[Security Solution] [Elastic AI Assistant] Adds support for plugin fe…

…ature registration (#174317)

## Summary

Resolves #172509

Adds ability to register feature capabilities through the assistant
server so they no longer need to be plumbed through the
`ElasticAssistantProvider`, which now also makes them available server
side.

Adds new `/internal/elastic_assistant/capabilities` route and
`useCapabilities()` UI hook for fetching capabilities.

### OpenAPI Codegen

Implemented using the new OpenAPI codegen and bundle packages:
* Includes OpenAPI codegen script and CI action as detailed in:
#166269
* Includes OpenAPI docs bundling script as detailed in:
#171526

To run codegen/bundling locally, cd to
`x-pack/plugins/elastic_assistant/` and run any of the following
commands:

```bash
yarn openapi:generate
yarn openapi:generate:debug
yarn openapi:bundle
```

> [!NOTE]
> At the moment `yarn openapi:bundle` will output an empty bundled
schema since `get_capabilities_route` is an internal route, this is to
be expected. Also, if you don't see the file in your IDE, it's probably
because `target` directories are ignored, so you may need to manually
find/open the bundled schema at it's target location:
`/x-pack/plugins/elastic_assistant/target/openapi/elastic_assistant.bundled.schema.yaml`

### Registering Capabilities 

To register a capability on plugin start, add the following in the
consuming plugin's `start()`:

```ts
plugins.elasticAssistant.registerFeatures(APP_UI_ID, {
  assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation,
  assistantStreamingEnabled: config.experimentalFeatures.assistantStreamingEnabled,
});
```

### Declaring Feature Capabilities
Feature capabilities are declared in
`x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts`:

```ts
/**
 * Interfaces for features available to the elastic assistant
 */
export type AssistantFeatures = { [K in keyof typeof assistantFeatures]: boolean };

export const assistantFeatures = Object.freeze({
  assistantModelEvaluation: false,
  assistantStreamingEnabled: false,
});
```
### Using Capabilities Client Side
And can be fetched client side using the `useCapabilities()` hook ala:

```ts
// Fetch assistant capabilities
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](#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`.

Perhaps this separation of concerns from one plugin to another isn't
necessary, but it was easy to add matching the behavior of registering
arbitrary tools. We can remove this granularity in favor of global
features if desired.


### Test Steps

* Verify `/internal/elastic_assistant/capabilities` route is called on
security solution page load in dev tools, and that by default the
`Evaluation` UI in setting does is not displayed and `404`'s if manually
called.
* Set the below experimental feature flag in your `kibana.dev.yml` and
observe the feature being enabled by inspecting the capabilities api
response, and that the evaluation feature becomes available:
```
xpack.securitySolution.enableExperimental: [ 'assistantModelEvaluation']
```
* Run the `yarn openapi:*` codegen scripts above and ensure they execute
as expected (code is generated/bundled)

### Checklist

- [x]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [X] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

---------

Co-authored-by: kibanamachine <[email protected]>

.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 defaultAssistantFeatures]: boolean };
+
+export const defaultAssistantFeatures = 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 ?? defaultAssistantFeatures;
+```
+
+### 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 defaultAssistantFeatures]: boolean };
+
+/**
+ * Default features available to the elastic assistant
+ */
+export const defaultAssistantFeatures = Object.freeze({
+  assistantModelEvaluation: false,
+  assistantStreamingEnabled: false,
+});

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

@@ -5,6 +5,9 @@
  * 2.0.
  */
 
+export { defaultAssistantFeatures } 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');
   });