diff --git a/docs/design/2026-07-20-background-agent-hot-continuation.md b/docs/design/2026-07-20-background-agent-hot-continuation.md new file mode 100644 index 00000000000..7d430061bb7 --- /dev/null +++ b/docs/design/2026-07-20-background-agent-hot-continuation.md @@ -0,0 +1,160 @@ +# Background Agent Hot Continuation + +## Context + +A completed background subagent currently loses its in-process runtime. A later +`send_message` reconstructs a new `AgentHeadless` from the JSONL transcript. +This preserves most visible conversation history, but recreates the chat, tool +surface, per-agent registries, and provider-side cache state. + +The launch path also constructs ordinary background agents twice: once with the +parent emitter and again with the dedicated background emitter. The first +instance is never executed or disposed. + +This design addresses only the in-session lifecycle. Discovery and continuity +after restoring the parent session are outside its scope. + +The distinction is behavioral, not just an implementation detail. Within one +session, transcript revival already preserves the model-visible conversation, +so hot continuation primarily avoids runtime reconstruction and preserves +provider/tool state. Across a parent-session restore, the original in-memory +runtime cannot survive process teardown; this change leaves that cold-start +behavior unchanged. + +## Goals + +- Create one runtime for a fresh ordinary background agent. +- Keep that runtime resident after a successful turn. +- Continue a completed task on the same chat and prepared tool surface. +- Preserve the current task row, task ID, per-turn start/completion events, and + terminal notifications. +- Keep transcript revival as the fallback when no compatible resident runtime + exists. +- Release resident resources on failure, cancellation, session shutdown/reset, + and terminal-entry eviction. + +## Non-goals + +- Persisting a live runtime across processes or parent-session restoration. +- Adding an `idle` value to the shared task-status union. +- Changing how messages sent to an actively running agent are injected between + tool rounds. +- Making fork agents persistent. +- Extending temporary worktree lifetime across completed turns. +- Making globally registered frontmatter hooks safe to leave installed while an + agent is idle. + +## Design + +### Reusable headless runtime + +`AgentHeadless` keeps its `GeminiChat` and prepared tool declarations as +instance state. Its public `execute()` remains a per-turn operation: + +- only one call may run at a time; +- final text and termination mode are reset at the start; statistics reset for + a new parent instruction but remain cumulative across internal stop-hook + retries of that instruction; +- the first call creates the chat and prepares tools; +- later calls append a new user turn to the same chat and emit an external + message event so the JSONL transcript remains complete. + +This keeps the existing `AgentHeadless` hooks, telemetry, external-message +drain, and terminal result contract. `AgentInteractive` is not used because its +queue API does not provide the per-turn completion result and notification +semantics required by background tasks. + +### Resident controller + +`BackgroundTaskRegistry` owns an in-memory controller table keyed by task ID. +The controller is intentionally separate from `AgentTask`, which remains a +serializable UI/status record. + +A controller can: + +- start a continuation from a completed row; +- abort and dispose its runtime. + +On a completed `send_message`, the tool first asks the registry for a resident +continuation. A hit synchronously changes the existing row back to `running`, +claims a normal background execution slot, and schedules the new turn after +the previous turn has fully settled. A miss uses the existing transcript +revival service. + +`completed` continues to mean “the latest turn completed.” Runtime residency is +an internal implementation fact, so the shared task status and UI do not gain a +new idle state. + +### Per-turn and resident resources + +Each continuation receives a fresh abort controller, SubagentStart/Stop hook +pair, task-start event, completion notification, and sidecar status transition. +Fresh-launch resident turns also keep the launch path's per-turn trace span. +Transcript-reconstructed residents preserve the existing resume path's +telemetry shape; adding spans to that path is outside this lifecycle change. A +runtime that would need a child-only AUTO permission lease is not retained +because those leases are not reference-counted across concurrent subagents. + +The chat, prepared tools, JSONL writer, event listeners, agent-scoped tool +registry, and per-agent MCP resources remain alive while the controller is +resident. Disposal is idempotent. + +The existing terminal-entry retention limit also bounds resident controllers. +Pruning a row disposes its controller. Registry reset and shutdown dispose all +controllers, including already-completed ones. + +### Compatibility exclusions + +The first version retains only ordinary named background agents that: + +- completed normally; +- do not use `isolation: "worktree"`; +- do not declare frontmatter hooks; +- do not require a child-only AUTO permission lease. + +Temporary worktrees are currently finalized after each turn, so retaining a +runtime would leave its Config pointing at a removed directory. Frontmatter +hooks are currently registered globally for their lifetime, so retaining them +while idle could affect unrelated work. Child-only AUTO leases mutate the +parent permission manager and are not reference-counted across concurrent +subagents, so reacquiring them per hot turn would be unsafe. Hooked, child-only +AUTO, and worktree-isolated agents continue through the existing JSONL revival +flow. + +## Races and failure handling + +- The registry performs the completed-to-running transition synchronously + before the continuation promise is scheduled. A second concurrent + `send_message` therefore observes `running` and uses the existing in-round + message queue. +- The next turn is chained after the prior turn promise, covering the window in + which the completion notification is emitted before the prior `finally` + block has finished. +- After the final in-loop message drain, a resident turn synchronously claims + any inputs queued while stop hooks or cleanup were running. Claimed inputs + start another turn on the same runtime; an empty claim is followed by the + terminal state transition without another asynchronous boundary. +- Pending inputs and cleanup are guarded by resident-controller identity so a + stale runtime cannot consume work from or unregister its replacement. +- Failed and cancelled turns remove and dispose the resident controller. +- A parent-session working-directory change invalidates the resident runtime; + the next continuation disposes it and falls back to transcript revival. +- If claiming a background slot fails, the row stays completed and the caller + can use the existing cold-revival error path. +- Disposal during an active turn aborts its controller and defers destructive + resource cleanup to the turn's finalizer. + +## Validation + +Unit tests must prove: + +- a fresh background launch creates exactly one `AgentHeadless`; +- two sequential turns use one `GeminiChat` and one prepared tool list; +- completed `send_message` prefers the resident controller; +- absence of a resident controller still invokes transcript revival; +- the second user instruction is present in JSONL; +- reset, shutdown/cancellation, and terminal pruning dispose exactly once. + +The E2E scenario uses one task ID for two completed phases and verifies that the +second phase remembers a nonce from the first. Physical runtime identity is +verified by unit tests because stream JSON does not expose constructor counts. diff --git a/packages/acp-bridge/src/bridge.test.ts b/packages/acp-bridge/src/bridge.test.ts index 1252ea893ef..187e5b95eff 100644 --- a/packages/acp-bridge/src/bridge.test.ts +++ b/packages/acp-bridge/src/bridge.test.ts @@ -49,6 +49,7 @@ import { import { BridgeChannelClosedError, BridgeTimeoutError, + SESSION_CLOSE_QUARANTINED_ERROR_KIND, SERVE_CONTROL_EXT_METHODS, SERVE_STATUS_EXT_METHODS, } from './status.js'; @@ -12909,6 +12910,31 @@ describe('createAcpSessionBridge', () => { await bridge.shutdown(); }); + it('kills the channel when the child quarantined the session before close failed', async () => { + const handle = makeChannel({ + extMethodImpl: (method) => { + if (method === SERVE_CONTROL_EXT_METHODS.sessionClose) { + throw new RequestError(-32603, 'dispose failed', { + errorKind: SESSION_CLOSE_QUARANTINED_ERROR_KIND, + }); + } + return {}; + }, + }); + const bridge = makeBridge({ + channelFactory: async () => handle.channel, + }); + const session = await bridge.spawnOrAttach({ workspaceCwd: WS_A }); + + await expect(bridge.closeSession(session.sessionId)).rejects.toThrow( + 'dispose failed', + ); + + expect(handle.killed).toBe(true); + await vi.waitFor(() => expect(bridge.sessionCount).toBe(0)); + await bridge.shutdown(); + }); + it('resolves pending permissions as cancelled', async () => { let capturedConn: AgentSideConnection | undefined; const factory: ChannelFactory = async () => { diff --git a/packages/acp-bridge/src/bridge.ts b/packages/acp-bridge/src/bridge.ts index 01096d4c262..95e6c54f9f1 100644 --- a/packages/acp-bridge/src/bridge.ts +++ b/packages/acp-bridge/src/bridge.ts @@ -48,6 +48,7 @@ import { BridgeTimeoutError, createIdleWorkspaceExtensionsStatus, createIdleWorkspaceHooksStatus, + SESSION_CLOSE_QUARANTINED_ERROR_KIND, SERVE_CONTROL_EXT_METHODS, SERVE_STATUS_EXT_METHODS, STATUS_SCHEMA_VERSION, @@ -202,6 +203,13 @@ function isRecord(value: unknown): value is Record { } function isDefinitiveAcpRequestError(error: unknown): boolean { + if ( + isRecord(error) && + isRecord(error['data']) && + error['data']['errorKind'] === SESSION_CLOSE_QUARANTINED_ERROR_KIND + ) { + return false; + } if (error instanceof RequestError) return true; if (!isRecord(error)) return false; return ( diff --git a/packages/acp-bridge/src/status.ts b/packages/acp-bridge/src/status.ts index 214dfb2b673..90472595e7c 100644 --- a/packages/acp-bridge/src/status.ts +++ b/packages/acp-bridge/src/status.ts @@ -10,6 +10,9 @@ import { SkillError } from '@qwen-code/qwen-code-core'; export const STATUS_SCHEMA_VERSION = 1 as const; +export const SESSION_CLOSE_QUARANTINED_ERROR_KIND = + 'session_close_quarantined' as const; + /** * Closed enumeration of structured error categories surfaced on diagnostic * status cells. Cells produced by `/workspace/preflight`, `/workspace/env`, diff --git a/packages/cli/src/acp-integration/acpAgent.test.ts b/packages/cli/src/acp-integration/acpAgent.test.ts index 6386580c4df..4d104e71f07 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -796,6 +796,7 @@ import { SessionTranscriptPageTooLargeError, encodeSessionTranscriptCursor, unregisterGoalHook, + uiTelemetryService, getActiveGoal, registerGoalHook, startEventLoopLagMonitor, @@ -803,6 +804,7 @@ import { SESSION_ARTIFACT_PERSISTENCE_VERSION, mcpServerRequiresOAuth, APPROVAL_MODES, + McpTransportPool, } from '@qwen-code/qwen-code-core'; import type { LoadSessionResponse, @@ -820,6 +822,7 @@ import { createLoadedSettingsAdapter } from '../config/loadedSettingsAdapter.js' import { AcpFileSystemService } from './service/filesystem.js'; import { Session, buildAvailableCommandsSnapshot } from './session/Session.js'; import { + SESSION_CLOSE_QUARANTINED_ERROR_KIND, SERVE_STATUS_EXT_METHODS, SERVE_CONTROL_EXT_METHODS, } from '@qwen-code/acp-bridge/status'; @@ -1418,6 +1421,7 @@ describe('QwenAgent MCP SSE/HTTP support', () => { getRewindableUserTurnCount: ReturnType; clearTodoStopGuardTrust: ReturnType; releaseTodoStopGuardQueuedPromptWait: ReturnType; + dispose: ReturnType; } | undefined; let processExitSpy: MockInstance; @@ -2271,7 +2275,7 @@ describe('QwenAgent MCP SSE/HTTP support', () => { waitForActiveTurnsToSettle: vi.fn().mockResolvedValue(undefined), cancelPendingPrompt: vi.fn().mockResolvedValue(undefined), assertCanStartTurn: vi.fn().mockResolvedValue(undefined), - dispose: vi.fn(), + dispose: vi.fn().mockResolvedValue(undefined), emitGoalStatus: vi.fn(), captureHistorySnapshot: vi .fn() @@ -9530,6 +9534,110 @@ describe('QwenAgent MCP SSE/HTTP support', () => { ); }); + it('waits for session disposal before draining the MCP pool on connection close', async () => { + await setupSessionMocks('session-shutdown-order'); + const { agent, agentPromise } = await bootAcpAgent(); + + await agent.newSession({ cwd: '/tmp', mcpServers: [] }); + const liveSession = vi.mocked(Session).mock.results.at(-1)?.value as + | { dispose: ReturnType } + | undefined; + const mcpPool = vi.mocked(McpTransportPool).mock.results.at(-1)?.value as + | { drainAll: ReturnType } + | undefined; + let releaseDisposal!: () => void; + liveSession?.dispose.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseDisposal = resolve; + }), + ); + + mockConnectionState.resolve(); + await vi.waitFor(() => expect(liveSession?.dispose).toHaveBeenCalledOnce()); + expect(mcpPool?.drainAll).not.toHaveBeenCalled(); + + releaseDisposal(); + await agentPromise; + expect(mcpPool?.drainAll).toHaveBeenCalledOnce(); + }); + + it('waits for an in-flight session lifecycle operation before connection teardown', async () => { + const innerConfig = await setupSessionMocks('session-in-flight-close'); + const sessionHookSystem = { + fireSessionEndEvent: vi.fn().mockResolvedValue(undefined), + }; + innerConfig.getHookSystem = vi.fn().mockReturnValue(sessionHookSystem); + innerConfig.getDisableAllHooks = vi.fn().mockReturnValue(false); + innerConfig.hasHooksForEvent = vi + .fn() + .mockImplementation((event: string) => event === 'SessionEnd'); + let releaseConfig!: (config: Config) => void; + vi.mocked(loadCliConfig).mockReturnValueOnce( + new Promise((resolve) => { + releaseConfig = resolve; + }), + ); + const { agent, agentPromise } = await bootAcpAgent(); + + const creating = agent.newSession({ cwd: '/tmp', mcpServers: [] }); + await vi.waitFor(() => expect(loadCliConfig).toHaveBeenCalledOnce()); + + let teardownSettled = false; + mockConnectionState.resolve(); + void agentPromise.finally(() => { + teardownSettled = true; + }); + await Promise.resolve(); + expect(teardownSettled).toBe(false); + expect(Session).not.toHaveBeenCalled(); + + releaseConfig(innerConfig as unknown as Config); + await creating; + const createdSession = lastSessionMock; + await agentPromise; + + expect(createdSession?.dispose).toHaveBeenCalledOnce(); + expect(innerConfig.shutdown).toHaveBeenCalledOnce(); + expect(sessionHookSystem.fireSessionEndEvent).toHaveBeenCalledWith( + SessionEndReason.PromptInputExit, + ); + }); + + it('coalesces concurrent session disposal', async () => { + const innerConfig = await setupSessionMocks('session-concurrent-disposal'); + const { agent, agentPromise } = await bootAcpAgent(); + await agent.newSession({ cwd: '/tmp', mcpServers: [] }); + const liveSession = vi.mocked(Session).mock.results.at(-1)?.value as + | { dispose: ReturnType } + | undefined; + let releaseDisposal!: () => void; + liveSession?.dispose.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseDisposal = resolve; + }), + ); + const disposableAgent = agent as unknown as { + disposeSessions: () => Promise; + }; + + const firstDisposal = disposableAgent.disposeSessions(); + await vi.waitFor(() => expect(liveSession?.dispose).toHaveBeenCalledOnce()); + const secondDisposal = disposableAgent.disposeSessions(); + expect(liveSession?.dispose).toHaveBeenCalledOnce(); + + releaseDisposal(); + await Promise.all([firstDisposal, secondDisposal]); + expect(liveSession?.dispose).toHaveBeenCalledOnce(); + expect(innerConfig.shutdown).toHaveBeenCalledOnce(); + + mockConnectionState.resolve(); + await agentPromise; + expect(liveSession?.dispose).toHaveBeenCalledOnce(); + expect(innerConfig.shutdown).toHaveBeenCalledOnce(); + }); + it('rewindSession extension method rewinds the active session', async () => { const sessionId = '11111111-1111-1111-1111-111111111111'; const innerConfig = await setupSessionMocks(sessionId); @@ -10860,6 +10968,13 @@ describe('QwenAgent extMethod renameSession routing', () => { expect(recording.flush).toHaveBeenCalledOnce(); expect(liveCancelPendingPrompt).not.toHaveBeenCalled(); expect(innerConfig.shutdown).not.toHaveBeenCalled(); + expect( + ( + vi.mocked(Session).mock.results.at(-1)?.value as + | { dispose: ReturnType } + | undefined + )?.dispose, + ).not.toHaveBeenCalled(); await expect( agent.extMethod('qwen/control/session/close', { @@ -10870,6 +10985,13 @@ describe('QwenAgent extMethod renameSession routing', () => { expect(recording.flush).toHaveBeenCalledTimes(2); expect(liveCancelPendingPrompt).not.toHaveBeenCalled(); expect(innerConfig.shutdown).not.toHaveBeenCalled(); + expect( + ( + vi.mocked(Session).mock.results.at(-1)?.value as + | { dispose: ReturnType } + | undefined + )?.dispose, + ).not.toHaveBeenCalled(); await expect( agent.extMethod('qwen/control/session/close', { @@ -10880,6 +11002,13 @@ describe('QwenAgent extMethod renameSession routing', () => { expect(recording.flush).toHaveBeenCalledTimes(3); expect(liveCancelPendingPrompt).toHaveBeenCalledOnce(); expect(innerConfig.shutdown).toHaveBeenCalledOnce(); + expect( + ( + vi.mocked(Session).mock.results.at(-1)?.value as + | { dispose: ReturnType } + | undefined + )?.dispose, + ).toHaveBeenCalledOnce(); expect( ( agent as unknown as { @@ -10894,6 +11023,119 @@ describe('QwenAgent extMethod renameSession routing', () => { await agentPromise; }); + it('quarantines a disposed session when the final strict flush fails', async () => { + const recording = makeRecordingService(); + let ownsWriter = true; + recording.hasWriteOwnership.mockImplementation(() => ownsWriter); + recording.close.mockImplementation(async () => { + ownsWriter = false; + }); + recording.flush + .mockResolvedValueOnce(undefined) + .mockRejectedValueOnce(new Error('flush failed')) + .mockResolvedValue(undefined); + const innerConfig = makeLiveSessionInnerConfig(recording); + const { agent, agentPromise } = await bootAgent(innerConfig); + + await agent.newSession({ cwd: '/tmp', mcpServers: [] }); + const liveSession = vi.mocked(Session).mock.results.at(-1)?.value as + | { dispose: ReturnType } + | undefined; + + const closeError = await agent + .extMethod('qwen/control/session/close', { + sessionId: liveSessionId, + requireFlush: true, + }) + .catch((error: unknown) => error); + expect(closeError).toMatchObject({ + message: 'flush failed', + data: { + errorKind: SESSION_CLOSE_QUARANTINED_ERROR_KIND, + sessionId: liveSessionId, + }, + }); + expect(recording.flush).toHaveBeenCalledTimes(2); + expect(recording.close).not.toHaveBeenCalled(); + expect(recording.hasWriteOwnership()).toBe(true); + expect(liveCancelPendingPrompt).toHaveBeenCalledOnce(); + expect(liveSession?.dispose).toHaveBeenCalledOnce(); + expect(innerConfig.shutdown).not.toHaveBeenCalled(); + expect( + ( + agent as unknown as { + getActiveSessions: () => Array<{ getId: () => string }>; + } + ) + .getActiveSessions() + .map((session) => session.getId()), + ).not.toContain(liveSessionId); + + await expect( + agent.extMethod('qwen/control/session/close', { + sessionId: liveSessionId, + requireFlush: true, + }), + ).resolves.toEqual({ sessionId: liveSessionId, closed: true }); + expect(recording.flush).toHaveBeenCalledTimes(3); + expect(recording.close).toHaveBeenCalledOnce(); + expect(recording.hasWriteOwnership()).toBe(false); + expect(liveSession?.dispose).toHaveBeenCalledTimes(2); + expect(innerConfig.shutdown).toHaveBeenCalledOnce(); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('keeps the disposed session quarantined while close retains writer ownership', async () => { + const recording = makeRecordingService(); + let ownsWriter = true; + recording.hasWriteOwnership.mockImplementation(() => ownsWriter); + recording.close + .mockRejectedValueOnce(new Error('lease release failed')) + .mockImplementation(async () => { + ownsWriter = false; + }); + const innerConfig = makeLiveSessionInnerConfig(recording); + const { agent, agentPromise } = await bootAgent(innerConfig); + await agent.newSession({ cwd: '/tmp', mcpServers: [] }); + const liveSession = vi.mocked(Session).mock.results.at(-1)?.value as + | { dispose: ReturnType } + | undefined; + + await expect( + agent.extMethod('qwen/control/session/close', { + sessionId: liveSessionId, + }), + ).rejects.toThrow('lease release failed'); + expect(liveSession?.dispose).toHaveBeenCalledOnce(); + expect(recording.close).toHaveBeenCalledOnce(); + expect(recording.hasWriteOwnership()).toBe(true); + expect(innerConfig.shutdown).not.toHaveBeenCalled(); + expect( + ( + agent as unknown as { + getActiveSessions: () => Array<{ getId: () => string }>; + } + ) + .getActiveSessions() + .map((session) => session.getId()), + ).not.toContain(liveSessionId); + + await expect( + agent.extMethod('qwen/control/session/close', { + sessionId: liveSessionId, + }), + ).resolves.toEqual({ sessionId: liveSessionId, closed: true }); + expect(liveSession?.dispose).toHaveBeenCalledTimes(2); + expect(recording.close).toHaveBeenCalledTimes(2); + expect(recording.hasWriteOwnership()).toBe(false); + expect(innerConfig.shutdown).toHaveBeenCalledOnce(); + + mockConnectionState.resolve(); + await agentPromise; + }); + it('does not abort an active generation when the close gate is unavailable', async () => { const recording = makeRecordingService(); const innerConfig = makeLiveSessionInnerConfig(recording); @@ -11050,6 +11292,72 @@ describe('QwenAgent extMethod renameSession routing', () => { mockConnectionState.resolve(); await agentPromise; }); + + it('quarantines a session until a failed disposal succeeds on retry', async () => { + const recording = makeRecordingService(); + const innerConfig = makeLiveSessionInnerConfig(recording); + const { agent, agentPromise } = await bootAgent(innerConfig); + await agent.newSession({ cwd: '/tmp', mcpServers: [] }); + const liveSession = vi.mocked(Session).mock.results.at(-1)?.value as + | { dispose: ReturnType } + | undefined; + liveSession?.dispose.mockRejectedValueOnce(new Error('dispose failed')); + const mcpPool = vi.mocked(McpTransportPool).mock.results.at(-1)?.value as + | { releaseSession: ReturnType } + | undefined; + const unregisterCalls = vi.mocked(unregisterGoalHook).mock.calls.length; + const telemetryCalls = vi.mocked(uiTelemetryService.removeSession).mock + .calls.length; + + await expect( + agent.extMethod('qwen/control/session/close', { + sessionId: liveSessionId, + requireFlush: false, + }), + ).rejects.toThrow('dispose failed'); + expect(recording.finalize).not.toHaveBeenCalled(); + expect(recording.flush).not.toHaveBeenCalled(); + expect(recording.close).not.toHaveBeenCalled(); + expect(innerConfig.shutdown).not.toHaveBeenCalled(); + expect(vi.mocked(unregisterGoalHook).mock.calls).toHaveLength( + unregisterCalls, + ); + expect(mcpPool?.releaseSession).not.toHaveBeenCalled(); + expect(vi.mocked(uiTelemetryService.removeSession).mock.calls).toHaveLength( + telemetryCalls, + ); + expect( + ( + agent as unknown as { + getActiveSessions: () => Array<{ getId: () => string }>; + } + ) + .getActiveSessions() + .map((session) => session.getId()), + ).not.toContain(liveSessionId); + + await expect( + agent.extMethod('qwen/control/session/close', { + sessionId: liveSessionId, + requireFlush: false, + }), + ).resolves.toEqual({ sessionId: liveSessionId, closed: true }); + + expect(liveSession?.dispose).toHaveBeenCalledTimes(2); + expect(recording.close).toHaveBeenCalledOnce(); + expect(liveSession!.dispose.mock.invocationCallOrder[1]).toBeLessThan( + recording.close.mock.invocationCallOrder[0]!, + ); + expect(innerConfig.shutdown).toHaveBeenCalledOnce(); + expect(unregisterGoalHook).toHaveBeenCalledWith(innerConfig, liveSessionId); + expect(mcpPool?.releaseSession).toHaveBeenCalledWith(liveSessionId); + expect(uiTelemetryService.removeSession).toHaveBeenCalledWith( + liveSessionId, + ); + + mockConnectionState.resolve(); + await agentPromise; + }); }); describe('QwenAgent unstable_listSessions cursor parsing', () => { @@ -11357,6 +11665,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { waitForCloseGateToRelease: ReturnType; waitForActiveTurnsToSettle: ReturnType; sendUpdate: ReturnType; + cancelPendingPrompt: ReturnType; dispose: ReturnType; } | undefined; @@ -11431,6 +11740,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { async (operation: () => Promise): Promise => operation(), ), }; + const toolRegistry = { stop: vi.fn().mockResolvedValue(undefined) }; return { initialize: vi.fn().mockResolvedValue(undefined), shutdown: vi.fn().mockResolvedValue(undefined), @@ -11475,6 +11785,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { .mockReturnValue('/tmp/qwen-runtime-test'), assertCanStartTurn: vi.fn().mockResolvedValue(undefined), getSessionService: vi.fn(), + getToolRegistry: vi.fn().mockReturnValue(toolRegistry), // load path reads back the persisted conversation here and feeds // it to `session.replayHistory`. resume path doesn't read this. getResumedSessionData: vi @@ -11519,11 +11830,11 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { loadSession, }) as unknown as InstanceType, ); - vi.mocked(Session).mockImplementation(() => { + vi.mocked(Session).mockImplementation((_sessionId, config) => { const releaseCloseGate = vi.fn(); const sessionMock = { getId: vi.fn().mockReturnValue('persisted-1'), - getConfig: vi.fn().mockReturnValue(innerConfig), + getConfig: vi.fn().mockReturnValue(config), sendAvailableCommandsUpdate: vi.fn().mockResolvedValue(undefined), replayHistory: vi .fn() @@ -11547,7 +11858,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { cancelPendingPrompt: vi.fn().mockResolvedValue(undefined), assertCanStartTurn: vi.fn().mockResolvedValue(undefined), sendUpdate: vi.fn().mockResolvedValue(undefined), - dispose: vi.fn(), + dispose: vi.fn().mockResolvedValue(undefined), }; lastSessionMock = sessionMock; return sessionMock as unknown as InstanceType; @@ -12359,7 +12670,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { expect(result).toBe(setupError); expect(recording.close).toHaveBeenCalledOnce(); - expect(innerConfig.shutdown).toHaveBeenCalledOnce(); + expect(innerConfig.shutdown).not.toHaveBeenCalled(); expect(recording.hasWriteOwnership()).toBe(true); expect(lastSessionMock?.dispose).toHaveBeenCalledOnce(); expect( @@ -12372,12 +12683,12 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { .map((session) => session.getId()), ).not.toContain('persisted-1'); - innerConfig.shutdown.mockImplementation(async () => { - recording.hasWriteOwnership.mockReturnValue(false); - }); + recording.close.mockResolvedValue(undefined); + recording.hasWriteOwnership.mockReturnValue(false); mockConnectionState.resolve(); await agentPromise; - expect(innerConfig.shutdown).toHaveBeenCalledTimes(2); + expect(recording.close).toHaveBeenCalledTimes(2); + expect(innerConfig.shutdown).toHaveBeenCalledOnce(); expect(recording.hasWriteOwnership()).toBe(false); }); @@ -12408,7 +12719,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { await agentPromise; }); - it('removes a stored session when replay and the first lease release fail', async () => { + it('quarantines a stored session when replay and the first lease release fail', async () => { const replayError = new Error('replay failed'); const innerConfig = bindRestoreMocks({ sessionExists: true, @@ -12451,11 +12762,14 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { .getActiveSessions() .map((session) => session.getId()), ).not.toContain('persisted-1'); - expect(innerConfig.shutdown).toHaveBeenCalledOnce(); - expect(recording.close).toHaveBeenCalledTimes(2); + expect(innerConfig.shutdown).not.toHaveBeenCalled(); + expect(recording.close).toHaveBeenCalledOnce(); mockConnectionState.resolve(); await agentPromise; + expect(failedSession.dispose).toHaveBeenCalledTimes(2); + expect(innerConfig.shutdown).toHaveBeenCalledOnce(); + expect(recording.close).toHaveBeenCalledTimes(3); }); it('cleans Config once when replay fails and lease release succeeds', async () => { @@ -12609,6 +12923,148 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { await agentPromise; }); + it('fully closes a quarantined owner before initializing its replacement', async () => { + const innerConfig = bindRestoreMocks({ + sessionExists: true, + resumedConversation: { messages: [] }, + }); + const { agent, agentPromise } = await spawnAgent(); + await agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + }); + const firstSession = lastSessionMock!; + firstSession.dispose.mockRejectedValueOnce(new Error('dispose failed')); + const closeStoredSession = ( + agent as unknown as { + closeStoredSession(sessionId: string): Promise; + } + ).closeStoredSession.bind(agent); + + await expect(closeStoredSession('persisted-1')).rejects.toThrow( + 'dispose failed', + ); + expect(firstSession.dispose).toHaveBeenCalledOnce(); + + const mcpPool = vi.mocked(McpTransportPool).mock.results.at(-1)?.value as + | { releaseSession: ReturnType } + | undefined; + const replacementConfig = makeRestoreInnerConfig({ + resumedConversation: { messages: [] }, + }); + vi.mocked(loadCliConfig).mockResolvedValueOnce( + replacementConfig as unknown as Config, + ); + let releaseDisposal!: () => void; + firstSession.dispose.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseDisposal = resolve; + }), + ); + const replacement = agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + }); + await vi.waitFor(() => + expect(firstSession.dispose).toHaveBeenCalledTimes(2), + ); + expect(loadCliConfig).toHaveBeenCalledOnce(); + expect(replacementConfig.initialize).not.toHaveBeenCalled(); + expect(Session).toHaveBeenCalledOnce(); + + await expect( + ( + agent as unknown as { + extMethod(method: string, params: unknown): Promise; + } + ).extMethod('qwen/control/session/close', { + sessionId: 'persisted-1', + }), + ).rejects.toMatchObject({ + data: { errorKind: 'session_busy', sessionId: 'persisted-1' }, + }); + + releaseDisposal(); + await expect(replacement).resolves.toMatchObject({ + modes: expect.anything(), + models: expect.anything(), + configOptions: expect.anything(), + }); + + expect(Session).toHaveBeenCalledTimes(2); + expect(lastSessionMock).not.toBe(firstSession); + expect(lastSessionMock?.getConfig()).toBe(replacementConfig); + expect(innerConfig.shutdown).toHaveBeenCalledOnce(); + expect(innerConfig.shutdown).toHaveBeenCalledWith({ + shutdownTelemetry: false, + }); + expect(mcpPool?.releaseSession).toHaveBeenCalledWith('persisted-1'); + expect(innerConfig.shutdown.mock.invocationCallOrder[0]).toBeLessThan( + replacementConfig.initialize.mock.invocationCallOrder[0]!, + ); + expect(mcpPool?.releaseSession.mock.invocationCallOrder[0]).toBeLessThan( + replacementConfig.initialize.mock.invocationCallOrder[0]!, + ); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('closes a quarantined owner before unstable resume initializes its replacement', async () => { + const innerConfig = bindRestoreMocks({ sessionExists: true }); + const { agent, agentPromise } = await spawnAgent(); + await agent.unstable_resumeSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + }); + const firstSession = lastSessionMock!; + firstSession.dispose.mockRejectedValueOnce(new Error('dispose failed')); + const closeStoredSession = ( + agent as unknown as { + closeStoredSession(sessionId: string): Promise; + } + ).closeStoredSession.bind(agent); + + await expect(closeStoredSession('persisted-1')).rejects.toThrow( + 'dispose failed', + ); + const mcpPool = vi.mocked(McpTransportPool).mock.results.at(-1)?.value as + | { releaseSession: ReturnType } + | undefined; + const replacementConfig = makeRestoreInnerConfig(); + vi.mocked(loadCliConfig).mockResolvedValueOnce( + replacementConfig as unknown as Config, + ); + + await expect( + agent.unstable_resumeSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + }), + ).resolves.toMatchObject({ + modes: expect.anything(), + models: expect.anything(), + configOptions: expect.anything(), + }); + + expect(innerConfig.shutdown).toHaveBeenCalledOnce(); + expect(mcpPool?.releaseSession).toHaveBeenCalledWith('persisted-1'); + expect(innerConfig.shutdown.mock.invocationCallOrder[0]).toBeLessThan( + replacementConfig.initialize.mock.invocationCallOrder[0]!, + ); + expect(mcpPool?.releaseSession.mock.invocationCallOrder[0]).toBeLessThan( + replacementConfig.initialize.mock.invocationCallOrder[0]!, + ); + + mockConnectionState.resolve(); + await agentPromise; + }); + it('times out a live load drain and releases its close gate', async () => { bindRestoreMocks({ sessionExists: true, diff --git a/packages/cli/src/acp-integration/acpAgent.ts b/packages/cli/src/acp-integration/acpAgent.ts index 56a07286e8b..71ea1387fcf 100644 --- a/packages/cli/src/acp-integration/acpAgent.ts +++ b/packages/cli/src/acp-integration/acpAgent.ts @@ -239,6 +239,7 @@ import { } from '../config/trustedFolders.js'; import { ACP_PREFLIGHT_KINDS, + SESSION_CLOSE_QUARANTINED_ERROR_KIND, STATUS_SCHEMA_VERSION, SERVE_CONTROL_EXT_METHODS, SERVE_STATUS_EXT_METHODS, @@ -2762,6 +2763,7 @@ export async function runAcpAgent( debugLogger.debug('[ACP] Shutdown signal received, closing streams'); try { + await agentInstance?.beginSessionShutdown(); // Fire SessionEnd hook for all active sessions (aligned with core path) await fireSessionEndOnce(SessionEndReason.Other); await agentInstance?.disposeSessions(); @@ -2800,12 +2802,13 @@ export async function runAcpAgent( try { await connection.closed; + await agentInstance?.beginSessionShutdown(); // Connection closed by IDE - fire SessionEnd hook (aligned with core path) await fireSessionEndOnce(SessionEndReason.PromptInputExit); // Mirror the SIGTERM handler's pool drain on the IDE-initiated // normal close path to avoid leaking shared MCP entries. - await drainPoolBeforeExit('ide_close'); await agentInstance?.disposeSessions(); + await drainPoolBeforeExit('ide_close'); } finally { process.off('SIGTERM', shutdownHandler); process.off('SIGINT', shutdownHandler); @@ -2977,6 +2980,11 @@ interface PendingMcpAuthentication { class QwenAgent implements Agent { private sessions: Map = new Map(); + private readonly closingSessions: Map = new Map(); + private readonly activeSessionOperations = new Set(); + private readonly activeSessionLifecycleOperations = new Set>(); + private disposingSessions = false; + private sessionDisposalPromise: Promise | undefined; private workspaceMcpDiscoveryConfig: Config | undefined; private workspaceMcpDiscoveryPromise: Promise | undefined; private workspaceMcpDiscoveryError: string | undefined; @@ -3266,27 +3274,61 @@ class QwenAgent implements Agent { } } + private quarantineStoredSession(sessionId: string, session: Session): void { + if (this.sessions.get(sessionId) === session) { + this.sessions.delete(sessionId); + this.closingSessions.set(sessionId, session); + } + } + private async removeStoredSessionEntry( sessionId: string, session: Session, cleanupErrors: unknown[] = [], - options: { shutdownConfig?: boolean } = {}, + options: { shutdownConfig?: boolean; requireFlush?: boolean } = {}, ): Promise { - if (this.sessions.get(sessionId) !== session) return; + if ( + this.sessions.get(sessionId) !== session && + this.closingSessions.get(sessionId) !== session + ) { + return; + } + this.quarantineStoredSession(sessionId, session); + await session.dispose(); + const config = session.getConfig(); + const recorder = config.getChatRecordingService(); + recorder?.finalize(); + let flushError: unknown; try { - session.dispose(); + await recorder?.flush(); } catch (error) { - cleanupErrors.push(error); + flushError = error; } + if (flushError !== undefined && options.requireFlush === true) { + throw flushError; + } + + let closeError: unknown; + try { + await recorder?.close(); + } catch (error) { + closeError = error; + } + if (recorder?.hasWriteOwnership()) { + throw closeError ?? new SessionWriterUnavailableError(); + } + if (flushError !== undefined) cleanupErrors.push(flushError); + if (closeError !== undefined) cleanupErrors.push(closeError); + if (options.shutdownConfig !== false) { try { - await session.getConfig().shutdown({ shutdownTelemetry: false }); + await config.shutdown({ shutdownTelemetry: false }); } catch (error) { cleanupErrors.push(error); } } try { - unregisterGoalHook(session.getConfig(), sessionId); + unregisterGoalHook(config, sessionId); } catch (error) { cleanupErrors.push(error); } @@ -3300,7 +3342,12 @@ class QwenAgent implements Agent { } catch (error) { cleanupErrors.push(error); } - this.sessions.delete(sessionId); + if (this.sessions.get(sessionId) === session) { + this.sessions.delete(sessionId); + } + if (this.closingSessions.get(sessionId) === session) { + this.closingSessions.delete(sessionId); + } if (cleanupErrors.length > 0) { debugLogger.warn( `Session ${sessionId} closed after ${cleanupErrors.length} cleanup failure(s): ${cleanupErrors @@ -3376,6 +3423,14 @@ class QwenAgent implements Agent { throw error; } + private ownsSessionConfig(config: Config): boolean { + const sessionId = config.getSessionId(); + return ( + this.sessions.get(sessionId)?.getConfig() === config || + this.closingSessions.get(sessionId)?.getConfig() === config + ); + } + private pendingConfigCleanupKey( runtimeBaseDir: string, sessionId: string, @@ -3420,6 +3475,14 @@ class QwenAgent implements Agent { waitForCloseGate?: boolean; }, ): Promise { + const closingSession = this.closingSessions.get(sessionId); + if (closingSession) { + await this.removeStoredSessionEntry(sessionId, closingSession, [], { + shutdownConfig: opts?.shutdownConfig, + requireFlush: opts?.requireFlush, + }); + return; + } const session = this.sessions.get(sessionId); if (!session) { this.mcpPool?.releaseSession(sessionId); @@ -3460,32 +3523,9 @@ class QwenAgent implements Agent { 'close', ); - recorder?.finalize(); - let flushError: unknown; - try { - await recorder?.flush(); - } catch (error) { - flushError = error; - } - if (flushError !== undefined && requireFlush) { - throw flushError; - } - - let closeError: unknown; - try { - await recorder?.close(); - } catch (error) { - closeError = error; - } - if (recorder?.hasWriteOwnership()) { - throw closeError ?? new SessionWriterUnavailableError(); - } - - const cleanupErrors: unknown[] = []; - if (flushError !== undefined) cleanupErrors.push(flushError); - if (closeError !== undefined) cleanupErrors.push(closeError); - await this.removeStoredSessionEntry(sessionId, session, cleanupErrors, { + await this.removeStoredSessionEntry(sessionId, session, [], { shutdownConfig: opts?.shutdownConfig, + requireFlush, }); removedFromStore = true; } finally { @@ -3509,7 +3549,63 @@ class QwenAgent implements Agent { await this.closeStoredSession(sessionId, opts); } - async disposeSessions(): Promise { + private acquireSessionOperation(sessionId: string): () => void { + if (this.activeSessionOperations.has(sessionId)) { + throw new RequestError( + -32000, + `Session operation already in progress: ${sessionId}`, + { errorKind: 'session_busy', sessionId }, + ); + } + this.activeSessionOperations.add(sessionId); + return () => { + this.activeSessionOperations.delete(sessionId); + }; + } + + private acquireSessionLifecycleOperation(): () => void { + if (this.disposingSessions) { + throw new RequestError(-32000, 'Session shutdown is in progress.', { + errorKind: 'session_busy', + }); + } + let resolve!: () => void; + const completion = new Promise((done) => { + resolve = done; + }); + this.activeSessionLifecycleOperations.add(completion); + let released = false; + return () => { + if (released) return; + released = true; + this.activeSessionLifecycleOperations.delete(completion); + resolve(); + }; + } + + private async runSessionLifecycleOperation( + operation: () => Promise, + ): Promise { + const release = this.acquireSessionLifecycleOperation(); + try { + return await operation(); + } finally { + release(); + } + } + + async beginSessionShutdown(): Promise { + this.disposingSessions = true; + await Promise.allSettled([...this.activeSessionLifecycleOperations]); + } + + disposeSessions(): Promise { + this.sessionDisposalPromise ??= this.disposeSessionsOnce(); + return this.sessionDisposalPromise; + } + + private async disposeSessionsOnce(): Promise { + await this.beginSessionShutdown(); for (const generation of this.generationControllers.values()) { generation.controller.abort(); } @@ -3521,6 +3617,11 @@ class QwenAgent implements Agent { }), ), ); + await Promise.allSettled( + [...this.closingSessions.entries()].map(([sessionId, session]) => + this.removeStoredSessionEntry(sessionId, session), + ), + ); await Promise.allSettled( [...this.pendingConfigCleanup.values()] .flatMap((configs) => [...configs]) @@ -3731,6 +3832,12 @@ class QwenAgent implements Agent { } async newSession(params: NewSessionRequest): Promise { + return this.runSessionLifecycleOperation(() => this.newSessionImpl(params)); + } + + private async newSessionImpl( + params: NewSessionRequest, + ): Promise { const { cwd, mcpServers } = params; const parentContext = extractDaemonTraceContext(params); return await withDaemonSpan( @@ -3762,9 +3869,7 @@ class QwenAgent implements Agent { ); } catch (error) { return this.cleanupAfterRequestFailure(error, async () => { - if ( - this.sessions.get(config.getSessionId())?.getConfig() !== config - ) { + if (!this.ownsSessionConfig(config)) { await this.cleanupUnstoredConfig(config); } }); @@ -3782,6 +3887,14 @@ class QwenAgent implements Agent { } async loadSession(params: LoadSessionRequest): Promise { + return this.runSessionLifecycleOperation(() => + this.loadSessionImpl(params), + ); + } + + private async loadSessionImpl( + params: LoadSessionRequest, + ): Promise { // Load per-request settings BEFORE the existence check: the check must // resolve `advanced.runtimeOutputDir` from THIS request's cwd, not from // whichever settings a concurrent handler loaded last. @@ -3854,138 +3967,144 @@ class QwenAgent implements Agent { if (!exists) { throw RequestError.resourceNotFound(`session:${params.sessionId}`); } - // Adopt into the "latest loaded" cache only once the session is - // confirmed — a failed probe for a stale id must not repoint - // agent-level readers at this request's workspace. - this.settings = settings; - - const config = await this.newSessionConfig( - params.cwd, - // `LoadSessionRequest.mcpServers` is required in today's ACP - // schema, but mirror `unstable_resumeSession` and tolerate a - // future loosening — `newSessionConfig` iterates the list, so - // a `null`/`undefined` would otherwise throw `TypeError`. - params.mcpServers ?? [], - settings, + const releaseSessionOperation = this.acquireSessionOperation( params.sessionId, - true, ); - const sessionData = config.getResumedSessionData(); - const bulkReplay = isBulkLoadReplayRequest(params); - const replayPageSize = bulkReplay - ? getLoadReplayPageSize(params) - : undefined; - let session: Session; try { - await this.ensureAuthenticated(config); - this.setupFileSystem(config); - session = await this.createAndStoreSession( - config, + if (this.closingSessions.has(params.sessionId)) { + await this.closeStoredSession(params.sessionId); + } + // Adopt into the "latest loaded" cache only once the session is + // confirmed — a failed probe for a stale id must not repoint + // agent-level readers at this request's workspace. + this.settings = settings; + + const config = await this.newSessionConfig( + params.cwd, + // `LoadSessionRequest.mcpServers` is required in today's ACP + // schema, but mirror `unstable_resumeSession` and tolerate a + // future loosening — `newSessionConfig` iterates the list, so + // a `null`/`undefined` would otherwise throw `TypeError`. + params.mcpServers ?? [], settings, - sessionData, - bulkReplay - ? { replayHistory: false, startPostReplayServices: false } - : {}, + params.sessionId, + true, ); - } catch (error) { - return this.cleanupAfterRequestFailure(error, async () => { - if (this.sessions.get(config.getSessionId())?.getConfig() !== config) { - await this.cleanupUnstoredConfig(config); - } - }); - } - let replayEnvelope: BridgeLoadReplayEnvelope | undefined; - if (bulkReplay) { + const sessionData = config.getResumedSessionData(); + const bulkReplay = isBulkLoadReplayRequest(params); + const replayPageSize = bulkReplay + ? getLoadReplayPageSize(params) + : undefined; + let session: Session; try { - const records = sessionData?.conversation.messages; - let replayUpdates: SessionUpdate[] = []; - if (records) { - session.primeTurnFromHistory(records); - const replayPage = selectRecentHistoryRecords( - records, - replayPageSize, - ); - const replayUsage = createReplayCumulativeUsage(); - const replay = await collectHistoryReplayUpdates({ - sessionId: params.sessionId, - config, - records: replayPage.records, - gaps: sessionData?.historyGaps, - cumulativeUsage: replayUsage, - // A resume: the goal restore runs right after this. - supersedeUnrestorableGoal: true, - logger: debugLogger, - }); - replayUpdates = replay.updates; - copyCumulativeUsage(session.cumulativeUsage, replayUsage); - if (replay.replayError !== undefined) { - replayEnvelope = { + await this.ensureAuthenticated(config); + this.setupFileSystem(config); + session = await this.createAndStoreSession( + config, + settings, + sessionData, + bulkReplay + ? { replayHistory: false, startPostReplayServices: false } + : {}, + ); + } catch (error) { + return this.cleanupAfterRequestFailure(error, async () => { + if (!this.ownsSessionConfig(config)) { + await this.cleanupUnstoredConfig(config); + } + }); + } + let replayEnvelope: BridgeLoadReplayEnvelope | undefined; + if (bulkReplay) { + try { + const records = sessionData?.conversation.messages; + let replayUpdates: SessionUpdate[] = []; + if (records) { + session.primeTurnFromHistory(records); + const replayPage = selectRecentHistoryRecords( + records, + replayPageSize, + ); + const replayUsage = createReplayCumulativeUsage(); + const replay = await collectHistoryReplayUpdates({ + sessionId: params.sessionId, + config, + records: replayPage.records, + gaps: sessionData?.historyGaps, + cumulativeUsage: replayUsage, + // A resume: the goal restore runs right after this. + supersedeUnrestorableGoal: true, + logger: debugLogger, + }); + replayUpdates = replay.updates; + copyCumulativeUsage(session.cumulativeUsage, replayUsage); + if (replay.replayError !== undefined) { + replayEnvelope = { + v: LOAD_REPLAY_VERSION, + updates: replayUpdates, + partial: true, + replayError: replay.replayError, + ...(replayPage.hasMore ? { hasMore: true } : {}), + }; + } + replayEnvelope ??= { v: LOAD_REPLAY_VERSION, updates: replayUpdates, - partial: true, - replayError: replay.replayError, ...(replayPage.hasMore ? { hasMore: true } : {}), }; } replayEnvelope ??= { v: LOAD_REPLAY_VERSION, updates: replayUpdates, - ...(replayPage.hasMore ? { hasMore: true } : {}), }; + session.installRewriter(); + session.startCronScheduler(); + } catch (err) { + return this.cleanupAfterRequestFailure(err, () => + this.discardStoredSessionIfCurrent(params.sessionId, session), + ); } - replayEnvelope ??= { - v: LOAD_REPLAY_VERSION, - updates: replayUpdates, - }; - session.installRewriter(); - session.startCronScheduler(); - } catch (err) { - return this.cleanupAfterRequestFailure(err, async () => { - try { - await this.discardStoredSessionIfCurrent(params.sessionId, session); - } catch (cleanupError) { - await this.removeStoredSessionEntry( - params.sessionId, - session, - [cleanupError], - { - shutdownConfig: false, - }, - ); - await this.cleanupUnstoredConfig(config); - } - }); } - } - await this.#restoreWorktreeOnResume(config, session); - this.#restoreGoalOnResume(config, session); + await this.#restoreWorktreeOnResume(config, session); + this.#restoreGoalOnResume(config, session); - const modesData = this.buildModesData(config); - const availableModels = this.buildAvailableModels(config); - const configOptions = this.buildConfigOptions(config); + const modesData = this.buildModesData(config); + const availableModels = this.buildAvailableModels(config); + const configOptions = this.buildConfigOptions(config); - const response: LoadSessionResponse = { - modes: modesData, - models: availableModels, - configOptions, - ...(sessionData?.artifactSnapshot - ? { artifactSnapshot: sessionData.artifactSnapshot } - : {}), - } as LoadSessionResponse; - if (!replayEnvelope) { - return response; + const response: LoadSessionResponse = { + modes: modesData, + models: availableModels, + configOptions, + ...(sessionData?.artifactSnapshot + ? { artifactSnapshot: sessionData.artifactSnapshot } + : {}), + } as LoadSessionResponse; + if (!replayEnvelope) { + return response; + } + return { + ...response, + _meta: { + [LOAD_REPLAY_META_KEY]: replayEnvelope, + }, + }; + } finally { + releaseSessionOperation(); } - return { - ...response, - _meta: { - [LOAD_REPLAY_META_KEY]: replayEnvelope, - }, - }; } async unstable_resumeSession( params: ResumeSessionRequest, + ): Promise { + return this.runSessionLifecycleOperation(() => + this.unstableResumeSessionImpl(params), + ); + } + + private async unstableResumeSessionImpl( + params: ResumeSessionRequest, ): Promise { // Same per-request settings discipline as `loadSession`. const settings = loadSettingsCached(params.cwd); @@ -4018,49 +4137,59 @@ class QwenAgent implements Agent { if (!exists) { throw RequestError.resourceNotFound(`session:${params.sessionId}`); } - this.settings = settings; - - const config = await this.newSessionConfig( - params.cwd, - params.mcpServers ?? [], - settings, + const releaseSessionOperation = this.acquireSessionOperation( params.sessionId, - true, ); - let session: Session; try { - await this.ensureAuthenticated(config); - this.setupFileSystem(config); - session = await this.createAndStoreSession( - config, + if (this.closingSessions.has(params.sessionId)) { + await this.closeStoredSession(params.sessionId); + } + this.settings = settings; + + const config = await this.newSessionConfig( + params.cwd, + params.mcpServers ?? [], settings, - config.getResumedSessionData(), - { replayHistory: false }, + params.sessionId, + true, ); - } catch (error) { - return this.cleanupAfterRequestFailure(error, async () => { - if (this.sessions.get(config.getSessionId())?.getConfig() !== config) { - await this.cleanupUnstoredConfig(config); - } - }); - } + let session: Session; + try { + await this.ensureAuthenticated(config); + this.setupFileSystem(config); + session = await this.createAndStoreSession( + config, + settings, + config.getResumedSessionData(), + { replayHistory: false }, + ); + } catch (error) { + return this.cleanupAfterRequestFailure(error, async () => { + if (!this.ownsSessionConfig(config)) { + await this.cleanupUnstoredConfig(config); + } + }); + } - await this.#restoreWorktreeOnResume(config, session); - this.#restoreGoalOnResume(config, session); + await this.#restoreWorktreeOnResume(config, session); + this.#restoreGoalOnResume(config, session); - const modesData = this.buildModesData(config); - const availableModels = this.buildAvailableModels(config); - const configOptions = this.buildConfigOptions(config); + const modesData = this.buildModesData(config); + const availableModels = this.buildAvailableModels(config); + const configOptions = this.buildConfigOptions(config); - const sessionData = config.getResumedSessionData(); - return { - modes: modesData, - models: availableModels, - configOptions, - ...(sessionData?.artifactSnapshot - ? { artifactSnapshot: sessionData.artifactSnapshot } - : {}), - } as ResumeSessionResponse; + const sessionData = config.getResumedSessionData(); + return { + modes: modesData, + models: availableModels, + configOptions, + ...(sessionData?.artifactSnapshot + ? { artifactSnapshot: sessionData.artifactSnapshot } + : {}), + } as ResumeSessionResponse; + } finally { + releaseSessionOperation(); + } } /** @@ -7964,13 +8093,39 @@ class QwenAgent implements Agent { 'Invalid session close drain timeout', ); } - await this.closeStoredSession(sessionId, { - requireFlush: params['requireFlush'] === true, - ...(typeof rawDrainTimeoutMs === 'number' - ? { drainTimeoutMs: rawDrainTimeoutMs } - : {}), - }); - return { sessionId, closed: true }; + const releaseLifecycleOperation = + this.acquireSessionLifecycleOperation(); + try { + const releaseSessionOperation = + this.acquireSessionOperation(sessionId); + try { + try { + await this.closeStoredSession(sessionId, { + requireFlush: params['requireFlush'] === true, + ...(typeof rawDrainTimeoutMs === 'number' + ? { drainTimeoutMs: rawDrainTimeoutMs } + : {}), + }); + } catch (error) { + if (this.closingSessions.has(sessionId)) { + throw new RequestError( + -32603, + error instanceof Error ? error.message : String(error), + { + errorKind: SESSION_CLOSE_QUARANTINED_ERROR_KIND, + sessionId, + }, + ); + } + throw error; + } + return { sessionId, closed: true }; + } finally { + releaseSessionOperation(); + } + } finally { + releaseLifecycleOperation(); + } } case SERVE_CONTROL_EXT_METHODS.sessionCd: { const sessionId = params['sessionId']; @@ -10176,7 +10331,7 @@ class QwenAgent implements Agent { await geminiClient.initialize(); } - if (this.sessions.has(sessionId)) { + if (this.sessions.has(sessionId) || this.closingSessions.has(sessionId)) { throw new Error(`Session ${sessionId} is already active.`); } @@ -10223,11 +10378,8 @@ class QwenAgent implements Agent { shutdownConfig: false, }); } catch (cleanupError) { - await this.removeStoredSessionEntry( - sessionId, - session, - [cleanupError], - { shutdownConfig: false }, + debugLogger.warn( + `Session ${sessionId} cleanup failed while preserving the creation error: ${cleanupError instanceof Error ? cleanupError.message : String(cleanupError)}`, ); } throw error; diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index c7458d707e1..5ef4a57cd03 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -336,6 +336,8 @@ describe('Session', () => { setNotificationCallback: ReturnType; hasUnfinalizedTasks: ReturnType; getAll: ReturnType; + abortAll: ReturnType; + abortAllAndWait: ReturnType; }; let mockMonitorRegistry: { setNotificationCallback: ReturnType; @@ -479,6 +481,8 @@ describe('Session', () => { setNotificationCallback: vi.fn(), hasUnfinalizedTasks: vi.fn().mockReturnValue(false), getAll: vi.fn().mockReturnValue([]), + abortAll: vi.fn(), + abortAllAndWait: vi.fn().mockResolvedValue(undefined), }; mockMonitorRegistry = { setNotificationCallback: vi.fn(), @@ -16055,6 +16059,9 @@ describe('Session', () => { expect(internals.notificationQueue).toHaveLength(0); expect(internals.cronQueue).toHaveLength(0); expect(internals.notificationProcessing).toBe(false); + expect(mockBackgroundTaskRegistry.abortAll).toHaveBeenCalledWith({ + notify: false, + }); expect( mockBackgroundTaskRegistry.setNotificationCallback, ).toHaveBeenLastCalledWith(undefined); @@ -16100,6 +16107,53 @@ describe('Session', () => { expect(internals.cronCompletion).toBeNull(); }); + it('waits for aborted background agents before disposal completes', async () => { + let release!: () => void; + mockBackgroundTaskRegistry.abortAllAndWait.mockReturnValueOnce( + new Promise((resolve) => { + release = resolve; + }), + ); + + let settled = false; + const disposal = session.dispose().then(() => { + settled = true; + }); + + await vi.waitFor(() => { + expect(mockBackgroundTaskRegistry.abortAllAndWait).toHaveBeenCalledWith( + { notify: false }, + ); + }); + expect(settled).toBe(false); + expect(mockBackgroundTaskRegistry.abortAll).toHaveBeenCalledOnce(); + + release(); + await disposal; + expect(settled).toBe(true); + }); + + it('rejects prompts once disposal has started', async () => { + let release!: () => void; + mockBackgroundTaskRegistry.abortAllAndWait.mockReturnValueOnce( + new Promise((resolve) => { + release = resolve; + }), + ); + + const disposal = session.dispose(); + await expect( + session.prompt({ + prompt: [], + sessionId: 'test-session-id', + }), + ).rejects.toThrow('Session is closing'); + expect(mockChat.sendMessageStream).not.toHaveBeenCalled(); + + release(); + await disposal; + }); + it('is idempotent — repeated dispose() calls do not throw or re-register', () => { const internals = session as unknown as SessionInternals; session.dispose(); diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index accd969fa2d..5c03136d5a6 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -1629,12 +1629,19 @@ export class Session implements SessionContext { return this.createdAt; } - dispose(): void { + async dispose(): Promise { + const pendingCompletions = [ + this.pendingPromptCompletion, + this.cronCompletion, + this.notificationCompletion, + ].filter((completion): completion is Promise => completion !== null); + this.pendingPrompt?.abort(); this.disposed = true; this.closing = true; this.resolveCloseGate?.(); this.resolveCloseGate = null; this.closeGateCompletion = null; + this.config.getBackgroundTaskRegistry().abortAll({ notify: false }); this.todoStopGuardQueuedPromptPriority = false; this.todoStopGuardDrainAutomaticQueuesWhenIdle = false; this.todoStopGuard.clearTrust(); @@ -1668,6 +1675,10 @@ export class Session implements SessionContext { this.unsubscribeChatRecordingFailure = undefined; this.config.setSubSessionSpawner(undefined); clearGoalTerminalObserver(this.sessionId); + await Promise.allSettled(pendingCompletions); + + const registry = this.config.getBackgroundTaskRegistry(); + await registry.abortAllAndWait({ notify: false }); } /** diff --git a/packages/cli/src/ui/hooks/useBranchCommand.test.ts b/packages/cli/src/ui/hooks/useBranchCommand.test.ts index fd1230931d1..684cd553b2e 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.test.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.test.ts @@ -33,6 +33,22 @@ describe('useBranchCommand', () => { let setSessionName: ReturnType; let remount: ReturnType; let addItem: ReturnType; + let backgroundTaskRegistry: { + hasRunningTasks: ReturnType; + reset: ReturnType; + }; + let monitorRegistry: { + getRunning: ReturnType; + reset: ReturnType; + }; + let backgroundShellRegistry: { + hasRunningEntries: ReturnType; + reset: ReturnType; + }; + let workflowRunRegistry: { + hasRunningEntries: ReturnType; + reset: ReturnType; + }; // Mock Config shape covers only what useBranchCommand touches. // eslint-disable-next-line @typescript-eslint/no-explicit-any let config: any; @@ -85,6 +101,22 @@ describe('useBranchCommand', () => { setSessionName = vi.fn(); remount = vi.fn(); addItem = vi.fn(); + backgroundTaskRegistry = { + hasRunningTasks: vi.fn().mockReturnValue(false), + reset: vi.fn(), + }; + monitorRegistry = { + getRunning: vi.fn().mockReturnValue([]), + reset: vi.fn(), + }; + backgroundShellRegistry = { + hasRunningEntries: vi.fn().mockReturnValue(false), + reset: vi.fn(), + }; + workflowRunRegistry = { + hasRunningEntries: vi.fn().mockReturnValue(false), + reset: vi.fn(), + }; config = { getSessionId: () => '12345678-aaaa-bbbb-cccc-dddddddddddd', getSessionService: () => ({ @@ -96,11 +128,106 @@ describe('useBranchCommand', () => { }), getChatRecordingService: () => ({ finalize, flush }), getGeminiClient: () => ({ initialize: vi.fn() }), + getBackgroundTaskRegistry: () => backgroundTaskRegistry, + getMonitorRegistry: () => monitorRegistry, + getBackgroundShellRegistry: () => backgroundShellRegistry, + getWorkflowRunRegistry: () => workflowRunRegistry, startNewSession: startNewSessionConfig, getDebugLogger: () => ({ warn: vi.fn() }), }; }); + it('does not branch while the current session has running background work', async () => { + backgroundTaskRegistry.hasRunningTasks.mockReturnValue(true); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('my-branch'); + }); + + expect(forkSession).not.toHaveBeenCalled(); + expect(startNewSessionConfig).not.toHaveBeenCalled(); + expect(backgroundTaskRegistry.reset).not.toHaveBeenCalled(); + expect(addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'error', + text: expect.stringContaining('Stop the current session'), + }), + expect.any(Number), + ); + }); + + it('clears terminal background state after branch initialization succeeds', async () => { + const order: string[] = []; + backgroundTaskRegistry.reset.mockImplementation(() => order.push('reset')); + startNewSessionConfig.mockImplementation(() => order.push('switch')); + config.getGeminiClient = () => ({ + initialize: vi.fn(async () => { + order.push('initialize'); + }), + }); + loadHistory.mockImplementation(() => { + order.push('ui'); + }); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('my-branch'); + }); + + expect(order).toEqual(['switch', 'initialize', 'ui', 'reset']); + expect(backgroundTaskRegistry.reset).toHaveBeenCalledOnce(); + expect(monitorRegistry.reset).toHaveBeenCalledOnce(); + expect(backgroundShellRegistry.reset).toHaveBeenCalledOnce(); + expect(workflowRunRegistry.reset).toHaveBeenCalledOnce(); + }); + + it('preserves terminal background state when branch initialization rolls back', async () => { + const initialize = vi + .fn() + .mockRejectedValueOnce(new Error('branch init failed')) + .mockResolvedValueOnce(undefined); + config.getGeminiClient = () => ({ initialize }); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('my-branch'); + }); + + expect(startNewSessionConfig).toHaveBeenCalledTimes(2); + expect(startNewSessionConfig.mock.calls[1]?.[0]).toBe( + '12345678-aaaa-bbbb-cccc-dddddddddddd', + ); + expect(backgroundTaskRegistry.reset).not.toHaveBeenCalled(); + expect(monitorRegistry.reset).not.toHaveBeenCalled(); + expect(backgroundShellRegistry.reset).not.toHaveBeenCalled(); + expect(workflowRunRegistry.reset).not.toHaveBeenCalled(); + }); + + it('preserves terminal background state when the UI swap throws', async () => { + loadHistory.mockImplementation(() => { + throw new Error('history load failed'); + }); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('my-branch'); + }); + + expect(startNewSessionConfig).toHaveBeenCalledTimes(2); + expect(backgroundTaskRegistry.reset).not.toHaveBeenCalled(); + expect(monitorRegistry.reset).not.toHaveBeenCalled(); + expect(backgroundShellRegistry.reset).not.toHaveBeenCalled(); + expect(workflowRunRegistry.reset).not.toHaveBeenCalled(); + expect(addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'error', + text: expect.stringContaining('history load failed'), + }), + expect.any(Number), + ); + }); + it('persists and reloads the title before switching core or UI', async () => { // The parent snapshot must come AFTER finalize(): finalize() appends a // trailing custom_title record to the parent JSONL, advancing the diff --git a/packages/cli/src/ui/hooks/useBranchCommand.ts b/packages/cli/src/ui/hooks/useBranchCommand.ts index b153834a7aa..eee3258bbb6 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.ts @@ -21,6 +21,10 @@ import { restoreGoalFromHistory } from '../utils/restoreGoal.js'; import type { UseHistoryManagerReturn } from './useHistoryManager.js'; import type { LoadedSettings } from '../../config/settings.js'; import { t } from '../../i18n/index.js'; +import { + hasBlockingBackgroundWork, + resetBackgroundStateForSessionSwitch, +} from '../utils/backgroundWorkUtils.js'; /** * Derives a short one-line title from the first *real* user message in the @@ -69,6 +73,9 @@ export interface UseBranchCommandResult { handleBranch: (name?: string) => Promise; } +const BACKGROUND_WORK_BRANCH_BLOCKED_MESSAGE = + "Stop the current session's running background tasks before branching the conversation."; + /** * Orchestrates `/branch`: * 1. Capture the current (soon-to-be-parent) sessionId for the resume hint. @@ -93,6 +100,17 @@ export function useBranchCommand( async (name?: string) => { if (!config) return; + if (hasBlockingBackgroundWork(config)) { + historyManager.addItem( + { + type: 'error', + text: BACKGROUND_WORK_BRANCH_BLOCKED_MESSAGE, + }, + Date.now(), + ); + return; + } + const oldSessionId = config.getSessionId(); const newSessionId = randomUUID(); const sessionService = config.getSessionService(); @@ -192,6 +210,7 @@ export function useBranchCommand( historyManager.clearItems(); historyManager.loadHistory(uiHistoryItems); uiSwapped = true; + resetBackgroundStateForSessionSwitch(config); // 9. Re-arm /goal under the fork's new sessionId. The branched JSONL // is a verbatim copy of the parent's, so an active goal sentinel diff --git a/packages/core/src/agents/background-agent-resume.test.ts b/packages/core/src/agents/background-agent-resume.test.ts index be71aaaa20c..e86f7cd5864 100644 --- a/packages/core/src/agents/background-agent-resume.test.ts +++ b/packages/core/src/agents/background-agent-resume.test.ts @@ -113,6 +113,7 @@ describe('BackgroundAgentResumeService', () => { isTrustedFolder: () => true, isInteractive: () => false, getProjectRoot: () => tempDir, + getWorkingDir: vi.fn().mockReturnValue(tempDir), getCliVersion: () => 'test-version', getGeminiClient: () => undefined, getSkillManager: () => undefined, @@ -124,6 +125,7 @@ describe('BackgroundAgentResumeService', () => { return { service: new BackgroundAgentResumeService(config), + config, subagentManager, hookSystem, monitorRegistry, @@ -1513,6 +1515,10 @@ describe('BackgroundAgentResumeService', () => { let releaseExecute: (() => void) | undefined; const execute = vi.fn( + (_context?: unknown, _signal?: AbortSignal, _options?: unknown) => + Promise.resolve(), + ); + execute.mockImplementationOnce( () => new Promise((resolve) => { releaseExecute = resolve; @@ -1549,12 +1555,17 @@ describe('BackgroundAgentResumeService', () => { await Promise.all([first, second]); await vi.waitFor(() => { expect(registry.get(agentId)?.status).toBe('completed'); + expect(execute).toHaveBeenCalledTimes(2); + }); + expect(execute.mock.calls[1]?.[2]).toEqual({ + resetStats: false, + initialExternalInputs: ['second message'], }); const provider = subagent.setExternalMessageProvider.mock.calls[0]?.[0] as | (() => string[]) | undefined; expect(provider).toBeDefined(); - expect(provider?.()).toEqual(['second message']); + expect(provider?.()).toEqual([]); }); it('routes owned monitor notifications into a resumed agent queue', async () => { @@ -1656,6 +1667,10 @@ describe('BackgroundAgentResumeService', () => { await expect(waitPromise).resolves.toEqual([]); releaseExecute?.(); await resume; + await vi.waitFor(() => { + expect(registry.get(agentId)?.status).toBe('completed'); + }); + expect(registry.disposeResidentAgent(agentId)).toBe(true); await vi.waitFor(() => { expect(monitorRegistry.setAgentNotificationCallback).toHaveBeenCalledWith( agentId, @@ -2348,7 +2363,7 @@ describe('BackgroundAgentResumeService', () => { ); }); - it('revives a completed background agent from its transcript and bumps resumeCount', async () => { + it('reconstructs a completed agent once, then reuses and disposes its resident runtime', async () => { const sessionId = 'session-revive'; const agentId = 'agent-revive'; const metaPath = getAgentMetaPath(tempDir, sessionId, agentId); @@ -2407,8 +2422,33 @@ describe('BackgroundAgentResumeService', () => { }); registry.complete(agentId, 'All done'); + let releaseStopHook: (() => void) | undefined; + const stopHookGate = new Promise((resolve) => { + releaseStopHook = resolve; + }); + let markStopHookStarted: (() => void) | undefined; + const stopHookStarted = new Promise((resolve) => { + markStopHookStarted = resolve; + }); + const hookSystem = { + fireSubagentStartEvent: vi.fn().mockResolvedValue(undefined), + fireSubagentStopEvent: vi + .fn() + .mockImplementationOnce(async () => { + markStopHookStarted?.(); + await stopHookGate; + return undefined; + }) + .mockResolvedValue(undefined), + }; + const notification = vi.fn(); + registry.setNotificationCallback(notification); const execute = vi.fn( - async (_context: { get: (key: string) => unknown }) => undefined, + async ( + _context: { get: (key: string) => unknown }, + _signal?: AbortSignal, + _options?: unknown, + ) => undefined, ); const subagent = { execute, @@ -2423,31 +2463,111 @@ describe('BackgroundAgentResumeService', () => { getFinalText: () => 'iterated', }; - const { service, subagentManager } = createService(); + let resolveDispose!: () => void; + const disposeGate = new Promise((resolve) => { + resolveDispose = resolve; + }); + const dispose = vi.fn(() => disposeGate); + const { service, subagentManager, config } = createService({ hookSystem }); + const trackAgentExecution = vi.spyOn(registry, 'trackAgentExecution'); + const getModel = vi.spyOn(config, 'getModel').mockReturnValue('model-a'); subagentManager.createAgentHeadless.mockResolvedValue({ subagent, - dispose: vi.fn().mockResolvedValue(undefined), + dispose, }); - const revived = await service.reviveCompletedBackgroundAgent( + const revive = service.reviveCompletedBackgroundAgent( agentId, 'now write the summary', ); + await stopHookStarted; + expect(registry.get(agentId)?.status).toBe('running'); + expect(notification).not.toHaveBeenCalled(); + expect(registry.queueMessage(agentId, 'late message')).toBe(true); + expect( + registry.queueExternalInput(agentId, { + kind: 'notification', + text: 'ready', + }), + ).toBe(true); + releaseStopHook?.(); + const revived = await revive; expect(revived).toBeDefined(); + expect(trackAgentExecution).toHaveBeenCalledOnce(); + expect(revived?.model).toBe('model-a'); expect(subagentManager.createAgentHeadless).toHaveBeenCalledTimes(1); - expect(execute).toHaveBeenCalledTimes(1); + const [, runtimeConfig, createOptions] = + subagentManager.createAgentHeadless.mock.calls[0]!; + expect((runtimeConfig as Config).getModel()).toBe('model-a'); + expect(createOptions).toEqual( + expect.objectContaining({ + modelConfigOverrides: { model: 'model-a' }, + }), + ); + await vi.waitFor(() => { + expect(execute).toHaveBeenCalledTimes(2); + expect(registry.get(agentId)?.status).toBe('completed'); + }); const contextArg = execute.mock.calls[0]?.[0]; expect(contextArg).toBeDefined(); expect(contextArg?.get('task_prompt')).toBe('now write the summary'); - await vi.waitFor(() => { - expect(registry.get(agentId)?.status).toBe('completed'); + expect(execute.mock.calls[1]?.[2]).toEqual({ + resetStats: false, + initialExternalInputs: [ + 'late message', + { + kind: 'notification', + text: 'ready', + }, + ], }); + expect(notification).toHaveBeenCalledTimes(1); const meta = JSON.parse(fs.readFileSync(metaPath, 'utf8')); expect(meta.resumeCount).toBe(1); + expect(meta.model).toBe('model-a'); + expect(meta.persistedCliFlags.model).toBe('model-a'); expect(fs.statSync(sessionDir).mtime.getTime()).toBeGreaterThan( oldSessionMtime.getTime(), ); + + getModel.mockReturnValue('model-b'); + expect(registry.continueResidentAgent(agentId, 'tighten the summary')).toBe( + true, + ); + expect(trackAgentExecution).toHaveBeenCalledTimes(2); + expect(registry.get(agentId)?.status).toBe('running'); + await vi.waitFor(() => { + expect(execute).toHaveBeenCalledTimes(3); + expect(registry.get(agentId)?.status).toBe('completed'); + }); + expect(subagentManager.createAgentHeadless).toHaveBeenCalledTimes(1); + expect((runtimeConfig as Config).getModel()).toBe('model-a'); + const hotContextArg = execute.mock.calls[2]?.[0]; + expect(hotContextArg?.get('task_prompt')).toBe('tighten the summary'); + expect(readAgentMeta(metaPath)?.resumeCount).toBe(2); + expect(dispose).not.toHaveBeenCalled(); + + vi.mocked(config.getWorkingDir).mockReturnValue( + path.join(tempDir, 'relocated'), + ); + expect(registry.continueResidentAgent(agentId, 'again')).toBe(false); + expect(dispose).toHaveBeenCalledTimes(1); + + let cleanupSettled = false; + const cleanup = registry.abortAllAndWait({ notify: false }).then(() => { + cleanupSettled = true; + }); + await Promise.resolve(); + expect(cleanupSettled).toBe(false); + resolveDispose(); + await cleanup; + expect(cleanupSettled).toBe(true); + + registry.reset(); + + expect(dispose).toHaveBeenCalledTimes(1); + expect(registry.continueResidentAgent(agentId, 'after reset')).toBe(false); }); it('does not revive non-completed or transcript-less entries', async () => { diff --git a/packages/core/src/agents/background-agent-resume.ts b/packages/core/src/agents/background-agent-resume.ts index 0cf813baa0b..e50f551fa7e 100644 --- a/packages/core/src/agents/background-agent-resume.ts +++ b/packages/core/src/agents/background-agent-resume.ts @@ -15,7 +15,10 @@ import { AgentEventType, type AgentToolCallEvent, } from './runtime/agent-events.js'; -import { AgentTerminateMode } from './runtime/agent-types.js'; +import { + AgentTerminateMode, + type AgentExternalInput, +} from './runtime/agent-types.js'; import { AgentHeadless, ContextState } from './runtime/agent-headless.js'; import { getSubagentSessionDir, @@ -47,6 +50,7 @@ import type { AgentCompletionStats, AgentTask, AgentTaskRegistration, + ResidentBackgroundAgent, } from './background-tasks.js'; import type { SubagentConfig } from '../subagents/types.js'; import { BUBBLE_APPROVAL_MODE } from '../subagents/types.js'; @@ -623,6 +627,12 @@ export class BackgroundAgentResumeService { return undefined; } + const resumeModel = + existing.model ?? + meta.model ?? + meta.persistedCliFlags?.model ?? + this.config.getModel(); + const bgAbortController = new AbortController(); try { @@ -637,6 +647,7 @@ export class BackgroundAgentResumeService { stats: undefined, recentActivities: [], pendingMessages: [...(existing.pendingMessages ?? [])], + model: resumeModel, }); } catch (error) { const errorMessage = @@ -654,7 +665,6 @@ export class BackgroundAgentResumeService { let cleanupOwnedMonitorNotifications: (() => void) | undefined; let cleanupJsonl: (() => void) | undefined; - try { const subagentName = meta.subagentName ?? meta.agentType; const target = await this.resolveResumeTarget(subagentName); @@ -694,7 +704,12 @@ export class BackgroundAgentResumeService { await createApprovalModeOverride( this.config, resolvedApprovalMode as ApprovalMode, - { persistedCliFlags: meta.persistedCliFlags }, + { + persistedCliFlags: { + ...meta.persistedCliFlags, + model: resumeModel, + }, + }, ); // Mirror the launch path's permission-bubbling gate (agent.ts): an // agent whose definition uses `approvalMode: bubble` surfaces @@ -791,6 +806,7 @@ export class BackgroundAgentResumeService { promptConfigOverrides: { initialMessages: resumeHistory, }, + modelConfigOverrides: { model: resumeModel }, }); subagent = result.subagent; subagentDispose = result.dispose; @@ -815,9 +831,14 @@ export class BackgroundAgentResumeService { status: 'running', lastUpdatedAt: new Date().toISOString(), resolvedApprovalMode, + persistedCliFlags: { + ...meta.persistedCliFlags, + model: resumeModel, + }, subagentName: target.agentName, agentColor: target.subagentConfig?.color ?? meta.agentColor, resumeCount: nextResumeCount, + model: resumeModel, lastError: undefined, }); @@ -838,6 +859,7 @@ export class BackgroundAgentResumeService { prompt: recovery.initialPrompt ?? existing.prompt, recentActivities: [], pendingMessages, + model: resumeModel, }; const entry = registry.register(registration, { suppressRegisterCallback: true, @@ -921,57 +943,164 @@ export class BackgroundAgentResumeService { ? registry.bridgeApprovalEvents(meta.agentId, bgEmitter) : undefined; - const runBody = async () => { - try { - await subagent.execute(contextState, bgAbortController.signal); + const canStayResident = + !target.isFork && + (!target.subagentConfig?.hooks || + Object.keys(target.subagentConfig.hooks).length === 0); + const needsAutoPermissionLease = () => + agentConfig.getApprovalMode() === 'auto' && + this.config.getApprovalMode() !== 'auto'; + const residentWorkingDir = agentConfig.getWorkingDir(); + let runtimeDisposed = false; + let disposeRequested = false; + let turnRunning = false; + let currentAbortController: AbortController | undefined = + bgAbortController; + let currentTurnPromise: Promise | undefined; + let runtimeCleanupPromise: Promise | undefined; + let hotResumeCount = nextResumeCount; + let residentRegistered = false; + + const runtimeCleanup = () => { + if (runtimeCleanupPromise) return runtimeCleanupPromise; + runtimeDisposed = true; + registry.unregisterResidentAgent(meta.agentId, residentController); + residentRegistered = false; + bgEmitter.off(AgentEventType.TOOL_CALL, onToolCall); + bgEmitter.off(AgentEventType.USAGE_METADATA, onUsageMetadata); + cleanupApprovalBridge?.(); + cleanupOwnedMonitorNotifications?.(); + cleanupJsonl?.(); + runtimeCleanupPromise = Promise.allSettled([ + agentConfig.getToolRegistry().stop(), + subagentDispose?.() ?? Promise.resolve(), + ]).then(() => undefined); + return runtimeCleanupPromise; + }; + + const requestRuntimeDisposal = () => { + if (disposeRequested || runtimeDisposed) { + return runtimeCleanupPromise; + } + disposeRequested = true; + registry.unregisterResidentAgent(meta.agentId, residentController); + residentRegistered = false; + currentAbortController?.abort(); + if (!turnRunning) { + return runtimeCleanup(); + } + return undefined; + }; - let stopHookWarning: string | undefined; - if (hookSystem && !bgAbortController.signal.aborted) { - stopHookWarning = await this.runSubagentStopHookLoop(subagent, { + const runBody = async ( + turnContextState: ContextState, + turnAbortController: AbortController, + fireStartHook: boolean, + ) => { + let keepResident = false; + turnRunning = true; + try { + if (fireStartHook) { + await this.applySubagentStartHook(turnContextState, { agentId: meta.agentId, agentType: meta.agentType, - transcriptPath: outputFile, resolvedMode, - signal: bgAbortController.signal, + signal: turnAbortController.signal, }); + const additionalContext = turnContextState.get('hook_context'); + if (additionalContext) { + turnContextState.set( + 'task_prompt', + `${String(turnContextState.get('task_prompt'))}\n\n${String(additionalContext)}`, + ); + } } - const terminateMode = subagent.getTerminateMode(); - const modelVisibleText = toModelVisibleSubagentResult( - subagent.getFinalText(), - terminateMode, - ); - const finalText = appendStopHookBlockingCapWarning( - terminateMode === AgentTerminateMode.GOAL - ? modelVisibleText || - '(subagent produced no model-visible output)' - : modelVisibleText, - stopHookWarning, - ); - const stats = getCompletionStats(subagent, liveToolCallCount); - if (terminateMode === AgentTerminateMode.GOAL) { - registry.complete(meta.agentId, finalText, stats); - patchAgentMeta(metaPath, { - status: 'completed', - lastUpdatedAt: new Date().toISOString(), - lastError: undefined, - }); - } else if (terminateMode === AgentTerminateMode.CANCELLED) { - registry.finalizeCancelled(meta.agentId, finalText, stats); - persistBackgroundCancellation( - metaPath, - registry.get(meta.agentId)?.persistedCancellationStatus ?? - 'cancelled', + let executionContextState = turnContextState; + let initialExternalInputs: readonly AgentExternalInput[] | undefined; + while (true) { + await subagent.execute( + executionContextState, + turnAbortController.signal, + initialExternalInputs + ? { resetStats: false, initialExternalInputs } + : undefined, ); - } else { - const failureText = - finalText || `Agent terminated with mode: ${terminateMode}`; - registry.fail(meta.agentId, failureText, stats); - patchAgentMeta(metaPath, { - status: 'failed', - lastUpdatedAt: new Date().toISOString(), - lastError: failureText, - }); + + let stopHookWarning: string | undefined; + if (hookSystem && !turnAbortController.signal.aborted) { + stopHookWarning = await this.runSubagentStopHookLoop(subagent, { + agentId: meta.agentId, + agentType: meta.agentType, + transcriptPath: outputFile, + resolvedMode, + signal: turnAbortController.signal, + }); + } + + const terminateMode = subagent.getTerminateMode(); + if ( + terminateMode === AgentTerminateMode.GOAL && + residentRegistered && + !needsAutoPermissionLease() + ) { + const claimedInputs = registry.claimPendingInputsForResident( + meta.agentId, + residentController, + ); + if (claimedInputs.length > 0) { + executionContextState = new ContextState(); + executionContextState.set('hook_context', ''); + initialExternalInputs = claimedInputs; + continue; + } + } + + const modelVisibleText = toModelVisibleSubagentResult( + subagent.getFinalText(), + terminateMode, + ); + const finalText = appendStopHookBlockingCapWarning( + terminateMode === AgentTerminateMode.GOAL + ? modelVisibleText || + '(subagent produced no model-visible output)' + : modelVisibleText, + stopHookWarning, + ); + const stats = getCompletionStats(subagent, liveToolCallCount); + if (terminateMode === AgentTerminateMode.GOAL) { + keepResident = residentRegistered && !needsAutoPermissionLease(); + if (!keepResident) { + registry.unregisterResidentAgent( + meta.agentId, + residentController, + ); + residentRegistered = false; + } + registry.complete(meta.agentId, finalText, stats); + patchAgentMeta(metaPath, { + status: 'completed', + lastUpdatedAt: new Date().toISOString(), + lastError: undefined, + }); + } else if (terminateMode === AgentTerminateMode.CANCELLED) { + registry.finalizeCancelled(meta.agentId, finalText, stats); + persistBackgroundCancellation( + metaPath, + registry.get(meta.agentId)?.persistedCancellationStatus ?? + 'cancelled', + ); + } else { + const failureText = + finalText || `Agent terminated with mode: ${terminateMode}`; + registry.fail(meta.agentId, failureText, stats); + patchAgentMeta(metaPath, { + status: 'failed', + lastUpdatedAt: new Date().toISOString(), + lastError: failureText, + }); + } + break; } } catch (error) { const errorMessage = @@ -979,7 +1108,7 @@ export class BackgroundAgentResumeService { debugLogger.error( `[BackgroundAgentResume] Background agent failed: ${errorMessage}`, ); - if (bgAbortController.signal.aborted) { + if (turnAbortController.signal.aborted) { registry.finalizeCancelled( meta.agentId, errorMessage, @@ -1003,45 +1132,122 @@ export class BackgroundAgentResumeService { }); } } finally { - bgEmitter.off(AgentEventType.TOOL_CALL, onToolCall); - bgEmitter.off(AgentEventType.USAGE_METADATA, onUsageMetadata); - cleanupApprovalBridge?.(); - cleanupOwnedMonitorNotifications?.(); - cleanupJsonl?.(); - // Release the per-subagent ToolRegistry the resumed agent's - // wrapper Config built in `createApprovalModeOverride` so any - // AgentTool / SkillTool the model instantiated during this - // run disposes its change-listeners on shared - // SubagentManager / SkillManager. Without this, every resume - // accumulates listeners for the rest of the session. - void agentConfig - .getToolRegistry() - .stop() - .catch(() => {}); - // Per-spawn cleanup from `createAgentHeadless`: releases agent- - // scope hook entries and stops the per-agent ToolRegistry that - // the force rebuild created for `mcpServers`. Distinct from the - // parent registry above (no-op when target.isFork). - void subagentDispose?.().catch(() => {}); - // Restore parent PermissionManager's dangerous allow rules if - // this override stripped them. See createApprovalModeOverride - // strip-lifecycle comment in agent.ts. + turnRunning = false; restoreParentPM(); + if (!keepResident || disposeRequested) { + await runtimeCleanup(); + } } }; - // Restore the persisted launch depth so a resumed nested agent keeps - // its original nesting level (and spawn eligibility) instead of - // recomputing to depth 0 from this top-level resume frame. Normalized - // because the sidecar is untrusted input — a tampered negative depth - // would otherwise mint unbounded spawn capacity. - const framedRunBody = () => - runWithAgentContext( - meta.agentId, - runBody, - normalizeResumedAgentDepth(meta.depth), + const runBackgroundTurn = ( + turnContextState: ContextState, + turnAbortController: AbortController, + fireStartHook: boolean, + ) => { + // Restore the persisted launch depth so a resumed nested agent keeps + // its original nesting level (and spawn eligibility) instead of + // recomputing to depth 0 from this top-level resume frame. + const framedRunBody = () => + runWithAgentContext( + meta.agentId, + () => runBody(turnContextState, turnAbortController, fireStartHook), + normalizeResumedAgentDepth(meta.depth), + ); + return target.isFork + ? runInForkContext(framedRunBody) + : framedRunBody(); + }; + + const reportUnexpectedBackgroundError = (error: unknown) => { + debugLogger.warn( + `[BackgroundAgentResume] Background agent ${meta.agentId} body raised unexpected rejection: ${ + error instanceof Error ? error.message : String(error) + }`, ); - void (target.isFork ? runInForkContext(framedRunBody) : framedRunBody()); + }; + + const residentController: ResidentBackgroundAgent = { + continue: (message) => { + if (!canStayResident || disposeRequested || runtimeDisposed) { + return false; + } + if (agentConfig.getWorkingDir() !== residentWorkingDir) { + registry.disposeResidentAgent(meta.agentId, residentController); + return false; + } + if (needsAutoPermissionLease()) { + registry.disposeResidentAgent(meta.agentId, residentController); + return false; + } + + const nextAbortController = new AbortController(); + let restarted; + try { + restarted = registry.restartCompletedAgent( + meta.agentId, + nextAbortController, + ); + } catch (error) { + debugLogger.warn( + `[BackgroundAgentResume] Could not continue resident background agent ${ + meta.agentId + }: ${error instanceof Error ? error.message : String(error)}`, + ); + return false; + } + if ( + !restarted || + disposeRequested || + runtimeDisposed || + registry.get(meta.agentId) !== restarted || + restarted.status !== 'running' + ) { + return false; + } + + liveToolCallCount = 0; + currentAbortController = nextAbortController; + hotResumeCount += 1; + patchAgentMeta(metaPath, { + status: 'running', + lastUpdatedAt: new Date().toISOString(), + lastError: undefined, + resumeCount: hotResumeCount, + }); + + const nextContextState = new ContextState(); + nextContextState.set('task_prompt', message); + nextContextState.set('hook_context', ''); + const previousTurn = currentTurnPromise ?? Promise.resolve(); + currentTurnPromise = previousTurn + .catch(reportUnexpectedBackgroundError) + .then(async () => { + if (disposeRequested || runtimeDisposed) return; + await runBackgroundTurn( + nextContextState, + nextAbortController, + true, + ); + }); + registry.trackAgentExecution(currentTurnPromise); + currentTurnPromise.catch(reportUnexpectedBackgroundError); + return true; + }, + dispose: requestRuntimeDisposal, + }; + if (canStayResident && !needsAutoPermissionLease()) { + registry.registerResidentAgent(meta.agentId, residentController); + residentRegistered = true; + } + + currentTurnPromise = runBackgroundTurn( + contextState, + bgAbortController, + false, + ); + registry.trackAgentExecution(currentTurnPromise); + currentTurnPromise.catch(reportUnexpectedBackgroundError); return entry; } catch (error) { cleanupOwnedMonitorNotifications?.(); @@ -1278,7 +1484,9 @@ export class BackgroundAgentResumeService { typedStopOutput.getEffectiveReason(), ); continueContext.set('hook_context', ''); - await subagent.execute(continueContext, signal); + await subagent.execute(continueContext, signal, { + resetStats: false, + }); if (signal?.aborted) return undefined; } catch (hookError) { diff --git a/packages/core/src/agents/background-tasks.test.ts b/packages/core/src/agents/background-tasks.test.ts index 5e590985cbe..70907104f7a 100644 --- a/packages/core/src/agents/background-tasks.test.ts +++ b/packages/core/src/agents/background-tasks.test.ts @@ -15,6 +15,7 @@ import { type AgentTaskRegistration, type BackgroundApproval, type BackgroundTaskEntry, + type ResidentBackgroundAgent, } from './background-tasks.js'; import { getCurrentAgentId, @@ -228,6 +229,266 @@ describe('BackgroundTaskRegistry', () => { expect(displayText).toContain('failed'); }); + describe('resident background agents', () => { + function makeResident( + overrides: Partial = {}, + ): ResidentBackgroundAgent { + return { + continue: vi.fn(() => true), + dispose: vi.fn(), + ...overrides, + }; + } + + it('continues only completed resident agents and supports guarded unregister', async () => { + registry.register(makeRegistration('resident-1')); + const resident = makeResident(); + registry.registerResidentAgent('resident-1', resident); + + expect(registry.continueResidentAgent('resident-1', 'too early')).toBe( + false, + ); + + registry.complete('resident-1', 'first result'); + expect(registry.continueResidentAgent('resident-1', 'keep going')).toBe( + true, + ); + expect(resident.continue).toHaveBeenCalledWith('keep going'); + + const staleHandle = makeResident(); + expect(registry.unregisterResidentAgent('resident-1', staleHandle)).toBe( + false, + ); + expect(registry.unregisterResidentAgent('resident-1', resident)).toBe( + true, + ); + expect(resident.dispose).not.toHaveBeenCalled(); + expect( + registry.continueResidentAgent('resident-1', 'after unregister'), + ).toBe(false); + }); + + it('claims pending inputs only for the current running resident', () => { + registry.register(makeRegistration('resident-1')); + const resident = makeResident(); + const staleResident = makeResident(); + registry.registerResidentAgent('resident-1', resident); + registry.queueMessage('resident-1', 'first'); + registry.queueExternalInput('resident-1', { + kind: 'notification', + text: '', + }); + registry.queueMessage('resident-1', 'third'); + + expect( + registry.claimPendingInputsForResident('resident-1', staleResident), + ).toEqual([]); + expect( + registry.claimPendingInputsForResident('resident-1', resident), + ).toEqual([ + 'first', + { kind: 'notification', text: '' }, + 'third', + ]); + expect( + registry.claimPendingInputsForResident('resident-1', resident), + ).toEqual([]); + + registry.complete('resident-1', 'done'); + expect( + registry.claimPendingInputsForResident('resident-1', resident), + ).toEqual([]); + }); + + it('disposes a replaced resident without letting its stale handle remove the replacement', () => { + registry.register(makeRegistration('resident-1')); + const first = makeResident(); + const second = makeResident(); + + registry.registerResidentAgent('resident-1', first); + registry.registerResidentAgent('resident-1', second); + + expect(first.dispose).toHaveBeenCalledOnce(); + expect(registry.disposeResidentAgent('resident-1', first)).toBe(false); + expect(second.dispose).not.toHaveBeenCalled(); + expect(registry.disposeResidentAgent('resident-1', second)).toBe(true); + expect(second.dispose).toHaveBeenCalledOnce(); + }); + + it('disposes the resident when a cold task registration replaces its entry', () => { + registry.register(makeRegistration('resident-1')); + const resident = makeResident(); + registry.registerResidentAgent('resident-1', resident); + registry.complete('resident-1', 'done'); + const completed = registry.get('resident-1')!; + + registry.register({ + ...completed, + status: 'paused', + abortController: new AbortController(), + }); + + expect(resident.dispose).toHaveBeenCalledOnce(); + }); + + it('restarts a completed entry with clean turn state and normal callbacks', () => { + const onRegister = vi.fn(); + const onStatusChange = vi.fn(); + registry.setRegisterCallback(onRegister); + registry.setStatusChangeCallback(onStatusChange); + const originalController = new AbortController(); + registry.register( + makeRegistration('resident-1', { + abortController: originalController, + recentActivities: [ + { name: 'Read', description: 'old activity', at: 1 }, + ], + pendingMessages: ['already queued'], + }), + ); + registry.complete('resident-1', 'first result', { + totalTokens: 10, + outputTokens: 4, + toolUses: 1, + durationMs: 20, + }); + const completed = registry.get('resident-1')!; + completed.error = 'stale error'; + completed.resumeBlockedReason = 'stale block'; + completed.persistedCancellationStatus = 'cancelled'; + const completedAt = completed.endTime; + const nextController = new AbortController(); + + const restarted = registry.restartCompletedAgent( + 'resident-1', + nextController, + ); + + expect(restarted).toBe(completed); + expect(restarted).toMatchObject({ + status: 'running', + abortController: nextController, + recentActivities: [], + pendingApprovals: [], + pendingMessages: ['already queued'], + notified: false, + outputOffset: 0, + }); + expect(restarted?.startTime).toBeGreaterThan(0); + expect(restarted?.endTime).toBeUndefined(); + expect(restarted?.result).toBeUndefined(); + expect(restarted?.error).toBeUndefined(); + expect(restarted?.resumeBlockedReason).toBeUndefined(); + expect(restarted?.stats).toBeUndefined(); + expect(restarted?.persistedCancellationStatus).toBeUndefined(); + expect(completedAt).toBeDefined(); + expect(onRegister).toHaveBeenCalledTimes(2); + expect(onStatusChange).toHaveBeenCalledTimes(3); + }); + + it('leaves a completed entry unchanged when restart capacity is full', () => { + registry = new BackgroundTaskRegistry({ + maxConcurrentBackgroundAgents: 1, + }); + registry.register(makeRegistration('resident-1')); + registry.complete('resident-1', 'first result', { + totalTokens: 10, + outputTokens: 4, + toolUses: 1, + durationMs: 20, + }); + const completed = registry.get('resident-1')!; + const completedController = completed.abortController; + const completedAt = completed.endTime; + registry.register(makeRegistration('busy')); + + expect(() => + registry.restartCompletedAgent('resident-1', new AbortController()), + ).toThrow('maximum concurrent background agents (1) reached'); + expect(completed).toMatchObject({ + status: 'completed', + abortController: completedController, + result: 'first result', + endTime: completedAt, + notified: true, + }); + expect(completed.stats).toEqual({ + totalTokens: 10, + outputTokens: 4, + toolUses: 1, + durationMs: 20, + }); + }); + + it('keeps a successful runtime resident but disposes failed and cancelled runtimes', () => { + registry.register(makeRegistration('completed')); + const completedResident = makeResident(); + registry.registerResidentAgent('completed', completedResident); + registry.complete('completed', 'done'); + expect(completedResident.dispose).not.toHaveBeenCalled(); + + registry.register(makeRegistration('failed')); + const failedResident = makeResident(); + registry.registerResidentAgent('failed', failedResident); + registry.fail('failed', 'boom'); + expect(failedResident.dispose).toHaveBeenCalledOnce(); + + registry.register(makeRegistration('cancelled')); + const cancelledResident = makeResident(); + registry.registerResidentAgent('cancelled', cancelledResident); + registry.cancel('cancelled'); + registry.finalizeCancelled('cancelled', 'partial'); + expect(cancelledResident.dispose).toHaveBeenCalledOnce(); + }); + + it('removes a cancelled resident before publishing a raced completion', () => { + registry.register(makeRegistration('cancelled-completion')); + const resident = makeResident(); + registry.registerResidentAgent('cancelled-completion', resident); + let continuation: boolean | undefined; + registry.setNotificationCallback(() => { + continuation = registry.continueResidentAgent( + 'cancelled-completion', + 'do not restart the cancelled runtime', + ); + }); + + registry.cancel('cancelled-completion'); + registry.complete('cancelled-completion', 'finished while cancelling'); + + expect(continuation).toBe(false); + expect(resident.continue).not.toHaveBeenCalled(); + expect(resident.dispose).toHaveBeenCalledOnce(); + }); + + it('disposes all resident runtimes on abortAll and reset', () => { + registry.register(makeRegistration('running')); + const runningResident = makeResident(); + registry.registerResidentAgent('running', runningResident); + registry.register( + makeRegistration('completed', { + status: 'completed', + }), + ); + const completedResident = makeResident(); + registry.registerResidentAgent('completed', completedResident); + + registry.abortAll({ notify: false }); + + expect(runningResident.dispose).toHaveBeenCalledOnce(); + expect(completedResident.dispose).toHaveBeenCalledOnce(); + + registry.register(makeRegistration('next')); + const nextResident = makeResident(); + registry.registerResidentAgent('next', nextResident); + registry.complete('next', 'done'); + + registry.reset(); + + expect(nextResident.dispose).toHaveBeenCalledOnce(); + }); + }); + it('cancels a running background agent without emitting a notification', () => { // cancel() is intent-only: it aborts the signal and marks the entry // cancelled, but does not emit a task-notification. The natural @@ -810,6 +1071,42 @@ describe('BackgroundTaskRegistry', () => { }), ).toThrow('invalidated by session reset'); }); + + it.each([ + ['default notifications', undefined], + ['suppressed notifications', { notify: false }], + ])('rejects queued launches when abortAll uses %s', async (_, options) => { + registry = new BackgroundTaskRegistry({ + maxConcurrentBackgroundAgents: 1, + }); + registry.register(makeRegistration('bg-1')); + const waiter = registry.waitForBackgroundSlot( + new AbortController().signal, + ); + + registry.abortAll(options); + + await expect(waiter).rejects.toThrow( + 'Agent launch cancelled while waiting for a background slot.', + ); + expect(registry.getQueuedCount()).toBe(0); + }); + + it('invalidates unconsumed slot reservations when aborting all agents', () => { + registry = new BackgroundTaskRegistry({ + maxConcurrentBackgroundAgents: 1, + }); + const reservation = registry.tryReserveBackgroundSlot()!; + + registry.abortAll({ notify: false }); + + expect(() => + registry.register(makeRegistration('bg-1'), { + slotReservation: reservation, + }), + ).toThrow('invalidated by session reset'); + expect(registry.get('bg-1')).toBeUndefined(); + }); }); describe('per-model background concurrency limit', () => { @@ -1017,6 +1314,115 @@ describe('BackgroundTaskRegistry', () => { expect(callback).toHaveBeenCalledTimes(2); }); + it('aborts residents and waits for every tracked agent execution to settle', async () => { + let resolveFirst!: () => void; + let rejectSecond!: (error: Error) => void; + const first = new Promise((resolve) => { + resolveFirst = resolve; + }); + const second = new Promise((_resolve, reject) => { + rejectSecond = reject; + }); + const resident: ResidentBackgroundAgent = { + continue: vi.fn(() => true), + dispose: vi.fn(), + }; + registry.register(makeRegistration('resident-1')); + registry.registerResidentAgent('resident-1', resident); + registry.trackAgentExecution(first); + const abortAll = vi.spyOn(registry, 'abortAll'); + + let settled = false; + const wait = registry.abortAllAndWait({ notify: false }).then(() => { + settled = true; + }); + await Promise.resolve(); + expect(settled).toBe(false); + expect(resident.dispose).toHaveBeenCalledOnce(); + + registry.trackAgentExecution(second); + resolveFirst(); + await vi.waitFor(() => expect(abortAll).toHaveBeenCalledTimes(2)); + expect(settled).toBe(false); + + rejectSecond(new Error('expected test rejection')); + await wait; + expect(settled).toBe(true); + expect(abortAll).toHaveBeenNthCalledWith(1, { notify: false }); + expect(abortAll).toHaveBeenNthCalledWith(2, { notify: false }); + }); + + it('waits for an idle resident async cleanup to settle', async () => { + let resolveCleanup!: () => void; + const cleanup = new Promise((resolve) => { + resolveCleanup = resolve; + }); + const resident: ResidentBackgroundAgent = { + continue: vi.fn(() => true), + dispose: vi.fn(() => cleanup), + }; + registry.register( + makeRegistration('idle-resident', { status: 'completed' }), + ); + registry.registerResidentAgent('idle-resident', resident); + + let settled = false; + const wait = registry.abortAllAndWait({ notify: false }).then(() => { + settled = true; + }); + await Promise.resolve(); + + expect(resident.dispose).toHaveBeenCalledOnce(); + expect(settled).toBe(false); + expect(registry.continueResidentAgent('idle-resident', 'continue')).toBe( + false, + ); + + resolveCleanup(); + await wait; + expect(settled).toBe(true); + }); + + it('disposes a resident registered while waiting for an execution to settle', async () => { + let resolveExecution!: () => void; + const execution = new Promise((resolve) => { + resolveExecution = resolve; + }); + registry.trackAgentExecution(execution); + + const wait = registry.abortAllAndWait({ notify: false }); + const resident: ResidentBackgroundAgent = { + continue: vi.fn(() => true), + dispose: vi.fn(), + }; + registry.register( + makeRegistration('late-resident', { status: 'completed' }), + ); + registry.registerResidentAgent('late-resident', resident); + + resolveExecution(); + await wait; + + expect(resident.dispose).toHaveBeenCalledOnce(); + }); + + it('bounds the wait when an aborted agent execution never settles', async () => { + vi.useFakeTimers(); + try { + registry.trackAgentExecution(new Promise(() => {})); + + const wait = registry.abortAllAndWait({ notify: false }); + const rejection = wait.catch((error: unknown) => error); + + await vi.advanceTimersByTimeAsync(5000); + expect(await rejection).toEqual( + new Error('Background agents did not stop within 5000ms.'), + ); + } finally { + vi.useRealTimers(); + } + }); + it('abortAll({ notify: false }) suppresses terminal notifications from old tasks', () => { const callback = vi.fn(); registry.setNotificationCallback(callback); @@ -1570,6 +1976,24 @@ describe('BackgroundTaskRegistry', () => { ).toBeDefined(); }); + it('disposes a resident runtime when its terminal entry is evicted', () => { + registry.register(makeRegisteredEntry('resident-oldest', 0)); + const resident = { + continue: vi.fn(() => true), + dispose: vi.fn(), + }; + registry.registerResidentAgent('resident-oldest', resident); + registry.complete('resident-oldest', 'done'); + + for (let i = 0; i < MAX_RETAINED_TERMINAL_AGENTS; i++) { + registry.register(makeRegisteredEntry(`newer-${i}`, 1000 + i)); + registry.complete(`newer-${i}`, 'done'); + } + + expect(registry.get('resident-oldest')).toBeUndefined(); + expect(resident.dispose).toHaveBeenCalledOnce(); + }); + it('never evicts running entries even when terminal entries blow past the cap', () => { // The user's only handle on a live subagent is its row in the // dialog; a prune that drops a running entry would silently diff --git a/packages/core/src/agents/background-tasks.ts b/packages/core/src/agents/background-tasks.ts index b3618e7b441..d9b0caf6a3a 100644 --- a/packages/core/src/agents/background-tasks.ts +++ b/packages/core/src/agents/background-tasks.ts @@ -437,6 +437,16 @@ export type BackgroundActivityChangeCallback = (entry: AgentTask) => void; */ export type BackgroundApprovalChangeCallback = (entry: AgentTask) => void; +/** + * Session-scoped handle for a background agent whose runtime remains alive + * after a completed turn. The handle is deliberately not part of AgentTask: + * task state is serializable, while the live runtime is process-local. + */ +export interface ResidentBackgroundAgent { + continue(message: string): boolean; + dispose(): void | Promise; +} + type MessageWaiter = () => void; export interface BackgroundTaskRegistryOptions { @@ -477,6 +487,8 @@ const BACKGROUND_SLOT_WAIT_CANCELLED = export class BackgroundTaskRegistry { private readonly agents = new Map(); + private readonly residentAgents = new Map(); + private readonly agentExecutions = new Set>(); private readonly messageWaiters = new Map>(); private readonly waitQueue: BackgroundSlotWaiter[] = []; // Maps each outstanding slot reservation to the concrete model ID it was @@ -639,6 +651,9 @@ export class BackgroundTaskRegistry { } } } + if (existing && existing !== registration) { + this.disposeResidentAgent(registration.agentId); + } // Mutate the registration in place to graduate it to an `AgentTask`. // Returning the same reference lets callers (e.g. the resume service) @@ -692,6 +707,108 @@ export class BackgroundTaskRegistry { return entry; } + /** + * Restart a completed background task for another turn while preserving its + * resident runtime. Capacity is checked before mutating the entry so a + * rejected restart leaves the completed task intact. + */ + restartCompletedAgent( + agentId: string, + abortController: AbortController, + ): AgentTask | undefined { + const entry = this.agents.get(agentId); + if (!entry || !entry.isBackgrounded || entry.status !== 'completed') { + return undefined; + } + + this.assertCanStartBackgroundAgent(entry.model); + + entry.status = 'running'; + entry.startTime = Date.now(); + entry.endTime = undefined; + entry.abortController = abortController; + entry.result = undefined; + entry.error = undefined; + entry.resumeBlockedReason = undefined; + entry.stats = undefined; + entry.recentActivities = []; + entry.pendingApprovals = []; + entry.persistedCancellationStatus = undefined; + + return this.register(entry); + } + + registerResidentAgent( + agentId: string, + resident: ResidentBackgroundAgent, + ): void { + const existing = this.residentAgents.get(agentId); + if (existing === resident) return; + if (existing) { + this.disposeResidentAgent(agentId, existing); + } + this.residentAgents.set(agentId, resident); + } + + continueResidentAgent(agentId: string, message: string): boolean { + const entry = this.agents.get(agentId); + const resident = this.residentAgents.get(agentId); + if (!resident || entry?.status !== 'completed') return false; + return resident.continue(message); + } + + claimPendingInputsForResident( + agentId: string, + resident: ResidentBackgroundAgent, + ): AgentExternalInput[] { + const entry = this.agents.get(agentId); + if ( + entry?.status !== 'running' || + this.residentAgents.get(agentId) !== resident || + !entry.pendingMessages?.length + ) { + return []; + } + return entry.pendingMessages.splice(0); + } + + unregisterResidentAgent( + agentId: string, + resident?: ResidentBackgroundAgent, + ): boolean { + const current = this.residentAgents.get(agentId); + if (!current || (resident && current !== resident)) return false; + return this.residentAgents.delete(agentId); + } + + disposeResidentAgent( + agentId: string, + resident?: ResidentBackgroundAgent, + ): boolean { + const current = this.residentAgents.get(agentId); + if (!current || (resident && current !== resident)) return false; + this.residentAgents.delete(agentId); + try { + const cleanup = current.dispose(); + if (cleanup) { + this.trackAgentExecution( + cleanup.catch((error) => { + debugLogger.error( + `Failed to dispose resident background agent ${agentId}:`, + error, + ); + }), + ); + } + } catch (error) { + debugLogger.error( + `Failed to dispose resident background agent ${agentId}:`, + error, + ); + } + return true; + } + // Transition a still-running entry to 'completed' and emit the terminal // notification. No-op if the entry is already terminal *and* has been // notified — protects against duplicate emission when cancel aborts the @@ -710,6 +827,7 @@ export class BackgroundTaskRegistry { if (entry.status !== 'running' && entry.status !== 'cancelled') return; if (entry.notified) return; + const wasCancelled = entry.status === 'cancelled'; entry.status = 'completed'; entry.endTime = Date.now(); entry.result = result; @@ -717,6 +835,9 @@ export class BackgroundTaskRegistry { debugLogger.info(`Background agent completed: ${agentId}`); this.rejectPendingApprovals(entry); + if (wasCancelled) { + this.disposeResidentAgent(agentId); + } this.emitNotification(entry); this.emitStatusChange(entry); this.drainWaitQueue(); @@ -749,7 +870,7 @@ export class BackgroundTaskRegistry { // complete/fail/cancel/finalize ordering on purpose — those // keep the entry around (terminal state) so callbacks can inspect // it on re-read; unregister removes it outright. - this.agents.delete(agentId); + this.deleteAgent(agentId); this.emitStatusChange(entry); debugLogger.info(`Unregistered foreground agent: ${agentId}`); this.drainWaitQueue(); @@ -771,6 +892,7 @@ export class BackgroundTaskRegistry { this.rejectPendingApprovals(entry); this.emitNotification(entry); this.emitStatusChange(entry); + this.disposeResidentAgent(agentId); this.drainWaitQueue(); } @@ -852,6 +974,7 @@ export class BackgroundTaskRegistry { debugLogger.info(`Abandoned paused background agent: ${agentId}`); this.rejectPendingApprovals(entry); this.emitStatusChange(entry); + this.disposeResidentAgent(agentId); this.drainWaitQueue(); } @@ -877,6 +1000,7 @@ export class BackgroundTaskRegistry { this.rejectPendingApprovals(entry); this.emitNotification(entry); this.emitStatusChange(entry); + this.disposeResidentAgent(agentId); this.drainWaitQueue(); } @@ -896,6 +1020,7 @@ export class BackgroundTaskRegistry { this.rejectPendingApprovals(entry); this.emitNotification(entry); this.emitStatusChange(entry); + this.disposeResidentAgent(agentId); this.drainWaitQueue(); } @@ -1071,6 +1196,51 @@ export class BackgroundTaskRegistry { return Array.from(this.agents.values()); } + trackAgentExecution(execution: Promise): void { + this.agentExecutions.add(execution); + void execution.then( + () => this.agentExecutions.delete(execution), + () => this.agentExecutions.delete(execution), + ); + } + + async abortAllAndWait( + options: BackgroundTaskCancelOptions = {}, + ): Promise { + const deadline = Date.now() + CANCEL_GRACE_MS; + while (true) { + this.abortAll(options); + if (this.agentExecutions.size === 0) return; + + const remainingMs = deadline - Date.now(); + if (remainingMs <= 0) { + throw new Error( + `Background agents did not stop within ${CANCEL_GRACE_MS}ms.`, + ); + } + + let timeout: ReturnType | undefined; + try { + await Promise.race([ + Promise.allSettled([...this.agentExecutions]), + new Promise((_resolve, reject) => { + timeout = setTimeout( + () => + reject( + new Error( + `Background agents did not stop within ${CANCEL_GRACE_MS}ms.`, + ), + ), + remainingMs, + ); + }), + ]); + } finally { + if (timeout) clearTimeout(timeout); + } + } + } + // Counts backgrounded agents that still occupy a slot: running, or // cancelled-but-not-yet-finalized. When `model` is given, only agents on // that model are counted (per-model cap); otherwise all of them (global). @@ -1215,6 +1385,7 @@ export class BackgroundTaskRegistry { | AgentTask | undefined; if (!firstEntry) { + this.disposeAllResidentAgents(); this.rejectWaitQueue(); return; } @@ -1229,6 +1400,7 @@ export class BackgroundTaskRegistry { } this.rejectWaitQueue(); this.agents.clear(); + this.disposeAllResidentAgents(); this.emitStatusChange(firstEntry); } @@ -1349,6 +1521,7 @@ export class BackgroundTaskRegistry { } abortAll(options: BackgroundTaskCancelOptions = {}): void { + this.rejectWaitQueue(); const cancelOptions: BackgroundTaskCancelOptions = { persistedStatus: 'running', ...options, @@ -1367,6 +1540,7 @@ export class BackgroundTaskRegistry { // notification here to honour the one-notification-per-agent contract. this.finalizeCancellationIfPending(entry.agentId); } + this.disposeAllResidentAgents(); debugLogger.info('Aborted all background agents'); } @@ -1506,11 +1680,22 @@ export class BackgroundTaskRegistry { while (evictable.length > MAX_RETAINED_TERMINAL_AGENTS) { const oldest = evictable.shift(); if (oldest) { - this.agents.delete(oldest.agentId); + this.deleteAgent(oldest.agentId); } } } + private deleteAgent(agentId: string): boolean { + this.disposeResidentAgent(agentId); + return this.agents.delete(agentId); + } + + private disposeAllResidentAgents(): void { + for (const agentId of Array.from(this.residentAgents.keys())) { + this.disposeResidentAgent(agentId); + } + } + private wakeMessageWaiters(agentId: string): void { const waiters = this.messageWaiters.get(agentId); if (!waiters) return; diff --git a/packages/core/src/agents/runtime/agent-core.ts b/packages/core/src/agents/runtime/agent-core.ts index b28eb434751..5acd6a8ddea 100644 --- a/packages/core/src/agents/runtime/agent-core.ts +++ b/packages/core/src/agents/runtime/agent-core.ts @@ -226,6 +226,8 @@ export interface ReasoningLoopOptions { maxTimeMinutes?: number; /** Start time in ms (for timeout calculation). Defaults to Date.now(). */ startTimeMs?: number; + /** Rounds already completed in the same logical turn. */ + roundOffset?: number; /** * Optional callback to drain external messages between model rounds. * Returned inputs are appended to the next model request as user-role @@ -293,6 +295,7 @@ export interface ExecutionStats { * or final result interpretation — those are the caller's responsibility. */ export class AgentCore { + private promptOrdinal = 0; readonly subagentId: string; readonly name: string; readonly runtimeContext: Config; @@ -787,7 +790,10 @@ export class AgentCore { } // Check termination conditions. - if (options?.maxTurns && turnCounter >= options.maxTurns) { + if ( + options?.maxTurns && + (options.roundOffset ?? 0) + turnCounter >= options.maxTurns + ) { terminateMode = AgentTerminateMode.MAX_TURNS; break; } @@ -805,7 +811,8 @@ export class AgentCore { const roundAbortController = createChildAbortController(abortController); try { - const promptId = `${this.runtimeContext.getSessionId()}#${this.subagentId}#${turnCounter++}`; + const promptId = `${this.runtimeContext.getSessionId()}#${this.subagentId}#${this.promptOrdinal++}`; + turnCounter += 1; const messageParams = { message: currentMessages[0]?.parts || [], @@ -1000,8 +1007,9 @@ export class AgentCore { } as AgentRoundTextEvent); } - this.executionStats.rounds = turnCounter; - this.stats.setRounds(turnCounter); + const cumulativeRounds = (options?.roundOffset ?? 0) + turnCounter; + this.executionStats.rounds = cumulativeRounds; + this.stats.setRounds(cumulativeRounds); durationMin = (Date.now() - startTime) / (1000 * 60); if (options?.maxTimeMinutes && durationMin >= options.maxTimeMinutes) { @@ -1200,7 +1208,10 @@ export class AgentCore { options: ReasoningLoopOptions | undefined, turnCounter: number, ): boolean { - return !options?.maxTurns || turnCounter < options.maxTurns; + return ( + !options?.maxTurns || + (options.roundOffset ?? 0) + turnCounter < options.maxTurns + ); } private getRemainingTimeMs( @@ -1835,6 +1846,22 @@ export class AgentCore { // ─── Stats & Events ─────────────────────────────────────── + resetExecutionStats(): void { + this.executionStats = { + startTimeMs: 0, + totalDurationMs: 0, + rounds: 0, + totalToolCalls: 0, + successfulToolCalls: 0, + failedToolCalls: 0, + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + }; + this.toolUsage.clear(); + this.stats.reset(); + } + getEventEmitter(): AgentEventEmitter { return this.eventEmitter; } diff --git a/packages/core/src/agents/runtime/agent-headless.test.ts b/packages/core/src/agents/runtime/agent-headless.test.ts index a990457f923..adaa1d76ceb 100644 --- a/packages/core/src/agents/runtime/agent-headless.test.ts +++ b/packages/core/src/agents/runtime/agent-headless.test.ts @@ -543,6 +543,315 @@ describe('subagent.ts', () => { ]); }); + it('should reuse chat and tools for sequential follow-up turns', async () => { + const { config, toolRegistry } = await createMockConfig(); + mockSendMessageStream.mockImplementation( + createMockStream(['stop', 'stop']), + ); + + const scope = await AgentHeadless.create( + 'test-agent', + config, + { systemPrompt: 'You are a test agent.' }, + defaultModelConfig, + defaultRunConfig, + ); + const externalMessages: string[] = []; + scope.getEventEmitter().on(AgentEventType.EXTERNAL_MESSAGE, (event) => { + externalMessages.push(event.text); + }); + + const initialContext = new ContextState(); + initialContext.set('task_prompt', 'Initial task'); + await scope.execute(initialContext); + + scope.getCore().recordToolCallStats('stale_tool', true, 25); + scope.getCore().stats.recordTokens(100, 50); + + const followUpContext = new ContextState(); + followUpContext.set('task_prompt', 'Follow-up task'); + await scope.execute(followUpContext); + + expect(GeminiChat).toHaveBeenCalledTimes(1); + expect(toolRegistry.warmAll).toHaveBeenCalledTimes(1); + expect(mockSendMessageStream).toHaveBeenCalledTimes(2); + expect(mockSendMessageStream.mock.calls[0][1].message).toEqual([ + { text: 'Initial task' }, + ]); + expect(mockSendMessageStream.mock.calls[1][1].message).toEqual([ + { text: '[Message from parent agent]: Follow-up task' }, + ]); + expect(mockSendMessageStream.mock.calls[0][2]).not.toBe( + mockSendMessageStream.mock.calls[1][2], + ); + expect(mockSendMessageStream.mock.calls[0][2]).toMatch(/#0$/); + expect(mockSendMessageStream.mock.calls[1][2]).toMatch(/#1$/); + expect(externalMessages).toEqual(['Follow-up task']); + expect(scope.getExecutionSummary()).toMatchObject({ + rounds: 1, + totalToolCalls: 0, + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + toolUsage: [], + }); + expect(scope.getStatistics()).toMatchObject({ + rounds: 1, + totalToolCalls: 0, + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, + toolUsage: [], + }); + }); + + it('should preserve claimed external input order and source kind on a continuation', async () => { + const { config, toolRegistry } = await createMockConfig(); + mockSendMessageStream.mockImplementation( + createMockStream(['stop', 'stop']), + ); + + const scope = await AgentHeadless.create( + 'test-agent', + config, + { systemPrompt: 'You are a test agent.' }, + defaultModelConfig, + defaultRunConfig, + ); + const externalMessages: Array<{ + kind?: 'message' | 'notification'; + text: string; + }> = []; + scope.getEventEmitter().on(AgentEventType.EXTERNAL_MESSAGE, (event) => { + externalMessages.push({ kind: event.kind, text: event.text }); + }); + + const initialContext = new ContextState(); + initialContext.set('task_prompt', 'Initial task'); + await scope.execute(initialContext); + scope.getCore().recordToolCallStats('first_attempt_tool', true, 25); + + const continuationContext = new ContextState(); + continuationContext.set('task_prompt', 'ignored fallback'); + await scope.execute(continuationContext, undefined, { + resetStats: false, + initialExternalInputs: [ + 'first message', + { + kind: 'notification', + text: 'ready', + }, + 'third message', + ], + }); + + expect(GeminiChat).toHaveBeenCalledTimes(1); + expect(toolRegistry.warmAll).toHaveBeenCalledTimes(1); + expect(mockSendMessageStream.mock.calls[1][1].message).toEqual([ + { text: '[Message from parent agent]: first message' }, + { text: 'ready' }, + { text: '[Message from parent agent]: third message' }, + ]); + expect(externalMessages).toEqual([ + { kind: 'message', text: 'first message' }, + { + kind: 'notification', + text: 'ready', + }, + { kind: 'message', text: 'third message' }, + ]); + expect(scope.getExecutionSummary()).toMatchObject({ + rounds: 2, + totalToolCalls: 1, + successfulToolCalls: 1, + }); + }); + + it('should preserve statistics for continuation work in the same logical turn', async () => { + const { config } = await createMockConfig(); + mockSendMessageStream.mockImplementation( + createMockStream(['stop', 'stop']), + ); + + const scope = await AgentHeadless.create( + 'test-agent', + config, + { systemPrompt: 'You are a test agent.' }, + defaultModelConfig, + defaultRunConfig, + ); + await scope.execute(new ContextState()); + scope.getCore().recordToolCallStats('first_attempt_tool', true, 25); + scope.getCore().stats.recordTokens(100, 50); + scope.getCore().executionStats.inputTokens = 100; + scope.getCore().executionStats.outputTokens = 50; + scope.getCore().executionStats.totalTokens = 150; + const logicalTurnStart = Date.now() - 10_000; + scope.getCore().executionStats.startTimeMs = logicalTurnStart; + scope.getCore().stats.start(logicalTurnStart); + + const continuationContext = new ContextState(); + continuationContext.set('task_prompt', 'Address the stop-hook reason'); + await scope.execute(continuationContext, undefined, { + resetStats: false, + }); + + expect(scope.getExecutionSummary()).toMatchObject({ + rounds: 2, + totalToolCalls: 1, + successfulToolCalls: 1, + inputTokens: 100, + outputTokens: 50, + }); + expect(scope.getCore().executionStats.startTimeMs).toBe( + logicalTurnStart, + ); + expect( + scope.getCore().executionStats.totalDurationMs, + ).toBeGreaterThanOrEqual(10_000); + expect(scope.getStatistics()).toMatchObject({ + rounds: 2, + totalDurationMs: expect.any(Number), + totalToolCalls: 1, + successfulToolCalls: 1, + inputTokens: 100, + outputTokens: 50, + }); + }); + + it('should enforce max turns before starting a same-turn continuation', async () => { + const { config } = await createMockConfig(); + const runConfig: RunConfig = { ...defaultRunConfig, max_turns: 1 }; + mockSendMessageStream.mockImplementation(createMockStream(['stop'])); + + const scope = await AgentHeadless.create( + 'test-agent', + config, + { systemPrompt: 'You are a test agent.' }, + defaultModelConfig, + runConfig, + ); + await scope.execute(new ContextState()); + + await scope.execute(new ContextState(), undefined, { + resetStats: false, + initialExternalInputs: ['continue'], + }); + + expect(mockSendMessageStream).toHaveBeenCalledOnce(); + expect(scope.getTerminateMode()).toBe(AgentTerminateMode.MAX_TURNS); + expect(scope.getExecutionSummary().rounds).toBe(1); + }); + + it('should not idle-wait after a continuation exhausts the logical turn budget', async () => { + const { config } = await createMockConfig(); + const runConfig: RunConfig = { ...defaultRunConfig, max_turns: 2 }; + mockSendMessageStream.mockImplementation( + createMockStream(['stop', 'stop', 'stop']), + ); + + const scope = await AgentHeadless.create( + 'test-agent', + config, + { systemPrompt: 'You are a test agent.' }, + defaultModelConfig, + runConfig, + ); + await scope.execute(new ContextState()); + + const waitForExternalMessages = vi.fn(async () => [ + { + kind: 'notification' as const, + text: 'late', + }, + ]); + scope.setExternalMessageProvider(() => []); + scope.setExternalMessageWaiter(waitForExternalMessages); + scope.setExternalMessageWaitPredicate(() => true); + + await scope.execute(new ContextState(), undefined, { + resetStats: false, + initialExternalInputs: ['continue'], + }); + + expect(mockSendMessageStream).toHaveBeenCalledTimes(2); + expect(waitForExternalMessages).not.toHaveBeenCalled(); + expect(scope.getTerminateMode()).toBe(AgentTerminateMode.MAX_TURNS); + expect(scope.getExecutionSummary().rounds).toBe(2); + }); + + it('should reject concurrent execute calls', async () => { + const { config } = await createMockConfig(); + let releaseResponse: (() => void) | undefined; + const responseGate = new Promise((resolve) => { + releaseResponse = resolve; + }); + mockSendMessageStream.mockImplementation(async () => + (async function* () { + await responseGate; + yield { + type: 'chunk', + value: { + candidates: [ + { + content: { + parts: [{ text: 'Done.' }], + }, + }, + ], + }, + }; + })(), + ); + + const scope = await AgentHeadless.create( + 'test-agent', + config, + { systemPrompt: 'You are a test agent.' }, + defaultModelConfig, + defaultRunConfig, + ); + const firstExecution = scope.execute(new ContextState()); + await vi.waitFor(() => + expect(mockSendMessageStream).toHaveBeenCalledTimes(1), + ); + + await expect(scope.execute(new ContextState())).rejects.toThrow( + 'AgentHeadless does not support concurrent execute() calls.', + ); + + releaseResponse?.(); + await firstExecution; + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + }); + + it('should clear the prior result before a failing follow-up turn', async () => { + const { config } = await createMockConfig(); + mockSendMessageStream + .mockImplementationOnce(createMockStream(['stop'])) + .mockRejectedValueOnce(new Error('follow-up failed')); + + const scope = await AgentHeadless.create( + 'test-agent', + config, + { systemPrompt: 'You are a test agent.' }, + defaultModelConfig, + defaultRunConfig, + ); + await scope.execute(new ContextState()); + expect(scope.getFinalText()).toBe('Done.'); + expect(scope.getTerminateMode()).toBe(AgentTerminateMode.GOAL); + + const followUpContext = new ContextState(); + followUpContext.set('task_prompt', 'Follow-up task'); + await expect(scope.execute(followUpContext)).rejects.toThrow( + 'follow-up failed', + ); + + expect(scope.getFinalText()).toBe(''); + expect(scope.getTerminateMode()).toBe(AgentTerminateMode.ERROR); + }); + it('should append userMemory to the system prompt when available', async () => { const { config } = await createMockConfig(); const userMemoryContent = diff --git a/packages/core/src/agents/runtime/agent-headless.ts b/packages/core/src/agents/runtime/agent-headless.ts index 106950201a3..925012c8efc 100644 --- a/packages/core/src/agents/runtime/agent-headless.ts +++ b/packages/core/src/agents/runtime/agent-headless.ts @@ -5,17 +5,18 @@ */ /** - * @fileoverview AgentHeadless — one-shot task execution wrapper around AgentCore. + * @fileoverview AgentHeadless — sequential task execution wrapper around AgentCore. * - * AgentHeadless manages - * the lifecycle of a single headless task: start → run → finish. + * AgentHeadless runs one headless task at a time while retaining its chat + * session for follow-up tasks. * It delegates all model reasoning and tool scheduling to AgentCore. * * For persistent interactive agents, see AgentInteractive (Phase 2). */ -import type { Content } from '@google/genai'; +import type { Content, FunctionDeclaration } from '@google/genai'; import type { Config } from '../../config/config.js'; +import type { GeminiChat } from '../../core/geminiChat.js'; import type { RuntimeContentGeneratorView } from './agent-context.js'; import { createChildAbortController } from '../../utils/abortController.js'; import { createDebugLogger } from '../../utils/debugLogger.js'; @@ -38,7 +39,7 @@ import type { import { AgentTerminateMode } from './agent-types.js'; import { logSubagentExecution } from '../../telemetry/loggers.js'; import { SubagentExecutionEvent } from '../../telemetry/types.js'; -import { AgentCore } from './agent-core.js'; +import { AgentCore, EXTERNAL_MESSAGE_PREFIX } from './agent-core.js'; import { DEFAULT_QWEN_MODEL } from '../../config/models.js'; const debugLogger = createDebugLogger('SUBAGENT'); @@ -129,17 +130,19 @@ export function templateString( // ─── AgentHeadless ────────────────────────────────────────── /** - * AgentHeadless — one-shot task executor. + * AgentHeadless — sequential task executor. * - * Takes a task, runs it through AgentCore's reasoning loop, and returns - * the result. - * - * Lifecycle: Born → execute() → die. + * Each execute() call runs one task through AgentCore's reasoning loop. Calls + * must be sequential; later calls reuse the same chat and prepared tools. */ export class AgentHeadless { private readonly core: AgentCore; private finalText: string = ''; private terminateMode: AgentTerminateMode = AgentTerminateMode.ERROR; + private chat?: GeminiChat; + private toolsList?: FunctionDeclaration[]; + private executing = false; + private hasStartedReasoning = false; private externalMessageProvider?: () => AgentExternalInput[]; private externalMessageWaiter?: ( signal: AbortSignal, @@ -203,10 +206,47 @@ export class AgentHeadless { async execute( context: ContextState, externalSignal?: AbortSignal, + options: { + resetStats?: boolean; + initialExternalInputs?: readonly AgentExternalInput[]; + } = {}, + ): Promise { + if (this.executing) { + throw new Error( + 'AgentHeadless does not support concurrent execute() calls.', + ); + } + + this.executing = true; + this.finalText = ''; + this.terminateMode = AgentTerminateMode.ERROR; + const resetStats = options.resetStats !== false; + if (resetStats) { + this.core.resetExecutionStats(); + } + + try { + await this.executeTurn( + context, + externalSignal, + !resetStats, + options.initialExternalInputs, + ); + } finally { + this.executing = false; + } + } + + private async executeTurn( + context: ContextState, + externalSignal?: AbortSignal, + preserveStats = false, + initialExternalInputs?: readonly AgentExternalInput[], ): Promise { const initialMessagesOverride = context.get('initial_messages_override') as | Content[] | undefined; + const isContinuation = this.hasStartedReasoning; // Record the initial user turn in the observable message log before // anything that can throw — createChat / prepareTools failures still // get a transcript showing the task that was asked, which is what @@ -215,11 +255,38 @@ export class AgentHeadless { const initialTaskText = String( (context.get('task_prompt') as string) ?? 'Get Started!', ); - if (!initialMessagesOverride || initialMessagesOverride.length === 0) { + const claimedExternalInputs = + isContinuation && initialExternalInputs?.length + ? initialExternalInputs + : undefined; + if (claimedExternalInputs) { + for (const input of claimedExternalInputs) { + this.core.eventEmitter.emit(AgentEventType.EXTERNAL_MESSAGE, { + subagentId: this.core.subagentId, + kind: typeof input === 'string' ? 'message' : input.kind, + text: typeof input === 'string' ? input : input.text, + timestamp: Date.now(), + }); + } + } else if (isContinuation) { + this.core.eventEmitter.emit(AgentEventType.EXTERNAL_MESSAGE, { + subagentId: this.core.subagentId, + kind: 'message', + text: initialTaskText, + timestamp: Date.now(), + }); + } else if ( + !initialMessagesOverride || + initialMessagesOverride.length === 0 + ) { this.core.pushMessage('user', initialTaskText); } - const chat = await this.core.createChat(context); + let chat = this.chat; + if (!chat) { + chat = await this.core.createChat(context); + this.chat = chat; + } if (!chat) { this.terminateMode = AgentTerminateMode.ERROR; @@ -231,16 +298,38 @@ export class AgentHeadless { const abortController = createChildAbortController(externalSignal); try { - const toolsList = await this.core.prepareTools(); - - const initialMessages = - initialMessagesOverride && initialMessagesOverride.length > 0 + if (!this.toolsList) { + this.toolsList = await this.core.prepareTools(); + } + const toolsList = this.toolsList; + + const initialMessages = isContinuation + ? [ + { + role: 'user' as const, + parts: claimedExternalInputs + ? claimedExternalInputs.map((input) => ({ + text: + typeof input === 'string' + ? `${EXTERNAL_MESSAGE_PREFIX} ${input}` + : input.text, + })) + : [{ text: `${EXTERNAL_MESSAGE_PREFIX} ${initialTaskText}` }], + }, + ] + : initialMessagesOverride && initialMessagesOverride.length > 0 ? initialMessagesOverride : [{ role: 'user' as const, parts: [{ text: initialTaskText }] }]; - const startTime = Date.now(); - this.core.executionStats.startTimeMs = startTime; - this.core.stats.start(startTime); + const startTime = + preserveStats && this.core.executionStats.startTimeMs > 0 + ? this.core.executionStats.startTimeMs + : Date.now(); + const roundOffset = preserveStats ? this.core.executionStats.rounds : 0; + if (!preserveStats || this.core.executionStats.startTimeMs === 0) { + this.core.executionStats.startTimeMs = startTime; + this.core.stats.start(startTime); + } try { // Emit start event @@ -265,6 +354,7 @@ export class AgentHeadless { logSubagentExecution(this.core.runtimeContext, startEvent); // Delegate to AgentCore's reasoning loop + this.hasStartedReasoning = true; const result = await this.core.runReasoningLoop( chat, initialMessages, @@ -274,6 +364,7 @@ export class AgentHeadless { maxTurns: this.core.runConfig.max_turns, maxTimeMinutes: this.core.runConfig.max_time_minutes, startTimeMs: startTime, + roundOffset, getExternalMessages: this.externalMessageProvider, waitForExternalMessages: this.externalMessageWaiter, shouldWaitForExternalMessages: this.externalMessageWaitPredicate, @@ -293,7 +384,8 @@ export class AgentHeadless { throw error; } finally { - this.core.executionStats.totalDurationMs = Date.now() - startTime; + this.core.executionStats.totalDurationMs = + Date.now() - this.core.executionStats.startTimeMs; const summary = this.core.stats.getSummary(Date.now()); this.core.eventEmitter?.emit(AgentEventType.FINISH, { subagentId: this.core.subagentId, diff --git a/packages/core/src/agents/runtime/agent-statistics.test.ts b/packages/core/src/agents/runtime/agent-statistics.test.ts index fee2e0b38a8..18b7fdd88a1 100644 --- a/packages/core/src/agents/runtime/agent-statistics.test.ts +++ b/packages/core/src/agents/runtime/agent-statistics.test.ts @@ -16,6 +16,30 @@ describe('AgentStatistics', () => { }); describe('basic statistics tracking', () => { + it('should reset all statistics for a new turn', () => { + stats.start(baseTime); + stats.setRounds(3); + stats.recordToolCall('file_read', false, 100, 'failed'); + stats.recordTokens(100, 50, 10, 5, 160); + + stats.reset(); + + expect(stats.getSummary(baseTime + 5000)).toEqual({ + rounds: 0, + totalDurationMs: 0, + totalToolCalls: 0, + successfulToolCalls: 0, + failedToolCalls: 0, + successRate: 0, + inputTokens: 0, + outputTokens: 0, + thoughtTokens: 0, + cachedTokens: 0, + totalTokens: 0, + toolUsage: [], + }); + }); + it('should track execution time', () => { stats.start(baseTime); const summary = stats.getSummary(baseTime + 5000); diff --git a/packages/core/src/agents/runtime/agent-statistics.ts b/packages/core/src/agents/runtime/agent-statistics.ts index 6128531b491..15d7b2f533e 100644 --- a/packages/core/src/agents/runtime/agent-statistics.ts +++ b/packages/core/src/agents/runtime/agent-statistics.ts @@ -42,6 +42,20 @@ export class AgentStatistics { private apiTotalTokens = 0; private toolUsage = new Map(); + reset(): void { + this.startTimeMs = 0; + this.rounds = 0; + this.totalToolCalls = 0; + this.successfulToolCalls = 0; + this.failedToolCalls = 0; + this.inputTokens = 0; + this.outputTokens = 0; + this.thoughtTokens = 0; + this.cachedTokens = 0; + this.apiTotalTokens = 0; + this.toolUsage.clear(); + } + start(now = Date.now()) { this.startTimeMs = now; } diff --git a/packages/core/src/agents/runtime/workflow-orchestrator.ts b/packages/core/src/agents/runtime/workflow-orchestrator.ts index c2caeeef58d..0e923ae2c39 100644 --- a/packages/core/src/agents/runtime/workflow-orchestrator.ts +++ b/packages/core/src/agents/runtime/workflow-orchestrator.ts @@ -535,12 +535,10 @@ function reportTokens( * "not available in this build" signal), and always disposes the * per-agent registry/hooks in `finally`. * - * Why opts.model goes into `SubagentConfig.model` (not - * `modelConfigOverrides`): `SubagentManager.buildRuntimeContentGeneratorView` - * (subagent-manager.ts:945) consults `SubagentConfig.model` to decide - * whether to build a dedicated ContentGenerator for a different provider - * — `modelConfigOverrides` would only swap the model name within the - * existing provider's runtime view. + * Why opts.model goes into `SubagentConfig.model`: the declarative selector + * remains the authoritative source for an explicitly requested provider and + * lets `SubagentManager.buildRuntimeContentGeneratorView` preserve that route + * when a caller also supplies effective model overrides. * * Why disallowed-floor is augmented on `SubagentConfig.disallowedTools` * (not via `toolConfigOverride`): augmenting before `convertToRuntimeConfig` diff --git a/packages/core/src/subagents/subagent-manager.test.ts b/packages/core/src/subagents/subagent-manager.test.ts index 32d700f0e59..d930e752854 100644 --- a/packages/core/src/subagents/subagent-manager.test.ts +++ b/packages/core/src/subagents/subagent-manager.test.ts @@ -2120,6 +2120,53 @@ bad`); expect(mockCreateContentGenerator).not.toHaveBeenCalled(); }); + it('should pin an inherited effective model to its launch provider', async () => { + const config = { ...agentConfig, model: 'inherit' }; + const launchGenerator = { generateContentStream: vi.fn() }; + mockCreateContentGenerator.mockResolvedValue(launchGenerator); + + await manager.createAgentHeadless(config, mockConfig, { + modelConfigOverrides: { model: 'parent-model' }, + }); + + expect(mockCreateContentGenerator).toHaveBeenCalledWith( + expect.objectContaining({ + model: 'parent-model', + authType: AuthType.USE_OPENAI, + apiKey: 'parent-key', + }), + mockConfig, + ); + const { modelConfig, runtimeView } = destructureAgentHeadlessCall( + mockAgentHeadlessCreate.mock.calls[0], + ); + expect(modelConfig).toEqual( + expect.objectContaining({ model: 'parent-model' }), + ); + expect(runtimeView).toEqual( + expect.objectContaining({ + contentGenerator: launchGenerator, + contentGeneratorConfig: expect.objectContaining({ + model: 'parent-model', + authType: AuthType.USE_OPENAI, + }), + }), + ); + + vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({ + model: 'model-b', + authType: AuthType.USE_ANTHROPIC, + apiKey: 'replacement-key', + }); + expect(runtimeView?.contentGenerator).toBe(launchGenerator); + expect(runtimeView?.contentGeneratorConfig).toEqual( + expect.objectContaining({ + model: 'parent-model', + authType: AuthType.USE_OPENAI, + }), + ); + }); + it('should NOT create a new ContentGenerator when model is omitted', async () => { await manager.createAgentHeadless(agentConfig, mockConfig); diff --git a/packages/core/src/subagents/subagent-manager.ts b/packages/core/src/subagents/subagent-manager.ts index 0d0041b7ef4..717d976cd17 100644 --- a/packages/core/src/subagents/subagent-manager.ts +++ b/packages/core/src/subagents/subagent-manager.ts @@ -815,13 +815,13 @@ export class SubagentManager { const toolConfig = options?.toolConfigOverride ?? runtimeConfig.toolConfig; - // When the model selector specifies a different provider, build a - // dedicated ContentGenerator + view so the subagent talks to the - // right API without affecting the parent process. The view is + // Build a dedicated ContentGenerator view for either a declarative + // model selector or a caller-pinned effective model. The view is // applied via AsyncLocalStorage when the agent runs. const runtimeView = await this.buildRuntimeContentGeneratorView( config, runtimeContext, + options?.modelConfigOverrides?.model, ); const { context: subagentContext, cleanup } = @@ -1004,11 +1004,10 @@ export class SubagentManager { } /** - * When a subagent's model selector resolves to a concrete model, build a - * dedicated ContentGenerator and the view the agent runtime should publish - * via AsyncLocalStorage during the run. Returns `undefined` when no - * override is needed — including `inherit`, an unset `fast` selector, or - * any selector that fails to resolve to a configured model. + * When a subagent's model selector or caller-pinned effective model resolves + * to a concrete model, build a dedicated ContentGenerator and the view the + * agent runtime should publish via AsyncLocalStorage during the run. Returns + * `undefined` when neither source resolves to a concrete model. * * FileReadCache isolation and tool-registry rebuilding are handled * separately in {@link buildSubagentContextOverride} — every subagent @@ -1018,8 +1017,16 @@ export class SubagentManager { private async buildRuntimeContentGeneratorView( config: SubagentConfig, base: Config, + effectiveModelOverride?: string, ): Promise { - const resolvedModel = this.resolveModelOverride(config.model, base); + const configuredModel = this.resolveModelOverride(config.model, base); + const overriddenModel = effectiveModelOverride?.trim() + ? resolveModelId(effectiveModelOverride, buildModelIdContext(base)) + : undefined; + const resolvedModel = + configuredModel?.modelId === overriddenModel?.modelId + ? configuredModel + : (overriddenModel ?? configuredModel); if (!resolvedModel) { return undefined; } diff --git a/packages/core/src/tools/agent/agent.test.ts b/packages/core/src/tools/agent/agent.test.ts index 05fac30211e..4dfdad02b7f 100644 --- a/packages/core/src/tools/agent/agent.test.ts +++ b/packages/core/src/tools/agent/agent.test.ts @@ -4254,6 +4254,7 @@ describe('AgentTool', () => { describe('Agent-level background: true', () => { let mockAgent: AgentHeadless; let mockContextState: ContextState; + let mockSubagentDispose: ReturnType; let mockRegistry: { assertCanStartBackgroundAgent: ReturnType; canStartBackgroundAgent: ReturnType; @@ -4261,6 +4262,7 @@ describe('AgentTool', () => { waitForBackgroundSlot: ReturnType; releaseBackgroundSlot: ReturnType; getQueuedCount: ReturnType; + get: ReturnType; register: ReturnType; unregisterForeground: ReturnType; complete: ReturnType; @@ -4271,6 +4273,12 @@ describe('AgentTool', () => { queueExternalInput: ReturnType; wakeExternalInputWaiters: ReturnType; appendActivity: ReturnType; + registerResidentAgent: ReturnType; + claimPendingInputsForResident: ReturnType; + unregisterResidentAgent: ReturnType; + disposeResidentAgent: ReturnType; + restartCompletedAgent: ReturnType; + trackAgentExecution: ReturnType; }; const bgSubagent: SubagentConfig = { @@ -4301,6 +4309,7 @@ describe('AgentTool', () => { mockContextState = { set: vi.fn() } as unknown as ContextState; MockedContextState.mockImplementation(() => mockContextState); + const restartedEntry = { status: 'running' }; mockRegistry = { assertCanStartBackgroundAgent: vi.fn(), canStartBackgroundAgent: vi.fn().mockReturnValue(true), @@ -4312,6 +4321,7 @@ describe('AgentTool', () => { .mockResolvedValue({ id: Symbol('background-slot') }), releaseBackgroundSlot: vi.fn(), getQueuedCount: vi.fn().mockReturnValue(0), + get: vi.fn().mockReturnValue(restartedEntry), register: vi.fn(), unregisterForeground: vi.fn(), complete: vi.fn(), @@ -4322,6 +4332,20 @@ describe('AgentTool', () => { queueExternalInput: vi.fn(), wakeExternalInputWaiters: vi.fn(), appendActivity: vi.fn(), + registerResidentAgent: vi.fn(), + claimPendingInputsForResident: vi.fn().mockReturnValue([]), + unregisterResidentAgent: vi.fn().mockReturnValue(true), + disposeResidentAgent: vi.fn( + ( + _agentId: string, + resident?: { dispose: () => void | Promise }, + ) => { + resident?.dispose(); + return true; + }, + ), + restartCompletedAgent: vi.fn().mockReturnValue(restartedEntry), + trackAgentExecution: vi.fn(), }; vi.mocked(config.getApprovalMode).mockReturnValue(ApprovalMode.DEFAULT); @@ -4345,9 +4369,10 @@ describe('AgentTool', () => { ] = vi.fn(); vi.mocked(mockSubagentManager.loadSubagent).mockResolvedValue(bgSubagent); + mockSubagentDispose = vi.fn().mockResolvedValue(undefined); vi.mocked(mockSubagentManager.createAgentHeadless).mockResolvedValue({ subagent: mockAgent, - dispose: vi.fn().mockResolvedValue(undefined), + dispose: mockSubagentDispose, }); }); @@ -4366,6 +4391,10 @@ describe('AgentTool', () => { const llmText = partToString(result.llmContent); expect(llmText).toContain('Background agent launched'); + expect(mockRegistry.trackAgentExecution).toHaveBeenCalledOnce(); + expect(mockRegistry.trackAgentExecution).toHaveBeenCalledWith( + expect.any(Promise), + ); expect(llmText).toContain( `Use ${ToolNames.SEND_MESSAGE} to continue this agent`, ); @@ -4404,10 +4433,55 @@ describe('AgentTool', () => { expect.any(String), expect.objectContaining({ persistedCliFlags: expect.objectContaining({ - model: 'subagent-model', + model: 'parent-model', }), }), ); + expect(mockSubagentManager.createAgentHeadless).toHaveBeenCalledTimes(1); + writeMetaSpy.mockRestore(); + }); + + it('pins an inherited model for resident continuations', async () => { + const writeMetaSpy = vi.spyOn(transcript, 'writeAgentMeta'); + vi.mocked(config.getModel).mockReturnValue('model-a'); + vi.mocked(mockAgent.getCore).mockReturnValue({ + modelConfig: {}, + getEventEmitter: () => ({ on: vi.fn(), off: vi.fn() }), + } as unknown as ReturnType); + + const invocation = ( + agentTool as AgentToolWithProtectedMethods + ).createInvocation({ + description: 'Start monitor', + prompt: 'Watch for changes', + subagent_type: 'monitor', + }); + await invocation.execute(); + + const runtimeConfig = vi.mocked(mockSubagentManager.createAgentHeadless) + .mock.calls[0]?.[1] as Config; + expect(runtimeConfig.getModel()).toBe('model-a'); + expect(mockSubagentManager.createAgentHeadless).toHaveBeenCalledWith( + expect.anything(), + runtimeConfig, + expect.objectContaining({ + modelConfigOverrides: { model: 'model-a' }, + }), + ); + expect(mockRegistry.register).toHaveBeenCalledWith( + expect.objectContaining({ model: 'model-a' }), + expect.any(Object), + ); + expect(writeMetaSpy).toHaveBeenCalledWith( + expect.any(String), + expect.objectContaining({ + model: 'model-a', + persistedCliFlags: expect.objectContaining({ model: 'model-a' }), + }), + ); + + vi.mocked(config.getModel).mockReturnValue('model-b'); + expect(runtimeConfig.getModel()).toBe('model-a'); writeMetaSpy.mockRestore(); }); @@ -4502,7 +4576,12 @@ describe('AgentTool', () => { ); }); - it('cleans up owned monitor routing when a background agent finishes', async () => { + it('keeps runtime resources while idle and cleans them when disposed', async () => { + let resolveDispose!: () => void; + const disposeGate = new Promise((resolve) => { + resolveDispose = resolve; + }); + mockSubagentDispose.mockReturnValueOnce(disposeGate); const params: AgentParams = { description: 'Start monitor', prompt: 'Watch for changes', @@ -4521,6 +4600,26 @@ describe('AgentTool', () => { cancelRunningForOwner: ReturnType; }; + await vi.waitFor(() => { + expect(mockRegistry.complete).toHaveBeenCalled(); + }); + expect( + monitorRegistry.setAgentNotificationCallback, + ).not.toHaveBeenCalledWith(agentId, undefined); + expect(mockSubagentDispose).not.toHaveBeenCalled(); + + const resident = mockRegistry.registerResidentAgent.mock.calls[0]?.[1] as + | { dispose: () => void | Promise } + | undefined; + expect(resident).toBeDefined(); + const cleanup = resident?.dispose(); + expect(cleanup).toBeInstanceOf(Promise); + const cleanupPromise = cleanup as Promise; + let cleanupSettled = false; + void cleanupPromise.then(() => { + cleanupSettled = true; + }); + await vi.waitFor(() => { expect( monitorRegistry.setAgentNotificationCallback, @@ -4534,6 +4633,232 @@ describe('AgentTool', () => { { notify: false }, ); }); + expect(mockSubagentDispose).toHaveBeenCalledOnce(); + expect(cleanupSettled).toBe(false); + + resolveDispose(); + await cleanupPromise; + expect(cleanupSettled).toBe(true); + }); + + it('continues a completed background agent on the same runtime', async () => { + const invocation = ( + agentTool as AgentToolWithProtectedMethods + ).createInvocation({ + description: 'Start monitor', + prompt: 'Watch for changes', + subagent_type: 'monitor', + }); + + await invocation.execute(); + await vi.waitFor(() => { + expect(mockRegistry.complete).toHaveBeenCalledTimes(1); + }); + + const resident = mockRegistry.registerResidentAgent.mock.calls[0]?.[1] as + | { continue: (message: string) => boolean } + | undefined; + expect(resident).toBeDefined(); + expect(resident?.continue('Now inspect the helper')).toBe(true); + expect(mockRegistry.trackAgentExecution).toHaveBeenCalledTimes(2); + + await vi.waitFor(() => { + expect(mockAgent.execute).toHaveBeenCalledTimes(2); + expect(mockRegistry.complete).toHaveBeenCalledTimes(2); + }); + expect(mockRegistry.restartCompletedAgent).toHaveBeenCalledWith( + expect.stringContaining('monitor-'), + expect.any(AbortController), + ); + expect(mockContextState.set).toHaveBeenCalledWith( + 'task_prompt', + 'Now inspect the helper', + ); + expect(mockSubagentManager.createAgentHeadless).toHaveBeenCalledTimes(1); + expect(mockSubagentDispose).not.toHaveBeenCalled(); + }); + + it('runs inputs queued during the finishing hook before publishing completion', async () => { + let releaseStopHook: (() => void) | undefined; + const stopHookGate = new Promise((resolve) => { + releaseStopHook = resolve; + }); + let markStopHookStarted: (() => void) | undefined; + const stopHookStarted = new Promise((resolve) => { + markStopHookStarted = resolve; + }); + const hookSystem = { + fireSubagentStartEvent: vi.fn().mockResolvedValue(undefined), + fireSubagentStopEvent: vi + .fn() + .mockImplementationOnce(async () => { + markStopHookStarted?.(); + await stopHookGate; + return undefined; + }) + .mockResolvedValue(undefined), + }; + (config as unknown as Record)['getHookSystem'] = vi + .fn() + .mockReturnValue(hookSystem); + mockRegistry.claimPendingInputsForResident + .mockReturnValueOnce([ + 'late message', + { + kind: 'notification', + text: 'ready', + }, + ]) + .mockReturnValue([]); + + const invocation = ( + agentTool as AgentToolWithProtectedMethods + ).createInvocation({ + description: 'Start monitor', + prompt: 'Watch for changes', + subagent_type: 'monitor', + }); + await invocation.execute(); + await stopHookStarted; + + expect(mockRegistry.complete).not.toHaveBeenCalled(); + releaseStopHook?.(); + + await vi.waitFor(() => { + expect(mockAgent.execute).toHaveBeenCalledTimes(2); + expect(mockRegistry.complete).toHaveBeenCalledTimes(1); + }); + const resident = mockRegistry.registerResidentAgent.mock.calls[0]?.[1]; + expect( + mockRegistry.claimPendingInputsForResident, + ).toHaveBeenNthCalledWith( + 1, + expect.stringContaining('monitor-'), + resident, + ); + expect(vi.mocked(mockAgent.execute).mock.calls[1]?.[2]).toEqual({ + resetStats: false, + initialExternalInputs: [ + 'late message', + { + kind: 'notification', + text: 'ready', + }, + ], + }); + expect(hookSystem.fireSubagentStopEvent).toHaveBeenCalledTimes(2); + expect(mockRegistry.restartCompletedAgent).not.toHaveBeenCalled(); + expect(mockSubagentManager.createAgentHeadless).toHaveBeenCalledTimes(1); + }); + + it('does not retain an agent whose frontmatter hooks are globally registered', async () => { + vi.mocked(mockSubagentManager.loadSubagent).mockResolvedValue({ + ...bgSubagent, + hooks: { PreToolUse: [] }, + }); + const invocation = ( + agentTool as AgentToolWithProtectedMethods + ).createInvocation({ + description: 'Start hooked monitor', + prompt: 'Watch for changes', + subagent_type: 'monitor', + }); + + await invocation.execute(); + await vi.waitFor(() => { + expect(mockRegistry.complete).toHaveBeenCalled(); + expect(mockSubagentDispose).toHaveBeenCalledOnce(); + }); + expect(mockRegistry.registerResidentAgent).not.toHaveBeenCalled(); + }); + + it('does not retain an agent that needs a child-only AUTO permission lease', async () => { + let releaseExecution: (() => void) | undefined; + vi.mocked(mockAgent.execute).mockImplementation( + () => + new Promise((resolve) => { + releaseExecution = resolve; + }), + ); + vi.mocked(mockSubagentManager.loadSubagent).mockResolvedValue({ + ...bgSubagent, + approvalMode: 'auto', + }); + const invocation = ( + agentTool as AgentToolWithProtectedMethods + ).createInvocation({ + description: 'Run classified work', + prompt: 'Inspect the helper', + subagent_type: 'monitor', + }); + + await invocation.execute(); + vi.mocked(config.getApprovalMode).mockReturnValue(ApprovalMode.AUTO); + releaseExecution?.(); + await vi.waitFor(() => { + expect(mockRegistry.complete).toHaveBeenCalled(); + expect(mockSubagentDispose).toHaveBeenCalledOnce(); + }); + expect(mockRegistry.registerResidentAgent).not.toHaveBeenCalled(); + }); + + it('disposes an idle AUTO resident if the parent leaves AUTO mode', async () => { + vi.mocked(config.getApprovalMode).mockReturnValue(ApprovalMode.AUTO); + vi.mocked(mockSubagentManager.loadSubagent).mockResolvedValue({ + ...bgSubagent, + approvalMode: 'auto', + }); + const invocation = ( + agentTool as AgentToolWithProtectedMethods + ).createInvocation({ + description: 'Run classified work', + prompt: 'Inspect the helper', + subagent_type: 'monitor', + }); + + await invocation.execute(); + await vi.waitFor(() => { + expect(mockRegistry.complete).toHaveBeenCalled(); + }); + const resident = mockRegistry.registerResidentAgent.mock.calls[0]?.[1] as + | { continue: (message: string) => boolean } + | undefined; + expect(resident).toBeDefined(); + expect(mockSubagentDispose).not.toHaveBeenCalled(); + + vi.mocked(config.getApprovalMode).mockReturnValue(ApprovalMode.DEFAULT); + expect(resident?.continue('Continue')).toBe(false); + + expect(mockRegistry.unregisterResidentAgent).toHaveBeenCalled(); + expect(mockSubagentDispose).toHaveBeenCalledOnce(); + }); + + it('disposes an idle resident after the session working directory changes', async () => { + const invocation = ( + agentTool as AgentToolWithProtectedMethods + ).createInvocation({ + description: 'Inspect the project', + prompt: 'Inspect the helper', + subagent_type: 'monitor', + }); + + await invocation.execute(); + await vi.waitFor(() => { + expect(mockRegistry.complete).toHaveBeenCalled(); + }); + const resident = mockRegistry.registerResidentAgent.mock.calls[0]?.[1] as + | { continue: (message: string) => boolean } + | undefined; + expect(resident).toBeDefined(); + + vi.mocked(config.getWorkingDir).mockReturnValue('/other/project'); + expect(resident?.continue('Continue')).toBe(false); + + expect(mockRegistry.restartCompletedAgent).not.toHaveBeenCalled(); + expect(mockRegistry.unregisterResidentAgent).toHaveBeenCalled(); + await vi.waitFor(() => { + expect(mockSubagentDispose).toHaveBeenCalledOnce(); + }); }); it('should run in background when run_in_background is true even without background config', async () => { diff --git a/packages/core/src/tools/agent/agent.ts b/packages/core/src/tools/agent/agent.ts index 00a34ec633f..ca8c611d777 100644 --- a/packages/core/src/tools/agent/agent.ts +++ b/packages/core/src/tools/agent/agent.ts @@ -114,7 +114,10 @@ import { writeAgentMeta, type AgentPersistedCliFlags, } from '../../agents/agent-transcript.js'; -import type { BackgroundSlotReservation } from '../../agents/background-tasks.js'; +import type { + BackgroundSlotReservation, + ResidentBackgroundAgent, +} from '../../agents/background-tasks.js'; import { getGitBranch } from '../../utils/gitUtils.js'; import { buildModelIdContext, resolveModelId } from '../../utils/modelId.js'; @@ -870,7 +873,6 @@ export class AgentTool extends BaseDeclarativeTool { const teamGuidance = this.config.isAgentTeamEnabled() ? `**For tasks requiring multiple agents to coordinate, communicate, or work as a team**: Use ${ToolNames.TEAM_CREATE} first to create a team, then spawn teammates using the Agent tool with the \`name\` parameter (the active team is selected automatically). Teams enable message passing between agents, shared task lists, and coordinated workflows. If the user asks for agents to collaborate, review each other's work, or produce a consolidated result — create a team.` : ''; - const baseDescription = `Launch a new agent to handle complex, multi-step tasks autonomously. The Agent tool launches specialized agents (subprocesses) that autonomously handle complex tasks. Each agent type has specific capabilities and tools available to it. @@ -1741,7 +1743,9 @@ class AgentToolInvocation extends BaseToolInvocation { typedStopOutput.getEffectiveReason(), ); continueContext.set('hook_context', ''); - await subagent.execute(continueContext, signal); + await subagent.execute(continueContext, signal, { + resetStats: false, + }); if (signal?.aborted) return undefined; } catch (hookError) { @@ -2749,6 +2753,7 @@ class AgentToolInvocation extends BaseToolInvocation { // Date.now() alone collides when two parallel background agents of the // same type land in the same ms; the registry is keyed by agentId. const agentIdSuffix = this.callId ?? randomUUID().slice(0, 8); + const launchDepth = childLaunchDepth(); const hookOpts = { agentId: `${subagentConfig.name}-${agentIdSuffix}`, // Resolved config name, not the raw requested type: a fork request @@ -2761,10 +2766,42 @@ class AgentToolInvocation extends BaseToolInvocation { updateOutput, }; - // Create the subagent. Fork bypasses SubagentManager because its - // runtime configs are synthesized from the parent's cache-safe params. + const shouldBubble = Boolean( + shouldRunInBackground && + subagentConfig.approvalMode === BUBBLE_APPROVAL_MODE && + this.config.isInteractive(), + ); + // Background agents have no inline UI. Preserve the resolved approval + // mode while overriding only the prompt-avoidance policy used by their + // scheduler. + const subagentRuntimeConfig = shouldRunInBackground + ? (Object.create(agentConfig) as Config) + : agentConfig; + if (shouldRunInBackground) { + subagentRuntimeConfig.getShouldAvoidPermissionPrompts = () => + !shouldBubble; + if (subagentModelId) { + const launchModel = subagentModelId; + subagentRuntimeConfig.getModel = () => launchModel; + } + } + + // Background agents need a dedicated emitter so their transcript never + // receives events from concurrent agents using the parent tool emitter. + // Choose it before construction so every launch creates exactly one + // runtime; the old background branch constructed a second runtime and + // leaked the first one. + const backgroundEventEmitter = shouldRunInBackground + ? new AgentEventEmitter() + : undefined; + + // Create the subagent. Fork bypasses SubagentManager because its runtime + // configs are synthesized from the parent's cache-safe params. let subagent: AgentHeadless; let taskPrompt: string; + let initialMessages: Content[] | undefined; + let promptConfig: PromptConfig | undefined; + let toolConfig: ToolConfig | undefined; // Per-spawn cleanup the subagent manager returns. The caller MUST // invoke this in the same `finally` block that wraps `execute()` — @@ -2775,14 +2812,25 @@ class AgentToolInvocation extends BaseToolInvocation { // dispose, so this stays undefined on the fork path. let subagentDispose: (() => Promise) | undefined; if (isFork) { - const fork = await this.createForkSubagent(agentConfig); + const fork = await this.createForkSubagent( + subagentRuntimeConfig as Config, + backgroundEventEmitter, + ); subagent = fork.subagent; taskPrompt = fork.taskPrompt; + initialMessages = fork.initialMessages; + promptConfig = fork.promptConfig; + toolConfig = fork.toolConfig; } else { const result = await this.subagentManager.createAgentHeadless( subagentConfig, - agentConfig, - { eventEmitter: this.eventEmitter }, + subagentRuntimeConfig as Config, + { + eventEmitter: backgroundEventEmitter ?? this.eventEmitter, + ...(shouldRunInBackground && subagentModelId + ? { modelConfigOverrides: { model: subagentModelId } } + : {}), + }, ); subagent = result.subagent; subagentDispose = result.dispose; @@ -2862,54 +2910,16 @@ class AgentToolInvocation extends BaseToolInvocation { // Non-interactive sessions can't answer, so they keep auto-deny. // (`bubble` resolves to `default` run behavior, so the resolved mode // already requires confirmation — this only flips deny → surface.) - const shouldBubble = Boolean( - subagentConfig.approvalMode === BUBBLE_APPROVAL_MODE && - this.config.isInteractive(), - ); - // Use Object.create so the resolved approval mode override (e.g. - // subagent-level `approvalMode: auto-edit`) is preserved. - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const bgConfig = Object.create(agentConfig) as any; - bgConfig.getShouldAvoidPermissionPrompts = () => !shouldBubble; - - // Register in the background task registry only AFTER init succeeds — if - // construction throws, a pre-registered phantom 'running' entry would hang - // the non-interactive hold-back loop forever. - // Dedicated emitter for this background agent so the transcript - // writer only sees *this* agent's events. Reusing the parent tool's - // UI emitter (this.eventEmitter) would mix events from every - // concurrent fork/subagent into the same transcript. - const bgEventEmitter = new AgentEventEmitter(); - let bgSubagent: AgentHeadless; - let bgInitialMessages: Content[] | undefined; - let bgTaskPrompt: string; - let bgPromptConfig: PromptConfig | undefined; - let bgToolConfig: ToolConfig | undefined; - // Per-spawn cleanup from `createAgentHeadless` (background path). - // The bg `finally` below invokes this alongside the existing - // parent-registry stop; see the foreground call site for the leak - // scenarios it covers. - let bgSubagentDispose: (() => Promise) | undefined; - if (isFork) { - const fork = await this.createForkSubagent( - bgConfig as Config, - bgEventEmitter, - ); - bgSubagent = fork.subagent; - bgInitialMessages = fork.initialMessages; - bgTaskPrompt = fork.taskPrompt; - bgPromptConfig = fork.promptConfig; - bgToolConfig = fork.toolConfig; - } else { - const bgResult = await this.subagentManager.createAgentHeadless( - subagentConfig, - bgConfig as Config, - { eventEmitter: bgEventEmitter }, - ); - bgSubagent = bgResult.subagent; - bgSubagentDispose = bgResult.dispose; - bgTaskPrompt = this.params.prompt; - } + // Register in the background task registry only AFTER init succeeds — + // if construction throws, a pre-registered phantom 'running' entry + // would hang the non-interactive hold-back loop forever. + const bgEventEmitter = backgroundEventEmitter!; + const bgSubagent = subagent; + const bgInitialMessages = initialMessages; + const bgTaskPrompt = taskPrompt; + const bgPromptConfig = promptConfig; + const bgToolConfig = toolConfig; + const bgSubagentDispose = subagentDispose; const registry = this.config.getBackgroundTaskRegistry(); @@ -2953,7 +2963,7 @@ class AgentToolInvocation extends BaseToolInvocation { // Nested-agent lineage (mirrors the meta sidecar); register() // resolves the parent's display name from parentAgentId. parentAgentId: getCurrentAgentId(), - depth: childLaunchDepth(), + depth: launchDepth, }, registerOptions, ); @@ -3007,6 +3017,8 @@ class AgentToolInvocation extends BaseToolInvocation { `[Agent] ToolRegistry stop after background registration failure failed: ${stopError}`, ); }); + void bgSubagentDispose?.().catch(() => {}); + restoreParentPM(); return { llmContent: `${errorMessage}${wtSuffix}`, returnDisplay: this.currentDisplay!, @@ -3052,14 +3064,14 @@ class AgentToolInvocation extends BaseToolInvocation { persistedCliFlags: capturePersistedCliFlags( this.config, resolvedApprovalMode, - bgSubagent.getCore().modelConfig.model, + subagentModelId, ), subagentName: subagentConfig.name, agentColor: subagentConfig.color, resumeCount: 0, // Persisted so resume restores the original nesting level; see // childLaunchDepth() for the rationale. - depth: childLaunchDepth(), + depth: launchDepth, model: subagentModelId, }); @@ -3140,110 +3152,246 @@ class AgentToolInvocation extends BaseToolInvocation { }; }; + // Some launch modes have resources whose lifecycle cannot safely span + // an idle turn yet. Forks carry inherited parent state, temporary + // worktrees are finalized after each turn, and frontmatter hooks are + // currently registered as global matchers. They retain the existing + // transcript-revival behavior. + const canStayResident = + !isFork && + this.params.isolation !== 'worktree' && + (!subagentConfig.hooks || + Object.keys(subagentConfig.hooks).length === 0); + const needsAutoPermissionLease = () => + agentConfig.getApprovalMode() === ApprovalMode.AUTO && + this.config.getApprovalMode() !== ApprovalMode.AUTO; + const residentWorkingDir = agentConfig.getWorkingDir(); + let runtimeDisposed = false; + let disposeRequested = false; + let turnRunning = false; + let currentAbortController: AbortController | undefined = + bgAbortController; + let currentTurnPromise: Promise | undefined; + let runtimeCleanupPromise: Promise | undefined; + let hotContinuationCount = 0; + let residentRegistered = false; + + const cleanupRuntime = () => { + if (runtimeCleanupPromise) return runtimeCleanupPromise; + runtimeDisposed = true; + registry.unregisterResidentAgent( + hookOpts.agentId, + residentController, + ); + residentRegistered = false; + bgEmitter.off(AgentEventType.TOOL_CALL, onToolCall); + bgEmitter.off(AgentEventType.USAGE_METADATA, onUsageMetadata); + cleanupApprovalBridge?.(); + cleanupOwnedMonitorNotifications(); + cleanupJsonl?.(); + runtimeCleanupPromise = Promise.allSettled([ + agentConfig.getToolRegistry().stop(), + bgSubagentDispose?.() ?? Promise.resolve(), + ]).then(() => undefined); + return runtimeCleanupPromise; + }; + + const requestRuntimeDisposal = () => { + if (disposeRequested || runtimeDisposed) { + return runtimeCleanupPromise; + } + disposeRequested = true; + registry.unregisterResidentAgent( + hookOpts.agentId, + residentController, + ); + residentRegistered = false; + currentAbortController?.abort(); + if (!turnRunning) { + return cleanupRuntime(); + } + return undefined; + }; + // Fire-and-forget: start the subagent without blocking the parent. // For forks, wrap the body in runInForkContext so the recursive-fork // guard in execute() fires if the fork child's model calls `agent` // again — otherwise background forks bypass the ALS marker and can // spawn nested forks. - const bgBody = async (recordSpanOutcome: SubagentOutcomeSink) => { + const bgBody = async ( + turnContextState: ContextState, + turnAbortController: AbortController, + recordSpanOutcome: SubagentOutcomeSink, + fireStartHook: boolean, + ) => { + let keepResident = false; + turnRunning = true; try { - await bgSubagent.execute(contextState, bgAbortController.signal); - - let stopHookWarning: string | undefined; - if (hookSystem && !bgAbortController.signal.aborted) { - stopHookWarning = await this.runSubagentStopHookLoop(bgSubagent, { - agentId: hookOpts.agentId, - agentType: hookOpts.agentType, - transcriptPath: jsonlPath, - resolvedMode, - signal: bgAbortController.signal, - }); + if (fireStartHook && hookSystem) { + try { + const startHookOutput = await hookSystem.fireSubagentStartEvent( + hookOpts.agentId, + hookOpts.agentType, + resolvedMode, + turnAbortController.signal, + ); + const additionalContext = + startHookOutput?.getAdditionalContext(); + if (additionalContext) { + turnContextState.set('hook_context', additionalContext); + // The resident chat's system instruction was rendered on + // its first turn, so make new hook context visible in the + // continuation user turn as well. + turnContextState.set( + 'task_prompt', + `${String(turnContextState.get('task_prompt'))}\n\n${additionalContext}`, + ); + } + } catch (hookError) { + debugLogger.warn( + `[Agent] SubagentStart hook failed, continuing execution: ${hookError}`, + ); + } } - // Report terminate mode: only GOAL counts as success. CANCELLED - // keeps the 'cancelled' status so the model sees task_stop's - // effect accurately (with any partial result attached). ERROR, - // MAX_TURNS, TIMEOUT, and SHUTDOWN are surfaced as failures so - // the parent model (and the UI) don't treat incomplete runs as - // completed. - // - // Snapshot the span-relevant terminal state and PUBLISH IT - // FIRST — if the worktree cleanup / registry update / patch - // throws, telemetry must still see the subagent's actual - // outcome (review wenshao @ #4410). - const terminateMode = bgSubagent.getTerminateMode(); - const subagentRawText = bgSubagent.getFinalText(); - recordSpanOutcome( - deriveSubagentOutcomeMetadata({ + let executionContextState = turnContextState; + let initialExternalInputs: + | readonly AgentExternalInput[] + | undefined; + while (true) { + await bgSubagent.execute( + executionContextState, + turnAbortController.signal, + initialExternalInputs + ? { resetStats: false, initialExternalInputs } + : undefined, + ); + + let stopHookWarning: string | undefined; + if (hookSystem && !turnAbortController.signal.aborted) { + stopHookWarning = await this.runSubagentStopHookLoop( + bgSubagent, + { + agentId: hookOpts.agentId, + agentType: hookOpts.agentType, + transcriptPath: jsonlPath, + resolvedMode, + signal: turnAbortController.signal, + }, + ); + } + + // Report terminate mode: only GOAL counts as success. CANCELLED + // keeps the 'cancelled' status so the model sees task_stop's + // effect accurately (with any partial result attached). ERROR, + // MAX_TURNS, TIMEOUT, and SHUTDOWN are surfaced as failures so + // the parent model (and the UI) don't treat incomplete runs as + // completed. + const terminateMode = bgSubagent.getTerminateMode(); + const subagentRawText = bgSubagent.getFinalText(); + const spanOutcome = deriveSubagentOutcomeMetadata({ terminateMode, - signalAborted: bgAbortController.signal.aborted, + signalAborted: turnAbortController.signal.aborted, resultSummaryPresent: Boolean( subagentRawText && subagentRawText.length > 0, ), - }), - ); - - const wtSuffix = formatWorktreeSuffix( - await cleanupWorktreeIsolation(), - ); - const modelVisibleText = toModelVisibleSubagentResult( - subagentRawText, - terminateMode, - ); - const finalText = - appendStopHookBlockingCapWarning( - terminateMode === AgentTerminateMode.GOAL - ? modelVisibleText || - '(subagent produced no model-visible output)' - : modelVisibleText, - stopHookWarning, - ) + wtSuffix; - const completionStats = getCompletionStats(); - if (terminateMode === AgentTerminateMode.GOAL) { - registry.complete(hookOpts.agentId, finalText, completionStats); - patchAgentMeta(metaPath, { - status: 'completed', - lastUpdatedAt: new Date().toISOString(), - lastError: undefined, }); - } else if ( - terminateMode === AgentTerminateMode.CANCELLED || - terminateMode === AgentTerminateMode.SHUTDOWN - ) { - // SHUTDOWN is grouped with CANCELLED in the span taxonomy - // (deriveSubagentOutcomeMetadata); align the registry side - // so dashboards don't see span=cancelled / registry=failed - // mismatch on graceful arena/team-session shutdown. - // wenshao @ #4410. - registry.finalizeCancelled( - hookOpts.agentId, - finalText, - completionStats, - ); - persistBackgroundCancellation( - metaPath, - registry.get(hookOpts.agentId)?.persistedCancellationStatus ?? - 'cancelled', + const mayRemainResident = + terminateMode === AgentTerminateMode.GOAL && residentRegistered; + + // Non-resident turns keep publishing telemetry before worktree + // cleanup so a cleanup failure cannot hide the model outcome. + if (!mayRemainResident) { + recordSpanOutcome(spanOutcome); + } + const wtSuffix = formatWorktreeSuffix( + await cleanupWorktreeIsolation(), ); - } else { - registry.fail( - hookOpts.agentId, - finalText || `Agent terminated with mode: ${terminateMode}`, - completionStats, + if (mayRemainResident) { + if (residentRegistered && !needsAutoPermissionLease()) { + const claimedInputs = registry.claimPendingInputsForResident( + hookOpts.agentId, + residentController, + ); + if (claimedInputs.length > 0) { + executionContextState = new ContextState(); + executionContextState.set('hook_context', ''); + initialExternalInputs = claimedInputs; + continue; + } + } + recordSpanOutcome(spanOutcome); + } + + const modelVisibleText = toModelVisibleSubagentResult( + subagentRawText, + terminateMode, ); - patchAgentMeta(metaPath, { - status: 'failed', - lastUpdatedAt: new Date().toISOString(), - lastError: + const finalText = + appendStopHookBlockingCapWarning( + terminateMode === AgentTerminateMode.GOAL + ? modelVisibleText || + '(subagent produced no model-visible output)' + : modelVisibleText, + stopHookWarning, + ) + wtSuffix; + const completionStats = getCompletionStats(); + if (terminateMode === AgentTerminateMode.GOAL) { + keepResident = + residentRegistered && !needsAutoPermissionLease(); + if (!keepResident) { + registry.unregisterResidentAgent( + hookOpts.agentId, + residentController, + ); + residentRegistered = false; + } + registry.complete(hookOpts.agentId, finalText, completionStats); + patchAgentMeta(metaPath, { + status: 'completed', + lastUpdatedAt: new Date().toISOString(), + lastError: undefined, + }); + } else if ( + terminateMode === AgentTerminateMode.CANCELLED || + terminateMode === AgentTerminateMode.SHUTDOWN + ) { + // SHUTDOWN is grouped with CANCELLED in the span taxonomy + // (deriveSubagentOutcomeMetadata); align the registry side + // so dashboards don't see span=cancelled / registry=failed + // mismatch on graceful arena/team-session shutdown. + // wenshao @ #4410. + registry.finalizeCancelled( + hookOpts.agentId, + finalText, + completionStats, + ); + persistBackgroundCancellation( + metaPath, + registry.get(hookOpts.agentId)?.persistedCancellationStatus ?? + 'cancelled', + ); + } else { + registry.fail( + hookOpts.agentId, finalText || `Agent terminated with mode: ${terminateMode}`, - }); + completionStats, + ); + patchAgentMeta(metaPath, { + status: 'failed', + lastUpdatedAt: new Date().toISOString(), + lastError: + finalText || `Agent terminated with mode: ${terminateMode}`, + }); + } + break; } } catch (error) { // Publish first — same reason as the success path. recordSpanOutcome( deriveSubagentExceptionMetadata( error, - bgAbortController.signal.aborted, + turnAbortController.signal.aborted, ), ); const baseErrorMsg = @@ -3270,7 +3418,7 @@ class AgentToolInvocation extends BaseToolInvocation { // If the error came from a cancellation, preserve the cancelled // status so the model's notification matches what task_stop // requested rather than reporting it as a generic failure. - if (bgAbortController.signal.aborted) { + if (turnAbortController.signal.aborted) { registry.finalizeCancelled( hookOpts.agentId, errorMsg, @@ -3290,73 +3438,139 @@ class AgentToolInvocation extends BaseToolInvocation { }); } } finally { - bgEmitter.off(AgentEventType.TOOL_CALL, onToolCall); - bgEmitter.off(AgentEventType.USAGE_METADATA, onUsageMetadata); - cleanupApprovalBridge?.(); - cleanupOwnedMonitorNotifications(); - cleanupJsonl?.(); - // Release the per-subagent ToolRegistry now that the - // background agent has finished — see the matching call in - // the foreground finally for why. Stopping here, after - // bgSubagent.execute resolves, is safe: by this point the - // detached body cannot invoke any more tool factories on - // this registry. - void agentConfig - .getToolRegistry() - .stop() - .catch(() => {}); - // Per-spawn cleanup from `SubagentManager.createAgentHeadless` - // (background path). Mirrors the foreground finally: releases - // agent-scope hook entries and stops the per-agent ToolRegistry - // owning MCP child processes; not redundant with the parent - // registry stop above. - void bgSubagentDispose?.().catch(() => {}); - // Restore parent PermissionManager's dangerous allow rules - // if this AUTO override stripped them. Background path: - // restore fires when the bg agent terminates (complete / - // fail / cancel), not when this outer execute() returns. + turnRunning = false; restoreParentPM(); + if (!keepResident || disposeRequested) { + await cleanupRuntime(); + } } }; - // Wrap in the agent-identity frame so nested `agent` tool calls - // from this subagent's model record this agent's id as their - // `parentAgentId` in the sidecar meta. Also wrap in - // qwen-code.subagent span (#3731 Phase 3) — background is - // fire-and-forget, so the span gets a new traceId + `Link` to the - // invoking AGENT tool span. `invocationKind` distinguishes a fork - // (subagent_type: "fork") from a named background agent; both are - // long-lived enough to qualify for the 4h TTL safety net. - const framedBgBody = () => - this.runWithSubagentSpan( - this.buildSubagentSpanSpec( - hookOpts, - subagentConfig, - isFork ? 'fork' : 'background', - ), - // bg uses the per-agent abort controller, not the parent turn - // signal — `task_stop` aborts the bg controller alone (silent - // failure: a task_stop'd bg agent was being reported as 'failed' - // because the wrapper saw an unaborted parent signal). - bgAbortController.signal, - (recordOutcome) => - runWithAgentContext(hookOpts.agentId, () => - bgBody(recordOutcome), + // Wrap every turn in a fresh span and the original agent-identity + // frame. The depth override is load-bearing for hot continuations: + // send_message runs from the top level, but the continued agent must + // retain the nesting budget it had when it was created. + const runBackgroundTurn = ( + turnContextState: ContextState, + turnAbortController: AbortController, + fireStartHook: boolean, + ) => { + const framedBgBody = () => + this.runWithSubagentSpan( + this.buildSubagentSpanSpec( + hookOpts, + subagentConfig, + isFork ? 'fork' : 'background', ), - ); - // Defensive `.catch`: bgBody is supposed to handle its own - // errors, but runWithSubagentSpan's `endSubagentSpan` finally - // call could theoretically throw if OTel internals break. - // Without this, such a throw becomes an unhandled rejection - // (Node ≥15 default = process termination). Review wenshao @ - // #4410 + silent-failure-hunter. - const bgPromise = isFork - ? runInForkContext(framedBgBody) - : framedBgBody(); - bgPromise.catch((err) => + turnAbortController.signal, + (recordOutcome) => + runWithAgentContext( + hookOpts.agentId, + () => + bgBody( + turnContextState, + turnAbortController, + recordOutcome, + fireStartHook, + ), + launchDepth, + ), + ); + return isFork ? runInForkContext(framedBgBody) : framedBgBody(); + }; + + const reportUnexpectedBackgroundError = (err: unknown) => { debugLogger.warn( `[Agent] background subagent ${hookOpts.agentId} body raised unexpected rejection: ${err instanceof Error ? err.message : String(err)}`, - ), + ); + }; + + const residentController: ResidentBackgroundAgent = { + continue: (message) => { + if (!canStayResident || disposeRequested || runtimeDisposed) { + return false; + } + if (agentConfig.getWorkingDir() !== residentWorkingDir) { + registry.disposeResidentAgent( + hookOpts.agentId, + residentController, + ); + return false; + } + if (needsAutoPermissionLease()) { + registry.disposeResidentAgent( + hookOpts.agentId, + residentController, + ); + return false; + } + + const nextAbortController = new AbortController(); + let restarted; + try { + restarted = registry.restartCompletedAgent( + hookOpts.agentId, + nextAbortController, + ); + } catch (error) { + debugLogger.warn( + `[Agent] Could not continue resident background agent ${hookOpts.agentId}: ${error instanceof Error ? error.message : String(error)}`, + ); + return false; + } + if ( + !restarted || + disposeRequested || + runtimeDisposed || + registry.get(hookOpts.agentId) !== restarted || + restarted.status !== 'running' + ) { + return false; + } + + liveToolCallCount = 0; + currentAbortController = nextAbortController; + hotContinuationCount += 1; + patchAgentMeta(metaPath, { + status: 'running', + lastUpdatedAt: new Date().toISOString(), + lastError: undefined, + resumeCount: hotContinuationCount, + }); + + const nextContextState = new ContextState(); + nextContextState.set('task_prompt', message); + nextContextState.set('hook_context', ''); + const previousTurn = currentTurnPromise ?? Promise.resolve(); + currentTurnPromise = previousTurn + .catch(reportUnexpectedBackgroundError) + .then(async () => { + if (disposeRequested || runtimeDisposed) return; + await runBackgroundTurn( + nextContextState, + nextAbortController, + true, + ); + }); + registry.trackAgentExecution(currentTurnPromise); + currentTurnPromise.catch(reportUnexpectedBackgroundError); + return true; + }, + dispose: requestRuntimeDisposal, + }; + if (canStayResident && !needsAutoPermissionLease()) { + registry.registerResidentAgent(hookOpts.agentId, residentController); + residentRegistered = true; + } + + // Defensive `.catch`: bgBody handles normal errors, but span teardown + // can still reject if telemetry internals fail. + currentTurnPromise = runBackgroundTurn( + contextState, + bgAbortController, + false, ); + registry.trackAgentExecution(currentTurnPromise); + currentTurnPromise.catch(reportUnexpectedBackgroundError); this.updateDisplay({ status: 'background' as const }, updateOutput); return { @@ -3608,7 +3822,7 @@ class AgentToolInvocation extends BaseToolInvocation { // Nested-agent lineage (mirrors the meta sidecar); register() // resolves the parent's display name from parentAgentId. parentAgentId: getCurrentAgentId(), - depth: childLaunchDepth(), + depth: launchDepth, }); writeAgentMeta(fgMetaPath, { agentId: hookOpts.agentId, @@ -3630,7 +3844,7 @@ class AgentToolInvocation extends BaseToolInvocation { resumeCount: 0, // Persisted so resume restores the original nesting level; see // childLaunchDepth() for the rationale. - depth: childLaunchDepth(), + depth: launchDepth, }); const stopHookWarning = await runFramed(); diff --git a/packages/core/src/tools/send-message.test.ts b/packages/core/src/tools/send-message.test.ts index 822676beaa2..350dc301e53 100644 --- a/packages/core/src/tools/send-message.test.ts +++ b/packages/core/src/tools/send-message.test.ts @@ -35,6 +35,17 @@ describe('SendMessageTool — team mode', () => { expect(tool.name).toBe('send_message'); }); + it('advertises resident-or-transcript continuation', () => { + const tool = new SendMessageTool(makeTeamConfig()); + + expect(tool.description).toContain( + 'a completed task continues on its resident runtime when available', + ); + expect(tool.description).toContain( + 'otherwise is revived from its transcript', + ); + }); + it('sends a message via TeamManager', async () => { const sendMessage = vi.fn().mockResolvedValue(undefined); const tool = new SendMessageTool( @@ -340,7 +351,37 @@ describe('SendMessageTool — background-task mode', () => { expect(result.llmContent).toContain('resumed'); }); - it('revives a completed task with the message as the next instruction', async () => { + it('continues a completed task on its resident runtime', async () => { + registry.register({ + agentId: 'agent-1', + description: 'test agent', + status: 'completed', + startTime: Date.now(), + abortController: new AbortController(), + isBackgrounded: true, + outputFile: '/tmp/test.jsonl', + metaPath: '/tmp/test.meta.json', + }); + const continueResident = vi.fn().mockReturnValue(true); + registry.registerResidentAgent('agent-1', { + continue: continueResident, + dispose: vi.fn(), + }); + + const result = await tool.validateBuildAndExecute( + { task_id: 'agent-1', message: 'now refactor the helper' }, + new AbortController().signal, + ); + + expect(continueResident).toHaveBeenCalledWith('now refactor the helper'); + expect(reviveCompletedBackgroundAgent).not.toHaveBeenCalled(); + expect(resumeBackgroundAgent).not.toHaveBeenCalled(); + expect(result.error).toBeUndefined(); + expect(result.llmContent).toContain('existing runtime'); + expect(result.returnDisplay).toContain('Continued'); + }); + + it('revives a completed task when no resident runtime is available', async () => { registry.register({ agentId: 'agent-1', description: 'test agent', diff --git a/packages/core/src/tools/send-message.ts b/packages/core/src/tools/send-message.ts index 37ad925e4ab..47b455dc6c0 100644 --- a/packages/core/src/tools/send-message.ts +++ b/packages/core/src/tools/send-message.ts @@ -130,10 +130,21 @@ class SendMessageInvocation extends BaseToolInvocation< }; } - // A completed background agent is revived from its persisted transcript - // and continued with this message — lets the model keep iterating on a - // finished sub-agent instead of spawning a fresh one. + // Prefer the same in-process runtime when the completed agent is still + // resident. This preserves its live chat and prepared tool surface. + // Fall back to the existing transcript-based revival path on a miss. if (entry.status === 'completed') { + const continued = registry.continueResidentAgent( + this.params.task_id, + this.params.message, + ); + if (continued) { + return { + llmContent: `Background task "${this.params.task_id}" continued on its existing runtime with your message as the next instruction.`, + returnDisplay: `Continued ${entry.description}`, + }; + } + const revived = await this.config.reviveCompletedBackgroundAgent( this.params.task_id, this.params.message, @@ -256,8 +267,8 @@ export class SendMessageTool extends BaseDeclarativeTool< ToolDisplayNames.SEND_MESSAGE, 'Send a message to a teammate (use "to") or to a running background task (use "task_id"). ' + 'For teams, set "to" to a bare teammate name (no @) or "*" to broadcast. ' + - 'For background tasks, set "task_id" to the id from the launch response, a recovered paused task, or a completed task to revive. ' + - 'Running tasks receive it at the next tool-round boundary; paused recovered tasks are resumed with the message as their first continuation instruction; a completed task is revived from its transcript and continued with your message. ' + + 'For background tasks, set "task_id" to the id from the launch response, a recovered paused task, or a completed task to continue. ' + + 'Running tasks receive it at the next tool-round boundary; paused recovered tasks are resumed with the message as their first continuation instruction; a completed task continues on its resident runtime when available and otherwise is revived from its transcript and continued with your message. ' + 'Your text output is NOT visible to other agents — use this tool to communicate.', Kind.Other, { @@ -270,7 +281,7 @@ export class SendMessageTool extends BaseDeclarativeTool< task_id: { type: 'string', description: - 'The ID of the background task (from the launch response, a recovered paused task, or a completed task to revive).', + 'The ID of the background task (from the launch response, a recovered paused task, or a completed task to continue).', }, message: { type: 'string',