github-mirrors/kibana
1
0
Fork 0
mirror of https://github.com/elastic/kibana.git synced 2025-04-24 01:38:56 -04:00
Code Issues Projects Releases Packages Wiki Activity

[8.16] [Security AI Assistant] Fixed license issue for Knowledge Base resources initialization (#198239) (#198451)

Browse source 
# Backport

This will backport the following commits from `main` to `8.16`:
- [[Security AI Assistant] Fixed license issue for Knowledge Base
resources initialization
(#198239)](https://github.com/elastic/kibana/pull/198239)

<!--- Backport version: 8.9.8 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sqren/backport)

<!--BACKPORT [{"author":{"name":"Yuliia
Naumenko","email":"jo.naumenko@gmail.com"},"sourceCommit":{"committedDate":"2024-10-30T21:13:33Z","message":"[Security
AI Assistant] Fixed license issue for Knowledge Base resources
initialization
(#198239)","sha":"ed81e4334f4d5608517dacdba28d46dfea966be0","branchLabelMapping":{"^v9.0.0$":"main","^v8.17.0$":"8.x","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["bug","release_note:skip","v9.0.0","v8.16.0","backport:version","v8.17.0","v8.15.4"],"number":198239,"url":"https://github.com/elastic/kibana/pull/198239","mergeCommit":{"message":"[Security
AI Assistant] Fixed license issue for Knowledge Base resources
initialization
(#198239)","sha":"ed81e4334f4d5608517dacdba28d46dfea966be0"}},"sourceBranch":"main","suggestedTargetBranches":["8.16"],"targetPullRequestStates":[{"branch":"main","label":"v9.0.0","labelRegex":"^v9.0.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/198239","number":198239,"mergeCommit":{"message":"[Security
AI Assistant] Fixed license issue for Knowledge Base resources
initialization
(#198239)","sha":"ed81e4334f4d5608517dacdba28d46dfea966be0"}},{"branch":"8.16","label":"v8.16.0","labelRegex":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"8.x","label":"v8.17.0","labelRegex":"^v8.17.0$","isSourceBranch":false,"url":"https://github.com/elastic/kibana/pull/198447","number":198447,"state":"OPEN"},{"branch":"8.15","label":"v8.15.4","labelRegex":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"url":"https://github.com/elastic/kibana/pull/198449","number":198449,"state":"OPEN"}]}]
BACKPORT-->

Co-authored-by: Yuliia Naumenko <jo.naumenko@gmail.com>
This commit is contained in:
Steph Milovic 2024-10-30 17:43:37 -06:00 • committed by GitHub
parent 1138c09886
commit f02d1303b5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
36 changed files with 291 additions and 223 deletions
Download patch file Download diff file Expand all files Collapse all files

3
x-pack/packages/kbn-elastic-assistant/impl/assistant/api/knowledge_base/use_knowledge_base_status.tsx
View file

@ -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<ReadKnowledgeBaseResponse, IHttpFetchError> => {
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

3
x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx
View file

@ -57,6 +57,9 @@ describe('use chat send', () => {
assistantTelemetry: {
reportAssistantMessageSent,
},
assistantAvailability: {
isAssistantEnabled: true,
},
});
});
it('handleOnChatCleared clears the conversation', async () => {

8
x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx
View file

@ -52,12 +52,16 @@ export const useChatSend = ({
setSelectedPromptContexts,
setCurrentConversation,
}: UseChatSendProps): UseChatSend => {
const { assistantTelemetry, toasts } = useAssistantContext();
const {
assistantTelemetry,
toasts,
assistantAvailability: { isAssistantEnabled },
} = useAssistantContext();
const [userPrompt, setUserPrompt] = useState<string | null>(null);
const { isLoading, sendMessage, abortStream } = useSendMessage();
const { clearConversation, removeLastMessage } = useConversation();
const { data: kbStatus } = useKnowledgeBaseStatus({ http });
const { data: kbStatus } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled });
const isSetupComplete =
kbStatus?.elser_exists &&
kbStatus?.index_exists &&

3
x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.test.tsx
View file

@ -26,6 +26,9 @@ const mockUseAssistantContext = {
},
setAllSystemPrompts: jest.fn(),
setConversations: jest.fn(),
assistantAvailability: {
isAssistantEnabled: true,
},
};
jest.mock('../assistant_context', () => {

12
x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx
View file

@ -44,8 +44,16 @@ interface Props {
*/
export const KnowledgeBaseSettings: React.FC<Props> = React.memo(
({ knowledgeBase, setUpdatedKnowledgeBaseSettings, modalMode = false }) => {
const { http, toasts } = useAssistantContext();
const { data: kbStatus, isLoading, isFetching } = useKnowledgeBaseStatus({ http });
const {
http,
toasts,
assistantAvailability: { isAssistantEnabled },
} = useAssistantContext();
const {
data: kbStatus,
isLoading,
isFetching,
} = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled });
const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts });
// Resource enabled state

9
x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx
View file

@ -74,12 +74,15 @@ interface Params {
export const KnowledgeBaseSettingsManagement: React.FC<Params> = React.memo(({ dataViews }) => {
const {
assistantFeatures: { assistantKnowledgeBaseByDefault: enableKnowledgeBaseByDefault },
assistantAvailability: { hasManageGlobalKnowledgeBase },
assistantAvailability: { hasManageGlobalKnowledgeBase, isAssistantEnabled },
http,
toasts,
} = useAssistantContext();
const [hasPendingChanges, setHasPendingChanges] = useState(false);
const { data: kbStatus, isFetched } = useKnowledgeBaseStatus({ http });
const { data: kbStatus, isFetched } = useKnowledgeBaseStatus({
http,
enabled: isAssistantEnabled,
});
const isKbSetup = isKnowledgeBaseSetup(kbStatus);
const [deleteKBItem, setDeleteKBItem] = useState<DocumentEntry | IndexEntry | null>(null);
@ -159,7 +162,7 @@ export const KnowledgeBaseSettingsManagement: React.FC<Params> = React.memo(({ d
} = useKnowledgeBaseEntries({
http,
toasts,
enabled: enableKnowledgeBaseByDefault,
enabled: enableKnowledgeBaseByDefault && isAssistantEnabled,
isRefetching: kbStatus?.is_setup_in_progress,
});

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

@ -23,9 +23,13 @@ interface Props {
*
*/
export const SetupKnowledgeBaseButton: React.FC<Props> = React.memo(({ display }: Props) => {
const { http, toasts } = useAssistantContext();
const {
http,
toasts,
assistantAvailability: { isAssistantEnabled },
} = useAssistantContext();
const { data: kbStatus } = useKnowledgeBaseStatus({ http });
const { data: kbStatus } = useKnowledgeBaseStatus({ http, enabled: isAssistantEnabled });
const { mutate: setupKB, isLoading: isSettingUpKB } = useSetupKnowledgeBase({ http, toasts });
const isSetupInProgress = kbStatus?.is_setup_in_progress || isSettingUpKB;

8
x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.test.ts
View file

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

4
x-pack/plugins/elastic_assistant/server/ai_assistant_service/create_resource_installation_helper.ts
View file

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

17
x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts
View file

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

9
x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts
View file

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

36
x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/bulk_actions_route.ts
View file

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

22
x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.test.ts
View file

@ -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<typeof serverMock.create>;
@ -21,19 +22,26 @@ describe('Find user anonymization fields route', () => {
beforeEach(async () => {
server = serverMock.create();
({ clients, context } = requestContextMock.createTools());
clients.elasticAssistant.getAIAssistantAnonymizationFieldsDataClient.findDocuments.mockResolvedValue(
Promise.resolve(getFindAnonymizationFieldsResultWithSingleHit())
);
clients.elasticAssistant.getCurrentUser.mockResolvedValue({
const mockUser1 = {
username: 'my_username',
authentication_realm: {
type: 'my_realm_type',
name: 'my_realm_name',
},
});
logger = loggingSystemMock.createLogger();
} as AuthenticatedUser;
clients.elasticAssistant.getAIAssistantAnonymizationFieldsDataClient.findDocuments.mockResolvedValue(
Promise.resolve(getFindAnonymizationFieldsResultWithSingleHit())
);
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);
});

18
x-pack/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.ts
View file

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

12
x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts
View file

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

6
x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts
View file

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

7
x-pack/plugins/elastic_assistant/server/routes/evaluate/get_evaluate.ts
View file

@ -48,15 +48,14 @@ export const getEvaluateRoute = (router: IRouter<ElasticAssistantRequestHandlerC
// Perform license, authenticated user and evaluation FF checks
const checkResponse = performChecks({
authenticatedUser: true,
capability: 'assistantModelEvaluation',
context: ctx,
license: true,
request,
response,
});
if (checkResponse) {
return checkResponse;
if (!checkResponse.isSuccess) {
return checkResponse.response;
}
// Fetch datasets from LangSmith // TODO: plumb apiKey so this will work in cloud w/o env vars

6
x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts
View file

@ -95,15 +95,13 @@ export const postEvaluateRoute = (
// Perform license, authenticated user and evaluation FF checks
const checkResponse = performChecks({
authenticatedUser: true,
capability: 'assistantModelEvaluation',
context: ctx,
license: true,
request,
response,
});
if (checkResponse) {
return checkResponse;
if (!checkResponse.isSuccess) {
return checkResponse.response;
}
try {

62
x-pack/plugins/elastic_assistant/server/routes/helpers.ts
View file

@ -7,6 +7,7 @@
import {
AnalyticsServiceSetup,
type AuthenticatedUser,
IKibanaResponse,
KibanaRequest,
KibanaResponseFactory,
@ -561,55 +562,66 @@ export const updateConversationWithUserInput = async ({
};
interface PerformChecksParams {
authenticatedUser?: boolean;
capability?: AssistantFeatureKey;
context: AwaitedProperties<
Pick<ElasticAssistantRequestHandlerContext, 'elasticAssistant' | 'licensing' | 'core'>
>;
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,
};
};
/**

11
x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts
View file

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

6
x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts
View file

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

8
x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts
View file

@ -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}` : '';

13
x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts
View file

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

15
x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts
View file

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

27
x-pack/plugins/elastic_assistant/server/routes/prompts/bulk_actions_route.ts
View file

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

22
x-pack/plugins/elastic_assistant/server/routes/prompts/find_route.test.ts
View file

@ -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<typeof serverMock.create>;
@ -21,19 +22,26 @@ describe('Find user prompts route', () => {
beforeEach(async () => {
server = serverMock.create();
({ clients, context } = requestContextMock.createTools());
clients.elasticAssistant.getAIAssistantPromptsDataClient.findDocuments.mockResolvedValue(
Promise.resolve(getFindPromptsResultWithSingleHit())
);
clients.elasticAssistant.getCurrentUser.mockResolvedValue({
const mockUser1 = {
username: 'my_username',
authentication_realm: {
type: 'my_realm_type',
name: 'my_realm_name',
},
});
logger = loggingSystemMock.createLogger();
} as AuthenticatedUser;
clients.elasticAssistant.getAIAssistantPromptsDataClient.findDocuments.mockResolvedValue(
Promise.resolve(getFindPromptsResultWithSingleHit())
);
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);
});

17
x-pack/plugins/elastic_assistant/server/routes/prompts/find_route.ts
View file

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

5
x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts
View file

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

24
x-pack/plugins/elastic_assistant/server/routes/user_conversations/append_conversation_messages_route.ts
View file

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

24
x-pack/plugins/elastic_assistant/server/routes/user_conversations/bulk_actions_route.ts
View file

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

6
x-pack/plugins/elastic_assistant/server/routes/user_conversations/create_route.ts
View file

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

25
x-pack/plugins/elastic_assistant/server/routes/user_conversations/delete_route.ts
View file

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

7
x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.test.ts
View file

@ -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<typeof serverMock.create>;
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);
});

19
x-pack/plugins/elastic_assistant/server/routes/user_conversations/find_route.ts
View file

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

24
x-pack/plugins/elastic_assistant/server/routes/user_conversations/read_route.ts
View file

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

8
x-pack/plugins/elastic_assistant/server/routes/user_conversations/update_route.ts
View file

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