[Obs AI Assistant] Replace Task Manager usage with Lock Manager

x-pack/platform/plugins/shared/observability_ai_assistant/server/plugin.ts

@@ -31,8 +31,8 @@ import { registerFunctions } from './functions';
 import { recallRankingEvent } from './analytics/recall_ranking';
 import { initLangtrace } from './service/client/instrumentation/init_langtrace';
 import { aiAssistantCapabilities } from '../common/capabilities';
-import { registerAndScheduleKbSemanticTextMigrationTask } from './service/task_manager_definitions/register_kb_semantic_text_migration_task';
-import { updateExistingIndexAssets } from './service/create_or_update_index_assets';
+import { populateMissingSemanticTextFieldMigration } from './service/startup_migrations/populate_missing_semantic_text_field_migration';
+import { updateExistingIndexAssets } from './service/startup_migrations/create_or_update_index_assets';
 
 export class ObservabilityAIAssistantPlugin
   implements
@@ -130,23 +130,19 @@ export class ObservabilityAIAssistantPlugin
     }));
 
     // Update existing index assets (mappings, templates, etc). This will not create assets if they do not exist.
-    const indexAssetsUpdatedPromise = updateExistingIndexAssets({
-      logger: this.logger.get('index_assets'),
-      core,
-    }).catch((e) => this.logger.error(`Index assets could not be updated: ${e.message}`));
-
-    // register task to migrate knowledge base entries to include semantic_text field
-    registerAndScheduleKbSemanticTextMigrationTask({
-      core,
-      taskManager: plugins.taskManager,
-      logger: this.logger.get('kb_semantic_text_migration_task'),
-      config: this.config,
-      indexAssetsUpdatedPromise,
-    }).catch((e) =>
-      this.logger.error(
-        `Knowledge base semantic_text migration task could not be registered: ${e.message}`
+    updateExistingIndexAssets({ logger: this.logger, core })
+      .then(() =>
+        populateMissingSemanticTextFieldMigration({
+          core,
+          logger: this.logger,
+          config: this.config,
+        })
       )
-    );
+      .catch((e) =>
+        this.logger.error(
+          `Error during knowledge base migration in AI Assistant plugin startup: ${e.message}`
+        )
+      );
 
     service.register(registerFunctions);
 

x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/knowledge_base/route.ts

@@ -100,8 +100,23 @@ const resetKnowledgeBase = createObservabilityAIAssistantServerRoute({
   },
 });
 
+const reIndexKnowledgeBase = createObservabilityAIAssistantServerRoute({
+  endpoint: 'POST /internal/observability_ai_assistant/kb/reindex',
+  security: {
+    authz: {
+      requiredPrivileges: ['ai_assistant'],
+    },
+  },
+  handler: async (resources): Promise<{ result: boolean }> => {
+    const client = await resources.service.getClient({ request: resources.request });
+    const result = await client.reIndexKnowledgeBaseWithLock();
+    return { result };
+  },
+});
+
 const semanticTextMigrationKnowledgeBase = createObservabilityAIAssistantServerRoute({
-  endpoint: 'POST /internal/observability_ai_assistant/kb/migrations/kb_semantic_text',
+  endpoint:
+    'POST /internal/observability_ai_assistant/kb/migrations/populate_missing_semantic_text_field',
   security: {
     authz: {
       requiredPrivileges: ['ai_assistant'],
@@ -320,6 +335,7 @@ const importKnowledgeBaseEntries = createObservabilityAIAssistantServerRoute({
 });
 
 export const knowledgeBaseRoutes = {
+  ...reIndexKnowledgeBase,
   ...semanticTextMigrationKnowledgeBase,
   ...setupKnowledgeBase,
   ...resetKnowledgeBase,

x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/top_level/route.ts

@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import { createOrUpdateIndexAssets } from '../../service/create_or_update_index_assets';
+import { createOrUpdateIndexAssets } from '../../service/startup_migrations/create_or_update_index_assets';
 import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route';
 
 const createOrUpdateIndexAssetsRoute = createObservabilityAIAssistantServerRoute({

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts

@@ -67,14 +67,12 @@ import { continueConversation } from './operators/continue_conversation';
 import { convertInferenceEventsToStreamingEvents } from './operators/convert_inference_events_to_streaming_events';
 import { extractMessages } from './operators/extract_messages';
 import { getGeneratedTitle } from './operators/get_generated_title';
-import {
-  reIndexKnowledgeBaseAndPopulateSemanticTextField,
-  scheduleKbSemanticTextMigrationTask,
-} from '../task_manager_definitions/register_kb_semantic_text_migration_task';
+import { populateMissingSemanticTextFieldMigration } from '../startup_migrations/populate_missing_semantic_text_field_migration';
 import { ObservabilityAIAssistantPluginStartDependencies } from '../../types';
 import { ObservabilityAIAssistantConfig } from '../../config';
 import { getElserModelId } from '../knowledge_base_service/get_elser_model_id';
 import { apmInstrumentation } from './operators/apm_instrumentation';
+import { reIndexKnowledgeBaseWithLock } from '../knowledge_base_service/reindex_knowledge_base';
 
 const MAX_FUNCTION_CALLS = 8;
 
@@ -680,14 +678,15 @@ export class ObservabilityAIAssistantClient {
     // setup the knowledge base
     const res = await knowledgeBaseService.setup(esClient, modelId);
 
-    core
-      .getStartServices()
-      .then(([_, pluginsStart]) =>
-        scheduleKbSemanticTextMigrationTask({ taskManager: pluginsStart.taskManager, logger })
-      )
-      .catch((error) => {
-        logger.error(`Failed to schedule semantic text migration task: ${error}`);
-      });
+    populateMissingSemanticTextFieldMigration({
+      core,
+      logger,
+      config: this.dependencies.config,
+    }).catch((e) => {
+      this.dependencies.logger.error(
+        `Failed to populate missing semantic text fields: ${e.message}`
+      );
+    });
 
     return res;
   };
@@ -697,14 +696,21 @@ export class ObservabilityAIAssistantClient {
     return this.dependencies.knowledgeBaseService.reset(esClient);
   };
 
-  reIndexKnowledgeBaseAndPopulateSemanticTextField = () => {
-    return reIndexKnowledgeBaseAndPopulateSemanticTextField({
+  reIndexKnowledgeBaseWithLock = () => {
+    return reIndexKnowledgeBaseWithLock({
+      core: this.dependencies.core,
       esClient: this.dependencies.esClient,
       logger: this.dependencies.logger,
-      config: this.dependencies.config,
     });
   };
 
+  reIndexKnowledgeBaseAndPopulateSemanticTextField = () => {
+    return populateMissingSemanticTextFieldMigration({
+      core: this.dependencies.core,
+      logger: this.dependencies.logger,
+      config: this.dependencies.config,
+    });
+  };
   addUserInstruction = async ({
     entry,
   }: {