v3: machine config (#978)

* add and use machine config

* assign tasks to worker nodes only

* add secure flag to zod connection

* changeset

* add pre stop hook

* don't use secure connection by default

* pass more identifiers to provider and apply labels

.changeset/clean-pianos-listen.md

@@ -0,0 +1,8 @@
+---
+"@trigger.dev/sdk": patch
+"@trigger.dev/core-apps": patch
+"trigger.dev": patch
+"@trigger.dev/core": patch
+---
+
+add machine config and secure zod connection

apps/coordinator/.env.example

@@ -1,3 +1,4 @@
 HTTP_SERVER_PORT=8020
 PLATFORM_ENABLED=true
-PLATFORM_WS_PORT=3030
+PLATFORM_WS_PORT=3030
+SECURE_CONNECTION=false

apps/coordinator/src/index.ts

@@ -26,6 +26,7 @@ const PLATFORM_ENABLED = ["1", "true"].includes(process.env.PLATFORM_ENABLED ??
 const PLATFORM_HOST = process.env.PLATFORM_HOST || "127.0.0.1";
 const PLATFORM_WS_PORT = process.env.PLATFORM_WS_PORT || 3030;
 const PLATFORM_SECRET = process.env.PLATFORM_SECRET || "coordinator-secret";
+const SECURE_CONNECTION = ["1", "true"].includes(process.env.SECURE_CONNECTION ?? "false");
 
 const logger = new SimpleLogger(`[${NODE_NAME}]`);
 
@@ -365,6 +366,7 @@ class TaskCoordinator {
       namespace: "coordinator",
       host: PLATFORM_HOST,
       port: Number(PLATFORM_WS_PORT),
+      secure: SECURE_CONNECTION,
       clientMessages: CoordinatorToPlatformMessages,
       serverMessages: PlatformToCoordinatorMessages,
       authToken: PLATFORM_SECRET,

apps/docker-provider/.env.example

@@ -2,6 +2,7 @@ HTTP_SERVER_PORT=8050
 
 PLATFORM_WS_PORT=3030
 PLATFORM_SECRET=provider-secret
+SECURE_CONNECTION=false
 
 # Use this if you are on macOS
 # COORDINATOR_HOST="host.docker.internal"

apps/docker-provider/src/index.ts

@@ -8,6 +8,7 @@ import {
   TaskOperationsIndexOptions,
 } from "@trigger.dev/core-apps";
 import { setTimeout } from "node:timers/promises";
+import { PostStartCauses, PreStopCauses } from "@trigger.dev/core/v3";
 
 const MACHINE_NAME = process.env.MACHINE_NAME || "local";
 const COORDINATOR_PORT = process.env.COORDINATOR_PORT || 8020;
@@ -190,6 +191,9 @@ class DockerTaskOperations implements TaskOperations {
   async delete(opts: { runId: string }) {
     await this.#initialize();
 
+    const containerName = this.#getRunContainerName(opts.runId);
+    await this.#sendPreStop(containerName);
+
     logger.log("noop: delete");
   }
 
@@ -208,6 +212,26 @@ class DockerTaskOperations implements TaskOperations {
   }
 
   async #sendPostStart(containerName: string): Promise<void> {
+    try {
+      const port = await this.#getHttpServerPort(containerName);
+      logger.debug(await this.#runLifecycleCommand(containerName, port, "postStart", "restore"));
+    } catch (error) {
+      logger.error("postStart error", { error });
+      throw new Error("postStart command failed");
+    }
+  }
+
+  async #sendPreStop(containerName: string): Promise<void> {
+    try {
+      const port = await this.#getHttpServerPort(containerName);
+      logger.debug(await this.#runLifecycleCommand(containerName, port, "preStop", "terminate"));
+    } catch (error) {
+      logger.error("preStop error", { error });
+      throw new Error("preStop command failed");
+    }
+  }
+
+  async #getHttpServerPort(containerName: string): Promise<number> {
     // We first get the correct port, which is random during dev as we run with host networking and need to avoid clashes
     // FIXME: Skip this in prod
     const logs = logger.debug(await $`docker logs ${containerName}`);
@@ -219,19 +243,14 @@ class DockerTaskOperations implements TaskOperations {
       throw new Error("failed to extract port from logs");
     }
 
-    try {
+    return port;
-      logger.debug(await this.#runLifecycleCommand(containerName, port, "postStart", "restore"));
-    } catch (error) {
-      logger.error("postStart error", { error });
-      throw new Error("postStart command failed");
-    }
   }
 
-  async #runLifecycleCommand(
+  async #runLifecycleCommand<THookType extends "postStart" | "preStop">(
     containerName: string,
     port: number,
-    type: "postStart" | "preStop",
+    type: THookType,
-    cause: "index" | "create" | "restore",
+    cause: THookType extends "postStart" ? PostStartCauses : PreStopCauses,
     retryCount = 0
   ): Promise<ExecaChildProcess> {
     try {
@@ -244,15 +263,15 @@ class DockerTaskOperations implements TaskOperations {
         `127.0.0.1:${port}/${type}?cause=${cause}`,
       ]);
     } catch (error: any) {
-      if (retryCount < 6) {
+      if (type === "postStart" && retryCount < 6) {
-        logger.debug("retriable postStart error", { retryCount, message: error?.message });
+        logger.debug(`retriable ${type} error`, { retryCount, message: error?.message });
         await setTimeout(exponentialBackoff(retryCount + 1, 2, 50, 1150, 50));
 
         return this.#runLifecycleCommand(containerName, port, type, cause, retryCount + 1);
       }
 
-      logger.error("final postStart error", { message: error?.message });
+      logger.error(`final ${type} error`, { message: error?.message });
-      throw new Error(`postStart command failed after ${retryCount - 1} retries`);
+      throw new Error(`${type} command failed after ${retryCount - 1} retries`);
     }
   }
 }

apps/kubernetes-provider/.env.example

@@ -2,6 +2,7 @@ HTTP_SERVER_PORT=8060
 
 PLATFORM_WS_PORT=3030
 PLATFORM_SECRET=provider-secret
+SECURE_CONNECTION=false
 
 # Use this if you are on macOS
 # COORDINATOR_HOST="host.docker.internal"

apps/kubernetes-provider/src/index.ts

@@ -7,6 +7,7 @@ import {
   TaskOperationsIndexOptions,
   TaskOperationsRestoreOptions,
 } from "@trigger.dev/core-apps";
+import { Machine, PostStartCauses, PreStopCauses, EnvironmentType } from "@trigger.dev/core/v3";
 import { randomUUID } from "crypto";
 
 const RUNTIME_ENV = process.env.KUBERNETES_PORT ? "kubernetes" : "local";
@@ -55,6 +56,12 @@ class KubernetesTaskOperations implements TaskOperations {
             metadata: {
               labels: {
                 app: "task-index",
+                "app.kubernetes.io/part-of": "trigger-worker",
+                "app.kubernetes.io/component": "index",
+                env: opts.envId,
+                envtype: this.#envTypeToLabelValue(opts.envType),
+                org: opts.orgId,
+                project: opts.projectId,
               },
             },
             spec: {
@@ -64,6 +71,9 @@ class KubernetesTaskOperations implements TaskOperations {
                   name: "registry-trigger",
                 },
               ],
+              nodeSelector: {
+                nodetype: "worker",
+              },
               containers: [
                 {
                   name: this.#getIndexContainerName(opts.shortCode),
@@ -79,6 +89,13 @@ class KubernetesTaskOperations implements TaskOperations {
                   //     memory: "50Mi",
                   //   },
                   // },
+                  lifecycle: {
+                    preStop: {
+                      exec: {
+                        command: this.#getLifecycleCommand("preStop", "terminate"),
+                      },
+                    },
+                  },
                   env: [
                     {
                       name: "DEBUG",
@@ -151,6 +168,13 @@ class KubernetesTaskOperations implements TaskOperations {
           namespace: this.#namespace.metadata.name,
           labels: {
             app: "task-run",
+            "app.kubernetes.io/part-of": "trigger-worker",
+            "app.kubernetes.io/component": "create",
+            env: opts.envId,
+            envtype: this.#envTypeToLabelValue(opts.envType),
+            org: opts.orgId,
+            project: opts.projectId,
+            run: opts.runId,
           },
         },
         spec: {
@@ -160,6 +184,9 @@ class KubernetesTaskOperations implements TaskOperations {
               name: "registry-trigger",
             },
           ],
+          nodeSelector: {
+            nodetype: "worker",
+          },
           containers: [
             {
               name: this.#getRunContainerName(opts.runId),
@@ -169,9 +196,9 @@ class KubernetesTaskOperations implements TaskOperations {
                   containerPort: 8000,
                 },
               ],
-              // resources: {
+              resources: {
-              //   limits: opts.machine,
+                limits: this.#getResourcesFromMachineConfig(opts.machine),
-              // },
+              },
               lifecycle: {
                 postStart: {
                   exec: {
@@ -180,7 +207,7 @@ class KubernetesTaskOperations implements TaskOperations {
                 },
                 preStop: {
                   exec: {
-                    command: this.#getLifecycleCommand("preStop", "create"),
+                    command: this.#getLifecycleCommand("preStop", "terminate"),
                   },
                 },
               },
@@ -262,6 +289,14 @@ class KubernetesTaskOperations implements TaskOperations {
           namespace: this.#namespace.metadata.name,
           labels: {
             app: "task-run",
+            "app.kubernetes.io/part-of": "trigger-worker",
+            "app.kubernetes.io/component": "restore",
+            env: opts.envId,
+            envtype: this.#envTypeToLabelValue(opts.envType),
+            org: opts.orgId,
+            project: opts.projectId,
+            run: opts.runId,
+            checkpoint: opts.checkpointId,
           },
         },
         spec: {
@@ -271,6 +306,9 @@ class KubernetesTaskOperations implements TaskOperations {
               name: "registry-trigger",
             },
           ],
+          nodeSelector: {
+            nodetype: "worker",
+          },
           initContainers: [
             {
               name: "pull-base-image",
@@ -309,9 +347,9 @@ class KubernetesTaskOperations implements TaskOperations {
                   containerPort: 8000,
                 },
               ],
-              // resources: {
+              resources: {
-              //   limits: opts.machine,
+                limits: this.#getResourcesFromMachineConfig(opts.machine),
-              // },
+              },
               lifecycle: {
                 postStart: {
                   exec: {
@@ -320,7 +358,7 @@ class KubernetesTaskOperations implements TaskOperations {
                 },
                 preStop: {
                   exec: {
-                    command: this.#getLifecycleCommand("preStop", "restore"),
+                    command: this.#getLifecycleCommand("preStop", "terminate"),
                   },
                 },
               },
@@ -355,7 +393,30 @@ class KubernetesTaskOperations implements TaskOperations {
     await this.#getPod(opts.runId, this.#namespace);
   }
 
-  #getLifecycleCommand(type: "postStart" | "preStop", cause: "index" | "create" | "restore") {
+  #envTypeToLabelValue(type: EnvironmentType) {
+    switch (type) {
+      case "PRODUCTION":
+        return "prod";
+      case "STAGING":
+        return "stg";
+      case "DEVELOPMENT":
+        return "dev";
+      case "PREVIEW":
+        return "preview";
+    }
+  }
+
+  #getResourcesFromMachineConfig(config: Machine) {
+    return {
+      cpu: `${config.cpu}`,
+      memory: `${config.memory}G`,
+    };
+  }
+
+  #getLifecycleCommand<THookType extends "postStart" | "preStop">(
+    type: THookType,
+    cause: THookType extends "postStart" ? PostStartCauses : PreStopCauses
+  ) {
     return ["/bin/sh", "-c", `sleep 1; wget -q -O- 127.0.0.1:8000/${type}?cause=${cause}`];
   }
 

apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts

@@ -1,5 +1,6 @@
 import { Context, ROOT_CONTEXT, Span, SpanKind, context, trace } from "@opentelemetry/api";
 import {
+  Machine,
   ProdTaskRunExecution,
   ProdTaskRunExecutionPayload,
   TaskRunError,
@@ -438,10 +439,27 @@ export class SharedQueueConsumer {
             queueId: queue.id,
             runtimeEnvironmentId: environment.id,
           },
+          include: {
+            backgroundWorkerTask: true,
+          },
         });
 
         const isRetry = taskRunAttempt.number > 1;
 
+        const { machineConfig } = taskRunAttempt.backgroundWorkerTask;
+        const machine = Machine.safeParse(machineConfig ?? {});
+
+        if (!machine.success) {
+          logger.error("Failed to parse machine config", {
+            queueMessage: message.data,
+            messageId: message.messageId,
+            attemptId: taskRunAttempt.id,
+            machineConfig,
+          });
+
+          await this.#ackAndDoMoreWork(message.messageId);
+          return;
+        }
         try {
           if (messageBody.data.checkpointEventId) {
             const restoreService = new RestoreCheckpointService();
@@ -470,11 +488,16 @@ export class SharedQueueConsumer {
               backgroundWorkerId: deployment.worker.friendlyId,
               data: {
                 type: "SCHEDULE_ATTEMPT",
-                id: taskRunAttempt.id,
                 image: deployment.imageReference,
-                envId: environment.id,
-                runId: taskRunAttempt.taskRunId,
                 version: deployment.version,
+                machine: machine.data,
+                // identifiers
+                id: taskRunAttempt.id,
+                envId: environment.id,
+                envType: environment.type,
+                orgId: environment.organizationId,
+                projectId: environment.projectId,
+                runId: taskRunAttempt.taskRunId,
               },
             });
           }

apps/webapp/app/v3/services/createBackgroundWorker.server.ts

@@ -101,6 +101,7 @@ export async function createBackgroundTasks(
           exportName: task.exportName,
           retryConfig: task.retry,
           queueConfig: task.queue,
+          machineConfig: task.machine,
         },
       });
 

apps/webapp/app/v3/services/indexDeployment.server.ts

@@ -49,9 +49,13 @@ export class IndexDeploymentService extends BaseService {
         version: "v1",
         shortCode: deployment.shortCode,
         imageTag: deployment.imageReference,
-        envId: deployment.environmentId,
         apiKey: deployment.environment.apiKey,
         apiUrl: env.APP_ORIGIN,
+        // identifiers
+        envId: deployment.environmentId,
+        envType: deployment.environment.type,
+        projectId: deployment.projectId,
+        orgId: deployment.environment.organizationId,
       });
 
       logger.debug("Index ACK received", { responses });

apps/webapp/app/v3/services/restoreCheckpoint.server.ts

@@ -3,6 +3,7 @@ import { logger } from "~/services/logger.server";
 import { socketIo } from "../handleSocketIo.server";
 import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server";
 import { BaseService } from "./baseService.server";
+import { Machine } from "@trigger.dev/core/v3";
 
 const RESTORABLE_RUN_STATUSES: TaskRunStatus[] = ["WAITING_TO_RESUME"];
 const RESTORABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["PAUSED"];
@@ -30,8 +31,14 @@ export class RestoreCheckpointService extends BaseService {
             attempt: {
               select: {
                 status: true,
+                backgroundWorkerTask: {
+                  select: {
+                    machineConfig: true,
+                  },
+                },
               },
             },
+            runtimeEnvironment: true,
           },
         },
       },
@@ -63,17 +70,34 @@ export class RestoreCheckpointService extends BaseService {
       return;
     }
 
+    const { machineConfig } = checkpoint.attempt.backgroundWorkerTask;
+    const machine = Machine.safeParse(machineConfig ?? {});
+
+    if (!machine.success) {
+      logger.error("Failed to parse machine config", {
+        attemptId: checkpoint.attemptId,
+        machineConfig: checkpoint.attempt.backgroundWorkerTask.machineConfig,
+      });
+      return;
+    }
+
     const eventService = new CreateCheckpointRestoreEventService(this._prisma);
     await eventService.restore({ checkpointId: checkpoint.id });
 
     socketIo.providerNamespace.emit("RESTORE", {
       version: "v1",
-      checkpointId: checkpoint.id,
-      runId: checkpoint.runId,
       type: checkpoint.type,
       location: checkpoint.location,
       reason: checkpoint.reason ?? undefined,
       imageRef: checkpoint.imageRef,
+      machine: machine.data,
+      // identifiers
+      checkpointId: checkpoint.id,
+      envId: checkpoint.runtimeEnvironment.id,
+      envType: checkpoint.runtimeEnvironment.type,
+      orgId: checkpoint.runtimeEnvironment.organizationId,
+      projectId: checkpoint.runtimeEnvironment.projectId,
+      runId: checkpoint.runId,
     });
 
     return checkpoint;

packages/cli-v3/src/workers/dev/worker-facade.ts

@@ -84,6 +84,7 @@ function getTasks(): Array<TaskMetadataWithFunctions> {
           filePath: (taskFile as any).filePath,
           queue: (task as any).__trigger.queue,
           retry: (task as any).__trigger.retry,
+          machine: (task as any).__trigger.machine,
           fns: (task as any).__trigger.fns,
         });
       }

packages/cli-v3/src/workers/prod/entry-point.ts

@@ -1,6 +1,8 @@
 import {
   Config,
   CoordinatorToProdWorkerMessages,
+  PostStartCauses,
+  PreStopCauses,
   ProdWorkerToCoordinatorMessages,
   TaskResource,
   WaitReason,
@@ -567,23 +569,15 @@ class ProdWorker {
           }
 
           case "/preStop": {
-            const schema = z.enum(["index", "create", "restore"]);
+            const cause = PreStopCauses.safeParse(url.searchParams.get("cause"));
-
-            const cause = schema.safeParse(url.searchParams.get("cause"));
 
             if (!cause.success) {
               logger.error("Failed to parse cause", { cause });
-              return;
+              return reply.text("Failed to parse cause", 400);
             }
 
             switch (cause.data) {
-              case "index": {
+              case "terminate": {
-                break;
-              }
-              case "create": {
-                break;
-              }
-              case "restore": {
                 break;
               }
               default: {
@@ -597,13 +591,11 @@ class ProdWorker {
           }
 
           case "/postStart": {
-            const schema = z.enum(["index", "create", "restore"]);
+            const cause = PostStartCauses.safeParse(url.searchParams.get("cause"));
-
-            const cause = schema.safeParse(url.searchParams.get("cause"));
 
             if (!cause.success) {
               logger.error("Failed to parse cause", { cause });
-              return;
+              return reply.text("Failed to parse cause", 400);
             }
 
             switch (cause.data) {
@@ -678,11 +670,7 @@ class ProdWorker {
     }
 
     for (const task of this.#backgroundWorker.tasks) {
-      taskResources.push({
+      taskResources.push(task);
-        id: task.id,
-        filePath: task.filePath,
-        exportName: task.exportName,
-      });
 
       packageVersion = task.packageVersion;
     }

packages/cli-v3/src/workers/prod/worker-facade.ts

@@ -78,6 +78,7 @@ function getTasks(): Array<TaskMetadataWithFunctions> {
           filePath: (taskFile as any).filePath,
           queue: (task as any).__trigger.queue,
           retry: (task as any).__trigger.retry,
+          machine: (task as any).__trigger.machine,
           fns: (task as any).__trigger.fns,
         });
       }

packages/core-apps/src/provider.ts

@@ -2,6 +2,7 @@ import { createServer } from "node:http";
 import {
   ClientToSharedQueueMessages,
   clientWebsocketMessages,
+  EnvironmentType,
   Machine,
   PlatformToProviderMessages,
   ProviderToPlatformMessages,
@@ -18,30 +19,46 @@ const MACHINE_NAME = process.env.MACHINE_NAME || "local";
 const PLATFORM_HOST = process.env.PLATFORM_HOST || "127.0.0.1";
 const PLATFORM_WS_PORT = process.env.PLATFORM_WS_PORT || 3030;
 const PLATFORM_SECRET = process.env.PLATFORM_SECRET || "provider-secret";
+const SECURE_CONNECTION = ["1", "true"].includes(process.env.SECURE_CONNECTION ?? "false");
 
 const logger = new SimpleLogger(`[${MACHINE_NAME}]`);
 
 export interface TaskOperationsIndexOptions {
   shortCode: string;
   imageRef: string;
-  envId: string;
   apiKey: string;
   apiUrl: string;
+  // identifiers
+  envId: string;
+  envType: EnvironmentType;
+  orgId: string;
+  projectId: string;
 }
 
 export interface TaskOperationsCreateOptions {
-  runId: string;
   image: string;
   machine: Machine;
-  envId: string;
   version: string;
+  // identifiers
+  envId: string;
+  envType: EnvironmentType;
+  orgId: string;
+  projectId: string;
+  runId: string;
+  attemptId: string;
 }
 
 export interface TaskOperationsRestoreOptions {
-  runId: string;
   imageRef: string;
   checkpointRef: string;
   machine: Machine;
+  // identifiers
+  envId: string;
+  envType: EnvironmentType;
+  orgId: string;
+  projectId: string;
+  runId: string;
+  checkpointId: string;
 }
 
 export interface TaskOperations {
@@ -87,6 +104,7 @@ export class ProviderShell implements Provider {
       namespace: "shared-queue",
       host: PLATFORM_HOST,
       port: Number(PLATFORM_WS_PORT),
+      secure: SECURE_CONNECTION,
       clientMessages: ClientToSharedQueueMessages,
       serverMessages: SharedQueueToClientMessages,
       authToken: PLATFORM_SECRET,
@@ -101,11 +119,16 @@ export class ProviderShell implements Provider {
           if (message.data.type === "SCHEDULE_ATTEMPT") {
             try {
               this.tasks.create({
-                envId: message.data.envId,
-                runId: message.data.runId,
                 image: message.data.image,
-                machine: {},
+                machine: message.data.machine,
-                version: message.version,
+                version: message.data.version,
+                // identifiers
+                envId: message.data.envId,
+                envType: message.data.envType,
+                orgId: message.data.orgId,
+                projectId: message.data.projectId,
+                runId: message.data.runId,
+                attemptId: message.data.id,
               });
             } catch (error) {
               logger.error("create failed", error);
@@ -138,6 +161,7 @@ export class ProviderShell implements Provider {
       namespace: "provider",
       host: PLATFORM_HOST,
       port: Number(PLATFORM_WS_PORT),
+      secure: SECURE_CONNECTION,
       clientMessages: ProviderToPlatformMessages,
       serverMessages: PlatformToProviderMessages,
       authToken: PLATFORM_SECRET,
@@ -165,9 +189,13 @@ export class ProviderShell implements Provider {
             await this.tasks.index({
               shortCode: message.shortCode,
               imageRef: message.imageTag,
-              envId: message.envId,
               apiKey: message.apiKey,
               apiUrl: message.apiUrl,
+              // identifiers
+              envId: message.envId,
+              envType: message.envType,
+              orgId: message.orgId,
+              projectId: message.projectId,
             });
           } catch (error) {
             logger.error("index failed", error);
@@ -206,13 +234,16 @@ export class ProviderShell implements Provider {
 
           try {
             await this.tasks.restore({
-              runId: message.runId,
               checkpointRef: message.location,
-              machine: {
+              machine: message.machine,
-                cpu: "1",
-                memory: "100Mi",
-              },
               imageRef: message.imageRef,
+              // identifiers
+              envId: message.envId,
+              envType: message.envType,
+              orgId: message.orgId,
+              projectId: message.projectId,
+              runId: message.runId,
+              checkpointId: message.checkpointId,
             });
           } catch (error) {
             logger.error("restore failed", error);
@@ -230,54 +261,34 @@ export class ProviderShell implements Provider {
 
       const reply = new HttpReply(res);
 
-      switch (req.url) {
+      try {
-        case "/health": {
+        const url = new URL(req.url ?? "", `http://${req.headers.host}`);
-          return reply.text("ok");
-        }
-        case "/whoami": {
-          return reply.text(`${MACHINE_NAME}`);
-        }
-        case "/close": {
-          this.#platformSocket.close();
-          return reply.text("platform socket closed");
-        }
-        case "/delete": {
-          const body = await getTextBody(req);
 
-          await this.tasks.delete({ runId: body });
+        switch (url.pathname) {
+          case "/health": {
+            return reply.text("ok");
+          }
+          case "/whoami": {
+            return reply.text(`${MACHINE_NAME}`);
+          }
+          case "/close": {
+            this.#platformSocket.close();
+            return reply.text("platform socket closed");
+          }
+          case "/delete": {
+            const body = await getTextBody(req);
 
-          return reply.text(`sent delete request: ${body}`);
+            await this.tasks.delete({ runId: body });
-        }
+
-        case "/invoke": {
+            return reply.text(`sent delete request: ${body}`);
-          const body = await getTextBody(req);
+          }
-
+          default: {
-          await this.tasks.create({
+            return reply.empty(404);
-            envId: "placeholder",
+          }
-            image: body,
-            machine: {
-              cpu: "1",
-              memory: "100Mi",
-            },
-            runId: "<missing>",
-            version: "<missing>",
-          });
-
-          return reply.text(`sent restore request: ${body}`);
-        }
-        case "/restore": {
-          const body = await getTextBody(req);
-
-          const items = body.split("&");
-          const image = items[0];
-          const baseImageTag = items[1] ?? image;
-
-          // await this.tasks.restore({});
-
-          return reply.text(`sent restore request: ${body}`);
-        }
-        default: {
-          return reply.empty(404);
         }
+      } catch (error) {
+        logger.error("HTTP server error", { error });
+        reply.empty(500);
       }
     });
 

packages/core/src/v3/schemas/messages.ts

@@ -1,6 +1,29 @@
 import { z } from "zod";
 import { TaskRunExecution, TaskRunExecutionResult } from "./common";
 
+export const EnvironmentType = z.enum(["PRODUCTION", "STAGING", "DEVELOPMENT", "PREVIEW"])
+export type EnvironmentType = z.infer<typeof EnvironmentType>;
+
+export const MachineCpu = z
+  .union([z.literal(0.25), z.literal(0.5), z.literal(1), z.literal(2), z.literal(4)])
+  .default(0.5);
+
+export type MachineCpu = z.infer<typeof MachineCpu>;
+
+export const MachineMemory = z
+  .union([z.literal(0.25), z.literal(0.5), z.literal(1), z.literal(2), z.literal(4), z.literal(8)])
+  .default(1);
+
+export type MachineMemory = z.infer<typeof MachineMemory>;
+
+export const Machine = z.object({
+  version: z.literal("v1").default("v1"),
+  cpu: MachineCpu,
+  memory: MachineMemory,
+});
+
+export type Machine = z.infer<typeof Machine>;
+
 export const TaskRunExecutionPayload = z.object({
   execution: TaskRunExecution,
   traceContext: z.record(z.unknown()),
@@ -39,11 +62,16 @@ export const BackgroundWorkerServerMessages = z.discriminatedUnion("type", [
   }),
   z.object({
     type: z.literal("SCHEDULE_ATTEMPT"),
-    id: z.string(),
     image: z.string(),
-    envId: z.string(),
-    runId: z.string(),
     version: z.string(),
+    machine: Machine,
+    // identifiers
+    id: z.string(), // attempt
+    envId: z.string(),
+    envType: EnvironmentType,
+    orgId: z.string(),
+    projectId: z.string(),
+    runId: z.string(),
   }),
 ]);
 
@@ -229,6 +257,7 @@ export const TaskMetadata = z.object({
   packageVersion: z.string(),
   queue: QueueOptions.optional(),
   retry: RetryOptions.optional(),
+  machine: Machine.partial().optional(),
 });
 
 export type TaskMetadata = z.infer<typeof TaskMetadata>;

packages/core/src/v3/schemas/resources.ts

@@ -1,5 +1,5 @@
 import { z } from "zod";
-import { QueueOptions, RetryOptions } from "./messages";
+import { Machine, QueueOptions, RetryOptions } from "./messages";
 
 export const TaskResource = z.object({
   id: z.string(),
@@ -7,6 +7,7 @@ export const TaskResource = z.object({
   exportName: z.string(),
   queue: QueueOptions.optional(),
   retry: RetryOptions.optional(),
+  machine: Machine.partial().optional(),
 });
 
 export type TaskResource = z.infer<typeof TaskResource>;

packages/core/src/v3/schemas/schemas.ts

@@ -7,9 +7,17 @@ import {
   ProdTaskRunExecution,
   ProdTaskRunExecutionPayload,
   RetryOptions,
+  Machine,
+  EnvironmentType,
 } from "./messages";
 import { TaskResource } from "./resources";
 
+export const PostStartCauses = z.enum(["index", "create", "restore"]);
+export type PostStartCauses = z.infer<typeof PostStartCauses>;
+
+export const PreStopCauses = z.enum(["terminate"]);
+export type PreStopCauses = z.infer<typeof PreStopCauses>;
+
 const RegexSchema = z.custom<RegExp>((val) => {
   try {
     // Check to see if val is a regex
@@ -42,13 +50,6 @@ export type ResolvedConfig = RequireKeys<
   "triggerDirectories" | "triggerUrl" | "projectDir" | "tsconfigPath"
 >;
 
-export const Machine = z.object({
-  cpu: z.string().default("1").optional(),
-  memory: z.string().default("500Mi").optional(),
-});
-
-export type Machine = z.infer<typeof Machine>;
-
 export const WaitReason = z.enum(["WAIT_FOR_DURATION", "WAIT_FOR_TASK", "WAIT_FOR_BATCH"]);
 
 export type WaitReason = z.infer<typeof WaitReason>;
@@ -85,9 +86,13 @@ export const PlatformToProviderMessages = {
       version: z.literal("v1").default("v1"),
       imageTag: z.string(),
       shortCode: z.string(),
-      envId: z.string(),
       apiKey: z.string(),
       apiUrl: z.string(),
+      // identifiers
+      envId: z.string(),
+      envType: EnvironmentType,
+      orgId: z.string(),
+      projectId: z.string(),
     }),
     callback: z.discriminatedUnion("success", [
       z.object({
@@ -103,22 +108,22 @@ export const PlatformToProviderMessages = {
       }),
     ]),
   },
-  INVOKE: {
+  // TODO: this should be a shared queue message instead
-    message: z.object({
-      version: z.literal("v1").default("v1"),
-      name: z.string(),
-      machine: Machine,
-    }),
-  },
   RESTORE: {
     message: z.object({
       version: z.literal("v1").default("v1"),
-      checkpointId: z.string(),
-      runId: z.string(),
       type: z.enum(["DOCKER", "KUBERNETES"]),
       location: z.string(),
       reason: z.string().optional(),
       imageRef: z.string(),
+      machine: Machine,
+      // identifiers
+      checkpointId: z.string(),
+      envId: z.string(),
+      envType: EnvironmentType,
+      orgId: z.string(),
+      projectId: z.string(),
+      runId: z.string(),
     }),
   },
   DELETE: {

packages/core/src/v3/zodSocket.ts

@@ -263,7 +263,8 @@ interface ZodSocketConnectionOptions<
   TServerMessages extends ZodSocketMessageCatalogSchema,
 > {
   host: string;
-  port: number;
+  port?: number;
+  secure?: boolean;
   namespace: string;
   clientMessages: TClientMessages;
   serverMessages: TServerMessages;
@@ -302,15 +303,24 @@ export class ZodSocketConnection<
   #logger: StructuredLogger;
 
   constructor(opts: ZodSocketConnectionOptions<TClientMessages, TServerMessages>) {
-    this.socket = io(`ws://${opts.host}:${opts.port}/${opts.namespace}`, {
+    const uri = `${opts.secure ? "wss" : "ws"}://${opts.host}:${
+      opts.port ?? (opts.secure ? "443" : "80")
+    }/${opts.namespace}`;
+
+    const logger = new SimpleStructuredLogger(opts.namespace, LogLevel.info);
+    logger.log("new zod socket", { uri });
+
+    this.socket = io(uri, {
       transports: ["websocket"],
       auth: {
         token: opts.authToken,
       },
       extraHeaders: opts.extraHeaders,
+      reconnectionDelay: 500,
+      reconnectionDelayMax: 1000,
     });
 
-    this.#logger = new SimpleStructuredLogger(opts.namespace, LogLevel.info, {
+    this.#logger = logger.child({
       socketId: this.socket.id,
     });
 

packages/database/prisma/migrations/20240327121557_add_machine_config_to_worker_task/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "BackgroundWorkerTask" ADD COLUMN     "machineConfig" JSONB;

packages/database/prisma/schema.prisma

@@ -1559,6 +1559,7 @@ model BackgroundWorkerTask {
 
   queueConfig Json?
   retryConfig Json?
+  machineConfig Json?
 
   @@unique([workerId, slug])
 }

packages/trigger-sdk/src/v3/shared.ts

@@ -9,6 +9,8 @@ import {
   HandleErrorResult,
   InitFnParams,
   InitOutput,
+  MachineCpu,
+  MachineMemory,
   MiddlewareFnParams,
   QueueOptions,
   RetryOptions,
@@ -114,8 +116,8 @@ export type TaskOptions<TPayload, TOutput = any, TInitOutput extends InitOutput
      * - 2
      * - 4
      */
-    cpu?: 0.25 | 0.5 | 1 | 2 | 4;
+    cpu?: MachineCpu;
-    /** In GBs of RAM. The default is 0.5.
+    /** In GBs of RAM. The default is 1.
      *
      * Possible values:
      * - 0.25
@@ -125,7 +127,7 @@ export type TaskOptions<TPayload, TOutput = any, TInitOutput extends InitOutput
      * - 4
      * - 8
      */
-    memory?: 0.25 | 0.5 | 1 | 2 | 4 | 8;
+    memory?: MachineMemory;
   };
   /** This gets called when a task is triggered. It's where you put the code you want to execute.
    *
@@ -489,6 +491,7 @@ export function createTask<TInput, TOutput, TInitOutput extends InitOutput>(
       packageVersion: packageJson.version,
       queue: params.queue,
       retry: params.retry ? { ...defaultRetryOptions, ...params.retry } : undefined,
+      machine: params.machine,
       fns: {
         run: params.run,
         init: params.init,