@fradser/pi-agent-teams 0.8.2 → 0.9.0

package/src/spawner.ts

@@ -11,7 +11,10 @@
  */
 
 import { randomUUID } from "node:crypto";
+import { WORKER_BUILTIN_TOOLS } from "./worker-tools.ts";
+export { WORKER_BUILTIN_TOOLS } from "./worker-tools.ts";
 import * as fs from "node:fs";
+import { tmpdir } from "node:os";
 import * as path from "node:path";
 import { fileURLToPath } from "node:url";
 import {
@@ -23,8 +26,16 @@ import {
 } from "@fradser/pi-kit";
 import type { ChildProcess } from "node:child_process";
 import type { WorkerUsage } from "./types.ts";
+import { writeWorkContext } from "./work-context.ts";
+import type { SessionContext } from "@earendil-works/pi-coding-agent";
 
 const OUTPUT_CAP = 16_000;
+/** A single JSONL record and an unterminated record may not exceed this many bytes. */
+export const MAX_JSONL_LINE_BYTES = 8 * 1024 * 1024;
+/** Live text, thinking, or tool arguments are bounded per assistant message. */
+export const MAX_TURN_OUTPUT_BYTES = 16 * 1024 * 1024;
+/** Diagnostic tail only; this is not a resident-process lifetime quota. */
+const DIAGNOSTIC_TAIL_BYTES = 64 * 1024;
 
 export interface WorkerProcessResult {
   pid: number;
@@ -48,8 +59,11 @@ export interface WorkerProgressUpdate {
    *  tool-call, or tool execution events). Independent of usage totals so
    *  providers that omit usage are not misclassified as silent. */
   modelOutputSeen?: boolean;
+  /** A prompt-less child acknowledged native startup readiness. */
+  ready?: boolean;
   /** Lifetime accumulated usage parsed from message_end events. */
   usage?: WorkerUsage;
+  controlError?: string;
 }
 
 /** A crashed teammate is one that closed without a normal zero exit. */
@@ -59,10 +73,8 @@ export function isCleanExit(result: Pick<WorkerProcessResult, "exitCode" | "sign
 
 /** Tools every teammate receives regardless of its role definition. */
 export const WORKER_CAPABILITY_TOOLS: readonly string[] = [
-  "send_message",
-  "task_list",
-  "task_claim",
-  "task_submit",
+  "agent_event",
+  "work",
 ];
 
 /** Effective tool allowlist for one teammate: the role's requested tools plus
@@ -78,16 +90,7 @@ export function resolveWorkerTools(requested?: string[]): string[] {
  *  --no-extensions, so this list (plus the capability set) is the complete
  *  grantable universe; anything else is silently dropped by the child's
  *  --tools allowlist filter. */
-export const WORKER_BUILTIN_TOOLS: readonly string[] = [
-  "read",
-  "bash",
-  "edit",
-  "write",
-  "grep",
-  "find",
-  "ls",
-  "powershell",
-];
+// Canonical IDs are shared with public schema and guidance in worker-tools.ts.
 
 /** Every tool id a teammate can actually receive: pi built-ins plus the capability set. */
 export const WORKER_TOOL_UNIVERSE: readonly string[] = [...WORKER_BUILTIN_TOOLS, ...WORKER_CAPABILITY_TOOLS];
@@ -122,10 +125,13 @@ export function isWorkerCloseObserved(name: string): boolean {
 export { terminateChildProcess };
 
 /** Terminate a living teammate and wait until its child process has closed. */
-export async function terminateTeammate(name: string, graceMs = DEFAULT_TERMINATION_GRACE_MS): Promise<boolean> {
+export type TeammateTermination = { outcome: "missing" } | { outcome: "closed" } | { outcome: "unconfirmed" };
+
+export async function terminateTeammate(name: string, graceMs = DEFAULT_TERMINATION_GRACE_MS): Promise<TeammateTermination> {
   const child = workers.get(name);
-  if (!child) return false;
-  return terminateChildProcess(child, graceMs);
+  if (!child) return { outcome: "missing" };
+  const closed = closedWorkers.has(child) || await terminateChildProcess(child, graceMs);
+  return { outcome: closed ? "closed" : "unconfirmed" };
 }
 
 export async function terminateAllTeammates(graceMs = DEFAULT_TERMINATION_GRACE_MS): Promise<Array<{ name: string; confirmedClosed: boolean }>> {
@@ -138,6 +144,30 @@ export async function terminateAllTeammates(graceMs = DEFAULT_TERMINATION_GRACE_
 
 // ── Control stream ────────────────────────────────────────────────
 
+interface PendingControlResponse {
+  child: ChildProcess;
+  command: string;
+  timeout: ReturnType<typeof setTimeout>;
+  resolve: (success: boolean) => void;
+}
+
+type PendingFreshDelivery = {
+  message: string;
+  streamingBehavior: "steer" | "followUp";
+};
+
+interface PendingFreshAssignment {
+  child: ChildProcess;
+  message: string;
+  resetId?: string;
+  queued: PendingFreshDelivery[];
+  resolve: (success: boolean) => void;
+}
+
+const CONTROL_RESPONSE_TIMEOUT_MS = 5_000;
+const pendingControlResponses = new Map<string, PendingControlResponse>();
+const pendingFreshAssignments = new Map<string, PendingFreshAssignment>();
+
 function writeToControlStream(child: ChildProcess, line: unknown): boolean {
   if (!child.stdin || child.stdin.destroyed || !child.stdin.writable) return false;
   child.stdin.write(`${JSON.stringify(line)}\n`);
@@ -151,7 +181,7 @@ function writeToControlStream(child: ChildProcess, line: unknown): boolean {
 export function deliverPrompt(name: string, message: string): boolean {
   const child = workers.get(name);
   if (!child) return false;
-  const sent = writeToControlStream(child, { type: "prompt", id: randomUUID(), message });
+  const sent = writeToControlStream(child, { type: "prompt", id: randomUUID(), message, streamingBehavior: "followUp" });
   if (sent) beginSequence(name);
   return sent;
 }
@@ -166,20 +196,118 @@ function beginSequence(name: string): void {
     state.text = "";
     state.thinking = "";
     clearActiveTools(state);
+    state.turnBytes = 0;
+    state.outputLimitError = undefined;
   }
   baselines.set(name, streamTurns.get(name) ?? 0);
 }
 
+function sendOrQueueDuringFreshReset(
+  name: string,
+  message: string,
+  streamingBehavior: PendingFreshDelivery["streamingBehavior"],
+): boolean {
+  const child = workers.get(name);
+  if (!child) return false;
+  const pending = pendingFreshAssignments.get(name);
+  if (pending?.child === child) {
+    pending.queued.push({ message, streamingBehavior });
+    return true;
+  }
+  return writeToControlStream(child, { type: "prompt", id: randomUUID(), message, streamingBehavior });
+}
+
 /** Send a mid-turn steer to a working teammate; no peer mailbox is involved. */
 export function sendWorkerSteer(name: string, message: string): boolean {
+  return sendOrQueueDuringFreshReset(name, message, "steer");
+}
+
+/** Peer traffic yields to leader steering without depending on roster freshness. */
+export function sendWorkerFollowUp(name: string, message: string): boolean {
+  return sendOrQueueDuringFreshReset(name, message, "followUp");
+}
+
+function settleFreshAssignment(name: string, success: boolean): void {
+  const pending = pendingFreshAssignments.get(name);
+  if (!pending) return;
+  pendingFreshAssignments.delete(name);
+  if (pending.resetId) {
+    const control = pendingControlResponses.get(pending.resetId);
+    if (control) clearTimeout(control.timeout);
+    pendingControlResponses.delete(pending.resetId);
+  }
+  pending.resolve(success);
+}
+
+function deliverFreshPrompts(name: string, pending: PendingFreshAssignment): boolean {
+  if (!writeToControlStream(pending.child, {
+    type: "prompt",
+    id: randomUUID(),
+    message: pending.message,
+    streamingBehavior: "followUp",
+  })) return false;
+  beginSequence(name);
+  for (const delivery of pending.queued) {
+    if (!writeToControlStream(pending.child, {
+      type: "prompt",
+      id: randomUUID(),
+      message: delivery.message,
+      streamingBehavior: delivery.streamingBehavior,
+    })) break;
+  }
+  return true;
+}
+
+function startFreshAssignmentReset(name: string): void {
+  const pending = pendingFreshAssignments.get(name);
+  const state = streamStates.get(name);
+  if (!pending || pending.resetId || state?.finalResponse !== true || workers.get(name) !== pending.child) return;
+  const id = randomUUID();
+  pending.resetId = id;
+  const timeout = setTimeout(() => settleFreshAssignment(name, false), CONTROL_RESPONSE_TIMEOUT_MS);
+  timeout.unref?.();
+  pendingControlResponses.set(id, {
+    child: pending.child,
+    command: "new_session",
+    timeout,
+    resolve: (reset) => {
+      const current = pendingFreshAssignments.get(name);
+      if (current !== pending || !reset || workers.get(name) !== pending.child) {
+        settleFreshAssignment(name, false);
+        return;
+      }
+      settleFreshAssignment(name, deliverFreshPrompts(name, pending));
+    },
+  });
+  if (!writeToControlStream(pending.child, { type: "new_session", id })) {
+    settleFreshAssignment(name, false);
+  }
+}
+
+/** True while a new Assignment Attempt is waiting for reset or prompt delivery. */
+export function isFreshAssignmentPending(name: string): boolean {
+  return pendingFreshAssignments.has(name);
+}
+
+/** Queue a fresh Pi session and deliver a new Assignment Attempt after settlement. */
+export function deliverFreshAssignment(name: string, message: string): Promise<boolean> {
   const child = workers.get(name);
-  return child ? writeToControlStream(child, { type: "steer", message }) : false;
+  if (!child || pendingFreshAssignments.has(name)) return Promise.resolve(false);
+  return new Promise((resolve) => {
+    pendingFreshAssignments.set(name, { child, message, queued: [], resolve });
+    startFreshAssignmentReset(name);
+  });
 }
 
 // ── Stream parsing ────────────────────────────────────────────────
 
 type JsonEvent = {
+  id?: string;
   type?: string;
+  command?: string;
+  success?: boolean;
+  error?: string;
+  data?: { cancelled?: boolean; isStreaming?: boolean; isCompacting?: boolean; pendingMessageCount?: number };
   toolCallId?: string;
   toolName?: string;
   args?: unknown;
@@ -215,10 +343,13 @@ interface StreamState {
    *  message_end artifact. The stall classifier uses this, not usage. */
   modelOutputSeen?: boolean;
   usage?: WorkerUsage;
+  controlError?: string;
+  turnBytes: number;
+  outputLimitError?: string;
 }
 
 function createStreamState(): StreamState {
-  return { text: "", thinking: "", toolcallArgs: "", activeTools: new Map(), turns: 0 };
+  return { text: "", thinking: "", toolcallArgs: "", activeTools: new Map(), turns: 0, turnBytes: 0 };
 }
 
 function clearActiveTools(state: StreamState): void {
@@ -229,7 +360,7 @@ function clearActiveTools(state: StreamState): void {
 
 function toolExecutionLabel(toolName: string | undefined, args: unknown): string {
   const serialized = typeof args === "string" ? args : JSON.stringify(args ?? {});
-  return toolcallLabel(serialized) ?? toolName ?? "tool";
+  return truncate(toolcallLabel(serialized) ?? toolName ?? "tool", OUTPUT_CAP);
 }
 
 /** Human-readable label from a partially streamed tool-call argument JSON. */
@@ -258,7 +389,18 @@ function normalizeInline(text: string): string {
   return text.replace(/\s+/g, " ").trim();
 }
 
-function applyStreamLine(state: StreamState, line: string): boolean {
+function appendTurnText(state: StreamState, field: "text" | "thinking" | "toolcallArgs", delta: string): boolean {
+  const nextBytes = state.turnBytes + Buffer.byteLength(delta, "utf8");
+  if (nextBytes > MAX_TURN_OUTPUT_BYTES) {
+    state.outputLimitError = `Resident worker ${field} exceeded ${MAX_TURN_OUTPUT_BYTES} bytes in one turn.`;
+    return false;
+  }
+  state.turnBytes = nextBytes;
+  state[field] += delta;
+  return true;
+}
+
+function applyStreamLine(state: StreamState, line: string, child?: ChildProcess): boolean {
   if (!line.trim()) return false;
   let event: JsonEvent;
   try {
@@ -266,6 +408,34 @@ function applyStreamLine(state: StreamState, line: string): boolean {
   } catch {
     return false;
   }
+  if (event.type === "response" && event.id) {
+    const pending = pendingControlResponses.get(event.id);
+    if (pending && child === pending.child && pending.command === event.command) {
+      clearTimeout(pending.timeout);
+      pendingControlResponses.delete(event.id);
+      pending.resolve(event.success === true && event.data?.cancelled !== true
+        && (pending.command !== "get_state" || (event.data?.isStreaming === false
+          && event.data?.isCompacting !== true && (event.data?.pendingMessageCount ?? 0) === 0)));
+    }
+  }
+  if (event.type === "response" && event.success === false) {
+    state.controlError = truncate(`RPC ${event.command ?? "command"} rejected: ${event.error ?? "unknown error"}`, 1000);
+    return true;
+  }
+  if (event.type === "agent_settled") {
+    state.finalResponse = true;
+    clearActiveTools(state);
+    return true;
+  }
+  if (event.type === "agent_start") {
+    state.finalResponse = false;
+    state.text = "";
+    state.thinking = "";
+    clearActiveTools(state);
+    state.turnBytes = 0;
+    state.outputLimitError = undefined;
+    return true;
+  }
   if (event.type === "tool_execution_start") {
     state.modelOutputSeen = true;
     state.activeTools.set(event.toolCallId ?? `tool-${state.activeTools.size}`, toolExecutionLabel(event.toolName, event.args));
@@ -284,9 +454,13 @@ function applyStreamLine(state: StreamState, line: string): boolean {
     state.turns++;
     clearActiveTools(state);
     state.thinking = "";
-    if (event.message.stopReason === "stop") state.finalResponse = true;
     const parts = extractTextContent(event.message.content, "");
+    if (Buffer.byteLength(parts, "utf8") > MAX_TURN_OUTPUT_BYTES) {
+      state.outputLimitError = `Resident worker message exceeded ${MAX_TURN_OUTPUT_BYTES} bytes.`;
+      return true;
+    }
     if (parts.trim()) state.text = parts;
+    state.turnBytes = 0;
     // Usage stays diagnostics: only streamed content counts as model output,
     // so input-only or failed responses cannot bypass the zero-output tier.
     const u = event.message.usage;
@@ -309,12 +483,12 @@ function applyStreamLine(state: StreamState, line: string): boolean {
     case "text_delta":
       state.modelOutputSeen = true;
       state.activeTool = undefined;
-      state.text += sub.delta ?? "";
+      appendTurnText(state, "text", sub.delta ?? "");
       return true;
     case "thinking_delta":
       state.modelOutputSeen = true;
       state.activeTool = undefined;
-      state.thinking += sub.delta ?? "";
+      appendTurnText(state, "thinking", sub.delta ?? "");
       return true;
     case "toolcall_start":
       state.modelOutputSeen = true;
@@ -322,7 +496,7 @@ function applyStreamLine(state: StreamState, line: string): boolean {
       return true;
     case "toolcall_delta": {
       state.modelOutputSeen = true;
-      state.toolcallArgs += sub.delta ?? "";
+      appendTurnText(state, "toolcallArgs", sub.delta ?? "");
       const label = toolcallLabel(state.toolcallArgs);
       if (label) state.activeTool = label;
       return true;
@@ -363,6 +537,10 @@ export interface ResidentSpawnOptions {
   model?: string;
   /** Execution-tool allowlist; capability tools are always appended. */
   tools?: string[];
+  /** Leader session's thinking level, forwarded as the child's default. */
+  thinking?: string;
+  /** Detached active-context snapshot; undefined starts with fresh history. */
+  context?: SessionContext["messages"];
   env?: Record<string, string | undefined>;
   cwd?: string;
   onUpdate?: (update: WorkerProgressUpdate) => void;
@@ -386,11 +564,11 @@ export function spawnResident(options: ResidentSpawnOptions): SpawnedResident |
   const args: string[] = [
     ...cli.args,
     "--mode", "rpc",
-    "--no-session",
     "--no-extensions",
     "--extension", WORKER_EXTENSION,
   ];
   if (options.model) args.push("--model", options.model);
+  if (options.thinking) args.push("--thinking", options.thinking);
   args.push("--tools", resolveWorkerTools(options.tools).join(","));
 
   let tempDir: string | undefined;
@@ -407,6 +585,12 @@ export function spawnResident(options: ResidentSpawnOptions): SpawnedResident |
 
   let child: ChildProcess;
   try {
+    if (options.context !== undefined) {
+      tempDir = fs.mkdtempSync(path.join(tmpdir(), "pi-work-context-"));
+      args.push("--session", writeWorkContext(options.cwd ?? process.cwd(), tempDir, options.context));
+    } else {
+      args.push("--no-session");
+    }
     child = spawnPiChild(cli.command, args, {
       cwd: options.cwd,
       env: { ...process.env, ...options.env },
@@ -428,6 +612,7 @@ export function spawnResident(options: ResidentSpawnOptions): SpawnedResident |
   const streamState = createStreamState();
   streamStates.set(options.workerName, streamState);
   let stdoutBuffer = "";
+  let failureReason: string | undefined;
   let settled = false;
 
   const emitProgress = () => options.onUpdate?.({
@@ -438,25 +623,54 @@ export function spawnResident(options: ResidentSpawnOptions): SpawnedResident |
     finalResponse: streamState.finalResponse,
     modelOutputSeen: streamState.modelOutputSeen,
     usage: streamState.usage,
+    controlError: streamState.controlError,
   });
 
+  let termination: Promise<boolean> | undefined;
+  const failWorker = (reason: string) => {
+    failureReason ??= reason;
+    stdoutBuffer = "";
+    streamState.outputLimitError ??= reason;
+    streamState.controlError = reason;
+    emitProgress();
+    termination ??= terminateChildProcess(child);
+  };
+
   child.stdout?.on("data", (chunk: Buffer) => {
+    if (failureReason) return;
     const text = chunk.toString();
-    appendCapped(stdoutChunks, text, OUTPUT_CAP * 4);
     stdoutBuffer += text;
     const lines = stdoutBuffer.split("\n");
     stdoutBuffer = lines.pop() ?? "";
     let changed = false;
     for (const line of lines) {
-      changed = applyStreamLine(streamState, line) || changed;
+      if (Buffer.byteLength(line, "utf8") > MAX_JSONL_LINE_BYTES) {
+        failWorker(`Resident worker JSONL line exceeded ${MAX_JSONL_LINE_BYTES} bytes.`);
+        return;
+      }
+      changed = applyStreamLine(streamState, line, child) || changed;
+      if (streamState.finalResponse) startFreshAssignmentReset(options.workerName);
+      if (streamState.outputLimitError) {
+        failWorker(streamState.outputLimitError);
+        return;
+      }
     }
+    if (Buffer.byteLength(stdoutBuffer, "utf8") > MAX_JSONL_LINE_BYTES) {
+      failWorker(`Resident worker JSONL line exceeded ${MAX_JSONL_LINE_BYTES} bytes.`);
+      return;
+    }
+    appendCapped(stdoutChunks, text, DIAGNOSTIC_TAIL_BYTES);
     streamTurns.set(options.workerName, streamState.turns);
     if (changed) emitProgress();
+    if (!failureReason) streamState.controlError = undefined;
+  });
+  child.stderr?.on("data", (chunk: Buffer) => {
+    if (failureReason) return;
+    appendCapped(stderrChunks, chunk.toString(), DIAGNOSTIC_TAIL_BYTES);
   });
-  child.stderr?.on("data", (chunk: Buffer) => appendCapped(stderrChunks, chunk.toString(), OUTPUT_CAP * 2));
 
   child.on("error", (error) => {
-    cleanupTempDir();
+    if (child.pid === undefined) cleanupTempDir();
     settled = true;
     // Keep the registry entry until close is observed so shutdown diagnostics
     // can distinguish an error/exit code from a confirmed close event.
@@ -467,30 +681,71 @@ export function spawnResident(options: ResidentSpawnOptions): SpawnedResident |
     cleanupTempDir();
     if (workers.get(options.workerName) === child) workers.delete(options.workerName);
     streamStates.delete(options.workerName);
+    const fresh = pendingFreshAssignments.get(options.workerName);
+    if (fresh?.child === child) settleFreshAssignment(options.workerName, false);
+    for (const [id, pending] of pendingControlResponses) {
+      if (pending.child !== child) continue;
+      clearTimeout(pending.timeout);
+      pendingControlResponses.delete(id);
+      pending.resolve(false);
+    }
     baselines.delete(options.workerName);
     streamTurns.delete(options.workerName);
     // A spawn failure was already reported via onError (Node fires error then
     // close) — do not double-report through onExit.
     if (settled) return;
-    const parsed = parseTeammateOutput(stdoutChunks.join(""));
+    const failed = failureReason !== undefined;
+    const parsed = failed ? { text: "", usage: undefined } : parseTeammateOutput(stdoutChunks.join(""));
     options.onExit({
       pid: child.pid ?? 0,
-      exitCode: code,
+      exitCode: code === 0 && failed ? 1 : code,
       signal,
-      stdout: truncate(parsed.text, OUTPUT_CAP),
-      stderr: truncate(stderrChunks.join("").trim(), OUTPUT_CAP),
+      stdout: failed ? "" : truncate(parsed.text, OUTPUT_CAP),
+      stderr: truncate([failureReason, stderrChunks.join("").trim()].filter(Boolean).join("\n"), OUTPUT_CAP),
       usage: parsed.usage,
     });
   });
 
+  if (!options.description?.trim()) {
+    const id = randomUUID();
+    const timeout = setTimeout(() => {
+      pendingControlResponses.delete(id);
+      streamState.controlError = "Resident startup readiness was not acknowledged.";
+      emitProgress();
+    }, CONTROL_RESPONSE_TIMEOUT_MS);
+    timeout.unref?.();
+    pendingControlResponses.set(id, {
+      child,
+      command: "get_state",
+      timeout,
+      resolve: (ready) => {
+        if (workers.get(options.workerName) !== child) return;
+        if (!ready) {
+          streamState.controlError = "Resident startup readiness acknowledgement was rejected or not idle.";
+          emitProgress();
+          return;
+        }
+        streamState.finalResponse = true;
+        options.onUpdate?.({ text: "", turns: 0, finalResponse: true, ready: true });
+        startFreshAssignmentReset(options.workerName);
+      },
+    });
+    if (!writeToControlStream(child, { type: "get_state", id })) {
+      clearTimeout(timeout);
+      pendingControlResponses.delete(id);
+      options.onError?.(new Error("Cannot request resident startup readiness."));
+    }
+  }
   return { pid: child.pid ?? 0 };
 }
 
 function appendCapped(chunks: string[], chunk: string, cap: number): void {
   chunks.push(chunk);
-  let total = chunks.reduce((sum, value) => sum + value.length, 0);
+  let total = chunks.reduce((sum, value) => sum + Buffer.byteLength(value, "utf8"), 0);
   while (total > cap && chunks.length > 1) {
-    total -= chunks.shift()?.length ?? 0;
+    total -= Buffer.byteLength(chunks.shift() ?? "", "utf8");
+  }
+  if (total > cap && chunks.length === 1) {
+    chunks[0] = Buffer.from(chunks[0], "utf8").subarray(-cap).toString("utf8");
   }
-  if (total > cap && chunks.length === 1) chunks[0] = chunks[0].slice(-cap);
 }