@aaroncql/pim-agent 0.0.1

package/src/extensions/bash/run.test.ts

@@ -0,0 +1,262 @@
+import { afterAll, beforeAll, describe, expect, test } from "bun:test";
+import { mkdtemp, rm, stat } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { SpillCache } from "../../shared/SpillCache";
+import { killAllActiveBashGroups, runBashCommand } from "./run";
+import {
+  DRAIN_GRACE_MS,
+  KILL_GRACE_MS,
+  STREAM_HEAD_BYTES,
+  STREAM_TAIL_BYTES,
+} from "./schema";
+
+let previousPimHomeDir: string | undefined;
+let testPimHomeDir: string | undefined;
+
+beforeAll(async () => {
+  previousPimHomeDir = process.env.PIM_HOME_DIR;
+  testPimHomeDir = await mkdtemp(join(tmpdir(), "pim-bash-home-"));
+  process.env.PIM_HOME_DIR = testPimHomeDir;
+});
+
+afterAll(async () => {
+  if (previousPimHomeDir === undefined) {
+    delete process.env.PIM_HOME_DIR;
+  } else {
+    process.env.PIM_HOME_DIR = previousPimHomeDir;
+  }
+  if (testPimHomeDir) {
+    await rm(testPimHomeDir, { recursive: true, force: true });
+  }
+});
+
+function shellQuote(value: string): string {
+  return `'${value.replaceAll("'", `'"'"'`)}'`;
+}
+
+async function waitForFile(path: string, timeoutMs: number): Promise<void> {
+  const deadline = Date.now() + timeoutMs;
+  while (Date.now() < deadline) {
+    if (await Bun.file(path).exists()) {
+      return;
+    }
+    await Bun.sleep(10);
+  }
+  throw new Error(`Timed out waiting for ${path}`);
+}
+
+async function waitForNoProcess(
+  marker: string,
+  timeoutMs: number
+): Promise<void> {
+  const deadline = Date.now() + timeoutMs;
+  while (true) {
+    const probe = Bun.spawnSync({ cmd: ["pgrep", "-f", marker] });
+    if (probe.exitCode !== 0) {
+      return;
+    }
+    if (Date.now() >= deadline) {
+      throw new Error(`Process still running for marker ${marker}`);
+    }
+    await Bun.sleep(25);
+  }
+}
+
+describe("runBashCommand (integration)", () => {
+  test("captures stdout from a successful command", async () => {
+    const r = await runBashCommand(
+      "echo hello",
+      5000,
+      undefined,
+      process.cwd()
+    );
+    expect(r.exitCode).toBe(0);
+    expect(r.aborted).toBe(false);
+    expect(r.timedOut).toBe(false);
+    expect(r.stdout.text.trim()).toBe("hello");
+    expect(r.stderr.totalBytes).toBe(0);
+  });
+
+  test("captures stderr and non-zero exit", async () => {
+    const r = await runBashCommand(
+      "echo oops 1>&2; exit 3",
+      5000,
+      undefined,
+      process.cwd()
+    );
+    expect(r.exitCode).toBe(3);
+    expect(r.stderr.text.trim()).toBe("oops");
+  });
+
+  test("respects cwd", async () => {
+    const r = await runBashCommand("pwd", 5000, undefined, "/tmp");
+    expect(r.exitCode).toBe(0);
+    expect(r.stdout.text.trim()).toBe("/tmp");
+  });
+
+  test("times out and reports timedOut", async () => {
+    const r = await runBashCommand("sleep 5", 25, undefined, process.cwd());
+    expect(r.timedOut).toBe(true);
+    expect(r.exitCode === null || r.exitCode !== 0).toBe(true);
+  });
+
+  test("aborts when signal fires", async () => {
+    const ctrl = new AbortController();
+    const promise = runBashCommand("sleep 5", 5000, ctrl.signal, process.cwd());
+    ctrl.abort();
+    const r = await promise;
+    expect(r.aborted).toBe(true);
+  });
+
+  test("returns promptly when a backgrounded child inherits the pipe", async () => {
+    const startedAt = Date.now();
+    const r = await runBashCommand(
+      "nohup sleep 47 > /dev/null 2>&1 & disown; echo done",
+      5000,
+      undefined,
+      process.cwd()
+    );
+    const elapsed = Date.now() - startedAt;
+    expect(r.exitCode).toBe(0);
+    expect(r.timedOut).toBe(false);
+    expect(r.stdout.text.trim()).toBe("done");
+    expect(elapsed).toBeLessThan(2000);
+    // clean up the orphaned sleep so it doesn't linger
+    try {
+      Bun.spawnSync({ cmd: ["pkill", "-f", "sleep 47"] });
+    } catch {}
+  });
+
+  test("timeout kills the whole process group", async () => {
+    const marker = `pim-test-timeout-${Date.now()}`;
+    const startedAt = Date.now();
+    const r = await runBashCommand(
+      `bash -c ${shellQuote(`exec -a ${marker} sleep 60`)}`,
+      50,
+      undefined,
+      process.cwd()
+    );
+    const elapsed = Date.now() - startedAt;
+    expect(r.timedOut).toBe(true);
+    expect(r.exitCode === null || r.exitCode !== 0).toBe(true);
+    expect(elapsed).toBeLessThan(KILL_GRACE_MS + 2000);
+    await waitForNoProcess(marker, KILL_GRACE_MS + 500);
+  });
+
+  test("does not crash on timeout while drains still hold readers", async () => {
+    // Regression: stream.cancel() on a locked stream rejects (Bun throws
+    // synchronously). Drains run fire-and-forget, so when a quiet command
+    // times out (no output → drain blocked on read), the finally hits
+    // cancel before drain has released. Unhandled rejection would crash Bun.
+    const rejections: unknown[] = [];
+    const onRejection = (err: unknown) => rejections.push(err);
+    process.on("unhandledRejection", onRejection);
+    try {
+      const r = await runBashCommand("sleep 5", 25, undefined, process.cwd());
+      expect(r.timedOut).toBe(true);
+      await Bun.sleep(25);
+      expect(rejections).toEqual([]);
+    } finally {
+      process.off("unhandledRejection", onRejection);
+    }
+  });
+
+  test("bounded drain returns even when a daemon escapes our process group", async () => {
+    // A child that calls setsid itself leaves our pgid and survives killGroup.
+    // If it keeps the pipe open, drain would block forever; the DRAIN_GRACE_MS
+    // bound forces us to return anyway. The marker lets us clean up after.
+    const marker = `pim-test-detached-${Date.now()}`;
+    const startedAt = Date.now();
+    const r = await runBashCommand(
+      `setsid bash -c 'sleep 60; echo ${marker}' > /tmp/${marker}.out 2>&1 < /dev/null & disown; echo done`,
+      5000,
+      undefined,
+      process.cwd()
+    );
+    const elapsed = Date.now() - startedAt;
+    expect(r.exitCode).toBe(0);
+    expect(r.stdout.text.trim()).toBe("done");
+    expect(elapsed).toBeLessThan(DRAIN_GRACE_MS + 2000);
+    try {
+      Bun.spawnSync({ cmd: ["pkill", "-f", marker] });
+      Bun.spawnSync({ cmd: ["rm", "-f", `/tmp/${marker}.out`] });
+    } catch {}
+  });
+
+  test("killAllActiveBashGroups sweeps in-flight subtrees", async () => {
+    const id = Date.now();
+    const marker = `/tmp/pim-test-active-${id}.marker`;
+    const ready = `/tmp/pim-test-active-${id}.ready`;
+    const processMarker = `pim-test-active-${id}`;
+    const pending = runBashCommand(
+      `touch ${shellQuote(ready)}; bash -c ${shellQuote(`exec -a ${processMarker} sleep 30`)} && touch ${shellQuote(marker)}`,
+      30_000,
+      undefined,
+      process.cwd()
+    );
+    try {
+      await waitForFile(ready, 1000);
+      killAllActiveBashGroups("SIGTERM");
+      const result = await pending;
+      expect(result.exitCode === null || result.exitCode !== 0).toBe(true);
+      await waitForNoProcess(processMarker, KILL_GRACE_MS + 500);
+      expect(await Bun.file(marker).exists()).toBe(false);
+    } finally {
+      Bun.spawnSync({ cmd: ["rm", "-f", marker, ready] });
+    }
+  });
+
+  test("truncates very large stdout", async () => {
+    const totalBytes = STREAM_HEAD_BYTES + STREAM_TAIL_BYTES + 1000;
+    const r = await runBashCommand(
+      `head -c ${totalBytes} /dev/zero | tr '\\0' 'A'`,
+      5000,
+      undefined,
+      process.cwd()
+    );
+    expect(r.exitCode).toBe(0);
+    expect(r.stdout.totalBytes).toBe(totalBytes);
+    expect(r.stdout.truncated).toBe(true);
+    expect(r.stdout.text).toContain("bytes truncated");
+  });
+
+  test("spills full stdout to ~/.pim/cache when truncated", async () => {
+    const totalBytes = STREAM_HEAD_BYTES + STREAM_TAIL_BYTES + 4096;
+    const r = await runBashCommand(
+      `head -c ${totalBytes} /dev/zero | tr '\\0' 'A'`,
+      5000,
+      undefined,
+      process.cwd()
+    );
+    try {
+      expect(r.exitCode).toBe(0);
+      expect(r.stdout.truncated).toBe(true);
+      expect(r.stdout.path).toBeTruthy();
+      expect(r.stdout.path!.startsWith(join(SpillCache.dir(), "bash-"))).toBe(
+        true
+      );
+      expect(r.stdout.path!.endsWith(".out")).toBe(true);
+      const cacheMode = (await stat(SpillCache.dir())).mode & 0o777;
+      const spillMode = (await stat(r.stdout.path!)).mode & 0o777;
+      expect(cacheMode).toBe(0o700);
+      expect(spillMode).toBe(0o600);
+      const spilled = await Bun.file(r.stdout.path!).text();
+      expect(spilled.length).toBe(totalBytes);
+      expect(spilled).toBe("A".repeat(totalBytes));
+    } finally {
+      if (r.stdout.path) {
+        try {
+          Bun.spawnSync({ cmd: ["rm", "-f", r.stdout.path] });
+        } catch {}
+      }
+    }
+  });
+
+  test("omits spill path when stream is empty", async () => {
+    const r = await runBashCommand("true", 5000, undefined, process.cwd());
+    expect(r.exitCode).toBe(0);
+    expect(r.stdout.path).toBeNull();
+    expect(r.stderr.path).toBeNull();
+  });
+});

package/src/extensions/bash/run.ts

@@ -0,0 +1,207 @@
+import { SpillCache } from "../../shared/SpillCache";
+import { StreamCapture } from "./capture";
+import {
+  type BashCommandResult,
+  DRAIN_GRACE_MS,
+  KILL_GRACE_MS,
+} from "./schema";
+
+type Reader = ReadableStreamDefaultReader<Uint8Array>;
+
+const activePids = new Set<number>();
+
+// Wired into the extension's signal handlers so a daemon that `setsid`s
+// out of our group (or harbor/parent SIGTERM) still tears down its subtree.
+export function killAllActiveBashGroups(sig: NodeJS.Signals = "SIGTERM"): void {
+  for (const pid of activePids) {
+    killGroup(pid, sig);
+  }
+  activePids.clear();
+}
+
+async function drain(reader: Reader | null, cap: StreamCapture): Promise<void> {
+  if (!reader) {
+    return;
+  }
+  try {
+    while (true) {
+      const { value, done } = await reader.read();
+      if (done) {
+        break;
+      }
+      if (value) {
+        cap.push(value);
+      }
+    }
+  } catch {
+    // reader cancelled; drop remaining bytes
+  } finally {
+    try {
+      reader.releaseLock();
+    } catch {}
+  }
+}
+
+async function spillIfTruncated(
+  cap: StreamCapture,
+  ext: "out" | "err"
+): Promise<string | null> {
+  if (!cap.truncated) {
+    return null;
+  }
+  return SpillCache.write("bash", ext, cap.full());
+}
+
+function killGroup(pid: number | undefined, sig: NodeJS.Signals): void {
+  if (pid === undefined) {
+    return;
+  }
+  try {
+    process.kill(-pid, sig);
+  } catch {
+    try {
+      process.kill(pid, sig);
+    } catch {}
+  }
+}
+
+function getReader(
+  stream: ReadableStream<Uint8Array> | undefined
+): Reader | null {
+  return stream ? stream.getReader() : null;
+}
+
+export async function runBashCommand(
+  command: string,
+  timeoutMs: number,
+  signal: AbortSignal | undefined,
+  cwd: string
+): Promise<BashCommandResult> {
+  const startedAt = Date.now();
+  const stdoutCap = new StreamCapture();
+  const stderrCap = new StreamCapture();
+
+  // setsid puts bash and its descendants into a fresh process group with
+  // pgid == proc.pid, so we can signal the whole tree on timeout/abort
+  // instead of leaving backgrounded grandchildren alive holding our pipes.
+  const proc = Bun.spawn({
+    cmd: ["setsid", "bash", "-lc", command],
+    cwd,
+    stdout: "pipe",
+    stderr: "pipe",
+    env: { ...process.env },
+  });
+  if (proc.pid !== undefined) {
+    activePids.add(proc.pid);
+  }
+
+  let timedOut = false;
+  let aborted = false;
+
+  // We own the readers so we can force-cancel them later even while the
+  // background drains are still mid-read. Cancelling via the held reader
+  // does not throw the way ReadableStream.cancel() on a locked stream does.
+  const stdoutReader = getReader(
+    proc.stdout as unknown as ReadableStream<Uint8Array>
+  );
+  const stderrReader = getReader(
+    proc.stderr as unknown as ReadableStream<Uint8Array>
+  );
+
+  // Fire-and-forget drains. A backgrounded child can inherit the subshell's
+  // fds and keep the pipes open after bash exits, so we can't block on EOF;
+  // we race proc.exited against a wall-clock timeout instead.
+  const stdoutDrain = drain(stdoutReader, stdoutCap);
+  const stderrDrain = drain(stderrReader, stderrCap);
+
+  let timeoutHandle: ReturnType<typeof setTimeout> | null = null;
+  const exitedPromise = proc.exited.then(() => "exited" as const);
+  const timeoutPromise = new Promise<"timeout">((resolve) => {
+    timeoutHandle = setTimeout(() => resolve("timeout"), timeoutMs);
+  });
+
+  let abortResolve: ((v: "aborted") => void) | null = null;
+  const abortPromise = new Promise<"aborted">((resolve) => {
+    abortResolve = resolve;
+  });
+  const onAbort = () => {
+    aborted = true;
+    abortResolve?.("aborted");
+  };
+  if (signal) {
+    if (signal.aborted) {
+      onAbort();
+    } else {
+      signal.addEventListener("abort", onAbort, { once: true });
+    }
+  }
+
+  let exitCode: number | null = null;
+  let signalCode: NodeJS.Signals | null = null;
+  try {
+    const result = await Promise.race([
+      exitedPromise,
+      timeoutPromise,
+      abortPromise,
+    ]);
+
+    if (result === "timeout") {
+      timedOut = true;
+    }
+
+    if (result !== "exited") {
+      killGroup(proc.pid, "SIGTERM");
+      const sigkillTimer = setTimeout(() => {
+        killGroup(proc.pid, "SIGKILL");
+      }, KILL_GRACE_MS);
+      try {
+        await proc.exited;
+      } finally {
+        clearTimeout(sigkillTimer);
+      }
+    }
+
+    exitCode = proc.exitCode ?? null;
+    signalCode = (proc.signalCode as NodeJS.Signals | null | undefined) ?? null;
+
+    // Bound the drain so a detached grandchild holding the pipe can't keep
+    // the drain promise + capture buffer alive past this call.
+    await Promise.race([
+      Promise.all([stdoutDrain, stderrDrain]),
+      Bun.sleep(DRAIN_GRACE_MS),
+    ]);
+  } finally {
+    if (timeoutHandle) {
+      clearTimeout(timeoutHandle);
+    }
+    if (signal) {
+      signal.removeEventListener("abort", onAbort);
+    }
+    for (const reader of [stdoutReader, stderrReader]) {
+      if (!reader) {
+        continue;
+      }
+      try {
+        void reader.cancel().catch(() => {});
+      } catch {}
+    }
+    if (proc.pid !== undefined) {
+      activePids.delete(proc.pid);
+    }
+  }
+
+  const [stdoutPath, stderrPath] = await Promise.all([
+    spillIfTruncated(stdoutCap, "out"),
+    spillIfTruncated(stderrCap, "err"),
+  ]);
+
+  return {
+    exitCode,
+    signal: signalCode,
+    stdout: { ...stdoutCap.snapshot(), path: stdoutPath },
+    stderr: { ...stderrCap.snapshot(), path: stderrPath },
+    timedOut,
+    aborted,
+    durationMs: Date.now() - startedAt,
+  };
+}

