tau-coding-agent 0.1.6 → 0.2.0

package/extensions/subagent/interrupt.ts

@@ -0,0 +1,117 @@
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import {
+  getKeybindings,
+  isKeyRelease,
+  type TuiMainScreen,
+  type Component,
+  type TUI,
+} from "@earendil-works/pi-tui";
+
+interface InterruptRequest {
+  sessionKey: string;
+  cancel?: boolean;
+  active: boolean;
+  confirming?: boolean;
+}
+
+/** Each extension owns its guard; optional peers contribute work through the event bus. */
+export function createInterruptGuard(
+  pi: ExtensionAPI,
+  ctx: ExtensionContext,
+  tui: TUI & Partial<Pick<TuiMainScreen, "getFocusedComponent">>,
+  isWorking: () => boolean,
+  cancel: () => void,
+  isPromptActive: () => boolean,
+) {
+  const sessionKey =
+    ctx.sessionManager.getSessionFile() ?? `session:${ctx.sessionManager.getSessionId()}`;
+  let closed = false;
+  let pending:
+    | { controller: AbortController; editor: Component; completion: Promise<void> }
+    | undefined;
+  const unsubscribeWork = pi.events.on("tau:interrupt", (data) => {
+    const request = data as InterruptRequest;
+    if (request.sessionKey !== sessionKey || closed) return;
+    if (pending) request.confirming = true;
+    if (!isWorking()) return;
+    request.active = true;
+    if (request.cancel) cancel();
+  });
+  const unsubscribeChanged = pi.events.on("tau:work-changed", (data) => {
+    if ((data as { sessionKey: string }).sessionKey === sessionKey && pending && !hasWork())
+      pending.controller.abort();
+  });
+  const unsubscribeInput = ctx.ui.onTerminalInput((data) => {
+    if (isKeyRelease(data) || !getKeybindings().matches(data, "app.interrupt")) return;
+    const editor = tui.getFocusedComponent?.();
+    // A dismissed native dialog restores focus before its promise settles. Swallow a rapid
+    // extra Escape in that gap instead of accidentally interrupting the parent.
+    if (pending) return editor === pending.editor ? { consume: true } : undefined;
+    if (
+      closed ||
+      isPromptActive() ||
+      !editor ||
+      !("getText" in editor) ||
+      !("setText" in editor) ||
+      ("isShowingAutocomplete" in editor &&
+        typeof editor.isShowingAutocomplete === "function" &&
+        editor.isShowingAutocomplete()) ||
+      !hasWork()
+    )
+      return;
+    const current = {
+      controller: new AbortController(),
+      editor,
+      completion: Promise.resolve(),
+    };
+    pending = current;
+    current.completion = (async () => {
+      try {
+        const request: {
+          run: (signal: AbortSignal) => Promise<void>;
+          signal: AbortSignal;
+          result?: Promise<unknown>;
+        } = {
+          signal: current.controller.signal,
+          async run(signal) {
+            const confirmed = await ctx.ui.confirm("Cancel all ongoing work?", "", { signal });
+            if (!confirmed || closed || signal.aborted || !hasWork()) return;
+            pi.events.emit("tau:interrupt", { sessionKey, active: false, cancel: true });
+            // Reuse the editor's native interrupt, including queue restoration, Bash and retries.
+            editor.handleInput?.(data);
+          },
+        };
+        // Keep queued approvals blocked until the confirmed cancellation has been applied.
+        pi.events.emit("subagent:permission", request);
+        await (request.result ?? request.run(request.signal));
+      } catch (error) {
+        if (!closed && !current.controller.signal.aborted)
+          ctx.ui.notify(`Could not confirm cancellation: ${String(error)}`, "warning");
+      } finally {
+        if (pending === current) pending = undefined;
+      }
+    })();
+    return { consume: true };
+  });
+
+  return {
+    refresh() {
+      if (!closed) pi.events.emit("tau:work-changed", { sessionKey });
+    },
+    async dispose() {
+      closed = true;
+      unsubscribeInput();
+      unsubscribeWork();
+      unsubscribeChanged();
+      pending?.controller.abort();
+      pi.events.emit("tau:work-changed", { sessionKey });
+      await pending?.completion;
+    },
+  };
+
+  function hasWork(): boolean {
+    const request: InterruptRequest = { sessionKey, active: false };
+    pi.events.emit("tau:interrupt", request);
+    return request.active;
+  }
+}

package/extensions/subagent/permissions.ts

