SnapOtter/packages/ai/src/bridge.ts

import { type ChildProcess, spawn } from "node:child_process";
import { randomUUID } from "node:crypto";
import { dirname, resolve } from "node:path";
import { fileURLToPath } from "node:url";

const __dirname = dirname(fileURLToPath(import.meta.url));
const PYTHON_DIR = resolve(__dirname, "../python");

/** Try venv first, then system python. */
function getPythonPath(): string {
  const venvPath = process.env.PYTHON_VENV_PATH || resolve(__dirname, "../../../.venv");
  return `${venvPath}/bin/python3`;
}

/**
 * Extract a user-friendly error from a Python process error.
 */
function extractPythonError(error: unknown): string {
  if (error && typeof error === "object") {
    const execError = error as {
      stderr?: string;
      stdout?: string;
      message?: string;
    };
    for (const output of [execError.stdout, execError.stderr]) {
      if (output) {
        try {
          const parsed = JSON.parse(output.trim());
          if (parsed.error) return parsed.error;
        } catch {
          const trimmed = output.trim();
          if (trimmed && !trimmed.startsWith("Traceback")) {
            return trimmed;
          }
          // Extract the last meaningful line from a Python Traceback
          if (trimmed) {
            const lines = trimmed
              .split("\n")
              .map((l) => l.trim())
              .filter(Boolean);
            const lastLine = lines[lines.length - 1];
            if (lastLine && lastLine !== "Traceback (most recent call last):") {
              return lastLine;
            }
          }
        }
      }
    }
    if (execError.message) return execError.message;
    // Return empty string for {stdout, stderr} objects with no useful content
    // so the caller's fallback message (e.g. exit code) kicks in.
    return "";
  }
  if (error instanceof Error) return error.message;
  return String(error);
}

export type ProgressCallback = (percent: number, stage: string) => void;

// ── Persistent dispatcher ───────────────────────────────────────────

interface PendingRequest {
  resolve: (result: { stdout: string; stderr: string }) => void;
  reject: (err: Error) => void;
  onProgress?: ProgressCallback;
  stderrLines: string[];
}

let dispatcher: ChildProcess | null = null;
let dispatcherReady = false;
let dispatcherFailed = false;
// biome-ignore lint/style/useConst: reassigned on dispatcher readiness signal
let dispatcherGpuAvailable = false;
const pendingRequests = new Map<string, PendingRequest>();
let stdoutBuffer = "";

// Crash recovery with exponential backoff
// biome-ignore lint/style/useConst: reassigned on crash events
let consecutiveCrashes = 0;
// biome-ignore lint/style/useConst: reassigned on crash events
let lastCrashTime = 0;
// biome-ignore lint/style/useConst: reassigned on crash events
let backoffUntil = 0;
const CRASH_WINDOW_MS = 60_000;
const MAX_CONSECUTIVE_CRASHES = 5;
const BASE_BACKOFF_MS = 1_000;

function recordCrash(): void {
  const now = Date.now();
  if (now - lastCrashTime > CRASH_WINDOW_MS) {
    consecutiveCrashes = 1;
  } else {
    consecutiveCrashes++;
  }
  lastCrashTime = now;

  if (consecutiveCrashes >= MAX_CONSECUTIVE_CRASHES) {
    console.error(
      `[bridge] Dispatcher crashed ${consecutiveCrashes} times in ${CRASH_WINDOW_MS / 1000}s, disabling permanently`,
    );
    dispatcherFailed = true;
    return;
  }

  const delay = BASE_BACKOFF_MS * 2 ** (consecutiveCrashes - 1);
  backoffUntil = now + delay;
  console.warn(
    `[bridge] Dispatcher crash #${consecutiveCrashes}, backing off ${delay}ms before restart`,
  );
}

