Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 49 additions & 0 deletions packages/core/src/tools/agent/agent.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -695,6 +695,55 @@ describe('AgentTool', () => {
undefined,
);
});

it('stops the per-subagent ToolRegistry after the fork body finishes', async () => {
// Regression: foreground-fork fires the body via
// `void runInForkContext(...)` and returns a placeholder
// synchronously. Without an inner try/finally, the per-subagent
// ToolRegistry built by `createApprovalModeOverride` would never
// be stopped, and any AgentTool / SkillTool the fork's model
// instantiates would leak its change-listener on shared
// SubagentManager / SkillManager. Other three spawn paths
// (foreground non-fork, background fork, background non-fork)
// already stop the registry in their finally blocks.
const stopSpy = vi.fn().mockResolvedValue(undefined);
const stubReg = {
copyDiscoveredToolsFrom: vi.fn(),
getAllTools: vi.fn().mockReturnValue([]),
getAllToolNames: vi.fn().mockReturnValue([]),
stop: stopSpy,
};
// The override Config built by `createApprovalModeOverride` calls
// `createToolRegistry` (returns the override's own registry) and
// `getToolRegistry` (during `copyDiscoveredToolsFrom(base...)`).
// The override's own getToolRegistry is then assigned to whatever
// `createToolRegistry` returned. Wire BOTH config getters so the
// post-override `agentConfig.getToolRegistry().stop()` reaches our
// spy.
vi.mocked(config.getToolRegistry).mockReturnValue(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
stubReg as any,
);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
vi.mocked((config as any).createToolRegistry).mockResolvedValue(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
stubReg as any,
);

const params: AgentParams = {
description: 'fork task',
prompt: 'do the thing',
};
const invocation = (
agentTool as AgentToolWithProtectedMethods
).createInvocation(params);
await invocation.execute();

// Drain the detached fork body so its finally block runs.
await vi.runAllTimersAsync();

expect(stopSpy).toHaveBeenCalledTimes(1);
});
});

describe('SubagentStart hook integration', () => {
Expand Down
23 changes: 20 additions & 3 deletions packages/core/src/tools/agent/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1434,10 +1434,27 @@ class AgentToolInvocation extends BaseToolInvocation<AgentParams, ToolResult> {
// Background fork execution. Run under an AsyncLocalStorage frame so
// nested `agent` tool calls by the fork's model can be detected.
// Forks run async (return a placeholder); skip foreground registration.
// Wrap the fork body in try/finally so the per-subagent ToolRegistry
// is stopped after the fork finishes — the other three spawn paths
// (foreground non-fork, background fork, background non-fork) already
// do this in their finally blocks. Without it, every AgentTool /
// SkillTool the fork's model instantiates from this registry leaks
// its change-listener on shared SubagentManager / SkillManager.
const runFramedFork = () =>
runWithAgentContext({ agentId: hookOpts.agentId }, () =>
this.runSubagentWithHooks(subagent, contextState, hookOpts),
);
runWithAgentContext({ agentId: hookOpts.agentId }, async () => {
try {
await this.runSubagentWithHooks(
subagent,
contextState,
hookOpts,
);
} finally {
void agentConfig
.getToolRegistry()
.stop()
.catch(() => {});
}
});
void runInForkContext(runFramedFork);
return {
llmContent: [{ text: FORK_PLACEHOLDER_RESULT }],
Expand Down
Loading