From d92ebefb9bc35e63c58dc806cdc096628dcc55cf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=B8ren=20Louv-Jansen?= <soren.louv@elastic.co>
Date: Fri, 9 May 2025 11:27:32 +0200
Subject: [PATCH 1/3] [Obs AI Assistant] Fix flaky `recall` tests

---
 .../service/knowledge_base_service/index.ts   |  4 +-
 .../index_write_block_utils.ts                | 41 +++++++++++++++++++
 .../reindex_knowledge_base.ts                 | 14 +++----
 .../update_knowledge_base_index_alias.ts      | 15 ++++++-
 .../complete/functions/recall.spec.ts         |  4 +-
 ...base_change_model_from_elser_to_e5.spec.ts |  4 +-
 .../knowledge_base_setup.spec.ts              | 13 +-----
 .../ai_assistant/utils/index_assets.ts        | 12 +++++-
 .../ai_assistant/utils/knowledge_base.ts      | 32 +++++++++++++--
 .../ai_assistant/utils/model_and_inference.ts |  9 +---
 10 files changed, 109 insertions(+), 39 deletions(-)
 create mode 100644 x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index_write_block_utils.ts

diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts
index 100e4782e4cbd..6f34edcf26c01 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts
@@ -416,7 +416,7 @@ export class KnowledgeBaseService {
     }
 
     try {
-      await this.dependencies.esClient.asInternalUser.index<
+      const indexResult = await this.dependencies.esClient.asInternalUser.index<
         Omit<KnowledgeBaseEntry, 'id'> & { namespace: string }
       >({
         index: resourceNames.writeIndexAlias.kb,
@@ -432,7 +432,7 @@ export class KnowledgeBaseService {
       });
 
       this.dependencies.logger.debug(
-        `Entry added to knowledge base. title = "${doc.title}", user = "${user?.name}, namespace = "${namespace}"`
+        `Entry added to knowledge base. title = "${doc.title}", user = "${user?.name}, namespace = "${namespace}", index = ${indexResult._index}, id = ${indexResult._id}`
       );
     } catch (error) {
       this.dependencies.logger.error(`Failed to add entry to knowledge base ${error}`);
diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index_write_block_utils.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index_write_block_utils.ts
new file mode 100644
index 0000000000000..b3057c666d1d3
--- /dev/null
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index_write_block_utils.ts
@@ -0,0 +1,41 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { errors } from '@elastic/elasticsearch';
+import { ElasticsearchClient } from '@kbn/core/server';
+import { resourceNames } from '..';
+
+export async function addIndexWriteBlock({
+  esClient,
+  index,
+}: {
+  esClient: { asInternalUser: ElasticsearchClient };
+  index: string;
+}) {
+  await esClient.asInternalUser.indices.addBlock({ index, block: 'write' });
+}
+
+export function removeIndexWriteBlock({
+  esClient,
+  index,
+}: {
+  esClient: { asInternalUser: ElasticsearchClient };
+  index: string;
+}) {
+  return esClient.asInternalUser.indices.putSettings({
+    index,
+    body: { 'index.blocks.write': false },
+  });
+}
+
+export function isKnowledgeBaseIndexWriteBlocked(error: any) {
+  return (
+    error instanceof errors.ResponseError &&
+    error.message.includes(`cluster_block_exception`) &&
+    error.message.includes(resourceNames.writeIndexAlias.kb)
+  );
+}
diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/reindex_knowledge_base.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/reindex_knowledge_base.ts
index 58b7b98d1e31d..8dd7cae3902fd 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/reindex_knowledge_base.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/reindex_knowledge_base.ts
@@ -59,13 +59,6 @@ async function reIndexKnowledgeBase({
     `Re-indexing knowledge base from "${currentWriteIndexName}" to index "${nextWriteIndexName}"...`
   );
 
-  const reindexResponse = await esClient.asInternalUser.reindex({
-    source: { index: currentWriteIndexName },
-    dest: { index: nextWriteIndexName },
-    refresh: true,
-    wait_for_completion: false,
-  });
-
   // Point write index alias to the new index
   await updateKnowledgeBaseWriteIndexAlias({
     esClient,
@@ -74,6 +67,13 @@ async function reIndexKnowledgeBase({
     currentWriteIndexName,
   });
 
+  const reindexResponse = await esClient.asInternalUser.reindex({
+    source: { index: currentWriteIndexName },
+    dest: { index: nextWriteIndexName },
+    refresh: true,
+    wait_for_completion: false,
+  });
+
   const taskId = reindexResponse.task?.toString();
   if (taskId) {
     await waitForReIndexTaskToComplete({ esClient, taskId, logger });
diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/update_knowledge_base_index_alias.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/update_knowledge_base_index_alias.ts
index 3d893d6c2a6f3..44d91626a7343 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/update_knowledge_base_index_alias.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/update_knowledge_base_index_alias.ts
@@ -8,6 +8,7 @@
 import { ElasticsearchClient } from '@kbn/core-elasticsearch-server';
 import { Logger } from '@kbn/logging';
 import { resourceNames } from '..';
+import { addIndexWriteBlock, removeIndexWriteBlock } from './index_write_block_utils';
 
 export async function updateKnowledgeBaseWriteIndexAlias({
   esClient,
@@ -25,6 +26,11 @@ export async function updateKnowledgeBaseWriteIndexAlias({
   );
   const alias = resourceNames.writeIndexAlias.kb;
   try {
+    await addIndexWriteBlock({ esClient, index: currentWriteIndexName });
+    logger.debug(
+      `Added write block to "${currentWriteIndexName}". It is now read-only and writes are temporarily blocked.`
+    );
+
     await esClient.asInternalUser.indices.updateAliases({
       actions: [
         { remove: { index: currentWriteIndexName, alias } },
@@ -32,7 +38,14 @@ export async function updateKnowledgeBaseWriteIndexAlias({
       ],
     });
   } catch (error) {
-    logger.error(`Failed to update write index alias: ${error.message}`);
+    await removeIndexWriteBlock({ esClient, index: currentWriteIndexName });
+    logger.error(
+      `Failed to update write index alias: ${error.message}. Reverting back to ${currentWriteIndexName}`
+    );
     throw error;
   }
+
+  logger.debug(
+    `Successfully updated write index alias to "${nextWriteIndexName}". Writes are now enabled again.`
+  );
 }
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/recall.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/recall.spec.ts
index db2218a3a76cb..ae3f8354029e8 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/recall.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/recall.spec.ts
@@ -25,7 +25,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
   const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
   const es = getService('es');
 
-  describe('recall', function () {
+  describe.only('recall', function () {
     before(async () => {
       await deployTinyElserAndSetupKb(getService);
       await addSampleDocsToInternalKb(getService, technicalSampleDocs);
@@ -40,7 +40,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     });
 
     describe('GET /internal/observability_ai_assistant/functions/recall', () => {
-      it('produces unique scores for each doc', async () => {
+      it.only('produces unique scores for each doc', async () => {
         const entries = await recall('What happened during the database outage?');
         const uniqueScores = uniq(entries.map(({ esScore }) => esScore));
         expect(uniqueScores.length).to.be.greaterThan(1);
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_change_model_from_elser_to_e5.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_change_model_from_elser_to_e5.spec.ts
index 48d141b46e8a5..4eac156b73868 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_change_model_from_elser_to_e5.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_change_model_from_elser_to_e5.spec.ts
@@ -56,7 +56,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     before(async () => {
       await importModel(getService, { modelId: TINY_ELSER_MODEL_ID });
       await createTinyElserInferenceEndpoint(getService, { inferenceId: TINY_ELSER_INFERENCE_ID });
-      await setupKnowledgeBase(observabilityAIAssistantAPIClient, TINY_ELSER_INFERENCE_ID);
+      await setupKnowledgeBase(getService, TINY_ELSER_INFERENCE_ID);
       await waitForKnowledgeBaseReady(getService);
 
       // ingest documents
@@ -76,7 +76,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
       await createTinyTextEmbeddingInferenceEndpoint(getService, {
         inferenceId: TINY_TEXT_EMBEDDING_INFERENCE_ID,
       });
-      await setupKnowledgeBase(observabilityAIAssistantAPIClient, TINY_TEXT_EMBEDDING_INFERENCE_ID);
+      await setupKnowledgeBase(getService, TINY_TEXT_EMBEDDING_INFERENCE_ID);
 
       await waitForKnowledgeBaseIndex(getService, '.kibana-observability-ai-assistant-kb-000002');
       await waitForKnowledgeBaseReady(getService);
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts
index 7e0a5c9466386..3674589ce78f7 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts
@@ -65,11 +65,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
         let body: Awaited<ReturnType<typeof setupKbAsAdmin>>['body'];
 
         before(async () => {
-          // setup KB initially
-
           await deployTinyElserAndSetupKb(getService);
-
-          // setup KB with custom inference endpoint
           await createTinyElserInferenceEndpoint(getService, {
             inferenceId: CUSTOM_TINY_ELSER_INFERENCE_ID,
           });
@@ -120,7 +116,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
         await createTinyElserInferenceEndpoint(getService, {
           inferenceId: customInferenceId,
         });
-        await setupKnowledgeBase(observabilityAIAssistantAPIClient, customInferenceId);
+        await setupKnowledgeBase(getService, customInferenceId);
         await waitForKnowledgeBaseReady(getService);
       });
 
@@ -154,12 +150,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
   }
 
   function setupKbAsAdmin(inferenceId: string) {
-    return observabilityAIAssistantAPIClient.admin({
-      endpoint: 'POST /internal/observability_ai_assistant/kb/setup',
-      params: {
-        query: { inference_id: inferenceId },
-      },
-    });
+    return setupKnowledgeBase(getService, inferenceId);
   }
 
   function setupKbAsViewer(inferenceId: string) {
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/index_assets.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/index_assets.ts
index 4cd1d9e4b0f3f..14ff7b8286b01 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/index_assets.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/index_assets.ts
@@ -14,6 +14,7 @@ import {
 import { DeploymentAgnosticFtrProviderContext } from '../../../../ftr_provider_context';
 import type { ObservabilityAIAssistantApiClient } from '../../../../services/observability_ai_assistant_api';
 import { TINY_ELSER_INFERENCE_ID } from './model_and_inference';
+import { getConcreteWriteIndexFromAlias } from './knowledge_base';
 
 export async function runStartupMigrations(
   observabilityAIAssistantAPIClient: ObservabilityAIAssistantApiClient
@@ -67,9 +68,16 @@ export async function restoreIndexAssets(
   getService: DeploymentAgnosticFtrProviderContext['getService']
 ) {
   const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+  const retry = getService('retry');
+  const es = getService('es');
+  const log = getService('log');
 
-  await deleteIndexAssets(getService);
-  await createOrUpdateIndexAssets(observabilityAIAssistantAPIClient);
+  await retry.try(async () => {
+    log.debug('Restoring index assets');
+    await deleteIndexAssets(getService);
+    await createOrUpdateIndexAssets(observabilityAIAssistantAPIClient);
+    expect(await getConcreteWriteIndexFromAlias(es)).to.be(resourceNames.concreteWriteIndexName.kb);
+  });
 }
 
 export async function getComponentTemplate(es: Client) {
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts
index 66a59bc93bc57..7332e62938c29 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts
@@ -40,6 +40,14 @@ export async function waitForKnowledgeBaseIndex(
   });
 }
 
+export async function getKnowledgeBaseStatus(
+  observabilityAIAssistantAPIClient: ObservabilityAIAssistantApiClient
+) {
+  return observabilityAIAssistantAPIClient.editor({
+    endpoint: 'GET /internal/observability_ai_assistant/kb/status',
+  });
+}
+
 export async function waitForKnowledgeBaseReady(
   getService: DeploymentAgnosticFtrProviderContext['getService']
 ) {
@@ -49,9 +57,7 @@ export async function waitForKnowledgeBaseReady(
 
   await retry.tryForTime(5 * 60 * 1000, async () => {
     log.debug(`Waiting for knowledge base to be ready...`);
-    const res = await observabilityAIAssistantAPIClient.editor({
-      endpoint: 'GET /internal/observability_ai_assistant/kb/status',
-    });
+    const res = await getKnowledgeBaseStatus(observabilityAIAssistantAPIClient);
     expect(res.status).to.be(200);
     expect(res.body.kbState).to.be(KnowledgeBaseState.READY);
     expect(res.body.isReIndexing).to.be(false);
@@ -60,9 +66,17 @@ export async function waitForKnowledgeBaseReady(
 }
 
 export async function setupKnowledgeBase(
-  observabilityAIAssistantAPIClient: ObservabilityAIAssistantApiClient,
+  getService: DeploymentAgnosticFtrProviderContext['getService'],
   inferenceId: string
 ) {
+  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+  const log = getService('log');
+
+  const statusResult = await getKnowledgeBaseStatus(observabilityAIAssistantAPIClient);
+
+  log.debug(
+    `Setting up knowledge base with inference endpoint = "${TINY_ELSER_INFERENCE_ID}", concreteWriteIndex = ${statusResult.body.concreteWriteIndex}, currentInferenceId = ${statusResult.body.currentInferenceId} ${statusResult.body.isReIndexing}`
+  );
   return observabilityAIAssistantAPIClient.admin({
     endpoint: 'POST /internal/observability_ai_assistant/kb/setup',
     params: {
@@ -77,6 +91,8 @@ export async function addSampleDocsToInternalKb(
 ) {
   const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
   const es = getService('es');
+  const log = getService('log');
+  const retry = getService('retry');
 
   await observabilityAIAssistantAPIClient.editor({
     endpoint: 'POST /internal/observability_ai_assistant/kb/entries/import',
@@ -88,6 +104,14 @@ export async function addSampleDocsToInternalKb(
   });
 
   await refreshKbIndex(es);
+
+  await retry.try(async () => {
+    const itemsInKb = await getKnowledgeBaseEntriesFromEs(es);
+    log.debug(
+      `Waiting for at least ${sampleDocs.length} docs to be available for search in KB. Currently ${itemsInKb.length} docs available.`
+    );
+    expect(itemsInKb.length >= sampleDocs.length).to.be(true);
+  });
 }
 
 // refresh the index to make sure the documents are searchable
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference.ts
index 2a610e0e45296..1d28bec819214 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference.ts
@@ -98,16 +98,9 @@ export function createTinyTextEmbeddingInferenceEndpoint(
 export async function deployTinyElserAndSetupKb(
   getService: DeploymentAgnosticFtrProviderContext['getService']
 ) {
-  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
-  const log = getService('log');
-
   await setupTinyElserModelAndInferenceEndpoint(getService);
 
-  log.debug(`Setting up knowledge base with inference endpoint "${TINY_ELSER_INFERENCE_ID}"`);
-  const { status, body } = await setupKnowledgeBase(
-    observabilityAIAssistantAPIClient,
-    TINY_ELSER_INFERENCE_ID
-  );
+  const { status, body } = await setupKnowledgeBase(getService, TINY_ELSER_INFERENCE_ID);
   await waitForKnowledgeBaseReady(getService);
 
   return { status, body };

From 159fd4d9259f8984c7611e0f3352d401327970e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=B8ren=20Louv-Jansen?= <soren.louv@elastic.co>
Date: Fri, 9 May 2025 12:07:08 +0200
Subject: [PATCH 2/3] Add `wait_until_complete` to /setup endpoint

---
 .../server/routes/knowledge_base/route.ts            | 12 +++++++-----
 .../server/service/client/index.ts                   |  9 +++++++--
 .../knowledge_base/knowledge_base_setup.spec.ts      |  2 +-
 .../ai_assistant/utils/knowledge_base.ts             |  2 +-
 4 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts
index b78dd82cc48e0..6596ca667dee5 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts
@@ -50,9 +50,10 @@ const getKnowledgeBaseStatus = createObservabilityAIAssistantServerRoute({
 const setupKnowledgeBase = createObservabilityAIAssistantServerRoute({
   endpoint: 'POST /internal/observability_ai_assistant/kb/setup',
   params: t.type({
-    query: t.type({
-      inference_id: t.string,
-    }),
+    query: t.intersection([
+      t.type({ inference_id: t.string }),
+      t.partial({ wait_until_complete: toBooleanRt }),
+    ]),
   }),
   security: {
     authz: {
@@ -67,8 +68,9 @@ const setupKnowledgeBase = createObservabilityAIAssistantServerRoute({
     nextInferenceId: string;
   }> => {
     const client = await resources.service.getClient({ request: resources.request });
-    const { inference_id: inferenceId } = resources.params.query;
-    return client.setupKnowledgeBase(inferenceId);
+    const { inference_id: inferenceId, wait_until_complete: waitUntilComplete } =
+      resources.params.query;
+    return client.setupKnowledgeBase(inferenceId, waitUntilComplete);
   },
 });
 
diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts
index 3d92b765b764a..4ad90ed30c37c 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts
@@ -664,7 +664,8 @@ export class ObservabilityAIAssistantClient {
   };
 
   setupKnowledgeBase = async (
-    nextInferenceId: string
+    nextInferenceId: string,
+    waitUntilComplete: boolean = false
   ): Promise<{
     reindex: boolean;
     currentInferenceId: string | undefined;
@@ -693,7 +694,7 @@ export class ObservabilityAIAssistantClient {
       inferenceId: nextInferenceId,
     });
 
-    waitForKbModel({
+    const kbSetupPromise = waitForKbModel({
       core: this.dependencies.core,
       esClient,
       logger,
@@ -728,6 +729,10 @@ export class ObservabilityAIAssistantClient {
         }
       });
 
+    if (waitUntilComplete) {
+      await kbSetupPromise;
+    }
+
     return { reindex: true, currentInferenceId, nextInferenceId };
   };
 
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts
index 3674589ce78f7..389bf18a952a9 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts
@@ -157,7 +157,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     return observabilityAIAssistantAPIClient.viewer({
       endpoint: 'POST /internal/observability_ai_assistant/kb/setup',
       params: {
-        query: { inference_id: inferenceId },
+        query: { inference_id: inferenceId, wait_until_complete: true },
       },
     });
   }
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts
index 7332e62938c29..dd0a77eafaa6a 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts
@@ -80,7 +80,7 @@ export async function setupKnowledgeBase(
   return observabilityAIAssistantAPIClient.admin({
     endpoint: 'POST /internal/observability_ai_assistant/kb/setup',
     params: {
-      query: { inference_id: inferenceId },
+      query: { inference_id: inferenceId, wait_until_complete: true },
     },
   });
 }

From d96b904eea82e6f8df2496b5d7b9d884e8debb1c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=B8ren=20Louv-Jansen?= <soren.louv@elastic.co>
Date: Fri, 9 May 2025 12:07:15 +0200
Subject: [PATCH 3/3] Remove focused test

---
 .../ai_assistant/complete/functions/recall.spec.ts            | 4 ++--
 .../apis/observability/ai_assistant/utils/knowledge_base.ts   | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/recall.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/recall.spec.ts
index ae3f8354029e8..db2218a3a76cb 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/recall.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/recall.spec.ts
@@ -25,7 +25,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
   const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
   const es = getService('es');
 
-  describe.only('recall', function () {
+  describe('recall', function () {
     before(async () => {
       await deployTinyElserAndSetupKb(getService);
       await addSampleDocsToInternalKb(getService, technicalSampleDocs);
@@ -40,7 +40,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     });
 
     describe('GET /internal/observability_ai_assistant/functions/recall', () => {
-      it.only('produces unique scores for each doc', async () => {
+      it('produces unique scores for each doc', async () => {
         const entries = await recall('What happened during the database outage?');
         const uniqueScores = uniq(entries.map(({ esScore }) => esScore));
         expect(uniqueScores.length).to.be.greaterThan(1);
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts
index dd0a77eafaa6a..6abc77856a1e5 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts
@@ -75,7 +75,7 @@ export async function setupKnowledgeBase(
   const statusResult = await getKnowledgeBaseStatus(observabilityAIAssistantAPIClient);
 
   log.debug(
-    `Setting up knowledge base with inference endpoint = "${TINY_ELSER_INFERENCE_ID}", concreteWriteIndex = ${statusResult.body.concreteWriteIndex}, currentInferenceId = ${statusResult.body.currentInferenceId} ${statusResult.body.isReIndexing}`
+    `Setting up knowledge base with inference endpoint = "${TINY_ELSER_INFERENCE_ID}", concreteWriteIndex = ${statusResult.body.concreteWriteIndex}, currentInferenceId = ${statusResult.body.currentInferenceId}, isReIndexing = ${statusResult.body.isReIndexing}`
   );
   return observabilityAIAssistantAPIClient.admin({
     endpoint: 'POST /internal/observability_ai_assistant/kb/setup',