Agent-System/deploy/pi-agent/runner.mjs

#!/usr/bin/env node
/**
 * Pi Agent Container HTTP Runner Server
 * Listens for task run requests, executes pi agent in RPC mode within workspace,
 * and returns generated artifact patches.
 */

import http from "node:http";
import { spawn } from "node:child_process";
import { readFileSync, writeFileSync, existsSync, mkdirSync } from "node:fs";
import { resolve, relative, isAbsolute, dirname } from "node:path";
import { TavilyKeyPool, searchTavily, formatTavilyResults } from "./extensions/tavily-search.mjs";
import { applyAgentEvent, classifyRunOutcome, newSettleState } from "./settle-policy.mjs";
import {
  withResultV1,
  capPatches,
  truncateExecutionLogContent,
  estimateWireBytes,
  MAX_RESULT_WIRE_BYTES,
} from "./result-v1.mjs";
import { buildPiLaunchArgs, TOOL_PROFILE_REASONING } from "./pi-args.mjs";

/** Keep in sync with reasoning-materialize.mjs (unit-tested; inlined so bind-mounted runner needs no extra mount). */
function stripThinkBlocks(text) {
  return String(text || "")
    .replace(/<think>[\s\S]*?<\/think>/gi, "")
    .replace(/```(?:json)?\s*/gi, "")
    .replace(/```/g, "")
    .trim();
}

function extractAssistantText(ev) {
  if (!ev || typeof ev !== "object") {
    return "";
  }
  const t = ev.type || "";
  if (t !== "message_end" && t !== "turn_end") {
    return "";
  }
  const msg = ev.message;
  if (!msg || typeof msg !== "object" || msg.role !== "assistant") {
    return "";
  }
  const parts = Array.isArray(msg.content) ? msg.content : [];
  const chunks = [];
  for (const part of parts) {
    if (part && typeof part === "object" && part.type === "text" && typeof part.text === "string") {
      chunks.push(part.text);
    }
  }
  return chunks.join("");
}

function materializeReasoningTargets({
  workspace,
  targets,
  assistantText,
  existsSync: exists,
  readFileSync: read,
  writeFileSync: write,
  mkdirSync: mkdir,
  resolve: pathResolve,
  dirname: pathDirname,
}) {
  const cleaned = stripThinkBlocks(assistantText);
  if (!cleaned || !workspace) {
    return [];
  }
  const written = [];
  for (const rel of Array.isArray(targets) ? targets : []) {
    if (!rel || typeof rel !== "string") {
      continue;
    }
    const abs = pathResolve(workspace, rel);
    if (exists(abs)) {
      try {
        if (String(read(abs, "utf8")).trim()) {
          continue;
        }
      } catch {
        // fall through
      }
    }
    try {
      mkdir(pathDirname(abs), { recursive: true });
      write(abs, cleaned, "utf8");
      written.push(rel);
    } catch {
      // leave missing
    }
  }
  return written;
}

const PORT = parseInt(process.env.PORT || process.env.PI_AGENT_RUNNER_PORT || "18766", 10);
const DATA_DIR = process.env.DATA_DIR || "/app/data/workspaces";
const RUNNER_AUTH_TOKEN = (process.env.RUNNER_AUTH_TOKEN || process.env.PI_AGENT_RUNNER_AUTH_TOKEN || "").trim();
const DEFAULT_MODEL = process.env.MINIMAX_MODEL_ID || "MiniMax-M3";
const DEFAULT_PROVIDER = process.env.MINIMAX_PROVIDER || "minimax";
const tavilyPool = TavilyKeyPool.fromEnv();
/** @type {Map<string, import('node:child_process').ChildProcess>} */
const activeChildren = new Map();

function safeRel(workspace, target) {
  const abs = isAbsolute(target) ? target : resolve(workspace, target);
  const rel = relative(workspace, abs);
  if (!rel || rel.startsWith("..") || isAbsolute(rel)) {
    return null;
  }
  return rel.replace(/\\/g, "/");
}

function resolveWorkspace(rawPath) {
  if (!rawPath) return DATA_DIR;

  // If path contains 'workspaces/', extract the relative part
  const wsIndex = rawPath.indexOf("workspaces/");
  let sub = "";
  if (wsIndex !== -1) {
    sub = rawPath.slice(wsIndex + "workspaces/".length);
  } else if (rawPath.startsWith(DATA_DIR)) {
    sub = relative(DATA_DIR, rawPath);
  } else {
    sub = rawPath.replace(/^[/\\]+/, "");
  }

  // Force all execution strictly inside DATA_DIR (which is mounted to host's controlled directory)
  const target = resolve(DATA_DIR, sub);
  const rel = relative(DATA_DIR, target);
  if (rel.startsWith("..") || isAbsolute(rel)) {
    return resolve(DATA_DIR, "sandboxes", "isolated_" + Date.now());
  }
  return target;
}

