src/worker/sessionWorker.ts

import { parentPort } from "node:worker_threads";
import {
  Bash,
  InMemoryFs,
  CommandCollectorPlugin,
} from "just-bash";
import { posix as path } from "node:path";
import { truncateUtf8Safe, sliceUtf8Bytes, isValidUtf8 } from "../utils/utf8.js";
import type {
  ParentMessage,
  WorkerMessage,
  InitMessage,
  ExecRequest,
  FilesWriteRequest,
  FilesReadRequest,
  FilesListRequest,
} from "./protocol.js";
import { validateAbsolutePath } from "../utils/paths.js";

if (!parentPort) {
  throw new Error("sessionWorker must run as a worker thread");
}

const port = parentPort;

let bash: Bash | null = null;
let fs: InMemoryFs | null = null;
let collector: CommandCollectorPlugin | null = null;
let initCfg: InitMessage | null = null;
const aborters = new Map<string, AbortController>();

function send(msg: WorkerMessage) {
  port.postMessage(msg);
}

function sendError(requestId: string, code: string, message: string) {
  send({ type: "error_result", requestId, code, message });
}

// Count only bytes inside the agent workspace roots. InMemoryFs ships with a
// large baseline (e.g. /etc, /usr/share) that we do not count toward the user quota.
const ACCOUNTED_ROOTS = ["/home/user", "/tmp"];

function isAccounted(p: string): boolean {
  return ACCOUNTED_ROOTS.some((r) => p === r || p.startsWith(`${r}/`));
}

async function totalFsBytes(): Promise<number> {
  if (!fs) return 0;
  let total = 0;
  const all = fs.getAllPaths();
  for (const p of all) {
    if (!isAccounted(p)) continue;
    try {
      const st = await fs.stat(p);
      if (st.isFile) total += st.size;
    } catch {
      // ignore
    }
  }
  return total;
}

async function handleInit(msg: InitMessage) {
  initCfg = msg;
  fs = new InMemoryFs();
  // Ensure cwd exists
  await fs.mkdir(msg.cwd, { recursive: true });
  // Apply initial files
  for (const [p, content] of Object.entries(msg.initialFiles)) {
    const norm = validateAbsolutePath(p);
    const parent = path.dirname(norm);
    await fs.mkdir(parent, { recursive: true });
    await fs.writeFile(norm, content, { encoding: "utf8" });
  }
  collector = new CommandCollectorPlugin();
  bash = new Bash({
    fs,
    cwd: msg.cwd,
    executionLimits: msg.executionLimits,
    network:
      msg.capabilities.network && msg.networkConfig
        ? {
            dangerouslyAllowFullInternetAccess:
              msg.networkConfig.dangerouslyAllowFullInternetAccess,
            allowedUrlPrefixes: msg.networkConfig.allowedUrlPrefixes,
            allowedMethods: msg.networkConfig.allowedMethods as any,
          }
        : undefined,
    python: msg.capabilities.python,
    javascript: msg.capabilities.javascript,
  });
  bash.registerTransformPlugin(collector);
  send({ type: "ready" });
}

