Handle large payloads and correct the trace ID propogation to child runs

Author: ericallam

apps/webapp/app/env.server.ts

@@ -528,6 +528,7 @@ const EnvironmentSchema = z
     MAXIMUM_TRACE_SUMMARY_VIEW_COUNT: z.coerce.number().int().default(25_000),
     MAXIMUM_TRACE_DETAILED_SUMMARY_VIEW_COUNT: z.coerce.number().int().default(10_000),
     TASK_PAYLOAD_OFFLOAD_THRESHOLD: z.coerce.number().int().default(524_288), // 512KB
+    BATCH_PAYLOAD_OFFLOAD_THRESHOLD: z.coerce.number().int().optional(), // Defaults to TASK_PAYLOAD_OFFLOAD_THRESHOLD if not set
     TASK_PAYLOAD_MAXIMUM_SIZE: z.coerce.number().int().default(3_145_728), // 3MB
     BATCH_TASK_PAYLOAD_MAXIMUM_SIZE: z.coerce.number().int().default(1_000_000), // 1MB
     TASK_RUN_METADATA_MAXIMUM_SIZE: z.coerce.number().int().default(262_144), // 256KB
@@ -539,7 +540,7 @@ const EnvironmentSchema = z
 
     // 2-phase batch API settings
     STREAMING_BATCH_MAX_ITEMS: z.coerce.number().int().default(1_000), // Max items in streaming batch
-    STREAMING_BATCH_ITEM_MAXIMUM_SIZE: z.coerce.number().int().default(512_000), // 512KB max per NDJSON line
+    STREAMING_BATCH_ITEM_MAXIMUM_SIZE: z.coerce.number().int().default(3_145_728),
 
     REALTIME_STREAM_VERSION: z.enum(["v1", "v2"]).default("v1"),
     REALTIME_STREAM_MAX_LENGTH: z.coerce.number().int().default(1000),

apps/webapp/app/routes/api.v2.batches.$batchId.ts