package/src/extensions/bash/schema.ts

@@ -0,0 +1,54 @@
+import { type Static, Type } from "typebox";
+
+export const STREAM_HEAD_BYTES = 8192;
+export const STREAM_TAIL_BYTES = 8192;
+export const DEFAULT_TIMEOUT_MS = 30_000;
+export const KILL_GRACE_MS = 2000;
+export const DRAIN_GRACE_MS = 1000;
+
+export const bashSchema = Type.Object({
+  command: Type.String({
+    description: "Runs via bash -lc, so login shell init applies.",
+  }),
+  timeoutMs: Type.Optional(
+    Type.Integer({
+      minimum: 1,
+      description: `Timeout in milliseconds. Default is ${DEFAULT_TIMEOUT_MS} (${DEFAULT_TIMEOUT_MS / 1000}s) — raise it for long-running commands like builds, test suites, training runs, or installs.`,
+    })
+  ),
+});
+export type BashInput = Static<typeof bashSchema>;
+
+export type CapturedStream = {
+  readonly text: string;
+  readonly totalBytes: number;
+  readonly truncated: boolean;
+  readonly path: string | null;
+  readonly nextStart: number | null;
+};
+
+export type BashCommandResult = {
+  readonly exitCode: number | null;
+  readonly signal: NodeJS.Signals | null;
+  readonly stdout: CapturedStream;
+  readonly stderr: CapturedStream;
+  readonly timedOut: boolean;
+  readonly aborted: boolean;
+  readonly durationMs: number;
+};
+
+export type BashStreamDetails = {
+  readonly totalBytes: number;
+  readonly truncated: boolean;
+  readonly path: string | null;
+};
+
+export type BashDetails = {
+  readonly exitCode: number | null;
+  readonly signal: NodeJS.Signals | null;
+  readonly durationMs: number;
+  readonly timedOut: boolean;
+  readonly aborted: boolean;
+  readonly stdout: BashStreamDetails;
+  readonly stderr: BashStreamDetails;
+};

