dotfiles/config/pi/extensions/worktree-agent/worker.ts

import { spawn } from "node:child_process";
import type { ProgressReporter, WorktreeAgentInput, WorktreeContext } from "./types.ts";

const MAX_CAPTURED_STDERR = 50 * 1024;

interface WorkerModel {
  provider: string;
  id: string;
}

interface RunWorkerOptions {
  input: WorktreeAgentInput;
  worktree: WorktreeContext;
  model?: WorkerModel;
  thinkingLevel: string;
  signal?: AbortSignal;
  report: ProgressReporter;
}

function assistantText(message: unknown): string {
  if (!message || typeof message !== "object") return "";
  const candidate = message as {
    role?: string;
    content?: Array<{ type?: string; text?: string }>;
  };
  if (candidate.role !== "assistant" || !Array.isArray(candidate.content)) return "";
  return candidate.content
    .filter((part) => part.type === "text" && typeof part.text === "string")
    .map((part) => part.text)
    .join("\n")
    .trim();
}

function appendCaptured(current: string, addition: string): string {
  const combined = current + addition;
  return combined.length <= MAX_CAPTURED_STDERR
    ? combined
    : combined.slice(combined.length - MAX_CAPTURED_STDERR);
}

function buildWorkerPrompt(input: WorktreeAgentInput, worktree: WorktreeContext): string {
  const criteria = input.acceptanceCriteria?.length
    ? input.acceptanceCriteria.map((criterion) => `- ${criterion}`).join("\n")
    : "- Derive acceptance criteria from the issue description and state any uncertainty.";

  return [
    "/skill:linear-issue-worker",
    "",
    `Issue ID: ${input.issueId.trim().toUpperCase()}`,
    `Title: ${input.title.trim()}`,
    input.issueUrl ? `URL: ${input.issueUrl}` : undefined,
    "",
    "Description:",
    input.description.trim() || "(No description supplied.)",
    "",
    "Acceptance criteria:",
    criteria,
    "",
    `Assigned branch: ${worktree.branchName}`,
    `Assigned worktree: ${worktree.worktreePath}`,
  ]
    .filter((line): line is string => line !== undefined)
    .join("\n");
}

export async function runWorker(options: RunWorkerOptions): Promise<string> {
  const { input, worktree, model, thinkingLevel, signal, report } = options;
  const sessionName = `${input.issueId.trim().toUpperCase()} worker`;
  const childArgs = ["--mode", "json", "-p", "--approve", "--name", sessionName];
  if (model) childArgs.push("--model", `${model.provider}/${model.id}`);
  childArgs.push("--thinking", thinkingLevel, buildWorkerPrompt(input, worktree));

  report(`Starting child Pi in ${worktree.worktreePath}`, {
    phase: "starting-worker",
    branchName: worktree.branchName,
    worktreePath: worktree.worktreePath,
  });

  let finalText = "";
  let stderr = "";
  let workerStopReason: string | undefined;
  let workerError: string | undefined;

  let exitCode: number;
  try {
    exitCode = await new Promise<number>((resolve, reject) => {
      const child = spawn("pi", childArgs, {
        cwd: worktree.worktreePath,
        env: {
          ...process.env,
          PI_WORKTREE_AGENT: "1",
          PI_WORKTREE_ROOT: worktree.worktreePath,
        },
        shell: false,
        stdio: ["ignore", "pipe", "pipe"],
      });

      let stdoutBuffer = "";
      let completed = false;
      let killTimer: ReturnType<typeof setTimeout> | undefined;

      const finish = (callback: () => void) => {
        if (completed) return;
        completed = true;
        if (killTimer) clearTimeout(killTimer);
        signal?.removeEventListener("abort", abortChild);
        callback();
      };

      const abortChild = () => {
        child.kill("SIGTERM");
        killTimer = setTimeout(() => child.kill("SIGKILL"), 5_000);
      };

      const processLine = (line: string) => {
        if (!line.trim()) return;
        let event: any;
        try {
          event = JSON.parse(line);
        } catch {
          return;
        }

        if (event.type === "message_end") {
          const text = assistantText(event.message);
          if (text) finalText = text;
          if (event.message?.role === "assistant") {
            workerStopReason = event.message.stopReason;
            workerError = event.message.errorMessage;
          }
        } else if (event.type === "tool_execution_start") {
          report(`Worker running: ${event.toolName}`, {
            phase: "worker-running",
            branchName: worktree.branchName,
            worktreePath: worktree.worktreePath,
            toolName: event.toolName,
          });
        }
      };

      child.stdout.on("data", (chunk) => {
        stdoutBuffer += chunk.toString();
        const lines = stdoutBuffer.split("\n");
        stdoutBuffer = lines.pop() ?? "";
        for (const line of lines) processLine(line.endsWith("\r") ? line.slice(0, -1) : line);
      });
      child.stderr.on("data", (chunk) => {
        stderr = appendCaptured(stderr, chunk.toString());
      });
      child.on("error", (error) => finish(() => reject(error)));
      child.on("close", (code) =>
        finish(() => {
          if (stdoutBuffer.trim()) processLine(stdoutBuffer);
          resolve(code ?? 1);
        }),
      );

      if (signal?.aborted) abortChild();
      else signal?.addEventListener("abort", abortChild, { once: true });
    });
  } catch (error) {
    const detail = error instanceof Error ? error.message : String(error);
    throw new Error(`Unable to run child Pi: ${detail}\nWorktree retained at ${worktree.worktreePath}`);
  }

  if (signal?.aborted) {
    throw new Error(`Worker aborted. Worktree retained at ${worktree.worktreePath}`);
  }
  if (exitCode !== 0 || workerStopReason === "error" || workerStopReason === "aborted") {
    const detail = workerError || stderr.trim() || finalText || `exit code ${exitCode}`;
    throw new Error(`Worker failed: ${detail}\nWorktree retained at ${worktree.worktreePath}`);
  }

  return finalText || "Worker exited successfully without a final text report.";
}