Skip to content
Closed
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
11 changes: 10 additions & 1 deletion CONTRIBUTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,16 @@ If you add a `workflow` tool parameter or a `~/.pi/workflows/settings.json` sett

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.

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.
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.

For workflow runtime reliability changes, exercise every affected item in this real-provider checklist:

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.
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.
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.
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.
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.
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.

## Style

Expand Down
74 changes: 53 additions & 21 deletions README.md

Large diffs are not rendered by default.

9 changes: 6 additions & 3 deletions extensions/workflow.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import {
saveWorkflowSettingsForCwd,
WorkflowManager,
} from "../src/index.js";
import { createWorkflowControlTool } from "../src/workflow-control-tool.js";

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

const workflowTool = createWorkflowTool({ cwd, manager, storage });
const workflowControlTool = createWorkflowControlTool({ manager });
pi.registerTool(workflowTool);
pi.registerTool(workflowControlTool);
// Standing /effort opt-in (off|high|ultra): auto-arms a workflow for substantive
// messages, like CC's ultracode. Shared with the editor's input hook below and
// with the explicit /workflows run <prompt> manual trigger.
Expand All @@ -58,9 +61,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
190 changes: 159 additions & 31 deletions src/agent.ts
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
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 {
type AgentSessionEvent,
AuthStorage,
type CreateAgentSessionOptions,
createAgentSession,
Expand All @@ -21,6 +22,7 @@ import { classifyProviderLimit, WorkflowError, WorkflowErrorCode } from "./error
import { canonicalModelSpec, resolveModelSpecWithThinking } from "./model-spec.js";
import { loadModelTierConfig, type ModelTierConfig, 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,
Expand Down Expand Up @@ -215,9 +217,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;
Expand All @@ -243,14 +244,77 @@ export function listAvailableModelSpecs(registry?: ModelRegistry): string[] {
}
}

/** 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;
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 text = event.message.content
.map((part) => (part.type === "text" ? part.text : part.type === "thinking" ? part.thinking : ""))
.join("");
const estimatedOutput = Math.ceil(text.length / 4);
if (estimatedOutput <= 0) return;
emit({
...exact,
output: exact.output + estimatedOutput,
total: exact.total + estimatedOutput,
estimated: true,
});
};
}

export interface AgentSessionCheckpoint {
/** File-backed Pi child session. Undefined means this invocation cannot resume in-place. */
sessionFile?: string;
/** True when the existing session was reopened rather than created fresh. */
resumed: boolean;
}

