omnirush 0.8.5 → 0.9.0

package/assets/extensions/omnirush/bgshell-lib.ts

@@ -0,0 +1,400 @@
+// bgshell-lib — shell commands that outlive their tool call.
+//
+// The bash tool (bgshell.ts) runs every command as a job here. A command
+// that finishes within its foreground window answers the call as before; one
+// that runs longer (or that the model started with run_in_background) is not
+// killed: it moves to the background, keeps running, its output goes to a
+// bounded buffer, and the call returns its id at once. bash_output /
+// bash_wait / bash_kill read, wait for and stop it; when it finishes, its
+// result is delivered to the main agent like a background sub-agent's
+// (deliveries.ts). Jobs still running when the session ends are stopped.
+//
+// Pure apart from the injected `exec` (pi's own local shell execution, which
+// runs the command in its own process group and kills the whole tree), so
+// node:test drives it with fakes.
+
+/** How long a command runs in the foreground before it moves to the background (default). */
+export const DEFAULT_FOREGROUND_SECONDS = 300;
+/** The longest foreground wait a call may ask for. */
+export const MAX_FOREGROUND_SECONDS = 3_600;
+/** Output kept per job (characters): the newest part of it. */
+export const JOB_BUFFER_CHARS = 1_000_000;
+/** Output a delivery or a read shows at most (characters): the newest part. */
+export const SHOW_CHARS = 30_000;
+/** Longest wait a bash_wait call may ask for (seconds). */
+export const MAX_WAIT_SECONDS = 3_600;
+
+export const BGSHELL_HOOK = Symbol.for("omnirush.bgshell");
+export interface BgShellHook {
+  /** Stop every running job (of one session, or all) and record their output; resolves once they exited. */
+  stopAll(sessionId?: string): Promise<Array<{ id: string; command: string }>>;
+}
+
+/** pi's shell execution (BashOperations.exec). */
+export type ShellExec = (
+  command: string,
+  cwd: string,
+  options: { onData: (data: Buffer | string) => void; signal?: AbortSignal; timeout?: number; env?: NodeJS.ProcessEnv },
+) => Promise<{ exitCode: number | null }>;
+
+export type JobStatus = "running" | "exited" | "killed" | "failed";
+
+/** The newest `max` characters of a stream, with absolute offsets. */
+export class OutputBuffer {
+  private text = "";
+  /** Characters appended in total. */
+  total = 0;
+  constructor(private readonly max: number = JOB_BUFFER_CHARS) {}
+
+  append(chunk: string): void {
+    if (!chunk) return;
+    this.text += chunk;
+    this.total += chunk.length;
+    if (this.text.length > this.max * 1.25) this.text = this.text.slice(this.text.length - this.max);
+  }
+
+  /** Offset of the oldest character still kept. */
+  get start(): number {
+    return this.total - this.text.length;
+  }
+
+  /** Everything after `offset` that is still kept, and how much of it was dropped. */
+  since(offset: number): { text: string; dropped: number } {
+    const from = Math.max(offset, this.start);
+    return { text: this.text.slice(from - this.start), dropped: Math.max(0, this.start - offset) };
+  }
+}
+
+// eslint-disable-next-line no-control-regex
+const ANSI = /\x1b\[[0-9;?]*[ -/]*[@-~]|\x1b\][^\x07\x1b]*(?:\x07|\x1b\\)|\x1b[@-Z\\-_]/g;
+// eslint-disable-next-line no-control-regex
+const CONTROL = /[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]/g;
+
+/** Terminal escapes and control characters out; \r\n and progress-bar \r become line ends. */
+export function cleanOutput(text: string): string {
+  return text.replace(ANSI, "").replace(/\r\n/g, "\n").replace(/\r(?!\n)/g, "\n").replace(CONTROL, "");
+}
+
+/** The newest `max` characters, starting at a line boundary when one is close. */
+export function tailText(text: string, max: number = SHOW_CHARS): { text: string; truncated: boolean } {
+  if (text.length <= max) return { text, truncated: false };
+  let tail = text.slice(text.length - max);
+  const newline = tail.indexOf("\n");
+  if (newline >= 0 && newline < 2_000) tail = tail.slice(newline + 1);
+  return { text: tail, truncated: true };
+}
+
+export interface ShellJob {
+  /** Short handle for the bash_* tools ("sh1", "sh2", ...). */
+  id: string;
+  session: string;
+  command: string;
+  cwd: string;
+  /** The bash call that started it. */
+  toolCallId: string | null;
+  startedAt: number;
+  finishedAt: number | null;
+  status: JobStatus;
+  exitCode: number | null;
+  error: string | null;
+  /** Moved to the background (its result is delivered when it finishes). */
+  background: boolean;
+  /** Started with run_in_background (the model did not wait for it). */
+  explicit: boolean;
+  backgroundedAt: number | null;
+  output: OutputBuffer;
+  /** Where bash_output continues ("new output"). */
+  readOffset: number;
+  /** Its final result reached the model (tool result, read, wait, kill or delivery). */
+  delivered: boolean;
+  /** Stopped because the session ended. */
+  interrupted: boolean;
+  controller: AbortController;
+  done: Promise<void>;
+}
+
+/** Tracks the shell jobs of this process per session. */
+export class ShellJobManager {
+  private readonly exec: ShellExec;
+  private readonly now: () => number;
+  private readonly onDeliverable: () => void;
+  private readonly jobs: ShellJob[] = [];
+  private readonly listeners = new Map<ShellJob, Set<(text: string) => void>>();
+  /** The foreground call's raw output sink, until the job moves to the background. */
+  private readonly forwards = new Map<ShellJob, (data: Buffer | string) => void>();
+  private changeWaiters: Array<() => void> = [];
+  private next = 1;
+  /** Jobs a bash_wait call is waiting for (it returns their result itself). */
+  private readonly awaited = new Map<ShellJob, number>();
+  /** Explicit background jobs the one-shot reminder already named. */
+  private readonly reminded = new Set<ShellJob>();
+
+  constructor(options: { exec: ShellExec; now?: () => number; onDeliverable?: () => void }) {
+    this.exec = options.exec;
+    this.now = options.now ?? Date.now;
+    this.onDeliverable = options.onDeliverable ?? (() => undefined);
+  }
+
+  start(input: {
+    session: string;
+    command: string;
+    cwd: string;
+    env?: NodeJS.ProcessEnv;
+    toolCallId?: string | null;
+    onData?: (data: Buffer | string) => void;
+  }): ShellJob {
+    const controller = new AbortController();
+    const job: ShellJob = {
+      id: `sh${this.next++}`,
+      session: input.session,
+      command: input.command,
+      cwd: input.cwd,
+      toolCallId: input.toolCallId ?? null,
+      startedAt: this.now(),
+      finishedAt: null,
+      status: "running",
+      exitCode: null,
+      error: null,
+      background: false,
+      explicit: false,
+      backgroundedAt: null,
+      output: new OutputBuffer(),
+      readOffset: 0,
+      delivered: false,
+      interrupted: false,
+      controller,
+      done: Promise.resolve(),
+    };
+    this.jobs.push(job);
+    const decoder = new TextDecoder();
+    const onData = (data: Buffer | string) => {
+      const text = cleanOutput(typeof data === "string" ? data : decoder.decode(data, { stream: true }));
+      job.output.append(text);
+      for (const listener of this.listeners.get(job) ?? []) listener(text);
+    };
+    if (input.onData) this.forwards.set(job, input.onData);
+    job.done = (async () => {
+      try {
+        const result = await this.exec(input.command, input.cwd, {
+          onData: (data) => {
+            this.forwards.get(job)?.(data);
+            onData(data);
+          },
+          signal: controller.signal,
+          ...(input.env ? { env: input.env } : {}),
+        });
+        this.finish(job, controller.signal.aborted ? "killed" : "exited", result?.exitCode ?? null, null);
+      } catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        if (controller.signal.aborted || message === "aborted") this.finish(job, "killed", null, null);
+        else this.finish(job, "failed", null, message);
+      }
+    })();
+    return job;
+  }
+
+  private finish(job: ShellJob, status: JobStatus, exitCode: number | null, error: string | null): void {
+    if (job.status !== "running") return;
+    job.status = status;
+    job.exitCode = exitCode;
+    job.error = error;
+    job.finishedAt = this.now();
+    this.listeners.delete(job);
+    this.forwards.delete(job);
+    // A waiting bash_wait hands the result over: no delivery message.
+    if ((this.awaited.get(job) ?? 0) > 0) job.delivered = true;
+    const waiters = this.changeWaiters;
+    this.changeWaiters = [];
+    waiters.forEach((resolve) => resolve());
+    if (job.background && !job.delivered) this.onDeliverable();
+  }
+
+  /** Hear the job's (cleaned) output from now on; returns the unsubscribe. */
+  listen(job: ShellJob, listener: (text: string) => void): () => void {
+    const set = this.listeners.get(job) ?? new Set();
+    set.add(listener);
+    this.listeners.set(job, set);
+    return () => set.delete(listener);
+  }
+
+  /** The foreground call stops waiting: the job keeps running in the background. */
+  toBackground(job: ShellJob, explicit = false): void {
+    if (job.background) return;
+    this.forwards.delete(job);
+    job.background = true;
+    job.explicit = explicit;
+    job.backgroundedAt = this.now();
+    // What the foreground call showed is not "new" for bash_output.
+    job.readOffset = job.output.total;
+    if (job.status !== "running" && !job.delivered) this.onDeliverable();
+  }
+
+  list(session?: string): ShellJob[] {
+    return this.jobs.filter((job) => session === undefined || job.session === session);
+  }
+
+  /** Jobs by id (every background job of the session when `ids` is empty or "all"); unknown ids apart. */
+  resolve(session: string, ids: readonly string[] | undefined): { jobs: ShellJob[]; unknown: string[] } {
+    const all = this.list(session).filter((job) => job.background);
+    if (!ids || ids.length === 0 || ids.includes("all")) return { jobs: all, unknown: [] };
+    const jobs: ShellJob[] = [];
+    const unknown: string[] = [];
+    for (const raw of ids) {
+      const id = String(raw).trim();
+      const job = all.find((candidate) => candidate.id === id);
+      if (job && !jobs.includes(job)) jobs.push(job);
+      else if (!job) unknown.push(id);
+    }
+    return { jobs, unknown };
+  }
+
+  running(session?: string): ShellJob[] {
+    return this.list(session).filter((job) => job.status === "running");
+  }
+
+  /** New output since the last read (marks it read; a finished job's result is then delivered). */
+  read(job: ShellJob): { text: string; dropped: number } {
+    const out = job.output.since(job.readOffset);
+    job.readOffset = job.output.total;
+    if (job.status !== "running") job.delivered = true;
+    return out;
+  }
+
+  /** Wait until the jobs finished, `timeoutMs` passed or `signal` fired (they keep running then). */
+  async wait(jobs: ShellJob[], options: { timeoutMs?: number; signal?: AbortSignal } = {}): Promise<void> {
+    for (const job of jobs) this.awaited.set(job, (this.awaited.get(job) ?? 0) + 1);
+    const all = Promise.all(jobs.map((job) => job.done)).then(() => undefined);
+    const stops: Array<Promise<void>> = [all];
+    let timer: ReturnType<typeof setTimeout> | null = null;
+    if (options.timeoutMs && options.timeoutMs > 0) stops.push(new Promise((resolve) => { timer = setTimeout(resolve, options.timeoutMs); }));
+    let onAbort: (() => void) | null = null;
+    if (options.signal) {
+      stops.push(new Promise((resolve) => {
+        onAbort = () => resolve();
+        if (options.signal!.aborted) resolve();
+        else options.signal!.addEventListener("abort", onAbort, { once: true });
+      }));
+    }
+    try {
+      await Promise.race(stops);
+    } finally {
+      if (timer) clearTimeout(timer);
+      if (onAbort) options.signal?.removeEventListener("abort", onAbort);
+      for (const job of jobs) {
+        const count = (this.awaited.get(job) ?? 1) - 1;
+        if (count > 0) this.awaited.set(job, count);
+        else this.awaited.delete(job);
+      }
+    }
+  }
+
+  /** Stop jobs (their process trees) and wait for them to exit. */
+  async kill(jobs: ShellJob[], options: { interrupted?: boolean; deliver?: boolean } = {}): Promise<ShellJob[]> {
+    const live = jobs.filter((job) => job.status === "running");
+    for (const job of live) {
+      if (options.interrupted) job.interrupted = true;
+      // The caller reports them itself (unless the user stopped them: then
+      // the model is told like for any finished background command).
+      if (!options.deliver) job.delivered = true;
+      job.controller.abort();
+    }
+    await Promise.all(live.map((job) => job.done));
+    return live;
+  }
+
+  /** Finished background jobs whose result the model has not seen (marked delivered). */
+  takeDeliverable(session: string): ShellJob[] {
+    const ready = this.list(session).filter((job) => job.background && job.status !== "running" && !job.delivered);
+    ready.forEach((job) => { job.delivered = true; });
+    return ready;
+  }
+
+  /**
+   * One-shot runs wait for: jobs that moved to the background while the model
+   * waited for them, and finished results not yet seen. Jobs the model
+   * started in the background on purpose get one reminder instead
+   * (`remindable`).
+   */
+  hasPending(session: string): boolean {
+    return this.list(session).some((job) => job.background && (job.status === "running" ? !job.explicit : !job.delivered));
+  }
+
+  /** Explicit background jobs still running that the model was not reminded of yet (marked reminded). */
+  remindable(session: string): ShellJob[] {
+    const jobs = this.running(session).filter((job) => job.background && job.explicit && !this.reminded.has(job));
+    jobs.forEach((job) => this.reminded.add(job));
+    return jobs;
+  }
+
+  changed(): Promise<void> {
+    return new Promise((resolve) => this.changeWaiters.push(resolve));
+  }
+}
+
+function minutes(ms: number): string {
+  const seconds = Math.max(0, Math.round(ms / 1000));
+  if (seconds < 90) return `${seconds} s`;
+  return `${Math.round((seconds / 60) * 10) / 10} min`;
+}
+
+export function shortCommand(command: string, max = 80): string {
+  const line = command.trim().replace(/\s+/g, " ");
+  return line.length > max ? `${line.slice(0, max - 1)}…` : line;
+}
+
+/** "exit 0", "killed", "failed: ..." */
+export function jobOutcome(job: ShellJob): string {
+  if (job.status === "running") return "running";
+  if (job.status === "killed") return job.interrupted ? "stopped (the session ended)" : "killed";
+  if (job.status === "failed") return `failed${job.error ? `: ${job.error}` : ""}`;
+  return `exit ${job.exitCode ?? "?"}`;
+}
+
+/** One status line per job. */
+export function renderJobStatus(jobs: ShellJob[], now: number = Date.now()): string {
+  if (jobs.length === 0) return "No background commands in this session.";
+  return jobs
+    .map((job) => `- ${job.id} [${jobOutcome(job)}] ${minutes((job.finishedAt ?? now) - job.startedAt)}, ${job.output.total} chars of output — ${shortCommand(job.command)}`)
+    .join("\n");
+}
+
+/** A job's result for the model: header, then the newest output. */
+export function renderJobResult(job: ShellJob, text: string, options: { dropped?: number; heading?: string; now?: number } = {}): { text: string; output: string; truncated: boolean } {
+  const shown = tailText(text);
+  const lost = (options.dropped ?? 0) > 0 || shown.truncated;
+  const heading = options.heading ?? `${job.id} [${jobOutcome(job)}] after ${minutes((job.finishedAt ?? options.now ?? Date.now()) - job.startedAt)}: ${shortCommand(job.command, 200)}`;
+  const body = shown.text.trim() ? shown.text.replace(/\s+$/, "") : "(no new output)";
+  return {
+    text: `${heading}\n${lost ? "[earlier output omitted — only the newest part is kept]\n" : ""}${body}`,
+    output: shown.text,
+    truncated: lost,
+  };
+}
+
+/** The message that brings finished background commands back to the model. */
+export function renderJobDelivery(jobs: Array<{ job: ShellJob; text: string; dropped: number }>): string {
+  const ids = jobs.map(({ job }) => job.id).join(", ");
+  return [
+    `Background command${jobs.length === 1 ? "" : "s"} finished (${ids}). Output since you last read it follows; continue your work with it.`,
+    "",
+    ...jobs.map(({ job, text, dropped }) => renderJobResult(job, text, { dropped }).text).join("\n\n---\n\n").split("\n"),
+  ].join("\n");
+}
+
+/** A job's record in a message's details (the trace turns it into a bash tool part). */
+export function jobDetails(job: ShellJob, output: string, truncated: boolean): Record<string, unknown> {
+  return {
+    id: job.id,
+    command: job.command,
+    cwd: job.cwd,
+    status: job.status,
+    exit_code: job.exitCode,
+    ...(job.error ? { error: job.error } : {}),
+    ...(job.interrupted ? { interrupted: true } : {}),
+    tool_call_id: job.toolCallId,
+    started_at: job.startedAt,
+    finished_at: job.finishedAt,
+    output,
+    truncated,
+  };
+}