async function handleExec(msg: ExecRequest) {
  if (!bash || !initCfg) {
    return sendError(msg.requestId, "NOT_INITIALIZED", "Worker not initialized");
  }
  const ac = new AbortController();
  aborters.set(msg.requestId, ac);
  const started = Date.now();
  let timedOut = false;
  const timer = setTimeout(() => {
    timedOut = true;
    ac.abort();
  }, msg.timeoutMs);
  try {
    const result = await bash.exec(msg.script, {
      stdin: msg.stdin,
      cwd: msg.cwd,
      env: msg.env,
      replaceEnv: msg.replaceEnv,
      signal: ac.signal,
    });
    clearTimeout(timer);
    aborters.delete(msg.requestId);

    const stdoutFull = result.stdout ?? "";
    const stderrFull = result.stderr ?? "";

    // Truncate for embedding
    const stdoutSafe = truncateUtf8Safe(stdoutFull, msg.maxStdoutBytes);
    const stderrSafe = truncateUtf8Safe(stderrFull, msg.maxStderrBytes);

    // Retained (capped) full output for resource links.
    const capStdoutFull = truncateUtf8Safe(
      stdoutFull,
      initCfg.maxCapturedOutputBytesExec,
    );
    const capStderrFull = truncateUtf8Safe(
      stderrFull,
      initCfg.maxCapturedOutputBytesExec,
    );

    const commands = extractCommandNames(result.metadata);

    send({
      type: "exec_result",
      requestId: msg.requestId,
      execId: msg.execId,
      exitCode: result.exitCode,
      stdout: stdoutSafe.text,
      stderr: stderrSafe.text,
      durationMs: Date.now() - started,
      timedOut,
      truncatedStdout: stdoutSafe.truncated,
      truncatedStderr: stderrSafe.truncated,
      capturedStdoutBytes: Buffer.byteLength(stdoutFull, "utf8"),
      capturedStderrBytes: Buffer.byteLength(stderrFull, "utf8"),
      retainedStdout: stdoutSafe.truncated ? capStdoutFull.text : undefined,
      retainedStderr: stderrSafe.truncated ? capStderrFull.text : undefined,
      commands,
      sessionFsBytes: await totalFsBytes(),
    });
  } catch (e: any) {
    clearTimeout(timer);
    aborters.delete(msg.requestId);
    // Treat as a tool error result so the parent can surface it.
    sendError(msg.requestId, timedOut ? "COOPERATIVE_TIMEOUT" : "EXEC_FAILED", e?.message ?? String(e));
  }
}

function extractCommandNames(meta: Record<string, unknown> | undefined): string[] {
  if (!meta) return [];
  const list = (meta as { commands?: unknown }).commands;
  if (!Array.isArray(list)) return [];
  const out: string[] = [];
  for (const c of list) {
    if (typeof c === "string" && c.length > 0) out.push(c);
  }
  return out;
}

async function handleFilesWrite(msg: FilesWriteRequest) {
  if (!fs || !initCfg) {
    return sendError(msg.requestId, "NOT_INITIALIZED", "Worker not initialized");
  }

  // Compute decoded sizes first
  const decoded: Array<{ path: string; bytes: Uint8Array }> = [];
  let totalDecodedBytes = 0;
  for (const f of msg.files) {
    let buf: Buffer;
    try {
      if (f.encoding === "base64") {
        buf = Buffer.from(f.content, "base64");
      } else {
        buf = Buffer.from(f.content, "utf8");
      }
    } catch (e: any) {
      return sendError(msg.requestId, "INVALID_INPUT", `Failed to decode ${f.path}: ${e.message}`);
    }
    if (buf.length > initCfg.maxSingleFileBytes) {
      return sendError(
        msg.requestId,
        "QUOTA_EXCEEDED",
        `File ${f.path} exceeds max single file bytes (${buf.length} > ${initCfg.maxSingleFileBytes})`,
      );
    }
    let norm: string;
    try {
      norm = validateAbsolutePath(f.path);
    } catch (e: any) {
      return sendError(msg.requestId, "INVALID_INPUT", `Invalid path ${f.path}: ${e.message}`);
    }
    totalDecodedBytes += buf.length;
    decoded.push({ path: norm, bytes: buf });
  }
  if (totalDecodedBytes > initCfg.maxWriteBatchBytes) {
    return sendError(
      msg.requestId,
      "QUOTA_EXCEEDED",
      `Batch decoded size ${totalDecodedBytes} exceeds max write batch ${initCfg.maxWriteBatchBytes}`,
    );
  }

  // Check session fs cap considering replacements.
  const currentBytesByPath = new Map<string, number>();
  for (const d of decoded) {
    if (await fs.exists(d.path)) {
      try {
        const st = await fs.stat(d.path);
        if (st.isFile) {
          if (!msg.overwrite) {
            return sendError(
              msg.requestId,
              "INVALID_INPUT",
              `File ${d.path} exists and overwrite=false`,
            );
          }
          currentBytesByPath.set(d.path, st.size);
        }
      } catch {
        // ignore
      }
    }
  }
  const currentTotal = await totalFsBytes();
  const replaceFreed = Array.from(currentBytesByPath.values()).reduce((a, b) => a + b, 0);
  const projected = currentTotal - replaceFreed + totalDecodedBytes;
  if (projected > initCfg.maxSessionFsBytes) {
    return sendError(
      msg.requestId,
      "QUOTA_EXCEEDED",
      `Session fs cap exceeded: ${projected} > ${initCfg.maxSessionFsBytes}`,
    );
  }

  // Write
  const written: Array<{ path: string; bytes: number }> = [];
  for (const d of decoded) {
    const parent = path.dirname(d.path);
    if (msg.createDirs) {
      await fs.mkdir(parent, { recursive: true });
    } else if (!(await fs.exists(parent))) {
      return sendError(msg.requestId, "INVALID_INPUT", `Parent dir does not exist: ${parent}`);
    }
    await fs.writeFile(d.path, d.bytes);
    written.push({ path: d.path, bytes: d.bytes.length });
  }
  send({
    type: "files_write_result",
    requestId: msg.requestId,
    written,
    sessionFsBytes: await totalFsBytes(),
  });
}