function parseJsonBody(req) {
  return new Promise((resolvePromise, rejectPromise) => {
    const chunks = [];
    req.on("data", (chunk) => chunks.push(chunk));
    req.on("end", () => {
      const raw = Buffer.concat(chunks).toString("utf8");
      if (!raw.trim()) {
        resolvePromise({});
        return;
      }
      try {
        resolvePromise(JSON.parse(raw));
      } catch (err) {
        rejectPromise(new Error("Invalid JSON: " + err.message));
      }
    });
    req.on("error", rejectPromise);
  });
}

function sendJson(res, statusCode, data) {
  // Compact JSON: pretty-print + oversized patches previously ballooned past Go's 8MB fail-closed.
  const body = JSON.stringify(data) + "\n";
  if (Buffer.byteLength(body, "utf8") > MAX_RESULT_WIRE_BYTES + 512 * 1024) {
    console.warn(
      `[runner] response body ${Buffer.byteLength(body, "utf8")} bytes exceeds soft wire budget`,
    );
  }
  res.writeHead(statusCode, { "Content-Type": "application/json; charset=utf-8" });
  res.end(body);
}

function authorizeRunner(req, res) {
  if (!RUNNER_AUTH_TOKEN) {
    return true;
  }
  const auth = req.headers.authorization || "";
  if (auth === `Bearer ${RUNNER_AUTH_TOKEN}`) {
    return true;
  }
  sendJson(res, 401, { ok: false, error: "unauthorized" });
  return false;
}

