WIP

Author: ericallam

apps/webapp/app/db.server.ts

@@ -107,14 +107,14 @@ function getClient() {
         emit: "stdout",
         level: "warn",
       },
-      {
-        emit: "stdout",
-        level: "query",
-      },
-      {
-        emit: "event",
-        level: "query",
-      },
+      // {
+      //   emit: "stdout",
+      //   level: "query",
+      // },
+      // {
+      //   emit: "event",
+      //   level: "query",
+      // },
     ],
   });
 

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

@@ -0,0 +1,35 @@
+import { ActionFunctionArgs } from "@remix-run/server-runtime";
+import { z } from "zod";
+import { validateJWTToken } from "~/services/apiAuth.server";
+import { logger } from "~/services/logger.server";
+
+const JWTPayloadSchema = z.object({
+  environment_id: z.string(),
+  org_id: z.string(),
+  project_id: z.string(),
+  run_id: z.string(),
+  machine_preset: z.string(),
+});
+
+export async function action({ request }: ActionFunctionArgs) {
+  // Ensure this is a POST request
+  if (request.method.toUpperCase() !== "POST") {
+    return { status: 405, body: "Method Not Allowed" };
+  }
+
+  const jwt = request.headers.get("x-trigger-jwt");
+
+  if (!jwt) {
+    return { status: 401, body: "Unauthorized" };
+  }
+
+  logger.debug("Validating JWT", { jwt });
+
+  const jwtPayload = await validateJWTToken(jwt, JWTPayloadSchema);
+
+  logger.debug("Validated JWT", { jwtPayload });
+
+  return new Response(null, {
+    status: 200,
+  });
+}

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

@@ -1,6 +1,7 @@
 import { Prettify } from "@trigger.dev/core";
 import { z } from "zod";
 import {
+  RuntimeEnvironment,
   findEnvironmentByApiKey,
   findEnvironmentByPublicApiKey,
 } from "~/models/runtimeEnvironment.server";
@@ -12,7 +13,8 @@ import {
 import { prisma } from "~/db.server";
 import { json } from "@remix-run/server-runtime";
 import { findProjectByRef } from "~/models/project.server";
-import { SignJWT } from "jose";
+import { SignJWT, jwtVerify } from "jose";
+import { env } from "~/env.server";
 
 type Optional<T, K extends keyof T> = Prettify<Omit<T, K> & Partial<Pick<T, K>>>;
 
@@ -211,14 +213,20 @@ export async function authenticatedEnvironmentForAuthentication(
   }
 }
 
-export async function generateJWTTokenForEnvironment(environment: AuthenticatedEnvironment) {
-  const secret = new TextEncoder().encode(
-    "cc7e0d44fd473002f1c42167459001140ec6389b7353f8088f4d9a95f2f596f2"
-  );
+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 })
+  const jwt = await new SignJWT({
+    environment_id: environment.id,
+    org_id: environment.organizationId,
+    project_id: environment.projectId,
+    ...payload,
+  })
     .setProtectedHeader({ alg })
     .setIssuedAt()
     .setIssuer("https://id.trigger.dev")
@@ -228,3 +236,20 @@ export async function generateJWTTokenForEnvironment(environment: AuthenticatedE
 
   return jwt;
 }
+
+export async function validateJWTToken<T extends z.ZodTypeAny>(
+  jwt: string,
+  payloadSchema: T
+): Promise<z.infer<T>> {
+  const secret = new TextEncoder().encode(env.SESSION_SECRET);
+
+  const { payload, protectedHeader } = await jwtVerify(jwt, secret, {
+    issuer: "https://id.trigger.dev",
+    audience: "https://api.trigger.dev",
+  });
+
+  console.log(protectedHeader);
+  console.log(payload);
+
+  return payloadSchema.parse(payload);
+}

apps/webapp/app/v3/environmentVariables/environmentVariablesRepository.server.ts

@@ -746,10 +746,6 @@ async function resolveBuiltInProdVariables(runtimeEnvironment: RuntimeEnvironmen
       key: "TRIGGER_ORG_ID",
       value: runtimeEnvironment.organizationId,
     },