async function handleFilesRead(msg: FilesReadRequest) {
  if (!fs || !initCfg) {
    return sendError(msg.requestId, "NOT_INITIALIZED", "Worker not initialized");
  }
  let norm: string;
  try {
    norm = validateAbsolutePath(msg.path);
  } catch (e: any) {
    return sendError(msg.requestId, "INVALID_INPUT", e.message);
  }
  if (!(await fs.exists(norm))) {
    return sendError(msg.requestId, "NOT_FOUND", `File not found: ${norm}`);
  }
  let st;
  try {
    st = await fs.stat(norm);
  } catch (e: any) {
    return sendError(msg.requestId, "NOT_FOUND", e.message);
  }
  if (!st.isFile) {
    return sendError(msg.requestId, "INVALID_INPUT", `Not a file: ${norm}`);
  }
  const maxBytes = Math.min(msg.maxBytes, initCfg.maxReadBytesPerCall);
  const buf = await fs.readFileBuffer(norm);
  const total = buf.length;
  if (msg.offset < 0 || msg.offset > total) {
    return sendError(msg.requestId, "INVALID_INPUT", `Offset out of range`);
  }
  const end = Math.min(msg.offset + maxBytes, total);

  // Decide encoding
  const slice = buf.subarray(msg.offset, end);
  let encoding: "utf8" | "base64";
  let content: string;
  let bytesRead: number;
  let nextOffset: number | null;

  if (msg.encoding === "base64") {
    content = Buffer.from(slice).toString("base64");
    bytesRead = slice.length;
    encoding = "base64";
    nextOffset = end < total ? end : null;
  } else {
    // utf8 or auto: try utf8-safe slice with codepoint trimming
    const decoded = sliceUtf8Bytes(Buffer.from(buf).toString("utf8"), msg.offset, maxBytes);
    // Actually we need to operate on the raw bytes, not on the re-encoded utf8 text.
    // If the bytes round-trip cleanly as UTF-8 within the slice, return them; otherwise fall back to base64.
    if (isValidUtf8(slice)) {
      content = Buffer.from(slice).toString("utf8");
      bytesRead = slice.length;
      encoding = "utf8";
      nextOffset = end < total ? end : null;
    } else if (msg.encoding === "auto") {
      content = Buffer.from(slice).toString("base64");
      bytesRead = slice.length;
      encoding = "base64";
      nextOffset = end < total ? end : null;
    } else {
      // For utf8 requested with binary data: trim at codepoint boundary using decoded.
      content = decoded.text;
      bytesRead = decoded.bytesRead;
      encoding = "utf8";
      nextOffset = decoded.nextOffset;
    }
  }

  send({
    type: "files_read_result",
    requestId: msg.requestId,
    path: norm,
    content,
    encoding,
    bytesRead,
    nextOffset,
  });
}