async function handleRun(body) {
  const rawWs = body.workspace || "";
  const workspace = resolveWorkspace(rawWs);
  const model = body.model || DEFAULT_MODEL;
  const provider = body.provider || DEFAULT_PROVIDER;
  const prompt = body.prompt || "";
  const instructions = body.instructions || "";
  const targets = Array.isArray(body.targets) ? body.targets : [];
  const timeoutMs = Number(body.timeout_ms) || 30 * 60 * 1000;
  const toolProfile = String(body.tool_profile || "").trim();

  if (!existsSync(workspace)) {
    mkdirSync(workspace, { recursive: true });
  }

  if (!prompt.trim()) {
    return { ok: false, run_id: String(body.run_id || "unknown").slice(0, 128), fallback_reason: "prompt_empty", patches: [] };
  }

  const allowed = new Set();
  for (const t of targets) {
    const rel = safeRel(workspace, t);
    if (rel) allowed.add(rel);
  }

  // 1. Write persona instructions and task into workspace
  if (instructions.trim()) {
    try {
      writeFileSync(resolve(workspace, "AGENTS.md"), instructions, "utf8");
      writeFileSync(resolve(workspace, "INSTRUCTIONS.md"), instructions, "utf8");
    } catch (e) {
      console.warn("[runner] Failed to write instructions to workspace:", e);
    }
  }

  if (prompt.trim()) {
    try {
      writeFileSync(resolve(workspace, "TASK.md"), prompt, "utf8");
    } catch (e) {
      console.warn("[runner] Failed to write TASK.md:", e);
    }
  }

  // 2. Build pi arguments. Unknown profiles fail closed before spawn.
  const piBin = process.env.PI_AGENT_BIN || "pi";
  const sysPromptFile = resolve(workspace, "INSTRUCTIONS.md");
  let args;
  try {
    args = buildPiLaunchArgs({
      toolProfile: body.tool_profile,
      provider,
      model,
      systemPromptFile: existsSync(sysPromptFile) ? "INSTRUCTIONS.md" : "",
      extension: "/app/extensions/tavily-search.mjs",
      workspaceTools: "read,write,edit,tavily_search,web_search",
    });
  } catch (e) {
    if (e && e.code === "UNKNOWN_TOOL_PROFILE") {
      return {
        ok: false,
        run_id: String(body.run_id || "unknown").slice(0, 128),
        fallback_reason: "unknown_tool_profile",
        patches: [],
      };
    }
    throw e;
  }

  const env = { ...process.env };
  env.PI_AGENT_MODEL = model;
  env.PI_AGENT_PROVIDER = provider;
  env.PI_AGENT_SYSTEM_PROMPT_FILE = "AGENTS.md";

  let agentId = "";
  let runId = String(body.run_id || "").slice(0, 128);
  let lastErr = "";
  let lastAssistantText = "";
  const settle = newSettleState();

  console.log(`[runner] Starting pi run in workspace: ${workspace}, model: ${model}`);

  const child = spawn(piBin, args, {
    cwd: workspace,
    env,
    stdio: ["pipe", "pipe", "pipe"],
  });
  activeChildren.set(workspace, child);

  const timer = setTimeout(() => {
    lastErr = `Execution timed out after ${timeoutMs}ms`;
    child.kill("SIGKILL");
  }, timeoutMs);

  const cmdId = "req-" + Date.now();
  child.stdin.write(
    JSON.stringify({
      id: cmdId,
      type: "prompt",
      message: instructions ? `[Persona & System Instructions]:\n${instructions}\n\n[User Task]:\n${prompt}` : prompt,
    }) + "\n",
  );

  let stdoutBuf = "";
  let fullStdout = "";
  child.stdout.on("data", (chunk) => {
    const text = chunk.toString("utf8");
    stdoutBuf += text;
    fullStdout += text;
    let nl;
    while ((nl = stdoutBuf.indexOf("\n")) >= 0) {
      const line = stdoutBuf.slice(0, nl).trim();
      stdoutBuf = stdoutBuf.slice(nl + 1);
      if (!line) continue;
      let ev;
      try { ev = JSON.parse(line); } catch { continue; }
      if (!ev || typeof ev !== "object") continue;

      const t = ev.type || "";
      if (t === "session" && ev.id) agentId = ev.id;
      const assistantChunk = extractAssistantText(ev);
      if (assistantChunk) {
        lastAssistantText = assistantChunk;
      }
      applyAgentEvent(settle, ev);
      if (t === "agent_settled") {
        child.stdin.end();
      }
    }
  });

  let stderrBuf = "";
  child.stderr.on("data", (chunk) => {
    stderrBuf += chunk.toString("utf8");
    lastErr = (stderrBuf + " " + lastErr).slice(-400);
  });

  await new Promise((resolveDone) => {
    child.on("close", () => {
      clearTimeout(timer);
      activeChildren.delete(workspace);
      resolveDone();
    });
    child.on("error", (err) => {
      clearTimeout(timer);
      activeChildren.delete(workspace);
      lastErr = err && err.message ? err.message : String(err);
      resolveDone();
    });
  });

  // Ensure all intermediate traces, logs and thinking process land on the host-mounted workspace
  try {
    const trace = [
      `=== PI-AGENT TASK EXECUTION LOG ===`,
      `Timestamp: ${new Date().toISOString()}`,
      `Workspace: ${workspace}`,
      `Model: ${model} (${provider})`,
      `Outcome: ${settle.succeeded ? "SUCCESS" : "FAILED"}`,
      `Error: ${settle.lastErr || lastErr || "None"}`,
      `\n--- RAW OUTPUT & EVENTS ---\n${fullStdout}`,
      `\n--- STDERR ---\n${stderrBuf}`,
    ].join("\n");
    writeFileSync(resolve(workspace, "execution.log"), trace, "utf8");
  } catch (logErr) {
    console.warn("[runner] Could not write execution.log:", logErr.message);
  }

  lastErr = settle.lastErr || lastErr;
  if (!settle.succeeded) {
    const classified = classifyRunOutcome({
      ...settle,
      lastErr,
      missing: [],
    });
    return {
      ok: false,
      agent_id: agentId,
      run_id: runId,
      retryable: classified.retryable,
      fallback_reason: classified.fallback_reason,
      patches: [],
    };
  }

  // Reasoning / --no-tools: materialize assistant text into missing targets (parity with pi-native).
  if (toolProfile === TOOL_PROFILE_REASONING && lastAssistantText) {
    materializeReasoningTargets({
      workspace,
      targets: [...allowed],
      assistantText: lastAssistantText,
      existsSync,
      readFileSync,
      writeFileSync,
      mkdirSync,
      resolve,
      dirname,
    });
  }

  // 3. Read generated files and deliverables
  const patches = [];
  const missing = [];
  for (const rel of allowed) {
    const abs = resolve(workspace, rel);
    if (!existsSync(abs)) {
      missing.push(rel);
      continue;
    }
    const content = readFileSync(abs, "utf8");
    if (!String(content).trim()) {
      missing.push(rel);
      continue;
    }
    patches.push({ rel_path: rel, content });
  }

  // execution.log full text stays on the shared volume; only a capped tail goes on the wire.
  const execLogPath = resolve(workspace, "execution.log");
  if (existsSync(execLogPath)) {
    try {
      const logContent = readFileSync(execLogPath, "utf8");
      patches.push({
        rel_path: "execution.log",
        content: truncateExecutionLogContent(logContent),
      });
    } catch {}
  }

  const cappedPatches = capPatches(patches);

  if (missing.length > 0) {
    const classified = classifyRunOutcome({
      ...settle,
      lastErr,
      missing,
    });
    return {
      ok: false,
      agent_id: agentId,
      run_id: runId,
      retryable: classified.retryable,
      fallback_reason: classified.fallback_reason,
      patches: cappedPatches,
    };
  }

  return {
    ok: true,
    agent_id: agentId,
    run_id: runId,
    patches: cappedPatches,
    fallback_reason: "",
  };
}