package/src/extensions/command-picker/index.ts

@@ -0,0 +1,52 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { rank } from "./ranker";
+
+const MAX_VISIBLE_ROWS = 10;
+const SLASH_PREFIX = /^\/([^\s]*)$/;
+const SLASH_LINES = ["/"];
+
+export default function (pi: ExtensionAPI): void {
+  pi.on("session_start", (_event, ctx) => {
+    ctx.ui.addAutocompleteProvider((current) => ({
+      async getSuggestions(lines, cursorLine, cursorCol, options) {
+        const line = lines[cursorLine] ?? "";
+        const beforeCursor = line.slice(0, cursorCol);
+
+        const slashMatch = beforeCursor.match(SLASH_PREFIX);
+        if (!slashMatch) {
+          return current.getSuggestions(lines, cursorLine, cursorCol, options);
+        }
+
+        const all = await current.getSuggestions(SLASH_LINES, 0, 1, options);
+        if (all === null) {
+          return null;
+        }
+
+        const items = rank(slashMatch[1] ?? "", all.items, {
+          limit: MAX_VISIBLE_ROWS,
+        });
+        if (items.length === 0) {
+          return null;
+        }
+        return { items, prefix: beforeCursor };
+      },
+
+      applyCompletion(lines, cursorLine, cursorCol, item, prefix) {
+        return current.applyCompletion(
+          lines,
+          cursorLine,
+          cursorCol,
+          item,
+          prefix
+        );
+      },
+
+      shouldTriggerFileCompletion(lines, cursorLine, cursorCol) {
+        return (
+          current.shouldTriggerFileCompletion?.(lines, cursorLine, cursorCol) ??
+          true
+        );
+      },
+    }));
+  });
+}

