feat(core): add telemetry for subagent execution (#10456)

packages/core/src/agents/executor.test.ts

@@ -6,13 +6,6 @@
 
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { AgentExecutor, type ActivityCallback } from './executor.js';
-import type {
-  AgentDefinition,
-  AgentInputs,
-  SubagentActivityEvent,
-  OutputConfig,
-} from './types.js';
-import { AgentTerminateMode } from './types.js';
 import { makeFakeConfig } from '../test-utils/config.js';
 import { ToolRegistry } from '../tools/tool-registry.js';
 import { LSTool } from '../tools/ls.js';
@@ -32,6 +25,16 @@ import type { Config } from '../config/config.js';
 import { MockTool } from '../test-utils/mock-tool.js';
 import { getDirectoryContextString } from '../utils/environmentContext.js';
 import { z } from 'zod';
+import { promptIdContext } from '../utils/promptIdContext.js';
+import { logAgentStart, logAgentFinish } from '../telemetry/loggers.js';
+import { AgentStartEvent, AgentFinishEvent } from '../telemetry/types.js';
+import type {
+  AgentDefinition,
+  AgentInputs,
+  SubagentActivityEvent,
+  OutputConfig,
+} from './types.js';
+import { AgentTerminateMode } from './types.js';
 
 const { mockSendMessageStream, mockExecuteToolCall } = vi.hoisted(() => ({
   mockSendMessageStream: vi.fn(),
@@ -54,8 +57,29 @@ vi.mock('../core/nonInteractiveToolExecutor.js', () => ({
 
 vi.mock('../utils/environmentContext.js');
 
+vi.mock('../telemetry/loggers.js', () => ({
+  logAgentStart: vi.fn(),
+  logAgentFinish: vi.fn(),
+}));
+
+vi.mock('../utils/promptIdContext.js', async (importOriginal) => {
+  const actual =
+    await importOriginal<typeof import('../utils/promptIdContext.js')>();
+  return {
+    ...actual,
+    promptIdContext: {
+      ...actual.promptIdContext,
+      getStore: vi.fn(),
+      run: vi.fn((_id, fn) => fn()),
+    },
+  };
+});
+
 const MockedGeminiChat = vi.mocked(GeminiChat);
 const mockedGetDirectoryContextString = vi.mocked(getDirectoryContextString);
+const mockedPromptIdContext = vi.mocked(promptIdContext);
+const mockedLogAgentStart = vi.mocked(logAgentStart);
+const mockedLogAgentFinish = vi.mocked(logAgentFinish);
 
 // Constants for testing
 const TASK_COMPLETE_TOOL_NAME = 'complete_task';
@@ -160,6 +184,10 @@ describe('AgentExecutor', () => {
     vi.resetAllMocks();
     mockSendMessageStream.mockReset();
     mockExecuteToolCall.mockReset();
+    mockedLogAgentStart.mockReset();
+    mockedLogAgentFinish.mockReset();
+    mockedPromptIdContext.getStore.mockReset();
+    mockedPromptIdContext.run.mockImplementation((_id, fn) => fn());
 
     MockedGeminiChat.mockImplementation(
       () =>
@@ -229,9 +257,52 @@ describe('AgentExecutor', () => {
       expect(agentRegistry.getAllToolNames()).toHaveLength(2);
       expect(agentRegistry.getTool(MOCK_TOOL_NOT_ALLOWED.name)).toBeUndefined();
     });
+
+    it('should use parentPromptId from context to create agentId', async () => {
+      const parentId = 'parent-id';
+      mockedPromptIdContext.getStore.mockReturnValue(parentId);
+
+      const definition = createTestDefinition();
+      const executor = await AgentExecutor.create(
+        definition,
+        mockConfig,
+        onActivity,
+      );
+
+      expect(executor['agentId']).toMatch(
+        new RegExp(`^${parentId}-${definition.name}-`),
+      );
+    });
   });
 
   describe('run (Execution Loop and Logic)', () => {
+    it('should log AgentFinish with error if run throws', async () => {
+      const definition = createTestDefinition();
+      // Make the definition invalid to cause an error during run
+      definition.inputConfig.inputs = {
+        goal: { type: 'string', required: true, description: 'goal' },
+      };
+      const executor = await AgentExecutor.create(
+        definition,
+        mockConfig,
+        onActivity,
+      );
+
+      // Run without inputs to trigger validation error
+      await expect(executor.run({}, signal)).rejects.toThrow(
+        /Missing required input parameters/,
+      );
+
+      expect(mockedLogAgentStart).toHaveBeenCalledTimes(1);
+      expect(mockedLogAgentFinish).toHaveBeenCalledTimes(1);
+      expect(mockedLogAgentFinish).toHaveBeenCalledWith(
+        mockConfig,
+        expect.objectContaining({
+          terminate_reason: AgentTerminateMode.ERROR,
+        }),
+      );
+    });
+
     it('should execute successfully when model calls complete_task with output (Happy Path with Output)', async () => {
       const definition = createTestDefinition();
       const executor = await AgentExecutor.create(
@@ -312,6 +383,34 @@ describe('AgentExecutor', () => {
       expect(output.result).toBe('Found file1.txt');
       expect(output.terminate_reason).toBe(AgentTerminateMode.GOAL);
 
+      // Telemetry checks
+      expect(mockedLogAgentStart).toHaveBeenCalledTimes(1);
+      expect(mockedLogAgentStart).toHaveBeenCalledWith(
+        mockConfig,
+        expect.any(AgentStartEvent),
+      );
+      expect(mockedLogAgentFinish).toHaveBeenCalledTimes(1);
+      expect(mockedLogAgentFinish).toHaveBeenCalledWith(
+        mockConfig,
+        expect.any(AgentFinishEvent),
+      );
+      const finishEvent = mockedLogAgentFinish.mock.calls[0][1];
+      expect(finishEvent.terminate_reason).toBe(AgentTerminateMode.GOAL);
+
+      // Context checks
+      expect(mockedPromptIdContext.run).toHaveBeenCalledTimes(2); // Two turns
+      const agentId = executor['agentId'];
+      expect(mockedPromptIdContext.run).toHaveBeenNthCalledWith(
+        1,
+        `${agentId}#0`,
+        expect.any(Function),
+      );
+      expect(mockedPromptIdContext.run).toHaveBeenNthCalledWith(
+        2,
+        `${agentId}#1`,
+        expect.any(Function),
+      );
+
       expect(activities).toEqual(
         expect.arrayContaining([
           expect.objectContaining({
@@ -425,6 +524,14 @@ describe('AgentExecutor', () => {
       expect(output.terminate_reason).toBe(AgentTerminateMode.ERROR);
       expect(output.result).toBe(expectedError);
 
+      // Telemetry check for error
+      expect(mockedLogAgentFinish).toHaveBeenCalledWith(
+        mockConfig,
+        expect.objectContaining({
+          terminate_reason: AgentTerminateMode.ERROR,
+        }),
+      );
+
       expect(activities).toContainEqual(
         expect.objectContaining({
           type: 'ERROR',

packages/core/src/agents/executor.ts

@@ -28,6 +28,9 @@ import { MemoryTool } from '../tools/memoryTool.js';
 import { ReadFileTool } from '../tools/read-file.js';
 import { ReadManyFilesTool } from '../tools/read-many-files.js';
 import { WebSearchTool } from '../tools/web-search.js';
+import { promptIdContext } from '../utils/promptIdContext.js';
+import { logAgentStart, logAgentFinish } from '../telemetry/loggers.js';
+import { AgentStartEvent, AgentFinishEvent } from '../telemetry/types.js';
 import type {
   AgentDefinition,
   AgentInputs,
@@ -104,10 +107,14 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
       await AgentExecutor.validateTools(agentToolRegistry, definition.name);
     }
 
+    // Get the parent prompt ID from context
+    const parentPromptId = promptIdContext.getStore();
+
     return new AgentExecutor(
       definition,
       runtimeContext,
       agentToolRegistry,
+      parentPromptId,
       onActivity,
     );
   }
@@ -122,6 +129,7 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
     definition: AgentDefinition<TOutput>,
     runtimeContext: Config,
     toolRegistry: ToolRegistry,
+    parentPromptId: string | undefined,
     onActivity?: ActivityCallback,
   ) {
     this.definition = definition;
@@ -130,7 +138,10 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
     this.onActivity = onActivity;
 
     const randomIdPart = Math.random().toString(36).slice(2, 8);
-    this.agentId = `${this.definition.name}-${randomIdPart}`;
+    // parentPromptId will be undefined if this agent is invoked directly
+    // (top-level), rather than as a sub-agent.
+    const parentPrefix = parentPromptId ? `${parentPromptId}-` : '';
+    this.agentId = `${parentPrefix}${this.definition.name}-${randomIdPart}`;
   }
 
   /**
@@ -143,12 +154,17 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
   async run(inputs: AgentInputs, signal: AbortSignal): Promise<OutputObject> {
     const startTime = Date.now();
     let turnCounter = 0;
+    let terminateReason: AgentTerminateMode = AgentTerminateMode.ERROR;
+    let finalResult: string | null = null;
+
+    logAgentStart(
+      this.runtimeContext,
+      new AgentStartEvent(this.agentId, this.definition.name),
+    );
 
     try {
       const chat = await this.createChatObject(inputs);
       const tools = this.prepareToolsList();
-      let terminateReason = AgentTerminateMode.ERROR;
-      let finalResult: string | null = null;
 
       const query = this.definition.promptConfig.query
         ? templateString(this.definition.promptConfig.query, inputs)
@@ -167,14 +183,12 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
           break;
         }
 
-        // Call model
-        const promptId = `${this.runtimeContext.getSessionId()}#${this.agentId}#${turnCounter++}`;
-        const { functionCalls } = await this.callModel(
-          chat,
-          currentMessage,
-          tools,
-          signal,
+        const promptId = `${this.agentId}#${turnCounter++}`;
+
+        const { functionCalls } = await promptIdContext.run(
           promptId,
+          async () =>
+            this.callModel(chat, currentMessage, tools, signal, promptId),
         );
 
         if (signal.aborted) {
@@ -220,6 +234,17 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
     } catch (error) {
       this.emitActivity('ERROR', { error: String(error) });
       throw error; // Re-throw the error for the parent context to handle.
+    } finally {
+      logAgentFinish(
+        this.runtimeContext,
+        new AgentFinishEvent(
+          this.agentId,
+          this.definition.name,
+          Date.now() - startTime,
+          turnCounter,
+          terminateReason,
+        ),
+      );
     }
   }