fix(webapp): correctly generate JWT tokens for preview branches after triggering a run (fix #2678) (#2695)

Author: ericallam

apps/webapp/app/routes/api.v1.tasks.$taskId.trigger.ts

@@ -11,6 +11,7 @@ import { prisma } from "~/db.server";
 import { env } from "~/env.server";
 import { ApiAuthenticationResultSuccess, getOneTimeUseToken } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
+import { extractJwtSigningSecretKey } from "~/services/realtime/jwtAuth.server";
 import { determineRealtimeStreamsVersion } from "~/services/realtime/v1StreamsGlobal.server";
 import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server";
@@ -85,7 +86,7 @@ const { action, loader } = createActionApiRoute(
         isCached: false,
       }),
       buildResponseHeaders: async (responseBody, cachedEntity) => {
-        return await responseHeaders(cachedEntity, authentication, triggerClient);
+        return await responseHeaders(cachedEntity, authentication);
       },
     });
 
@@ -140,7 +141,12 @@ const { action, loader } = createActionApiRoute(
 
       await saveRequestIdempotency(requestIdempotencyKey, "trigger", result.run.id);
 
-      const $responseHeaders = await responseHeaders(result.run, authentication, triggerClient);
+      const $responseHeaders = await responseHeaders(result.run, authentication);
+
+      logger.debug("responseHeaders authentication", {
+        authentication,
+        responseHeaders: $responseHeaders,
+      });
 
       return json(
         {
@@ -170,39 +176,26 @@ const { action, loader } = createActionApiRoute(
 
 async function responseHeaders(
   run: Pick<TaskRun, "friendlyId">,
-  authentication: ApiAuthenticationResultSuccess,
-  triggerClient?: string | null
+  authentication: ApiAuthenticationResultSuccess
 ): Promise<Record<string, string>> {
   const { environment, realtime } = authentication;
 
-  const claimsHeader = JSON.stringify({
+  const claims = {
     sub: environment.id,
     pub: true,
+    scopes: [`read:runs:${run.friendlyId}`],
     realtime,
-  });
-
-  if (triggerClient === "browser") {
-    const claims = {
-      sub: environment.id,
-      pub: true,
-      scopes: [`read:runs:${run.friendlyId}`],
-      realtime,
-    };
-
-    const jwt = await internal_generateJWT({
-      secretKey: environment.apiKey,
-      payload: claims,
-      expirationTime: "1h",
-    });
+  };
 
-    return {
-      "x-trigger-jwt-claims": claimsHeader,
-      "x-trigger-jwt": jwt,
-    };
-  }
+  const jwt = await internal_generateJWT({
+    secretKey: extractJwtSigningSecretKey(environment),
+    payload: claims,
+    expirationTime: "1h",
+  });
 
   return {
-    "x-trigger-jwt-claims": claimsHeader,
+    "x-trigger-jwt-claims": JSON.stringify(claims),
+    "x-trigger-jwt": jwt,
   };
 }
 

apps/webapp/app/routes/api.v1.tasks.batch.ts

@@ -17,6 +17,7 @@ import {
 import { OutOfEntitlementError } from "~/v3/services/triggerTask.server";
 import { HeadersSchema } from "./api.v1.tasks.$taskId.trigger";
 import { determineRealtimeStreamsVersion } from "~/services/realtime/v1StreamsGlobal.server";
+import { extractJwtSigningSecretKey } from "~/services/realtime/jwtAuth.server";
 
 const { action, loader } = createActionApiRoute(
   {
@@ -163,7 +164,7 @@ async function responseHeaders(
     };
 
     const jwt = await generateJWT({
-      secretKey: environment.apiKey,
+      secretKey: extractJwtSigningSecretKey(environment),
       payload: claims,
       expirationTime: "1h",
     });

apps/webapp/app/routes/api.v2.tasks.batch.ts

@@ -19,6 +19,7 @@ import { BatchProcessingStrategy } from "~/v3/services/batchTriggerV3.server";
 import { OutOfEntitlementError } from "~/v3/services/triggerTask.server";
 import { HeadersSchema } from "./api.v1.tasks.$taskId.trigger";
 import { determineRealtimeStreamsVersion } from "~/services/realtime/v1StreamsGlobal.server";
+import { extractJwtSigningSecretKey } from "~/services/realtime/jwtAuth.server";
 
 const { action, loader } = createActionApiRoute(
   {
@@ -178,7 +179,7 @@ async function responseHeaders(
     };
 
     const jwt = await generateJWT({
-      secretKey: environment.apiKey,
+      secretKey: extractJwtSigningSecretKey(environment),
       payload: claims,
       expirationTime: "1h",
     });

apps/webapp/app/services/apiAuth.server.ts

@@ -236,6 +236,8 @@ async function authenticateApiKeyWithFailure(
     case "PUBLIC_JWT": {
       const validationResults = await validatePublicJwtKey(result.apiKey);
 
+      logger.debug("validatePublicJwtKey", { validationResults });
+
       if (!validationResults.ok) {
         return validationResults;
       }

apps/webapp/app/services/realtime/jwtAuth.server.ts

@@ -2,6 +2,7 @@ import { json } from "@remix-run/server-runtime";
 import { validateJWT } from "@trigger.dev/core/v3/jwt";
 import { findEnvironmentById } from "~/models/runtimeEnvironment.server";
 import { AuthenticatedEnvironment } from "../apiAuth.server";
+import { logger } from "../logger.server";
 
 export type ValidatePublicJwtKeySuccess = {
   ok: true;
@@ -38,6 +39,8 @@ export async function validatePublicJwtKey(token: string): Promise<ValidatePubli
     environment.parentEnvironment?.apiKey ?? environment.apiKey
   );
 
+  logger.debug("validateJWT result", { result });
+
   if (!result.ok) {
     switch (result.code) {
       case "ERR_JWT_EXPIRED": {
@@ -89,6 +92,12 @@ export function isPublicJWT(token: string): boolean {
   }
 }
 
+export function extractJwtSigningSecretKey(
+  environment: AuthenticatedEnvironment & { parentEnvironment?: { apiKey: string } }
+) {
+  return environment.parentEnvironment?.apiKey ?? environment.apiKey;
+}
+
 function extractJWTSub(token: string): string | undefined {
   // Split the token
   const parts = token.split(".");