import { type StreamBatchItemsResponse, BatchItemNDJSON as BatchItemNDJSONSchema, } from "@trigger.dev/core/v3"; import { BatchId } from "@trigger.dev/core/v3/isomorphic"; import type { BatchItem, RunEngine } from "@internal/run-engine"; import pMap from "p-map"; import type { BatchTaskRunStatus } from "@trigger.dev/database"; 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"; /** * Phase 2 retry idempotency check. * * Returns true when the batch is in a state that means the Phase 2 stream's * job has already been done — every item has a TaskRun record (real or * pre-failed) for the customer to monitor. A retry, or the original call * racing against a fast-completing BatchQueue, should return sealed:true * in these states so the SDK stops retrying. * * Three "work is done" shapes: * - status moved out of PENDING into PROCESSING/COMPLETED/PARTIAL_FAILED * (PROCESSING via our seal, COMPLETED via tryCompleteBatch, PARTIAL_FAILED * via the V2 batchCompletionCallback). * - status stuck at PENDING but `sealed=true`: another concurrent * streamBatchItems call sealed the batch and then the callback's * happy-path branch reset status to PENDING ("all runs created"). * - status stuck at PENDING with `sealed=false` but `processingCompletedAt` * set: the cleanup-race. BatchQueue rushed through all items, callback * fired (setting processingCompletedAt), cleanup deleted the Redis * metadata — all before our service got the chance to seal. The work * is done; the discriminator is processingCompletedAt which is set * exclusively by the V2 completion callback. * * ABORTED is excluded — it means ZERO TaskRun records were created (every * per-item attempt failed AND the pre-failed-TaskRun fallback also failed, * or queue-overload on every item). The customer has nothing to monitor * at the run level, so the trigger call must throw to give their retry/ * error handling a chance to create a fresh batch. */ export function isIdempotentRetrySuccess( status: BatchTaskRunStatus | null | undefined, sealed: boolean | null | undefined, processingCompletedAt: Date | null | undefined ): boolean { return ( status === "PROCESSING" || status === "COMPLETED" || status === "PARTIAL_FAILED" || (status === "PENDING" && (sealed === true || processingCompletedAt != null)) ); } export type StreamBatchItemsServiceOptions = { maxItemBytes: number; /** Max items processed concurrently. The route wires this to STREAMING_BATCH_INGEST_CONCURRENCY. */ concurrency: number; }; export type OversizedItemMarker = { __batchItemError: "OVERSIZED"; index: number; task: string; actualSize: number; maxSize: number; }; export type StreamBatchItemsServiceConstructorOptions = { prisma?: PrismaClientOrTransaction; engine?: RunEngine; /** Override the payload processor (used in tests to observe ingest concurrency). */ payloadProcessor?: BatchPayloadProcessor; }; /** * Stream Batch Items Service (Phase 2 of 2-phase batch API). * * This service handles Phase 2 of the streaming batch API: * 1. Validates batch exists and is in PENDING status * 2. Processes NDJSON stream item by item * 3. Calls engine.enqueueBatchItem() for each item * 4. Tracks accepted/deduplicated counts * 5. On completion: validates count, seals the batch * * The service is designed for streaming and processes items as they arrive, * providing backpressure through the async iterator pattern. */ export class StreamBatchItemsService extends WithRunEngine { private readonly payloadProcessor: BatchPayloadProcessor; constructor(opts: StreamBatchItemsServiceConstructorOptions = {}) { super({ prisma: opts.prisma ?? prisma, engine: opts.engine }); this.payloadProcessor = opts.payloadProcessor ?? new BatchPayloadProcessor(); } /** * Parse a batch friendly ID to its internal ID format. * Throws a ServiceValidationError with 400 status if the ID is malformed. */ private parseBatchFriendlyId(friendlyId: string): string { try { return BatchId.fromFriendlyId(friendlyId); } catch { throw new ServiceValidationError(`Invalid batchFriendlyId: ${friendlyId}`, 400); } } /** * Process a stream of batch items from an async iterator. * Each item is validated and enqueued to the BatchQueue. * The batch is sealed when the stream completes. */ public async call( environment: AuthenticatedEnvironment, batchFriendlyId: string, itemsIterator: AsyncIterable, options: StreamBatchItemsServiceOptions ): Promise { return this.traceWithEnv( "streamBatchItems()", environment, async (span) => { span.setAttribute("batchId", batchFriendlyId); // Convert friendly ID to internal ID const batchId = this.parseBatchFriendlyId(batchFriendlyId); // Validate batch exists and belongs to this environment. Routed by batch id so a // run-ops id (NEW-resident) batch is found on the owning DB; the env-ownership check that // was in the where clause is enforced app-side below. const batch = await this._engine.runStore.findBatchTaskRunById(batchId); if (!batch || batch.runtimeEnvironmentId !== environment.id) { throw new ServiceValidationError(`Batch ${batchFriendlyId} not found`); } if (isIdempotentRetrySuccess(batch.status, batch.sealed, batch.processingCompletedAt)) { logger.info("Batch already sealed/completed - treating Phase 2 retry as success", { batchId: batchFriendlyId, batchSealed: batch.sealed, batchStatus: batch.status, processingCompletedAt: batch.processingCompletedAt, }); return { id: batchFriendlyId, itemsAccepted: 0, itemsDeduplicated: 0, sealed: true, runCount: batch.runCount, }; } if (batch.status !== "PENDING") { // ABORTED or any other unexpected non-PENDING state — surface as an error. // For ABORTED specifically, throwing is required so the customer's // batchTrigger() retries (a new batch) can recreate the runs. throw new ServiceValidationError( `Batch ${batchFriendlyId} is not in PENDING status (current: ${batch.status})` ); } // Process items from the stream with bounded concurrency. // // Ordering and idempotency do NOT depend on processing order: // - The BatchQueue derives run order from each item's index // (enqueue timestamp = batch.createdAt + itemIndex), not enqueue order. // - enqueueBatchItem() dedups atomically per index. // We cap concurrency to bound peak in-flight memory (≈ concurrency × // maxItemBytes) and to keep backpressure on the request body stream. // p-map pulls lazily from the async iterator — at most `concurrency` // items are read and in flight at once. stopOnError aborts ingestion on // the first failure (the batch is left unsealed; the SDK's retry // re-streams and dedups already-enqueued items). const outcomes = await pMap( itemsIterator, (rawItem) => this.#processItem(rawItem, batchId, environment, batch.runCount), { concurrency: options.concurrency, stopOnError: true } ); let itemsAccepted = 0; let itemsDeduplicated = 0; for (const outcome of outcomes) { if (outcome === "accepted") { itemsAccepted++; } else { itemsDeduplicated++; } } // Get the actual enqueued count from Redis const enqueuedCount = await this._engine.getBatchEnqueuedCount(batchId); // Validate we received the expected number of items if (enqueuedCount !== batch.runCount) { // The batch queue consumers may have already processed all items and // cleaned up the Redis keys before we got here. This happens when all // runs complete fast enough that cleanup() deletes the enqueuedItemsKey // before we read it — typically when the last item executes in the // milliseconds between the loop ending and getBatchEnqueuedCount() being called. // Check both sealed (sealed by this endpoint on a concurrent request) and // COMPLETED (sealed by the BatchQueue completion path before we got here). const currentBatch = await this._engine.runStore.findBatchTaskRunById(batchId); if ( isIdempotentRetrySuccess( currentBatch?.status, currentBatch?.sealed, currentBatch?.processingCompletedAt ) ) { logger.info("Batch already sealed before count check (fast completion)", { batchId: batchFriendlyId, itemsAccepted, itemsDeduplicated, enqueuedCount, expectedCount: batch.runCount, batchStatus: currentBatch?.status, processingCompletedAt: currentBatch?.processingCompletedAt, }); return { id: batchFriendlyId, itemsAccepted, itemsDeduplicated, sealed: true, runCount: batch.runCount, }; } if (currentBatch?.status === "ABORTED") { // Zero TaskRuns exist — the count-mismatch sealed:false semantics // ("retry with missing items") would mislead the SDK. Throw so the // customer's batchTrigger() retry creates a fresh batch. throw new ServiceValidationError( `Batch ${batchFriendlyId} is not in PENDING status (current: ABORTED)` ); } logger.warn("Batch item count mismatch", { batchId: batchFriendlyId, expected: batch.runCount, received: enqueuedCount, itemsAccepted, itemsDeduplicated, }); // Don't seal the batch if count doesn't match // Return sealed: false so client knows to retry with missing items return { id: batchFriendlyId, itemsAccepted, itemsDeduplicated, sealed: false, enqueuedCount, expectedCount: batch.runCount, runCount: batch.runCount, }; } // Seal the batch - use conditional update to prevent TOCTOU race // Another concurrent request may have already sealed this batch const now = new Date(); const sealResult = await this._engine.runStore.updateManyBatchTaskRun({ where: { id: batchId, sealed: false, status: "PENDING", }, data: { sealed: true, sealedAt: now, status: "PROCESSING", processingStartedAt: now, }, }); // Check if we won the race to seal the batch if (sealResult.count === 0) { // The conditional update failed because the batch was no longer in // PENDING status. Re-query to determine which path got there first: // - A concurrent streaming request already sealed and moved it to // PROCESSING. // - The BatchQueue completion path finished all runs and set it to // COMPLETED (without setting sealed=true — that's this endpoint's // job). This window exists between completionCallback (which calls // tryCompleteBatch) and cleanup() in BatchQueue — see // batch-queue/index.ts. // Either way the goal — a durable batch that the SDK stops retrying — // has been achieved, so we return sealed: true. const currentBatch = await this._engine.runStore.findBatchTaskRunById(batchId); if ( isIdempotentRetrySuccess( currentBatch?.status, currentBatch?.sealed, currentBatch?.processingCompletedAt ) ) { logger.info("Batch already sealed/completed by concurrent path", { batchId: batchFriendlyId, itemsAccepted, itemsDeduplicated, envId: environment.id, batchStatus: currentBatch?.status, batchSealed: currentBatch?.sealed, processingCompletedAt: currentBatch?.processingCompletedAt, }); span.setAttribute("itemsAccepted", itemsAccepted); span.setAttribute("itemsDeduplicated", itemsDeduplicated); span.setAttribute("sealedByConcurrentRequest", true); return { id: batchFriendlyId, itemsAccepted, itemsDeduplicated, sealed: true, runCount: batch.runCount, }; } // Batch is in an unexpected state - fail with error const actualStatus = currentBatch?.status ?? "unknown"; const actualSealed = currentBatch?.sealed ?? "unknown"; logger.error("Batch seal race condition: unexpected state", { batchId: batchFriendlyId, expectedStatus: "PENDING", actualStatus, expectedSealed: false, actualSealed, envId: environment.id, }); throw new ServiceValidationError( `Batch ${batchFriendlyId} is in unexpected state (status: ${actualStatus}, sealed: ${actualSealed}). Cannot seal batch.` ); } logger.info("Batch sealed and ready for processing", { batchId: batchFriendlyId, itemsAccepted, itemsDeduplicated, totalEnqueued: enqueuedCount, envId: environment.id, }); span.setAttribute("itemsAccepted", itemsAccepted); span.setAttribute("itemsDeduplicated", itemsDeduplicated); return { id: batchFriendlyId, itemsAccepted, itemsDeduplicated, sealed: true, runCount: batch.runCount, }; } ); } /** * Process a single streamed batch item: validate it, offload its payload to * object storage if oversized, and enqueue it. Returns whether the item was * newly enqueued ("accepted") or was a duplicate ("deduplicated"). Throws * ServiceValidationError for invalid items, which aborts the stream. * * Safe to run concurrently: enqueueBatchItem() is atomic and order-independent * per item index, and each item carries its own index (real items from the * SDK; oversized markers are stamped by the NDJSON parser). */ async #processItem( rawItem: unknown, batchId: string, environment: AuthenticatedEnvironment, runCount: number ): Promise<"accepted" | "deduplicated"> { // Oversized item marker emitted by the NDJSON parser if (rawItem && typeof rawItem === "object" && "__batchItemError" in rawItem) { const marker = rawItem as OversizedItemMarker; // Same out-of-range guard as normal items: an oversized item with an // out-of-range index must 4xx rather than create a stray pre-failed run. if (marker.index >= runCount) { throw new ServiceValidationError( `Item index ${marker.index} exceeds batch runCount ${runCount}` ); } const errorMessage = `Batch item payload is too large (${(marker.actualSize / 1024).toFixed( 1 )} KB). Maximum allowed size is ${(marker.maxSize / 1024).toFixed( 1 )} KB. Reduce the payload size or offload large data to external storage.`; // Enqueue with __error metadata - processItemCallback will detect this // and use TriggerFailedTaskService to create a pre-failed run const batchItem: BatchItem = { task: marker.task, payload: "{}", payloadType: "application/json", options: { __error: errorMessage, __errorCode: "PAYLOAD_TOO_LARGE", }, }; const result = await this._engine.enqueueBatchItem( batchId, environment.id, marker.index, batchItem ); return result.enqueued ? "accepted" : "deduplicated"; } // Parse and validate the item const parseResult = BatchItemNDJSONSchema.safeParse(rawItem); if (!parseResult.success) { const rawIndex = (rawItem as { index?: unknown } | null)?.index; const where = typeof rawIndex === "number" ? `index ${rawIndex}` : "unknown index"; throw new ServiceValidationError(`Invalid item at ${where}: ${parseResult.error.message}`); } const item = parseResult.data; // Validate index is within expected range if (item.index >= runCount) { throw new ServiceValidationError( `Item index ${item.index} exceeds batch runCount ${runCount}` ); } // Get the original payload type const originalPayloadType = (item.options?.payloadType as string) ?? "application/json"; // Process payload - offload to object storage 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: processedPayload.payload, payloadType: processedPayload.payloadType, options: item.options, }; // Enqueue the item const result = await this._engine.enqueueBatchItem( batchId, environment.id, item.index, batchItem ); return result.enqueued ? "accepted" : "deduplicated"; } } /** * Extract `index` and `task` from raw JSON bytes without decoding the full line. * Scans at most 512 bytes, tracking JSON nesting depth to only match top-level keys. */ export function extractIndexAndTask(bytes: Uint8Array): { index: number; task: string } { let index = -1; let task = "unknown"; let depth = 0; let foundIndex = false; let foundTask = false; const limit = Math.min(bytes.byteLength, 512); const QUOTE = 0x22; // " const COLON = 0x3a; // : const LBRACE = 0x7b; // { const RBRACE = 0x7d; // } const LBRACKET = 0x5b; // [ const RBRACKET = 0x5d; // ] const BACKSLASH = 0x5c; // \ // Byte patterns for "index" and "task" (without quotes) const INDEX_BYTES = [0x69, 0x6e, 0x64, 0x65, 0x78]; // index const TASK_BYTES = [0x74, 0x61, 0x73, 0x6b]; // task let i = 0; while (i < limit && !(foundIndex && foundTask)) { const b = bytes[i]; if (b === LBRACE || b === LBRACKET) { depth++; i++; continue; } if (b === RBRACE || b === RBRACKET) { depth--; i++; continue; } // Only match keys at depth 1 (top-level object) if (b === QUOTE && depth === 1) { // Read the key inside quotes const keyStart = i + 1; let keyEnd = keyStart; while (keyEnd < limit && bytes[keyEnd] !== QUOTE) { if (bytes[keyEnd] === BACKSLASH) keyEnd++; // skip escaped char keyEnd++; } const keyLen = keyEnd - keyStart; // Check if this key matches "index" or "task" const isIndex = !foundIndex && keyLen === INDEX_BYTES.length && INDEX_BYTES.every((b, j) => bytes[keyStart + j] === b); const isTask = !foundTask && keyLen === TASK_BYTES.length && TASK_BYTES.every((b, j) => bytes[keyStart + j] === b); if (isIndex || isTask) { // Skip past closing quote and find colon let pos = keyEnd + 1; while (pos < limit && bytes[pos] !== COLON) pos++; pos++; // skip colon // Skip whitespace while (pos < limit && (bytes[pos] === 0x20 || bytes[pos] === 0x09)) pos++; if (isIndex) { // Parse digits let num = 0; let hasDigit = false; while (pos < limit && bytes[pos] >= 0x30 && bytes[pos] <= 0x39) { num = num * 10 + (bytes[pos] - 0x30); hasDigit = true; pos++; } if (hasDigit) { index = num; foundIndex = true; } } else { // Parse quoted string value if (pos < limit && bytes[pos] === QUOTE) { const valStart = pos + 1; let valEnd = valStart; while (valEnd < limit && bytes[valEnd] !== QUOTE) { if (bytes[valEnd] === BACKSLASH) valEnd++; valEnd++; } // Decode just this slice try { task = new TextDecoder("utf-8", { fatal: true }).decode( bytes.slice(valStart, valEnd) ); foundTask = true; } catch { // Leave as "unknown" } } } } // Skip past the key's closing quote i = keyEnd + 1; continue; } i++; } return { index, task }; } /** * Create an NDJSON parser transform stream. * * Converts a stream of Uint8Array chunks into parsed JSON objects. * Each line in the NDJSON is parsed independently. * * Uses byte-buffer accumulation to: * - Prevent OOM from unbounded string buffers * - Properly handle multibyte UTF-8 characters across chunk boundaries * - Check size limits on raw bytes before decoding * * @param maxItemBytes - Maximum allowed bytes per line (item) * @returns TransformStream that outputs parsed JSON objects */ export function createNdjsonParserStream( maxItemBytes: number ): TransformStream { // Single decoder instance, reused for all lines const decoder = new TextDecoder("utf-8", { fatal: true }); // Byte buffer: array of chunks with tracked total length let chunks: Uint8Array[] = []; let totalBytes = 0; let lineNumber = 0; // 0-based position of the next object we emit (parsed item or oversized // marker). The parser is the single sequential point in the pipeline, so this // is the authoritative source of item ordering — downstream consumers can // process items concurrently and must not rely on processing order to derive // an item's index. Used to back-fill an oversized marker's index when it // couldn't be extracted from the (truncated) raw bytes. let emittedCount = 0; // When an oversized incomplete line is detected (Case 2), we must discard // all remaining bytes of that line until the next newline delimiter. let skipUntilNewline = false; const NEWLINE_BYTE = 0x0a; // '\n' /** * Emit a parsed object or marker downstream and advance the emit position. * Every emitted object MUST go through here so `emittedCount` stays aligned * with item position (empty/skipped lines never emit, so they don't count). */ function emit(controller: TransformStreamDefaultController, obj: unknown): void { controller.enqueue(obj); emittedCount++; } /** * Concatenate all chunks into a single Uint8Array */ function concatenateChunks(): Uint8Array { if (chunks.length === 0) { return new Uint8Array(0); } if (chunks.length === 1) { return chunks[0]; } const result = new Uint8Array(totalBytes); let offset = 0; for (const chunk of chunks) { result.set(chunk, offset); offset += chunk.byteLength; } return result; } /** * Find the index of the first newline byte in the buffer. * Returns -1 if not found. */ function findNewlineIndex(): number { let globalIndex = 0; for (const chunk of chunks) { for (let i = 0; i < chunk.byteLength; i++) { if (chunk[i] === NEWLINE_BYTE) { return globalIndex + i; } } globalIndex += chunk.byteLength; } return -1; } /** * Extract bytes from the buffer up to (but not including) the given index, * and remove those bytes plus the delimiter from the buffer. */ function extractLine(newlineIndex: number): Uint8Array { const fullBuffer = concatenateChunks(); const lineBytes = fullBuffer.slice(0, newlineIndex); const remaining = fullBuffer.slice(newlineIndex + 1); // Skip the newline // Reset buffer with remaining bytes if (remaining.byteLength > 0) { chunks = [remaining]; totalBytes = remaining.byteLength; } else { chunks = []; totalBytes = 0; } return lineBytes; } /** * Parse a line from bytes, handling whitespace trimming. * Returns the parsed object or null for empty lines. */ function parseLine( lineBytes: Uint8Array, controller: TransformStreamDefaultController ): void { lineNumber++; // Decode the line bytes (stream: false since this is a complete line) let lineText: string; try { lineText = decoder.decode(lineBytes, { stream: false }); } catch (err) { throw new Error(`Invalid UTF-8 at line ${lineNumber}: ${(err as Error).message}`); } const trimmed = lineText.trim(); if (!trimmed) { return; // Skip empty lines } try { const obj = JSON.parse(trimmed); emit(controller, obj); } catch (err) { throw new Error(`Invalid JSON at line ${lineNumber}: ${(err as Error).message}`); } } return new TransformStream({ transform(chunk, controller) { // If we're skipping the remainder of an oversized line, scan for the // next newline in this chunk and discard everything before it. if (skipUntilNewline) { const nlPos = chunk.indexOf(NEWLINE_BYTE); if (nlPos === -1) { // Entire chunk is still part of the oversized line — discard it return; } // Found the newline — keep everything after it skipUntilNewline = false; const remaining = chunk.slice(nlPos + 1); if (remaining.byteLength === 0) { return; } // Replace chunk with the remainder and fall through to normal processing chunk = remaining; } // Append chunk to buffer chunks.push(chunk); totalBytes += chunk.byteLength; // Process all complete lines in the buffer let newlineIndex: number; while ((newlineIndex = findNewlineIndex()) !== -1) { // Check size limit BEFORE extracting/decoding (bytes up to newline) if (newlineIndex > maxItemBytes) { // Case 1: Complete line exceeds limit - emit marker instead of throwing const lineBytes = extractLine(newlineIndex); const extracted = extractIndexAndTask(lineBytes); const marker: OversizedItemMarker = { __batchItemError: "OVERSIZED", index: extracted.index >= 0 ? extracted.index : emittedCount, task: extracted.task, actualSize: newlineIndex, maxSize: maxItemBytes, }; emit(controller, marker); lineNumber++; continue; } const lineBytes = extractLine(newlineIndex); parseLine(lineBytes, controller); } // Check if the remaining buffer (incomplete line) exceeds the limit // This prevents OOM from a single huge line without newlines if (totalBytes > maxItemBytes) { // Case 2: Incomplete line exceeds limit - emit marker instead of throwing const extracted = extractIndexAndTask(concatenateChunks()); const marker: OversizedItemMarker = { __batchItemError: "OVERSIZED", index: extracted.index >= 0 ? extracted.index : emittedCount, task: extracted.task, actualSize: totalBytes, maxSize: maxItemBytes, }; emit(controller, marker); lineNumber++; // Clear buffer and skip remaining bytes of this oversized line // until the next newline delimiter is found in a subsequent chunk chunks = []; totalBytes = 0; skipUntilNewline = true; return; } }, flush(controller) { // Flush any remaining bytes from the decoder's internal state // This handles multibyte characters that may have been split across chunks decoder.decode(new Uint8Array(0), { stream: false }); // Process any remaining buffered data (no trailing newline case) if (totalBytes === 0) { return; } // Check size limit before processing final line if (totalBytes > maxItemBytes) { // Case 3: Flush with oversized remaining - emit marker instead of throwing const extracted = extractIndexAndTask(concatenateChunks()); const marker: OversizedItemMarker = { __batchItemError: "OVERSIZED", index: extracted.index >= 0 ? extracted.index : emittedCount, task: extracted.task, actualSize: totalBytes, maxSize: maxItemBytes, }; emit(controller, marker); return; } const finalBytes = concatenateChunks(); parseLine(finalBytes, controller); }, }); } /** * Convert a ReadableStream into an AsyncIterable. * Useful for processing streams with for-await-of loops. */ export async function* streamToAsyncIterable(stream: ReadableStream): AsyncIterable { const reader = stream.getReader(); try { while (true) { const { done, value } = await reader.read(); if (done) break; yield value; } } finally { reader.releaseLock(); } }