@fradser/pi-agent-teams 0.5.1 → 0.5.3

package/src/run-machine.ts

@@ -2,25 +2,23 @@
 
 import { randomUUID } from "node:crypto";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
-import { truncateTail, DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES } from "@earendil-works/pi-coding-agent";
 import { resolveAgent } from "./agents";
 import { isWorkerEvent } from "./worker";
 import {
   cancelBlockedDependents, cancelNode, cancelRun, clearWorkerRunEvents, clearStateDirty, isStateDirty,
-  deliverToLeader, failRunTimeout, findSharedWorkspaceWriteConflict, getNode, getNodeByWorkerKey,
+  deliverToLeader, findSharedWorkspaceWriteConflict, getNode,
   getRun, getState, listNodes, listRuns, markNodeRunning, markStateDirty,
-  readyPendingNodes, receiveWorkerMessage, retryRun, runningNodeCount, settleRun, setNodeSpawnInfo,
-  SUMMARY_NODE_ID, updateNodeSpawnProgress, updateNodeStatus,
+  readyPendingNodes, receiveWorkerMessage, runningNodeCount, settleRun, setNodeSpawnInfo,
+  SUMMARY_NODE_ID, updateNodeSpawnProgress, updateNodeStatus, validateStructuredOutput,
 } from "./state";
 import {
-  CancellationIntents, buildAutonomousPrompt, finishReportedWorker, isCompletedWorkerExit,
-  POST_REPORT_GRACE_MS, spawnPiWorker, terminateAllWorkers, terminateWorker,
+  CancellationIntents, buildAutonomousPrompt, DEFAULT_TURN_BUDGET, finishReportedWorker, isCompletedWorkerExit,
+  POST_REPORT_GRACE_MS, spawnPiWorker, terminateWorker,
   type WorkerProcessResult,
 } from "./spawner";
 import { buildNodeTerminalResult } from "./terminal";
 import { captureWorktreeDiff, cleanupWorktree, createWorktree, discardWorktree } from "./worktree";
 import { readWorkerEvents, removeWorkerOutbox, stateFilePath, workerOutboxPath, writeStateFile } from "./statefile";
