AgentHub/server/__tests__/server-mode-proxy-gateway.test.js

/**
 * @file server-mode-proxy-gateway.test.js
 * @description The AI Proxy key gateway (S3 plan, P1b). In server mode a pane
 * never gets a proxy's key: its CLI talks to `/api/proxy-gw/<name>/v1/…` on
 * loopback with its own pane token, and the gateway forwards to the profile's
 * endpoint with the real key, streaming both ways. Refused: local mode, a
 * non-loopback request, an unknown/dead pane token, a profile the pane was not
 * started on, a path outside `/v1/`. Also what a spawn writes in server mode:
 * no key in any file, the gateway as the endpoint, the pane token as the key.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const { describe, it, before, after } = require("node:test");
const assert = require("node:assert/strict");
const crypto = require("node:crypto");
const http = require("node:http");
const path = require("node:path");
const fs = require("node:fs");
const os = require("node:os");
const { execFileSync } = require("node:child_process");

const TMP = fs.mkdtempSync(path.join(os.tmpdir(), "server-mode-proxy-gw-"));
process.env.DASHBOARD_DB_PATH = path.join(TMP, "dashboard.db");
process.env.DASHBOARD_LIVENESS_PROBE = "0";
process.env.LANES_ROOT = path.join(TMP, "lanes");
process.env.CLAUDE_HOME = path.join(TMP, ".claude");
process.env.CODEX_HOME = path.join(TMP, ".codex");
process.env.DASHBOARD_DATA_DIR = path.join(TMP, "data");

const { createApp, startServer } = require("../index");
const auth = require("../lib/auth");
const proxies = require("../lib/ai-proxies");
const { proxyForSpawn, resolveProxy } = require("../lib/proxy-apply");
const { db } = require("../db");

const KEY = "sk-upstream-real-key";
const sha256 = (s) => crypto.createHash("sha256").update(s).digest("hex");

let server;
let BASE;
let upstream;
let seen = []; // what the fake upstream received
const LIVE_TOKEN = "ahpane_live-token-for-tests-0001";
const DEAD_TOKEN = "ahpane_dead-token-for-tests-0002";
const LIVE_RUN = "agenthub-lane-900";
const DEAD_RUN = "agenthub-lane-901";
let profileName;
let otherName;

function gw(p, { token = LIVE_TOKEN, headers = {}, body, method = "POST" } = {}) {
  return fetch(`${BASE}/api/proxy-gw/${p}`, {
    method,
    headers: {
      "content-type": "application/json",
      ...(token ? { authorization: `Bearer ${token}` } : {}),
      ...headers,
    },
    body,
  });
}

before(async () => {
  upstream = http.createServer((req, res) => {
    let raw = "";
    req.on("data", (c) => (raw += c));
    req.on("end", () => {
      seen.push({ url: req.url, headers: req.headers, body: raw });
      if (req.url.endsWith("/stream")) {
        res.writeHead(200, { "content-type": "text/event-stream" });
        res.write("event: a\ndata: 1\n\n");
        setTimeout(() => res.end("event: b\ndata: 2\n\n"), 30);
        return;
      }
      res.writeHead(201, { "content-type": "application/json", "x-upstream": "yes" });
      res.end(JSON.stringify({ ok: true }));
    });
  });
  await new Promise((r) => upstream.listen(0, "127.0.0.1", r));
  const upBase = `http://127.0.0.1:${upstream.address().port}/root`;

  server = await startServer(createApp(), 0);
  BASE = `http://127.0.0.1:${server.address().port}`;
  const mia = await auth.createUser("mia", "member-pass-1", "developer");
  profileName = proxies.createProxy({
    label: "Team GW",
    baseUrl: upBase,
    apiKey: KEY,
    model: "m",
  }).name;
  otherName = proxies.createProxy({
    label: "Other GW",
    baseUrl: upBase,
    apiKey: KEY,
    model: "m",
  }).name;
  const laneId = db
    .prepare("INSERT INTO lanes (cwd, tool, owner_id) VALUES ('/w/gw', 'claude', ?)")
    .run(mia.id).lastInsertRowid;
  const run = db.prepare(
    `INSERT INTO dashboard_runs (id, lane_id, cwd, mode, status, proxy, pane_token_hash)
     VALUES (?, ?, '/w/gw', 'interactive', 'running', ?, ?)`
  );
  run.run(LIVE_RUN, laneId, profileName, sha256(LIVE_TOKEN));
  run.run(DEAD_RUN, laneId, profileName, sha256(DEAD_TOKEN));
  // A live pane for LIVE_RUN only, on the suite's private tmux socket.
  execFileSync("tmux", ["new-session", "-d", "-s", LIVE_RUN, "sleep 300"]);
  process.env.AGENTHUB_MODE = "server";
});

after(() => {
  delete process.env.AGENTHUB_MODE;
  try {
    execFileSync("tmux", ["kill-session", "-t", LIVE_RUN]);
  } catch {
    /* already gone */
  }
  server?.close();
  upstream?.close();
  fs.rmSync(TMP, { recursive: true, force: true });
});

