[9.1] [Security Solution] [AI assistant] Make Security AI assistant tools available even when the chat completion endpoint is called from outside of security. (#226206) (#226444)

kibanamachine · KDKHD · web-flow · commit 268b5e68d101 · 2025-07-03T17:59:46.000+02:00
# Backport This will backport the following commits from `main` to `9.1`: - [[Security Solution] [AI assistant] Make Security AI assistant tools available even when the chat completion endpoint is called from outside of security. (#226206)](#226206)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sorenlouv/backport)  Co-authored-by: Kenneth Kreindler <42113355+KDKHD@users.noreply.github.com>
diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/plugin.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/plugin.ts
@@ -194,7 +194,7 @@ export class ElasticAssistantPlugin
       getRegisteredFeatures: (pluginName: string) => {
         return appContextService.getRegisteredFeatures(pluginName);
       },
-      getRegisteredTools: (pluginName: string) => {
+      getRegisteredTools: (pluginName: string | string[]) => {
         return appContextService.getRegisteredTools(pluginName);
       },
     };
@@ -220,7 +220,7 @@ export class ElasticAssistantPlugin
       getRegisteredFeatures: (pluginName: string) => {
         return appContextService.getRegisteredFeatures(pluginName);
       },
-      getRegisteredTools: (pluginName: string) => {
+      getRegisteredTools: (pluginName: string | string[]) => {
         return appContextService.getRegisteredTools(pluginName);
       },
       registerFeatures: (pluginName: string, features: Partial<AssistantFeatures>) => {
diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts
@@ -293,8 +293,10 @@ export const langChainExecute = async ({
   const assistantContext = context.elasticAssistant;
   // We don't (yet) support invoking these tools interactively
   const unsupportedTools = new Set(['attack-discovery', DEFEND_INSIGHTS_ID]);
+  const pluginNames = Array.from(new Set([pluginName, DEFAULT_PLUGIN_NAME]));
+
   const assistantTools = assistantContext
-    .getRegisteredTools(pluginName)
+    .getRegisteredTools(pluginNames)
     .filter((tool) => !unsupportedTools.has(tool.id));
 
   // get a scoped esClient for assistant memory
diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/request_context_factory.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/request_context_factory.ts
@@ -107,7 +107,7 @@ export class RequestContextFactory implements IRequestContextFactory {
 
       getCurrentUser,
 
-      getRegisteredTools: (pluginName: string) => {
+      getRegisteredTools: (pluginName: string | string[]) => {
         return appContextService.getRegisteredTools(pluginName);
       },
 
diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/services/app_context.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/services/app_context.test.ts
@@ -35,6 +35,14 @@ describe('AppContextService', () => {
     isSupported: jest.fn(),
     getTool: jest.fn(),
   };
+  const toolThree: AssistantTool = {
+    id: 'tool-three',
+    name: 'ToolThree',
+    description: 'Description 3',
+    sourceRegister: 'Source3',
+    isSupported: jest.fn(),
+    getTool: jest.fn(),
+  };
 
   beforeEach(() => {
     appContextService.stop();
@@ -99,6 +107,17 @@ describe('AppContextService', () => {
     });
   });
 
+  it('get tools for multiple plugins', () => {
+    const pluginName1 = 'pluginName1';
+    const pluginName2 = 'pluginName2';
+
+    appContextService.start(mockAppContext);
+    appContextService.registerTools(pluginName2, [toolOne, toolThree]);
+    appContextService.registerTools(pluginName1, [toolOne]);
+
+    expect(appContextService.getRegisteredTools([pluginName1, pluginName2]).length).toEqual(2);
+  });
+
   describe('registering features', () => {
     it('should register and get features for a single plugin', () => {
       const pluginName = 'pluginName';
diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/services/app_context.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/services/app_context.ts
@@ -12,7 +12,7 @@ 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 GetRegisteredTools = (pluginName: string | string[]) => AssistantTool[];
 export type GetRegisteredFeatures = (pluginName: string) => AssistantFeatures;
 export interface ElasticAssistantAppContext {
   logger: Logger;
@@ -67,8 +67,16 @@ class AppContextService {
    *
    * @param pluginName
    */
-  public getRegisteredTools(pluginName: string): AssistantTool[] {
-    const tools = Array.from(this.registeredTools?.get(pluginName) ?? new Set<AssistantTool>());
+  public getRegisteredTools(pluginName: string | string[]): AssistantTool[] {
+    const pluginNames = Array.isArray(pluginName) ? pluginName : [pluginName];
+
+    const tools = [
+      ...new Set(
+        pluginNames
+          .map((name) => this.registeredTools?.get(name) ?? new Set<AssistantTool>())
+          .flatMap((set) => [...set])
+      ),
+    ];
 
     this.logger?.debug('AppContextService:getRegisteredTools');
     this.logger?.debug(`pluginName: ${pluginName}`);
