Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`, `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 |
| --- | --- |
Expand Down Expand Up @@ -184,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.

</details>
Expand Down
9 changes: 6 additions & 3 deletions extensions/workflow.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import {
createEffortState,
createWorkflowControlTool,
createWorkflowStorage,
createWorkflowTool,
installResultDelivery,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down
84 changes: 73 additions & 11 deletions src/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -300,14 +301,72 @@ 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;
cacheRead: number;
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<object>();
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,
});
};
}

/**
Expand Down Expand Up @@ -346,10 +405,9 @@ export interface AgentRunOptions<TSchemaDef extends TSchema | undefined = undefi
instructions?: string;
signal?: AbortSignal;
/**
* Called once with this subagent's real usage, read from the session right
* before disposal. Fires on both the success and error paths so partial
* usage is never lost — but NOT when the provider reported no usage at all
* (all-zero stats), so consumers keep their scalar fallback.
* Called with absolute cumulative usage during the run and once more with
* authoritative session totals before disposal. Streaming estimates have
* `estimated: true`; message-boundary and terminal updates are exact.
*/
onUsage?: (usage: AgentUsage) => void;
/**
Expand Down Expand Up @@ -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 = () => {
Expand All @@ -593,15 +651,19 @@ 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) {
const onAbort = () => void session.abort();
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<any>, Boolean(options.schema)));
Expand Down Expand Up @@ -630,17 +692,17 @@ export class WorkflowAgent {
return text as AgentRunResult<TSchemaDef>;
} 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.
}
Expand Down
13 changes: 11 additions & 2 deletions src/display.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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. */
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -338,6 +345,8 @@ export function statusIcon(status: WorkflowAgentStatus): string {
return "○";
case "running":
return "●";
case "paused":
return "⏸";
case "done":
return "✓";
case "error":
Expand Down
6 changes: 6 additions & 0 deletions src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading