src/runtime/WorkerClient.ts

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<string, Pending>();
  private readyPromise: Promise<void>;
  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<void>((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<InitMessage, "type">): Promise<void> {
    this.worker.postMessage({ type: "init", ...msg } satisfies InitMessage);
    await this.readyPromise;
  }

  private send<T extends WorkerMessage>(
    msg: ParentMessage & { requestId: string },
    expected: T["type"],
  ): Promise<T> {
    if (this.terminated) {
      return Promise.reject(new Error(`Worker terminated: ${this.terminatedReason}`));
    }
    return new Promise<T>((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<Extract<ParentMessage, { type: "exec" }>, "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<ExecResult>((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<Extract<ParentMessage, { type: "files_write" }>, "type" | "requestId">,
  ): Promise<FilesWriteResult> {
    const requestId = newId("req");
    return this.send<FilesWriteResult>(
      { type: "files_write", requestId, ...args } as ParentMessage & { requestId: string },
      "files_write_result",
    );
  }

  filesRead(
    args: Omit<Extract<ParentMessage, { type: "files_read" }>, "type" | "requestId">,
  ): Promise<FilesReadResult> {
    const requestId = newId("req");
    return this.send<FilesReadResult>(
      { type: "files_read", requestId, ...args } as ParentMessage & { requestId: string },
      "files_read_result",
    );
  }

  filesList(
    args: Omit<Extract<ParentMessage, { type: "files_list" }>, "type" | "requestId">,
  ): Promise<FilesListResult> {
    const requestId = newId("req");
    return this.send<FilesListResult>(
      { type: "files_list", requestId, ...args } as ParentMessage & { requestId: string },
      "files_list_result",
    );
  }

  fsBytes(): Promise<FsBytesResult> {
    const requestId = newId("req");
    return this.send<FsBytesResult>(
      { type: "fs_bytes", requestId } as ParentMessage & { requestId: string },
      "fs_bytes_result",
    );
  }

  async terminate(reason?: string): Promise<void> {
    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;