Use x-trigger-request-idempotency-key header instead

Plus cleanup shared logic in the routes

Author: ericallam

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

@@ -9,15 +9,14 @@ import { z } from "zod";
 import { prisma } from "~/db.server";
 import { env } from "~/env.server";
 import { EngineServiceValidationError } from "~/runEngine/concerns/errors";
-import {
-  ApiAuthenticationResultSuccess,
-  AuthenticatedEnvironment,
-  getOneTimeUseToken,
-} from "~/services/apiAuth.server";
+import { ApiAuthenticationResultSuccess, getOneTimeUseToken } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
-import { requestIdempotency } from "~/services/requestIdempotencyInstance.server";
 import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server";
+import {
+  handleRequestIdempotency,
+  saveRequestIdempotency,
+} from "~/utils/requestIdempotency.server";
 import { ServiceValidationError } from "~/v3/services/baseService.server";
 import { OutOfEntitlementError, TriggerTaskService } from "~/v3/services/triggerTask.server";
 
@@ -33,7 +32,7 @@ export const HeadersSchema = z.object({
   "x-trigger-worker": z.string().nullish(),
   "x-trigger-client": z.string().nullish(),
   "x-trigger-engine-version": RunEngineVersionSchema.nullish(),
-  "x-trigger-request-id": z.string().nullish(),
+  "x-trigger-request-idempotency-key": z.string().nullish(),
   traceparent: z.string().optional(),
   tracestate: z.string().optional(),
 });
@@ -63,51 +62,32 @@ const { action, loader } = createActionApiRoute(
       "x-trigger-worker": isFromWorker,
       "x-trigger-client": triggerClient,
       "x-trigger-engine-version": engineVersion,
-      "x-trigger-request-id": requestId,
+      "x-trigger-request-idempotency-key": requestIdempotencyKey,
     } = headers;
 
