[Security Assistant] Fix Product documentation installation banner (#…

…212463)

## Summary

Fixes logic on fresh cluster where the ELSER was not started yet, in
this case API reports `status` as `uninstalled`, but it doesn't mean
that the Product documentation was actually uninstall, but rather it's a
default state.
Added internal `product_documentation_status` to KB status API to make
sure we keep track of the status internally and present the banner only
if the docs were intentionally uninstalled

---------

Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
(cherry picked from commit f0d6669)

oas_docs/output/kibana.serverless.yaml

@@ -37898,6 +37898,8 @@ paths:
                     type: boolean
                   pipeline_exists:
                     type: boolean
+                  product_documentation_status:
+                    type: string
                   security_labs_exists:
                     type: boolean
                   user_data_exists:

oas_docs/output/kibana.yaml

@@ -40450,6 +40450,8 @@ paths:
                     type: boolean
                   pipeline_exists:
                     type: boolean
+                  product_documentation_status:
+                    type: string
                   security_labs_exists:
                     type: boolean
                   user_data_exists:

x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml

@@ -445,6 +445,8 @@ paths:
                     type: boolean
                   pipeline_exists:
                     type: boolean
+                  product_documentation_status:
+                    type: string
                   security_labs_exists:
                     type: boolean
                   user_data_exists:

x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml

@@ -445,6 +445,8 @@ paths:
                     type: boolean
                   pipeline_exists:
                     type: boolean
+                  product_documentation_status:
+                    type: string
                   security_labs_exists:
                     type: boolean
                   user_data_exists:

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.gen.ts

@@ -73,4 +73,5 @@ export const ReadKnowledgeBaseResponse = z.object({
   pipeline_exists: z.boolean().optional(),
   security_labs_exists: z.boolean().optional(),
   user_data_exists: z.boolean().optional(),
+  product_documentation_status: z.string().optional(),
 });

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/knowledge_base/crud_kb_route.schema.yaml

@@ -87,6 +87,8 @@ paths:
                     type: boolean
                   user_data_exists:
                     type: boolean
+                  product_documentation_status:
+                    type: string
         400:
           description: Generic Error
           content:

x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx

@@ -12,6 +12,7 @@ import type { IToasts } from '@kbn/core-notifications-browser';
 import { i18n } from '@kbn/i18n';
 import { useCallback } from 'react';
 import { ReadKnowledgeBaseResponse } from '@kbn/elastic-assistant-common';
+import { InstallationStatus } from '@kbn/product-doc-base-plugin/common/install_status';
 import { getKnowledgeBaseStatus } from './api';
 
 const KNOWLEDGE_BASE_STATUS_QUERY_KEY = ['elastic-assistant', 'knowledge-base-status'];
