[Security Solution] [Elastic AI Assistant] Adds support for plugin feature registration (#174317)

## Summary

Resolves https://github.com/elastic/kibana/issues/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:
https://github.com/elastic/kibana/pull/166269
* Includes OpenAPI docs bundling script as detailed in:
https://github.com/elastic/kibana/pull/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](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`.

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 <42973632+kibanamachine@users.noreply.github.com>

.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');
   });

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

@@ -13,6 +13,7 @@ import type { IToasts } from '@kbn/core-notifications-browser';
 import { ActionTypeRegistryContract } from '@kbn/triggers-actions-ui-plugin/public';
 import { useLocalStorage } from 'react-use';
 import type { DocLinksStart } from '@kbn/core-doc-links-browser';
+import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common';
 import { WELCOME_CONVERSATION_TITLE } from '../assistant/use_conversation/translations';
 import { updatePromptContexts } from './helpers';
 import type {
@@ -37,6 +38,7 @@ import {
 } from './constants';
 import { CONVERSATIONS_TAB, SettingsTabs } from '../assistant/settings/assistant_settings';
 import { AssistantAvailability, AssistantTelemetry } from './types';
+import { useCapabilities } from '../assistant/api/capabilities/use_capabilities';
 
 export interface ShowAssistantOverlayProps {
   showOverlay: boolean;
@@ -53,7 +55,6 @@ export interface AssistantProviderProps {
   actionTypeRegistry: ActionTypeRegistryContract;
   alertsIndexPattern?: string;
   assistantAvailability: AssistantAvailability;
-  assistantStreamingEnabled?: boolean;
   assistantTelemetry?: AssistantTelemetry;
   augmentMessageCodeBlocks: (currentConversation: Conversation) => CodeBlockDetails[][];
   baseAllow: string[];
@@ -87,7 +88,6 @@ export interface AssistantProviderProps {
   }) => EuiCommentProps[];
   http: HttpSetup;
   getInitialConversations: () => Record<string, Conversation>;
-  modelEvaluatorEnabled?: boolean;
   nameSpace?: string;
   setConversations: React.Dispatch<React.SetStateAction<Record<string, Conversation>>>;
   setDefaultAllow: React.Dispatch<React.SetStateAction<string[]>>;
@@ -163,7 +163,6 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
   actionTypeRegistry,
   alertsIndexPattern,
   assistantAvailability,
-  assistantStreamingEnabled = false,
   assistantTelemetry,
   augmentMessageCodeBlocks,
   baseAllow,
@@ -179,7 +178,6 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
   getComments,
   http,
   getInitialConversations,
-  modelEvaluatorEnabled = false,
   nameSpace = DEFAULT_ASSISTANT_NAMESPACE,
   setConversations,
   setDefaultAllow,
@@ -298,6 +296,11 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
     [localStorageLastConversationId]
   );
 
+  // Fetch assistant capabilities
+  const { data: capabilities } = useCapabilities({ http, toasts });
+  const { assistantModelEvaluation: modelEvaluatorEnabled, assistantStreamingEnabled } =
+    capabilities ?? defaultAssistantFeatures;
+
   const value = useMemo(
     () => ({
       actionTypeRegistry,

x-pack/packages/security-solution/ecs_data_quality_dashboard/impl/data_quality/mock/test_providers/test_providers.tsx

@@ -13,6 +13,7 @@ import { euiDarkVars } from '@kbn/ui-theme';
 import React from 'react';
 import { ThemeProvider } from 'styled-components';
 
+import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 import { DataQualityProvider } from '../../data_quality_panel/data_quality_context';
 
 interface Props {
@@ -39,38 +40,52 @@ export const TestProvidersComponent: React.FC<Props> = ({ children, isILMAvailab
     hasConnectorsReadPrivilege: true,
     isAssistantEnabled: true,
   };
+  const queryClient = new QueryClient({
+    defaultOptions: {
+      queries: {
+        retry: false,
+      },
+    },
+    logger: {
+      log: jest.fn(),
+      warn: jest.fn(),
+      error: () => {},
+    },
+  });
 
   return (
     <I18nProvider>
       <ThemeProvider theme={() => ({ eui: euiDarkVars, darkMode: true })}>
-        <AssistantProvider
-          actionTypeRegistry={actionTypeRegistry}
-          assistantAvailability={mockAssistantAvailability}
-          augmentMessageCodeBlocks={jest.fn()}
-          baseAllow={[]}
-          baseAllowReplacement={[]}
-          basePath={'https://localhost:5601/kbn'}
-          defaultAllow={[]}
-          defaultAllowReplacement={[]}
-          docLinks={{
-            ELASTIC_WEBSITE_URL: 'https://www.elastic.co/',
-            DOC_LINK_VERSION: 'current',
-          }}
-          getComments={mockGetComments}
-          getInitialConversations={mockGetInitialConversations}
-          setConversations={jest.fn()}
-          setDefaultAllow={jest.fn()}
-          setDefaultAllowReplacement={jest.fn()}
-          http={mockHttp}
-        >
-          <DataQualityProvider
-            httpFetch={http.fetch}
-            isILMAvailable={isILMAvailable}
-            telemetryEvents={mockTelemetryEvents}
+        <QueryClientProvider client={queryClient}>
+          <AssistantProvider
+            actionTypeRegistry={actionTypeRegistry}
+            assistantAvailability={mockAssistantAvailability}
+            augmentMessageCodeBlocks={jest.fn()}
+            baseAllow={[]}
+            baseAllowReplacement={[]}
+            basePath={'https://localhost:5601/kbn'}
+            defaultAllow={[]}
+            defaultAllowReplacement={[]}
+            docLinks={{
+              ELASTIC_WEBSITE_URL: 'https://www.elastic.co/',
+              DOC_LINK_VERSION: 'current',
+            }}
+            getComments={mockGetComments}
+            getInitialConversations={mockGetInitialConversations}
+            setConversations={jest.fn()}
+            setDefaultAllow={jest.fn()}
+            setDefaultAllowReplacement={jest.fn()}
+            http={mockHttp}
           >
-            {children}
-          </DataQualityProvider>
-        </AssistantProvider>
+            <DataQualityProvider
+              httpFetch={http.fetch}
+              isILMAvailable={isILMAvailable}
+              telemetryEvents={mockTelemetryEvents}
+            >
+              {children}
+            </DataQualityProvider>
+          </AssistantProvider>
+        </QueryClientProvider>
       </ThemeProvider>
     </I18nProvider>
   );

x-pack/plugins/elastic_assistant/common/constants.ts

@@ -17,3 +17,6 @@ export const KNOWLEDGE_BASE = `${BASE_PATH}/knowledge_base/{resource?}`;
 
 // Model Evaluation
 export const EVALUATE = `${BASE_PATH}/evaluate`;
+
+// Capabilities
+export const CAPABILITIES = `${BASE_PATH}/capabilities`;

x-pack/plugins/elastic_assistant/package.json

@@ -5,6 +5,9 @@
   "private": true,
   "license": "Elastic License 2.0",
   "scripts": {
-    "evaluate-model": "node ./scripts/model_evaluator"
+    "evaluate-model": "node ./scripts/model_evaluator",
+    "openapi:generate": "node scripts/openapi/generate",
+    "openapi:generate:debug": "node --inspect-brk scripts/openapi/generate",
+    "openapi:bundle": "node scripts/openapi/bundle"
   }
 }

x-pack/plugins/elastic_assistant/scripts/openapi/bundle.js

@@ -0,0 +1,18 @@
+/*
+ * 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.
+ */
+
+require('../../../../../src/setup_node_env');
+const { bundle } = require('@kbn/openapi-bundler');
+const { resolve } = require('path');
+
+const ELASTIC_ASSISTANT_ROOT = resolve(__dirname, '../..');
+
+bundle({
+  rootDir: ELASTIC_ASSISTANT_ROOT,
+  sourceGlob: './server/schemas/**/*.schema.yaml',
+  outputFilePath: './target/openapi/elastic_assistant.bundled.schema.yaml',
+});

x-pack/plugins/elastic_assistant/scripts/openapi/generate.js

@@ -0,0 +1,18 @@
+/*
+ * 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.
+ */
+
+require('../../../../../src/setup_node_env');
+const { generate } = require('@kbn/openapi-generator');
+const { resolve } = require('path');
+
+const ELASTIC_ASSISTANT_ROOT = resolve(__dirname, '../..');
+
+generate({
+  rootDir: ELASTIC_ASSISTANT_ROOT,
+  sourceGlob: './server/schemas/**/*.schema.yaml',
+  templateName: 'zod_operation_schema',
+});

x-pack/plugins/elastic_assistant/server/__mocks__/request.ts

@@ -5,7 +5,11 @@
  * 2.0.
  */
 import { httpServerMock } from '@kbn/core/server/mocks';
-import { KNOWLEDGE_BASE } from '../../common/constants';
+import { CAPABILITIES, EVALUATE, KNOWLEDGE_BASE } from '../../common/constants';
+import {
+  PostEvaluateBodyInputs,
+  PostEvaluatePathQueryInputs,
+} from '../schemas/evaluate/post_evaluate';
 
 export const requestMock = {
   create: httpServerMock.createKibanaRequest,
@@ -31,3 +35,23 @@ export const getDeleteKnowledgeBaseRequest = (resource?: string) =>
     path: KNOWLEDGE_BASE,
     query: { resource },
   });
+
+export const getGetCapabilitiesRequest = () =>
+  requestMock.create({
+    method: 'get',
+    path: CAPABILITIES,
+  });
+
+export const getPostEvaluateRequest = ({
+  body,
+  query,
+}: {
+  body: PostEvaluateBodyInputs;
+  query: PostEvaluatePathQueryInputs;
+}) =>
+  requestMock.create({
+    body,
+    method: 'post',
+    path: EVALUATE,
+    query,
+  });

x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts

@@ -24,6 +24,7 @@ export const createMockClients = () => {
     clusterClient: core.elasticsearch.client,
     elasticAssistant: {
       actions: actionsClientMock.create(),
+      getRegisteredFeatures: jest.fn(),
       getRegisteredTools: jest.fn(),
       logger: loggingSystemMock.createLogger(),
       telemetry: coreMock.createSetup().analytics,
@@ -74,6 +75,7 @@ const createElasticAssistantRequestContextMock = (
 ): jest.Mocked<ElasticAssistantApiRequestHandlerContext> => {
   return {
     actions: clients.elasticAssistant.actions as unknown as ActionsPluginStart,
+    getRegisteredFeatures: jest.fn(),
     getRegisteredTools: jest.fn(),
     logger: clients.elasticAssistant.logger,
     telemetry: clients.elasticAssistant.telemetry,

x-pack/plugins/elastic_assistant/server/__mocks__/server.ts

@@ -5,35 +5,59 @@
  * 2.0.
  */
 import { httpServiceMock } from '@kbn/core/server/mocks';
-import type { RequestHandler, RouteConfig, KibanaRequest } from '@kbn/core/server';
-import type { RequestHandlerContext } from '@kbn/core-http-request-handler-context-server';
+import type {
+  RequestHandler,
+  RouteConfig,
+  KibanaRequest,
+  RequestHandlerContext,
+} from '@kbn/core/server';
 
 import { requestMock } from './request';
 import { responseMock as responseFactoryMock } from './response';
 import { requestContextMock } from './request_context';
 import { responseAdapter } from './test_adapters';
+import type { RegisteredVersionedRoute } from '@kbn/core-http-router-server-mocks';
 
 interface Route {
-  config: RouteConfig<unknown, unknown, unknown, 'get' | 'post' | 'delete' | 'patch' | 'put'>;
+  validate: RouteConfig<
+    unknown,
+    unknown,
+    unknown,
+    'get' | 'post' | 'delete' | 'patch' | 'put'
+  >['validate'];
   handler: RequestHandler;
 }
 
-const getRoute = (routerMock: MockServer['router']): Route => {
-  const routeCalls = [
-    ...routerMock.get.mock.calls,
-    ...routerMock.post.mock.calls,
-    ...routerMock.put.mock.calls,
-    ...routerMock.patch.mock.calls,
-    ...routerMock.delete.mock.calls,
-  ];
+const HTTP_METHODS = ['get', 'post', 'put', 'patch', 'delete'] as const;
 
-  const [route] = routeCalls;
-  if (!route) {
-    throw new Error('No route registered!');
+const getClassicRoute = (routerMock: MockServer['router']): Route | undefined => {
+  const method = HTTP_METHODS.find((m) => routerMock[m].mock.calls.length > 0);
+  if (!method) {
+    return undefined;
   }
 
-  const [config, handler] = route;
-  return { config, handler };
+  const [config, handler] = routerMock[method].mock.calls[0];
+  return { validate: config.validate, handler };
+};
+
+const getVersionedRoute = (router: MockServer['router']): Route => {
+  const method = HTTP_METHODS.find((m) => router.versioned[m].mock.calls.length > 0);
+  if (!method) {
+    throw new Error('No route registered!');
+  }
+  const config = router.versioned[method].mock.calls[0][0];
+  const routePath = config.path;
+
+  const route: RegisteredVersionedRoute = router.versioned.getRoute(method, routePath);
+  const firstVersion = Object.values(route.versions)[0];
+
+  return {
+    validate:
+      firstVersion.config.validate === false
+        ? false
+        : firstVersion.config.validate.request || false,
+    handler: firstVersion.handler,
+  };
 };
 
 const buildResultMock = () => ({ ok: jest.fn((x) => x), badRequest: jest.fn((x) => x) });
@@ -63,7 +87,7 @@ class MockServer {
   }
 
   private getRoute(): Route {
-    return getRoute(this.router);
+    return getClassicRoute(this.router) ?? getVersionedRoute(this.router);
   }
 
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -72,7 +96,7 @@ class MockServer {
   }
 
   private validateRequest(request: KibanaRequest): KibanaRequest {
-    const validations = this.getRoute().config.validate;
+    const validations = this.getRoute().validate;
     if (!validations) {
       return request;
     }
@@ -88,6 +112,7 @@ class MockServer {
     return validatedRequest;
   }
 }
+
 const createMockServer = () => new MockServer();
 
 export const serverMock = {

x-pack/plugins/elastic_assistant/server/__mocks__/test_adapters.ts

@@ -43,6 +43,8 @@ const buildResponses = (method: Method, calls: MockCall[]): ResponseCall[] => {
         status: call.statusCode,
         body: call.body,
       }));
+    case 'notFound':
+      return calls.map(() => ({ status: 404, body: undefined }));
     default:
       throw new Error(`Encountered unexpected call to response.${method}`);
   }

x-pack/plugins/elastic_assistant/server/plugin.ts

@@ -18,6 +18,7 @@ import {
 } from '@kbn/core/server';
 import { once } from 'lodash';
 
+import { AssistantFeatures } from '@kbn/elastic-assistant-common';
 import { events } from './lib/telemetry/event_based_telemetry';
 import {
   AssistantTool,
@@ -36,11 +37,17 @@ import {
   postEvaluateRoute,
   postKnowledgeBaseRoute,
 } from './routes';
-import { appContextService, GetRegisteredTools } from './services/app_context';
+import {
+  appContextService,
+  GetRegisteredFeatures,
+  GetRegisteredTools,
+} from './services/app_context';
+import { getCapabilitiesRoute } from './routes/capabilities/get_capabilities_route';
 
 interface CreateRouteHandlerContextParams {
   core: CoreSetup<ElasticAssistantPluginStart, unknown>;
   logger: Logger;
+  getRegisteredFeatures: GetRegisteredFeatures;
   getRegisteredTools: GetRegisteredTools;
   telemetry: AnalyticsServiceSetup;
 }
@@ -63,6 +70,7 @@ export class ElasticAssistantPlugin
   private createRouteHandlerContext = ({
     core,
     logger,
+    getRegisteredFeatures,
     getRegisteredTools,
     telemetry,
   }: CreateRouteHandlerContextParams): IContextProvider<
@@ -74,6 +82,7 @@ export class ElasticAssistantPlugin
 
       return {
         actions: pluginsStart.actions,
+        getRegisteredFeatures,
         getRegisteredTools,
         logger,
         telemetry,
@@ -89,6 +98,9 @@ export class ElasticAssistantPlugin
       this.createRouteHandlerContext({
         core: core as CoreSetup<ElasticAssistantPluginStart, unknown>,
         logger: this.logger,
+        getRegisteredFeatures: (pluginName: string) => {
+          return appContextService.getRegisteredFeatures(pluginName);
+        },
         getRegisteredTools: (pluginName: string) => {
           return appContextService.getRegisteredTools(pluginName);
         },
@@ -112,40 +124,37 @@ export class ElasticAssistantPlugin
     postActionsConnectorExecuteRoute(router, getElserId);
     // Evaluate
     postEvaluateRoute(router, getElserId);
+    // Capabilities
+    getCapabilitiesRoute(router);
     return {
       actions: plugins.actions,
+      getRegisteredFeatures: (pluginName: string) => {
+        return appContextService.getRegisteredFeatures(pluginName);
+      },
       getRegisteredTools: (pluginName: string) => {
         return appContextService.getRegisteredTools(pluginName);
       },
     };
   }
 
-  public start(core: CoreStart, plugins: ElasticAssistantPluginStartDependencies) {
+  public start(
+    core: CoreStart,
+    plugins: ElasticAssistantPluginStartDependencies
+  ): ElasticAssistantPluginStart {
     this.logger.debug('elasticAssistant: Started');
     appContextService.start({ logger: this.logger });
 
     return {
-      /**
-       * Actions plugin start contract
-       */
       actions: plugins.actions,
-
-      /**
-       * Get the registered tools for a given plugin name.
-       * @param pluginName
-       */
+      getRegisteredFeatures: (pluginName: string) => {
+        return appContextService.getRegisteredFeatures(pluginName);
+      },
       getRegisteredTools: (pluginName: string) => {
         return appContextService.getRegisteredTools(pluginName);
       },
-
-      /**
-       * Register tools to be used by the Elastic Assistant for a given plugin. Use the plugin name that
-       * corresponds to your application as defined in the `x-kbn-context` header of requests made from your
-       * application.
-       *
-       * @param pluginName
-       * @param tools
-       */
+      registerFeatures: (pluginName: string, features: Partial<AssistantFeatures>) => {
+        return appContextService.registerFeatures(pluginName, features);
+      },
       registerTools: (pluginName: string, tools: AssistantTool[]) => {
         return appContextService.registerTools(pluginName, tools);
       },

x-pack/plugins/elastic_assistant/server/routes/capabilities/get_capabilities_route.test.ts

@@ -0,0 +1,47 @@
+/*
+ * 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 { getCapabilitiesRoute } from './get_capabilities_route';
+import { serverMock } from '../../__mocks__/server';
+import { requestContextMock } from '../../__mocks__/request_context';
+import { getGetCapabilitiesRequest } from '../../__mocks__/request';
+import { getPluginNameFromRequest } from '../helpers';
+
+jest.mock('../helpers');
+
+describe('Get Capabilities Route', () => {
+  let server: ReturnType<typeof serverMock.create>;
+  let { context } = requestContextMock.createTools();
+
+  beforeEach(() => {
+    server = serverMock.create();
+    ({ context } = requestContextMock.createTools());
+
+    getCapabilitiesRoute(server.router);
+  });
+
+  describe('Status codes', () => {
+    it('returns 200 with capabilities', async () => {
+      const response = await server.inject(
+        getGetCapabilitiesRequest(),
+        requestContextMock.convertContext(context)
+      );
+      expect(response.status).toEqual(200);
+    });
+
+    it('returns 500 if an error is thrown in fetching capabilities', async () => {
+      (getPluginNameFromRequest as jest.Mock).mockImplementation(() => {
+        throw new Error('Mocked error');
+      });
+      const response = await server.inject(
+        getGetCapabilitiesRequest(),
+        requestContextMock.convertContext(context)
+      );
+      expect(response.status).toEqual(500);
+    });
+  });
+});

x-pack/plugins/elastic_assistant/server/routes/capabilities/get_capabilities_route.ts

@@ -0,0 +1,60 @@
+/*
+ * 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 { IKibanaResponse, IRouter } from '@kbn/core/server';
+import { transformError } from '@kbn/securitysolution-es-utils';
+
+import { CAPABILITIES } from '../../../common/constants';
+import { ElasticAssistantRequestHandlerContext } from '../../types';
+
+import { GetCapabilitiesResponse } from '../../schemas/capabilities/get_capabilities_route.gen';
+import { buildResponse } from '../../lib/build_response';
+import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers';
+
+/**
+ * Get the assistant capabilities for the requesting plugin
+ *
+ * @param router IRouter for registering routes
+ */
+export const getCapabilitiesRoute = (router: IRouter<ElasticAssistantRequestHandlerContext>) => {
+  router.versioned
+    .get({
+      access: 'internal',
+      path: CAPABILITIES,
+      options: {
+        tags: ['access:elasticAssistant'],
+      },
+    })
+    .addVersion(
+      {
+        version: '1',
+        validate: {},
+      },
+      async (context, request, response): Promise<IKibanaResponse<GetCapabilitiesResponse>> => {
+        const resp = buildResponse(response);
+        const assistantContext = await context.elasticAssistant;
+        const logger = assistantContext.logger;
+
+        try {
+          const pluginName = getPluginNameFromRequest({
+            request,
+            defaultPluginName: DEFAULT_PLUGIN_NAME,
+            logger,
+          });
+          const registeredFeatures = assistantContext.getRegisteredFeatures(pluginName);
+
+          return response.ok({ body: registeredFeatures });
+        } catch (err) {
+          const error = transformError(err);
+          return resp.error({
+            body: error.message,
+            statusCode: error.statusCode,
+          });
+        }
+      }
+    );
+};

x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.test.ts

@@ -0,0 +1,59 @@
+/*
+ * 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 { postEvaluateRoute } from './post_evaluate';
+import { serverMock } from '../../__mocks__/server';
+import { requestContextMock } from '../../__mocks__/request_context';
+import { getPostEvaluateRequest } from '../../__mocks__/request';
+import {
+  PostEvaluateBodyInputs,
+  PostEvaluatePathQueryInputs,
+} from '../../schemas/evaluate/post_evaluate';
+
+const defaultBody: PostEvaluateBodyInputs = {
+  dataset: undefined,
+  evalPrompt: undefined,
+};
+
+const defaultQueryParams: PostEvaluatePathQueryInputs = {
+  agents: 'agents',
+  datasetName: undefined,
+  evaluationType: undefined,
+  evalModel: undefined,
+  models: 'models',
+  outputIndex: '.kibana-elastic-ai-assistant-',
+  projectName: undefined,
+  runName: undefined,
+};
+
+describe('Post Evaluate Route', () => {
+  let server: ReturnType<typeof serverMock.create>;
+  let { context } = requestContextMock.createTools();
+  const mockGetElser = jest.fn().mockResolvedValue('.elser_model_2');
+
+  beforeEach(() => {
+    server = serverMock.create();
+    ({ context } = requestContextMock.createTools());
+
+    postEvaluateRoute(server.router, mockGetElser);
+  });
+
+  describe('Capabilities', () => {
+    it('returns a 404 if evaluate feature is not registered', async () => {
+      context.elasticAssistant.getRegisteredFeatures.mockReturnValueOnce({
+        assistantModelEvaluation: false,
+        assistantStreamingEnabled: false,
+      });
+
+      const response = await server.inject(
+        getPostEvaluateRequest({ body: defaultBody, query: defaultQueryParams }),
+        requestContextMock.convertContext(context)
+      );
+      expect(response.status).toEqual(404);
+    });
+  });
+});

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

@@ -29,6 +29,7 @@ import {
 } from '../../lib/model_evaluator/output_index/utils';
 import { fetchLangSmithDataset, getConnectorName, getLangSmithTracer, getLlmType } from './utils';
 import { RequestBody } from '../../lib/langchain/types';
+import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers';
 
 /**
  * To support additional Agent Executors from the UI, add them to this map
@@ -53,11 +54,22 @@ export const postEvaluateRoute = (
         query: buildRouteValidation(PostEvaluatePathQuery),
       },
     },
-    // TODO: Limit route based on experimental feature
     async (context, request, response) => {
       const assistantContext = await context.elasticAssistant;
       const logger = assistantContext.logger;
       const telemetry = assistantContext.telemetry;
+
+      // Validate evaluation feature is enabled
+      const pluginName = getPluginNameFromRequest({
+        request,
+        defaultPluginName: DEFAULT_PLUGIN_NAME,
+        logger,
+      });
+      const registeredFeatures = assistantContext.getRegisteredFeatures(pluginName);
+      if (!registeredFeatures.assistantModelEvaluation) {
+        return response.notFound();
+      }
+
       try {
         const evaluationId = uuidv4();
         const {

x-pack/plugins/elastic_assistant/server/routes/helpers.ts

@@ -7,10 +7,9 @@
 
 import { KibanaRequest } from '@kbn/core-http-server';
 import { Logger } from '@kbn/core/server';
-import { RequestBody } from '../lib/langchain/types';
 
 interface GetPluginNameFromRequestParams {
-  request: KibanaRequest<unknown, unknown, RequestBody>;
+  request: KibanaRequest;
   defaultPluginName: string;
   logger?: Logger;
 }

x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts

@@ -22,6 +22,7 @@ import { ESQL_DOCS_LOADED_QUERY, ESQL_RESOURCE, KNOWLEDGE_BASE_INDEX_PATTERN } f
  * Get the status of the Knowledge Base index, pipeline, and resources (collection of documents)
  *
  * @param router IRouter for registering routes
+ * @param getElser Function to get the default Elser ID
  */
 export const getKnowledgeBaseStatusRoute = (
   router: IRouter<ElasticAssistantRequestHandlerContext>,

x-pack/plugins/elastic_assistant/server/schemas/capabilities/get_capabilities_route.gen.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.
+ */
+
+import { z } from 'zod';
+
+/*
+ * NOTICE: Do not edit this file manually.
+ * This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator.
+ */
+
+export type GetCapabilitiesResponse = z.infer<typeof GetCapabilitiesResponse>;
+export const GetCapabilitiesResponse = z.object({
+  assistantModelEvaluation: z.boolean(),
+  assistantStreamingEnabled: z.boolean(),
+});

x-pack/plugins/elastic_assistant/server/schemas/capabilities/get_capabilities_route.schema.yaml

@@ -0,0 +1,41 @@
+openapi: 3.0.0
+info:
+  title: Get Capabilities API endpoint
+  version: '1'
+paths:
+  /internal/elastic_assistant/capabilities:
+    get:
+      operationId: GetCapabilities
+      x-codegen-enabled: true
+      description: Get Elastic Assistant capabilities for the requesting plugin
+      summary: Get Elastic Assistant capabilities
+      tags:
+        - Capabilities API
+      responses:
+        '200':
+          description: Successful response
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  assistantModelEvaluation:
+                    type: boolean
+                  assistantStreamingEnabled:
+                    type: boolean
+                required:
+                  - assistantModelEvaluation
+                  - assistantStreamingEnabled
+        '400':
+          description: Generic Error
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  statusCode:
+                    type: number
+                  error:
+                    type: string
+                  message:
+                    type: string

x-pack/plugins/elastic_assistant/server/schemas/evaluate/post_evaluate.ts

@@ -35,6 +35,8 @@ export const PostEvaluatePathQuery = t.type({
   runName: t.union([t.string, t.undefined]),
 });
 
+export type PostEvaluatePathQueryInputs = t.TypeOf<typeof PostEvaluatePathQuery>;
+
 export type DatasetItem = t.TypeOf<typeof DatasetItem>;
 export const DatasetItem = t.type({
   id: t.union([t.string, t.undefined]),

x-pack/plugins/elastic_assistant/server/services/app_context.test.ts

@@ -8,6 +8,7 @@
 import { appContextService, ElasticAssistantAppContext } from './app_context';
 import { loggerMock } from '@kbn/logging-mocks';
 import { AssistantTool } from '../types';
+import { AssistantFeatures, defaultAssistantFeatures } from '@kbn/elastic-assistant-common';
 
 // Mock Logger
 const mockLogger = loggerMock.create();
@@ -48,6 +49,19 @@ describe('AppContextService', () => {
 
       expect(appContextService.getRegisteredTools('super').length).toBe(0);
     });
+
+    it('should return default registered features when stopped ', () => {
+      appContextService.start(mockAppContext);
+      appContextService.registerFeatures('super', {
+        assistantModelEvaluation: true,
+        assistantStreamingEnabled: true,
+      });
+      appContextService.stop();
+
+      expect(appContextService.getRegisteredFeatures('super')).toEqual(
+        expect.objectContaining(defaultAssistantFeatures)
+      );
+    });
   });
 
   describe('registering tools', () => {
@@ -84,4 +98,81 @@ describe('AppContextService', () => {
       expect(appContextService.getRegisteredTools(pluginName).length).toEqual(1);
     });
   });
+
+  describe('registering features', () => {
+    it('should register and get features for a single plugin', () => {
+      const pluginName = 'pluginName';
+      const features: AssistantFeatures = {
+        assistantModelEvaluation: true,
+        assistantStreamingEnabled: true,
+      };
+
+      appContextService.start(mockAppContext);
+      appContextService.registerFeatures(pluginName, features);
+
+      // Check if getRegisteredFeatures returns the correct tools
+      const retrievedFeatures = appContextService.getRegisteredFeatures(pluginName);
+      expect(retrievedFeatures).toEqual(features);
+    });
+
+    it('should register and get features for multiple plugins', () => {
+      const pluginOne = 'plugin1';
+      const featuresOne: AssistantFeatures = {
+        assistantModelEvaluation: true,
+        assistantStreamingEnabled: false,
+      };
+      const pluginTwo = 'plugin2';
+      const featuresTwo: AssistantFeatures = {
+        assistantModelEvaluation: false,
+        assistantStreamingEnabled: true,
+      };
+
+      appContextService.start(mockAppContext);
+      appContextService.registerFeatures(pluginOne, featuresOne);
+      appContextService.registerFeatures(pluginTwo, featuresTwo);
+
+      expect(appContextService.getRegisteredFeatures(pluginOne)).toEqual(featuresOne);
+      expect(appContextService.getRegisteredFeatures(pluginTwo)).toEqual(featuresTwo);
+    });
+
+    it('should update features if registered again', () => {
+      const pluginName = 'pluginName';
+      const featuresOne: AssistantFeatures = {
+        assistantModelEvaluation: true,
+        assistantStreamingEnabled: false,
+      };
+      const featuresTwo: AssistantFeatures = {
+        assistantModelEvaluation: false,
+        assistantStreamingEnabled: true,
+      };
+
+      appContextService.start(mockAppContext);
+      appContextService.registerFeatures(pluginName, featuresOne);
+      appContextService.registerFeatures(pluginName, featuresTwo);
+
+      expect(appContextService.getRegisteredFeatures(pluginName)).toEqual(featuresTwo);
+    });
+
+    it('should return default features if pluginName not present', () => {
+      appContextService.start(mockAppContext);
+
+      expect(appContextService.getRegisteredFeatures('super')).toEqual(
+        expect.objectContaining(defaultAssistantFeatures)
+      );
+    });
+
+    it('allows registering a subset of all available features', () => {
+      const pluginName = 'pluginName';
+      const featuresSubset: Partial<AssistantFeatures> = {
+        assistantModelEvaluation: true,
+      };
+
+      appContextService.start(mockAppContext);
+      appContextService.registerFeatures(pluginName, featuresSubset);
+
+      expect(appContextService.getRegisteredFeatures(pluginName)).toEqual(
+        expect.objectContaining({ ...defaultAssistantFeatures, ...featuresSubset })
+      );
+    });
+  });
 });

x-pack/plugins/elastic_assistant/server/services/app_context.ts

@@ -6,11 +6,14 @@
  */
 
 import type { Logger } from '@kbn/core/server';
-import type { AssistantTool } from '../types';
+import { defaultAssistantFeatures, AssistantFeatures } from '@kbn/elastic-assistant-common';
+import { AssistantTool } from '../types';
 
 export type PluginName = string;
 export type RegisteredToolsStorage = Map<PluginName, Set<AssistantTool>>;
+export type RegisteredFeaturesStorage = Map<PluginName, AssistantFeatures>;
 export type GetRegisteredTools = (pluginName: string) => AssistantTool[];
+export type GetRegisteredFeatures = (pluginName: string) => AssistantFeatures;
 export interface ElasticAssistantAppContext {
   logger: Logger;
 }
@@ -23,6 +26,7 @@ export interface ElasticAssistantAppContext {
 class AppContextService {
   private logger: Logger | undefined;
   private registeredTools: RegisteredToolsStorage = new Map<PluginName, Set<AssistantTool>>();
+  private registeredFeatures: RegisteredFeaturesStorage = new Map<PluginName, AssistantFeatures>();
 
   public start(appContext: ElasticAssistantAppContext) {
     this.logger = appContext.logger;
@@ -30,6 +34,7 @@ class AppContextService {
 
   public stop() {
     this.registeredTools.clear();
+    this.registeredFeatures.clear();
   }
 
   /**
@@ -44,7 +49,7 @@ class AppContextService {
     this.logger?.debug(`tools: ${tools.map((tool) => tool.name).join(', ')}`);
 
     if (!this.registeredTools.has(pluginName)) {
-      this.logger?.debug('plugin has no tools, making new set');
+      this.logger?.debug('plugin has no tools, initializing...');
       this.registeredTools.set(pluginName, new Set<AssistantTool>());
     }
     tools.forEach((tool) => this.registeredTools.get(pluginName)?.add(tool));
@@ -64,6 +69,51 @@ class AppContextService {
 
     return tools;
   }
+
+  /**
+   * Register features to be used by the Elastic Assistant
+   *
+   * @param pluginName
+   * @param features
+   */
+  public registerFeatures(pluginName: string, features: Partial<AssistantFeatures>) {
+    this.logger?.debug('AppContextService:registerFeatures');
+    this.logger?.debug(`pluginName: ${pluginName}`);
+    this.logger?.debug(
+      `features: ${Object.entries(features)
+        .map(([feature, enabled]) => `${feature}:${enabled}`)
+        .join(', ')}`
+    );
+
+    if (!this.registeredFeatures.has(pluginName)) {
+      this.logger?.debug('plugin has no features, initializing...');
+      this.registeredFeatures.set(pluginName, defaultAssistantFeatures);
+    }
+
+    const registeredFeatures = this.registeredFeatures.get(pluginName);
+    if (registeredFeatures != null) {
+      this.registeredFeatures.set(pluginName, { ...registeredFeatures, ...features });
+    }
+  }
+
+  /**
+   * Get the registered features
+   *
+   * @param pluginName
+   */
+  public getRegisteredFeatures(pluginName: string): AssistantFeatures {
+    const features = this.registeredFeatures?.get(pluginName) ?? defaultAssistantFeatures;
+
+    this.logger?.debug('AppContextService:getRegisteredFeatures');
+    this.logger?.debug(`pluginName: ${pluginName}`);
+    this.logger?.debug(
+      `features: ${Object.entries(features)
+        .map(([feature, enabled]) => `${feature}:${enabled}`)
+        .join(', ')}`
+    );
+
+    return features;
+  }
 }
 
 export const appContextService = new AppContextService();

x-pack/plugins/elastic_assistant/server/types.ts

@@ -20,8 +20,9 @@ import { type MlPluginSetup } from '@kbn/ml-plugin/server';
 import { Tool } from 'langchain/dist/tools/base';
 import { RetrievalQAChain } from 'langchain/chains';
 import { ElasticsearchClient } from '@kbn/core/server';
+import { AssistantFeatures } from '@kbn/elastic-assistant-common';
 import { RequestBody } from './lib/langchain/types';
-import type { GetRegisteredTools } from './services/app_context';
+import type { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context';
 
 export const PLUGIN_ID = 'elasticAssistant' as const;
 
@@ -32,15 +33,37 @@ export interface ElasticAssistantPluginSetup {
 
 /** The plugin start interface */
 export interface ElasticAssistantPluginStart {
+  /**
+   * Actions plugin start contract.
+   */
   actions: ActionsPluginStart;
   /**
-   * Register tools to be used by the elastic assistant
+   * Register features to be used by the elastic assistant.
+   *
+   * Note: Be sure to use the pluginName that is sent in the request headers by your plugin to ensure it is extracted
+   * and the correct features are available. See {@link getPluginNameFromRequest} for more details.
+   *
+   * @param pluginName Name of the plugin the features should be registered to
+   * @param features Partial<AssistantFeatures> to be registered with for the given plugin
+   */
+  registerFeatures: (pluginName: string, features: Partial<AssistantFeatures>) => void;
+  /**
+   * Get the registered features for a given plugin name.
+   * @param pluginName Name of the plugin to get the features for
+   */
+  getRegisteredFeatures: GetRegisteredFeatures;
+  /**
+   * Register tools to be used by the elastic assistant.
+   *
+   * Note: Be sure to use the pluginName that is sent in the request headers by your plugin to ensure it is extracted
+   * and the correct tools are selected. See {@link getPluginNameFromRequest} for more details.
+   *
    * @param pluginName Name of the plugin the tool should be registered to
    * @param tools AssistantTools to be registered with for the given plugin
    */
   registerTools: (pluginName: string, tools: AssistantTool[]) => void;
   /**
-   * Get the registered tools
+   * Get the registered tools for a given plugin name.
    * @param pluginName Name of the plugin to get the tools for
    */
   getRegisteredTools: GetRegisteredTools;
@@ -56,6 +79,7 @@ export interface ElasticAssistantPluginStartDependencies {
 
 export interface ElasticAssistantApiRequestHandlerContext {
   actions: ActionsPluginStart;
+  getRegisteredFeatures: GetRegisteredFeatures;
   getRegisteredTools: GetRegisteredTools;
   logger: Logger;
   telemetry: AnalyticsServiceSetup;

x-pack/plugins/elastic_assistant/tsconfig.json

@@ -17,7 +17,6 @@
     "@kbn/core",
     "@kbn/core-http-server",
     "@kbn/licensing-plugin",
-    "@kbn/core-http-request-handler-context-server",
     "@kbn/securitysolution-es-utils",
     "@kbn/securitysolution-io-ts-utils",
     "@kbn/actions-plugin",
@@ -34,6 +33,8 @@
     "@kbn/ml-plugin",
     "@kbn/apm-utils",
     "@kbn/core-analytics-server",
+    "@kbn/elastic-assistant-common",
+    "@kbn/core-http-router-server-mocks",
   ],
   "exclude": [
     "target/**/*",

x-pack/plugins/security_solution/public/assistant/provider.tsx

@@ -22,7 +22,6 @@ import { useAnonymizationStore } from './use_anonymization_store';
 import { useAssistantAvailability } from './use_assistant_availability';
 import { APP_ID } from '../../common/constants';
 import { useAppToasts } from '../common/hooks/use_app_toasts';
-import { useIsExperimentalFeatureEnabled } from '../common/hooks/use_experimental_features';
 import { useSignalIndex } from '../detections/containers/detection_engine/alerts/use_signal_index';
 
 const ASSISTANT_TITLE = i18n.translate('xpack.securitySolution.assistant.title', {
@@ -39,8 +38,6 @@ export const AssistantProvider: React.FC = ({ children }) => {
     docLinks: { ELASTIC_WEBSITE_URL, DOC_LINK_VERSION },
   } = useKibana().services;
   const basePath = useBasePath();
-  const isModelEvaluationEnabled = useIsExperimentalFeatureEnabled('assistantModelEvaluation');
-  const assistantStreamingEnabled = useIsExperimentalFeatureEnabled('assistantStreamingEnabled');
 
   const { conversations, setConversations } = useConversationStore();
   const getInitialConversation = useCallback(() => {
@@ -78,8 +75,6 @@ export const AssistantProvider: React.FC = ({ children }) => {
       getInitialConversations={getInitialConversation}
       getComments={getComments}
       http={http}
-      assistantStreamingEnabled={assistantStreamingEnabled}
-      modelEvaluatorEnabled={isModelEvaluationEnabled}
       nameSpace={nameSpace}
       setConversations={setConversations}
       setDefaultAllow={setDefaultAllow}

x-pack/plugins/security_solution/public/detections/components/rules/rule_execution_status/rule_status_failed_callout.test.tsx

@@ -15,6 +15,7 @@ import { AssistantProvider } from '@kbn/elastic-assistant';
 import type { AssistantAvailability } from '@kbn/elastic-assistant';
 import { httpServiceMock } from '@kbn/core-http-browser-mocks';
 import { actionTypeRegistryMock } from '@kbn/triggers-actions-ui-plugin/public/application/action_type_registry.mock';
+import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 
 jest.mock('../../../../common/lib/kibana');
 
@@ -32,29 +33,44 @@ const mockAssistantAvailability: AssistantAvailability = {
   hasConnectorsReadPrivilege: true,
   isAssistantEnabled: true,
 };
+const queryClient = new QueryClient({
+  defaultOptions: {
+    queries: {
+      retry: false,
+    },
+  },
+  logger: {
+    log: jest.fn(),
+    warn: jest.fn(),
+    error: () => {},
+  },
+});
+
 const ContextWrapper: React.FC = ({ children }) => (
-  <AssistantProvider
-    actionTypeRegistry={actionTypeRegistry}
-    assistantAvailability={mockAssistantAvailability}
-    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>
+  <QueryClientProvider client={queryClient}>
+    <AssistantProvider
+      actionTypeRegistry={actionTypeRegistry}
+      assistantAvailability={mockAssistantAvailability}
+      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>
+  </QueryClientProvider>
 );
 
 describe('RuleStatusFailedCallOut', () => {

x-pack/plugins/security_solution/server/plugin.ts

@@ -514,6 +514,10 @@ export class Plugin implements ISecuritySolutionPlugin {
 
     // Assistant Tool and Feature Registration
     plugins.elasticAssistant.registerTools(APP_UI_ID, getAssistantTools());
+    plugins.elasticAssistant.registerFeatures(APP_UI_ID, {
+      assistantModelEvaluation: config.experimentalFeatures.assistantModelEvaluation,
+      assistantStreamingEnabled: config.experimentalFeatures.assistantStreamingEnabled,
+    });
 
     if (this.lists && plugins.taskManager && plugins.fleet) {
       // Exceptions, Artifacts and Manifests start