AgentHub/server/__tests__/remote-client.test.js

/**
 * @file remote-client.test.js
 * @description A LOCAL AgentHub's side of the remote-clients architecture
 * (server/lib/remote-client.js, /api/remote-client), against a fake hub:
 * enrollment keeps the device token in a 0600 file and never answers it; the
 * Company account is Claude-only, enrolled-only, and points a pane's
 * `--settings` at the hub's gateway; a resume stays on it; only sessions
 * started on it are uploaded — whole lines, secret-shaped strings replaced,
 * from where the hub stands (a 409 resyncs), the end sent once; a pane that
 * gets a new session id registers it; logout revokes on the hub.
 * @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 TMP = fs.mkdtempSync(path.join(os.tmpdir(), "remote-client-"));
process.env.DASHBOARD_DB_PATH = path.join(TMP, "dashboard.db");
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");
process.env.DASHBOARD_LIVENESS_PROBE = "0";
delete process.env.AGENTHUB_MODE;

const { createApp, startServer } = require("../index");
const { db } = require("../db");
const remote = require("../lib/remote-client");
const { resolveProxy, resumeProxy, proxyForSpawn } = require("../lib/proxy-apply");

const TOKEN = `ahdev_${"t".repeat(43)}`;
let hub;
let HUB;
let server;
let BASE;
const hubSeen = { registered: [], uploads: [], logout: 0 };
const consumed = new Map();
let forceOffset = null;
/** Sessions the hub answers ENOSESSION for until they are registered again. */
const forgotten = new Set();
/** Sessions the hub already purged. */
const purged = new Set();

function readBody(req) {
  return new Promise((r) => {
    let raw = "";
    req.on("data", (c) => (raw += c));
    req.on("end", () => r(raw ? JSON.parse(raw) : {}));
  });
}

before(async () => {
  hub = http.createServer(async (req, res) => {
    const body = await readBody(req);
    const json = (status, obj) => {
      res.writeHead(status, { "content-type": "application/json" });
      res.end(JSON.stringify(obj));
    };
    if (req.url === "/api/remote/device/login") {
      if (body.password !== "right-pass") {
        return json(401, { error: { code: "EBADLOGIN", message: "wrong username or password" } });
      }
      return json(200, {
        token: TOKEN,
        device: {
          id: 7,
          name: body.deviceName,
          expiresAt: new Date(Date.now() + 90 * 864e5).toISOString(),
        },
        user: { id: 3, username: body.username, role: "developer" },
        retainDays: 30,
      });
    }
    if (req.headers.authorization !== `Bearer ${TOKEN}`) {
      return json(401, { error: { code: "EDEVICE", message: "no" } });
    }
    if (req.url === "/api/remote/device/me") return json(200, { device: { id: 7 } });
    if (req.url === "/api/remote/device/logout") {
      hubSeen.logout++;
      return json(200, { ok: true });
    }
    if (req.url === "/api/remote/device/sessions") {
      hubSeen.registered.push(body);
      forgotten.delete(body.sessionId);
      return json(200, { ok: true, sessionId: body.sessionId });
    }
    if (req.url === "/api/remote/device/upload") {
      hubSeen.uploads.push(body);
      if (forgotten.has(body.sessionId)) {
        return json(404, { error: { code: "ENOSESSION", message: "not this device's" } });
      }
      if (purged.has(body.sessionId)) {
        return json(410, { error: { code: "EPURGED", message: "removed" } });
      }
      const files = [];
      for (const f of body.files) {
        const key = `${body.sessionId}/${f.path}`;
        const at = forceOffset ?? consumed.get(key) ?? 0;
        if (f.from !== at) {
          forceOffset = null;
          return json(409, {
            error: { code: "EOFFSET", message: "offset", expected: { path: f.path, consumed: at } },
          });
        }
        consumed.set(key, f.to);
        files.push({ path: f.path, consumed: f.to });
      }
      return json(200, { files, imported: true });
    }
    json(404, { error: { code: "ENOTFOUND", message: req.url } });
  });
  await new Promise((r) => hub.listen(0, "127.0.0.1", r));
  HUB = `http://127.0.0.1:${hub.address().port}`;
  server = await startServer(createApp(), 0);
  BASE = `http://127.0.0.1:${server.address().port}`;
});

