[Security solution] Implement dashboard to track Gen AI Token Usage

x-pack/plugins/actions/server/sub_action_framework/sub_action_connector.ts

@@ -9,6 +9,8 @@ import { isPlainObject, isEmpty } from 'lodash';
 import { Type } from '@kbn/config-schema';
 import { Logger } from '@kbn/logging';
 import axios, { AxiosInstance, AxiosResponse, AxiosError, AxiosRequestHeaders } from 'axios';
+import { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server';
+import { ElasticsearchClient } from '@kbn/core-elasticsearch-server';
 import { assertURL } from './helpers/validators';
 import { ActionsConfigurationUtilities } from '../actions_config';
 import { SubAction, SubActionRequestParams } from './types';
@@ -28,6 +30,8 @@ export abstract class SubActionConnector<Config, Secrets> {
   private subActions: Map<string, SubAction> = new Map();
   private configurationUtilities: ActionsConfigurationUtilities;
   protected logger: Logger;
+  protected esClient: ElasticsearchClient;
+  protected savedObjectsClient: SavedObjectsClientContract;
   protected connector: ServiceParams<Config, Secrets>['connector'];
   protected config: Config;
   protected secrets: Secrets;
@@ -37,6 +41,8 @@ export abstract class SubActionConnector<Config, Secrets> {
     this.logger = params.logger;
     this.config = params.config;
     this.secrets = params.secrets;
+    this.savedObjectsClient = params.services.savedObjectsClient;
+    this.esClient = params.services.scopedClusterClient;
     this.configurationUtilities = params.configurationUtilities;
     this.axiosInstance = axios.create();
   }

x-pack/plugins/actions/tsconfig.json

@@ -36,6 +36,8 @@
     "@kbn/core-http-server-mocks",
     "@kbn/tinymath",
     "@kbn/core-saved-objects-utils-server",
+    "@kbn/core-saved-objects-api-server",
+    "@kbn/core-elasticsearch-server",
   ],
   "exclude": [
     "target/**/*",

x-pack/plugins/stack_connectors/common/gen_ai/constants.ts

@@ -16,6 +16,7 @@ export const GEN_AI_TITLE = i18n.translate(
 export const GEN_AI_CONNECTOR_ID = '.gen-ai';
 export enum SUB_ACTION {
   RUN = 'run',
+  DASHBOARD = 'getDashboard',
   TEST = 'test',
 }
 export enum OpenAiProviderType {

x-pack/plugins/stack_connectors/common/gen_ai/schema.ts

@@ -44,3 +44,12 @@ export const GenAiRunActionResponseSchema = schema.object(
   },
   { unknowns: 'ignore' }
 );
+
+// Run action schema
+export const GenAiDashboardActionParamsSchema = schema.object({
+  dashboardId: schema.string(),
+});
+
+export const GenAiDashboardActionResponseSchema = schema.object({
+  available: schema.boolean(),
+});

x-pack/plugins/stack_connectors/common/gen_ai/types.ts

@@ -11,9 +11,13 @@ import {
   GenAiSecretsSchema,
   GenAiRunActionParamsSchema,
   GenAiRunActionResponseSchema,
+  GenAiDashboardActionParamsSchema,
+  GenAiDashboardActionResponseSchema,
 } from './schema';
 
 export type GenAiConfig = TypeOf<typeof GenAiConfigSchema>;
 export type GenAiSecrets = TypeOf<typeof GenAiSecretsSchema>;
 export type GenAiRunActionParams = TypeOf<typeof GenAiRunActionParamsSchema>;
 export type GenAiRunActionResponse = TypeOf<typeof GenAiRunActionResponseSchema>;
+export type GenAiDashboardActionParams = TypeOf<typeof GenAiDashboardActionParamsSchema>;
+export type GenAiDashboardActionResponse = TypeOf<typeof GenAiDashboardActionResponseSchema>;

x-pack/plugins/stack_connectors/public/connector_types/gen_ai/api.test.ts

@@ -0,0 +1,37 @@
+/*
+ * 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 { httpServiceMock } from '@kbn/core-http-browser-mocks';
+import { getDashboard } from './api';
+import { SUB_ACTION } from '../../../common/gen_ai/constants';
+const response = {
+  available: true,
+};
+
+describe('Gen AI Dashboard API', () => {
+  const http = httpServiceMock.createStartContract();
+
+  beforeEach(() => jest.resetAllMocks());
+  describe('getDashboard', () => {
+    test('should call get dashboard API', async () => {
+      const abortCtrl = new AbortController();
+      http.post.mockResolvedValueOnce(response);
+      const res = await getDashboard({
+        http,
+        signal: abortCtrl.signal,
+        connectorId: 'te/st',
+        dashboardId: 'cool-dashboard',
+      });
+
+      expect(res).toEqual(response);
+      expect(http.post).toHaveBeenCalledWith('/api/actions/connector/te%2Fst/_execute', {
+        body: `{"params":{"subAction":"${SUB_ACTION.DASHBOARD}","subActionParams":{"dashboardId":"cool-dashboard"}}}`,
+        signal: abortCtrl.signal,
+      });
+    });
+  });
+});

x-pack/plugins/stack_connectors/public/connector_types/gen_ai/api.ts

@@ -0,0 +1,34 @@
+/*
+ * 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 { ActionTypeExecutorResult, BASE_ACTION_API_PATH } from '@kbn/actions-plugin/common';
+import { SUB_ACTION } from '../../../common/gen_ai/constants';
+import { ConnectorExecutorResult, rewriteResponseToCamelCase } from '../lib/rewrite_response_body';
+
+export async function getDashboard({
+  http,
+  signal,
+  dashboardId,
+  connectorId,
+}: {
+  http: HttpSetup;
+  signal: AbortSignal;
+  connectorId: string;
+  dashboardId: string;
+}): Promise<ActionTypeExecutorResult<{ available: boolean }>> {
+  const res = await http.post<ConnectorExecutorResult<{ available: boolean }>>(
+    `${BASE_ACTION_API_PATH}/connector/${encodeURIComponent(connectorId)}/_execute`,
+    {
+      body: JSON.stringify({
+        params: { subAction: SUB_ACTION.DASHBOARD, subActionParams: { dashboardId } },
+      }),
+      signal,
+    }
+  );
+  return rewriteResponseToCamelCase(res);
+}

x-pack/plugins/stack_connectors/public/connector_types/gen_ai/connector.test.tsx

@@ -8,27 +8,54 @@
 import React from 'react';
 import GenerativeAiConnectorFields from './connector';
 import { ConnectorFormTestProvider } from '../lib/test_utils';
-import { act, render, waitFor } from '@testing-library/react';
+import { act, fireEvent, render, waitFor } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import { OpenAiProviderType } from '../../../common/gen_ai/constants';
+import { useKibana } from '@kbn/triggers-actions-ui-plugin/public';
+import { useGetDashboard } from './use_get_dashboard';
+
+jest.mock('@kbn/triggers-actions-ui-plugin/public/common/lib/kibana');
+jest.mock('./use_get_dashboard');
+
+const useKibanaMock = useKibana as jest.Mocked<typeof useKibana>;
+const mockDashboard = useGetDashboard as jest.Mock;
+const openAiConnector = {
+  actionTypeId: '.gen-ai',
+  name: 'genAi',
+  id: '123',
+  config: {
+    apiUrl: 'https://openaiurl.com',
+    apiProvider: OpenAiProviderType.OpenAi,
+  },
+  secrets: {
+    apiKey: 'thats-a-nice-looking-key',
+  },
+  isDeprecated: false,
+};
+const azureConnector = {
+  ...openAiConnector,
+  config: {
+    apiUrl: 'https://azureaiurl.com',
+    apiProvider: OpenAiProviderType.AzureAi,
+  },
+  secrets: {
+    apiKey: 'thats-a-nice-looking-key',
+  },
+};
+
+const navigateToUrl = jest.fn();
 
 describe('GenerativeAiConnectorFields renders', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    useKibanaMock().services.application.navigateToUrl = navigateToUrl;
+    mockDashboard.mockImplementation(({ connectorId }) => ({
+      dashboardUrl: `https://dashboardurl.com/${connectorId}`,
+    }));
+  });
   test('open ai connector fields are rendered', async () => {
-    const actionConnector = {
-      actionTypeId: '.gen-ai',
-      name: 'genAi',
-      config: {
-        apiUrl: 'https://openaiurl.com',
-        apiProvider: OpenAiProviderType.OpenAi,
-      },
-      secrets: {
-        apiKey: 'thats-a-nice-looking-key',
-      },
-      isDeprecated: false,
-    };
-
     const { getAllByTestId } = render(
-      <ConnectorFormTestProvider connector={actionConnector}>
+      <ConnectorFormTestProvider connector={openAiConnector}>
         <GenerativeAiConnectorFields
           readOnly={false}
           isEdit={false}
@@ -37,71 +64,99 @@ describe('GenerativeAiConnectorFields renders', () => {
       </ConnectorFormTestProvider>
     );
     expect(getAllByTestId('config.apiUrl-input')[0]).toBeInTheDocument();
-    expect(getAllByTestId('config.apiUrl-input')[0]).toHaveValue(actionConnector.config.apiUrl);
+    expect(getAllByTestId('config.apiUrl-input')[0]).toHaveValue(openAiConnector.config.apiUrl);
     expect(getAllByTestId('config.apiProvider-select')[0]).toBeInTheDocument();
     expect(getAllByTestId('config.apiProvider-select')[0]).toHaveValue(
-      actionConnector.config.apiProvider
+      openAiConnector.config.apiProvider
     );
     expect(getAllByTestId('open-ai-api-doc')[0]).toBeInTheDocument();
     expect(getAllByTestId('open-ai-api-keys-doc')[0]).toBeInTheDocument();
   });
 
   test('azure ai connector fields are rendered', async () => {
-    const actionConnector = {
-      actionTypeId: '.gen-ai',
-      name: 'genAi',
-      config: {
-        apiUrl: 'https://azureaiurl.com',
-        apiProvider: OpenAiProviderType.AzureAi,
-      },
-      secrets: {
-        apiKey: 'thats-a-nice-looking-key',
-      },
-      isDeprecated: false,
-    };
-
     const { getAllByTestId } = render(
-      <ConnectorFormTestProvider connector={actionConnector}>
+      <ConnectorFormTestProvider connector={azureConnector}>
         <GenerativeAiConnectorFields
           readOnly={false}
           isEdit={false}
           registerPreSubmitValidator={() => {}}
         />
       </ConnectorFormTestProvider>
     );
-
     expect(getAllByTestId('config.apiUrl-input')[0]).toBeInTheDocument();
-    expect(getAllByTestId('config.apiUrl-input')[0]).toHaveValue(actionConnector.config.apiUrl);
+    expect(getAllByTestId('config.apiUrl-input')[0]).toHaveValue(azureConnector.config.apiUrl);
     expect(getAllByTestId('config.apiProvider-select')[0]).toBeInTheDocument();
     expect(getAllByTestId('config.apiProvider-select')[0]).toHaveValue(
-      actionConnector.config.apiProvider
+      azureConnector.config.apiProvider
     );
     expect(getAllByTestId('azure-ai-api-doc')[0]).toBeInTheDocument();
     expect(getAllByTestId('azure-ai-api-keys-doc')[0]).toBeInTheDocument();
   });
 
+  describe('Dashboard link', () => {
+    it('Does not render if isEdit is false and dashboardUrl is defined', async () => {
+      const { queryByTestId } = render(
+        <ConnectorFormTestProvider connector={openAiConnector}>
+          <GenerativeAiConnectorFields
+            readOnly={false}
+            isEdit={false}
+            registerPreSubmitValidator={() => {}}
+          />
+        </ConnectorFormTestProvider>
+      );
+      expect(queryByTestId('link-gen-ai-token-dashboard')).not.toBeInTheDocument();
+    });
+    it('Does not render if isEdit is true and dashboardUrl is null', async () => {
+      mockDashboard.mockImplementation((id: string) => ({
+        dashboardUrl: null,
+      }));
+      const { queryByTestId } = render(
+        <ConnectorFormTestProvider connector={openAiConnector}>
+          <GenerativeAiConnectorFields
+            readOnly={false}
+            isEdit={false}
+            registerPreSubmitValidator={() => {}}
+          />
+        </ConnectorFormTestProvider>
+      );
+      expect(queryByTestId('link-gen-ai-token-dashboard')).not.toBeInTheDocument();
+    });
+    it('Renders if isEdit is true and dashboardUrl is defined', async () => {
+      const { getByTestId } = render(
+        <ConnectorFormTestProvider connector={openAiConnector}>
+          <GenerativeAiConnectorFields
+            readOnly={false}
+            isEdit={true}
+            registerPreSubmitValidator={() => {}}
+          />
+        </ConnectorFormTestProvider>
+      );
+      expect(getByTestId('link-gen-ai-token-dashboard')).toBeInTheDocument();
+    });
+    it('On click triggers redirect with correct saved object id', async () => {
+      const { getByTestId } = render(
+        <ConnectorFormTestProvider connector={openAiConnector}>
+          <GenerativeAiConnectorFields
+            readOnly={false}
+            isEdit={true}
+            registerPreSubmitValidator={() => {}}
+          />
+        </ConnectorFormTestProvider>
+      );
+      fireEvent.click(getByTestId('link-gen-ai-token-dashboard'));
+      expect(navigateToUrl).toHaveBeenCalledWith(`https://dashboardurl.com/123`);
+    });
+  });
   describe('Validation', () => {
     const onSubmit = jest.fn();
-    const actionConnector = {
-      actionTypeId: '.gen-ai',
-      name: 'genAi',
-      config: {
-        apiUrl: 'https://openaiurl.com',
-        apiProvider: OpenAiProviderType.OpenAi,
-      },
-      secrets: {
-        apiKey: 'thats-a-nice-looking-key',
-      },
-      isDeprecated: false,
-    };
 
     beforeEach(() => {
       jest.clearAllMocks();
     });
 
     it('connector validation succeeds when connector config is valid', async () => {
       const { getByTestId } = render(
-        <ConnectorFormTestProvider connector={actionConnector} onSubmit={onSubmit}>
+        <ConnectorFormTestProvider connector={openAiConnector} onSubmit={onSubmit}>
           <GenerativeAiConnectorFields
             readOnly={false}
             isEdit={false}
@@ -119,16 +174,16 @@ describe('GenerativeAiConnectorFields renders', () => {
       });
 
       expect(onSubmit).toBeCalledWith({
-        data: actionConnector,
+        data: openAiConnector,
         isValid: true,
       });
     });
 
     it('validates correctly if the apiUrl is empty', async () => {
       const connector = {
-        ...actionConnector,
+        ...openAiConnector,
         config: {
-          ...actionConnector.config,
+          ...openAiConnector.config,
           apiUrl: '',
         },
       };
@@ -159,9 +214,9 @@ describe('GenerativeAiConnectorFields renders', () => {
     ];
     it.each(tests)('validates correctly %p', async (field, value) => {
       const connector = {
-        ...actionConnector,
+        ...openAiConnector,
         config: {
-          ...actionConnector.config,
+          ...openAiConnector.config,
           headers: [],
         },
       };