AgentHub/server/hub/remote-sessions.js

/**
 * @file remote-sessions.js
 * @description The company-account sessions remote clients registered with the
 * hub (server mode only). A client registers a session ONCE, when it starts it
 * on the "Company account" (`POST /api/remote/device/sessions`), with its id,
 * folder and tool — no content. The gateway (remote-gateway.js) relays a
 * request only for a session THAT device registered; everything a session
 * says arrives later in the hourly upload (remote-ingest.js).
 *
 * A session id is claimed by its first device: another device (or a session id
 * the hub already holds from anywhere else) is refused with EFOREIGN, so one
 * client can never write into another user's session. The pool account the
 * gateway picks is kept on the row for the session's whole life: the prompt
 * cache belongs to one account, and switching would re-write it all.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const sql = (q) => require("../db").db.prepare(q);
const err = (code, message, status) => Object.assign(new Error(message), { code, status });

const UUID_RE = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;
const TOOLS = new Set(["claude"]);
const MAX_CWD = 1024;

/** A client-sent folder, as text only: absolute-looking, printable, capped. */
function cleanCwd(cwd) {
  const s = String(cwd || "").replace(/[\u0000-\u001f\u007f]/g, "");
  if (!s || s.length > MAX_CWD) throw err("EBADCWD", "cwd must be a folder path", 400);
  if (!s.startsWith("/") && !/^[A-Za-z]:[\\/]/.test(s)) {
    throw err("EBADCWD", "cwd must be absolute", 400);
  }
  return s;
}

/**
 * Register (or re-register, on a resume) `sessionId` for `device`. Refuses a
 * session another device registered, and one the hub already knows that did
 * not come from this device (a local session, another user's upload).
 */
function registerSession(device, { sessionId, cwd, tool = "claude" } = {}) {
  if (!UUID_RE.test(String(sessionId || ""))) {
    throw err("EBADSESSION", "sessionId must be a UUID", 400);
  }
  if (!TOOLS.has(tool)) throw err("EBADTOOL", "only claude sessions use the company account", 400);
  const folder = cleanCwd(cwd);
  const id = String(sessionId).toLowerCase();
  let row = sql(
    "SELECT device_id, user_id, cwd, purged_at FROM remote_sessions WHERE session_id = ?"
  ).get(id);
  if (row?.purged_at) {
    throw err(
      "EPURGED",
      "this session's conversation was removed after the retention period; start a new session",
      410
    );
  }
  if (row && row.device_id !== device.id) {
    // The same user on a new device (enrolled again, a new machine) resumes
    // their own session: it moves to that device. Anyone else's never does.
    if (row.user_id !== device.userId) {
      throw err("EFOREIGN", "this session belongs to another device", 409);
    }
    require("./remote-ingest").moveSession(id, row.device_id, device.id, device.userId);
    row = { ...row, device_id: device.id, user_id: device.userId };
  }
  if (!row) {
    const known = sql("SELECT source FROM sessions WHERE id = ?").get(id);
    if (known && known.source !== `device:${device.id}`) {
      throw err("EFOREIGN", "this session id is already taken", 409);
    }
    sql(
      "INSERT INTO remote_sessions (session_id, device_id, user_id, cwd, tool) VALUES (?, ?, ?, ?, ?)"
    ).run(id, device.id, device.userId, folder, tool);
  } else {
    // The folder is kept from the first registration: the hub's copy lives
    // under it, and a later one (a resume, a hook) would split that copy.
    // A registration that lost the race to its first request (the gateway
    // waited REGISTER_GRACE_MS and recorded it unmonitored) clears the flag:
    // its device did register it, which is all "monitored" means.
    sql(
      "UPDATE remote_sessions SET ended_at = NULL, unmonitored = 0, cwd = COALESCE(NULLIF(cwd, ''), ?) WHERE session_id = ?"
    ).run(folder, id);
  }
  return getSession(id);
}

function getSession(sessionId) {
  return (
    sql("SELECT * FROM remote_sessions WHERE session_id = ?").get(
      String(sessionId || "").toLowerCase()
    ) || null
  );
}

/** The registered session `sessionId` of `deviceId`, or null. */
function sessionOfDevice(sessionId, deviceId) {
  const row = getSession(sessionId);
  return row && row.device_id === deviceId ? row : null;
}

function markEnded(sessionId, deviceId) {
  return (
    sql(
      "UPDATE remote_sessions SET ended_at = strftime('%Y-%m-%dT%H:%M:%fZ','now') WHERE session_id = ? AND device_id = ?"
    ).run(String(sessionId).toLowerCase(), deviceId).changes > 0
  );
}

