[Security solution] Data stream resource initialization helper

x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.test.ts

@@ -11,7 +11,7 @@ import { loggingSystemMock } from '@kbn/core/server/mocks';
 import {
   createResourceInstallationHelper,
   errorResult,
-  InitializationPromise,
+  InitializationState,
   ResourceInstallationHelper,
   successResult,
   calculateDelay,
@@ -54,7 +54,7 @@ const getCommonInitPromise = async (
   resolution: boolean,
   timeoutMs: number = 1,
   customLogString: string = ''
-): Promise<InitializationPromise> => {
+): Promise<InitializationState> => {
   if (timeoutMs < 0) {
     throw new Error('fail');
   }

x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.ts

@@ -8,7 +8,7 @@
 import { DEFAULT_NAMESPACE_STRING } from '@kbn/core-saved-objects-utils-server';
 import { Logger } from '@kbn/core/server';
 
-export interface InitializationPromise {
+export interface InitializationState {
   result: boolean;
   error?: string;
 }
@@ -23,10 +23,10 @@ export interface ResourceInstallationHelper {
   add: (namespace?: string, timeoutMs?: number) => void;
   retry: (
     namespace?: string,
-    initPromise?: Promise<InitializationPromise>,
+    initPromise?: Promise<InitializationState>,
     timeoutMs?: number
   ) => void;
-  getInitializedResources: (namespace: string) => Promise<InitializationPromise>;
+  getInitializedResources: (namespace: string) => Promise<InitializationState>;
 }
 
 /**
@@ -42,17 +42,17 @@ export interface ResourceInstallationHelper {
  */
 export function createResourceInstallationHelper(
   logger: Logger,
-  commonResourcesInitPromise: Promise<InitializationPromise>,
+  commonResourcesInitPromise: Promise<InitializationState>,
   installFn: (namespace: string, timeoutMs?: number) => Promise<void>
 ): ResourceInstallationHelper {
-  let commonInitPromise: Promise<InitializationPromise> = commonResourcesInitPromise;
-  const initializedResources: Map<string, Promise<InitializationPromise>> = new Map();
+  let commonInitPromise: Promise<InitializationState> = commonResourcesInitPromise;
+  const initializedResources: Map<string, Promise<InitializationState>> = new Map();
   const lastRetry: Map<string, Retry> = new Map();
 
   const waitUntilResourcesInstalled = async (
     namespace: string = DEFAULT_NAMESPACE_STRING,
     timeoutMs?: number
-  ): Promise<InitializationPromise> => {
+  ): Promise<InitializationState> => {
     try {
       const { result: commonInitResult, error: commonInitError } = await commonInitPromise;
       if (commonInitResult) {
@@ -81,7 +81,7 @@ export function createResourceInstallationHelper(
     },
     retry: (
       namespace: string = DEFAULT_NAMESPACE_STRING,
-      initPromise?: Promise<InitializationPromise>,
+      initPromise?: Promise<InitializationState>,
       timeoutMs?: number
     ) => {
       const key = namespace;
@@ -109,19 +109,23 @@ export function createResourceInstallationHelper(
         );
       }
     },
-    getInitializedResources: async (namespace: string): Promise<InitializationPromise> => {
-      const key = namespace;
-      return (
-        initializedResources.has(key)
-          ? initializedResources.get(key)
-          : errorResult(`Unrecognized spaceId ${key}`)
-      ) as InitializationPromise;
+    getInitializedResources: async (namespace: string): Promise<InitializationState> => {
+      try {
+        const key = namespace;
+        if (!initializedResources.has(key)) {
+          return errorResult(`Unrecognized spaceId ${key}`);
+        }
+        const initializationState = await initializedResources.get(key);
+        return initializationState ?? errorResult(`Initialize state returned undefined`);
+      } catch (e) {
+        return errorResult(`Initialize resources encountered an error: ${e.message}`);
+      }
     },
   };
 }
 
 export const successResult = () => ({ result: true });
-export const errorResult = (error?: string) => ({ result: false, error });
+export const errorResult = (error?: string): InitializationState => ({ result: false, error });
 
 export const getShouldRetry = ({ time, attempts }: Retry) => {
   const now = new Date().valueOf();

x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts

@@ -16,7 +16,7 @@ import { getDefaultAnonymizationFields } from '../../common/anonymization';
 import { AssistantResourceNames, GetElser } from '../types';
 import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations';
 import {
-  InitializationPromise,
+  InitializationState,
   ResourceInstallationHelper,
   createResourceInstallationHelper,
   errorResult,
@@ -80,7 +80,7 @@ export class AIAssistantService {
   private anonymizationFieldsDataStream: DataStreamSpacesAdapter;
   private attackDiscoveryDataStream: DataStreamSpacesAdapter;
   private resourceInitializationHelper: ResourceInstallationHelper;
-  private initPromise: Promise<InitializationPromise>;
+  private initPromise: Promise<InitializationState>;
   private isKBSetupInProgress: boolean = false;
   // Temporary 'feature flag' to determine if we should initialize the new kb mappings, toggled when accessing kbDataClient
   private v2KnowledgeBaseEnabled: boolean = false;
@@ -168,7 +168,7 @@ export class AIAssistantService {
     return newDataStream;
   };
 
-  private async initializeResources(): Promise<InitializationPromise> {
+  private async initializeResources(): Promise<InitializationState> {
     this.isInitializing = true;
     try {
       this.options.logger.debug(`Initializing resources for AIAssistantService`);
@@ -283,13 +283,13 @@ export class AIAssistantService {
 
   private async checkResourcesInstallation(opts: CreateAIAssistantClientParams) {
     // Check if resources installation has succeeded
-    const { result: initialized, error } = await this.getSpaceResourcesInitializationPromise(
+    const { result: initialized, error } = await this.getSpaceResourcesInitializationState(
       opts.spaceId
     );
 
     // If space level resources initialization failed, retry
     if (!initialized && error) {
-      let initRetryPromise: Promise<InitializationPromise> | undefined;
+      let initRetryPromise: Promise<InitializationState> | undefined;
 
       // If !this.initialized, we know that resource initialization failed
       // and we need to retry this before retrying the spaceId specific resources
@@ -453,9 +453,9 @@ export class AIAssistantService {
     });
   }
 
-  public async getSpaceResourcesInitializationPromise(
+  public async getSpaceResourcesInitializationState(
     spaceId: string | undefined = DEFAULT_NAMESPACE_STRING
-  ): Promise<InitializationPromise> {
+  ): Promise<InitializationState> {
     const result = await this.resourceInitializationHelper.getInitializedResources(spaceId);
     // If the spaceId is unrecognized and spaceId is not the default, we
     // need to kick off resource installation and return the promise