diff --git a/.cspell-project-words.txt b/.cspell-project-words.txt
new file mode 100644
index 0000000000000000000000000000000000000000..719d3e4d017d7d1cf67ced46757292ad000c634f
--- /dev/null
+++ b/.cspell-project-words.txt
@@ -0,0 +1,11 @@
+andrii
+autoreference
+autoreferences
+categorising
+chunker
+collpsible
+entityreference
+euser
+getbatch
+nnevill
+sakhaniuk
diff --git a/ai_auto_reference.info.yml b/ai_auto_reference.info.yml
index 091d4a8dc1524b065a7ac6be2651be03448df9fd..37180d6a0e5fc8d0a853bcf4bf6b39ea364aeab2 100644
--- a/ai_auto_reference.info.yml
+++ b/ai_auto_reference.info.yml
@@ -2,7 +2,7 @@ name: AI Auto-reference
 type: module
 description: Automatically sets or suggests references to other content entities in the site.
 package: AI Auto-reference
-core_version_requirement: ^9 || ^10
+core_version_requirement: ^10 || ^11
 configure: ai_auto_reference.ai_auto_reference_settings_form
 dependencies:
   - drupal:ai
diff --git a/ai_auto_reference.install b/ai_auto_reference.install
index e6191b181629f72e5f7de4a58d6103fd8235e089..bf5641ed83c92774fdc7d5530f23638e66e817c7 100644
--- a/ai_auto_reference.install
+++ b/ai_auto_reference.install
@@ -5,10 +5,56 @@
  * Update scripts for the AI Auto Reference module.
  */
 
+use Drupal\node\Entity\NodeType;
+
 /**
- * Add auto_apply_relevance_levels to the default configuration for ai_auto_reference.settings.
+ * Add relevance levels to the default configuration settings.
  */
 function ai_auto_reference_update_10001() {
   $config = \Drupal::configFactory()->getEditable('ai_auto_reference.settings');
   $config->set('auto_apply_relevance_levels', ['high', 'medium'])->save();
 }
