[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/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.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/use_capabilities.tsx

@@ -0,0 +1,57 @@
+/*
+ * 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(
+    CAPABILITIES_QUERY_KEY,
+    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 { assistantFeatures } 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/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 ?? assistantFeatures;
+
   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`;