AgentHub/server/hub/remote-gateway.js

/**
 * @file remote-gateway.js
 * @description The hub's gateway for remote clients' company-account sessions
 * (server mode only; docs/plans/remote-clients/README.md). A client's Claude
 * Code runs with `ANTHROPIC_BASE_URL = <hub>/api/remote/gw` and its device
 * token as `ANTHROPIC_AUTH_TOKEN`; this handler checks the token, checks that
 * the request's `x-claude-code-session-id` is a session THAT device
 * registered, swaps in the pool account's OAuth token and relays the request
 * to Anthropic, streaming the answer back untouched.
 *
 * The pool credential never leaves the hub: the client sends a device token
 * and gets back only Anthropic's answer. An unknown session id is a device
 * token used outside AgentHub — `AGENTHUB_REMOTE_GATE` decides: `warn` (the
 * default) relays it and records the session as unmonitored, `block` answers
 * 403. A message with no session id at all is always refused (Claude Code
 * always sends one), and so is a user over their spend quota (EQUOTA, the
 * same rule `spawnRun` applies). Each answer's `usage` is counted onto the
 * session as it streams past.
 *
 * Only `/v1/…` is relayed, to `AGENTHUB_REMOTE_UPSTREAM` (default
 * https://api.anthropic.com — set only from the server's own environment,
 * never from a request).
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const http = require("node:http");
const https = require("node:https");
const zlib = require("node:zlib");

const devices = require("./devices");
const sessions = require("./remote-sessions");
const { GATEWAY_MARK } = require("../lib/proxy-gateway");

const IDLE_MS = 10 * 60 * 1000;
/** A request may race its session's registration by this long. */
const REGISTER_GRACE_MS = Number(process.env.AGENTHUB_REMOTE_REGISTER_GRACE_MS) || 5000;
const OAUTH_BETA = "oauth-2025-04-20";
/** How much of a non-streamed answer is kept to read its `usage`. */
const MAX_TAP = 2 * 1024 * 1024;

const DROP_REQUEST = new Set([
  GATEWAY_MARK,
  "host",
  "connection",
  "keep-alive",
  "proxy-authorization",
  "proxy-connection",
  "te",
  "trailer",
  "transfer-encoding",
  "upgrade",
  "authorization",
  "x-api-key",
  "cookie",
  "x-agenthub-pane-token",
  "x-forwarded-for",
  "x-forwarded-host",
  "x-forwarded-proto",
  "forwarded",
]);
const DROP_RESPONSE = new Set(["connection", "keep-alive", "transfer-encoding"]);

const fail = (res, status, code, message) =>
  res.status(status).json({ type: "error", error: { type: code, code, message } });

const gate = () => (process.env.AGENTHUB_REMOTE_GATE === "block" ? "block" : "warn");
const upstream = () => new URL(process.env.AGENTHUB_REMOTE_UPSTREAM || "https://api.anthropic.com");
const sleep = (ms) => new Promise((r) => setTimeout(r, ms));

// A user's quota status, re-read at most every QUOTA_TTL_MS: computing spend
// on every relayed request would be an aggregate per request.
const QUOTA_TTL_MS = 60_000;
const quotaSeen = new Map();
function overQuota(userId) {
  const hit = quotaSeen.get(userId);
  if (hit && Date.now() - hit.at < QUOTA_TTL_MS) return hit.error;
  let error = null;
  try {
    require("../lib/user-quota").assertSpawnAllowed(userId);
  } catch (e) {
    if (e.code === "EQUOTA") error = e.message;
  }
  quotaSeen.set(userId, { at: Date.now(), error });
  return error;
}

/** The registered session, waiting out a registration still on its way. */
async function waitForSession(sessionId, deviceId) {
  const until = Date.now() + REGISTER_GRACE_MS;
  for (;;) {
    const row = sessions.sessionOfDevice(sessionId, deviceId);
    if (row || Date.now() >= until) return row;
    await sleep(200);
  }
}

/**
 * Reads `usage` out of an answer as it streams past, without holding it up:
 * SSE line by line (`message_start` carries input + cache, `message_delta`
 * the output), or a JSON body read whole up to `MAX_TAP`.
 */
function usageTap(contentType) {
  const usage = {};
  const add = (u) => {
    if (!u || typeof u !== "object") return;
    for (const [k, v] of Object.entries(u)) {
      if (typeof v === "number") usage[k] = v; // later events carry the running total
    }
  };
  const sse = /text\/event-stream/i.test(contentType || "");
  let carry = "";
  let size = 0;
  const chunks = [];
  return {
    write(buf) {
      if (sse) {
        const text = carry + buf.toString("utf8");
        const lines = text.split("\n");
        carry = lines.pop() || "";
        if (carry.length > 1024 * 1024) carry = "";
        for (const line of lines) {
          if (!line.startsWith("data:") || !line.includes('"usage"')) continue;
          try {
            const ev = JSON.parse(line.slice(5));
            add(ev.message?.usage);
            add(ev.usage);
          } catch {
            /* not JSON: skip */
          }
        }
      } else if (size < MAX_TAP) {
        chunks.push(buf);
        size += buf.length;
      }
    },
    end() {
      if (!sse && size > 0 && size <= MAX_TAP) {
        try {
          add(JSON.parse(Buffer.concat(chunks).toString("utf8")).usage);
        } catch {
          /* not JSON */
        }
      }
      return usage;
    },
  };
}

