@donkeylabs/server 0.4.7 → 0.4.8

package/src/core/jobs.ts

@@ -1,7 +1,30 @@
 // Core Jobs Service
 // Background job queue with scheduling
+// Supports both in-process handlers and external processes (Python, Go, Shell, etc.)
 
 import type { Events } from "./events";
+import type {
+  ExternalJobConfig,
+  ExternalJob,
+  ExternalJobProcessState,
+  ExternalJobsConfig,
+  AnyExternalJobMessage,
+} from "./external-jobs";
+import {
+  isProcessAlive,
+  isExternalJob,
+  isProgressMessage,
+  isHeartbeatMessage,
+  isLogMessage,
+  isCompletedMessage,
+  isFailedMessage,
+  isStartedMessage,
+  createInitialPayload,
+} from "./external-jobs";
+import {
+  createExternalJobSocketServer,
+  type ExternalJobSocketServer,
+} from "./external-job-socket";
 
 export type JobStatus = "pending" | "running" | "completed" | "failed" | "scheduled";
 
@@ -18,6 +41,19 @@ export interface Job {
   error?: string;
   attempts: number;
   maxAttempts: number;
+  // External job fields (null/undefined for in-process jobs)
+  /** Flag indicating this is an external job */
+  external?: boolean;
+  /** Process ID of the external process */
+  pid?: number;
+  /** Unix socket path for communication */
+  socketPath?: string;
+  /** TCP port for Windows fallback */
+  tcpPort?: number;
+  /** Timestamp of last heartbeat */
+  lastHeartbeat?: Date;
+  /** Current process state */
+  processState?: ExternalJobProcessState;
 }
 
 export interface JobHandler<T = any, R = any> {
@@ -32,6 +68,10 @@ export interface JobAdapter {
   getPending(limit?: number): Promise<Job[]>;
   getScheduledReady(now: Date): Promise<Job[]>;
   getByName(name: string, status?: JobStatus): Promise<Job[]>;
+  /** Get all running external jobs */
+  getRunningExternal(): Promise<Job[]>;
+  /** Get external jobs that need reconnection after server restart */
+  getOrphanedExternal(): Promise<Job[]>;
 }
 
 export interface JobsConfig {
@@ -40,16 +80,30 @@ export interface JobsConfig {
   concurrency?: number; // Max concurrent jobs, default 5
   pollInterval?: number; // ms, default 1000
   maxAttempts?: number; // Default retry attempts, default 3
+  /** External jobs configuration */
+  external?: ExternalJobsConfig;
 }
 
 export interface Jobs {
+  /** Register an in-process job handler */
   register<T = any, R = any>(name: string, handler: JobHandler<T, R>): void;
+  /** Register an external job (Python, Go, Shell, etc.) */
+  registerExternal(name: string, config: ExternalJobConfig): void;
+  /** Enqueue a job (works for both in-process and external jobs) */
   enqueue<T = any>(name: string, data: T, options?: { maxAttempts?: number }): Promise<string>;
+  /** Schedule a job to run at a specific time */
   schedule<T = any>(name: string, data: T, runAt: Date, options?: { maxAttempts?: number }): Promise<string>;
+  /** Get a job by ID */
   get(jobId: string): Promise<Job | null>;
+  /** Cancel a pending job */
   cancel(jobId: string): Promise<boolean>;
+  /** Get jobs by name and optionally filter by status */
   getByName(name: string, status?: JobStatus): Promise<Job[]>;
+  /** Get all running external jobs */
+  getRunningExternal(): Promise<Job[]>;
+  /** Start the job processing loop */
   start(): void;
+  /** Stop the job processing and cleanup */
   stop(): Promise<void>;
 }
 
@@ -110,6 +164,26 @@ export class MemoryJobAdapter implements JobAdapter {
     }
     return results;
   }
+
+  async getRunningExternal(): Promise<Job[]> {
+    const results: Job[] = [];
+    for (const job of this.jobs.values()) {
+      if (job.external && job.status === "running") {
+        results.push(job);
+      }
+    }
+    return results;
+  }
+
+  async getOrphanedExternal(): Promise<Job[]> {
+    const results: Job[] = [];
+    for (const job of this.jobs.values()) {
+      if (job.external && (job.processState === "orphaned" || job.processState === "running")) {
+        results.push(job);
+      }
+    }
+    return results;
+  }
 }
 
 class JobsImpl implements Jobs {
@@ -118,28 +192,49 @@ class JobsImpl implements Jobs {
   private handlers = new Map<string, JobHandler>();
   private running = false;
   private timer: ReturnType<typeof setInterval> | null = null;
+  private heartbeatTimer: ReturnType<typeof setInterval> | null = null;
   private activeJobs = 0;
   private concurrency: number;
   private pollInterval: number;
   private defaultMaxAttempts: number;
 
+  // External jobs support
+  private externalConfigs = new Map<string, ExternalJobConfig>();
+  private externalConfig: ExternalJobsConfig;
+  private socketServer: ExternalJobSocketServer | null = null;
+  private externalProcesses = new Map<string, { pid: number; timeout?: ReturnType<typeof setTimeout> }>();
+
   constructor(config: JobsConfig = {}) {
     this.adapter = config.adapter ?? new MemoryJobAdapter();
     this.events = config.events;
     this.concurrency = config.concurrency ?? 5;
     this.pollInterval = config.pollInterval ?? 1000;
     this.defaultMaxAttempts = config.maxAttempts ?? 3;
+    this.externalConfig = config.external ?? {};
   }
 
   register<T = any, R = any>(name: string, handler: JobHandler<T, R>): void {
-    if (this.handlers.has(name)) {
+    if (this.handlers.has(name) || this.externalConfigs.has(name)) {
       throw new Error(`Job handler "${name}" is already registered`);
     }
     this.handlers.set(name, handler);
   }
 
+  registerExternal(name: string, config: ExternalJobConfig): void {
+    if (this.handlers.has(name) || this.externalConfigs.has(name)) {
+      throw new Error(`Job handler "${name}" is already registered`);
+    }
+    this.externalConfigs.set(name, config);
+  }
+
+  private isExternalJob(name: string): boolean {
+    return this.externalConfigs.has(name);
+  }
+
   async enqueue<T = any>(name: string, data: T, options: { maxAttempts?: number } = {}): Promise<string> {
-    if (!this.handlers.has(name)) {
+    const isExternal = this.isExternalJob(name);
+
+    if (!isExternal && !this.handlers.has(name)) {
       throw new Error(`No handler registered for job "${name}"`);
     }
 
@@ -150,6 +245,8 @@ class JobsImpl implements Jobs {
       createdAt: new Date(),
       attempts: 0,
       maxAttempts: options.maxAttempts ?? this.defaultMaxAttempts,
+      external: isExternal || undefined,
+      processState: isExternal ? "spawning" : undefined,
     });
 
     return job.id;
@@ -161,7 +258,9 @@ class JobsImpl implements Jobs {
     runAt: Date,
     options: { maxAttempts?: number } = {}
   ): Promise<string> {
-    if (!this.handlers.has(name)) {
+    const isExternal = this.isExternalJob(name);
+
+    if (!isExternal && !this.handlers.has(name)) {
       throw new Error(`No handler registered for job "${name}"`);
     }
 
@@ -173,6 +272,8 @@ class JobsImpl implements Jobs {
       runAt,
       attempts: 0,
       maxAttempts: options.maxAttempts ?? this.defaultMaxAttempts,
+      external: isExternal || undefined,
+      processState: isExternal ? "spawning" : undefined,
     });
 
     return job.id;
@@ -187,8 +288,18 @@ class JobsImpl implements Jobs {
     if (!job) return false;
 
     if (job.status === "running") {
-      // Can't cancel running job
-      return false;
+      // For external jobs, we can try to kill the process
+      if (job.external && job.pid) {
+        try {
+          process.kill(job.pid, "SIGTERM");
+        } catch {
+          // Process may already be dead
+        }
+        await this.cleanupExternalJob(jobId);
+      } else {
+        // Can't cancel running in-process job
+        return false;
+      }
     }
 
     return this.adapter.delete(jobId);
@@ -198,10 +309,22 @@ class JobsImpl implements Jobs {
     return this.adapter.getByName(name, status);
   }
 
+  async getRunningExternal(): Promise<Job[]> {
+    return this.adapter.getRunningExternal();
+  }
+
   start(): void {
     if (this.running) return;
     this.running = true;
 
+    // Initialize socket server for external jobs
+    if (this.externalConfigs.size > 0) {
+      this.initializeSocketServer();
+      this.startHeartbeatMonitor();
+      // Attempt to reconnect to orphaned jobs from previous run
+      this.reconnectOrphanedJobs();
+    }
+
     this.timer = setInterval(() => this.tick(), this.pollInterval);
     // Run immediately too
     this.tick();
@@ -209,19 +332,322 @@ class JobsImpl implements Jobs {
 
   async stop(): Promise<void> {
     this.running = false;
+
     if (this.timer) {
       clearInterval(this.timer);
       this.timer = null;
     }
 
-    // Wait for active jobs to complete (with timeout)
+    if (this.heartbeatTimer) {
+      clearInterval(this.heartbeatTimer);
+      this.heartbeatTimer = null;
+    }
+
+    // Cleanup external job processes
+    for (const [jobId, procInfo] of this.externalProcesses) {
+      if (procInfo.timeout) {
+        clearTimeout(procInfo.timeout);
+      }
+      try {
+        process.kill(procInfo.pid, "SIGTERM");
+      } catch {
+        // Process may already be dead
+      }
+    }
+    this.externalProcesses.clear();
+
+    // Shutdown socket server
+    if (this.socketServer) {
+      await this.socketServer.shutdown();
+      this.socketServer = null;
+    }
+
+    // Wait for active in-process jobs to complete (with timeout)
     const maxWait = 30000; // 30 seconds
-    const start = Date.now();
-    while (this.activeJobs > 0 && Date.now() - start < maxWait) {
+    const startTime = Date.now();
+    while (this.activeJobs > 0 && Date.now() - startTime < maxWait) {
       await new Promise(resolve => setTimeout(resolve, 100));
     }
   }
 
+  private initializeSocketServer(): void {
+    this.socketServer = createExternalJobSocketServer(this.externalConfig, {
+      onMessage: (msg) => this.handleExternalMessage(msg),
+      onConnect: (jobId) => {
+        console.log(`[Jobs] External job ${jobId} connected`);
+      },
+      onDisconnect: (jobId) => {
+        console.log(`[Jobs] External job ${jobId} disconnected`);
+        // Check if the job is still running and mark as orphaned if so
+        this.handleExternalDisconnect(jobId);
+      },
+      onError: (err, jobId) => {
+        console.error(`[Jobs] External job socket error:`, err, jobId);
+      },
+    });
+  }
+
+  private startHeartbeatMonitor(): void {
+    const checkInterval = this.externalConfig.heartbeatCheckInterval ?? 10000;
+
+    this.heartbeatTimer = setInterval(async () => {
+      await this.checkHeartbeats();
+    }, checkInterval);
+  }
+
+  private async checkHeartbeats(): Promise<void> {
+    try {
+      const runningExternal = await this.adapter.getRunningExternal();
+      const now = Date.now();
+
+      for (const job of runningExternal) {
+        if (!job.external || !job.lastHeartbeat) continue;
+
+        const config = this.externalConfigs.get(job.name);
+        const heartbeatTimeout = config?.heartbeatTimeout ?? this.externalConfig.defaultHeartbeatTimeout ?? 30000;
+        const timeSinceHeartbeat = now - job.lastHeartbeat.getTime();
+
+        if (timeSinceHeartbeat > heartbeatTimeout) {
+          // Job is stale
+          console.warn(`[Jobs] External job ${job.id} is stale (no heartbeat for ${timeSinceHeartbeat}ms)`);
+
+          if (this.events) {
+            await this.events.emit("job.stale", {
+              jobId: job.id,
+              name: job.name,
+              timeSinceHeartbeat,
+            });
+          }
+
+          // If stale for 2x timeout, kill the process
+          if (timeSinceHeartbeat > heartbeatTimeout * 2) {
+            console.error(`[Jobs] Killing stale external job ${job.id}`);
+
+            if (job.pid) {
+              try {
+                process.kill(job.pid, "SIGKILL");
+              } catch {
+                // Process may already be dead
+              }
+            }
+
+            await this.adapter.update(job.id, {
+              status: "failed",
+              error: "Heartbeat timeout - job process unresponsive",
+              completedAt: new Date(),
+              processState: "orphaned",
+            });
+
+            await this.cleanupExternalJob(job.id);
+
+            if (this.events) {
+              await this.events.emit("job.failed", {
+                jobId: job.id,
+                name: job.name,
+                error: "Heartbeat timeout",
+              });
+            }
+          }
+        }
+      }
+    } catch (err) {
+      console.error("[Jobs] Heartbeat check error:", err);
+    }
+  }
+
+  private async reconnectOrphanedJobs(): Promise<void> {
+    try {
+      const orphaned = await this.adapter.getOrphanedExternal();
+      const activeJobIds = new Set<string>();
+
+      for (const job of orphaned) {
+        if (!job.pid) {
+          // No PID, mark as failed
+          await this.adapter.update(job.id, {
+            status: "failed",
+            error: "Lost job state - no PID available",
+            completedAt: new Date(),
+          });
+          continue;
+        }
+
+        // Check if process is still alive
+        if (isProcessAlive(job.pid)) {
+          console.log(`[Jobs] Found orphaned job ${job.id} with PID ${job.pid}, attempting reconnect`);
+          activeJobIds.add(job.id);
+
+          // Try to reconnect to the socket
+          const reconnected = await this.socketServer?.reconnect(
+            job.id,
+            job.socketPath,
+            job.tcpPort
+          );
+
+          if (reconnected) {
+            await this.adapter.update(job.id, {
+              processState: "running",
+              lastHeartbeat: new Date(),
+            });
+
+            if (this.events) {
+              await this.events.emit("job.reconnected", {
+                jobId: job.id,
+                name: job.name,
+              });
+            }
+          } else {
+            // Mark as orphaned, but keep tracking
+            await this.adapter.update(job.id, { processState: "orphaned" });
+
+            if (this.events) {
+              await this.events.emit("job.lost", {
+                jobId: job.id,
+                name: job.name,
+              });
+            }
+          }
+        } else {
+          // Process is dead, mark job as failed
+          console.log(`[Jobs] Orphaned job ${job.id} process (PID ${job.pid}) is dead`);
+          await this.adapter.update(job.id, {
+            status: "failed",
+            error: "Process died unexpectedly",
+            completedAt: new Date(),
+          });
+
+          if (this.events) {
+            await this.events.emit("job.failed", {
+              jobId: job.id,
+              name: job.name,
+              error: "Process died unexpectedly",
+            });
+          }
+        }
+      }
+
+      // Clean orphaned socket files
+      await this.socketServer?.cleanOrphanedSockets(activeJobIds);
+    } catch (err) {
+      console.error("[Jobs] Orphan reconnection error:", err);
+    }
+  }
+
+  private async handleExternalMessage(message: AnyExternalJobMessage): Promise<void> {
+    const job = await this.adapter.get(message.jobId);
+    if (!job) {
+      console.warn(`[Jobs] Received message for unknown job: ${message.jobId}`);
+      return;
+    }
+
+    if (isStartedMessage(message)) {
+      await this.adapter.update(message.jobId, {
+        processState: "running",
+        lastHeartbeat: new Date(message.timestamp),
+      });
+
+      if (this.events) {
+        await this.events.emit("job.external.spawned", {
+          jobId: message.jobId,
+          name: job.name,
+        });
+      }
+    } else if (isHeartbeatMessage(message)) {
+      await this.adapter.update(message.jobId, {
+        lastHeartbeat: new Date(message.timestamp),
+      });
+    } else if (isProgressMessage(message)) {
+      await this.adapter.update(message.jobId, {
+        lastHeartbeat: new Date(message.timestamp),
+      });
+
+      if (this.events) {
+        await this.events.emit("job.external.progress", {
+          jobId: message.jobId,
+          name: job.name,
+          percent: message.percent,
+          message: message.message,
+          data: message.data,
+        });
+      }
+    } else if (isLogMessage(message)) {
+      await this.adapter.update(message.jobId, {
+        lastHeartbeat: new Date(message.timestamp),
+      });
+
+      if (this.events) {
+        await this.events.emit("job.external.log", {
+          jobId: message.jobId,
+          name: job.name,
+          level: message.level,
+          message: message.message,
+          data: message.data,
+        });
+      }
+    } else if (isCompletedMessage(message)) {
+      await this.adapter.update(message.jobId, {
+        status: "completed",
+        result: message.result,
+        completedAt: new Date(message.timestamp),
+      });
+
+      await this.cleanupExternalJob(message.jobId);
+
+      if (this.events) {
+        await this.events.emit("job.completed", {
+          jobId: message.jobId,
+          name: job.name,
+          result: message.result,
+        });
+        await this.events.emit(`job.${job.name}.completed`, {
+          jobId: message.jobId,
+          result: message.result,
+        });
+      }
+    } else if (isFailedMessage(message)) {
+      await this.adapter.update(message.jobId, {
+        status: "failed",
+        error: message.error,
+        completedAt: new Date(message.timestamp),
+      });
+
+      await this.cleanupExternalJob(message.jobId);
+
+      if (this.events) {
+        await this.events.emit("job.failed", {
+          jobId: message.jobId,
+          name: job.name,
+          error: message.error,
+          stack: message.stack,
+        });
+        await this.events.emit(`job.${job.name}.failed`, {
+          jobId: message.jobId,
+          error: message.error,
+          stack: message.stack,
+        });
+      }
+    }
+  }
+
+  private async handleExternalDisconnect(jobId: string): Promise<void> {
+    const job = await this.adapter.get(jobId);
+    if (!job || job.status !== "running") return;
+
+    // Mark as orphaned - the job might still be running
+    await this.adapter.update(jobId, { processState: "orphaned" });
+  }
+
+  private async cleanupExternalJob(jobId: string): Promise<void> {
+    // Clear any timeout
+    const procInfo = this.externalProcesses.get(jobId);
+    if (procInfo?.timeout) {
+      clearTimeout(procInfo.timeout);
+    }
+    this.externalProcesses.delete(jobId);
+
+    // Close the socket
+    await this.socketServer?.closeSocket(jobId);
+  }
+
   private async tick(): Promise<void> {
     if (!this.running) return;
 
@@ -240,13 +666,230 @@ class JobsImpl implements Jobs {
       const pending = await this.adapter.getPending(availableSlots);
       for (const job of pending) {
         if (this.activeJobs >= this.concurrency) break;
-        this.processJob(job);
+
+        if (job.external) {
+          this.processExternalJob(job);
+        } else {
+          this.processJob(job);
+        }
       }
     } catch (err) {
       console.error("[Jobs] Tick error:", err);
     }
   }
 
+  private async processExternalJob(job: Job): Promise<void> {
+    const config = this.externalConfigs.get(job.name);
+    if (!config) {
+      await this.adapter.update(job.id, {
+        status: "failed",
+        error: `No external config registered for job "${job.name}"`,
+        completedAt: new Date(),
+      });
+      return;
+    }
+
+    this.activeJobs++;
+    const startedAt = new Date();
+
+    try {
+      // Create socket for this job
+      if (!this.socketServer) {
+        this.initializeSocketServer();
+      }
+
+      const { socketPath, tcpPort } = await this.socketServer!.createSocket(job.id);
+
+      // Update job with socket info
+      await this.adapter.update(job.id, {
+        status: "running",
+        startedAt,
+        attempts: job.attempts + 1,
+        socketPath,
+        tcpPort,
+        processState: "spawning",
+        lastHeartbeat: startedAt,
+      });
+
+      // Create initial payload
+      const payload = createInitialPayload(
+        job.id,
+        job.name,
+        job.data,
+        socketPath ?? `tcp://127.0.0.1:${tcpPort}`
+      );
+
+      // Spawn the external process
+      const env = {
+        ...process.env,
+        ...config.env,
+        DONKEYLABS_JOB_ID: job.id,
+        DONKEYLABS_JOB_NAME: job.name,
+        DONKEYLABS_SOCKET_PATH: socketPath ?? "",
+        DONKEYLABS_TCP_PORT: String(tcpPort ?? ""),
+      };
+
+      const proc = Bun.spawn([config.command, ...(config.args ?? [])], {
+        cwd: config.cwd,
+        env,
+        stdin: "pipe",
+        stdout: "pipe",
+        stderr: "pipe",
+      });
+
+      // Store process info
+      this.externalProcesses.set(job.id, { pid: proc.pid });
+
+      // Update job with PID
+      await this.adapter.update(job.id, { pid: proc.pid });
+
+      // Send payload via stdin
+      proc.stdin.write(payload + "\n");
+      proc.stdin.end();
+
+      // Set up process timeout if configured
+      if (config.timeout) {
+        const timeout = setTimeout(async () => {
+          console.warn(`[Jobs] External job ${job.id} timed out after ${config.timeout}ms`);
+          try {
+            process.kill(proc.pid, "SIGTERM");
+          } catch {
+            // Process may already be dead
+          }
+
+          await this.adapter.update(job.id, {
+            status: "failed",
+            error: `Job timed out after ${config.timeout}ms`,
+            completedAt: new Date(),
+          });
+
+          await this.cleanupExternalJob(job.id);
+
+          if (this.events) {
+            await this.events.emit("job.failed", {
+              jobId: job.id,
+              name: job.name,
+              error: "Timeout",
+            });
+          }
+        }, config.timeout);
+
+        const procInfo = this.externalProcesses.get(job.id);
+        if (procInfo) {
+          procInfo.timeout = timeout;
+        }
+      }
+
+      // Handle process exit
+      proc.exited.then(async (code) => {
+        // Only handle exit if job is still running (not already completed/failed via message)
+        const currentJob = await this.adapter.get(job.id);
+        if (currentJob?.status === "running") {
+          if (code === 0) {
+            // Process exited cleanly but didn't send completion message
+            // This might be ok, or might indicate an issue
+            console.warn(`[Jobs] External job ${job.id} exited with code 0 but no completion message`);
+          } else {
+            // Process failed
+            await this.adapter.update(job.id, {
+              status: "failed",
+              error: `Process exited with code ${code}`,
+              completedAt: new Date(),
+            });
+
+            if (this.events) {
+              await this.events.emit("job.failed", {
+                jobId: job.id,
+                name: job.name,
+                error: `Process exited with code ${code}`,
+              });
+            }
+          }
+
+          await this.cleanupExternalJob(job.id);
+        }
+
+        this.activeJobs--;
+      });
+
+      // Stream stdout/stderr to logs
+      this.streamProcessOutput(job.id, job.name, proc);
+    } catch (err) {
+      const error = err instanceof Error ? err.message : String(err);
+
+      await this.adapter.update(job.id, {
+        status: "failed",
+        error,
+        completedAt: new Date(),
+      });
+
+      await this.cleanupExternalJob(job.id);
+
+      if (this.events) {
+        await this.events.emit("job.failed", {
+          jobId: job.id,
+          name: job.name,
+          error,
+        });
+      }
+
+      this.activeJobs--;
+    }
+  }
+
+  private streamProcessOutput(
+    jobId: string,
+    jobName: string,
+    proc: ReturnType<typeof Bun.spawn>
+  ): void {
+    const decoder = new TextDecoder();
+    const events = this.events;
+
+    // Helper to stream a ReadableStream
+    const streamOutput = async (
+      stream: ReadableStream<Uint8Array> | undefined,
+      level: "info" | "error"
+    ) => {
+      if (!stream) return;
+
+      try {
+        const reader = stream.getReader();
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+
+          const text = decoder.decode(value);
+          if (level === "error") {
+            console.error(`[Jobs:${jobId}] stderr:`, text.trim());
+          } else {
+            console.log(`[Jobs:${jobId}] stdout:`, text.trim());
+          }
+
+          if (events) {
+            await events.emit("job.external.log", {
+              jobId,
+              name: jobName,
+              level,
+              message: text.trim(),
+            });
+          }
+        }
+      } catch {
+        // Stream may be closed
+      }
+    };
+
+    // Stream stdout
+    if (proc.stdout && typeof proc.stdout !== "number") {
+      streamOutput(proc.stdout as ReadableStream<Uint8Array>, "info");
+    }
+
+    // Stream stderr
+    if (proc.stderr && typeof proc.stderr !== "number") {
+      streamOutput(proc.stderr as ReadableStream<Uint8Array>, "error");
+    }
+  }
+
   private async processJob(job: Job): Promise<void> {
     const handler = this.handlers.get(job.name);
     if (!handler) {