@@ -38,10 +38,7 @@ export const loader = createLoaderApiRoute(
       updatedAt: batch.updatedAt,
       runCount: batch.runCount,
       runs: batch.runIds,
-      // Include error details for PARTIAL_FAILED batches
-      successfulRunCount: batch.successfulRunCount ?? undefined,
-      failedRunCount: batch.failedRunCount ?? undefined,
-      errors:
+      processingErrors:
         batch.errors.length > 0
           ? batch.errors.map((err) => ({
               index: err.index,

apps/webapp/app/runEngine/concerns/batchPayloads.server.ts

@@ -0,0 +1,165 @@
+import { IOPacket, packetRequiresOffloading, tryCatch } from "@trigger.dev/core/v3";
+import { env } from "~/env.server";
+import { startActiveSpan } from "~/v3/tracer.server";
+import { uploadPacketToObjectStore, r2 } from "~/v3/r2.server";
+import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
+import { logger } from "~/services/logger.server";
+
+export type BatchPayloadProcessResult = {
+  /** The processed payload - either the original or an R2 path */
+  payload: unknown;
+  /** The payload type - "application/store" if offloaded to R2 */
+  payloadType: string;
+  /** Whether the payload was offloaded to R2 */
+  wasOffloaded: boolean;
+  /** Size of the payload in bytes */
+  size: number;
+};
+
+/**
+ * BatchPayloadProcessor handles payload offloading for batch items.
+ *
+ * When a batch item's payload exceeds the configured threshold, it's uploaded
+ * to object storage (R2) and the payload is replaced with the storage path.
+ * This aligns with how single task triggers work via DefaultPayloadProcessor.
+ *
+ * Path format: batch_{batchId}/item_{index}/payload.json
+ */
+export class BatchPayloadProcessor {
+  /**
+   * Check if object storage is available for payload offloading.
+   * If not available, large payloads will be stored inline (which may fail for very large payloads).
+   */
+  isObjectStoreAvailable(): boolean {
+    return r2 !== undefined && env.OBJECT_STORE_BASE_URL !== undefined;
+  }
+
+  /**
+   * Process a batch item payload, offloading to R2 if it exceeds the threshold.
+   *
+   * @param payload - The raw payload from the batch item
+   * @param payloadType - The payload type (e.g., "application/json")
+   * @param batchId - The batch ID (internal format)
+   * @param itemIndex - The item index within the batch
+   * @param environment - The authenticated environment for R2 path construction
+   * @returns The processed result with potentially offloaded payload
+   */
+  async process(
+    payload: unknown,
+    payloadType: string,
+    batchId: string,
+    itemIndex: number,
+    environment: AuthenticatedEnvironment
+  ): Promise<BatchPayloadProcessResult> {
+    return startActiveSpan("BatchPayloadProcessor.process()", async (span) => {
+      span.setAttribute("batchId", batchId);
+      span.setAttribute("itemIndex", itemIndex);
+      span.setAttribute("payloadType", payloadType);
+
+      // Create the packet for size checking
+      const packet = this.#createPayloadPacket(payload, payloadType);
+
+      if (!packet.data) {
+        return {
+          payload,
+          payloadType,
+          wasOffloaded: false,
+          size: 0,
+        };
+      }
+
+      const threshold = env.BATCH_PAYLOAD_OFFLOAD_THRESHOLD ?? env.TASK_PAYLOAD_OFFLOAD_THRESHOLD;
+      const { needsOffloading, size } = packetRequiresOffloading(packet, threshold);
+
+      span.setAttribute("payloadSize", size);
+      span.setAttribute("needsOffloading", needsOffloading);
+      span.setAttribute("threshold", threshold);
+
+      if (!needsOffloading) {
+        return {
+          payload,
+          payloadType,
+          wasOffloaded: false,
+          size,
+        };
+      }
+
+      // Check if object store is available
+      if (!this.isObjectStoreAvailable()) {
+        logger.warn("Payload exceeds threshold but object store is not available", {
+          batchId,
+          itemIndex,
+          size,
+          threshold,
+        });
+
+        // Return without offloading - the payload will be stored inline
+        // This may fail downstream for very large payloads
+        return {
+          payload,
+          payloadType,
+          wasOffloaded: false,
+          size,
+        };
+      }
+
+      // Upload to R2
+      const filename = `batch_${batchId}/item_${itemIndex}/payload.json`;
+
+      const [uploadError] = await tryCatch(
+        uploadPacketToObjectStore(filename, packet.data, packet.dataType, environment)
+      );
+
+      if (uploadError) {
+        logger.error("Failed to upload batch item payload to object store", {
+          batchId,
+          itemIndex,
+          error: uploadError instanceof Error ? uploadError.message : String(uploadError),
+        });
+
+        // Throw to fail this item - SDK can retry
+        throw new Error(
+          `Failed to upload large payload to object store: ${uploadError instanceof Error ? uploadError.message : String(uploadError)}`
+        );
+      }
+
+      logger.debug("Batch item payload offloaded to R2", {
+        batchId,
+        itemIndex,
+        filename,
+        size,
+      });
+
+      span.setAttribute("wasOffloaded", true);
+      span.setAttribute("offloadPath", filename);
+
+      return {
+        payload: filename,
+        payloadType: "application/store",
+        wasOffloaded: true,
+        size,
+      };
+    });
+  }
+
+  /**
+   * Create an IOPacket from payload for size checking.
+   */
+  #createPayloadPacket(payload: unknown, payloadType: string): IOPacket {
+    if (payloadType === "application/json") {
+      return { data: JSON.stringify(payload), dataType: "application/json" };
+    }
+
+    if (typeof payload === "string") {
+      return { data: payload, dataType: payloadType };
+    }
+
+    // For other types, try to stringify
+    try {
+      return { data: JSON.stringify(payload), dataType: payloadType };
+    } catch {
+      return { dataType: payloadType };
+    }
+  }
+}
+

apps/webapp/app/runEngine/services/streamBatchItems.server.ts

@@ -9,6 +9,7 @@ import { prisma, type PrismaClientOrTransaction } from "~/db.server";
 import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { ServiceValidationError, WithRunEngine } from "../../v3/services/baseService.server";
+import { BatchPayloadProcessor } from "../concerns/batchPayloads.server";
 
 export type StreamBatchItemsServiceOptions = {
   maxItemBytes: number;
@@ -28,8 +29,11 @@ export type StreamBatchItemsServiceOptions = {
  * providing backpressure through the async iterator pattern.
  */
 export class StreamBatchItemsService extends WithRunEngine {
+  private readonly payloadProcessor: BatchPayloadProcessor;
+
   constructor(protected readonly _prisma: PrismaClientOrTransaction = prisma) {
     super({ prisma });
+    this.payloadProcessor = new BatchPayloadProcessor();
   }
 
   /**
@@ -108,11 +112,23 @@ export class StreamBatchItemsService extends WithRunEngine {
             );
           }
 
-          // Convert to BatchItem format
+          // Get the original payload type
+          const originalPayloadType = (item.options?.payloadType as string) ?? "application/json";
+
+          // Process payload - offload to R2 if it exceeds threshold
+          const processedPayload = await this.payloadProcessor.process(
+            item.payload,
+            originalPayloadType,
+            batchId,
+            item.index,
+            environment
+          );
+
+          // Convert to BatchItem format with potentially offloaded payload
           const batchItem: BatchItem = {
             task: item.task,
-            payload: item.payload,
-            payloadType: (item.options?.payloadType as string) ?? "application/json",
+            payload: processedPayload.payload,
+            payloadType: processedPayload.payloadType,
             options: item.options,
           };
 

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

@@ -160,26 +160,23 @@ function createRunEngine() {
     },
     // BatchQueue with DRR scheduling for fair batch processing
     // Consumers are controlled by options.worker.disabled (same as main worker)
-    batchQueue:
-      env.BATCH_TRIGGER_WORKER_ENABLED === "true"
-        ? {
-            redis: {
-              keyPrefix: "engine:",
-              port: env.BATCH_TRIGGER_WORKER_REDIS_PORT ?? undefined,
-              host: env.BATCH_TRIGGER_WORKER_REDIS_HOST ?? undefined,
-              username: env.BATCH_TRIGGER_WORKER_REDIS_USERNAME ?? undefined,
-              password: env.BATCH_TRIGGER_WORKER_REDIS_PASSWORD ?? undefined,
-              enableAutoPipelining: true,
-              ...(env.BATCH_TRIGGER_WORKER_REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
-            },
-            drr: {
-              quantum: env.BATCH_QUEUE_DRR_QUANTUM,
-              maxDeficit: env.BATCH_QUEUE_MAX_DEFICIT,
-            },
-            consumerCount: env.BATCH_QUEUE_CONSUMER_COUNT,
-            consumerIntervalMs: env.BATCH_QUEUE_CONSUMER_INTERVAL_MS,
-          }
-        : undefined,
+    batchQueue: {
+      redis: {
+        keyPrefix: "engine:",
+        port: env.BATCH_TRIGGER_WORKER_REDIS_PORT ?? undefined,
+        host: env.BATCH_TRIGGER_WORKER_REDIS_HOST ?? undefined,
+        username: env.BATCH_TRIGGER_WORKER_REDIS_USERNAME ?? undefined,
+        password: env.BATCH_TRIGGER_WORKER_REDIS_PASSWORD ?? undefined,
+        enableAutoPipelining: true,
+        ...(env.BATCH_TRIGGER_WORKER_REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
+      },
+      drr: {
+        quantum: env.BATCH_QUEUE_DRR_QUANTUM,
+        maxDeficit: env.BATCH_QUEUE_MAX_DEFICIT,
+      },
+      consumerCount: env.BATCH_QUEUE_CONSUMER_COUNT,
+      consumerIntervalMs: env.BATCH_QUEUE_CONSUMER_INTERVAL_MS,
+    },
   });
 
   // Set up BatchQueue callbacks if enabled
@@ -192,17 +189,24 @@ function createRunEngine() {
 
 /**
  * Normalize the payload from BatchQueue.
- * The payload might be a JSON string if the SDK sent it pre-serialized.
- * If it's a JSON string and payloadType is "application/json", parse it
- * to avoid double-stringification in DefaultPayloadProcessor.
+ *
+ * Handles different payload types:
+ * - "application/store": Already offloaded to R2, payload is the path - pass through as-is
+ * - "application/json": May be a pre-serialized JSON string - parse to avoid double-stringification
+ * - Other types: Pass through as-is
+ *
+ * @param payload - The raw payload from the batch item
+ * @param payloadType - The payload type (e.g., "application/json", "application/store")
  */
 function normalizePayload(payload: unknown, payloadType?: string): unknown {
-  // Only normalize for JSON payloads
+  // For non-JSON payloads (including application/store for R2-offloaded payloads),
+  // return as-is - no normalization needed
   if (payloadType !== "application/json" && payloadType !== undefined) {
     return payload;
   }
 
-  // If payload is a string, try to parse it as JSON
+  // For JSON payloads, if payload is a string, try to parse it
+  // This handles pre-serialized JSON from the SDK
   if (typeof payload === "string") {
     try {
       return JSON.parse(payload);
@@ -218,14 +222,19 @@ function normalizePayload(payload: unknown, payloadType?: string): unknown {
 /**
  * Set up the BatchQueue processing callbacks.
  * These handle creating runs from batch items and completing batches.
+ *
+ * Payload handling:
+ * - If payloadType is "application/store", the payload is an R2 path (already offloaded)
+ * - DefaultPayloadProcessor in TriggerTaskService will pass it through without re-offloading
+ * - The run engine will download from R2 when the task executes
  */
 function setupBatchQueueCallbacks(engine: RunEngine) {
   // Item processing callback - creates a run for each batch item
   engine.setBatchProcessItemCallback(async ({ batchId, friendlyId, itemIndex, item, meta }) => {
     try {
       const triggerTaskService = new TriggerTaskService();
 
-      // Normalize payload to avoid double-stringification
+      // Normalize payload - for application/store (R2 paths), this passes through as-is
       const payload = normalizePayload(item.payload, item.payloadType);
 
       const result = await triggerTaskService.call(