Cache worker group authentication and remove old self-hosted worker code (only managed is currently supported)

Author: ericallam

apps/webapp/app/routes/engine.v1.worker-actions.dequeue.ts

@@ -9,7 +9,10 @@ export const action = createActionWorkerApiRoute(
   {
     body: WorkerApiDequeueRequestBody, // Even though we don't use it, we need to keep it for backwards compatibility
   },
-  async ({ authenticatedWorker }): Promise<TypedResponse<WorkerApiDequeueResponseBody>> => {
-    return json(await authenticatedWorker.dequeue());
+  async ({
+    authenticatedWorker,
+    runnerId,
+  }): Promise<TypedResponse<WorkerApiDequeueResponseBody>> => {
+    return json(await authenticatedWorker.dequeue({ runnerId }));
   }
 );

apps/webapp/app/routes/engine.v1.worker-actions.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.attempts.complete.ts

@@ -18,6 +18,7 @@ export const action = createActionWorkerApiRoute(
     authenticatedWorker,
     body,
     params,
+    runnerId,
   }): Promise<TypedResponse<WorkerApiRunAttemptCompleteResponseBody>> => {
     const { completion } = body;
     const { runFriendlyId, snapshotFriendlyId } = params;
@@ -26,6 +27,7 @@ export const action = createActionWorkerApiRoute(
       runFriendlyId,
       snapshotFriendlyId,
       completion,
+      runnerId,
     });
 
     return json({ result: completeResult });

apps/webapp/app/routes/engine.v1.worker-actions.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.attempts.start.ts

@@ -18,13 +18,15 @@ export const action = createActionWorkerApiRoute(
     authenticatedWorker,
     body,
     params,
+    runnerId,
   }): Promise<TypedResponse<WorkerApiRunAttemptStartResponseBody>> => {
     const { runFriendlyId, snapshotFriendlyId } = params;
 
     const runExecutionData = await authenticatedWorker.startRunAttempt({
       runFriendlyId,
       snapshotFriendlyId,
       isWarmStart: body.isWarmStart,
+      runnerId,
     });
 
     return json(runExecutionData);

apps/webapp/app/routes/engine.v1.worker-actions.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.continue.ts

@@ -14,6 +14,7 @@ export const loader = createLoaderWorkerApiRoute(
   async ({
     authenticatedWorker,
     params,
+    runnerId,
   }): Promise<TypedResponse<WorkerApiContinueRunExecutionRequestBody>> => {
     const { runFriendlyId, snapshotFriendlyId } = params;
 
@@ -23,6 +24,7 @@ export const loader = createLoaderWorkerApiRoute(
       const continuationResult = await authenticatedWorker.continueRunExecution({
         runFriendlyId,
         snapshotFriendlyId,
+        runnerId,
       });
 
       return json(continuationResult);

apps/webapp/app/routes/engine.v1.worker-actions.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.heartbeat.ts

@@ -13,12 +13,14 @@ export const action = createActionWorkerApiRoute(
   async ({
     authenticatedWorker,
     params,
+    runnerId,
   }): Promise<TypedResponse<WorkloadHeartbeatResponseBody>> => {
     const { runFriendlyId, snapshotFriendlyId } = params;
 
     await authenticatedWorker.heartbeatRun({
       runFriendlyId,
       snapshotFriendlyId,
+      runnerId,
     });
 
     return json({ ok: true });

apps/webapp/app/routes/engine.v1.worker-actions.runs.$runFriendlyId.snapshots.$snapshotFriendlyId.suspend.ts

@@ -19,6 +19,7 @@ export const action = createActionWorkerApiRoute(
     authenticatedWorker,
     params,
     body,
+    runnerId,
   }): Promise<TypedResponse<WorkerApiSuspendRunResponseBody>> => {
     const { runFriendlyId, snapshotFriendlyId } = params;
 
@@ -39,6 +40,7 @@ export const action = createActionWorkerApiRoute(
         runFriendlyId,
         snapshotFriendlyId,
         checkpoint: body.checkpoint,
+        runnerId,
       });
 
       return json({ ok: true });

apps/webapp/app/services/routeBuilders/apiBuilder.server.ts

@@ -22,6 +22,7 @@ import {
   WorkerGroupTokenService,
 } from "~/v3/services/worker/workerGroupTokenService.server";
 import { API_VERSIONS, getApiVersion } from "~/api/versions";
+import { WORKER_HEADERS } from "@trigger.dev/core/v3/runEngineWorker";
 
 type AnyZodSchema = z.ZodFirstPartySchemaTypes | z.ZodDiscriminatedUnion<any, any>;
 
@@ -795,6 +796,7 @@ type WorkerLoaderHandlerFunction<
   headers: THeadersSchema extends z.ZodFirstPartySchemaTypes | z.ZodDiscriminatedUnion<any, any>
     ? z.infer<THeadersSchema>
     : undefined;
+  runnerId?: string;
 }) => Promise<Response>;
 
 export function createLoaderWorkerApiRoute<
@@ -858,12 +860,15 @@ export function createLoaderWorkerApiRoute<
         parsedHeaders = headers.data;
       }
 
+      const runnerId = request.headers.get(WORKER_HEADERS.RUNNER_ID) ?? undefined;
+
       const result = await handler({
         params: parsedParams,
         searchParams: parsedSearchParams,
         authenticatedWorker: authenticationResult,
         request,
         headers: parsedHeaders,
+        runnerId,
       });
       return result;
     } catch (error) {
@@ -924,6 +929,7 @@ type WorkerActionHandlerFunction<
   body: TBodySchema extends z.ZodFirstPartySchemaTypes | z.ZodDiscriminatedUnion<any, any>
     ? z.infer<TBodySchema>
     : undefined;
+  runnerId?: string;
 }) => Promise<Response>;
 
 export function createActionWorkerApiRoute<
@@ -1021,13 +1027,16 @@ export function createActionWorkerApiRoute<
         parsedBody = parsed.data;
       }
 
+      const runnerId = request.headers.get(WORKER_HEADERS.RUNNER_ID) ?? undefined;
+
       const result = await handler({
         params: parsedParams,
         searchParams: parsedSearchParams,
         authenticatedWorker: authenticationResult,
         request,
         body: parsedBody,
         headers: parsedHeaders,
+        runnerId,
       });
       return result;
     } catch (error) {