From ea2cee57b384661f93c69aedf4c8fec352ba3d6b Mon Sep 17 00:00:00 2001 From: Yuliia Naumenko Date: Wed, 30 Oct 2024 14:13:33 -0700 Subject: [PATCH 1/3] [Security AI Assistant] Fixed license issue for Knowledge Base resources initialization (#198239) (cherry picked from commit ed81e4334f4d5608517dacdba28d46dfea966be0) # 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 --- .../use_knowledge_base_status.tsx | 3 + .../chat_send/use_chat_send.test.tsx | 3 + .../assistant/chat_send/use_chat_send.tsx | 8 ++- .../knowledge_base_settings.test.tsx | 3 + .../knowledge_base_settings.tsx | 8 ++- .../setup_knowledge_base_button.tsx | 8 ++- ...reate_resource_installation_helper.test.ts | 8 +-- .../create_resource_installation_helper.ts | 4 +- .../server/ai_assistant_service/index.test.ts | 19 +++++- .../server/ai_assistant_service/index.ts | 9 ++- .../bulk_actions_route.ts | 36 +++++------ .../anonymization_fields/find_route.test.ts | 14 ++++- .../routes/anonymization_fields/find_route.ts | 18 +++--- .../routes/chat/chat_complete_route.test.ts | 12 +++- .../server/routes/chat/chat_complete_route.ts | 6 +- .../server/routes/evaluate/get_evaluate.ts | 7 +-- .../server/routes/evaluate/post_evaluate.ts | 6 +- .../server/routes/helpers.ts | 62 ++++++++++++------- .../entries/bulk_actions_route.ts | 11 ++-- .../knowledge_base/entries/create_route.ts | 6 +- .../knowledge_base/entries/find_route.ts | 8 +-- .../post_actions_connector_execute.test.ts | 13 +++- .../routes/post_actions_connector_execute.ts | 15 +++-- .../routes/prompts/bulk_actions_route.ts | 27 ++++---- .../server/routes/prompts/find_route.test.ts | 14 ++++- .../server/routes/prompts/find_route.ts | 17 ++--- .../server/routes/request_context_factory.ts | 5 ++ .../append_conversation_messages_route.ts | 24 +++---- .../user_conversations/bulk_actions_route.ts | 24 +++---- .../routes/user_conversations/create_route.ts | 6 +- .../routes/user_conversations/delete_route.ts | 25 +++----- .../user_conversations/find_route.test.ts | 7 +-- .../routes/user_conversations/find_route.ts | 19 +++--- .../routes/user_conversations/read_route.ts | 24 +++---- .../routes/user_conversations/update_route.ts | 8 +-- 35 files changed, 274 insertions(+), 213 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx index ba6317329d350..4b7a127c9f8a6 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx +++ b/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,6 +37,7 @@ export const useKnowledgeBaseStatus = ({ http, resource, toasts, + enabled, }: UseKnowledgeBaseStatusParams): UseQueryResult => { return useQuery( KNOWLEDGE_BASE_STATUS_QUERY_KEY, @@ -43,6 +45,7 @@ export const useKnowledgeBaseStatus = ({ 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 index a3532fac8d57f..62dd0ceaa16ea 100644 --- 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 index 7dfe3748bbae2..d09a67da296d5 100644 --- 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(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 index 67b48ac9354d7..b382a39a99bae 100644 --- 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 index 873296ea57840..b2ec440ce78ff 100644 --- 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 = 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 index 008597c1e8243..582cf28be9e70 100644 --- 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/x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.test.ts index 481211309123e..df1f937ed4e83 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.test.ts +++ b/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'); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.ts index 39e0e69a8fc49..e8d1f1eb1d85d 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.ts +++ b/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; }, 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 index a11ad418c7bad..23a1a55564415 100644 --- 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; +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 index 6fff5f767154b..9119b29210edd 100644 --- 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; } 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 index f7b87b49dd2a7..491051cdaa970 100644 --- 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,7 +191,7 @@ export const bulkActionAnonymizationFieldsRoute = ( } const writer = await dataClient?.getWriter(); - const changedAt = new Date().toISOString(); + const createdAt = new Date().toISOString(); const { errors, docs_created: docsCreated, @@ -203,12 +199,12 @@ export const bulkActionAnonymizationFieldsRoute = ( 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 }), diff --git a/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.test.ts index 4659503261a6e..7c2b1d330a3db 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.test.ts @@ -12,6 +12,7 @@ import { requestContextMock } from '../../__mocks__/request_context'; import { getFindAnonymizationFieldsResultWithSingleHit } from '../../__mocks__/response'; import { findAnonymizationFieldsRoute } from './find_route'; import { loggingSystemMock } from '@kbn/core-logging-server-mocks'; +import type { AuthenticatedUser } from '@kbn/core-security-common'; describe('Find user anonymization fields route', () => { let server: ReturnType; @@ -21,19 +22,26 @@ describe('Find user anonymization fields route', () => { beforeEach(async () => { server = serverMock.create(); ({ clients, context } = requestContextMock.createTools()); + const mockUser1 = { + username: 'my_username', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, + } as AuthenticatedUser; clients.elasticAssistant.getAIAssistantAnonymizationFieldsDataClient.findDocuments.mockResolvedValue( Promise.resolve(getFindAnonymizationFieldsResultWithSingleHit()) ); - clients.elasticAssistant.getCurrentUser.mockResolvedValue({ + context.elasticAssistant.getCurrentUser.mockReturnValue({ username: 'my_username', authentication_realm: { type: 'my_realm_type', name: 'my_realm_name', }, - }); + } as AuthenticatedUser); logger = loggingSystemMock.createLogger(); - + context.elasticAssistant.getCurrentUser.mockReturnValue(mockUser1); findAnonymizationFieldsRoute(server.router, logger); }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.ts b/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.ts index 904a80d6a3ea4..061dd9ff3eac6 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.ts @@ -22,7 +22,7 @@ import { ElasticAssistantPluginRouter } from '../../types'; import { buildResponse } from '../utils'; import { EsAnonymizationFieldsSchema } from '../../ai_assistant_data_clients/anonymization_fields/types'; import { transformESSearchToAnonymizationFields } from '../../ai_assistant_data_clients/anonymization_fields/helpers'; -import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../helpers'; +import { performChecks } from '../helpers'; export const findAnonymizationFieldsRoute = ( router: ElasticAssistantPluginRouter, @@ -55,14 +55,16 @@ export const findAnonymizationFieldsRoute = ( try { const { query } = request; 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, + }); + if (!checkResponse.isSuccess) { + return checkResponse.response; } + const dataClient = await ctx.elasticAssistant.getAIAssistantAnonymizationFieldsDataClient(); diff --git a/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts index 0b05bb2875cb6..4aca370aa700f 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts @@ -32,7 +32,17 @@ const actionsClient = actionsClientMock.create(); jest.mock('../../lib/build_response', () => ({ buildResponse: jest.fn().mockImplementation((x) => x), })); -jest.mock('../helpers'); + +jest.mock('../helpers', () => { + const original = jest.requireActual('../helpers'); + + return { + ...original, + appendAssistantMessageToConversation: jest.fn(), + createConversationWithUserInput: jest.fn(), + langChainExecute: jest.fn(), + }; +}); const mockAppendAssistantMessageToConversation = appendAssistantMessageToConversation as jest.Mock; const mockLangChainExecute = langChainExecute as jest.Mock; diff --git a/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts b/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts index b8f75e2376863..88e352eba98bd 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts @@ -69,14 +69,12 @@ export const chatCompleteRoute = ( // Perform license and authenticated user checks const checkResponse = performChecks({ - authenticatedUser: true, context: ctx, - license: true, request, response, }); - if (checkResponse) { - return checkResponse; + if (!checkResponse.isSuccess) { + return checkResponse.response; } const conversationsDataClient = diff --git a/x-pack/plugins/elastic_assistant/server/routes/evaluate/get_evaluate.ts b/x-pack/plugins/elastic_assistant/server/routes/evaluate/get_evaluate.ts index 41b455a73598b..dd7462696621b 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/evaluate/get_evaluate.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/evaluate/get_evaluate.ts @@ -48,15 +48,14 @@ export const getEvaluateRoute = (router: IRouter >; - license?: boolean; request: KibanaRequest; response: KibanaResponseFactory; } /** - * Helper to perform checks for authenticated user, capability, and license. Perform all or one - * of the checks by providing relevant optional params. Check order is license, authenticated user, - * then capability. + * Helper to perform checks for authenticated user, license, and optionally capability. + * Check order is license, authenticated user, then capability. + * + * Returns either a successful check with an AuthenticatedUser or + * an unsuccessful check with an error IKibanaResponse. * - * @param authenticatedUser - Whether to check for an authenticated user * @param capability - Specific capability to check if enabled, e.g. `assistantModelEvaluation` * @param context - Route context - * @param license - Whether to check for a valid license * @param request - Route KibanaRequest * @param response - Route KibanaResponseFactory + * @returns PerformChecks */ + +type PerformChecks = + | { + isSuccess: true; + currentUser: AuthenticatedUser; + } + | { + isSuccess: false; + response: IKibanaResponse; + }; export const performChecks = ({ - authenticatedUser, capability, context, - license, request, response, -}: PerformChecksParams): IKibanaResponse | undefined => { +}: PerformChecksParams): PerformChecks => { const assistantResponse = buildResponse(response); - if (license) { - if (!hasAIAssistantLicense(context.licensing.license)) { - return response.forbidden({ + if (!hasAIAssistantLicense(context.licensing.license)) { + return { + isSuccess: false, + response: response.forbidden({ body: { message: UPGRADE_LICENSE_MESSAGE, }, - }); - } + }), + }; } - if (authenticatedUser) { - if (context.elasticAssistant.getCurrentUser() == null) { - return assistantResponse.error({ + const currentUser = context.elasticAssistant.getCurrentUser(); + + if (currentUser == null) { + return { + isSuccess: false, + response: assistantResponse.error({ body: `Authenticated user not found`, statusCode: 401, - }); - } + }), + }; } if (capability) { @@ -585,11 +597,17 @@ export const performChecks = ({ }); const registeredFeatures = context.elasticAssistant.getRegisteredFeatures(pluginName); if (!registeredFeatures[capability]) { - return response.notFound(); + return { + isSuccess: false, + response: response.notFound(), + }; } } - return undefined; + return { + isSuccess: true, + currentUser, + }; }; /** diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts index 76ea154fe203f..93de3ad460b51 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts @@ -6,7 +6,7 @@ */ import moment from 'moment'; -import type { AuthenticatedUser, IKibanaResponse, KibanaResponseFactory } from '@kbn/core/server'; +import type { IKibanaResponse, KibanaResponseFactory } from '@kbn/core/server'; import { transformError } from '@kbn/securitysolution-es-utils'; import { @@ -132,15 +132,13 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug // Perform license, authenticated user and FF checks const checkResponse = performChecks({ - authenticatedUser: true, capability: 'assistantKnowledgeBaseByDefault', context: ctx, - license: true, request, response, }); - if (checkResponse) { - return checkResponse; + if (!checkResponse.isSuccess) { + return checkResponse.response; } logger.debug( @@ -167,8 +165,7 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug request.events.completed$.subscribe(() => abortController.abort()); const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient(); const spaceId = ctx.elasticAssistant.getSpaceId(); - // Authenticated user null check completed in `performChecks()` above - const authenticatedUser = ctx.elasticAssistant.getCurrentUser() as AuthenticatedUser; + const authenticatedUser = checkResponse.currentUser; if (body.create && body.create.length > 0) { const result = await kbDataClient?.findDocuments({ diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts index fe21bfc043c78..a39dcee6d2487 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts @@ -49,15 +49,13 @@ export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout // Perform license, authenticated user and FF checks const checkResponse = performChecks({ - authenticatedUser: true, capability: 'assistantKnowledgeBaseByDefault', context: ctx, - license: true, request, response, }); - if (checkResponse) { - return checkResponse; + if (!checkResponse.isSuccess) { + return checkResponse.response; } logger.debug(`Creating KB Entry:\n${JSON.stringify(request.body)}`); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts index e939eae1eb3ad..f787825c31172 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts @@ -52,19 +52,17 @@ export const findKnowledgeBaseEntriesRoute = (router: ElasticAssistantPluginRout // Perform license, authenticated user and FF checks const checkResponse = performChecks({ - authenticatedUser: true, capability: 'assistantKnowledgeBaseByDefault', context: ctx, - license: true, request, response, }); - if (checkResponse) { - return checkResponse; + if (!checkResponse.isSuccess) { + return checkResponse.response; } const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient(); - const currentUser = ctx.elasticAssistant.getCurrentUser(); + const currentUser = checkResponse.currentUser; const additionalFilter = query.filter ? ` AND ${query.filter}` : ''; const result = await kbDataClient?.findDocuments({ diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts index 12e854d5d0bf3..8053068da1fc0 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts @@ -46,7 +46,18 @@ jest.mock('../lib/executor', () => ({ const mockStream = jest.fn().mockImplementation(() => new PassThrough()); const mockLangChainExecute = langChainExecute as jest.Mock; const mockAppendAssistantMessageToConversation = appendAssistantMessageToConversation as jest.Mock; -jest.mock('./helpers'); +jest.mock('./helpers', () => { + const original = jest.requireActual('./helpers'); + + return { + ...original, + getIsKnowledgeBaseEnabled: jest.fn(), + appendAssistantMessageToConversation: jest.fn(), + langChainExecute: jest.fn(), + getPluginNameFromRequest: jest.fn(), + getSystemPromptFromUserConversation: jest.fn(), + }; +}); const existingConversation = getConversationResponseMock(); const reportEvent = jest.fn(); const appendConversationMessages = jest.fn(); diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index 12dce3fca30ef..191220d592dc1 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -25,6 +25,7 @@ import { appendAssistantMessageToConversation, getIsKnowledgeBaseEnabled, langChainExecute, + performChecks, } from './helpers'; export const postActionsConnectorExecuteRoute = ( @@ -62,12 +63,16 @@ export const postActionsConnectorExecuteRoute = ( let onLlmResponse; try { - const authenticatedUser = assistantContext.getCurrentUser(); - if (authenticatedUser == null) { - return response.unauthorized({ - body: `Authenticated user not found`, - }); + const checkResponse = performChecks({ + context: ctx, + request, + response, + }); + + if (!checkResponse.isSuccess) { + return checkResponse.response; } + let latestReplacements: Replacements = request.body.replacements; const onNewReplacements = (newReplacements: Replacements) => { latestReplacements = { ...latestReplacements, ...newReplacements }; diff --git a/x-pack/plugins/elastic_assistant/server/routes/prompts/bulk_actions_route.ts b/x-pack/plugins/elastic_assistant/server/routes/prompts/bulk_actions_route.ts index c2713f28088d8..88178851126b5 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/prompts/bulk_actions_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/prompts/bulk_actions_route.ts @@ -35,7 +35,7 @@ import { transformESSearchToPrompts, } from '../../ai_assistant_data_clients/prompts/helpers'; import { EsPromptsSchema, UpdatePromptSchema } from '../../ai_assistant_data_clients/prompts/types'; -import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../helpers'; +import { performChecks } from '../helpers'; export interface BulkOperationError { message: string; @@ -152,22 +152,17 @@ export const bulkPromptsRoute = (router: ElasticAssistantPluginRouter, logger: L 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, + }); + if (!checkResponse.isSuccess) { + return checkResponse.response; } + const authenticatedUser = checkResponse.currentUser; - const authenticatedUser = ctx.elasticAssistant.getCurrentUser(); - if (authenticatedUser == null) { - return assistantResponse.error({ - body: `Authenticated user not found`, - statusCode: 401, - }); - } const dataClient = await ctx.elasticAssistant.getAIAssistantPromptsDataClient(); if (body.create && body.create.length > 0) { @@ -207,7 +202,7 @@ export const bulkPromptsRoute = (router: ElasticAssistantPluginRouter, logger: L ), getUpdateScript: (document: UpdatePromptSchema) => getUpdateScript({ prompt: document, isPatch: true }), - authenticatedUser, + authenticatedUser: authenticatedUser ?? undefined, }); const created = docsCreated.length > 0 diff --git a/x-pack/plugins/elastic_assistant/server/routes/prompts/find_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/prompts/find_route.test.ts index 68ce67d842a0f..151c1622d0219 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/prompts/find_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/prompts/find_route.test.ts @@ -12,6 +12,7 @@ import { requestContextMock } from '../../__mocks__/request_context'; import { getFindPromptsResultWithSingleHit } from '../../__mocks__/response'; import { findPromptsRoute } from './find_route'; import { loggingSystemMock } from '@kbn/core-logging-server-mocks'; +import type { AuthenticatedUser } from '@kbn/core-security-common'; describe('Find user prompts route', () => { let server: ReturnType; @@ -21,19 +22,26 @@ describe('Find user prompts route', () => { beforeEach(async () => { server = serverMock.create(); ({ clients, context } = requestContextMock.createTools()); + const mockUser1 = { + username: 'my_username', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, + } as AuthenticatedUser; clients.elasticAssistant.getAIAssistantPromptsDataClient.findDocuments.mockResolvedValue( Promise.resolve(getFindPromptsResultWithSingleHit()) ); - clients.elasticAssistant.getCurrentUser.mockResolvedValue({ + context.elasticAssistant.getCurrentUser.mockReturnValue({ username: 'my_username', authentication_realm: { type: 'my_realm_type', name: 'my_realm_name', }, - }); + } as AuthenticatedUser); logger = loggingSystemMock.createLogger(); - + context.elasticAssistant.getCurrentUser.mockReturnValue(mockUser1); findPromptsRoute(server.router, logger); }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/prompts/find_route.ts b/x-pack/plugins/elastic_assistant/server/routes/prompts/find_route.ts index 142b63a7d04b5..d188c615b1282 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/prompts/find_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/prompts/find_route.ts @@ -18,7 +18,7 @@ import { ElasticAssistantPluginRouter } from '../../types'; import { buildResponse } from '../utils'; import { EsPromptsSchema } from '../../ai_assistant_data_clients/prompts/types'; import { transformESSearchToPrompts } from '../../ai_assistant_data_clients/prompts/helpers'; -import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../helpers'; +import { performChecks } from '../helpers'; export const findPromptsRoute = (router: ElasticAssistantPluginRouter, logger: Logger) => { router.versioned @@ -44,13 +44,14 @@ export const findPromptsRoute = (router: ElasticAssistantPluginRouter, logger: L try { const { query } = request; 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, + }); + if (!checkResponse.isSuccess) { + return checkResponse.response; } const dataClient = await ctx.elasticAssistant.getAIAssistantPromptsDataClient(); diff --git a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts index b98ddfc6ac8f9..a71ddd84653bf 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts @@ -85,6 +85,7 @@ export class RequestContextFactory implements IRequestContextFactory { const currentUser = getCurrentUser(); return this.assistantService.createAIAssistantKnowledgeBaseDataClient({ spaceId: getSpaceId(), + licensing: context.licensing, logger: this.logger, currentUser, }); @@ -94,6 +95,7 @@ export class RequestContextFactory implements IRequestContextFactory { const currentUser = getCurrentUser(); return this.assistantService.createAttackDiscoveryDataClient({ spaceId: getSpaceId(), + licensing: context.licensing, logger: this.logger, currentUser, }); @@ -103,6 +105,7 @@ export class RequestContextFactory implements IRequestContextFactory { const currentUser = getCurrentUser(); return this.assistantService.createAIAssistantPromptsDataClient({ spaceId: getSpaceId(), + licensing: context.licensing, logger: this.logger, currentUser, }); @@ -112,6 +115,7 @@ export class RequestContextFactory implements IRequestContextFactory { const currentUser = getCurrentUser(); return this.assistantService.createAIAssistantAnonymizationFieldsDataClient({ spaceId: getSpaceId(), + licensing: context.licensing, logger: this.logger, currentUser, }); @@ -121,6 +125,7 @@ export class RequestContextFactory implements IRequestContextFactory { const currentUser = getCurrentUser(); return this.assistantService.createAIAssistantConversationsDataClient({ spaceId: getSpaceId(), + licensing: context.licensing, logger: this.logger, currentUser, }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/append_conversation_messages_route.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/append_conversation_messages_route.ts index 796c0d617fe5d..06bfa023136d9 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/append_conversation_messages_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/append_conversation_messages_route.ts @@ -17,7 +17,7 @@ import { import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; import { buildResponse } from '../utils'; import { ElasticAssistantPluginRouter } from '../../types'; -import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../helpers'; +import { performChecks } from '../helpers'; export const appendConversationMessageRoute = (router: ElasticAssistantPluginRouter) => { router.versioned @@ -43,22 +43,16 @@ export const appendConversationMessageRoute = (router: ElasticAssistantPluginRou const { id } = request.params; 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, - }, - }); + const checkResponse = performChecks({ + context: ctx, + request, + response, + }); + if (!checkResponse.isSuccess) { + return checkResponse.response; } const dataClient = await ctx.elasticAssistant.getAIAssistantConversationsDataClient(); - const authenticatedUser = ctx.elasticAssistant.getCurrentUser(); - if (authenticatedUser == null) { - return assistantResponse.error({ - body: `Authenticated user not found`, - statusCode: 401, - }); - } + const authenticatedUser = checkResponse.currentUser; const existingConversation = await dataClient?.getConversation({ id, authenticatedUser }); if (existingConversation == null) { diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.ts index 6e30acb1a47c7..9c353997f1d46 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.ts @@ -35,7 +35,7 @@ import { transformToUpdateScheme, } from '../../ai_assistant_data_clients/conversations/update_conversation'; import { EsConversationSchema } from '../../ai_assistant_data_clients/conversations/types'; -import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../helpers'; +import { performChecks } from '../helpers'; export interface BulkOperationError { message: string; @@ -156,23 +156,17 @@ export const bulkActionConversationsRoute = ( 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, - }, - }); + const checkResponse = performChecks({ + context: ctx, + request, + response, + }); + if (!checkResponse.isSuccess) { + return checkResponse.response; } + const authenticatedUser = checkResponse.currentUser; const dataClient = await ctx.elasticAssistant.getAIAssistantConversationsDataClient(); const spaceId = ctx.elasticAssistant.getSpaceId(); - const authenticatedUser = ctx.elasticAssistant.getCurrentUser(); - if (authenticatedUser == null) { - return assistantResponse.error({ - body: `Authenticated user not found`, - statusCode: 401, - }); - } if (body.create && body.create.length > 0) { const userFilter = authenticatedUser?.username diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.ts index f997121f95aa8..6421014934d14 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.ts @@ -43,14 +43,12 @@ export const createConversationRoute = (router: ElasticAssistantPluginRouter): v const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']); // Perform license and authenticated user checks const checkResponse = performChecks({ - authenticatedUser: true, context: ctx, - license: true, request, response, }); - if (checkResponse) { - return checkResponse; + if (!checkResponse.isSuccess) { + return checkResponse.response; } const dataClient = await ctx.elasticAssistant.getAIAssistantConversationsDataClient(); diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/delete_route.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/delete_route.ts index 5d761c09f682c..9c974fdb78de8 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/delete_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/delete_route.ts @@ -14,7 +14,7 @@ import { import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; import { ElasticAssistantPluginRouter } from '../../types'; import { buildResponse } from '../utils'; -import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../helpers'; +import { performChecks } from '../helpers'; export const deleteConversationRoute = (router: ElasticAssistantPluginRouter) => { router.versioned @@ -40,23 +40,18 @@ export const deleteConversationRoute = (router: ElasticAssistantPluginRouter) => const { id } = request.params; const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']); - const license = ctx.licensing.license; - if (!hasAIAssistantLicense(license)) { - return response.forbidden({ - body: { - message: UPGRADE_LICENSE_MESSAGE, - }, - }); + const checkResponse = performChecks({ + context: ctx, + request, + response, + }); + if (!checkResponse.isSuccess) { + return checkResponse.response; } const dataClient = await ctx.elasticAssistant.getAIAssistantConversationsDataClient(); - const authenticatedUser = ctx.elasticAssistant.getCurrentUser(); - if (authenticatedUser == null) { - return assistantResponse.error({ - body: `Authenticated user not found`, - statusCode: 401, - }); - } + const authenticatedUser = checkResponse.currentUser; + const existingConversation = await dataClient?.getConversation({ id, authenticatedUser }); if (existingConversation == null) { return assistantResponse.error({ diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.test.ts index 63141fe5475a6..2b20ab03371f6 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.test.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.test.ts @@ -4,7 +4,7 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ - +import { type AuthenticatedUser } from '@kbn/core/server'; import { getCurrentUserFindRequest, requestMock } from '../../__mocks__/request'; import { ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL_FIND } from '@kbn/elastic-assistant-common'; import { serverMock } from '../../__mocks__/server'; @@ -15,7 +15,6 @@ import { findUserConversationsRoute } from './find_route'; describe('Find user conversations route', () => { let server: ReturnType; let { clients, context } = requestContextMock.createTools(); - beforeEach(async () => { server = serverMock.create(); ({ clients, context } = requestContextMock.createTools()); @@ -23,13 +22,13 @@ describe('Find user conversations route', () => { clients.elasticAssistant.getAIAssistantConversationsDataClient.findDocuments.mockResolvedValue( Promise.resolve(getFindConversationsResultWithSingleHit()) ); - clients.elasticAssistant.getCurrentUser.mockResolvedValue({ + context.elasticAssistant.getCurrentUser.mockReturnValue({ username: 'my_username', authentication_realm: { type: 'my_realm_type', name: 'my_realm_name', }, - }); + } as AuthenticatedUser); findUserConversationsRoute(server.router); }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts index e7ce80039beb0..07ba23710b12c 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts @@ -21,7 +21,7 @@ import { ElasticAssistantPluginRouter } from '../../types'; import { buildResponse } from '../utils'; import { EsConversationSchema } from '../../ai_assistant_data_clients/conversations/types'; import { transformESSearchToConversations } from '../../ai_assistant_data_clients/conversations/transforms'; -import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../helpers'; +import { performChecks } from '../helpers'; export const findUserConversationsRoute = (router: ElasticAssistantPluginRouter) => { router.versioned @@ -46,16 +46,17 @@ export const findUserConversationsRoute = (router: ElasticAssistantPluginRouter) try { const { query } = request; 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, + }); + if (!checkResponse.isSuccess) { + return checkResponse.response; } const dataClient = await ctx.elasticAssistant.getAIAssistantConversationsDataClient(); - const currentUser = ctx.elasticAssistant.getCurrentUser(); + const currentUser = checkResponse.currentUser; const additionalFilter = query.filter ? ` AND ${query.filter}` : ''; const userFilter = currentUser?.username diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/read_route.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/read_route.ts index dd540897b0ece..ab69dc20999a2 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/read_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/read_route.ts @@ -16,7 +16,7 @@ import { ReadConversationRequestParams } from '@kbn/elastic-assistant-common/imp import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; import { buildResponse } from '../utils'; import { ElasticAssistantPluginRouter } from '../../types'; -import { UPGRADE_LICENSE_MESSAGE, hasAIAssistantLicense } from '../helpers'; +import { performChecks } from '../helpers'; export const readConversationRoute = (router: ElasticAssistantPluginRouter) => { router.versioned @@ -43,21 +43,15 @@ export const readConversationRoute = (router: ElasticAssistantPluginRouter) => { 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, - }, - }); - } - const authenticatedUser = ctx.elasticAssistant.getCurrentUser(); - if (authenticatedUser == null) { - return assistantResponse.error({ - body: `Authenticated user not found`, - statusCode: 401, - }); + const checkResponse = performChecks({ + context: ctx, + request, + response, + }); + if (!checkResponse.isSuccess) { + return checkResponse.response; } + const authenticatedUser = checkResponse.currentUser; const dataClient = await ctx.elasticAssistant.getAIAssistantConversationsDataClient(); const conversation = await dataClient?.getConversation({ id, authenticatedUser }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/update_route.ts b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/update_route.ts index 4ad819ef0caa0..41956b9bc80f7 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/user_conversations/update_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/user_conversations/update_route.ts @@ -45,18 +45,16 @@ export const updateConversationRoute = (router: ElasticAssistantPluginRouter) => const { id } = request.params; try { const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']); - const authenticatedUser = ctx.elasticAssistant.getCurrentUser(); // Perform license and authenticated user checks const checkResponse = performChecks({ - authenticatedUser: true, context: ctx, - license: true, request, response, }); - if (checkResponse) { - return checkResponse; + if (!checkResponse.isSuccess) { + return checkResponse.response; } + const authenticatedUser = checkResponse.currentUser; const dataClient = await ctx.elasticAssistant.getAIAssistantConversationsDataClient(); From bb68a30f9955641e23ca743d66827bfb7bc8fe85 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Wed, 30 Oct 2024 21:58:40 +0000 Subject: [PATCH 2/3] [CI] Auto-commit changed files from 'node scripts/eslint --no-cache --fix' --- .../impl/assistant/chat_send/use_chat_send.tsx | 6 +++++- .../impl/knowledge_base/setup_knowledge_base_button.tsx | 6 +++++- .../server/ai_assistant_service/index.test.ts | 2 +- .../elastic_assistant/server/ai_assistant_service/index.ts | 2 +- 4 files changed, 12 insertions(+), 4 deletions(-) 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 index d09a67da296d5..ed57d3d3bcf5c 100644 --- 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 @@ -66,7 +66,11 @@ export const useChatSend = ({ const { isLoading, sendMessage, abortStream } = useSendMessage(); const { clearConversation, removeLastMessage } = useConversation(); - const { data: kbStatus } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled , 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/setup_knowledge_base_button.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx index 582cf28be9e70..3fcbac893f477 100644 --- 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 @@ -26,7 +26,11 @@ export const SetupKnowledgeBaseButton: React.FC = React.memo(() => { assistantAvailability: { isAssistantEnabled }, } = useAssistantContext(); - const { data: kbStatus } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled , 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/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 index 23a1a55564415..30b293761c944 100644 --- 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 @@ -29,7 +29,7 @@ const licensing = Promise.resolve( license: { type: 'enterprise' }, }) ); -let logger: ReturnType<(typeof loggingSystemMock)['createLogger']>; +let logger: ReturnType; const clusterClient = elasticsearchServiceMock.createClusterClient().asInternalUser; const SimulateTemplateResponse = { 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 index 9119b29210edd..aa02dbbd4b4da 100644 --- 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,8 +11,8 @@ 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 { attackDiscoveryFieldMap } from '../ai_assistant_data_clients/attack_discovery/field_maps_configuration'; import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server'; +import { attackDiscoveryFieldMap } from '../ai_assistant_data_clients/attack_discovery/field_maps_configuration'; import { getDefaultAnonymizationFields } from '../../common/anonymization'; import { AssistantResourceNames, GetElser } from '../types'; import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations'; From c3eb7d8bc387eb0a5f15ffc813bce73687c2110d Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Wed, 6 Nov 2024 13:33:02 -0700 Subject: [PATCH 3/3] fix type --- .../knowledge_base/knowledge_base_settings_management.tsx | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management.tsx index e2c9c07aee4c3..f60576f77dfc6 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management.tsx @@ -43,7 +43,11 @@ const KNOWLEDGE_BASE_INDEX_PATTERN = '.kibana-elastic-ai-assistant-knowledge-bas * Knowledge Base Settings -- set up the Knowledge Base and configure RAG on alerts */ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => { - const { http, toasts } = useAssistantContext(); + const { + http, + toasts, + assistantAvailability: { isAssistantEnabled }, + } = useAssistantContext(); const [hasPendingChanges, setHasPendingChanges] = useState(false); const { knowledgeBase, setUpdatedKnowledgeBaseSettings, resetSettings, saveSettings } = @@ -88,7 +92,7 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => { data: kbStatus, isLoading, isFetching, - } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE }); + } = useKnowledgeBaseStatus({ enabled: isAssistantEnabled, http, resource: ESQL_RESOURCE }); const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts }); // Resource enabled state