@@ -0,0 +1,101 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+
+/** Optional integration for sharing the parent's approval and confirmation queue. */
+interface PermissionRequest {
+  run: (signal: AbortSignal) => Promise<unknown>;
+  signal?: AbortSignal;
+  result?: Promise<unknown>;
+}
+
+export function registerPermissions(pi: ExtensionAPI) {
+  let queue = new PermissionQueue();
+  pi.events.on("subagent:permission", (data) => {
+    const request = data as PermissionRequest;
+    request.result = queue.run(request.run, request.signal);
+  });
+  pi.on("ui_prompt_start", () => queue.setPromptActive(true));
+  pi.on("ui_prompt_end", () => queue.setPromptActive(false));
+  return {
+    get active() {
+      return queue.active;
+    },
+    run<T>(run: (signal: AbortSignal) => Promise<T>, signal: AbortSignal) {
+      return queue.run(run, signal);
+    },
+    async reset() {
+      await queue.close();
+      queue = new PermissionQueue();
+    },
+    close: () => queue.close(),
+  };
+}
+
+/** Cancel queued requests independently; stopping one child must not wait for another child's user. */
+class PermissionQueue {
+  private readonly controller = new AbortController();
+  private readonly pending: (() => Promise<void>)[] = [];
+  private current?: Promise<void>;
+  private promptActive = false;
+
+  get active(): boolean {
+    return this.promptActive || this.current !== undefined;
+  }
+
+  run<T>(run: (signal: AbortSignal) => Promise<T>, signal?: AbortSignal): Promise<T> {
+    const combined = signal
+      ? AbortSignal.any([signal, this.controller.signal])
+      : this.controller.signal;
+    return new Promise<T>((resolve, reject) => {
+      if (combined.aborted) {
+        reject(combined.reason);
+        return;
+      }
+      const start = async () => {
+        try {
+          combined.throwIfAborted();
+          resolve(await run(combined));
+        } catch (error) {
+          reject(error);
+        } finally {
+          combined.removeEventListener("abort", cancel);
+        }
+      };
+      const cancel = () => {
+        const index = this.pending.indexOf(start);
+        if (index === -1) return;
+        this.pending.splice(index, 1);
+        combined.removeEventListener("abort", cancel);
+        reject(combined.reason);
+      };
+      combined.addEventListener("abort", cancel, { once: true });
+      this.pending.push(start);
+      this.advance();
+    });
+  }
+
+  setPromptActive(active: boolean): void {
+    this.promptActive = active;
+    this.advance();
+  }
+
+  async close(): Promise<void> {
+    this.controller.abort();
+    await this.current;
+  }
+
+  private advance(): void {
+    if (this.current || this.promptActive || this.controller.signal.aborted) return;
+    const start = this.pending.shift();
+    if (!start) return;
+    let begin!: () => void;
+    this.current = new Promise<void>((resolve, reject) => {
+      begin = () => {
+        void start().then(resolve, reject);
+      };
+    }).finally(() => {
+      this.current = undefined;
+      this.advance();
+    });
+    begin();
+  }
+}

package/extensions/subagent/rpc.ts

