assistme 0.7.0 → 0.8.2

package/src/workers/manager.ts

@@ -0,0 +1,506 @@
+/**
+ * WorkerManager — runs in the main process (orchestrator).
+ *
+ * Manages the lifecycle of worker child processes:
+ * - Spawns workers via child_process.fork()
+ * - Routes tasks to the appropriate worker (non-blocking)
+ * - Monitors worker health and restarts on crash
+ * - Tracks busy count via reference counting
+ * - Handles graceful shutdown
+ */
+
+import { fork, type ChildProcess } from "child_process";
+import { existsSync } from "fs";
+import { fileURLToPath } from "url";
+import { dirname, join } from "path";
+import { randomUUID } from "crypto";
+import type {
+  WorkerType,
+  WorkerConfig,
+  WorkerInfo,
+  WorkerStatus,
+  OrchestratorMessage,
+  WorkerMessage,
+} from "./types.js";
+import type { ConversationMessage } from "../db/types.js";
+import { log, setLogConversationId } from "../utils/logger.js";
+
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = dirname(__filename);
+
+/**
+ * Resolve the worker entry point path.
+ * tsup may bundle this into dist/index.js or dist/chunk-*.js,
+ * so we check both sibling and nested paths.
+ */
+function resolveWorkerEntry(): string {
+  const sibling = join(__dirname, "entry.js");
+  if (existsSync(sibling)) return sibling;
+
+  const nested = join(__dirname, "workers", "entry.js");
+  if (existsSync(nested)) return nested;
+
+  return nested;
+}
+
+const WORKER_ENTRY = resolveWorkerEntry();
+
+/** Auto-terminate idle conversation workers after this duration. */
+const IDLE_TIMEOUT_MS = 5 * 60_000; // 5 minutes
+
+interface ManagedWorker {
+  id: string;
+  type: WorkerType;
+  process: ChildProcess;
+  status: WorkerStatus;
+  config: WorkerConfig;
+  conversationId?: string;
+  startedAt: Date;
+  lastActivityAt: Date;
+  tasksProcessed: number;
+  idleTimer: ReturnType<typeof setTimeout> | null;
+  /** Resolves when the worker reports "ready" after init. Rejects if worker exits first. */
+  readyPromise: Promise<void>;
+  readyResolve: () => void;
+  readyReject: (err: Error) => void;
+  /** Task IDs currently being processed by this worker. */
+  pendingTaskIds: Set<string>;
+}
+
+export type BusyChangeCallback = (busyCount: number) => void;
+
+/** Pending completion tracker for a specific task. */
+interface TaskCompletionTracker {
+  resolve: () => void;
+  reject: (err: Error) => void;
+}
+
+export class WorkerManager {
+  private workers = new Map<string, ManagedWorker>();
+  /** Maps conversation ID → worker ID for fast lookup. */
+  private conversationWorkers = new Map<string, string>();
+  private userId: string;
+  private sessionId: string;
+  private running = true;
+
+  /** Number of conversation workers currently processing a task. */
+  private busyCount = 0;
+  /** Called whenever busyCount changes. */
+  private onBusyChange: BusyChangeCallback | null = null;
+
+  /** Maps task ID → completion promise, for callers that need to await a specific task. */
+  private taskCompletions = new Map<string, TaskCompletionTracker>();
+
+  constructor(userId: string, sessionId: string) {
+    this.userId = userId;
+    this.sessionId = sessionId;
+  }
+
+  /**
+   * Register a callback for busy count changes.
+   * Called with the new count whenever a worker starts or finishes a task.
+   */
+  setBusyChangeCallback(cb: BusyChangeCallback): void {
+    this.onBusyChange = cb;
+  }
+
+  // ── Public API ──────────────────────────────────────────────────
+
+  /**
+   * Dispatch a task to a conversation worker. Non-blocking — returns
+   * immediately after sending the task. The orchestrator can continue
+   * polling for more tasks.
+   */
+  async dispatchTask(task: ConversationMessage): Promise<void> {
+    const conversationId = task.conversation_id;
+    let worker = this.getConversationWorker(conversationId);
+
+    if (!worker || worker.status === "stopped" || worker.status === "error") {
+      worker = await this.spawnWorker("conversation", { conversationId });
+    }
+
+    // Wait for ready, but not for task completion
+    await worker.readyPromise;
+
+    // If this worker is still busy with a previous task, log a warning.
+    // The worker will queue it internally (single-threaded).
+    if (worker.status === "busy") {
+      log.warn(`[worker:${worker.id}] Still busy — task ${task.id.slice(0, 8)} will queue`);
+    }
+
+    this.sendToWorker(worker, { type: "process_task", task });
+  }
+
+  /**
+   * Dispatch a task and wait for completion.
+   * Use this when you need to know when the task finishes
+   * (e.g. job runs that need completion tracking).
+   */
+  async dispatchAndWait(task: ConversationMessage): Promise<void> {
+    const completionPromise = new Promise<void>((resolve, reject) => {
+      this.taskCompletions.set(task.id, { resolve, reject });
+    });
+
+    try {
+      await this.dispatchTask(task);
+    } catch (err) {
+      // Clean up tracker if dispatch fails
+      this.taskCompletions.delete(task.id);
+      throw err;
+    }
+
+    await completionPromise;
+  }
+
+  /**
+   * Graceful shutdown: stop all workers.
+   */
+  async shutdown(timeoutMs = 5_000): Promise<void> {
+    this.running = false;
+
+    const shutdownPromises = Array.from(this.workers.values()).map((worker) =>
+      this.stopWorker(worker, timeoutMs)
+    );
+
+    await Promise.allSettled(shutdownPromises);
+
+    // Reject any remaining task completions
+    const shutdownError = new Error("WorkerManager shutting down");
+    for (const tracker of this.taskCompletions.values()) {
+      tracker.reject(shutdownError);
+    }
+    this.taskCompletions.clear();
+
+    this.workers.clear();
+    this.conversationWorkers.clear();
+  }
+
+  /**
+   * Get info about all active workers.
+   */
+  getWorkerInfos(): WorkerInfo[] {
+    return Array.from(this.workers.values()).map((w) => ({
+      id: w.id,
+      type: w.type,
+      status: w.status,
+      pid: w.process.pid,
+      conversationId: w.conversationId,
+      startedAt: w.startedAt,
+      lastActivityAt: w.lastActivityAt,
+      tasksProcessed: w.tasksProcessed,
+    }));
+  }
+
+  /**
+   * Number of conversation workers currently processing a task.
+   */
+  getBusyCount(): number {
+    return this.busyCount;
+  }
+
+  // ── Worker Lifecycle ────────────────────────────────────────────
+
+  private async spawnWorker(
+    type: WorkerType,
+    opts?: { conversationId?: string }
+  ): Promise<ManagedWorker> {
+    const id = randomUUID().slice(0, 8);
+    const config: WorkerConfig = {
+      type,
+      userId: this.userId,
+      sessionId: this.sessionId,
+      conversationId: opts?.conversationId,
+    };
+
+    let readyResolve!: () => void;
+    let readyReject!: (err: Error) => void;
+    const readyPromise = new Promise<void>((resolve, reject) => {
+      readyResolve = resolve;
+      readyReject = reject;
+    });
+
+    const child = fork(WORKER_ENTRY, ["--type", type], {
+      stdio: ["pipe", "pipe", "pipe", "ipc"],
+      env: { ...process.env },
+    });
+
+    const worker: ManagedWorker = {
+      id,
+      type,
+      process: child,
+      status: "starting",
+      config,
+      conversationId: opts?.conversationId,
+      startedAt: new Date(),
+      lastActivityAt: new Date(),
+      tasksProcessed: 0,
+      idleTimer: null,
+      readyPromise,
+      readyResolve,
+      readyReject,
+      pendingTaskIds: new Set(),
+    };
+
+    this.workers.set(id, worker);
+
+    if (opts?.conversationId) {
+      this.conversationWorkers.set(opts.conversationId, id);
+    }
+
+    // Handle IPC messages from worker
+    child.on("message", (raw: unknown) => {
+      this.handleWorkerMessage(worker, raw as WorkerMessage);
+    });
+
+    // Handle worker exit
+    child.on("exit", (code, signal) => {
+      this.handleWorkerExit(worker, code, signal);
+    });
+
+    // Forward stderr for debugging
+    child.stderr?.on("data", (data: Buffer) => {
+      log.debug(`[worker:${id}] ${data.toString().trim()}`);
+    });
+
+    // Send init message
+    const initMsg: OrchestratorMessage = { type: "init", config };
+    child.send(initMsg);
+
+    log.info(`Worker spawned: ${type} (${id}, pid=${child.pid})`);
+    return worker;
+  }
+
+  private handleWorkerMessage(worker: ManagedWorker, message: WorkerMessage): void {
+    worker.lastActivityAt = new Date();
+
+    switch (message.type) {
+      case "ready":
+        worker.status = "ready";
+        worker.readyResolve();
+        break;
+
+      case "task_completed": {
+        worker.tasksProcessed++;
+        worker.pendingTaskIds.delete(message.taskId);
+        log.info(`[worker:${worker.id}] Task completed: ${message.taskId.slice(0, 8)}`);
+        this.workerBecameIdle(worker);
+        const completedTracker = this.taskCompletions.get(message.taskId);
+        if (completedTracker) {
+          this.taskCompletions.delete(message.taskId);
+          completedTracker.resolve();
+        }
+        break;
+      }
+
+      case "task_failed": {
+        worker.tasksProcessed++;
+        worker.pendingTaskIds.delete(message.taskId);
+        log.error(`[worker:${worker.id}] Task failed: ${message.error}`);
+        this.workerBecameIdle(worker);
+        const failedTracker = this.taskCompletions.get(message.taskId);
+        if (failedTracker) {
+          this.taskCompletions.delete(message.taskId);
+          failedTracker.reject(new Error(message.error));
+        }
+        break;
+      }
+
+      case "idle":
+        this.workerBecameIdle(worker);
+        break;
+
+      case "log":
+        switch (message.level) {
+          case "error":
+            log.error(`[worker:${worker.id}] ${message.message}`);
+            break;
+          case "warn":
+            log.warn(`[worker:${worker.id}] ${message.message}`);
+            break;
+          case "info":
+            log.info(`[worker:${worker.id}] ${message.message}`);
+            break;
+          default:
+            log.debug(`[worker:${worker.id}] ${message.message}`);
+        }
+        break;
+
+      case "log_forward": {
+        // Replay worker process logs through the main logger (with conversation context).
+        // This feeds both the console output and the SessionLogEmitter (Supabase persistence).
+        // Set conversationId context so logHook passes it to SessionLogEmitter.
+        const convId = worker.conversationId || null;
+        setLogConversationId(convId);
+
+        const convTag = convId
+          ? `[conv:${convId.slice(0, 8)}]`
+          : `[worker:${worker.id}]`;
+        switch (message.method) {
+          case "agent":
+            log.agent(`${convTag} ${message.message}`);
+            break;
+          case "tool":
+            log.tool(message.extra || "unknown", `${convTag} ${message.message}`);
+            break;
+          case "result":
+            log.result(`${convTag} ${message.message}`);
+            break;
+          case "success":
+            log.success(`${convTag} ${message.message}`);
+            break;
+          case "error":
+            log.error(`${convTag} ${message.message}`);
+            break;
+          case "warn":
+            log.warn(`${convTag} ${message.message}`);
+            break;
+          case "info":
+            log.info(`${convTag} ${message.message}`);
+            break;
+          default:
+            log.debug(`${convTag} ${message.message}`);
+        }
+
+        setLogConversationId(null);
+        break;
+      }
+
+      case "error":
+        log.error(`[worker:${worker.id}] Error: ${message.error}`);
+        break;
+    }
+  }
+
+  /**
+   * Transition a worker from busy → idle and update busy count.
+   */
+  private workerBecameIdle(worker: ManagedWorker): void {
+    if (worker.status === "busy" && worker.type === "conversation") {
+      this.busyCount = Math.max(0, this.busyCount - 1);
+      this.onBusyChange?.(this.busyCount);
+    }
+    worker.status = "idle";
+    this.scheduleIdleTimeout(worker);
+  }
+
+  /**
+   * Send a message to a worker and mark it busy (for conversation workers).
+   */
+  private sendToWorker(worker: ManagedWorker, message: OrchestratorMessage): void {
+    if (worker.type === "conversation" && worker.status !== "busy") {
+      worker.status = "busy";
+      this.busyCount++;
+      this.onBusyChange?.(this.busyCount);
+    }
+
+    // Track task ID for cleanup on worker exit
+    if (message.type === "process_task") {
+      worker.pendingTaskIds.add(message.task.id);
+    }
+
+    worker.lastActivityAt = new Date();
+
+    // Cancel idle timer
+    if (worker.idleTimer) {
+      clearTimeout(worker.idleTimer);
+      worker.idleTimer = null;
+    }
+
+    worker.process.send(message);
+  }
+
+  private handleWorkerExit(
+    worker: ManagedWorker,
+    code: number | null,
+    signal: string | null
+  ): void {
+    // If worker was busy, decrement count
+    if (worker.status === "busy" && worker.type === "conversation") {
+      this.busyCount = Math.max(0, this.busyCount - 1);
+      this.onBusyChange?.(this.busyCount);
+    }
+
+    worker.status = "stopped";
+
+    if (worker.idleTimer) {
+      clearTimeout(worker.idleTimer);
+      worker.idleTimer = null;
+    }
+
+    if (worker.conversationId) {
+      this.conversationWorkers.delete(worker.conversationId);
+    }
+
+    if (code !== 0 && code !== null) {
+      log.warn(`Worker ${worker.id} (${worker.type}) exited with code ${code}`);
+      worker.status = "error";
+    } else if (signal) {
+      log.debug(`Worker ${worker.id} (${worker.type}) killed by signal ${signal}`);
+    } else {
+      log.debug(`Worker ${worker.id} (${worker.type}) exited normally`);
+    }
+
+    // Reject readyPromise if worker exited before becoming ready
+    const exitError = new Error(`Worker ${worker.id} exited (code=${code}, signal=${signal})`);
+    worker.readyReject(exitError);
+
+    // Reject all pending task completions for this worker
+    for (const taskId of worker.pendingTaskIds) {
+      const tracker = this.taskCompletions.get(taskId);
+      if (tracker) {
+        this.taskCompletions.delete(taskId);
+        tracker.reject(exitError);
+      }
+    }
+    worker.pendingTaskIds.clear();
+
+    this.workers.delete(worker.id);
+  }
+
+  private async stopWorker(worker: ManagedWorker, timeoutMs: number): Promise<void> {
+    if (worker.status === "stopped") return;
+
+    worker.status = "stopping";
+
+    if (worker.idleTimer) {
+      clearTimeout(worker.idleTimer);
+      worker.idleTimer = null;
+    }
+
+    try {
+      const shutdownMsg: OrchestratorMessage = { type: "shutdown" };
+      worker.process.send(shutdownMsg);
+
+      await Promise.race([
+        new Promise<void>((resolve) => {
+          worker.process.on("exit", () => resolve());
+        }),
+        new Promise<void>((_, reject) =>
+          setTimeout(() => reject(new Error("Shutdown timeout")), timeoutMs)
+        ),
+      ]);
+    } catch {
+      worker.process.kill("SIGKILL");
+    }
+  }
+
+  private scheduleIdleTimeout(worker: ManagedWorker): void {
+    if (worker.type !== "conversation" || !this.running) return;
+
+    if (worker.idleTimer) {
+      clearTimeout(worker.idleTimer);
+    }
+
+    worker.idleTimer = setTimeout(() => {
+      if (worker.status === "idle" || worker.status === "ready") {
+        log.debug(`Worker ${worker.id} idle timeout — terminating`);
+        this.stopWorker(worker, 3_000).catch(() => {});
+      }
+    }, IDLE_TIMEOUT_MS);
+  }
+
+  private getConversationWorker(conversationId: string): ManagedWorker | undefined {
+    const workerId = this.conversationWorkers.get(conversationId);
+    return workerId ? this.workers.get(workerId) : undefined;
+  }
+}

