expose a useTriggerChatTransport hook

2026-02-21 15:18:22 +00:00
parent b8d0483729
commit bf229d6faa
6 changed files with 384 additions and 14 deletions
@@ -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<string, unknown> | 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`
@@ -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",
@@ -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<typeof chat>({
+ *     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<typeof myTask>({ task: "my-task", ... })
+ * ```
+ */
+export type UseTriggerChatTransportOptions<TTask extends AnyTask = AnyTask> = Omit<
+  TriggerChatTransportOptions,
+  "task"
+> & {
+  /** The task ID. Strongly typed when a task type parameter is provided. */
+  task: TaskIdentifier<TTask>;
+};
+
+/**
+ * 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<typeof chat>({
+ *     task: "ai-chat",
+ *     accessToken: () => fetchToken(),
+ *   });
+ *
+ *   const { messages, sendMessage } = useChat({ transport });
+ * }
+ * ```
+ */
+export function useTriggerChatTransport<TTask extends AnyTask = AnyTask>(
+  options: UseTriggerChatTransportOptions<TTask>
+): TriggerChatTransport {
+  const ref = useRef<TriggerChatTransport | null>(null);
+  if (ref.current === null) {
+    ref.current = new TriggerChatTransport(options);
+  }
+  return ref.current;
+}
@@ -391,3 +391,4 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
 export function createChatTransport(options: TriggerChatTransportOptions): TriggerChatTransport {
  return new TriggerChatTransport(options);
 }
+
@@ -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: {}
@@ -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<typeof chat>({
+    task: "ai-chat",
+    accessToken: getChatToken,
+    baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL,
+  });

  const { messages, sendMessage, status, error } = useChat({
    transport,