const server = http.createServer(async (req, res) => {
  const url = new URL(req.url, `http://${req.headers.host || "localhost"}`);
  const method = req.method.toUpperCase();

  // 1. Healthcheck
  if (method === "GET" && (url.pathname === "/health" || url.pathname === "/")) {
    sendJson(res, 200, {
      ok: true,
      service: "shuimu-pi-agent-runner",
      endpoint_url: process.env.MINIMAX_ENDPOINT_URL || "https://api.minimaxi.com/v1",
      model_id: DEFAULT_MODEL,
      api: process.env.MINIMAX_API_TYPE || "openai-completions",
      has_token: Boolean(process.env.MINIMAX_API_KEY || process.env.MINIMAX_API_TOKEN),
      tavily: {
        active: tavilyPool.size() > 0,
        keys_count: tavilyPool.size(),
      },
      data_dir: DATA_DIR,
      uptime: process.uptime(),
    });
    return;
  }

  // 2. Direct search test endpoint
  if (method === "POST" && url.pathname === "/v1/search") {
    try {
      const body = await parseJsonBody(req);
      const query = (body.query || "").trim();
      if (!query) {
        sendJson(res, 400, { error: "Missing query parameter" });
        return;
      }
      const data = await tavilyPool.executeWithPool(async (apiKey) => {
        return await searchTavily(apiKey, query, {
          search_depth: body.search_depth,
          max_results: body.max_results,
        });
      });
      const markdown = formatTavilyResults(data, query);
      sendJson(res, 200, { ok: true, data, markdown });
    } catch (err) {
      sendJson(res, 500, { ok: false, error: err.message });
    }
    return;
  }

  // 3. Cancel active child process by workspace key
  if (method === "POST" && url.pathname === "/v1/cancel") {
    if (!authorizeRunner(req, res)) {
      return;
    }
    try {
      const body = await parseJsonBody(req);
      const workspace = resolveWorkspace(body.workspace || "");
      const child = activeChildren.get(workspace);
      if (child) {
        child.kill("SIGKILL");
        activeChildren.delete(workspace);
        sendJson(res, 200, { ok: true, cancelled: true, workspace });
      } else {
        sendJson(res, 200, { ok: true, cancelled: false, workspace });
      }
    } catch (err) {
      sendJson(res, 500, { ok: false, error: err.message });
    }
    return;
  }

  // 4. Agent Task Run Endpoint
  if (method === "POST" && url.pathname === "/v1/run") {
    if (!authorizeRunner(req, res)) {
      return;
    }
    try {
      const body = await parseJsonBody(req);
      const result = await handleRun(body);
      const typed = withResultV1(result);
      const wireBytes = estimateWireBytes(typed);
      if (wireBytes > MAX_RESULT_WIRE_BYTES) {
        console.warn(`[runner] typed result wire ${wireBytes} bytes; dropping non-target patches`);
        typed.patches = capPatches(
          (typed.patches || []).filter((p) => p.rel_path !== "execution.log"),
        );
        const shrunk = withResultV1({ ...typed, patches: typed.patches });
        sendJson(res, 200, shrunk);
        return;
      }
      sendJson(res, 200, typed);
    } catch (err) {
      sendJson(res, 500, withResultV1({ ok: false, error: "runner_failed", patches: [], run_id: "unknown" }));
    }
    return;
  }

  sendJson(res, 404, { error: "Not Found", path: url.pathname });
});

server.listen(PORT, "0.0.0.0", () => {
  console.log(`[pi-agent-runner] Listening on 0.0.0.0:${PORT}`);
  console.log(`[pi-agent-runner] Endpoint: ${process.env.MINIMAX_ENDPOINT_URL || "https://api.minimaxi.com/v1"}`);
  console.log(`[pi-agent-runner] Model: ${DEFAULT_MODEL}, Provider: ${DEFAULT_PROVIDER}`);
  console.log(`[pi-agent-runner] Tavily API keys pool size: ${tavilyPool.size()}`);
});