Large waitpoints now get offloaded

matt-aitken · matt-aitken · commit adab2252c860 · 2026-03-31T16:36:12.000+01:00
diff --git a/apps/webapp/app/routes/api.v1.waitpoints.tokens.$waitpointFriendlyId.callback.$hash.ts b/apps/webapp/app/routes/api.v1.waitpoints.tokens.$waitpointFriendlyId.callback.$hash.ts
@@ -1,13 +1,11 @@
 import { type ActionFunctionArgs, json } from "@remix-run/server-runtime";
-import {
-  type CompleteWaitpointTokenResponseBody,
-  conditionallyExportPacket,
-  stringifyIO,
-} from "@trigger.dev/core/v3";
+import { type CompleteWaitpointTokenResponseBody, stringifyIO } from "@trigger.dev/core/v3";
 import { WaitpointId } from "@trigger.dev/core/v3/isomorphic";
 import { z } from "zod";
 import { $replica } from "~/db.server";
 import { env } from "~/env.server";
+import { processWaitpointCompletionPacket } from "~/runEngine/concerns/waitpointCompletionPacket.server";
+import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { verifyHttpCallbackHash } from "~/services/httpCallback.server";
 import { logger } from "~/services/logger.server";
 import { engine } from "~/v3/runEngine.server";