@@ -38,7 +39,10 @@ export const useKnowledgeBaseStatus = ({
   resource,
   toasts,
   enabled,
-}: UseKnowledgeBaseStatusParams): UseQueryResult<ReadKnowledgeBaseResponse, IHttpFetchError> => {
+}: UseKnowledgeBaseStatusParams): UseQueryResult<
+  ReadKnowledgeBaseResponse & { product_documentation_status: InstallationStatus },
+  IHttpFetchError
+> => {
   return useQuery(
     KNOWLEDGE_BASE_STATUS_QUERY_KEY,
     async ({ signal }) => {
@@ -49,7 +53,10 @@ export const useKnowledgeBaseStatus = ({
       retry: false,
       keepPreviousData: true,
       // Polling interval for Knowledge Base setup in progress
-      refetchInterval: (data) => (data?.is_setup_in_progress ? 30000 : false),
+      refetchInterval: (data) =>
+        data?.is_setup_in_progress || data?.product_documentation_status === 'installing'
+          ? 30000
+          : false,
       // Deprecated, hoist to `queryCache` w/in `QueryClient. See: https://stackoverflow.com/a/76961109
       onError: (error: IHttpFetchError<ResponseErrorBody>) => {
         if (error.name !== 'AbortError') {

x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/product_documentation/index.test.tsx

@@ -9,7 +9,6 @@ import { render, screen, fireEvent, waitFor } from '@testing-library/react';
 import { ProductDocumentationManagement } from '.';
 import * as i18n from './translations';
 import { useInstallProductDoc } from '../../api/product_docs/use_install_product_doc';
-import { useGetProductDocStatus } from '../../api/product_docs/use_get_product_doc_status';
 
 jest.mock('../../api/product_docs/use_install_product_doc');
 jest.mock('../../api/product_docs/use_get_product_doc_status');
@@ -18,69 +17,64 @@ describe('ProductDocumentationManagement', () => {
   const mockInstallProductDoc = jest.fn().mockResolvedValue({});
 
   beforeEach(() => {
-    (useInstallProductDoc as jest.Mock).mockReturnValue({ mutateAsync: mockInstallProductDoc });
-    (useGetProductDocStatus as jest.Mock).mockReturnValue({ status: null, isLoading: false });
-    jest.clearAllMocks();
-  });
-
-  it('renders loading spinner when status is loading', async () => {
-    (useGetProductDocStatus as jest.Mock).mockReturnValue({
-      status: { overall: 'not_installed' },
-      isLoading: true,
+    (useInstallProductDoc as jest.Mock).mockReturnValue({
+      mutateAsync: mockInstallProductDoc,
+      isLoading: false,
+      isSuccess: false,
     });
-    render(<ProductDocumentationManagement />);
-    expect(screen.getByTestId('statusLoading')).toBeInTheDocument();
+    jest.clearAllMocks();
   });
 
   it('renders install button when not installed', () => {
-    (useGetProductDocStatus as jest.Mock).mockReturnValue({
-      status: { overall: 'not_installed' },
-      isLoading: false,
-    });
-    render(<ProductDocumentationManagement />);
+    render(<ProductDocumentationManagement status="uninstalled" />);
     expect(screen.getByText(i18n.INSTALL)).toBeInTheDocument();
   });
 
   it('does not render anything when already installed', () => {
-    (useGetProductDocStatus as jest.Mock).mockReturnValue({
-      status: { overall: 'installed' },
+    const { container } = render(<ProductDocumentationManagement status="installed" />);
+    expect(container).toBeEmptyDOMElement();
+  });
+
+  it('does not render anything when the installation was started by the plugin', () => {
+    (useInstallProductDoc as jest.Mock).mockReturnValue({
+      mutateAsync: mockInstallProductDoc,
       isLoading: false,
+      isSuccess: false,
     });
-    const { container } = render(<ProductDocumentationManagement />);
+    const { container } = render(<ProductDocumentationManagement status="installing" />);
     expect(container).toBeEmptyDOMElement();
   });
 
   it('shows installing spinner and text when installing', async () => {
-    (useGetProductDocStatus as jest.Mock).mockReturnValue({
-      status: { overall: 'not_installed' },
-      isLoading: false,
-    });
-    render(<ProductDocumentationManagement />);
-    fireEvent.click(screen.getByText(i18n.INSTALL));
-    await waitFor(() => {
-      expect(screen.getByTestId('installing')).toBeInTheDocument();
-      expect(screen.getByText(i18n.INSTALLING)).toBeInTheDocument();
+    (useInstallProductDoc as jest.Mock).mockReturnValue({
+      mutateAsync: mockInstallProductDoc,
+      isLoading: true,
+      isSuccess: false,
     });
+    render(<ProductDocumentationManagement status="uninstalled" />);
+    expect(screen.getByTestId('installing')).toBeInTheDocument();
+    expect(screen.getByText(i18n.INSTALLING)).toBeInTheDocument();
   });
 
   it('sets installed state to true after successful installation', async () => {
-    (useGetProductDocStatus as jest.Mock).mockReturnValue({
-      status: { overall: 'not_installed' },
+    (useInstallProductDoc as jest.Mock).mockReturnValue({
+      mutateAsync: mockInstallProductDoc,
       isLoading: false,
+      isSuccess: true,
     });
     mockInstallProductDoc.mockResolvedValueOnce({});
-    render(<ProductDocumentationManagement />);
-    fireEvent.click(screen.getByText(i18n.INSTALL));
-    await waitFor(() => expect(screen.queryByText(i18n.INSTALL)).not.toBeInTheDocument());
+    render(<ProductDocumentationManagement status="uninstalled" />);
+    expect(screen.queryByText(i18n.INSTALL)).not.toBeInTheDocument();
   });
 
   it('sets installed state to false after failed installation', async () => {
-    (useGetProductDocStatus as jest.Mock).mockReturnValue({
-      status: { overall: 'not_installed' },
+    (useInstallProductDoc as jest.Mock).mockReturnValue({
+      mutateAsync: mockInstallProductDoc,
       isLoading: false,
+      isSuccess: false,
     });
     mockInstallProductDoc.mockRejectedValueOnce(new Error('Installation failed'));
-    render(<ProductDocumentationManagement />);
+    render(<ProductDocumentationManagement status="uninstalled" />);
     fireEvent.click(screen.getByText(i18n.INSTALL));
     await waitFor(() => expect(screen.getByText(i18n.INSTALL)).toBeInTheDocument());
   });

x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/product_documentation/index.tsx

@@ -14,43 +14,25 @@ import {
   EuiSpacer,
   EuiText,
 } from '@elastic/eui';
-import React, { useEffect, useState, useCallback, useMemo } from 'react';
+import React, { useCallback, useMemo } from 'react';
+import { InstallationStatus } from '@kbn/product-doc-base-plugin/common/install_status';
 import { useInstallProductDoc } from '../../api/product_docs/use_install_product_doc';
-import { useGetProductDocStatus } from '../../api/product_docs/use_get_product_doc_status';
 import * as i18n from './translations';
 
-export const ProductDocumentationManagement: React.FC = React.memo(() => {
-  const [{ isInstalled, isInstalling }, setState] = useState({
-    isInstalled: true,
-    isInstalling: false,
-  });
+export const ProductDocumentationManagement = React.memo<{
+  status?: InstallationStatus;
+}>(({ status }) => {
+  const {
+    mutateAsync: installProductDoc,
+    isSuccess: isInstalled,
+    isLoading: isInstalling,
+  } = useInstallProductDoc();
 
-  const { mutateAsync: installProductDoc } = useInstallProductDoc();
-  const { status, isLoading: isStatusLoading } = useGetProductDocStatus();
-
-  useEffect(() => {
-    if (status) {
-      setState((prevState) => ({
-        ...prevState,
-        isInstalled: status.overall === 'installed',
-      }));
-    }
-  }, [status]);
-
-  const onClickInstall = useCallback(async () => {
-    setState((prevState) => ({ ...prevState, isInstalling: true }));
-    try {
-      await installProductDoc();
-      setState({ isInstalled: true, isInstalling: false });
-    } catch {
-      setState({ isInstalled: false, isInstalling: false });
-    }
+  const onClickInstall = useCallback(() => {
+    installProductDoc();
   }, [installProductDoc]);
 
   const content = useMemo(() => {
-    if (isStatusLoading) {
-      return <EuiLoadingSpinner data-test-subj="statusLoading" size="m" />;
-    }
     if (isInstalling) {
       return (
         <EuiFlexGroup justifyContent="flexStart" alignItems="center">
@@ -72,11 +54,18 @@ export const ProductDocumentationManagement: React.FC = React.memo(() => {
         </EuiFlexItem>
       </EuiFlexGroup>
     );
-  }, [isInstalling, isStatusLoading, onClickInstall]);
+  }, [isInstalling, onClickInstall]);
 
-  if (isInstalled) {
+  // The last condition means that the installation was started by the plugin
+  if (
+    !status ||
+    status === 'installed' ||
+    isInstalled ||
+    (status === 'installing' && !isInstalling)
+  ) {
     return null;
   }
+
   return (
     <>
       <EuiCallOut title={i18n.LABEL} iconType="iInCircle">

x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx

@@ -340,7 +340,7 @@ export const KnowledgeBaseSettingsManagement: React.FC<Params> = React.memo(({ d
 
   return (
     <>
-      <ProductDocumentationManagement />
+      <ProductDocumentationManagement status={kbStatus?.product_documentation_status} />
       <EuiPanel hasShadow={false} hasBorder paddingSize="l">
         <EuiText size={'m'}>
           <FormattedMessage

x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/data_clients.mock.ts

@@ -15,7 +15,9 @@ type ConversationsDataClientContract = PublicMethodsOf<AIAssistantConversationsD
 export type ConversationsDataClientMock = jest.Mocked<ConversationsDataClientContract>;
 type AttackDiscoveryDataClientContract = PublicMethodsOf<AttackDiscoveryDataClient>;
 export type AttackDiscoveryDataClientMock = jest.Mocked<AttackDiscoveryDataClientContract>;
-type KnowledgeBaseDataClientContract = PublicMethodsOf<AIAssistantKnowledgeBaseDataClient>;
+type KnowledgeBaseDataClientContract = PublicMethodsOf<AIAssistantKnowledgeBaseDataClient> & {
+  isSetupInProgress: AIAssistantKnowledgeBaseDataClient['isSetupInProgress'];
+};
 export type KnowledgeBaseDataClientMock = jest.Mocked<KnowledgeBaseDataClientContract>;
 
 const createConversationsDataClientMock = () => {
@@ -73,9 +75,11 @@ const createKnowledgeBaseDataClientMock = () => {
     isModelInstalled: jest.fn(),
     isSecurityLabsDocsLoaded: jest.fn(),
     isSetupAvailable: jest.fn(),
+    isSetupInProgress: jest.fn().mockReturnValue(false)(),
     isUserDataExists: jest.fn(),
     setupKnowledgeBase: jest.fn(),
     getLoadedSecurityLabsDocsCount: jest.fn(),
+    getProductDocumentationStatus: jest.fn(),
   };
   return mocked;
 };

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts

@@ -64,6 +64,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => {
       ml,
       getElserId: getElserId.mockResolvedValue('elser-id'),
       getIsKBSetupInProgress: mockGetIsKBSetupInProgress.mockReturnValue(false),
+      getProductDocumentationStatus: jest.fn().mockResolvedValue('installed'),
       ingestPipelineResourceName: 'something',
       setIsKBSetupInProgress: jest.fn().mockImplementation(() => {}),
       manageGlobalKnowledgeBaseAIAssistant: true,

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts

@@ -33,6 +33,7 @@ import {
 } from '@kbn/core/server';
 import { IndexPatternsFetcher } from '@kbn/data-views-plugin/server';
 import { map } from 'lodash';
+import { InstallationStatus } from '@kbn/product-doc-base-plugin/common/install_status';
 import type { TrainedModelsProvider } from '@kbn/ml-plugin/server/shared_services/providers';
 import { getMlNodeCount } from '@kbn/ml-plugin/server/lib/node_utils';
 import { AIAssistantDataClient, AIAssistantDataClientParams } from '..';
@@ -86,6 +87,7 @@ export interface KnowledgeBaseDataClientParams extends AIAssistantDataClientPara
   ml: MlPluginSetup;
   getElserId: GetElser;
   getIsKBSetupInProgress: (spaceId: string) => boolean;
+  getProductDocumentationStatus: () => Promise<InstallationStatus>;
   ingestPipelineResourceName: string;
   setIsKBSetupInProgress: (spaceId: string, isInProgress: boolean) => void;
   manageGlobalKnowledgeBaseAIAssistant: boolean;
@@ -100,6 +102,11 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
   public get isSetupInProgress() {
     return this.options.getIsKBSetupInProgress(this.spaceId);
   }
+
+  public getProductDocumentationStatus = async () => {
+    return (await this.options.getProductDocumentationStatus()) ?? 'uninstalled';
+  };
+
   /**
    * Returns whether setup of the Knowledge Base can be performed (essentially an ML features check)
    *

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/helpers.test.ts

@@ -26,7 +26,11 @@ describe('helpers', () => {
       mockProductDocManager.getStatus.mockResolvedValue({ status: 'uninstalled' });
       mockProductDocManager.install.mockResolvedValue(null);
 
-      await ensureProductDocumentationInstalled(mockProductDocManager, mockLogger);
+      await ensureProductDocumentationInstalled({
+        productDocManager: mockProductDocManager,
+        setIsProductDocumentationInProgress: jest.fn(),
+        logger: mockLogger,
+      });
 
       expect(mockProductDocManager.getStatus).toHaveBeenCalled();
       expect(mockLogger.debug).toHaveBeenCalledWith(
@@ -42,7 +46,11 @@ describe('helpers', () => {
     it('should not install product documentation if already installed', async () => {
       mockProductDocManager.getStatus.mockResolvedValue({ status: 'installed' });
 
-      await ensureProductDocumentationInstalled(mockProductDocManager, mockLogger);
+      await ensureProductDocumentationInstalled({
+        productDocManager: mockProductDocManager,
+        setIsProductDocumentationInProgress: jest.fn(),
+        logger: mockLogger,
+      });
 
       expect(mockProductDocManager.getStatus).toHaveBeenCalled();
       expect(mockProductDocManager.install).not.toHaveBeenCalled();
@@ -54,7 +62,11 @@ describe('helpers', () => {
       mockProductDocManager.getStatus.mockResolvedValue({ status: 'not_installed' });
       mockProductDocManager.install.mockRejectedValue(new Error('Install failed'));
 
-      await ensureProductDocumentationInstalled(mockProductDocManager, mockLogger);
+      await ensureProductDocumentationInstalled({
+        productDocManager: mockProductDocManager,
+        setIsProductDocumentationInProgress: jest.fn(),
+        logger: mockLogger,
+      });
 
       expect(mockProductDocManager.getStatus).toHaveBeenCalled();
       expect(mockProductDocManager.install).toHaveBeenCalled();
@@ -67,7 +79,11 @@ describe('helpers', () => {
     it('should log a warning if getStatus fails', async () => {
       mockProductDocManager.getStatus.mockRejectedValue(new Error('Status check failed'));
 
-      await ensureProductDocumentationInstalled(mockProductDocManager, mockLogger);
+      await ensureProductDocumentationInstalled({
+        productDocManager: mockProductDocManager,
+        setIsProductDocumentationInProgress: jest.fn(),
+        logger: mockLogger,
+      });
 
       expect(mockProductDocManager.getStatus).toHaveBeenCalled();
       expect(mockLogger.warn).toHaveBeenCalledWith(