From fbb5725f463707850539944bdeb8bab23a1b2471 Mon Sep 17 00:00:00 2001 From: Alexander Penkin Date: Thu, 16 Jul 2026 10:23:50 +0300 Subject: [PATCH 1/4] chore: initialize durable resume split From 46678846a288f6d05f00a97e412d0149aa78c48a Mon Sep 17 00:00:00 2001 From: Alexander Penkin Date: Thu, 16 Jul 2026 10:51:53 +0300 Subject: [PATCH 2/4] feat: add workflow control and observability --- README.md | 4 +- extensions/workflow.ts | 9 +- src/index.ts | 6 + src/workflow-control-tool.ts | 216 +++++++++++++++++++++++++ tests/agent-registry.test.ts | 4 + tests/task-panel.test.ts | 5 +- tests/usage-limit-integration.test.ts | 4 +- tests/workflow-control-tool.test.ts | 188 +++++++++++++++++++++ tests/workflow-tools-available.test.ts | 59 ++++++- 9 files changed, 486 insertions(+), 9 deletions(-) create mode 100644 src/workflow-control-tool.ts create mode 100644 tests/workflow-control-tool.test.ts diff --git a/README.md b/README.md index 3619a1c7..b6ec5fd9 100644 --- a/README.md +++ b/README.md @@ -107,7 +107,9 @@ return await agent( For an always-on exhaustive mode, use `/ultracode`; `/effort high` is the lighter standing option. -## Commands +## Commands and run control + +Pi can manage background runs directly with the `workflow_control` tool instead of asking you to type a command. It supports `list`, `status`, `pause`, `resume`, and `stop`; run-specific actions use the canonical run ID returned when the workflow starts. Status output includes the run state, current phase, agent counts, active labels, and recorded token total. | Command | Purpose | | --- | --- | diff --git a/extensions/workflow.ts b/extensions/workflow.ts index 90282040..25c6922a 100644 --- a/extensions/workflow.ts +++ b/extensions/workflow.ts @@ -1,6 +1,7 @@ import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent"; import { createEffortState, + createWorkflowControlTool, createWorkflowStorage, createWorkflowTool, installResultDelivery, @@ -33,7 +34,9 @@ export default function extension(pi: ExtensionAPI) { }); const workflowTool = createWorkflowTool({ cwd, manager, storage }); + const workflowControlTool = createWorkflowControlTool({ manager }); pi.registerTool(workflowTool); + pi.registerTool(workflowControlTool); // Auto-resume runs that paused on a provider usage limit once the quota is // likely refilled. Standalone: only consumes the manager's public surface, so // it stays decoupled from manager/persistence internals. Its constructor also @@ -68,9 +71,9 @@ export default function extension(pi: ExtensionAPI) { // advertise the shared registry's models. manager.setModelRegistry(ctx.modelRegistry); const active = pi.getActiveTools(); - if (!active.includes(workflowTool.name)) { - pi.setActiveTools([...active, workflowTool.name]); - } + const workflowTools = [workflowTool.name, workflowControlTool.name]; + const missing = workflowTools.filter((name) => !active.includes(name)); + if (missing.length) pi.setActiveTools([...active, ...missing]); // Scope the /workflows history to this session: runs persist on disk across // sessions, but the navigator/task panel show only the current session's runs. // Switching back to a previous session re-shows that session's runs. diff --git a/src/index.ts b/src/index.ts index 534493ab..a47f5944 100644 --- a/src/index.ts +++ b/src/index.ts @@ -96,6 +96,12 @@ export type { } from "./workflow.js"; export { parseWorkflowScript, runWorkflow } from "./workflow.js"; export { registerWorkflowCommands } from "./workflow-commands.js"; +export type { + WorkflowControlInput, + WorkflowControlRunDetails, + WorkflowControlToolOptions, +} from "./workflow-control-tool.js"; +export { createWorkflowControlTool } from "./workflow-control-tool.js"; export { buildForcedWorkflowPrompt, colorizeWorkflow, diff --git a/src/workflow-control-tool.ts b/src/workflow-control-tool.ts new file mode 100644 index 00000000..5e76cdc6 --- /dev/null +++ b/src/workflow-control-tool.ts @@ -0,0 +1,216 @@ +import { defineTool, type ToolDefinition } from "@earendil-works/pi-coding-agent"; +import { type Static, Type } from "typebox"; +import { aggregateAgentUsage, tokenFigures, type WorkflowAgentSnapshot, type WorkflowSnapshot } from "./display.js"; +import type { PersistedRunState, RunStatus } from "./run-persistence.js"; +import type { WorkflowManager } from "./workflow-manager.js"; + +const runActionSchema = Type.Union([ + Type.Literal("status"), + Type.Literal("pause"), + Type.Literal("resume"), + Type.Literal("stop"), +]); + +const workflowControlSchema = Type.Union([ + Type.Object( + { action: Type.Literal("list", { description: "List workflow runs." }) }, + { additionalProperties: false }, + ), + Type.Object( + { + action: runActionSchema, + runId: Type.String({ minLength: 1, description: "Canonical workflow run ID." }), + }, + { additionalProperties: false }, + ), +]); + +export type WorkflowControlInput = Static; + +export interface WorkflowControlToolOptions { + manager: WorkflowManager; +} + +export interface WorkflowControlRunDetails { + runId: string; + workflowName: string; + status: RunStatus; + phase: string | null; + counts: { + total: number; + done: number; + running: number; + queued: number; + error: number; + skipped: number; + }; + activeLabels: string[]; + tokenTotal: number; +} + +type ControlResult = { + content: Array<{ type: "text"; text: string }>; + details: Record; +}; + +export function createWorkflowControlTool( + options: WorkflowControlToolOptions, +): ToolDefinition> { + const manager = options.manager; + return defineTool({ + name: "workflow_control", + label: "Workflow Control", + description: + "List and inspect workflow runs, or pause, resume, and stop them without asking the user to run slash commands.", + promptSnippet: "Inspect and manage workflow runs directly by canonical run ID.", + promptGuidelines: [ + "Use workflow_control for workflow lifecycle management; do not ask the user to type /workflows when this tool can perform the action.", + "Use stop to terminate or quit a run. Closing the navigator does not stop a run.", + ], + parameters: workflowControlSchema, + prepareArguments: normalizeInput, + async execute(_toolCallId, params) { + if (params.action === "list") { + const runs = manager.listRuns(); + const summaries = runs.map((run) => summarizeRun(run, manager.getSnapshot(run.runId))); + return result( + summaries.length + ? `action=list result=ok runs=${summaries.length}\n${summaries.map(formatRun).join("\n")}` + : "action=list result=ok runs=0", + { action: "list", result: "ok", runs: summaries }, + ); + } + + const run = findRun(manager, params.runId); + if (!run) return controlError(params.action, params.runId, "run not found", ["list"]); + + switch (params.action) { + case "status": { + const summary = summarizeRun(run, manager.getSnapshot(run.runId)); + return result(`action=status result=ok ${formatRun(summary)}`, { + action: "status", + result: "ok", + run: summary, + }); + } + case "pause": + if (!manager.pause(run.runId)) return invalidTransition("pause", run); + return actionSuccess("pause", "paused", currentSummary(manager, run)); + case "resume": + if (!(await manager.resume(run.runId))) return invalidTransition("resume", run); + return actionSuccess("resume", "resumed", currentSummary(manager, run)); + case "stop": + if (!manager.stop(run.runId)) return invalidTransition("stop", run); + return actionSuccess("stop", "stopped", currentSummary(manager, run)); + } + }, + }); +} + +function normalizeInput(value: unknown): WorkflowControlInput { + if (!value || typeof value !== "object" || Array.isArray(value)) { + throw new Error("workflow_control requires an object argument"); + } + const input = value as Record; + const actions = new Set(["list", "status", "pause", "resume", "stop"]); + if (typeof input.action !== "string" || !actions.has(input.action)) { + throw new Error("workflow_control requires action: list|status|pause|resume|stop"); + } + + const allowedKeys = input.action === "list" ? new Set(["action"]) : new Set(["action", "runId"]); + const extraKey = Object.keys(input).find((key) => !allowedKeys.has(key)); + if (extraKey) throw new Error(`workflow_control action "${input.action}" does not accept ${extraKey}`); + + if (input.action !== "list" && (typeof input.runId !== "string" || !input.runId.trim())) { + throw new Error(`workflow_control action "${input.action}" requires runId`); + } + return input as WorkflowControlInput; +} + +function result(text: string, details: Record): ControlResult { + return { content: [{ type: "text", text }], details }; +} + +function findRun(manager: WorkflowManager, runId: string): PersistedRunState | undefined { + return manager.listRuns().find((candidate) => candidate.runId === runId); +} + +function currentSummary(manager: WorkflowManager, fallback: PersistedRunState): WorkflowControlRunDetails { + const current = findRun(manager, fallback.runId) ?? fallback; + return summarizeRun(current, manager.getSnapshot(current.runId)); +} + +function actionSuccess(action: string, actionResult: string, run: WorkflowControlRunDetails): ControlResult { + return result(`action=${action} result=${actionResult} ${formatRun(run)}`, { + action, + result: actionResult, + run, + }); +} + +function invalidTransition(action: string, run: PersistedRunState): ControlResult { + return controlError(action, run.runId, `cannot ${action} run with status ${run.status}`, allowedActions(run.status)); +} + +function controlError(action: string, runId: string, message: string, allowed: string[]): ControlResult { + return result( + `action=${action} result=error runId=${runId} error=${message} allowed=${allowed.join(",") || "none"}`, + { action, result: "error", runId, error: message, allowedActions: allowed }, + ); +} + +function allowedActions(status: RunStatus): string[] { + switch (status) { + case "running": + return ["status", "pause", "stop"]; + case "paused": + return ["status", "resume", "stop"]; + case "failed": + case "pending": + return ["status", "resume"]; + case "completed": + case "aborted": + return ["status"]; + } +} + +function summarizeRun(run: PersistedRunState, live?: WorkflowSnapshot | null): WorkflowControlRunDetails { + const agents = live?.agents ?? run.agents; + const counts = countAgents(agents); + const liveUsage = tokenFigures(live?.tokenUsage); + const persistedUsage = tokenFigures(run.tokenUsage); + const agentUsage = aggregateAgentUsage(agents); + return { + runId: run.runId, + workflowName: live?.name ?? run.workflowName, + status: run.status, + phase: live?.currentPhase ?? run.currentPhase ?? null, + counts, + activeLabels: agents.filter((agent) => agent.status === "running").map((agent) => agent.label), + tokenTotal: Math.max( + liveUsage.fresh + liveUsage.cacheRead, + persistedUsage.fresh + persistedUsage.cacheRead, + agentUsage.fresh + agentUsage.cacheRead, + ), + }; +} + +function countAgents(agents: Array>): WorkflowControlRunDetails["counts"] { + return { + total: agents.length, + done: agents.filter((agent) => agent.status === "done").length, + running: agents.filter((agent) => agent.status === "running").length, + queued: agents.filter((agent) => agent.status === "queued").length, + error: agents.filter((agent) => agent.status === "error").length, + skipped: agents.filter((agent) => agent.status === "skipped").length, + }; +} + +function formatRun(run: WorkflowControlRunDetails): string { + const active = run.activeLabels.join(",") || "-"; + return `runId=${run.runId} name=${quote(run.workflowName)} status=${run.status} phase=${quote(run.phase ?? "-")} total=${run.counts.total} done=${run.counts.done} running=${run.counts.running} queued=${run.counts.queued} error=${run.counts.error} skipped=${run.counts.skipped} active=${quote(active)} tokens=${run.tokenTotal}`; +} + +function quote(value: string): string { + return JSON.stringify(value); +} diff --git a/tests/agent-registry.test.ts b/tests/agent-registry.test.ts index 66143e69..aa188e85 100644 --- a/tests/agent-registry.test.ts +++ b/tests/agent-registry.test.ts @@ -147,9 +147,11 @@ describe("loadAgentRegistry", () => { it("default userDir resolution uses getAgentDir() (~/.pi/agent/agents) with no injected opts", () => { const tmpHome = mkdtempSync(join(tmpdir(), "pi-home-")); const originalHome = process.env.HOME; + const originalUserProfile = process.env.USERPROFILE; const originalAgentDirEnv = process.env.PI_CODING_AGENT_DIR; delete process.env.PI_CODING_AGENT_DIR; process.env.HOME = tmpHome; + process.env.USERPROFILE = tmpHome; try { const expectedUserDir = join(getAgentDir(), "agents"); assert.equal(expectedUserDir, join(tmpHome, ".pi", "agent", "agents"), "sanity: HOME override took effect"); @@ -163,6 +165,8 @@ describe("loadAgentRegistry", () => { } finally { if (originalHome === undefined) delete process.env.HOME; else process.env.HOME = originalHome; + if (originalUserProfile === undefined) delete process.env.USERPROFILE; + else process.env.USERPROFILE = originalUserProfile; if (originalAgentDirEnv === undefined) delete process.env.PI_CODING_AGENT_DIR; else process.env.PI_CODING_AGENT_DIR = originalAgentDirEnv; rmSync(tmpHome, { recursive: true, force: true }); diff --git a/tests/task-panel.test.ts b/tests/task-panel.test.ts index 39efeffa..d03a8b05 100644 --- a/tests/task-panel.test.ts +++ b/tests/task-panel.test.ts @@ -1,5 +1,6 @@ import assert from "node:assert/strict"; import { EventEmitter } from "node:events"; +import { join } from "node:path"; import { before, describe, it } from "node:test"; import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; import { visibleWidth } from "@earendil-works/pi-tui"; @@ -242,7 +243,7 @@ describe("installResultDelivery", () => { const content = (pi as unknown as { _calls: { content: string }[] })._calls[0].content; assert.ok(content.includes("Full result:"), "should include the pointer label"); - assert.ok(content.includes("/runs/test-run-1.json"), "should point at /.json"); + assert.ok(content.includes(join("/runs", "test-run-1.json")), "should point at /.json"); // The verdict summary itself is unchanged apart from the appended pointer. assert.ok(content.includes("All tests passed"), "verdict text preserved"); }); @@ -275,7 +276,7 @@ describe("installResultDelivery", () => { const content = (pi as unknown as { _calls: { content: string }[] })._calls[0].content; assert.ok(/…\(truncated [\d.]+ (B|KB|MB)\)/.test(content), "the 50-char setting truncates a sub-400 dump"); assert.ok(!content.includes("z".repeat(200)), "the body is cut at the configured threshold"); - assert.ok(content.includes("/runs/test-run-1.json"), "pointer still appended"); + assert.ok(content.includes(join("/runs", "test-run-1.json")), "pointer still appended"); }); // ── installResultDelivery: guard / stale ctx ── diff --git a/tests/usage-limit-integration.test.ts b/tests/usage-limit-integration.test.ts index 03c8483e..354c9a66 100644 --- a/tests/usage-limit-integration.test.ts +++ b/tests/usage-limit-integration.test.ts @@ -15,7 +15,7 @@ import { existsSync, mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import test from "node:test"; -import { fileURLToPath } from "node:url"; +import { fileURLToPath, pathToFileURL } from "node:url"; import { WorkflowAgent } from "../src/agent.js"; import { WorkflowErrorCode } from "../src/errors.js"; import { WorkflowManager } from "../src/workflow-manager.js"; @@ -38,7 +38,7 @@ async function loadFaux(): Promise; } diff --git a/tests/workflow-control-tool.test.ts b/tests/workflow-control-tool.test.ts new file mode 100644 index 00000000..31d0374b --- /dev/null +++ b/tests/workflow-control-tool.test.ts @@ -0,0 +1,188 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { Check } from "typebox/value"; +import type { WorkflowSnapshot } from "../src/display.js"; +import type { PersistedRunState, RunStatus } from "../src/run-persistence.js"; +import { createWorkflowControlTool } from "../src/workflow-control-tool.js"; +import type { WorkflowManager } from "../src/workflow-manager.js"; + +function run(status: RunStatus = "running", runId = "audit-abc123"): PersistedRunState { + return { + runId, + workflowName: "audit", + script: "export const meta = { name: 'audit', description: 'audit' }; return await agent('x')", + status, + phases: ["Inspect"], + currentPhase: "Inspect", + agents: [ + { id: 1, label: "active scan", prompt: "scan", status: status === "running" ? "running" : "done", tokens: 30 }, + { id: 2, label: "queued check", prompt: "check", status: "queued" }, + { id: 3, label: "failed check", prompt: "fail", status: "error" }, + { id: 4, label: "optional check", prompt: "optional", status: "skipped" }, + ], + logs: [], + startedAt: "2026-07-14T00:00:00.000Z", + updatedAt: "2026-07-14T00:00:01.000Z", + tokenUsage: { input: 20, output: 10, total: 30 }, + }; +} + +function fakeManager(initial: PersistedRunState[], liveSnapshots: Record = {}) { + const runs = new Map(initial.map((item) => [item.runId, item])); + const calls: Array<{ action: string; runId: string }> = []; + const manager = { + listRuns: () => [...runs.values()], + getSnapshot: (runId: string) => liveSnapshots[runId] ?? null, + pause(runId: string) { + calls.push({ action: "pause", runId }); + const item = runs.get(runId); + if (item?.status !== "running") return false; + item.status = "paused"; + return true; + }, + async resume(runId: string) { + calls.push({ action: "resume", runId }); + const item = runs.get(runId); + if (!item || (item.status !== "paused" && item.status !== "failed" && item.status !== "pending")) return false; + item.status = "running"; + return true; + }, + stop(runId: string) { + calls.push({ action: "stop", runId }); + const item = runs.get(runId); + if (!item || (item.status !== "running" && item.status !== "paused")) return false; + item.status = "aborted"; + return true; + }, + } as unknown as WorkflowManager; + return { manager, calls }; +} + +async function execute(manager: WorkflowManager, params: Record) { + const tool = createWorkflowControlTool({ manager }); + return (tool.execute as any)("control-call", params, undefined, undefined, {}); +} + +function text(result: Awaited>): string { + return result.content[0].text; +} + +test("workflow_control exposes only list, status, pause, resume, and stop in a strict schema", () => { + const { manager } = fakeManager([]); + const tool = createWorkflowControlTool({ manager }); + + assert.equal(tool.name, "workflow_control"); + assert.equal(Check(tool.parameters, { action: "list" }), true); + assert.equal(Check(tool.parameters, { action: "status", runId: "abc" }), true); + assert.equal(Check(tool.parameters, { action: "pause", runId: "abc" }), true); + assert.equal(Check(tool.parameters, { action: "resume", runId: "abc" }), true); + assert.equal(Check(tool.parameters, { action: "stop", runId: "abc" }), true); + assert.equal(Check(tool.parameters, { action: "restart", runId: "abc" }), false); + assert.equal(Check(tool.parameters, { action: "remove", runId: "abc" }), false); + assert.equal(Check(tool.parameters, { action: "set_concurrency", runId: "abc", concurrency: 2 }), false); + assert.equal(Check(tool.parameters, { action: "status" }), false); + assert.equal(Check(tool.parameters, { action: "list", runId: "abc" }), false); + assert.equal(Check(tool.parameters, { action: "status", runId: "abc", extra: true }), false); + + const prepare = tool.prepareArguments as (value: unknown) => unknown; + assert.throws(() => prepare({ action: "pause" }), /requires runId/); + assert.throws(() => prepare({ action: "status", runId: "abc", extra: true }), /does not accept extra/); + assert.throws(() => prepare({ action: "restart", runId: "abc" }), /requires action/); +}); + +test("list and status return stable lifecycle and observability fields", async () => { + const { manager } = fakeManager([run()]); + + const listed = await execute(manager, { action: "list" }); + assert.match(text(listed), /^action=list result=ok runs=1\n/); + assert.match(text(listed), /runId=audit-abc123 name="audit" status=running phase="Inspect"/); + assert.match(text(listed), /total=4 done=0 running=1 queued=1 error=1 skipped=1/); + assert.match(text(listed), /active="active scan" tokens=30/); + assert.deepEqual(listed.details, { + action: "list", + result: "ok", + runs: [ + { + runId: "audit-abc123", + workflowName: "audit", + status: "running", + phase: "Inspect", + counts: { total: 4, done: 0, running: 1, queued: 1, error: 1, skipped: 1 }, + activeLabels: ["active scan"], + tokenTotal: 30, + }, + ], + }); + + const status = await execute(manager, { action: "status", runId: "audit-abc123" }); + assert.match(text(status), /^action=status result=ok /); + assert.equal(status.details.action, "status"); + assert.equal((status.details.run as { runId: string }).runId, "audit-abc123"); + assert.doesNotMatch(text(status), /\/workflows/); +}); + +test("status uses agent usage when the live run aggregate is lagging", async () => { + const live: WorkflowSnapshot = { + name: "audit", + phases: ["Inspect"], + currentPhase: "Inspect", + logs: [], + agents: [ + { id: 1, label: "estimated", prompt: "scan", status: "running", tokens: 80 }, + { + id: 2, + label: "reported", + prompt: "check", + status: "done", + tokens: 40, + tokenUsage: { input: 15, output: 5, total: 40, cacheRead: 20, cacheWrite: 0, cost: 0 }, + }, + ], + agentCount: 2, + runningCount: 1, + doneCount: 1, + errorCount: 0, + tokenUsage: { input: 0, output: 0, total: 0 }, + }; + const { manager } = fakeManager([run()], { "audit-abc123": live }); + + const status = await execute(manager, { action: "status", runId: "audit-abc123" }); + + assert.match(text(status), /tokens=120$/); + assert.equal((status.details.run as { tokenTotal: number }).tokenTotal, 120); +}); + +test("list reports an explicit empty result", async () => { + const { manager } = fakeManager([]); + const response = await execute(manager, { action: "list" }); + assert.equal(text(response), "action=list result=ok runs=0"); + assert.deepEqual(response.details, { action: "list", result: "ok", runs: [] }); +}); + +test("pause, resume, and stop call the shared manager lifecycle methods", async () => { + const fixture = fakeManager([run()]); + + assert.match(text(await execute(fixture.manager, { action: "pause", runId: "audit-abc123" })), /result=paused/); + assert.match(text(await execute(fixture.manager, { action: "resume", runId: "audit-abc123" })), /result=resumed/); + assert.match(text(await execute(fixture.manager, { action: "stop", runId: "audit-abc123" })), /result=stopped/); + assert.deepEqual( + fixture.calls.map((call) => call.action), + ["pause", "resume", "stop"], + ); +}); + +test("unknown IDs and illegal transitions return explicit errors with allowed actions", async () => { + const fixture = fakeManager([run("completed"), run("running", "live-123")]); + + const unknown = text(await execute(fixture.manager, { action: "status", runId: "missing" })); + assert.match(unknown, /result=error runId=missing error=run not found allowed=list/); + + const pauseCompleted = text(await execute(fixture.manager, { action: "pause", runId: "audit-abc123" })); + assert.match(pauseCompleted, /cannot pause run with status completed/); + assert.match(pauseCompleted, /allowed=status/); + + await execute(fixture.manager, { action: "stop", runId: "live-123" }); + const stopAborted = text(await execute(fixture.manager, { action: "stop", runId: "live-123" })); + assert.match(stopAborted, /cannot stop run with status aborted/); + assert.match(stopAborted, /allowed=status/); +}); diff --git a/tests/workflow-tools-available.test.ts b/tests/workflow-tools-available.test.ts index 539a5b83..4456e630 100644 --- a/tests/workflow-tools-available.test.ts +++ b/tests/workflow-tools-available.test.ts @@ -14,9 +14,13 @@ */ import assert from "node:assert/strict"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; import { describe, it, mock } from "node:test"; import type { ExtensionAPI, ExtensionUIContext } from "@earendil-works/pi-coding-agent"; import { buildForcedWorkflowPrompt, WORKFLOW_TOOL_NAME, type WorkflowModeState } from "../src/workflow-editor.js"; +import { withFakeHomeAsync } from "./helpers/fake-home.js"; // --------------------------------------------------------------------------- // Default Pi tools that every Pi install provides (plugin-independent) @@ -33,6 +37,7 @@ const DEFAULT_PI_TOOLS = [ "advisor", "subagent", "workflow", + "workflow_control", ]; // Additional tools from context-mode plugin (common but not guaranteed) @@ -141,7 +146,7 @@ describe("installWorkflowEditor - tool availability", () => { const { installWorkflowEditor } = await import("../src/workflow-editor.js"); // Add a bonus tool to simulate a plugin adding a tool - const originalTools = ["bash", "read", "edit", "write", "custom-plugin-tool", "workflow"]; + const originalTools = ["bash", "read", "edit", "write", "custom-plugin-tool", "workflow", "workflow_control"]; const mockPi = createMockPi(originalTools); const ui = { @@ -450,3 +455,55 @@ describe("installWorkflowEditor - tool availability", () => { assert.equal(state.active, false); }); }); + +describe("workflow extension - control tool availability", () => { + it("registers and activates workflow and workflow_control together", async () => { + const fakeHome = mkdtempSync(join(tmpdir(), "pi-dw-control-extension-")); + try { + await withFakeHomeAsync(fakeHome, async () => { + const registeredTools: string[] = []; + const activeTools = ["bash", "read"]; + const handlers: Record any>> = {}; + const pi = { + registerTool: (tool: { name: string }) => registeredTools.push(tool.name), + registerCommand: () => {}, + getCommands: () => [], + on: (event: string, handler: (...args: any[]) => any) => { + if (!handlers[event]) handlers[event] = []; + handlers[event].push(handler); + }, + getActiveTools: () => [...activeTools], + setActiveTools: (tools: string[]) => { + activeTools.splice(0, activeTools.length, ...tools); + }, + sendMessage: () => {}, + } as unknown as ExtensionAPI; + const { default: installExtension } = await import("../extensions/workflow.js"); + + installExtension(pi); + + assert.deepEqual(registeredTools.slice(0, 2), ["workflow", "workflow_control"]); + assert.equal(handlers.session_start.length, 1); + handlers.session_start[0]( + {}, + { + model: undefined, + modelRegistry: {}, + sessionManager: { getSessionId: () => "session-1" }, + ui: { + setWidget: () => {}, + getEditorComponent: () => undefined, + setEditorComponent: () => {}, + }, + }, + ); + + assert.ok(activeTools.includes("workflow")); + assert.ok(activeTools.includes("workflow_control")); + handlers.session_shutdown?.[0]?.(); + }); + } finally { + rmSync(fakeHome, { recursive: true, force: true }); + } + }); +}); From e2bb4e6736672c7555946238dffd1295202c8fc7 Mon Sep 17 00:00:00 2001 From: Alexander Penkin Date: Thu, 16 Jul 2026 14:49:25 +0300 Subject: [PATCH 3/4] feat: resume workflow runs durably --- README.md | 4 +- src/agent.ts | 84 ++- src/display.ts | 13 +- src/run-persistence.ts | 391 ++++++++++-- src/task-panel.ts | 5 +- src/workflow-control-tool.ts | 55 +- src/workflow-manager.ts | 874 +++++++++++++++++++++------ src/workflow-ui.ts | 28 +- src/workflow.ts | 396 ++++++++---- tests/agent.test.ts | 114 +++- tests/run-persistence.test.ts | 253 +++++++- tests/workflow-control-tool.test.ts | 86 ++- tests/workflow-manager-abort.test.ts | 47 +- tests/workflow-manager.test.ts | 323 ++++++++-- tests/workflow-resume-state.test.ts | 515 ++++++++++++++++ tests/workflow-runtime.test.ts | 294 ++++++++- 16 files changed, 3008 insertions(+), 474 deletions(-) create mode 100644 tests/workflow-resume-state.test.ts diff --git a/README.md b/README.md index b6ec5fd9..ce84e893 100644 --- a/README.md +++ b/README.md @@ -109,7 +109,7 @@ For an always-on exhaustive mode, use `/ultracode`; `/effort high` is the lighte ## Commands and run control -Pi can manage background runs directly with the `workflow_control` tool instead of asking you to type a command. It supports `list`, `status`, `pause`, `resume`, and `stop`; run-specific actions use the canonical run ID returned when the workflow starts. Status output includes the run state, current phase, agent counts, active labels, and recorded token total. +Pi can manage background runs directly with the `workflow_control` tool instead of asking you to type a command. It supports `list`, `status`, `pause`, `resume`, `stop`, `restart`, and `remove`; run-specific actions use the canonical run ID returned when the workflow starts. Status output includes the run state, current phase, agent counts, active labels, and recorded token total. `remove` accepts only terminal runs, so stop running or paused work first. | Command | Purpose | | --- | --- | @@ -186,6 +186,8 @@ Extension state lives outside the repository under `~/.pi/workflows`: Subagents are in-memory by default. Set `persistAgentSessions: true` to retain full transcripts in Pi's standard session directory. This creates one file per agent and may store sensitive material that an agent read, so enable it deliberately. +Run files use a versioned, backward-compatible state model with stable execution IDs, exact terminal usage, and crash-safe temp/rename writes plus backup recovery. After a crash, orphaned running work is recovered as paused and waits for an explicit resume. Resume replays the longest unchanged completed prefix—including nested workflows—then starts incomplete work fresh, without reopening child transcripts or double-charging completed usage. + Completed background runs persist their full result in the project run JSON. The conversation delivery includes a pointer to that file when the visible summary is shortened. diff --git a/src/agent.ts b/src/agent.ts index 7a8b4999..8841f472 100644 --- a/src/agent.ts +++ b/src/agent.ts @@ -3,6 +3,7 @@ import { unlinkSync, writeFileSync } from "node:fs"; import { join } from "node:path"; import type { AssistantMessage, Model, TextContent } from "@earendil-works/pi-ai"; import { + type AgentSessionEvent, AuthStorage, type CreateAgentSessionOptions, createAgentSession, @@ -300,7 +301,7 @@ function warnPersistSecretsOnce(sessionDir: string): void { ); } -/** Real token/cost usage for a single subagent run, read from the SDK session. */ +/** Token/cost usage for a single subagent run. */ export interface AgentUsage { input: number; output: number; @@ -308,6 +309,64 @@ export interface AgentUsage { cacheWrite: number; total: number; cost: number; + /** True only for an in-progress output-token estimate. */ + estimated?: boolean; +} + +/** + * Convert session events into absolute cumulative usage. Exact message usage is + * emitted at message_end; throttled message_update events add only a temporary + * output estimate, which the next exact event replaces. + */ +export function createAgentUsageEventHandler( + onUsage: (usage: AgentUsage) => void, + now: () => number = Date.now, +): (event: AgentSessionEvent) => void { + const exact: AgentUsage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0, cost: 0 }; + const endedMessages = new WeakSet(); + let lastEstimateEmit = Number.NEGATIVE_INFINITY; + const emit = (usage: AgentUsage) => { + try { + onUsage(usage); + } catch { + // Telemetry is best-effort; never interrupt the child session. + } + }; + + return (event) => { + if (event.type === "message_end" && event.message.role === "assistant") { + if (endedMessages.has(event.message)) return; + endedMessages.add(event.message); + const usage = event.message.usage; + exact.input += usage.input; + exact.output += usage.output; + exact.cacheRead += usage.cacheRead; + exact.cacheWrite += usage.cacheWrite; + exact.total += usage.totalTokens; + exact.cost += usage.cost.total; + lastEstimateEmit = Number.NEGATIVE_INFINITY; + if (exact.total > 0 || exact.cost > 0) emit({ ...exact, estimated: false }); + return; + } + + if (event.type !== "message_update" || event.message.role !== "assistant") return; + const timestamp = now(); + if (timestamp - lastEstimateEmit < 250) return; + lastEstimateEmit = timestamp; + const textLength = event.message.content.reduce( + (total, part) => + total + (part.type === "text" ? part.text.length : part.type === "thinking" ? part.thinking.length : 0), + 0, + ); + const estimatedOutput = Math.ceil(textLength / 4); + if (estimatedOutput <= 0) return; + emit({ + ...exact, + output: exact.output + estimatedOutput, + total: exact.total + estimatedOutput, + estimated: true, + }); + }; } /** @@ -346,10 +405,9 @@ export interface AgentRunOptions void; /** @@ -583,7 +641,7 @@ export class WorkflowAgent { } let removeAbortListener: (() => void) | undefined; - let removeHistoryListener: (() => void) | undefined; + let removeSessionListener: (() => void) | undefined; let lastHistoryEmit = 0; const emitHistory = () => options.onHistory?.(compactAgentHistory(session.messages)); const maybeEmitHistory = () => { @@ -593,6 +651,7 @@ export class WorkflowAgent { lastHistoryEmit = now; emitHistory(); }; + const handleUsageEvent = options.onUsage ? createAgentUsageEventHandler(options.onUsage) : undefined; try { if (options.signal?.aborted) throw new Error("Subagent was aborted"); if (options.signal) { @@ -600,8 +659,11 @@ export class WorkflowAgent { options.signal.addEventListener("abort", onAbort, { once: true }); removeAbortListener = () => options.signal?.removeEventListener("abort", onAbort); } - if (options.onHistory) { - removeHistoryListener = session.subscribe(() => maybeEmitHistory()); + if (options.onHistory || handleUsageEvent) { + removeSessionListener = session.subscribe((event) => { + maybeEmitHistory(); + handleUsageEvent?.(event); + }); } await session.prompt(this.buildPrompt(prompt, options as AgentRunOptions, Boolean(options.schema))); @@ -630,17 +692,17 @@ export class WorkflowAgent { return text as AgentRunResult; } finally { removeAbortListener?.(); - removeHistoryListener?.(); + removeSessionListener?.(); try { emitHistory(); } catch { // History is diagnostic only; never let it mask the real result/error. } - // Read real usage before disposing — dispose tears down the session state. + // Emit authoritative terminal usage before disposing the session state. if (options.onUsage) { try { const usage = usageFromStats(session.getSessionStats()); - if (usage) options.onUsage(usage); + if (usage) options.onUsage({ ...usage, estimated: false }); } catch { // Usage is best-effort; never let stats failure mask the real result/error. } diff --git a/src/display.ts b/src/display.ts index 848ae175..d8e3f360 100644 --- a/src/display.ts +++ b/src/display.ts @@ -4,10 +4,13 @@ import type { AgentHistoryEntry } from "./agent-history.js"; import type { WorkflowErrorCode } from "./errors.js"; import type { WorkflowMeta } from "./workflow.js"; -export type WorkflowAgentStatus = "queued" | "running" | "done" | "error" | "skipped"; +export type WorkflowAgentStatus = "queued" | "running" | "paused" | "done" | "error" | "skipped"; export interface WorkflowAgentSnapshot { id: number; + /** Stable invocation identity (`runId:callIndex`). */ + executionId?: string; + callIndex?: number; label: string; phase?: string; prompt: string; @@ -19,6 +22,10 @@ export interface WorkflowAgentSnapshot { history?: AgentHistoryEntry[]; /** Tokens used by this agent (a scalar estimate when the provider reports no usage). */ tokens?: number; + /** Whether tokens is a live streaming estimate. */ + tokensEstimated?: boolean; + /** Exact or provisional cumulative usage for this invocation. */ + usage?: AgentUsage; /** Per-agent token usage breakdown (fresh input+output vs cached), when known. */ tokenUsage?: AgentUsage; /** The model this agent ran on (provider/id), when known. */ @@ -247,7 +254,7 @@ const NO_THEME: ThemeLike = { fg: (_c, t) => t, bold: (t) => t }; /** The bracketed per-agent token cell (" [89 tok · 3,000 cached]"), or "" when nothing is known yet. */ function agentTokenCell(agent: WorkflowAgentSnapshot, theme: ThemeLike): string { const segment = fmtTokenSegment(tokenFigures(agent.tokenUsage, agent.tokens), fmtFull); - return segment ? theme.fg("dim", ` [${segment}]`) : ""; + return segment ? theme.fg("dim", ` [${agent.tokensEstimated ? "~" : ""}${segment}]`) : ""; } export function renderWorkflowLines( @@ -338,6 +345,8 @@ export function statusIcon(status: WorkflowAgentStatus): string { return "○"; case "running": return "●"; + case "paused": + return "⏸"; case "done": return "✓"; case "error": diff --git a/src/run-persistence.ts b/src/run-persistence.ts index e9a73291..d1c4ec8c 100644 --- a/src/run-persistence.ts +++ b/src/run-persistence.ts @@ -2,39 +2,63 @@ * Workflow run state persistence for pause/resume support. */ -import { existsSync, mkdirSync, readdirSync, readFileSync, renameSync, unlinkSync, writeFileSync } from "node:fs"; +import { + existsSync, + linkSync, + mkdirSync, + readdirSync, + readFileSync, + renameSync, + unlinkSync, + writeFileSync, +} from "node:fs"; import { join } from "node:path"; import type { AgentUsage } from "./agent.js"; import type { AgentHistoryEntry } from "./agent-history.js"; import type { WorkflowErrorCode } from "./errors.js"; +import type { JournalEntry } from "./workflow.js"; import { workflowProjectPaths } from "./workflow-paths.js"; +export const RUN_STATE_VERSION = 2; + export type RunStatus = "pending" | "running" | "paused" | "completed" | "failed" | "aborted"; export interface PersistedAgentState { id: number; + /** Stable runtime identity (`runId:callIndex`). Optional for legacy save callers. */ + executionId?: string; + callIndex?: number; label: string; phase?: string; prompt: string; - status: "queued" | "running" | "done" | "error" | "skipped"; + status: "queued" | "running" | "paused" | "done" | "error" | "skipped"; result?: unknown; + resultPreview?: string; error?: string; errorCode?: WorkflowErrorCode; recoverable?: boolean; history?: AgentHistoryEntry[]; + /** Exact cumulative usage. Streaming estimates are never persisted. */ + usage?: AgentUsage; + /** Legacy per-agent usage field retained for backward-compatible readers. */ + tokenUsage?: AgentUsage; + tokens?: number; startedAt?: string; endedAt?: string; - /** Tokens used by this agent (a scalar estimate when the provider reports no usage). */ - tokens?: number; - /** Per-agent token usage breakdown, when the provider reported one. */ - tokenUsage?: AgentUsage; /** The model this agent ran on (provider/id), when known. */ model?: string; } +export interface LoadedPersistedAgentState extends PersistedAgentState { + executionId: string; + callIndex: number; +} + export interface PersistedRunState { + version?: number; runId: string; workflowName: string; + workflowDescription?: string; script: string; args?: unknown; /** The pi session this run belongs to. Runs persist on disk across sessions but @@ -54,6 +78,12 @@ export interface PersistedRunState { updatedAt: string; completedAt?: string; durationMs?: number; + /** Run controls required to resume with equivalent execution behavior. */ + concurrency?: number; + maxAgents?: number; + agentRetries?: number; + agentTimeoutMs?: number | null; + tokenBudget?: number | null; tokenUsage?: { input: number; output: number; @@ -63,27 +93,24 @@ export interface PersistedRunState { cacheWrite?: number; }; /** Cached agent results for resume, keyed by deterministic call index. */ - journal?: Array<{ index: number; hash: string; result: unknown }>; - /** - * Opt-out of auto-resume for this run (default true, i.e. eligible unless - * explicitly set to false via ExecOptions.autoResume). Set once at run start - * and carried through resumes; see UsageLimitScheduler. - */ + journal?: JournalEntry[]; + /** Auto-resume eligibility for provider usage-limit pauses. Undefined means enabled. */ autoResume?: boolean; - /** - * Auto-resume attempt counter for the current usage_limit pause-cycle, owned - * and persisted by UsageLimitScheduler (best-effort). Absent/0 means no - * auto-resume attempt has been recorded yet. - */ + /** Best-effort scheduler attempt count for the current usage-limit pause cycle. */ autoResumeAttempts?: number; } +export interface LoadedPersistedRunState extends Omit { + version: number; + agents: LoadedPersistedAgentState[]; +} + export interface RunPersistence { /** Save current run state. */ save(state: PersistedRunState): void; - /** Load a persisted run by ID. */ + /** Load a persisted run by ID. Built-in persistence returns normalized state. */ load(runId: string): PersistedRunState | null; - /** List all persisted runs. */ + /** List persisted runs. Built-in persistence returns normalized state. */ list(): PersistedRunState[]; /** Delete a persisted run. */ delete(runId: string): boolean; @@ -103,6 +130,11 @@ export interface RunLease { token: string; } +interface NormalizedRunPersistence extends RunPersistence { + load(runId: string): LoadedPersistedRunState | null; + list(): LoadedPersistedRunState[]; +} + interface LockFile { runId: string; runPath: string; @@ -111,12 +143,219 @@ interface LockFile { token: string; } +const RUN_STATUSES = new Set(["pending", "running", "paused", "completed", "failed", "aborted"]); +const AGENT_STATUSES = new Set([ + "queued", + "running", + "paused", + "done", + "error", + "skipped", +]); +const warnedMigrationSources = new Set(); + +function warnMigrationOnce(source: string, message: string): void { + if (warnedMigrationSources.has(source)) return; + warnedMigrationSources.add(source); + console.warn(message); +} + +function finiteNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) ? value : undefined; +} + +function exactUsage(value: unknown): AgentUsage | undefined { + if (!value || typeof value !== "object" || Array.isArray(value)) return undefined; + const usage = value as Record; + return { + input: finiteNumber(usage.input) ?? 0, + output: finiteNumber(usage.output) ?? 0, + cacheRead: finiteNumber(usage.cacheRead) ?? 0, + cacheWrite: finiteNumber(usage.cacheWrite) ?? 0, + total: finiteNumber(usage.total) ?? 0, + cost: finiteNumber(usage.cost) ?? 0, + estimated: false, + }; +} + +function migrateAgent( + raw: unknown, + runId: string, + fallbackIndex: number, + legacyIdentity: boolean, +): LoadedPersistedAgentState | undefined { + if (!raw || typeof raw !== "object" || Array.isArray(raw)) return undefined; + const agent = raw as Record; + const persistedId = finiteNumber(agent.id); + const callIndex = Math.max(0, Math.floor(finiteNumber(agent.callIndex) ?? (persistedId ?? fallbackIndex + 1) - 1)); + const id = Math.max(1, Math.floor(persistedId ?? callIndex + 1)); + const status = AGENT_STATUSES.has(agent.status as PersistedAgentState["status"]) + ? (agent.status as PersistedAgentState["status"]) + : "queued"; + const usage = exactUsage(agent.usage ?? agent.tokenUsage); + return { + id, + executionId: !legacyIdentity && typeof agent.executionId === "string" ? agent.executionId : `${runId}:${callIndex}`, + callIndex, + label: typeof agent.label === "string" ? agent.label : `agent-${id}`, + phase: typeof agent.phase === "string" ? agent.phase : undefined, + prompt: typeof agent.prompt === "string" ? agent.prompt : "", + status, + result: agent.result, + resultPreview: typeof agent.resultPreview === "string" ? agent.resultPreview : undefined, + error: typeof agent.error === "string" ? agent.error : undefined, + errorCode: typeof agent.errorCode === "string" ? (agent.errorCode as WorkflowErrorCode) : undefined, + recoverable: typeof agent.recoverable === "boolean" ? agent.recoverable : undefined, + history: Array.isArray(agent.history) ? (agent.history as AgentHistoryEntry[]) : undefined, + usage, + tokens: finiteNumber(agent.tokens) ?? usage?.total, + startedAt: typeof agent.startedAt === "string" ? agent.startedAt : undefined, + endedAt: typeof agent.endedAt === "string" ? agent.endedAt : undefined, + model: typeof agent.model === "string" ? agent.model : undefined, + }; +} + +function migrateRunState(raw: unknown, source: string): LoadedPersistedRunState | null { + if (!raw || typeof raw !== "object" || Array.isArray(raw)) return null; + const state = raw as Record; + if (typeof state.runId !== "string" || !state.runId) return null; + const runId = state.runId; + const now = new Date().toISOString(); + const legacyIdentity = + state.version === undefined || finiteNumber(state.version) === undefined || Number(state.version) < 2; + let malformedNested = false; + const malformedUsage = (value: unknown): boolean => { + if (value === undefined) return false; + if (!value || typeof value !== "object" || Array.isArray(value)) return true; + const usage = value as Record; + return ["input", "output", "cacheRead", "cacheWrite", "total", "cost"].some( + (field) => usage[field] !== undefined && finiteNumber(usage[field]) === undefined, + ); + }; + const agents = Array.isArray(state.agents) + ? state.agents.flatMap((agent, index) => { + if (!agent || typeof agent !== "object" || Array.isArray(agent)) { + malformedNested = true; + return []; + } + const agentRecord = agent as Record; + if (malformedUsage(agentRecord.usage ?? agentRecord.tokenUsage)) malformedNested = true; + const migrated = migrateAgent(agent, runId, index, legacyIdentity); + return migrated ? [migrated] : []; + }) + : []; + const journal = Array.isArray(state.journal) + ? state.journal.flatMap((rawEntry) => { + if (!rawEntry || typeof rawEntry !== "object" || Array.isArray(rawEntry)) { + malformedNested = true; + return []; + } + const entry = rawEntry as Record; + const index = finiteNumber(entry.index); + if (index === undefined || typeof entry.hash !== "string") { + malformedNested = true; + return []; + } + if ( + (!legacyIdentity && typeof entry.executionId !== "string") || + malformedUsage(entry.usage) || + (entry.storeDelta !== undefined && + (!entry.storeDelta || typeof entry.storeDelta !== "object" || Array.isArray(entry.storeDelta))) + ) { + malformedNested = true; + } + const usage = exactUsage(entry.usage); + return [ + { + index: Math.floor(index), + executionId: + !legacyIdentity && typeof entry.executionId === "string" + ? entry.executionId + : `${runId}:${Math.floor(index)}`, + hash: entry.hash, + result: entry.result, + usage, + storeDelta: + entry.storeDelta && typeof entry.storeDelta === "object" && !Array.isArray(entry.storeDelta) + ? (entry.storeDelta as Record) + : undefined, + }, + ]; + }) + : undefined; + const rawTokenUsage = + state.tokenUsage && typeof state.tokenUsage === "object" && !Array.isArray(state.tokenUsage) + ? (state.tokenUsage as Record) + : undefined; + const tokenUsage = rawTokenUsage + ? { + input: finiteNumber(rawTokenUsage.input) ?? 0, + output: finiteNumber(rawTokenUsage.output) ?? 0, + total: finiteNumber(rawTokenUsage.total) ?? 0, + ...(finiteNumber(rawTokenUsage.cost) !== undefined ? { cost: finiteNumber(rawTokenUsage.cost) } : {}), + ...(finiteNumber(rawTokenUsage.cacheRead) !== undefined + ? { cacheRead: finiteNumber(rawTokenUsage.cacheRead) } + : {}), + ...(finiteNumber(rawTokenUsage.cacheWrite) !== undefined + ? { cacheWrite: finiteNumber(rawTokenUsage.cacheWrite) } + : {}), + } + : undefined; + const status = RUN_STATUSES.has(state.status as RunStatus) ? (state.status as RunStatus) : "paused"; + const malformed = + (state.phases !== undefined && !Array.isArray(state.phases)) || + (Array.isArray(state.phases) && state.phases.some((phase) => typeof phase !== "string")) || + (state.agents !== undefined && !Array.isArray(state.agents)) || + (state.logs !== undefined && !Array.isArray(state.logs)) || + (Array.isArray(state.logs) && state.logs.some((log) => typeof log !== "string")) || + malformedNested || + (state.tokenBudget !== undefined && state.tokenBudget !== null && finiteNumber(state.tokenBudget) === undefined); + if (state.version !== RUN_STATE_VERSION) { + warnMigrationOnce(source, `[run-persistence] Migrated legacy workflow run ${runId} from ${source}`); + } else if (malformed) { + warnMigrationOnce(source, `[run-persistence] Ignored malformed fields in workflow run ${runId} from ${source}`); + } + return { + version: RUN_STATE_VERSION, + runId, + workflowName: typeof state.workflowName === "string" ? state.workflowName : runId, + workflowDescription: typeof state.workflowDescription === "string" ? state.workflowDescription : undefined, + script: typeof state.script === "string" ? state.script : "", + args: state.args, + sessionId: typeof state.sessionId === "string" ? state.sessionId : undefined, + status, + pauseReason: typeof state.pauseReason === "string" ? state.pauseReason : undefined, + resetHint: typeof state.resetHint === "string" ? state.resetHint : undefined, + phases: Array.isArray(state.phases) + ? state.phases.filter((phase): phase is string => typeof phase === "string") + : [], + currentPhase: typeof state.currentPhase === "string" ? state.currentPhase : undefined, + agents, + logs: Array.isArray(state.logs) ? state.logs.filter((log): log is string => typeof log === "string") : [], + result: state.result, + startedAt: typeof state.startedAt === "string" ? state.startedAt : now, + updatedAt: typeof state.updatedAt === "string" ? state.updatedAt : now, + completedAt: typeof state.completedAt === "string" ? state.completedAt : undefined, + durationMs: finiteNumber(state.durationMs), + concurrency: finiteNumber(state.concurrency), + maxAgents: finiteNumber(state.maxAgents), + agentRetries: finiteNumber(state.agentRetries), + agentTimeoutMs: state.agentTimeoutMs === null ? null : finiteNumber(state.agentTimeoutMs), + tokenBudget: state.tokenBudget === null ? null : finiteNumber(state.tokenBudget), + tokenUsage, + journal, + autoResume: typeof state.autoResume === "boolean" ? state.autoResume : undefined, + autoResumeAttempts: finiteNumber(state.autoResumeAttempts), + }; +} + /** * Filesystem operations used by run persistence. * Exposed for testing – pass overrides to inject mock implementations. */ export type FsLayer = { existsSync: typeof existsSync; + linkSync: typeof linkSync; mkdirSync: typeof mkdirSync; readdirSync: typeof readdirSync; readFileSync: typeof readFileSync; @@ -125,8 +364,9 @@ export type FsLayer = { writeFileSync: typeof writeFileSync; }; -export function createRunPersistence(cwd: string, fsOverride?: Partial): RunPersistence { +export function createRunPersistence(cwd: string, fsOverride?: Partial): NormalizedRunPersistence { const _existsSync = fsOverride?.existsSync ?? existsSync; + const _linkSync = fsOverride?.linkSync ?? linkSync; const _mkdirSync = fsOverride?.mkdirSync ?? mkdirSync; const _readdirSync = fsOverride?.readdirSync ?? readdirSync; const _readFileSync = fsOverride?.readFileSync ?? readFileSync; @@ -185,9 +425,26 @@ export function createRunPersistence(cwd: string, fsOverride?: Partial) return true; }; + const loadRun = (runId: string): LoadedPersistedRunState | null => { + // Try the primary, then the .bak — so a corrupt primary doesn't lose the run. + for (const path of candidateRunPaths(runId)) { + for (const candidate of [path, `${path}.bak`]) { + try { + if (!_existsSync(candidate)) continue; + const migrated = migrateRunState(JSON.parse(_readFileSync(candidate, "utf-8")), candidate); + if (migrated) return migrated; + } catch { + // corrupt candidate -> fall through to the next candidate + } + } + } + return null; + }; + return { save(state: PersistedRunState) { ensureDir(); + state.version = RUN_STATE_VERSION; state.updatedAt = new Date().toISOString(); const path = primaryRunPath(state.runId); const json = JSON.stringify(state, null, 2); @@ -203,34 +460,25 @@ export function createRunPersistence(cwd: string, fsOverride?: Partial) } }, - load(runId: string): PersistedRunState | null { - // Try the primary, then the .bak — so a corrupt primary doesn't lose the run. - for (const path of candidateRunPaths(runId)) { - for (const candidate of [path, `${path}.bak`]) { - try { - if (!_existsSync(candidate)) continue; - return JSON.parse(_readFileSync(candidate, "utf-8")) as PersistedRunState; - } catch { - // corrupt candidate -> fall through to the next candidate - } - } - } - return null; - }, + load: loadRun, - list(): PersistedRunState[] { - const byRunId = new Map(); + list(): LoadedPersistedRunState[] { + const byRunId = new Map(); for (const dir of [runsDir, legacyRunsDir]) { try { if (!_existsSync(dir)) continue; - const files = _readdirSync(dir).filter((f) => f.endsWith(".json")); - for (const file of files) { - try { - const state = JSON.parse(_readFileSync(join(dir, file), "utf-8")) as PersistedRunState; - if (!byRunId.has(state.runId)) byRunId.set(state.runId, state); - } catch { - // Skip corrupted files - } + const runIds = new Set( + _readdirSync(dir).flatMap((file) => + file.endsWith(".json") + ? [file.slice(0, -".json".length)] + : file.endsWith(".json.bak") + ? [file.slice(0, -".json.bak".length)] + : [], + ), + ); + for (const runId of runIds) { + const state = loadRun(runId); + if (state && !byRunId.has(state.runId)) byRunId.set(state.runId, state); } } catch { // Skip unreadable directories; another storage location may still work. @@ -243,11 +491,14 @@ export function createRunPersistence(cwd: string, fsOverride?: Partial) let deleted = false; try { for (const path of candidateRunPaths(runId)) { - const dir = path === primaryRunPath(runId) ? runsDir : legacyRunsDir; - // Best-effort cleanup of the sidecar files alongside the primary. - for (const sidecar of [`${path}.bak`, `${path}.tmp`, lockPath(dir, runId)]) { + // Lease files are ownership records, not artifact sidecars. Only the + // lease owner may remove them through releaseRunLease(). + for (const sidecar of [`${path}.bak`, `${path}.tmp`]) { try { - if (_existsSync(sidecar)) _unlinkSync(sidecar); + if (_existsSync(sidecar)) { + _unlinkSync(sidecar); + deleted = true; + } } catch { // ignore sidecar cleanup failures } @@ -271,8 +522,22 @@ export function createRunPersistence(cwd: string, fsOverride?: Partial) ensureDir(); const path = primaryRunPath(runId); const lock = primaryLockPath(runId); + const takeover = `${lock}.takeover`; if (!removeStaleLegacyLock(runId)) return null; - for (let attempt = 0; attempt < 2; attempt++) { + for (let attempt = 0; attempt < 3; attempt++) { + if (_existsSync(takeover)) { + const claim = readLockAt(takeover); + if (!claim || pidIsAlive(claim.pid)) return null; + const currentClaim = readLockAt(takeover); + if (!currentClaim || currentClaim.token !== claim.token || pidIsAlive(currentClaim.pid)) return null; + try { + _unlinkSync(takeover); + } catch { + return null; + } + continue; + } + const token = `${process.pid}-${Date.now().toString(36)}-${Math.random().toString(36).slice(2)}`; const payload: LockFile = { runId, @@ -288,13 +553,33 @@ export function createRunPersistence(cwd: string, fsOverride?: Partial) const code = (err as { code?: string }).code; if (code !== "EEXIST") throw err; const existing = readLock(runId); - if (existing && existing.runPath === path && pidIsAlive(existing.pid)) { - return null; + if (existing && existing.runPath === path && pidIsAlive(existing.pid)) return null; + } + + const takeoverCandidate = `${takeover}.${token}.tmp`; + try { + _writeFileSync(takeoverCandidate, JSON.stringify(payload, null, 2), { flag: "wx" }); + _linkSync(takeoverCandidate, takeover); + } catch { + return null; + } finally { + try { + if (_existsSync(takeoverCandidate)) _unlinkSync(takeoverCandidate); + } catch { + // The fixed takeover claim, not its temporary source, is authoritative. } + } + try { + const existing = readLock(runId); + if (existing && existing.runPath === path && pidIsAlive(existing.pid)) return null; + if (_existsSync(lock)) _unlinkSync(lock); + _writeFileSync(lock, JSON.stringify(payload, null, 2), { flag: "wx" }); + return { runId, token }; + } finally { try { - _unlinkSync(lock); + if (readLockAt(takeover)?.token === token) _unlinkSync(takeover); } catch { - return null; + // A dead claimant is recovered by the next acquirer through its PID. } } } diff --git a/src/task-panel.ts b/src/task-panel.ts index 9611193f..d2f49be0 100644 --- a/src/task-panel.ts +++ b/src/task-panel.ts @@ -28,7 +28,10 @@ import { shortModel } from "./workflow-ui.js"; // as tokens accrue (not only on agent start/end). It is harmless in compact mode — // it redraws identical content. const RUN_EVENTS = [ + "agentQueued", "agentStart", + "agentUsage", + "agentHistory", "agentEnd", "phase", "log", @@ -335,7 +338,7 @@ function renderRunBody( const visible = phaseAgents.slice(-maxAgents); for (const a of visible) { const segment = fmtTokenSegment(tokenFigures(a.tokenUsage, a.tokens), fmtTokensShort); - const tok = segment ? dim(` ${segment}`) : ""; + const tok = segment ? dim(` ${a.tokensEstimated ? "~" : ""}${segment}`) : ""; const mdl = shortModel(a.model); const model = mdl ? dim(` · ${mdl}`) : ""; lines.push(` [${a.id}] ${statusIcon(a.status)} ${shorten(a.label, 40)}${tok}${model}`); diff --git a/src/workflow-control-tool.ts b/src/workflow-control-tool.ts index 5e76cdc6..183b3090 100644 --- a/src/workflow-control-tool.ts +++ b/src/workflow-control-tool.ts @@ -9,6 +9,8 @@ const runActionSchema = Type.Union([ Type.Literal("pause"), Type.Literal("resume"), Type.Literal("stop"), + Type.Literal("restart"), + Type.Literal("remove"), ]); const workflowControlSchema = Type.Union([ @@ -40,6 +42,7 @@ export interface WorkflowControlRunDetails { total: number; done: number; running: number; + paused: number; queued: number; error: number; skipped: number; @@ -61,7 +64,7 @@ export function createWorkflowControlTool( name: "workflow_control", label: "Workflow Control", description: - "List and inspect workflow runs, or pause, resume, and stop them without asking the user to run slash commands.", + "List and inspect workflow runs, or pause, resume, stop, restart, and remove them without asking the user to run slash commands.", promptSnippet: "Inspect and manage workflow runs directly by canonical run ID.", promptGuidelines: [ "Use workflow_control for workflow lifecycle management; do not ask the user to type /workflows when this tool can perform the action.", @@ -102,6 +105,43 @@ export function createWorkflowControlTool( case "stop": if (!manager.stop(run.runId)) return invalidTransition("stop", run); return actionSuccess("stop", "stopped", currentSummary(manager, run)); + case "restart": { + if (run.status === "running") { + return controlError("restart", run.runId, "cannot restart a running run", allowedActions(run.status)); + } + if (!run.script) { + return controlError("restart", run.runId, "run has no saved script", allowedActions(run.status)); + } + const restarted = manager.restart(run.runId); + if (!restarted) { + return controlError("restart", run.runId, "run could not be restarted", allowedActions(run.status)); + } + const next = findRun(manager, restarted.runId); + const suffix = next + ? formatRun(summarizeRun(next, manager.getSnapshot(next.runId))) + : `runId=${restarted.runId} status=running`; + return result(`action=restart result=restarted sourceRunId=${run.runId} ${suffix}`, { + action: "restart", + result: "restarted", + sourceRunId: run.runId, + runId: restarted.runId, + }); + } + case "remove": + if (run.status === "running" || run.status === "paused") { + return controlError("remove", run.runId, `cannot remove a ${run.status} run; stop it first`, [ + "status", + "stop", + ]); + } + if (!manager.deleteRun(run.runId)) { + return controlError("remove", run.runId, "run could not be removed", ["list"]); + } + return result(`action=remove result=removed runId=${run.runId}`, { + action: "remove", + result: "removed", + runId: run.runId, + }); } }, }); @@ -112,9 +152,9 @@ function normalizeInput(value: unknown): WorkflowControlInput { throw new Error("workflow_control requires an object argument"); } const input = value as Record; - const actions = new Set(["list", "status", "pause", "resume", "stop"]); + const actions = new Set(["list", "status", "pause", "resume", "stop", "restart", "remove"]); if (typeof input.action !== "string" || !actions.has(input.action)) { - throw new Error("workflow_control requires action: list|status|pause|resume|stop"); + throw new Error("workflow_control requires action: list|status|pause|resume|stop|restart|remove"); } const allowedKeys = input.action === "list" ? new Set(["action"]) : new Set(["action", "runId"]); @@ -164,13 +204,13 @@ function allowedActions(status: RunStatus): string[] { case "running": return ["status", "pause", "stop"]; case "paused": - return ["status", "resume", "stop"]; + return ["status", "resume", "stop", "restart"]; case "failed": case "pending": - return ["status", "resume"]; + return ["status", "resume", "restart", "remove"]; case "completed": case "aborted": - return ["status"]; + return ["status", "restart", "remove"]; } } @@ -200,6 +240,7 @@ function countAgents(agents: Array>): Work total: agents.length, done: agents.filter((agent) => agent.status === "done").length, running: agents.filter((agent) => agent.status === "running").length, + paused: agents.filter((agent) => agent.status === "paused").length, queued: agents.filter((agent) => agent.status === "queued").length, error: agents.filter((agent) => agent.status === "error").length, skipped: agents.filter((agent) => agent.status === "skipped").length, @@ -208,7 +249,7 @@ function countAgents(agents: Array>): Work function formatRun(run: WorkflowControlRunDetails): string { const active = run.activeLabels.join(",") || "-"; - return `runId=${run.runId} name=${quote(run.workflowName)} status=${run.status} phase=${quote(run.phase ?? "-")} total=${run.counts.total} done=${run.counts.done} running=${run.counts.running} queued=${run.counts.queued} error=${run.counts.error} skipped=${run.counts.skipped} active=${quote(active)} tokens=${run.tokenTotal}`; + return `runId=${run.runId} name=${quote(run.workflowName)} status=${run.status} phase=${quote(run.phase ?? "-")} total=${run.counts.total} done=${run.counts.done} running=${run.counts.running} paused=${run.counts.paused} queued=${run.counts.queued} error=${run.counts.error} skipped=${run.counts.skipped} active=${quote(active)} tokens=${run.tokenTotal}`; } function quote(value: string): string { diff --git a/src/workflow-manager.ts b/src/workflow-manager.ts index b511abc8..dfce3cca 100644 --- a/src/workflow-manager.ts +++ b/src/workflow-manager.ts @@ -4,18 +4,41 @@ import { EventEmitter } from "node:events"; import type { ModelRegistry } from "@earendil-works/pi-coding-agent"; -import type { WorkflowAgent } from "./agent.js"; -import { preview, type WorkflowSnapshot } from "./display.js"; +import type { AgentUsage, WorkflowAgent } from "./agent.js"; +import { MAX_AGENTS_PER_RUN } from "./config.js"; +import { preview, recomputeWorkflowSnapshot, type WorkflowSnapshot } from "./display.js"; import { WorkflowError, WorkflowErrorCode } from "./errors.js"; import { createRunPersistence, generateRunId, - type PersistedRunState, + type LoadedPersistedRunState, type RunLease, type RunPersistence, type RunStatus, } from "./run-persistence.js"; -import { type JournalEntry, parseWorkflowScript, runWorkflow, type WorkflowRunResult } from "./workflow.js"; +import { + type JournalEntry, + normalizeConcurrency, + parseWorkflowScript, + runWorkflow, + WORKFLOW_PAUSE_ABORT_REASON, + type WorkflowRunResult, +} from "./workflow.js"; + +interface DurableAgentMetadata { + result?: unknown; + usage?: AgentUsage; + startedAt?: string; + endedAt?: string; +} + +interface DurableExecutionControls { + concurrency: number; + maxAgents: number; + agentRetries: number; + agentTimeoutMs: number | null; + tokenBudget: number | null; +} export interface ManagedRun { runId: string; @@ -30,8 +53,22 @@ export interface ManagedRun { args?: unknown; /** Accumulated agent results for resume (deterministic call index -> result). */ journal: JournalEntry[]; + /** Serializable execution controls retained across resume. */ + execution: DurableExecutionControls; + /** Exact usage already consumed before a cold resume. */ + usageBaseline: NonNullable; + /** Result, exact usage, and timestamps not carried by the display snapshot. */ + durableAgents: Map; + sessionId?: string; + persistTimer?: ReturnType; + /** Current execution settlement, retained across pause so resume can serialize behind it. */ + executionPromise?: Promise; + /** Prevent a stopped run that is still unwinding from recreating its deleted artifact. */ + deleted?: boolean; /** Cross-process execution lease for this run, when it is actively executing. */ lease?: RunLease; + /** First required persistence failure; aborts execution and is surfaced to the caller. */ + durabilityError?: WorkflowError; /** * True when the run was started in the background (or resumed) and the caller is * not awaiting its result inline. Only background runs deliver their result back @@ -39,18 +76,14 @@ export interface ManagedRun { * result, so re-delivering would duplicate it. */ background: boolean; - /** - * Auto-resume eligibility for this run (see ExecOptions.autoResume). Set once - * at creation and carried through resume() so it survives pause/resume cycles. - * Undefined means eligible (default-on); false opts out. - */ + /** Undefined means eligible; false explicitly disables usage-limit auto-resume. */ autoResume?: boolean; } /** Per-execution options shared by sync, background, and resume runs. */ export interface ExecOptions { - /** Replay these journaled agent results for the unchanged prefix (resume). */ - resumeJournal?: Map; + /** Replay these journaled results for the unchanged prefix (resume). */ + resumeJournal?: ReadonlyMap; /** Cap on total agents for this run. */ maxAgents?: number; /** Per-agent timeout in milliseconds. null/omitted means no hard timeout. */ @@ -67,12 +100,7 @@ export interface ExecOptions { agentRetries?: number; /** Resolve a checkpoint() question with a human reply (only for UI-bearing runs). */ confirm?: (promptText: string, options: unknown) => Promise; - /** - * Whether this run is eligible for auto-resume when it pauses on a provider - * usage limit. Default-on: omit or pass true to stay eligible, pass false to - * opt out. Persisted on the run so a cold-start UsageLimitScheduler respects - * it too. See usage-limit-scheduler.ts. - */ + /** Whether provider usage-limit pauses are eligible for scheduler auto-resume. */ autoResume?: boolean; } @@ -97,16 +125,13 @@ export interface WorkflowManagerOptions { defaultAgentTimeoutMs?: number | null; /** Default retry attempts after recoverable agent failures. */ defaultAgentRetries?: number; - /** - * Persist each subagent transcript as a real pi session file under the - * standard sessions directory. Default false (in-memory, discarded). - */ + /** Persist subagent transcripts through the existing opt-in agent setting. Default false. */ persistAgentSessions?: boolean; } export class WorkflowManager extends EventEmitter { private runs = new Map(); - private persistence: RunPersistence; + private persistence: ReturnType; private cwd: string; private concurrency: number; private loadSavedWorkflow?: (name: string) => string | undefined; @@ -124,7 +149,7 @@ export class WorkflowManager extends EventEmitter { constructor(options: WorkflowManagerOptions = {}) { super(); this.cwd = options.cwd ?? process.cwd(); - this.concurrency = options.concurrency ?? 8; + this.concurrency = normalizeConcurrency(options.concurrency ?? 8); this.loadSavedWorkflow = options.loadSavedWorkflow; this.agent = options.agent; this.mainModel = options.mainModel; @@ -143,27 +168,42 @@ export class WorkflowManager extends EventEmitter { this.sessionId = id; } - /** - * On startup, any persisted run still marked "running" belongs to a process - * that died mid-run (this fresh manager has it nowhere in memory). Reconcile it - * to "paused" — never "failed" — so its journal is preserved and resume() can - * replay the completed prefix and finish the rest. - */ + /** Reconcile persisted work not owned by a live process into restart-safe rows. */ private recoverStaleRuns(): void { + let persistedRuns: LoadedPersistedRunState[]; try { - for (const p of this.listAllRuns()) { - if (p.status === "running" && !this.runs.has(p.runId)) { - const lease = this.persistence.acquireRunLease(p.runId); - if (!lease) continue; - try { - this.persistence.save({ ...p, status: "paused" }); - } finally { - this.persistence.releaseRunLease(lease); - } + persistedRuns = this.listAllRuns(); + } catch { + return; + } + + for (const persisted of persistedRuns) { + if (this.runs.has(persisted.runId)) continue; + let lease: RunLease | null = null; + try { + lease = this.persistence.acquireRunLease(persisted.runId); + if (!lease) continue; + const current = this.persistence.load(persisted.runId) ?? persisted; + const status = current.status === "running" ? "paused" : current.status; + const resumable = status === "pending" || status === "paused" || status === "failed"; + const agents = current.agents.map((agent) => { + if (agent.status !== "running") return agent; + return resumable + ? { ...agent, status: "paused" as const, endedAt: undefined } + : { + ...agent, + status: "skipped" as const, + endedAt: agent.endedAt ?? current.completedAt ?? current.updatedAt, + }; + }); + if (status !== current.status || agents.some((agent, index) => agent !== current.agents[index])) { + this.persistence.save({ ...current, status, agents }); } + } catch { + // Recovery is best-effort per run; one bad artifact must not block the rest. + } finally { + if (lease) this.persistence.releaseRunLease(lease); } - } catch { - // Recovery is best-effort; never let it block manager construction. } } @@ -178,8 +218,10 @@ export class WorkflowManager extends EventEmitter { } /** - * Expose the host session's model registry to integrations sharing this - * manager. Workflow execution reads the same registry internally. + * The host session's model registry, when set. Read lazily (e.g. by the + * workflow tool's model routing guideline) since `setModelRegistry` is called + * from `session_start`, which runs after the tool is created — a snapshot + * taken at tool-creation time would miss it. */ getModelRegistry(): ModelRegistry | undefined { return this.modelRegistry; @@ -206,6 +248,7 @@ export class WorkflowManager extends EventEmitter { const controller = new AbortController(); const lease = this.persistence.acquireRunLease(runId); if (!lease) throw new Error(`Could not acquire workflow run lease for ${runId}`); + const controls = this.resolveExecutionControls(exec); const managed: ManagedRun = { runId, @@ -226,29 +269,19 @@ export class WorkflowManager extends EventEmitter { script, args, journal: [], + execution: controls, + usageBaseline: this.zeroUsage(), + durableAgents: new Map(), + sessionId: this.sessionId, background: true, - lease, autoResume: exec.autoResume, + lease, }; this.runs.set(runId, managed); try { - // Persist initial state - this.persistence.save({ - runId, - workflowName: parsed.meta.name, - script, - args, - sessionId: this.sessionId, - status: "running", - phases: managed.snapshot.phases, - agents: [], - logs: [], - startedAt: managed.startedAt.toISOString(), - updatedAt: managed.startedAt.toISOString(), - autoResume: managed.autoResume, - }); + this.persistRun(managed, true); } catch (err) { this.releaseRunLease(managed); this.runs.delete(runId); @@ -260,7 +293,7 @@ export class WorkflowManager extends EventEmitter { // when a workflow is aborted/paused/stopped — executeRun()'s catch block // already records status/event/persist, but the promise still rejects. // The original promise is returned so callers can await it in try/catch. - const promise = this.executeRun(managed, script, args, exec); + const promise = this.beginExecution(managed, script, args, exec); promise.catch(() => {}); return { runId, promise }; @@ -273,20 +306,44 @@ export class WorkflowManager extends EventEmitter { * a caller (e.g. the workflow tool) drive its own inline display. */ async runSync(script: string, args?: unknown, exec: ExecOptions = {}): Promise { - const managed = this.createManaged(script, args); + const managed = this.createManaged(script, args, exec); const lease = this.persistence.acquireRunLease(managed.runId); if (!lease) throw new Error(`Could not acquire workflow run lease for ${managed.runId}`); managed.lease = lease; - managed.autoResume = exec.autoResume; this.runs.set(managed.runId, managed); // Persist the initial state immediately so listRuns()/the task panel can see // the run the moment it starts, not only after the first agent journals. - this.persistRun(managed); - return this.executeRun(managed, script, args, exec); + try { + this.persistRun(managed, true); + } catch (err) { + this.releaseRunLease(managed); + this.runs.delete(managed.runId); + throw err; + } + return this.beginExecution(managed, script, args, exec); + } + + private beginExecution( + managed: ManagedRun, + script: string, + args?: unknown, + exec: ExecOptions = {}, + ): Promise { + const promise = this.executeRun(managed, script, args, exec); + managed.executionPromise = promise; + void promise.then( + () => { + if (managed.executionPromise === promise) managed.executionPromise = undefined; + }, + () => { + if (managed.executionPromise === promise) managed.executionPromise = undefined; + }, + ); + return promise; } /** Build a fresh managed run with an empty snapshot. */ - private createManaged(script: string, args?: unknown): ManagedRun { + private createManaged(script: string, args?: unknown, exec: ExecOptions = {}): ManagedRun { const parsed = parseWorkflowScript(script); const slug = parsed.meta.name ? parsed.meta.name @@ -296,6 +353,7 @@ export class WorkflowManager extends EventEmitter { .slice(0, 40) || "workflow" : ""; const runId = slug ? `${slug}-${generateRunId()}` : generateRunId(); + const controls = this.resolveExecutionControls(exec); return { runId, status: "running", @@ -315,7 +373,38 @@ export class WorkflowManager extends EventEmitter { script, args, journal: [], + execution: controls, + usageBaseline: this.zeroUsage(), + durableAgents: new Map(), + sessionId: this.sessionId, background: false, + autoResume: exec.autoResume, + }; + } + + private zeroUsage(): NonNullable { + return { input: 0, output: 0, total: 0, cost: 0, cacheRead: 0, cacheWrite: 0 }; + } + + private addUsage(left: AgentUsage | undefined, right: AgentUsage): AgentUsage { + return { + input: (left?.input ?? 0) + right.input, + output: (left?.output ?? 0) + right.output, + cacheRead: (left?.cacheRead ?? 0) + right.cacheRead, + cacheWrite: (left?.cacheWrite ?? 0) + right.cacheWrite, + total: (left?.total ?? 0) + right.total, + cost: (left?.cost ?? 0) + right.cost, + estimated: right.estimated === true, + }; + } + + private resolveExecutionControls(exec: ExecOptions): DurableExecutionControls { + return { + concurrency: normalizeConcurrency(exec.concurrency ?? this.concurrency), + maxAgents: exec.maxAgents ?? MAX_AGENTS_PER_RUN, + agentRetries: exec.agentRetries ?? this.defaultAgentRetries, + agentTimeoutMs: exec.agentTimeoutMs !== undefined ? exec.agentTimeoutMs : this.defaultAgentTimeoutMs, + tokenBudget: exec.tokenBudget ?? null, }; } @@ -325,21 +414,38 @@ export class WorkflowManager extends EventEmitter { args?: unknown, exec: ExecOptions = {}, ): Promise { - const { - resumeJournal, - maxAgents, - agentTimeoutMs, - externalSignal, - onProgress, - tokenBudget, - concurrency, - agentRetries, - confirm, - } = exec; - const resolvedAgentTimeoutMs = agentTimeoutMs !== undefined ? agentTimeoutMs : this.defaultAgentTimeoutMs; - const resolvedConcurrency = concurrency ?? this.concurrency; - const resolvedAgentRetries = agentRetries ?? this.defaultAgentRetries; + const { resumeJournal, externalSignal, onProgress, confirm } = exec; + const controls = managed.execution; + const runtimeTokenBudget = + resumeJournal && controls.tokenBudget !== null + ? Math.max(0, controls.tokenBudget - managed.usageBaseline.total) + : controls.tokenBudget; const progress = () => onProgress?.(managed.snapshot); + const refresh = () => { + managed.snapshot = recomputeWorkflowSnapshot(managed.snapshot); + progress(); + }; + const agentsByExecutionId = new Map( + managed.snapshot.agents.flatMap((agent) => (agent.executionId ? [[agent.executionId, agent] as const] : [])), + ); + const findAgent = (executionId: string) => agentsByExecutionId.get(executionId); + const invocationBaselines = new Map( + managed.snapshot.agents.flatMap((agent) => { + const usage = + agent.executionId && agent.status !== "done" + ? managed.durableAgents.get(agent.executionId)?.usage + : undefined; + return agent.executionId && usage ? [[agent.executionId, { ...usage, estimated: false }] as const] : []; + }), + ); + const resumePhaseUsage = new Map(); + if (resumeJournal) { + for (const agent of managed.snapshot.agents) { + if (!agent.phase || !agent.executionId) continue; + const total = managed.durableAgents.get(agent.executionId)?.usage?.total ?? 0; + resumePhaseUsage.set(agent.phase, (resumePhaseUsage.get(agent.phase) ?? 0) + total); + } + } // Let a host abort (e.g. Esc during a blocking tool call) cancel this run. if (externalSignal) { if (externalSignal.aborted) managed.controller.abort(); @@ -359,20 +465,44 @@ export class WorkflowManager extends EventEmitter { modelRegistry: this.modelRegistry, persistAgentSessions: this.persistAgentSessions, signal: managed.controller.signal, - concurrency: resolvedConcurrency, - agentRetries: resolvedAgentRetries, - maxAgents, - agentTimeoutMs: resolvedAgentTimeoutMs, - tokenBudget, + concurrency: controls.concurrency, + agentRetries: controls.agentRetries, + maxAgents: controls.maxAgents, + agentTimeoutMs: controls.agentTimeoutMs, + tokenBudget: runtimeTokenBudget, confirm, loadSavedWorkflow: this.loadSavedWorkflow, resumeJournal, resumeFromRunId: resumeJournal ? managed.runId : undefined, + resumePhaseUsage, + onResumeMiss: (replayedExecutionIds) => { + if (!resumeJournal) return; + managed.journal = managed.journal.filter((entry) => + replayedExecutionIds.has(entry.executionId ?? `${managed.runId}:${entry.index}`), + ); + this.requirePersist(managed); + }, onAgentJournal: (entry) => { - // Append (crash-safe-ish): keep the latest entry per index, then persist. - managed.journal = managed.journal.filter((e) => e.index !== entry.index); - managed.journal.push(entry); - this.persistRun(managed); + const baseline = entry.executionId ? invocationBaselines.get(entry.executionId) : undefined; + const usage = entry.usage + ? entry.usage.estimated + ? baseline + : this.addUsage(baseline, entry.usage) + : baseline; + const reconciled = { ...entry, usage: usage ? { ...usage, estimated: false } : undefined }; + managed.journal = managed.journal.filter( + (existing) => + (existing.executionId ?? `${managed.runId}:${existing.index}`) !== + (reconciled.executionId ?? `${managed.runId}:${reconciled.index}`), + ); + managed.journal.push(reconciled); + managed.journal.sort( + (left, right) => + left.index - right.index || (left.executionId ?? "").localeCompare(right.executionId ?? ""), + ); + // Persist the completed result before the separate terminal-row callback; + // a crash between callbacks must replay rather than rerun and recharge it. + this.requirePersist(managed); }, onLog: (message) => { managed.snapshot.logs.push(message); @@ -387,74 +517,174 @@ export class WorkflowManager extends EventEmitter { this.emit("phase", { runId: managed.runId, title }); progress(); }, + onAgentQueued: (event) => { + let agent = findAgent(event.executionId); + if (!agent) { + agent = { + id: managed.snapshot.agents.length + 1, + executionId: event.executionId, + callIndex: event.callIndex, + label: event.label, + phase: event.phase, + prompt: event.prompt, + status: "queued", + model: event.model, + }; + managed.snapshot.agents.push(agent); + agentsByExecutionId.set(event.executionId, agent); + } else if (!event.replayed) { + agent.status = "queued"; + agent.label = event.label; + agent.phase = event.phase; + agent.prompt = event.prompt; + agent.resultPreview = undefined; + agent.error = undefined; + agent.errorCode = undefined; + agent.recoverable = undefined; + agent.tokensEstimated = false; + const metadata = managed.durableAgents.get(event.executionId) ?? {}; + metadata.result = undefined; + metadata.endedAt = undefined; + managed.durableAgents.set(event.executionId, metadata); + if (event.model) agent.model = event.model; + } + this.schedulePersist(managed); + this.emit("agentQueued", { runId: managed.runId, ...event }); + refresh(); + }, onAgentStart: (event) => { - managed.snapshot.agents.push({ - id: managed.snapshot.agents.length + 1, - label: event.label, - phase: event.phase, - prompt: event.prompt, - status: "running", - model: event.model, - }); + const agent = findAgent(event.executionId); + if (agent && !event.replayed) { + agent.status = "running"; + if (event.model) agent.model = event.model; + const metadata = managed.durableAgents.get(event.executionId) ?? {}; + metadata.startedAt = new Date().toISOString(); + metadata.endedAt = undefined; + managed.durableAgents.set(event.executionId, metadata); + } + this.schedulePersist(managed); this.emit("agentStart", { runId: managed.runId, ...event }); + refresh(); + }, + onAgentUsage: (event) => { + const usage = event.replayed + ? { ...event.usage, estimated: false } + : this.addUsage(invocationBaselines.get(event.executionId), event.usage); + const agent = findAgent(event.executionId); + if (agent) { + agent.usage = usage; + agent.tokenUsage = usage; + agent.tokens = usage.total; + agent.tokensEstimated = usage.estimated === true; + if (!usage.estimated) { + const metadata = managed.durableAgents.get(event.executionId) ?? {}; + metadata.usage = { ...usage, estimated: false }; + managed.durableAgents.set(event.executionId, metadata); + } + } + this.emit("agentUsage", { runId: managed.runId, ...event, usage }); progress(); }, onAgentEnd: (event) => { - const agent = [...managed.snapshot.agents] - .reverse() - .find((a) => a.label === event.label && a.status === "running"); + const agent = findAgent(event.executionId); + const baseline = invocationBaselines.get(event.executionId); + const usage = event.replayed + ? event.usage + ? { ...event.usage, estimated: false } + : baseline + : event.usage + ? this.addUsage(baseline, event.usage) + : baseline; if (agent) { - agent.status = event.result === null ? "error" : "done"; + const alreadyCompleted = agent.status === "done"; + const terminalStatus = managed.status === "aborted" ? "skipped" : event.status; + agent.status = terminalStatus; agent.resultPreview = preview(event.result); agent.error = event.error; agent.errorCode = event.errorCode; agent.recoverable = event.recoverable; - agent.tokens = event.tokens; - if (event.tokenUsage) agent.tokenUsage = event.tokenUsage; + if (!alreadyCompleted || usage) agent.tokens = usage?.total ?? event.tokens; + agent.tokensEstimated = usage?.estimated === true; + if (usage) { + agent.usage = usage; + agent.tokenUsage = usage; + } if (event.model) agent.model = event.model; + const metadata = managed.durableAgents.get(event.executionId) ?? {}; + metadata.result = event.result; + if (terminalStatus === "paused") metadata.endedAt = undefined; + else metadata.endedAt ??= new Date().toISOString(); + if (usage && !usage.estimated) metadata.usage = { ...usage, estimated: false }; + managed.durableAgents.set(event.executionId, metadata); } - this.emit("agentEnd", { runId: managed.runId, ...event }); - progress(); + this.requirePersist(managed); + this.emit("agentEnd", { runId: managed.runId, ...event, usage }); + refresh(); }, onAgentHistory: (event) => { - const agent = [...managed.snapshot.agents] - .reverse() - .find((a) => a.label === event.label && a.status === "running"); - if (agent) { - agent.history = event.history; - } + const agent = findAgent(event.executionId); + if (agent) agent.history = event.history; this.emit("agentHistory", { runId: managed.runId, ...event }); progress(); }, onTokenUsage: (usage) => { - managed.snapshot.tokenUsage = usage; - this.emit("tokenUsage", { runId: managed.runId, usage }); + const baseline = managed.usageBaseline; + managed.snapshot.tokenUsage = { + input: baseline.input + usage.input, + output: baseline.output + usage.output, + total: baseline.total + usage.total, + cost: (baseline.cost ?? 0) + usage.cost, + cacheRead: (baseline.cacheRead ?? 0) + (usage.cacheRead ?? 0), + cacheWrite: (baseline.cacheWrite ?? 0) + (usage.cacheWrite ?? 0), + }; + this.requirePersist(managed); + this.emit("tokenUsage", { runId: managed.runId, usage: managed.snapshot.tokenUsage }); progress(); }, }); managed.status = "completed"; - managed.result = result; - this.emit("complete", { runId: managed.runId, result }); - - // Persist final state - this.persistRun(managed); + const finalUsage = managed.snapshot.tokenUsage; + const completedResult: WorkflowRunResult = { + ...result, + tokenUsage: finalUsage + ? { + input: finalUsage.input, + output: finalUsage.output, + total: finalUsage.total, + cost: finalUsage.cost ?? 0, + cacheRead: finalUsage.cacheRead, + cacheWrite: finalUsage.cacheWrite, + } + : result.tokenUsage, + }; + managed.result = completedResult; + managed.snapshot.result = result.result; + managed.snapshot.durationMs = result.durationMs; + this.requirePersist(managed); this.releaseRunLease(managed); + this.emit("complete", { runId: managed.runId, result: managed.result }); - return result; + return completedResult; } catch (error) { + const durableError = managed.durabilityError; const workflowError = - error instanceof WorkflowError + durableError ?? + (error instanceof WorkflowError ? error : new WorkflowError( error instanceof Error ? error.message : String(error), WorkflowErrorCode.WORKFLOW_ABORTED, { recoverable: true }, - ); + )); + const lifecycleAbort = + managed.controller.signal.aborted && (managed.status === "paused" || managed.status === "aborted"); const usageLimitPaused = !managed.controller.signal.aborted && workflowError.code === WorkflowErrorCode.PROVIDER_USAGE_LIMIT; - if (managed.controller.signal.aborted) { + if (durableError) { + managed.status = "failed"; + } else if (managed.controller.signal.aborted) { // Intentional abort (pause/stop/Esc) — preserve status set by pause()/stop() if (managed.status === "running") { managed.status = "aborted"; @@ -475,12 +705,21 @@ export class WorkflowManager extends EventEmitter { error: workflowError, resetHint: workflowError.resetHint, }); - } else { + } else if (!lifecycleAbort && this.listenerCount("error") > 0) { + // pause()/stop() already emitted their lifecycle event. Only unexpected + // failures and external aborts reach the error channel. this.emit("error", { runId: managed.runId, error: workflowError }); } - // Persist final state - this.persistRun(managed); + // Persist final state. A durability failure cannot safely expose a stale, + // resumable artifact that would rerun already-paid work. + if (durableError) { + managed.deleted = true; + this.persistence.delete(managed.runId); + this.runs.delete(managed.runId); + } else { + this.persistRun(managed); + } this.releaseRunLease(managed); throw workflowError; @@ -493,24 +732,53 @@ export class WorkflowManager extends EventEmitter { managed.lease = undefined; } - private persistRun(managed: ManagedRun) { + private requirePersist(managed: ManagedRun): void { + try { + this.persistRun(managed, true); + } catch (error) { + const durabilityError = + error instanceof WorkflowError + ? error + : new WorkflowError( + error instanceof Error ? error.message : String(error), + WorkflowErrorCode.PERSISTENCE_ERROR, + { recoverable: false, details: error }, + ); + managed.durabilityError ??= durabilityError; + managed.controller.abort(durabilityError); + throw durabilityError; + } + } + + private schedulePersist(managed: ManagedRun): void { + if (managed.deleted || managed.persistTimer) return; + managed.persistTimer = setTimeout(() => this.persistRun(managed), 100); + } + + private persistRun(managed: ManagedRun, required = false): void { + if (managed.persistTimer) { + clearTimeout(managed.persistTimer); + managed.persistTimer = undefined; + } + if (managed.deleted) return; try { this.persistence.save({ runId: managed.runId, workflowName: managed.snapshot.name, - // Persist the real script + journal so the run can be resumed. Runs live - // in workflow run storage — protect via directory permissions, not blanking. + workflowDescription: managed.snapshot.description, script: managed.script, args: managed.args, - sessionId: this.sessionId, - journal: managed.journal, + sessionId: managed.sessionId, + journal: managed.journal.map((entry) => ({ + index: entry.index, + executionId: entry.executionId, + hash: entry.hash, + result: entry.result, + usage: entry.usage ? { ...entry.usage, estimated: false } : undefined, + storeDelta: entry.storeDelta, + })), status: managed.status, - // Persisted every write (not just at pause) so a stale read during the - // "paused" event race (see UsageLimitScheduler) is still correct — this - // is fixed at run-start and doesn't change over the run's lifetime. autoResume: managed.autoResume, - // Why a usage-limit pause happened, so the navigator / a future cold start - // can show it and (eventually) re-arm resume after the budget refills. pauseReason: managed.status === "paused" && managed.error?.code === WorkflowErrorCode.PROVIDER_USAGE_LIMIT ? "usage_limit" @@ -519,14 +787,33 @@ export class WorkflowManager extends EventEmitter { managed.status === "paused" && managed.error?.code === WorkflowErrorCode.PROVIDER_USAGE_LIMIT ? managed.error.resetHint : undefined, - phases: managed.snapshot.phases, + phases: [...managed.snapshot.phases], currentPhase: managed.snapshot.currentPhase, - agents: managed.snapshot.agents.map((a) => ({ - ...a, - startedAt: managed.startedAt.toISOString(), - endedAt: new Date().toISOString(), - })), - logs: managed.snapshot.logs, + agents: managed.snapshot.agents.map((agent) => { + const executionId = agent.executionId ?? `${managed.runId}:${agent.callIndex ?? agent.id - 1}`; + const metadata = managed.durableAgents.get(executionId); + return { + id: agent.id, + executionId, + callIndex: agent.callIndex ?? agent.id - 1, + label: agent.label, + phase: agent.phase, + prompt: agent.prompt, + status: agent.status, + result: metadata?.result, + resultPreview: agent.resultPreview, + error: agent.error, + errorCode: agent.errorCode, + recoverable: agent.recoverable, + history: agent.history, + usage: metadata?.usage ? { ...metadata.usage, estimated: false } : undefined, + tokens: metadata?.usage?.total ?? (agent.tokensEstimated ? undefined : agent.tokens), + startedAt: metadata?.startedAt, + endedAt: metadata?.endedAt, + model: agent.model, + }; + }), + logs: [...managed.snapshot.logs], result: managed.result?.result, tokenUsage: managed.snapshot.tokenUsage ? { @@ -542,11 +829,19 @@ export class WorkflowManager extends EventEmitter { updatedAt: new Date().toISOString(), completedAt: managed.status === "completed" ? new Date().toISOString() : undefined, durationMs: managed.result?.durationMs, + concurrency: managed.execution.concurrency, + maxAgents: managed.execution.maxAgents, + agentRetries: managed.execution.agentRetries, + agentTimeoutMs: managed.execution.agentTimeoutMs, + tokenBudget: managed.execution.tokenBudget, }); } catch (err) { - // Persistence is best-effort: the run is still healthy in memory. - // Log so an operator debugging state-loss has a lead, but never crash - // the workflow over a disk-full situation. + if (required) { + throw new WorkflowError(err instanceof Error ? err.message : String(err), WorkflowErrorCode.PERSISTENCE_ERROR, { + recoverable: false, + details: err, + }); + } console.warn("[workflow-manager] Persist run failed:", err); } } @@ -558,11 +853,14 @@ export class WorkflowManager extends EventEmitter { const managed = this.runs.get(runId); if (managed?.status !== "running") return false; - managed.controller.abort(); managed.status = "paused"; + for (const agent of managed.snapshot.agents) { + if (agent.status === "running") agent.status = "paused"; + } + managed.snapshot = recomputeWorkflowSnapshot(managed.snapshot); + managed.controller.abort(WORKFLOW_PAUSE_ABORT_REASON); this.emit("paused", { runId }); this.persistRun(managed); - this.releaseRunLease(managed); return true; } @@ -586,68 +884,256 @@ export class WorkflowManager extends EventEmitter { const active = this.runs.get(runId); if (active?.status === "running") return false; if (active?.status === "aborted") return false; + if (active?.executionPromise) await active.executionPromise.catch(() => {}); - const persisted = this.persistence.load(runId); - if (!persisted?.script || persisted.status === "completed" || persisted.status === "aborted") return false; + const settled = this.runs.get(runId); + if (settled?.status === "running" || settled?.status === "aborted") return false; + const preflight = this.persistence.load(runId); + if (!preflight?.script || preflight.status === "completed" || preflight.status === "aborted") return false; const lease = this.persistence.acquireRunLease(runId); if (!lease) return false; + const persisted = this.persistence.load(runId); + if (!persisted?.script || persisted.status === "completed" || persisted.status === "aborted") { + this.persistence.releaseRunLease(lease); + return false; + } // Use the edited script when supplied, else the persisted one (backward-compat). const script = opts?.script ?? persisted.script; const args = opts?.args !== undefined ? opts.args : persisted.args; - - const controller = new AbortController(); + const controls = this.resolveExecutionControls({ + concurrency: persisted.concurrency, + maxAgents: persisted.maxAgents, + agentRetries: persisted.agentRetries, + agentTimeoutMs: persisted.agentTimeoutMs !== undefined ? persisted.agentTimeoutMs : this.defaultAgentTimeoutMs, + tokenBudget: persisted.tokenBudget !== undefined ? persisted.tokenBudget : null, + autoResume: persisted.autoResume, + }); + const agents = persisted.agents.map((agent) => ({ + id: agent.id, + executionId: agent.executionId, + callIndex: agent.callIndex, + label: agent.label, + phase: agent.phase, + prompt: agent.prompt, + status: agent.status === "running" ? ("paused" as const) : agent.status, + resultPreview: agent.resultPreview ?? (agent.result !== undefined ? preview(agent.result) : undefined), + error: agent.error, + errorCode: agent.errorCode, + recoverable: agent.recoverable, + history: agent.history, + tokens: agent.usage?.total ?? agent.tokens, + tokensEstimated: false, + usage: agent.usage ? { ...agent.usage, estimated: false } : undefined, + tokenUsage: agent.usage ? { ...agent.usage, estimated: false } : undefined, + model: agent.model, + })); + const aggregateUsage = agents.reduce((total, agent) => { + const usage = agent.usage; + if (usage) { + total.input += usage.input; + total.output += usage.output; + total.total += usage.total; + total.cost = (total.cost ?? 0) + usage.cost; + total.cacheRead = (total.cacheRead ?? 0) + usage.cacheRead; + total.cacheWrite = (total.cacheWrite ?? 0) + usage.cacheWrite; + } else if (agent.status === "done" && agent.tokens !== undefined) { + // v1 rows carried only a terminal token total. + total.total += agent.tokens; + } + return total; + }, this.zeroUsage()); + const usageBaseline = persisted.tokenUsage ? { ...persisted.tokenUsage } : aggregateUsage; + const startedAt = new Date(persisted.startedAt); const managed: ManagedRun = { runId, status: "running", - snapshot: { + snapshot: recomputeWorkflowSnapshot({ name: persisted.workflowName, - phases: persisted.phases ?? [], - logs: persisted.logs ?? [], - agents: [], - agentCount: 0, + description: persisted.workflowDescription, + phases: [...persisted.phases], + currentPhase: persisted.currentPhase, + logs: [...persisted.logs], + agents, + agentCount: agents.length, runningCount: 0, doneCount: 0, errorCount: 0, - }, - controller, - startedAt: new Date(), + result: persisted.result, + durationMs: persisted.durationMs, + tokenUsage: usageBaseline, + runId, + }), + controller: new AbortController(), + startedAt: Number.isNaN(startedAt.getTime()) ? new Date() : startedAt, // The (possibly edited) script + args become the run's own — persistRun() - // writes them below, so a later resume of this run sees the edited script. + // writes them below, so a later resume sees the edited script. script, args, journal: persisted.journal ?? [], + execution: controls, + usageBaseline, + durableAgents: new Map( + persisted.agents.map((agent) => [ + agent.executionId, + { + result: agent.result, + usage: agent.usage ? { ...agent.usage, estimated: false } : undefined, + startedAt: agent.startedAt, + endedAt: agent.status === "running" || agent.status === "paused" ? undefined : agent.endedAt, + }, + ]), + ), + sessionId: persisted.sessionId, background: true, - lease, - // Carry the original opt-out forward across resumes; it's fixed at - // run-start and persistRun() re-persists it on every subsequent write. autoResume: persisted.autoResume, + lease, }; this.runs.set(runId, managed); - // Persist before notifying renderers: listRuns() is their source of truth for - // lifecycle status, while getRun() supplies the live in-memory snapshot. - this.persistRun(managed); + try { + this.persistRun(managed, true); + } catch (err) { + this.releaseRunLease(managed); + this.runs.delete(runId); + throw err; + } - const resumeJournal = new Map((persisted.journal ?? []).map((e) => [e.index, e] as const)); + const resumeJournal = new Map(managed.journal.map((entry) => [entry.executionId ?? entry.index, entry] as const)); this.emit("resumed", { runId }); // Run in the background; executeRun records status/errors on the managed run. - void this.executeRun(managed, script, args, { resumeJournal }).catch(() => {}); + void this.beginExecution(managed, script, args, { resumeJournal }).catch(() => {}); return true; } + /** Restart a saved run with the same durable execution controls. */ + restart(runId: string): { runId: string; promise: Promise } | null { + const active = this.runs.get(runId); + if (active?.status === "running" || active?.executionPromise) return null; + const sourceLease = this.persistence.acquireRunLease(runId); + if (!sourceLease) return null; + try { + const persisted = this.persistence.load(runId); + if (!persisted?.script || !["paused", "completed", "failed", "aborted"].includes(persisted.status)) return null; + if (persisted.status === "paused") { + const endedAt = new Date().toISOString(); + try { + this.persistence.save({ + ...persisted, + status: "aborted", + autoResume: false, + pauseReason: undefined, + resetHint: undefined, + agents: persisted.agents.map((agent) => + agent.status === "queued" || agent.status === "running" || agent.status === "paused" + ? { ...agent, status: "skipped" as const, endedAt: agent.endedAt ?? endedAt } + : agent, + ), + }); + } catch (error) { + throw new WorkflowError( + error instanceof Error ? error.message : String(error), + WorkflowErrorCode.PERSISTENCE_ERROR, + { recoverable: false, details: error }, + ); + } + if (active) { + active.status = "aborted"; + active.autoResume = false; + for (const agent of active.snapshot.agents) { + if (agent.status === "queued" || agent.status === "running" || agent.status === "paused") { + agent.status = "skipped"; + } + } + active.snapshot = recomputeWorkflowSnapshot(active.snapshot); + } + this.emit("stopped", { runId }); + } + return this.startInBackground(persisted.script, persisted.args, { + concurrency: persisted.concurrency, + maxAgents: persisted.maxAgents, + agentRetries: persisted.agentRetries, + agentTimeoutMs: persisted.agentTimeoutMs !== undefined ? persisted.agentTimeoutMs : this.defaultAgentTimeoutMs, + tokenBudget: persisted.tokenBudget !== undefined ? persisted.tokenBudget : null, + autoResume: persisted.autoResume, + }); + } finally { + this.persistence.releaseRunLease(sourceLease); + } + } + /** * Stop a running workflow. */ stop(runId: string): boolean { const managed = this.runs.get(runId); - if (!managed || (managed.status !== "running" && managed.status !== "paused")) return false; + if (managed) { + if (managed.status !== "running" && managed.status !== "paused") return false; + if (!managed.lease) { + managed.lease = this.persistence.acquireRunLease(runId) ?? undefined; + if (!managed.lease) return false; + } + const endedAt = new Date().toISOString(); + for (const agent of managed.snapshot.agents) { + if (agent.status !== "queued" && agent.status !== "running" && agent.status !== "paused") continue; + agent.status = "skipped"; + const executionId = agent.executionId; + if (executionId) { + const metadata = managed.durableAgents.get(executionId) ?? {}; + metadata.endedAt ??= endedAt; + managed.durableAgents.set(executionId, metadata); + } + } + managed.snapshot = recomputeWorkflowSnapshot(managed.snapshot); + managed.status = "aborted"; + managed.controller.abort(); + try { + this.requirePersist(managed); + } catch (error) { + managed.deleted = true; + this.persistence.delete(runId); + this.runs.delete(runId); + throw error; + } + this.emit("stopped", { runId }); + if (!managed.executionPromise) this.releaseRunLease(managed); + return true; + } - managed.controller.abort(); - managed.status = "aborted"; - this.emit("stopped", { runId }); - this.persistRun(managed); - this.releaseRunLease(managed); - return true; + const persisted = this.persistence.load(runId); + if (!persisted || (persisted.status !== "running" && persisted.status !== "paused")) return false; + const lease = this.persistence.acquireRunLease(runId); + if (!lease) return false; + try { + const current = this.persistence.load(runId); + if (!current || (current.status !== "running" && current.status !== "paused")) return false; + try { + this.persistence.save({ + ...current, + status: "aborted", + pauseReason: undefined, + resetHint: undefined, + agents: current.agents.map((agent) => + agent.status === "queued" || agent.status === "running" || agent.status === "paused" + ? { + ...agent, + status: "skipped" as const, + endedAt: agent.endedAt ?? new Date().toISOString(), + } + : agent, + ), + }); + } catch (error) { + throw new WorkflowError( + error instanceof Error ? error.message : String(error), + WorkflowErrorCode.PERSISTENCE_ERROR, + { recoverable: false, details: error }, + ); + } + this.emit("stopped", { runId }); + return true; + } finally { + this.persistence.releaseRunLease(lease); + } } /** @@ -665,13 +1151,13 @@ export class WorkflowManager extends EventEmitter { * that session's runs are returned — runs from other sessions stay on disk and * reappear when you switch back. Unbound (tests/legacy) returns everything. */ - listRuns(): PersistedRunState[] { + listRuns(): LoadedPersistedRunState[] { const all = this.persistence.list(); return this.sessionId ? all.filter((r) => r.sessionId === this.sessionId) : all; } /** All persisted runs regardless of session (used by cross-session recovery). */ - listAllRuns(): PersistedRunState[] { + listAllRuns(): LoadedPersistedRunState[] { return this.persistence.list(); } @@ -682,14 +1168,40 @@ export class WorkflowManager extends EventEmitter { return this.runs.get(runId)?.snapshot ?? null; } - /** - * Delete a persisted run. - */ + /** Delete a terminal persisted run. Running and paused work must be stopped first. */ deleteRun(runId: string): boolean { const managed = this.runs.get(runId); - if (managed) this.releaseRunLease(managed); - this.runs.delete(runId); - return this.persistence.delete(runId); + const persisted = this.persistence.load(runId); + const status = managed?.status ?? persisted?.status; + if (status === "running" || status === "paused" || (!managed && !persisted)) return false; + + const acquiredLease = managed?.lease ? null : this.persistence.acquireRunLease(runId); + const lease = managed?.lease ?? acquiredLease; + if (!lease) return false; + + try { + const current = this.persistence.load(runId); + const currentStatus = managed?.status ?? current?.status; + if (currentStatus === "running" || currentStatus === "paused" || (!managed && !current)) return false; + + if (managed) { + managed.deleted = true; + if (managed.persistTimer) { + clearTimeout(managed.persistTimer); + managed.persistTimer = undefined; + } + } + const deleted = this.persistence.delete(runId); + if (!deleted) { + if (managed) managed.deleted = false; + return false; + } + this.runs.delete(runId); + if (managed?.lease?.token === lease.token) this.releaseRunLease(managed); + return true; + } finally { + if (acquiredLease) this.persistence.releaseRunLease(acquiredLease); + } } /** diff --git a/src/workflow-ui.ts b/src/workflow-ui.ts index 3a51b597..d099f8a2 100644 --- a/src/workflow-ui.ts +++ b/src/workflow-ui.ts @@ -81,6 +81,7 @@ interface AgentRow { status: string; phase?: string; tokens?: number; + tokensEstimated?: boolean; tokenUsage?: AgentUsage; model?: string; } @@ -188,6 +189,7 @@ export class NavigatorModel { status: a.status, phase: a.phase, tokens: a.tokens, + tokensEstimated: a.tokensEstimated, tokenUsage: a.tokenUsage, model: a.model, })); @@ -215,6 +217,8 @@ function persistedToSnapshot(p: PersistedRunState): WorkflowSnapshot { logs: p.logs, agents: p.agents.map((a) => ({ id: a.id, + executionId: a.executionId, + callIndex: a.callIndex, label: a.label, phase: a.phase, prompt: a.prompt, @@ -225,8 +229,9 @@ function persistedToSnapshot(p: PersistedRunState): WorkflowSnapshot { errorCode: a.errorCode, recoverable: a.recoverable, history: a.history, - tokens: a.tokens, - tokenUsage: a.tokenUsage, + tokens: a.usage?.total ?? a.tokens, + usage: a.usage ?? a.tokenUsage, + tokenUsage: a.usage ?? a.tokenUsage, model: a.model, })), agentCount: p.agents.length, @@ -478,7 +483,8 @@ function rightAgentRow( theme: ThemeLike, ): string { const dotColor = AGENT_DOT_COLOR[a.status] ?? "dim"; - const stats = fmtTokenSegment(tokenFigures(a.tokenUsage, a.tokens), compactTokens); + const tokenSegment = fmtTokenSegment(tokenFigures(a.tokenUsage, a.tokens), compactTokens); + const stats = tokenSegment ? `${a.tokensEstimated ? "~" : ""}${tokenSegment}` : ""; const model = shortModel(a.model) ?? ""; // Stable 2-cell marker so columns never shift on selection: "› " | " ". @@ -1028,7 +1034,21 @@ export function openWorkflowNavigator( return ui.custom( (tui: TUI, theme: Theme, _keybindings, done: (r: undefined) => void) => { const rerender = () => tui.requestRender(); - const events = ["agentStart", "agentEnd", "phase", "log", "complete", "error", "stopped", "paused", "resumed"]; + const events = [ + "agentQueued", + "agentStart", + "agentUsage", + "agentHistory", + "agentEnd", + "tokenUsage", + "phase", + "log", + "complete", + "error", + "stopped", + "paused", + "resumed", + ]; const onEvent = () => rerender(); for (const ev of events) manager.on(ev, onEvent); const cleanup = () => { diff --git a/src/workflow.ts b/src/workflow.ts index bf9e2b18..c4ecbc59 100644 --- a/src/workflow.ts +++ b/src/workflow.ts @@ -20,6 +20,9 @@ import { parseModelRoutingFromMeta, resolveModelForPhase } from "./model-routing import { createAgentStoreTools, SharedStore } from "./shared-store.js"; import { createWorktree, removeWorktree, type Worktree } from "./worktree.js"; +/** AbortSignal reason used for a resumable workflow pause (distinct from stop/Esc). */ +export const WORKFLOW_PAUSE_ABORT_REASON = "workflow-paused"; + export interface WorkflowMetaPhase { title: string; detail?: string; @@ -37,9 +40,15 @@ export interface WorkflowMeta { /** One cached agent() result, keyed by its deterministic call index. */ export interface JournalEntry { index: number; + /** Stable invocation identity (`runId:index`). */ + executionId?: string; + /** Runtime-only origin; checkpoints have no following agent-end persistence callback. */ + source?: "agent" | "checkpoint"; /** sha256 of the call's identity (prompt + model + phase + agentType + schema). */ hash: string; result: unknown; + /** Exact cumulative usage for the completed invocation, when available. */ + usage?: AgentUsage; /** * Per-agent write delta (keys set by this agent) for additive replay on resume. * Replaces the former full-map snapshot to fix parallel-agent ordering: applying @@ -50,9 +59,8 @@ export interface JournalEntry { } /** - * Global resources shared across a run and any workflow() nested inside it, so - * the 16-concurrent / 1000-total caps and the token budget hold across nesting - * instead of each level getting its own limiter and counters. + * Global resources shared across a run and any nested workflow() call, so + * concurrency, maxAgents, and token budgets apply to the whole tree. */ export interface SharedRuntime { limiter: (fn: () => Promise) => Promise; @@ -60,6 +68,10 @@ export interface SharedRuntime { spent: number; tokenUsage: { input: number; output: number; total: number; cost: number; cacheRead: number; cacheWrite: number }; depth: number; + /** Scoped execution IDs replayed before the first cache miss. */ + replayedExecutionIds?: Set; + /** Shared resume boundary: after any parent/child miss, all later scoped calls run live. */ + resumeMissed?: boolean; } export interface WorkflowRunOptions extends WorkflowAgentOptions { @@ -87,12 +99,16 @@ export interface WorkflowRunOptions extends WorkflowAgentOptions { persistLogs?: boolean; /** Run ID for persistence. Auto-generated if not provided. */ runId?: string; - /** Resume: cached agent results keyed by deterministic call index. */ - resumeJournal?: Map; + /** Resume: cached results keyed by scoped execution ID (numeric keys remain supported for legacy callers). */ + resumeJournal?: ReadonlyMap; /** Resume: the run being resumed (informational; enables resume mode). */ resumeFromRunId?: string; /** Called after each live agent completes so the caller can persist the journal. */ onAgentJournal?: (entry: JournalEntry) => void; + /** Called once at the first replay miss with the execution IDs safe to retain. */ + onResumeMiss?: (replayedExecutionIds: ReadonlySet) => void; + /** Exact historical usage by phase, used to restore phase sub-budget gates. */ + resumePhaseUsage?: ReadonlyMap; /** Internal: shared runtime inherited by a nested workflow() call. */ sharedRuntime?: SharedRuntime; /** @@ -111,20 +127,58 @@ export interface WorkflowRunOptions extends WorkflowAgentOptions { confirm?: (promptText: string, options: CheckpointOptions) => Promise; onLog?: (message: string) => void; onPhase?: (title: string) => void; - onAgentStart?: (event: { label: string; phase?: string; prompt: string; model?: string }) => void; + onAgentQueued?: (event: { + executionId: string; + callIndex: number; + label: string; + phase?: string; + prompt: string; + model?: string; + /** True only when this callback represents a journal cache hit. */ + replayed?: boolean; + }) => void; + onAgentStart?: (event: { + executionId: string; + callIndex: number; + label: string; + phase?: string; + prompt: string; + model?: string; + /** True only when this callback represents a journal cache hit. */ + replayed?: boolean; + }) => void; onAgentEnd?: (event: { + executionId: string; + callIndex: number; label: string; phase?: string; result: unknown; + status: "done" | "error" | "skipped" | "paused"; tokens?: number; - tokenUsage?: AgentUsage; - worktree?: string; + usage?: AgentUsage; model?: string; error?: string; errorCode?: WorkflowErrorCode; recoverable?: boolean; + /** True only when this callback represents a journal cache hit. */ + replayed?: boolean; + }) => void; + onAgentHistory?: (event: { + executionId: string; + callIndex: number; + label: string; + phase?: string; + history: AgentHistoryEntry[]; + }) => void; + onAgentUsage?: (event: { + executionId: string; + callIndex: number; + label: string; + phase?: string; + usage: AgentUsage; + /** True only when this callback represents a journal cache hit. */ + replayed?: boolean; }) => void; - onAgentHistory?: (event: { label: string; phase?: string; history: AgentHistoryEntry[] }) => void; onTokenUsage?: (usage: { input: number; output: number; @@ -220,6 +274,8 @@ interface RuntimeState { * callIndex < firstMiss; once a call misses, it AND everything after run live. */ firstMiss: number; + /** Deterministic nested workflow invocation sequence within this script. */ + nestedSeq: number; } type AnyNode = Node & { [key: string]: any; start: number; end: number }; @@ -294,21 +350,31 @@ export async function runWorkflow( phaseBudgets: new Map(), callSeq: 0, firstMiss: Number.POSITIVE_INFINITY, + nestedSeq: 0, }; const agentRunner = options.agent ?? new WorkflowAgent(options); const concurrency = normalizeConcurrency( options.concurrency ?? Math.max(1, (globalThis.navigator?.hardwareConcurrency ?? 8) - 2), ); - // Global caps + budget are shared with any nested workflow() so they hold across nesting. + // Global controls are shared with any nested workflow() so they hold across nesting. const shared: SharedRuntime = options.sharedRuntime ?? { limiter: createLimiter(concurrency), agentCount: 0, spent: 0, tokenUsage: { input: 0, output: 0, total: 0, cost: 0, cacheRead: 0, cacheWrite: 0 }, depth: 0, + replayedExecutionIds: new Set(), + resumeMissed: false, }; const limiter = shared.limiter; + if (!shared.replayedExecutionIds) shared.replayedExecutionIds = new Set(); + const replayedExecutionIds = shared.replayedExecutionIds; + const markResumeMiss = () => { + if (shared.resumeMissed) return; + shared.resumeMissed = true; + options.onResumeMiss?.(new Set(replayedExecutionIds)); + }; // One store instance per run; nested workflow() calls inherit the parent's store // so all agents across nesting levels share the same key-value space. @@ -327,7 +393,12 @@ export async function runWorkflow( // Re-declaring re-bases from the current spent (idempotent across resume: the // script re-runs phase() and the ceiling is recomputed from live spent). if (typeof phaseOptions?.budget === "number" && phaseOptions.budget > 0) { - state.phaseBudgets.set(title, { budget: phaseOptions.budget, startSpent: shared.spent, warned: false }); + const historicalUsage = options.resumePhaseUsage?.get(title) ?? 0; + state.phaseBudgets.set(title, { + budget: phaseOptions.budget, + startSpent: shared.spent - historicalUsage, + warned: false, + }); } options.onPhase?.(title); }; @@ -344,6 +415,14 @@ export async function runWorkflow( } }; + const cachedJournalEntry = (executionId: string, callIndex: number): JournalEntry | undefined => { + const scoped = options.resumeJournal?.get(executionId); + if (scoped) return scoped; + // Numeric journals predate scoped nested execution IDs. Only the root may + // consume them; otherwise child index 0 can replay the root's index 0 entry. + return options.sharedRuntime ? undefined : options.resumeJournal?.get(callIndex); + }; + const agent = async (prompt: string, agentOptions: AgentOptions = {}) => { throwIfAborted(); @@ -356,36 +435,8 @@ export async function runWorkflow( ); } - if (budget.total !== null && budget.remaining() <= 0) { - throw new WorkflowError("workflow token budget exhausted", WorkflowErrorCode.TOKEN_BUDGET_EXHAUSTED, { - recoverable: false, - }); - } - const assignedPhase = agentOptions.phase ?? state.currentPhase; - // Per-phase soft sub-budget gate: a noisy phase can exhaust its own ceiling - // without touching the run's overall budget. Soft (spent accrues post-agent), - // warns once at ~80%, throws at 100%. Scripts can try/catch around a phase's - // work so later phases still proceed. - if (assignedPhase) { - const pb = state.phaseBudgets.get(assignedPhase); - if (pb) { - const phaseSpent = shared.spent - pb.startSpent; - if (phaseSpent >= pb.budget) { - throw new WorkflowError( - `phase "${assignedPhase}" token sub-budget exhausted (${pb.budget})`, - WorkflowErrorCode.TOKEN_BUDGET_EXHAUSTED, - { recoverable: false }, - ); - } - if (!pb.warned && phaseSpent >= pb.budget * 0.8) { - pb.warned = true; - log(`phase "${assignedPhase}" at ${Math.round((phaseSpent / pb.budget) * 100)}% of its token sub-budget`); - } - } - } - const requestedLabel = agentOptions.label?.trim(); // Resolve a named agentType to its bound definition (tools/model/prompt). @@ -416,45 +467,86 @@ export async function runWorkflow( // running nested-run agent can both get callIndex 0 and collide in // SharedStore.agentDeltas — whichever commits last steals/overwrites the // other's journaled delta. Composing the run's own runId (unique per - // top-level run AND per nested run, see `${runId}-nested${shared.depth}` - // below) with callIndex makes the key unique across the whole store. - const deltaKey = `${runId}:${callIndex}`; - - // Reserve the agent slot synchronously — atomic with the limit/budget gate - // above (no await in between) — so a parallel() fan-out can't all observe the - // same agentCount and overshoot maxAgents. (Token budget stays a soft gate: - // spent accrues after each agent, matching Claude Code; in-flight agents may - // push slightly past total, then further agent() calls throw.) - shared.agentCount++; - const label = requestedLabel || defaultAgentLabel(assignedPhase, shared.agentCount); - - // Longest-unchanged-prefix resume: replay a cached result only while the - // prefix is still intact — this call's index is before the first changed/new - // call. Once any call misses, it AND everything after it run live (matching - // Claude Code's contract), so an edited upstream call never leaves stale - // downstream results served from the journal. - const cached = options.resumeJournal?.get(callIndex); + // top-level and deterministic nested invocation) with callIndex makes the + // key unique across the whole store. + const executionId = `${runId}:${callIndex}`; + const deltaKey = executionId; + + const label = requestedLabel || defaultAgentLabel(assignedPhase, shared.agentCount + 1); + const invocation = { executionId, callIndex, label, phase: assignedPhase }; + + // Longest-unchanged-prefix resume: replay a cached result before live-work + // budget gates. Historical work is free even when the remaining budget is 0. + const cached = cachedJournalEntry(executionId, callIndex); const hashMatches = cached != null && cached.hash === callHash; const cachedEmptyOutput = hashMatches && isEmptyTextAgentResult(cached.result, agentOptions.schema); - if (hashMatches && !cachedEmptyOutput && callIndex < state.firstMiss) { - options.onAgentStart?.({ label, phase: assignedPhase, prompt, model: displayModel }); - options.onAgentEnd?.({ label, phase: assignedPhase, result: cached.result, tokens: 0, model: displayModel }); - // Apply this agent's write delta so live agents later in the run see a - // consistent store. Additive apply preserves parallel-agent writes that - // came from higher-callIndex agents finishing before this one. + if (hashMatches && !cachedEmptyOutput && !shared.resumeMissed && callIndex < state.firstMiss) { + shared.agentCount++; + replayedExecutionIds.add(executionId); + options.onAgentQueued?.({ ...invocation, prompt, model: displayModel, replayed: true }); + options.onAgentStart?.({ ...invocation, prompt, model: displayModel, replayed: true }); + if (cached.usage) + options.onAgentUsage?.({ ...invocation, usage: { ...cached.usage, estimated: false }, replayed: true }); + options.onAgentEnd?.({ + ...invocation, + result: cached.result, + status: "done", + tokens: cached.usage?.total ?? 0, + usage: cached.usage, + model: displayModel, + replayed: true, + }); if (cached.storeDelta) store.applyDelta(cached.storeDelta); return cached.result; } - // A genuine miss (no journal entry, or the hash changed) marks where the - // unchanged prefix ends; this call and every later one then run live. - if (!hashMatches || cachedEmptyOutput) state.firstMiss = Math.min(state.firstMiss, callIndex); + if (!hashMatches || cachedEmptyOutput) { + state.firstMiss = Math.min(state.firstMiss, callIndex); + markResumeMiss(); + } + + if (budget.total !== null && budget.remaining() <= 0) { + throw new WorkflowError("workflow token budget exhausted", WorkflowErrorCode.TOKEN_BUDGET_EXHAUSTED, { + recoverable: false, + }); + } + if (assignedPhase) { + const pb = state.phaseBudgets.get(assignedPhase); + if (pb) { + const phaseSpent = shared.spent - pb.startSpent; + if (phaseSpent >= pb.budget) { + throw new WorkflowError( + `phase "${assignedPhase}" token sub-budget exhausted (${pb.budget})`, + WorkflowErrorCode.TOKEN_BUDGET_EXHAUSTED, + { recoverable: false }, + ); + } + if (!pb.warned && phaseSpent >= pb.budget * 0.8) { + pb.warned = true; + log(`phase "${assignedPhase}" at ${Math.round((phaseSpent / pb.budget) * 100)}% of its token sub-budget`); + } + } + } + + // Reserve synchronously before the limiter so parallel fan-out cannot exceed maxAgents. + shared.agentCount++; + options.onAgentQueued?.({ + ...invocation, + prompt, + model: displayModel, + }); return limiter(async () => { + // Calls that never entered the limiter before pause remain visibly queued. + throwIfAborted(); const timeout = agentOptions.timeoutMs !== undefined ? agentOptions.timeoutMs : agentTimeoutMs; const retryAttempts = normalizeAgentRetries(agentOptions.retries ?? options.agentRetries ?? 0); const maxAttempts = retryAttempts + 1; - options.onAgentStart?.({ label, phase: assignedPhase, prompt, model: displayModel }); + options.onAgentStart?.({ + ...invocation, + prompt, + model: displayModel, + }); // Optional per-agent worktree isolation (deterministic name -> stable resume keys). // Precedence: explicit call-site isolation > agentDef isolation. @@ -469,27 +561,81 @@ export async function runWorkflow( } const runCwd = worktree?.isolated ? worktree.cwd : undefined; - // Captured from the subagent's real session usage; falls back to an - // estimate when the provider reports no usage (total === 0). Usage is reset - // per retry attempt so a failed attempt does not double-count the next one. - let usage: AgentUsage | undefined; - const recordTokens = (result: unknown): number => { - const tokens = usage && usage.total > 0 ? usage.total : estimateTokens(result) + estimateTokens(prompt); - if (usage) { - shared.tokenUsage.input += usage.input; - shared.tokenUsage.output += usage.output; - shared.tokenUsage.cost += usage.cost; - shared.tokenUsage.cacheRead += usage.cacheRead; - shared.tokenUsage.cacheWrite += usage.cacheWrite; + // Provider callbacks are absolute within an attempt. Only exact callbacks + // enter durable/run accounting; estimates are provisional display values. + const zeroUsage = (): AgentUsage => ({ + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + total: 0, + cost: 0, + estimated: false, + }); + const addUsage = (left: AgentUsage, right: AgentUsage): AgentUsage => ({ + input: left.input + right.input, + output: left.output + right.output, + cacheRead: left.cacheRead + right.cacheRead, + cacheWrite: left.cacheWrite + right.cacheWrite, + total: left.total + right.total, + cost: left.cost + right.cost, + estimated: left.estimated === true || right.estimated === true, + }); + let completedExactUsage = zeroUsage(); + let attemptExactUsage: AgentUsage | undefined; + let accountedUsage = zeroUsage(); + const emitUsage = (usage: AgentUsage) => options.onAgentUsage?.({ ...invocation, usage }); + const applyExactUsage = (usage: AgentUsage) => { + const input = usage.input - accountedUsage.input; + const output = usage.output - accountedUsage.output; + const cacheRead = usage.cacheRead - accountedUsage.cacheRead; + const cacheWrite = usage.cacheWrite - accountedUsage.cacheWrite; + const total = usage.total - accountedUsage.total; + const cost = usage.cost - accountedUsage.cost; + shared.tokenUsage.input += input; + shared.tokenUsage.output += output; + shared.tokenUsage.cacheRead += cacheRead; + shared.tokenUsage.cacheWrite += cacheWrite; + shared.tokenUsage.total += total; + shared.tokenUsage.cost += cost; + shared.spent += total; + accountedUsage = { ...usage, estimated: false }; + if (input || output || cacheRead || cacheWrite || total || cost) options.onTokenUsage?.(shared.tokenUsage); + }; + const receiveUsage = (usage: AgentUsage) => { + if (usage.estimated) { + emitUsage(addUsage(completedExactUsage, { ...usage, estimated: true })); + return; + } + attemptExactUsage = { ...usage, estimated: false }; + const cumulative = addUsage(completedExactUsage, attemptExactUsage); + emitUsage(cumulative); + applyExactUsage(cumulative); + }; + const finishAttempt = (result: unknown, allowEstimate = true): { display?: AgentUsage; exact?: AgentUsage } => { + if (attemptExactUsage) { + completedExactUsage = addUsage(completedExactUsage, attemptExactUsage); + completedExactUsage.estimated = false; + emitUsage(completedExactUsage); + applyExactUsage(completedExactUsage); + attemptExactUsage = undefined; + return { display: completedExactUsage, exact: completedExactUsage }; } - shared.tokenUsage.total += tokens; - shared.spent += tokens; - return tokens; + if (!allowEstimate) return { exact: completedExactUsage.total > 0 ? completedExactUsage : undefined }; + const display = addUsage(completedExactUsage, { + ...zeroUsage(), + total: estimateTokens(result) + estimateTokens(prompt), + estimated: true, + }); + emitUsage(display); + return { display, exact: completedExactUsage.total > 0 ? completedExactUsage : undefined }; }; + let activeAttempt = 0; try { for (let attempt = 1; attempt <= maxAttempts; attempt++) { - usage = undefined; + activeAttempt = attempt; + attemptExactUsage = undefined; try { throwIfAborted(); @@ -520,11 +666,11 @@ export async function runWorkflow( // Make the silent degrade visible in /workflows, not just console. log(`${label}: model "${spec}" unavailable — using the session default`); }, - onUsage: (u: AgentUsage) => { - usage = u; + onUsage: (usage: AgentUsage) => { + if (activeAttempt === attempt) receiveUsage(usage); }, onHistory: (history: AgentHistoryEntry[]) => { - options.onAgentHistory?.({ label, phase: assignedPhase, history }); + options.onAgentHistory?.({ ...invocation, history }); }, }), timeout, @@ -539,29 +685,49 @@ export async function runWorkflow( }); } - const tokens = recordTokens(result); + const usage = finishAttempt(result); + activeAttempt = 0; options.onAgentJournal?.({ index: callIndex, + executionId, + source: "agent", hash: callHash, result, + usage: usage.exact, storeDelta: store.commitDelta(deltaKey), }); options.onAgentEnd?.({ - label, - phase: assignedPhase, + ...invocation, result, - tokens, - tokenUsage: usage, - worktree: runCwd, + status: "done", + tokens: usage.display?.total ?? 0, + usage: usage.display, model: displayModel, }); return result; } catch (error) { - if (options.signal?.aborted) throw error; + if (options.signal?.aborted) { + const usage = finishAttempt(null, false); + activeAttempt = 0; + const paused = options.signal.reason === WORKFLOW_PAUSE_ABORT_REASON; + options.onAgentEnd?.({ + ...invocation, + result: null, + status: paused ? "paused" : "skipped", + tokens: usage.display?.total ?? usage.exact?.total ?? 0, + usage: usage.display ?? usage.exact, + model: displayModel, + error: paused ? "Subagent paused; incomplete work will restart on resume" : "Subagent was aborted", + errorCode: WorkflowErrorCode.WORKFLOW_ABORTED, + recoverable: true, + }); + throw error; + } const workflowError = wrapError(error, { agentLabel: label }); logger.error(`agent ${label} attempt ${attempt}/${maxAttempts} failed: ${workflowError.message}`); - const tokens = recordTokens(null); + const usage = finishAttempt(null); + activeAttempt = 0; if (workflowError.recoverable && attempt < maxAttempts) { log( @@ -570,19 +736,20 @@ export async function runWorkflow( continue; } + const providerPaused = workflowError.code === WorkflowErrorCode.PROVIDER_USAGE_LIMIT; options.onAgentEnd?.({ - label, - phase: assignedPhase, + ...invocation, result: null, - tokens, - tokenUsage: usage, - worktree: runCwd, + status: providerPaused ? "paused" : "error", + tokens: usage.display?.total ?? usage.exact?.total ?? 0, + usage: usage.display ?? usage.exact, model: displayModel, error: workflowError.message, errorCode: workflowError.code, recoverable: workflowError.recoverable, }); + if (providerPaused) throw workflowError; if (workflowError.recoverable) { log( `agent "${label}" exhausted ${maxAttempts} attempt${maxAttempts === 1 ? "" : "s"}: ${workflowError.code} ${workflowError.message}`, @@ -594,7 +761,6 @@ export async function runWorkflow( } return null; } finally { - // Always tear down the worktree, even on timeout/abort. if (worktree?.isolated) await removeWorktree(worktree); } }); @@ -666,6 +832,7 @@ export async function runWorkflow( } const resolved = options.loadSavedWorkflow?.(String(nameOrScript)); const childScript = resolved ?? String(nameOrScript); + const childRunId = `${runId}-nested${state.nestedSeq++}`; shared.depth++; try { const child = await runWorkflow(childScript, { @@ -674,10 +841,10 @@ export async function runWorkflow( sharedRuntime: shared, // Propagate the parent's store so nested agents share the same key-value space. sharedStore: store, - // A nested run is its own script; never reuse the parent's resume journal. - resumeJournal: undefined, - resumeFromRunId: undefined, - runId: `${runId}-nested${shared.depth}`, + // Scoped execution IDs let parent and child replay from one durable journal. + resumeJournal: options.resumeJournal, + resumeFromRunId: options.resumeFromRunId, + runId: childRunId, persistLogs: false, }); return child.result; @@ -856,13 +1023,18 @@ export async function runWorkflow( ); } const callIndex = state.callSeq++; + const executionId = `${runId}:${callIndex}`; const callHash = hashCheckpoint(promptText, checkpointOptions); - const cached = options.resumeJournal?.get(callIndex); - if (cached != null && cached.hash === callHash && callIndex < state.firstMiss) { + const cached = cachedJournalEntry(executionId, callIndex); + if (cached != null && cached.hash === callHash && !shared.resumeMissed && callIndex < state.firstMiss) { shared.agentCount++; + replayedExecutionIds.add(executionId); return cached.result; // replay the journaled human reply } - if (cached == null || cached.hash !== callHash) state.firstMiss = Math.min(state.firstMiss, callIndex); + if (cached == null || cached.hash !== callHash) { + state.firstMiss = Math.min(state.firstMiss, callIndex); + markResumeMiss(); + } shared.agentCount++; let reply: unknown; @@ -878,7 +1050,13 @@ export async function runWorkflow( reply = checkpointOptions.default ?? true; } throwIfAborted(); - options.onAgentJournal?.({ index: callIndex, hash: callHash, result: reply }); + options.onAgentJournal?.({ + index: callIndex, + executionId, + source: "checkpoint", + hash: callHash, + result: reply, + }); return reply; }; @@ -1146,7 +1324,7 @@ function estimateTokens(value: unknown): number { return Math.ceil(JSON.stringify(value ?? "").length / 4); } -function normalizeConcurrency(value: unknown): number { +export function normalizeConcurrency(value: unknown): number { if (typeof value !== "number" || !Number.isFinite(value) || value < 1) return 1; return Math.min(MAX_CONCURRENCY, Math.floor(value)); } diff --git a/tests/agent.test.ts b/tests/agent.test.ts index e1bdb124..d08cf81a 100644 --- a/tests/agent.test.ts +++ b/tests/agent.test.ts @@ -4,7 +4,13 @@ import { tmpdir } from "node:os"; import { dirname, join } from "node:path"; import test from "node:test"; import type { AgentRunOptions, AgentUsage } from "../src/agent.js"; -import { listAvailableModelSpecs, resolveAgentModelSpec, usageFromStats, WorkflowAgent } from "../src/agent.js"; +import { + createAgentUsageEventHandler, + listAvailableModelSpecs, + resolveAgentModelSpec, + usageFromStats, + WorkflowAgent, +} from "../src/agent.js"; import { WorkflowError, WorkflowErrorCode } from "../src/errors.js"; import { resolveModelSpecWithThinking } from "../src/model-spec.js"; import type { ModelTierConfig } from "../src/model-tier-config.js"; @@ -388,6 +394,104 @@ test("lastAssistantText picks the last assistant message, not first", () => { assert.equal(text, "final"); }); +test("usage events emit throttled estimates then exact cumulative message usage", () => { + const events: AgentUsage[] = []; + let timestamp = 0; + const handle = createAgentUsageEventHandler( + (usage) => events.push(usage), + () => timestamp, + ); + const updatingMessage = { + role: "assistant", + content: [{ type: "text", text: "12345678" }], + }; + + handle({ type: "message_update", message: updatingMessage, assistantMessageEvent: {} } as never); + timestamp = 100; + updatingMessage.content[0].text = "123456789012"; + handle({ type: "message_update", message: updatingMessage, assistantMessageEvent: {} } as never); + timestamp = 250; + handle({ type: "message_update", message: updatingMessage, assistantMessageEvent: {} } as never); + + const endedMessage = { + role: "assistant", + content: [{ type: "text", text: "done" }], + usage: { + input: 7, + output: 3, + cacheRead: 2, + cacheWrite: 1, + totalTokens: 13, + cost: { input: 0.01, output: 0.02, cacheRead: 0, cacheWrite: 0, total: 0.03 }, + }, + }; + handle({ type: "message_end", message: endedMessage } as never); + handle({ type: "message_end", message: endedMessage } as never); + + assert.deepEqual( + events.map((usage) => ({ total: usage.total, output: usage.output, estimated: usage.estimated })), + [ + { total: 2, output: 2, estimated: true }, + { total: 3, output: 3, estimated: true }, + { total: 13, output: 3, estimated: false }, + ], + ); + assert.equal(events[2].input, 7); + assert.equal(events[2].cacheRead, 2); + assert.equal(events[2].cost, 0.03); +}); + +test("all-zero message usage keeps the streaming estimate fallback", () => { + const events: AgentUsage[] = []; + const handle = createAgentUsageEventHandler((usage) => events.push(usage)); + const message = { + role: "assistant", + content: [{ type: "text", text: "12345678" }], + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }; + + handle({ type: "message_update", message, assistantMessageEvent: {} } as never); + handle({ type: "message_end", message } as never); + + assert.deepEqual( + events.map((usage) => usage.estimated), + [true], + ); + assert.ok(events[0].total > 0); +}); + +test("usage message_end events accumulate across assistant turns", () => { + const events: AgentUsage[] = []; + const handle = createAgentUsageEventHandler((usage) => events.push(usage)); + const message = (input: number, output: number) => ({ + role: "assistant", + content: [{ type: "text", text: "done" }], + usage: { + input, + output, + cacheRead: 0, + cacheWrite: 0, + totalTokens: input + output, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0.001 }, + }, + }); + + handle({ type: "message_end", message: message(10, 4) } as never); + handle({ type: "message_end", message: message(6, 2) } as never); + + assert.equal(events.at(-1)?.input, 16); + assert.equal(events.at(-1)?.output, 6); + assert.equal(events.at(-1)?.total, 22); + assert.equal(events.at(-1)?.cost, 0.002); +}); + // ═══════════════════════════════════════════════════════════════════════════ // Full agent() pipeline inside runWorkflow — verifies the agent() function // in workflow.ts correctly invokes the runner with all options. @@ -638,8 +742,8 @@ test("agent() in workflow fires onTokenUsage after run", async () => { onTokenUsage: (u) => usageEvents.push({ input: u.input, output: u.output, total: u.total }), }, ); - assert.equal(usageEvents.length, 1, "should fire onTokenUsage once"); - assert.equal(usageEvents[0].total, 30, "should accumulate from agent usage"); + assert.ok(usageEvents.length >= 1, "should fire onTokenUsage during the run"); + assert.equal(usageEvents.at(-1)?.total, 30, "should finish with accumulated agent usage"); }); test("agent() passes onModelResolved callback for display model updates", async () => { @@ -673,8 +777,8 @@ test("agent() accumulates usage across multiple agents", async () => { onTokenUsage: (u) => usageEvents.push({ total: u.total }), }, ); - assert.equal(usageEvents.length, 1, "one final usage event"); - assert.equal(usageEvents[0].total, 60, "two agents × 30 tokens each"); + assert.ok(usageEvents.length >= 2, "usage should update as each agent reports"); + assert.equal(usageEvents.at(-1)?.total, 60, "two agents × 30 tokens each"); }); test("agent() with timeout should handle gracefully (timeout returns null)", async () => { diff --git a/tests/run-persistence.test.ts b/tests/run-persistence.test.ts index af56a426..8c34af1f 100644 --- a/tests/run-persistence.test.ts +++ b/tests/run-persistence.test.ts @@ -4,7 +4,13 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import test from "node:test"; import { WORKFLOW_RUNS_DIR } from "../src/config.js"; -import { createRunPersistence, generateRunId, type PersistedRunState } from "../src/run-persistence.js"; +import { + createRunPersistence, + generateRunId, + type PersistedRunState, + RUN_STATE_VERSION, + type RunPersistence, +} from "../src/run-persistence.js"; import { WorkflowManager } from "../src/workflow-manager.js"; import { workflowProjectPaths } from "../src/workflow-paths.js"; import { withFakeHomeAsync } from "./helpers/fake-home.js"; @@ -359,6 +365,157 @@ test( }), ); +test( + "save accepts legacy agent rows without identity and load returns normalized identity", + withTempCwd(async (cwd) => { + const persistence = createRunPersistence(cwd); + persistence.save({ + runId: "legacy-save-shape", + workflowName: "legacy", + script: "return 1", + status: "paused", + phases: [], + agents: [{ id: 1, label: "old", prompt: "work", status: "done", tokens: 4 }], + logs: [], + startedAt: "2024-01-01T00:00:00.000Z", + updatedAt: "2024-01-01T00:00:00.000Z", + }); + + const loaded = persistence.load("legacy-save-shape"); + assert.equal(loaded?.agents[0].executionId, "legacy-save-shape:0"); + assert.equal(loaded?.agents[0].callIndex, 0); + }), +); + +test("RunPersistence remains compatible with legacy PersistedRunState return types", () => { + const state = { + runId: "legacy-implementation", + workflowName: "legacy", + script: "return 1", + status: "completed", + phases: [], + agents: [], + logs: [], + startedAt: "2024-01-01T00:00:00.000Z", + updatedAt: "2024-01-01T00:00:00.000Z", + } satisfies PersistedRunState; + const legacyImplementation: RunPersistence = { + save() {}, + load: () => state, + list: () => [state], + delete: () => false, + acquireRunLease: () => null, + releaseRunLease() {}, + getRunsDir: () => ".", + }; + + assert.equal(legacyImplementation.load(state.runId)?.runId, state.runId); +}); + +test( + "legacy run migration fills stable identity and ignores malformed optional fields", + withTempCwd(async (cwd) => { + const runsDir = workflowProjectPaths(cwd).runsDir; + mkdirSync(runsDir, { recursive: true }); + writeFileSync( + join(runsDir, "legacy-v1.json"), + JSON.stringify({ + runId: "legacy-v1", + workflowName: "legacy", + script: "return 1", + status: "paused", + phases: ["Work", 7], + agents: [ + { + id: 1, + label: "old agent", + prompt: "work", + status: "done", + tokens: 12, + usage: { input: 8, output: 4, total: 12, cost: 0.01 }, + }, + ], + logs: ["ok", null], + tokenBudget: "bad", + journal: [{ index: 0, hash: "hash", result: "done", usage: { input: 8, output: 4, total: 12 } }], + startedAt: "2024-01-01T00:00:00.000Z", + updatedAt: "2024-01-01T00:01:00.000Z", + }), + ); + + const loaded = createRunPersistence(cwd).load("legacy-v1"); + assert.equal(loaded?.version, RUN_STATE_VERSION); + assert.deepEqual(loaded?.phases, ["Work"]); + assert.deepEqual(loaded?.logs, ["ok"]); + assert.equal(loaded?.agents[0].executionId, "legacy-v1:0"); + assert.equal(loaded?.agents[0].callIndex, 0); + assert.equal(loaded?.agents[0].usage?.total, 12); + assert.equal(loaded?.journal?.[0].executionId, "legacy-v1:0"); + assert.equal(loaded?.journal?.[0].usage?.total, 12); + assert.equal(loaded?.tokenBudget, undefined); + }), +); + +test( + "current loader warns when malformed nested usage and journal fields are normalized or dropped", + withTempCwd(async (cwd) => { + const runsDir = workflowProjectPaths(cwd).runsDir; + mkdirSync(runsDir, { recursive: true }); + writeFileSync( + join(runsDir, "malformed-nested.json"), + JSON.stringify({ + version: RUN_STATE_VERSION, + runId: "malformed-nested", + workflowName: "malformed", + script: "return 1", + status: "paused", + phases: [], + agents: [ + { + id: 1, + executionId: "malformed-nested:0", + callIndex: 0, + label: "agent", + prompt: "work", + status: "done", + usage: { input: "bad", output: 2, total: 2 }, + }, + ], + logs: [], + journal: [ + null, + { + index: 0, + executionId: 42, + hash: "hash", + result: "done", + usage: { total: "bad" }, + storeDelta: [], + }, + ], + startedAt: "2024-01-01T00:00:00.000Z", + updatedAt: "2024-01-01T00:01:00.000Z", + }), + ); + + const warnings: string[] = []; + const originalWarn = console.warn; + console.warn = (...args: unknown[]) => warnings.push(args.map(String).join(" ")); + try { + const loaded = createRunPersistence(cwd).load("malformed-nested"); + assert.equal(loaded?.agents[0].usage?.input, 0); + assert.equal(loaded?.journal?.length, 1); + assert.equal(loaded?.journal?.[0].executionId, "malformed-nested:0"); + assert.equal(loaded?.journal?.[0].usage?.total, 0); + assert.equal(loaded?.journal?.[0].storeDelta, undefined); + assert.equal(warnings.length, 1); + assert.match(warnings[0], /Ignored malformed fields.*malformed-nested/); + } finally { + console.warn = originalWarn; + } + }), +); + test("generateRunId returns a string with timestamp and random parts", () => { const id = generateRunId(); assert.equal(typeof id, "string"); @@ -483,7 +640,16 @@ test( assert.equal(loaded.agents[1].status, "running"); assert.deepEqual(loaded.logs, ["started", "phase: Scan", "phase: Analyze"]); assert.deepEqual(loaded.tokenUsage, { input: 500, output: 200, total: 700 }); - assert.deepEqual(loaded.journal, [{ index: 0, hash: "abc", result: { ok: true } }]); + assert.deepEqual(loaded.journal, [ + { + index: 0, + executionId: "concurrent-test:0", + hash: "abc", + result: { ok: true }, + usage: undefined, + storeDelta: undefined, + }, + ]); }), ); @@ -524,6 +690,31 @@ test( writeFileSync(join(workflowProjectPaths(cwd).runsDir, "r1.json"), "{ truncated", "utf-8"); const loaded = rp.load("r1"); assert.equal(loaded?.runId, "r1", "load falls back to the intact .bak"); + assert.deepEqual( + rp.list().map((run) => run.runId), + ["r1"], + "list discovers the backup so startup recovery can reconcile the run", + ); + }), +); + +test( + "delete reports success for a backup-only run", + withTempCwd(async (cwd) => { + const rp = createRunPersistence(cwd); + rp.save({ + runId: "backup-only", + workflowName: "w", + status: "completed", + phases: [], + agents: [], + logs: [], + } as PersistedRunState); + const runsDir = workflowProjectPaths(cwd).runsDir; + rmSync(join(runsDir, "backup-only.json")); + + assert.equal(rp.delete("backup-only"), true); + assert.equal(rp.load("backup-only"), null); }), ); @@ -674,7 +865,43 @@ test( ); test( - "delete removes the lock sidecar too", + "a live takeover claim blocks contenders and a dead claim is recoverable", + withTempCwd(async (cwd) => { + const runsDir = workflowProjectPaths(cwd).runsDir; + const rp = createRunPersistence(cwd); + rp.save({ + runId: "takeover-race", + workflowName: "w", + status: "paused", + phases: [], + agents: [], + logs: [], + } as PersistedRunState); + const lockPath = join(runsDir, "takeover-race.lock"); + const takeoverPath = `${lockPath}.takeover`; + const claim = (pid: number, token: string) => ({ + runId: "takeover-race", + runPath: join(runsDir, "takeover-race.json"), + pid, + startedAt: "2024-01-01T00:00:00.000Z", + token, + }); + writeFileSync(lockPath, JSON.stringify(claim(2147483647, "stale")), "utf-8"); + writeFileSync(takeoverPath, JSON.stringify(claim(process.pid, "live-claim")), "utf-8"); + + assert.equal(rp.acquireRunLease("takeover-race"), null); + assert.equal(JSON.parse(readFileSync(takeoverPath, "utf-8")).token, "live-claim"); + + writeFileSync(takeoverPath, JSON.stringify(claim(2147483647, "dead-claim")), "utf-8"); + const lease = rp.acquireRunLease("takeover-race"); + assert.ok(lease); + assert.equal(existsSync(takeoverPath), false); + rp.releaseRunLease(lease); + }), +); + +test( + "artifact deletion leaves the lease for its owner to release", withTempCwd(async (cwd) => { const rp = createRunPersistence(cwd); rp.save({ @@ -688,7 +915,10 @@ test( const lease = rp.acquireRunLease("delete-lock"); assert.ok(lease, "lease exists before delete"); rp.delete("delete-lock"); - assert.equal(existsSync(join(workflowProjectPaths(cwd).runsDir, "delete-lock.lock")), false, "lock cleaned up"); + const lockPath = join(workflowProjectPaths(cwd).runsDir, "delete-lock.lock"); + assert.equal(existsSync(lockPath), true, "artifact deletion must not remove an ownership record"); + rp.releaseRunLease(lease); + assert.equal(existsSync(lockPath), false, "the owner explicitly releases its lease"); }), ); @@ -702,12 +932,23 @@ test( status: "running", script: "export const meta = { name: 'w', description: 'd' }\nawait agent('x',{label:'x'})\nreturn 1", phases: [], - agents: [], + agents: [ + { + id: 1, + executionId: "stale:0", + callIndex: 0, + label: "orphan", + prompt: "x", + status: "running", + }, + ], logs: [], } as PersistedRunState); // A fresh manager (the previous process died) should recover the orphan. new WorkflowManager({ cwd }); - assert.equal(rp.load("stale")?.status, "paused", "stale running -> paused (journal preserved for resume)"); + const recovered = rp.load("stale"); + assert.equal(recovered?.status, "paused", "stale running -> paused (journal preserved for resume)"); + assert.equal(recovered?.agents[0].status, "paused", "orphaned running invocation is visibly paused mid-run"); }), ); diff --git a/tests/workflow-control-tool.test.ts b/tests/workflow-control-tool.test.ts index 31d0374b..f1e4a872 100644 --- a/tests/workflow-control-tool.test.ts +++ b/tests/workflow-control-tool.test.ts @@ -1,10 +1,14 @@ import assert from "node:assert/strict"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; import test from "node:test"; import { Check } from "typebox/value"; import type { WorkflowSnapshot } from "../src/display.js"; import type { PersistedRunState, RunStatus } from "../src/run-persistence.js"; import { createWorkflowControlTool } from "../src/workflow-control-tool.js"; -import type { WorkflowManager } from "../src/workflow-manager.js"; +import { WorkflowManager } from "../src/workflow-manager.js"; +import { withFakeHomeAsync } from "./helpers/fake-home.js"; function run(status: RunStatus = "running", runId = "audit-abc123"): PersistedRunState { return { @@ -54,6 +58,20 @@ function fakeManager(initial: PersistedRunState[], liveSnapshots: Record>): string { return result.content[0].text; } -test("workflow_control exposes only list, status, pause, resume, and stop in a strict schema", () => { +test("workflow_control exposes durable lifecycle actions in a strict schema", () => { const { manager } = fakeManager([]); const tool = createWorkflowControlTool({ manager }); @@ -77,8 +95,8 @@ test("workflow_control exposes only list, status, pause, resume, and stop in a s assert.equal(Check(tool.parameters, { action: "pause", runId: "abc" }), true); assert.equal(Check(tool.parameters, { action: "resume", runId: "abc" }), true); assert.equal(Check(tool.parameters, { action: "stop", runId: "abc" }), true); - assert.equal(Check(tool.parameters, { action: "restart", runId: "abc" }), false); - assert.equal(Check(tool.parameters, { action: "remove", runId: "abc" }), false); + assert.equal(Check(tool.parameters, { action: "restart", runId: "abc" }), true); + assert.equal(Check(tool.parameters, { action: "remove", runId: "abc" }), true); assert.equal(Check(tool.parameters, { action: "set_concurrency", runId: "abc", concurrency: 2 }), false); assert.equal(Check(tool.parameters, { action: "status" }), false); assert.equal(Check(tool.parameters, { action: "list", runId: "abc" }), false); @@ -87,7 +105,7 @@ test("workflow_control exposes only list, status, pause, resume, and stop in a s const prepare = tool.prepareArguments as (value: unknown) => unknown; assert.throws(() => prepare({ action: "pause" }), /requires runId/); assert.throws(() => prepare({ action: "status", runId: "abc", extra: true }), /does not accept extra/); - assert.throws(() => prepare({ action: "restart", runId: "abc" }), /requires action/); + assert.throws(() => prepare({ action: "restart" }), /requires runId/); }); test("list and status return stable lifecycle and observability fields", async () => { @@ -96,7 +114,7 @@ test("list and status return stable lifecycle and observability fields", async ( const listed = await execute(manager, { action: "list" }); assert.match(text(listed), /^action=list result=ok runs=1\n/); assert.match(text(listed), /runId=audit-abc123 name="audit" status=running phase="Inspect"/); - assert.match(text(listed), /total=4 done=0 running=1 queued=1 error=1 skipped=1/); + assert.match(text(listed), /total=4 done=0 running=1 paused=0 queued=1 error=1 skipped=1/); assert.match(text(listed), /active="active scan" tokens=30/); assert.deepEqual(listed.details, { action: "list", @@ -107,7 +125,7 @@ test("list and status return stable lifecycle and observability fields", async ( workflowName: "audit", status: "running", phase: "Inspect", - counts: { total: 4, done: 0, running: 1, queued: 1, error: 1, skipped: 1 }, + counts: { total: 4, done: 0, running: 1, paused: 0, queued: 1, error: 1, skipped: 1 }, activeLabels: ["active scan"], tokenTotal: 30, }, @@ -171,6 +189,53 @@ test("pause, resume, and stop call the shared manager lifecycle methods", async ); }); +test("restart creates a new run and remove deletes terminal history", async () => { + const fixture = fakeManager([run("completed")]); + + const restarted = await execute(fixture.manager, { action: "restart", runId: "audit-abc123" }); + assert.match(text(restarted), /action=restart result=restarted sourceRunId=audit-abc123/); + assert.equal(restarted.details.runId, "audit-new456"); + + const removed = await execute(fixture.manager, { action: "remove", runId: "audit-abc123" }); + assert.match(text(removed), /action=remove result=removed runId=audit-abc123/); + assert.deepEqual( + fixture.calls.map((call) => call.action), + ["restart", "remove"], + ); +}); + +test("cold paused run can be stopped and then removed", async () => { + const cwd = mkdtempSync(join(tmpdir(), "pi-dw-control-cold-")); + const fakeHome = mkdtempSync(join(tmpdir(), "pi-dw-control-home-")); + try { + await withFakeHomeAsync(fakeHome, async () => { + const bootstrap = new WorkflowManager({ cwd }); + const runId = "cold-paused-control"; + bootstrap.getPersistence().save({ + runId, + workflowName: "cold_control", + script: "export const meta = { name: 'cold_control', description: 'cold' }; return true", + status: "paused", + phases: [], + agents: [], + logs: [], + startedAt: "2026-07-14T00:00:00.000Z", + updatedAt: "2026-07-14T00:00:01.000Z", + }); + + const manager = new WorkflowManager({ cwd }); + assert.equal(manager.getRun(runId), undefined); + assert.match(text(await execute(manager, { action: "stop", runId })), /result=stopped/); + assert.equal(manager.getPersistence().load(runId)?.status, "aborted"); + assert.match(text(await execute(manager, { action: "remove", runId })), /result=removed/); + assert.equal(manager.getPersistence().load(runId), null); + }); + } finally { + rmSync(cwd, { recursive: true, force: true }); + rmSync(fakeHome, { recursive: true, force: true }); + } +}); + test("unknown IDs and illegal transitions return explicit errors with allowed actions", async () => { const fixture = fakeManager([run("completed"), run("running", "live-123")]); @@ -179,10 +244,13 @@ test("unknown IDs and illegal transitions return explicit errors with allowed ac const pauseCompleted = text(await execute(fixture.manager, { action: "pause", runId: "audit-abc123" })); assert.match(pauseCompleted, /cannot pause run with status completed/); - assert.match(pauseCompleted, /allowed=status/); + assert.match(pauseCompleted, /allowed=status,restart,remove/); + + const removeRunning = text(await execute(fixture.manager, { action: "remove", runId: "live-123" })); + assert.match(removeRunning, /cannot remove a running run; stop it first/); await execute(fixture.manager, { action: "stop", runId: "live-123" }); const stopAborted = text(await execute(fixture.manager, { action: "stop", runId: "live-123" })); assert.match(stopAborted, /cannot stop run with status aborted/); - assert.match(stopAborted, /allowed=status/); + assert.match(stopAborted, /allowed=status,restart,remove/); }); diff --git a/tests/workflow-manager-abort.test.ts b/tests/workflow-manager-abort.test.ts index 100975fe..ebbdb627 100644 --- a/tests/workflow-manager-abort.test.ts +++ b/tests/workflow-manager-abort.test.ts @@ -350,21 +350,20 @@ test( assert.equal(paused, true); assert.equal(manager.getRun(runId)?.status, "paused"); - // Resume — replays journal (empty for single-agent that never completed) and - // re-runs the live agent with a fresh (non-aborted) controller. - const resumed = await manager.resume(runId); - assert.equal(resumed, true, "resume should succeed"); - - // The resumed run should be running - assert.equal(manager.getRun(runId)?.status, "running", "resumed run should be running"); + // Resume serializes behind the interrupted execution before it starts a + // replacement, so an invocation can never overlap with its prior attempt. + const resumePromise = manager.resume(runId); + await new Promise((r) => setTimeout(r, 20)); + assert.equal(manager.getRun(runId)?.status, "paused"); - // Resolve the deferred agent so the resumed run's agent completes + // Let the interrupted runner settle; the replacement then starts with a + // fresh controller. This test runner reuses an already-resolved promise, so + // the replacement completes immediately with the same terminal result. da.resolve("resumed-done"); - - // The original promise will reject (its controller was aborted). Suppress it. await origPromise.catch(() => {}); + assert.equal(await resumePromise, true, "resume should succeed"); - // Wait for the resumed run to complete + // Wait for the resumed run to complete. await new Promise((r) => setTimeout(r, 50)); const finalRun = manager.getRun(runId); @@ -513,7 +512,7 @@ test( // ─── deleteRun tests (2 tests) ───────────────────────────────────────────────── test( - "deleteRun can delete a running run (removes from memory and persistence)", + "deleteRun refuses active work until it is stopped", withTempCwd(async (cwd) => { const da = deferredAgent(); const manager = new WorkflowManager({ cwd, agent: da.runner }); @@ -521,18 +520,16 @@ test( const { runId, promise } = manager.startInBackground(oneAgentScript); await new Promise((r) => setTimeout(r, 20)); - // Delete while running — should succeed (removes from tracking) - const deleted = manager.deleteRun(runId); - assert.equal(deleted, true); + assert.equal(manager.deleteRun(runId), false, "running artifacts must remain owned by the execution"); + assert.ok(manager.getRun(runId)); + assert.ok(manager.listRuns().some((run) => run.runId === runId)); - // Should not be in memory + assert.equal(manager.stop(runId), true); + assert.equal(manager.deleteRun(runId), true); assert.equal(manager.getRun(runId), undefined); - - // Should not be in persistence - const runs = manager.listRuns(); assert.equal( - runs.find((r) => r.runId === runId), - undefined, + manager.listRuns().some((run) => run.runId === runId), + false, ); da.resolve("done"); @@ -674,14 +671,14 @@ test( const { runId, promise: origPromise } = manager.startInBackground(oneAgentScript); await new Promise((r) => setTimeout(r, 20)); manager.pause(runId); - await manager.resume(runId); + const resumePromise = manager.resume(runId); + da.resolve("done"); + await origPromise.catch(() => {}); + assert.equal(await resumePromise, true); assert.ok(resumedEvent, "resumed event should fire on resume"); assert.equal(resumedEvent?.runId, runId); - da.resolve("done"); - await origPromise.catch(() => {}); - // Now test error event on abort let capturedError: { runId: string; error: WorkflowError } | null = null; const da2 = deferredAgent(); diff --git a/tests/workflow-manager.test.ts b/tests/workflow-manager.test.ts index 3aa9fd3c..158995c7 100644 --- a/tests/workflow-manager.test.ts +++ b/tests/workflow-manager.test.ts @@ -5,6 +5,7 @@ import { join } from "node:path"; import test from "node:test"; import type { AgentUsage } from "../src/agent.js"; import { WorkflowError, WorkflowErrorCode } from "../src/errors.js"; +import { UsageLimitScheduler } from "../src/usage-limit-scheduler.js"; import { WorkflowManager } from "../src/workflow-manager.js"; import { withFakeHomeAsync } from "./helpers/fake-home.js"; @@ -108,6 +109,41 @@ test( }), ); +test( + "initial persistence failures prevent background and synchronous execution", + withTempCwd(async (cwd) => { + const manager = new WorkflowManager({ cwd, agent: fakeAgent() }); + manager.getPersistence().save = () => { + throw new Error("disk unavailable"); + }; + + assert.throws(() => manager.startInBackground(oneAgentScript), /disk unavailable/); + await assert.rejects(manager.runSync(oneAgentScript), /disk unavailable/); + assert.equal(manager.listRuns().length, 0); + }), +); + +test( + "a post-start persistence failure rejects completion and removes the unsafe resume artifact", + withTempCwd(async (cwd) => { + const manager = new WorkflowManager({ cwd, agent: fakeAgent({}, "charged-result") }); + const persistence = manager.getPersistence(); + const save = persistence.save.bind(persistence); + let saves = 0; + persistence.save = (state) => { + saves++; + if (saves > 1) throw new Error("disk failed after initial save"); + save(state); + }; + + await assert.rejects( + manager.runSync(oneAgentScript), + (error: unknown) => error instanceof WorkflowError && error.code === WorkflowErrorCode.PERSISTENCE_ERROR, + ); + assert.equal(manager.listRuns().length, 0); + }), +); + test( "manager defaultAgentTimeoutMs applies when run options omit agentTimeoutMs", withTempCwd(async (cwd) => { @@ -645,6 +681,141 @@ return { a, b }`; // ─── Stop tests ──────────────────────────────────────────────────────────────── +test( + "restarting a usage-limit pause aborts the source and cancels its auto-resume timer", + withTempCwd(async (cwd) => { + const seeder = new WorkflowManager({ cwd }); + seeder.getPersistence().save({ + runId: "restart-paused-source", + workflowName: "paused", + script: oneAgentScript, + status: "paused", + pauseReason: "usage_limit", + resetHint: "Resets in 1h", + autoResume: true, + phases: [], + agents: [], + logs: [], + startedAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }); + + const manager = new WorkflowManager({ cwd, agent: fakeAgent() }); + const scheduler = new UsageLimitScheduler(manager, { + setTimer: () => 1 as never, + clearTimer: () => {}, + }); + assert.equal(scheduler.hasArmedTimer("restart-paused-source"), true); + + const restarted = manager.restart("restart-paused-source"); + assert.ok(restarted); + const source = manager.getPersistence().load("restart-paused-source"); + assert.equal(source?.status, "aborted"); + assert.equal(source?.autoResume, false); + assert.equal(scheduler.hasArmedTimer("restart-paused-source"), false); + await restarted.promise; + scheduler.dispose(); + }), +); + +test( + "active stop surfaces persistence failure and removes the unsafe artifact", + withTempCwd(async (cwd) => { + const agent = deferredAgent(); + const manager = new WorkflowManager({ cwd, agent: agent.runner }); + manager.on("error", () => {}); + const { runId, promise } = manager.startInBackground(oneAgentScript); + await new Promise((resolve) => setTimeout(resolve, 20)); + manager.getPersistence().save = () => { + throw new Error("stop disk failure"); + }; + + assert.throws( + () => manager.stop(runId), + (error: unknown) => error instanceof WorkflowError && error.code === WorkflowErrorCode.PERSISTENCE_ERROR, + ); + assert.equal(manager.getPersistence().load(runId), null); + agent.resolve(); + await promise.catch(() => {}); + }), +); + +test( + "cold stop wraps persistence failures without reporting success", + withTempCwd(async (cwd) => { + const manager = new WorkflowManager({ cwd }); + const persistence = manager.getPersistence(); + persistence.save({ + runId: "cold-stop-failure", + workflowName: "cold", + script: oneAgentScript, + status: "paused", + phases: [], + agents: [], + logs: [], + startedAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }); + const save = persistence.save.bind(persistence); + persistence.save = () => { + throw new Error("cold stop disk failure"); + }; + + assert.throws( + () => manager.stop("cold-stop-failure"), + (error: unknown) => error instanceof WorkflowError && error.code === WorkflowErrorCode.PERSISTENCE_ERROR, + ); + persistence.save = save; + assert.equal(persistence.load("cold-stop-failure")?.status, "paused"); + }), +); + +test( + "stop retains its lease until abort-ignoring execution settles", + withTempCwd(async (cwd) => { + const agent = deferredAgent(); + const owner = new WorkflowManager({ cwd, agent: agent.runner }); + owner.on("error", () => {}); + const { runId, promise } = owner.startInBackground(oneAgentScript); + await new Promise((resolve) => setTimeout(resolve, 20)); + + assert.equal(owner.stop(runId), true); + const contender = new WorkflowManager({ cwd }); + assert.equal(contender.deleteRun(runId), false, "another process cannot delete while execution unwinds"); + + agent.resolve(); + await promise.catch(() => {}); + assert.equal(contender.deleteRun(runId), true); + assert.equal(contender.getPersistence().load(runId), null); + }), +); + +test( + "a stale manager cannot stop a paused run resumed by another process", + withTempCwd(async (cwd) => { + const firstAgent = deferredAgent(); + const first = new WorkflowManager({ cwd, agent: firstAgent.runner }); + first.on("error", () => {}); + const { runId, promise } = first.startInBackground(oneAgentScript); + await new Promise((resolve) => setTimeout(resolve, 20)); + assert.equal(first.pause(runId), true); + firstAgent.resolve(); + await promise.catch(() => {}); + + const secondAgent = deferredAgent(); + const second = new WorkflowManager({ cwd, agent: secondAgent.runner }); + second.on("error", () => {}); + assert.equal(await second.resume(runId), true); + await new Promise((resolve) => setTimeout(resolve, 20)); + + assert.equal(first.stop(runId), false); + assert.equal(second.getRun(runId)?.status, "running"); + second.pause(runId); + secondAgent.resolve(); + await second.getRun(runId)?.executionPromise?.catch(() => {}); + }), +); + test( "stop on paused run transitions to aborted", withTempCwd(async (cwd) => { @@ -666,6 +837,7 @@ test( da.resolve("done"); await promise.catch(() => {}); + assert.equal(manager.getPersistence().load(runId)?.agents[0]?.status, "skipped"); }), ); @@ -790,19 +962,13 @@ test( assert.equal(paused, true); assert.equal(manager.getRun(runId)?.status, "paused"); - // Resume — replays journal (empty for single-agent that never completed) and - // re-runs the live agent with a fresh (non-aborted) controller. - const resumed = await manager.resume(runId); - assert.equal(resumed, true, "resume should succeed"); - - // The resumed run should be running - assert.equal(manager.getRun(runId)?.status, "running", "resumed run should be running"); - - // Resolve the deferred agent so the resumed run's agent completes + // Resume serializes behind the interrupted execution so replacement work + // cannot overlap with the old invocation. + const resumePromise = manager.resume(runId); da.resolve("resumed-done"); - - // The original promise will reject (its controller was aborted). Suppress it. await origPromise.catch(() => {}); + const resumed = await resumePromise; + assert.equal(resumed, true, "resume should succeed"); // Wait for the resumed run to complete await new Promise((r) => setTimeout(r, 50)); @@ -956,6 +1122,48 @@ test( }), ); +test( + "resume revalidates terminal state after acquiring the lease", + withTempCwd(async (cwd) => { + let calls = 0; + const manager = new WorkflowManager({ + cwd, + agent: { + async run() { + calls++; + return "live"; + }, + }, + }); + const persistence = manager.getPersistence(); + persistence.save({ + runId: "resume-lease-race", + workflowName: "race", + script: oneAgentScript, + status: "paused", + phases: [], + agents: [], + logs: [], + startedAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }); + const acquire = persistence.acquireRunLease.bind(persistence); + persistence.acquireRunLease = (runId) => { + const current = persistence.load(runId); + assert.ok(current); + persistence.save({ ...current, status: "completed", result: "already done" }); + return acquire(runId); + }; + + assert.equal(await manager.resume("resume-lease-race"), false); + assert.equal(calls, 0); + assert.equal(persistence.load("resume-lease-race")?.status, "completed"); + const lease = persistence.acquireRunLease("resume-lease-race"); + assert.ok(lease, "resume released the lease after revalidation"); + persistence.releaseRunLease(lease); + }), +); + // ─── Cold-start resume tests ──────────────────────────────────────────────────── // These tests manually persist runs via the persistence layer (as though the // process was restarted) and then resume them from disk — no in-memory state. @@ -1225,7 +1433,7 @@ test( // ─── deleteRun tests ─────────────────────────────────────────────────────────── test( - "deleteRun can delete a running run (removes from memory and persistence)", + "deleteRun refuses active work until it is stopped", withTempCwd(async (cwd) => { const da = deferredAgent(); const manager = new WorkflowManager({ cwd, agent: da.runner }); @@ -1233,18 +1441,16 @@ test( const { runId, promise } = manager.startInBackground(oneAgentScript); await new Promise((r) => setTimeout(r, 20)); - // Delete while running — should succeed (removes from tracking) - const deleted = manager.deleteRun(runId); - assert.equal(deleted, true); + assert.equal(manager.deleteRun(runId), false); + assert.ok(manager.getRun(runId)); + assert.ok(manager.listRuns().some((run) => run.runId === runId)); - // Should not be in memory + assert.equal(manager.stop(runId), true); + assert.equal(manager.deleteRun(runId), true); assert.equal(manager.getRun(runId), undefined); - - // Should not be in persistence - const runs = manager.listRuns(); assert.equal( - runs.find((r) => r.runId === runId), - undefined, + manager.listRuns().some((run) => run.runId === runId), + false, ); da.resolve("done"); @@ -1382,14 +1588,14 @@ test( await new Promise((resolve) => setTimeout(resolve, 20)); manager.pause(runId); - const resumed = await manager.resume(runId); + const resumePromise = manager.resume(runId); + da.resolve("done"); + await promise.catch(() => {}); + const resumed = await resumePromise; const persisted = manager.listRuns().find((run) => run.runId === runId); assert.equal(resumed, true); assert.equal(persisted?.status, "running", "listRuns should show running status after resume"); - - da.resolve("done"); - await promise.catch(() => {}); }), ); @@ -1408,13 +1614,13 @@ test( const { runId, promise } = manager.startInBackground(oneAgentScript); await new Promise((r) => setTimeout(r, 20)); manager.pause(runId); - await manager.resume(runId); + const resumePromise = manager.resume(runId); + da.resolve("done"); + await promise.catch(() => {}); + assert.equal(await resumePromise, true); assert.ok(resumedEvent, "resumed event should fire"); assert.equal(resumedEvent?.runId, runId); - - da.resolve("done"); - await promise.catch(() => {}); }), ); @@ -1466,13 +1672,11 @@ test( assert.equal(manager.pause(runId), true); assert.equal(manager.getRun(runId)?.status, "paused", "should be paused after pause"); - const resumed = await manager.resume(runId); - assert.equal(resumed, true); - assert.equal(manager.getRun(runId)?.status, "running", "should be running after resume"); - - // Complete the resumed run + const resumePromise = manager.resume(runId); da.resolve("resumed-done"); await origPromise.catch(() => {}); + const resumed = await resumePromise; + assert.equal(resumed, true); await new Promise((r) => setTimeout(r, 30)); assert.equal(manager.getRun(runId)?.status, "completed", "should complete after resume finishes"); @@ -1597,16 +1801,16 @@ test( assert.equal(manager.pause(runId), true); assert.equal(manager.getRun(runId)?.status, "paused"); - // First resume should succeed - const firstResume = await manager.resume(runId); + // The first resume waits for the interrupted execution to settle. + const firstResumePromise = manager.resume(runId); + da.resolve("done"); + await origPromise.catch(() => {}); + const firstResume = await firstResumePromise; assert.equal(firstResume, true, "first resume should succeed"); - // The resumed run is now running; second resume should return false + // The resumed run is now running; second resume should return false. const secondResume = await manager.resume(runId); assert.equal(secondResume, false, "second resume should return false when the resumed run is already running"); - - da.resolve("done"); - await origPromise.catch(() => {}); }), ); @@ -1660,8 +1864,11 @@ test( }); try { - // Resume — executeRun calls runWorkflow which calls the mocked runner - const resumed = await manager.resume(runId); + // Resume waits for the interrupted execution before invoking the mocked runner. + const resumePromise = manager.resume(runId); + da.resolve("done"); + await origPromise.catch(() => {}); + const resumed = await resumePromise; assert.equal(resumed, true, "resume should schedule the run"); // Wait for the background executed run to process the agent error @@ -1676,10 +1883,7 @@ test( "error code should be AGENT_EXECUTION_ERROR", ); } finally { - // Resolve the original deferred promise so the first executeRun settles da.runner.run = async (_prompt: string) => "done"; - da.resolve("done"); - await origPromise.catch(() => {}); } }), ); @@ -1732,8 +1936,11 @@ test( }); try { - // Resume — the run will fail because the mocked agent throws - const resumed = await manager.resume(runId); + // Resume — settle the interrupted execution, then the mocked agent fails. + const resumePromise = manager.resume(runId); + da.resolve("done"); + await origPromise.catch(() => {}); + const resumed = await resumePromise; assert.equal(resumed, true, "resume should schedule the run"); await new Promise((r) => setTimeout(r, 100)); @@ -1748,8 +1955,6 @@ test( assert.equal(manager.getRun(runId)?.status, "failed", "status should remain failed after rejected pause"); } finally { da.runner.run = async (_prompt: string) => "done"; - da.resolve("done"); - await origPromise.catch(() => {}); } }), ); @@ -1774,8 +1979,11 @@ test( }); try { - // Resume — the run will fail - const resumed = await manager.resume(runId); + // Resume — settle the interrupted execution, then the mocked agent fails. + const resumePromise = manager.resume(runId); + da.resolve("done"); + await origPromise.catch(() => {}); + const resumed = await resumePromise; assert.equal(resumed, true, "resume should schedule the run"); await new Promise((r) => setTimeout(r, 100)); @@ -1790,8 +1998,6 @@ test( assert.equal(manager.getRun(runId)?.status, "failed", "status should remain failed after rejected stop"); } finally { da.runner.run = async (_prompt: string) => "done"; - da.resolve("done"); - await origPromise.catch(() => {}); } }), ); @@ -1816,8 +2022,11 @@ test( }); try { - // Resume — the run will fail - await manager.resume(runId); + // Resume — settle the interrupted execution, then the mocked agent fails. + const firstResume = manager.resume(runId); + da.resolve("done"); + await origPromise.catch(() => {}); + assert.equal(await firstResume, true); await new Promise((r) => setTimeout(r, 100)); // Verify the run is now in failed state @@ -1825,10 +2034,8 @@ test( assert.equal(failedRun?.status, "failed", "run should be in failed state"); assert.ok(failedRun?.error instanceof WorkflowError, "error should be a WorkflowError"); } finally { - // Restore the runner so the resumed run's agent call succeeds + // Restore the runner so the resumed run's agent call succeeds. da.runner.run = async (_prompt: string) => "done"; - da.resolve("done"); - await origPromise.catch(() => {}); } // Resume the failed run — resume() allows failed status diff --git a/tests/workflow-resume-state.test.ts b/tests/workflow-resume-state.test.ts new file mode 100644 index 00000000..95fc8708 --- /dev/null +++ b/tests/workflow-resume-state.test.ts @@ -0,0 +1,515 @@ +import assert from "node:assert/strict"; +import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import test from "node:test"; +import type { AgentUsage } from "../src/agent.js"; +import { WorkflowManager } from "../src/workflow-manager.js"; +import { workflowProjectPaths } from "../src/workflow-paths.js"; +import { withFakeHomeAsync } from "./helpers/fake-home.js"; + +const script = `export const meta = { name: 'resume_42', description: 'durable restart fixture' } +phase('Work') +const results = await parallel(Array.from({ length: 42 }, (_, i) => () => agent('task-' + i, { label: 'worker' }))) +return results`; + +function usage(total: number): AgentUsage { + return { input: total - 1, output: 1, cacheRead: 0, cacheWrite: 0, total, cost: total / 1000 }; +} + +function agentWithUsage(total: number, activity?: { active: number; maximum: number; calls: number }) { + return { + async run(prompt: string, options: { onUsage?: (value: AgentUsage) => void }) { + if (activity) { + activity.active++; + activity.maximum = Math.max(activity.maximum, activity.active); + activity.calls++; + await new Promise((resolve) => setTimeout(resolve, 2)); + } + options.onUsage?.(usage(total)); + if (activity) activity.active--; + return `result:${prompt}`; + }, + }; +} + +async function waitForCompletion(manager: WorkflowManager, runId: string): Promise { + for (let attempt = 0; attempt < 200; attempt++) { + if (manager.getRun(runId)?.status === "completed") return; + await new Promise((resolve) => setTimeout(resolve, 10)); + } + assert.fail("resumed workflow did not complete"); +} + +test("cold resume preserves completed rows and execution controls without double charging", async () => { + const cwd = mkdtempSync(join(tmpdir(), "pi-dw-resume-")); + const fakeHome = mkdtempSync(join(tmpdir(), "pi-dw-home-")); + try { + await withFakeHomeAsync(fakeHome, async () => { + const first = new WorkflowManager({ cwd, agent: agentWithUsage(10) }); + await first.runSync(script, undefined, { + concurrency: 1, + maxAgents: 42, + agentRetries: 1, + agentTimeoutMs: null, + tokenBudget: 10_000, + }); + const seeded = first.listRuns()[0]; + assert.equal(seeded.agents.length, 42); + assert.equal(seeded.journal?.length, 42); + + first.getPersistence().save({ + ...seeded, + status: "paused", + agents: seeded.agents.map((agent, index) => + index < 27 + ? agent + : { + ...agent, + status: "queued", + result: undefined, + resultPreview: undefined, + usage: undefined, + tokens: undefined, + endedAt: undefined, + }, + ), + journal: seeded.journal?.slice(0, 27), + tokenUsage: { input: 243, output: 27, total: 270, cost: 0.27, cacheRead: 0, cacheWrite: 0 }, + result: undefined, + completedAt: undefined, + durationMs: undefined, + }); + + const activity = { active: 0, maximum: 0, calls: 0 }; + const resumed = new WorkflowManager({ cwd, concurrency: 8, agent: agentWithUsage(20, activity) }); + assert.equal(await resumed.resume(seeded.runId), true); + await waitForCompletion(resumed, seeded.runId); + + const final = resumed.listRuns().find((run) => run.runId === seeded.runId); + assert.ok(final); + assert.equal(final.agents.length, 42); + assert.equal(new Set(final.agents.map((agent) => agent.executionId)).size, 42); + assert.deepEqual( + final.agents.slice(0, 27).map((agent) => agent.status), + Array(27).fill("done"), + ); + assert.deepEqual( + final.agents.slice(0, 27).map((agent) => agent.tokens), + Array(27).fill(10), + ); + assert.deepEqual( + final.agents.slice(27).map((agent) => agent.tokens), + Array(15).fill(20), + ); + assert.equal(activity.calls, 15, "journaled prefix is replayed, not re-run"); + assert.equal(activity.maximum, 1, "cold resume restores the original concurrency instead of manager defaults"); + assert.equal(final.tokenUsage?.total, 570); + assert.equal(final.concurrency, 1); + assert.equal(final.maxAgents, 42); + assert.equal(final.agentRetries, 1); + assert.equal(final.agentTimeoutMs, null); + assert.equal(final.tokenBudget, 10_000); + assert.equal(new Date(final.startedAt).getTime(), new Date(seeded.startedAt).getTime()); + }); + } finally { + rmSync(cwd, { recursive: true, force: true }); + rmSync(fakeHome, { recursive: true, force: true }); + } +}); + +test("exact message usage is debounced to disk while streaming estimates remain live-only", async () => { + const cwd = mkdtempSync(join(tmpdir(), "pi-dw-exact-usage-")); + const fakeHome = mkdtempSync(join(tmpdir(), "pi-dw-home-")); + let release: (() => void) | undefined; + try { + await withFakeHomeAsync(fakeHome, async () => { + const manager = new WorkflowManager({ + cwd, + agent: { + async run(_prompt: string, options: { onUsage?: (value: AgentUsage) => void }) { + options.onUsage?.(usage(11)); + options.onUsage?.({ ...usage(20), estimated: true }); + await new Promise((resolve) => { + release = resolve; + }); + return "done"; + }, + }, + }); + manager.on("error", () => {}); + const one = `export const meta = { name: 'usage', description: 'message usage' } +return await agent('one', { label: 'one' })`; + const { runId, promise } = manager.startInBackground(one); + while (!release) await new Promise((resolve) => setTimeout(resolve, 0)); + + const live = manager.getSnapshot(runId)?.agents[0]; + const durable = manager.getPersistence().load(runId)?.agents[0]; + assert.equal(live?.tokens, 20); + assert.equal(live?.tokensEstimated, true); + assert.equal(durable?.usage?.total, 11); + assert.equal(durable?.tokens, 11); + assert.equal(manager.getPersistence().load(runId)?.tokenUsage?.total, 11); + + manager.pause(runId); + release?.(); + await promise.catch(() => {}); + }); + } finally { + release?.(); + rmSync(cwd, { recursive: true, force: true }); + rmSync(fakeHome, { recursive: true, force: true }); + } +}); + +test("resuming an incomplete invocation keeps its exact usage durable until replacement usage arrives", async () => { + const cwd = mkdtempSync(join(tmpdir(), "pi-dw-resume-baseline-")); + const fakeHome = mkdtempSync(join(tmpdir(), "pi-dw-home-")); + let release: (() => void) | undefined; + try { + await withFakeHomeAsync(fakeHome, async () => { + const one = `export const meta = { name: 'baseline', description: 'usage baseline' } +return await agent('one', { label: 'one' })`; + const first = new WorkflowManager({ cwd, agent: agentWithUsage(5) }); + await first.runSync(one); + const seeded = first.listRuns()[0]; + first.getPersistence().save({ + ...seeded, + status: "paused", + agents: seeded.agents.map((agent) => ({ ...agent, status: "paused", result: undefined, endedAt: undefined })), + journal: [], + result: undefined, + completedAt: undefined, + }); + + const resumed = new WorkflowManager({ + cwd, + agent: { + async run() { + await new Promise((resolve) => { + release = resolve; + }); + return "replacement"; + }, + }, + }); + assert.equal(await resumed.resume(seeded.runId), true); + while (!release) await new Promise((resolve) => setTimeout(resolve, 0)); + await new Promise((resolve) => setTimeout(resolve, 150)); + + assert.equal(resumed.getSnapshot(seeded.runId)?.agents[0]?.usage?.total, 5); + assert.equal(resumed.getPersistence().load(seeded.runId)?.agents[0]?.usage?.total, 5); + resumed.pause(seeded.runId); + release(); + await resumed.getRun(seeded.runId)?.executionPromise?.catch(() => {}); + }); + } finally { + release?.(); + rmSync(cwd, { recursive: true, force: true }); + rmSync(fakeHome, { recursive: true, force: true }); + } +}); + +test("the first replay miss durably removes every unreplayed suffix entry", async () => { + const cwd = mkdtempSync(join(tmpdir(), "pi-dw-resume-suffix-")); + const fakeHome = mkdtempSync(join(tmpdir(), "pi-dw-home-")); + let release: (() => void) | undefined; + try { + await withFakeHomeAsync(fakeHome, async () => { + const two = `export const meta = { name: 'suffix', description: 'suffix invalidation' } +const a = await agent('a', { label: 'a' }) +const b = await agent('b', { label: 'b' }) +return { a, b }`; + const first = new WorkflowManager({ cwd, agent: agentWithUsage(3) }); + await first.runSync(two); + const seeded = first.listRuns()[0]; + first.getPersistence().save({ + ...seeded, + status: "paused", + journal: seeded.journal?.map((entry, index) => (index === 0 ? { ...entry, hash: "stale" } : entry)), + result: undefined, + completedAt: undefined, + }); + + const resumed = new WorkflowManager({ + cwd, + agent: { + async run() { + await new Promise((resolve) => { + release = resolve; + }); + return "replacement"; + }, + }, + }); + assert.equal(await resumed.resume(seeded.runId), true); + while (!release) await new Promise((resolve) => setTimeout(resolve, 0)); + + assert.deepEqual(resumed.getPersistence().load(seeded.runId)?.journal, []); + resumed.pause(seeded.runId); + release(); + await resumed.getRun(seeded.runId)?.executionPromise?.catch(() => {}); + }); + } finally { + release?.(); + rmSync(cwd, { recursive: true, force: true }); + rmSync(fakeHome, { recursive: true, force: true }); + } +}); + +test("cold resume restores historical phase usage before admitting live work", async () => { + const cwd = mkdtempSync(join(tmpdir(), "pi-dw-phase-budget-")); + const fakeHome = mkdtempSync(join(tmpdir(), "pi-dw-home-")); + try { + await withFakeHomeAsync(fakeHome, async () => { + const original = `export const meta = { name: 'phase_budget', description: 'phase budget' } +phase('Work', { budget: 100 }) +const a = await agent('a', { label: 'a' }) +const b = await agent('b', { label: 'b' }) +return { a, b }`; + const first = new WorkflowManager({ cwd, agent: agentWithUsage(10) }); + await first.runSync(original); + const seeded = first.listRuns()[0]; + first.getPersistence().save({ + ...seeded, + script: original.replace("budget: 100", "budget: 5"), + status: "paused", + agents: seeded.agents.map((agent, index) => + index === 0 + ? agent + : { + ...agent, + status: "queued", + result: undefined, + resultPreview: undefined, + usage: undefined, + tokens: undefined, + endedAt: undefined, + }, + ), + journal: seeded.journal?.slice(0, 1), + tokenUsage: usage(10), + result: undefined, + completedAt: undefined, + }); + + let calls = 0; + const resumed = new WorkflowManager({ + cwd, + agent: { + async run() { + calls++; + return "unexpected"; + }, + }, + }); + resumed.on("error", () => {}); + assert.equal(await resumed.resume(seeded.runId), true); + for (let attempt = 0; attempt < 100 && resumed.getRun(seeded.runId)?.status === "running"; attempt++) { + await new Promise((resolve) => setTimeout(resolve, 10)); + } + + assert.equal(calls, 0); + assert.equal(resumed.getRun(seeded.runId)?.status, "failed"); + }); + } finally { + rmSync(cwd, { recursive: true, force: true }); + rmSync(fakeHome, { recursive: true, force: true }); + } +}); + +test("agent completion persists its journal before the terminal-row callback", async () => { + const cwd = mkdtempSync(join(tmpdir(), "pi-dw-journal-boundary-")); + const fakeHome = mkdtempSync(join(tmpdir(), "pi-dw-home-")); + try { + await withFakeHomeAsync(fakeHome, async () => { + const manager = new WorkflowManager({ cwd, agent: agentWithUsage(7) }); + const persistence = manager.getPersistence(); + const save = persistence.save.bind(persistence); + let journalWasDurableWhileRunning = false; + persistence.save = (state) => { + if (state.journal?.length === 1 && state.agents[0]?.status === "running") { + journalWasDurableWhileRunning = true; + } + save(state); + }; + const one = `export const meta = { name: 'journal_boundary', description: 'journal boundary' } +return await agent('one', { label: 'one' })`; + + await manager.runSync(one); + assert.equal(journalWasDurableWhileRunning, true); + }); + } finally { + rmSync(cwd, { recursive: true, force: true }); + rmSync(fakeHome, { recursive: true, force: true }); + } +}); + +test("agent completion persists journal, terminal row, result, and exact usage together", async () => { + const cwd = mkdtempSync(join(tmpdir(), "pi-dw-crash-window-")); + const fakeHome = mkdtempSync(join(tmpdir(), "pi-dw-home-")); + try { + await withFakeHomeAsync(fakeHome, async () => { + const one = `export const meta = { name: 'atomic', description: 'atomic completion' } +return await agent('one', { label: 'one' })`; + const manager = new WorkflowManager({ cwd, agent: agentWithUsage(33) }); + let durableAtEvent = false; + manager.on("agentEnd", ({ runId }: { runId: string }) => { + const saved = manager.getPersistence().load(runId); + durableAtEvent = + saved?.journal?.[0]?.usage?.total === 33 && + saved.agents[0]?.status === "done" && + saved.agents[0]?.usage?.total === 33 && + saved.agents[0]?.result === "result:one"; + }); + + await manager.runSync(one); + assert.equal(durableAtEvent, true); + }); + } finally { + rmSync(cwd, { recursive: true, force: true }); + rmSync(fakeHome, { recursive: true, force: true }); + } +}); + +test("nested cold resume replays scoped journal rows and preserves completed timestamps at zero remaining budget", async () => { + const cwd = mkdtempSync(join(tmpdir(), "pi-dw-nested-resume-")); + const fakeHome = mkdtempSync(join(tmpdir(), "pi-dw-home-")); + try { + await withFakeHomeAsync(fakeHome, async () => { + const child = `export const meta = { name: 'child', description: 'child' } +return await agent('child work', { label: 'child' })`; + const parent = `export const meta = { name: 'parent', description: 'parent' } +const parentResult = await agent('parent work', { label: 'parent' }) +const childResult = await workflow('child') +return { parentResult, childResult }`; + const loadSavedWorkflow = (name: string) => (name === "child" ? child : undefined); + const first = new WorkflowManager({ cwd, agent: agentWithUsage(10), loadSavedWorkflow }); + await first.runSync(parent, undefined, { tokenBudget: 20 }); + const seeded = first.listRuns()[0]; + assert.equal(seeded.journal?.length, 2); + assert.equal(new Set(seeded.journal?.map((entry) => entry.executionId)).size, 2); + const timestamps = new Map(seeded.agents.map((agent) => [agent.executionId, [agent.startedAt, agent.endedAt]])); + first.getPersistence().save({ + ...seeded, + status: "paused", + result: undefined, + completedAt: undefined, + durationMs: undefined, + }); + + const activity = { active: 0, maximum: 0, calls: 0 }; + const resumed = new WorkflowManager({ cwd, agent: agentWithUsage(99, activity), loadSavedWorkflow }); + assert.equal(await resumed.resume(seeded.runId), true); + await waitForCompletion(resumed, seeded.runId); + + const final = resumed.listRuns().find((run) => run.runId === seeded.runId); + assert.ok(final); + assert.equal(activity.calls, 0, "both scoped rows replay without live work"); + assert.equal(final.tokenUsage?.total, 20); + assert.equal(new Set(final.agents.map((agent) => agent.executionId)).size, 2); + for (const agent of final.agents) { + assert.deepEqual([agent.startedAt, agent.endedAt], timestamps.get(agent.executionId)); + } + }); + } finally { + rmSync(cwd, { recursive: true, force: true }); + rmSync(fakeHome, { recursive: true, force: true }); + } +}); + +test("tokens-only v1 completed rows survive cold replay and enter aggregate accounting once", async () => { + const cwd = mkdtempSync(join(tmpdir(), "pi-dw-v1-resume-")); + const fakeHome = mkdtempSync(join(tmpdir(), "pi-dw-home-")); + try { + await withFakeHomeAsync(fakeHome, async () => { + const one = `export const meta = { name: 'legacy_tokens', description: 'legacy tokens' } +return await agent('work', { label: 'legacy' })`; + const first = new WorkflowManager({ cwd, agent: agentWithUsage(10) }); + await first.runSync(one); + const seeded = first.listRuns()[0]; + const runsDir = workflowProjectPaths(cwd).runsDir; + mkdirSync(runsDir, { recursive: true }); + writeFileSync( + join(runsDir, `${seeded.runId}.json`), + JSON.stringify({ + version: 1, + runId: seeded.runId, + workflowName: seeded.workflowName, + script: seeded.script, + status: "paused", + phases: seeded.phases, + agents: seeded.agents.map( + ({ executionId: _executionId, callIndex: _callIndex, usage: _usage, ...agent }) => ({ + ...agent, + tokens: 10, + }), + ), + logs: seeded.logs, + journal: seeded.journal?.map(({ executionId: _executionId, usage: _usage, ...entry }) => entry), + startedAt: seeded.startedAt, + updatedAt: seeded.updatedAt, + }), + ); + + const activity = { active: 0, maximum: 0, calls: 0 }; + const resumed = new WorkflowManager({ cwd, agent: agentWithUsage(50, activity) }); + assert.equal(await resumed.resume(seeded.runId), true); + await waitForCompletion(resumed, seeded.runId); + + const final = resumed.listRuns().find((run) => run.runId === seeded.runId); + assert.ok(final); + assert.equal(activity.calls, 0); + assert.equal(final.agents[0].tokens, 10); + assert.equal(final.agents[0].usage, undefined); + assert.equal(final.tokenUsage?.total, 10); + }); + } finally { + rmSync(cwd, { recursive: true, force: true }); + rmSync(fakeHome, { recursive: true, force: true }); + } +}); + +test("unfinished exact usage is a per-invocation resume baseline included exactly once", async () => { + const cwd = mkdtempSync(join(tmpdir(), "pi-dw-usage-baseline-")); + const fakeHome = mkdtempSync(join(tmpdir(), "pi-dw-home-")); + try { + await withFakeHomeAsync(fakeHome, async () => { + const one = `export const meta = { name: 'usage_baseline', description: 'usage baseline' } +return await agent('work', { label: 'worker' })`; + const first = new WorkflowManager({ cwd, agent: agentWithUsage(7) }); + await first.runSync(one); + const seeded = first.listRuns()[0]; + first.getPersistence().save({ + ...seeded, + status: "paused", + agents: seeded.agents.map((agent) => ({ + ...agent, + status: "queued", + result: undefined, + resultPreview: undefined, + endedAt: undefined, + })), + journal: [], + tokenUsage: { input: 6, output: 1, total: 7, cost: 0.007, cacheRead: 0, cacheWrite: 0 }, + result: undefined, + completedAt: undefined, + durationMs: undefined, + }); + + const resumed = new WorkflowManager({ cwd, agent: agentWithUsage(5) }); + assert.equal(await resumed.resume(seeded.runId), true); + await waitForCompletion(resumed, seeded.runId); + + const final = resumed.listRuns().find((run) => run.runId === seeded.runId); + assert.ok(final); + assert.equal(final.agents[0].usage?.total, 12); + assert.equal(final.agents[0].tokens, 12); + assert.equal(final.journal?.[0].usage?.total, 12); + assert.equal(final.tokenUsage?.total, 12); + }); + } finally { + rmSync(cwd, { recursive: true, force: true }); + rmSync(fakeHome, { recursive: true, force: true }); + } +}); diff --git a/tests/workflow-runtime.test.ts b/tests/workflow-runtime.test.ts index 3e290ef8..1ad624dd 100644 --- a/tests/workflow-runtime.test.ts +++ b/tests/workflow-runtime.test.ts @@ -79,6 +79,171 @@ return xs`; assert.equal(result.agentCount, 4); }); +test("queued and running callbacks keep duplicate labels separated by executionId", async () => { + const release = createDeferred(); + let runnerStarts = 0; + const queued: Array<{ executionId: string; callIndex: number; label: string }> = []; + const started: Array<{ executionId: string; callIndex: number; label: string }> = []; + const histories: Array<{ executionId: string; history: unknown[] }> = []; + const usages: Array<{ executionId: string; usage: AgentUsage }> = []; + const ended: Array<{ executionId: string; result: unknown }> = []; + const runner = { + async run( + prompt: string, + options: { onHistory?: (history: unknown[]) => void; onUsage?: (usage: AgentUsage) => void }, + ) { + runnerStarts++; + options.onHistory?.([{ role: "assistant", kind: "text", text: prompt }]); + await release.promise; + const total = prompt.charCodeAt(0) - 96; + options.onUsage?.({ input: total, output: 0, cacheRead: 0, cacheWrite: 0, total, cost: 0 }); + return `done:${prompt}`; + }, + }; + const script = `export const meta = { name: 'identities', description: 'duplicate labels' } +const xs = await parallel(['a','b','c','d'].map((p) => () => agent(p, { label: 'worker' }))) +return xs`; + + const run = runWorkflow(script, { + runId: "identity-run", + agent: runner, + concurrency: 2, + persistLogs: false, + onAgentQueued: (event) => queued.push(event), + onAgentStart: (event) => started.push(event), + onAgentHistory: (event) => histories.push(event), + onAgentUsage: (event) => usages.push(event), + onAgentEnd: (event) => ended.push(event), + }); + while (runnerStarts < 2) await new Promise((resolve) => setTimeout(resolve, 0)); + + assert.equal(queued.length, 4, "all logical invocations should be visible before limiter admission"); + assert.equal(started.length, 2, "only admitted invocations should be running"); + assert.deepEqual( + queued.map((event) => event.executionId), + ["identity-run:0", "identity-run:1", "identity-run:2", "identity-run:3"], + ); + assert.deepEqual( + started.map((event) => event.executionId), + ["identity-run:0", "identity-run:1"], + ); + + release.resolve(); + await run; + assert.equal(started.length, 4); + assert.equal(new Set(histories.map((event) => event.executionId)).size, 4); + assert.equal(new Set(usages.map((event) => event.executionId)).size, 4); + assert.equal(new Set(ended.map((event) => event.executionId)).size, 4); +}); + +test("aborted queued agents without exact usage finish skipped at zero tokens", async () => { + const release = createDeferred(); + const controller = new AbortController(); + let starts = 0; + const queued: string[] = []; + const ended: Array<{ status: string; tokens?: number; usage?: AgentUsage }> = []; + const run = runWorkflow( + `export const meta = { name: 'abort_queue', description: 'abort queue' } +return await parallel(['a','b','c'].map((p) => () => agent(p, { label: p })))`, + { + agent: { + async run() { + starts++; + await release.promise; + return "done"; + }, + }, + concurrency: 1, + signal: controller.signal, + persistLogs: false, + onAgentQueued: (event) => queued.push(event.executionId), + onAgentEnd: (event) => ended.push(event), + }, + ); + while (starts < 1) await new Promise((resolve) => setTimeout(resolve, 0)); + controller.abort(); + release.resolve(); + await assert.rejects(run, /aborted/i); + + assert.equal(starts, 1, "queued work never reaches the runner"); + assert.equal(queued.length, 3, "all logical invocations remain visible"); + assert.equal(ended.length, 1, "only the invocation that actually started can be aborted"); + assert.ok(ended.every((event) => event.status === "skipped" && event.tokens === 0 && event.usage === undefined)); +}); + +test("incremental estimates are display-only and exact usage replaces them before completion", async () => { + const reported = createDeferred(); + const release = createDeferred(); + const agentUsage: AgentUsage[] = []; + const runUsage: number[] = []; + const runner = { + async run(_prompt: string, options: { onUsage?: (usage: AgentUsage) => void }) { + options.onUsage?.({ + input: 0, + output: 20, + cacheRead: 0, + cacheWrite: 0, + total: 20, + cost: 0, + estimated: true, + }); + options.onUsage?.({ input: 8, output: 4, cacheRead: 0, cacheWrite: 0, total: 12, cost: 0.01 }); + reported.resolve(); + await release.promise; + return "done"; + }, + }; + + const run = runWorkflow( + `export const meta = { name: 'live_usage', description: 'usage before completion' } +return await agent('work', { label: 'worker' })`, + { + runId: "usage-run", + agent: runner, + persistLogs: false, + onAgentUsage: (event) => agentUsage.push(event.usage), + onTokenUsage: (usage) => runUsage.push(usage.total), + }, + ); + await reported.promise; + + assert.deepEqual( + agentUsage.slice(0, 2).map((usage) => ({ total: usage.total, estimated: usage.estimated })), + [ + { total: 20, estimated: true }, + { total: 12, estimated: false }, + ], + ); + assert.deepEqual(runUsage, [12], "the streaming estimate must not enter budget/run accounting"); + + release.resolve(); + const result = await run; + assert.equal(result.tokenUsage?.total, 12); + assert.equal(agentUsage.at(-1)?.total, 12, "terminal reconciliation replaces rather than adds usage"); +}); + +test("repeated absolute usage updates are idempotent", async () => { + const result = await runWorkflow( + `export const meta = { name: 'idempotent_usage', description: 'usage' } +return await agent('work', { label: 'worker' })`, + { + agent: { + async run(_prompt: string, options: { onUsage?: (usage: AgentUsage) => void }) { + const usage = { input: 6, output: 4, cacheRead: 0, cacheWrite: 0, total: 10, cost: 0.01 }; + options.onUsage?.(usage); + options.onUsage?.(usage); + return "done"; + }, + }, + persistLogs: false, + }, + ); + + assert.equal(result.tokenUsage?.total, 10); + assert.equal(result.tokenUsage?.input, 6); + assert.equal(result.tokenUsage?.cost, 0.01); +}); + test("runWorkflow retries recoverable empty output then succeeds", async () => { let calls = 0; const journal: JournalEntry[] = []; @@ -216,7 +381,7 @@ return 1`; assert.equal(seenModel, "meta/default-model", "an agent with no model/tier/phase route uses meta.model"); }); -test("runWorkflow falls back to an estimate when provider reports total === 0", async () => { +test("an exact provider total of zero is not replaced by an estimate", async () => { const result = await runWorkflow(twoAgentScript, { agent: fakeAgent({ total: 0 }, "a result string"), persistLogs: false, @@ -224,10 +389,34 @@ test("runWorkflow falls back to an estimate when provider reports total === 0", assert.equal(result.tokenUsage?.input, 0); assert.equal(result.tokenUsage?.output, 0); - assert.ok((result.tokenUsage?.total ?? 0) > 0, "estimate should be positive"); + assert.equal(result.tokenUsage?.total, 0); assert.equal(result.tokenUsage?.cost, 0); }); +test("missing provider usage emits an estimate without charging or journaling it", async () => { + const journal: JournalEntry[] = []; + const displayUsage: AgentUsage[] = []; + const result = await runWorkflow( + `export const meta = { name: 'estimated', description: 'missing usage' } +return await agent('work', { label: 'worker' })`, + { + agent: { + async run() { + return "a result string"; + }, + }, + persistLogs: false, + onAgentUsage: (event) => displayUsage.push(event.usage), + onAgentJournal: (entry) => journal.push(entry), + }, + ); + + assert.equal(displayUsage.at(-1)?.estimated, true); + assert.ok((displayUsage.at(-1)?.total ?? 0) > 0); + assert.equal(result.tokenUsage?.total, 0); + assert.equal(journal[0]?.usage, undefined); +}); + test("agents default to the first declared phase when the script omits phase()", async () => { // Regression for the "(no phase) has agents, declared phase 0/0" bug: a script // that declares meta.phases but never calls phase() should still group its @@ -475,6 +664,76 @@ return { a, nested }`; assert.equal(result.result.nested.child, "ran:child task"); }); +test("a changed child invalidates later parent resume results", async () => { + const childV1 = `export const meta = { name: 'child', description: 'c' } +return await agent('child original', { label: 'child' })`; + const childV2 = `export const meta = { name: 'child', description: 'c' } +return await agent('child changed', { label: 'child' })`; + const parent = `export const meta = { name: 'parent', description: 'p' } +const before = await agent('parent before', { label: 'before' }) +const child = await workflow('child') +const after = await agent('parent after', { label: 'after' }) +return { before, child, after }`; + const journal: JournalEntry[] = []; + await runWorkflow(parent, { + runId: "nested-miss", + agent: { + async run(prompt: string) { + return `first:${prompt}`; + }, + }, + loadSavedWorkflow: () => childV1, + persistLogs: false, + onAgentJournal: (entry) => journal.push(entry), + }); + + const livePrompts: string[] = []; + const resumed = await runWorkflow<{ before: string; child: string; after: string }>(parent, { + runId: "nested-miss", + agent: { + async run(prompt: string) { + livePrompts.push(prompt); + return `second:${prompt}`; + }, + }, + loadSavedWorkflow: () => childV2, + persistLogs: false, + resumeJournal: new Map(journal.map((entry) => [entry.executionId ?? entry.index, entry] as const)), + }); + + assert.deepEqual(livePrompts, ["child changed", "parent after"]); + assert.equal(resumed.result.before, "first:parent before"); + assert.equal(resumed.result.child, "second:child changed"); + assert.equal(resumed.result.after, "second:parent after"); +}); + +test("legacy numeric journals are never consumed by nested workflows", async () => { + const child = `export const meta = { name: 'child', description: 'child' } +return await agent('same', { label: 'same' })`; + const seedJournal: JournalEntry[] = []; + await runWorkflow(child, { + agent: countingAgent().runner, + runId: "seed", + persistLogs: false, + onAgentJournal: (entry) => seedJournal.push(entry), + }); + const legacyEntry = { ...seedJournal[0], executionId: undefined }; + const counted = countingAgent(); + const parent = `export const meta = { name: 'parent', description: 'parent' } +return await workflow('child')`; + + const result = await runWorkflow(parent, { + agent: counted.runner, + runId: "root", + persistLogs: false, + loadSavedWorkflow: (name) => (name === "child" ? child : undefined), + resumeJournal: new Map([[0, legacyEntry]]), + }); + + assert.equal(counted.state.calls, 1); + assert.equal(result.result, "ran:same"); +}); + test("workflow() nesting is one level deep (second level throws)", async () => { const map: Record = { gc: `export const meta = { name: 'gc', description: 'g' } @@ -497,6 +756,37 @@ return { err }`; assert.match(result.result.err, /one level deep/); }); +test("cached replay succeeds when the live-work token budget is already exhausted", async () => { + const journal: JournalEntry[] = []; + const script = `export const meta = { name: 'cached_budget', description: 'cached budget' } +return await agent('cached', { label: 'cached' })`; + await runWorkflow(script, { + runId: "cached-budget", + agent: fakeAgent({ total: 10 }), + persistLogs: false, + onAgentJournal: (entry) => journal.push(entry), + }); + + let liveCalls = 0; + const resumed = await runWorkflow(script, { + runId: "cached-budget", + agent: { + async run() { + liveCalls++; + return "live"; + }, + }, + tokenBudget: 0, + persistLogs: false, + resumeJournal: new Map( + journal.flatMap((entry) => (entry.executionId ? [[entry.executionId, entry] as const] : [])), + ), + }); + + assert.equal(liveCalls, 0); + assert.equal(resumed.result, "ok"); +}); + test("runWorkflow budget gates on accumulated tokens", async () => { const script = `export const meta = { name: 'budget_demo', description: 'budget' } const a = await agent('first', { label: 'a' }) From 72ac76013b197ce87311f884d016818900b777d9 Mon Sep 17 00:00:00 2001 From: Alexander Penkin Date: Thu, 16 Jul 2026 15:15:38 +0300 Subject: [PATCH 4/4] feat: resume workflow child sessions --- CONTRIBUTING.md | 2 + README.md | 8 +- src/agent.ts | 114 ++++++++++++++++++----- src/run-persistence.ts | 30 ++++++- src/workflow-manager.ts | 76 +++++++++++++++- src/workflow-paths.ts | 3 + src/workflow-settings.ts | 7 +- src/workflow.ts | 97 ++++++++++++++++++-- tests/agent.test.ts | 117 +++++++++++++++++++++--- tests/run-persistence.test.ts | 43 +++++++++ tests/workflow-manager-abort.test.ts | 123 ++++++++++++++++++++++++- tests/workflow-manager.test.ts | 117 ++++++++++++++++++++++++ tests/workflow-paths.test.ts | 1 + tests/workflow-runtime.test.ts | 129 +++++++++++++++++++++++++++ 14 files changed, 815 insertions(+), 52 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 847e84af..53e34134 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -27,6 +27,8 @@ Fake-agent unit tests are necessary but not sufficient. Any change to how agents A throwaway harness for this should live in the repo root (not `/tmp`, whose symlink breaks relative imports), import from `./src`, and be deleted before commit — don't commit harnesses. +Child-session persistence changes must keep the default off and test the private storage path, missing/corrupt-session fallback, durable turn-boundary continuation, and paused-worktree cleanup. Tests must use an isolated fake home and never write transcripts into the developer's normal Pi session index. + ## Style Formatting and linting are handled by Biome (`npm run format`, `npm run lint`). Match the existing code; don't reformat files you aren't otherwise changing. diff --git a/README.md b/README.md index ce84e893..ad36d53f 100644 --- a/README.md +++ b/README.md @@ -77,7 +77,7 @@ return await agent( - **Real parallel orchestration** — fan out up to 16 concurrent and 1000 total subagents from one orchestration script. - **Per-agent model routing** — use `small`, `medium`, or `big` tiers, or choose an exact provider/model and thinking level. -- **Journaled resume** — replay completed agents after interruption without rerunning them or spending their tokens again. The orchestrator can also resume with an **edited script** (`resumeFromRunId`): unchanged `agent()` calls replay from cache and only edited/new ones re-run — so a single bad prompt no longer means paying to re-run the whole workflow. +- **Journaled resume** — replay completed agents without rerunning or repaying them, resume edited scripts from the unchanged prefix, and optionally continue interrupted child sessions from their last durable turn boundary. - **Git worktree isolation** — let parallel agents edit safely on throwaway branches with `isolation: "worktree"`. - **Measured usage** — report real tokens and cost from each subagent session; add run, phase, or agent budgets only when you want them. - **Visible background runs** — track phases, agents, models, fresh/cache tokens, cost, and live tok/s from the progress panel or `/workflows` navigator. @@ -184,9 +184,11 @@ Extension state lives outside the repository under `~/.pi/workflows`: - project runs, journals, locks, and saved overrides: `~/.pi/workflows/projects//` - older project-local `.pi/workflows/runs` and `.pi/workflows/saved` remain readable as fallbacks -Subagents are in-memory by default. Set `persistAgentSessions: true` to retain full transcripts in Pi's standard session directory. This creates one file per agent and may store sensitive material that an agent read, so enable it deliberately. +Subagents are in-memory by default. Set `persistAgentSessions: true` in `~/.pi/workflows/settings.json` (or a project override) to opt in to durable child transcripts. They are stored under the project's private workflow state directory (`~/.pi/workflows/projects//agent-sessions/`), outside Pi's normal `/resume` picker. This creates one file per agent and may store secrets or other sensitive material that an agent read, so enable it deliberately. -Run files use a versioned, backward-compatible state model with stable execution IDs, exact terminal usage, and crash-safe temp/rename writes plus backup recovery. After a crash, orphaned running work is recovered as paused and waits for an explicit resume. Resume replays the longest unchanged completed prefix—including nested workflows—then starts incomplete work fresh, without reopening child transcripts or double-charging completed usage. +Run files use a versioned, backward-compatible state model with stable execution IDs, exact terminal usage, and crash-safe temp/rename writes plus backup recovery. After a crash, orphaned running work is recovered as paused and waits for an explicit resume. Resume replays the longest unchanged completed prefix—including nested workflows—without double-charging completed usage. With child persistence enabled, an interrupted agent reopens its private transcript and compatible isolated worktree; otherwise it starts fresh. Missing, corrupt, or unwritable child sessions also fall back to a fresh session and are reported in run logs. + +Continuation begins at the last durable Pi message/tool-result boundary. Provider streams cannot resume mid-token, and a tool interrupted before its result was durably recorded may have uncertain side effects; the continuation prompt tells the agent to inspect existing state before acting. Completed background runs persist their full result in the project run JSON. The conversation delivery includes a pointer to that file when the visible summary is shortened. diff --git a/src/agent.ts b/src/agent.ts index 8841f472..58d66c08 100644 --- a/src/agent.ts +++ b/src/agent.ts @@ -1,5 +1,5 @@ import { randomUUID } from "node:crypto"; -import { unlinkSync, writeFileSync } from "node:fs"; +import { existsSync, unlinkSync, writeFileSync } from "node:fs"; import { join } from "node:path"; import type { AssistantMessage, Model, TextContent } from "@earendil-works/pi-ai"; import { @@ -28,6 +28,7 @@ import { resolveTierModel, } from "./model-tier-config.js"; import { createStructuredOutputTool, type StructuredOutputCapture } from "./structured-output.js"; +import { workflowProjectPaths } from "./workflow-paths.js"; /** * Find a JSON object/array in free-form text: a fenced ```json block if present, @@ -226,9 +227,8 @@ export interface WorkflowAgentOptions { */ modelRegistry?: ModelRegistry; /** - * Persist each subagent transcript as a real pi session file under the - * standard sessions directory (keyed by the runner's project cwd), instead - * of the default in-memory session that is discarded when the run ends. + * Persist each subagent transcript as a private Pi session file under the + * workflow project's state directory, outside Pi's normal /resume picker. * Default: false (current behavior). */ persistAgentSessions?: boolean; @@ -391,12 +391,19 @@ export function usageFromStats(stats: { }; } +export interface AgentSessionCheckpoint { + /** File-backed Pi child session. Undefined means this invocation cannot resume in place. */ + sessionFile?: string; + /** True when an existing session was reopened rather than created fresh. */ + resumed: boolean; +} + export interface AgentRunOptions { label?: string; /** - * Display name recorded on the persisted session (session_info entry) when - * `persistAgentSessions` is enabled, so transcripts are identifiable in - * session pickers (e.g. `workflow: