[Security AI Assistant] Fixed license issue for Knowledge Base resour…

…ces initialization (elastic#198239) (cherry picked from commit ed81e43) # Conflicts: # x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx # x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx # x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx # x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx # x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts # x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts # x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts # x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts # x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts
stephmilovic · Oct 30, 2024 · ea2cee5 · ea2cee5
1 parent 3219dc8
commit ea2cee5
Show file tree

Hide file tree

Showing 35 changed files with 274 additions and 213 deletions.
diff --git a/...ges/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx b/...ges/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx
@@ -20,6 +20,7 @@ export interface UseKnowledgeBaseStatusParams {
   http: HttpSetup;
   resource?: string;
   toasts?: IToasts;
+  enabled: boolean;
 }
 
 /**
@@ -36,13 +37,15 @@ export const useKnowledgeBaseStatus = ({
   http,
   resource,
   toasts,
+  enabled,
 }: UseKnowledgeBaseStatusParams): UseQueryResult<ReadKnowledgeBaseResponse, IHttpFetchError> => {
   return useQuery(
     KNOWLEDGE_BASE_STATUS_QUERY_KEY,
     async ({ signal }) => {
       return getKnowledgeBaseStatus({ http, resource, signal });
     },
     {
+      enabled,
       retry: false,
       keepPreviousData: true,
       // Deprecated, hoist to `queryCache` w/in `QueryClient. See: https://stackoverflow.com/a/76961109

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx
@@ -60,6 +60,9 @@ describe('use chat send', () => {
       assistantTelemetry: {
         reportAssistantMessageSent,
       },
+      assistantAvailability: {
+        isAssistantEnabled: true,
+      },
     });
   });
   it('handleOnChatCleared clears the conversation', async () => {

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx
@@ -57,12 +57,16 @@ export const useChatSend = ({
   setSelectedPromptContexts,
   setCurrentConversation,
 }: UseChatSendProps): UseChatSend => {
-  const { assistantTelemetry, toasts } = useAssistantContext();
+  const {
+    assistantTelemetry,
+    toasts,
+    assistantAvailability: { isAssistantEnabled },
+  } = useAssistantContext();
   const [userPrompt, setUserPrompt] = useState<string | null>(null);
 
   const { isLoading, sendMessage, abortStream } = useSendMessage();
   const { clearConversation, removeLastMessage } = useConversation();
-  const { data: kbStatus } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE });
+  const { data: kbStatus } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled , resource: ESQL_RESOURCE });
   const isSetupComplete =
     kbStatus?.elser_exists &&
     kbStatus?.index_exists &&

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx
@@ -26,6 +26,9 @@ const mockUseAssistantContext = {
   },
   setAllSystemPrompts: jest.fn(),
   setConversations: jest.fn(),
+  assistantAvailability: {
+    isAssistantEnabled: true,
+  },
 };
 
 jest.mock('../assistant_context', () => {

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx
@@ -44,12 +44,16 @@ interface Props {
  */
 export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
   ({ knowledgeBase, setUpdatedKnowledgeBaseSettings }) => {
-    const { http, toasts } = useAssistantContext();
+    const {
+      http,
+      toasts,
+      assistantAvailability: { isAssistantEnabled },
+    } = useAssistantContext();
     const {
       data: kbStatus,
       isLoading,
       isFetching,
-    } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE });
+    } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled, resource: ESQL_RESOURCE });
     const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts });
 
     // Resource enabled state

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx
@@ -20,9 +20,13 @@ export const ESQL_RESOURCE = 'esql';
  *
  */
 export const SetupKnowledgeBaseButton: React.FC = React.memo(() => {
-  const { http, toasts } = useAssistantContext();
+  const {
+    http,
+    toasts,
+    assistantAvailability: { isAssistantEnabled },
+  } = useAssistantContext();
 
-  const { data: kbStatus } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE });
+  const { data: kbStatus } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled , resource: ESQL_RESOURCE });
   const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts });
 
   const isSetupInProgress = kbStatus?.is_setup_in_progress || isSettingUpKB;

