@chrok/pi-braid 0.4.0 → 0.5.1

package/dist/jobs.js

@@ -1,14 +1,21 @@
 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;
+    onReplanRequested;
     jobs = new Map();
     assistants = new Map();
     cwd = process.cwd();
@@ -16,9 +23,25 @@ export class BraidJobs {
     nextHandle = 1;
     listeners = new Set();
     disposed = false;
-    constructor(onFinished = () => { }, onNodeFinished = () => { }) {
+    constructor(onFinished = () => { }, onNodeFinished = () => { }, onControl = () => { }, onTimeoutPause = () => { }, onReplanRequested = () => { }) {
         this.onFinished = onFinished;
         this.onNodeFinished = onNodeFinished;
+        this.onControl = onControl;
+        this.onTimeoutPause = onTimeoutPause;
+        this.onReplanRequested = onReplanRequested;
+    }
+    /** 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 +59,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 +72,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 +82,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 +114,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 +126,7 @@ export class BraidJobs {
                             this.changed();
                         },
                     });
-                    return async (request) => {
+                    return Object.assign(async (request) => {
                         try {
                             await nextTurn();
                             return await invoke(request);
@@ -116,9 +139,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 +156,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 +187,42 @@ 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. */ }
+                        }
+                    }
+                    // A replan request wakes the parent now, while the node keeps working.
+                    if (!this.disposed && event.type === "replan_requested" && event.executionId) {
+                        try {
+                            void Promise.resolve(this.onReplanRequested({
+                                jobId: job.jobId, handle: job.handle, nodeId: event.nodeId, executionId: event.executionId,
+                                ...(event.iteration ? { iteration: event.iteration } : {}),
+                                reason: event.reason, suggested: event.suggestion !== undefined, 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,
@@ -165,6 +234,7 @@ export class BraidJobs {
                                 status: instance.status === "completed" ? "completed" : "failed",
                                 eventSequence: event.sequence,
                                 ...(instance.error ? { errorCode: instance.error.code } : {}),
+                                ...(instance.replan ? { replanRequested: true } : {}),
                             })).catch(() => { });
                         }
                         catch {
@@ -188,12 +258,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 +284,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 +363,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 +433,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 +446,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 +463,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 +475,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);