Skip to content

Commit af55f2d

Browse files
committed
feat: make workflow runs observable and resumable
1 parent b587566 commit af55f2d

38 files changed

Lines changed: 4811 additions & 588 deletions

CONTRIBUTING.md

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,16 @@ If you add a `workflow` tool parameter or a `~/.pi/workflows/settings.json` sett
2525

2626
Fake-agent unit tests are necessary but not sufficient. Any change to how agents actually run — retries, timeouts, model routing, token accounting, concurrency, resume — must also be verified **end-to-end against a real Pi subagent session** (real `createAgentSession` → real model), because the real SDK path behaves differently than a mock. If you don't have a real-provider environment, say so in the PR and a maintainer will run it before merge.
2727

28-
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.
28+
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, credentials, or provider output. Record the provider/model and observed results in the PR instead.
29+
30+
For workflow runtime reliability changes, exercise every affected item in this real-provider checklist:
31+
32+
1. **Safe typing:** Submit ordinary prompts containing `workflow`, `workflows`, mixed case, punctuation, paths, and slash commands. With fresh/default settings, none may rewrite the prompt or emit a forced-tool message. Explicit `/workflows-trigger on` must restore only word-bounded activation.
33+
2. **Autonomous management:** Let Pi start a background run with `workflow`, then have Pi use `workflow_control` to `list`, inspect `status`, `pause`, `resume`, and `stop` it without asking the user to type `/workflows` commands.
34+
3. **Active fleet:** Start 42 independent agents with Pi selecting `concurrency: 8`. Confirm no more than eight run simultaneously, all active labels appear in `Running now (N/8)`, and remaining work is shown as queued.
35+
4. **Live usage:** Run a multi-turn child. Confirm `~N tok` moves during streaming, exact tokens replace the estimate at assistant-message boundaries, and terminal reconciliation remains exact.
36+
5. **Restart/resume:** Interrupt after at least 27 completed calls, restart Pi, and resume the same run. Confirm completed agents retain status and tokens, aggregate usage remains monotonic, no rows duplicate, and the original concurrency/options remain in effect. Repeat with an older run artifact when changing persistence migration.
37+
6. **Human compatibility:** Confirm `/workflows` commands, navigator pause/stop/restart, `q` close, `x` stop, trigger opt-in, saved workflows, and stale-run recovery still work.
2938

3039
## Style
3140

README.md

Lines changed: 53 additions & 21 deletions
Large diffs are not rendered by default.

extensions/workflow.ts

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import {
1515
saveWorkflowSettingsForCwd,
1616
WorkflowManager,
1717
} from "../src/index.js";
18+
import { createWorkflowControlTool } from "../src/workflow-control-tool.js";
1819

1920
export default function extension(pi: ExtensionAPI) {
2021
// Single manager/storage shared by the workflow tool and the /workflows command,
@@ -32,7 +33,9 @@ export default function extension(pi: ExtensionAPI) {
3233
});
3334