+
+/**
+ * Convert to third party settings for config schema validation.
+ */
+function ai_auto_reference_update_10002() {
+  $bundles = NodeType::loadMultiple();
+  if (!$bundles) {
+    return;
+  }
+  foreach ($bundles as $bundle) {
+    $changed = FALSE;
+    $machine_name = $bundle->id();
+    /** @var \Drupal\Core\Entity\Display\EntityDisplayInterface $form_display_entity */
+    $form_display_entity = \Drupal::entityTypeManager()
+      ->getStorage('entity_form_display')
+      ->load('node.' . $bundle->id() . '.default');
+    $components = $form_display_entity->getComponents();
+    foreach ($components as $field_name => $component) {
+      if (
+        is_array($component)
+        && isset($component['third_party_settings'])
+        && is_array($component['third_party_settings'])
+        && !empty($component['third_party_settings']['ai_auto_reference'])
+        && !empty($component['third_party_settings']['ai_auto_reference']['view_mode'])
+      ) {
+        $changed = TRUE;
+
+        // Set third party setting instead.
+        $view_mode = $component['third_party_settings']['ai_auto_reference']['view_mode'];
+        $form_display_entity->setThirdPartySetting('ai_auto_reference', $field_name, [
+          'view_mode' => $view_mode,
+        ]);
+
+        // Remove the original within the component.
+        unset($component['third_party_settings']['ai_auto_reference']);
+        $form_display_entity->setComponent($field_name, $component);
+      }
+    }
+
+    if ($changed) {
+      $form_display_entity->save();
+    }
+  }
+}
diff --git a/ai_auto_reference.module b/ai_auto_reference.module
index 313d45fd14ce485af1929b0a10ae909ce6538848..0a81104d80e75651001be43e1b337a921f1c928a 100644
--- a/ai_auto_reference.module
+++ b/ai_auto_reference.module
@@ -16,7 +16,8 @@ use Drupal\node\NodeInterface;
 function ai_auto_reference_form_node_form_alter(&$form, FormStateInterface $form_state, $form_id) {
   $node = $form_state->getFormObject()->getEntity();
   // Turning off the functionality on the 'create' page of nodes.
-  // The modules refreshes the page, which would just delete the progress on such pages.
+  // The modules refreshes the page, which would just delete the progress on
+  // such pages.
   if (
     $node instanceof NodeInterface
     && !$node->isNew()
@@ -28,13 +29,15 @@ function ai_auto_reference_form_node_form_alter(&$form, FormStateInterface $form
 
     // To prevent messages duplications.
     if (!empty($apply_form['container'])) {
+      $apply_form_rendered = \Drupal::service('renderer')->renderRoot($apply_form);
       \Drupal::messenger()->deleteByType('status');
-      \Drupal::messenger()->addMessage($apply_form);
+      \Drupal::messenger()->addMessage($apply_form_rendered);
     }
 
     $build_info = $form_state->getBuildInfo();
     $node = $build_info['callback_object']->getEntity();
     if ($node instanceof NodeInterface && $node->id()) {
+      $form['#attached']['library'][] = 'ai_auto_reference/ai-auto-reference-admin';
       $form['actions']['ai_auto_reference'] = [
         '#type' => 'submit',
         '#value' => t('Generate references with AI'),
diff --git a/assets/css/ai-auto-reference.css b/assets/css/ai-auto-reference.css
index 6ac1105cd72b9033743470b24b8008fe5f5b93e3..ad88aadf94553a44cb506663d8fab5f38efaaeb9 100644
--- a/assets/css/ai-auto-reference.css
+++ b/assets/css/ai-auto-reference.css
@@ -1,5 +1,7 @@
 .fieldset__ai-auto-reference,
-.fieldset__ai-auto-reference:not(.fieldgroup) {
+.fieldset__ai-auto-reference:not(.fieldgroup),
+.fieldset.fieldset__ai-auto-reference,
+.fieldset.fieldset__ai-auto-reference:not(.fieldgroup) {
   /** Prevent Gin theme compatibility issue for the apply form */
   background-color: white;
 }
diff --git a/composer.json b/composer.json
index 2c9c40f7d191f749bb5f8957cbb1274225e226c6..65fb8fadd0afdd0bea5af5953fd7f182ab930095 100644
--- a/composer.json
+++ b/composer.json
@@ -1,19 +1,22 @@
 {
-    "name": "drupal/ai_auto_reference",
-    "description": "Automatically sets or suggests references to other content entities in the site.",
-    "keywords": ["openai", "GPT-3", "GPT-4", "content generation", "AI content", "gpt", "content moderation"],
-    "type": "drupal-module",
-    "license": "GPL-2.0-or-later",
-    "authors": [
-        {
-            "name": "Scott Euser",
-            "homepage": "https://scotteuser.com",
-            "role": "Maintainer"
-        },
-        {
-            "name": "Andrii Sakhaniuk",
-            "homepage": "https://dev-branch.com",
-            "role": "Maintainer"
-        }
-    ]
+  "name": "drupal/ai_auto_reference",
+  "description": "Automatically sets or suggests references to other content entities in the site.",
+  "keywords": ["openai", "GPT-3", "GPT-4", "content generation", "AI content", "gpt", "content moderation"],
+  "type": "drupal-module",
+  "license": "GPL-2.0-or-later",
+  "require": {
+    "drupal/ai": "^1.0.4"
+  },
+  "authors": [
+    {
+      "name": "Scott Euser",
+      "homepage": "https://scotteuser.com",
+      "role": "Maintainer"
+    },
+    {
+      "name": "Andrii Sakhaniuk",
+      "homepage": "https://dev-branch.com",
+      "role": "Maintainer"
+    }
+  ]
 }
diff --git a/config/install/ai_auto_reference.settings.yml b/config/install/ai_auto_reference.settings.yml
index 64d0bd96b34bf3df3b382617299f0a2266662783..44b6c5b198f18c952559c86ff682f308ff089982 100644
--- a/config/install/ai_auto_reference.settings.yml
+++ b/config/install/ai_auto_reference.settings.yml
@@ -1,6 +1,6 @@
 service: "gpt"
 service_model: "gpt-3.5-turbo"
 service_api_key: ""
-token_limit: "4000"
+token_limit: 4000
 auto_apply_suggestions: false
 auto_apply_relevance_levels: ["high","medium"]
diff --git a/config/schema/ai_auto_reference.entity_display.schema.yml b/config/schema/ai_auto_reference.entity_display.schema.yml
new file mode 100644
index 0000000000000000000000000000000000000000..d63d68243b0f35ddc41d9f4e91b663c88b551c71
--- /dev/null
+++ b/config/schema/ai_auto_reference.entity_display.schema.yml
@@ -0,0 +1,21 @@
+core.entity_view_display.*.*.*.third_party.ai_auto_reference:
+  type: sequence
+  label: 'AI Auto-Reference field settings'
+  sequence:
+    type: mapping
+    label: An AI Auto-Reference
+    mapping:
+      view_mode:
+        type: string
+        label: 'The view mode'
+
+core.entity_form_display.*.*.*.third_party.ai_auto_reference:
+  type: sequence
+  label: 'AI Auto-Reference field settings'
+  sequence:
+    type: mapping
+    label: An AI Auto-Reference
+    mapping:
+      view_mode:
+        type: string
+        label: 'The view mode'
diff --git a/config/schema/ai_auto_reference.schema.yml b/config/schema/ai_auto_reference.schema.yml
index b0e75d115c2cf67f884ec7017380b3051ed0f421..28475fbd1ee55c333c0d50204feb0ef1d1238538 100644
--- a/config/schema/ai_auto_reference.schema.yml
+++ b/config/schema/ai_auto_reference.schema.yml
@@ -19,4 +19,6 @@ ai_auto_reference.settings:
       type: boolean
     auto_apply_relevance_levels:
       label: 'Automatically save the suggestions at the specified relevance without editor review'
-      type: array
\ No newline at end of file
+      type: sequence
+      sequence:
+        type: string
diff --git a/src/AiReferenceGenerator.php b/src/AiReferenceGenerator.php
index 0f4d67b93c7b18c2de435e23acbafcb4eafe9306..46aa30559eca83781b7943c8479f64528ef05efc 100644
--- a/src/AiReferenceGenerator.php
+++ b/src/AiReferenceGenerator.php
@@ -5,8 +5,6 @@ namespace Drupal\ai_auto_reference;
 use Drupal\ai\AiProviderPluginManager;
 use Drupal\ai\OperationType\Chat\ChatInput;
 use Drupal\ai\OperationType\Chat\ChatMessage;
-use Drupal\ai\OperationType\GenericType\AudioFile;
-use Drupal\ai\OperationType\SpeechToText\SpeechToTextInput;
 use Drupal\ai\Utility\TextChunkerInterface;
 use Drupal\ai\Utility\TokenizerInterface;
 use Drupal\Component\Serialization\Json;
@@ -56,7 +54,7 @@ class AiReferenceGenerator {
    *   The AI provider plugin manager.
    * @param \Drupal\ai\Utility\TokenizerInterface $tokenizer
    *   The tokenizer.
-   * @param \Drupal\ai\Utility\TextChunkerInterface
+   * @param \Drupal\ai\Utility\TextChunkerInterface $textChunker
    *   The text chunker.
    */
   public function __construct(
@@ -84,17 +82,22 @@ class AiReferenceGenerator {
    *   Array with configurations.
    */
   public function getBundleAiReferencesConfiguration($bundle) {
-    $form_display = $this->entityTypeManager
+    /** @var \Drupal\Core\Entity\Display\EntityDisplayInterface $form_display_entity */
+    $form_display_entity = $this->entityTypeManager
       ->getStorage('entity_form_display')
       ->load("node.{$bundle}.default");
+    $settings = $form_display_entity->getThirdPartySettings('ai_auto_reference');
+    if (empty($settings)) {
+      return $settings;
+    }
 
     $ai_references = [];
-    foreach ($form_display->getComponents() as $field_name => $component) {
+    foreach ($settings as $field_name => $setting) {
       // Add AI autogenerate button only if at least one field is configured.
-      if (!empty(($component['third_party_settings']['ai_auto_reference']))) {
+      if (!empty($setting['view_mode'])) {
         $ai_references[] = [
           'field_name' => $field_name,
-          'view_mode' => $component['third_party_settings']['ai_auto_reference']['view_mode'],
+          'view_mode' => $setting['view_mode'],
         ];
       }
     }
diff --git a/src/Batch/AiReferenceBatch.php b/src/Batch/AiReferenceBatch.php
index 10cc595513aa6130dad5bda7f1657d20db55fd93..e9d75db71d2d45fabe0593057db792dd35d483cc 100644
--- a/src/Batch/AiReferenceBatch.php
+++ b/src/Batch/AiReferenceBatch.php
@@ -42,7 +42,7 @@ class AiReferenceBatch {
   public function __construct(
     EntityTypeManagerInterface $entity_type_manager,
     RendererInterface $renderer,
-    ConfigFactoryInterface $config
+    ConfigFactoryInterface $config,
   ) {
     $this->entityTypeManager = $entity_type_manager;
     $this->renderer = $renderer;
diff --git a/src/Form/AutoReferenceApplyForm.php b/src/Form/AutoReferenceApplyForm.php
index eb6b759880faae0dd6554459d8e212dc53d882fb..ea970754418479015d1896eeb4c6054adba1881f 100644
--- a/src/Form/AutoReferenceApplyForm.php
+++ b/src/Form/AutoReferenceApplyForm.php
@@ -36,7 +36,7 @@ class AutoReferenceApplyForm extends FormBase {
    */
   public function __construct(
     EntityTypeManagerInterface $entity_type_manager,
-    RequestStack $request_stack
+    RequestStack $request_stack,
   ) {
     $this->entityTypeManager = $entity_type_manager;
     $this->requestStack = $request_stack;
@@ -74,7 +74,7 @@ class AutoReferenceApplyForm extends FormBase {
    * @return array
    *   The form structure.
    */
-  public function buildForm(array $form, FormStateInterface $form_state, NodeInterface $node = NULL, array $configuration = NULL) {
+  public function buildForm(array $form, FormStateInterface $form_state, ?NodeInterface $node = NULL, ?array $configuration = NULL) {
     $form = [];
     $query = $this->requestStack->getCurrentRequest()->query->all();
     $auto_apply = !empty($query['auto-apply']);
diff --git a/src/Form/AutoReferenceDeleteForm.php b/src/Form/AutoReferenceDeleteForm.php
index da4da83a9532fe82a2cdf4e856f2c20672a15a25..f53be376f0abdfd5f33dcc9c6253d80703d2b14d 100644
--- a/src/Form/AutoReferenceDeleteForm.php
+++ b/src/Form/AutoReferenceDeleteForm.php
@@ -96,17 +96,18 @@ class AutoReferenceDeleteForm extends ConfirmFormBase {
     if (empty($this->nodeType) || empty($this->fieldName)) {
       return;
     }
-    $entity = $this->entityTypeManager
+
+    /** @var \Drupal\Core\Entity\Display\EntityDisplayInterface $form_display_entity */
+    $form_display_entity = $this->entityTypeManager
       ->getStorage('entity_form_display')
-      ->load("node.{$this->nodeType}.default");
+      ->load('node.' . $this->nodeType . '.default');
 
-    if (!$entity) {
+    if (!$form_display_entity) {
       return;
     }
-    $field = $entity->getComponent($this->fieldName);
-    unset($field['third_party_settings']['ai_auto_reference']);
-    $entity->setComponent($this->fieldName, $field);
-    $entity->save();
+
+    $form_display_entity->unsetThirdPartySetting('ai_auto_reference', $this->fieldName);
+    $form_display_entity->save();
     $form_state->setRedirectUrl($this->getCancelUrl());
   }
 
diff --git a/src/Form/AutoReferenceEditForm.php b/src/Form/AutoReferenceEditForm.php
index 35def5f6943698ea0a9c20feffcf0517b4ecfcbd..34807f45955c8c9358659f36dab45ad274beb9e8 100644
--- a/src/Form/AutoReferenceEditForm.php
+++ b/src/Form/AutoReferenceEditForm.php
@@ -57,7 +57,7 @@ class AutoReferenceEditForm extends FormBase {
   public function __construct(
     EntityDisplayRepositoryInterface $entity_display_repository,
     EntityTypeManagerInterface $entity_type_manager,
-    EntityFieldManagerInterface $entity_field_manager
+    EntityFieldManagerInterface $entity_field_manager,
   ) {
     $this->entityDisplayRepository = $entity_display_repository;
     $this->entityTypeManager = $entity_type_manager;
@@ -91,6 +91,7 @@ class AutoReferenceEditForm extends FormBase {
       return $form;
     }
 
+    /** @var \Drupal\Core\Entity\Display\EntityDisplayInterface $form_display_entity */
     $form_display_entity = $this->entityTypeManager
       ->getStorage('entity_form_display')
       ->load("node.{$node_type}.default");
@@ -99,9 +100,8 @@ class AutoReferenceEditForm extends FormBase {
       return $form;
     }
 
-    $field = $form_display_entity->getComponent($field_name);
-
-    if (empty($field['third_party_settings']['ai_auto_reference']['view_mode'])) {
+    $setting = $form_display_entity->getThirdPartySetting('ai_auto_reference', $field_name);
+    if (!is_array($setting) || empty($setting['view_mode'])) {
       return $form;
     }
 
@@ -112,7 +112,7 @@ class AutoReferenceEditForm extends FormBase {
       '#title' => $this->t('View mode'),
       '#description' => $this->t('View mode that will be used as a source of content for AI analysis.'),
       '#options' => $view_modes,
-      '#default_value' => $field['third_party_settings']['ai_auto_reference']['view_mode'],
+      '#default_value' => $setting['view_mode'],
       '#required' => TRUE,
     ];
     $form['bundle'] = [
@@ -141,18 +141,18 @@ class AutoReferenceEditForm extends FormBase {
     $bundle = $values['bundle'];
     $field_name = $values['field_name'];
 
-    /** @var \Drupal\Core\Entity\Entity\EntityFormDisplay $entity */
-    $entity = $this->entityTypeManager
+    /** @var \Drupal\Core\Entity\Display\EntityDisplayInterface $form_display_entity */
+    $form_display_entity = $this->entityTypeManager
       ->getStorage('entity_form_display')
       ->load("{$this->entityType}.{$bundle}.default");
-    if (!$entity) {
+    if (!$form_display_entity) {
       return;
     }
 
-    $field = $entity->getComponent($field_name);
-    $field['third_party_settings']['ai_auto_reference']['view_mode'] = $values['view_mode'];
-    $entity->setComponent($field_name, $field);
-    $entity->save();
+    $form_display_entity->setThirdPartySetting('ai_auto_reference', $field_name, [
+      'view_mode' => $values['view_mode'],
+    ]);
+    $form_display_entity->save();
     $form_state->setRedirect($this->redirectPath, ['node_type' => $bundle]);
   }
 
diff --git a/src/Form/EntityBundleSettingsForm.php b/src/Form/EntityBundleSettingsForm.php
index 2f5f7a5be17b51638c1146fb1db5112936ecb05c..2460d8f4dc6229733c210b1a721af1c873b6b1fe 100644
--- a/src/Form/EntityBundleSettingsForm.php
+++ b/src/Form/EntityBundleSettingsForm.php
@@ -81,7 +81,7 @@ class EntityBundleSettingsForm extends FormBase {
   public function __construct(
     EntityFieldManagerInterface $entity_field_manager,
     EntityTypeManagerInterface $entity_type_manager,
-    EntityDisplayRepositoryInterface $entity_display_repository
+    EntityDisplayRepositoryInterface $entity_display_repository,
   ) {
     $this->entityFieldManager = $entity_field_manager;
     $this->entityTypeManager = $entity_type_manager;
@@ -139,11 +139,9 @@ class EntityBundleSettingsForm extends FormBase {
     }
 
     $field_name = $table['add_new_autoreference']['field_name'];
-    $field = $entity_form_display->getComponent($field_name);
-    $field['third_party_settings']['ai_auto_reference'] = [
+    $entity_form_display->setThirdPartySetting('ai_auto_reference', $field_name, [
       'view_mode' => $table['add_new_autoreference']['view_mode'],
-    ];
-    $entity_form_display->setComponent($field_name, $field);
+    ]);
     $entity_form_display->save();
   }
 
@@ -185,6 +183,7 @@ class EntityBundleSettingsForm extends FormBase {
       }
     }
 
+    /** @var \Drupal\Core\Entity\Display\EntityDisplayInterface $form_display_entity */
     $form_display_entity = $this->entityTypeManager
       ->getStorage('entity_form_display')
       ->load("$this->entityType.$this->bundleName.default");
@@ -196,9 +195,9 @@ class EntityBundleSettingsForm extends FormBase {
     $view_modes = $this->entityDisplayRepository->getViewModeOptionsByBundle($this->entityType, $this->bundleName);
 
     // Displaying already added AI autoreference blocks.
+    $settings = $form_display_entity->getThirdPartySettings('ai_auto_reference');
     foreach ($reference_fields as $field_name => $label) {
-      $field = $form_display_entity->getComponent($field_name);
-      if (empty($field['third_party_settings']['ai_auto_reference']['view_mode'])) {
+      if (empty($settings[$field_name]['view_mode'])) {
         continue;
       }
       $parameters = [
@@ -207,7 +206,7 @@ class EntityBundleSettingsForm extends FormBase {
       ];
       $table[] = [
         'field_name' => ['#markup' => "{$label} ({$field_name})"],
-        'view_mode' => ['#markup' => $view_modes[$field['third_party_settings']['ai_auto_reference']['view_mode']]],
+        'view_mode' => ['#markup' => $view_modes[$settings[$field_name]['view_mode']]],
         'actions' => [
           '#type' => 'operations',
           '#links' => [
diff --git a/src/Form/SettingsForm.php b/src/Form/SettingsForm.php
index 415422f28a56dcfb4c7b1d5fc46f650afbd4a2cb..277bf9dbb8a68202e8b2e3c2e4284c8004fabf85 100644
--- a/src/Form/SettingsForm.php
+++ b/src/Form/SettingsForm.php
@@ -21,14 +21,14 @@ class SettingsForm extends ConfigFormBase {
    *
    * @param \Drupal\Core\Config\ConfigFactoryInterface $config_factory
    *   The factory for configuration objects.
-   * @param \Drupal\Core\Config\TypedConfigManagerInterface|null $typedConfigManager
+   * @param \Drupal\Core\Config\TypedConfigManagerInterface $typedConfigManager
    *   The typed config manager.
    * @param \Drupal\ai\AiProviderPluginManager $aiProviderManager
    *   The AI provider manager.
    */
   public function __construct(
     ConfigFactoryInterface $config_factory,
-    protected $typedConfigManager,
+    protected TypedConfigManagerInterface $typedConfigManager,
     protected AiProviderPluginManager $aiProviderManager,
   ) {
     parent::__construct($config_factory, $typedConfigManager);
@@ -135,7 +135,7 @@ class SettingsForm extends ConfigFormBase {
     $relevance_levels = array_filter($form_state->getValue('auto_apply_relevance_levels'));
     $config = $this->config('ai_auto_reference.settings')
       ->set('provider', $form_state->getValue('provider'))
-      ->set('token_limit', $form_state->getValue('token_limit'))
+      ->set('token_limit', (int) $form_state->getValue('token_limit'))
       ->set('auto_apply_suggestions', $form_state->getValue('auto_apply_suggestions'))
       ->set('auto_apply_relevance_levels', array_keys($relevance_levels));
     $config->save();
diff --git a/tests/src/Functional/AiAutoReferenceUiTest.php b/tests/src/Functional/AiAutoReferenceUiTest.php
new file mode 100644
index 0000000000000000000000000000000000000000..bdb1c9a46cf9dbab37cf211969015f0b5b81d2bd
--- /dev/null
+++ b/tests/src/Functional/AiAutoReferenceUiTest.php
@@ -0,0 +1,260 @@
+<?php
+
+namespace Drupal\Tests\ai_auto_reference\Functional;
+
+use Drupal\Core\Url;
+use Drupal\Tests\BrowserTestBase;
+
+/**
+ * Contains AI Auto-Reference UI setup functional tests.
+ *
+ * @group ai_auto_reference
+ */
+class AiAutoReferenceUiTest extends BrowserTestBase {
+
+  /**
+   * {@inheritdoc}
+   */
+  protected static $modules = [
+    'ai',
+    'ai_auto_reference',
+    'node',
+    'taxonomy',
+    'user',
+    'system',
+    'field_ui',
+    'views_ui',
+  ];
+
+  /**
+   * {@inheritdoc}
+   */
+  protected $defaultTheme = 'stark';
+
+  /**
+   * A user with permission to bypass access content.
+   *
+   * @var \Drupal\user\UserInterface
+   */
+  protected $adminUser;
+
+  /**
+   * Nodes for testing the indexing.
+   *
+   * @var array
+   *   An array of nodes for testing.
+   */
+  protected $nodes = [];
+
+  /**
+   * {@inheritdoc}
+   */
+  protected function setUp(): void {
+    parent::setUp();
+
+    $this->drupalCreateContentType([
+      'type' => 'article',
+      'name' => 'Article',
+    ]);
+
+    $this->adminUser = $this->drupalCreateUser([
+      'access administration pages',
+      'administer content types',
+      'access content overview',
+      'administer nodes',
+      'administer node fields',
+      'bypass node access',
+      'administer ai',
+      'administer ai autoreference',
+    ]);
+
+    // Create an entity reference field "Related food" for articles.
+    $this->createEntityReferenceField(
+      'node',
+      'article',
+      'field_related_food',
+      'Related food',
+      'node',
+      ['article'],
+    );
+
+    $this->createSampleContent();
+  }
+
+  /**
+   * Helper function to create an entity reference field.
+   *
+   * @param string $entity_type
+   *   The entity type (e.g., 'node').
+   * @param string $bundle
+   *   The bundle (content type).
+   * @param string $field_name
+   *   The machine name of the field.
+   * @param string $field_label
+   *   The human-readable label.
+   * @param string $target_type
+   *   The entity type being referenced.
+   * @param array $target_bundles
+   *   Allowed target bundles.
+   */
+  protected function createEntityReferenceField($entity_type, $bundle, $field_name, $field_label, $target_type, array $target_bundles) {
+    $entity_type_manager = \Drupal::entityTypeManager();
+    $display_repository = \Drupal::service('entity_display.repository');
+
+    // Create the field storage.
+    $field_storage = $entity_type_manager->getStorage('field_storage_config')->create([
+      'field_name' => $field_name,
+      'entity_type' => $entity_type,
+      'type' => 'entity_reference',
+      'settings' => [
+        'target_type' => $target_type,
+      ],
+      'cardinality' => -1,
+      'translatable' => FALSE,
+    ]);
+    $field_storage->save();
+
+    // Create the field configuration.
+    $field = $entity_type_manager->getStorage('field_config')->create([
+      'field_name' => $field_name,
+      'entity_type' => $entity_type,
+      'bundle' => $bundle,
+      'label' => $field_label,
+      'settings' => [
+        'handler' => 'default',
+        'handler_settings' => [
+          'target_bundles' => array_combine($target_bundles, $target_bundles),
+        ],
+      ],
+    ]);
+    $field->save();
+
+    // Set form display.
+    $form_display = $display_repository->getFormDisplay($entity_type, $bundle, 'default');
+    $form_display->setComponent($field_name, [
+      'type' => 'entity_reference_autocomplete',
+    ])->save();
+
+    // Set view display.
+    $view_display = $display_repository->getViewDisplay($entity_type, $bundle, 'default');
+    $view_display->setComponent($field_name, [
+      'type' => 'entity_reference_label',
+    ])->save();
+  }
+
+  /**
+   * Create sample content to check index.
+   */
+  public function createSampleContent(): void {
+    $this->nodes[] = $this->drupalCreateNode([
+      'type' => 'article',
+      'title' => 'Chocolate Cake',
+      'field_body' => [
+        'value' => 'A delicious chocolate dessert made with cocoa powder and dark chocolate.',
+        'format' => 'plain_text',
+      ],
+    ]);
+    $this->nodes[] = $this->drupalCreateNode([
+      'type' => 'article',
+      'title' => 'Strawberry Cheese Cake',
+      'field_body' => [
+        'value' => 'A sweet cheese based dessert make with strawberries on a pie-like crust.',
+        'format' => 'plain_text',
+      ],
+      'status' => 0,
+    ]);
+    $this->nodes[] = $this->drupalCreateNode([
+      'type' => 'article',
+      'title' => 'Vanilla Ice Cream',
+      'field_body' => [
+        'value' => 'A creamy vanilla dessert made with milk, cream, and vanilla extract.',
+        'format' => 'plain_text',
+      ],
+    ]);
+    $this->nodes[] = $this->drupalCreateNode([
+      'type' => 'article',
+      'title' => 'Tomato Soup',
+      'field_body' => [
+        'value' => 'A warm starter made with fresh tomatoes, garlic, and basil.',
+        'format' => 'plain_text',
+      ],
+    ]);
+  }
+
+  /**
+   * Test the access.
+   */
+  public function testSettingsPageAccess() {
+    $this->drupalGet('admin/config/content/ai-autoreference-settings');
+    $this->assertSession()->statusCodeEquals(403);
+
+    $this->drupalGet('admin/structure/types/manage/article/ai-autoreference');
+    $this->assertSession()->statusCodeEquals(403);
+  }
+
+  /**
+   * Test the settings.
+   */
+  public function testNodeSettings() {
+    $this->drupalLogin($this->adminUser);
+
+    // Check that no button exists yet.
+    $this->drupalGet('node/1/edit');
+    $this->assertSession()->buttonNotExists('Generate references with AI');
+
+    // Create a new auto-reference.
+    $this->drupalGet('admin/structure/types/manage/article/ai-autoreference');
+    $this->submitForm([
+      'table[add_new_autoreference][field_name]' => 'field_related_food',
+      'table[add_new_autoreference][view_mode]' => 'teaser',
+    ], 'Add autoreference');
+
+    // Now edit it.
+    $this->clickLink('Edit');
+
+    // Re-save.
+    $this->submitForm([], 'Save settings');
+
+    // Delete it.
+    $this->drupalGet('admin/structure/types/manage/article/ai-autoreference/field_related_food/delete');
+    $this->submitForm([], 'Confirm');
+
+    // Add it again.
+    $this->submitForm([
+      'table[add_new_autoreference][field_name]' => 'field_related_food',
+      'table[add_new_autoreference][view_mode]' => 'teaser',
+    ], 'Add autoreference');
+
+    // Edit a node.
+    $this->drupalGet('node/1/edit');
+    $this->assertSession()->buttonExists('Generate references with AI');
+
+    // Pretend we have generated references.
+    $url = Url::fromRoute('entity.node.edit_form', ['node' => 1], [
+      'query' => [
+        'field_related_food[h]' => '1,2',
+        'field_related_food[m]' => '3',
+        'auto-apply' => '',
+      ],
+    ]);
+    $this->drupalGet($url);
+    // High results.
+    $this->assertSession()->elementTextContains(
+      'css',
+      '.js-form-item-container-field-related-food-h-1 label',
+      'Chocolate Cake',
+    );
+    $this->assertSession()->elementTextContains(
+      'css',
+      '.js-form-item-container-field-related-food-h-2 label',
+      'Strawberry Cheese Cake',
+    );
+    // Medium results.
+    $this->assertSession()->elementTextContains(
+      'css',
+      '.js-form-item-container-field-related-food-m-3 label',
+      'Vanilla Ice Cream',
+    );
+  }
+
+}