@chrok/pi-braid 0.4.0 → 0.5.0

package/dist/jobs.js

@@ -1,14 +1,20 @@
 import { setImmediate as nextTurn } from "node:timers/promises";
-import { ExecutionActivityTracker, startBraid, validateGraph, } from "@chrok/braid";
+import { ExecutionActivityTracker, formatActivityStatus, startBraid, validateGraph, } from "@chrok/braid";
 import { applyEvent, applyProgress, createLiveState, } from "./display.js";
 import { createPiRunner, sumPiUsage } from "./runner.js";
 import { PiTranscripts } from "./transcript.js";
 import { PiHumanInputs } from "./human-input.js";
 import { BraidAssistant } from "./assistant.js";
+function timedOut(state) {
+    const ids = (state?.suspendedExecutionIds ?? []).filter(id => state.executions[id]?.suspension?.source === "node_timeout");
+    return ids.length ? { timedOutExecutionIds: ids } : {};
+}
 /** Jobs belong to one extension/session lifetime, independently of foreground turns. */
 export class BraidJobs {
     onFinished;
     onNodeFinished;
+    onControl;
+    onTimeoutPause;
     jobs = new Map();
     assistants = new Map();
     cwd = process.cwd();
@@ -16,9 +22,24 @@ export class BraidJobs {
     nextHandle = 1;
     listeners = new Set();
     disposed = false;
-    constructor(onFinished = () => { }, onNodeFinished = () => { }) {
+    constructor(onFinished = () => { }, onNodeFinished = () => { }, onControl = () => { }, onTimeoutPause = () => { }) {
         this.onFinished = onFinished;
         this.onNodeFinished = onNodeFinished;
+        this.onControl = onControl;
+        this.onTimeoutPause = onTimeoutPause;
+    }
+    /** The parent made the change itself; anyone else's change becomes a notice for it. */
+    notify(job, action, source, executionIds) {
+        if (source === "agent" || this.disposed)
+            return;
+        const state = job.run?.snapshot();
+        try {
+            this.onControl({
+                jobId: job.jobId, handle: job.handle, action, source,
+                ...(executionIds ? { executions: executionIds.map(executionId => ({ executionId, nodeId: state?.executions[executionId]?.id ?? "?" })) } : {}),
+            });
+        }
+        catch { /* Notices are advisory. */ }
     }
     subscribe(listener) {
         this.listeners.add(listener);
@@ -36,12 +57,10 @@ export class BraidJobs {
             }
         }
     }
-    start(input, options, ctx, blocking = false) {
+    start(input, options, ctx) {
         if (this.disposed)
             throw new Error("Braid session has closed");
         validateGraph(input);
-        if (blocking && input.nodes.some(node => node.pauseAfter))
-            throw new Error("Blocking Braid jobs cannot use pauseAfter; remove it or disable blocking with /braid:blocking off.");
         const snapshot = structuredClone(input);
         const settings = { ...options };
         const model = ctx.model
@@ -51,7 +70,6 @@ export class BraidJobs {
         const cwd = ctx.cwd;
         this.cwd = cwd;
         const job = {
-            blocking,
             handle: `job-${this.nextHandle++}`,
             jobId: crypto.randomUUID(),
             goal: snapshot.goal,
@@ -62,6 +80,8 @@ export class BraidJobs {
             done: Promise.resolve(),
             usageClaimed: false,
             runningExecutions: new Map(),
+            suspendedExecutions: new Map(),
+            timeoutActivity: new Map(),
             activity: new ExecutionActivityTracker({ onChange: () => this.changed() }),
             transcripts: new PiTranscripts(() => this.changed()),
             humanInputs: new PiHumanInputs(() => this.changed()),
@@ -92,6 +112,7 @@ export class BraidJobs {
                         humanInputs: job.humanInputs,
                         maxToolRounds: options.maxToolRounds ?? Infinity,
                         maxToolCalls: options.maxToolCalls ?? Infinity,
+                        ...(options.maxModelRetries === undefined ? {} : { maxModelRetries: options.maxModelRetries }),
                         onUsage: (usage) => {
                             if (job.status === "running" && !job.controller.signal.aborted)
                                 reports.push(usage);
@@ -103,7 +124,7 @@ export class BraidJobs {
                             this.changed();
                         },
                     });
-                    return async (request) => {
+                    return Object.assign(async (request) => {
                         try {
                             await nextTurn();
                             return await invoke(request);
@@ -116,9 +137,15 @@ export class BraidJobs {
                             }
                             job.transcripts.finish(request.execution.executionId, request.signal.aborted ? "Execution was stopped before this step finished" : undefined);
                         }
-                    };
+                    }, { pausable: invoke.pausable === true, interruptible: invoke.interruptible === true });
                 })(),
                 onEvent: (event) => {
+                    // Capture what the worker was doing before the interruption changes it.
+                    if (event.type === "execution_suspending" && event.source === "node_timeout" && event.executionId) {
+                        const activity = job.activity.get(event.executionId);
+                        if (activity)
+                            job.timeoutActivity.set(event.executionId, formatActivityStatus(activity, event.timestamp));
+                    }
                     job.activity.observe(event);
                     if (event.type === "workspace_updated") {
                         job.workspaces ??= {};
@@ -127,14 +154,25 @@ export class BraidJobs {
                         });
                     }
                     // Definition edits can remove nodes whose executions are still running.
+                    // A suspended execution is not running: drain can then see a stalled job.
                     if (event.type === "node_started" && event.executionId) {
                         job.runningExecutions.set(event.executionId, event.nodeId);
                     }
+                    else if (event.type === "execution_suspended" && event.executionId) {
+                        job.runningExecutions.delete(event.executionId);
+                        job.suspendedExecutions.set(event.executionId, event.nodeId);
+                    }
+                    else if (event.type === "execution_resumed" && event.executionId && job.suspendedExecutions.has(event.executionId)) {
+                        job.runningExecutions.set(event.executionId, event.nodeId);
+                        job.suspendedExecutions.delete(event.executionId);
+                    }
                     else if ((event.type === "node_completed" || event.type === "node_failed") && event.executionId) {
                         job.runningExecutions.delete(event.executionId);
+                        job.suspendedExecutions.delete(event.executionId);
                     }
                     else if (event.type === "graph_completed" || event.type === "graph_failed") {
                         job.runningExecutions.clear();
+                        job.suspendedExecutions.clear();
                     }
                     applyEvent(job.live, event);
                     if ("nodeId" in event && Object.hasOwn(job.live.nodes, event.nodeId)) {
@@ -147,13 +185,31 @@ export class BraidJobs {
                         });
                     }
                     this.changed();
+                    if (event.type === "execution_suspended" && event.source === "node_timeout" && event.executionId) {
+                        const instance = job.run?.snapshot().executions[event.executionId];
+                        const timeout = instance?.suspension?.timeout;
+                        const activity = job.timeoutActivity.get(event.executionId);
+                        job.timeoutActivity.delete(event.executionId);
+                        // A timeout pause always wakes the parent, which must resume or cancel it.
+                        if (!this.disposed && instance && timeout) {
+                            try {
+                                void Promise.resolve(this.onTimeoutPause({
+                                    jobId: job.jobId, handle: job.handle, nodeId: event.nodeId, executionId: event.executionId,
+                                    ...(instance.iteration ? { iteration: instance.iteration } : {}),
+                                    nodeTimeoutMs: timeout.nodeTimeoutMs, activeMs: timeout.activeMs,
+                                    ...(activity ? { activity } : {}), eventSequence: event.sequence,
+                                })).catch(() => { });
+                            }
+                            catch { /* Notification failures must not affect scheduling. */ }
+                        }
+                    }
                     const terminal = event.type === "execution_paused" || event.type === "node_completed" || event.type === "node_failed";
                     const state = terminal ? job.run?.snapshot() : undefined;
                     const instance = event.executionId ? state?.executions[event.executionId] : undefined;
                     const shouldNotify = event.type === "execution_paused" ||
                         ((event.type === "node_completed" || event.type === "node_failed") && instance?.node.notifyOnCompletion &&
-                            (!instance.node.pauseAfter || state?.error));
-                    if (!this.disposed && !job.blocking && shouldNotify && "nodeId" in event && instance) {
+                            (!instance.holdAfter || state?.error));
+                    if (!this.disposed && shouldNotify && "nodeId" in event && instance) {
                         try {
                             void Promise.resolve(this.onNodeFinished({
                                 jobId: job.jobId,
@@ -188,12 +244,13 @@ export class BraidJobs {
             job.live.status = "failed";
             job.live.pausedExecutionIds = [];
             job.runningExecutions.clear();
+            job.suspendedExecutions.clear();
             job.error = error instanceof Error ? error.message : String(error);
             if (reports.length)
                 job.usage = sumPiUsage(reports);
         }
         this.changed();
-        if (!this.disposed && (!job.blocking || job.status === "cancelled")) {
+        if (!this.disposed) {
             try {
                 this.onFinished(this.get(job.jobId));
             }
@@ -213,7 +270,7 @@ export class BraidJobs {
         const job = this.lookup(jobId);
         if (!job)
             return undefined;
-        const { blocking: _blocking, run: _run, controller: _controller, done: _done, usageClaimed: _claimed, nodeResults: _nodeResults, runningExecutions: _runningExecutions, activity: _activity, transcripts: _transcripts, humanInputs: _humanInputs, ...snapshot } = job;
+        const { run: _run, controller: _controller, done: _done, usageClaimed: _claimed, nodeResults: _nodeResults, runningExecutions: _runningExecutions, suspendedExecutions: _suspendedExecutions, activity: _activity, transcripts: _transcripts, humanInputs: _humanInputs, ...snapshot } = job;
         const copy = structuredClone(snapshot);
         if (job.run)
             copy.execution = job.run.snapshot();
@@ -292,13 +349,60 @@ export class BraidJobs {
         this.changed();
         return this.get(jobId);
     }
-    resume(jobId, executionIds, expectedRevision) {
+    /** Release held executions or resume paused ones. */
+    resume(jobId, executionIds, expectedRevision, source = "agent", options) {
         const job = this.lookup(jobId);
         if (!job)
             throw this.unknownJob(jobId);
         if (!job.run)
             throw new Error("Job is not accepting updates");
-        job.run.resume(executionIds, expectedRevision);
+        job.run.resume(executionIds, expectedRevision, options);
+        this.notify(job, "resumed", source, executionIds);
+        this.changed();
+        return this.get(jobId);
+    }
+    /**
+     * Pause running executions at their next model request, or the whole job
+     * when executionIds is omitted. source records who asked, e.g. "agent" or "user".
+     */
+    pause(jobId, executionIds, source) {
+        const job = this.lookup(jobId);
+        if (!job)
+            throw this.unknownJob(jobId);
+        if (!job.run || job.status !== "running")
+            throw new Error("Job is not accepting updates");
+        if (executionIds)
+            job.run.pause(executionIds, { source });
+        else
+            job.run.pauseRun({ source });
+        this.notify(job, "paused", source, executionIds);
+        this.changed();
+        return this.get(jobId);
+    }
+    /**
+     * Hold running executions' results once they finish (hold = true), or remove
+     * that hold, without suspending them. source records who asked.
+     */
+    holdAfter(jobId, executionIds, hold, source) {
+        const job = this.lookup(jobId);
+        if (!job)
+            throw this.unknownJob(jobId);
+        if (!job.run || job.status !== "running")
+            throw new Error("Job is not accepting updates");
+        job.run.holdAfter(executionIds, hold, { source });
+        this.notify(job, hold ? "held" : "unheld", source, executionIds);
+        this.changed();
+        return this.get(jobId);
+    }
+    /** Admit work again and resume the executions the job pause paused. */
+    resumeJob(jobId, source = "agent") {
+        const job = this.lookup(jobId);
+        if (!job)
+            throw this.unknownJob(jobId);
+        if (!job.run || job.status !== "running")
+            throw new Error("Job is not accepting updates");
+        job.run.resumeRun();
+        this.notify(job, "resumed", source);
         this.changed();
         return this.get(jobId);
     }
@@ -315,7 +419,8 @@ export class BraidJobs {
     }
     /**
      * Jobs that have not reached a terminal status. A job is stalled when only
-     * paused executions remain, so it cannot progress without resume or cancel.
+     * held, suspended, or job-paused work remains, so it cannot progress without
+     * resume or cancel.
      */
     outstanding() {
         return [...this.jobs.values()].filter(job => job.status === "running").map(job => {
@@ -327,6 +432,9 @@ export class BraidJobs {
                 stalled: state?.status === "waiting",
                 ...(state ? { revision: state.revision } : {}),
                 pausedExecutionIds: state?.pausedExecutionIds ?? [],
+                suspendedExecutionIds: state?.suspendedExecutionIds ?? [],
+                ...timedOut(state),
+                ...(state?.pause ? { paused: true } : {}),
             };
         });
     }
@@ -341,7 +449,7 @@ export class BraidJobs {
             done: nodes.filter(node => node.status === "completed" || node.status === "failed" || node.status === "skipped").length,
             failed: nodes.filter(node => node.status === "failed").length,
             running: [...new Set(job.runningExecutions.values())],
-            paused: job.live.pausedExecutionIds?.length ?? 0,
+            paused: (job.live.pausedExecutionIds?.length ?? 0) + job.suspendedExecutions.size,
         };
     }
     cancel(jobId) {
@@ -353,6 +461,18 @@ export class BraidJobs {
         job.controller.abort();
         return true;
     }
+    /** Fail running or paused executions only; their successors follow the usual failure semantics. */
+    cancelExecutions(jobId, executionIds, source) {
+        const job = this.lookup(jobId);
+        if (!job)
+            throw this.unknownJob(jobId);
+        if (!job.run || job.status !== "running")
+            throw new Error("Job is not accepting updates");
+        job.run.cancelExecutions(executionIds, { source });
+        this.notify(job, "cancelled", source, executionIds);
+        this.changed();
+        return this.get(jobId);
+    }
     /** Pi accounts usage only once, on the first retrieval of a terminal result. */
     claimUsage(jobId) {
         const job = this.lookup(jobId);

package/dist/runner.js

@@ -1,10 +1,28 @@
 import { resolve } from "node:path";
-import { StringEnum, Type, validateToolCall, } from "@earendil-works/pi-ai";
-import { formatBudgetReminder, gitToolDefinition, finishMergeToolDefinition, mergeInstructions, parseGitToolArguments, parseFinishMergeArguments, } from "@chrok/braid";
+import { StringEnum, Type, isRetryableAssistantError, retryDelayMs, validateToolCall, } from "@earendil-works/pi-ai";
+import { formatBudgetReminder, interruptedToolResult, isStepInterrupted, notRunToolResult, gitToolDefinition, finishMergeToolDefinition, mergeInstructions, parseGitToolArguments, parseFinishMergeArguments, } from "@chrok/braid";
 import { createWorktreeWriteTools } from "./write-tools.js";
 import { createAvailableReadTools } from "./read-tools.js";
 import { createWorkspaceShellTools } from "./shell-tools.js";
 import { formatHumanInput } from "./human-input.js";
+const DEFAULT_MODEL_RETRIES = 3;
+/** Caps both Braid's backoff and a provider's Retry-After; longer requested waits fail immediately. */
+const MAX_MODEL_RETRY_DELAY_MS = 60_000;
+function abortableDelay(ms, signal) {
+    return new Promise((resolve, reject) => {
+        if (signal.aborted)
+            return reject(signal.reason);
+        const onAbort = () => {
+            clearTimeout(timer);
+            reject(signal.reason);
+        };
+        const timer = setTimeout(() => {
+            signal.removeEventListener("abort", onAbort);
+            resolve();
+        }, ms);
+        signal.addEventListener("abort", onAbort, { once: true });
+    });
+}
 /**
  * Present a Git result the way Pi's bash tool presents command output. Output
  * stays byte-exact, so a diff can be passed back to git apply.
@@ -30,6 +48,14 @@ export function createPiRunner(registry, onUsageOrOptions, cwd = process.cwd())
             throw new TypeError(`${name} must be Infinity or a positive safe integer`);
         }
     }
+    const maxModelRetries = options.maxModelRetries ?? DEFAULT_MODEL_RETRIES;
+    const modelRetryDelayMs = options.modelRetryDelayMs ?? 2000;
+    if (!Number.isSafeInteger(maxModelRetries) || maxModelRetries < 0) {
+        throw new TypeError("maxModelRetries must be a non-negative safe integer");
+    }
+    if (!Number.isFinite(modelRetryDelayMs) || modelRetryDelayMs < 0) {
+        throw new TypeError("modelRetryDelayMs must be a non-negative finite number");
+    }
     const onUsage = options.onUsage;
     const onProgress = options.onProgress;
     const reportProgress = (progress) => {
@@ -45,7 +71,7 @@ export function createPiRunner(registry, onUsageOrOptions, cwd = process.cwd())
         messages: context.messages,
         tools: context.tools,
     }).length / 4);
-    return async (request) => {
+    const runner = async (request) => {
         const name = request.model;
         const slash = name?.indexOf("/") ?? -1;
         if (!name || slash < 1 || slash === name.length - 1) {
@@ -182,8 +208,10 @@ export function createPiRunner(registry, onUsageOrOptions, cwd = process.cwd())
         let decided = false;
         let actualModel = name;
         let reportedContextTokens;
-        const complete = async () => {
+        const prepareRequest = async () => {
             request.signal.throwIfAborted();
+            // A pause takes effect here; messages sent while suspended are delivered after resume.
+            await request.checkpoint?.();
             deliverHumanInputs();
             const budgets = [];
             if (Number.isFinite(maxToolRounds)) {
@@ -195,6 +223,7 @@ export function createPiRunner(registry, onUsageOrOptions, cwd = process.cwd())
             if (budgets.length > 0) {
                 budgets.push("When a tool budget reaches zero, make no further tool calls and return your final answer. Reserve budget for decide or finish_merge if required.");
             }
+            // Rebuilt per attempt: a suspension moves the node deadline.
             context.systemPrompt = systemPrompt + formatBudgetReminder(request, budgets);
             reportProgress({
                 nodeId: request.node.id,
@@ -206,82 +235,125 @@ export function createPiRunner(registry, onUsageOrOptions, cwd = process.cwd())
                 toolRounds,
                 phase: "model",
             });
-            // The registry resolves stored API keys, OAuth, custom headers and endpoints.
-            // complete() is stream().result(); streaming only adds observation.
-            const requestOptions = {
-                signal: request.signal,
-                sessionId,
-                cacheRetention: "none",
-                transport: "sse",
-                maxRetries: 0,
-            };
-            activity?.modelRequest();
-            let response;
-            try {
-                if (registry.stream) {
-                    const stream = registry.stream(model, context, requestOptions);
-                    for await (const event of stream) {
-                        if (event.type !== "done" && event.type !== "error")
-                            transcripts?.streaming(executionId, event.partial);
-                        if (event.type === "text_delta")
-                            activity?.modelStream("text", event.delta);
-                        else if (event.type === "thinking_delta")
-                            activity?.modelStream("reasoning", event.delta);
-                        else if (event.type === "toolcall_start" || event.type === "toolcall_delta") {
-                            const block = event.partial.content[event.contentIndex];
-                            activity?.modelStream("tool_call", event.type === "toolcall_delta" ? event.delta : "", block?.type === "toolCall" ? block.name : undefined);
+        };
+        const complete = async () => {
+            let retries = 0;
+            while (true) {
+                // Retries pass the checkpoint too, so a pause during backoff stops before the resend.
+                await prepareRequest();
+                // A node timeout pause interrupts only this step; the request is sent again after resume.
+                const step = request.stepSignal?.() ?? request.signal;
+                // The registry resolves stored API keys, OAuth, custom headers and endpoints.
+                // complete() is stream().result(); streaming only adds observation.
+                // pi-ai retries rejected requests before a response starts, honoring Retry-After.
+                const requestOptions = {
+                    signal: step,
+                    sessionId,
+                    cacheRetention: "none",
+                    transport: "sse",
+                    maxRetries: maxModelRetries,
+                    maxRetryDelayMs: MAX_MODEL_RETRY_DELAY_MS,
+                };
+                activity?.modelRequest();
+                let response;
+                // Without a stream we cannot tell where a request failed, so Braid may retry it too.
+                let started = !registry.stream;
+                try {
+                    if (registry.stream) {
+                        const stream = registry.stream(model, context, requestOptions);
+                        for await (const event of stream) {
+                            if (event.type === "start")
+                                started = true;
+                            if (event.type !== "done" && event.type !== "error")
+                                transcripts?.streaming(executionId, event.partial);
+                            if (event.type === "text_delta")
+                                activity?.modelStream("text", event.delta);
+                            else if (event.type === "thinking_delta")
+                                activity?.modelStream("reasoning", event.delta);
+                            else if (event.type === "toolcall_start" || event.type === "toolcall_delta") {
+                                const block = event.partial.content[event.contentIndex];
+                                activity?.modelStream("tool_call", event.type === "toolcall_delta" ? event.delta : "", block?.type === "toolCall" ? block.name : undefined);
+                            }
                         }
+                        response = await stream.result();
+                    }
+                    else {
+                        response = await registry.complete(model, context, requestOptions);
                     }
-                    response = await stream.result();
                 }
-                else {
-                    response = await registry.complete(model, context, requestOptions);
+                catch (error) {
+                    if (isStepInterrupted(request, step)) {
+                        activity?.modelError("Model request interrupted by the node timeout");
+                        continue;
+                    }
+                    activity?.modelError(error instanceof Error ? error.message : String(error));
+                    throw error;
                 }
-            }
-            catch (error) {
-                activity?.modelError(error instanceof Error ? error.message : String(error));
-                throw error;
-            }
-            transcripts?.assistant(executionId, response);
-            activity?.modelResponse({
-                text: response.content.flatMap(block => block.type === "text" ? [block.text] : []).join("\n\n"),
-                reasoning: response.content.flatMap(block => block.type === "thinking" ? [block.thinking] : []).join("\n\n"),
-                toolCalls: response.content.flatMap(block => block.type === "toolCall" ? [{ id: block.id, name: block.name }] : []),
-                stopReason: response.stopReason,
-                inputTokens: response.usage.input + response.usage.cacheRead + response.usage.cacheWrite,
-                outputTokens: response.usage.output,
-            });
-            if (response.stopReason !== "stop" && response.stopReason !== "toolUse") {
-                activity?.modelError(response.errorMessage || `Pi model ended with '${response.stopReason}'`);
-            }
-            actualModel = `${response.provider}/${response.responseModel ?? response.model}`;
-            const providerContextTokens = response.usage.input +
-                response.usage.cacheRead +
-                response.usage.cacheWrite;
-            if (providerContextTokens)
-                reportedContextTokens = providerContextTokens;
-            reportProgress({
-                nodeId: request.node.id,
-                ...(request.execution.executionId ? { executionId: request.execution.executionId } : {}),
-                contextTokens: reportedContextTokens ?? estimateContextTokens(context),
-                contextWindow: model.contextWindow,
-                contextSource: reportedContextTokens !== undefined ? "reported" : "estimate",
-                toolCalls,
-                toolRounds,
-                phase: "model",
-            });
-            onUsage?.(structuredClone(response.usage));
-            usage.inputTokens +=
-                response.usage.input +
+                // A failed attempt stays visible in the session view but never enters the context.
+                transcripts?.assistant(executionId, response);
+                activity?.modelResponse({
+                    text: response.content.flatMap(block => block.type === "text" ? [block.text] : []).join("\n\n"),
+                    reasoning: response.content.flatMap(block => block.type === "thinking" ? [block.thinking] : []).join("\n\n"),
+                    toolCalls: response.content.flatMap(block => block.type === "toolCall" ? [{ id: block.id, name: block.name }] : []),
+                    stopReason: response.stopReason,
+                    inputTokens: response.usage.input + response.usage.cacheRead + response.usage.cacheWrite,
+                    outputTokens: response.usage.output,
+                });
+                if (response.stopReason !== "stop" && response.stopReason !== "toolUse") {
+                    activity?.modelError(response.errorMessage || `Pi model ended with '${response.stopReason}'`);
+                }
+                actualModel = `${response.provider}/${response.responseModel ?? response.model}`;
+                const providerContextTokens = response.usage.input +
                     response.usage.cacheRead +
                     response.usage.cacheWrite;
-            usage.outputTokens += response.usage.output;
-            request.signal.throwIfAborted();
-            if (response.stopReason !== "stop" && response.stopReason !== "toolUse") {
-                throw new Error(response.errorMessage ||
-                    `Pi model ended with '${response.stopReason}'`);
+                if (providerContextTokens)
+                    reportedContextTokens = providerContextTokens;
+                reportProgress({
+                    nodeId: request.node.id,
+                    ...(request.execution.executionId ? { executionId: request.execution.executionId } : {}),
+                    contextTokens: reportedContextTokens ?? estimateContextTokens(context),
+                    contextWindow: model.contextWindow,
+                    contextSource: reportedContextTokens !== undefined ? "reported" : "estimate",
+                    toolCalls,
+                    toolRounds,
+                    phase: "model",
+                });
+                // Failed attempts are billed too.
+                onUsage?.(structuredClone(response.usage));
+                usage.inputTokens +=
+                    response.usage.input +
+                        response.usage.cacheRead +
+                        response.usage.cacheWrite;
+                usage.outputTokens += response.usage.output;
+                request.signal.throwIfAborted();
+                // The partial response never enters the context.
+                if (isStepInterrupted(request, step))
+                    continue;
+                // A stream that failed after it started is restarted from the same context. Tools run
+                // only after a complete response, so nothing from the failed attempt has executed.
+                if (response.stopReason === "error" && started && retries < maxModelRetries && isRetryableAssistantError(response)) {
+                    const delayMs = retryDelayMs({ baseDelayMs: modelRetryDelayMs, maxAgentDelayMs: MAX_MODEL_RETRY_DELAY_MS }, retries + 1);
+                    const deadline = Math.min(request.deadlines?.node ?? Infinity, request.deadlines?.graph ?? Infinity);
+                    // Report the provider error rather than a timeout when the wait cannot finish in time.
+                    if (performance.now() + delayMs < deadline) {
+                        retries++;
+                        activity?.modelError(`Retrying model request in ${Math.ceil(delayMs / 1000)}s (${retries}/${maxModelRetries})`);
+                        try {
+                            await abortableDelay(delayMs, step);
+                        }
+                        catch (error) {
+                            if (!isStepInterrupted(request, step))
+                                throw error;
+                        }
+                        continue;
+                    }
+                }
+                if (response.stopReason !== "stop" && response.stopReason !== "toolUse") {
+                    throw new Error(response.errorMessage ||
+                        `Pi model ended with '${response.stopReason}'`);
+                }
+                return response;
             }
-            return response;
         };
         const toolResult = (call, content, isError) => ({
             role: "toolResult",
@@ -291,11 +363,12 @@ export function createPiRunner(registry, onUsageOrOptions, cwd = process.cwd())
             isError,
             timestamp: Date.now(),
         });
-        const executeFileTool = async (call) => {
+        const executeFileTool = async (call, signal) => {
             const tool = toolByName.get(call.name);
             if (!tool && !(call.name === "git" && request.git) && !(call.name === "finish_merge" && request.merge)) {
                 return toolResult(call, `Tool '${call.name}' is unavailable in this Braid node. Available tools: ${allToolDefinitions.map(tool => tool.name).join(", ")}.`, true);
             }
+            let partialOutput = "";
             try {
                 request.signal.throwIfAborted();
                 if (call.name === "git" && request.git) {
@@ -309,17 +382,19 @@ export function createPiRunner(registry, onUsageOrOptions, cwd = process.cwd())
                     return toolResult(call, "Merge dispositions recorded. Return your final answer.", false);
                 }
                 const args = validateToolCall(allToolDefinitions, call);
-                const execute = () => tool.execute(call.id, args, request.signal, 
+                const execute = () => tool.execute(call.id, args, signal, 
                 // Bash sends cumulative, throttled output snapshots.
                 (partial) => {
-                    activity?.toolOutput(call.id, partial.content
-                        .flatMap(block => block.type === "text" ? [block.text] : []).join(""), "replace");
+                    partialOutput = partial.content.flatMap(block => block.type === "text" ? [block.text] : []).join("");
+                    activity?.toolOutput(call.id, partialOutput, "replace");
                     transcripts?.toolUpdate(executionId, call.id, { content: partial.content, details: partial.details });
                 });
                 const result = ["write", "edit", "bash", "powershell"].includes(call.name) && request.withWorkspaceWrite
                     ? await request.withWorkspaceWrite(execute)
                     : await execute();
                 request.signal.throwIfAborted();
+                if (isStepInterrupted(request, signal))
+                    return toolResult(call, interruptedToolResult(partialOutput), true);
                 return {
                     role: "toolResult",
                     toolCallId: call.id,
@@ -332,12 +407,14 @@ export function createPiRunner(registry, onUsageOrOptions, cwd = process.cwd())
             }
             catch (error) {
                 request.signal.throwIfAborted();
+                if (isStepInterrupted(request, signal))
+                    return toolResult(call, interruptedToolResult(partialOutput), true);
                 return toolResult(call, error instanceof Error ? error.message : String(error), true);
             }
         };
-        const executeCall = async (call) => {
+        const executeCall = async (call, signal) => {
             if (call.name !== "decide")
-                return executeFileTool(call);
+                return executeFileTool(call, signal);
             if (request.node.type !== "decision") {
                 return toolResult(call, "decide is available only on decision nodes", true);
             }
@@ -391,11 +468,19 @@ export function createPiRunner(registry, onUsageOrOptions, cwd = process.cwd())
             const results = [];
             for (const call of calls) {
                 request.signal.throwIfAborted();
+                const step = request.stepSignal?.() ?? request.signal;
+                // After an interruption the rest of the batch is not run, but every call still gets a result.
+                if (isStepInterrupted(request, step)) {
+                    const result = toolResult(call, notRunToolResult(), true);
+                    transcripts?.toolResult(executionId, result);
+                    results.push(result);
+                    continue;
+                }
                 activity?.toolStart(call.id, call.name, call.arguments);
                 transcripts?.toolStart(executionId, call.id);
                 let result;
                 try {
-                    result = await executeCall(call);
+                    result = await executeCall(call, step);
                 }
                 catch (error) {
                     activity?.toolEnd(call.id, error instanceof Error ? error.message : String(error), true);
@@ -421,6 +506,7 @@ export function createPiRunner(registry, onUsageOrOptions, cwd = process.cwd())
             usage,
         };
     };
+    return Object.assign(runner, { pausable: true, interruptible: true });
 }
 /** Includes cache/cost fields for Pi's tool-usage accounting, separate from core token totals. */
 export function sumPiUsage(reports) {