diff --git a/.scratch/plan-graceful-oversized-batch-items.md b/.scratch/plan-graceful-oversized-batch-items.md new file mode 100644 index 000000000..cb463b962 --- /dev/null +++ b/.scratch/plan-graceful-oversized-batch-items.md @@ -0,0 +1,257 @@ +# Graceful handling of oversized batch items + +## Prerequisites + +This plan builds on top of PR #2980 which provides: +- `TriggerFailedTaskService` at `apps/webapp/app/runEngine/services/triggerFailedTask.server.ts` - creates pre-failed TaskRuns with proper trace events, waitpoint connections, and parent run associations +- `engine.createFailedTaskRun()` on RunEngine - creates a SYSTEM_FAILURE run with associated waitpoints +- Retry support in `processItemCallback` with `attempt` and `isFinalAttempt` params +- The callback already uses `TriggerFailedTaskService` for items that fail after retries + +## Problem + +When the NDJSON parser in `createNdjsonParserStream` detects an oversized line, it throws inside the TransformStream's `transform()` method. This aborts the request body stream (due to `pipeThrough` coupling), causing the client's `fetch()` to see `TypeError: fetch failed` instead of the server's 400 response. The SDK treats this as a connection error and retries with exponential backoff (~25s wasted). + +## Goal + +Instead of throwing, treat oversized items as per-item failures that flow through the existing batch failure pipeline. The batch seals normally, other items process fine, and the user sees a clear failure for the specific oversized item(s). + +## Approach + +The NDJSON parser emits an error marker object instead of throwing. `StreamBatchItemsService` detects these markers and enqueues the item to the FairQueue with error metadata in its options. The `processItemCallback` (already enhanced with `TriggerFailedTaskService` in PR #2980) detects the error metadata and creates a pre-failed run via `TriggerFailedTaskService`, which handles all the waitpoint/trace machinery. + +## Changes + +### 1. Byte-level key extractor for oversized lines + +**`apps/webapp/app/runEngine/services/streamBatchItems.server.ts`** - new function + +Add `extractIndexAndTask(bytes: Uint8Array): { index: number; task: string }` - a state machine that extracts top-level `"index"` and `"task"` values from raw bytes without decoding the full line. + +How it works: +- Scan bytes tracking JSON nesting depth (count `{`/`[` vs `}`/`]`) +- At depth 1 (inside the top-level object), look for byte sequences matching `"index"` and `"task"` key patterns +- For `"index"`: after the `:`, parse the digit sequence as a number +- For `"task"`: after the `:`, find opening `"`, read bytes until closing `"`, decode just that slice +- Stop when both found, or after scanning 512 bytes (whichever comes first) +- Fallback: `index = -1`, `task = "unknown"` if not found + +This avoids decoding/allocating the full 3MB line - only the first few hundred bytes are examined. + +### 2. Modify `createNdjsonParserStream` to emit error markers + +**`apps/webapp/app/runEngine/services/streamBatchItems.server.ts`** + +Define a marker type: +```typescript +type OversizedItemMarker = { + __batchItemError: "OVERSIZED"; + index: number; + task: string; + actualSize: number; + maxSize: number; +}; +``` + +**Case 1 - Complete line exceeds limit** (newline found, `newlineIndex > maxItemBytes`): +- Call `extractLine(newlineIndex)` to consume the line from the buffer +- Call `extractIndexAndTask(lineBytes)` on the extracted bytes +- `controller.enqueue(marker)` instead of throwing +- Increment `lineNumber` and continue + +**Case 2 - Incomplete line exceeds limit** (no newline, `totalBytes > maxItemBytes`): +- Call `extractIndexAndTask(concatenateChunks())` on current buffer +- `controller.enqueue(marker)` +- Clear the buffer (`chunks = []; totalBytes = 0`) +- Return from transform (don't throw) + +**Case 3 - Flush with oversized remaining** (`totalBytes > maxItemBytes` in flush): +- Same as case 2 but in `flush()`. + +### 3. Handle markers in `StreamBatchItemsService` + +**`apps/webapp/app/runEngine/services/streamBatchItems.server.ts`** - in the `for await` loop + +Before the existing `BatchItemNDJSONSchema.safeParse(rawItem)`, check for the marker: + +```typescript +if (rawItem && typeof rawItem === "object" && (rawItem as any).__batchItemError === "OVERSIZED") { + const marker = rawItem as OversizedItemMarker; + const itemIndex = marker.index >= 0 ? marker.index : lastIndex + 1; + + 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 the item normally but with error metadata in options. + // The processItemCallback will detect __error and use TriggerFailedTaskService + // to create a pre-failed run with proper waitpoint connections. + 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, itemIndex, batchItem + ); + + if (result.enqueued) { + itemsAccepted++; + } else { + itemsDeduplicated++; + } + lastIndex = itemIndex; + continue; +} +``` + +### 4. Handle `__error` items in `processItemCallback` + +**`apps/webapp/app/v3/runEngineHandlers.server.ts`** - in the `setupBatchQueueCallbacks` function + +In the `processItemCallback`, before the `TriggerTaskService.call()`, check for `__error` in `item.options`: + +```typescript +const itemError = item.options?.__error as string | undefined; +if (itemError) { + const errorCode = (item.options?.__errorCode as string) ?? "ITEM_ERROR"; + + // Use TriggerFailedTaskService to create a pre-failed run. + // This creates a proper TaskRun with waitpoint connections so the + // parent's batchTriggerAndWait resolves correctly for this item. + const failedRunId = await triggerFailedTaskService.call({ + taskId: item.task, + environment, + payload: item.payload ?? "{}", + payloadType: item.payloadType, + errorMessage: itemError, + errorCode: errorCode as TaskRunErrorCodes, + parentRunId: meta.parentRunId, + resumeParentOnCompletion: meta.resumeParentOnCompletion, + batch: { id: batchId, index: itemIndex }, + traceContext: meta.traceContext as Record | undefined, + spanParentAsLink: meta.spanParentAsLink, + }); + + if (failedRunId) { + span.setAttribute("batch.result.pre_failed", true); + span.setAttribute("batch.result.run_id", failedRunId); + span.end(); + return { success: true as const, runId: failedRunId }; + } + + // Fallback if TriggerFailedTaskService fails + span.end(); + return { success: false as const, error: itemError, errorCode }; +} +``` + +Note: this returns `{ success: true, runId }` because the pre-failed run IS a real run. The BatchQueue records it as a success (run was created). The run itself is already in SYSTEM_FAILURE status, so the batch completion flow handles it correctly. + +If `environment` is null (environment not found), fall through to the existing environment-not-found handling which already uses `triggerFailedTaskService.callWithoutTraceEvents()` on `isFinalAttempt`. + +### 5. Handle undefined/null payload in BatchQueue serialization + +**`internal-packages/run-engine/src/batch-queue/index.ts`** - in `#handleMessage` + +Both payload serialization blocks (in the `success: false` branch and the `catch` block) do: +```typescript +const str = typeof item.payload === "string" ? item.payload : JSON.stringify(item.payload); +innerSpan?.setAttribute("batch.payloadSize", str.length); +``` + +`JSON.stringify(undefined)` returns `undefined`, causing `str.length` to crash. Fix both: +```typescript +const str = + item.payload === undefined || item.payload === null + ? "{}" + : typeof item.payload === "string" + ? item.payload + : JSON.stringify(item.payload); +``` + +### 6. Remove stale error handling in route + +**`apps/webapp/app/routes/api.v3.batches.$batchId.items.ts`** + +The `error.message.includes("exceeds maximum size")` branch is no longer reachable since oversized items don't throw. Remove that condition, keep the `"Invalid JSON"` check. + +### 7. Remove `BatchItemTooLargeError` and SDK pre-validation + +**`packages/core/src/v3/apiClient/errors.ts`** - remove `BatchItemTooLargeError` class + +**`packages/core/src/v3/apiClient/index.ts`**: +- Remove `BatchItemTooLargeError` import +- Remove `instanceof BatchItemTooLargeError` check in the retry catch block +- Remove `MAX_BATCH_ITEM_BYTES` constant +- Remove size validation from `createNdjsonStream` (revert `encodeAndValidate` to simple encode) + +**`packages/trigger-sdk/src/v3/shared.ts`** - remove `BatchItemTooLargeError` import and handling in `buildBatchErrorMessage` + +**`packages/trigger-sdk/src/v3/index.ts`** - remove `BatchItemTooLargeError` re-export + +### 8. Update tests + +**`apps/webapp/test/engine/streamBatchItems.test.ts`**: +- Update "should reject lines exceeding maxItemBytes" to assert `OversizedItemMarker` emission instead of throw +- Update "should reject unbounded accumulation without newlines" similarly +- Update the emoji byte-size test to assert marker instead of throw + +### 9. Update reference project test task + +**`references/hello-world/src/trigger/batches.ts`**: +- Remove `BatchItemTooLargeError` import +- Update `batchSealFailureOversizedPayload` task to test the new behavior: + - Send 2 items: one normal, one oversized (~3.2MB) + - Assert `batchTriggerAndWait` returns (doesn't throw) + - Assert `results.runs[0].ok === true` (normal item succeeded) + - Assert `results.runs[1].ok === false` (oversized item failed) + - Assert error message contains "too large" + +## Data flow + +``` +NDJSON bytes arrive + | +createNdjsonParserStream + |-- Line <= limit --> parse JSON --> enqueue object + `-- Line > limit --> extractIndexAndTask(bytes) --> enqueue OversizedItemMarker + | +StreamBatchItemsService for-await loop + |-- OversizedItemMarker --> engine.enqueueBatchItem() with __error in options + `-- Normal item --> validate --> engine.enqueueBatchItem() + | +FairQueue consumer (#handleMessage) + |-- __error in options --> processItemCallback detects it + | --> TriggerFailedTaskService.call() + | --> Creates pre-failed TaskRun with SYSTEM_FAILURE status + | --> Proper waitpoint + TaskRunWaitpoint connections created + | --> Returns { success: true, runId: failedRunFriendlyId } + `-- Normal item --> TriggerTaskService.call() --> creates normal run + | +Batch sealing: enqueuedCount === runCount (all items go through enqueueBatchItem) +Batch completion: all items have runs (real or pre-failed), waitpoints resolve normally +Parent run: batchTriggerAndWait resolves with per-item results +``` + +## Why this works + +The key insight is that `TriggerFailedTaskService` (from PR #2980) creates a real `TaskRun` in `SYSTEM_FAILURE` status. This means: +1. A RUN waitpoint is created and connected to the parent via `TaskRunWaitpoint` with correct `batchId`/`batchIndex` +2. The run is immediately completed, which completes the waitpoint +3. The SDK's `waitForBatch` resolver for that index fires with the error result +4. The batch completion flow counts this as a processed item (it's a real run) +5. No special-casing needed in the batch completion callback + +## Verification + +1. Rebuild `@trigger.dev/core`, `@trigger.dev/sdk`, `@internal/run-engine` +2. Restart webapp + trigger dev +3. Trigger `batch-seal-failure-oversized` task - should complete in ~2-3s with: + - Normal item: `ok: true` + - Oversized item: `ok: false` with "too large" error +4. Run NDJSON parser tests: updated tests assert marker emission instead of throws +5. Run `pnpm run build --filter @internal/run-engine --filter @trigger.dev/core --filter @trigger.dev/sdk` diff --git a/packages/trigger-sdk/package.json b/packages/trigger-sdk/package.json index f9a122724..7fe2e18a7 100644 --- a/packages/trigger-sdk/package.json +++ b/packages/trigger-sdk/package.json @@ -25,7 +25,8 @@ ".": "./src/v3/index.ts", "./v3": "./src/v3/index.ts", "./ai": "./src/v3/ai.ts", - "./chat": "./src/v3/chat.ts" + "./chat": "./src/v3/chat.ts", + "./chat/react": "./src/v3/chat-react.ts" }, "sourceDialects": [ "@triggerdotdev/source" @@ -41,6 +42,9 @@ ], "chat": [ "dist/commonjs/v3/chat.d.ts" + ], + "chat/react": [ + "dist/commonjs/v3/chat-react.d.ts" ] } }, @@ -70,6 +74,7 @@ "devDependencies": { "@arethetypeswrong/cli": "^0.15.4", "@types/debug": "^4.1.7", + "@types/react": "^19.2.14", "@types/slug": "^5.0.3", "@types/uuid": "^9.0.0", "@types/ws": "^8.5.3", @@ -82,12 +87,16 @@ "zod": "3.25.76" }, "peerDependencies": { - "zod": "^3.0.0 || ^4.0.0", - "ai": "^5.0.0 || ^6.0.0" + "ai": "^5.0.0 || ^6.0.0", + "react": "^18.0 || ^19.0", + "zod": "^3.0.0 || ^4.0.0" }, "peerDependenciesMeta": { "ai": { "optional": true + }, + "react": { + "optional": true } }, "engines": { @@ -138,6 +147,17 @@ "types": "./dist/commonjs/v3/chat.d.ts", "default": "./dist/commonjs/v3/chat.js" } + }, + "./chat/react": { + "import": { + "@triggerdotdev/source": "./src/v3/chat-react.ts", + "types": "./dist/esm/v3/chat-react.d.ts", + "default": "./dist/esm/v3/chat-react.js" + }, + "require": { + "types": "./dist/commonjs/v3/chat-react.d.ts", + "default": "./dist/commonjs/v3/chat-react.js" + } } }, "main": "./dist/commonjs/v3/index.js", diff --git a/packages/trigger-sdk/src/v3/chat-react.ts b/packages/trigger-sdk/src/v3/chat-react.ts new file mode 100644 index 000000000..a62496463 --- /dev/null +++ b/packages/trigger-sdk/src/v3/chat-react.ts @@ -0,0 +1,84 @@ +"use client"; + +/** + * @module @trigger.dev/sdk/chat/react + * + * React hooks for AI SDK chat transport integration. + * Use alongside `@trigger.dev/sdk/chat` for a type-safe, ergonomic DX. + * + * @example + * ```tsx + * import { useChat } from "@ai-sdk/react"; + * import { useTriggerChatTransport } from "@trigger.dev/sdk/chat/react"; + * import type { chat } from "@/trigger/chat"; + * + * function Chat() { + * const transport = useTriggerChatTransport({ + * task: "ai-chat", + * accessToken: () => fetchToken(), + * }); + * + * const { messages, sendMessage } = useChat({ transport }); + * } + * ``` + */ + +import { useRef } from "react"; +import { + TriggerChatTransport, + type TriggerChatTransportOptions, +} from "./chat.js"; +import type { AnyTask, TaskIdentifier } from "@trigger.dev/core/v3"; + +/** + * Options for `useTriggerChatTransport`, with a type-safe `task` field. + * + * Pass a task type parameter to get compile-time validation of the task ID: + * ```ts + * useTriggerChatTransport({ task: "my-task", ... }) + * ``` + */ +export type UseTriggerChatTransportOptions = Omit< + TriggerChatTransportOptions, + "task" +> & { + /** The task ID. Strongly typed when a task type parameter is provided. */ + task: TaskIdentifier; +}; + +/** + * React hook that creates and memoizes a `TriggerChatTransport` instance. + * + * The transport is created once on first render and reused for the lifetime + * of the component. This avoids the need for `useMemo` and ensures the + * transport's internal session state (waitpoint tokens, lastEventId, etc.) + * is preserved across re-renders. + * + * For dynamic access tokens, pass a function — it will be called on each + * request without needing to recreate the transport. + * + * @example + * ```tsx + * import { useChat } from "@ai-sdk/react"; + * import { useTriggerChatTransport } from "@trigger.dev/sdk/chat/react"; + * import type { chat } from "@/trigger/chat"; + * + * function Chat() { + * const transport = useTriggerChatTransport({ + * task: "ai-chat", + * accessToken: () => fetchToken(), + * }); + * + * const { messages, sendMessage } = useChat({ transport }); + * } + * ``` + */ +export function useTriggerChatTransport( + options: UseTriggerChatTransportOptions +): TriggerChatTransport { + const ref = useRef(null); + if (ref.current === null) { + ref.current = new TriggerChatTransport(options); + } + return ref.current; +} diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 60c3445a7..e36c57618 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -391,3 +391,4 @@ export class TriggerChatTransport implements ChatTransport { export function createChatTransport(options: TriggerChatTransportOptions): TriggerChatTransport { return new TriggerChatTransport(options); } + diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 094e9af14..3514ce7f0 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2130,6 +2130,9 @@ importers: '@types/debug': specifier: ^4.1.7 version: 4.1.7 + '@types/react': + specifier: ^19.2.14 + version: 19.2.14 '@types/slug': specifier: ^5.0.3 version: 5.0.3 @@ -11018,6 +11021,9 @@ packages: '@types/react@19.2.14': resolution: {integrity: sha512-ilcTH/UniCkMdtexkoCN0bI7pMcJDvmQFPvuPvmEaYA/NSfFTAgdUSLAoVjaRJm7+6PvcM+q1zYOwS4wTYMF9w==} + '@types/react@19.2.14': + resolution: {integrity: sha512-ilcTH/UniCkMdtexkoCN0bI7pMcJDvmQFPvuPvmEaYA/NSfFTAgdUSLAoVjaRJm7+6PvcM+q1zYOwS4wTYMF9w==} + '@types/readable-stream@4.0.14': resolution: {integrity: sha512-xZn/AuUbCMShGsqH/ehZtGDwQtbx00M9rZ2ENLe4tOjFZ/JFeWMhEZkk2fEe1jAUqqEAURIkFJ7Az/go8mM1/w==} @@ -12683,6 +12689,9 @@ packages: csstype@3.2.3: resolution: {integrity: sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==} + csstype@3.2.3: + resolution: {integrity: sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==} + csv-generate@3.4.3: resolution: {integrity: sha512-w/T+rqR0vwvHqWs/1ZyMDWtHHSJaN06klRqJXBEpDJaM/+dZkso0OKh1VcuuYvK3XM53KysVNq8Ko/epCK8wOw==} @@ -33343,6 +33352,8 @@ snapshots: csstype@3.2.3: {} + csstype@3.2.3: {} + csv-generate@3.4.3: {} csv-parse@4.16.3: {} diff --git a/references/ai-chat/src/components/chat.tsx b/references/ai-chat/src/components/chat.tsx index aa09d530b..9755e15f6 100644 --- a/references/ai-chat/src/components/chat.tsx +++ b/references/ai-chat/src/components/chat.tsx @@ -1,9 +1,10 @@ "use client"; import { useChat } from "@ai-sdk/react"; -import { TriggerChatTransport } from "@trigger.dev/sdk/chat"; -import { useMemo, useState } from "react"; +import { useTriggerChatTransport } from "@trigger.dev/sdk/chat/react"; +import { useState } from "react"; import { getChatToken } from "@/app/actions"; +import type { chat } from "@/trigger/chat"; function ToolInvocation({ part }: { part: any }) { const [expanded, setExpanded] = useState(false); @@ -72,15 +73,11 @@ function ToolInvocation({ part }: { part: any }) { export function Chat() { const [input, setInput] = useState(""); - const transport = useMemo( - () => - new TriggerChatTransport({ - task: "ai-chat", - accessToken: getChatToken, - baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL, - }), - [] - ); + const transport = useTriggerChatTransport({ + task: "ai-chat", + accessToken: getChatToken, + baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL, + }); const { messages, sendMessage, status, error } = useChat({ transport,