after(() => {
  server?.close();
  hub?.close();
  fs.rmSync(TMP, { recursive: true, force: true });
});

const local = (method, p, body) =>
  fetch(BASE + p, {
    method,
    headers: { "content-type": "application/json", origin: BASE },
    body: body !== undefined ? JSON.stringify(body) : undefined,
  }).then(async (r) => ({ status: r.status, body: await r.json().catch(() => null) }));

/** A local session whose transcript is `lines`, optionally a Company-account one. */
function makeSession({ company, lines, runStatus = "running" }) {
  const id = crypto.randomUUID();
  const dir = path.join(process.env.CLAUDE_HOME, "projects", "-w-p");
  fs.mkdirSync(dir, { recursive: true });
  const file = path.join(dir, `${id}.jsonl`);
  fs.writeFileSync(file, lines);
  db.prepare(
    "INSERT INTO sessions (id, name, status, cwd, transcript_path) VALUES (?, 'x', 'active', '/w/p', ?)"
  ).run(id, file);
  db.prepare(
    `INSERT INTO dashboard_runs (id, session_id, cwd, mode, status, proxy) VALUES (?, ?, '/w/p', 'interactive', ?, ?)`
  ).run(`run-${id}`, id, runStatus, company ? "agenthub.company" : null);
  if (company) remote.registerSoon({ sessionId: id, cwd: "/w/p" });
  return { id, file };
}

describe("enrollment", () => {
  it("is off until login; a wrong password stores nothing", async () => {
    const st = await local("GET", "/api/remote-client");
    assert.equal(st.body.enrolled, false);
    const bad = await local("POST", "/api/remote-client/login", {
      hub: HUB,
      username: "mia",
      password: "wrong",
    });
    assert.equal(bad.status, 401);
    assert.equal(remote.readState(), null);
  });

  it("refuses a hub address with a login in it or another scheme", () => {
    assert.throws(() => remote.normalizeHub("http://u:p@10.1.13.7:4820"), { code: "EBADHUB" });
    assert.throws(() => remote.normalizeHub("ftp://10.1.13.7"), { code: "EBADHUB" });
    assert.equal(remote.normalizeHub("http://10.1.13.7:4820/"), "http://10.1.13.7:4820");
  });

  it("keeps the token in a 0600 file and never answers it", async () => {
    const r = await local("POST", "/api/remote-client/login", {
      hub: `${HUB}/`,
      username: "mia",
      password: "right-pass",
      deviceName: "mia-pc",
    });
    assert.equal(r.status, 200);
    assert.equal(r.body.enrolled, true);
    assert.equal(r.body.user.username, "mia");
    assert.ok(!JSON.stringify(r.body).includes(TOKEN));
    const file = path.join(process.env.DASHBOARD_DATA_DIR, "runtime", "remote-client.json");
    assert.equal(fs.statSync(file).mode & 0o777, 0o600);
    assert.ok(!fs.readFileSync(file, "utf8").includes("right-pass"), "the password is not kept");
    const st = await local("GET", "/api/remote-client");
    assert.ok(!JSON.stringify(st.body).includes(TOKEN));
  });
});

describe("the Company account", () => {
  it("is Claude-only and becomes a --settings file pointing at the gateway", () => {
    assert.throws(() => resolveProxy("codex", "agenthub.company"), { code: "EUNSUPPORTED" });
    const profile = resolveProxy("claude", "agenthub.company");
    assert.equal(profile.company, true);
    const { argv, env } = proxyForSpawn("claude", profile);
    assert.deepEqual(env, {}, "nothing secret goes through the pane env");
    const file = argv[0].replace("--settings=", "");
    assert.equal(fs.statSync(file).mode & 0o777, 0o600);
    const s = JSON.parse(fs.readFileSync(file, "utf8"));
    assert.equal(s.env.ANTHROPIC_BASE_URL, `${HUB}/api/remote/gw`);
    assert.equal(s.env.ANTHROPIC_AUTH_TOKEN, TOKEN);
    assert.equal(s.env.ENABLE_PROMPT_CACHING_1H, "1");
  });

  it("a resume stays on it", () => {
    const { id } = makeSession({ company: true, lines: "" });
    assert.equal(resumeProxy(id), "agenthub.company");
  });

  it("no profile may take its name", () => {
    const proxies = require("../lib/ai-proxies");
    const p = proxies.createProxy({
      label: "Company",
      baseUrl: "http://127.0.0.1:9",
      apiKey: "k",
      model: "m",
    });
    assert.notEqual(p.name, "agenthub.company");
    // A profile made before the Company account existed keeps working as itself.
    assert.equal(p.name, "company");
    const resolved = resolveProxy("claude", "company");
    assert.equal(resolved.company, undefined);
    assert.equal(resolved.name, "company");
    proxies.removeProxy(p.name);
  });
});

describe("the upload", () => {
  it("sends whole lines of Company sessions only, secrets replaced, then only what is new", async () => {
    hubSeen.uploads = [];
    const secret = `sk-ant-api03-${"s".repeat(40)}`;
    const line1 = JSON.stringify({ type: "user", message: { content: `my key is ${secret}` } });
    const line2 = JSON.stringify({
      type: "assistant",
      message: { usage: { input_tokens: 5, output_tokens: 6 } },
    });
    const company = makeSession({ company: true, lines: `${line1}\n${line2}\n{"partial":` });
    const own = makeSession({ company: false, lines: `${line1}\n` });
    const r = await remote.uploadAll();
    assert.deepEqual(r.errors, []);
    const mine = hubSeen.uploads.filter((u) => u.sessionId === company.id);
    const sent = mine.filter((u) => u.files.length);
    assert.equal(sent.length, 1);
    assert.equal(sent[0].more, true, "a chunk leaves the import to the round's close");
    assert.equal(mine.length, 2);
    assert.deepEqual(mine[1], { sessionId: company.id, files: [] }, "the close: import now");
    const data = sent[0].files[0].data;
    assert.ok(!data.includes(secret), "the key never leaves the machine");
    assert.ok(data.includes("[redacted]"));
    assert.ok(data.includes('"output_tokens":6'), "usage numbers survive");
    assert.ok(!data.includes("partial"), "a line still being written waits");
    assert.equal(sent[0].files[0].to, Buffer.byteLength(`${line1}\n${line2}\n`));
    assert.equal(
      hubSeen.uploads.filter((u) => u.sessionId === own.id).length,
      0,
      "an own-login session is never sent"
    );

    fs.appendFileSync(company.file, ' "done"}\n');
    hubSeen.uploads = [];
    await remote.uploadAll();
    const next = hubSeen.uploads.filter((u) => u.sessionId === company.id && u.files.length);
    assert.equal(next.length, 1);
    assert.equal(next[0].files[0].from, Buffer.byteLength(`${line1}\n${line2}\n`));
    assert.equal(next[0].files[0].data, '{"partial":"done"}\n');
  });

  it("resyncs from where the hub stands after a 409", async () => {
    const line = JSON.stringify({ type: "user", message: { content: "hi" } });
    const s = makeSession({ company: true, lines: `${line}\n${line}\n` });
    // The hub already holds the first line (an earlier answer got lost).
    consumed.set(`${s.id}/${s.id}.jsonl`, Buffer.byteLength(`${line}\n`));
    hubSeen.uploads = [];
    await remote.uploadAll();
    const tries = hubSeen.uploads.filter((u) => u.sessionId === s.id && u.files.length);
    assert.equal(tries.length, 2);
    assert.equal(tries[1].files[0].from, Buffer.byteLength(`${line}\n`));
    assert.equal(consumed.get(`${s.id}/${s.id}.jsonl`), Buffer.byteLength(`${line}\n${line}\n`));
  });

  it("sends the end once, when no pane holds the session any more", async () => {
    const s = makeSession({ company: true, lines: '{"a":1}\n', runStatus: "exited" });
    db.prepare("UPDATE sessions SET updated_at = '2000-01-01T00:00:00Z' WHERE id = ?").run(s.id);
    hubSeen.uploads = [];
    await remote.uploadAll();
    await remote.uploadAll();
    const ends = hubSeen.uploads.filter((u) => u.sessionId === s.id && u.ended === true);
    assert.equal(ends.length, 1);
  });
});