-    {
-      key: "TRIGGER_MACHINE_PRESET",
-      value: "tiny-1x",
-    },
   ];
 
   if (env.PROD_OTEL_BATCH_PROCESSING_ENABLED === "1") {

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

@@ -15,6 +15,8 @@ import { ZodMessageSender } from "@trigger.dev/core/v3/zodMessageHandler";
 import {
   BackgroundWorker,
   BackgroundWorkerTask,
+  RuntimeEnvironment,
+  TaskRun,
   TaskRunAttemptStatus,
   TaskRunStatus,
 } from "@trigger.dev/database";
@@ -37,6 +39,8 @@ import { CrashTaskRunService } from "../services/crashTaskRun.server";
 import { CreateTaskRunAttemptService } from "../services/createTaskRunAttempt.server";
 import { RestoreCheckpointService } from "../services/restoreCheckpoint.server";
 import { tracer } from "../tracer.server";
+import { generateJWTTokenForEnvironment } from "~/services/apiAuth.server";
+import { EnvironmentVariable } from "../environmentVariables/repository";
 
 const WithTraceContext = z.object({
   traceparent: z.string().optional(),
@@ -1063,7 +1067,7 @@ class SharedQueueTasks {
       },
     };
 
-    const variables = await resolveVariablesForEnvironment(attempt.runtimeEnvironment);
+    const variables = await this.#buildEnvironmentVariables(attempt.runtimeEnvironment, taskRun);
 
     const payload: ProdTaskRunExecutionPayload = {
       execution,
@@ -1129,7 +1133,7 @@ class SharedQueueTasks {
       return;
     }
 
-    const variables = await resolveVariablesForEnvironment(environment);
+    const variables = await this.#buildEnvironmentVariables(environment, run);
 
     return {
       traceContext: run.traceContext as Record<string, unknown>,
@@ -1170,6 +1174,30 @@ class SharedQueueTasks {
 
     await service.call(completion.id, completion);
   }
+
+  async #buildEnvironmentVariables(
+    environment: RuntimeEnvironment,
+    run: TaskRun
+  ): Promise<Array<EnvironmentVariable>> {
+    const variables = await resolveVariablesForEnvironment(environment);
+
+    const jwt = await generateJWTTokenForEnvironment(environment, {
+      run_id: run.id,
+      machine_present: "tiny-1x",
+    });
+
+    return [
+      ...variables,
+      ...[
+        { key: "TRIGGER_JWT", value: jwt },
+        { key: "TRIGGER_RUN_ID", value: run.id },
+        {
+          key: "TRIGGER_MACHINE_PRESET",
+          value: "tiny-1x",
+        },
+      ],
+    ];
+  }
 }
 
 export const sharedQueueTasks = singleton("sharedQueueTasks", () => new SharedQueueTasks());

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

@@ -49,22 +49,13 @@ import {
 import { ProdRuntimeManager } from "@trigger.dev/core/v3/prod";
 
 const heartbeatIntervalMs = getEnvVar("USAGE_HEARTBEAT_INTERVAL_MS");
-const subject = getEnvVar("TRIGGER_ORG_ID");
-const machinePreset = getEnvVar("TRIGGER_MACHINE_PRESET");
-const openMeterApiKey = getEnvVar("USAGE_OPEN_METER_API_KEY");
-const openMeterBaseUrl = getEnvVar("USAGE_OPEN_METER_BASE_URL");
+const usageEventUrl = getEnvVar("USAGE_EVENT_URL");
+const triggerJWT = getEnvVar("TRIGGER_JWT");
 
 const prodUsageManager = new ProdUsageManager(new DevUsageManager(), {
   heartbeatIntervalMs: heartbeatIntervalMs ? parseInt(heartbeatIntervalMs, 10) : undefined,
-  subject: subject!,
-  machinePreset: machinePreset,
-  client:
-    openMeterApiKey && openMeterBaseUrl
-      ? {
-          token: openMeterApiKey,
-          baseUrl: openMeterBaseUrl,
-        }
-      : undefined,
+  url: usageEventUrl,
+  jwt: triggerJWT,
 });
 
 usage.setGlobalUsageManager(prodUsageManager);

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

@@ -1,12 +1,11 @@
 import { setInterval } from "node:timers/promises";
 import { UsageManager, UsageMeasurement, UsageSample } from "./types";
-import { UsageClient, UsageClientOptions } from "./usageClient";
+import { UsageClient } from "./usageClient";
 
 export type ProdUsageManagerOptions = {
   heartbeatIntervalMs?: number;
-  client?: UsageClientOptions;
-  subject: string;
-  machinePreset?: string;
+  url?: string;
+  jwt?: string;
 };
 
 export class ProdUsageManager implements UsageManager {
@@ -19,13 +18,13 @@ export class ProdUsageManager implements UsageManager {
     private readonly delegageUsageManager: UsageManager,
     private readonly options: ProdUsageManagerOptions
   ) {
-    if (typeof this.options.client !== "undefined") {
-      this._usageClient = new UsageClient(this.options.client);
+    if (this.options.url && this.options.jwt) {
+      this._usageClient = new UsageClient(this.options.url, this.options.jwt);
     }
   }
 
   get isReportingEnabled() {
-    return typeof this.options.client !== "undefined";
+    return typeof this._usageClient !== "undefined";
   }
 
   disable(): void {
@@ -95,9 +94,6 @@ export class ProdUsageManager implements UsageManager {
     }
 
     const sample = this._measurement.sample();
-    const wallTimeSinceLastSample = this._lastSample
-      ? sample.wallTime - this._lastSample.wallTime
-      : sample.wallTime;
 
     const cpuTimeSinceLastSample = this._lastSample
       ? sample.cpuTime - this._lastSample.cpuTime
@@ -106,10 +102,7 @@ export class ProdUsageManager implements UsageManager {
     this._lastSample = sample;
 
     console.log("Reporting usage", {
-      wallTimeSinceLastSample,
       cpuTimeSinceLastSample,
-      subject: this.options.subject,
-      machine: this.options.machinePreset,
     });
 
     if (cpuTimeSinceLastSample <= 0) {
@@ -118,24 +111,12 @@ export class ProdUsageManager implements UsageManager {
 
     const now = performance.now();
 
-    const event = {
-      source: "prod-usage-manager",
-      type: "usage",
-      subject: this.options.subject,
-      data: {
-        durationMs: cpuTimeSinceLastSample,
-        wallTimeInMs: wallTimeSinceLastSample,
-        machinePreset: this.options.machinePreset ?? "unknown",
-      },
-    };
-
-    await client.sendUsageEvent(event);
+    await client.sendUsageEvent({ durationMs: cpuTimeSinceLastSample });
 
     const durationInMs = performance.now() - now;
 
     console.log("Reported usage", {
       durationInMs,
-      event,
     });
   }
 }

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

@@ -4,49 +4,28 @@ export type UsageClientOptions = {
 };
 
 export type UsageEvent = {
-  source: string;
-  subject: string;
-  type: string;
-  id?: string;
-  time?: Date;
-  data?: Record<string, unknown>;
+  durationMs: number;
 };
 
 export class UsageClient {
-  constructor(private readonly options: UsageClientOptions) {}
+  constructor(
+    private readonly url: string,
+    private readonly jwt: string
+  ) {}
 
   async sendUsageEvent(event: UsageEvent): Promise<void> {
-    const body = {
-      specversion: "1.0",
-      id: event.id ?? globalThis.crypto.randomUUID(),
-      source: event.source,
-      type: event.type,
-      time: (event.time ?? new Date()).toISOString(),
-      subject: event.subject,
-      datacontenttype: "application/json",
-      data: event.data,
-    };
-
-    const url = `${this.baseUrl}/api/v1/events`;
-
     try {
-      await fetch(url, {
+      await fetch(this.url, {
         method: "POST",
-        body: JSON.stringify(body),
+        body: JSON.stringify(event),
         headers: {
-          "Content-Type": "application/cloudevents+json",
-          Authorization: `Bearer ${this.token}`,
-          Accept: "application/json",
+          "content-type": "application/json",
+          "x-trigger-jwt": `Bearer ${this.jwt}`,
+          accept: "application/json",
         },
       });
-    } catch {}
-  }
-
-  get baseUrl(): string {
-    return this.options.baseUrl;
-  }
-
-  private get token(): string {
-    return this.options.token;
+    } catch (error) {
+      console.error(`Failed to send usage event: ${error}`);
+    }
   }
 }