@@ -0,0 +1,177 @@
+import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process";
+import type {
+  JsonAgentSessionEvent,
+  RpcCommand,
+  RpcExtensionUIRequest,
+  RpcExtensionUIResponse,
+  RpcResponse,
+} from "@earendil-works/pi-coding-agent";
+
+export type ChildEvent = JsonAgentSessionEvent | RpcExtensionUIRequest;
+
+/** Own the native RPC process, its LF-delimited protocol, and all pending requests through close. */
+export class SubagentProcess {
+  private readonly process: ChildProcessWithoutNullStreams;
+  private readonly pending = new Map<
+    string,
+    {
+      command: RpcCommand["type"];
+      resolve: (response: RpcResponse) => void;
+      reject: (error: Error) => void;
+      timer: ReturnType<typeof setTimeout>;
+    }
+  >();
+  private sequence = 0;
+  private buffer = "";
+  private stderr = "";
+  private failure?: Error;
+  private stopping = false;
+  private closed = false;
+  private escalation?: ReturnType<typeof setTimeout>;
+  private shutdownRequest?: Promise<void>;
+  readonly completion: Promise<void>;
+
+  constructor(
+    cwd: string,
+    args: string[],
+    env: NodeJS.ProcessEnv,
+    onEvent: (event: ChildEvent) => void,
+    onClose: (error?: Error) => void,
+  ) {
+    this.process = spawn("pi", ["--mode", "rpc", ...args], {
+      cwd,
+      env,
+      detached: true,
+      stdio: ["pipe", "pipe", "pipe"],
+    });
+    this.process.stdout.setEncoding("utf8");
+    this.process.stderr.setEncoding("utf8");
+    this.process.stdout.on("data", (chunk: string) => {
+      try {
+        this.buffer += chunk;
+        let newline: number;
+        while ((newline = this.buffer.indexOf("\n")) !== -1) {
+          const line = this.buffer.slice(0, newline).replace(/\r$/, "");
+          this.buffer = this.buffer.slice(newline + 1);
+          if (!line) continue;
+          if (Buffer.byteLength(line) > 32 * 1024 * 1024)
+            throw new Error("Subagent RPC record exceeds 32 MiB");
+          const event = JSON.parse(line) as RpcResponse | ChildEvent;
+          if (!event || typeof event !== "object" || typeof event.type !== "string") {
+            throw new Error("Invalid subagent RPC record");
+          }
+          if (event.type === "response") {
+            const request = event.id ? this.pending.get(event.id) : undefined;
+            if (!request) continue;
+            if (event.command !== request.command || typeof event.success !== "boolean")
+              throw new Error("Invalid subagent RPC acknowledgement");
+            clearTimeout(request.timer);
+            this.pending.delete(event.id!);
+            if (event.success) request.resolve(event);
+            else request.reject(new Error(event.error));
+          } else if (!this.stopping) {
+            onEvent(event);
+          }
+        }
+        if (Buffer.byteLength(this.buffer) > 32 * 1024 * 1024)
+          throw new Error("Subagent RPC record exceeds 32 MiB");
+      } catch (error) {
+        this.fail(error);
+      }
+    });
+    this.process.stderr.on("data", (chunk: string) => {
+      this.stderr = (this.stderr + chunk).slice(-8192);
+    });
+    this.process.on("error", (error) => this.fail(error));
+    this.process.stdin.on("error", (error) => this.fail(error));
+    this.completion = new Promise<void>((resolve) => {
+      this.process.once("close", (code, signal) => {
+        this.closed = true;
+        clearTimeout(this.escalation);
+        const error =
+          this.failure ??
+          new Error(
+            `Subagent exited (${signal ?? code}).${this.stderr.trim() ? ` ${this.stderr.trim()}` : ""}`,
+          );
+        for (const request of this.pending.values()) {
+          clearTimeout(request.timer);
+          request.reject(error);
+        }
+        this.pending.clear();
+        try {
+          onClose(this.stopping && !this.failure ? undefined : error);
+        } finally {
+          resolve();
+        }
+      });
+    });
+  }
+
+  async request(command: RpcCommand): Promise<RpcResponse> {
+    if (this.closed || this.stopping) throw this.failure ?? new Error("Subagent has stopped");
+    return this.send(command);
+  }
+
+  respond(response: RpcExtensionUIResponse): void {
+    if (!this.closed && !this.stopping) this.write(response);
+  }
+
+  async stop(): Promise<void> {
+    if (!this.closed && !this.stopping) {
+      this.stopping = true;
+      let terminating = false;
+      const terminate = () => {
+        if (terminating || this.closed) return;
+        terminating = true;
+        clearTimeout(this.escalation);
+        // Pi's SIGTERM handler also stops separately detached Bash children.
+        this.signal("SIGTERM");
+        this.escalation = setTimeout(() => this.signal("SIGKILL"), 1000);
+      };
+      this.escalation = setTimeout(terminate, 250);
+      this.shutdownRequest = (
+        this.failure
+          ? Promise.resolve()
+          : this.send({ type: "clear_queue" }).then(() => this.send({ type: "abort" }))
+      ).then(terminate, terminate);
+    }
+    await this.completion;
+    await this.shutdownRequest;
+  }
+
+  private send(command: RpcCommand): Promise<RpcResponse> {
+    if (this.closed) return Promise.reject(this.failure ?? new Error("Subagent has stopped"));
+    const id = String(++this.sequence);
+    return new Promise<RpcResponse>((resolve, reject) => {
+      const timer = setTimeout(
+        () =>
+          this.fail(new Error(`Subagent did not acknowledge ${command.type} within 30 seconds`)),
+        30_000,
+      );
+      this.pending.set(id, { command: command.type, resolve, reject, timer });
+      this.write({ ...command, id });
+    });
+  }
+
+  private write(message: RpcCommand | RpcExtensionUIResponse): void {
+    try {
+      this.process.stdin.write(`${JSON.stringify(message)}\n`);
+    } catch (error) {
+      this.fail(error);
+    }
+  }
+
+  private fail(error: unknown): void {
+    this.failure ??= error instanceof Error ? error : new Error(String(error));
+    void this.stop();
+  }
+
+  private signal(signal: NodeJS.Signals): void {
+    if (this.closed) return;
+    try {
+      if (this.process.pid) process.kill(-this.process.pid, signal);
+    } catch (error) {
+      if ((error as NodeJS.ErrnoException).code !== "ESRCH") this.process.kill(signal);
+    }
+  }
+}