3435
const workflowTool = createWorkflowTool({ cwd, manager, storage });
36+
const workflowControlTool = createWorkflowControlTool({ manager });
3537
pi.registerTool(workflowTool);
38+
pi.registerTool(workflowControlTool);
3639
// Standing /effort opt-in (off|high|ultra): auto-arms a workflow for substantive
3740
// messages, like CC's ultracode. Shared with the editor's input hook below and
3841
// with the explicit /workflows run <prompt> manual trigger.
@@ -58,9 +61,9 @@ export default function extension(pi: ExtensionAPI) {
5861
// advertise the shared registry's models.
5962
manager.setModelRegistry(ctx.modelRegistry);
6063
const active = pi.getActiveTools();
61-
if (!active.includes(workflowTool.name)) {
62-
pi.setActiveTools([...active, workflowTool.name]);
63-
}
64+
const workflowTools = [workflowTool.name, workflowControlTool.name];
65+
const missing = workflowTools.filter((name) => !active.includes(name));
66+
if (missing.length) pi.setActiveTools([...active, ...missing]);
6467
// Scope the /workflows history to this session: runs persist on disk across
6568
// sessions, but the navigator/task panel show only the current session's runs.
6669
// Switching back to a previous session re-shows that session's runs.

src/agent.ts

Lines changed: 159 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
import { randomUUID } from "node:crypto";
2-
import { unlinkSync, writeFileSync } from "node:fs";
2+
import { existsSync, unlinkSync, writeFileSync } from "node:fs";
33
import { join } from "node:path";
44
import type { AssistantMessage, Model, TextContent } from "@earendil-works/pi-ai";
55
import {
6+
type AgentSessionEvent,
67
AuthStorage,
78
type CreateAgentSessionOptions,
89
createAgentSession,
@@ -21,6 +22,7 @@ import { classifyProviderLimit, WorkflowError, WorkflowErrorCode } from "./error
2122
import { canonicalModelSpec, resolveModelSpecWithThinking } from "./model-spec.js";
2223
import { loadModelTierConfig, type ModelTierConfig, resolveTierModel } from "./model-tier-config.js";
2324
import { createStructuredOutputTool, type StructuredOutputCapture } from "./structured-output.js";
25+
import { workflowProjectPaths } from "./workflow-paths.js";
2426

2527
/**
2628
* Find a JSON object/array in free-form text: a fenced ```json block if present,
@@ -215,9 +217,8 @@ export interface WorkflowAgentOptions {
215217
*/
216218
modelRegistry?: ModelRegistry;
217219
/**
218-
* Persist each subagent transcript as a real pi session file under the
219-
* standard sessions directory (keyed by the runner's project cwd), instead
220-
* of the default in-memory session that is discarded when the run ends.
220+
* Persist each subagent transcript as a private pi session file under the
221+
* workflow project's state directory, outside Pi's normal /resume picker.
221222
* Default: false (current behavior).
222223
*/
223224
persistAgentSessions?: boolean;
@@ -243,14 +244,77 @@ export function listAvailableModelSpecs(registry?: ModelRegistry): string[] {
243244
}
244245
}
245246

246-
/** Real token/cost usage for a single subagent run, read from the SDK session. */
247+
/** Token/cost usage for a single subagent run. */
247248
export interface AgentUsage {
248249
input: number;
249250
output: number;
250251
cacheRead: number;
251252
cacheWrite: number;
252253
total: number;
253254
cost: number;
255+
/** True only for an in-progress output-token estimate. */
256+
estimated?: boolean;
257+
}
258+
259+
/**
260+
* Convert session events into absolute cumulative usage. Exact message usage is
261+
* emitted at message_end; throttled message_update events add only a temporary
262+
* output estimate, which the next exact event replaces.
263+
*/
264+
export function createAgentUsageEventHandler(
265+
onUsage: (usage: AgentUsage) => void,
266+
now: () => number = Date.now,
267+
): (event: AgentSessionEvent) => void {
268+
const exact: AgentUsage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0, cost: 0 };
269+
const endedMessages = new WeakSet<object>();
270+
let lastEstimateEmit = Number.NEGATIVE_INFINITY;
271+
const emit = (usage: AgentUsage) => {
272+
try {
273+
onUsage(usage);
274+
} catch {
275+
// Telemetry is best-effort; never interrupt the child session.
276+
}
277+
};
278+
279+
return (event) => {
280+
if (event.type === "message_end" && event.message.role === "assistant") {
281+
if (endedMessages.has(event.message)) return;
282+
endedMessages.add(event.message);
283+
const usage = event.message.usage;
284+
exact.input += usage.input;
285+
exact.output += usage.output;
286+
exact.cacheRead += usage.cacheRead;
287+
exact.cacheWrite += usage.cacheWrite;
288+
exact.total += usage.totalTokens;
289+
exact.cost += usage.cost.total;
290+
lastEstimateEmit = Number.NEGATIVE_INFINITY;
291+
emit({ ...exact, estimated: false });
292+
return;
293+
}
294+
295+
if (event.type !== "message_update" || event.message.role !== "assistant") return;
296+
const timestamp = now();
297+
if (timestamp - lastEstimateEmit < 250) return;
298+
lastEstimateEmit = timestamp;
299+
const text = event.message.content
300+
.map((part) => (part.type === "text" ? part.text : part.type === "thinking" ? part.thinking : ""))
301+
.join("");
302+
const estimatedOutput = Math.ceil(text.length / 4);
303+
if (estimatedOutput <= 0) return;
304+
emit({
305+
...exact,
306+
output: exact.output + estimatedOutput,
307+
total: exact.total + estimatedOutput,
308+
estimated: true,
309+
});
310+
};
311+
}
312+
313+
export interface AgentSessionCheckpoint {
314+
/** File-backed Pi child session. Undefined means this invocation cannot resume in-place. */
315+
sessionFile?: string;
316+
/** True when the existing session was reopened rather than created fresh. */
317+
resumed: boolean;
254318
}
255319