diff --git a/...elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.test.ts b/...elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.test.ts
@@ -141,7 +141,7 @@ describe('createResourceInstallationHelper', () => {
       async () => (await getContextInitialized(helper)) === false
     );
 
-    expect(logger.error).toHaveBeenCalledWith(`Error initializing resources test1 - fail`);
+    expect(logger.warn).toHaveBeenCalledWith(`Error initializing resources test1 - fail`);
     expect(await helper.getInitializedResources('test1')).toEqual({
       result: false,
       error: `fail`,
@@ -204,7 +204,7 @@ describe('createResourceInstallationHelper', () => {
       async () => (await getContextInitialized(helper)) === false
     );
 
-    expect(logger.error).toHaveBeenCalledWith(`Error initializing resources default - first error`);
+    expect(logger.warn).toHaveBeenCalledWith(`Error initializing resources default - first error`);
     expect(await helper.getInitializedResources(DEFAULT_NAMESPACE_STRING)).toEqual({
       result: false,
       error: `first error`,
@@ -221,9 +221,7 @@ describe('createResourceInstallationHelper', () => {
       return logger.error.mock.calls.length === 1;
     });
 
-    expect(logger.error).toHaveBeenCalledWith(
-      `Error initializing resources default - second error`
-    );
+    expect(logger.warn).toHaveBeenCalledWith(`Error initializing resources default - second error`);
 
     // the second retry is throttled so this is never called
     expect(logger.info).not.toHaveBeenCalledWith('test1_default successfully retried');

diff --git a/...gins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.ts b/...gins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.ts
@@ -65,7 +65,7 @@ export function createResourceInstallationHelper(
         return errorResult(commonInitError);
       }
     } catch (err) {
-      logger.error(`Error initializing resources ${namespace} - ${err.message}`);
+      logger.warn(`Error initializing resources ${namespace} - ${err.message}`);
       return errorResult(err.message);
     }
   };
@@ -113,7 +113,7 @@ export function createResourceInstallationHelper(
       const key = namespace;
       return (
         initializedResources.has(key)
-          ? initializedResources.get(key)
+          ? await initializedResources.get(key)
           : errorResult(`Unrecognized spaceId ${key}`)
       ) as InitializationPromise;
     },

diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts
@@ -18,12 +18,18 @@ import { AIAssistantService, AIAssistantServiceOpts } from '.';
 import { retryUntil } from './create_resource_installation_helper.test';
 import { mlPluginMock } from '@kbn/ml-plugin/public/mocks';
 import type { MlPluginSetup } from '@kbn/ml-plugin/server';
+import { licensingMock } from '@kbn/licensing-plugin/server/mocks';
 
 jest.mock('../ai_assistant_data_clients/conversations', () => ({
   AIAssistantConversationsDataClient: jest.fn(),
 }));
 
-let logger: ReturnType<typeof loggingSystemMock['createLogger']>;
+const licensing = Promise.resolve(
+  licensingMock.createRequestHandlerContext({
+    license: { type: 'enterprise' },
+  })
+);
+let logger: ReturnType<(typeof loggingSystemMock)['createLogger']>;
 const clusterClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
 
 const SimulateTemplateResponse = {
@@ -191,6 +197,7 @@ describe('AI Assistant Service', () => {
         logger,
         spaceId: 'default',
         currentUser: mockUser1,
+        licensing,
       });
 
       expect(AIAssistantConversationsDataClient).toHaveBeenCalledWith({
@@ -221,6 +228,7 @@ describe('AI Assistant Service', () => {
         logger,
         spaceId: 'default',
         currentUser: mockUser1,
+        licensing,
       });
 
       expect(clusterClient.indices.putIndexTemplate).toHaveBeenCalled();
@@ -274,11 +282,13 @@ describe('AI Assistant Service', () => {
           logger,
           spaceId: 'default',
           currentUser: mockUser1,
+          licensing,
         }),
         assistantService.createAIAssistantConversationsDataClient({
           logger,
           spaceId: 'default',
           currentUser: mockUser1,
+          licensing,
         }),
       ]);
 
@@ -340,6 +350,7 @@ describe('AI Assistant Service', () => {
         logger,
         spaceId: 'default',
         currentUser: mockUser1,
+        licensing,
       });
 
       expect(AIAssistantConversationsDataClient).toHaveBeenCalledWith({
@@ -400,6 +411,7 @@ describe('AI Assistant Service', () => {
           logger,
           spaceId: 'default',
           currentUser: mockUser1,
+          licensing,
         });
       };
 
@@ -472,6 +484,7 @@ describe('AI Assistant Service', () => {
           logger,
           spaceId: 'default',
           currentUser: mockUser1,
+          licensing,
         });
       };
 
@@ -513,6 +526,7 @@ describe('AI Assistant Service', () => {
         logger,
         spaceId: 'test',
         currentUser: mockUser1,
+        licensing,
       });
 
       expect(clusterClient.indices.putIndexTemplate).not.toHaveBeenCalled();
@@ -560,6 +574,7 @@ describe('AI Assistant Service', () => {
         logger,
         spaceId: 'test',
         currentUser: mockUser1,
+        licensing,
       });
 
       expect(clusterClient.indices.putIndexTemplate).not.toHaveBeenCalled();
@@ -607,6 +622,7 @@ describe('AI Assistant Service', () => {
         logger,
         spaceId: 'test',
         currentUser: mockUser1,
+        licensing,
       });
 
       expect(AIAssistantConversationsDataClient).not.toHaveBeenCalled();
@@ -752,6 +768,7 @@ describe('AI Assistant Service', () => {
         logger,
         spaceId: 'default',
         currentUser: mockUser1,
+        licensing,
       });
 
       await retryUntil(

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
@@ -12,6 +12,7 @@ import type { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server';
 import type { MlPluginSetup } from '@kbn/ml-plugin/server';
 import { Subject } from 'rxjs';
 import { attackDiscoveryFieldMap } from '../ai_assistant_data_clients/attack_discovery/field_maps_configuration';
+import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server';
 import { getDefaultAnonymizationFields } from '../../common/anonymization';
 import { AssistantResourceNames, GetElser } from '../types';
 import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations';
@@ -30,6 +31,7 @@ import { knowledgeBaseFieldMap } from '../ai_assistant_data_clients/knowledge_ba
 import { AIAssistantKnowledgeBaseDataClient } from '../ai_assistant_data_clients/knowledge_base';
 import { AttackDiscoveryDataClient } from '../ai_assistant_data_clients/attack_discovery';
 import { createGetElserId, createPipeline, pipelineExists } from './helpers';
+import { hasAIAssistantLicense } from '../routes/helpers';
 
 const TOTAL_FIELDS_LIMIT = 2500;
 
@@ -50,6 +52,7 @@ export interface CreateAIAssistantClientParams {
   logger: Logger;
   spaceId: string;
   currentUser: AuthenticatedUser | null;
+  licensing: Promise<LicensingApiRequestHandlerContext>;
 }
 
 export type CreateDataStream = (params: {
@@ -217,7 +220,7 @@ export class AIAssistantService {
         pluginStop$: this.options.pluginStop$,
       });
     } catch (error) {
-      this.options.logger.error(`Error initializing AI assistant resources: ${error.message}`);
+      this.options.logger.warn(`Error initializing AI assistant resources: ${error.message}`);
       this.initialized = false;
       this.isInitializing = false;
       return errorResult(error.message);
@@ -262,6 +265,8 @@ export class AIAssistantService {
   };
 
   private async checkResourcesInstallation(opts: CreateAIAssistantClientParams) {
+    const licensing = await opts.licensing;
+    if (!hasAIAssistantLicense(licensing.license)) return null;
     // Check if resources installation has succeeded
     const { result: initialized, error } = await this.getSpaceResourcesInitializationPromise(
       opts.spaceId
@@ -456,7 +461,7 @@ export class AIAssistantService {
         await this.createDefaultAnonymizationFields(spaceId);
       }
     } catch (error) {
-      this.options.logger.error(
+      this.options.logger.warn(
         `Error initializing AI assistant namespace level resources: ${error.message}`
       );
       throw error;

diff --git a/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/bulk_actions_route.ts b/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/bulk_actions_route.ts
@@ -38,7 +38,7 @@ import {
   EsAnonymizationFieldsSchema,
   UpdateAnonymizationFieldSchema,
 } from '../../ai_assistant_data_clients/anonymization_fields/types';
-import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../helpers';
+import { performChecks } from '../helpers';
 
 export interface BulkOperationError {
   message: string;
@@ -158,22 +158,18 @@ export const bulkActionAnonymizationFieldsRoute = (
         request.events.completed$.subscribe(() => abortController.abort());
         try {
           const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
-          const license = ctx.licensing.license;
-          if (!hasAIAssistantLicense(license)) {
-            return response.forbidden({
-              body: {
-                message: UPGRADE_LICENSE_MESSAGE,
-              },
-            });
-          }
+          // Perform license and authenticated user checks
+          const checkResponse = performChecks({
+            context: ctx,
+            request,
+            response,
+          });
 
-          const authenticatedUser = ctx.elasticAssistant.getCurrentUser();
-          if (authenticatedUser == null) {
-            return assistantResponse.error({
-              body: `Authenticated user not found`,
-              statusCode: 401,
-            });
+          if (!checkResponse.isSuccess) {
+            return checkResponse.response;
           }
+          const authenticatedUser = checkResponse.currentUser;
+
           const dataClient =
             await ctx.elasticAssistant.getAIAssistantAnonymizationFieldsDataClient();
 
@@ -195,20 +191,20 @@ export const bulkActionAnonymizationFieldsRoute = (
           }
 
           const writer = await dataClient?.getWriter();
-          const changedAt = new Date().toISOString();
+          const createdAt = new Date().toISOString();
           const {
             errors,
             docs_created: docsCreated,
             docs_updated: docsUpdated,
             docs_deleted: docsDeleted,
             // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
           } = await writer!.bulk({
-            documentsToCreate: body.create?.map((f) =>
-              transformToCreateScheme(authenticatedUser, changedAt, f)
+            documentsToCreate: body.create?.map((doc) =>
+              transformToCreateScheme(authenticatedUser, createdAt, doc)
             ),
             documentsToDelete: body.delete?.ids,
-            documentsToUpdate: body.update?.map((f) =>
-              transformToUpdateScheme(authenticatedUser, changedAt, f)
+            documentsToUpdate: body.update?.map((doc) =>
+              transformToUpdateScheme(authenticatedUser, createdAt, doc)
             ),
             getUpdateScript: (document: UpdateAnonymizationFieldSchema) =>
               getUpdateScript({ anonymizationField: document, isPatch: true }),