/**
 * A session the gateway relayed for although its device never registered it
 * (remote_gate = warn): kept as UNMONITORED so an admin sees it. A session id
 * already held by another device stays refused (null).
 */
function recordUnmonitored(device, sessionId) {
  const id = String(sessionId || "").toLowerCase();
  if (!UUID_RE.test(id)) return null;
  const row = getSession(id);
  if (row) return row.device_id === device.id ? row : null;
  // Never for a session the hub already knows from elsewhere (a hub pane's,
  // another device's upload): that would hand it to this device's user.
  const known = sql("SELECT source FROM sessions WHERE id = ?").get(id);
  if (known && known.source !== `device:${device.id}`) return null;
  sql(
    "INSERT INTO remote_sessions (session_id, device_id, user_id, unmonitored) VALUES (?, ?, ?, 1)"
  ).run(id, device.id, device.userId);
  scheduleUnmonitoredAlert(id);
  return getSession(id);
}

/** How long a session may stay unregistered before the alert fires. */
const ALERT_AFTER_MS = () => Number(process.env.AGENTHUB_REMOTE_UNMONITORED_ALERT_MS) || 60_000;

/**
 * Fire the `remote_unmonitored` alert rules for `sessionId` once it has stayed
 * unregistered for ALERT_AFTER_MS — a registration that was only late (the
 * hub was briefly unreachable) clears the flag first and alerts nobody.
 */
function scheduleUnmonitoredAlert(sessionId) {
  setTimeout(() => {
    try {
      const row = sql(
        `SELECT r.*, d.name AS device_name, u.username
           FROM remote_sessions r JOIN devices d ON d.id = r.device_id
           LEFT JOIN users u ON u.id = r.user_id
          WHERE r.session_id = ?`
      ).get(sessionId);
      if (!row || !row.unmonitored) return;
      require("../lib/alerts").fireRemoteUnmonitored({
        sessionId,
        username: row.username || `#${row.user_id}`,
        device: row.device_name,
        deviceId: row.device_id,
        userId: row.user_id,
      });
    } catch (e) {
      console.warn(`[remote] unmonitored alert failed: ${e.message}`);
    }
  }, ALERT_AFTER_MS()).unref();
}

/** Keep the pool account the gateway picked for this session. */
function setAccount(sessionId, account) {
  sql("UPDATE remote_sessions SET account = ? WHERE session_id = ?").run(account, sessionId);
}

/** Count one relayed request and the usage its answer reported. */
function countRequest(sessionId, usage = {}) {
  const n = (v) => (Number.isFinite(v) && v > 0 ? Math.round(v) : 0);
  sql(
    `UPDATE remote_sessions SET requests = requests + 1,
            last_request_at = strftime('%Y-%m-%dT%H:%M:%fZ','now'),
            gw_input_tokens = gw_input_tokens + ?, gw_output_tokens = gw_output_tokens + ?,
            gw_cache_read_tokens = gw_cache_read_tokens + ?,
            gw_cache_write_tokens = gw_cache_write_tokens + ?
      WHERE session_id = ?`
  ).run(
    n(usage.input_tokens),
    n(usage.output_tokens),
    n(usage.cache_read_input_tokens),
    n(usage.cache_creation_input_tokens),
    sessionId
  );
}

/**
 * The registered sessions a user (or every user, `userId` null) has, newest
 * first, with their device's name — the "device › folder" view's rows.
 */
function listSessions(userId = null, limit = 200) {
  const where = userId == null ? "" : "WHERE r.user_id = ?";
  const args = userId == null ? [] : [Number(userId)];
  return sql(
    `SELECT r.session_id AS sessionId, r.device_id AS deviceId, d.name AS deviceName,
            r.user_id AS userId, r.cwd, r.tool, r.account, r.registered_at AS registeredAt,
            r.ended_at AS endedAt, r.last_request_at AS lastRequestAt, r.requests,
            r.gw_input_tokens AS inputTokens, r.gw_output_tokens AS outputTokens,
            r.gw_cache_read_tokens AS cacheReadTokens, r.gw_cache_write_tokens AS cacheWriteTokens,
            r.last_upload_at AS lastUploadAt, r.purged_at AS purgedAt,
            r.unmonitored,
            s.name AS title
       FROM remote_sessions r
       JOIN devices d ON d.id = r.device_id
       LEFT JOIN sessions s ON s.id = r.session_id
       ${where}
      ORDER BY COALESCE(r.last_request_at, r.registered_at) DESC
      LIMIT ?`
  ).all(...args, limit);
}

module.exports = {
  UUID_RE,
  registerSession,
  getSession,
  sessionOfDevice,
  markEnded,
  recordUnmonitored,
  setAccount,
  countRequest,
  listSessions,
};