tau-coding-agent 0.1.5 → 0.2.0

package/extensions/review/runner.ts

@@ -1,4 +1,4 @@
-import { spawn, type ChildProcess } from "node:child_process";
+import { spawn } from "node:child_process";
 import { promises as fs } from "node:fs";
 import os from "node:os";
 import path from "node:path";
@@ -20,9 +20,9 @@ export type TaskErrorKind =
   | "unsupported_reasoning"
   | "other";
 
-export type TaskExecutionControl = {
-  isCancelled: () => boolean;
-  registerProcess: (proc: ChildProcess) => () => void;
+export type ReviewCommandContext = {
+  cwd: string;
+  signal?: AbortSignal;
 };
 
 export type PiTaskStatus =
@@ -47,7 +47,7 @@ export type PiTaskOptions = {
   prompt: string;
   cwd: string;
   timeoutMs: number;
-  control?: TaskExecutionControl;
+  signal: AbortSignal;
   submitTool?: string;
 };
 
@@ -347,12 +347,12 @@ export async function runPiSubmitToolTask(options: PiSubmitToolTaskOptions): Pro
     if (
       firstResult.status !== "ok" ||
       firstResult.submittedPayloads.length > 0 ||
-      options.control?.isCancelled()
+      options.signal.aborted
     ) {
       return firstResult;
     }
 
-    return runPiOneShotTask({
+    return await runPiOneShotTask({
       ...options,
       args,
       prompt: SUBMIT_TOOL_RETRY_PROMPT.replaceAll("{SUBMIT_TOOL}", options.submitTool),
@@ -367,156 +367,148 @@ export async function runPiOneShotTask({
   prompt,
   cwd,
   timeoutMs,
-  control,
+  signal,
   submitTool,
 }: PiTaskOptions): Promise<PiTaskResult> {
-  if (control?.isCancelled()) {
-    return {
-      status: "cancelled",
-      assistantOutput: "",
-      stderr: "",
-      submittedPayloads: [],
-    };
+  if (signal.aborted) {
+    return { status: "cancelled", assistantOutput: "", stderr: "", submittedPayloads: [] };
   }
 
-  return new Promise<PiTaskResult>((resolve) => {
-    const proc = spawn("pi", [...args, prompt], {
-      cwd,
-      shell: false,
-      stdio: ["ignore", "pipe", "pipe"],
-      env: process.env,
-    });
-    const unregisterProcess = control?.registerProcess(proc);
-
-    let stdoutBuffer = "";
-    let latestAssistantOutput = "";
-    let latestAssistantError = "";
-    const submittedPayloads: unknown[] = [];
-    let stderr = "";
-    let settled = false;
-    let timeoutId: ReturnType<typeof setTimeout> | undefined;
-
-    const finish = (result: Omit<PiTaskResult, "submittedPayloads">) => {
-      if (settled) return;
-      settled = true;
-      if (timeoutId) {
-        clearTimeout(timeoutId);
-        timeoutId = undefined;
-      }
-      unregisterProcess?.();
-      resolve({
-        ...result,
-        submittedPayloads,
-      });
-    };
-
-    const processLine = (line: string) => {
-      const trimmed = line.trim();
-      if (!trimmed) return;
-      try {
-        const event = JSON.parse(trimmed);
-        const submittedPayload = submitTool
-          ? extractSubmitToolPayloadFromEvent(event, submitTool)
-          : null;
-        if (submittedPayload !== null) {
-          submittedPayloads.push(submittedPayload);
-        }
-
-        const message = extractAssistantMessageFromEvent(event);
-        if (message) {
-          const text = extractTextContent(message.content);
-          if (text) latestAssistantOutput = text;
-          latestAssistantError =
-            message.stopReason === "error" && typeof message.errorMessage === "string"
-              ? message.errorMessage
-              : "";
-        }
-      } catch {
-        // Ignore non-JSON lines.
-      }
-    };
+  const child = spawnReviewProcess("pi", [...args, prompt], { cwd, signal });
+  let stdoutBuffer = "";
+  let latestAssistantOutput = "";
+  let latestAssistantError = "";
+  const submittedPayloads: unknown[] = [];
+  let stderr = "";
+  let timedOut = false;
+  const timeout = setTimeout(() => {
+    timedOut = true;
+    child.stop();
+  }, timeoutMs);
+
+  child.proc.stdout.on("data", (chunk: string) => {
+    stdoutBuffer += chunk;
+    const lines = stdoutBuffer.split("\n");
+    stdoutBuffer = lines.pop() ?? "";
+    for (const line of lines) processLine(line);
+  });
+  child.proc.stderr.on("data", (chunk: string) => {
+    stderr += chunk;
+  });
+  child.proc.stdin.end();
 
-    if (control?.isCancelled()) {
-      try {
-        proc.kill("SIGKILL");
-      } catch {
-        // Best effort.
-      }
-      finish({
-        status: "cancelled",
-        assistantOutput: latestAssistantOutput,
-        stderr,
-      });
-      return;
+  try {
+    const { code, error } = await child.completion;
+    if (stdoutBuffer.trim()) processLine(stdoutBuffer);
+    const output = { assistantOutput: latestAssistantOutput, stderr, submittedPayloads };
+    if (signal.aborted) return { ...output, status: "cancelled" };
+    if (timedOut) return { ...output, status: "timeout" };
+    if (error) return { ...output, status: "spawn_error", error: error.message };
+    if (code !== 0) return { ...output, status: "non_zero_exit", exitCode: code ?? 1 };
+    if (latestAssistantError) {
+      return { ...output, status: "assistant_error", error: latestAssistantError, exitCode: 0 };
     }
+    return { ...output, status: "ok", exitCode: 0 };
+  } finally {
+    clearTimeout(timeout);
+  }
 
-    timeoutId = setTimeout(() => {
-      try {
-        proc.kill("SIGKILL");
-      } catch {
-        // Best effort.
-      }
-      finish({
-        status: "timeout",
-        assistantOutput: latestAssistantOutput,
-        stderr,
-      });
-    }, timeoutMs);
-
-    proc.stdout.on("data", (chunk) => {
-      stdoutBuffer += chunk.toString();
-      const lines = stdoutBuffer.split("\n");
-      stdoutBuffer = lines.pop() ?? "";
-      for (const line of lines) {
-        processLine(line);
-      }
-    });
-
-    proc.stderr.on("data", (chunk) => {
-      stderr += chunk.toString();
-    });
-
-    proc.on("error", (error) => {
-      finish({
-        status: "spawn_error",
-        assistantOutput: latestAssistantOutput,
-        stderr,
-        error: error.message,
-      });
-    });
-
-    proc.on("close", (code) => {
-      if (stdoutBuffer.trim()) {
-        processLine(stdoutBuffer);
-      }
-
-      if ((code ?? 1) !== 0) {
-        finish({
-          status: "non_zero_exit",
-          assistantOutput: latestAssistantOutput,
-          stderr,
-          exitCode: code ?? 1,
-        });
-        return;
+  function processLine(line: string) {
+    const trimmed = line.trim();
+    if (!trimmed) return;
+    try {
+      const event = JSON.parse(trimmed);
+      const submittedPayload = submitTool
+        ? extractSubmitToolPayloadFromEvent(event, submitTool)
+        : null;
+      if (submittedPayload !== null) submittedPayloads.push(submittedPayload);
+      const message = extractAssistantMessageFromEvent(event);
+      if (message) {
+        const text = extractTextContent(message.content);
+        if (text) latestAssistantOutput = text;
+        latestAssistantError =
+          message.stopReason === "error" && typeof message.errorMessage === "string"
+            ? message.errorMessage
+            : "";
       }
+    } catch {
+      // Ignore non-JSON lines.
+    }
+  }
+}
 
-      if (latestAssistantError) {
-        finish({
-          status: "assistant_error",
-          assistantOutput: latestAssistantOutput,
-          stderr,
-          error: latestAssistantError,
-          exitCode: 0,
-        });
-        return;
-      }
+export async function runReviewCommand(
+  context: ReviewCommandContext,
+  command: string,
+  args: string[],
+  input?: string,
+): Promise<{ stdout: string; stderr: string; code: number }> {
+  const child = spawnReviewProcess(command, args, context);
+  let stdout = "";
+  let stderr = "";
+  child.proc.stdout.on("data", (chunk: string) => {
+    stdout += chunk;
+  });
+  child.proc.stderr.on("data", (chunk: string) => {
+    stderr += chunk;
+  });
+  child.proc.stdin.end(input);
+  const { code, error } = await child.completion;
+  context.signal?.throwIfAborted();
+  return { stdout, stderr: error?.message ?? stderr, code: error ? 1 : (code ?? 1) };
+}
 
-      finish({
-        status: "ok",
-        assistantOutput: latestAssistantOutput,
-        stderr,
-        exitCode: 0,
-      });
+/** Join process and pipes after graceful cancellation, escalating only if they have not closed. */
+function spawnReviewProcess(
+  command: string,
+  args: string[],
+  { cwd, signal }: ReviewCommandContext,
+) {
+  signal?.throwIfAborted();
+  const proc = spawn(command, args, {
+    cwd,
+    detached: true,
+    stdio: ["pipe", "pipe", "pipe"],
+    env: process.env,
+  });
+  let closed = false;
+  let stopping = false;
+  let error: Error | undefined;
+  let escalation: NodeJS.Timeout | undefined;
+  proc.stdout.setEncoding("utf8");
+  proc.stderr.setEncoding("utf8");
+  proc.on("error", (cause: Error) => {
+    error = cause;
+  });
+  proc.stdin.on("error", () => {
+    /* A cancelled child may close stdin before it drains. */
+  });
+  const completion = new Promise<{ code: number | null; error?: Error }>((resolve) => {
+    proc.once("close", (code) => {
+      closed = true;
+      clearTimeout(escalation);
+      signal?.removeEventListener("abort", stop);
+      resolve({ code, error });
     });
   });
+  signal?.addEventListener("abort", stop, { once: true });
+  if (signal?.aborted) stop();
+  return { proc, completion, stop };
+
+  function stop() {
+    if (closed || stopping) return;
+    stopping = true;
+    // Pi handles SIGTERM by cancelling its agent and cleaning up detached Bash tools.
+    sendSignal("SIGTERM");
+    escalation = setTimeout(() => sendSignal("SIGKILL"), 1_000);
+  }
+
+  function sendSignal(value: "SIGTERM" | "SIGKILL") {
+    if (closed || proc.pid === undefined) return;
+    try {
+      process.kill(-proc.pid, value);
+    } catch (cause) {
+      if ((cause as NodeJS.ErrnoException).code !== "ESRCH") proc.kill(value);
+    }
+  }
 }

package/extensions/review/runtime.ts

@@ -3,9 +3,8 @@ import type {
   ExtensionCommandContext,
   ExtensionContext,
 } from "@earendil-works/pi-coding-agent";
-import { matchesKey } from "@earendil-works/pi-tui";
-import type { ChildProcess } from "node:child_process";
-
+import { truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
+import { createInterruptGuard } from "./interrupt.js";
 import type { ReviewRunOutcome, ReviewRunSource } from "./schema.js";
 
 export const REVIEW_CANCELLED_ERROR = "Review aborted";
@@ -13,23 +12,12 @@ export const REVIEW_CANCELLED_ERROR = "Review aborted";
 export const REVIEW_EVENT_START = "review:start";
 export const REVIEW_EVENT_END = "review:end";
 
-export const REVIEW_STATUS_KEY = "0-review";
+export const REVIEW_PROGRESS_WIDGET_KEY = "review-progress";
 
 export const STATUS_SPINNER_INTERVAL_MS = 80;
 
 export const STATUS_SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
 
-export type ReviewExecutionControl = {
-  isCancelled: () => boolean;
-  registerProcess: (proc: ChildProcess) => () => void;
-};
-
-export type ManagedReviewRun = {
-  control: ReviewExecutionControl;
-  markSuccessful: () => void;
-  markCancelled: () => void;
-};
-
 export type AgentEndMessage = {
   role?: string;
   stopReason?: string;
@@ -56,79 +44,125 @@ export type AgentRunTracker = FixPassAgentTracker & {
   reset: () => void;
 };
 
-const runtimeState = {
-  activeReviewRuns: new Set<string>(),
-  activeReviewCancels: new Map<string, () => void>(),
-  activePromptCount: 0,
-};
+export type ReviewRuntime = ReturnType<typeof createReviewRuntime>;
 
-export async function withManagedReviewRun<T>(
+export function createReviewRuntime(
   pi: ExtensionAPI,
-  ctx: ExtensionCommandContext,
-  source: ReviewRunSource,
-  run: (managed: ManagedReviewRun) => Promise<T>,
-): Promise<T> {
-  const sessionKey = getReviewSessionKey(ctx);
-  const activeProcesses = new Set<ChildProcess>();
-  let cancelRequested = false;
-  let outcome: ReviewRunOutcome = "failed";
+  onInterrupt: (ctx: ExtensionContext) => void,
+) {
+  let busy = false;
+  let closed = false;
+  let promptActive = false;
+  let active: { controller: AbortController; completion: Promise<void> } | undefined;
 
-  const cancelActiveProcesses = () => {
-    for (const proc of activeProcesses) {
+  return {
+    get closed() {
+      return closed;
+    },
+    acquire(ctx: ExtensionContext, busyMessage: string): boolean {
+      if (closed) return false;
+      if (busy) {
+        notify(ctx, busyMessage, "warning");
+        return false;
+      }
+      busy = true;
+      return true;
+    },
+    release() {
+      busy = false;
+    },
+    setPromptActive(value: boolean) {
+      promptActive = value;
+    },
+    async run<T extends { ok: boolean }>(
+      ctx: ExtensionCommandContext,
+      source: ReviewRunSource,
+      run: (signal: AbortSignal) => Promise<T>,
+    ): Promise<T | { ok: false; error: string }> {
+      if (closed) return { ok: false, error: REVIEW_CANCELLED_ERROR };
+      const sessionKey = getReviewSessionKey(ctx);
+      const controller = new AbortController();
+      const { signal } = controller;
+      let resolveCompletion!: () => void;
+      const completion = new Promise<void>((resolve) => {
+        resolveCompletion = resolve;
+      });
+      active = { controller, completion };
+      let outcome: ReviewRunOutcome = "failed";
+      let started = false;
+      let interrupt: ReturnType<typeof createInterruptGuard> | undefined;
       try {
-        proc.kill("SIGKILL");
-      } catch {
-        // Best effort.
+        publishReviewActivity(pi, ctx, "review preparing");
+        if (ctx.mode === "tui") {
+          ctx.ui.setWidget(
+            REVIEW_PROGRESS_WIDGET_KEY,
+            (tui, theme) => {
+              interrupt = createInterruptGuard(
+                pi,
+                ctx,
+                tui,
+                () => !signal.aborted,
+                () => {
+                  onInterrupt(ctx);
+                  controller.abort();
+                },
+                () => promptActive,
+              );
+              return {
+                invalidate() {},
+                render: (width) => {
+                  const handled = publishReviewActivity(pi, ctx, "review preparing");
+                  return !ctx.ui.getToolsExpanded() && handled
+                    ? []
+                    : [renderReviewProgressHeader("preparing", theme, width)];
+                },
+              };
+            },
+            { placement: "aboveEditor" },
+          );
+        }
+        started = true;
+        pi.events.emit(REVIEW_EVENT_START, { sessionKey, source });
+        signal.throwIfAborted();
+        const result = await run(signal);
+        signal.throwIfAborted();
+        outcome = result.ok ? "success" : "failed";
+        return result;
+      } catch (error) {
+        if (!signal.aborted) throw error;
+        outcome = "cancelled";
+        return { ok: false, error: REVIEW_CANCELLED_ERROR };
+      } finally {
+        try {
+          controller.abort();
+          await interrupt?.dispose();
+          if (ctx.mode === "tui") ctx.ui.setWidget(REVIEW_PROGRESS_WIDGET_KEY, undefined);
+          active = undefined;
+          if (started) pi.events.emit(REVIEW_EVENT_END, { sessionKey, source, outcome });
+        } finally {
+          publishReviewActivity(pi, ctx, undefined);
+          resolveCompletion();
+        }
       }
-    }
-    activeProcesses.clear();
-  };
-
-  const requestCancellation = () => {
-    if (cancelRequested) return;
-    cancelRequested = true;
-    cancelActiveProcesses();
-  };
-
-  const control: ReviewExecutionControl = {
-    isCancelled: () => cancelRequested,
-    registerProcess: (proc) => {
-      activeProcesses.add(proc);
-      return () => {
-        activeProcesses.delete(proc);
-      };
+    },
+    async shutdown() {
+      closed = true;
+      const current = active;
+      current?.controller.abort();
+      await current?.completion;
     },
   };
+}
 
-  pi.events.emit(REVIEW_EVENT_START, { sessionKey, source });
-  runtimeState.activeReviewCancels.set(sessionKey, requestCancellation);
-  const unsubscribeInterrupt = ctx.hasUI
-    ? ctx.ui.onTerminalInput((data) => {
-        if (!matchesKey(data, "escape")) return undefined;
-        if (runtimeState.activePromptCount > 0) return undefined;
-        requestCancellation();
-        return { consume: true };
-      })
-    : undefined;
-
-  try {
-    return await run({
-      control,
-      markSuccessful: () => {
-        outcome = "success";
-      },
-      markCancelled: () => {
-        outcome = "cancelled";
-      },
-    });
-  } finally {
-    unsubscribeInterrupt?.();
-    cancelActiveProcesses();
-    if (runtimeState.activeReviewCancels.get(sessionKey) === requestCancellation) {
-      runtimeState.activeReviewCancels.delete(sessionKey);
-    }
-    pi.events.emit(REVIEW_EVENT_END, { sessionKey, source, outcome });
-  }
+/** Join every sibling before propagating failure, so cancellation cannot leave background work behind. */
+export async function joinAll<T extends readonly unknown[] | []>(
+  tasks: T,
+): Promise<{ -readonly [P in keyof T]: Awaited<T[P]> }> {
+  const results = await Promise.allSettled(tasks);
+  return results.map((result) => {
+    if (result.status === "rejected") throw result.reason;
+    return result.value;
+  }) as { -readonly [P in keyof T]: Awaited<T[P]> };
 }
 
 export function notify(
@@ -144,7 +178,19 @@ export function getReviewSessionKey(ctx: ExtensionContext): string {
   return ctx.sessionManager.getSessionFile() ?? `session:${ctx.sessionManager.getSessionId()}`;
 }
 
+export function publishReviewActivity(
+  pi: ExtensionAPI,
+  ctx: ExtensionContext,
+  text: string | undefined,
+): boolean {
+  if (ctx.mode !== "tui") return false;
+  const request = { sessionKey: getReviewSessionKey(ctx), source: "review", text, handled: false };
+  pi.events.emit("tau:activity", request);
+  return request.handled === true;
+}
+
 export async function withSpinner<T>(
+  pi: ExtensionAPI,
   ctx: ExtensionContext,
   buildStatusText: () => string,
   run: () => Promise<T>,
@@ -152,9 +198,39 @@ export async function withSpinner<T>(
   if (!ctx.hasUI) return run();
 
   let frame = 0;
+  let requestRender: (() => void) | undefined;
+  if (ctx.mode === "tui") {
+    ctx.ui.setWidget(
+      REVIEW_PROGRESS_WIDGET_KEY,
+      (tui, theme) => {
+        requestRender = () => tui.requestRender();
+        return {
+          invalidate() {},
+          render(width) {
+            const handled = publishReviewActivity(pi, ctx, `review ${buildStatusText()}`);
+            if (!ctx.ui.getToolsExpanded() && handled) return [];
+            const spinner = theme.fg("accent", STATUS_SPINNER_FRAMES[frame]!);
+            return [
+              renderReviewProgressHeader(
+                `${spinner} ${theme.fg("muted", buildStatusText())}`,
+                theme,
+                width,
+              ),
+            ];
+          },
+        };
+      },
+      { placement: "aboveEditor" },
+    );
+  }
   const render = () => {
-    const spinner = STATUS_SPINNER_FRAMES[frame % STATUS_SPINNER_FRAMES.length];
-    ctx.ui.setStatus(REVIEW_STATUS_KEY, `${spinner} ${buildStatusText()}`);
+    if (ctx.mode === "tui") {
+      publishReviewActivity(pi, ctx, `review ${buildStatusText()}`);
+      requestRender?.();
+    } else
+      ctx.ui.setWidget(REVIEW_PROGRESS_WIDGET_KEY, [`Review · ${buildStatusText()}`], {
+        placement: "aboveEditor",
+      });
   };
 
   render();
@@ -167,23 +243,24 @@ export async function withSpinner<T>(
     return await run();
   } finally {
     clearInterval(timer);
-    ctx.ui.setStatus(REVIEW_STATUS_KEY, undefined);
+    publishReviewActivity(pi, ctx, "review");
+    ctx.ui.setWidget(REVIEW_PROGRESS_WIDGET_KEY, undefined);
   }
 }
 
-export function acquireReviewRunLock(ctx: ExtensionContext, busyMessage: string): string | null {
-  const sessionKey = getReviewSessionKey(ctx);
-  if (runtimeState.activeReviewRuns.has(sessionKey)) {
-    notify(ctx, busyMessage, "warning");
-    return null;
-  }
-
-  runtimeState.activeReviewRuns.add(sessionKey);
-  return sessionKey;
-}
-
-export function releaseReviewRunLock(sessionKey: string): void {
-  runtimeState.activeReviewRuns.delete(sessionKey);
+export function renderReviewProgressHeader(
+  status: string,
+  theme: ExtensionContext["ui"]["theme"],
+  width: number,
+  hint = "",
+): string {
+  const label = `Review · ${status}`;
+  const innerWidth = Math.max(0, width - 2);
+  const gap = innerWidth - visibleWidth(label) - visibleWidth(hint);
+  return truncateToWidth(
+    ` ${theme.fg("muted", truncateToWidth(label, innerWidth))}${hint && gap >= 2 ? " ".repeat(gap) + theme.fg("dim", hint) : ""}`,
+    width,
+  );
 }
 
 export function createAgentRunTracker(): AgentRunTracker {
@@ -249,27 +326,3 @@ export function createAgentRunTracker(): AgentRunTracker {
     },
   };
 }
-
-export function recordPromptStart(): void {
-  runtimeState.activePromptCount += 1;
-}
-
-export function recordPromptEnd(): void {
-  runtimeState.activePromptCount = Math.max(0, runtimeState.activePromptCount - 1);
-}
-
-export function handleReviewSessionStart(ctx: ExtensionContext): void {
-  runtimeState.activePromptCount = 0;
-  const sessionKey = getReviewSessionKey(ctx);
-  for (const [key, cancel] of runtimeState.activeReviewCancels) {
-    if (key !== sessionKey) cancel();
-  }
-}
-
-export function handleReviewSessionShutdown(ctx: ExtensionContext): void {
-  const sessionKey = getReviewSessionKey(ctx);
-  runtimeState.activeReviewCancels.get(sessionKey)?.();
-  runtimeState.activeReviewCancels.delete(sessionKey);
-  runtimeState.activeReviewRuns.delete(sessionKey);
-  runtimeState.activePromptCount = 0;
-}