AgentHub/server/lib/event-payload.js

/**
 * @file event-payload.js
 * @description Trims a hook payload before it is stored as `events.data`. A
 * PostToolUse payload carries whole files — Edit's `originalFile`, Read's file
 * `content`/`base64`, Write's `content` — and long command output, none of which
 * any analysis reads: the server reads only `tool_input.skill`, the event views
 * only preview. The full text stays in the session's transcript (and AgentHub's
 * snapshot of it under `<data dir>/transcripts/`), which the Conversation tab
 * reads. Kept: every field outside `tool_input`/`tool_response`, every number and
 * flag, `structuredPatch` (the +/- line counts), and each long string's head and
 * tail. Dropped fields are named in `_trimmed`, so a view can say so.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

// Strings inside tool_input / tool_response longer than this keep head + tail.
const STR_CAP = 8192;
const HEAD = 6000;
const TAIL = 2000;
// A payload this short cannot hold anything worth trimming, so it is not parsed.
const SKIP_UNDER = STR_CAP;

// Per tool, the tool_response fields that mirror a whole file or repeat
// tool_input. Other tools' fields (Grep's `content` is its matches) only get capped.
const EDIT_DROP = ["originalFile", "oldString", "newString"];
const DROP_RESPONSE = {
  Edit: EDIT_DROP,
  MultiEdit: EDIT_DROP,
  NotebookEdit: EDIT_DROP,
  Write: [...EDIT_DROP, "content"],
};
const DROP_RESPONSE_FILE = { Read: ["content", "base64"] };

function capString(s, hit) {
  if (s.length <= STR_CAP) return s;
  hit.n++;
  return `${s.slice(0, HEAD)}\n…[${s.length - HEAD - TAIL} chars trimmed]…\n${s.slice(-TAIL)}`;
}

// Caps every string at any depth; `structuredPatch` is left whole.
function capDeep(v, key, hit) {
  if (typeof v === "string") return capString(v, hit);
  if (key === "structuredPatch" || v === null || typeof v !== "object") return v;
  if (Array.isArray(v)) return v.map((x) => capDeep(x, undefined, hit));
  const out = {};
  for (const [k, x] of Object.entries(v)) out[k] = capDeep(x, k, hit);
  return out;
}

/**
 * @param {object} data A parsed hook payload.
 * @param {{n: number}} [hit] Counts every drop and cap made, so a caller can tell
 *   "nothing to trim" apart.
 * @returns {object} A trimmed copy (the input is not mutated).
 */
function trimHookPayload(data, hit = { n: 0 }) {
  if (!data || typeof data !== "object" || Array.isArray(data)) return data;
  const out = { ...data };
  const dropped = new Set(Array.isArray(data._trimmed) ? data._trimmed : []);
  const res = out.tool_response;
  if (res && typeof res === "object" && !Array.isArray(res)) {
    const r = { ...res };
    for (const k of DROP_RESPONSE[data.tool_name] || []) {
      if (typeof r[k] === "string" && r[k].length > 0) {
        delete r[k];
        hit.n++;
        dropped.add(`tool_response.${k}`);
      }
    }
    const fileDrop = DROP_RESPONSE_FILE[data.tool_name];
    if (fileDrop && r.file && typeof r.file === "object") {
      const f = { ...r.file };
      for (const k of fileDrop) {
        if (typeof f[k] === "string" && f[k].length > 0) {
          delete f[k];
          hit.n++;
          dropped.add(`tool_response.file.${k}`);
        }
      }
      r.file = f;
    }
    out.tool_response = r;
  }
  if ("tool_response" in out) out.tool_response = capDeep(out.tool_response, undefined, hit);
  if ("tool_input" in out) out.tool_input = capDeep(out.tool_input, undefined, hit);
  if (dropped.size > 0) out._trimmed = [...dropped];
  return out;
}

/**
 * The SQL-side entry point (`trim_event_data(?)`, registered in `server/db.js`):
 * a JSON string in, a JSON string out. Short, null or unparseable input is
 * returned unchanged, and so is a payload with nothing to trim — re-serializing
 * it would round integers past 2^53. Trimming is idempotent.
 * @param {string|null} json
 * @returns {string|null}
 */
function trimEventData(json) {
  if (typeof json !== "string" || json.length <= SKIP_UNDER) return json;
  let parsed;
  try {
    parsed = JSON.parse(json);
  } catch {
    return json;
  }
  if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) return json;
  const hit = { n: 0 };
  const out = trimHookPayload(parsed, hit);
  // ponytail: a payload that IS trimmed is re-serialized whole, so a huge integer
  // elsewhere in it still rounds; splice the JSON text instead if one ever matters.
  return hit.n ? JSON.stringify(out) : json;
}

/**
 * Trims the rows stored before trimming existed, a few rows per write
 * transaction so a hook never waits long for the lock: ids are picked first and
 * each row's payload is read alone, and a transaction ends after `chunk` rows or
 * once the payloads it read pass `maxBytes` — a Read of an image can hold MBs of
 * base64, and 200 of those in one go stalled the event loop. Idempotent: a crash
 * just resumes from the start on the next run.
 * @param {import("better-sqlite3").Database} db
 * @param {{chunk?: number, maxBytes?: number}} [opts]
 * @returns {Promise<{rows: number, bytesSaved: number}>}
 */
function sweepEventPayloads(db, { chunk = 200, maxBytes = 8 * 1024 * 1024 } = {}) {
  const pick = db
    .prepare("SELECT id FROM events WHERE id > ? AND length(data) > ? ORDER BY id LIMIT ?")
    .pluck();
  const read = db.prepare("SELECT data FROM events WHERE id = ?").pluck();
  const update = db.prepare("UPDATE events SET data = ? WHERE id = ?");
  const out = { rows: 0, bytesSaved: 0 };
  // Returns the last id it handled; the next pick starts after it.
  const trimChunk = db.transaction((ids) => {
    let bytes = 0;
    let last = ids[0];
    for (const id of ids) {
      last = id;
      const data = read.get(id);
      if (typeof data !== "string") continue;
      bytes += data.length;
      const t = trimEventData(data);
      if (t !== data) {
        update.run(t, id);
        out.rows++;
        out.bytesSaved += data.length - t.length;
      }
      if (bytes >= maxBytes) break;
    }
    return last;
  });
  let lastId = 0;
  return new Promise((resolve, reject) => {
    const step = () => {
      try {
        const ids = pick.all(lastId, SKIP_UNDER, chunk);
        if (ids.length === 0) return resolve(out);
        lastId = trimChunk(ids);
        setImmediate(step);
      } catch (err) {
        reject(err);
      }
    };
    step();
  });
}

module.exports = { trimHookPayload, trimEventData, sweepEventPayloads, STR_CAP };