From ed81e4334f4d5608517dacdba28d46dfea966be0 Mon Sep 17 00:00:00 2001 From: Yuliia Naumenko Date: Wed, 30 Oct 2024 14:13:33 -0700 Subject: [PATCH] [Security AI Assistant] Fixed license issue for Knowledge Base resources initialization (#198239) --- .../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 | 12 +++- .../index.tsx | 9 ++- .../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 | 17 +++++ .../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 +-- 36 files changed, 283 insertions(+), 215 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 45c6d011b46d4..3ae89edc2a912 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, // Polling interval for Knowledge Base setup in progress diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx index 0de7adc484fc1..f72f85892d379 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 @@ -57,6 +57,9 @@ describe('use chat send', () => { assistantTelemetry: { reportAssistantMessageSent, }, + assistantAvailability: { + isAssistantEnabled: true, + }, }); }); it('handleOnChatCleared clears the conversation', async () => { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx index 4ea376518b5a7..c240d5ac6b60b 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 @@ -52,12 +52,16 @@ export const useChatSend = ({ setSelectedPromptContexts, setCurrentConversation, }: UseChatSendProps): UseChatSend => { - const { assistantTelemetry, toasts } = useAssistantContext(); + const { + assistantTelemetry, + toasts, + assistantAvailability: { isAssistantEnabled }, + } = useAssistantContext(); const [userPrompt, setUserPrompt] = useState(null); const { isLoading, sendMessage, abortStream } = useSendMessage(); const { clearConversation, removeLastMessage } = useConversation(); - const { data: kbStatus } = useKnowledgeBaseStatus({ http }); + const { data: kbStatus } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled }); const isSetupComplete = kbStatus?.elser_exists && kbStatus?.index_exists && diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx index 3d18885902326..763a2578ee273 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 7041bf909601f..18bc0cbe5a384 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,8 +44,16 @@ interface Props { */ export const KnowledgeBaseSettings: React.FC = React.memo( ({ knowledgeBase, setUpdatedKnowledgeBaseSettings, modalMode = false }) => { - const { http, toasts } = useAssistantContext(); - const { data: kbStatus, isLoading, isFetching } = useKnowledgeBaseStatus({ http }); + const { + http, + toasts, + assistantAvailability: { isAssistantEnabled }, + } = useAssistantContext(); + const { + data: kbStatus, + isLoading, + isFetching, + } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled }); const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts }); // Resource enabled state diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx index bc2d60941679a..092cc7e36689e 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx @@ -74,12 +74,15 @@ interface Params { export const KnowledgeBaseSettingsManagement: React.FC = React.memo(({ dataViews }) => { const { assistantFeatures: { assistantKnowledgeBaseByDefault: enableKnowledgeBaseByDefault }, - assistantAvailability: { hasManageGlobalKnowledgeBase }, + assistantAvailability: { hasManageGlobalKnowledgeBase, isAssistantEnabled }, http, toasts, } = useAssistantContext(); const [hasPendingChanges, setHasPendingChanges] = useState(false); - const { data: kbStatus, isFetched } = useKnowledgeBaseStatus({ http }); + const { data: kbStatus, isFetched } = useKnowledgeBaseStatus({ + http, + enabled: isAssistantEnabled, + }); const isKbSetup = isKnowledgeBaseSetup(kbStatus); const [deleteKBItem, setDeleteKBItem] = useState(null); @@ -159,7 +162,7 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(({ d } = useKnowledgeBaseEntries({ http, toasts, - enabled: enableKnowledgeBaseByDefault, + enabled: enableKnowledgeBaseByDefault && isAssistantEnabled, isRefetching: kbStatus?.is_setup_in_progress, }); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx index 948e45232028c..41656c968d38e 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 @@ -23,9 +23,13 @@ interface Props { * */ export const SetupKnowledgeBaseButton: React.FC = React.memo(({ display }: Props) => { - const { http, toasts } = useAssistantContext(); + const { + http, + toasts, + assistantAvailability: { isAssistantEnabled }, + } = useAssistantContext(); - const { data: kbStatus } = useKnowledgeBaseStatus({ http }); + const { data: kbStatus } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled }); const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts }); const isSetupInProgress = kbStatus?.is_setup_in_progress || isSettingUpKB; diff --git a/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 8e34581332ff6..85f6de83592ae 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 8bd1173e93d89..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,11 +18,17 @@ import { AIAssistantService, AIAssistantServiceOpts } from '.'; import { retryUntil } from './create_resource_installation_helper.test'; import { mlPluginMock } from '@kbn/ml-plugin/public/mocks'; import type { MlPluginSetup } from '@kbn/ml-plugin/server'; +import { licensingMock } from '@kbn/licensing-plugin/server/mocks'; jest.mock('../ai_assistant_data_clients/conversations', () => ({ AIAssistantConversationsDataClient: jest.fn(), })); +const licensing = Promise.resolve( + licensingMock.createRequestHandlerContext({ + license: { type: 'enterprise' }, + }) +); let logger: ReturnType<(typeof loggingSystemMock)['createLogger']>; const clusterClient = elasticsearchServiceMock.createClusterClient().asInternalUser; @@ -191,6 +197,7 @@ describe('AI Assistant Service', () => { logger, spaceId: 'default', currentUser: mockUser1, + licensing, }); expect(AIAssistantConversationsDataClient).toHaveBeenCalledWith({ @@ -221,6 +228,7 @@ describe('AI Assistant Service', () => { logger, spaceId: 'default', currentUser: mockUser1, + licensing, }); expect(clusterClient.indices.putIndexTemplate).toHaveBeenCalled(); @@ -274,11 +282,13 @@ describe('AI Assistant Service', () => { logger, spaceId: 'default', currentUser: mockUser1, + licensing, }), assistantService.createAIAssistantConversationsDataClient({ logger, spaceId: 'default', currentUser: mockUser1, + licensing, }), ]); @@ -340,6 +350,7 @@ describe('AI Assistant Service', () => { logger, spaceId: 'default', currentUser: mockUser1, + licensing, }); expect(AIAssistantConversationsDataClient).toHaveBeenCalledWith({ @@ -400,6 +411,7 @@ describe('AI Assistant Service', () => { logger, spaceId: 'default', currentUser: mockUser1, + licensing, }); }; @@ -472,6 +484,7 @@ describe('AI Assistant Service', () => { logger, spaceId: 'default', currentUser: mockUser1, + licensing, }); }; @@ -513,6 +526,7 @@ describe('AI Assistant Service', () => { logger, spaceId: 'test', currentUser: mockUser1, + licensing, }); expect(clusterClient.indices.putIndexTemplate).not.toHaveBeenCalled(); @@ -560,6 +574,7 @@ describe('AI Assistant Service', () => { logger, spaceId: 'test', currentUser: mockUser1, + licensing, }); expect(clusterClient.indices.putIndexTemplate).not.toHaveBeenCalled(); @@ -607,6 +622,7 @@ describe('AI Assistant Service', () => { logger, spaceId: 'test', currentUser: mockUser1, + licensing, }); expect(AIAssistantConversationsDataClient).not.toHaveBeenCalled(); @@ -752,6 +768,7 @@ describe('AI Assistant Service', () => { logger, spaceId: 'default', currentUser: mockUser1, + licensing, }); await retryUntil( diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index a7b54dd5ca4be..15274f2323259 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,6 +11,7 @@ import type { AuthenticatedUser, Logger, ElasticsearchClient } from '@kbn/core/s import type { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server'; import type { MlPluginSetup } from '@kbn/ml-plugin/server'; import { Subject } from 'rxjs'; +import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server'; import { attackDiscoveryFieldMap } from '../lib/attack_discovery/persistence/field_maps_configuration/field_maps_configuration'; import { getDefaultAnonymizationFields } from '../../common/anonymization'; import { AssistantResourceNames, GetElser } from '../types'; @@ -36,6 +37,7 @@ import { } from '../ai_assistant_data_clients/knowledge_base'; import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence'; import { createGetElserId, createPipeline, pipelineExists } from './helpers'; +import { hasAIAssistantLicense } from '../routes/helpers'; const TOTAL_FIELDS_LIMIT = 2500; @@ -56,6 +58,7 @@ export interface CreateAIAssistantClientParams { logger: Logger; spaceId: string; currentUser: AuthenticatedUser | null; + licensing: Promise; } export type CreateDataStream = (params: { @@ -245,7 +248,7 @@ export class AIAssistantService { pluginStop$: this.options.pluginStop$, }); } catch (error) { - this.options.logger.error(`Error initializing AI assistant resources: ${error.message}`); + this.options.logger.warn(`Error initializing AI assistant resources: ${error.message}`); this.initialized = false; this.isInitializing = false; return errorResult(error.message); @@ -290,6 +293,8 @@ export class AIAssistantService { }; private async checkResourcesInstallation(opts: CreateAIAssistantClientParams) { + const licensing = await opts.licensing; + if (!hasAIAssistantLicense(licensing.license)) return null; // Check if resources installation has succeeded const { result: initialized, error } = await this.getSpaceResourcesInitializationPromise( opts.spaceId @@ -510,7 +515,7 @@ export class AIAssistantService { await this.createDefaultAnonymizationFields(spaceId); } } catch (error) { - this.options.logger.error( + this.options.logger.warn( `Error initializing AI assistant namespace level resources: ${error.message}` ); throw error; diff --git a/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/bulk_actions_route.ts b/x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/bulk_actions_route.ts index 5464756739c08..9aedffae5cfb5 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; @@ -162,22 +162,18 @@ export const bulkActionAnonymizationFieldsRoute = ( request.events.completed$.subscribe(() => abortController.abort()); try { const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']); - const license = ctx.licensing.license; - if (!hasAIAssistantLicense(license)) { - return response.forbidden({ - body: { - message: UPGRADE_LICENSE_MESSAGE, - }, - }); - } + // Perform license and authenticated user checks + const checkResponse = performChecks({ + context: ctx, + request, + response, + }); - const authenticatedUser = ctx.elasticAssistant.getCurrentUser(); - if (authenticatedUser == null) { - return assistantResponse.error({ - body: `Authenticated user not found`, - statusCode: 401, - }); + if (!checkResponse.isSuccess) { + return checkResponse.response; } + const authenticatedUser = checkResponse.currentUser; + const dataClient = await ctx.elasticAssistant.getAIAssistantAnonymizationFieldsDataClient(); @@ -199,7 +195,7 @@ export const bulkActionAnonymizationFieldsRoute = ( } const writer = await dataClient?.getWriter(); - const changedAt = new Date().toISOString(); + const createdAt = new Date().toISOString(); const { errors, docs_created: docsCreated, @@ -207,12 +203,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 47f6f1a486957..f2365f0320967 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 @@ -71,14 +71,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) { @@ -619,11 +631,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 cfb2303010756..fbe73525578b0 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 { @@ -143,15 +143,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( @@ -181,8 +179,7 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug v2KnowledgeBaseEnabled: true, }); const spaceId = ctx.elasticAssistant.getSpaceId(); - // Authenticated user null check completed in `performChecks()` above - const authenticatedUser = ctx.elasticAssistant.getCurrentUser() as AuthenticatedUser; + const authenticatedUser = checkResponse.currentUser; const userFilter = getKBUserFilter(authenticatedUser); const manageGlobalKnowledgeBaseAIAssistant = kbDataClient?.options.manageGlobalKnowledgeBaseAIAssistant; 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 96753bdd690bd..0bfe9de269f7c 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 @@ -47,15 +47,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; } // Check mappings and upgrade if necessary -- this route only supports v2 KB, so always `true` 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 356d5d9150a67..13334d0d829b1 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 @@ -58,21 +58,19 @@ 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({ v2KnowledgeBaseEnabled: true, }); - const currentUser = ctx.elasticAssistant.getCurrentUser(); + const currentUser = checkResponse.currentUser; const userFilter = getKBUserFilter(currentUser); const systemFilter = ` AND (kb_resource:"user" OR type:"index")`; const additionalFilter = query.filter ? ` AND ${query.filter}` : ''; 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 d19127be0d7e8..998790f332c45 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 4b65b5bb3f1e5..acf2dd32a060b 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 @@ -28,6 +28,7 @@ import { getPluginNameFromRequest, getSystemPromptFromUserConversation, langChainExecute, + performChecks, } from './helpers'; import { isOpenSourceModel } from './utils'; @@ -66,12 +67,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 44a949cd22eeb..d3ee47854e7a0 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; @@ -156,22 +156,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) { @@ -211,7 +206,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 848680be662a3..a2980b173d76a 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 eeb1a5564d1cf..7d97029e7252a 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 @@ -101,6 +101,7 @@ export class RequestContextFactory implements IRequestContextFactory { return this.assistantService.createAIAssistantKnowledgeBaseDataClient({ spaceId: getSpaceId(), logger: this.logger, + licensing: context.licensing, currentUser, modelIdOverride, v2KnowledgeBaseEnabled, @@ -114,6 +115,7 @@ export class RequestContextFactory implements IRequestContextFactory { const currentUser = getCurrentUser(); return this.assistantService.createAttackDiscoveryDataClient({ spaceId: getSpaceId(), + licensing: context.licensing, logger: this.logger, currentUser, }); @@ -123,6 +125,7 @@ export class RequestContextFactory implements IRequestContextFactory { const currentUser = getCurrentUser(); return this.assistantService.createAIAssistantPromptsDataClient({ spaceId: getSpaceId(), + licensing: context.licensing, logger: this.logger, currentUser, }); @@ -132,6 +135,7 @@ export class RequestContextFactory implements IRequestContextFactory { const currentUser = getCurrentUser(); return this.assistantService.createAIAssistantAnonymizationFieldsDataClient({ spaceId: getSpaceId(), + licensing: context.licensing, logger: this.logger, currentUser, }); @@ -141,6 +145,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 b92ad5462963e..9955494b5f294 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 @@ -44,14 +44,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();