Handle revalidating JWT tokens

Author: ericallam

apps/webapp/app/routes/api.v1.usage.ingest.ts

@@ -2,7 +2,7 @@ import { ActionFunctionArgs } from "@remix-run/server-runtime";
 import { MachinePresetName } from "@trigger.dev/core/v3";
 import { z } from "zod";
 import { prisma } from "~/db.server";
-import { validateJWTToken } from "~/services/apiAuth.server";
+import { validateJWTTokenAndRenew } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { workerQueue } from "~/services/worker.server";
 import { machinePresetFromName } from "~/v3/machinePresets.server";
@@ -26,16 +26,12 @@ export async function action({ request }: ActionFunctionArgs) {
     return { status: 405, body: "Method Not Allowed" };
   }
 
-  const jwt = request.headers.get("x-trigger-jwt");
+  const jwtResult = await validateJWTTokenAndRenew(request, JWTPayloadSchema);
 
-  if (!jwt) {
+  if (!jwtResult) {
     return { status: 401, body: "Unauthorized" };
   }
 
-  logger.debug("Validating JWT", { jwt });
-
-  const jwtPayload = await validateJWTToken(jwt, JWTPayloadSchema);
-
   const rawJson = await request.json();
 
   const json = BodySchema.safeParse(rawJson);
@@ -46,16 +42,16 @@ export async function action({ request }: ActionFunctionArgs) {
     return { status: 400, body: "Bad Request" };
   }
 
-  const preset = machinePresetFromName(jwtPayload.machine_preset as MachinePresetName);
+  const preset = machinePresetFromName(jwtResult.payload.machine_preset as MachinePresetName);
 
-  logger.debug("Validated JWT", { jwtPayload, json: json.data, preset });
+  logger.debug("[/api/v1/usage/ingest] Reporting usage", { jwtResult, json: json.data, preset });
 
-  if (json.data.durationMs > 10) {
+  if (json.data.durationMs > 0) {
     const costInCents = json.data.durationMs * preset.centsPerMs;
 
     await prisma.taskRun.update({
       where: {
-        id: jwtPayload.run_id,
+        id: jwtResult.payload.run_id,
       },
       data: {
         usageDurationMs: {
@@ -71,7 +67,7 @@ export async function action({ request }: ActionFunctionArgs) {
       await reportUsageEvent({
         source: "webapp",
         type: "usage",
-        subject: jwtPayload.org_id,
+        subject: jwtResult.payload.org_id,
         data: {
           durationMs: json.data.durationMs,
           costInCents: String(costInCents),
@@ -81,7 +77,7 @@ export async function action({ request }: ActionFunctionArgs) {
       logger.error("Failed to report usage event, enqueing v3.reportUsage", { error: e });
 
       await workerQueue.enqueue("v3.reportUsage", {
-        orgId: jwtPayload.org_id,
+        orgId: jwtResult.payload.org_id,
         data: {
           costInCents: String(costInCents),
         },
@@ -94,5 +90,8 @@ export async function action({ request }: ActionFunctionArgs) {
 
   return new Response(null, {
     status: 200,
+    headers: {
+      "x-trigger-jwt": jwtResult.jwt,
+    },
   });
 }

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

@@ -13,8 +13,9 @@ import {
 import { prisma } from "~/db.server";
 import { json } from "@remix-run/server-runtime";
 import { findProjectByRef } from "~/models/project.server";
-import { SignJWT, jwtVerify } from "jose";
+import { SignJWT, jwtVerify, errors } from "jose";
 import { env } from "~/env.server";
+import { logger } from "./logger.server";
 
 type Optional<T, K extends keyof T> = Prettify<Omit<T, K> & Partial<Pick<T, K>>>;
 
@@ -213,40 +214,126 @@ export async function authenticatedEnvironmentForAuthentication(
   }
 }
 
+const JWT_SECRET = new TextEncoder().encode(env.SESSION_SECRET);
+const JWT_ALGORITHM = "HS256";
+const DEFAULT_JWT_EXPIRATION_IN_MS = 1000 * 60 * 60; // 1 hour
+
 export async function generateJWTTokenForEnvironment(
   environment: RuntimeEnvironment,
   payload: Record<string, string>
 ) {
-  const secret = new TextEncoder().encode(env.SESSION_SECRET);
-
-  const alg = "HS256";
-
   const jwt = await new SignJWT({
     environment_id: environment.id,
     org_id: environment.organizationId,
     project_id: environment.projectId,
     ...payload,
   })
-    .setProtectedHeader({ alg })
+    .setProtectedHeader({ alg: JWT_ALGORITHM })
     .setIssuedAt()
     .setIssuer("https://id.trigger.dev")
     .setAudience("https://api.trigger.dev")
-    .setExpirationTime("24h")
-    .sign(secret);
+    .setExpirationTime(calculateJWTExpiration())
+    .sign(JWT_SECRET);
 
   return jwt;
 }
 
-export async function validateJWTToken<T extends z.ZodTypeAny>(
-  jwt: string,
+export async function validateJWTTokenAndRenew<T extends z.ZodTypeAny>(
+  request: Request,
   payloadSchema: T
-): Promise<z.infer<T>> {
-  const secret = new TextEncoder().encode(env.SESSION_SECRET);
+): Promise<{ payload: z.infer<T>; jwt: string } | undefined> {
+  try {
+    const jwt = request.headers.get("x-trigger-jwt");
+
+    if (!jwt) {
+      logger.debug("Missing JWT token in request", {
+        headers: Object.fromEntries(request.headers),
+      });
+
+      return;
+    }
+
+    const { payload: rawPayload } = await jwtVerify(jwt, JWT_SECRET, {
+      issuer: "https://id.trigger.dev",
+      audience: "https://api.trigger.dev",
+    });
+
+    const payload = payloadSchema.safeParse(rawPayload);
+
+    if (!payload.success) {
+      logger.error("Failed to validate JWT", { payload: rawPayload, issues: payload.error.issues });
+
+      return;
+    }
+
+    const renewedJwt = await renewJWTToken(payload.data);
+
+    return {
+      payload: payload.data,
+      jwt: renewedJwt,
+    };
+  } catch (error) {
+    if (error instanceof errors.JWTExpired) {
+      // Now we need to try and renew the token using the API key auth
+      const authenticatedEnv = await authenticateApiRequest(request);
+
+      if (!authenticatedEnv) {
+        logger.error("Failed to renew JWT token, missing or invalid Authorization header", {
+          error: error.message,
+        });
+
+        return;
+      }
+
+      const payload = payloadSchema.safeParse(error.payload);
 
-  const { payload, protectedHeader } = await jwtVerify(jwt, secret, {
-    issuer: "https://id.trigger.dev",
-    audience: "https://api.trigger.dev",
-  });
+      if (!payload.success) {
+        logger.error("Failed to parse jwt payload after expired", {
+          payload: error.payload,
+          issues: payload.error.issues,
+        });
+
+        return;
+      }
+
+      const renewedJwt = await generateJWTTokenForEnvironment(authenticatedEnv.environment, {
+        ...payload.data,
+      });
+
+      logger.debug("Renewed JWT token from Authorization header API Key", {
+        environment: authenticatedEnv.environment,
+        payload: payload.data,
+      });
+
+      return {
+        payload: payload.data,
+        jwt: renewedJwt,
+      };
+    }
+
+    logger.error("Failed to validate JWT token", { error });
+  }
+}
+
+async function renewJWTToken(payload: Record<string, string>) {
+  const jwt = await new SignJWT(payload)
+    .setProtectedHeader({ alg: JWT_ALGORITHM })
+    .setIssuedAt()
+    .setIssuer("https://id.trigger.dev")
+    .setAudience("https://api.trigger.dev")
+    .setExpirationTime(calculateJWTExpiration())
+    .sign(JWT_SECRET);
+
+  return jwt;
+}
+
+function calculateJWTExpiration() {
+  if (env.PROD_USAGE_HEARTBEAT_INTERVAL_MS) {
+    return (
+      (Date.now() + Math.max(DEFAULT_JWT_EXPIRATION_IN_MS, env.PROD_USAGE_HEARTBEAT_INTERVAL_MS)) /
+      1000
+    );
+  }
 
-  return payloadSchema.parse(payload);
+  return (Date.now() + DEFAULT_JWT_EXPIRATION_IN_MS) / 1000;
 }

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

@@ -29,7 +29,7 @@ export async function reportUsageEvent(event: UsageEvent) {
 
   const url = `${env.USAGE_OPEN_METER_BASE_URL}/api/v1/events`;
 
-  logger.debug("Reporting usage event", { url, body });
+  logger.debug("Reporting usage event to OpenMeter", { url, body });
 
   const response = await fetch(url, {
     method: "POST",
@@ -42,6 +42,6 @@ export async function reportUsageEvent(event: UsageEvent) {
   });
 
   if (!response.ok) {
-    throw new Error(`Failed to report usage event: ${response.status} ${response.statusText}`);
+    logger.error(`Failed to report usage event: ${response.status} ${response.statusText}`);
   }
 }

packages/core/src/v3/usage/usageClient.ts

@@ -1,3 +1,5 @@
+import { apiClientManager } from "../apiClientManager-api";
+
 export type UsageClientOptions = {
   token: string;
   baseUrl: string;
@@ -10,20 +12,29 @@ export type UsageEvent = {
 export class UsageClient {
   constructor(
     private readonly url: string,
-    private readonly jwt: string
+    private jwt: string
   ) {}
 
   async sendUsageEvent(event: UsageEvent): Promise<void> {
     try {
-      await fetch(this.url, {
+      const response = await fetch(this.url, {
         method: "POST",
         body: JSON.stringify(event),
         headers: {
           "content-type": "application/json",
           "x-trigger-jwt": this.jwt,
           accept: "application/json",
+          authorization: `Bearer ${apiClientManager.accessToken}`, // this is used to renew the JWT
         },
       });
+
+      if (response.ok) {
+        const renewedJwt = response.headers.get("x-trigger-jwt");
+
+        if (renewedJwt) {
+          this.jwt = renewedJwt;
+        }
+      }
     } catch (error) {
       console.error(`Failed to send usage event: ${error}`);
     }

references/v3-catalog/src/trigger/longRunning.ts

@@ -3,12 +3,13 @@ import { logger, task, wait } from "@trigger.dev/sdk/v3";
 export const longRunning = task({
   id: "long-running",
   run: async (payload: { message: string }, { ctx }) => {
-    logger.info("Long running payloadddd", { payload });
+    logger.info("Long running", { payload });
 
-    // Wait for 3 minutes
-    await new Promise((resolve) => setTimeout(resolve, 5000));
+    await new Promise((resolve) => setTimeout(resolve, 20000));
+
+    await wait.for({ seconds: 10 });
 
-    await wait.for({ seconds: 5 });
+    await new Promise((resolve) => setTimeout(resolve, 20000));
   },
 });