[Security AI Assistant] Fixed license issue for Knowledge Base resources initialization

x-pack/packages/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

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 () => {

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 &&

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', () => {

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

x-pack/packages/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,
   });
 
   // Flyout Save/Cancel Actions

x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx

@@ -22,9 +22,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;

x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/anonymization_fields/helpers.ts

@@ -62,30 +62,30 @@ export const transformESSearchToAnonymizationFields = (
 };
 
 export const transformToUpdateScheme = (
-  user: AuthenticatedUser,
+  user: AuthenticatedUser | null,
   updatedAt: string,
   { allowed, anonymized, id }: AnonymizationFieldUpdateProps
 ): UpdateAnonymizationFieldSchema => {
   return {
     id,
     updated_at: updatedAt,
-    updated_by: user.username,
+    updated_by: user?.username,
     allowed,
     anonymized,
   };
 };
 
 export const transformToCreateScheme = (
-  user: AuthenticatedUser,
   createdAt: string,
-  { allowed, anonymized, field }: AnonymizationFieldCreateProps
+  { allowed, anonymized, field }: AnonymizationFieldCreateProps,
+  user: AuthenticatedUser | null
 ): CreateAnonymizationFieldSchema => {
   return {
     '@timestamp': createdAt,
     updated_at: createdAt,
     field,
     created_at: createdAt,
-    created_by: user.username,
+    created_by: user?.username,
     allowed,
     anonymized,
   };

x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/prompts/helpers.ts

@@ -81,7 +81,7 @@ export const transformESSearchToPrompts = (
 };
 
 export const transformToUpdateScheme = (
-  user: AuthenticatedUser,
+  user: AuthenticatedUser | null,
   updatedAt: string,
   { content, isNewConversationDefault, categories, color, id }: PromptUpdateProps
 ): UpdatePromptSchema => {
@@ -92,17 +92,19 @@ export const transformToUpdateScheme = (
     is_new_conversation_default: isNewConversationDefault,
     categories,
     color,
-    users: [
-      {
-        id: user.profile_uid,
-        name: user.username,
-      },
-    ],
+    users: user
+      ? [
+          {
+            id: user.profile_uid,
+            name: user.username,
+          },
+        ]
+      : undefined,
   };
 };
 
 export const transformToCreateScheme = (
-  user: AuthenticatedUser,
+  user: AuthenticatedUser | null,
   updatedAt: string,
   {
     content,
@@ -126,12 +128,14 @@ export const transformToCreateScheme = (
     name,
     is_default: isDefault,
     prompt_type: promptType,
-    users: [
-      {
-        id: user.profile_uid,
-        name: user.username,
-      },
-    ],
+    users: user
+      ? [
+          {
+            id: user.profile_uid,
+            name: user.username,
+          },
+        ]
+      : undefined,
   };
 };
 

x-pack/plugins/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');

x-pack/plugins/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;
     },

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(