package/src/extensions/command-picker/ranker.test.ts

@@ -0,0 +1,46 @@
+import { expect, test } from "bun:test";
+import type { AutocompleteItem } from "@earendil-works/pi-tui";
+import { rank } from "./ranker";
+
+const item = (name: string, description?: string): AutocompleteItem => ({
+  value: name,
+  label: name,
+  ...(description !== undefined && { description }),
+});
+
+test("empty query returns items alphabetically by label", () => {
+  const items = rank("", [
+    item("rename", "Rename the session."),
+    item("clear", "Clear the session."),
+    item("help", "Show help."),
+  ]);
+
+  expect(items.map((i) => i.value)).toEqual(["clear", "help", "rename"]);
+});
+
+test("query ranks fuzzy matches by score", () => {
+  const items = rank("cl", [
+    item("rename", "Rename the session."),
+    item("clear", "Clear the session."),
+    item("help", "Show help."),
+  ]);
+
+  expect(items[0]?.value).toBe("clear");
+});
+
+test("matches against description when label doesn't contain query", () => {
+  const items = rank("rename", [
+    item("noop", "fully unrelated"),
+    item("x", "rename the session"),
+  ]);
+
+  expect(items[0]?.value).toBe("x");
+});
+
+test("limit caps the returned items", () => {
+  const items = rank("", [item("c"), item("a"), item("b"), item("d")], {
+    limit: 2,
+  });
+
+  expect(items.map((i) => i.value)).toEqual(["a", "b"]);
+});

package/src/extensions/command-picker/ranker.ts

@@ -0,0 +1,17 @@
+import type { AutocompleteItem } from "@earendil-works/pi-tui";
+import { FuzzyMatcher, type FuzzyCandidate } from "../../shared/FuzzyMatcher";
+
+export function rank(
+  query: string,
+  items: readonly AutocompleteItem[],
+  options: { readonly limit?: number } = {}
+): AutocompleteItem[] {
+  const candidates: FuzzyCandidate<AutocompleteItem>[] = items.map((item) => ({
+    item,
+    haystacks: [item.label, item.description ?? ""],
+  }));
+
+  const hits = FuzzyMatcher.rank(query, candidates, options);
+
+  return hits.map((hit) => hit.item);
+}