diff --git a/README.md b/README.md index 878cddc5..a3a75cbf 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,7 @@ https://github.com/user-attachments/assets/8685261b-9338-4fea-8dfe-1c590d5df543 - **Custom agent types** — define agents in `.pi/agents/.md` or `.agents/agents/.md` (project) or globally, with YAML frontmatter: custom system prompts, model selection, thinking levels, tool restrictions - **Nested subagents** — opt-in, default-off delegation: a custom agent that sets `allowed_subagents` gets its own ownership-scoped `Agent`, `get_subagent_result`, and `steer_subagent` tools, depth-capped from the main session (default 2). It can control only its own children, they are stopped when it finishes, and their transcripts and token spend roll up to it. The allowlist is a privilege boundary — a child runs with its own tools, so pick it as carefully as `tools:` itself - **Mid-run steering** — inject messages into running agents to redirect their work without restarting -- **Session resume** — pick up where an agent left off, preserving full conversation context +- **Session resume** — pick up where an agent left off, preserving full conversation context. Resumes in the foreground by default, or pass `run_in_background: true` to resume detached and be notified on completion, just like a background spawn - **Graceful turn limits** — agents get a "wrap up" warning before hard abort, producing clean partial results instead of cut-off output - **Case-insensitive agent types** — `"explore"`, `"Explore"`, `"EXPLORE"` all work. Unknown types fall back to general-purpose with a note - **Fuzzy model selection** — specify models by name (`"haiku"`, `"sonnet"`) instead of full IDs, with automatic filtering to only available/configured models diff --git a/src/agent-manager.ts b/src/agent-manager.ts index bfc3cc93..46462eb1 100644 --- a/src/agent-manager.ts +++ b/src/agent-manager.ts @@ -121,6 +121,23 @@ interface SpawnOptions { rootSessionId?: string; } +interface ResumeOptions { + /** + * Run the resumed turn detached in the background: return immediately with + * the record still "running" (or "queued" at the concurrency limit) and + * notify on completion via onComplete, exactly like a background spawn. + * Default (false/undefined) runs the resume inline and returns the settled + * record — the historical behavior. + */ + isBackground?: boolean; + /** Called on tool start/end with activity info (for streaming progress to UI). */ + onToolActivity?: (activity: ToolActivity) => void; + /** Called once per assistant message_end with that message's usage delta. */ + onAssistantUsage?: (usage: { input: number; output: number; cacheWrite: number }) => void; + /** Called when the session successfully compacts. */ + onCompaction?: (info: CompactionInfo) => void; +} + export class AgentManager { private agents = new Map(); private cleanupInterval: ReturnType; @@ -133,7 +150,7 @@ export class AgentManager { private worktreeRepos = new Set(); /** Queue of background agents waiting to start. */ - private queue: { id: string; args: SpawnArgs }[] = []; + private queue: { id: string; start: () => void }[] = []; /** Number of currently running background agents. */ private runningBackground = 0; @@ -209,7 +226,7 @@ export class AgentManager { if (occupiesPoolSlot(record) && !options.bypassQueue && this.runningBackground >= this.maxConcurrent) { // Queue it — will be started when a running agent completes - this.queue.push({ id, args }); + this.queue.push({ id, start: () => this.startAgent(id, record, args) }); return id; } @@ -441,7 +458,7 @@ export class AgentManager { const record = this.agents.get(next.id); if (!record || record.status !== "queued") continue; try { - this.startAgent(next.id, record, next.args); + next.start(); } catch (err) { // Late failure (e.g. strict worktree-isolation) — surface on the record // so the user/agent can see it via /agents, then keep draining. @@ -500,10 +517,36 @@ export class AgentManager { id: string, prompt: string, signal?: AbortSignal, + options?: ResumeOptions, ): Promise { const record = this.agents.get(id); if (!record?.session) return undefined; + // Background resume: settle asynchronously and notify on completion exactly + // like a background spawn, returning immediately with the record still + // "running" — or "queued" when at the concurrency limit. Previously + // run_in_background was ignored on resume (the Agent tool's resume branch + // returned before its background branch, and resume() only ever awaited + // inline), so a resumed agent always blocked the caller until it finished. + if (options?.isBackground) { + record.isBackground = true; + record.resultConsumed = false; + record.result = undefined; + record.error = undefined; + record.completedAt = undefined; + record.status = "queued"; + + const start = () => this.startResume(id, record, prompt, signal, options); + if (occupiesPoolSlot(record) && this.runningBackground >= this.maxConcurrent) { + // At the concurrency limit — queue it, drains when a slot frees. + this.queue.push({ id, start }); + } else { + start(); + } + return record; + } + + // Foreground resume: run inline and return the settled record. record.status = "running"; record.startedAt = Date.now(); record.completedAt = undefined; @@ -514,13 +557,16 @@ export class AgentManager { const { text, failure } = await resumeAgent(record.session, prompt, { onToolActivity: (activity) => { if (activity.type === "end") record.toolUses++; + options?.onToolActivity?.(activity); }, onAssistantUsage: (usage) => { addUsage(record.lifetimeUsage, usage); + options?.onAssistantUsage?.(usage); }, onCompaction: (info) => { record.compactionCount++; this.onCompact?.(record, info); + options?.onCompaction?.(info); }, signal, }); @@ -543,6 +589,96 @@ export class AgentManager { return record; } + /** + * Start a background resume run: detached, settling and notifying like + * startAgent's background path. Invoked immediately, or from drainQueue when + * a concurrency slot frees. The session already exists (resume reuses it), so + * output-file streaming is wired by the caller against record.session rather + * than through onSessionCreated. + */ + private startResume( + id: string, + record: AgentRecord, + prompt: string, + parentSignal: AbortSignal | undefined, + options: ResumeOptions, + ) { + if (!record.session) return; + + record.status = "running"; + record.startedAt = Date.now(); + if (occupiesPoolSlot(record)) this.runningBackground++; + this.onStart?.(record); + + // Fresh abort controller so /agents stop and steering target THIS run — and + // so the detached run isn't tied to the tool-call signal, which resolves the + // moment the tool returns. Wire the parent signal to abort the run. + const abortController = new AbortController(); + record.abortController = abortController; + let detachParentSignal: (() => void) | undefined; + if (parentSignal) { + const onParentAbort = () => this.abort(id); + parentSignal.addEventListener("abort", onParentAbort, { once: true }); + detachParentSignal = () => parentSignal.removeEventListener("abort", onParentAbort); + } + + const settle = () => { + detachParentSignal?.(); + detachParentSignal = undefined; + // Final flush of streaming output file + if (record.outputCleanup) { + try { record.outputCleanup(); } catch { /* ignore */ } + record.outputCleanup = undefined; + } + // Children spawned during the resumed turn must not outlive it. + this.abortOwnedChildren(id); + if (occupiesPoolSlot(record)) this.runningBackground--; + try { this.onComplete?.(record); } catch { /* ignore completion side-effect errors */ } + this.drainQueue(); + }; + + const promise = resumeAgent(record.session, prompt, { + onToolActivity: (activity) => { + if (activity.type === "end") record.toolUses++; + options.onToolActivity?.(activity); + }, + onAssistantUsage: (usage) => { + addUsage(record.lifetimeUsage, usage); + options.onAssistantUsage?.(usage); + }, + onCompaction: (info) => { + record.compactionCount++; + this.onCompact?.(record, info); + options.onCompaction?.(info); + }, + signal: abortController.signal, + }) + .then(({ text, failure }) => { + // Don't overwrite status if externally stopped via abort(). + if (record.status !== "stopped") { + // Same contract as the spawn path (#144): a failed final turn is an + // error, not a completion — but the resumed text stays available. + record.status = failure ? "error" : "completed"; + if (failure) record.error = failure; + } + record.result = text; + record.completedAt ??= Date.now(); + settle(); + return text; + }) + .catch((err) => { + if (record.status !== "stopped") { + record.status = "error"; + record.error = err instanceof Error ? err.message : String(err); + } + record.completedAt ??= Date.now(); + settle(); + return ""; + }); + + record.promise = promise; + } + /** * Send a steering message to an agent from the UI (mirrors the steer_subagent * tool). A live session delivers it now — it interrupts the agent after its diff --git a/src/index.ts b/src/index.ts index 8736dd23..d261b116 100644 --- a/src/index.ts +++ b/src/index.ts @@ -940,7 +940,7 @@ Terse command-style prompts produce shallow, generic work. ), resume: Type.Optional( Type.String({ - description: "Optional agent ID to resume from. Continues from previous context.", + description: "Optional agent ID to resume from. Continues from previous context. Combine with run_in_background to resume detached and be notified on completion.", }), ), isolated: Type.Optional( @@ -1215,6 +1215,65 @@ Terse command-style prompts produce shallow, generic work. if (!existing.session) { return textResult(`Agent "${params.resume}" has no active session to resume.`); } + + // Background resume: detached run that notifies on completion, mirroring + // a background spawn. Previously run_in_background was silently ignored + // on resume (this branch returned before the background branch below), + // so a resumed agent always blocked the main loop until it finished. + if (runInBackground) { + const id = existing.id; + const joinMode = resolveJoinMode(defaultJoinMode, true); + existing.toolCallId = toolCallId; + if (joinMode) existing.joinMode = joinMode; + // Session already exists — attach a fresh transcript for THIS run; + // streaming is wired directly below (no onSessionCreated fires). + attachTranscript(existing, id); + + const { state: bgState, callbacks: bgCallbacks } = createActivityTracker(effectiveMaxTurns); + const record = await manager.resume(params.resume, params.prompt, signal, { + isBackground: true, + onToolActivity: bgCallbacks.onToolActivity, + onAssistantUsage: bgCallbacks.onAssistantUsage, + }); + if (!record) { + return textResult(`Failed to resume agent "${params.resume}".`); + } + if (record.session && record.outputFile) { + record.outputCleanup = streamToOutputFile(record.session, record.outputFile, id, ctx.cwd); + } + + if (joinMode != null && joinMode !== 'async') { + currentBatchAgents.push({ id, joinMode }); + if (batchFinalizeTimer) clearTimeout(batchFinalizeTimer); + batchFinalizeTimer = setTimeout(finalizeBatch, 100); + } + + agentActivity.set(id, bgState); + widget.ensureTimer(); + widget.update(); + fleet.ensureTimer(); + fleet.update(); + + pi.events.emit("subagents:created", { + id, + type: subagentType, + description: params.description, + isBackground: true, + }); + + const isQueued = record.status === "queued"; + return textResult( + `Agent ${isQueued ? "queued" : "resumed"} in background.\n` + + `Agent ID: ${id}\n` + + `Type: ${displayName}\n` + + (record.outputFile ? `Output file: ${record.outputFile}\n` : "") + + (isQueued ? `Position: queued (max ${manager.getMaxConcurrent()} concurrent)\n` : "") + + `\nYou will be notified when this agent completes.\n` + + `Use get_subagent_result to retrieve full results, or steer_subagent to send it messages.`, + { ...detailBase, toolUses: record.toolUses, tokens: "", durationMs: 0, status: "background" as const, agentId: id }, + ); + } + const record = await manager.resume(params.resume, params.prompt, signal); if (!record) { return textResult(`Failed to resume agent "${params.resume}".`); diff --git a/test/agent-manager.test.ts b/test/agent-manager.test.ts index 35e20563..bc050510 100644 --- a/test/agent-manager.test.ts +++ b/test/agent-manager.test.ts @@ -1210,3 +1210,134 @@ describe("AgentManager — resolved runs with a failed final turn map to error ( expect(record.result).toBe("new partial progress"); // salvageable, this-run text }); }); + +describe("AgentManager — background resume", () => { + let manager: AgentManager; + + afterEach(() => { + manager?.dispose(); + }); + + // Spawn a background agent and let it settle so it holds a session to resume. + async function spawnSettled(mgr: AgentManager): Promise { + vi.mocked(runAgent).mockResolvedValue({ + responseText: "first", + session: mockSession(), + aborted: false, + steered: false, + }); + const id = mgr.spawn(mockPi, mockCtx, "general-purpose", "task", { + description: "task", + isBackground: true, + }); + await mgr.getRecord(id)!.promise; + return id; + } + + it("returns immediately with a running record + promise, then settles and fires onComplete", async () => { + const onComplete = vi.fn(); + manager = new AgentManager(onComplete); + const id = await spawnSettled(manager); + onComplete.mockClear(); // drop the spawn's own completion + + // Deferred resumeAgent so we can observe the mid-flight state. + let finish!: (v: { text: string; failure?: string }) => void; + vi.mocked(resumeAgent).mockImplementation( + () => new Promise((resolve) => { finish = resolve; }), + ); + + const record = await manager.resume(id, "keep going", undefined, { isBackground: true }); + // Returned immediately: still running, with a tracked promise, no notify yet. + expect(record?.status).toBe("running"); + expect(record?.promise).toBeDefined(); + expect(onComplete).not.toHaveBeenCalled(); + + finish({ text: "second" }); + await record!.promise; + + expect(manager.getRecord(id)!.status).toBe("completed"); + expect(manager.getRecord(id)!.result).toBe("second"); + expect(onComplete).toHaveBeenCalledTimes(1); + }); + + it("a failed final turn on a background resume maps to error and still notifies", async () => { + const onComplete = vi.fn(); + manager = new AgentManager(onComplete); + const id = await spawnSettled(manager); + onComplete.mockClear(); + + vi.mocked(resumeAgent).mockResolvedValue({ + text: "partial", + failure: "provider exploded", + } as any); + + const record = await manager.resume(id, "again", undefined, { isBackground: true }); + await record!.promise; + + expect(manager.getRecord(id)!.status).toBe("error"); + expect(manager.getRecord(id)!.error).toBe("provider exploded"); + expect(manager.getRecord(id)!.result).toBe("partial"); // #144: keep this-run text + expect(onComplete).toHaveBeenCalledTimes(1); + }); + + it("forwards activity/usage callbacks to the resumed run", async () => { + manager = new AgentManager(); + const id = await spawnSettled(manager); + + const onToolActivity = vi.fn(); + const onAssistantUsage = vi.fn(); + vi.mocked(resumeAgent).mockImplementation(async (_session, _prompt, opts: any) => { + opts.onToolActivity?.({ type: "end", toolName: "grep" }); + opts.onAssistantUsage?.({ input: 5, output: 3, cacheWrite: 0 }); + return { text: "ok" }; + }); + + const record = await manager.resume(id, "go", undefined, { + isBackground: true, + onToolActivity, + onAssistantUsage, + }); + await record!.promise; + + expect(onToolActivity).toHaveBeenCalledWith({ type: "end", toolName: "grep" }); + expect(onAssistantUsage).toHaveBeenCalledWith({ input: 5, output: 3, cacheWrite: 0 }); + // Internal record bookkeeping still runs alongside the forwarded callbacks. + expect(manager.getRecord(id)!.toolUses).toBe(1); + expect(manager.getRecord(id)!.lifetimeUsage).toEqual({ input: 5, output: 3, cacheWrite: 0 }); + }); + + it("queues a background resume when the concurrency pool is full", async () => { + manager = new AgentManager(undefined, 1); // maxConcurrent = 1 + const id = await spawnSettled(manager); + + // Occupy the single slot with a never-settling background spawn. + vi.mocked(runAgent).mockImplementation(() => new Promise(() => {})); + const blockerId = manager.spawn(mockPi, mockCtx, "general-purpose", "blocker", { + description: "blocker", + isBackground: true, + }); + expect(manager.getRecord(blockerId)!.status).toBe("running"); + + vi.mocked(resumeAgent).mockImplementation(() => new Promise(() => {})); + vi.mocked(resumeAgent).mockClear(); // drop call history from earlier tests + const record = await manager.resume(id, "later", undefined, { isBackground: true }); + + expect(record?.status).toBe("queued"); + expect(resumeAgent).not.toHaveBeenCalled(); + }); + + it("foreground resume is unchanged: awaits inline and does not fire onComplete", async () => { + const onComplete = vi.fn(); + manager = new AgentManager(onComplete); + const id = await spawnSettled(manager); + onComplete.mockClear(); + + vi.mocked(resumeAgent).mockResolvedValue({ text: "inline result" } as any); + const record = await manager.resume(id, "sync"); + + expect(record?.status).toBe("completed"); + expect(record?.result).toBe("inline result"); + // Foreground resume returns its result inline and never notified (historical). + expect(onComplete).not.toHaveBeenCalled(); + }); +});