From 606c78203539353f4f58531f44c67ff33dd62efd Mon Sep 17 00:00:00 2001 From: wenshao Date: Thu, 7 May 2026 14:06:35 +0800 Subject: [PATCH 1/2] fix(core): close bound-tool gap on runForkedAgent's YOLO wrapper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to #3873 review (#3 of the three flagged adjacent Config-wrapper sites). `runForkedAgent`'s AgentHeadless path used to build its YOLO override via a local `Object.create(parent) + getApprovalMode = YOLO` helper that did NOT rebuild the tool registry, so: 1. The YOLO approval mode was silently ignored on the bound-tool path — parent's already-bound `EditTool` / `WriteFileTool` / `ReadFileTool` resolved `this.config.getApprovalMode()` back to the parent. 2. The fork's reads / mutations went through the parent's `FileReadCache` instead of a per-fork cache. 3. Memory-extraction and dream-agent paths stack the YOLO wrapper over a `getPermissionManager`-overriding scoped wrapper. Since the bound tools resolved to the parent, BOTH overrides — the YOLO approval mode AND the scoped permission manager — were bypassed. The fix routes through the existing `createApprovalModeOverride` helper, which: - rebuilds the tool registry on the wrapper (so bound tools resolve `this.config` to the wrapper), - copies discovered tools from the upstream registry, - sets the `TOOL_REGISTRY_REBUILT` Symbol marker so any further downstream wrapper layer recognises the rebuild and skips redundant work. The memory-extraction / dream-agent composition now resolves correctly via prototype walk — the YOLO wrapper sits above the scoped wrapper, so bound tools observe `getApprovalMode() = YOLO` on the wrapper itself and `getPermissionManager() = scopedPm` one prototype level up. Adds a try/finally around the AgentHeadless run so the per-fork ToolRegistry is stopped after execution — same shape as the spawn finallys in `agent.ts` and `background-agent-resume.ts`. Without this, every AgentTool / SkillTool the fork's model later instantiates leaks its change-listener on shared SubagentManager / SkillManager. Adds `forkedAgent.agent.test.ts` covering: marker + YOLO + distinct registry on the wrapper passed to AgentHeadless.create; bound EditTool resolves to the wrapper; memory-scoped composition preserves both YOLO and scopedPm; `stop()` fires after the AgentHeadless body finishes. Uses `vi.spyOn(AgentHeadless, 'create')` rather than module mocking so the real `ContextState` / `AgentEventEmitter` keep working. `npx vitest run packages/core/src` — 269 files / 6992 passed. --- .../core/src/utils/forkedAgent.agent.test.ts | 229 ++++++++++++++++++ packages/core/src/utils/forkedAgent.ts | 113 +++++---- 2 files changed, 294 insertions(+), 48 deletions(-) create mode 100644 packages/core/src/utils/forkedAgent.agent.test.ts diff --git a/packages/core/src/utils/forkedAgent.agent.test.ts b/packages/core/src/utils/forkedAgent.agent.test.ts new file mode 100644 index 00000000000..d2f4ae839f5 --- /dev/null +++ b/packages/core/src/utils/forkedAgent.agent.test.ts @@ -0,0 +1,229 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi } from 'vitest'; +import type { Config } from '../config/config.js'; +import { Config as ConfigImpl, ApprovalMode } from '../config/config.js'; +import { AgentHeadless } from '../agents/runtime/agent-headless.js'; +import { AgentTerminateMode } from '../agents/runtime/agent-types.js'; +import { runForkedAgent } from './forkedAgent.js'; +import { ToolNames } from '../tools/tool-names.js'; +import { EditTool } from '../tools/edit.js'; +import { + hasRebuiltToolRegistry, + TOOL_REGISTRY_REBUILT, +} from '../tools/agent/agent.js'; + +/** + * Regression: `runForkedAgent` (AgentHeadless path) used to produce its + * YOLO wrapper via `Object.create(parent) + getApprovalMode = YOLO`, + * which left the parent's already-bound `EditTool` / `WriteFileTool` / + * `ReadFileTool` reachable through the wrapper's prototype chain. Bound + * tools then read `this.config.getApprovalMode()` from the parent + * (silently ignoring the YOLO override) and `this.config.getFileReadCache()` + * from the parent's cache. + * + * The fix: route through `createApprovalModeOverride`, which rebuilds + * the tool registry on the wrapper so bound tools resolve `this.config` + * to the wrapper. + */ +describe('runForkedAgent (AgentHeadless path) bound-tool isolation', () => { + // Bare mode keeps the registry small (ReadFile / Edit / Shell only) so + // the rebuild covers the file tools we actually care about. + const baseParams = { + cwd: '/tmp', + targetDir: '/tmp', + debugMode: false, + model: 'test-model', + usageStatisticsEnabled: false, + bareMode: true, + }; + + // Spy on AgentHeadless.create at the source module rather than mocking + // the re-export layer in `agents/index.js` — vitest's module-mock layer + // doesn't reliably forward `export *` re-exports through `...actual`, + // and stubbing the full surface manually is brittle. + function captureAgentHeadlessConfig(): { + captured: { config: Config | undefined }; + restore: () => void; + } { + const captured: { config: Config | undefined } = { config: undefined }; + const spy = vi + .spyOn(AgentHeadless, 'create') + .mockImplementation( + async ( + _name: string, + config: Config, + ..._rest: unknown[] + ): Promise => { + captured.config = config; + return { + execute: vi.fn().mockResolvedValue(undefined), + getTerminateMode: vi.fn().mockReturnValue(AgentTerminateMode.GOAL), + getFinalText: vi.fn().mockReturnValue('done'), + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } as any; + }, + ); + return { captured, restore: () => spy.mockRestore() }; + } + + it('passes a Config with the rebuilt-registry marker and YOLO approval mode to AgentHeadless.create', async () => { + const parent = new ConfigImpl(baseParams); + const parentRegistry = await parent.createToolRegistry(undefined, { + skipDiscovery: true, + }); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (parent as any).toolRegistry = parentRegistry; + + const { captured, restore } = captureAgentHeadlessConfig(); + try { + const result = await runForkedAgent({ + name: 'test-fork', + systemPrompt: 'You are a test fork.', + taskPrompt: 'do the task', + config: parent, + }); + expect(result.status).toBe('completed'); + } finally { + restore(); + } + + expect(captured.config).toBeDefined(); + // The wrapper passed to AgentHeadless must: + // 1. Have its own rebuilt registry (Symbol marker propagation) + expect(hasRebuiltToolRegistry(captured.config!)).toBe(true); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + expect((captured.config as any)[TOOL_REGISTRY_REBUILT]).toBe(true); + // 2. Resolve approval mode to YOLO (the override) + expect(captured.config!.getApprovalMode()).toBe(ApprovalMode.YOLO); + // 3. Hand out a different ToolRegistry instance from the parent + expect(captured.config!.getToolRegistry()).not.toBe(parentRegistry); + }); + + it('binds EditTool from the wrapper registry to the wrapper Config (not the parent)', async () => { + const parent = new ConfigImpl(baseParams); + const parentRegistry = await parent.createToolRegistry(undefined, { + skipDiscovery: true, + }); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (parent as any).toolRegistry = parentRegistry; + + const { captured, restore } = captureAgentHeadlessConfig(); + try { + await runForkedAgent({ + name: 'test-fork', + systemPrompt: 'You are a test fork.', + taskPrompt: 'do the task', + config: parent, + }); + } finally { + restore(); + } + + expect(captured.config).toBeDefined(); + const wrapperRegistry = captured.config!.getToolRegistry(); + const editTool = await wrapperRegistry.ensureTool(ToolNames.EDIT); + expect(editTool).toBeInstanceOf(EditTool); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + expect((editTool as any).config).toBe(captured.config); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const boundConfig = (editTool as any).config as Config; + expect(boundConfig.getApprovalMode()).toBe(ApprovalMode.YOLO); + expect(boundConfig.getFileReadCache()).toBe( + captured.config!.getFileReadCache(), + ); + expect(boundConfig.getFileReadCache()).not.toBe(parent.getFileReadCache()); + }); + + it('preserves an upstream getPermissionManager override (memory-scoped composition)', async () => { + // The memory extraction / dream agent path stacks two wrappers: + // parent + // └── scopedConfig (Object.create + getPermissionManager override) + // └── yoloConfig (createApprovalModeOverride, sets registry + marker) + // Bound tools must see: + // - approval mode = YOLO (from yoloConfig's own override) + // - permission manager = scopedPm (walks proto past yoloConfig to scopedConfig) + const parent = new ConfigImpl(baseParams); + const parentRegistry = await parent.createToolRegistry(undefined, { + skipDiscovery: true, + }); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (parent as any).toolRegistry = parentRegistry; + + const scopedPm = { id: 'scoped-pm-marker' } as never; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const scopedConfig = Object.create(parent) as any; + scopedConfig.getPermissionManager = () => scopedPm; + + const { captured, restore } = captureAgentHeadlessConfig(); + try { + await runForkedAgent({ + name: 'test-fork', + systemPrompt: 'You are a test fork.', + taskPrompt: 'do the task', + config: scopedConfig as Config, + }); + } finally { + restore(); + } + + expect(captured.config).toBeDefined(); + const editTool = await captured + .config!.getToolRegistry() + .ensureTool(ToolNames.EDIT); + expect(editTool).toBeInstanceOf(EditTool); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const boundConfig = (editTool as any).config as Config; + // YOLO from yoloConfig's own override + expect(boundConfig.getApprovalMode()).toBe(ApprovalMode.YOLO); + // Scoped PM from scopedConfig (one prototype level up) + expect(boundConfig.getPermissionManager?.()).toBe(scopedPm); + }); + + it('stops the per-fork ToolRegistry after the AgentHeadless body finishes', async () => { + const parent = new ConfigImpl(baseParams); + const parentRegistry = await parent.createToolRegistry(undefined, { + skipDiscovery: true, + }); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (parent as any).toolRegistry = parentRegistry; + + // Wrap parent.createToolRegistry so the registry it returns to + // `createApprovalModeOverride` carries a stop spy. The wrapper's + // own getToolRegistry is then assigned this same instance. + const stopSpy = vi.fn().mockResolvedValue(undefined); + const originalCreate = parent.createToolRegistry.bind(parent); + vi.spyOn(parent, 'createToolRegistry').mockImplementation( + async (...args) => { + const reg = await originalCreate(...args); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (reg as any).stop = stopSpy; + return reg; + }, + ); + + const { restore } = captureAgentHeadlessConfig(); + try { + await runForkedAgent({ + name: 'test-fork', + systemPrompt: 'You are a test fork.', + taskPrompt: 'do the task', + config: parent, + }); + } finally { + restore(); + } + + // stop() is fire-and-forget inside the runForkedAgent finally — + // it is awaited by the runtime via the resolved promise chain, so + // by the time `await runForkedAgent` returns the stop call has + // already started; flush microtasks for the catch handler. + await new Promise((resolve) => setImmediate(resolve)); + + expect(stopSpy).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/core/src/utils/forkedAgent.ts b/packages/core/src/utils/forkedAgent.ts index 2f0de662a33..162a51f1a11 100644 --- a/packages/core/src/utils/forkedAgent.ts +++ b/packages/core/src/utils/forkedAgent.ts @@ -31,6 +31,7 @@ import type { } from '@google/genai'; import { ApprovalMode, type Config } from '../config/config.js'; import { GeminiChat, StreamEventType } from '../core/geminiChat.js'; +import { createApprovalModeOverride } from '../tools/agent/agent.js'; import { AgentHeadless, AgentEventEmitter, @@ -256,17 +257,6 @@ export interface ForkedAgentResult { filesTouched: string[]; } -/** - * Returns a shallow clone of config with ApprovalMode forced to YOLO. - * Background agents must never block on permission prompts — there is - * no user present to answer them. - */ -function createYoloConfig(config: Config): Config { - const yoloConfig = Object.create(config) as Config; - yoloConfig.getApprovalMode = () => ApprovalMode.YOLO; - return yoloConfig; -} - /** * Extracts file paths from a tool call's args object. * Matches any arg key that contains "path", "file", or "target". @@ -376,7 +366,23 @@ export async function runForkedAgent( } // ── AgentHeadless path ──────────────────────────────────────────────────── - const yoloConfig = createYoloConfig(params.config); + // `createApprovalModeOverride` rebuilds the tool registry on the YOLO + // wrapper Config so core file tools (`EditTool` / `WriteFileTool` / + // `ReadFileTool`) resolve `this.config` to the wrapper, not to the + // parent. Without that rebuild the YOLO override is silently ignored + // on the bound-tool path (parent's pre-bound tool instances keep + // reading the parent's approval mode), and the wrapper's own + // `FileReadCache` lazy-init is bypassed too. + // + // Consumers that pre-wrap with `createMemoryScopedAgentConfig` + // (memory extraction / dream agent) compose correctly: the YOLO + // wrapper's bound tools resolve `this.config.getPermissionManager()` + // through the prototype chain to the scoped wrapper's own override, + // while `this.config.getApprovalMode()` lands on YOLO. + const yoloConfig = await createApprovalModeOverride( + params.config, + ApprovalMode.YOLO, + ); const filesTouched = new Set(); const emitter = new AgentEventEmitter(); @@ -401,47 +407,58 @@ export async function runForkedAgent( const toolConfig: ToolConfig | undefined = params.tools !== undefined ? { tools: params.tools } : undefined; - const headless = await AgentHeadless.create( - params.name, - yoloConfig, - promptConfig, - modelConfig, - runConfig, - toolConfig, - emitter, - ); - - const context = new ContextState(); - context.set('task_prompt', params.taskPrompt); - await headless.execute(context, params.abortSignal); - - const terminateReason = headless.getTerminateMode(); - const finalText = headless.getFinalText() || undefined; - const touched = [...filesTouched]; + try { + const headless = await AgentHeadless.create( + params.name, + yoloConfig, + promptConfig, + modelConfig, + runConfig, + toolConfig, + emitter, + ); - if (terminateReason === AgentTerminateMode.CANCELLED) { - return { - status: 'cancelled', - terminateReason, - finalText, - filesTouched: touched, - }; - } - if ( - terminateReason === AgentTerminateMode.ERROR || - terminateReason === AgentTerminateMode.TIMEOUT - ) { + const context = new ContextState(); + context.set('task_prompt', params.taskPrompt); + await headless.execute(context, params.abortSignal); + + const terminateReason = headless.getTerminateMode(); + const finalText = headless.getFinalText() || undefined; + const touched = [...filesTouched]; + + if (terminateReason === AgentTerminateMode.CANCELLED) { + return { + status: 'cancelled', + terminateReason, + finalText, + filesTouched: touched, + }; + } + if ( + terminateReason === AgentTerminateMode.ERROR || + terminateReason === AgentTerminateMode.TIMEOUT + ) { + return { + status: 'failed', + terminateReason, + finalText, + filesTouched: touched, + }; + } return { - status: 'failed', + status: 'completed', terminateReason, finalText, filesTouched: touched, }; + } finally { + // Release the per-fork ToolRegistry so AgentTool / SkillTool + // instances dispose their change-listeners on shared + // SubagentManager / SkillManager. Same shape as the spawn-path + // finallys in `agent.ts` and `background-agent-resume.ts`. + void yoloConfig + .getToolRegistry() + .stop() + .catch(() => {}); } - return { - status: 'completed', - terminateReason, - finalText, - filesTouched: touched, - }; } From 8f266fa7ba43b8c5806c34637e61dc43f0bf47a6 Mon Sep 17 00:00:00 2001 From: wenshao Date: Thu, 7 May 2026 15:50:53 +0800 Subject: [PATCH 2/2] test(core): cover stop() lifecycle on AgentHeadless.create + execute failure paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Self-review feedback on #3892: the stop lifecycle test only covered the success path. A future refactor could move the stop() out of the `finally` block and onto the success branch, reintroducing listener leaks when create or execute rejects, while every existing test still passes. Two new tests pin the cleanup to the `finally`: 1. `stops the per-fork ToolRegistry even when AgentHeadless.create rejects` — make `AgentHeadless.create` return a rejected promise; assert the rejection propagates and the stop spy still fires once. 2. `stops the per-fork ToolRegistry even when headless.execute rejects` — return a headless object whose `execute` rejects; same shape. Together with the success-path test these three cases cover every exit edge of the AgentHeadless body. `npx vitest run packages/core/src` — 269 files / 6994 passed. --- .../core/src/utils/forkedAgent.agent.test.ts | 100 ++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/packages/core/src/utils/forkedAgent.agent.test.ts b/packages/core/src/utils/forkedAgent.agent.test.ts index d2f4ae839f5..92d57e17b10 100644 --- a/packages/core/src/utils/forkedAgent.agent.test.ts +++ b/packages/core/src/utils/forkedAgent.agent.test.ts @@ -226,4 +226,104 @@ describe('runForkedAgent (AgentHeadless path) bound-tool isolation', () => { expect(stopSpy).toHaveBeenCalledTimes(1); }); + + it('stops the per-fork ToolRegistry even when AgentHeadless.create rejects', async () => { + // Failure-path regression: a future refactor could accidentally + // move the stop() out of the `finally` and onto the success path + // while every other test still passes. This test pins that the + // cleanup runs when `AgentHeadless.create` rejects before any + // body executes. + const parent = new ConfigImpl(baseParams); + const parentRegistry = await parent.createToolRegistry(undefined, { + skipDiscovery: true, + }); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (parent as any).toolRegistry = parentRegistry; + + const stopSpy = vi.fn().mockResolvedValue(undefined); + const originalCreate = parent.createToolRegistry.bind(parent); + vi.spyOn(parent, 'createToolRegistry').mockImplementation( + async (...args) => { + const reg = await originalCreate(...args); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (reg as any).stop = stopSpy; + return reg; + }, + ); + + const createSpy = vi + .spyOn(AgentHeadless, 'create') + .mockRejectedValue(new Error('agent-headless-create-blew-up')); + + try { + await expect( + runForkedAgent({ + name: 'test-fork', + systemPrompt: 'You are a test fork.', + taskPrompt: 'do the task', + config: parent, + }), + ).rejects.toThrow('agent-headless-create-blew-up'); + } finally { + createSpy.mockRestore(); + } + + await new Promise((resolve) => setImmediate(resolve)); + expect(stopSpy).toHaveBeenCalledTimes(1); + }); + + it('stops the per-fork ToolRegistry even when headless.execute rejects', async () => { + // Same shape as the create-rejects test, but for the execute + // failure path. Together they pin the lifecycle stop to the + // `finally` block rather than any specific success branch. + const parent = new ConfigImpl(baseParams); + const parentRegistry = await parent.createToolRegistry(undefined, { + skipDiscovery: true, + }); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (parent as any).toolRegistry = parentRegistry; + + const stopSpy = vi.fn().mockResolvedValue(undefined); + const originalCreate = parent.createToolRegistry.bind(parent); + vi.spyOn(parent, 'createToolRegistry').mockImplementation( + async (...args) => { + const reg = await originalCreate(...args); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (reg as any).stop = stopSpy; + return reg; + }, + ); + + const createSpy = vi + .spyOn(AgentHeadless, 'create') + .mockImplementation( + async (..._args: unknown[]): Promise => + ({ + execute: vi + .fn() + .mockRejectedValue(new Error('headless-execute-blew-up')), + getTerminateMode: vi + .fn() + .mockReturnValue(AgentTerminateMode.GOAL), + getFinalText: vi.fn().mockReturnValue(''), + // eslint-disable-next-line @typescript-eslint/no-explicit-any + }) as any, + ); + + try { + await expect( + runForkedAgent({ + name: 'test-fork', + systemPrompt: 'You are a test fork.', + taskPrompt: 'do the task', + config: parent, + }), + ).rejects.toThrow('headless-execute-blew-up'); + } finally { + createSpy.mockRestore(); + } + + await new Promise((resolve) => setImmediate(resolve)); + expect(stopSpy).toHaveBeenCalledTimes(1); + }); });