/** A decompressor for a `content-encoding`, or null for an uncompressed answer. */
function decoderFor(encoding) {
  const e = String(encoding || "")
    .trim()
    .toLowerCase();
  if (e === "gzip" || e === "x-gzip") return zlib.createGunzip();
  if (e === "br") return zlib.createBrotliDecompress();
  if (e === "deflate") return zlib.createInflate();
  if (e === "zstd" && typeof zlib.createZstdDecompress === "function") {
    return zlib.createZstdDecompress();
  }
  return null;
}

/** Express handler, mounted at `/api/remote/gw` BEFORE the body parsers. */
async function remoteGateway(req, res) {
  const who = devices.deviceFromRequest(req);
  if (!who) return fail(res, 401, "EAUTH", "not a valid AgentHub device token");
  const m = /^(\/v1\/[^?]*)(\?.*)?$/.exec(req.url);
  if (!m) return fail(res, 400, "EBADPATH", "expected /api/remote/gw/v1/…");

  const deviceRow = { id: who.device.id, userId: who.user.id };
  const quota = overQuota(who.user.id);
  if (quota) return fail(res, 429, "EQUOTA", quota);
  const sid = String(req.headers["x-claude-code-session-id"] || "").toLowerCase();
  let session = null;
  if (sid) {
    session = await waitForSession(sid, deviceRow.id);
    if (!session) {
      if (gate() === "block") {
        return fail(
          res,
          403,
          "EUNREGISTERED",
          "this session was not started from AgentHub; company accounts are used only through AgentHub"
        );
      }
      session = sessions.recordUnmonitored(deviceRow, sid);
      if (!session) return fail(res, 409, "EFOREIGN", "this session belongs to another device");
      console.warn(
        `[remote-gw] unmonitored session ${sid} from device ${deviceRow.id} (user ${who.user.username})`
      );
    }
  } else if (m[1].startsWith("/v1/messages") && !m[1].startsWith("/v1/messages/count_tokens")) {
    // A message without a session id did not come from Claude Code, and could
    // be neither counted nor shown: refused whatever the gate.
    return fail(res, 403, "EUNREGISTERED", "a message needs its Claude Code session id");
  }
  if (session?.purged_at) {
    return fail(
      res,
      410,
      "EPURGED",
      "this session is older than the hub keeps conversations; start a new session"
    );
  }

  const acct = require("../lib/claude-pool").forRemote(
    who.user.id,
    session ? session.account : null
  );
  if (!acct) return fail(res, 503, "ENOACCOUNT", "no company account is available right now");
  if (session && session.account !== acct.name) sessions.setAccount(session.session_id, acct.name);

  const base = upstream();
  const target = new URL(`${base.origin}${m[1]}${m[2] || ""}`);
  if (target.origin !== base.origin || !target.pathname.startsWith("/v1/")) {
    return fail(res, 400, "EBADPATH", "path leaves /v1/");
  }
  const headers = {};
  for (const [k, v] of Object.entries(req.headers)) {
    if (!DROP_REQUEST.has(k.toLowerCase())) headers[k] = v;
  }
  headers.authorization = `Bearer ${acct.bearer}`;
  const betas = String(headers["anthropic-beta"] || "")
    .split(",")
    .map((s) => s.trim())
    .filter(Boolean);
  if (!betas.includes(OAUTH_BETA)) betas.push(OAUTH_BETA);
  headers["anthropic-beta"] = betas.join(",");
  headers[GATEWAY_MARK] = "1";

  const mod = target.protocol === "https:" ? https : http;
  const up = mod.request(target, { method: req.method, headers }, (answer) => {
    const out = {};
    for (const [k, v] of Object.entries(answer.headers)) {
      if (!DROP_RESPONSE.has(k.toLowerCase())) out[k] = v;
    }
    res.writeHead(answer.statusCode || 502, out);
    const tap = session ? usageTap(answer.headers["content-type"]) : null;
    // The answer is relayed exactly as it came (compressed, if it was); the
    // usage is read from a decompressed COPY of it.
    const inflate = tap ? decoderFor(answer.headers["content-encoding"]) : null;
    const count = () => {
      try {
        sessions.countRequest(session.session_id, tap.end());
      } catch (e) {
        console.warn(`[remote-gw] usage not counted: ${e.message}`);
      }
    };
    if (inflate) {
      inflate.on("data", (buf) => tap.write(buf));
      inflate.on("end", count);
      inflate.on("error", count); // a broken copy only costs the count's detail
    }
    answer.on("data", (buf) => {
      if (inflate) inflate.write(buf);
      else if (tap) tap.write(buf);
      res.write(buf);
    });
    answer.on("end", () => {
      res.end();
      if (inflate) inflate.end();
      else if (tap) count();
    });
    answer.on("error", () => res.destroy());
  });
  up.setTimeout(IDLE_MS, () => up.destroy(new Error("upstream went quiet")));
  up.on("error", (e) => {
    if (!res.headersSent) fail(res, 502, "EUPSTREAM", e.message);
    else res.destroy();
  });
  res.on("close", () => {
    if (!res.writableFinished) up.destroy();
  });
  req.pipe(up);
}

module.exports = {
  remoteGateway,
  usageTap,
  __gate: gate,
  __resetQuota: () => quotaSeen.clear(),
};