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..14a11a5019a --- /dev/null +++ b/docs/design/2026-07-20-background-agent-hot-continuation.md @@ -0,0 +1,169 @@ +# 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 the in-session lifecycle. Logical discovery and +continuation after restoring the parent session are handled separately by the +background-agent roster restore design. + +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. Logical continuity therefore comes +from restoring the task identity and transcript into the new session, followed +by one cold reconstruction. + +## 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, + terminal-entry eviction, working-directory changes, branch switches, and ACP + session close/disposal. +- Atomically claim input queued in the finishing window before publishing a + successful completion. + +## 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, trace span, task-start event, completion notification, and sidecar status +transition. 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, +worktree-isolated, and child-only AUTO agents continue through the existing +JSONL revival flow. The reconstructed worktree agent runs from the current +parent working directory because its temporary launch worktree has already +been finalized. + +## Races and failure handling + +- Immediately before a compatible runtime publishes successful completion, the + active turn drains the registry queue without yielding. If it claims input, + the same headless runtime executes that input and the task remains running. + If the queue is empty, transcript persistence and the running-to-completed + transition happen synchronously, so a later `send_message` observes the + completed row and uses the resident-continuation path instead of receiving a + misleading queued acknowledgement. Worktree-isolated turns perform their + final drain before teardown because their runtime is intentionally not + continuable afterward. +- 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. +- Failed and cancelled turns remove and dispose the resident controller. +- 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. +- `/branch` refuses running background work and disposes terminal residents only + after the branch has initialized successfully; +- working-directory changes and ACP session disposal release resident runtimes. + +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/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index c7458d707e1..14ae0417ef7 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -333,6 +333,7 @@ describe('Session', () => { tryCompressChat: ReturnType; }; let mockBackgroundTaskRegistry: { + abortAll: ReturnType; setNotificationCallback: ReturnType; hasUnfinalizedTasks: ReturnType; getAll: ReturnType; @@ -476,6 +477,7 @@ describe('Session', () => { }), }; mockBackgroundTaskRegistry = { + abortAll: vi.fn(), setNotificationCallback: vi.fn(), hasUnfinalizedTasks: vi.fn().mockReturnValue(false), getAll: vi.fn().mockReturnValue([]), @@ -16055,6 +16057,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); diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index accd969fa2d..2573db0ffbe 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -1660,6 +1660,7 @@ export class Session implements SessionContext { this.#stopCronSchedulerInRuntime(); } + this.config.getBackgroundTaskRegistry().abortAll({ notify: false }); this.config.getBackgroundTaskRegistry().setNotificationCallback(undefined); this.config.getMonitorRegistry().setNotificationCallback(undefined); this.config.getBackgroundShellRegistry().setNotificationCallback(undefined); diff --git a/packages/cli/src/ui/hooks/useBranchCommand.test.ts b/packages/cli/src/ui/hooks/useBranchCommand.test.ts index fd1230931d1..e431f96cde3 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,50 @@ describe('useBranchCommand', () => { }), getChatRecordingService: () => ({ finalize, flush }), getGeminiClient: () => ({ initialize: vi.fn() }), + getBackgroundTaskRegistry: () => backgroundTaskRegistry, + getMonitorRegistry: () => monitorRegistry, + getBackgroundShellRegistry: () => backgroundShellRegistry, + getWorkflowRunRegistry: () => workflowRunRegistry, startNewSession: startNewSessionConfig, getDebugLogger: () => ({ warn: vi.fn() }), }; }); + it('refuses to branch while background work is running', async () => { + backgroundTaskRegistry.hasRunningTasks.mockReturnValue(true); + + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('blocked'); + }); + + expect(finalize).not.toHaveBeenCalled(); + expect(forkSession).not.toHaveBeenCalled(); + expect(startNewSessionConfig).not.toHaveBeenCalled(); + expect(addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'error', + text: expect.stringContaining('running background tasks'), + }), + expect.any(Number), + ); + }); + + it('clears terminal background state after the branch initializes', async () => { + const { result } = renderHook(() => useBranchCommand(makeOptions())); + await act(async () => { + await result.current.handleBranch('ready'); + }); + + expect(backgroundTaskRegistry.reset).toHaveBeenCalledOnce(); + expect(monitorRegistry.reset).toHaveBeenCalledOnce(); + expect(backgroundShellRegistry.reset).toHaveBeenCalledOnce(); + expect(workflowRunRegistry.reset).toHaveBeenCalledOnce(); + expect(startNewSessionUI.mock.invocationCallOrder[0]).toBeLessThan( + backgroundTaskRegistry.reset.mock.invocationCallOrder[0]!, + ); + }); + 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 @@ -505,6 +576,7 @@ describe('useBranchCommand', () => { expect(startNewSessionUI).not.toHaveBeenCalled(); expect(setSessionName).not.toHaveBeenCalled(); expect(removeSession).toHaveBeenCalledTimes(1); + expect(backgroundTaskRegistry.reset).not.toHaveBeenCalled(); // User sees the failure. expect(addItem).toHaveBeenCalledWith( expect.objectContaining({ diff --git a/packages/cli/src/ui/hooks/useBranchCommand.ts b/packages/cli/src/ui/hooks/useBranchCommand.ts index b153834a7aa..983ab767ea2 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.ts @@ -21,6 +21,13 @@ 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'; + +const BACKGROUND_WORK_BRANCH_BLOCKED_MESSAGE = + "Stop the current session's running background tasks before branching the conversation."; /** * Derives a short one-line title from the first *real* user message in the @@ -93,6 +100,17 @@ export function useBranchCommand( async (name?: string) => { if (!config) return; + if (hasBlockingBackgroundWork(config)) { + historyManager.addItem( + { + type: 'error', + text: t(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/agent-transcript.ts b/packages/core/src/agents/agent-transcript.ts index 689bc4cc97c..23aea30cf29 100644 --- a/packages/core/src/agents/agent-transcript.ts +++ b/packages/core/src/agents/agent-transcript.ts @@ -110,6 +110,8 @@ export interface AgentMeta { * `running` as resumable work that was interrupted by process exit. */ status?: 'running' | 'completed' | 'failed' | 'cancelled' | 'paused'; + /** Whether the original launch used temporary worktree isolation. */ + isolation?: 'worktree'; /** ISO 8601 timestamp of the latest lifecycle transition. */ lastUpdatedAt?: string; /** Resolved approval mode used when the agent was launched. */ @@ -144,6 +146,8 @@ export interface AgentPersistedCliFlags { sandbox?: SandboxConfig | null; screenReader?: boolean; model?: string; + authType?: string; + baseUrl?: string; maxSessionTurns?: number; maxToolCalls?: number; /** diff --git a/packages/core/src/agents/background-agent-resume.test.ts b/packages/core/src/agents/background-agent-resume.test.ts index be71aaaa20c..3ff81ba7595 100644 --- a/packages/core/src/agents/background-agent-resume.test.ts +++ b/packages/core/src/agents/background-agent-resume.test.ts @@ -89,6 +89,10 @@ describe('BackgroundAgentResumeService', () => { isDeferredToolRevealed: vi.fn().mockReturnValue(false), getMcpServerInstructions: vi.fn().mockReturnValue(new Map()), }; + const permissionManager = { + stripDangerousRulesForAutoMode: vi.fn(), + restoreDangerousRules: vi.fn(), + }; const monitorRegistry = { setAgentNotificationCallback: vi.fn(), setAgentLifecycleCallback: vi.fn(), @@ -120,6 +124,7 @@ describe('BackgroundAgentResumeService', () => { getTranscriptPath: () => path.join(tempDir, 'session.jsonl'), getToolRegistry: () => stubToolRegistry, createToolRegistry: vi.fn().mockResolvedValue(stubToolRegistry), + getPermissionManager: () => permissionManager, } as unknown as Config; return { @@ -127,6 +132,9 @@ describe('BackgroundAgentResumeService', () => { subagentManager, hookSystem, monitorRegistry, + config, + permissionManager, + stubToolRegistry, }; } @@ -1230,6 +1238,8 @@ describe('BackgroundAgentResumeService', () => { sandbox: { command: 'docker', image: 'qwen-code-sandbox' }, screenReader: true, model: 'agent-model', + authType: 'anthropic', + baseUrl: 'https://launch-provider.example.com', maxSessionTurns: 7, maxToolCalls: 11, // Deliberately out of range: the resume path must re-normalize @@ -1287,7 +1297,8 @@ describe('BackgroundAgentResumeService', () => { expect(resumed).toBeDefined(); expect(createAgentHeadless).toHaveBeenCalledTimes(1); - const [, overriddenConfig] = createAgentHeadless.mock.calls[0]!; + const [, overriddenConfig, createOptions] = + createAgentHeadless.mock.calls[0]!; expect(overriddenConfig.getApprovalMode()).toBe('auto-edit'); expect(overriddenConfig.getBareMode()).toBe(true); expect(overriddenConfig.getSandbox()).toEqual({ @@ -1296,6 +1307,13 @@ describe('BackgroundAgentResumeService', () => { }); expect(overriddenConfig.getScreenReader()).toBe(true); expect(overriddenConfig.getModel()).toBe('agent-model'); + expect(createOptions.modelConfigOverrides).toEqual({ + model: 'agent-model', + }); + expect(createOptions.runtimeAuthOverrides).toEqual({ + authType: 'anthropic', + baseUrl: 'https://launch-provider.example.com', + }); expect(overriddenConfig.getMaxSessionTurns()).toBe(7); expect(overriddenConfig.getMaxToolCalls()).toBe(11); expect(overriddenConfig.getMaxSubagentDepth()).toBe(100); @@ -1518,8 +1536,10 @@ describe('BackgroundAgentResumeService', () => { releaseExecute = resolve; }), ); + const executeExternalInputs = vi.fn().mockResolvedValue(undefined); const subagent = { execute, + executeExternalInputs, setExternalMessageProvider: vi.fn(), getCore: () => ({ getEventEmitter: () => new AgentEventEmitter() }), getExecutionSummary: () => ({ @@ -1554,7 +1574,12 @@ describe('BackgroundAgentResumeService', () => { | (() => string[]) | undefined; expect(provider).toBeDefined(); - expect(provider?.()).toEqual(['second message']); + expect(executeExternalInputs).toHaveBeenCalledWith( + ['second message'], + expect.any(AbortSignal), + { resetStats: false }, + ); + expect(provider?.()).toEqual([]); }); it('routes owned monitor notifications into a resumed agent queue', async () => { @@ -1656,6 +1681,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, @@ -1735,10 +1764,12 @@ describe('BackgroundAgentResumeService', () => { getTerminateMode: () => AgentTerminateMode.GOAL, getFinalText: () => 'done', }; - const { service, subagentManager, monitorRegistry } = createService(); + const { service, subagentManager, monitorRegistry, stubToolRegistry } = + createService(); + const dispose = vi.fn().mockResolvedValue(undefined); subagentManager.createAgentHeadless.mockResolvedValue({ subagent, - dispose: vi.fn().mockResolvedValue(undefined), + dispose, }); await expect( @@ -1746,6 +1777,9 @@ describe('BackgroundAgentResumeService', () => { ).resolves.toBeUndefined(); expect(subagent.execute).not.toHaveBeenCalled(); + expect(registry.get(agentId)?.status).toBe('paused'); + expect(stubToolRegistry.stop).toHaveBeenCalledTimes(1); + expect(dispose).toHaveBeenCalledTimes(1); expect(monitorRegistry.setAgentNotificationCallback).toHaveBeenCalledWith( agentId, expect.any(Function), @@ -1927,7 +1961,7 @@ describe('BackgroundAgentResumeService', () => { createdAt: '2026-04-20T00:00:00.000Z', status: 'running', subagentName: FORK_SUBAGENT_TYPE, - resolvedApprovalMode: 'default', + resolvedApprovalMode: 'auto', }); fs.writeFileSync( outputFile, @@ -1956,7 +1990,7 @@ describe('BackgroundAgentResumeService', () => { }); const createSpy = vi.spyOn(AgentHeadless, 'create'); - const { service } = createService(); + const { service, permissionManager, stubToolRegistry } = createService(); const resumed = await service.resumeBackgroundAgent(agentId, 'continue'); expect(resumed).toBeUndefined(); @@ -1966,6 +2000,11 @@ describe('BackgroundAgentResumeService', () => { ); expect(registry.get(agentId)?.error).toBeUndefined(); expect(createSpy).not.toHaveBeenCalled(); + expect(stubToolRegistry.stop).toHaveBeenCalledTimes(1); + expect( + permissionManager.stripDangerousRulesForAutoMode, + ).toHaveBeenCalledTimes(1); + expect(permissionManager.restoreDangerousRules).toHaveBeenCalledTimes(1); createSpy.mockRestore(); }); @@ -2348,7 +2387,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); @@ -2423,10 +2462,11 @@ describe('BackgroundAgentResumeService', () => { getFinalText: () => 'iterated', }; + const dispose = vi.fn().mockResolvedValue(undefined); const { service, subagentManager } = createService(); subagentManager.createAgentHeadless.mockResolvedValue({ subagent, - dispose: vi.fn().mockResolvedValue(undefined), + dispose, }); const revived = await service.reviveCompletedBackgroundAgent( @@ -2448,6 +2488,97 @@ describe('BackgroundAgentResumeService', () => { expect(fs.statSync(sessionDir).mtime.getTime()).toBeGreaterThan( oldSessionMtime.getTime(), ); + + expect(registry.continueResidentAgent(agentId, 'tighten the summary')).toBe( + true, + ); + expect(registry.get(agentId)?.status).toBe('running'); + await vi.waitFor(() => { + expect(execute).toHaveBeenCalledTimes(2); + expect(registry.get(agentId)?.status).toBe('completed'); + }); + expect(subagentManager.createAgentHeadless).toHaveBeenCalledTimes(1); + const hotContextArg = execute.mock.calls[1]?.[0]; + expect(hotContextArg?.get('task_prompt')).toBe('tighten the summary'); + expect(readAgentMeta(metaPath)?.resumeCount).toBe(2); + expect(dispose).not.toHaveBeenCalled(); + + registry.reset(); + + expect(dispose).toHaveBeenCalledTimes(1); + expect(registry.continueResidentAgent(agentId, 'again')).toBe(false); + }); + + it('cold-revives a completed worktree-isolated agent without retaining it', async () => { + const agentId = 'completed-isolated'; + const metaPath = path.join(tempDir, `${agentId}.meta.json`); + const outputFile = path.join(tempDir, `${agentId}.jsonl`); + writeAgentMeta(metaPath, { + agentId, + agentType: 'researcher', + description: 'Isolated result', + parentSessionId: 'session-isolated-revive', + parentAgentId: null, + createdAt: '2026-04-20T00:00:00.000Z', + status: 'completed', + isolation: 'worktree', + subagentName: 'researcher', + }); + fs.writeFileSync( + outputFile, + JSON.stringify({ + uuid: 'isolated-result', + parentUuid: null, + sessionId: 'session-isolated-revive', + timestamp: '2026-04-20T00:00:00.000Z', + type: 'user', + message: { role: 'user', parts: [{ text: 'Isolated result' }] }, + }) + '\n', + 'utf8', + ); + registry.register({ + agentId, + description: 'Isolated result', + subagentType: 'researcher', + isBackgrounded: true, + status: 'running', + startTime: Date.now(), + abortController: new AbortController(), + outputFile, + metaPath, + }); + registry.complete(agentId, 'done'); + + const subagent = { + execute: vi.fn().mockResolvedValue(undefined), + executeExternalInputs: vi.fn().mockResolvedValue(undefined), + setExternalMessageProvider: vi.fn(), + getCore: () => ({ getEventEmitter: () => new AgentEventEmitter() }), + getExecutionSummary: () => ({ + totalTokens: 0, + outputTokens: 0, + totalDurationMs: 0, + }), + getTerminateMode: () => AgentTerminateMode.GOAL, + getFinalText: () => 'continued from transcript', + }; + const dispose = vi.fn().mockResolvedValue(undefined); + const { service, subagentManager } = createService(); + subagentManager.createAgentHeadless.mockResolvedValue({ + subagent, + dispose, + }); + + await expect( + service.reviveCompletedBackgroundAgent(agentId, 'continue'), + ).resolves.toBeDefined(); + await vi.waitFor(() => { + expect(registry.get(agentId)?.status).toBe('completed'); + }); + + expect(subagentManager.createAgentHeadless).toHaveBeenCalledOnce(); + expect(registry.continueResidentAgent(agentId, 'again')).toBe(false); + expect(dispose).toHaveBeenCalledOnce(); }); it('does not revive non-completed or transcript-less entries', async () => { @@ -2659,7 +2790,6 @@ describe('BackgroundAgentResumeService', () => { }); } }); - const { service, subagentManager } = createService(); subagentManager.createAgentHeadless.mockRejectedValue( new Error('setup failed'), diff --git a/packages/core/src/agents/background-agent-resume.ts b/packages/core/src/agents/background-agent-resume.ts index 0cf813baa0b..d5f796946bf 100644 --- a/packages/core/src/agents/background-agent-resume.ts +++ b/packages/core/src/agents/background-agent-resume.ts @@ -47,12 +47,17 @@ 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'; import { EXCLUDED_TOOLS_FOR_SUBAGENTS } from './runtime/agent-core.js'; import { ToolNames } from '../tools/tool-names.js'; -import type { PromptConfig, ToolConfig } from './runtime/agent-types.js'; +import type { + AgentExternalInput, + PromptConfig, + ToolConfig, +} from './runtime/agent-types.js'; import type { AgentBootstrapRecordPayload, NotificationRecordPayload, @@ -654,6 +659,31 @@ export class BackgroundAgentResumeService { let cleanupOwnedMonitorNotifications: (() => void) | undefined; let cleanupJsonl: (() => void) | undefined; + let agentConfig: Config | undefined; + let restoreParentPM: (() => void) | undefined; + let subagentDispose: (() => Promise) | undefined; + let cleanupRuntime: (() => void) | undefined; + let runtimeLifecycleOwned = false; + let setupCleaned = false; + + const cleanupPreparedRuntime = () => { + if (runtimeLifecycleOwned || setupCleaned) return; + setupCleaned = true; + if (cleanupRuntime) { + cleanupRuntime(); + } else { + cleanupOwnedMonitorNotifications?.(); + cleanupJsonl?.(); + if (agentConfig) { + void agentConfig + .getToolRegistry() + .stop() + .catch(() => {}); + } + void subagentDispose?.().catch(() => {}); + } + restoreParentPM?.(); + }; try { const subagentName = meta.subagentName ?? meta.agentType; @@ -690,12 +720,15 @@ export class BackgroundAgentResumeService { // continuing to read the parent's. Reusing `this.config` // directly here would short-circuit that isolation. See the // matching wrapper in `agent.ts:createApprovalModeOverride`. - const { config: agentConfig, cleanup: restoreParentPM } = - await createApprovalModeOverride( - this.config, - resolvedApprovalMode as ApprovalMode, - { persistedCliFlags: meta.persistedCliFlags }, - ); + const approvalOverride = await createApprovalModeOverride( + this.config, + resolvedApprovalMode as ApprovalMode, + { persistedCliFlags: meta.persistedCliFlags }, + ); + const activeAgentConfig = approvalOverride.config; + const activeRestoreParentPM = approvalOverride.cleanup; + agentConfig = activeAgentConfig; + restoreParentPM = activeRestoreParentPM; // Mirror the launch path's permission-bubbling gate (agent.ts): an // agent whose definition uses `approvalMode: bubble` surfaces // confirmations to the parent UI instead of auto-denying, in @@ -706,7 +739,7 @@ export class BackgroundAgentResumeService { this.config.isInteractive(), ); // eslint-disable-next-line @typescript-eslint/no-explicit-any - const bgConfig = Object.create(agentConfig) as any; + const bgConfig = Object.create(activeAgentConfig) as any; bgConfig.getShouldAvoidPermissionPrompts = () => !shouldBubble; const records = await jsonl.read(outputFile); @@ -743,6 +776,7 @@ export class BackgroundAgentResumeService { lastUpdatedAt: new Date().toISOString(), }); this.restorePausedEntry(agentId, { resumeBlockedReason: reason }); + cleanupPreparedRuntime(); return undefined; } if (target.isFork && !recovery.forkBootstrap) { @@ -752,6 +786,7 @@ export class BackgroundAgentResumeService { lastUpdatedAt: new Date().toISOString(), }); this.restorePausedEntry(agentId, { resumeBlockedReason: reason }); + cleanupPreparedRuntime(); return undefined; } if ( @@ -765,16 +800,12 @@ export class BackgroundAgentResumeService { lastUpdatedAt: new Date().toISOString(), }); this.restorePausedEntry(agentId, { resumeBlockedReason: reason }); + cleanupPreparedRuntime(); return undefined; } const bgEventEmitter = new AgentEventEmitter(); - // Per-spawn cleanup from `SubagentManager.createAgentHeadless` — - // the resume `finally` invokes this so per-agent hook entries and - // the force-rebuilt ToolRegistry don't leak across the resume - // boundary. Stays undefined on the fork-resume path (forks share - // the parent's registry + hook lifecycle). - let subagentDispose: (() => Promise) | undefined; + const launchModel = meta.model ?? meta.persistedCliFlags?.model; let subagent: AgentHeadless; if (target.isFork) { subagent = await this.createResumedForkSubagent( @@ -791,8 +822,28 @@ export class BackgroundAgentResumeService { promptConfigOverrides: { initialMessages: resumeHistory, }, + ...(launchModel + ? { + modelConfigOverrides: { + model: launchModel, + }, + } + : {}), + ...(meta.persistedCliFlags?.authType + ? { + runtimeAuthOverrides: { + authType: meta.persistedCliFlags.authType, + baseUrl: meta.persistedCliFlags.baseUrl, + }, + } + : {}), }); subagent = result.subagent; + // Per-spawn cleanup from `SubagentManager.createAgentHeadless` — + // the resume `finally` invokes this so per-agent hook entries and + // the force-rebuilt ToolRegistry don't leak across the resume + // boundary. Stays undefined on the fork-resume path (forks share + // the parent's registry + hook lifecycle). subagentDispose = result.dispose; } @@ -921,57 +972,157 @@ export class BackgroundAgentResumeService { ? registry.bridgeApprovalEvents(meta.agentId, bgEmitter) : undefined; - const runBody = async () => { - try { - await subagent.execute(contextState, bgAbortController.signal); - - let stopHookWarning: string | undefined; - if (hookSystem && !bgAbortController.signal.aborted) { - stopHookWarning = await this.runSubagentStopHookLoop(subagent, { - agentId: meta.agentId, - agentType: meta.agentType, - transcriptPath: outputFile, - resolvedMode, - signal: bgAbortController.signal, - }); - } + const canStayResident = + !target.isFork && + meta.isolation !== 'worktree' && + (!target.subagentConfig?.hooks || + Object.keys(target.subagentConfig.hooks).length === 0); + const needsAutoPermissionLease = () => + activeAgentConfig.getApprovalMode() === 'auto' && + this.config.getApprovalMode() !== 'auto'; + let runtimeDisposed = false; + let disposeRequested = false; + let turnRunning = false; + let currentAbortController: AbortController | undefined = + bgAbortController; + let currentTurnPromise: Promise | undefined; + let hotResumeCount = nextResumeCount; + let residentRegistered = false; + + const runtimeCleanup = () => { + if (runtimeDisposed) return; + runtimeDisposed = true; + registry.unregisterResidentAgent(meta.agentId, residentController); + residentRegistered = false; + bgEmitter.off(AgentEventType.TOOL_CALL, onToolCall); + bgEmitter.off(AgentEventType.USAGE_METADATA, onUsageMetadata); + cleanupApprovalBridge?.(); + cleanupOwnedMonitorNotifications?.(); + cleanupJsonl?.(); + void activeAgentConfig + .getToolRegistry() + .stop() + .catch(() => {}); + void subagentDispose?.().catch(() => {}); + }; + cleanupRuntime = runtimeCleanup; + + const requestRuntimeDisposal = () => { + if (disposeRequested || runtimeDisposed) return; + disposeRequested = true; + registry.unregisterResidentAgent(meta.agentId, residentController); + residentRegistered = false; + currentAbortController?.abort(); + if (!turnRunning) { + runtimeCleanup(); + } + }; - 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', + const runBody = async ( + turnContextState: ContextState, + turnAbortController: AbortController, + fireStartHook: boolean, + ) => { + let keepResident = false; + let finishingInputs: AgentExternalInput[] | undefined; + let shouldFireStartHook = fireStartHook; + turnRunning = true; + try { + while (true) { + if (shouldFireStartHook) { + await this.applySubagentStartHook(turnContextState, { + agentId: meta.agentId, + agentType: meta.agentType, + resolvedMode, + signal: turnAbortController.signal, + }); + const additionalContext = turnContextState.get('hook_context'); + if (additionalContext) { + turnContextState.set( + 'task_prompt', + `${String(turnContextState.get('task_prompt'))}\n\n${String(additionalContext)}`, + ); + } + } + shouldFireStartHook = false; + + if (finishingInputs) { + await subagent.executeExternalInputs( + finishingInputs, + turnAbortController.signal, + { resetStats: false }, + ); + finishingInputs = undefined; + } else { + await subagent.execute( + turnContextState, + turnAbortController.signal, + ); + } + + 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(); + const modelVisibleText = toModelVisibleSubagentResult( + subagent.getFinalText(), + terminateMode, ); - } 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, - }); + const finalText = appendStopHookBlockingCapWarning( + terminateMode === AgentTerminateMode.GOAL + ? modelVisibleText || + '(subagent produced no model-visible output)' + : modelVisibleText, + stopHookWarning, + ); + const stats = getCompletionStats(subagent, liveToolCallCount); + if (terminateMode === AgentTerminateMode.GOAL) { + const pending = registry.drainMessages(meta.agentId); + if (pending.length > 0) { + finishingInputs = pending; + continue; + } + + keepResident = residentRegistered && !needsAutoPermissionLease(); + if (!keepResident) { + registry.unregisterResidentAgent( + meta.agentId, + residentController, + ); + residentRegistered = false; + } + patchAgentMeta(metaPath, { + status: 'completed', + lastUpdatedAt: new Date().toISOString(), + lastError: undefined, + }); + registry.complete(meta.agentId, finalText, stats); + } 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 +1130,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,49 +1154,120 @@ 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. - restoreParentPM(); + turnRunning = false; + activeRestoreParentPM(); + if (!keepResident || disposeRequested) { + 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 (needsAutoPermissionLease()) { + requestRuntimeDisposal(); + 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, + ); + }); + currentTurnPromise.catch(reportUnexpectedBackgroundError); + return true; + }, + dispose: requestRuntimeDisposal, + }; + if (canStayResident && !needsAutoPermissionLease()) { + registry.registerResidentAgent(meta.agentId, residentController); + residentRegistered = true; + } + + currentTurnPromise = runBackgroundTurn( + contextState, + bgAbortController, + false, + ); + runtimeLifecycleOwned = true; + currentTurnPromise.catch(reportUnexpectedBackgroundError); return entry; } catch (error) { - cleanupOwnedMonitorNotifications?.(); - cleanupJsonl?.(); + cleanupPreparedRuntime(); const errorMessage = error instanceof Error ? error.message : String(error); debugLogger.warn( @@ -1278,7 +1500,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..53ed6ef4fd4 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,235 @@ 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('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'); + expect(cancelledResident.dispose).toHaveBeenCalledOnce(); + 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 @@ -1570,6 +1800,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 @@ -1729,6 +1977,30 @@ describe('BackgroundTaskRegistry', () => { it('returns empty array for non-existent agent', () => { expect(registry.drainMessages('nope')).toEqual([]); }); + + it('rejects new messages after finalization begins and releases waiters on completion', async () => { + registry.register({ + agentId: 'test-1', + description: 'test agent', + status: 'running', + startTime: Date.now(), + abortController: new AbortController(), + isBackgrounded: true, + outputFile: '/tmp/test.jsonl', + }); + + expect(registry.beginFinishing('test-1')).toBe(true); + expect(registry.queueMessage('test-1', 'late correction')).toBe(false); + + const settled = registry.waitForFinishing( + 'test-1', + new AbortController().signal, + ); + registry.complete('test-1', 'done'); + + await expect(settled).resolves.toBe(true); + expect(registry.get('test-1')!.pendingMessages).toEqual([]); + }); }); describe('waitForMessages', () => { diff --git a/packages/core/src/agents/background-tasks.ts b/packages/core/src/agents/background-tasks.ts index b3618e7b441..1b2cc371809 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; +} + type MessageWaiter = () => void; export interface BackgroundTaskRegistryOptions { @@ -477,7 +487,13 @@ const BACKGROUND_SLOT_WAIT_CANCELLED = export class BackgroundTaskRegistry { private readonly agents = new Map(); + private readonly residentAgents = new Map(); private readonly messageWaiters = new Map>(); + private readonly finishingAgents = new Set(); + private readonly finishingWaiters = new Map< + string, + Set<(settled: boolean) => void> + >(); private readonly waitQueue: BackgroundSlotWaiter[] = []; // Maps each outstanding slot reservation to the concrete model ID it was // reserved for (undefined when unresolved). A Map rather than a Set so the @@ -639,6 +655,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) @@ -663,6 +682,7 @@ export class BackgroundTaskRegistry { entry.parentName = this.agents.get(entry.parentAgentId)?.subagentType; } this.agents.set(entry.agentId, entry); + this.releaseFinishingWaiters(entry.agentId, true); debugLogger.info(`Registered background agent: ${entry.agentId}`); if ( wasRunningBackground && @@ -692,6 +712,87 @@ 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); + } + + 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 { + current.dispose(); + } catch (error) { + debugLogger.error( + `Failed to dispose resident background agent ${agentId}:`, + error, + ); + } + return true; + } + + disposeResidentAgents(): void { + this.disposeAllResidentAgents(); + } + // 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,13 +811,18 @@ 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; entry.stats = stats; + this.releaseFinishingWaiters(agentId, true); debugLogger.info(`Background agent completed: ${agentId}`); this.rejectPendingApprovals(entry); + if (wasCancelled) { + this.disposeResidentAgent(agentId); + } this.emitNotification(entry); this.emitStatusChange(entry); this.drainWaitQueue(); @@ -749,7 +855,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(); @@ -766,11 +872,13 @@ export class BackgroundTaskRegistry { entry.endTime = Date.now(); entry.error = error; entry.stats = stats; + this.releaseFinishingWaiters(agentId, true); debugLogger.info(`Background agent failed: ${agentId}`); this.rejectPendingApprovals(entry); this.emitNotification(entry); this.emitStatusChange(entry); + this.disposeResidentAgent(agentId); this.drainWaitQueue(); } @@ -807,6 +915,8 @@ export class BackgroundTaskRegistry { entry.status = 'cancelled'; entry.endTime = Date.now(); entry.persistedCancellationStatus = persistedStatus; + this.releaseFinishingWaiters(agentId, true); + this.disposeResidentAgent(agentId); if (entry.metaPath) { patchAgentMeta(entry.metaPath, { status: persistedStatus, @@ -849,9 +959,11 @@ export class BackgroundTaskRegistry { entry.status = 'cancelled'; entry.endTime = Date.now(); entry.notified = true; + this.releaseFinishingWaiters(agentId, true); debugLogger.info(`Abandoned paused background agent: ${agentId}`); this.rejectPendingApprovals(entry); this.emitStatusChange(entry); + this.disposeResidentAgent(agentId); this.drainWaitQueue(); } @@ -874,9 +986,11 @@ export class BackgroundTaskRegistry { entry.endTime ??= Date.now(); if (partialResult) entry.result = partialResult; entry.stats = stats; + this.releaseFinishingWaiters(agentId, true); this.rejectPendingApprovals(entry); this.emitNotification(entry); this.emitStatusChange(entry); + this.disposeResidentAgent(agentId); this.drainWaitQueue(); } @@ -896,6 +1010,7 @@ export class BackgroundTaskRegistry { this.rejectPendingApprovals(entry); this.emitNotification(entry); this.emitStatusChange(entry); + this.disposeResidentAgent(agentId); this.drainWaitQueue(); } @@ -1215,6 +1330,8 @@ export class BackgroundTaskRegistry { | AgentTask | undefined; if (!firstEntry) { + this.releaseAllFinishingWaiters(false); + this.disposeAllResidentAgents(); this.rejectWaitQueue(); return; } @@ -1228,7 +1345,9 @@ export class BackgroundTaskRegistry { this.wakeMessageWaiters(entry.agentId); } this.rejectWaitQueue(); + this.releaseAllFinishingWaiters(false); this.agents.clear(); + this.disposeAllResidentAgents(); this.emitStatusChange(firstEntry); } @@ -1247,7 +1366,13 @@ export class BackgroundTaskRegistry { */ queueExternalInput(agentId: string, input: AgentExternalInput): boolean { const entry = this.agents.get(agentId); - if (!entry || entry.status !== 'running') return false; + if ( + !entry || + entry.status !== 'running' || + this.finishingAgents.has(agentId) + ) { + return false; + } const queue = entry.pendingMessages!; queue.push(input); debugLogger.info( @@ -1257,6 +1382,40 @@ export class BackgroundTaskRegistry { return true; } + /** Close the input queue after its final drain but before async teardown. */ + beginFinishing(agentId: string): boolean { + const entry = this.agents.get(agentId); + if (!entry || entry.status !== 'running') return false; + this.finishingAgents.add(agentId); + return true; + } + + isFinishing(agentId: string): boolean { + return this.finishingAgents.has(agentId); + } + + /** Wait until a finishing task publishes its terminal state. */ + waitForFinishing(agentId: string, signal: AbortSignal): Promise { + if (!this.finishingAgents.has(agentId)) return Promise.resolve(true); + if (signal.aborted) return Promise.resolve(false); + + return new Promise((resolve) => { + const settle = (settled: boolean) => { + signal.removeEventListener('abort', onAbort); + const waiters = this.finishingWaiters.get(agentId); + waiters?.delete(settle); + if (waiters?.size === 0) this.finishingWaiters.delete(agentId); + resolve(settled); + }; + const onAbort = () => settle(false); + const waiters = this.finishingWaiters.get(agentId) ?? new Set(); + waiters.add(settle); + this.finishingWaiters.set(agentId, waiters); + signal.addEventListener('abort', onAbort, { once: true }); + if (signal.aborted) onAbort(); + }); + } + /** * Drain all pending messages for an agent. Returns the messages * and clears the queue. Called by the agent's reasoning loop. @@ -1367,6 +1526,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 +1666,37 @@ 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.releaseFinishingWaiters(agentId, false); + this.disposeResidentAgent(agentId); + return this.agents.delete(agentId); + } + + private releaseFinishingWaiters(agentId: string, settled: boolean): void { + this.finishingAgents.delete(agentId); + const waiters = this.finishingWaiters.get(agentId); + if (!waiters) return; + this.finishingWaiters.delete(agentId); + for (const resolve of waiters) resolve(settled); + } + + private releaseAllFinishingWaiters(settled: boolean): void { + for (const agentId of Array.from(this.finishingAgents)) { + this.releaseFinishingWaiters(agentId, settled); + } + } + + 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..98df340d5ae 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; @@ -805,7 +808,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 +1004,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) { @@ -1835,6 +1840,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..36440f33104 100644 --- a/packages/core/src/agents/runtime/agent-headless.test.ts +++ b/packages/core/src/agents/runtime/agent-headless.test.ts @@ -543,6 +543,233 @@ 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 continue with atomically claimed finishing inputs', 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, + ); + const externalEvents: Array<{ + kind: string | undefined; + text: string; + }> = []; + scope.getEventEmitter().on(AgentEventType.EXTERNAL_MESSAGE, (event) => { + externalEvents.push({ kind: event.kind, text: event.text }); + }); + + const initialContext = new ContextState(); + initialContext.set('task_prompt', 'Initial task'); + await scope.execute(initialContext); + await scope.executeExternalInputs( + ['late correction', { kind: 'notification', text: 'monitor fired' }], + undefined, + { resetStats: false }, + ); + + expect(mockSendMessageStream.mock.calls[1][1].message).toEqual([ + { text: '[Message from parent agent]: late correction' }, + { text: 'monitor fired' }, + ]); + expect(externalEvents).toEqual([ + { kind: 'message', text: 'late correction' }, + { kind: 'notification', text: 'monitor fired' }, + ]); + expect(scope.getExecutionSummary()).toMatchObject({ rounds: 2 }); + }); + + 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 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..2870d4bc7be 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,54 @@ export class AgentHeadless { async execute( context: ContextState, externalSignal?: AbortSignal, + options: { resetStats?: boolean } = {}, + ): 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); + } finally { + this.executing = false; + } + } + + async executeExternalInputs( + inputs: AgentExternalInput[], + externalSignal?: AbortSignal, + options: { resetStats?: boolean } = {}, + ): Promise { + if (inputs.length === 0) return; + const context = new ContextState(); + context.set('external_inputs_override', inputs); + await this.execute(context, externalSignal, options); + } + + private async executeTurn( + context: ContextState, + externalSignal?: AbortSignal, + preserveStats = false, ): Promise { const initialMessagesOverride = context.get('initial_messages_override') as | Content[] | undefined; + const isContinuation = this.hasStartedReasoning; + const externalInputsOverride = isContinuation + ? (context.get('external_inputs_override') as + | AgentExternalInput[] + | undefined) + : undefined; // 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 +262,28 @@ export class AgentHeadless { const initialTaskText = String( (context.get('task_prompt') as string) ?? 'Get Started!', ); - if (!initialMessagesOverride || initialMessagesOverride.length === 0) { + if (isContinuation) { + const transcriptInputs = externalInputsOverride ?? [initialTaskText]; + for (const input of transcriptInputs) { + 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 ( + !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 +295,45 @@ export class AgentHeadless { const abortController = createChildAbortController(externalSignal); try { - const toolsList = await this.core.prepareTools(); - - const initialMessages = - 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); + if (!this.toolsList) { + this.toolsList = await this.core.prepareTools(); + } + const toolsList = this.toolsList; + + const initialMessages = externalInputsOverride + ? [ + { + role: 'user' as const, + parts: externalInputsOverride.map((input) => ({ + text: + typeof input === 'string' + ? `${EXTERNAL_MESSAGE_PREFIX} ${input}` + : input.text, + })), + }, + ] + : isContinuation + ? [ + { + role: 'user' as const, + parts: [ + { text: `${EXTERNAL_MESSAGE_PREFIX} ${initialTaskText}` }, + ], + }, + ] + : initialMessagesOverride && initialMessagesOverride.length > 0 + ? initialMessagesOverride + : [{ role: 'user' as const, parts: [{ text: initialTaskText }] }]; + + 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 +358,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 +368,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 +388,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/config/config.test.ts b/packages/core/src/config/config.test.ts index ca4212f11f7..749af6b3888 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -4191,6 +4191,10 @@ describe('Server Config (config.ts)', () => { it('relocateWorkingDirectory should update the session working roots', async () => { const config = new Config(baseParams); + const disposeResidentAgents = vi.spyOn( + config.getBackgroundTaskRegistry(), + 'disposeResidentAgents', + ); const newDir = path.resolve('/path/to/other'); const workspaceContext = config.getWorkspaceContext(); const directoriesChanged = vi.fn(); @@ -4210,6 +4214,7 @@ describe('Server Config (config.ts)', () => { expect(config.getWorkspaceContext()).toBe(workspaceContext); expect(config.getWorkspaceContext().getDirectories()[0]).toBe(newDir); expect(config.storage.getProjectRoot()).toBe(newDir); + expect(disposeResidentAgents).toHaveBeenCalledOnce(); expect(directoriesChanged).toHaveBeenCalled(); expect(loadServerHierarchicalMemory).toHaveBeenCalledWith( newDir, @@ -4413,6 +4418,10 @@ describe('Server Config (config.ts)', () => { it('relocateWorkingDirectory should reject and roll back when session artifact migration fails', async () => { const config = new Config({ ...baseParams, chatRecording: true }); + const disposeResidentAgents = vi.spyOn( + config.getBackgroundTaskRegistry(), + 'disposeResidentAgents', + ); const oldDir = config.getTargetDir(); const sessionId = config.getSessionId(); const newDir = path.resolve('/path/to/other'); @@ -4466,6 +4475,7 @@ describe('Server Config (config.ts)', () => { expect(config.getTargetDir()).toBe(oldDir); expect(config.storage.getProjectRoot()).toBe(oldDir); expect(config.getTranscriptPath()).toBe(oldTranscriptPath); + expect(disposeResidentAgents).not.toHaveBeenCalled(); chdirSpy.mockRestore(); cwdSpy.mockRestore(); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index e7b25278527..8485bd748b2 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -4240,6 +4240,7 @@ export class Config { this.chatRecordingService?.resetStoragePaths(); } + this.backgroundTaskRegistry.disposeResidentAgents(); this.targetDir = expected; this.cwd = expected; await this.refreshCurrentRuntimeStatus(expected); diff --git a/packages/core/src/subagents/subagent-manager.test.ts b/packages/core/src/subagents/subagent-manager.test.ts index 32d700f0e59..03b79b961ea 100644 --- a/packages/core/src/subagents/subagent-manager.test.ts +++ b/packages/core/src/subagents/subagent-manager.test.ts @@ -2120,6 +2120,27 @@ bad`); expect(mockCreateContentGenerator).not.toHaveBeenCalled(); }); + it('should snapshot the launch provider when inherit receives a concrete model override', async () => { + const config = { ...agentConfig, model: 'inherit' }; + + await manager.createAgentHeadless(config, mockConfig, { + modelConfigOverrides: { model: 'launch-model' }, + runtimeAuthOverrides: { + authType: AuthType.USE_ANTHROPIC, + baseUrl: 'https://launch-provider.example.com', + }, + }); + + expect(mockCreateContentGenerator).toHaveBeenCalledWith( + expect.objectContaining({ + model: 'launch-model', + authType: AuthType.USE_ANTHROPIC, + baseUrl: 'https://launch-provider.example.com', + }), + mockConfig, + ); + }); + 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..76133d4bb7d 100644 --- a/packages/core/src/subagents/subagent-manager.ts +++ b/packages/core/src/subagents/subagent-manager.ts @@ -40,7 +40,10 @@ import type { Config, MCPServerConfig } from '../config/config.js'; import { APPROVAL_MODES } from '../config/config.js'; import type { HookDefinition, HookEventName } from '../hooks/types.js'; import type { RuntimeContentGeneratorView } from '../agents/runtime/agent-context.js'; -import { createRuntimeContentGeneratorView } from '../models/content-generator-config.js'; +import { + createRuntimeContentGeneratorView, + type AuthOverrides, +} from '../models/content-generator-config.js'; import { createDebugLogger } from '../utils/debugLogger.js'; import { normalizeContent } from '../utils/textUtils.js'; import { @@ -760,6 +763,7 @@ export class SubagentManager { hooks?: AgentHooks; promptConfigOverrides?: Partial; modelConfigOverrides?: Partial; + runtimeAuthOverrides?: AuthOverrides; runConfigOverrides?: Partial; toolConfigOverride?: ToolConfig; }, @@ -822,6 +826,8 @@ export class SubagentManager { const runtimeView = await this.buildRuntimeContentGeneratorView( config, runtimeContext, + modelConfig.model, + options?.runtimeAuthOverrides, ); const { context: subagentContext, cleanup } = @@ -1018,22 +1024,30 @@ export class SubagentManager { private async buildRuntimeContentGeneratorView( config: SubagentConfig, base: Config, + fallbackModelId?: string, + runtimeAuthOverrides?: AuthOverrides, ): Promise { const resolvedModel = this.resolveModelOverride(config.model, base); - if (!resolvedModel) { + const modelId = resolvedModel?.modelId ?? fallbackModelId; + if (!modelId) { return undefined; } const authType = - resolvedModel.authType ?? base.getContentGeneratorConfig().authType; - const authOverrides = { - authType: authType as string, - }; + resolvedModel?.authType ?? + runtimeAuthOverrides?.authType ?? + base.getContentGeneratorConfig().authType; + const authOverrides: AuthOverrides = resolvedModel + ? { authType: authType as string } + : { + ...runtimeAuthOverrides, + authType: authType as string, + }; const view = await createRuntimeContentGeneratorView( base, base, - resolvedModel.modelId, + modelId, authOverrides, ); diff --git a/packages/core/src/tools/agent/agent.test.ts b/packages/core/src/tools/agent/agent.test.ts index 05fac30211e..c0cc480dcac 100644 --- a/packages/core/src/tools/agent/agent.test.ts +++ b/packages/core/src/tools/agent/agent.test.ts @@ -37,6 +37,7 @@ import type { AgentEventEmitter, } from '../../agents/runtime/agent-events.js'; import { partToString } from '../../utils/partUtils.js'; +import { AuthType } from '../../core/contentGenerator.js'; import type { HookSystem } from '../../hooks/hookSystem.js'; import { PermissionMode } from '../../hooks/types.js'; import { runWithAgentContext } from '../../agents/runtime/agent-context.js'; @@ -151,6 +152,7 @@ describe('AgentTool', () => { get: vi.fn(), getAll: vi.fn().mockReturnValue([]), drainMessages: vi.fn().mockReturnValue([]), + beginFinishing: vi.fn().mockReturnValue(true), queueMessage: vi.fn(), queueExternalInput: vi.fn(), wakeExternalInputWaiters: vi.fn(), @@ -190,6 +192,10 @@ describe('AgentTool', () => { isAgentTeamEnabled: vi.fn().mockReturnValue(false), getApprovalMode: vi.fn().mockReturnValue('default'), getModel: vi.fn().mockReturnValue('parent-model'), + getContentGeneratorConfig: vi.fn().mockReturnValue({ + model: 'parent-model', + authType: 'openai', + }), getBareMode: vi.fn().mockReturnValue(false), isSafeMode: vi.fn().mockReturnValue(false), getSandbox: vi.fn().mockReturnValue(undefined), @@ -378,7 +384,10 @@ describe('AgentTool', () => { 'paused agents resume with it as their first continuation instruction', ); expect(tool.description).toContain( - 'completed agents are revived from their retained transcript', + 'completed agents continue on their resident runtime when available', + ); + expect(tool.description).toContain( + 'otherwise revive from their retained transcript', ); expect(tool.description).toContain('return to their direct parent'); expect(tool.description).not.toContain('Top-level one-shot agents'); @@ -4254,6 +4263,7 @@ describe('AgentTool', () => { describe('Agent-level background: true', () => { let mockAgent: AgentHeadless; let mockContextState: ContextState; + let mockSubagentDispose: ReturnType; let mockRegistry: { assertCanStartBackgroundAgent: ReturnType; canStartBackgroundAgent: ReturnType; @@ -4261,16 +4271,21 @@ describe('AgentTool', () => { waitForBackgroundSlot: ReturnType; releaseBackgroundSlot: ReturnType; getQueuedCount: ReturnType; + get: ReturnType; register: ReturnType; unregisterForeground: ReturnType; complete: ReturnType; fail: ReturnType; finalizeCancelled: ReturnType; drainMessages: ReturnType; + beginFinishing: ReturnType; waitForMessages: ReturnType; queueExternalInput: ReturnType; wakeExternalInputWaiters: ReturnType; appendActivity: ReturnType; + registerResidentAgent: ReturnType; + unregisterResidentAgent: ReturnType; + restartCompletedAgent: ReturnType; }; const bgSubagent: SubagentConfig = { @@ -4285,6 +4300,7 @@ describe('AgentTool', () => { beforeEach(() => { mockAgent = { execute: vi.fn().mockResolvedValue(undefined), + executeExternalInputs: vi.fn().mockResolvedValue(undefined), getFinalText: vi.fn().mockReturnValue('Monitor done'), getTerminateMode: vi.fn().mockReturnValue(AgentTerminateMode.GOAL), getExecutionSummary: vi.fn().mockReturnValue({}), @@ -4301,6 +4317,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,16 +4329,21 @@ 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(), fail: vi.fn(), finalizeCancelled: vi.fn(), drainMessages: vi.fn().mockReturnValue([]), + beginFinishing: vi.fn().mockReturnValue(true), waitForMessages: vi.fn().mockResolvedValue([]), queueExternalInput: vi.fn(), wakeExternalInputWaiters: vi.fn(), appendActivity: vi.fn(), + registerResidentAgent: vi.fn(), + unregisterResidentAgent: vi.fn().mockReturnValue(true), + restartCompletedAgent: vi.fn().mockReturnValue(restartedEntry), }; vi.mocked(config.getApprovalMode).mockReturnValue(ApprovalMode.DEFAULT); @@ -4345,9 +4367,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, }); }); @@ -4369,6 +4392,7 @@ describe('AgentTool', () => { expect(llmText).toContain( `Use ${ToolNames.SEND_MESSAGE} to continue this agent`, ); + expect(llmText).toContain('agentId: monitor-'); expect(llmText).toContain(`or ${ToolNames.TASK_STOP} to cancel.`); expect(llmText).not.toContain('with to:'); expect(llmText).not.toContain('Use send_message with task_id:'); @@ -4405,9 +4429,48 @@ describe('AgentTool', () => { expect.objectContaining({ persistedCliFlags: expect.objectContaining({ model: 'subagent-model', + authType: 'openai', }), }), ); + expect(mockSubagentManager.createAgentHeadless).toHaveBeenCalledTimes(1); + writeMetaSpy.mockRestore(); + }); + + it('does not persist the parent base URL for a cross-provider runtime', async () => { + const writeMetaSpy = vi.spyOn(transcript, 'writeAgentMeta'); + vi.mocked(config.getContentGeneratorConfig).mockReturnValue({ + model: 'parent-model', + authType: AuthType.USE_OPENAI, + baseUrl: 'https://parent-provider.example.com', + }); + vi.mocked(mockAgent.getCore).mockReturnValue({ + modelConfig: { model: 'subagent-model' }, + runtimeView: { + contentGenerator: {}, + contentGeneratorConfig: { + model: 'subagent-model', + authType: AuthType.USE_ANTHROPIC, + }, + }, + getEventEmitter: () => ({ on: vi.fn(), off: vi.fn() }), + } as never); + + const invocation = ( + agentTool as AgentToolWithProtectedMethods + ).createInvocation({ + description: 'Start monitor', + prompt: 'Watch for changes', + subagent_type: 'monitor', + }); + await invocation.execute(); + + const persistedFlags = writeMetaSpy.mock.calls[0]?.[1].persistedCliFlags; + expect(persistedFlags).toMatchObject({ + model: 'subagent-model', + authType: 'anthropic', + }); + expect(persistedFlags).toHaveProperty('baseUrl', undefined); writeMetaSpy.mockRestore(); }); @@ -4502,7 +4565,7 @@ 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 () => { const params: AgentParams = { description: 'Start monitor', prompt: 'Watch for changes', @@ -4521,6 +4584,20 @@ 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 } + | undefined; + expect(resident).toBeDefined(); + resident?.dispose(); + await vi.waitFor(() => { expect( monitorRegistry.setAgentNotificationCallback, @@ -4534,6 +4611,188 @@ describe('AgentTool', () => { { notify: false }, ); }); + expect(mockSubagentDispose).toHaveBeenCalledOnce(); + }); + + 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); + + 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(mockSubagentManager.createAgentHeadless).toHaveBeenCalledWith( + expect.any(Object), + expect.any(Object), + expect.objectContaining({ + modelConfigOverrides: { model: 'parent-model' }, + runtimeAuthOverrides: expect.objectContaining({ + authType: 'openai', + }), + }), + ); + expect(mockSubagentDispose).not.toHaveBeenCalled(); + }); + + it('claims finishing-window input before publishing completion', async () => { + mockRegistry.drainMessages + .mockReturnValueOnce(['late correction']) + .mockReturnValue([]); + 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).toHaveBeenCalledOnce(); + }); + + expect(mockAgent.execute).toHaveBeenCalledOnce(); + expect(mockAgent.executeExternalInputs).toHaveBeenCalledWith( + ['late correction'], + expect.any(AbortSignal), + { resetStats: false }, + ); + expect( + vi.mocked(mockAgent.executeExternalInputs).mock.invocationCallOrder[0], + ).toBeLessThan(mockRegistry.complete.mock.invocationCallOrder[0]!); + }); + + it('persists completion before publishing the terminal notification', async () => { + const patchMetaSpy = vi.spyOn(transcript, 'patchAgentMeta'); + 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).toHaveBeenCalled(); + }); + + const completedPatchIndex = patchMetaSpy.mock.calls.findIndex( + ([, update]) => update.status === 'completed', + ); + expect(completedPatchIndex).toBeGreaterThanOrEqual(0); + expect( + patchMetaSpy.mock.invocationCallOrder[completedPatchIndex], + ).toBeLessThan(mockRegistry.complete.mock.invocationCallOrder[0]!); + patchMetaSpy.mockRestore(); + }); + + 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('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..dd6a086b122 100644 --- a/packages/core/src/tools/agent/agent.ts +++ b/packages/core/src/tools/agent/agent.ts @@ -114,9 +114,13 @@ 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'; +import type { AuthOverrides } from '../../models/content-generator-config.js'; // Memoize git branch per cwd for the agent-launch path. `getGitBranch` // shells out to `git rev-parse` synchronously; caching avoids the per-launch @@ -602,7 +606,9 @@ function capturePersistedCliFlags( config: Config, resolvedApprovalMode: ApprovalMode, modelOverride?: string, + runtimeAuthOverrides?: { authType?: string; baseUrl?: string }, ): AgentPersistedCliFlags { + const contentGeneratorConfig = config.getContentGeneratorConfig(); return { approvalMode: resolvedApprovalMode, bare: config.getBareMode(), @@ -610,6 +616,10 @@ function capturePersistedCliFlags( sandbox: config.getSandbox() ?? null, screenReader: config.getScreenReader(), model: modelOverride ?? config.getModel(), + authType: runtimeAuthOverrides?.authType ?? contentGeneratorConfig.authType, + baseUrl: runtimeAuthOverrides + ? runtimeAuthOverrides.baseUrl + : contentGeneratorConfig.baseUrl, maxSessionTurns: config.getMaxSessionTurns(), maxToolCalls: config.getMaxToolCalls(), maxSubagentDepth: config.getMaxSubagentDepth(), @@ -870,7 +880,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. @@ -899,7 +908,7 @@ Usage notes: - Run agents concurrently only when their tasks are independent. For code changes, give concurrent agents disjoint write scopes; launch them in a single message with multiple tool uses. - A background agent reports its result through a completion notification in a later turn. A foreground agent returns its result inline. Agent results are not visible to the user, so relay the relevant outcome in your response. - While background agents run, continue meaningful non-overlapping work. Wait for an agent only when its result blocks the next required step. -- Reuse an existing background agent for related follow-up work instead of launching a duplicate: call ${ToolNames.SEND_MESSAGE} with the \`agentId\` from its launch result as its \`task_id\`. Running agents receive the message at the next tool-round boundary; paused agents resume with it as their first continuation instruction; completed agents are revived from their retained transcript. If the task is no longer retained or cannot be resumed or revived, launch a new agent. +- Reuse an existing background agent for related follow-up work instead of launching a duplicate: call ${ToolNames.SEND_MESSAGE} with the \`agentId\` from its launch result as its \`task_id\`. Running agents receive the message at the next tool-round boundary; paused agents resume with it as their first continuation instruction; completed agents continue on their resident runtime when available and otherwise revive from their retained transcript. If the task is no longer retained or cannot be resumed or revived, launch a new agent. - Provide clear, detailed prompts so the agent can work autonomously and return exactly the information you need. - Regular subagents and named teammates start without parent conversation history. Only fork agents accept \`fork_turns\`; omit it for the full conversation or use a positive integer string such as \`"3"\` for a bounded recent window. - Treat the agent's output as evidence, not as automatically correct. Verify factual claims, review code changes, and run relevant checks before integrating or relaying the result. @@ -1741,7 +1750,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) { @@ -2334,6 +2345,7 @@ class AgentToolInvocation extends BaseToolInvocation { // selector once subagentConfig is loaded. Used to enforce per-model // background-agent concurrency caps (agents.maxParallelAgentsByModel). let subagentModelId: string | undefined; + let subagentRuntimeAuthOverrides: AuthOverrides | undefined; const releaseBackgroundSlotReservation = () => { if (backgroundSlotReservation && !backgroundSlotReservationConsumed) { this.config @@ -2487,10 +2499,25 @@ class AgentToolInvocation extends BaseToolInvocation { // resolveModelId maps it to the actual model ID, falling back to the // parent's current model when the sub-agent inherits (forks always // inherit, since FORK_AGENT has no model selector). - subagentModelId = resolveModelId( + const resolvedSubagentModel = resolveModelId( subagentConfig.model, buildModelIdContext(this.config), - )?.modelId; + ); + subagentModelId = resolvedSubagentModel?.modelId; + subagentModelId ??= this.config.getModel(); + const parentContentGeneratorConfig = + this.config.getContentGeneratorConfig(); + const authType = + resolvedSubagentModel?.authType ?? + parentContentGeneratorConfig.authType; + subagentRuntimeAuthOverrides = authType + ? { + authType, + ...(authType === parentContentGeneratorConfig.authType + ? { baseUrl: parentContentGeneratorConfig.baseUrl } + : {}), + } + : undefined; const registry = this.config.getBackgroundTaskRegistry(); backgroundSlotReservation = registry.tryReserveBackgroundSlot(subagentModelId); @@ -2749,6 +2776,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 +2789,38 @@ 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; + } + + // 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 +2831,28 @@ 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 } } + : {}), + ...(shouldRunInBackground && subagentRuntimeAuthOverrides + ? { runtimeAuthOverrides: subagentRuntimeAuthOverrides } + : {}), + }, ); subagent = result.subagent; subagentDispose = result.dispose; @@ -2862,54 +2932,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 +2985,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 +3039,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!, @@ -3047,19 +3081,22 @@ class AgentToolInvocation extends BaseToolInvocation { parentAgentId: getCurrentAgentId(), createdAt: new Date().toISOString(), status: 'running', + isolation: this.params.isolation, lastUpdatedAt: new Date().toISOString(), resolvedApprovalMode, persistedCliFlags: capturePersistedCliFlags( this.config, resolvedApprovalMode, bgSubagent.getCore().modelConfig.model, + bgSubagent.getCore().runtimeView?.contentGeneratorConfig ?? + subagentRuntimeAuthOverrides, ), 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 +3177,270 @@ 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; + let runtimeDisposed = false; + let disposeRequested = false; + let turnRunning = false; + let currentAbortController: AbortController | undefined = + bgAbortController; + let currentTurnPromise: Promise | undefined; + let hotContinuationCount = 0; + let residentRegistered = false; + + const cleanupRuntime = () => { + if (runtimeDisposed) return; + runtimeDisposed = true; + registry.unregisterResidentAgent( + hookOpts.agentId, + residentController, + ); + residentRegistered = false; + bgEmitter.off(AgentEventType.TOOL_CALL, onToolCall); + bgEmitter.off(AgentEventType.USAGE_METADATA, onUsageMetadata); + cleanupApprovalBridge?.(); + cleanupOwnedMonitorNotifications(); + cleanupJsonl?.(); + void agentConfig + .getToolRegistry() + .stop() + .catch(() => {}); + void bgSubagentDispose?.().catch(() => {}); + }; + + const requestRuntimeDisposal = () => { + if (disposeRequested || runtimeDisposed) return; + disposeRequested = true; + registry.unregisterResidentAgent( + hookOpts.agentId, + residentController, + ); + residentRegistered = false; + currentAbortController?.abort(); + if (!turnRunning) { + cleanupRuntime(); + } + }; + // 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; + let finishingInputs: AgentExternalInput[] | undefined; + let shouldFireStartHook = fireStartHook; + 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, - }); - } + while (true) { + if (shouldFireStartHook && 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}`, + ); + } + } + shouldFireStartHook = false; + + if (finishingInputs) { + await bgSubagent.executeExternalInputs( + finishingInputs, + turnAbortController.signal, + { resetStats: false }, + ); + finishingInputs = undefined; + } else { + await bgSubagent.execute( + turnContextState, + 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. - // - // 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({ - terminateMode, - signalAborted: bgAbortController.signal.aborted, - resultSummaryPresent: Boolean( - subagentRawText && subagentRawText.length > 0, - ), - }), - ); + 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, + }, + ); + } - 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', + // 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 hadWorktreeIsolation = worktreeIsolation !== null; + const recordTerminalOutcome = () => + recordSpanOutcome( + deriveSubagentOutcomeMetadata({ + terminateMode, + signalAborted: turnAbortController.signal.aborted, + resultSummaryPresent: Boolean( + subagentRawText && subagentRawText.length > 0, + ), + }), + ); + if ( + terminateMode === AgentTerminateMode.GOAL && + hadWorktreeIsolation + ) { + const pending = registry.drainMessages(hookOpts.agentId); + if (pending.length > 0) { + finishingInputs = pending; + continue; + } + registry.beginFinishing(hookOpts.agentId); + } + if (hadWorktreeIsolation) { + recordTerminalOutcome(); + } + + const wtSuffix = formatWorktreeSuffix( + hadWorktreeIsolation ? await cleanupWorktreeIsolation() : {}, ); - } else { - registry.fail( - hookOpts.agentId, - finalText || `Agent terminated with mode: ${terminateMode}`, - completionStats, + 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 && + !hadWorktreeIsolation + ) { + const pending = registry.drainMessages(hookOpts.agentId); + if (pending.length > 0) { + finishingInputs = pending; + continue; + } + // Mirror the worktree path: close the input queue before + // publishing completion so a send_message racing the terminal + // transition is rejected (queueExternalInput checks + // finishingAgents) rather than accepted and silently orphaned. + registry.beginFinishing(hookOpts.agentId); + } + + if (!hadWorktreeIsolation) { + recordTerminalOutcome(); + } + + if (terminateMode === AgentTerminateMode.GOAL) { + keepResident = + residentRegistered && !needsAutoPermissionLease(); + if (!keepResident) { + registry.unregisterResidentAgent( + hookOpts.agentId, + residentController, + ); + residentRegistered = false; + } + patchAgentMeta(metaPath, { + status: 'completed', + lastUpdatedAt: new Date().toISOString(), + lastError: undefined, + }); + registry.complete(hookOpts.agentId, finalText, completionStats); + } 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) { + // A resident runtime is only safe to keep for a cleanly completed + // agent. If completion bookkeeping (patchAgentMeta / + // registry.complete) threw after keepResident was set, the entry is + // finalized as failed/cancelled below and can never be continued — + // so release keepResident here to let the finally block dispose the + // runtime instead of leaking a zombie resident. + keepResident = false; // Publish first — same reason as the success path. recordSpanOutcome( deriveSubagentExceptionMetadata( error, - bgAbortController.signal.aborted, + turnAbortController.signal.aborted, ), ); const baseErrorMsg = @@ -3270,7 +3467,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 +3487,127 @@ 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) { + 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 (needsAutoPermissionLease()) { + requestRuntimeDisposal(); + 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, + ); + }); + 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, ); + currentTurnPromise.catch(reportUnexpectedBackgroundError); this.updateDisplay({ status: 'background' as const }, updateOutput); return { @@ -3608,7 +3859,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 +3881,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..ea1526e20b9 100644 --- a/packages/core/src/tools/send-message.test.ts +++ b/packages/core/src/tools/send-message.test.ts @@ -258,6 +258,39 @@ describe('SendMessageTool — background-task mode', () => { ]); }); + it('revives a task when it finishes while a message waits at the finalization boundary', async () => { + registry.register({ + agentId: 'agent-1', + description: 'test agent', + status: 'running', + startTime: Date.now(), + abortController: new AbortController(), + isBackgrounded: true, + outputFile: '/tmp/test.jsonl', + }); + registry.beginFinishing('agent-1'); + reviveCompletedBackgroundAgent.mockResolvedValue(registry.get('agent-1')); + + const resultPromise = tool.validateBuildAndExecute( + { task_id: 'agent-1', message: 'late correction' }, + new AbortController().signal, + ); + await Promise.resolve(); + + expect(registry.get('agent-1')!.pendingMessages).toEqual([]); + expect(reviveCompletedBackgroundAgent).not.toHaveBeenCalled(); + + registry.complete('agent-1', 'done'); + const result = await resultPromise; + + expect(result.error).toBeUndefined(); + expect(result.llmContent).toContain('revived it with your message'); + expect(reviveCompletedBackgroundAgent).toHaveBeenCalledWith( + 'agent-1', + 'late correction', + ); + }); + it('returns error for non-existent task', async () => { const result = await tool.validateBuildAndExecute( { task_id: 'nope', message: 'hello' }, @@ -340,7 +373,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..714867e5911 100644 --- a/packages/core/src/tools/send-message.ts +++ b/packages/core/src/tools/send-message.ts @@ -80,7 +80,7 @@ class SendMessageInvocation extends BaseToolInvocation< return 'ask'; } - async execute(_signal: AbortSignal): Promise { + async execute(signal: AbortSignal): Promise { if (isPlanRequiredTeammateAwaitingApproval(this.config)) { const msg = getPlanRequiredTeammatePreApprovalMessage( ToolNames.SEND_MESSAGE, @@ -130,10 +130,22 @@ 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. A + // compatible runtime is not retained across session restore, so the + // persisted transcript remains the cold fallback for resumable agents. 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, @@ -166,7 +178,27 @@ class SendMessageInvocation extends BaseToolInvocation< }; } - registry.queueMessage(this.params.task_id, this.params.message); + if ( + registry.isFinishing(this.params.task_id) || + !registry.queueMessage(this.params.task_id, this.params.message) + ) { + const settled = await registry.waitForFinishing( + this.params.task_id, + signal, + ); + if (!settled) { + const message = `Message delivery to background task "${this.params.task_id}" was cancelled.`; + return { + llmContent: `Error: ${message}`, + returnDisplay: message, + error: { + message, + type: ToolErrorType.SEND_MESSAGE_NOT_RUNNING, + }, + }; + } + return this.execute(signal); + } return { llmContent: `Message queued for delivery to background task "${this.params.task_id}". The task will receive it at the next tool-round boundary.`, @@ -256,8 +288,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. ' + + 'Running tasks receive it at the next tool-round boundary; paused recovered tasks resume with the message as their first continuation instruction; completed tasks continue on their resident runtime when available and otherwise revive from their transcript. ' + 'Your text output is NOT visible to other agents — use this tool to communicate.', Kind.Other, { @@ -270,7 +302,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',