AgentHub/server/lib/upload-staging.js

/**
 * @file Resumable uploads for the File Explorer: an upload's bytes are sent in
 * parts and kept on the HOST's disk under
 * `<data dir>/upload-staging/<user key>/<id>/` (0700/0600; never in a member's
 * home), so a restart of the dashboard — an update, a manual restart, a crash —
 * only costs the part in flight: the browser asks how much arrived (`status`)
 * and sends the rest. Only when every byte is here is the file written where
 * it goes, by the SAME `upload`/`uploadBatch` code (and guards) as a one-shot
 * upload, reading the staged file instead of the request (`commit`).
 *
 * `<id>` is the browser's own random name for one upload; ids live per user
 * key, so one user can never read, extend or commit another's. A part must
 * start exactly where the staged bytes end (`EOFFSET` with `received`), one
 * request per id at a time (`EBUSY`), never past the `total` announced first.
 * A committed upload keeps its answer (`result.json`) and drops its bytes, so
 * a commit whose answer a restart swallowed is answered again, never written
 * twice; a failed commit (a taken name, say) keeps the bytes for a retry with
 * another `conflict`. At most `MAX_OPEN` unfinished uploads per user; anything
 * untouched for `KEEP_MS` is swept.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

"use strict";

const fs = require("node:fs");
const path = require("node:path");
const { Transform } = require("node:stream");
const { pipeline } = require("node:stream/promises");

const { getDataDir } = require("./claude-home");

/** The largest upload (the browser sends a whole file — or a batch — under one id). */
const TOTAL_MAX = 512 * 1024 * 1024;
/** The largest part one request may carry. */
const PART_MAX = 16 * 1024 * 1024;
const MAX_OPEN = 8;
const KEEP_MS = 24 * 60 * 60 * 1000;

const ID_RE = /^[A-Za-z0-9_-]{16,64}$/;
const KEY_RE = /^[A-Za-z0-9_-]{1,64}$/;

const bad = (code, message, extra = {}) => Object.assign(new Error(message), { code, ...extra });

/** Ids (`<key>/<id>`) a request in THIS process is working on. */
const busy = new Set();

const root = () => path.join(getDataDir(), "upload-staging");

function dirOf(userKey, id) {
  const key = String(userKey ?? "");
  if (!KEY_RE.test(key)) throw bad("EBADBODY", "no user for this upload");
  if (typeof id !== "string" || !ID_RE.test(id)) {
    throw bad("EBADBODY", "an upload id is 16–64 letters, digits, '-' or '_'");
  }
  return { key, id, dir: path.join(root(), key, id), lock: `${key}/${id}` };
}

const dataFile = (d) => path.join(d.dir, "data");
const metaFile = (d) => path.join(d.dir, "meta.json");
const resultFile = (d) => path.join(d.dir, "result.json");

function readJson(file) {
  try {
    return JSON.parse(fs.readFileSync(file, "utf8"));
  } catch {
    return null;
  }
}

function writeJson(file, value) {
  const tmp = `${file}.tmp`;
  fs.writeFileSync(tmp, JSON.stringify(value), { mode: 0o600 });
  fs.renameSync(tmp, file);
}

const sizeOf = (d) => {
  try {
    return fs.statSync(dataFile(d)).size;
  } catch {
    return 0;
  }
};

function int(v, name, max) {
  const n = typeof v === "string" && /^\d+$/.test(v) ? Number(v) : v;
  if (!Number.isSafeInteger(n) || n < 0 || n > max) throw bad("EBADBODY", `bad ${name}`);
  return n;
}

/** Drop uploads untouched for KEEP_MS (every user's — the store is one). */
function sweep(now = Date.now()) {
  let keys = [];
  try {
    keys = fs.readdirSync(root());
  } catch {
    return;
  }
  for (const key of keys) {
    let ids = [];
    try {
      ids = fs.readdirSync(path.join(root(), key));
    } catch {
      continue;
    }
    for (const id of ids) {
      if (busy.has(`${key}/${id}`)) continue;
      const dir = path.join(root(), key, id);
      try {
        if (now - fs.statSync(dir).mtimeMs > KEEP_MS)
          fs.rmSync(dir, { recursive: true, force: true });
      } catch {
        /* gone meanwhile */
      }
    }
  }
}

/** Unfinished uploads of `key` (a committed one keeps only its answer). */
function openCount(key) {
  let ids = [];
  try {
    ids = fs.readdirSync(path.join(root(), key));
  } catch {
    return 0;
  }
  return ids.filter((id) => !fs.existsSync(path.join(root(), key, id, "result.json"))).length;
}