export interface AgentRunOptions<TSchemaDef extends TSchema | undefined = undefined> {
Expand All @@ -267,11 +331,15 @@ 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. `total === 0` means the provider reported no usage.
* 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;
/** Reopen this persisted Pi child session and continue from its last durable turn boundary. */
resumeSessionFile?: string;
/** Called before prompting so workflow persistence can durably link the invocation to its child session. */
onSession?: (checkpoint: AgentSessionCheckpoint) => void;
/**
* Model spec for this subagent: either `provider/modelId` (unambiguous) or a
* bare `modelId`. When it can't be resolved, the session default is used and
Expand Down Expand Up @@ -376,9 +444,9 @@ export class WorkflowAgent {
}

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

/** Reopen a durable child session when possible; otherwise create the configured fresh session. */
private resolveSessionManager(
resumeSessionFile: string | undefined,
runCwd: string,
): {
manager: SessionManager;
resumed: boolean;
} {
if (resumeSessionFile && existsSync(resumeSessionFile)) {
try {
return { manager: SessionManager.open(resumeSessionFile, undefined, runCwd), resumed: true };
} catch (error) {
console.warn(
`[workflow] could not reopen child session ${resumeSessionFile} (${
error instanceof Error ? error.message : String(error)
}); restarting this agent from a fresh session`,
);
}
}
return { manager: this.createSessionManager(), resumed: false };
}

/** Best-effort write probe: throws if the session directory isn't actually writable. */
private assertSessionDirWritable(dir: string): void {
const probePath = join(dir, `.write-probe-${randomUUID()}`);
Expand Down Expand Up @@ -463,11 +553,13 @@ export class WorkflowAgent {
}

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

// Name the persisted session so it's identifiable in session pickers.
// Skip when an injected session.sessionManager override won (tests/embedders).
if (this.persistAgentSessions && !this.sessionOptions.sessionManager && options.sessionName) {
// Persist the child-session link before the model starts, so a crash can recover it.
try {
options.onSession?.({ sessionFile: sessionManager.getSessionFile(), resumed: resolvedSession.resumed });
} catch {
// Session-link telemetry is best-effort; never block the child.
}

// Name a newly persisted session so it's identifiable in session pickers.
// Skip reopened sessions and injected managers to avoid duplicate session_info entries.
if (
this.persistAgentSessions &&
!resolvedSession.resumed &&
!this.sessionOptions.sessionManager &&
options.sessionName
) {
try {
sessionManager.appendSessionInfo(options.sessionName);
} catch {
// Naming is best-effort; never fail the run over it.
}
}

const initialStats = resolvedSession.resumed ? session.getSessionStats() : undefined;
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 @@ -510,18 +615,26 @@ 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)));
await session.prompt(
resolvedSession.resumed
? this.buildResumePrompt(Boolean(options.schema))
: this.buildPrompt(prompt, options as AgentRunOptions<any>, Boolean(options.schema)),
);

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

Expand All @@ -547,7 +660,7 @@ export class WorkflowAgent {
return text as AgentRunResult<TSchemaDef>;
} finally {
removeAbortListener?.();
removeHistoryListener?.();
removeSessionListener?.();
try {
emitHistory();
} catch {
Expand All @@ -558,12 +671,13 @@ export class WorkflowAgent {
try {
const { tokens, cost } = session.getSessionStats();
options.onUsage({
input: tokens.input,
output: tokens.output,
cacheRead: tokens.cacheRead,
cacheWrite: tokens.cacheWrite,
total: tokens.total,
cost,
input: Math.max(0, tokens.input - (initialStats?.tokens.input ?? 0)),
output: Math.max(0, tokens.output - (initialStats?.tokens.output ?? 0)),
cacheRead: Math.max(0, tokens.cacheRead - (initialStats?.tokens.cacheRead ?? 0)),
cacheWrite: Math.max(0, tokens.cacheWrite - (initialStats?.tokens.cacheWrite ?? 0)),
total: Math.max(0, tokens.total - (initialStats?.tokens.total ?? 0)),
cost: Math.max(0, cost - (initialStats?.cost ?? 0)),
estimated: false,
});
} catch {
// Usage is best-effort; never let stats failure mask the real result/error.
Expand All @@ -573,6 +687,20 @@ export class WorkflowAgent {
}
}

private buildResumePrompt(structured: boolean): string {
const parts = [
"Continue the interrupted task from the existing session at the last durable message/tool-result boundary.",
"Review the prior messages and current filesystem state before acting. Do not repeat completed side effects.",
"Finish the original task and return its requested final result.",
];
if (structured) {
parts.push(
"When finished, call structured_output with the required schema, even if it was called before interruption.",
);
}
return parts.join("\n\n");
}

private buildPrompt(prompt: string, options: AgentRunOptions<any>, structured: boolean): string {
const parts = [
this.instructions,
Expand Down
4 changes: 2 additions & 2 deletions src/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@ export const MAX_AGENTS_PER_RUN = 1000;
/** Default timeout for a single agent in milliseconds. null means no hard timeout. */
export const DEFAULT_AGENT_TIMEOUT_MS = null;

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

/** Maximum automatic retry attempts after a recoverable agent failure. */
export const MAX_AGENT_RETRIES = 3;
Expand Down
Loading