diff --git a/packages/core/src/agents/backends/InProcessBackend.test.ts b/packages/core/src/agents/backends/InProcessBackend.test.ts index a7fd17726..822d818ad 100644 --- a/packages/core/src/agents/backends/InProcessBackend.test.ts +++ b/packages/core/src/agents/backends/InProcessBackend.test.ts @@ -81,6 +81,28 @@ vi.mock('../runtime/agent-core.js', () => ({ }), })); +// Mirrors the positional AgentCore constructor parameters so tests can +// destructure by name instead of indexing — adding new parameters can't +// silently shift assertions onto the wrong slot. +function destructureAgentCoreCall(call: unknown[]) { + return { + name: call[0] as string, + runtimeContext: call[1] as Record, + promptConfig: call[2], + modelConfig: call[3], + runConfig: call[4], + toolConfig: call[5], + eventEmitter: call[6], + hooks: call[7], + runtimeView: call[8] as + | { + contentGenerator: unknown; + contentGeneratorConfig: { authType: string; model?: string }; + } + | undefined, + }; +} + function createMockToolRegistry() { return { getFunctionDeclarations: vi.fn().mockReturnValue([]), @@ -388,8 +410,8 @@ describe('InProcessBackend', () => { const lastCall = MockAgentCore.mock.calls.at(-1); expect(lastCall).toBeDefined(); - // Second arg is the runtime context (Config) - const agentContext = lastCall![1] as { + const { runtimeContext } = destructureAgentCoreCall(lastCall!); + const agentContext = runtimeContext as unknown as { getWorkingDir: () => string; getTargetDir: () => string; getToolRegistry: () => unknown; @@ -560,6 +582,13 @@ describe('InProcessBackend', () => { await backend.spawnAgent(config); const mockCreate = createContentGenerator as ReturnType; + // Owner must be the per-agent override Config (the same instance + // AgentCore receives as runtimeContext) — NOT the parent. Asserting + // that match exactly catches a regression where `base` slips in. + const MockAgentCore = AgentCore as unknown as ReturnType; + const { runtimeContext: agentContext } = destructureAgentCoreCall( + MockAgentCore.mock.calls.at(-1)!, + ); expect(mockCreate).toHaveBeenCalledWith( expect.objectContaining({ authType: 'anthropic', @@ -567,11 +596,11 @@ describe('InProcessBackend', () => { baseUrl: 'https://agent.example.com', model: 'test-model', }), - expect.anything(), + agentContext, ); }); - it('should override getContentGenerator on per-agent config', async () => { + it('should pass per-agent ContentGenerator via runtimeView', async () => { const agentGenerator = { generateContentStream: vi.fn() }; const mockCreate = createContentGenerator as ReturnType; mockCreate.mockResolvedValueOnce(agentGenerator); @@ -588,14 +617,11 @@ describe('InProcessBackend', () => { const MockAgentCore = AgentCore as unknown as ReturnType; const lastCall = MockAgentCore.mock.calls.at(-1); - const agentContext = lastCall![1] as { - getContentGenerator: () => unknown; - getAuthType: () => string | undefined; - getModel: () => string; - }; + const { runtimeView } = destructureAgentCoreCall(lastCall!); - expect(agentContext.getContentGenerator()).toBe(agentGenerator); - expect(agentContext.getAuthType()).toBe('anthropic'); + expect(runtimeView).toBeDefined(); + expect(runtimeView!.contentGenerator).toBe(agentGenerator); + expect(runtimeView!.contentGeneratorConfig.authType).toBe('anthropic'); expect(backend.getAgentContentGenerator('agent-1')).toBe(agentGenerator); }); @@ -629,12 +655,9 @@ describe('InProcessBackend', () => { const MockAgentCore = AgentCore as unknown as ReturnType; const lastCall = MockAgentCore.mock.calls.at(-1); - const agentContext = lastCall![1] as { - getContentGenerator: () => unknown; - }; - // Falls back to parent's content generator - expect(agentContext.getContentGenerator()).toBe(mockContentGenerator); + // No runtimeView when per-agent creation failed; agent inherits parent. + expect(destructureAgentCoreCall(lastCall!).runtimeView).toBeUndefined(); expect(backend.getAgentContentGenerator('agent-1')).toBeUndefined(); }); @@ -665,16 +688,12 @@ describe('InProcessBackend', () => { const MockAgentCore = AgentCore as unknown as ReturnType; const calls = MockAgentCore.mock.calls; - const ctx1 = calls.at(-2)![1] as { - getContentGenerator: () => unknown; - }; - const ctx2 = calls.at(-1)![1] as { - getContentGenerator: () => unknown; - }; + const view1 = calls.at(-2)![8] as { contentGenerator: unknown }; + const view2 = calls.at(-1)![8] as { contentGenerator: unknown }; - expect(ctx1.getContentGenerator()).toBe(gen1); - expect(ctx2.getContentGenerator()).toBe(gen2); - expect(ctx1.getContentGenerator()).not.toBe(ctx2.getContentGenerator()); + expect(view1.contentGenerator).toBe(gen1); + expect(view2.contentGenerator).toBe(gen2); + expect(view1.contentGenerator).not.toBe(view2.contentGenerator); }); }); }); diff --git a/packages/core/src/agents/backends/InProcessBackend.ts b/packages/core/src/agents/backends/InProcessBackend.ts index 7860dfe82..d5912400d 100644 --- a/packages/core/src/agents/backends/InProcessBackend.ts +++ b/packages/core/src/agents/backends/InProcessBackend.ts @@ -13,16 +13,12 @@ import { createDebugLogger } from '../../utils/debugLogger.js'; import type { Config } from '../../config/config.js'; -import { - type AuthType, - type ContentGenerator, - type ContentGeneratorConfig, - createContentGenerator, -} from '../../core/contentGenerator.js'; +import { type ContentGenerator } from '../../core/contentGenerator.js'; +import type { RuntimeContentGeneratorView } from '../runtime/agent-context.js'; import type { ToolRegistry } from '../../tools/tool-registry.js'; import { WorkspaceContext } from '../../utils/workspaceContext.js'; import { FileDiscoveryService } from '../../services/fileDiscoveryService.js'; -import { buildAgentContentGeneratorConfig } from '../../models/content-generator-config.js'; +import { createRuntimeContentGeneratorView } from '../../models/content-generator-config.js'; import { AgentStatus, isTerminalStatus } from '../runtime/agent-types.js'; import { AgentCore } from '../runtime/agent-core.js'; import { AgentEventEmitter } from '../runtime/agent-events.js'; @@ -122,6 +118,8 @@ export class InProcessBackend implements Backend { runConfig, toolConfig, eventEmitter, + undefined, + perAgent.runtimeView, ); const interactive = new AgentInteractive( @@ -382,19 +380,26 @@ export class InProcessBackend implements Backend { * - `getFileService()` → FileDiscoveryService rooted at agent's cwd * - `getToolRegistry()` → per-agent tool registry with core tools bound to * the agent Config - * - `getContentGenerator()` / `getContentGeneratorConfig()` / `getAuthType()` - * → per-agent ContentGenerator when `authOverrides` is provided - * - returned `contentGenerator` → the generator safe to use for summaries + * + * When `authOverrides` is provided, also returns a `runtimeView` describing + * the per-agent ContentGenerator. The agent runtime publishes the view via + * AsyncLocalStorage so the CG-related Config getters resolve to the + * agent's values during the run. */ async function createPerAgentConfig( base: Config, cwd: string, modelId?: string, authOverrides?: InProcessSpawnConfig['authOverrides'], -): Promise<{ config: Config; contentGenerator?: ContentGenerator }> { +): Promise<{ + config: Config; + contentGenerator?: ContentGenerator; + runtimeView?: RuntimeContentGeneratorView; +}> { // eslint-disable-next-line @typescript-eslint/no-explicit-any const override = Object.create(base) as any; let dedicatedContentGenerator: ContentGenerator | undefined; + let runtimeView: RuntimeContentGeneratorView | undefined; override.getWorkingDir = () => cwd; override.getTargetDir = () => cwd; @@ -415,25 +420,16 @@ async function createPerAgentConfig( if (authOverrides?.authType) { try { - const agentGeneratorConfig = buildAgentContentGeneratorConfig( + runtimeView = await createRuntimeContentGeneratorView( base, + override as Config, modelId, authOverrides, ); - const agentGenerator = await createContentGenerator( - agentGeneratorConfig, - override as Config, - ); - dedicatedContentGenerator = agentGenerator; - override.getContentGenerator = (): ContentGenerator => agentGenerator; - override.getContentGeneratorConfig = (): ContentGeneratorConfig => - agentGeneratorConfig; - override.getAuthType = (): AuthType | undefined => - agentGeneratorConfig.authType; - override.getModel = (): string => agentGeneratorConfig.model; + dedicatedContentGenerator = runtimeView.contentGenerator; debugLogger.info( - `Created per-agent ContentGenerator: authType=${authOverrides.authType}, model=${agentGeneratorConfig.model}`, + `Created per-agent ContentGenerator: authType=${authOverrides.authType}, model=${runtimeView.contentGeneratorConfig.model}`, ); } catch (error) { debugLogger.error( @@ -448,5 +444,6 @@ async function createPerAgentConfig( contentGenerator: dedicatedContentGenerator ?? (authOverrides?.authType ? undefined : base.getContentGenerator()), + runtimeView, }; } diff --git a/packages/core/src/agents/background-agent-resume.ts b/packages/core/src/agents/background-agent-resume.ts index f7631b31c..77e534795 100644 --- a/packages/core/src/agents/background-agent-resume.ts +++ b/packages/core/src/agents/background-agent-resume.ts @@ -27,7 +27,7 @@ import type { ChatRecord } from '../services/chatRecordingService.js'; import { getInitialChatHistory } from '../utils/environmentContext.js'; import { getGitBranch } from '../utils/gitUtils.js'; import { PermissionMode, type StopHookOutput } from '../hooks/types.js'; -import { runWithAgentContext } from '../tools/agent/agent-context.js'; +import { runWithAgentContext } from './runtime/agent-context.js'; import { createApprovalModeOverride } from '../tools/agent/agent.js'; import type { ApprovalMode } from '../config/config.js'; import { @@ -784,8 +784,7 @@ export class BackgroundAgentResumeService { } }; - const framedRunBody = () => - runWithAgentContext({ agentId: meta.agentId }, runBody); + const framedRunBody = () => runWithAgentContext(meta.agentId, runBody); void (target.isFork ? runInForkContext(framedRunBody) : framedRunBody()); return entry; } catch (error) { diff --git a/packages/core/src/agents/runtime/agent-context.test.ts b/packages/core/src/agents/runtime/agent-context.test.ts new file mode 100644 index 000000000..f1b713f6a --- /dev/null +++ b/packages/core/src/agents/runtime/agent-context.test.ts @@ -0,0 +1,163 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import { + getCurrentAgentId, + getRuntimeContentGenerator, + runWithAgentContext, + runWithRuntimeContentGenerator, + type RuntimeContentGeneratorView, +} from './agent-context.js'; +import { + AuthType, + type ContentGenerator, + type ContentGeneratorConfig, +} from '../../core/contentGenerator.js'; + +function makeView(model: string): RuntimeContentGeneratorView { + return { + contentGenerator: { tag: model } as unknown as ContentGenerator, + contentGeneratorConfig: { + model, + authType: AuthType.USE_OPENAI, + } as ContentGeneratorConfig, + }; +} + +describe('agent-context (agentId)', () => { + it('returns null outside any frame', () => { + expect(getCurrentAgentId()).toBeNull(); + }); + + it('exposes the agentId inside a frame', async () => { + await runWithAgentContext('explore-abc', async () => { + expect(getCurrentAgentId()).toBe('explore-abc'); + }); + expect(getCurrentAgentId()).toBeNull(); + }); + + it('propagates across awaits', async () => { + await runWithAgentContext('outer-1', async () => { + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(getCurrentAgentId()).toBe('outer-1'); + }); + }); + + it('nested frames shadow the outer agentId', async () => { + await runWithAgentContext('outer-1', async () => { + expect(getCurrentAgentId()).toBe('outer-1'); + await runWithAgentContext('inner-2', async () => { + expect(getCurrentAgentId()).toBe('inner-2'); + }); + expect(getCurrentAgentId()).toBe('outer-1'); + }); + }); + + it('isolates concurrent frames', async () => { + const results: string[] = []; + await Promise.all([ + runWithAgentContext('a', async () => { + await new Promise((resolve) => setTimeout(resolve, 5)); + results.push(getCurrentAgentId() ?? 'null'); + }), + runWithAgentContext('b', async () => { + await new Promise((resolve) => setTimeout(resolve, 1)); + results.push(getCurrentAgentId() ?? 'null'); + }), + ]); + expect(results.sort()).toEqual(['a', 'b']); + }); +}); + +describe('agent-context (runtimeView)', () => { + it('returns undefined outside any frame', () => { + expect(getRuntimeContentGenerator()).toBeUndefined(); + }); + + it('exposes the view to code running inside the frame', async () => { + const view = makeView('qwen3.6-plus'); + const inner = await runWithRuntimeContentGenerator(view, async () => + getRuntimeContentGenerator(), + ); + expect(inner).toBe(view); + expect(getRuntimeContentGenerator()).toBeUndefined(); + }); + + it('isolates sibling runs', async () => { + const v1 = makeView('model-a'); + const v2 = makeView('model-b'); + const [seen1, seen2] = await Promise.all([ + runWithRuntimeContentGenerator(v1, async () => + getRuntimeContentGenerator(), + ), + runWithRuntimeContentGenerator(v2, async () => + getRuntimeContentGenerator(), + ), + ]); + expect(seen1).toBe(v1); + expect(seen2).toBe(v2); + }); + + it('propagates through await chains', async () => { + const view = makeView('chained'); + const seen = await runWithRuntimeContentGenerator(view, async () => { + await Promise.resolve(); + await Promise.resolve(); + return getRuntimeContentGenerator(); + }); + expect(seen).toBe(view); + }); + + it('lets a nested run shadow the outer view', async () => { + const outer = makeView('outer'); + const inner = makeView('inner'); + const [seenOuter, seenInner] = await runWithRuntimeContentGenerator( + outer, + async () => { + const before = getRuntimeContentGenerator(); + const after = await runWithRuntimeContentGenerator(inner, async () => + getRuntimeContentGenerator(), + ); + return [before, after]; + }, + ); + expect(seenOuter).toBe(outer); + expect(seenInner).toBe(inner); + const outerAgain = await runWithRuntimeContentGenerator(outer, async () => { + await runWithRuntimeContentGenerator(inner, async () => undefined); + return getRuntimeContentGenerator(); + }); + expect(outerAgain).toBe(outer); + }); +}); + +describe('agent-context (merging)', () => { + it('runtimeView wrap preserves agentId from outer frame', async () => { + const view = makeView('inner-model'); + await runWithAgentContext('outer-agent', async () => { + await runWithRuntimeContentGenerator(view, async () => { + expect(getCurrentAgentId()).toBe('outer-agent'); + expect(getRuntimeContentGenerator()).toBe(view); + }); + // After the inner run resolves, runtimeView is gone but agentId stays. + expect(getCurrentAgentId()).toBe('outer-agent'); + expect(getRuntimeContentGenerator()).toBeUndefined(); + }); + }); + + it('agentId wrap preserves runtimeView from outer frame', async () => { + const view = makeView('outer-model'); + await runWithRuntimeContentGenerator(view, async () => { + await runWithAgentContext('inner-agent', async () => { + expect(getRuntimeContentGenerator()).toBe(view); + expect(getCurrentAgentId()).toBe('inner-agent'); + }); + expect(getRuntimeContentGenerator()).toBe(view); + expect(getCurrentAgentId()).toBeNull(); + }); + }); +}); diff --git a/packages/core/src/agents/runtime/agent-context.ts b/packages/core/src/agents/runtime/agent-context.ts new file mode 100644 index 000000000..285e47430 --- /dev/null +++ b/packages/core/src/agents/runtime/agent-context.ts @@ -0,0 +1,63 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * @fileoverview Per-run AsyncLocalStorage frame for agent execution. + * + * Tools capture `this.config` at construction time, so a sub-agent running + * with a different model cannot rely on the constructor-bound Config to + * report the right ContentGenerator or modalities. This frame lets + * `Config.getContentGenerator{,Config}()` resolve to the active sub-agent + * view, and lets nested `agent` tool launches discover their parent's id — + * both without threading extra parameters through every call site. + * + * Helpers patch one field at a time and merge with whatever is already on + * the stack, so wrapping at different layers preserves every set field. + */ + +import { AsyncLocalStorage } from 'node:async_hooks'; +import type { + ContentGenerator, + ContentGeneratorConfig, +} from '../../core/contentGenerator.js'; + +export interface RuntimeContentGeneratorView { + readonly contentGenerator: ContentGenerator; + readonly contentGeneratorConfig: ContentGeneratorConfig; +} + +interface AgentContext { + readonly agentId?: string; + readonly runtimeView?: RuntimeContentGeneratorView; +} + +const storage = new AsyncLocalStorage(); + +export function runWithAgentContext( + agentId: string, + fn: () => Promise, +): Promise { + const current = storage.getStore() ?? {}; + return storage.run({ ...current, agentId }, fn); +} + +export function runWithRuntimeContentGenerator( + view: RuntimeContentGeneratorView, + fn: () => Promise, +): Promise { + const current = storage.getStore() ?? {}; + return storage.run({ ...current, runtimeView: view }, fn); +} + +export function getCurrentAgentId(): string | null { + return storage.getStore()?.agentId ?? null; +} + +export function getRuntimeContentGenerator(): + | RuntimeContentGeneratorView + | undefined { + return storage.getStore()?.runtimeView; +} diff --git a/packages/core/src/agents/runtime/agent-core.test.ts b/packages/core/src/agents/runtime/agent-core.test.ts new file mode 100644 index 000000000..1abc33212 --- /dev/null +++ b/packages/core/src/agents/runtime/agent-core.test.ts @@ -0,0 +1,211 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import { AgentCore } from './agent-core.js'; +import { + getRuntimeContentGenerator, + runWithRuntimeContentGenerator, + type RuntimeContentGeneratorView, +} from './agent-context.js'; +import { subagentNameContext } from '../../utils/subagentNameContext.js'; +import type { Config } from '../../config/config.js'; +import type { ModelConfig, PromptConfig, RunConfig } from './agent-types.js'; +import type { + ContentGenerator, + ContentGeneratorConfig, +} from '../../core/contentGenerator.js'; + +describe('AgentCore.runInAgentFrames', () => { + // The deferred-approval `respond` callback that AgentCore hands to the + // UI must restore both ALS frames the agent normally runs under, so any + // tool body resumed via approval — including ones that trigger LLM + // calls — sees the agent's ContentGenerator (modalities, auth) and is + // attributed to the agent in token stats. + // + // The reasoning loop uses the same wrap, so anything that breaks here + // also breaks the synchronous path. These tests pin the contract. + + function makeCore(name: string, runtimeView?: RuntimeContentGeneratorView) { + const promptConfig: PromptConfig = { systemPrompt: '' }; + const modelConfig: ModelConfig = { model: 'test-model' }; + const runConfig: RunConfig = { max_turns: 1 }; + return new AgentCore( + name, + {} as unknown as Config, + promptConfig, + modelConfig, + runConfig, + undefined, + undefined, + undefined, + runtimeView, + ); + } + + it('publishes both the runtime view and the agent name when invoked from outside any frame', async () => { + const view: RuntimeContentGeneratorView = { + contentGenerator: { + generateContentStream: () => Promise.resolve(), + } as unknown as ContentGenerator, + contentGeneratorConfig: { + model: 'agent-model', + authType: 'anthropic', + } as ContentGeneratorConfig, + }; + const core = makeCore('image-agent', view); + + let observedView: RuntimeContentGeneratorView | undefined; + let observedName: string | undefined; + await core.runInAgentFrames(async () => { + observedView = getRuntimeContentGenerator(); + observedName = subagentNameContext.getStore(); + }); + + expect(observedView).toBe(view); + expect(observedName).toBe('image-agent'); + }); + + it('restores frames even when called from a fresh async chain (deferred-approval path)', async () => { + // Simulates the UI's async-input handler invoking the captured + // `respond` callback after the reasoning-loop frame has unwound. + // Without `runInAgentFrames` re-entering, the body would see the + // top-level (parent) view. + const view: RuntimeContentGeneratorView = { + contentGenerator: { + generateContentStream: () => Promise.resolve(), + } as unknown as ContentGenerator, + contentGeneratorConfig: { + model: 'agent-model', + authType: 'anthropic', + } as ContentGeneratorConfig, + }; + const core = makeCore('approval-agent', view); + + // Capture a thunk equivalent to the `respond` closure that AgentCore + // emits with TOOL_WAITING_APPROVAL — the wrap is identical. + let capturedRespond: (() => Promise) | undefined; + const onConfirmInvocations: Array<{ + view: RuntimeContentGeneratorView | undefined; + name: string | undefined; + }> = []; + const onConfirm = async () => { + onConfirmInvocations.push({ + view: getRuntimeContentGenerator(), + name: subagentNameContext.getStore(), + }); + }; + + await core.runInAgentFrames(async () => { + // Inside the reasoning-loop frame the agent would build the + // closure that the UI later invokes — same shape as line 938 of + // agent-core.ts. + capturedRespond = () => core.runInAgentFrames(onConfirm); + }); + + // After the loop frame has unwound, neither frame is active. + expect(getRuntimeContentGenerator()).toBeUndefined(); + expect(subagentNameContext.getStore()).toBeUndefined(); + + // Hop to a brand-new microtask chain to be sure no parent ALS frame + // is in scope, then invoke the captured callback. + await new Promise((resolve) => setImmediate(resolve)); + await capturedRespond!(); + + expect(onConfirmInvocations).toHaveLength(1); + expect(onConfirmInvocations[0]!.view).toBe(view); + expect(onConfirmInvocations[0]!.name).toBe('approval-agent'); + }); + + it('still publishes the agent name when no runtime view is set (inheriting agent)', async () => { + const core = makeCore('inherit-agent'); + + let observedView: RuntimeContentGeneratorView | undefined; + let observedName: string | undefined; + await core.runInAgentFrames(async () => { + observedView = getRuntimeContentGenerator(); + observedName = subagentNameContext.getStore(); + }); + + expect(observedView).toBeUndefined(); + expect(observedName).toBe('inherit-agent'); + }); + + it('uses inheritedView for deferred-approval continuation when the agent owns no view', async () => { + // A nested `model: inherit` child under a runtime-view-bearing parent + // owns no view of its own, but its tool bodies (e.g. `read_file` + // checking modalities) need the parent's view. The reasoning loop + // sees it via ALS, but the deferred-approval `respond` callback runs + // from a fresh async chain where that frame is gone — so the agent + // must capture it at emit time and pass it back through. + const parentView: RuntimeContentGeneratorView = { + contentGenerator: { + generateContentStream: () => Promise.resolve(), + } as unknown as ContentGenerator, + contentGeneratorConfig: { + model: 'parent-model', + authType: 'anthropic', + } as ContentGeneratorConfig, + }; + const inheritingCore = makeCore('inherit-agent'); + + let respondClosure: (() => Promise) | undefined; + let observedView: RuntimeContentGeneratorView | undefined; + let observedName: string | undefined; + const onConfirm = async () => { + observedView = getRuntimeContentGenerator(); + observedName = subagentNameContext.getStore(); + }; + + // Simulate the parent's loop frame being live at emit time. + await runWithRuntimeContentGenerator(parentView, async () => { + const inheritedView = getRuntimeContentGenerator(); + respondClosure = () => + inheritingCore.runInAgentFrames(onConfirm, inheritedView); + }); + + // Parent frame is gone; jump to a fresh microtask chain to be sure. + expect(getRuntimeContentGenerator()).toBeUndefined(); + await new Promise((resolve) => setImmediate(resolve)); + + await respondClosure!(); + + expect(observedView).toBe(parentView); + expect(observedName).toBe('inherit-agent'); + }); + + it("prefers the agent's own view over inheritedView when both are present", async () => { + // Defensive: if a future caller wires both, the agent's explicit view + // wins — we never want a captured snapshot to override the agent's + // declared view. + const ownView: RuntimeContentGeneratorView = { + contentGenerator: { + generateContentStream: () => Promise.resolve(), + } as unknown as ContentGenerator, + contentGeneratorConfig: { + model: 'own-model', + authType: 'anthropic', + } as ContentGeneratorConfig, + }; + const otherView: RuntimeContentGeneratorView = { + contentGenerator: { + generateContentStream: () => Promise.resolve(), + } as unknown as ContentGenerator, + contentGeneratorConfig: { + model: 'other-model', + authType: 'openai', + } as ContentGeneratorConfig, + }; + const core = makeCore('own-view-agent', ownView); + + let observed: RuntimeContentGeneratorView | undefined; + await core.runInAgentFrames(async () => { + observed = getRuntimeContentGenerator(); + }, otherView); + + expect(observed).toBe(ownView); + }); +}); diff --git a/packages/core/src/agents/runtime/agent-core.ts b/packages/core/src/agents/runtime/agent-core.ts index bb959f4d0..910e3475d 100644 --- a/packages/core/src/agents/runtime/agent-core.ts +++ b/packages/core/src/agents/runtime/agent-core.ts @@ -19,6 +19,11 @@ import { reportError } from '../../utils/errorReporting.js'; import { subagentNameContext } from '../../utils/subagentNameContext.js'; import type { Config } from '../../config/config.js'; +import { + getRuntimeContentGenerator, + runWithRuntimeContentGenerator, + type RuntimeContentGeneratorView, +} from './agent-context.js'; import { type ToolCallRequestInfo } from '../../core/turn.js'; import { CoreToolScheduler, @@ -183,6 +188,14 @@ export class AgentCore { readonly eventEmitter: AgentEventEmitter; readonly hooks?: AgentHooks; readonly stats = new AgentStatistics(); + /** + * When the agent runs with a model different from the parent session, + * this view is published via AsyncLocalStorage during execution so any + * `Config.getContentGenerator{,Config}()` call inside the run resolves + * to the agent's values — even from tools that captured the parent + * Config at construction. + */ + readonly runtimeView?: RuntimeContentGeneratorView; // Observable state lives on Core (not a wrapper) so headless and // background agents can be observed with the same accessors as @@ -236,6 +249,7 @@ export class AgentCore { toolConfig?: ToolConfig, eventEmitter?: AgentEventEmitter, hooks?: AgentHooks, + runtimeView?: RuntimeContentGeneratorView, ) { const randomPart = Math.random().toString(36).slice(2, 8); this.subagentId = `${name}-${randomPart}`; @@ -247,6 +261,7 @@ export class AgentCore { this.toolConfig = toolConfig; this.eventEmitter = eventEmitter ?? new AgentEventEmitter(); this.hooks = hooks; + this.runtimeView = runtimeView; this.setupStateListeners(); } @@ -433,21 +448,66 @@ export class AgentCore { abortController: AbortController, options?: ReasoningLoopOptions, ): Promise { - // Tag every API call emitted from this loop with the owning subagent's - // name so the `/stats` panel can attribute tokens/requests to the - // originating subagent. The store is read inside - // `LoggingContentGenerator` via `subagentNameContext.getStore()`. - return subagentNameContext.run(this.name, () => + const inner = () => this._runReasoningLoopInner( chat, initialMessages, toolsList, abortController, options, - ), + ); + return this.runInAgentFrames(inner); + } + + /** + * Run `fn` inside both ALS frames this agent owns: + * 1. {@link subagentNameContext} so token-attribution code resolves to + * this agent's name. + * 2. The per-agent runtime ContentGenerator view (when set) so + * `Config.getContentGenerator{,Config}()` calls inside resolve to + * the agent rather than to the parent Config tools captured at + * construction time. + * + * Used both around the reasoning loop and around the deferred-approval + * `onConfirm` continuation — the latter runs from the parent UI's input + * handler, on a different async chain than the loop, so without this + * re-entry the resumed tool body would fall back to the parent's view + * and mis-attribute its tokens. + * + * `inheritedView` lets a caller pass an ambient view captured earlier + * (e.g. at approval-emit time, when the parent's ALS frame is still + * live) for inheriting agents that own no view themselves. Without it, + * a nested `model: inherit` agent under a runtime-view-bearing parent + * would lose that view across the deferred-approval boundary, since + * the UI invokes `respond` from a fresh async chain where the parent's + * ALS frame is gone. + * + * Exposed (rather than inlined twice) so the contract stays testable in + * isolation; see `agent-core.test.ts`. + */ + runInAgentFrames( + fn: () => Promise, + inheritedView?: RuntimeContentGeneratorView, + ): Promise { + return subagentNameContext.run(this.name, () => + this.withRuntimeView(fn, inheritedView), ); } + /** + * Wraps `fn` in the effective runtime view: this agent's own view if + * set, else `inheritedView` if the caller captured one. Internal — + * public callers should use {@link runInAgentFrames}, which also + * restores the subagent-name frame. + */ + private withRuntimeView( + fn: () => Promise, + inheritedView?: RuntimeContentGeneratorView, + ): Promise { + const view = this.runtimeView ?? inheritedView; + return view ? runWithRuntimeContentGenerator(view, fn) : fn(); + } + private async _runReasoningLoopInner( chat: GeminiChat, initialMessages: Content[], @@ -901,6 +961,12 @@ export class AgentCore { try { const { confirmationDetails } = waiting; const { onConfirm: _onConfirm, ...rest } = confirmationDetails; + // Snapshot the ambient runtime view here, while the loop frame + // is still live. For inheriting agents (no own runtimeView) + // this captures the parent's view so the deferred-approval + // continuation — invoked later from the UI's async chain — can + // restore it. See `runInAgentFrames` for the wiring. + const inheritedView = getRuntimeContentGenerator(); this.eventEmitter?.emit(AgentEventType.TOOL_WAITING_APPROVAL, { subagentId: this.subagentId, round: currentRound, @@ -919,7 +985,14 @@ export class AgentCore { ) => { if (responded.has(waiting.request.callId)) return; responded.add(waiting.request.callId); - await waiting.confirmationDetails.onConfirm(outcome, payload); + // UI invokes this from its own async chain (outside the + // reasoning-loop ALS frames), so re-enter both the agent's + // runtime view AND its name context before the resumed + // tool body runs. See `runInAgentFrames` for rationale. + await this.runInAgentFrames( + () => waiting.confirmationDetails.onConfirm(outcome, payload), + inheritedView, + ); }, timestamp: Date.now(), }); diff --git a/packages/core/src/agents/runtime/agent-headless.ts b/packages/core/src/agents/runtime/agent-headless.ts index 8c85c46aa..49dc93022 100644 --- a/packages/core/src/agents/runtime/agent-headless.ts +++ b/packages/core/src/agents/runtime/agent-headless.ts @@ -16,6 +16,7 @@ import type { Content } from '@google/genai'; import type { Config } from '../../config/config.js'; +import type { RuntimeContentGeneratorView } from './agent-context.js'; import { createDebugLogger } from '../../utils/debugLogger.js'; import type { AgentEventEmitter, @@ -164,6 +165,7 @@ export class AgentHeadless { toolConfig?: ToolConfig, eventEmitter?: AgentEventEmitter, hooks?: AgentHooks, + runtimeView?: RuntimeContentGeneratorView, ): Promise { const core = new AgentCore( name, @@ -174,6 +176,7 @@ export class AgentHeadless { toolConfig, eventEmitter, hooks, + runtimeView, ); return new AgentHeadless(core); } diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index c1d9ef2a9..55b5d34f7 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -2197,4 +2197,102 @@ describe('Model Switching and Config Updates', () => { expect(mockHasHooksForEvent).toHaveBeenCalledWith('Stop'); }); }); + + describe('runtime ContentGenerator view (AsyncLocalStorage)', () => { + // The Config getters consult the per-run ALS view published by the + // agent runtime when a sub-agent runs on a different model than the + // parent. These tests pin that integration: tools that captured the + // parent Config at construction must still resolve to the agent's + // values when called inside the agent's runtime frame. + function setInstanceFields( + config: Config, + contentGenerator: ContentGenerator, + generatorConfig: ContentGeneratorConfig, + ): void { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (config as any).contentGenerator = contentGenerator; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (config as any).contentGeneratorConfig = generatorConfig; + } + + it('resolves getters to the runtime view inside the frame, instance fields outside', async () => { + const { runWithRuntimeContentGenerator } = await import( + '../agents/runtime/agent-context.js' + ); + const config = new Config(baseParams); + const parentGenerator = { + generateContentStream: vi.fn(), + } as unknown as ContentGenerator; + const parentGeneratorConfig: ContentGeneratorConfig = { + model: 'parent-model', + authType: AuthType.QWEN_OAUTH, + apiKey: 'parent-key', + }; + setInstanceFields(config, parentGenerator, parentGeneratorConfig); + + const agentGenerator = { + generateContentStream: vi.fn(), + } as unknown as ContentGenerator; + const agentGeneratorConfig: ContentGeneratorConfig = { + model: 'agent-model', + authType: AuthType.USE_OPENAI, + apiKey: 'agent-key', + }; + + // Outside the frame, getters resolve to the parent's instance fields. + expect(config.getContentGenerator()).toBe(parentGenerator); + expect(config.getContentGeneratorConfig()).toBe(parentGeneratorConfig); + expect(config.getModel()).toBe('parent-model'); + expect(config.getAuthType()).toBe(AuthType.QWEN_OAUTH); + + // Inside the frame, every getter resolves to the agent's view. + await runWithRuntimeContentGenerator( + { + contentGenerator: agentGenerator, + contentGeneratorConfig: agentGeneratorConfig, + }, + async () => { + expect(config.getContentGenerator()).toBe(agentGenerator); + expect(config.getContentGeneratorConfig()).toBe(agentGeneratorConfig); + expect(config.getModel()).toBe('agent-model'); + expect(config.getAuthType()).toBe(AuthType.USE_OPENAI); + }, + ); + + // Frame exit restores resolution to the parent's instance fields. + expect(config.getContentGenerator()).toBe(parentGenerator); + expect(config.getModel()).toBe('parent-model'); + }); + + it('falls back to the parent model id when the runtime view config has no model', async () => { + const { runWithRuntimeContentGenerator } = await import( + '../agents/runtime/agent-context.js' + ); + const config = new Config(baseParams); + setInstanceFields( + config, + { generateContentStream: vi.fn() } as unknown as ContentGenerator, + { + model: 'parent-model', + authType: AuthType.QWEN_OAUTH, + } as ContentGeneratorConfig, + ); + + await runWithRuntimeContentGenerator( + { + contentGenerator: { + generateContentStream: vi.fn(), + } as unknown as ContentGenerator, + contentGeneratorConfig: { + model: '', + authType: AuthType.USE_OPENAI, + } as ContentGeneratorConfig, + }, + async () => { + // Empty model on the runtime view falls through to modelsConfig. + expect(config.getModel()).toBe(baseParams.model); + }, + ); + }); + }); }); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index f635ea1b8..3a104fb11 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -33,6 +33,7 @@ import { createContentGenerator, resolveContentGeneratorConfigWithSources, } from '../core/contentGenerator.js'; +import { getRuntimeContentGenerator } from '../agents/runtime/agent-context.js'; // Services import { FileDiscoveryService } from '../services/fileDiscoveryService.js'; @@ -1285,7 +1286,9 @@ export class Config { } getContentGenerator(): ContentGenerator { - return this.contentGenerator; + return ( + getRuntimeContentGenerator()?.contentGenerator ?? this.contentGenerator + ); } /** @@ -1471,7 +1474,10 @@ export class Config { } getContentGeneratorConfig(): ContentGeneratorConfig { - return this.contentGeneratorConfig; + return ( + getRuntimeContentGenerator()?.contentGeneratorConfig ?? + this.contentGeneratorConfig + ); } getContentGeneratorConfigSources(): ContentGeneratorConfigSources { @@ -1487,7 +1493,9 @@ export class Config { } getModel(): string { - return this.contentGeneratorConfig?.model || this.modelsConfig.getModel(); + return ( + this.getContentGeneratorConfig()?.model || this.modelsConfig.getModel() + ); } onModelChange(listener: (model: string) => void): () => void { @@ -2455,7 +2463,7 @@ export class Config { } getAuthType(): AuthType | undefined { - return this.contentGeneratorConfig?.authType; + return this.getContentGeneratorConfig()?.authType; } getCliVersion(): string | undefined { diff --git a/packages/core/src/models/content-generator-config.test.ts b/packages/core/src/models/content-generator-config.test.ts index e090d0daf..68c6df0c2 100644 --- a/packages/core/src/models/content-generator-config.test.ts +++ b/packages/core/src/models/content-generator-config.test.ts @@ -7,12 +7,23 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { buildAgentContentGeneratorConfig, + createRuntimeContentGeneratorView, resolveCredentialField, } from './content-generator-config.js'; +import { createContentGenerator } from '../core/contentGenerator.js'; import type { ContentGeneratorConfig } from '../core/contentGenerator.js'; import type { Config } from '../config/config.js'; import type { ResolvedModelConfig } from './types.js'; +vi.mock('../core/contentGenerator.js', async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createContentGenerator: vi.fn(), + }; +}); + function createMockConfig( parentConfig: ContentGeneratorConfig, resolvedModel?: ResolvedModelConfig, @@ -211,6 +222,44 @@ describe('buildAgentContentGeneratorConfig', () => { }); }); +describe('createRuntimeContentGeneratorView', () => { + const parentConfig: ContentGeneratorConfig = { + model: 'parent-model', + authType: 'openai' as ContentGeneratorConfig['authType'], + apiKey: 'parent-key', + baseUrl: 'https://parent.example.com', + }; + + beforeEach(() => { + vi.mocked(createContentGenerator).mockReset(); + }); + + it('should bind the new ContentGenerator to contentGeneratorOwner, not base', async () => { + const baseConfig = createMockConfig(parentConfig); + // Distinct instance — represents the per-agent override Config. + const ownerConfig = createMockConfig(parentConfig); + const fakeGenerator = { generateContentStream: vi.fn() }; + vi.mocked(createContentGenerator).mockResolvedValueOnce( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + fakeGenerator as any, + ); + + const view = await createRuntimeContentGeneratorView( + baseConfig, + ownerConfig, + 'custom-model', + { authType: 'openai' }, + ); + + expect(createContentGenerator).toHaveBeenCalledTimes(1); + const [, ownerArg] = vi.mocked(createContentGenerator).mock.calls[0]; + expect(ownerArg).toBe(ownerConfig); + expect(ownerArg).not.toBe(baseConfig); + expect(view.contentGenerator).toBe(fakeGenerator); + expect(view.contentGeneratorConfig.model).toBe('custom-model'); + }); +}); + describe('resolveCredentialField', () => { afterEach(() => { vi.unstubAllEnvs(); diff --git a/packages/core/src/models/content-generator-config.ts b/packages/core/src/models/content-generator-config.ts index c2ef1a242..2dca50dbd 100644 --- a/packages/core/src/models/content-generator-config.ts +++ b/packages/core/src/models/content-generator-config.ts @@ -14,9 +14,11 @@ import type { Config } from '../config/config.js'; import { + createContentGenerator, type AuthType, type ContentGeneratorConfig, } from '../core/contentGenerator.js'; +import type { RuntimeContentGeneratorView } from '../agents/runtime/agent-context.js'; import { AUTH_ENV_MAPPINGS, MODEL_GENERATION_CONFIG_FIELDS, @@ -97,6 +99,34 @@ export function buildAgentContentGeneratorConfig( return nextConfig; } +/** + * Compose `buildAgentContentGeneratorConfig` + `createContentGenerator` into + * a single {@link RuntimeContentGeneratorView}. Both InProcessBackend and + * SubagentManager need the same three-step recipe; this helper centralizes + * it so the two paths can't drift. + * + * `contentGeneratorOwner` is the Config instance the new ContentGenerator + * should bind to for cwd / workspace / telemetry purposes — typically the + * per-agent override Config when one exists, or the parent Config otherwise. + */ +export async function createRuntimeContentGeneratorView( + base: Config, + contentGeneratorOwner: Config, + modelId: string | undefined, + authOverrides: AuthOverrides, +): Promise { + const contentGeneratorConfig = buildAgentContentGeneratorConfig( + base, + modelId, + authOverrides, + ); + const contentGenerator = await createContentGenerator( + contentGeneratorConfig, + contentGeneratorOwner, + ); + return { contentGenerator, contentGeneratorConfig }; +} + function applyResolvedModelConfig( targetConfig: ContentGeneratorConfig, resolvedModel: ResolvedModelConfig, diff --git a/packages/core/src/models/modelRegistry.test.ts b/packages/core/src/models/modelRegistry.test.ts index 9005dd52a..f97441043 100644 --- a/packages/core/src/models/modelRegistry.test.ts +++ b/packages/core/src/models/modelRegistry.test.ts @@ -162,6 +162,61 @@ describe('ModelRegistry', () => { }); }); + describe('modalities auto-fill', () => { + // Sub-agents that read straight from the registry (e.g. via + // getResolvedModel) need the registry to populate modalities for them; + // otherwise they inherit the parent session's modalities and fail the + // image/pdf/video gates set on tools like ReadFile. + it('populates modalities from the model name when not provided', () => { + const registry = new ModelRegistry({ + openai: [ + { + id: 'gpt-4-turbo', + name: 'GPT-4 Turbo', + baseUrl: 'https://api.openai.com/v1', + generationConfig: {}, + }, + ], + }); + + const model = registry.getModel(AuthType.USE_OPENAI, 'gpt-4-turbo'); + expect(model?.generationConfig.modalities).toEqual({ image: true }); + }); + + it('preserves caller-provided modalities verbatim', () => { + const explicitModalities = { image: true, pdf: true }; + const registry = new ModelRegistry({ + openai: [ + { + id: 'gpt-4-turbo', + name: 'GPT-4 Turbo', + baseUrl: 'https://api.openai.com/v1', + generationConfig: { modalities: explicitModalities }, + }, + ], + }); + + const model = registry.getModel(AuthType.USE_OPENAI, 'gpt-4-turbo'); + expect(model?.generationConfig.modalities).toEqual(explicitModalities); + }); + + it('returns text-only ({}) for models with no multimodal default', () => { + const registry = new ModelRegistry({ + openai: [ + { + id: 'qwen3-coder-plus', + name: 'Qwen3 Coder Plus', + baseUrl: 'https://example.invalid', + generationConfig: {}, + }, + ], + }); + + const model = registry.getModel(AuthType.USE_OPENAI, 'qwen3-coder-plus'); + expect(model?.generationConfig.modalities).toEqual({}); + }); + }); + describe('hasModel', () => { let registry: ModelRegistry; diff --git a/packages/core/src/models/modelRegistry.ts b/packages/core/src/models/modelRegistry.ts index c2815fb32..d580a1226 100644 --- a/packages/core/src/models/modelRegistry.ts +++ b/packages/core/src/models/modelRegistry.ts @@ -125,8 +125,9 @@ export class ModelRegistry { isVision: model.capabilities?.vision ?? false, contextWindowSize: model.generationConfig.contextWindowSize ?? tokenLimit(model.id), - modalities: - model.generationConfig.modalities ?? defaultModalities(model.id), + // `modalities` is auto-filled in `resolveModelConfig`, so it is + // always defined on `ResolvedModelConfig` — no fallback needed here. + modalities: model.generationConfig.modalities, baseUrl: model.baseUrl, envKey: model.envKey, })); @@ -176,12 +177,21 @@ export class ModelRegistry { ): ResolvedModelConfig { this.validateModelConfig(config, authType); + const generationConfig = { ...(config.generationConfig ?? {}) }; + // Auto-fill modalities from the model name when the provider didn't set + // them explicitly. Without this, downstream consumers that read straight + // from the registry (e.g. sub-agents via getResolvedModel) would inherit + // the parent session's modalities instead of the agent's own. + if (generationConfig.modalities === undefined) { + generationConfig.modalities = defaultModalities(config.id); + } + return { ...config, authType, name: config.name || config.id, baseUrl: config.baseUrl || this.getDefaultBaseUrl(authType), - generationConfig: config.generationConfig ?? {}, + generationConfig, capabilities: config.capabilities || {}, }; } diff --git a/packages/core/src/subagents/subagent-manager-override.test.ts b/packages/core/src/subagents/subagent-manager-override.test.ts index ec795d667..53e2428aa 100644 --- a/packages/core/src/subagents/subagent-manager-override.test.ts +++ b/packages/core/src/subagents/subagent-manager-override.test.ts @@ -4,58 +4,27 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi } from 'vitest'; +import { describe, it, expect } from 'vitest'; import { Config, ApprovalMode } from '../config/config.js'; import { SubagentManager } from './subagent-manager.js'; -import type { SubagentConfig } from './types.js'; import { ToolNames } from '../tools/tool-names.js'; import { EditTool } from '../tools/edit.js'; import { ReadFileTool } from '../tools/read-file.js'; import { createApprovalModeOverride } from '../tools/agent/agent.js'; -// The non-inherit (explicit-model) branch in maybeOverrideContentGenerator -// builds a fresh ContentGenerator. We don't want the test to actually -// reach the OpenAI / Anthropic SDK — replacing the factory with a stub -// is enough to exercise the code path. -vi.mock('../core/contentGenerator.js', async () => { - const actual = await vi.importActual< - typeof import('../core/contentGenerator.js') - >('../core/contentGenerator.js'); - return { - ...actual, - createContentGenerator: vi.fn().mockResolvedValue({ - generateContent: vi.fn(), - generateContentStream: vi.fn(), - }), - }; -}); - -vi.mock('../models/content-generator-config.js', async () => { - const actual = await vi.importActual< - typeof import('../models/content-generator-config.js') - >('../models/content-generator-config.js'); - return { - ...actual, - buildAgentContentGeneratorConfig: vi.fn().mockReturnValue({ - model: 'override-model', - authType: 'openai', - apiKey: 'override-key', - }), - }; -}); - /** * Companion to `tools/agent/agent-override.test.ts`. Same regression: * Object.create(parent) by itself is not enough to isolate a subagent's * core tools from the parent's bound `EditTool` / `WriteFileTool` / - * `ReadFileTool`. The subagent path that flows through - * `SubagentManager.maybeOverrideContentGenerator` must rebuild the - * tool registry on the override Config so bound tools resolve - * `this.config` to the subagent rather than the parent — otherwise - * mutations executed via the bound tool reach the parent's - * FileReadCache and silently weaken prior-read enforcement. + * `ReadFileTool`. The subagent path (which flows through + * `SubagentManager.createAgentHeadless` → + * `buildSubagentContextOverride`) must rebuild the tool registry on + * the override Config so bound tools resolve `this.config` to the + * subagent rather than the parent — otherwise mutations executed via + * the bound tool reach the parent's FileReadCache and silently weaken + * prior-read enforcement. */ -describe('SubagentManager.maybeOverrideContentGenerator bound-tool isolation', () => { +describe('SubagentManager.buildSubagentContextOverride bound-tool isolation', () => { // Bare mode keeps the registry small (ReadFile / Edit / Shell only) and // avoids needing extra setup for optional tools. const baseParams = { @@ -70,23 +39,19 @@ describe('SubagentManager.maybeOverrideContentGenerator bound-tool isolation', ( // The method is `private`. Cast via `unknown` to invoke it directly — // testing through the public `createAgentHeadless` pathway would also // work but pulls in a much larger graph (file IO, hooks, etc.). - function callMaybeOverride( + function callBuildOverride( manager: SubagentManager, - config: SubagentConfig, base: Config, ): Promise { const fn = ( manager as unknown as { - maybeOverrideContentGenerator: ( - c: SubagentConfig, - b: Config, - ) => Promise; + buildSubagentContextOverride: (b: Config) => Promise; } - ).maybeOverrideContentGenerator.bind(manager); - return fn(config, base); + ).buildSubagentContextOverride.bind(manager); + return fn(base); } - it('inherits branch: returns a Config whose registry is distinct from the parent and binds Edit/Read to the override', async () => { + it('returns a Config whose registry is distinct from the parent and binds Edit/Read to the override', async () => { const parent = new Config(baseParams); const parentRegistry = await parent.createToolRegistry(undefined, { skipDiscovery: true, @@ -96,16 +61,7 @@ describe('SubagentManager.maybeOverrideContentGenerator bound-tool isolation', ( const manager = new SubagentManager(parent); - const subagentConfig: SubagentConfig = { - name: 'inheriting-agent', - description: 'Inherits parent model', - systemPrompt: 'You are a helpful assistant.', - level: 'project', - filePath: '/test/project/.qwen/agents/inheriting-agent.md', - // model omitted -> inherits=true branch - }; - - const child = await callMaybeOverride(manager, subagentConfig, parent); + const child = await callBuildOverride(manager, parent); expect(child).not.toBe(parent); expect(child.getToolRegistry()).not.toBe(parentRegistry); @@ -130,7 +86,7 @@ describe('SubagentManager.maybeOverrideContentGenerator bound-tool isolation', ( expect(boundConfig.getFileReadCache()).not.toBe(parent.getFileReadCache()); }); - it('inherits branch: parent and child caches are independent', async () => { + it('parent and child caches are independent', async () => { const parent = new Config(baseParams); const parentRegistry = await parent.createToolRegistry(undefined, { skipDiscovery: true, @@ -139,15 +95,8 @@ describe('SubagentManager.maybeOverrideContentGenerator bound-tool isolation', ( (parent as any).toolRegistry = parentRegistry; const manager = new SubagentManager(parent); - const subagentConfig: SubagentConfig = { - name: 'inheriting-agent', - description: 'Inherits parent model', - systemPrompt: 'You are a helpful assistant.', - level: 'project', - filePath: '/test/project/.qwen/agents/inheriting-agent.md', - }; - const child = await callMaybeOverride(manager, subagentConfig, parent); + const child = await callBuildOverride(manager, parent); // Record a read on parent. Child must not see it. const fakeStats = { @@ -166,12 +115,12 @@ describe('SubagentManager.maybeOverrideContentGenerator bound-tool isolation', ( expect(child.getFileReadCache().size()).toBe(0); }); - it('inherits branch: skips rebuild and inherits registry via prototype when the base already has its own registry (real-world chained-override case)', async () => { + it('skips rebuild and inherits registry via prototype when an upstream wrapper has already rebuilt the registry (real-world chained-override case)', async () => { // This mirrors the real-world flow: agent.ts wraps the parent in // `createApprovalModeOverride` (which builds R1 on the wrapper), // then passes that wrapper — sometimes wrapped one more level in // `bgConfig = Object.create(agentConfig)` for the background path — - // through `createAgentHeadless` → `maybeOverrideContentGenerator`. + // through `createAgentHeadless` → `buildSubagentContextOverride`. // We do NOT want the second layer to build a redundant R2 — that // would (a) waste work, (b) leak listeners on every later // AgentTool/SkillTool factory invocation, and (c) split the cache @@ -205,19 +154,8 @@ describe('SubagentManager.maybeOverrideContentGenerator bound-tool isolation', ( bgWrapper.getShouldAvoidPermissionPrompts = () => true; const manager = new SubagentManager(parent); - const subagentConfig: SubagentConfig = { - name: 'inheriting-agent', - description: 'Inherits parent model', - systemPrompt: 'You are a helpful assistant.', - level: 'project', - filePath: '/test/project/.qwen/agents/inheriting-agent.md', - }; - const child = await callMaybeOverride( - manager, - subagentConfig, - bgWrapper as Config, - ); + const child = await callBuildOverride(manager, bgWrapper as Config); // child is still a distinct instance (Object.create) so the // FileReadCache lazy-init still works, but its registry must @@ -238,113 +176,7 @@ describe('SubagentManager.maybeOverrideContentGenerator bound-tool isolation', ( expect((childEdit as any).config).toBe(upstreamWrapper); }); - it('non-inherit branch (explicit-model selector): rebuilds registry and binds Edit/Read to the override Config', async () => { - // The non-inherit branch swaps the ContentGenerator (so the - // subagent talks to the model the selector requests). It must - // ALSO rebuild the tool registry — without that step explicit-model - // subagents would still resolve their core tools' `this.config` to - // the parent and read the parent's FileReadCache. - const parent = new Config(baseParams); - // Even though bare mode skips most tools, the non-inherit branch - // requires getContentGeneratorConfig() to return something for the - // authType fallback. Stub it minimally. - vi.spyOn(parent, 'getContentGeneratorConfig').mockReturnValue({ - model: 'parent-model', - authType: 'openai', - // eslint-disable-next-line @typescript-eslint/no-explicit-any - } as any); - const parentRegistry = await parent.createToolRegistry(undefined, { - skipDiscovery: true, - }); - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (parent as any).toolRegistry = parentRegistry; - - const manager = new SubagentManager(parent); - const subagentConfig: SubagentConfig = { - name: 'explicit-model-agent', - description: 'Uses an explicit model selector', - systemPrompt: 'You are a helpful assistant.', - level: 'project', - filePath: '/test/project/.qwen/agents/explicit-model-agent.md', - // Bare model ID -> non-inherits branch (parses to {modelId, - // inherits:false}). - model: 'override-model', - }; - - const child = await callMaybeOverride(manager, subagentConfig, parent); - - expect(child).not.toBe(parent); - expect(child.getToolRegistry()).not.toBe(parentRegistry); - - const childEdit = await child.getToolRegistry().ensureTool(ToolNames.EDIT); - const childRead = await child - .getToolRegistry() - .ensureTool(ToolNames.READ_FILE); - - expect(childEdit).toBeInstanceOf(EditTool); - expect(childRead).toBeInstanceOf(ReadFileTool); - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - expect((childEdit as any).config).toBe(child); - // eslint-disable-next-line @typescript-eslint/no-explicit-any - expect((childRead as any).config).toBe(child); - - // The bound EditTool's FileReadCache must be the override's, not - // the parent's. - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const boundConfig = (childEdit as any).config as Config; - expect(boundConfig.getFileReadCache()).toBe(child.getFileReadCache()); - expect(boundConfig.getFileReadCache()).not.toBe(parent.getFileReadCache()); - }); - - it('non-inherit branch: skips rebuild when an upstream wrapper has already rebuilt the registry', async () => { - const parent = new Config(baseParams); - vi.spyOn(parent, 'getContentGeneratorConfig').mockReturnValue({ - model: 'parent-model', - authType: 'openai', - // eslint-disable-next-line @typescript-eslint/no-explicit-any - } as any); - const parentRegistry = await parent.createToolRegistry(undefined, { - skipDiscovery: true, - }); - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (parent as any).toolRegistry = parentRegistry; - - const upstreamWrapper = await createApprovalModeOverride( - parent, - ApprovalMode.AUTO_EDIT, - ); - const upstreamRegistry = upstreamWrapper.getToolRegistry(); - - const manager = new SubagentManager(parent); - const subagentConfig: SubagentConfig = { - name: 'explicit-model-agent', - description: 'Uses an explicit model selector', - systemPrompt: 'You are a helpful assistant.', - level: 'project', - filePath: '/test/project/.qwen/agents/explicit-model-agent.md', - model: 'override-model', - }; - - const child = await callMaybeOverride( - manager, - subagentConfig, - upstreamWrapper, - ); - - // Upstream rebuild was detected via the symbol marker, so the - // override has no own registry — it inherits via the prototype. - expect(child.getToolRegistry()).toBe(upstreamRegistry); - - // Bound tools resolve to upstreamWrapper, not the second-layer - // child — same as the inherits branch's chained-override case. - const childEdit = await child.getToolRegistry().ensureTool(ToolNames.EDIT); - expect(childEdit).toBeInstanceOf(EditTool); - // eslint-disable-next-line @typescript-eslint/no-explicit-any - expect((childEdit as any).config).toBe(upstreamWrapper); - }); - - it('inherits branch: the override approval mode (inherited via prototype) still resolves via the override Config', async () => { + it('the override approval mode (inherited via prototype) still resolves via the override Config', async () => { const parent = new Config(baseParams); const parentRegistry = await parent.createToolRegistry(undefined, { skipDiscovery: true, @@ -353,15 +185,8 @@ describe('SubagentManager.maybeOverrideContentGenerator bound-tool isolation', ( (parent as any).toolRegistry = parentRegistry; const manager = new SubagentManager(parent); - const subagentConfig: SubagentConfig = { - name: 'inheriting-agent', - description: 'Inherits parent model', - systemPrompt: 'You are a helpful assistant.', - level: 'project', - filePath: '/test/project/.qwen/agents/inheriting-agent.md', - }; - const child = await callMaybeOverride(manager, subagentConfig, parent); + const child = await callBuildOverride(manager, parent); // Child has no own getApprovalMode; falls through prototype to parent. // Verify mutating parent's mode via setter is observed by child. diff --git a/packages/core/src/subagents/subagent-manager.test.ts b/packages/core/src/subagents/subagent-manager.test.ts index 06c37672b..90b843e14 100644 --- a/packages/core/src/subagents/subagent-manager.test.ts +++ b/packages/core/src/subagents/subagent-manager.test.ts @@ -48,6 +48,28 @@ vi.mock('../agents/runtime/agent-headless.js', () => ({ ContextState: class {}, })); +// Mirrors the positional AgentHeadless.create parameters so tests can +// destructure by name instead of indexing — adding new parameters can't +// silently shift assertions onto the wrong slot. +function destructureAgentHeadlessCall(call: unknown[]) { + return { + name: call[0] as string, + runtimeContext: call[1], + promptConfig: call[2], + modelConfig: call[3], + runConfig: call[4], + toolConfig: call[5], + eventEmitter: call[6], + hooks: call[7], + runtimeView: call[8] as + | { + contentGenerator: unknown; + contentGeneratorConfig: { authType?: string; model?: string }; + } + | undefined, + }; +} + // Mock createContentGenerator for model override tests const mockCreateContentGenerator = vi.hoisted(() => vi.fn()); vi.mock('../core/contentGenerator.js', async (importOriginal) => { @@ -72,7 +94,7 @@ describe('SubagentManager', () => { { name: 'write_file', displayName: 'Write File' }, { name: 'grep', displayName: 'Search Files' }, ]), - // `maybeOverrideContentGenerator` now rebuilds the tool registry on + // `buildSubagentContextOverride` now rebuilds the tool registry on // its override and copies discovered tools from this parent // registry. The real implementation iterates `source.tools.values()`, // so the stub needs a `tools` Map to avoid a TypeError. @@ -1446,9 +1468,12 @@ System prompt 3`); await manager.createAgentHeadless(config, mockConfig); + // Owner is the runtimeContext passed to createAgentHeadless — assert + // the exact instance so a regression that swaps in a different Config + // (e.g. the override) gets caught. expect(mockCreateContentGenerator).toHaveBeenCalledWith( expect.objectContaining({ model: 'custom-model' }), - expect.anything(), + mockConfig, ); }); @@ -1462,7 +1487,7 @@ System prompt 3`); model: 'claude-sonnet', authType: 'anthropic', }), - expect.anything(), + mockConfig, ); }); @@ -1480,16 +1505,23 @@ System prompt 3`); expect(mockCreateContentGenerator).not.toHaveBeenCalled(); }); - it('should pass the overridden Config to AgentHeadless.create', async () => { + it('should pass the agent runtimeView to AgentHeadless.create', async () => { const config = { ...agentConfig, model: 'custom-model' }; const fakeGenerator = { generateContentStream: vi.fn() }; mockCreateContentGenerator.mockResolvedValue(fakeGenerator); await manager.createAgentHeadless(config, mockConfig); - const passedConfig = mockAgentHeadlessCreate.mock.calls[0][1]; - expect(passedConfig.getContentGenerator()).toBe(fakeGenerator); - expect(passedConfig.getModel()).toBe('custom-model'); + const { runtimeContext, runtimeView } = destructureAgentHeadlessCall( + mockAgentHeadlessCreate.mock.calls[0], + ); + // Subagents always get an `Object.create(parent)` wrapper for + // FileReadCache isolation — distinct instance, prototype === parent. + expect(runtimeContext).not.toBe(mockConfig); + expect(Object.getPrototypeOf(runtimeContext)).toBe(mockConfig); + expect(runtimeView).toBeDefined(); + expect(runtimeView!.contentGenerator).toBe(fakeGenerator); + expect(runtimeView!.contentGeneratorConfig.model).toBe('custom-model'); }); }); }); diff --git a/packages/core/src/subagents/subagent-manager.ts b/packages/core/src/subagents/subagent-manager.ts index d18b519e5..860aabeb3 100644 --- a/packages/core/src/subagents/subagent-manager.ts +++ b/packages/core/src/subagents/subagent-manager.ts @@ -35,13 +35,8 @@ import type { } from '../agents/runtime/agent-events.js'; import type { Config } from '../config/config.js'; import { APPROVAL_MODES } from '../config/config.js'; -import { - type AuthType, - type ContentGenerator, - type ContentGeneratorConfig, - createContentGenerator, -} from '../core/contentGenerator.js'; -import { buildAgentContentGeneratorConfig } from '../models/content-generator-config.js'; +import type { RuntimeContentGeneratorView } from '../agents/runtime/agent-context.js'; +import { createRuntimeContentGeneratorView } from '../models/content-generator-config.js'; import { createDebugLogger } from '../utils/debugLogger.js'; import { normalizeContent } from '../utils/textUtils.js'; import { parseSubagentModelSelection } from './model-selection.js'; @@ -660,22 +655,27 @@ export class SubagentManager { options?.toolConfigOverride ?? runtimeConfig.toolConfig; // When the model selector specifies a different provider, build a - // per-agent Config with a dedicated ContentGenerator so the subagent - // talks to the right API without affecting the parent process. - const agentContext = await this.maybeOverrideContentGenerator( + // dedicated ContentGenerator + view so the subagent talks to the + // right API without affecting the parent process. The view is + // applied via AsyncLocalStorage when the agent runs. + const runtimeView = await this.buildRuntimeContentGeneratorView( config, runtimeContext, ); + const subagentContext = + await this.buildSubagentContextOverride(runtimeContext); + return await AgentHeadless.create( config.name, - agentContext, + subagentContext, promptConfig, modelConfig, runConfig, toolConfig, options?.eventEmitter, options?.hooks, + runtimeView, ); } catch (error) { if (error instanceof Error) { @@ -690,53 +690,55 @@ export class SubagentManager { } /** - * When a subagent's model selector specifies a model (bare ID or - * authType-prefixed), build a Config override with a dedicated - * ContentGenerator so the model actually reaches the API. - * For inherit selectors we still build a thin Object.create - * override so the subagent gets an isolated FileReadCache via - * the per-Config own-property machinery — returning `base` - * directly would let the subagent share the parent's read entries - * and silently weaken prior-read enforcement on its mutation - * paths. + * Build the per-subagent Config override used as the AgentHeadless + * runtime context. The override is a thin prototype-delegation wrapper + * (`Object.create(runtimeContext)`): no method changes, but a distinct + * instance triggers the lazy own-property init in + * `Config.getFileReadCache()` so the subagent gets its own cache + * rather than inheriting the parent's recorded reads — which would + * silently weaken prior-read enforcement on its mutation paths. + * + * The tool registry is also rebuilt on the override so `EditTool` / + * `WriteFileTool` / `ReadFileTool` resolve `this.config` to the + * subagent — without that step, the parent's cached tool instances + * still reach the parent's FileReadCache. The rebuild is skipped when + * a wrapper above `runtimeContext` already rebuilt one (typically + * `agent.ts:createApprovalModeOverride`, which marks itself via a + * Symbol-keyed flag — Symbol lookup walks the prototype chain, so + * this also catches wrapper-on-wrapper layering like + * `bgConfig = Object.create(agentConfig)` from the background path). + * Rebuilding twice would waste work, leak listeners on shared + * managers, and split caches across registry layers. */ - private async maybeOverrideContentGenerator( + private async buildSubagentContextOverride( + runtimeContext: Config, + ): Promise { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const subagentContext = Object.create(runtimeContext) as any as Config; + if (!hasRebuiltToolRegistry(runtimeContext)) { + await rebuildToolRegistryOnOverride(subagentContext, runtimeContext); + } + return subagentContext; + } + + /** + * When a subagent's model selector specifies a model (bare ID or + * authType-prefixed), build a dedicated ContentGenerator and the view + * the agent runtime should publish via AsyncLocalStorage during the + * run. Returns `undefined` for inherit selectors (no override needed). + * + * FileReadCache isolation and tool-registry rebuilding are handled + * separately in {@link buildSubagentContextOverride} — every subagent + * (inherit or explicit) gets that, regardless of whether a runtime + * view is built here. + */ + private async buildRuntimeContentGeneratorView( config: SubagentConfig, base: Config, - ): Promise { + ): Promise { const selection = parseSubagentModelSelection(config.model); - // Skip the registry rebuild if any wrapper above `base` already - // rebuilt one (typically `agent.ts:createApprovalModeOverride`, - // which marks itself via Symbol-keyed flag — Symbol property lookup - // walks the prototype chain, so this also catches - // wrapper-on-wrapper layering like - // `bgConfig = Object.create(agentConfig)` passed in from the - // background path). Rebuilding a second time would waste work, - // leak listeners on shared managers (any AgentTool / SkillTool the - // second registry later instantiates registers a change-listener - // and the short-lived registry has no explicit stop() site), and - // split the cache so client-level cache clears target an empty - // second-layer cache while bound tools (still in the upstream - // layer's registry) keep using the upstream cache. - const upstreamRebuilt = hasRebuiltToolRegistry(base); - if (selection.inherits) { - // Thin prototype-delegation override: no method changes, but a - // distinct instance triggers the lazy-init in - // `Config.getFileReadCache()` so the subagent gets its own - // cache rather than inheriting the parent's. - // - // When no upstream rebuild has happened, also rebuild the tool - // registry so `EditTool` / `WriteFileTool` / `ReadFileTool` are - // bound to the override and resolve `this.config` to the subagent - // — without that step, the parent's cached tool instances still - // reach the parent's FileReadCache and silently weaken prior-read - // enforcement on the subagent's mutation paths. - const isolated = Object.create(base) as Config; - if (!upstreamRebuilt) { - await rebuildToolRegistryOnOverride(isolated, base); - } - return isolated; + return undefined; } const authType = @@ -745,39 +747,18 @@ export class SubagentManager { authType: authType as string, }; - const agentGeneratorConfig = buildAgentContentGeneratorConfig( + const view = await createRuntimeContentGeneratorView( + base, base, selection.modelId, authOverrides, ); - const agentGenerator = await createContentGenerator( - agentGeneratorConfig, - base, - ); - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const override = Object.create(base) as any; - override.getContentGenerator = (): ContentGenerator => agentGenerator; - override.getContentGeneratorConfig = (): ContentGeneratorConfig => - agentGeneratorConfig; - override.getAuthType = (): AuthType | undefined => - agentGeneratorConfig.authType; - override.getModel = (): string => agentGeneratorConfig.model; - - // Rebuild the tool registry on the override so core tools resolve - // `this.config` to the subagent — but only if the upstream caller - // did not already build one. See the comment at the top of this - // function for the reasoning. - if (!upstreamRebuilt) { - await rebuildToolRegistryOnOverride(override as Config, base); - } - debugLogger.info( - `Created per-agent ContentGenerator for subagent "${config.name}": authType=${authType}, model=${agentGeneratorConfig.model}`, + `Created per-agent ContentGenerator for subagent "${config.name}": authType=${authType}, model=${view.contentGeneratorConfig.model}`, ); - return override as Config; + return view; } /** diff --git a/packages/core/src/tools/agent/agent-context.test.ts b/packages/core/src/tools/agent/agent-context.test.ts deleted file mode 100644 index e486d9b1f..000000000 --- a/packages/core/src/tools/agent/agent-context.test.ts +++ /dev/null @@ -1,53 +0,0 @@ -/** - * @license - * Copyright 2025 Qwen - * SPDX-License-Identifier: Apache-2.0 - */ - -import { describe, expect, it } from 'vitest'; -import { getCurrentAgentId, runWithAgentContext } from './agent-context.js'; - -describe('agent-context', () => { - it('returns null outside any frame', () => { - expect(getCurrentAgentId()).toBeNull(); - }); - - it('exposes the agentId inside a frame', async () => { - await runWithAgentContext({ agentId: 'explore-abc' }, async () => { - expect(getCurrentAgentId()).toBe('explore-abc'); - }); - expect(getCurrentAgentId()).toBeNull(); - }); - - it('propagates across awaits', async () => { - await runWithAgentContext({ agentId: 'outer-1' }, async () => { - await new Promise((resolve) => setTimeout(resolve, 0)); - expect(getCurrentAgentId()).toBe('outer-1'); - }); - }); - - it('nested frames shadow the outer agentId', async () => { - await runWithAgentContext({ agentId: 'outer-1' }, async () => { - expect(getCurrentAgentId()).toBe('outer-1'); - await runWithAgentContext({ agentId: 'inner-2' }, async () => { - expect(getCurrentAgentId()).toBe('inner-2'); - }); - expect(getCurrentAgentId()).toBe('outer-1'); - }); - }); - - it('isolates concurrent frames', async () => { - const results: string[] = []; - await Promise.all([ - runWithAgentContext({ agentId: 'a' }, async () => { - await new Promise((resolve) => setTimeout(resolve, 5)); - results.push(getCurrentAgentId() ?? 'null'); - }), - runWithAgentContext({ agentId: 'b' }, async () => { - await new Promise((resolve) => setTimeout(resolve, 1)); - results.push(getCurrentAgentId() ?? 'null'); - }), - ]); - expect(results.sort()).toEqual(['a', 'b']); - }); -}); diff --git a/packages/core/src/tools/agent/agent-context.ts b/packages/core/src/tools/agent/agent-context.ts deleted file mode 100644 index 87643aefb..000000000 --- a/packages/core/src/tools/agent/agent-context.ts +++ /dev/null @@ -1,45 +0,0 @@ -/** - * @license - * Copyright 2025 Qwen - * SPDX-License-Identifier: Apache-2.0 - */ - -/** - * @fileoverview Ambient agent-identity context for nested `agent` tool calls. - * - * When a subagent's model calls the `agent` tool, the resulting - * AgentToolInvocation's `this.config` is the main process Config (see - * comment in `fork-subagent.ts`) — it has no way to know which subagent - * made the call. We carry the launching agent's id via AsyncLocalStorage - * so nested launches can record `parentAgentId` in their sidecar meta. - */ - -import { AsyncLocalStorage } from 'node:async_hooks'; - -interface AgentContext { - readonly agentId: string; -} - -const agentContextStorage = new AsyncLocalStorage(); - -/** - * Runs `fn` with an ambient agent-identity frame. - * - * Wrap the subagent's execution (headless run loop and any hook-driven - * continuations) so every nested `agent` tool invocation inside it reads - * the launching agent's id via {@link getCurrentAgentId}. - */ -export function runWithAgentContext( - context: AgentContext, - fn: () => Promise, -): Promise { - return agentContextStorage.run(context, fn); -} - -/** - * Returns the id of the subagent whose execution is currently on the call - * stack, or `null` at the top-level user session. - */ -export function getCurrentAgentId(): string | null { - return agentContextStorage.getStore()?.agentId ?? null; -} diff --git a/packages/core/src/tools/agent/agent.test.ts b/packages/core/src/tools/agent/agent.test.ts index d7b6ec189..3f3de78d7 100644 --- a/packages/core/src/tools/agent/agent.test.ts +++ b/packages/core/src/tools/agent/agent.test.ts @@ -32,7 +32,7 @@ import type { import { partToString } from '../../utils/partUtils.js'; import type { HookSystem } from '../../hooks/hookSystem.js'; import { PermissionMode } from '../../hooks/types.js'; -import { runWithAgentContext } from './agent-context.js'; +import { runWithAgentContext } from '../../agents/runtime/agent-context.js'; import * as fs from 'node:fs'; import * as os from 'node:os'; import * as path from 'node:path'; @@ -1845,12 +1845,9 @@ describe('AgentTool', () => { invocation as unknown as { setCallId: (id: string) => void } ).setCallId('nested-1'); - await runWithAgentContext( - { agentId: 'explore-parent-42' }, - async () => { - await invocation.execute(); - }, - ); + await runWithAgentContext('explore-parent-42', async () => { + await invocation.execute(); + }); const meta = readSidecar('monitor-nested-1'); expect(meta.parentAgentId).toBe('explore-parent-42'); diff --git a/packages/core/src/tools/agent/agent.ts b/packages/core/src/tools/agent/agent.ts index 1f45a045e..1ba760990 100644 --- a/packages/core/src/tools/agent/agent.ts +++ b/packages/core/src/tools/agent/agent.ts @@ -39,7 +39,10 @@ import { isInForkExecution, runInForkContext, } from './fork-subagent.js'; -import { getCurrentAgentId, runWithAgentContext } from './agent-context.js'; +import { + getCurrentAgentId, + runWithAgentContext, +} from '../../agents/runtime/agent-context.js'; import { AgentEventEmitter, AgentEventType, @@ -187,7 +190,7 @@ export const TOOL_REGISTRY_REBUILT: unique symbol = Symbol.for( * rebuilt its tool registry via {@link rebuildToolRegistryOnOverride}. * * Used by spawn sites that may be called with a wrapper-on-wrapper - * argument (e.g. `subagent-manager.ts:maybeOverrideContentGenerator` + * argument (e.g. `subagent-manager.ts:buildSubagentContextOverride` * receiving `bgConfig = Object.create(agentConfig)` from the * background-agent path) to skip a redundant rebuild. */ @@ -200,7 +203,7 @@ export function hasRebuiltToolRegistry(config: Config): boolean { * Rebuilds the tool registry on `override` so core tools resolve * `this.config` to `override` instead of `base`. Used by both * {@link createApprovalModeOverride} and - * `subagent-manager.ts:maybeOverrideContentGenerator` to avoid + * `subagent-manager.ts:buildSubagentContextOverride` to avoid * duplicated rebuild logic. * * - `override.createToolRegistry(...)` runs on the override (so the @@ -1409,7 +1412,7 @@ class AgentToolInvocation extends BaseToolInvocation { // from this subagent's model record this agent's id as their // `parentAgentId` in the sidecar meta. const framedBgBody = () => - runWithAgentContext({ agentId: hookOpts.agentId }, bgBody); + runWithAgentContext(hookOpts.agentId, bgBody); void (isFork ? runInForkContext(framedBgBody) : framedBgBody()); this.updateDisplay({ status: 'background' as const }, updateOutput); @@ -1441,13 +1444,9 @@ class AgentToolInvocation extends BaseToolInvocation { // SkillTool the fork's model instantiates from this registry leaks // its change-listener on shared SubagentManager / SkillManager. const runFramedFork = () => - runWithAgentContext({ agentId: hookOpts.agentId }, async () => { + runWithAgentContext(hookOpts.agentId, async () => { try { - await this.runSubagentWithHooks( - subagent, - contextState, - hookOpts, - ); + await this.runSubagentWithHooks(subagent, contextState, hookOpts); } finally { void agentConfig .getToolRegistry() @@ -1477,7 +1476,7 @@ class AgentToolInvocation extends BaseToolInvocation { const fgHookOpts = { ...hookOpts, signal: fgAbortController.signal }; const runFramed = () => - runWithAgentContext({ agentId: hookOpts.agentId }, () => + runWithAgentContext(hookOpts.agentId, () => this.runSubagentWithHooks(subagent, contextState, fgHookOpts), );