[8.18] [Security Assistant] Fix telemetry tool reporting (#223832) (#224344)

# Backport

This will backport the following commits from `main` to `8.18`:
- [[Security Assistant] Fix telemetry tool reporting
(#223832)](https://github.com/elastic/kibana/pull/223832)

<!--- Backport version: 10.0.1 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sorenlouv/backport)

<!--BACKPORT [{"author":{"name":"Steph
Milovic","email":"stephanie.milovic@elastic.co"},"sourceCommit":{"committedDate":"2025-06-12T22:52:33Z","message":"[Security
Assistant] Fix telemetry tool reporting
(#223832)","sha":"602f96d46a3332471437e5e79faa5454293cbe14","branchLabelMapping":{"^v9.1.0$":"main","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:skip","backport
missing","Team: SecuritySolution","Team:Security Generative
AI","backport:version","v9.1.0","v8.19.0","v9.0.3","v8.18.3"],"title":"[Security
Assistant] Fix telemetry tool
reporting","number":223832,"url":"https://github.com/elastic/kibana/pull/223832","mergeCommit":{"message":"[Security
Assistant] Fix telemetry tool reporting
(#223832)","sha":"602f96d46a3332471437e5e79faa5454293cbe14"}},"sourceBranch":"main","suggestedTargetBranches":["9.0","8.18"],"targetPullRequestStates":[{"branch":"main","label":"v9.1.0","branchLabelMappingKey":"^v9.1.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/223832","number":223832,"mergeCommit":{"message":"[Security
Assistant] Fix telemetry tool reporting
(#223832)","sha":"602f96d46a3332471437e5e79faa5454293cbe14"}},{"branch":"8.19","label":"v8.19.0","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"url":"https://github.com/elastic/kibana/pull/223845","number":223845,"state":"OPEN"},{"branch":"9.0","label":"v9.0.3","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"8.18","label":"v8.18.3","branchLabelMappingKey":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"}]}]
BACKPORT-->

x-pack/platform/packages/shared/kbn-langchain/server/tracers/telemetry/telemetry_tracer.test.ts

@@ -162,7 +162,6 @@ describe('TelemetryTracer', () => {
         elasticTools,
         telemetry,
         telemetryParams,
-        totalTools: 9,
       },
       logger
     );
@@ -173,7 +172,6 @@ describe('TelemetryTracer', () => {
     expect(telemetryTracer.elasticTools).toEqual(elasticTools);
     expect(telemetryTracer.telemetry).toBe(telemetry);
     expect(telemetryTracer.telemetryParams).toBe(telemetryParams);
-    expect(telemetryTracer.totalTools).toBe(9);
   });
 
   it('should not log and report event on chain end if parent_run_id exists', async () => {

x-pack/platform/packages/shared/kbn-langchain/server/tracers/telemetry/telemetry_tracer.ts

@@ -21,7 +21,6 @@ export interface LangChainTracerFields extends BaseCallbackHandlerInput {
   elasticTools: string[];
   telemetry: AnalyticsServiceSetup;
   telemetryParams: TelemetryParams;
-  totalTools: number;
 }
 interface ToolRunStep {
   action: {
@@ -37,14 +36,12 @@ export class TelemetryTracer extends BaseTracer implements LangChainTracerFields
   elasticTools: string[];
   telemetry: AnalyticsServiceSetup;
   telemetryParams: TelemetryParams;
-  totalTools: number;
   constructor(fields: LangChainTracerFields, logger: Logger) {
     super(fields);
     this.logger = logger.get('telemetryTracer');
     this.elasticTools = fields.elasticTools;
     this.telemetry = fields.telemetry;
     this.telemetryParams = fields.telemetryParams;
-    this.totalTools = fields.totalTools;
   }
 
   async onChainEnd(run: Run): Promise<void> {

x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.test.ts

@@ -10,6 +10,7 @@ import { callAssistantGraph } from '.';
 import { getDefaultAssistantGraph } from './graph';
 import { invokeGraph, streamGraph } from './helpers';
 import { loggerMock } from '@kbn/logging-mocks';
+import { TelemetryTracer } from '@kbn/langchain/server/tracers/telemetry';
 import { AgentExecutorParams, AssistantDataClients } from '../../executors/types';
 import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks';
 import { getPrompt, resolveProviderAndModel } from '@kbn/security-ai-prompts';
@@ -30,6 +31,16 @@ jest.mock('@kbn/security-ai-prompts');
 const getDefaultAssistantGraphMock = getDefaultAssistantGraph as jest.Mock;
 const resolveProviderAndModelMock = resolveProviderAndModel as jest.Mock;
 const getPromptMock = getPrompt as jest.Mock;
+const telemetryTracerMock = TelemetryTracer as unknown as jest.Mock;
+const getTool = jest.fn();
+const mockTool: AssistantTool = {
+  id: 'id',
+  name: 'name',
+  description: 'description',
+  sourceRegister: 'sourceRegister',
+  isSupported: (params: AssistantToolParams) => true,
+  getTool: getTool.mockReturnValue({ name: 'name' }),
+};
 describe('callAssistantGraph', () => {
   const mockDataClients = {
     anonymizationFieldsDataClient: {
@@ -37,7 +48,7 @@ describe('callAssistantGraph', () => {
     },
     kbDataClient: {
       isInferenceEndpointExists: jest.fn(),
-      getAssistantTools: jest.fn(),
+      getAssistantTools: jest.fn().mockReturnValue([{ name: 'MyKBTool' }]),
     },
   } as unknown as AssistantDataClients;
 
@@ -54,6 +65,7 @@ describe('callAssistantGraph', () => {
     total: 0,
     saved_objects: [],
   });
+  const mockLogger = loggerMock.create();
   const defaultParams = {
     actionsClient: actionsClientMock.create(),
     alertsIndexPattern: 'test-pattern',
@@ -67,7 +79,7 @@ describe('callAssistantGraph', () => {
     llmTasks: { retrieveDocumentationAvailable: jest.fn(), retrieveDocumentation: jest.fn() },
     llmType: 'openai',
     isOssModel: false,
-    logger: loggerMock.create(),
+    logger: mockLogger,
     isStream: false,
     onLlmResponse: jest.fn(),
     onNewReplacements: jest.fn(),
@@ -176,15 +188,6 @@ describe('callAssistantGraph', () => {
   });
 
   it('calls getPrompt for each tool and the default system prompt', async () => {
-    const getTool = jest.fn();
-    const mockTool: AssistantTool = {
-      id: 'id',
-      name: 'name',
-      description: 'description',
-      sourceRegister: 'sourceRegister',
-      isSupported: (params: AssistantToolParams) => true,
-      getTool,
-    };
     const params = {
       ...defaultParams,
       assistantTools: [
@@ -227,6 +230,31 @@ describe('callAssistantGraph', () => {
     );
   });
 
+  it('Passes only Elastic tools, not custom, to Telemetry tracer', async () => {
+    await callAssistantGraph({
+      ...defaultParams,
+      assistantTools: [
+        { ...mockTool, name: 'test-tool', getTool: getTool.mockReturnValue({ name: 'test-tool' }) },
+        {
+          ...mockTool,
+          name: 'test-tool2',
+          getTool: getTool.mockReturnValue({ name: 'test-tool2' }),
+        },
+      ],
+    });
+    expect(telemetryTracerMock).toHaveBeenCalledWith(
+      {
+        elasticTools: ['test-tool2', 'test-tool2'],
+        telemetry: {},
+        telemetryParams: {},
+      },
+      {
+        ...mockLogger,
+        context: ['defaultAssistantGraph'],
+      }
+    );
+  });
+
   describe('agentRunnable', () => {
     it('creates OpenAIToolsAgent for openai llmType', async () => {
       const params = { ...defaultParams, llmType: 'openai' };

x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts

@@ -156,6 +156,19 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
     )
   ).filter((e) => e != null) as StructuredTool[];
 
+  const apmTracer = new APMTracer({ projectName: traceOptions?.projectName ?? 'default' }, logger);
+  const telemetryTracer = telemetryParams
+    ? new TelemetryTracer(
+        {
+          // this line MUST come before kbTools are added
+          elasticTools: tools.map(({ name }) => name),
+          telemetry,
+          telemetryParams,
+        },
+        logger
+      )
+    : undefined;
+
   // If KB enabled, fetch for any KB IndexEntries and generate a tool for each
   if (isEnabledKnowledgeBase) {
     const kbTools = await dataClients?.kbDataClient?.getAssistantTools({
@@ -191,18 +204,6 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
     prompt: chatPromptTemplate,
   });
 
-  const apmTracer = new APMTracer({ projectName: traceOptions?.projectName ?? 'default' }, logger);
-  const telemetryTracer = telemetryParams
-    ? new TelemetryTracer(
-        {
-          elasticTools: tools.map(({ name }) => name),
-          totalTools: tools.length,
-          telemetry,
-          telemetryParams,
-        },
-        logger
-      )
-    : undefined;
   const { provider } =
     !llmType || llmType === 'inference'
       ? await resolveProviderAndModel({