-import type { Run } from "./types";
 
 export interface DispatchCtx {
   ui: ExtensionContext["ui"];
@@ -29,25 +27,21 @@ export interface DispatchCtx {
 }
 
 export const MAX_SESSION_WORKERS = 8;
-const DEFAULT_NODE_TIMEOUT_MS = 30 * 60 * 1000;
 const LIVE_POLL_MS = 500;
 const cancellationIntents = new CancellationIntents();
 const reportedWorkerShutdowns = new Set<string>();
 let liveStateFile: string | undefined;
 let livePollTimer: ReturnType<typeof setInterval> | undefined;
-let machineCtx: DispatchCtx | undefined;
 let sendUpdate: (subject: string, body: string, runId?: string) => void = () => {};
 let notifyChange: () => void = () => {};
 
-export function initRunMachine(ctx: DispatchCtx, stateFile: string, hooks: { sendUpdate: typeof sendUpdate; notifyChange: () => void }): void {
-  machineCtx = ctx;
+export function initRunMachine(_ctx: DispatchCtx, stateFile: string, hooks: { sendUpdate: typeof sendUpdate; notifyChange: () => void }): void {
   liveStateFile = stateFile;
   sendUpdate = hooks.sendUpdate;
   notifyChange = hooks.notifyChange;
 }
 
 export function ensureRunContext(ctx: DispatchCtx): void {
-  machineCtx = ctx;
   if (!liveStateFile && ctx.sessionManager) {
     liveStateFile = stateFilePath(ctx.sessionManager.getSessionFile(), ctx.cwd ?? process.cwd());
   }
@@ -57,16 +51,47 @@ export function shutdownRunMachine(): void {
   if (livePollTimer) clearInterval(livePollTimer);
   livePollTimer = undefined;
   liveStateFile = undefined;
-  machineCtx = undefined;
   sendUpdate = () => {};
   notifyChange = () => {};
 }
 
+function resolveInputBinding(run: import("./types").Run, node: import("./types").Node, source: string): string {
+  const separator = source.indexOf("#");
+  const dependencyId = separator >= 0 ? source.slice(0, separator) : source;
+  if (!node.dependsOn.includes(dependencyId)) return "(rejected: source is not a dependency)";
+  const dependency = run.nodes[dependencyId];
+  if (!dependency) return "(rejected: dependency not found)";
+  const value = dependency.structuredOutput ?? dependency.namedOutputs ?? dependency.result;
+  if (separator < 0) return JSON.stringify(value);
+  const pointer = source.slice(separator + 1).replace(/^\/json/, "");
+  if (!pointer || pointer === "") return JSON.stringify(value);
+  let current: unknown = value;
+  for (const part of pointer.split("/").filter(Boolean).map((item) => item.replaceAll("~1", "/").replaceAll("~0", "~"))) {
+    if (!current || typeof current !== "object") return "(missing)";
+    current = (current as Record<string, unknown>)[part];
+  }
+  return JSON.stringify(current);
+}
+
 function currentStateFile(): string {
   if (!liveStateFile) throw new Error("Agent Teams session state is unavailable.");
   return liveStateFile;
 }
 
+function sendNodeFollowUp(run: import("./types").Run, node: import("./types").Node, subject: string, body: string): void {
+  if (!run.background || node.id === SUMMARY_NODE_ID || node.nodeFollowUpSent) return;
+  node.nodeFollowUpSent = true;
+  markStateDirty();
+  sendUpdate(`Teammate ${node.id} ${terminalVerb(subject)}`, body, run.id);
+}
+
+function terminalVerb(subject: string): string {
+  if (subject === "Node completed") return "completed";
+  if (subject === "Node failed") return "failed";
+  if (subject === "Node cancelled") return "was cancelled";
+  return subject.toLowerCase();
+}
+
 function flushStateSnapshot(): void {
   if (!liveStateFile || !isStateDirty()) return;
   try {
@@ -108,13 +133,25 @@ export function applyWorkerEvents(): void {
             type: "message",
             subject: event.subject,
             body: event.body,
+            status: event.status,
+            data: event.data,
           });
+          if (event.data?.kind === "named_output" && event.data.name && typeof event.data.value === "string") {
+            node.namedOutputs = { ...node.namedOutputs, [event.data.name]: event.data.value };
+            markStateDirty();
+          }
+          if (event.data?.kind === "output" && event.data.output !== undefined && !validateStructuredOutput(event.data.output)) {
+            node.structuredOutput = event.data.output;
+            markStateDirty();
+          }
           if (event.status && !["completed", "failed", "cancelled"].includes(node.status)) {
             if (event.status === "completed") {
               updateNodeStatus(run.id, node.id, "completed", event.body, undefined);
+              sendNodeFollowUp(run, node, "Node completed", event.body);
               requestReportedWorkerShutdown(node.workerKey, spawnId);
             } else if (event.status === "failed") {
               updateNodeStatus(run.id, node.id, "failed", undefined, event.body);
+              sendNodeFollowUp(run, node, "Node failed", event.body);
               requestReportedWorkerShutdown(node.workerKey, spawnId);
             }
           }
@@ -158,7 +195,6 @@ export function ensureLivePoll(): void {
     livePollTimer = setInterval(() => {
       try {
         applyWorkerEvents();
-        enforceRunTimeouts();
         flushStateSnapshot();
         notifyChange();
         ensureLivePoll();
@@ -172,32 +208,8 @@ export function ensureLivePoll(): void {
   }
 }
 
-/** Fail runs whose run-level hard wall-clock cap was exceeded. */
-function enforceRunTimeouts(): void {
-  for (const run of listRuns()) {
-    if (run.status !== "running" || !run.deadlineAt || Date.now() < run.deadlineAt) continue;
-    const failed = failRunTimeout(run.id, `Run timed out after ${Math.round((run.timeoutMs ?? 0) / 1000)}s.`);
-    if (!failed.ok) continue;
-    for (const nodeId of failed.runningNodeIds) {
-      // Spawn was cleared by failRunTimeout, so terminate by the stable workerKey
-      // (a no-op when the child already closed).
-      const node = getNode(run.id, nodeId);
-      if (node) void terminateWorker(node.workerKey).catch(() => false);
-    }
-    onRunSettled(run.id, machineCtx);
-    // failRunTimeout clears spawns, so those nodes' close events skip their
-    // finalize path — release write-deferred nodes in OTHER runs here.
-    if (machineCtx) scheduleAllRuns(machineCtx);
-  }
-}
-
-
 // ── Run dispatch machinery ────────────────────────────────────────
 
-function sleep(ms: number): Promise<void> {
-  return new Promise((resolve) => setTimeout(resolve, ms));
-}
-
 /** Compact run summary for tool returns and follow-ups. When the run has a
  * synthesized __summary node result, that is shown instead of per-node
  * headlines (no truncation heuristics). */
@@ -209,11 +221,8 @@ export function buildRunSummary(runId: string): string {
     acc[node.status] = (acc[node.status] ?? 0) + 1;
     return acc;
   }, {});
-  const lines = [
-    `## Run [${run.id}] ${run.status}`,
-    `${Object.keys(run.nodes).length} node(s): ${Object.entries(counts).map(([status, n]) => `${status} ${n}`).join(", ")}`,
-    "",
-  ];
+  const countsStr = Object.entries(counts).map(([status, n]) => `${status} ${n}`).join(", ");
+  const lines = [`Run [${run.id}] ${countsStr}`, ""];
   if (run.summary) {
     lines.push(run.summary, "");
   } else if (nodes.length === 1) {
@@ -228,17 +237,13 @@ export function buildRunSummary(runId: string): string {
       }
     }
   }
-  for (const node of nodes) {
-    lines.push(`- [${node.id}] ${node.status} (${node.agent})`);
-  }
-  lines.push("", "Console: /teammate");
   return lines.join("\n");
 }
 
 /** Called once a run reaches a terminal status. Idempotent: the leader summary
  * and follow-up fire only on the first settled observation (a run can transition
  * through settleRun once per node close). */
-export function onRunSettled(runId: string, ctx?: DispatchCtx): void {
+export function onRunSettled(runId: string): void {
   const run = getRun(runId);
   if (!run) return;
   if (run.settledMessageSent) return;
@@ -266,7 +271,7 @@ export function scheduleRun(runId: string, ctx: DispatchCtx): void {
   if (!run || run.status !== "running") return;
   const settled = settleRun(runId);
   if (settled !== "running") {
-    onRunSettled(runId, ctx);
+    onRunSettled(runId);
     return;
   }
   const sessionBudget = MAX_SESSION_WORKERS - listNodes().filter((node) => node.status === "running").length;
@@ -333,9 +338,6 @@ export async function cancelRunAndTerminate(
   if (run.status !== "running") return { ok: false, error: `Run "${runId}" is already ${run.status}.` };
   const cancellation = cancelRun(runId);
   if (!cancellation.ok) return { ok: false, error: cancellation.error ?? "Failed to cancel run." };
-  run.status = "cancelled";
-  run.finishedAt = Date.now();
-  run.updatedAt = Date.now();
   run.completionNotified = true;
   markStateDirty();
   for (const nodeId of cancellation.runningNodeIds) {
@@ -346,7 +348,13 @@ export async function cancelRunAndTerminate(
     cancellationIntents.resolve(spawnId, terminated);
   }
   publishStateSnapshot();
-  onRunSettled(runId, ctx);
+  if (!cancellation.runningNodeIds.some((nodeId) => getNode(runId, nodeId)?.status === "running")) {
+    run.status = "cancelled";
+    run.finishedAt = Date.now();
+    run.updatedAt = Date.now();
+    markStateDirty();
+    onRunSettled(runId);
+  }
   scheduleAllRuns(ctx);
   notifyChange();
   return { ok: true };
@@ -360,10 +368,18 @@ export function startNode(runId: string, nodeId: string, ctx: DispatchCtx): void
   if (!run || !node || node.status !== "pending") return;
   const agent = resolveAgent(node.agent, run.cwd);
   const stateFile = liveStateFile ?? (ctx.sessionManager ? stateFilePath(ctx.sessionManager.getSessionFile(), run.cwd) : "");
-  if (!agent || !stateFile) {
-    updateNodeStatus(runId, nodeId, "failed", undefined, agent ? "Shared state file unavailable." : `Agent "${node.agent}" not found.`);
+  const failBeforeSpawn = (error: string): void => {
+    updateNodeStatus(runId, nodeId, "failed", undefined, error);
+    const failedNode = getNode(runId, nodeId);
+    const subject = "Node failed";
+    const body = `Node [${runId}/${nodeId}] could not start.\nError: ${error}`;
+    deliverToLeader({ from: node.workerKey, subject, body, runId });
+    if (failedNode) sendNodeFollowUp(run, failedNode, subject, body);
     cancelBlockedDependents(runId, nodeId);
     scheduleRun(runId, ctx);
+  };
+  if (!agent || !stateFile) {
+    failBeforeSpawn(agent ? "Shared state file unavailable." : `Agent "${node.agent}" not found.`);
     return;
   }
 
@@ -375,9 +391,7 @@ export function startNode(runId: string, nodeId: string, ctx: DispatchCtx): void
   if (run.worktree) {
     worktree = createWorktree(run.cwd, `${runId}-${nodeId}`);
     if ("error" in worktree) {
-      updateNodeStatus(runId, nodeId, "failed", undefined, `Cannot isolate node: ${worktree.error}`);
-      cancelBlockedDependents(runId, nodeId);
-      scheduleRun(runId, ctx);
+      failBeforeSpawn(`Cannot isolate node: ${worktree.error}`);
       return;
     }
   }
@@ -396,15 +410,24 @@ export function startNode(runId: string, nodeId: string, ctx: DispatchCtx): void
     status: "running",
     startedAt: Date.now(),
     isolation: run.worktree ? "worktree" : "none",
+    mode: node.mode,
   });
 
-  const timeoutMs = node.timeoutMs ?? DEFAULT_NODE_TIMEOUT_MS;
+  const turnBudget = node.turnBudget ?? DEFAULT_TURN_BUDGET;
+  const bindingLines = node.inputBindings
+    ? Object.entries(node.inputBindings).map(([name, source]) => `Input ${name}: ${resolveInputBinding(run, node, source)}`)
+    : [];
   const upstream = node.dependsOn
     .map((depId) => run.nodes[depId])
     .filter((dep): dep is NonNullable<typeof dep> => Boolean(dep))
+    .filter((dep) => !node.forkContext || node.forkContext.includes(dep.id))
     .map((dep) => {
       const body = dep.result?.trim() || dep.errorMessage?.trim() || `${dep.status} with no written result.`;
-      return `--- ${dep.id} (${dep.agent}, ${dep.status}) ---\n${body}`;
+      const named = dep.namedOutputs && Object.keys(dep.namedOutputs).length > 0
+        ? `\nNamed outputs: ${JSON.stringify(dep.namedOutputs)}`
+        : "";
+      const structured = dep.structuredOutput === undefined ? "" : `\nStructured output: ${JSON.stringify(dep.structuredOutput)}`;
+      return `--- ${dep.id} (${dep.agent}, ${dep.status}) ---\n${body}${named}${structured}`;
     });
   const description = [
     buildAutonomousPrompt({
@@ -412,7 +435,7 @@ export function startNode(runId: string, nodeId: string, ctx: DispatchCtx): void
       role: node.agent,
       prompt: agent.prompt,
       taskId: nodeId,
-      timeoutSec: Math.round(timeoutMs / 1000),
+      turnBudget,
     }),
     "",
     "=== TASK ===",
@@ -420,6 +443,7 @@ export function startNode(runId: string, nodeId: string, ctx: DispatchCtx): void
     `Paths: ${node.paths.join(", ")}`,
     node.prompt,
     ...(upstream.length > 0 ? ["", "=== UPSTREAM HANDOFF ===", ...upstream] : []),
+    ...(bindingLines.length > 0 ? ["", "=== NAMED INPUTS ===", ...bindingLines] : []),
   ].join("\n");
 
   const finalizeNode = (result: WorkerProcessResult, cancelled = false) => {
@@ -445,7 +469,7 @@ export function startNode(runId: string, nodeId: string, ctx: DispatchCtx): void
     const completedAfterFinalResponse = nodeNow?.spawn?.finalResponse === true;
     const workerReportedFailure = reportedTerminalStatus === "failed";
     const completedAfterShutdown = (reportedTerminalStatus === "completed" || completedAfterFinalResponse)
-      && (result.signal === "SIGTERM" || result.exitCode === 128 + 15 || result.timedOut);
+      && result.signal === "SIGTERM";
     const ok = isCompletedWorkerExit(
       result,
       reportedTerminalStatus === "completed" || completedAfterFinalResponse,
@@ -460,12 +484,11 @@ export function startNode(runId: string, nodeId: string, ctx: DispatchCtx): void
       stdout: (result.stdout + patchText).trim() ? result.stdout + patchText : undefined,
       stderr: ok ? undefined : result.stderr,
       usage: result.usage,
-      timedOut: result.timedOut,
       isolation: run.worktree ? "worktree" : "none",
       error: ok
         ? undefined
-        : result.timedOut
-          ? `Worker timed out after ${Math.round(timeoutMs / 1000)}s.`
+        : result.turnBudgetExceeded
+          ? `Worker exceeded its turn budget of ${turnBudget} turn(s).`
           : result.signal
             ? `Worker was terminated by ${result.signal}.`
             : workerReportedFailure
@@ -486,7 +509,7 @@ export function startNode(runId: string, nodeId: string, ctx: DispatchCtx): void
       runId,
       nodeId,
       agent: node.agent,
-      result,
+      result: { ...result, timedOut: result.turnBudgetExceeded ?? false },
       nodeResult: nodeNow?.result,
       nodeError: nodeNow?.errorMessage,
       cancelled,
@@ -494,6 +517,8 @@ export function startNode(runId: string, nodeId: string, ctx: DispatchCtx): void
       patchText,
     });
     deliverToLeader({ from: workerKey, subject: terminalSubject, body: terminalBody, runId: runId });
+    const settledNode = getNode(runId, nodeId);
+    if (!cancelled && settledNode) sendNodeFollowUp(run, settledNode, terminalSubject, terminalBody);
     if (cancelled) {
       // A cancelled node keeps its process outcome but not a misleading error.
       const cleared = getNode(runId, nodeId);
@@ -525,12 +550,11 @@ export function startNode(runId: string, nodeId: string, ctx: DispatchCtx): void
       error: typeof error === "string" ? error : error.message,
     });
     reportedWorkerShutdowns.delete(spawnId);
-    deliverToLeader({
-      from: workerKey,
-            subject: "Node failed",
-      body: `Node [${runId}/${nodeId}] could not start.\nError: ${typeof error === "string" ? error : error.message}`,
-      runId: runId,
-    });
+    const subject = "Node failed";
+    const body = `Node [${runId}/${nodeId}] could not start.\nError: ${typeof error === "string" ? error : error.message}`;
+    deliverToLeader({ from: workerKey, subject, body, runId: runId });
+    const failedNode = getNode(runId, nodeId);
+    if (failedNode) sendNodeFollowUp(run, failedNode, subject, body);
     cancelBlockedDependents(runId, nodeId);
     compactFinishedNodeRun(stateFile, workerKey, spawnId);
     if (worktree && !("error" in worktree)) discardWorktree(worktree);
@@ -540,12 +564,13 @@ export function startNode(runId: string, nodeId: string, ctx: DispatchCtx): void
 
   const started = spawnPiWorker({
     workerName: workerKey,
+    mode: node.mode,
     description,
     model: node.model ?? agent.model,
     tools: agent.tools,
     cwd: workerCwd,
     env: workerEnv,
-    timeoutMs,
+    turnBudget,
     onUpdate: (progress) => {
       updateNodeSpawnProgress(runId, nodeId, spawnId, {
         liveText: progress.text,
@@ -572,6 +597,7 @@ export function startNode(runId: string, nodeId: string, ctx: DispatchCtx): void
     status: "running",
     startedAt: node.spawn?.startedAt ?? Date.now(),
     isolation: run.worktree ? "worktree" : "none",
+    mode: node.mode,
   });
   ensureLivePoll();
   publishStateSnapshot();