256320
export interface AgentRunOptions<TSchemaDef extends TSchema | undefined = undefined> {
@@ -267,11 +331,15 @@ export interface AgentRunOptions<TSchemaDef extends TSchema | undefined = undefi
267331
instructions?: string;
268332
signal?: AbortSignal;
269333
/**
270-
* Called once with this subagent's real usage, read from the session right
271-
* before disposal. Fires on both the success and error paths so partial
272-
* usage is never lost. `total === 0` means the provider reported no usage.
334+
* Called with absolute cumulative usage during the run and once more with
335+
* authoritative session totals before disposal. Streaming estimates have
336+
* `estimated: true`; message-boundary and terminal updates are exact.
273337
*/
274338
onUsage?: (usage: AgentUsage) => void;
339+
/** Reopen this persisted Pi child session and continue from its last durable turn boundary. */
340+
resumeSessionFile?: string;
341+
/** Called before prompting so workflow persistence can durably link the invocation to its child session. */
342+
onSession?: (checkpoint: AgentSessionCheckpoint) => void;
275343
/**
276344
* Model spec for this subagent: either `provider/modelId` (unambiguous) or a
277345
* bare `modelId`. When it can't be resolved, the session default is used and
@@ -376,9 +444,9 @@ export class WorkflowAgent {
376444
}
377445

378446
/**
379-
* Session manager for one subagent run. File-backed (persisted under the
380-
* standard sessions dir, keyed by the runner's project cwd — never a
381-
* per-call worktree cwd) when persistAgentSessions is on; in-memory otherwise.
447+
* Session manager for one subagent run. File-backed in the workflow project's
448+
* private agent-sessions directory (never the normal Pi /resume directory and
449+
* never a per-call worktree cwd) when persistence is on; in-memory otherwise.
382450
*
383451
* SessionManager.create() only creates the session directory — the SDK writes
384452
* the session file lazily (synchronous fs calls, uncaught) on the first
@@ -391,7 +459,7 @@ export class WorkflowAgent {
391459
private createSessionManager(): SessionManager {
392460
if (!this.persistAgentSessions) return SessionManager.inMemory();
393461
try {
394-
const manager = SessionManager.create(this.cwd);
462+
const manager = SessionManager.create(this.cwd, workflowProjectPaths(this.cwd).agentSessionsDir);
395463
this.assertSessionDirWritable(manager.getSessionDir());
396464
return manager;
397465
} catch (error) {
@@ -404,6 +472,28 @@ export class WorkflowAgent {
404472
}
405473
}
406474

475+
/** Reopen a durable child session when possible; otherwise create the configured fresh session. */
476+
private resolveSessionManager(
477+
resumeSessionFile: string | undefined,
478+
runCwd: string,
479+
): {
480+
manager: SessionManager;
481+
resumed: boolean;
482+
} {
483+
if (resumeSessionFile && existsSync(resumeSessionFile)) {
484+
try {
485+
return { manager: SessionManager.open(resumeSessionFile, undefined, runCwd), resumed: true };
486+
} catch (error) {
487+
console.warn(
488+
`[workflow] could not reopen child session ${resumeSessionFile} (${
489+
error instanceof Error ? error.message : String(error)
490+
}); restarting this agent from a fresh session`,
491+
);
492+
}
493+
}
494+
return { manager: this.createSessionManager(), resumed: false };
495+
}
496+
407497
/** Best-effort write probe: throws if the session directory isn't actually writable. */
408498
private assertSessionDirWritable(dir: string): void {
409499
const probePath = join(dir, `.write-probe-${randomUUID()}`);
@@ -463,11 +553,13 @@ export class WorkflowAgent {
463553
}
464554

465555
const agentDir = getAgentDir();
466-
// Key persisted sessions by the runner's project cwd (this.cwd), NOT the
467-
// per-call runCwd: agents working in short-lived git worktrees should still
468-
// group under the project's session dir instead of scattering across
469-
// temporary worktree paths.
470-
const sessionManager = this.createSessionManager();
556+
// Key new persisted sessions by the runner's project cwd (this.cwd), NOT the
557+
// per-call runCwd. A resumed session is reopened with runCwd as its cwd override
558+
// so coding tools continue in the same shared tree or preserved worktree.
559+
const resolvedSession = this.sessionOptions.sessionManager
560+
? { manager: this.sessionOptions.sessionManager, resumed: false }
561+
: this.resolveSessionManager(options.resumeSessionFile, runCwd);
562+
const sessionManager = resolvedSession.manager;
471563
const { session } = await createAgentSession({
472564
cwd: runCwd,
473565
agentDir,
@@ -489,18 +581,31 @@ export class WorkflowAgent {
489581
...(resolvedThinkingLevel ? { thinkingLevel: resolvedThinkingLevel } : {}),
490582
});
491583

492-
// Name the persisted session so it's identifiable in session pickers.
493-
// Skip when an injected session.sessionManager override won (tests/embedders).
494-
if (this.persistAgentSessions && !this.sessionOptions.sessionManager && options.sessionName) {
584+
// Persist the child-session link before the model starts, so a crash can recover it.
585+
try {
586+
options.onSession?.({ sessionFile: sessionManager.getSessionFile(), resumed: resolvedSession.resumed });
587+
} catch {
588+
// Session-link telemetry is best-effort; never block the child.
589+
}
590+
591+
// Name a newly persisted session so it's identifiable in session pickers.
592+
// Skip reopened sessions and injected managers to avoid duplicate session_info entries.
593+
if (
594+
this.persistAgentSessions &&
595+
!resolvedSession.resumed &&
596+
!this.sessionOptions.sessionManager &&
597+
options.sessionName
598+
) {
495599
try {
496600
sessionManager.appendSessionInfo(options.sessionName);
497601
} catch {
498602
// Naming is best-effort; never fail the run over it.
499603
}
500604
}
501605

606+
const initialStats = resolvedSession.resumed ? session.getSessionStats() : undefined;
502607
let removeAbortListener: (() => void) | undefined;
503-
let removeHistoryListener: (() => void) | undefined;
608+
let removeSessionListener: (() => void) | undefined;
504609
let lastHistoryEmit = 0;
505610
const emitHistory = () => options.onHistory?.(compactAgentHistory(session.messages));
506611
const maybeEmitHistory = () => {
@@ -510,18 +615,26 @@ export class WorkflowAgent {
510615
lastHistoryEmit = now;
511616
emitHistory();
512617
};
618+
const handleUsageEvent = options.onUsage ? createAgentUsageEventHandler(options.onUsage) : undefined;
513619
try {
514620
if (options.signal?.aborted) throw new Error("Subagent was aborted");
515621
if (options.signal) {
516622
const onAbort = () => void session.abort();
517623
options.signal.addEventListener("abort", onAbort, { once: true });
518624
removeAbortListener = () => options.signal?.removeEventListener("abort", onAbort);
519625
}
520-
if (options.onHistory) {
521-
removeHistoryListener = session.subscribe(() => maybeEmitHistory());
626+
if (options.onHistory || handleUsageEvent) {
627+
removeSessionListener = session.subscribe((event) => {
628+
maybeEmitHistory();
629+
handleUsageEvent?.(event);
630+
});
522631
}
523632

524-
await session.prompt(this.buildPrompt(prompt, options as AgentRunOptions<any>, Boolean(options.schema)));
633+
await session.prompt(
634+
resolvedSession.resumed
635+
? this.buildResumePrompt(Boolean(options.schema))
636+
: this.buildPrompt(prompt, options as AgentRunOptions<any>, Boolean(options.schema)),
637+
);
525638

526639
if (options.signal?.aborted) throw new Error("Subagent was aborted");
527640

@@ -547,7 +660,7 @@ export class WorkflowAgent {
547660
return text as AgentRunResult<TSchemaDef>;
548661
} finally {
549662
removeAbortListener?.();
550-
removeHistoryListener?.();
663+
removeSessionListener?.();
551664
try {
552665
emitHistory();
553666
} catch {
@@ -558,12 +671,13 @@ export class WorkflowAgent {
558671
try {
559672
const { tokens, cost } = session.getSessionStats();
560673
options.onUsage({
561-
input: tokens.input,
562-
output: tokens.output,
563-
cacheRead: tokens.cacheRead,
564-
cacheWrite: tokens.cacheWrite,
565-
total: tokens.total,
566-
cost,
674+
input: Math.max(0, tokens.input - (initialStats?.tokens.input ?? 0)),
675+
output: Math.max(0, tokens.output - (initialStats?.tokens.output ?? 0)),
676+
cacheRead: Math.max(0, tokens.cacheRead - (initialStats?.tokens.cacheRead ?? 0)),
677+
cacheWrite: Math.max(0, tokens.cacheWrite - (initialStats?.tokens.cacheWrite ?? 0)),
678+
total: Math.max(0, tokens.total - (initialStats?.tokens.total ?? 0)),
679+
cost: Math.max(0, cost - (initialStats?.cost ?? 0)),
680+
estimated: false,
567681
});
568682
} catch {
569683
// Usage is best-effort; never let stats failure mask the real result/error.
@@ -573,6 +687,20 @@ export class WorkflowAgent {
573687
}
574688
}
575689

690+
private buildResumePrompt(structured: boolean): string {
691+
const parts = [
692+
"Continue the interrupted task from the existing session at the last durable message/tool-result boundary.",
693+
"Review the prior messages and current filesystem state before acting. Do not repeat completed side effects.",
694+
"Finish the original task and return its requested final result.",
695+
];
696+
if (structured) {
697+
parts.push(
698+
"When finished, call structured_output with the required schema, even if it was called before interruption.",
699+
);
700+
}
701+
return parts.join("\n\n");
702+
}
703+
576704
private buildPrompt(prompt: string, options: AgentRunOptions<any>, structured: boolean): string {
577705
const parts = [
578706
this.instructions,

src/config.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@ export const MAX_AGENTS_PER_RUN = 1000;
88
/** Default timeout for a single agent in milliseconds. null means no hard timeout. */
99
export const DEFAULT_AGENT_TIMEOUT_MS = null;
1010

11-
/** Maximum concurrent agents (matches Claude Code limit). */
12-
export const MAX_CONCURRENCY = 16;
11+
/** Concurrent agents used when a run does not specify a limit. */
12+
export const DEFAULT_CONCURRENCY = 16;
1313

1414
/** Maximum automatic retry attempts after a recoverable agent failure. */
1515
export const MAX_AGENT_RETRIES = 3;

0 commit comments

Comments
 (0)