function startDispatcher(): ChildProcess | null {
  if (dispatcherFailed) return null;

  try {
    const child = spawn(getPythonPath(), [resolve(PYTHON_DIR, "dispatcher.py")], {
      stdio: ["pipe", "pipe", "pipe"],
    });

    let stderrBuffer = "";

    child.stderr?.on("data", (chunk: Buffer) => {
      stderrBuffer += chunk.toString();
      const lines = stderrBuffer.split("\n");
      stderrBuffer = lines.pop() ?? "";

      for (const line of lines) {
        const trimmed = line.trim();
        if (!trimmed) continue;

        try {
          const parsed = JSON.parse(trimmed);

          // Readiness signal
          if (parsed.ready === true) {
            dispatcherReady = true;
            dispatcherGpuAvailable = parsed.gpu === true;
            consecutiveCrashes = 0;
            console.log(`[bridge] Python dispatcher ready (GPU: ${parsed.gpu === true})`);
            continue;
          }

          // Progress event - route to the currently active request
          if (typeof parsed.progress === "number" && typeof parsed.stage === "string") {
            // Progress goes to all pending requests (only one should be active at a time
            // since Python processes synchronously)
            for (const req of pendingRequests.values()) {
              req.onProgress?.(parsed.progress, parsed.stage);
            }
          }
        } catch {
          // Not JSON - forward diagnostic messages to Node.js logger,
          // collect the rest as error output for pending requests.
          if (trimmed.startsWith("[")) {
            console.log(`[python] ${trimmed}`);
          }
          for (const req of pendingRequests.values()) {
            req.stderrLines.push(trimmed);
          }
        }
      }
    });

    child.stdout?.on("data", (chunk: Buffer) => {
      stdoutBuffer += chunk.toString();
      const lines = stdoutBuffer.split("\n");
      stdoutBuffer = lines.pop() ?? "";

      for (const line of lines) {
        const trimmed = line.trim();
        if (!trimmed) continue;

        try {
          const response = JSON.parse(trimmed);
          const reqId = response.id;
          const pending = pendingRequests.get(reqId);
          if (pending) {
            pendingRequests.delete(reqId);
            if (response.exitCode !== 0) {
              const errText =
                extractPythonError({
                  stdout: response.stdout,
                  stderr: pending.stderrLines.join("\n"),
                }) ||
                (response.exitCode === 137
                  ? "Process killed (out of memory) — try a lighter model or smaller image"
                  : response.exitCode === 139
                    ? "Process crashed (segmentation fault)"
                    : `Python script exited with code ${response.exitCode}`);
              pending.reject(new Error(errText));
            } else {
              pending.resolve({
                stdout: response.stdout || "",
                stderr: pending.stderrLines.join("\n"),
              });
            }
          }
        } catch {
          // Not a valid response line
        }
      }
    });

    child.on("error", (err: NodeJS.ErrnoException) => {
      console.error(`[bridge] Dispatcher error: ${err.message} (code: ${err.code})`);
      if (err.code === "ENOENT") {
        dispatcherFailed = true;
      } else {
        recordCrash();
      }
      for (const [id, req] of pendingRequests.entries()) {
        req.reject(new Error(extractPythonError(err)));
        pendingRequests.delete(id);
      }
      dispatcher = null;
      dispatcherReady = false;
    });

    child.on("close", (code) => {
      for (const [id, req] of pendingRequests.entries()) {
        req.reject(new Error("Python dispatcher exited unexpectedly"));
        pendingRequests.delete(id);
      }
      if (code !== 0) {
        recordCrash();
      }
      dispatcher = null;
      dispatcherReady = false;
    });

    return child;
  } catch {
    dispatcherFailed = true;
    return null;
  }
}

function getDispatcher(): ChildProcess | null {
  if (dispatcherFailed) return null;
  if (!dispatcher || dispatcher.killed) {
    if (Date.now() < backoffUntil) return null;
    dispatcher = startDispatcher();
  }
  return dispatcher;
}

/**
 * Send a request to the persistent Python dispatcher.
 * Returns null if the dispatcher is unavailable (caller should fall back).
 */
