[Obs AI Assistant] Improve flaky recall tests

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);
   },
 });
 

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 };
   };
 

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}`);

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)
+  );
+}

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 });

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,14 +26,26 @@ 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 } },
         { add: { index: nextWriteIndexName, alias, is_write_index: true } },
       ],
     });
   } 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.`
+  );
 }

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);

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,19 +150,14 @@ 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) {
     return observabilityAIAssistantAPIClient.viewer({
       endpoint: 'POST /internal/observability_ai_assistant/kb/setup',
       params: {
-        query: { inference_id: inferenceId },
+        query: { inference_id: inferenceId, wait_until_complete: true },
       },
     });
   }

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) {

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,13 +66,21 @@ 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}, isReIndexing = ${statusResult.body.isReIndexing}`
+  );
   return observabilityAIAssistantAPIClient.admin({
     endpoint: 'POST /internal/observability_ai_assistant/kb/setup',
     params: {
-      query: { inference_id: inferenceId },
+      query: { inference_id: inferenceId, wait_until_complete: true },
     },
   });
 }
@@ -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

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 };