feat(core): Thread AgentLoopContext through core. (#21944)

Author: joshualitt

packages/a2a-server/src/agent/task.ts

@@ -542,7 +542,7 @@ export class Task {
     const messageBus = this.config.getMessageBus();
     const scheduler = new Scheduler({
       schedulerId: this.id,
-      config: this.config,
+      context: this.config,
       messageBus,
       getPreferredEditor: () => DEFAULT_GUI_EDITOR,
     });

packages/cli/src/nonInteractiveCli.ts

@@ -211,7 +211,7 @@ export async function runNonInteractive({
 
       const geminiClient = config.getGeminiClient();
       const scheduler = new Scheduler({
-        config,
+        context: config,
         messageBus: config.getMessageBus(),
         getPreferredEditor: () => undefined,
         schedulerId: ROOT_SCHEDULER_ID,

packages/cli/src/ui/hooks/useToolScheduler.ts

@@ -97,7 +97,7 @@ export function useToolScheduler(
   const scheduler = useMemo(
     () =>
       new Scheduler({
-        config,
+        context: config,
         messageBus,
         getPreferredEditor: () => getPreferredEditorRef.current(),
         schedulerId: ROOT_SCHEDULER_ID,

packages/core/src/agents/agent-scheduler.test.ts

@@ -20,6 +20,7 @@ vi.mock('../scheduler/scheduler.js', () => ({
 
 describe('agent-scheduler', () => {
   let mockToolRegistry: Mocked<ToolRegistry>;
+  let mockConfig: Mocked<Config>;
   let mockMessageBus: Mocked<MessageBus>;
 
   beforeEach(() => {
@@ -29,6 +30,14 @@ describe('agent-scheduler', () => {
       getTool: vi.fn(),
       getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
     } as unknown as Mocked<ToolRegistry>;
+    mockConfig = {
+      getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
+      toolRegistry: mockToolRegistry,
+    } as unknown as Mocked<Config>;
+    (mockConfig as unknown as { messageBus: MessageBus }).messageBus =
+      mockMessageBus;
+    (mockConfig as unknown as { toolRegistry: ToolRegistry }).toolRegistry =
+      mockToolRegistry;
   });
 
   it('should create a scheduler with agent-specific config', async () => {
@@ -69,7 +78,8 @@ describe('agent-scheduler', () => {
       }),
     );
 
-    const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].config;
+    // Verify that the scheduler's context has the overridden tool registry
+    const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].context;
     expect(schedulerConfig.toolRegistry).toBe(mockToolRegistry);
   });
 
@@ -106,9 +116,8 @@ describe('agent-scheduler', () => {
       },
     );
 
-    const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].config;
+    const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].context;
     expect(schedulerConfig.toolRegistry).toBe(agentRegistry);
     expect(schedulerConfig.toolRegistry).not.toBe(mainRegistry);
-    expect(schedulerConfig.getToolRegistry()).toBe(agentRegistry);
   });
 });

packages/core/src/agents/agent-scheduler.ts

@@ -65,7 +65,7 @@ export async function scheduleAgentTools(
   });
 
   const scheduler = new Scheduler({
-    config: agentConfig,
+    context: agentConfig,
     messageBus: toolRegistry.getMessageBus(),
     getPreferredEditor: getPreferredEditor ?? (() => undefined),
     schedulerId,

packages/core/src/agents/browser/browserAgentInvocation.ts

@@ -16,6 +16,7 @@
 
 import { randomUUID } from 'node:crypto';
 import type { Config } from '../../config/config.js';
+import { type AgentLoopContext } from '../../config/agent-loop-context.js';
 import { LocalAgentExecutor } from '../local-executor.js';
 import { safeJsonToMarkdown } from '../../utils/markdownUtils.js';
 import {
@@ -179,7 +180,7 @@ export class BrowserAgentInvocation extends BaseToolInvocation<
   ToolResult
 > {
   constructor(
-    private readonly config: Config,
+    private readonly context: AgentLoopContext,
     params: AgentInputs,
     messageBus: MessageBus,
     _toolName?: string,
@@ -194,6 +195,10 @@ export class BrowserAgentInvocation extends BaseToolInvocation<
     );
   }
 
+  private get config(): Config {
+    return this.context.config;
+  }
+
   /**
    * Returns a concise, human-readable description of the invocation.
    */
@@ -409,7 +414,7 @@ export class BrowserAgentInvocation extends BaseToolInvocation<
       // Create and run executor with the configured definition
       const executor = await LocalAgentExecutor.create(
         definition,
-        this.config,
+        this.context,
         onActivity,
       );
 

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

@@ -307,6 +307,11 @@ describe('LocalAgentExecutor', () => {
     vi.useFakeTimers();
 
     mockConfig = makeFakeConfig();
+    // .config is already set correctly by the getter on the instance.
+    Object.defineProperty(mockConfig, 'promptId', {
+      get: () => 'test-prompt-id',
+      configurable: true,
+    });
     parentToolRegistry = new ToolRegistry(
       mockConfig,
       mockConfig.getMessageBus(),
@@ -319,7 +324,9 @@ describe('LocalAgentExecutor', () => {
     );
     parentToolRegistry.registerTool(MOCK_TOOL_NOT_ALLOWED);
 
-    vi.spyOn(mockConfig, 'getToolRegistry').mockReturnValue(parentToolRegistry);
+    vi.spyOn(mockConfig, 'toolRegistry', 'get').mockReturnValue(
+      parentToolRegistry,
+    );
     vi.spyOn(mockConfig, 'getAgentRegistry').mockReturnValue({
       getAllAgentNames: () => [],
     } as unknown as AgentRegistry);
@@ -382,7 +389,10 @@ describe('LocalAgentExecutor', () => {
 
     it('should use parentPromptId from context to create agentId', async () => {
       const parentId = 'parent-id';
-      mockedPromptIdContext.getStore.mockReturnValue(parentId);
+      Object.defineProperty(mockConfig, 'promptId', {
+        get: () => parentId,
+        configurable: true,
+      });
 
       const definition = createTestDefinition();
       const executor = await LocalAgentExecutor.create(
@@ -2052,7 +2062,7 @@ describe('LocalAgentExecutor', () => {
         vi.spyOn(configWithHints, 'getAgentRegistry').mockReturnValue({
           getAllAgentNames: () => [],
         } as unknown as AgentRegistry);
-        vi.spyOn(configWithHints, 'getToolRegistry').mockReturnValue(
+        vi.spyOn(configWithHints, 'toolRegistry', 'get').mockReturnValue(
           parentToolRegistry,
         );
       });

packages/core/src/agents/local-executor.ts

@@ -5,6 +5,7 @@
  */
 
 import type { Config } from '../config/config.js';
+import { type AgentLoopContext } from '../config/agent-loop-context.js';
 import { reportError } from '../utils/errorReporting.js';
 import { GeminiChat, StreamEventType } from '../core/geminiChat.js';
 import {
@@ -92,29 +93,33 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
 
   private readonly agentId: string;
   private readonly toolRegistry: ToolRegistry;
-  private readonly runtimeContext: Config;
+  private readonly context: AgentLoopContext;
   private readonly onActivity?: ActivityCallback;
   private readonly compressionService: ChatCompressionService;
   private readonly parentCallId?: string;
   private hasFailedCompressionAttempt = false;
 
+  private get config(): Config {
+    return this.context.config;
+  }
+
   /**
    * Creates and validates a new `AgentExecutor` instance.
    *
    * This method ensures that all tools specified in the agent's definition are
    * safe for non-interactive use before creating the executor.
    *
    * @param definition The definition object for the agent.
-   * @param runtimeContext The global runtime configuration.
+   * @param context The execution context.
    * @param onActivity An optional callback to receive activity events.
    * @returns A promise that resolves to a new `LocalAgentExecutor` instance.
    */
   static async create<TOutput extends z.ZodTypeAny>(
     definition: LocalAgentDefinition<TOutput>,
-    runtimeContext: Config,
+    context: AgentLoopContext,
     onActivity?: ActivityCallback,
   ): Promise<LocalAgentExecutor<TOutput>> {
-    const parentMessageBus = runtimeContext.getMessageBus();
+    const parentMessageBus = context.messageBus;
 
     // Create an override object to inject the subagent name into tool confirmation requests
     // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
@@ -133,12 +138,12 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
 
     // Create an isolated tool registry for this agent instance.
     const agentToolRegistry = new ToolRegistry(
-      runtimeContext,
+      context.config,
       subagentMessageBus,
     );
-    const parentToolRegistry = runtimeContext.getToolRegistry();
+    const parentToolRegistry = context.toolRegistry;
     const allAgentNames = new Set(
-      runtimeContext.getAgentRegistry().getAllAgentNames(),
+      context.config.getAgentRegistry().getAllAgentNames(),
     );
 
     const registerToolByName = (toolName: string) => {
@@ -190,15 +195,15 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
     agentToolRegistry.sortTools();
 
     // Get the parent prompt ID from context
-    const parentPromptId = promptIdContext.getStore();
+    const parentPromptId = context.promptId;
 
     // Get the parent tool call ID from context
     const toolContext = getToolCallContext();
     const parentCallId = toolContext?.callId;
 
     return new LocalAgentExecutor(
       definition,
-      runtimeContext,
+      context,
       agentToolRegistry,
       parentPromptId,
       parentCallId,
@@ -214,14 +219,14 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
    */
   private constructor(
     definition: LocalAgentDefinition<TOutput>,
-    runtimeContext: Config,
+    context: AgentLoopContext,
     toolRegistry: ToolRegistry,
     parentPromptId: string | undefined,
     parentCallId: string | undefined,
     onActivity?: ActivityCallback,
   ) {
     this.definition = definition;
-    this.runtimeContext = runtimeContext;
+    this.context = context;
     this.toolRegistry = toolRegistry;
     this.onActivity = onActivity;
     this.compressionService = new ChatCompressionService();
@@ -418,7 +423,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
     } finally {
       clearTimeout(graceTimeoutId);
       logRecoveryAttempt(
-        this.runtimeContext,
+        this.config,
         new RecoveryAttemptEvent(
           this.agentId,
           this.definition.name,
@@ -466,7 +471,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
     const combinedSignal = AbortSignal.any([signal, deadlineTimer.signal]);
 
     logAgentStart(
-      this.runtimeContext,
+      this.config,
       new AgentStartEvent(this.agentId, this.definition.name),
     );
 
@@ -477,7 +482,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
       const augmentedInputs = {
         ...inputs,
         cliVersion: await getVersion(),
-        activeModel: this.runtimeContext.getActiveModel(),
+        activeModel: this.config.getActiveModel(),
         today: new Date().toLocaleDateString(),
       };
 
@@ -494,13 +499,12 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
       // Capture the index of the last hint before starting to avoid re-injecting old hints.
       // NOTE: Hints added AFTER this point will be broadcast to all currently running
       // local agents via the listener below.
-      const startIndex =
-        this.runtimeContext.userHintService.getLatestHintIndex();
-      this.runtimeContext.userHintService.onUserHint(hintListener);
+      const startIndex = this.config.userHintService.getLatestHintIndex();
+      this.config.userHintService.onUserHint(hintListener);
 
       try {
         const initialHints =
-          this.runtimeContext.userHintService.getUserHintsAfter(startIndex);
+          this.config.userHintService.getUserHintsAfter(startIndex);
         const formattedInitialHints = formatUserHintsForModel(initialHints);
 
         let currentMessage: Content = formattedInitialHints
@@ -561,7 +565,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
           }
         }
       } finally {
-        this.runtimeContext.userHintService.offUserHint(hintListener);
+        this.config.userHintService.offUserHint(hintListener);
       }
 
       // === UNIFIED RECOVERY BLOCK ===
@@ -674,7 +678,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
     } finally {
       deadlineTimer.abort();
       logAgentFinish(
-        this.runtimeContext,
+        this.config,
         new AgentFinishEvent(
           this.agentId,
           this.definition.name,
@@ -697,7 +701,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
       prompt_id,
       false,
       model,
-      this.runtimeContext,
+      this.config,
       this.hasFailedCompressionAttempt,
     );
 
@@ -735,11 +739,10 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
     const modelConfigAlias = getModelConfigAlias(this.definition);
 
     // Resolve the model config early to get the concrete model string (which may be `auto`).
-    const resolvedConfig =
-      this.runtimeContext.modelConfigService.getResolvedConfig({
-        model: modelConfigAlias,
-        overrideScope: this.definition.name,
-      });
+    const resolvedConfig = this.config.modelConfigService.getResolvedConfig({
+      model: modelConfigAlias,
+      overrideScope: this.definition.name,
+    });
     const requestedModel = resolvedConfig.model;
 
     let modelToUse: string;
@@ -756,7 +759,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
           signal,
           requestedModel,
         };
-        const router = this.runtimeContext.getModelRouterService();
+        const router = this.config.getModelRouterService();
         const decision = await router.route(routingContext);
         modelToUse = decision.model;
       } catch (error) {
@@ -844,7 +847,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
 
     try {
       return new GeminiChat(
-        this.runtimeContext,
+        this.config,
         systemInstruction,
         [{ functionDeclarations: tools }],
         startHistory,
@@ -1092,7 +1095,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
     // Execute standard tool calls using the new scheduler
     if (toolRequests.length > 0) {
       const completedCalls = await scheduleAgentTools(
-        this.runtimeContext,
+        this.config,
         toolRequests,
         {
           schedulerId: this.agentId,
@@ -1240,7 +1243,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
     let finalPrompt = templateString(promptConfig.systemPrompt, inputs);
 
     // Append environment context (CWD and folder structure).
-    const dirContext = await getDirectoryContextString(this.runtimeContext);
+    const dirContext = await getDirectoryContextString(this.config);
     finalPrompt += `\n\n# Environment Context\n${dirContext}`;
 
     // Append standard rules for non-interactive execution.

packages/core/src/agents/local-invocation.test.ts

@@ -67,6 +67,11 @@ describe('LocalSubagentInvocation', () => {
   beforeEach(() => {
     vi.clearAllMocks();
     mockConfig = makeFakeConfig();
+    // .config is already set correctly by the getter on the instance.
+    Object.defineProperty(mockConfig, 'promptId', {
+      get: () => 'test-prompt-id',
+      configurable: true,
+    });
     mockMessageBus = createMockMessageBus();
 
     mockExecutorInstance = {