Merge remote-tracking branch 'origin/run-engine-2' into run-engine-batch-trigger

Author: matt-aitken

.npmrc

@@ -1,3 +1,4 @@
 link-workspace-packages=false
 public-hoist-pattern[]=*prisma*
-prefer-workspace-packages=true
+prefer-workspace-packages=true
+update-notifier=false

apps/webapp/app/routes/admin.api.v1.workers.ts

@@ -0,0 +1,65 @@
+import { ActionFunctionArgs, json } from "@remix-run/server-runtime";
+import { z } from "zod";
+import { prisma } from "~/db.server";
+import { authenticateApiRequestWithPersonalAccessToken } from "~/services/personalAccessToken.server";
+import { WorkerGroupService } from "~/v3/services/worker/workerGroupService.server";
+
+const RequestBodySchema = z.object({
+  name: z.string().optional(),
+  description: z.string().optional(),
+  projectId: z.string().optional(),
+  makeDefault: z.boolean().optional(),
+});
+
+export async function action({ request }: ActionFunctionArgs) {
+  // Next authenticate the request
+  const authenticationResult = await authenticateApiRequestWithPersonalAccessToken(request);
+
+  if (!authenticationResult) {
+    return json({ error: "Invalid or Missing API key" }, { status: 401 });
+  }
+
+  const user = await prisma.user.findUnique({
+    where: {
+      id: authenticationResult.userId,
+    },
+  });
+
+  if (!user) {
+    return json({ error: "Invalid or Missing API key" }, { status: 401 });
+  }
+
+  if (!user.admin) {
+    return json({ error: "You must be an admin to perform this action" }, { status: 403 });
+  }
+
+  try {
+    const rawBody = await request.json();
+    const { name, description, projectId, makeDefault } = RequestBodySchema.parse(rawBody ?? {});
+
+    const service = new WorkerGroupService();
+    const { workerGroup, token } = await service.createWorkerGroup({
+      name,
+      description,
+    });
+
+    if (makeDefault && projectId) {
+      await prisma.project.update({
+        where: {
+          id: projectId,
+        },
+        data: {
+          defaultWorkerGroupId: workerGroup.id,
+          engine: "V2",
+        },
+      });
+    }
+
+    return json({
+      token,
+      workerGroup,
+    });
+  } catch (error) {
+    return json({ error: error instanceof Error ? error.message : error }, { status: 400 });
+  }
+}

apps/webapp/app/routes/api.v1.worker-actions.connect.ts

@@ -1,8 +1,5 @@
 import { json, TypedResponse } from "@remix-run/server-runtime";