function dispatcherRun(
  scriptName: string,
  args: string[],
  options: { onProgress?: ProgressCallback; timeout?: number } = {},
): Promise<{ stdout: string; stderr: string }> | null {
  const proc = getDispatcher();
  if (!proc || !proc.stdin || !dispatcherReady) return null;

  const id = randomUUID();
  const timeout =
    options.timeout ??
    (process.env.PROCESSING_TIMEOUT_S && parseInt(process.env.PROCESSING_TIMEOUT_S, 10) > 0
      ? parseInt(process.env.PROCESSING_TIMEOUT_S, 10) * 1000
      : 600000);

  return new Promise((resolvePromise, rejectPromise) => {
    const timer = setTimeout(() => {
      pendingRequests.delete(id);
      rejectPromise(new Error("Python script timed out"));
    }, timeout);

    const wrappedResolve = (result: { stdout: string; stderr: string }) => {
      clearTimeout(timer);
      resolvePromise(result);
    };

    const wrappedReject = (err: Error) => {
      clearTimeout(timer);
      rejectPromise(err);
    };

    pendingRequests.set(id, {
      resolve: wrappedResolve,
      reject: wrappedReject,
      onProgress: options.onProgress,
      stderrLines: [],
    });

    const request = JSON.stringify({ id, script: scriptName.replace(".py", ""), args });
    proc.stdin!.write(request + "\n");
  });
}

/**
 * Whether the Python dispatcher detected a CUDA GPU at startup.
 */
export function isGpuAvailable(): boolean {
  return dispatcherGpuAvailable;
}

export interface DispatcherStatus {
  running: boolean;
  ready: boolean;
  failed: boolean;
  gpu: boolean;
  pid: number | null;
  consecutiveCrashes: number;
}

export function getDispatcherStatus(): DispatcherStatus {
  return {
    running: dispatcher !== null && !dispatcher.killed,
    ready: dispatcherReady,
    failed: dispatcherFailed,
    gpu: dispatcherGpuAvailable,
    pid: dispatcher?.pid ?? null,
    consecutiveCrashes,
  };
}

/**
 * Shut down the persistent dispatcher process.
 */
export function shutdownDispatcher(): void {
  if (dispatcher && !dispatcher.killed) {
    dispatcher.stdin?.end();
    dispatcher.kill("SIGTERM");
    dispatcher = null;
    dispatcherReady = false;
  }
}

/**
 * Eagerly start the Python dispatcher and wait for its readiness signal.
 * Returns the GPU status once ready, or {ready: false} on timeout/failure.
 * Safe to call multiple times -- idempotent if the dispatcher is already running.
 */
export function initDispatcher(timeoutMs = 30_000): Promise<{ ready: boolean; gpu: boolean }> {
  if (dispatcherReady) {
    return Promise.resolve({ ready: true, gpu: dispatcherGpuAvailable });
  }
  if (dispatcherFailed) {
    return Promise.resolve({ ready: false, gpu: false });
  }

  const proc = getDispatcher();
  if (!proc) {
    return Promise.resolve({ ready: false, gpu: false });
  }

  return new Promise((resolve) => {
    const timer = setTimeout(() => {
      clearInterval(poll);
      resolve({ ready: false, gpu: false });
    }, timeoutMs);

    const poll = setInterval(() => {
      if (dispatcherReady) {
        clearTimeout(timer);
        clearInterval(poll);
        resolve({ ready: true, gpu: dispatcherGpuAvailable });
      } else if (dispatcherFailed) {
        clearTimeout(timer);
        clearInterval(poll);
        resolve({ ready: false, gpu: false });
      }
    }, 50);
  });
}

// ── Per-request fallback (original implementation) ──────────────────

