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 } from "node:path";
import { TavilyKeyPool, searchTavily, formatTavilyResults } from "./extensions/tavily-search.mjs";
import { applyAgentEvent, classifyRunOutcome, newSettleState } from "./settle-policy.mjs";
import { withResultV1 } from "./result-v1.mjs";

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) {
  res.writeHead(statusCode, { "Content-Type": "application/json; charset=utf-8" });
  res.end(JSON.stringify(data, null, 2) + "\n");
}

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;

  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
  const piBin = process.env.PI_AGENT_BIN || "pi";
  const args = [
    "--mode", "rpc",
    "-e", "/app/extensions/tavily-search.mjs",
    "--no-session",
    "--provider", provider,
    "--model", model,
  ];

  const sysPromptFile = resolve(workspace, "INSTRUCTIONS.md");
  if (existsSync(sysPromptFile)) {
    args.push("--append-system-prompt", "INSTRUCTIONS.md");
  }

  args.push(
    "--tools", "read,write,edit,tavily_search,web_search",
    "--approve",
  );

  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 = "";
  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;
      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: [],
    };
  }

  // 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 });
  }

  // Also include execution.log as patch if present so host definitely receives it
  const execLogPath = resolve(workspace, "execution.log");
  if (existsSync(execLogPath)) {
    try {
      const logContent = readFileSync(execLogPath, "utf8");
      patches.push({ rel_path: "execution.log", content: logContent });
    } catch {}
  }

  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,
    };
  }

  return {
    ok: true,
    agent_id: agentId,
    run_id: runId,
    patches,
    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);
      sendJson(res, 200, withResultV1(result));
    } 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()}`);
});