@chrok/pi-braid 0.3.6 → 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);
     }
@@ -313,6 +417,27 @@ export class BraidJobs {
         }))
             .reverse();
     }
+    /**
+     * Jobs that have not reached a terminal status. A job is stalled when only
+     * 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 => {
+            // Skip the snapshot copy while executions are visibly running.
+            const state = job.runningExecutions.size ? undefined : job.run?.snapshot();
+            return {
+                jobId: job.jobId,
+                handle: job.handle,
+                stalled: state?.status === "waiting",
+                ...(state ? { revision: state.revision } : {}),
+                pausedExecutionIds: state?.pausedExecutionIds ?? [],
+                suspendedExecutionIds: state?.suspendedExecutionIds ?? [],
+                ...timedOut(state),
+                ...(state?.pause ? { paused: true } : {}),
+            };
+        });
+    }
     /** Editor progress without copying outputs, transcripts, or execution snapshots. */
     progress(jobId) {
         const job = this.lookup(jobId);
@@ -324,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) {
@@ -336,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);