-import {
-  WorkerApiConnectRequestBody,
-  WorkerApiConnectResponseBody,
-} from "@trigger.dev/worker";
+import { WorkerApiConnectRequestBody, WorkerApiConnectResponseBody } from "@trigger.dev/worker";
 import { createActionWorkerApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 
 export const action = createActionWorkerApiRoute(
@@ -11,6 +8,12 @@ export const action = createActionWorkerApiRoute(
   },
   async ({ authenticatedWorker, body }): Promise<TypedResponse<WorkerApiConnectResponseBody>> => {
     await authenticatedWorker.connect(body.metadata);
-    return json({ ok: true });
+    return json({
+      ok: true,
+      workerGroup: {
+        type: authenticatedWorker.type,
+        name: authenticatedWorker.name,
+      },
+    });
   }
 );

apps/webapp/app/routes/api.v1.worker-actions.heartbeat.ts

@@ -1,15 +1,12 @@
 import { json, TypedResponse } from "@remix-run/server-runtime";
-import {
-  WorkerApiConnectResponseBody,
-  WorkerApiHeartbeatRequestBody,
-} from "@trigger.dev/worker";
+import { WorkerApiHeartbeatResponseBody, WorkerApiHeartbeatRequestBody } from "@trigger.dev/worker";
 import { createActionWorkerApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 
 export const action = createActionWorkerApiRoute(
   {
     body: WorkerApiHeartbeatRequestBody,
   },
-  async ({ authenticatedWorker }): Promise<TypedResponse<WorkerApiConnectResponseBody>> => {
+  async ({ authenticatedWorker }): Promise<TypedResponse<WorkerApiHeartbeatResponseBody>> => {
     await authenticatedWorker.heartbeatWorkerInstance();
     return json({ ok: true });
   }

apps/webapp/app/v3/runEngine.server.ts

@@ -27,7 +27,7 @@ function createRunEngine() {
       pollIntervalMs: env.RUN_ENGINE_WORKER_POLL_INTERVAL,
     },
     machines: {
-      defaultMachine: defaultMachine,
+      defaultMachine,
       machines: allMachines(),
       baseCostInCents: env.CENTS_PER_RUN,
     },

apps/webapp/app/v3/services/worker/workerGroupTokenService.server.ts

@@ -25,6 +25,8 @@ import { CURRENT_UNMANAGED_DEPLOYMENT_LABEL } from "~/consts";
 import { resolveVariablesForEnvironment } from "~/v3/environmentVariables/environmentVariablesRepository.server";
 import { generateJWTTokenForEnvironment } from "~/services/apiAuth.server";
 import { fromFriendlyId } from "@trigger.dev/core/v3/apps";
+import { machinePresetFromName } from "~/v3/machinePresets.server";
+import { defaultMachine } from "@trigger.dev/platform/v3";
 
 export class WorkerGroupTokenService extends WithRunEngine {
   private readonly tokenPrefix = "tr_wgt_";
@@ -205,6 +207,7 @@ export class WorkerGroupTokenService extends WithRunEngine {
         prisma: this._prisma,
         engine: this._engine,
         type: WorkerInstanceGroupType.MANAGED,
+        name: workerGroup.name,
         workerGroupId: workerGroup.id,
         workerInstanceId: workerInstance.id,
         masterQueue: workerGroup.masterQueue,
@@ -240,6 +243,7 @@ export class WorkerGroupTokenService extends WithRunEngine {
       prisma: this._prisma,
       engine: this._engine,
       type: WorkerInstanceGroupType.UNMANAGED,
+      name: workerGroup.name,
       workerGroupId: workerGroup.id,
       workerInstanceId: workerInstance.id,
       masterQueue: workerGroup.masterQueue,
@@ -479,6 +483,7 @@ export type WorkerInstanceEnv = z.infer<typeof WorkerInstanceEnv>;
 
 export type AuthenticatedWorkerInstanceOptions = WithRunEngineOptions<{
   type: WorkerInstanceGroupType;
+  name: string;
   workerGroupId: string;
   workerInstanceId: string;
   masterQueue: string;
@@ -490,20 +495,22 @@ export type AuthenticatedWorkerInstanceOptions = WithRunEngineOptions<{
 
 export class AuthenticatedWorkerInstance extends WithRunEngine {
   readonly type: WorkerInstanceGroupType;
+  readonly name: string;
   readonly workerGroupId: string;
   readonly workerInstanceId: string;
   readonly masterQueue: string;
   readonly environment: RuntimeEnvironment | null;
   readonly deploymentId?: string;
   readonly backgroundWorkerId?: string;
 
-  // FIXME
+  // FIXME: Required for unmanaged workers
   readonly isLatestDeployment = true;
 
   constructor(opts: AuthenticatedWorkerInstanceOptions) {
     super({ prisma: opts.prisma, engine: opts.engine });
 
     this.type = opts.type;
+    this.name = opts.name;
     this.workerGroupId = opts.workerGroupId;
     this.workerInstanceId = opts.workerInstanceId;
     this.masterQueue = opts.masterQueue;
@@ -647,12 +654,7 @@ export class AuthenticatedWorkerInstance extends WithRunEngine {
       isWarmStart,
     });
 
-    const defaultMachinePreset = {
-      name: "small-1x",
-      cpu: 1,
-      memory: 1,
-      centsPerMs: 0,
-    } satisfies MachinePreset;
+    const defaultMachinePreset = machinePresetFromName(defaultMachine);
 
     const environment =
       this.environment ??
@@ -718,6 +720,7 @@ export class AuthenticatedWorkerInstance extends WithRunEngine {
     if (this.type === WorkerInstanceGroupType.MANAGED) {
       return {
         type: WorkerInstanceGroupType.MANAGED,
+        name: this.name,
         workerGroupId: this.workerGroupId,
         workerInstanceId: this.workerInstanceId,
         masterQueue: this.masterQueue,
@@ -726,6 +729,7 @@ export class AuthenticatedWorkerInstance extends WithRunEngine {
 
     return {
       type: WorkerInstanceGroupType.UNMANAGED,
+      name: this.name,
       workerGroupId: this.workerGroupId,
       workerInstanceId: this.workerInstanceId,
       masterQueue: this.masterQueue,
@@ -764,12 +768,14 @@ export class AuthenticatedWorkerInstance extends WithRunEngine {
 export type WorkerGroupTokenAuthenticationResponse =
   | {
       type: typeof WorkerInstanceGroupType.MANAGED;
+      name: string;
       workerGroupId: string;
       workerInstanceId: string;
       masterQueue: string;
     }
   | {
       type: typeof WorkerInstanceGroupType.UNMANAGED;
+      name: string;
       workerGroupId: string;
       workerInstanceId: string;
       masterQueue: string;

internal-packages/database/prisma/schema.prisma

@@ -423,7 +423,7 @@ model RuntimeEnvironment {
   taskRunNumberCounter TaskRunNumberCounter[]
   taskRunCheckpoints   TaskRunCheckpoint[]
   waitpoints           Waitpoint[]
-  workerInstance       WorkerInstance[]
+  workerInstances      WorkerInstance[]
 
   @@unique([projectId, slug, orgMemberId])
   @@unique([projectId, shortcode])

packages/worker/src/supervisor/schemas.ts

@@ -33,6 +33,10 @@ export type WorkerApiConnectRequestBody = z.infer<typeof WorkerApiConnectRequest
 
 export const WorkerApiConnectResponseBody = z.object({
   ok: z.literal(true),
+  workerGroup: z.object({
+    type: z.string(),
+    name: z.string(),
+  }),
 });
 export type WorkerApiConnectResponseBody = z.infer<typeof WorkerApiConnectResponseBody>;
 

packages/worker/src/supervisor/session.ts

@@ -100,17 +100,17 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
       extraHeaders: getDefaultWorkerHeaders(this.opts),
     });
     this.socket.on("run:notify", ({ version, run }) => {
-      console.log("[WorkerSession] Received run notification", { version, run });
+      console.log("[WorkerSession][WS] Received run notification", { version, run });
       this.emit("runNotification", { time: new Date(), run });
     });
     this.socket.on("connect", () => {
-      console.log("[WorkerSession] Connected to platform");
+      console.log("[WorkerSession][WS] Connected to platform");
     });
     this.socket.on("connect_error", (error) => {
-      console.error("[WorkerSession] Connection error", { error });
+      console.error("[WorkerSession][WS] Connection error", { error });
     });
     this.socket.on("disconnect", (reason, description) => {
-      console.log("[WorkerSession] Disconnected from platform", { reason, description });
+      console.log("[WorkerSession][WS] Disconnected from platform", { reason, description });
     });
   }
 
@@ -122,10 +122,17 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
     });
 
     if (!connect.success) {
-      console.error("[WorkerSession] Failed to connect via HTTP client", { error: connect.error });
-      throw new Error("[WorkerSession] Failed to connect via HTTP client");
+      console.error("[WorkerSession][HTTP] Failed to connect", { error: connect.error });
+      throw new Error("[WorkerSession][HTTP] Failed to connect");
     }
 
+    const { workerGroup } = connect.data;
+
+    console.log("[WorkerSession][HTTP] Connected to platform", {
+      type: workerGroup.type,
+      name: workerGroup.name,
+    });
+
     this.queueConsumer.start();
     this.heartbeatService.start();
     this.createSocket();