import { Worker } from "node:worker_threads"; import { fileURLToPath } from "node:url"; import type { ParentMessage, WorkerMessage, InitMessage, ExecResult, FilesWriteResult, FilesReadResult, FilesListResult, FsBytesResult, } from "../worker/protocol.js"; import { newId } from "../utils/ids.js"; const isCompiled = import.meta.url.endsWith(".js"); const WORKER_URL = isCompiled ? new URL("../worker/sessionWorker.js", import.meta.url) : new URL("../worker/bootstrap.mjs", import.meta.url); interface Pending { resolve: (value: WorkerMessage) => void; reject: (err: Error) => void; } export interface WorkerClientOptions { workerMaxOldGenMb: number; } export class WorkerClient { readonly id: string; private worker: Worker; private pending = new Map(); private readyPromise: Promise; private terminated = false; private terminatedReason: string | null = null; constructor(opts: WorkerClientOptions) { this.id = newId("w"); this.worker = new Worker(WORKER_URL, { resourceLimits: { maxOldGenerationSizeMb: opts.workerMaxOldGenMb, maxYoungGenerationSizeMb: Math.max(8, Math.floor(opts.workerMaxOldGenMb / 4)), }, }); let resolveReady!: () => void; let rejectReady!: (err: Error) => void; this.readyPromise = new Promise((resolve, reject) => { resolveReady = resolve; rejectReady = reject; }); this.worker.on("message", (msg: WorkerMessage) => { if (msg.type === "ready") { resolveReady(); return; } const reqId = (msg as { requestId?: string }).requestId; if (!reqId) return; const p = this.pending.get(reqId); if (!p) return; this.pending.delete(reqId); p.resolve(msg); }); this.worker.on("error", (err) => { rejectReady(err); this.failAllPending(err); }); this.worker.on("exit", (code) => { this.terminated = true; this.terminatedReason = `worker exited code=${code}`; rejectReady(new Error(`Worker exited before init code=${code}`)); this.failAllPending(new Error(`Worker exited code=${code}`)); }); } private failAllPending(err: Error) { for (const [id, p] of this.pending) { p.reject(err); this.pending.delete(id); } } isTerminated(): boolean { return this.terminated; } async init(msg: Omit): Promise { this.worker.postMessage({ type: "init", ...msg } satisfies InitMessage); await this.readyPromise; } private send( msg: ParentMessage & { requestId: string }, expected: T["type"], ): Promise { if (this.terminated) { return Promise.reject(new Error(`Worker terminated: ${this.terminatedReason}`)); } return new Promise((resolve, reject) => { this.pending.set(msg.requestId, { resolve: (m) => { if (m.type === "error_result") { const e = new Error(m.message); (e as any).code = m.code; reject(e); return; } if (m.type !== expected) { reject(new Error(`Unexpected worker reply type: ${m.type}`)); return; } resolve(m as T); }, reject, }); this.worker.postMessage(msg); }); } exec( args: Omit, "type" | "requestId">, hardTimeoutMs: number, ): Promise<{ result?: ExecResult; hardKilled: boolean; aborted: boolean; reason?: string }> { if (this.terminated) { return Promise.resolve({ hardKilled: false, aborted: true, reason: "worker_terminated" }); } const requestId = newId("req"); let watchdog: NodeJS.Timeout | null = null; const wait = new Promise((resolve, reject) => { this.pending.set(requestId, { resolve: (m) => { if (m.type === "error_result") { const e = new Error(m.message); (e as any).code = m.code; reject(e); return; } if (m.type !== "exec_result") { reject(new Error(`Unexpected reply: ${m.type}`)); return; } resolve(m); }, reject, }); this.worker.postMessage({ type: "exec", requestId, ...args }); }); let timedOutHard = false; const watchdogPromise = new Promise<{ hardKilled: true }>((resolve) => { watchdog = setTimeout(() => { timedOutHard = true; // Force terminate worker this.worker.terminate().catch(() => { /* ignored */ }); resolve({ hardKilled: true }); }, hardTimeoutMs); }); return Promise.race([ wait.then((res) => { if (watchdog) clearTimeout(watchdog); return { result: res, hardKilled: false, aborted: false }; }), watchdogPromise.then(() => ({ hardKilled: true, aborted: true, reason: "hard_timeout" })), ]).catch((err) => { if (watchdog) clearTimeout(watchdog); if (timedOutHard) { return { hardKilled: true, aborted: true, reason: "hard_timeout" }; } throw err; }); } filesWrite( args: Omit, "type" | "requestId">, ): Promise { const requestId = newId("req"); return this.send( { type: "files_write", requestId, ...args } as ParentMessage & { requestId: string }, "files_write_result", ); } filesRead( args: Omit, "type" | "requestId">, ): Promise { const requestId = newId("req"); return this.send( { type: "files_read", requestId, ...args } as ParentMessage & { requestId: string }, "files_read_result", ); } filesList( args: Omit, "type" | "requestId">, ): Promise { const requestId = newId("req"); return this.send( { type: "files_list", requestId, ...args } as ParentMessage & { requestId: string }, "files_list_result", ); } fsBytes(): Promise { const requestId = newId("req"); return this.send( { type: "fs_bytes", requestId } as ParentMessage & { requestId: string }, "fs_bytes_result", ); } async terminate(reason?: string): Promise { if (this.terminated) return; this.terminatedReason = reason ?? "terminated"; try { await this.worker.terminate(); } catch { /* ignored */ } this.terminated = true; this.failAllPending(new Error(`Worker terminated: ${this.terminatedReason}`)); } } // Reference fileURLToPath to satisfy ts-node-style imports (kept for future use) void fileURLToPath;