[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>
This commit is contained in:
Patryk Kopyciński 2025-03-04 02:28:35 +01:00 committed by GitHub
parent 750e156c26
commit f0d66691b8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 161 additions and 96 deletions

View file

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

View file

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

View file

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

View file

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

View file

@ -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(),
});

View file

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

View file

@ -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') {

View file

@ -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 });
(useInstallProductDoc as jest.Mock).mockReturnValue({
mutateAsync: mockInstallProductDoc,
isLoading: false,
isSuccess: false,
});
jest.clearAllMocks();
});
it('renders loading spinner when status is loading', async () => {
(useGetProductDocStatus as jest.Mock).mockReturnValue({
status: { overall: 'not_installed' },
isLoading: true,
});
render(<ProductDocumentationManagement />);
expect(screen.getByTestId('statusLoading')).toBeInTheDocument();
});
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());
});

View file

@ -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">

View file

@ -338,7 +338,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

View file

@ -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;
};

View file

@ -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,

View file

@ -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)
*

View file

@ -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(

View file

@ -144,19 +144,27 @@ const ESQL_QUERY_GENERATION_TITLE = i18n.translate(
}
);
export const ensureProductDocumentationInstalled = async (
productDocManager: ProductDocBaseStartContract['management'],
logger: Logger
) => {
export const ensureProductDocumentationInstalled = async ({
productDocManager,
setIsProductDocumentationInProgress,
logger,
}: {
productDocManager: ProductDocBaseStartContract['management'];
setIsProductDocumentationInProgress: (value: boolean) => void;
logger: Logger;
}) => {
try {
const { status } = await productDocManager.getStatus();
if (status !== 'installed') {
logger.debug(`Installing product documentation for AIAssistantService`);
setIsProductDocumentationInProgress(true);
try {
await productDocManager.install();
await productDocManager.install({ wait: true });
logger.debug(`Successfully installed product documentation for AIAssistantService`);
} catch (e) {
logger.warn(`Failed to install product documentation for AIAssistantService: ${e.message}`);
} finally {
setIsProductDocumentationInProgress(false);
}
}
} catch (e) {

View file

@ -18,6 +18,7 @@ import {
IndicesIndexSettings,
} from '@elastic/elasticsearch/lib/api/types';
import { omit } from 'lodash';
import { InstallationStatus } from '@kbn/product-doc-base-plugin/common/install_status';
import { TrainedModelsProvider } from '@kbn/ml-plugin/server/shared_services/providers';
import { attackDiscoveryFieldMap } from '../lib/attack_discovery/persistence/field_maps_configuration/field_maps_configuration';
import { defendInsightsFieldMap } from '../ai_assistant_data_clients/defend_insights/field_maps_configuration';
@ -105,6 +106,7 @@ export class AIAssistantService {
private isKBSetupInProgress: Map<string, boolean> = new Map();
private hasInitializedV2KnowledgeBase: boolean = false;
private productDocManager?: ProductDocBaseStartContract['management'];
private isProductDocumentationInProgress: boolean = false;
// Temporary 'feature flag' to determine if we should initialize the new knowledge base mappings
private assistantDefaultInferenceEndpoint: boolean = false;
@ -170,6 +172,14 @@ export class AIAssistantService {
this.isKBSetupInProgress.set(spaceId, isInProgress);
}
public getIsProductDocumentationInProgress() {
return this.isProductDocumentationInProgress;
}
public setIsProductDocumentationInProgress(isInProgress: boolean) {
this.isProductDocumentationInProgress = isInProgress;
}
private createDataStream: CreateDataStream = ({
resource,
kibanaVersion,
@ -220,7 +230,11 @@ export class AIAssistantService {
if (this.productDocManager) {
// install product documentation without blocking other resources
void ensureProductDocumentationInstalled(this.productDocManager, this.options.logger);
void ensureProductDocumentationInstalled({
productDocManager: this.productDocManager,
logger: this.options.logger,
setIsProductDocumentationInProgress: this.setIsProductDocumentationInProgress.bind(this),
});
}
await this.conversationsDataStream.install({
@ -469,6 +483,16 @@ export class AIAssistantService {
}
}
public async getProductDocumentationStatus(): Promise<InstallationStatus> {
const status = await this.productDocManager?.getStatus();
if (!status) {
return 'uninstalled';
}
return this.isProductDocumentationInProgress ? 'installing' : status.status;
}
public async createAIAssistantConversationsDataClient(
opts: CreateAIAssistantClientParams & GetAIAssistantConversationsDataClientParams
): Promise<AIAssistantConversationsDataClient | null> {
@ -523,6 +547,7 @@ export class AIAssistantService {
ingestPipelineResourceName: this.resourceNames.pipelines.knowledgeBase,
getElserId: this.getElserId,
getIsKBSetupInProgress: this.getIsKBSetupInProgress.bind(this),
getProductDocumentationStatus: this.getProductDocumentationStatus.bind(this),
kibanaVersion: this.options.kibanaVersion,
ml: this.options.ml,
setIsKBSetupInProgress: this.setIsKBSetupInProgress.bind(this),

View file

@ -10,6 +10,7 @@ import { serverMock } from '../../__mocks__/server';
import { requestContextMock } from '../../__mocks__/request_context';
import { getGetKnowledgeBaseStatusRequest } from '../../__mocks__/request';
import { AuthenticatedUser } from '@kbn/core-security-common';
import { knowledgeBaseDataClientMock } from '../../__mocks__/data_clients.mock';
describe('Get Knowledge Base Status Route', () => {
let server: ReturnType<typeof serverMock.create>;
@ -28,19 +29,18 @@ describe('Get Knowledge Base Status Route', () => {
server = serverMock.create();
({ context } = requestContextMock.createTools());
context.elasticAssistant.getCurrentUser.mockResolvedValue(mockUser);
context.elasticAssistant.getAIAssistantKnowledgeBaseDataClient = jest.fn().mockResolvedValue({
getKnowledgeBaseDocuments: jest.fn().mockResolvedValue([]),
indexTemplateAndPattern: {
alias: 'knowledge-base-alias',
},
isModelInstalled: jest.fn().mockResolvedValue(true),
isSetupAvailable: jest.fn().mockResolvedValue(true),
isInferenceEndpointExists: jest.fn().mockResolvedValue(true),
isSetupInProgress: false,
isSecurityLabsDocsLoaded: jest.fn().mockResolvedValue(true),
isUserDataExists: jest.fn().mockResolvedValue(true),
getLoadedSecurityLabsDocsCount: jest.fn().mockResolvedValue(0),
});
const kbDataClient = knowledgeBaseDataClientMock.create();
context.elasticAssistant.getAIAssistantKnowledgeBaseDataClient = jest
.fn()
.mockResolvedValue(kbDataClient);
kbDataClient.isInferenceEndpointExists.mockResolvedValue(true);
kbDataClient.isModelInstalled.mockResolvedValue(true);
kbDataClient.isSetupAvailable.mockResolvedValue(true);
kbDataClient.getProductDocumentationStatus.mockResolvedValue('installed');
kbDataClient.isSecurityLabsDocsLoaded.mockResolvedValue(true);
kbDataClient.isUserDataExists.mockResolvedValue(true);
kbDataClient.isSetupInProgress = false;
getKnowledgeBaseStatusRoute(server.router);
});
@ -61,6 +61,7 @@ describe('Get Knowledge Base Status Route', () => {
pipeline_exists: true,
security_labs_exists: true,
user_data_exists: true,
product_documentation_status: 'installed',
});
});
});

View file

@ -61,6 +61,7 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter
const securityLabsExists = await kbDataClient.isSecurityLabsDocsLoaded();
const loadedSecurityLabsDocsCount = await kbDataClient.getLoadedSecurityLabsDocsCount();
const userDataExists = await kbDataClient.isUserDataExists();
const productDocumentationStatus = await kbDataClient.getProductDocumentationStatus();
return response.ok({
body: {
@ -72,6 +73,7 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter
// If user data exists, we should have at least one document in the Security Labs index
user_data_exists: userDataExists || !!loadedSecurityLabsDocsCount,
pipeline_exists: pipelineExists,
product_documentation_status: productDocumentationStatus,
},
});
} catch (err) {