function matchGlob(p: string, glob: string | undefined): boolean {
  if (!glob) return true;
  // Tiny glob → regex: ** = any chars including /; * = any chars except /; {a,b} = alternation; ? = single
  let re = "";
  let i = 0;
  while (i < glob.length) {
    const c = glob[i];
    if (c === "*") {
      if (glob[i + 1] === "*") {
        re += ".*";
        i += 2;
        if (glob[i] === "/") i++;
        continue;
      }
      re += "[^/]*";
      i++;
      continue;
    }
    if (c === "?") {
      re += "[^/]";
      i++;
      continue;
    }
    if (c === "{") {
      const end = glob.indexOf("}", i);
      if (end > 0) {
        const opts = glob.slice(i + 1, end).split(",");
        re += `(?:${opts.map((o) => o.replace(/[.+^$()|[\]\\]/g, "\\$&")).join("|")})`;
        i = end + 1;
        continue;
      }
    }
    re += c.replace(/[.+^$()|[\]\\]/g, "\\$&");
    i++;
  }
  return new RegExp(`^${re}$`).test(p);
}

async function handleFilesList(msg: FilesListRequest) {
  if (!fs || !initCfg) {
    return sendError(msg.requestId, "NOT_INITIALIZED", "Worker not initialized");
  }
  let root: string;
  try {
    root = validateAbsolutePath(msg.path);
  } catch (e: any) {
    return sendError(msg.requestId, "INVALID_INPUT", e.message);
  }
  if (!(await fs.exists(root))) {
    return sendError(msg.requestId, "NOT_FOUND", `Path not found: ${root}`);
  }
  const startStat = await fs.stat(root);
  const collected: Array<{
    path: string;
    type: "file" | "directory" | "symlink";
    bytes?: number;
  }> = [];

  const startIndex = msg.cursor ? Number(msg.cursor) : 0;
  let traversed = 0;

  async function walk(dir: string) {
    const names = await fs!.readdir(dir);
    for (const name of names) {
      const child = path.join(dir, name);
      let st;
      try {
        st = await fs!.lstat(child);
      } catch {
        continue;
      }
      const type: "file" | "directory" | "symlink" = st.isSymbolicLink
        ? "symlink"
        : st.isDirectory
          ? "directory"
          : "file";

      const relMatch = matchGlob(child, msg.glob);
      if (relMatch) {
        traversed++;
        if (traversed > startIndex) {
          collected.push({
            path: child,
            type,
            bytes: type === "file" ? st.size : undefined,
          });
          if (collected.length >= msg.maxEntries) return;
        }
      }
      if (msg.recursive && type === "directory") {
        await walk(child);
        if (collected.length >= msg.maxEntries) return;
      }
    }
  }

  if (startStat.isFile) {
    if (matchGlob(root, msg.glob)) {
      collected.push({ path: root, type: "file", bytes: startStat.size });
    }
  } else {
    await walk(root);
  }
  // Determine nextCursor
  let nextCursor: string | null = null;
  if (collected.length >= msg.maxEntries) {
    nextCursor = String(startIndex + collected.length);
  }
  send({
    type: "files_list_result",
    requestId: msg.requestId,
    entries: collected,
    nextCursor,
  });
}

port.on("message", async (raw: ParentMessage) => {
  try {
    switch (raw.type) {
      case "init":
        await handleInit(raw);
        break;
      case "exec":
        await handleExec(raw);
        break;
      case "abort": {
        const ac = aborters.get(raw.requestId);
        if (ac) ac.abort();
        break;
      }
      case "files_write":
        await handleFilesWrite(raw);
        break;
      case "files_read":
        await handleFilesRead(raw);
        break;
      case "files_list":
        await handleFilesList(raw);
        break;
      case "fs_bytes":
        send({
          type: "fs_bytes_result",
          requestId: raw.requestId,
          sessionFsBytes: await totalFsBytes(),
        });
        break;
      default: {
        const _exhaustive: never = raw;
        void _exhaustive;
      }
    }
  } catch (e: any) {
    const reqId = (raw as any)?.requestId ?? "unknown";
    sendError(reqId, "WORKER_ERROR", e?.message ?? String(e));
  }
});