Merge branch 'fix/checkpoints' into fix/v4-misc

Author: nicktrn

apps/supervisor/src/env.ts

@@ -26,6 +26,7 @@ const Env = z.object({
     .transform((s) => z.enum(["http", "https"]).parse(s.toLowerCase()))
     .default("http"),
   TRIGGER_WORKLOAD_API_DOMAIN: z.string().optional(), // If unset, will use orchestrator-specific default
+  TRIGGER_WORKLOAD_API_HOST_INTERNAL: z.string().default("0.0.0.0"),
   TRIGGER_WORKLOAD_API_PORT_INTERNAL: z.coerce.number().default(8020), // This is the port the workload API listens on
   TRIGGER_WORKLOAD_API_PORT_EXTERNAL: z.coerce.number().default(8020), // This is the exposed port passed to the run controller
 
@@ -53,7 +54,10 @@ const Env = z.object({
   EPHEMERAL_STORAGE_SIZE_REQUEST: z.string().default("2Gi"),
 
   // Metrics
+  METRICS_ENABLED: BoolEnv.default(true),
   METRICS_COLLECT_DEFAULTS: BoolEnv.default(true),
+  METRICS_HOST: z.string().default("127.0.0.1"),
+  METRICS_PORT: z.coerce.number().int().default(9090),
 
   // Pod cleaner
   POD_CLEANER_ENABLED: BoolEnv.default(true),

apps/supervisor/src/index.ts

@@ -30,7 +30,7 @@ if (env.METRICS_COLLECT_DEFAULTS) {
 
 class ManagedSupervisor {
   private readonly workerSession: SupervisorSession;
-  private readonly httpServer: HttpServer;
+  private readonly metricsServer?: HttpServer;
   private readonly workloadServer: WorkloadServer;
   private readonly workloadManager: WorkloadManager;
   private readonly logger = new SimpleStructuredLogger("managed-worker");
@@ -61,6 +61,7 @@ class ManagedSupervisor {
         intervalMs: env.POD_CLEANER_INTERVAL_MS,
       });
       this.podCleaner = new PodCleaner({
+        register,
         namespace: env.KUBERNETES_NAMESPACE,
         batchSize: env.POD_CLEANER_BATCH_SIZE,
         intervalMs: env.POD_CLEANER_INTERVAL_MS,
@@ -75,6 +76,7 @@ class ManagedSupervisor {
         reconnectIntervalMs: env.FAILED_POD_HANDLER_RECONNECT_INTERVAL_MS,
       });
       this.failedPodHandler = new FailedPodHandler({
+        register,
         namespace: env.KUBERNETES_NAMESPACE,
         reconnectIntervalMs: env.FAILED_POD_HANDLER_RECONNECT_INTERVAL_MS,
       });
@@ -243,16 +245,21 @@ class ManagedSupervisor {
       }
     });
 
-    // Used for health checks and metrics
-    this.httpServer = new HttpServer({ port: 8080, host: "0.0.0.0" }).route("/health", "GET", {
-      handler: async ({ reply }) => {
-        reply.text("OK");
-      },
-    });
+    if (env.METRICS_ENABLED) {
+      this.metricsServer = new HttpServer({
+        port: env.METRICS_PORT,
+        host: env.METRICS_HOST,
+        metrics: {
+          register,
+          expose: true,
+        },
+      });
+    }
 
     // Responds to workload requests only
     this.workloadServer = new WorkloadServer({
       port: env.TRIGGER_WORKLOAD_API_PORT_INTERNAL,
+      host: env.TRIGGER_WORKLOAD_API_HOST_INTERNAL,
       workerClient: this.workerSession.httpClient,
       checkpointClient: this.checkpointClient,
     });
@@ -321,6 +328,7 @@ class ManagedSupervisor {
     // Optional services
     await this.podCleaner?.start();
     await this.failedPodHandler?.start();
+    await this.metricsServer?.start();
 
     if (env.TRIGGER_WORKLOAD_API_ENABLED) {
       this.logger.log("[ManagedWorker] Workload API enabled", {
@@ -334,16 +342,16 @@ class ManagedSupervisor {
     }
 
     await this.workerSession.start();
-    await this.httpServer.start();
   }
 
   async stop() {
     this.logger.log("[ManagedWorker] Shutting down");
-    await this.httpServer.stop();
+    await this.workerSession.stop();
 
     // Optional services
     await this.podCleaner?.stop();
     await this.failedPodHandler?.stop();
+    await this.metricsServer?.stop();
   }
 }
 

apps/supervisor/src/workloadServer/index.ts

@@ -22,6 +22,7 @@ import {
 } from "@trigger.dev/core/v3/workers";
 import { HttpServer, type CheckpointClient } from "@trigger.dev/core/v3/serverOnly";
 import { type IncomingMessage } from "node:http";
+import { register } from "../metrics.js";
 
 // Use the official export when upgrading to socket.io@4.8.0
 interface DefaultEventsMap {
@@ -121,7 +122,19 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
   }
 
   private createHttpServer({ host, port }: { host: string; port: number }) {
-    return new HttpServer({ port, host })
+    return new HttpServer({
+      port,
+      host,
+      metrics: {
+        register,
+        expose: false,
+      },
+    })
+      .route("/health", "GET", {
+        handler: async ({ reply }) => {
+          reply.text("OK");
+        },
+      })
       .route(
         "/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/attempts/start",
         "POST",