index.ts

import { randomBytes } from "node:crypto";
import { createWriteStream, mkdirSync } from "node:fs";
import { join } from "node:path";
import { tmpdir } from "node:os";
import { spawn } from "node:child_process";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import {
  createLocalBashOperations,
  DEFAULT_MAX_BYTES,
  DEFAULT_MAX_LINES,
  formatSize,
  truncateTail,
  getShellConfig,
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

const LOG_DIR = join(tmpdir(), "ppb");

// ---------------------------------------------------------------------------
// Background process tracking
// ---------------------------------------------------------------------------

const trackedPids = new Set<number>();

function killProcessTree(pid: number): void {
  if (process.platform === "win32") {
    try {
      spawn("taskkill", ["/F", "/T", "/PID", String(pid)], {
        stdio: "ignore",
        detached: true,
        windowsHide: true,
      });
    } catch {
      /* ignore */
    }
  } else {
    try {
      process.kill(-pid, "SIGKILL");
    } catch {
      try {
        process.kill(pid, "SIGKILL");
      } catch {
        /* already dead */
      }
    }
  }
}

function killAllTracked(): void {
  for (const pid of trackedPids) {
    killProcessTree(pid);
  }
  trackedPids.clear();
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

function generateId(): string {
  return randomBytes(3).toString("hex");
}

function ensureLogDir(): void {
  mkdirSync(LOG_DIR, { recursive: true });
}

function logPath(id: string): string {
  return join(LOG_DIR, `${id}.log`);
}

// ---------------------------------------------------------------------------
// Simple output collector (foreground mode)
// ---------------------------------------------------------------------------

interface SnapshotResult {
  content: string;
  truncated: boolean;
  truncatedBy: "lines" | "bytes" | null;
  totalLines: number;
  totalBytes: number;
  fullOutputPath: string | undefined;
}

class SimpleCollector {
  private chunks: Buffer[] = [];
  private text = "";
  private totalBytes = 0;
  private totalLines = 0;
  private finished = false;
  private tempFilePath: string | undefined;
  private tempFileStream: ReturnType<typeof createWriteStream> | undefined;
  private decoder = new TextDecoder();

  append(data: Buffer): void {
    if (this.finished) return;
    this.totalBytes += data.length;
    this.text += this.decoder.decode(data, { stream: true });
    this.countLines(data);
    this.chunks.push(data);
    this.ensureTempFileIfNeeded();
  }

  finish(): void {
    this.finished = true;
    this.text += this.decoder.decode();
    this.ensureTempFileIfNeeded();
  }

  snapshot(): SnapshotResult {
    const truncation = truncateTail(this.text, {
      maxLines: DEFAULT_MAX_LINES,
      maxBytes: DEFAULT_MAX_BYTES,
    });

    const truncated =
      this.totalLines > DEFAULT_MAX_LINES || this.totalBytes > DEFAULT_MAX_BYTES;

    return {
      content: truncation.content,
      truncated,
      truncatedBy: truncation.truncated
        ? (truncation.truncatedBy ??
          (this.totalBytes > DEFAULT_MAX_BYTES ? "bytes" : "lines"))
        : null,
      totalLines: this.totalLines,
      totalBytes: this.totalBytes,
      fullOutputPath: this.tempFilePath,
    };
  }

  async close(): Promise<void> {
    if (!this.tempFileStream) return;
    const stream = this.tempFileStream;
    this.tempFileStream = undefined;
    await new Promise<void>((resolve, reject) => {
      stream.once("error", reject);
      stream.once("finish", resolve);
      stream.end();
    });
  }

  private countLines(data: Buffer): void {
    for (let i = 0; i < data.length; i++) {
      if (data[i] === 0x0a) this.totalLines++;
    }
  }

  private ensureTempFileIfNeeded(): void {
    if (this.tempFilePath || this.totalBytes <= DEFAULT_MAX_BYTES) return;
    this.tempFilePath = join(tmpdir(), `pi-bash-${randomBytes(8).toString("hex")}.log`);
    this.tempFileStream = createWriteStream(this.tempFilePath);
    for (const chunk of this.chunks) {
      this.tempFileStream.write(chunk);
    }
    this.chunks = [];
  }

  private writeTemp(data: Buffer): void {
    if (this.tempFileStream) {
      this.tempFileStream.write(data);
    }
  }
}

// ---------------------------------------------------------------------------
// Schema
// ---------------------------------------------------------------------------

const bashSchema = Type.Object({
  command: Type.String({ description: "Bash command to execute" }),
  timeout: Type.Optional(
    Type.Number({
      description: "Timeout in seconds (optional, no default timeout)",
    }),
  ),
  background: Type.Optional(
    Type.Boolean({
      description:
        "If true, run the command in the background. Returns immediately with PID and log file path. The timeout is still respected — the process is killed when it expires.",
    }),
  ),
});

// ---------------------------------------------------------------------------
// Tool definition
// ---------------------------------------------------------------------------

export default function (pi: ExtensionAPI) {
  const ops = createLocalBashOperations();

  pi.on("session_shutdown", () => killAllTracked());

  pi.registerTool({
    name: "bash",
    label: "bash",
    description: `Execute a bash command in the current working directory. Returns stdout and stderr. Output is truncated to last ${DEFAULT_MAX_LINES} lines or ${DEFAULT_MAX_BYTES / 1024}KB (whichever is hit first). If truncated, full output is saved to a temp file. Optionally provide a timeout in seconds.

When background is true, the command runs in the background. The tool returns immediately with the PID and log file path (/tmp/ppb/{id}.log). Use the PID to manage the process (kill, check status) and read the log file to inspect output. The timeout parameter is still respected — the process is automatically killed when it expires.`,
    promptSnippet: "Execute bash commands (ls, grep, find, etc.); supports background mode",
    parameters: bashSchema,
    async execute(_toolCallId, { command, timeout, background }, signal, onUpdate, ctx) {
      // Resolve the working directory per execution from the extension context.
      // ctx.cwd is the session's cwd, which is correct both in the pi TUI and
      // in multi-session hosts like pi-web (where process.cwd() is the server
      // launch directory, not the session's project).
      const cwd = ctx?.cwd ?? process.cwd();
      if (background) {
        return executeBackground(command, cwd, timeout);
      }
      return executeForeground(command, cwd, ops, timeout, signal, onUpdate);
    },
  });
}

// ---------------------------------------------------------------------------
// Foreground execution
// ---------------------------------------------------------------------------

async function executeForeground(
  command: string,
  cwd: string,
  ops: ReturnType<typeof createLocalBashOperations>,
  timeout: number | undefined,
  signal: AbortSignal | undefined,
  onUpdate: ((r: { content: Array<{ type: string; text?: string }>; details?: object }) => void) | undefined,
): Promise<{ content: Array<{ type: string; text: string }>; details?: object }> {
  const collector = new SimpleCollector();

  const handleData = (data: Buffer) => {
    collector.append(data);
    if (onUpdate) {
      const snap = collector.snapshot();
      onUpdate({ content: [{ type: "text", text: snap.content || "" }] });
    }
  };

  try {
    const result = await ops.exec(command, cwd, {
      onData: handleData,
      signal,
      timeout,
    });

    collector.finish();
    const snap = collector.snapshot();
    await collector.close();

    let text = snap.content || "(no output)";

    if (snap.truncated) {
      const startLine = snap.totalLines - DEFAULT_MAX_LINES + 1;
      if (snap.truncatedBy === "lines") {
        text += `\n\n[Showing lines ${startLine}-${snap.totalLines} of ${snap.totalLines}. Full output: ${snap.fullOutputPath}]`;
      } else {
        text += `\n\n[Showing ${DEFAULT_MAX_LINES} lines (${formatSize(DEFAULT_MAX_BYTES)} limit). Full output: ${snap.fullOutputPath}]`;
      }
    }

    if (result.exitCode !== 0 && result.exitCode !== null) {
      throw new Error(
        `${text ? `${text}\n\n` : ""}Command exited with code ${result.exitCode}`,
      );
    }

    const details = snap.truncated
      ? {
          truncation: {
            truncated: snap.truncated,
            truncatedBy: snap.truncatedBy,
            totalLines: snap.totalLines,
            outputLines: DEFAULT_MAX_LINES,
          },
          fullOutputPath: snap.fullOutputPath,
        }
      : undefined;

    return { content: [{ type: "text", text }], details };
  } catch (err) {
    collector.finish();
    await collector.close();

    const snap = collector.snapshot();
    let text = snap.content || "";

    if (err instanceof Error && err.message === "aborted") {
      throw new Error(`${text ? `${text}\n\n` : ""}Command aborted`);
    }
    if (err instanceof Error && err.message.startsWith("timeout:")) {
      const timeoutSecs = err.message.split(":")[1];
      throw new Error(
        `${text ? `${text}\n\n` : ""}Command timed out after ${timeoutSecs} seconds`,
      );
    }
    throw err;
  }
}

// ---------------------------------------------------------------------------
// Background execution
// ---------------------------------------------------------------------------

function executeBackground(
  command: string,
  cwd: string,
  timeout: number | undefined,
): { content: Array<{ type: string; text: string }>; details?: object } {
  ensureLogDir();
  const id = generateId();
  const filePath = logPath(id);
  const stream = createWriteStream(filePath, { flags: "w" });

  const shellConfig = getShellConfig();
  const commandFromStdin = shellConfig.commandTransport === "stdin";
  const child = spawn(
    shellConfig.shell,
    commandFromStdin ? shellConfig.args : [...shellConfig.args, command],
    {
      cwd,
      detached: process.platform !== "win32",
      env: process.env,
      stdio: [commandFromStdin ? "pipe" : "ignore", "pipe", "pipe"],
      windowsHide: true,
    },
  );

  if (commandFromStdin) {
    child.stdin?.on("error", () => {});
    child.stdin?.end(command);
  }

  const pid = child.pid;
  if (pid) trackedPids.add(pid);

  let timeoutHandle: ReturnType<typeof setTimeout> | undefined;
  if (timeout !== undefined) {
    timeoutHandle = setTimeout(() => {
      if (pid) killProcessTree(pid);
    }, timeout * 1000);
  }

  // Stream output to log file.
  child.stdout?.on("data", (chunk: Buffer) => stream.write(chunk));
  child.stderr?.on("data", (chunk: Buffer) => stream.write(chunk));

  // Cleanup on exit or error.
  const cleanup = () => {
    trackedPids.delete(pid!);
    if (timeoutHandle) clearTimeout(timeoutHandle);
    stream.end();
  };
  child.on("exit", cleanup);
  child.on("error", cleanup);

  const timeoutInfo = timeout ? `\nTimeout: ${timeout}s (auto-kill)` : "";

  return {
    content: [
      {
        type: "text",
        text: `Process started in background.\nPID: ${pid}\nLog: ${filePath}${timeoutInfo}`,
      },
    ],
  };
}