diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.test.ts index 8e34581332ff6..0316a33b1d7c2 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.test.ts +++ b/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 => { +): Promise => { if (timeoutMs < 0) { throw new Error('fail'); } diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.ts index 39e0e69a8fc49..371974ac9a20a 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.ts +++ b/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, + initPromise?: Promise, timeoutMs?: number ) => void; - getInitializedResources: (namespace: string) => Promise; + getInitializedResources: (namespace: string) => Promise; } /** @@ -42,17 +42,17 @@ export interface ResourceInstallationHelper { */ export function createResourceInstallationHelper( logger: Logger, - commonResourcesInitPromise: Promise, + commonResourcesInitPromise: Promise, installFn: (namespace: string, timeoutMs?: number) => Promise ): ResourceInstallationHelper { - let commonInitPromise: Promise = commonResourcesInitPromise; - const initializedResources: Map> = new Map(); + let commonInitPromise: Promise = commonResourcesInitPromise; + const initializedResources: Map> = new Map(); const lastRetry: Map = new Map(); const waitUntilResourcesInstalled = async ( namespace: string = DEFAULT_NAMESPACE_STRING, timeoutMs?: number - ): Promise => { + ): Promise => { 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, + initPromise?: Promise, timeoutMs?: number ) => { const key = namespace; @@ -109,19 +109,23 @@ export function createResourceInstallationHelper( ); } }, - getInitializedResources: async (namespace: string): Promise => { - const key = namespace; - return ( - initializedResources.has(key) - ? initializedResources.get(key) - : errorResult(`Unrecognized spaceId ${key}`) - ) as InitializationPromise; + getInitializedResources: async (namespace: string): Promise => { + 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(); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index bfdf8b96f44b0..f26ce8b7c1189 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/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; + private initPromise: Promise; 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 { + private async initializeResources(): Promise { 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 | undefined; + let initRetryPromise: Promise | 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 { + ): Promise { 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