feat(agents): migrate subagents to event-driven scheduler (#17567)

packages/core/src/scheduler/scheduler.test.ts

@@ -70,6 +70,10 @@ import { ROOT_SCHEDULER_ID } from './types.js';
 import { ToolErrorType } from '../tools/tool-error.js';
 import * as ToolUtils from '../utils/tool-utils.js';
 import type { EditorType } from '../utils/editor.js';
+import {
+  getToolCallContext,
+  type ToolCallContext,
+} from '../utils/toolCallContext.js';
 
 describe('Scheduler (Orchestrator)', () => {
   let scheduler: Scheduler;
@@ -1010,4 +1014,68 @@ describe('Scheduler (Orchestrator)', () => {
       expect(mockStateManager.finalizeCall).toHaveBeenCalledWith('call-1');
     });
   });
+
+  describe('Tool Call Context Propagation', () => {
+    it('should propagate context to the tool executor', async () => {
+      const schedulerId = 'custom-scheduler';
+      const parentCallId = 'parent-call';
+      const customScheduler = new Scheduler({
+        config: mockConfig,
+        messageBus: mockMessageBus,
+        getPreferredEditor,
+        schedulerId,
+        parentCallId,
+      });
+
+      const validatingCall: ValidatingToolCall = {
+        status: 'validating',
+        request: req1,
+        tool: mockTool,
+        invocation: mockInvocation as unknown as AnyToolInvocation,
+      };
+
+      // Mock queueLength to run the loop once
+      Object.defineProperty(mockStateManager, 'queueLength', {
+        get: vi.fn().mockReturnValueOnce(1).mockReturnValue(0),
+        configurable: true,
+      });
+
+      vi.mocked(mockStateManager.dequeue).mockReturnValue(validatingCall);
+      Object.defineProperty(mockStateManager, 'firstActiveCall', {
+        get: vi.fn().mockReturnValue(validatingCall),
+        configurable: true,
+      });
+      vi.mocked(mockStateManager.getToolCall).mockReturnValue(validatingCall);
+
+      mockToolRegistry.getTool.mockReturnValue(mockTool);
+      mockPolicyEngine.check.mockResolvedValue({
+        decision: PolicyDecision.ALLOW,
+      });
+
+      let capturedContext: ToolCallContext | undefined;
+      mockExecutor.execute.mockImplementation(async () => {
+        capturedContext = getToolCallContext();
+        return {
+          status: 'success',
+          request: req1,
+          tool: mockTool,
+          invocation: mockInvocation as unknown as AnyToolInvocation,
+          response: {
+            callId: req1.callId,
+            responseParts: [],
+            resultDisplay: 'ok',
+            error: undefined,
+            errorType: undefined,
+          },
+        } as unknown as SuccessfulToolCall;
+      });
+
+      await customScheduler.schedule(req1, signal);
+
+      expect(capturedContext).toBeDefined();
+      expect(capturedContext!.callId).toBe(req1.callId);
+      expect(capturedContext!.schedulerId).toBe(schedulerId);
+      expect(capturedContext!.parentCallId).toBe(parentCallId);
+    });
+  });
 });

packages/core/src/scheduler/scheduler.ts

@@ -36,6 +36,7 @@ import {
   type SerializableConfirmationDetails,
   type ToolConfirmationRequest,
 } from '../confirmation-bus/types.js';
+import { runWithToolCallContext } from '../utils/toolCallContext.js';
 
 interface SchedulerQueueItem {
   requests: ToolCallRequestInfo[];
@@ -256,6 +257,7 @@ export class Scheduler {
       return this.state.completedBatch;
     } finally {
       this.isProcessing = false;
+      this.state.clearBatch();
       this._processNextInRequestQueue();
     }
   }
@@ -282,30 +284,39 @@ export class Scheduler {
     request: ToolCallRequestInfo,
     tool: AnyDeclarativeTool,
   ): ValidatingToolCall | ErroredToolCall {
-    try {
-      const invocation = tool.build(request.args);
-      return {
-        status: 'validating',
-        request,
-        tool,
-        invocation,
-        startTime: Date.now(),
+    return runWithToolCallContext(
+      {
+        callId: request.callId,
         schedulerId: this.schedulerId,
-      };
-    } catch (e) {
-      return {
-        status: 'error',
-        request,
-        tool,
-        response: createErrorResponse(
-          request,
-          e instanceof Error ? e : new Error(String(e)),
-          ToolErrorType.INVALID_TOOL_PARAMS,
-        ),
-        durationMs: 0,
-        schedulerId: this.schedulerId,
-      };
-    }
+        parentCallId: this.parentCallId,
+      },
+      () => {
+        try {
+          const invocation = tool.build(request.args);
+          return {
+            status: 'validating',
+            request,
+            tool,
+            invocation,
+            startTime: Date.now(),
+            schedulerId: this.schedulerId,
+          };
+        } catch (e) {
+          return {
+            status: 'error',
+            request,
+            tool,
+            response: createErrorResponse(
+              request,
+              e instanceof Error ? e : new Error(String(e)),
+              ToolErrorType.INVALID_TOOL_PARAMS,
+            ),
+            durationMs: 0,
+            schedulerId: this.schedulerId,
+          };
+        }
+      },
+    );
   }
 
   // --- Phase 2: Processing Loop ---
@@ -460,17 +471,29 @@ export class Scheduler {
     if (signal.aborted) throw new Error('Operation cancelled');
     this.state.updateStatus(callId, 'executing');
 
-    const result = await this.executor.execute({
-      call: this.state.firstActiveCall as ExecutingToolCall,
-      signal,
-      outputUpdateHandler: (id, out) =>
-        this.state.updateStatus(id, 'executing', { liveOutput: out }),
-      onUpdateToolCall: (updated) => {
-        if (updated.status === 'executing' && updated.pid) {
-          this.state.updateStatus(callId, 'executing', { pid: updated.pid });
-        }
+    const activeCall = this.state.firstActiveCall as ExecutingToolCall;
+
+    const result = await runWithToolCallContext(
+      {
+        callId: activeCall.request.callId,
+        schedulerId: this.schedulerId,
+        parentCallId: this.parentCallId,
       },
-    });
+      () =>
+        this.executor.execute({
+          call: activeCall,
+          signal,
+          outputUpdateHandler: (id, out) =>
+            this.state.updateStatus(id, 'executing', { liveOutput: out }),
+          onUpdateToolCall: (updated) => {
+            if (updated.status === 'executing' && updated.pid) {
+              this.state.updateStatus(callId, 'executing', {
+                pid: updated.pid,
+              });
+            }
+          },
+        }),
+    );
 
     if (result.status === 'success') {
       this.state.updateStatus(callId, 'success', result.response);