describe("upload edges", () => {
  it("sends a line too big for the hub as a placeholder and moves past it", async () => {
    const big = JSON.stringify({
      type: "user",
      message: { content: "x".repeat(remote.MAX_LINE_BYTES + 10) },
    });
    const small = JSON.stringify({ type: "user", message: { content: "after" } });
    const s = makeSession({ company: true, lines: `${big}\n${small}\n` });
    hubSeen.uploads = [];
    const r = await remote.uploadAll({ onlySession: s.id });
    assert.deepEqual(r.errors, []);
    const data = hubSeen.uploads
      .filter((u) => u.sessionId === s.id)
      .map((u) => u.files.map((f) => f.data).join(""))
      .join("");
    assert.ok(data.length < 1024, "the big line is not sent");
    assert.match(data, /"type":"agenthub-omitted","bytes":\d+/);
    assert.match(data, /after/);
    assert.equal(
      consumed.get(`${s.id}/${s.id}.jsonl`),
      Buffer.byteLength(`${big}\n${small}\n`),
      "the hub's offset still covers the whole file"
    );
  });

  it("registers a session again after this machine logged in again (ENOSESSION)", async () => {
    const s = makeSession({ company: true, lines: '{"a":1}\n' });
    await new Promise((r) => setTimeout(r, 100)); // registerSoon's own registration
    forgotten.add(s.id);
    hubSeen.registered = [];
    hubSeen.uploads = [];
    const r = await remote.uploadAll({ onlySession: s.id });
    assert.deepEqual(r.errors, []);
    assert.equal(hubSeen.registered.filter((b) => b.sessionId === s.id).length, 1);
    assert.equal(hubSeen.registered[0].cwd, "/w/p");
    assert.equal(consumed.get(`${s.id}/${s.id}.jsonl`), 8);
  });

  it("also registers again when only the end is left to send", async () => {
    const s = makeSession({ company: true, lines: '{"a":1}\n', runStatus: "exited" });
    db.prepare("UPDATE sessions SET updated_at = '2000-01-01T00:00:00Z' WHERE id = ?").run(s.id);
    await new Promise((r) => setTimeout(r, 100));
    consumed.set(`${s.id}/${s.id}.jsonl`, 8);
    db.prepare("INSERT INTO remote_uploads (session_id, rel_path, sent) VALUES (?, ?, 8)").run(
      s.id,
      `${s.id}.jsonl`
    );
    forgotten.add(s.id);
    hubSeen.uploads = [];
    const r = await remote.uploadAll({ onlySession: s.id });
    assert.deepEqual(r.errors, []);
    assert.ok(hubSeen.uploads.some((u) => u.sessionId === s.id && u.ended === true));
    assert.equal(forgotten.has(s.id), false, "registered again");
  });

  it("stops sending a session the hub already purged", async () => {
    const s = makeSession({ company: true, lines: '{"a":1}\n' });
    await new Promise((r) => setTimeout(r, 100));
    purged.add(s.id);
    const r = await remote.uploadAll({ onlySession: s.id });
    assert.deepEqual(r.errors, []);
    assert.equal(remote.isCompany(s.id), false);
    hubSeen.uploads = [];
    await remote.uploadAll({ onlySession: s.id });
    assert.equal(hubSeen.uploads.length, 0);
  });
});

describe("a new session id in a Company pane", () => {
  it("is registered from its hook; another pane's is not", async () => {
    const { id } = makeSession({ company: true, lines: "" });
    const run = `run-${id}`;
    await new Promise((r) => setTimeout(r, 100)); // its own registration lands first
    hubSeen.registered = [];
    const fresh = crypto.randomUUID();
    remote.onHookSession({ sessionId: fresh, runId: run, cwd: "/w/p" });
    const other = makeSession({ company: false, lines: "" });
    remote.onHookSession({ sessionId: crypto.randomUUID(), runId: `run-${other.id}`, cwd: "/w/p" });
    await new Promise((r) => setTimeout(r, 100));
    assert.deepEqual(
      hubSeen.registered.map((x) => x.sessionId),
      [fresh]
    );
    assert.equal(remote.isCompany(fresh), true);
  });
});

describe("logout", () => {
  it("revokes on the hub and forgets the token; the Company account is then refused", async () => {
    const r = await local("POST", "/api/remote-client/logout", {});
    assert.equal(r.status, 200);
    assert.equal(r.body.revokedOnHub, true);
    assert.equal(hubSeen.logout, 1);
    assert.equal(remote.readState(), null);
    assert.throws(() => resolveProxy("claude", "agenthub.company"), { code: "ENOHUB" });
  });
});