/** Touch the folder so `sweep` counts from the last part, not the first. */
function touch(d) {
  const now = new Date();
  try {
    fs.utimesSync(d.dir, now, now);
  } catch {
    /* best effort */
  }
}

/** `{received, total, done}` of one upload (`received: 0`, `total: null` when unknown). */
function status(userKey, id) {
  const d = dirOf(userKey, id);
  const meta = readJson(metaFile(d));
  if (!meta) return { received: 0, total: null, done: false };
  const done = fs.existsSync(resultFile(d));
  return { received: done ? meta.total : sizeOf(d), total: meta.total, done };
}

/**
 * Append the request `stream` at `offset` of upload `id` (of `total` bytes in
 * all). Returns `{received, total}`.
 */
async function putPart(userKey, id, { offset, total, stream }) {
  const d = dirOf(userKey, id);
  const at = int(offset, "offset", TOTAL_MAX);
  const all = int(total, "total", TOTAL_MAX);
  if (busy.has(d.lock)) throw bad("EBUSY", "this upload is already receiving a part");
  busy.add(d.lock);
  try {
    let meta = readJson(metaFile(d));
    if (!meta) {
      sweep();
      if (openCount(d.key) >= MAX_OPEN) {
        throw bad("ETOOMANY", `at most ${MAX_OPEN} unfinished uploads at a time`);
      }
      fs.mkdirSync(d.dir, { recursive: true, mode: 0o700 });
      meta = { total: all, createdAt: new Date().toISOString() };
      writeJson(metaFile(d), meta);
    }
    if (meta.total !== all) throw bad("EBADBODY", `this upload is ${meta.total} bytes`);
    if (fs.existsSync(resultFile(d))) return { received: all, total: all };
    const have = sizeOf(d);
    if (at !== have) {
      throw bad("EOFFSET", `this upload has ${have} bytes, not ${at}`, { received: have });
    }
    const room = Math.min(PART_MAX, all - have);
    let n = 0;
    const cap = new Transform({
      transform(chunk, _enc, cb) {
        n += chunk.length;
        if (n > room) return cb(bad("ETOOBIG", "this part goes past the upload's end"));
        cb(null, chunk);
      },
    });
    // A part cut short leaves what arrived: the next one starts from there.
    await pipeline(stream, cap, fs.createWriteStream(dataFile(d), { flags: "a", mode: 0o600 }));
    touch(d);
    return { received: sizeOf(d), total: all };
  } finally {
    busy.delete(d.lock);
  }
}

/**
 * Write a fully staged upload: `write(stream, size)` is the one-shot upload's
 * own code, given the staged bytes; it answers `{status, body}`. A commit
 * already made answers its stored answer again; a failed one keeps the bytes.
 */
async function commit(userKey, id, write) {
  const d = dirOf(userKey, id);
  const meta = readJson(metaFile(d));
  if (!meta) throw bad("ENOFILE", "no such upload — send it again");
  if (busy.has(d.lock)) throw bad("EBUSY", "this upload is still being written");
  busy.add(d.lock);
  try {
    const stored = readJson(resultFile(d));
    if (stored) return stored;
    const have = sizeOf(d);
    if (have !== meta.total) {
      throw bad("EINCOMPLETE", `this upload has ${have} of ${meta.total} bytes`, {
        received: have,
      });
    }
    const stream = fs.createReadStream(dataFile(d));
    let out;
    try {
      out = await write(stream, have);
    } finally {
      stream.destroy();
    }
    writeJson(resultFile(d), out);
    fs.rmSync(dataFile(d), { force: true });
    touch(d);
    return out;
  } finally {
    busy.delete(d.lock);
  }
}

/**
 * Forget one upload (the browser stopped it, or has its answer). Also while a
 * part is still arriving: a Stop drops the request, and the part writes on
 * into a file no longer listed — never refused, or the stopped upload would
 * hold one of MAX_OPEN for a day.
 */
function drop(userKey, id) {
  const d = dirOf(userKey, id);
  fs.rmSync(d.dir, { recursive: true, force: true });
  return { dropped: true };
}

module.exports = {
  status,
  putPart,
  commit,
  drop,
  sweep,
  root,
  TOTAL_MAX,
  PART_MAX,
  MAX_OPEN,
  KEEP_MS,
};