[Security Solution] [Elastic AI Assistant] Adds support for plugin feature registration

.buildkite/scripts/steps/checks.sh

@@ -23,6 +23,7 @@ export DISABLE_BOOTSTRAP_VALIDATION=false
 .buildkite/scripts/steps/checks/ftr_configs.sh
 .buildkite/scripts/steps/checks/saved_objects_compat_changes.sh
 .buildkite/scripts/steps/checks/saved_objects_definition_change.sh
+.buildkite/scripts/steps/code_generation/elastic_assistant_codegen.sh
 .buildkite/scripts/steps/code_generation/security_solution_codegen.sh
 .buildkite/scripts/steps/code_generation/osquery_codegen.sh
 .buildkite/scripts/steps/checks/yarn_deduplicate.sh

.buildkite/scripts/steps/code_generation/elastic_assistant_codegen.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+source .buildkite/scripts/common/util.sh
+
+echo --- Elastic Assistant OpenAPI Code Generation
+
+(cd x-pack/plugins/elastic_assistant && yarn openapi:generate)
+check_for_changed_files "yarn openapi:generate" true

x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/README.md

@@ -0,0 +1,51 @@
+### Feature Capabilities
+
+Feature capabilities are an object describing specific capabilities of the assistant, like whether a feature like streaming is enabled, and are defined in the sibling `./index.ts` file within this `kbn-elastic-assistant-common` package. These capabilities can be registered for a given plugin through the assistant server, and so do not need to be plumbed through the `ElasticAssistantProvider`.
+
+Storage and accessor functions are made available via the `AppContextService`, and exposed to clients via the`/internal/elastic_assistant/capabilities` route, which can be fetched by clients using the `useCapabilities()` UI hook.
+
+### Registering Capabilities
+
+To register a capability on plugin start, add the following in the consuming plugin's `start()`, specifying any number of capabilities you would like to explicitly declare:
+
+```ts
+plugins.elasticAssistant.registerFeatures(APP_UI_ID, {
+  assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation,
+  assistantStreamingEnabled: config.experimentalFeatures.assistantStreamingEnabled,
+});
+```
+
+### Declaring Feature Capabilities
+Default feature capabilities are declared in `x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts`:
+
+```ts
+export type AssistantFeatures = { [K in keyof typeof assistantFeatures]: boolean };
+
+export const assistantFeatures = Object.freeze({
+  assistantModelEvaluation: false,
+  assistantStreamingEnabled: false,
+});
+```
+
+### Using Capabilities Client Side
+Capabilities can be fetched client side using the `useCapabilities()` hook ala:
+
+```ts
+const { data: capabilities } = useCapabilities({ http, toasts });
+const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } = capabilities ?? assistantFeatures;
+```
+
+### Using Capabilities Server Side
+Or server side within a route (or elsewhere) via the `assistantContext`:
+
+```ts
+const assistantContext = await context.elasticAssistant;
+const pluginName = getPluginNameFromRequest({ request, logger });
+const registeredFeatures = assistantContext.getRegisteredFeatures(pluginName);
+if (!registeredFeatures.assistantModelEvaluation) {
+  return response.notFound();
+}
+```
+
+> [!NOTE]
+> Note, just as with [registering arbitrary tools](https://github.com/elastic/kibana/pull/172234), features are registered for a specific plugin, where the plugin name that corresponds to your application is defined in the `x-kbn-context` header of requests made from your application, which may be different than your plugin's registered `APP_ID`.

x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts

@@ -0,0 +1,19 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+/**
+ * Interface for features available to the elastic assistant
+ */
+export type AssistantFeatures = { [K in keyof typeof assistantFeatures]: boolean };
+
+/**
+ * Default features available to the elastic assistant
+ */
+export const assistantFeatures = Object.freeze({
+  assistantModelEvaluation: false,
+  assistantStreamingEnabled: false,
+});

x-pack/packages/kbn-elastic-assistant-common/index.ts

@@ -5,6 +5,9 @@
  * 2.0.
  */
 
+export { assistantFeatures } from './impl/capabilities';
+export type { AssistantFeatures } from './impl/capabilities';
+
 export { getAnonymizedValue } from './impl/data_anonymization/get_anonymized_value';
 
 export {

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.test.tsx

@@ -0,0 +1,43 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { HttpSetup } from '@kbn/core-http-browser';
+
+import { getCapabilities } from './capabilities';
+import { API_ERROR } from '../../translations';
+
+jest.mock('@kbn/core-http-browser');
+
+const mockHttp = {
+  fetch: jest.fn(),
+} as unknown as HttpSetup;
+
+describe('Capabilities API tests', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  describe('getCapabilities', () => {
+    it('calls the internal assistant API for fetching assistant capabilities', async () => {
+      await getCapabilities({ http: mockHttp });
+
+      expect(mockHttp.fetch).toHaveBeenCalledWith('/internal/elastic_assistant/capabilities', {
+        method: 'GET',
+        signal: undefined,
+        version: '1',
+      });
+    });
+
+    it('returns API_ERROR when the response status is error', async () => {
+      (mockHttp.fetch as jest.Mock).mockResolvedValue({ status: API_ERROR });
+
+      const result = await getCapabilities({ http: mockHttp });
+
+      expect(result).toEqual({ status: API_ERROR });
+    });
+  });
+});

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/capabilities.tsx

@@ -0,0 +1,44 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { HttpSetup, IHttpFetchError } from '@kbn/core-http-browser';
+import { AssistantFeatures } from '@kbn/elastic-assistant-common';
+
+export interface GetCapabilitiesParams {
+  http: HttpSetup;
+  signal?: AbortSignal | undefined;
+}
+
+export type GetCapabilitiesResponse = AssistantFeatures;
+
+/**
+ * API call for fetching assistant capabilities
+ *
+ * @param {Object} options - The options object.
+ * @param {HttpSetup} options.http - HttpSetup
+ * @param {AbortSignal} [options.signal] - AbortSignal
+ *
+ * @returns {Promise<GetCapabilitiesResponse | IHttpFetchError>}
+ */
+export const getCapabilities = async ({
+  http,
+  signal,
+}: GetCapabilitiesParams): Promise<GetCapabilitiesResponse | IHttpFetchError> => {
+  try {
+    const path = `/internal/elastic_assistant/capabilities`;
+
+    const response = await http.fetch(path, {
+      method: 'GET',
+      signal,
+      version: '1',
+    });
+
+    return response as GetCapabilitiesResponse;
+  } catch (error) {
+    return error as IHttpFetchError;
+  }
+};

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.test.tsx

@@ -0,0 +1,49 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { renderHook } from '@testing-library/react-hooks';
+
+import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
+import type { ReactNode } from 'react';
+import React from 'react';
+import { useCapabilities, UseCapabilitiesParams } from './use_capabilities';
+
+const statusResponse = { assistantModelEvaluation: true, assistantStreamingEnabled: false };
+
+const http = {
+  fetch: jest.fn().mockResolvedValue(statusResponse),
+};
+const toasts = {
+  addError: jest.fn(),
+};
+const defaultProps = { http, toasts } as unknown as UseCapabilitiesParams;
+
+const createWrapper = () => {
+  const queryClient = new QueryClient();
+  // eslint-disable-next-line react/display-name
+  return ({ children }: { children: ReactNode }) => (
+    <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+  );
+};
+
+describe('useFetchRelatedCases', () => {
+  it(`should make http request to fetch capabilities`, () => {
+    renderHook(() => useCapabilities(defaultProps), {
+      wrapper: createWrapper(),
+    });
+
+    expect(defaultProps.http.fetch).toHaveBeenCalledWith(
+      '/internal/elastic_assistant/capabilities',
+      {
+        method: 'GET',
+        version: '1',
+        signal: new AbortController().signal,
+      }
+    );
+    expect(toasts.addError).not.toHaveBeenCalled();
+  });
+});

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/capabilities/use_capabilities.tsx

@@ -0,0 +1,52 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { UseQueryResult } from '@tanstack/react-query';
+import { useQuery } from '@tanstack/react-query';
+import type { HttpSetup, IHttpFetchError, ResponseErrorBody } from '@kbn/core-http-browser';
+import type { IToasts } from '@kbn/core-notifications-browser';
+import { i18n } from '@kbn/i18n';
+import { getCapabilities, GetCapabilitiesResponse } from './capabilities';
+
+const CAPABILITIES_QUERY_KEY = ['elastic-assistant', 'capabilities'];
+
+export interface UseCapabilitiesParams {
+  http: HttpSetup;
+  toasts?: IToasts;
+}
+/**
+ * Hook for getting the feature capabilities of the assistant
+ *
+ * @param {Object} options - The options object.
+ * @param {HttpSetup} options.http - HttpSetup
+ * @param {IToasts} options.toasts - IToasts
+ *
+ * @returns {useQuery} hook for getting the status of the Knowledge Base
+ */
+export const useCapabilities = ({
+  http,
+  toasts,
+}: UseCapabilitiesParams): UseQueryResult<GetCapabilitiesResponse, IHttpFetchError> => {
+  return useQuery({
+    queryKey: CAPABILITIES_QUERY_KEY,
+    queryFn: async ({ signal }) => {
+      return getCapabilities({ http, signal });
+    },
+    retry: false,
+    keepPreviousData: true,
+    // Deprecated, hoist to `queryCache` w/in `QueryClient. See: https://stackoverflow.com/a/76961109
+    onError: (error: IHttpFetchError<ResponseErrorBody>) => {
+      if (error.name !== 'AbortError') {
+        toasts?.addError(error.body && error.body.message ? new Error(error.body.message) : error, {
+          title: i18n.translate('xpack.elasticAssistant.capabilities.statusError', {
+            defaultMessage: 'Error fetching capabilities',
+          }),
+        });
+      }
+    },
+  });
+};

x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.test.tsx

@@ -6,53 +6,14 @@
  */
 
 import { renderHook } from '@testing-library/react-hooks';
-import React from 'react';
 
-import { AssistantProvider, useAssistantContext } from '.';
-import { httpServiceMock } from '@kbn/core-http-browser-mocks';
-import { actionTypeRegistryMock } from '@kbn/triggers-actions-ui-plugin/public/application/action_type_registry.mock';
-import { AssistantAvailability } from '../..';
+import { useAssistantContext } from '.';
 import { useLocalStorage } from 'react-use';
+import { TestProviders } from '../mock/test_providers/test_providers';
 
 jest.mock('react-use', () => ({
   useLocalStorage: jest.fn().mockReturnValue(['456', jest.fn()]),
 }));
-const actionTypeRegistry = actionTypeRegistryMock.create();
-const mockGetInitialConversations = jest.fn(() => ({}));
-const mockGetComments = jest.fn(() => []);
-const mockHttp = httpServiceMock.createStartContract({ basePath: '/test' });
-const mockAssistantAvailability: AssistantAvailability = {
-  hasAssistantPrivilege: false,
-  hasConnectorsAllPrivilege: true,
-  hasConnectorsReadPrivilege: true,
-  isAssistantEnabled: true,
-};
-
-const ContextWrapper: React.FC = ({ children }) => (
-  <AssistantProvider
-    actionTypeRegistry={actionTypeRegistry}
-    assistantAvailability={mockAssistantAvailability}
-    assistantStreamingEnabled
-    augmentMessageCodeBlocks={jest.fn()}
-    baseAllow={[]}
-    baseAllowReplacement={[]}
-    basePath={'https://localhost:5601/kbn'}
-    defaultAllow={[]}
-    defaultAllowReplacement={[]}
-    docLinks={{
-      ELASTIC_WEBSITE_URL: 'https://www.elastic.co/',
-      DOC_LINK_VERSION: 'current',
-    }}
-    getInitialConversations={mockGetInitialConversations}
-    getComments={mockGetComments}
-    http={mockHttp}
-    setConversations={jest.fn()}
-    setDefaultAllow={jest.fn()}
-    setDefaultAllowReplacement={jest.fn()}
-  >
-    {children}
-  </AssistantProvider>
-);
 
 describe('AssistantContext', () => {
   beforeEach(() => jest.clearAllMocks());
@@ -66,30 +27,29 @@ describe('AssistantContext', () => {
   });
 
   test('it should return the httpFetch function', async () => {
-    const { result } = renderHook(useAssistantContext, { wrapper: ContextWrapper });
-    const http = await result.current.http;
+    const { result } = renderHook(useAssistantContext, { wrapper: TestProviders });
 
     const path = '/path/to/resource';
-    await http.fetch(path);
+    await result.current.http.fetch(path);
 
-    expect(mockHttp.fetch).toBeCalledWith(path);
+    expect(result.current.http.fetch).toBeCalledWith(path);
   });
 
   test('getConversationId defaults to provided id', async () => {
-    const { result } = renderHook(useAssistantContext, { wrapper: ContextWrapper });
+    const { result } = renderHook(useAssistantContext, { wrapper: TestProviders });
     const id = result.current.getConversationId('123');
     expect(id).toEqual('123');
   });
 
   test('getConversationId uses local storage id when no id is provided ', async () => {
-    const { result } = renderHook(useAssistantContext, { wrapper: ContextWrapper });
+    const { result } = renderHook(useAssistantContext, { wrapper: TestProviders });
     const id = result.current.getConversationId();
     expect(id).toEqual('456');
   });
 
   test('getConversationId defaults to Welcome when no local storage id and no id is provided ', async () => {
     (useLocalStorage as jest.Mock).mockReturnValue([undefined, jest.fn()]);
-    const { result } = renderHook(useAssistantContext, { wrapper: ContextWrapper });
+    const { result } = renderHook(useAssistantContext, { wrapper: TestProviders });
     const id = result.current.getConversationId();
     expect(id).toEqual('Welcome');
   });