package/src/workers/types.ts

@@ -0,0 +1,61 @@
+/**
+ * IPC message protocol between orchestrator (main process) and workers (child processes).
+ *
+ * Workers are spawned via child_process.fork() and communicate via JSON messages.
+ */
+
+import type { ConversationMessage } from "../db/types.js";
+
+// ── Worker Types ──────────────────────────────────────────────────
+
+export type WorkerType = "conversation";
+
+export interface WorkerConfig {
+  type: WorkerType;
+  userId: string;
+  sessionId: string;
+  /** For conversation workers: the conversation this worker handles. */
+  conversationId?: string;
+}
+
+// ── Orchestrator → Worker Messages ────────────────────────────────
+
+export type OrchestratorMessage =
+  | { type: "init"; config: WorkerConfig }
+  | { type: "process_task"; task: ConversationMessage }
+  | { type: "shutdown" };
+
+// ── Worker → Orchestrator Messages ────────────────────────────────
+
+export type LogMethod = "debug" | "info" | "success" | "warn" | "error" | "agent" | "tool" | "result";
+
+export type WorkerMessage =
+  | { type: "ready" }
+  | { type: "task_completed"; taskId: string }
+  | { type: "task_failed"; taskId: string; error: string }
+  | { type: "error"; error: string }
+  | { type: "idle" }
+  | { type: "log"; level: string; message: string }
+  | { type: "log_forward"; method: LogMethod; message: string; extra?: string };
+
+// ── Worker Status ─────────────────────────────────────────────────
+
+export type WorkerStatus =
+  | "starting"
+  | "ready"
+  | "busy"
+  | "idle"
+  | "stopping"
+  | "stopped"
+  | "error";
+
+export interface WorkerInfo {
+  id: string;
+  type: WorkerType;
+  status: WorkerStatus;
+  pid?: number;
+  conversationId?: string;
+  startedAt: Date;
+  lastActivityAt: Date;
+  tasksProcessed: number;
+}