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

…ces initialization (#198239) (cherry picked from commit ed81e43)
elastic · Oct 30, 2024 · e81d234 · e81d234
1 parent 2064a2c
commit e81d234
Show file tree

Hide file tree

Showing 36 changed files with 283 additions and 215 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,
       // Polling interval for Knowledge Base setup in progress

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
@@ -57,6 +57,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
@@ -52,12 +52,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 });
+  const { data: kbStatus } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled });
   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,8 +44,16 @@ interface Props {
  */
 export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
   ({ knowledgeBase, setUpdatedKnowledgeBaseSettings, modalMode = false }) => {
-    const { http, toasts } = useAssistantContext();
-    const { data: kbStatus, isLoading, isFetching } = useKnowledgeBaseStatus({ http });
+    const {
+      http,
+      toasts,
+      assistantAvailability: { isAssistantEnabled },
+    } = useAssistantContext();
+    const {
+      data: kbStatus,
+      isLoading,
+      isFetching,
+    } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled });
     const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts });
 
     // Resource enabled state

diff --git a/...es/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx b/...es/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx
@@ -74,12 +74,15 @@ interface Params {
 export const KnowledgeBaseSettingsManagement: React.FC<Params> = React.memo(({ dataViews }) => {
   const {
     assistantFeatures: { assistantKnowledgeBaseByDefault: enableKnowledgeBaseByDefault },
-    assistantAvailability: { hasManageGlobalKnowledgeBase },
+    assistantAvailability: { hasManageGlobalKnowledgeBase, isAssistantEnabled },
     http,
     toasts,
   } = useAssistantContext();
   const [hasPendingChanges, setHasPendingChanges] = useState(false);
-  const { data: kbStatus, isFetched } = useKnowledgeBaseStatus({ http });
+  const { data: kbStatus, isFetched } = useKnowledgeBaseStatus({
+    http,
+    enabled: isAssistantEnabled,
+  });
   const isKbSetup = isKnowledgeBaseSetup(kbStatus);
 
   const [deleteKBItem, setDeleteKBItem] = useState<DocumentEntry | IndexEntry | null>(null);
@@ -159,7 +162,7 @@ export const KnowledgeBaseSettingsManagement: React.FC<Params> = React.memo(({ d
   } = useKnowledgeBaseEntries({
     http,
     toasts,
-    enabled: enableKnowledgeBaseByDefault,
+    enabled: enableKnowledgeBaseByDefault && isAssistantEnabled,
     isRefetching: kbStatus?.is_setup_in_progress,
   });
 

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
@@ -23,9 +23,13 @@ interface Props {
  *
  */
 export const SetupKnowledgeBaseButton: React.FC<Props> = React.memo(({ display }: Props) => {
-  const { http, toasts } = useAssistantContext();
+  const {
+    http,
+    toasts,
+    assistantAvailability: { isAssistantEnabled },
+  } = useAssistantContext();
 
-  const { data: kbStatus } = useKnowledgeBaseStatus({ http });
+  const { data: kbStatus } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled });
   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,11 +18,17 @@ 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(),
 }));
 
+const licensing = Promise.resolve(
+  licensingMock.createRequestHandlerContext({
+    license: { type: 'enterprise' },
+  })
+);
 let logger: ReturnType<(typeof loggingSystemMock)['createLogger']>;
 const clusterClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
 
@@ -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
@@ -11,6 +11,7 @@ import type { AuthenticatedUser, Logger, ElasticsearchClient } from '@kbn/core/s
 import type { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server';
 import type { MlPluginSetup } from '@kbn/ml-plugin/server';
 import { Subject } from 'rxjs';
+import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server';
 import { attackDiscoveryFieldMap } from '../lib/attack_discovery/persistence/field_maps_configuration/field_maps_configuration';
 import { getDefaultAnonymizationFields } from '../../common/anonymization';
 import { AssistantResourceNames, GetElser } from '../types';
@@ -36,6 +37,7 @@ import {
 } from '../ai_assistant_data_clients/knowledge_base';
 import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence';
 import { createGetElserId, createPipeline, pipelineExists } from './helpers';
+import { hasAIAssistantLicense } from '../routes/helpers';
 
 const TOTAL_FIELDS_LIMIT = 2500;
 
@@ -56,6 +58,7 @@ export interface CreateAIAssistantClientParams {
   logger: Logger;
   spaceId: string;
   currentUser: AuthenticatedUser | null;
+  licensing: Promise<LicensingApiRequestHandlerContext>;
 }
 
 export type CreateDataStream = (params: {
@@ -245,7 +248,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);
@@ -290,6 +293,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
@@ -510,7 +515,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;
@@ -162,22 +162,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();
 
@@ -199,20 +195,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 }),