@@ -41,8 +39,10 @@ export async function action({ request, params }: ActionFunctionArgs) {
       },
       include: {
         environment: {
-          select: {
-            apiKey: true,
+          include: {
+            project: true,
+            organization: true,
+            orgMember: true,
             parentEnvironment: {
               select: {
                 apiKey: true,
@@ -77,9 +77,10 @@ export async function action({ request, params }: ActionFunctionArgs) {
     const body = await request.json().catch(() => ({}));
 
     const stringifiedData = await stringifyIO(body);
-    const finalData = await conditionallyExportPacket(
+    const finalData = await processWaitpointCompletionPacket(
       stringifiedData,
-      `${waitpointId}/waitpoint/http-callback`
+      waitpoint.environment,
+      `${WaitpointId.toFriendlyId(waitpointId)}/http-callback`
     );
 
     const result = await engine.completeWaitpoint({
diff --git a/apps/webapp/app/routes/api.v1.waitpoints.tokens.$waitpointFriendlyId.complete.ts b/apps/webapp/app/routes/api.v1.waitpoints.tokens.$waitpointFriendlyId.complete.ts
@@ -2,14 +2,14 @@ import { json } from "@remix-run/server-runtime";
 import {
   CompleteWaitpointTokenRequestBody,
   type CompleteWaitpointTokenResponseBody,
-  conditionallyExportPacket,
   stringifyIO,
 } from "@trigger.dev/core/v3";
 import { WaitpointId } from "@trigger.dev/core/v3/isomorphic";
 import { z } from "zod";
 import { $replica } from "~/db.server";
 import { env } from "~/env.server";
 import { logger } from "~/services/logger.server";
+import { processWaitpointCompletionPacket } from "~/runEngine/concerns/waitpointCompletionPacket.server";
 import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 import { engine } from "~/v3/runEngine.server";
 
@@ -52,9 +52,10 @@ const { action, loader } = createActionApiRoute(
       }
 
       const stringifiedData = await stringifyIO(body.data);
-      const finalData = await conditionallyExportPacket(
+      const finalData = await processWaitpointCompletionPacket(
         stringifiedData,
-        `${waitpointId}/waitpoint/token`
+        authentication.environment,
+        `${WaitpointId.toFriendlyId(waitpointId)}/token`
       );
 
       const result = await engine.completeWaitpoint({
diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.waitpoints.$waitpointFriendlyId.complete/route.tsx b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.waitpoints.$waitpointFriendlyId.complete/route.tsx
@@ -2,13 +2,7 @@ import { env } from "~/env.server";
 import { parse } from "@conform-to/zod";
 import { Form, useLocation, useNavigation, useSubmit } from "@remix-run/react";
 import { type ActionFunctionArgs, json } from "@remix-run/server-runtime";
-import {
-  conditionallyExportPacket,
-  IOPacket,
-  stringifyIO,
-  timeoutError,
-  WaitpointTokenStatus,
-} from "@trigger.dev/core/v3";
+import { stringifyIO, timeoutError, WaitpointTokenStatus } from "@trigger.dev/core/v3";
 import { WaitpointId } from "@trigger.dev/core/v3/isomorphic";
 import type { Waitpoint } from "@trigger.dev/database";
 import { useCallback, useRef } from "react";
@@ -24,6 +18,8 @@ import { $replica } from "~/db.server";
 import { useOrganization } from "~/hooks/useOrganizations";
 import { useProject } from "~/hooks/useProject";
 import { redirectWithErrorMessage, redirectWithSuccessMessage } from "~/models/message.server";
+import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server";
+import { processWaitpointCompletionPacket } from "~/runEngine/concerns/waitpointCompletionPacket.server";
 import { logger } from "~/services/logger.server";
 import { requireUserId } from "~/services/session.server";
 import { EnvironmentParamSchema, ProjectParamSchema, v3RunsPath } from "~/utils/pathBuilder";
@@ -86,6 +82,7 @@ export const action = async ({ request, params }: ActionFunctionArgs) => {
     const waitpoint = await $replica.waitpoint.findFirst({
       select: {
         projectId: true,
+        environmentId: true,
       },
       where: {
         id: waitpointId,
@@ -150,11 +147,29 @@ export const action = async ({ request, params }: ActionFunctionArgs) => {
             );
           }
 
+          const environment = await findEnvironmentBySlug(project.id, envParam, userId);
+          if (!environment) {
+            return redirectWithErrorMessage(
+              submission.value.failureRedirect,
+              request,
+              "Environment not found"
+            );
+          }
+
+          if (environment.id !== waitpoint.environmentId) {
+            return redirectWithErrorMessage(
+              submission.value.failureRedirect,
+              request,
+              "No waitpoint found"
+            );
+          }
+
           const data = submission.value.payload ? JSON.parse(submission.value.payload) : {};
           const stringifiedData = await stringifyIO(data);
-          const finalData = await conditionallyExportPacket(
+          const finalData = await processWaitpointCompletionPacket(
             stringifiedData,
-            `${waitpointId}/waitpoint/token`
+            environment,
+            `${WaitpointId.toFriendlyId(waitpointId)}/token`
           );
 
           const result = await engine.completeWaitpoint({
diff --git a/apps/webapp/app/runEngine/concerns/waitpointCompletionPacket.server.ts b/apps/webapp/app/runEngine/concerns/waitpointCompletionPacket.server.ts
@@ -0,0 +1,69 @@
+import { IOPacket, packetRequiresOffloading, tryCatch } from "@trigger.dev/core/v3";
+import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
+import { env } from "~/env.server";
+import { uploadPacketToObjectStore } from "~/v3/objectStore.server";
+import { ServiceValidationError } from "~/v3/services/common.server";
+import { startActiveSpan } from "~/v3/tracer.server";
+
+function packetExtensionForDataType(dataType: string): string {
+  switch (dataType) {
+    case "application/json":
+    case "application/super+json":
+      return "json";
+    case "text/plain":
+      return "txt";
+    default:
+      return "txt";
+  }
+}
+
+/**
+ * Offloads large waitpoint completion payloads to object store (same threshold and
+ * upload path pattern as DefaultPayloadProcessor). Object key prefix should use the
+ * waitpoint friendly id folder, e.g. `${WaitpointId.toFriendlyId(internalId)}/token`.
+ * Replaces no-op conditionallyExportPacket usage in webapp routes where apiClientManager is unset.
+ */
+export async function processWaitpointCompletionPacket(
+  packet: IOPacket,
+  environment: AuthenticatedEnvironment,
+  pathPrefix: string
+): Promise<IOPacket> {
+  return await startActiveSpan("processWaitpointCompletionPacket()", async (span) => {
+    if (!packet.data) {
+      return packet;
+    }
+
+    const { needsOffloading, size } = packetRequiresOffloading(
+      packet,
+      env.TASK_PAYLOAD_OFFLOAD_THRESHOLD
+    );
+
+    span.setAttribute("needsOffloading", needsOffloading);
+    span.setAttribute("size", size);
+
+    if (!needsOffloading) {
+      return packet;
+    }
+
+    const filename = `${pathPrefix}.${packetExtensionForDataType(packet.dataType)}`;
+
+    const [uploadError, uploadedFilename] = await tryCatch(
+      uploadPacketToObjectStore(
+        filename,
+        packet.data,
+        packet.dataType,
+        environment,
+        env.OBJECT_STORE_DEFAULT_PROTOCOL
+      )
+    );
+
+    if (uploadError) {
+      throw new ServiceValidationError("Failed to upload large waitpoint to object store", 500);
+    }
+
+    return {
+      data: uploadedFilename!,
+      dataType: "application/store",
+    };
+  });
+}
diff --git a/apps/webapp/test/objectStore.test.ts b/apps/webapp/test/objectStore.test.ts
@@ -3,6 +3,7 @@ import { type IOPacket } from "@trigger.dev/core/v3";
 import { type PrismaClient } from "@trigger.dev/database";
 import { afterAll, describe, expect, it, vi } from "vitest";
 import { env } from "~/env.server";
+import { processWaitpointCompletionPacket } from "~/runEngine/concerns/waitpointCompletionPacket.server";
 import {
   downloadPacketFromObjectStore,
   formatStorageUri,
@@ -65,6 +66,7 @@ const originalEnvObj = {
   OBJECT_STORE_SECRET_ACCESS_KEY: env.OBJECT_STORE_SECRET_ACCESS_KEY,
   OBJECT_STORE_REGION: env.OBJECT_STORE_REGION,
   OBJECT_STORE_DEFAULT_PROTOCOL: env.OBJECT_STORE_DEFAULT_PROTOCOL,
+  TASK_PAYLOAD_OFFLOAD_THRESHOLD: env.TASK_PAYLOAD_OFFLOAD_THRESHOLD,
 };
 
 describe("Object Storage", () => {
@@ -344,6 +346,66 @@ describe("Object Storage", () => {
     }
   );
 
+  postgresAndMinioTest(
+    "processWaitpointCompletionPacket offloads above threshold and round-trips download",
+    async ({ minioConfig, prisma }) => {
+      env.OBJECT_STORE_BASE_URL = minioConfig.baseUrl;
+      env.OBJECT_STORE_ACCESS_KEY_ID = minioConfig.accessKeyId;
+      env.OBJECT_STORE_SECRET_ACCESS_KEY = minioConfig.secretAccessKey;
+      env.OBJECT_STORE_REGION = minioConfig.region;
+      env.OBJECT_STORE_DEFAULT_PROTOCOL = undefined;
+
+      const savedThreshold = env.TASK_PAYLOAD_OFFLOAD_THRESHOLD;
+      env.TASK_PAYLOAD_OFFLOAD_THRESHOLD = 256;
+
+      const environment = await createTestEnvironment(prisma);
+      const pathPrefix = `waitpoint_completiontest/token`;
+
+      try {
+        const smallPacket: IOPacket = {
+          data: JSON.stringify({ ok: true }),
+          dataType: "application/json",
+        };
+        const smallResult = await processWaitpointCompletionPacket(
+          smallPacket,
+          environment as any,
+          pathPrefix
+        );
+        expect(smallResult).toEqual(smallPacket);
+
+        const largeBody = "x".repeat(400);
+        const largePacket: IOPacket = {
+          data: largeBody,
+          dataType: "text/plain",
+        };
+        const largeResult = await processWaitpointCompletionPacket(
+          largePacket,
+          environment as any,
+          pathPrefix
+        );
+
+        expect(largeResult.dataType).toBe("application/store");
+        expect(largeResult.data).toBe(`${pathPrefix}.txt`);
+
+        const downloadedPacket = await downloadPacketFromObjectStore(
+          {
+            data: largeResult.data,
+            dataType: "application/store",
+          },
+          environment as any
+        );
+
+        expect(downloadedPacket.data).toBe(largeBody);
+      } finally {
+        env.TASK_PAYLOAD_OFFLOAD_THRESHOLD = savedThreshold;
+
+        await prisma.runtimeEnvironment.delete({ where: { id: environment.id } });
+        await prisma.project.delete({ where: { id: environment.projectId } });
+        await prisma.organization.delete({ where: { id: environment.organizationId } });
+      }
+    }
+  );
+
   describe("hasObjectStoreClient", () => {
     it("returns false when no store is configured", () => {
       env.OBJECT_STORE_BASE_URL = undefined;
@@ -582,5 +644,6 @@ describe("Object Storage", () => {
     env.OBJECT_STORE_SECRET_ACCESS_KEY = originalEnvObj.OBJECT_STORE_SECRET_ACCESS_KEY;
     env.OBJECT_STORE_REGION = originalEnvObj.OBJECT_STORE_REGION;
     env.OBJECT_STORE_DEFAULT_PROTOCOL = originalEnvObj.OBJECT_STORE_DEFAULT_PROTOCOL;
+    env.TASK_PAYLOAD_OFFLOAD_THRESHOLD = originalEnvObj.TASK_PAYLOAD_OFFLOAD_THRESHOLD;
   });
 });