function runPythonPerRequest(
  scriptName: string,
  args: string[],
  options: {
    onProgress?: ProgressCallback;
    timeout?: number;
  } = {},
): Promise<{ stdout: string; stderr: string }> {
  const scriptPath = resolve(PYTHON_DIR, scriptName);
  const timeout =
    options.timeout ??
    (process.env.PROCESSING_TIMEOUT_S && parseInt(process.env.PROCESSING_TIMEOUT_S, 10) > 0
      ? parseInt(process.env.PROCESSING_TIMEOUT_S, 10) * 1000
      : 600000);

  return new Promise((resolvePromise, rejectPromise) => {
    const trySpawn = (pythonBin: string, isFallback: boolean) => {
      const child = spawn(pythonBin, [scriptPath, ...args], {
        stdio: ["ignore", "pipe", "pipe"],
      });

      let stdout = "";
      const stderrLines: string[] = [];
      let stderrBuffer = "";
      let timedOut = false;

      const timer = setTimeout(() => {
        timedOut = true;
        child.kill("SIGTERM");
      }, timeout);

      child.stdout.on("data", (chunk: Buffer) => {
        stdout += chunk.toString();
      });

      child.stderr.on("data", (chunk: Buffer) => {
        stderrBuffer += chunk.toString();
        const lines = stderrBuffer.split("\n");
        stderrBuffer = lines.pop() ?? "";

        for (const line of lines) {
          const trimmed = line.trim();
          if (!trimmed) continue;

          try {
            const parsed = JSON.parse(trimmed);
            if (typeof parsed.progress === "number" && typeof parsed.stage === "string") {
              options.onProgress?.(parsed.progress, parsed.stage);
              continue;
            }
          } catch {
            // Not JSON - collect as regular stderr
          }
          stderrLines.push(trimmed);
        }
      });

      child.on("error", (err: NodeJS.ErrnoException) => {
        clearTimeout(timer);
        if (err.code === "ENOENT" && !isFallback) {
          trySpawn("python3", true);
        } else {
          rejectPromise(new Error(extractPythonError(err)));
        }
      });

      child.on("close", (code, signal) => {
        clearTimeout(timer);

        if (stderrBuffer.trim()) {
          stderrLines.push(stderrBuffer.trim());
        }

        if (timedOut) {
          rejectPromise(new Error("Python script timed out"));
          return;
        }

        const stderr = stderrLines.join("\n");

        if (code !== 0) {
          // When the process was killed by a signal, use a clear message
          // instead of surfacing unrelated stderr (e.g. CUDA warnings).
          const signalMsg =
            signal === "SIGKILL" || code === 137
              ? "Process killed (out of memory) — try a lighter model or smaller image"
              : signal === "SIGSEGV" || code === 139
                ? "Process crashed (segmentation fault)"
                : null;
          const errorText =
            signalMsg ||
            extractPythonError({ stdout: stdout.trim(), stderr }) ||
            `Python script exited with code ${code}`;
          rejectPromise(new Error(errorText));
          return;
        }

        resolvePromise({ stdout: stdout.trim(), stderr });
      });
    };

    trySpawn(getPythonPath(), false);
  });
}

// ── Public API (unchanged signature) ────────────────────────────────

/**
 * Run a Python script with real-time progress streaming via stderr.
 *
 * Tries the persistent dispatcher first for warm-start performance.
 * Falls back to per-request spawning if the dispatcher is unavailable.
 */
export function runPythonWithProgress(
  scriptName: string,
  args: string[],
  options: {
    onProgress?: ProgressCallback;
    timeout?: number;
  } = {},
): Promise<{ stdout: string; stderr: string }> {
  // Try persistent dispatcher first
  const dispatcherPromise = dispatcherRun(scriptName, args, options);
  if (dispatcherPromise) {
    return dispatcherPromise.catch((err: Error) => {
      if (err.message === "Python dispatcher exited unexpectedly") {
        console.warn(
          `[bridge] Dispatcher crashed during ${scriptName}, retrying with per-request process`,
        );
        return runPythonPerRequest(scriptName, args, options).then((result) => ({
          ...result,
          stderr: `${result.stderr}\n[bridge] retried after dispatcher crash`,
        }));
      }
      throw err;
    });
  }

  // Fall back to per-request spawning
  return runPythonPerRequest(scriptName, args, options);
}

// biome-ignore lint/suspicious/noExplicitAny: matches JSON.parse return type
export function parseStdoutJson(stdout: string): any {
  const match = stdout.match(/\{[\s\S]*\}$/);
  if (!match) throw new Error("No JSON response from Python script");
  return JSON.parse(match[0]);
}