describe("the proxy gateway", () => {
  it("forwards with the real key, never the pane token, body and status intact", async () => {
    seen = [];
    const res = await gw(`${profileName}/v1/messages?beta=true`, {
      headers: { "anthropic-version": "2023-06-01", "x-api-key": LIVE_TOKEN },
      body: JSON.stringify({ hello: "world" }),
    });
    assert.equal(res.status, 201);
    assert.equal(res.headers.get("x-upstream"), "yes");
    assert.deepEqual(await res.json(), { ok: true });
    assert.equal(seen.length, 1);
    const got = seen[0];
    assert.equal(got.url, "/root/v1/messages?beta=true");
    assert.equal(got.headers.authorization, `Bearer ${KEY}`);
    assert.equal(got.headers["x-api-key"], KEY);
    assert.equal(got.headers["anthropic-version"], "2023-06-01");
    assert.ok(!JSON.stringify(got.headers).includes("ahpane_"));
    assert.equal(got.body, JSON.stringify({ hello: "world" }));
  });

  it("streams an SSE answer through", async () => {
    const res = await gw(`${profileName}/v1/stream`, { body: "{}" });
    assert.equal(res.status, 200);
    assert.equal(await res.text(), "event: a\ndata: 1\n\nevent: b\ndata: 2\n\n");
  });

  it("refuses a missing, unknown or dead pane token", async () => {
    for (const token of [null, "ahpane_nope", DEAD_TOKEN]) {
      seen = [];
      const res = await gw(`${profileName}/v1/messages`, { token, body: "{}" });
      assert.equal(res.status, 401, String(token));
      assert.equal(seen.length, 0);
    }
  });

  it("refuses a profile the pane was not started on", async () => {
    seen = [];
    const res = await gw(`${otherName}/v1/messages`, { body: "{}" });
    assert.equal(res.status, 403);
    assert.equal(seen.length, 0);
  });

  it("refuses a path that leaves /v1/", async () => {
    seen = [];
    for (const p of [`${profileName}/v1/../../admin`, `${profileName}/v2/x`, `${profileName}`]) {
      const res = await gw(p, { body: "{}" });
      assert.equal(res.status, 400, p);
    }
    assert.equal(seen.length, 0);
  });

  it("refuses a raw `..` path the client did not normalise", async () => {
    seen = [];
    for (const raw of [`${profileName}/v1/../../admin`, `${profileName}/v1/%2e%2e/%2e%2e/admin`]) {
      const status = await new Promise((resolve, reject) => {
        const req = http.request(
          `${BASE}/api/proxy-gw/x`,
          {
            method: "POST",
            path: `/api/proxy-gw/${raw}`,
            headers: { authorization: `Bearer ${LIVE_TOKEN}`, "content-type": "application/json" },
          },
          (res) => {
            res.resume();
            resolve(res.statusCode);
          }
        );
        req.on("error", reject);
        req.end("{}");
      });
      assert.equal(status, 400, raw);
    }
    assert.equal(seen.length, 0);
  });

  it("refuses a request that came through a reverse proxy", async () => {
    const res = await gw(`${profileName}/v1/messages`, {
      headers: { "x-forwarded-for": "10.0.0.9" },
      body: "{}",
    });
    assert.equal(res.status, 403);
  });

  it("is not there in local mode", async () => {
    delete process.env.AGENTHUB_MODE;
    try {
      const res = await gw(`${profileName}/v1/messages`, { body: "{}" });
      assert.equal(res.status, 404);
    } finally {
      process.env.AGENTHUB_MODE = "server";
    }
  });
});

describe("a spawn on a proxy in server mode", () => {
  it("puts no key in Claude's settings file — the pane's own, with its pane token", () => {
    const spawn = proxyForSpawn(
      "claude",
      resolveProxy("claude", profileName),
      LIVE_TOKEN,
      null,
      "agenthub-lane-3"
    );
    const file = spawn.argv[0].replace(/^--settings=/, "");
    assert.equal(path.basename(file), `${profileName}--agenthub-lane-3.json`);
    assert.equal(fs.statSync(file).mode & 0o777, 0o600);
    const text = fs.readFileSync(file, "utf8");
    assert.ok(!text.includes(KEY));
    const env = JSON.parse(text).env;
    assert.match(env.ANTHROPIC_BASE_URL, new RegExp(`/api/proxy-gw/${profileName}$`));
    // The command-line layer: a user's or project's env.ANTHROPIC_* cannot override it.
    assert.equal(env.ANTHROPIC_AUTH_TOKEN, LIVE_TOKEN);
    assert.deepEqual(spawn.env, { ANTHROPIC_AUTH_TOKEN: LIVE_TOKEN });
  });

  it("points Codex's layer at the gateway, keyed by the pane token", () => {
    const spawn = proxyForSpawn("codex", resolveProxy("codex", profileName), LIVE_TOKEN);
    const layer = fs.readFileSync(
      path.join(process.env.CODEX_HOME, `${profileName}.config.toml`),
      "utf8"
    );
    assert.ok(!layer.includes(KEY));
    assert.match(layer, new RegExp(`base_url = ".*/api/proxy-gw/${profileName}/v1"`));
    assert.match(layer, /env_key = "AGENTHUB_PROXY_KEY"/);
    assert.deepEqual(spawn.env, { AGENTHUB_PROXY_KEY: LIVE_TOKEN });
  });
});