import { describe, expect, it, vi } from "vitest"; // Mock the db prisma client - needs to be before other imports vi.mock("~/db.server", () => ({ prisma: {}, $replica: {}, })); vi.mock("~/services/platform.v3.server", async (importOriginal) => { const actual = (await importOriginal()) as Record; return { ...actual, getEntitlement: vi.fn(), }; }); import { RunEngine } from "@internal/run-engine"; import { PostgresRunStore } from "@internal/run-store"; import { setupAuthenticatedEnvironment } from "@internal/run-engine/tests"; // Per-test redis isolation: each test runs its own RunEngine whose background work outlives the test // body. NoClickhouse because this suite never touches ClickHouse - skips the worker-scoped boot+migrate. import { containerTestWithIsolatedRedisNoClickhouse as containerTest } from "@internal/testcontainers"; import { trace } from "@opentelemetry/api"; import type { PrismaClient } from "@trigger.dev/database"; import { BatchId } from "@trigger.dev/core/v3/isomorphic"; import { StreamBatchItemsService, createNdjsonParserStream, streamToAsyncIterable, extractIndexAndTask, type OversizedItemMarker, } from "../../app/runEngine/services/streamBatchItems.server"; import { ServiceValidationError } from "../../app/v3/services/baseService.server"; import { BatchPayloadProcessor, type BatchPayloadProcessResult, } from "../../app/runEngine/concerns/batchPayloads.server"; import { setTimeout as sleep } from "node:timers/promises"; // 120s: a cold per-test container boot counts against the test's own timeout, and under CI Docker // contention 30s was too tight. Matches the run-engine convention for this footprint. vi.setConfig({ testTimeout: 120_000 }); describe("StreamBatchItemsService", () => { async function createBatch( prisma: PrismaClient, environmentId: string, options: { runCount: number; status?: "PENDING" | "PROCESSING" | "COMPLETED" | "PARTIAL_FAILED" | "ABORTED"; sealed?: boolean; processingCompletedAt?: Date | null; } ) { const { id, friendlyId } = BatchId.generate(); const batch = await prisma.batchTaskRun.create({ data: { id, friendlyId, runtimeEnvironmentId: environmentId, status: options.status ?? "PENDING", runCount: options.runCount, expectedCount: options.runCount, runIds: [], batchVersion: "runengine:v2", sealed: options.sealed ?? false, processingCompletedAt: options.processingCompletedAt ?? null, }, }); return batch; } async function* itemsToAsyncIterable( items: Array<{ task: string; payload: string; index: number }> ) { for (const item of items) { yield item; } } function makeItems(count: number, taskId = "test-task") { return Array.from({ length: count }, (_, index) => ({ task: taskId, payload: JSON.stringify({ i: index }), index, })); } /** * Build a RunEngine for tests. The worker is disabled and no item-processing * callback is wired, so streamed items stay enqueued (the BatchQueue consumer * doesn't drain them) — keeping the enqueued-count assertions deterministic. */ function buildEngine(prisma: PrismaClient, redisOptions: any) { return new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions }, runLock: { redis: redisOptions }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions }, tracer: trace.getTracer("test", "0.0.0"), }); } containerTest( "should seal batch successfully when no race condition", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); // Create a batch const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 2, status: "PENDING", sealed: false, }); // Initialize the batch in Redis await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: authenticatedEnvironment.id, environmentType: authenticatedEnvironment.type, organizationId: authenticatedEnvironment.organizationId, projectId: authenticatedEnvironment.projectId, runCount: 2, processingConcurrency: 10, }); // Enqueue items directly to Redis (bypassing the service's item processing) await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, { task: "test-task", payload: JSON.stringify({ data: "item1" }), payloadType: "application/json", }); await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, { task: "test-task", payload: JSON.stringify({ data: "item2" }), payloadType: "application/json", }); // Create service with our test engine and prisma const service = new StreamBatchItemsService({ prisma, engine, }); // Create an empty items iterator since items are already enqueued const items = itemsToAsyncIterable([]); const result = await service.call(authenticatedEnvironment, batch.friendlyId, items, { maxItemBytes: 1024 * 1024, concurrency: 1, }); expect(result.sealed).toBe(true); expect(result.id).toBe(batch.friendlyId); // Verify the batch is sealed in the database const updatedBatch = await prisma.batchTaskRun.findUnique({ where: { id: batch.id }, }); expect(updatedBatch?.sealed).toBe(true); expect(updatedBatch?.status).toBe("PROCESSING"); expect(updatedBatch?.sealedAt).toBeDefined(); expect(updatedBatch?.processingStartedAt).toBeDefined(); await engine.quit(); } ); containerTest( "should return sealed=true when batch is already sealed and PROCESSING (Phase 2 retry idempotency)", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); // Simulate the SDK retrying Phase 2 after the original request succeeded: // the original request already sealed the batch and moved it to PROCESSING. const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 2, status: "PROCESSING", sealed: true, }); const service = new StreamBatchItemsService({ prisma, engine, }); const result = await service.call( authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, } ); // The retry should be treated as success — the original request already // enqueued every item, so the SDK should stop retrying. expect(result.sealed).toBe(true); expect(result.id).toBe(batch.friendlyId); expect(result.itemsAccepted).toBe(0); expect(result.itemsDeduplicated).toBe(0); expect(result.runCount).toBe(2); await engine.quit(); } ); containerTest( "should return sealed=true when batch is COMPLETED before Phase 2 retry arrives (TRI-9944)", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); // The customer-reported scenario: single-item batch where the original // Phase 2 request succeeded server-side, the run executed fast, the batch // flipped to COMPLETED, then the lost-response SDK retry hits us. // Note: tryCompleteBatch sets status=COMPLETED but does NOT set sealed=true. const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 1, status: "COMPLETED", sealed: false, }); const service = new StreamBatchItemsService({ prisma, engine, }); const result = await service.call( authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, } ); expect(result.sealed).toBe(true); expect(result.id).toBe(batch.friendlyId); expect(result.itemsAccepted).toBe(0); expect(result.itemsDeduplicated).toBe(0); await engine.quit(); } ); containerTest( "should throw when batch is in ABORTED status", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 2, status: "ABORTED", sealed: false, }); const service = new StreamBatchItemsService({ prisma, engine, }); await expect( service.call(authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, }) ).rejects.toThrow(ServiceValidationError); await engine.quit(); } ); containerTest( "should throw when batch is sealed but ABORTED (callback aborted post-seal must surface as error)", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); // V2 batch completion callback sets status=ABORTED (failedRunCount > 0 && // successfulRunCount === 0) without touching sealed=true that the seal // step previously set. The Phase 2 retry must NOT mask this terminal // failure as success — every run failed. const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 2, status: "ABORTED", sealed: true, }); const service = new StreamBatchItemsService({ prisma, engine, }); await expect( service.call(authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, }) ).rejects.toThrow(ServiceValidationError); await engine.quit(); } ); containerTest( "should return sealed=true when batch is PARTIAL_FAILED (Phase 2 retry idempotency)", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); // V2 completion callback sets PARTIAL_FAILED when some run-creation // attempts failed but at least one succeeded. The Phase 2 stream itself // did its job (items were enqueued and processed), so a retry should // see this as terminal success — the per-item failures are visible on // the individual run records. const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 2, status: "PARTIAL_FAILED", sealed: false, }); const service = new StreamBatchItemsService({ prisma, engine, }); const result = await service.call( authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, } ); expect(result.sealed).toBe(true); expect(result.id).toBe(batch.friendlyId); expect(result.itemsAccepted).toBe(0); expect(result.itemsDeduplicated).toBe(0); await engine.quit(); } ); containerTest( "should return sealed=true when batch is PARTIAL_FAILED by callback before seal attempt", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 2, status: "PENDING", sealed: false, }); await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: authenticatedEnvironment.id, environmentType: authenticatedEnvironment.type, organizationId: authenticatedEnvironment.organizationId, projectId: authenticatedEnvironment.projectId, runCount: 2, processingConcurrency: 10, }); await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, { task: "test-task", payload: JSON.stringify({ data: "item1" }), payloadType: "application/json", }); await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, { task: "test-task", payload: JSON.stringify({ data: "item2" }), payloadType: "application/json", }); // Simulate the race where V2's batchCompletionCallback runs between // getEnqueuedCount and the seal updateMany — some runs failed to create // but at least one succeeded, so callback sets status=PARTIAL_FAILED // without setting sealed=true. const racingPrisma = { ...prisma, batchTaskRun: { ...prisma.batchTaskRun, findFirst: prisma.batchTaskRun.findFirst.bind(prisma.batchTaskRun), updateMany: async () => { await prisma.batchTaskRun.update({ where: { id: batch.id }, data: { status: "PARTIAL_FAILED", }, }); return { count: 0 }; }, findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun), }, } as unknown as PrismaClient; // The batch find + seal updateMany now route through the engine's run-store // (route-by-batch-id under the run-ops split), so the racing client must back the // store the service reads through. engine.runStore = new PostgresRunStore({ prisma: racingPrisma, readOnlyPrisma: racingPrisma, }); const service = new StreamBatchItemsService({ prisma: racingPrisma, engine, }); const result = await service.call( authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, } ); expect(result.sealed).toBe(true); expect(result.id).toBe(batch.friendlyId); const updatedBatch = await prisma.batchTaskRun.findUnique({ where: { id: batch.id }, }); expect(updatedBatch?.status).toBe("PARTIAL_FAILED"); expect(updatedBatch?.sealed).toBe(false); await engine.quit(); } ); containerTest( "should return sealed=true when concurrent request already sealed the batch during seal attempt", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); // Create a batch in PENDING state const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 2, status: "PENDING", sealed: false, }); // Initialize the batch in Redis await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: authenticatedEnvironment.id, environmentType: authenticatedEnvironment.type, organizationId: authenticatedEnvironment.organizationId, projectId: authenticatedEnvironment.projectId, runCount: 2, processingConcurrency: 10, }); // Enqueue items await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, { task: "test-task", payload: JSON.stringify({ data: "item1" }), payloadType: "application/json", }); await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, { task: "test-task", payload: JSON.stringify({ data: "item2" }), payloadType: "application/json", }); // Create a custom prisma client that simulates a race condition: // When updateMany is called on batchTaskRun, it returns count: 0 (as if another request beat us) // but the subsequent findUnique shows the batch is sealed and PROCESSING const racingPrisma = { ...prisma, batchTaskRun: { ...prisma.batchTaskRun, findFirst: prisma.batchTaskRun.findFirst.bind(prisma.batchTaskRun), updateMany: async () => { // Simulate another request winning the race - seal the batch first await prisma.batchTaskRun.update({ where: { id: batch.id }, data: { sealed: true, sealedAt: new Date(), status: "PROCESSING", processingStartedAt: new Date(), }, }); // Return 0 as if the conditional update failed return { count: 0 }; }, findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun), }, } as unknown as PrismaClient; // The batch find + seal updateMany now route through the engine's run-store // (route-by-batch-id under the run-ops split), so the racing client must back the // store the service reads through. engine.runStore = new PostgresRunStore({ prisma: racingPrisma, readOnlyPrisma: racingPrisma, }); const service = new StreamBatchItemsService({ prisma: racingPrisma, engine, }); // Call the service - it should detect the race and return success since batch is sealed const result = await service.call( authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, } ); // Should return sealed=true because the batch was sealed (by the "other" request) expect(result.sealed).toBe(true); expect(result.id).toBe(batch.friendlyId); // Verify the batch is sealed in the database const updatedBatch = await prisma.batchTaskRun.findUnique({ where: { id: batch.id }, }); expect(updatedBatch?.sealed).toBe(true); expect(updatedBatch?.status).toBe("PROCESSING"); await engine.quit(); } ); containerTest( "should return sealed=true when batch is COMPLETED by BatchQueue before seal attempt", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); // Create a batch in PENDING state const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 2, status: "PENDING", sealed: false, }); // Initialize the batch in Redis await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: authenticatedEnvironment.id, environmentType: authenticatedEnvironment.type, organizationId: authenticatedEnvironment.organizationId, projectId: authenticatedEnvironment.projectId, runCount: 2, processingConcurrency: 10, }); // Enqueue items - the enqueued count check passes but the seal updateMany // will race with tryCompleteBatch moving status to COMPLETED. await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, { task: "test-task", payload: JSON.stringify({ data: "item1" }), payloadType: "application/json", }); await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, { task: "test-task", payload: JSON.stringify({ data: "item2" }), payloadType: "application/json", }); // Simulate the race where BatchQueue's completionCallback runs // tryCompleteBatch between getEnqueuedCount and the seal updateMany. // tryCompleteBatch sets status=COMPLETED but NOT sealed=true. const racingPrisma = { ...prisma, batchTaskRun: { ...prisma.batchTaskRun, findFirst: prisma.batchTaskRun.findFirst.bind(prisma.batchTaskRun), updateMany: async () => { await prisma.batchTaskRun.update({ where: { id: batch.id }, data: { status: "COMPLETED", }, }); // The conditional updateMany(where: status="PENDING") would now fail return { count: 0 }; }, findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun), }, } as unknown as PrismaClient; // The batch find + seal updateMany now route through the engine's run-store // (route-by-batch-id under the run-ops split), so the racing client must back the // store the service reads through. engine.runStore = new PostgresRunStore({ prisma: racingPrisma, readOnlyPrisma: racingPrisma, }); const service = new StreamBatchItemsService({ prisma: racingPrisma, engine, }); const result = await service.call( authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, } ); // The endpoint should accept the COMPLETED state as a success case so the // SDK does not retry a batch whose child runs have already finished. expect(result.sealed).toBe(true); expect(result.id).toBe(batch.friendlyId); const updatedBatch = await prisma.batchTaskRun.findUnique({ where: { id: batch.id }, }); expect(updatedBatch?.status).toBe("COMPLETED"); // sealed stays false because the BatchQueue completion path does not set // it - that's fine, the batch is terminal. expect(updatedBatch?.sealed).toBe(false); await engine.quit(); } ); containerTest( "should throw error when race condition leaves batch in unexpected state", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); // Create a batch in PENDING state const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 2, status: "PENDING", sealed: false, }); // Initialize the batch in Redis await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: authenticatedEnvironment.id, environmentType: authenticatedEnvironment.type, organizationId: authenticatedEnvironment.organizationId, projectId: authenticatedEnvironment.projectId, runCount: 2, processingConcurrency: 10, }); // Enqueue items await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, { task: "test-task", payload: JSON.stringify({ data: "item1" }), payloadType: "application/json", }); await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, { task: "test-task", payload: JSON.stringify({ data: "item2" }), payloadType: "application/json", }); // Create a custom prisma client that simulates a race condition where // the batch ends up in an unexpected state (ABORTED instead of PROCESSING) const racingPrisma = { ...prisma, batchTaskRun: { ...prisma.batchTaskRun, findFirst: prisma.batchTaskRun.findFirst.bind(prisma.batchTaskRun), updateMany: async () => { // Simulate the batch being aborted by another process await prisma.batchTaskRun.update({ where: { id: batch.id }, data: { sealed: true, status: "ABORTED", }, }); // Return 0 as if the conditional update failed return { count: 0 }; }, findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun), }, } as unknown as PrismaClient; // The batch find + seal updateMany now route through the engine's run-store // (route-by-batch-id under the run-ops split), so the racing client must back the // store the service reads through. engine.runStore = new PostgresRunStore({ prisma: racingPrisma, readOnlyPrisma: racingPrisma, }); const service = new StreamBatchItemsService({ prisma: racingPrisma, engine, }); // Call the service - it should throw because the batch is in an unexpected state await expect( service.call(authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, }) ).rejects.toThrow(/unexpected state/); await engine.quit(); } ); containerTest( "should return sealed=false when item count does not match", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); // Create a batch expecting 3 items const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 3, status: "PENDING", sealed: false, }); // Initialize the batch in Redis await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: authenticatedEnvironment.id, environmentType: authenticatedEnvironment.type, organizationId: authenticatedEnvironment.organizationId, projectId: authenticatedEnvironment.projectId, runCount: 3, processingConcurrency: 10, }); // Only enqueue 2 items (1 short of expected) await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, { task: "test-task", payload: JSON.stringify({ data: "item1" }), payloadType: "application/json", }); await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, { task: "test-task", payload: JSON.stringify({ data: "item2" }), payloadType: "application/json", }); const service = new StreamBatchItemsService({ prisma, engine, }); const result = await service.call( authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, } ); // Should return sealed=false because item count doesn't match expect(result.sealed).toBe(false); expect(result.enqueuedCount).toBe(2); expect(result.expectedCount).toBe(3); // Verify the batch is NOT sealed in the database const updatedBatch = await prisma.batchTaskRun.findUnique({ where: { id: batch.id }, }); expect(updatedBatch?.sealed).toBe(false); expect(updatedBatch?.status).toBe("PENDING"); await engine.quit(); } ); containerTest( "should return sealed=true when seal-failed race produces sealed=true + PENDING (post-callback all-created)", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 2, status: "PENDING", sealed: false, }); await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: authenticatedEnvironment.id, environmentType: authenticatedEnvironment.type, organizationId: authenticatedEnvironment.organizationId, projectId: authenticatedEnvironment.projectId, runCount: 2, processingConcurrency: 10, }); await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, { task: "test-task", payload: JSON.stringify({ data: "item1" }), payloadType: "application/json", }); await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, { task: "test-task", payload: JSON.stringify({ data: "item2" }), payloadType: "application/json", }); // Simulate the race where a concurrent path seals the batch (sealed=true, // PROCESSING), then the V2 batchCompletionCallback fires with all runs // created successfully and resets status to PENDING (sealed stays true). // Our seal updateMany then fails the conditional (sealed=false no longer // matches), and the re-query sees sealed=true + PENDING — a perfectly // valid post-callback state that the SDK retry should treat as success. const racingPrisma = { ...prisma, batchTaskRun: { ...prisma.batchTaskRun, findFirst: prisma.batchTaskRun.findFirst.bind(prisma.batchTaskRun), updateMany: async () => { await prisma.batchTaskRun.update({ where: { id: batch.id }, data: { sealed: true, sealedAt: new Date(), // Intentionally leave status as PENDING — that's exactly what // the V2 batchCompletionCallback does after all runs are // created (status PROCESSING → PENDING). }, }); return { count: 0 }; }, findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun), }, } as unknown as PrismaClient; // The batch find + seal updateMany now route through the engine's run-store // (route-by-batch-id under the run-ops split), so the racing client must back the // store the service reads through. engine.runStore = new PostgresRunStore({ prisma: racingPrisma, readOnlyPrisma: racingPrisma, }); const service = new StreamBatchItemsService({ prisma: racingPrisma, engine, }); const result = await service.call( authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, } ); expect(result.sealed).toBe(true); expect(result.id).toBe(batch.friendlyId); const updatedBatch = await prisma.batchTaskRun.findUnique({ where: { id: batch.id }, }); expect(updatedBatch?.sealed).toBe(true); expect(updatedBatch?.status).toBe("PENDING"); await engine.quit(); } ); containerTest( "should throw when count-mismatch race produces sealed=true + ABORTED (no TaskRuns created)", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 3, status: "PENDING", sealed: false, }); await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: authenticatedEnvironment.id, environmentType: authenticatedEnvironment.type, organizationId: authenticatedEnvironment.organizationId, projectId: authenticatedEnvironment.projectId, runCount: 3, processingConcurrency: 10, }); // Only enqueue 2 items so the post-loop count check trips into the // mismatch handler. The race we're simulating: between our pre-loop // findFirst and the count-mismatch re-query, a concurrent path sealed // the batch, runs were attempted, every run-creation failed AND the // pre-failed-TaskRun fallback also failed → callback sets ABORTED. // The customer has zero TaskRun records to monitor, so the retry must // throw rather than silently succeed. await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 0, { task: "test-task", payload: JSON.stringify({ data: "item1" }), payloadType: "application/json", }); await engine.enqueueBatchItem(batch.id, authenticatedEnvironment.id, 1, { task: "test-task", payload: JSON.stringify({ data: "item2" }), payloadType: "application/json", }); // Override findFirst to flip the batch to sealed=true + ABORTED on the // re-query that happens INSIDE the count-mismatch branch. The first // findFirst (pre-loop) must still see PENDING + sealed=false so we // pass through and reach the count-mismatch branch. let findFirstCallCount = 0; const racingPrisma = { ...prisma, batchTaskRun: { ...prisma.batchTaskRun, findFirst: async (args: Parameters[0]) => { findFirstCallCount++; if (findFirstCallCount >= 2) { await prisma.batchTaskRun.update({ where: { id: batch.id }, data: { sealed: true, sealedAt: new Date(), status: "ABORTED", completedAt: new Date(), }, }); } return prisma.batchTaskRun.findFirst.call(prisma.batchTaskRun, args); }, updateMany: prisma.batchTaskRun.updateMany.bind(prisma.batchTaskRun), findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun), }, } as unknown as PrismaClient; // The batch find + seal updateMany now route through the engine's run-store // (route-by-batch-id under the run-ops split), so the racing client must back the // store the service reads through. engine.runStore = new PostgresRunStore({ prisma: racingPrisma, readOnlyPrisma: racingPrisma, }); const service = new StreamBatchItemsService({ prisma: racingPrisma, engine, }); await expect( service.call(authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, }) ).rejects.toThrow(ServiceValidationError); await engine.quit(); } ); containerTest( "should return sealed=true when batch is sealed=false + PENDING + processingCompletedAt set (pre-loop post-callback)", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); // The V2 batchCompletionCallback set processingCompletedAt without // touching sealed (sealed gets set by streamBatchItems, not the callback). // Status stays PENDING because every run was created successfully (the // callback's "all created, waiting for completion" branch). A Phase 2 // retry arriving in this state must treat it as success — every item // has a TaskRun record for the customer to monitor. const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 4, status: "PENDING", sealed: false, processingCompletedAt: new Date(), }); const service = new StreamBatchItemsService({ prisma, engine, }); const result = await service.call( authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, } ); expect(result.sealed).toBe(true); expect(result.id).toBe(batch.friendlyId); expect(result.itemsAccepted).toBe(0); expect(result.itemsDeduplicated).toBe(0); await engine.quit(); } ); containerTest( "should return sealed=true on count-mismatch when callback fired before our getEnqueuedCount (cleanup race — customer scenario)", async ({ prisma, redisOptions }) => { const engine = new RunEngine({ prisma, worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100, disabled: true, }, queue: { redis: redisOptions, }, runLock: { redis: redisOptions, }, machines: { defaultMachine: "small-1x", machines: { "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001, }, }, baseCostInCents: 0.0005, }, batchQueue: { redis: redisOptions, }, tracer: trace.getTracer("test", "0.0.0"), }); const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); // The customer's exact case: 4-item batch, BatchQueue rushes through // all items before our service finishes its loop, callback fires // (setting processingCompletedAt; status stays PENDING because all // runs created cleanly), cleanup deletes Redis state. Our service // hits the count-mismatch branch because getBatchEnqueuedCount returns // 0 (cleaned). The re-query sees sealed=false + PENDING but // processingCompletedAt is set — the work is done. const batch = await createBatch(prisma, authenticatedEnvironment.id, { runCount: 4, status: "PENDING", sealed: false, }); await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: authenticatedEnvironment.id, environmentType: authenticatedEnvironment.type, organizationId: authenticatedEnvironment.organizationId, projectId: authenticatedEnvironment.projectId, runCount: 4, processingConcurrency: 10, }); // The pre-loop validate-find must see "initial" state (so it passes through to the // loop), and the count-mismatch re-query must see "post-callback" state. Use a findFirst // counter to flip the DB between those two reads, matching production timing where the // callback fires while our loop runs. Both reads route through the engine's run-store // (route-by-batch-id under the split), so the racing client backs the store. let findFirstCallCount = 0; const racingPrisma = { ...prisma, batchTaskRun: { ...prisma.batchTaskRun, findFirst: async (args: Parameters[0]) => { findFirstCallCount++; if (findFirstCallCount === 2) { await prisma.batchTaskRun.update({ where: { id: batch.id }, data: { processingCompletedAt: new Date(), successfulRunCount: 4, }, }); } return prisma.batchTaskRun.findFirst.call(prisma.batchTaskRun, args); }, updateMany: prisma.batchTaskRun.updateMany.bind(prisma.batchTaskRun), findUnique: prisma.batchTaskRun.findUnique.bind(prisma.batchTaskRun), }, } as unknown as PrismaClient; engine.runStore = new PostgresRunStore({ prisma: racingPrisma, readOnlyPrisma: racingPrisma, }); const service = new StreamBatchItemsService({ prisma: racingPrisma, engine, }); const result = await service.call( authenticatedEnvironment, batch.friendlyId, itemsToAsyncIterable([]), { maxItemBytes: 1024 * 1024, concurrency: 1, } ); // The retry must be treated as success — every item's TaskRun was // created by the original Phase 2 call. Returning sealed:false here // (the previous behavior) made the SDK retry the stream against a // cleaned-up batch, which then 5xx'd, exhausted SDK retries, and // surfaced as BatchTriggerError despite all runs succeeding. expect(result.sealed).toBe(true); expect(result.id).toBe(batch.friendlyId); await engine.quit(); } ); containerTest( "ingests a large batch with bounded concurrency and seals it", async ({ prisma, redisOptions }) => { const engine = buildEngine(prisma, redisOptions); const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); const runCount = 100; const batch = await createBatch(prisma, environment.id, { runCount, status: "PENDING", sealed: false, }); await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: environment.id, environmentType: environment.type, organizationId: environment.organizationId, projectId: environment.projectId, runCount, processingConcurrency: 10, }); const service = new StreamBatchItemsService({ prisma, engine }); const result = await service.call( environment, batch.friendlyId, itemsToAsyncIterable(makeItems(runCount)), { maxItemBytes: 1024 * 1024, concurrency: 10 } ); expect(result.sealed).toBe(true); expect(result.itemsAccepted).toBe(runCount); expect(result.itemsDeduplicated).toBe(0); // Every item was enqueued exactly once despite concurrent ingestion. const enqueuedCount = await engine.getBatchEnqueuedCount(batch.id); expect(enqueuedCount).toBe(runCount); const updatedBatch = await prisma.batchTaskRun.findUnique({ where: { id: batch.id } }); expect(updatedBatch?.sealed).toBe(true); expect(updatedBatch?.status).toBe("PROCESSING"); await engine.quit(); } ); containerTest( "bounds in-flight item processing to the configured concurrency", async ({ prisma, redisOptions }) => { const engine = buildEngine(prisma, redisOptions); const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); const runCount = 40; const concurrency = 5; const batch = await createBatch(prisma, environment.id, { runCount, status: "PENDING", sealed: false, }); await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: environment.id, environmentType: environment.type, organizationId: environment.organizationId, projectId: environment.projectId, runCount, processingConcurrency: 10, }); // A real payload processor that holds each slot briefly so we can observe // how many items are processed simultaneously. This guards against an // accidental regression that buffers the whole stream / ignores the cap. let inFlight = 0; let maxInFlight = 0; class TrackingPayloadProcessor extends BatchPayloadProcessor { async process(payload: unknown, payloadType: string): Promise { inFlight++; maxInFlight = Math.max(maxInFlight, inFlight); try { await sleep(20); return { payload, payloadType, wasOffloaded: false, size: 0 }; } finally { inFlight--; } } } const service = new StreamBatchItemsService({ prisma, engine, payloadProcessor: new TrackingPayloadProcessor(), }); const result = await service.call( environment, batch.friendlyId, itemsToAsyncIterable(makeItems(runCount)), { maxItemBytes: 1024 * 1024, concurrency } ); expect(result.sealed).toBe(true); expect(result.itemsAccepted).toBe(runCount); // Actually ran concurrently, but never exceeded the configured cap. expect(maxInFlight).toBeGreaterThan(1); expect(maxInFlight).toBeLessThanOrEqual(concurrency); await engine.quit(); } ); containerTest( "perf: higher ingest concurrency processes a batch proportionally faster", async ({ prisma, redisOptions }) => { const engine = buildEngine(prisma, redisOptions); const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); const runCount = 150; // A payload processor that holds each slot for a fixed duration, modelling // the per-item object-store offload. perItemLatencyMs=0 models a small // payload that never offloads, leaving the real per-item enqueueBatchItem // (a Redis round-trip) as the only serialized cost. Local MinIO/Redis // latency is sub-millisecond and too noisy to compare directly, so the // offload case uses a fixed hold instead. class FixedLatencyPayloadProcessor extends BatchPayloadProcessor { constructor(private readonly perItemLatencyMs: number) { super(); } async process(payload: unknown, payloadType: string): Promise { if (this.perItemLatencyMs > 0) { await sleep(this.perItemLatencyMs); } return { payload, payloadType, wasOffloaded: this.perItemLatencyMs > 0, size: 0 }; } } async function timeIngest(concurrency: number, perItemLatencyMs: number): Promise { // Each run needs its own batch: sealing mutates state and a re-stream // of the same batch would dedup every item. const batch = await createBatch(prisma, environment.id, { runCount, status: "PENDING", sealed: false, }); await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: environment.id, environmentType: environment.type, organizationId: environment.organizationId, projectId: environment.projectId, runCount, processingConcurrency: 10, }); const service = new StreamBatchItemsService({ prisma, engine, payloadProcessor: new FixedLatencyPayloadProcessor(perItemLatencyMs), }); const start = performance.now(); const result = await service.call( environment, batch.friendlyId, itemsToAsyncIterable(makeItems(runCount)), { maxItemBytes: 1024 * 1024, concurrency } ); const elapsedMs = performance.now() - start; // Correctness holds at every concurrency: all items accepted and sealed. expect(result.sealed).toBe(true); expect(result.itemsAccepted).toBe(runCount); expect(result.itemsDeduplicated).toBe(0); expect(await engine.getBatchEnqueuedCount(batch.id)).toBe(runCount); return elapsedMs; } // Scenario A: large payloads, where each item offloads to object storage. const offloadLatencyMs = 10; const offloadSeqMs = await timeIngest(1, offloadLatencyMs); const offload10Ms = await timeIngest(10, offloadLatencyMs); const offload50Ms = await timeIngest(50, offloadLatencyMs); // Scenario B: small payloads (no offload). The only per-item cost is the // real Redis enqueue, so this is the floor case that proves the speedup // applies to all batch ingest, not just large-payload batches. const enqueueSeqMs = await timeIngest(1, 0); const enqueue10Ms = await timeIngest(10, 0); // eslint-disable-next-line no-console console.log( `\n[streamBatchItems perf] runCount=${runCount}\n` + ` large payloads (${offloadLatencyMs}ms/item offload):\n` + ` concurrency=1 ${offloadSeqMs.toFixed(0)}ms\n` + ` concurrency=10 ${offload10Ms.toFixed(0)}ms (${(offloadSeqMs / offload10Ms).toFixed( 1 )}x faster)\n` + ` concurrency=50 ${offload50Ms.toFixed(0)}ms (${(offloadSeqMs / offload50Ms).toFixed( 1 )}x faster)\n` + ` small payloads (Redis enqueue only, no offload):\n` + ` concurrency=1 ${enqueueSeqMs.toFixed(0)}ms\n` + ` concurrency=10 ${enqueue10Ms.toFixed(0)}ms (${(enqueueSeqMs / enqueue10Ms).toFixed( 1 )}x faster)\n` ); // Sequential floor: N items each held for offloadLatencyMs cannot finish // faster than N x latency. Parallel ingest must beat that floor decisively. const sequentialFloorMs = runCount * offloadLatencyMs; expect(offloadSeqMs).toBeGreaterThan(sequentialFloorMs * 0.8); // 10x concurrency on a latency-bound workload should be well over 2x faster. expect(offload10Ms).toBeLessThan(offloadSeqMs / 2); // More concurrency keeps helping (or at least never regresses). expect(offload50Ms).toBeLessThanOrEqual(offload10Ms * 1.2); // Even with no offload, overlapping the per-item Redis enqueue is strictly // faster than doing them one at a time. expect(enqueue10Ms).toBeLessThan(enqueueSeqMs); await engine.quit(); } ); containerTest( "deduplicates already-enqueued items during concurrent ingest (Phase 2 retry)", async ({ prisma, redisOptions }) => { const engine = buildEngine(prisma, redisOptions); const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); const runCount = 20; const batch = await createBatch(prisma, environment.id, { runCount, status: "PENDING", sealed: false, }); await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: environment.id, environmentType: environment.type, organizationId: environment.organizationId, projectId: environment.projectId, runCount, processingConcurrency: 10, }); // Pre-enqueue the first half (simulating a partially-completed prior stream // that failed before sealing). The retry re-streams everything. for (let index = 0; index < 10; index++) { await engine.enqueueBatchItem(batch.id, environment.id, index, { task: "test-task", payload: JSON.stringify({ i: index }), payloadType: "application/json", }); } const service = new StreamBatchItemsService({ prisma, engine }); const result = await service.call( environment, batch.friendlyId, itemsToAsyncIterable(makeItems(runCount)), { maxItemBytes: 1024 * 1024, concurrency: 8 } ); expect(result.sealed).toBe(true); expect(result.itemsDeduplicated).toBe(10); expect(result.itemsAccepted).toBe(10); const enqueuedCount = await engine.getBatchEnqueuedCount(batch.id); expect(enqueuedCount).toBe(runCount); await engine.quit(); } ); containerTest( "rejects an oversized item whose index exceeds runCount", async ({ prisma, redisOptions }) => { const engine = buildEngine(prisma, redisOptions); const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); const batch = await createBatch(prisma, environment.id, { runCount: 2, status: "PENDING", sealed: false, }); await engine.initializeBatch({ batchId: batch.id, friendlyId: batch.friendlyId, environmentId: environment.id, environmentType: environment.type, organizationId: environment.organizationId, projectId: environment.projectId, runCount: 2, processingConcurrency: 10, }); const service = new StreamBatchItemsService({ prisma, engine }); // An oversized marker must hit the same out-of-range guard as a normal // item rather than slipping through to create a stray pre-failed run. async function* oversizedOutOfRange() { yield { __batchItemError: "OVERSIZED", index: 5, task: "test-task", actualSize: 9999, maxSize: 1000, }; } await expect( service.call(environment, batch.friendlyId, oversizedOutOfRange(), { maxItemBytes: 1024, concurrency: 4, }) ).rejects.toThrow(ServiceValidationError); await engine.quit(); } ); }); describe("createNdjsonParserStream", () => { async function collectStream(stream: ReadableStream): Promise { const results: T[] = []; for await (const item of streamToAsyncIterable(stream)) { results.push(item); } return results; } function chunksToStream(chunks: Uint8Array[]): ReadableStream { let index = 0; return new ReadableStream({ pull(controller) { if (index < chunks.length) { controller.enqueue(chunks[index++]); } else { controller.close(); } }, }); } it("should parse basic NDJSON correctly", async () => { const ndjson = '{"name":"alice"}\n{"name":"bob"}\n{"name":"charlie"}\n'; const encoder = new TextEncoder(); const stream = chunksToStream([encoder.encode(ndjson)]); const parser = createNdjsonParserStream(1024); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toEqual([{ name: "alice" }, { name: "bob" }, { name: "charlie" }]); }); it("should handle NDJSON without trailing newline", async () => { const ndjson = '{"id":1}\n{"id":2}'; const encoder = new TextEncoder(); const stream = chunksToStream([encoder.encode(ndjson)]); const parser = createNdjsonParserStream(1024); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toEqual([{ id: 1 }, { id: 2 }]); }); it("should handle escaped newlines in JSON string values", async () => { // JSON.stringify escapes newlines as \n (two chars: backslash + n), // so they don't break NDJSON line boundaries. This is the normal case // when the SDK serializes payloads containing newlines. const item1 = JSON.stringify({ payload: "line1\nline2\nline3" }); const item2 = JSON.stringify({ payload: "no newlines" }); const ndjson = item1 + "\n" + item2 + "\n"; const encoder = new TextEncoder(); const stream = chunksToStream([encoder.encode(ndjson)]); const parser = createNdjsonParserStream(1024); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toEqual([{ payload: "line1\nline2\nline3" }, { payload: "no newlines" }]); }); it("should skip empty lines", async () => { const ndjson = '{"a":1}\n\n{"b":2}\n \n{"c":3}\n'; const encoder = new TextEncoder(); const stream = chunksToStream([encoder.encode(ndjson)]); const parser = createNdjsonParserStream(1024); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toEqual([{ a: 1 }, { b: 2 }, { c: 3 }]); }); it("should handle chunks split mid-line", async () => { // Split '{"split":"value"}\n' across multiple chunks const encoder = new TextEncoder(); const chunks = [encoder.encode('{"spl'), encoder.encode('it":"va'), encoder.encode('lue"}\n')]; const stream = chunksToStream(chunks); const parser = createNdjsonParserStream(1024); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toEqual([{ split: "value" }]); }); it("should handle multibyte UTF-8 characters split across chunks", async () => { // Test with emoji and other multibyte characters // The emoji "😀" is 4 bytes: 0xF0 0x9F 0x98 0x80 const json = '{"emoji":"😀"}\n'; const fullBytes = new TextEncoder().encode(json); // Split in the middle of the emoji (between byte 2 and 3 of the 4-byte sequence) // Find where the emoji starts const emojiStart = fullBytes.indexOf(0xf0); expect(emojiStart).toBeGreaterThan(0); // Split after first 2 bytes of the emoji const chunk1 = fullBytes.slice(0, emojiStart + 2); const chunk2 = fullBytes.slice(emojiStart + 2); const stream = chunksToStream([chunk1, chunk2]); const parser = createNdjsonParserStream(1024); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toEqual([{ emoji: "😀" }]); }); it("should handle multiple multibyte characters across chunks", async () => { // Japanese text: "こんにちは" (each hiragana is 3 bytes in UTF-8) const json = '{"greeting":"こんにちは"}\n'; const fullBytes = new TextEncoder().encode(json); // Split into many small chunks to stress test UTF-8 handling const chunkSize = 3; // Deliberately misaligned with UTF-8 boundaries const chunks: Uint8Array[] = []; for (let i = 0; i < fullBytes.length; i += chunkSize) { chunks.push(fullBytes.slice(i, i + chunkSize)); } const stream = chunksToStream(chunks); const parser = createNdjsonParserStream(1024); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toEqual([{ greeting: "こんにちは" }]); }); it("should emit OversizedItemMarker for lines exceeding maxItemBytes", async () => { const maxBytes = 50; // Create a line that exceeds the limit with index and task fields const largeJson = JSON.stringify({ index: 3, task: "my-task", data: "x".repeat(100) }) + "\n"; const encoder = new TextEncoder(); const stream = chunksToStream([encoder.encode(largeJson)]); const parser = createNdjsonParserStream(maxBytes); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toHaveLength(1); const marker = results[0] as OversizedItemMarker; expect(marker.__batchItemError).toBe("OVERSIZED"); expect(marker.index).toBe(3); expect(marker.task).toBe("my-task"); expect(marker.maxSize).toBe(maxBytes); expect(marker.actualSize).toBeGreaterThan(maxBytes); }); it("should emit OversizedItemMarker for unbounded accumulation without newlines", async () => { const maxBytes = 50; // Send data without any newlines that exceeds the buffer limit const encoder = new TextEncoder(); const chunks = [ encoder.encode('{"index":7,"task":"big-task","start":"'), encoder.encode("x".repeat(60)), // This will push buffer over 50 bytes ]; const stream = chunksToStream(chunks); const parser = createNdjsonParserStream(maxBytes); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toHaveLength(1); const marker = results[0] as OversizedItemMarker; expect(marker.__batchItemError).toBe("OVERSIZED"); expect(marker.index).toBe(7); expect(marker.task).toBe("big-task"); expect(marker.maxSize).toBe(maxBytes); }); it("should skip remaining bytes of oversized line arriving in subsequent chunks", async () => { const maxBytes = 50; const encoder = new TextEncoder(); // Simulate a normal item, then an oversized item split across many chunks, // then another normal item after the newline. // The oversized line is: {"index":1,"task":"t","data":"xxxx...120 x's...xxxx"}\n const normalItem1 = '{"index":0,"task":"t","x":1}\n'; const oversizedStart = '{"index":1,"task":"t","data":"'; const oversizedMiddle = "x".repeat(120); // way over 50 bytes const oversizedEnd = '"}\n'; const normalItem2 = '{"index":2,"task":"t","x":2}\n'; // Send as separate chunks to trigger Case 2 (no newline, buffer > limit) const chunks = [ encoder.encode(normalItem1 + oversizedStart), encoder.encode(oversizedMiddle.slice(0, 60)), encoder.encode(oversizedMiddle.slice(60)), encoder.encode(oversizedEnd + normalItem2), ]; const stream = chunksToStream(chunks); const parser = createNdjsonParserStream(maxBytes); const results = await collectStream(stream.pipeThrough(parser)); // Should get: normal item 1, oversized marker, normal item 2 expect(results).toHaveLength(3); expect(results[0]).toEqual({ index: 0, task: "t", x: 1 }); expect((results[1] as OversizedItemMarker).__batchItemError).toBe("OVERSIZED"); expect((results[1] as OversizedItemMarker).index).toBe(1); expect(results[2]).toEqual({ index: 2, task: "t", x: 2 }); }); it("stamps an oversized marker's index by emit position when it can't be extracted", async () => { // Concurrency-safety: when the `index` field can't be recovered from the // (truncated) raw bytes, the parser must back-fill it from the emit // position rather than leaving it for the consumer to infer sequentially. const maxBytes = 50; const encoder = new TextEncoder(); // item 0: small + valid const item0 = '{"index":0,"task":"t","x":1}\n'; // item 1: oversized AND missing an `index` field → extractIndexAndTask returns -1 const item1 = `{"task":"big","data":"${"x".repeat(100)}"}\n`; // item 2: small + valid, must keep its own index const item2 = '{"index":2,"task":"t","x":2}\n'; const stream = chunksToStream([encoder.encode(item0 + item1 + item2)]); const parser = createNdjsonParserStream(maxBytes); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toHaveLength(3); expect(results[0]).toEqual({ index: 0, task: "t", x: 1 }); const marker = results[1] as OversizedItemMarker; expect(marker.__batchItemError).toBe("OVERSIZED"); expect(marker.index).toBe(1); // emit-position fallback, not -1 expect(marker.task).toBe("big"); expect(results[2]).toEqual({ index: 2, task: "t", x: 2 }); }); it("should check byte size before decoding to prevent OOM", async () => { // This test verifies that size is checked on raw bytes, not decoded string length // Unicode characters like "🎉" are 4 bytes but 2 UTF-16 code units (string length 2) const maxBytes = 30; // Create a line with emojis - each emoji is 4 bytes // {"e":"🎉🎉🎉🎉🎉"} = 5 + 20 (5 emojis * 4 bytes) + 2 = 27 bytes - should pass const smallJson = '{"e":"🎉🎉🎉🎉🎉"}\n'; const smallBytes = new TextEncoder().encode(smallJson); expect(smallBytes.length).toBeLessThan(maxBytes); // {"e":"🎉🎉🎉🎉🎉🎉🎉"} = 7 emojis * 4 bytes + overhead = 35 bytes - should fail const largeJson = '{"e":"🎉🎉🎉🎉🎉🎉🎉"}\n'; const largeBytes = new TextEncoder().encode(largeJson); expect(largeBytes.length).toBeGreaterThan(maxBytes); // Small one should succeed const stream1 = chunksToStream([smallBytes]); const parser1 = createNdjsonParserStream(maxBytes); const results1 = await collectStream(stream1.pipeThrough(parser1)); expect(results1).toHaveLength(1); // Large one should emit an OversizedItemMarker const stream2 = chunksToStream([largeBytes]); const parser2 = createNdjsonParserStream(maxBytes); const results2 = await collectStream(stream2.pipeThrough(parser2)); expect(results2).toHaveLength(1); expect((results2[0] as OversizedItemMarker).__batchItemError).toBe("OVERSIZED"); }); it("should handle final line in flush without trailing newline", async () => { const ndjson = '{"first":1}\n{"second":2}'; // No trailing newline const encoder = new TextEncoder(); const stream = chunksToStream([encoder.encode(ndjson)]); const parser = createNdjsonParserStream(1024); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toEqual([{ first: 1 }, { second: 2 }]); }); it("should reject invalid JSON", async () => { const ndjson = '{"valid":true}\n{invalid json}\n'; const encoder = new TextEncoder(); const stream = chunksToStream([encoder.encode(ndjson)]); const parser = createNdjsonParserStream(1024); await expect(collectStream(stream.pipeThrough(parser))).rejects.toThrow( /Invalid JSON at line 2/ ); }); it("should reject invalid UTF-8 sequences", async () => { // Invalid UTF-8: 0xFF is never valid in UTF-8 const invalidBytes = new Uint8Array([ 0x7b, 0x22, 0x78, 0x22, 0x3a, 0xff, 0x7d, 0x0a, // {"x":�}\n with invalid byte ]); const stream = chunksToStream([invalidBytes]); const parser = createNdjsonParserStream(1024); await expect(collectStream(stream.pipeThrough(parser))).rejects.toThrow(/Invalid UTF-8/); }); it("should handle many small chunks efficiently", async () => { // Simulate streaming byte-by-byte const json = '{"test":"value"}\n'; const bytes = new TextEncoder().encode(json); const chunks = Array.from(bytes).map((b) => new Uint8Array([b])); const stream = chunksToStream(chunks); const parser = createNdjsonParserStream(1024); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toEqual([{ test: "value" }]); }); it("should handle multiple lines per chunk", async () => { const ndjson = '{"a":1}\n{"b":2}\n{"c":3}\n{"d":4}\n{"e":5}\n'; const encoder = new TextEncoder(); const stream = chunksToStream([encoder.encode(ndjson)]); const parser = createNdjsonParserStream(1024); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toEqual([{ a: 1 }, { b: 2 }, { c: 3 }, { d: 4 }, { e: 5 }]); }); it("should handle empty stream", async () => { const stream = chunksToStream([]); const parser = createNdjsonParserStream(1024); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toEqual([]); }); it("should pass normal items and emit markers for oversized items in the same stream", async () => { const maxBytes = 50; const encoder = new TextEncoder(); // Normal item, then oversized item, then another normal item const normalItem1 = '{"index":0,"task":"t","x":1}\n'; const oversizedItem = JSON.stringify({ index: 1, task: "t", data: "x".repeat(100) }) + "\n"; const normalItem2 = '{"index":2,"task":"t","x":2}\n'; const stream = chunksToStream([encoder.encode(normalItem1 + oversizedItem + normalItem2)]); const parser = createNdjsonParserStream(maxBytes); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toHaveLength(3); // First: normal parsed object expect(results[0]).toEqual({ index: 0, task: "t", x: 1 }); // Second: oversized marker expect((results[1] as OversizedItemMarker).__batchItemError).toBe("OVERSIZED"); expect((results[1] as OversizedItemMarker).index).toBe(1); // Third: normal parsed object expect(results[2]).toEqual({ index: 2, task: "t", x: 2 }); }); it("should handle stream with only whitespace", async () => { const encoder = new TextEncoder(); const stream = chunksToStream([encoder.encode(" \n\n \n")]); const parser = createNdjsonParserStream(1024); const results = await collectStream(stream.pipeThrough(parser)); expect(results).toEqual([]); }); }); describe("extractIndexAndTask", () => { const encoder = new TextEncoder(); it("should extract index and task from JSON bytes", () => { const bytes = encoder.encode('{"index":42,"task":"my-task","data":"x"}'); const result = extractIndexAndTask(bytes); expect(result.index).toBe(42); expect(result.task).toBe("my-task"); }); it("should return defaults for empty or malformed bytes", () => { const result = extractIndexAndTask(new Uint8Array(0)); expect(result.index).toBe(-1); expect(result.task).toBe("unknown"); }); it("should handle keys in any order", () => { const bytes = encoder.encode('{"task":"other-task","data":"y","index":99}'); const result = extractIndexAndTask(bytes); expect(result.index).toBe(99); expect(result.task).toBe("other-task"); }); it("should not match nested keys", () => { const bytes = encoder.encode( '{"nested":{"index":999,"task":"inner"},"index":5,"task":"outer"}' ); const result = extractIndexAndTask(bytes); expect(result.index).toBe(5); expect(result.task).toBe("outer"); }); });