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

packages/core/src/utils/toolCallContext.test.ts

@@ -0,0 +1,84 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect } from 'vitest';
+import {
+  runWithToolCallContext,
+  getToolCallContext,
+} from './toolCallContext.js';
+
+describe('toolCallContext', () => {
+  it('should store and retrieve tool call context', () => {
+    const context = {
+      callId: 'test-call-id',
+      schedulerId: 'test-scheduler-id',
+    };
+
+    runWithToolCallContext(context, () => {
+      const storedContext = getToolCallContext();
+      expect(storedContext).toEqual(context);
+    });
+  });
+
+  it('should return undefined when no context is set', () => {
+    expect(getToolCallContext()).toBeUndefined();
+  });
+
+  it('should support nested contexts', () => {
+    const parentContext = {
+      callId: 'parent-call-id',
+      schedulerId: 'parent-scheduler-id',
+    };
+
+    const childContext = {
+      callId: 'child-call-id',
+      schedulerId: 'child-scheduler-id',
+      parentCallId: 'parent-call-id',
+    };
+
+    runWithToolCallContext(parentContext, () => {
+      expect(getToolCallContext()).toEqual(parentContext);
+
+      runWithToolCallContext(childContext, () => {
+        expect(getToolCallContext()).toEqual(childContext);
+      });
+
+      expect(getToolCallContext()).toEqual(parentContext);
+    });
+  });
+
+  it('should maintain isolation between parallel executions', async () => {
+    const context1 = {
+      callId: 'call-1',
+      schedulerId: 'scheduler-1',
+    };
+
+    const context2 = {
+      callId: 'call-2',
+      schedulerId: 'scheduler-2',
+    };
+
+    const promise1 = new Promise<void>((resolve) => {
+      runWithToolCallContext(context1, () => {
+        setTimeout(() => {
+          expect(getToolCallContext()).toEqual(context1);
+          resolve();
+        }, 10);
+      });
+    });
+
+    const promise2 = new Promise<void>((resolve) => {
+      runWithToolCallContext(context2, () => {
+        setTimeout(() => {
+          expect(getToolCallContext()).toEqual(context2);
+          resolve();
+        }, 5);
+      });
+    });
+
+    await Promise.all([promise1, promise2]);
+  });
+});

packages/core/src/utils/toolCallContext.ts

@@ -0,0 +1,47 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { AsyncLocalStorage } from 'node:async_hooks';
+
+/**
+ * Contextual information for a tool call execution.
+ */
+export interface ToolCallContext {
+  /** The unique ID of the tool call. */
+  callId: string;
+  /** The ID of the scheduler managing the execution. */
+  schedulerId: string;
+  /** The ID of the parent tool call, if this is a nested execution (e.g., in a subagent). */
+  parentCallId?: string;
+}
+
+/**
+ * AsyncLocalStorage instance for tool call context.
+ */
+export const toolCallContext = new AsyncLocalStorage<ToolCallContext>();
+
+/**
+ * Runs a function within a tool call context.
+ *
+ * @param context The context to set.
+ * @param fn The function to run.
+ * @returns The result of the function.
+ */
+export function runWithToolCallContext<T>(
+  context: ToolCallContext,
+  fn: () => T,
+): T {
+  return toolCallContext.run(context, fn);
+}
+
+/**
+ * Retrieves the current tool call context.
+ *
+ * @returns The current ToolCallContext, or undefined if not in a context.
+ */
+export function getToolCallContext(): ToolCallContext | undefined {
+  return toolCallContext.getStore();
+}