-    if (requestId) {
-      logger.debug("request-idempotency: checking for cached trigger request", {
-        requestId,
-      });
-
-      const cachedRequest = await requestIdempotency.checkRequest("trigger", requestId);
-
-      if (cachedRequest) {
-        logger.info("request-idempotency: found cached trigger request", {
-          requestId,
-          cachedRequest,
-        });
-
-        const cachedRun = await prisma.taskRun.findFirst({
+    const cachedResponse = await handleRequestIdempotency(requestIdempotencyKey, {
+      requestType: "trigger",
+      findCachedEntity: async (cachedRequestId) => {
+        return await prisma.taskRun.findFirst({
           where: {
-            id: cachedRequest.id,
+            id: cachedRequestId,
           },
           select: {
             friendlyId: true,
           },
         });
+      },
+      buildResponse: (cachedRun) => ({
+        id: cachedRun.friendlyId,
+        isCached: false,
+      }),
+      buildResponseHeaders: async (responseBody, cachedEntity) => {
+        return await responseHeaders(cachedEntity, authentication, triggerClient);
+      },
+    });
 
-        if (cachedRun) {
-          logger.info("request-idempotency: found cached trigger run", {
-            requestId,
-            cachedRun,
-          });
-
-          const $responseHeaders = await responseHeaders(cachedRun, authentication, triggerClient);
-
-          return json(
-            {
-              id: cachedRun.friendlyId,
-              isCached: false,
-            },
-            {
-              headers: $responseHeaders,
-              status: 200,
-            }
-          );
-        }
-      }
+    if (cachedResponse) {
+      return cachedResponse;
     }
 
     const service = new TriggerTaskService();
@@ -152,11 +132,7 @@ const { action, loader } = createActionApiRoute(
         return json({ error: "Task not found" }, { status: 404 });
       }
 
-      if (requestId) {
-        await requestIdempotency.saveRequest("trigger", requestId, {
-          id: result.run.id,
-        });
-      }
+      await saveRequestIdempotency(requestIdempotencyKey, "trigger", result.run.id);
 
       const $responseHeaders = await responseHeaders(result.run, authentication, triggerClient);
 
@@ -167,7 +143,7 @@ const { action, loader } = createActionApiRoute(
         },
         {
           headers: $responseHeaders,
-          status: 200,
+          status: 408,
         }
       );
     } catch (error) {

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

@@ -4,18 +4,20 @@ import {
   BatchTriggerTaskV3Response,
   generateJWT,
 } from "@trigger.dev/core/v3";
+import { prisma } from "~/db.server";
 import { env } from "~/env.server";
+import { RunEngineBatchTriggerService } from "~/runEngine/services/batchTrigger.server";
 import { AuthenticatedEnvironment, getOneTimeUseToken } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
+import {
+  handleRequestIdempotency,
+  saveRequestIdempotency,
+} from "~/utils/requestIdempotency.server";
 import { ServiceValidationError } from "~/v3/services/baseService.server";
 import { BatchProcessingStrategy } from "~/v3/services/batchTriggerV3.server";
 import { OutOfEntitlementError } from "~/v3/services/triggerTask.server";
 import { HeadersSchema } from "./api.v1.tasks.$taskId.trigger";
-import { RunEngineBatchTriggerService } from "~/runEngine/services/batchTrigger.server";
-import { z } from "zod";
-import { requestIdempotency } from "~/services/requestIdempotencyInstance.server";
-import { prisma } from "~/db.server";
 
 const { action, loader } = createActionApiRoute(
   {
@@ -56,7 +58,7 @@ const { action, loader } = createActionApiRoute(
       "x-trigger-client": triggerClient,
       "x-trigger-engine-version": engineVersion,
       "batch-processing-strategy": batchProcessingStrategy,
-      "x-trigger-request-id": requestId,
+      "x-trigger-request-idempotency-key": requestIdempotencyKey,
       traceparent,
       tracestate,
     } = headers;
@@ -71,55 +73,34 @@ const { action, loader } = createActionApiRoute(
       traceparent,
       tracestate,
       batchProcessingStrategy,
-      requestId,
+      requestIdempotencyKey,
     });
 
-    if (requestId) {
-      logger.debug("request-idempotency: checking for cached batch-trigger request", {
-        requestId,
-      });
-
-      const cachedRequest = await requestIdempotency.checkRequest("batch-trigger", requestId);
-
-      if (cachedRequest) {
-        logger.info("request-idempotency: found cached batch-trigger request", {
-          requestId,
-          cachedRequest,
-        });
-
-        // Find the batch and return it as a 200 response
-        const cachedBatch = await prisma.batchTaskRun.findFirst({
+    const cachedResponse = await handleRequestIdempotency(requestIdempotencyKey, {
+      requestType: "batch-trigger",
+      findCachedEntity: async (cachedRequestId) => {
+        return await prisma.batchTaskRun.findFirst({
           where: {
-            id: cachedRequest.id,
+            id: cachedRequestId,
             runtimeEnvironmentId: authentication.environment.id,
           },
           select: {
             friendlyId: true,
             runCount: true,
           },
         });
+      },
+      buildResponse: (cachedBatch) => ({
+        id: cachedBatch.friendlyId,
+        runCount: cachedBatch.runCount,
+      }),
+      buildResponseHeaders: async (responseBody, cachedEntity) => {
+        return await responseHeaders(responseBody, authentication.environment, triggerClient);
+      },
+    });
 
-        if (cachedBatch) {
-          logger.info("request-idempotency: found cached batch-trigger request", {
-            requestId,
-            cachedRequest,
-            cachedBatch,
-          });
-
-          const responseBody = {
-            id: cachedBatch.friendlyId,
-            runCount: cachedBatch.runCount,
-          };
-
-          const $responseHeaders = await responseHeaders(
-            responseBody,
-            authentication.environment,
-            triggerClient
-          );
-
-          return json(responseBody, { status: 200, headers: $responseHeaders });
-        }
-      }
+    if (cachedResponse) {
+      return cachedResponse;
     }
 
     const traceContext =
@@ -130,11 +111,7 @@ const { action, loader } = createActionApiRoute(
     const service = new RunEngineBatchTriggerService(batchProcessingStrategy ?? undefined);
 
     service.onBatchTaskRunCreated.attachOnce(async (batch) => {
-      if (requestId) {
-        await requestIdempotency.saveRequest("batch-trigger", requestId, {
-          id: batch.id,
-        });
-      }
+      await saveRequestIdempotency(requestIdempotencyKey, "batch-trigger", batch.id);
     });
 
     try {

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

@@ -4,6 +4,7 @@ import { MemoryStore } from "@unkey/cache/stores";
 import { RedisCacheStore } from "./unkey/redisCacheStore.server";
 import { RedisWithClusterOptions } from "~/redis.server";
 import { validate as uuidValidate, version as uuidVersion } from "uuid";
+import { startActiveSpan } from "~/v3/tracer.server";
 
 export type RequestIdempotencyServiceOptions<TTypes extends string> = {
   types: TTypes[];
@@ -53,37 +54,46 @@ export class RequestIdempotencyService<TTypes extends string> {
     this.cache = cache;
   }
 
-  async checkRequest(type: TTypes, requestId: string) {
-    if (!this.#validateRequestId(requestId)) {
-      this.logger.warn("RequestIdempotency: invalid requestId", {
-        requestId,
+  async checkRequest(type: TTypes, requestIdempotencyKey: string) {
+    if (!this.#validateRequestId(requestIdempotencyKey)) {
+      this.logger.warn("RequestIdempotency: invalid requestIdempotencyKey", {
+        requestIdempotencyKey,
       });
 
       return undefined;
     }
 
-    const key = `${type}:${requestId}`;
-    const result = await this.cache.requests.get(key);
+    return startActiveSpan("RequestIdempotency.checkRequest()", async (span) => {
+      span.setAttribute("request_id", requestIdempotencyKey);
+      span.setAttribute("type", type);
 
-    this.logger.debug("RequestIdempotency: checking request", {
-      type,
-      requestId,
-      key,
-      result,
-    });
+      const key = `${type}:${requestIdempotencyKey}`;
+      const result = await this.cache.requests.get(key);
+
+      this.logger.debug("RequestIdempotency: checking request", {
+        type,
+        requestIdempotencyKey,
+        key,
+        result,
+      });
 
-    return result.val ? result.val : undefined;
+      return result.val ? result.val : undefined;
+    });
   }
 
-  async saveRequest(type: TTypes, requestId: string, value: RequestIdempotencyCacheEntry) {
-    if (!this.#validateRequestId(requestId)) {
-      this.logger.warn("RequestIdempotency: invalid requestId", {
-        requestId,
+  async saveRequest(
+    type: TTypes,
+    requestIdempotencyKey: string,
+    value: RequestIdempotencyCacheEntry
+  ) {
+    if (!this.#validateRequestId(requestIdempotencyKey)) {
+      this.logger.warn("RequestIdempotency: invalid requestIdempotencyKey", {
+        requestIdempotencyKey,
       });
       return undefined;
     }
 
-    const key = `${type}:${requestId}`;
+    const key = `${type}:${requestIdempotencyKey}`;
     const result = await this.cache.requests.set(key, value);
 
     if (result.err) {
@@ -94,7 +104,7 @@ export class RequestIdempotencyService<TTypes extends string> {
     } else {
       this.logger.debug("RequestIdempotency: saved request", {
         type,
-        requestId,
+        requestIdempotencyKey,
         key,
         value,
       });
@@ -103,9 +113,9 @@ export class RequestIdempotencyService<TTypes extends string> {
     return result;
   }
 
-  // The requestId should be a valid UUID
-  #validateRequestId(requestId: string): boolean {
-    return isValidV4UUID(requestId);
+  // The requestIdempotencyKey should be a valid UUID
+  #validateRequestId(requestIdempotencyKey: string): boolean {
+    return isValidV4UUID(requestIdempotencyKey);
   }
 }