Add support for toUIMessageStream() options

Author: ericallam

docs/ai-chat/backend.mdx

@@ -504,6 +504,89 @@ run: async ({ messages, signal }) => {
   Longer warm timeout means faster responses but more compute usage. Set to `0` to suspend immediately after each turn (minimum latency cost, slight delay on next message).
 </Info>
 
+#### Stream options
+
+Control how `streamText` results are converted to the frontend stream via `toUIMessageStream()`. Set static defaults on the task, or override per-turn.
+
+##### Error handling with onError
+
+When `streamText` encounters an error mid-stream (rate limits, API failures, network errors), the `onError` callback converts it to a string that's sent to the frontend as an `{ type: "error", errorText }` chunk. The AI SDK's `useChat` receives this via its `onError` callback.
+
+By default, the raw error message is sent to the frontend. Use `onError` to sanitize errors and avoid leaking internal details:
+
+```ts
+export const myChat = chat.task({
+  id: "my-chat",
+  uiMessageStreamOptions: {
+    onError: (error) => {
+      // Log the full error server-side for debugging
+      console.error("Stream error:", error);
+      // Return a sanitized message — this is what the frontend sees
+      if (error instanceof Error && error.message.includes("rate limit")) {
+        return "Rate limited — please wait a moment and try again.";
+      }
+      return "Something went wrong. Please try again.";
+    },
+  },
+  run: async ({ messages, signal }) => {
+    return streamText({ model: openai("gpt-4o"), messages, abortSignal: signal });
+  },
+});
+```
+
+`onError` is also called for tool execution errors, so a single handler covers both LLM errors and tool failures.
+
+On the frontend, handle the error in `useChat`:
+
+```tsx
+const { messages, sendMessage } = useChat({
+  transport,
+  onError: (error) => {
+    // error.message contains the string returned by your onError handler
+    toast.error(error.message);
+  },
+});
+```
+
+##### Reasoning and sources
+
+Control which AI SDK features are forwarded to the frontend:
+
+```ts
+export const myChat = chat.task({
+  id: "my-chat",
+  uiMessageStreamOptions: {
+    sendReasoning: true,  // Forward model reasoning (default: true)
+    sendSources: true,    // Forward source citations (default: false)
+  },
+  run: async ({ messages, signal }) => {
+    return streamText({ model: openai("gpt-4o"), messages, abortSignal: signal });
+  },
+});
+```
+
+##### Per-turn overrides
+
+Override per-turn with `chat.setUIMessageStreamOptions()` — per-turn values merge with the static config (per-turn wins on conflicts). The override is cleared automatically after each turn.
+
+```ts
+run: async ({ messages, clientData, signal }) => {
+  // Enable reasoning only for certain models
+  if (clientData.model?.includes("claude")) {
+    chat.setUIMessageStreamOptions({ sendReasoning: true });
+  }
+  return streamText({ model: openai(clientData.model ?? "gpt-4o"), messages, abortSignal: signal });
+},
+```
+
+`chat.setUIMessageStreamOptions()` works across all abstraction levels — `chat.task()`, `chat.createSession()` / `turn.complete()`, and `chat.pipeAndCapture()`.
+
+See [ChatUIMessageStreamOptions](/ai-chat/reference#chatuimessagestreamoptions) for the full reference.
+
+<Note>
+  `onFinish` is managed internally for response capture and cannot be overridden here. Use `streamText`'s `onFinish` callback for custom finish handling, or use [raw task mode](#raw-task-with-primitives) for full control over `toUIMessageStream()`.
+</Note>
+
 ### Manual mode with task()
 
 If you need full control over task options, use the standard `task()` with `ChatTaskPayload` and `chat.pipe()`:
@@ -647,6 +730,8 @@ for await (const turn of session) {
 
 For full control, use a standard `task()` with the composable primitives from the `chat` namespace. You manage everything: the turn loop, stop signals, message accumulation, and turn-complete signaling.
 
+Raw task mode also lets you call `.toUIMessageStream()` yourself with any options — including `onFinish` and `originalMessages`. This is the right choice when you need complete control over the stream conversion beyond what `chat.setUIMessageStreamOptions()` provides.
+
 ### Primitives
 
 | Primitive | Description |

docs/ai-chat/reference.mdx

@@ -23,6 +23,7 @@ Options for `chat.task()`.
 | `chatAccessTokenTTL` | `string` | `"1h"` | How long the scoped access token remains valid |
 | `preloadWarmTimeoutInSeconds` | `number` | Same as `warmTimeoutInSeconds` | Warm timeout after `onPreload` fires |
 | `preloadTimeout` | `string` | Same as `turnTimeout` | Suspend timeout for preloaded runs |
+| `uiMessageStreamOptions` | `ChatUIMessageStreamOptions` | — | Default options for `toUIMessageStream()`. Per-turn override via `chat.setUIMessageStreamOptions()` |
 
 Plus all standard [TaskOptions](/tasks/overview) — `retry`, `queue`, `machine`, `maxDuration`, etc.
 
@@ -156,12 +157,28 @@ All methods available on the `chat` object from `@trigger.dev/sdk/ai`.
 | `chat.setTurnTimeout(duration)` | Override turn timeout at runtime (e.g. `"2h"`) |
 | `chat.setTurnTimeoutInSeconds(seconds)` | Override turn timeout at runtime (in seconds) |
 | `chat.setWarmTimeoutInSeconds(seconds)` | Override warm timeout at runtime |
+| `chat.setUIMessageStreamOptions(options)` | Override `toUIMessageStream()` options for the current turn |
 | `chat.defer(promise)` | Run background work in parallel with streaming, awaited before `onTurnComplete` |
 | `chat.isStopped()` | Check if the current turn was stopped by the user |
 | `chat.cleanupAbortedParts(message)` | Remove incomplete parts from a stopped response message |
 | `chat.stream` | Typed chat output stream — use `.writer()`, `.pipe()`, `.append()`, `.read()` |
 | `chat.MessageAccumulator` | Class that accumulates conversation messages across turns |
 
+## ChatUIMessageStreamOptions
+
+Options for customizing `toUIMessageStream()`. Set as static defaults via `uiMessageStreamOptions` on `chat.task()`, or override per-turn via `chat.setUIMessageStreamOptions()`. See [Stream options](/ai-chat/backend#stream-options) for usage examples.
+
+Derived from the AI SDK's `UIMessageStreamOptions` with `onFinish`, `originalMessages`, and `generateMessageId` omitted (managed internally).
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `onError` | `(error: unknown) => string` | Raw error message | Called on LLM errors and tool execution errors. Return a sanitized string — sent as `{ type: "error", errorText }` to the frontend. |
+| `sendReasoning` | `boolean` | `true` | Send reasoning parts to the client |
+| `sendSources` | `boolean` | `false` | Send source parts to the client |
+| `sendFinish` | `boolean` | `true` | Send the finish event. Set to `false` when chaining multiple `streamText` calls. |
+| `sendStart` | `boolean` | `true` | Send the message start event. Set to `false` when chaining. |
+| `messageMetadata` | `(options: { part }) => metadata` | — | Extract message metadata to send to the client. Called on `start` and `finish` events. |
+
 ## TriggerChatTransport options
 
 Options for the frontend transport constructor and `useTriggerChatTransport` hook.

packages/trigger-sdk/src/v3/ai.ts

@@ -14,7 +14,7 @@ import {
   type TaskSchema,
   type TaskWithSchema,
 } from "@trigger.dev/core/v3";
-import type { ModelMessage, UIMessage, UIMessageChunk } from "ai";
+import type { ModelMessage, UIMessage, UIMessageChunk, UIMessageStreamOptions } from "ai";
 import type { StreamWriteResult } from "@trigger.dev/core/v3";
 import { convertToModelMessages, dynamicTool, generateId as generateMessageId, jsonSchema, JSONSchema7, Schema, Tool, ToolCallOptions, zodSchema } from "ai";
 import { type Attributes, trace } from "@opentelemetry/api";
@@ -399,6 +399,10 @@ const chatDeferKey = locals.create<Set<Promise<unknown>>>("chat.defer");
  */
 const chatPipeCountKey = locals.create<number>("chat.pipeCount");
 const chatStopControllerKey = locals.create<AbortController>("chat.stopController");
+/** Static (task-level) UIMessageStream options, set once during chatTask setup. @internal */
+const chatUIStreamStaticKey = locals.create<ChatUIMessageStreamOptions>("chat.uiMessageStreamOptions.static");
+/** Per-turn UIMessageStream options, set via chat.setUIMessageStreamOptions(). @internal */
+const chatUIStreamPerTurnKey = locals.create<ChatUIMessageStreamOptions>("chat.uiMessageStreamOptions.perTurn");
 
 /**
  * Options for `pipeChat`.
@@ -423,6 +427,23 @@ export type PipeChatOptions = {
   spanName?: string;
 };
 
+/**
+ * Options for customizing the `toUIMessageStream()` call used when piping
+ * `streamText` results to the frontend.
+ *
+ * Set static defaults via `uiMessageStreamOptions` on `chat.task()`, or
+ * override per-turn via `chat.setUIMessageStreamOptions()`.
+ *
+ * `onFinish`, `originalMessages`, and `generateMessageId` are omitted because
+ * they are managed internally for response capture and message accumulation.
+ * Use `streamText`'s `onFinish` for custom finish handling, or drop down to
+ * raw task mode with `chat.pipe()` for full control.
+ */
+export type ChatUIMessageStreamOptions = Omit<
+  UIMessageStreamOptions<UIMessage>,
+  "onFinish" | "originalMessages" | "generateMessageId"
+>;
+
 /**
  * An object with a `toUIMessageStream()` method (e.g. `StreamTextResult` from `streamText()`).
  */
@@ -803,6 +824,35 @@ export type ChatTaskOptions<
    * @default Same as `turnTimeout`
    */
   preloadTimeout?: string;
+
+  /**
+   * Default options for `toUIMessageStream()` when auto-piping or using
+   * `turn.complete()` / `chat.pipeAndCapture()`.
+   *
+   * Controls how the `StreamTextResult` is converted to a `UIMessageChunk`
+   * stream — error handling, reasoning/source visibility, metadata, etc.
+   *
+   * Can be overridden per-turn by calling `chat.setUIMessageStreamOptions()`
+   * inside `run()` or lifecycle hooks. Per-turn values are merged on top
+   * of these defaults (per-turn wins on conflicts).
+   *
+   * `onFinish`, `originalMessages`, and `generateMessageId` are managed
+   * internally and cannot be overridden here. Use `streamText`'s `onFinish`
+   * for custom finish handling, or drop to raw task mode for full control.
+   *
+   * @example
+   * ```ts
+   * chat.task({
+   *   id: "my-chat",
+   *   uiMessageStreamOptions: {
+   *     sendReasoning: true,
+   *     onError: (error) => error instanceof Error ? error.message : "An error occurred.",
+   *   },
+   *   run: async ({ messages, signal }) => { ... },
+   * });
+   * ```
+   */
+  uiMessageStreamOptions?: ChatUIMessageStreamOptions;
 };
 
 /**
@@ -851,6 +901,7 @@ function chatTask<
     chatAccessTokenTTL = "1h",
     preloadWarmTimeoutInSeconds,
     preloadTimeout,
+    uiMessageStreamOptions,
     ...restOptions
   } = options;
 
@@ -867,6 +918,11 @@ function chatTask<
         activeSpan.setAttribute("gen_ai.conversation.id", payload.chatId);
       }
 
+      // Store static UIMessageStream options in locals so resolveUIMessageStreamOptions() can read them
+      if (uiMessageStreamOptions) {
+        locals.set(chatUIStreamStaticKey, uiMessageStreamOptions);
+      }
+
       let currentWirePayload = payload;
       const continuation = payload.continuation ?? false;
       const previousRunId = payload.previousRunId;
@@ -1192,6 +1248,7 @@ function chatTask<
                 if ((locals.get(chatPipeCountKey) ?? 0) === 0 && isUIMessageStreamable(result)) {
                   onFinishAttached = true;
                   const uiStream = result.toUIMessageStream({
+                    ...resolveUIMessageStreamOptions(),
                     onFinish: ({ responseMessage }: { responseMessage: UIMessage }) => {
                       capturedResponseMessage = responseMessage;
                       resolveOnFinish!();
@@ -1447,6 +1504,48 @@ function setWarmTimeoutInSeconds(seconds: number): void {
   metadata.set(WARM_TIMEOUT_METADATA_KEY, seconds);
 }
 
+/**
+ * Override the `toUIMessageStream()` options for the current turn.
+ *
+ * These options control how the `StreamTextResult` is converted to a
+ * `UIMessageChunk` stream — error handling, reasoning/source visibility,
+ * message metadata, etc.
+ *
+ * Per-turn options are merged on top of the static `uiMessageStreamOptions`
+ * set on `chat.task()`. Per-turn values win on conflicts.
+ *
+ * @example
+ * ```ts
+ * run: async ({ messages, signal }) => {
+ *   chat.setUIMessageStreamOptions({
+ *     sendReasoning: true,
+ *     onError: (error) => error instanceof Error ? error.message : "An error occurred.",
+ *   });
+ *   return streamText({ model, messages, abortSignal: signal });
+ * }
+ * ```
+ */
+function setUIMessageStreamOptions(options: ChatUIMessageStreamOptions): void {
+  locals.set(chatUIStreamPerTurnKey, options);
+}
+
+/**
+ * Resolve the effective UIMessageStream options by merging:
+ * 1. Static task-level options (from `chat.task({ uiMessageStreamOptions })`)
+ * 2. Per-turn overrides (from `chat.setUIMessageStreamOptions()`)
+ *
+ * Per-turn values win on conflicts. Clears the per-turn override after reading
+ * so it doesn't leak into subsequent turns.
+ * @internal
+ */
+function resolveUIMessageStreamOptions(): ChatUIMessageStreamOptions {
+  const staticOptions = locals.get(chatUIStreamStaticKey) ?? {};
+  const perTurnOptions = locals.get(chatUIStreamPerTurnKey) ?? {};
+  // Clear per-turn override so it doesn't leak into subsequent turns
+  locals.set(chatUIStreamPerTurnKey, undefined);
+  return { ...staticOptions, ...perTurnOptions };
+}
+
 // ---------------------------------------------------------------------------
 // Stop detection
 // ---------------------------------------------------------------------------
@@ -1641,6 +1740,7 @@ async function pipeChatAndCapture(
   const onFinishPromise = new Promise<void>((r) => { resolveOnFinish = r; });
 
   const uiStream = source.toUIMessageStream({
+    ...resolveUIMessageStreamOptions(),
     onFinish: ({ responseMessage }: { responseMessage: UIMessage }) => {
       captured = responseMessage;
       resolveOnFinish!();
@@ -2180,6 +2280,8 @@ export const chat = {
   setTurnTimeoutInSeconds,
   /** Override the warm timeout at runtime. See {@link setWarmTimeoutInSeconds}. */
   setWarmTimeoutInSeconds,
+  /** Override toUIMessageStream() options for the current turn. See {@link setUIMessageStreamOptions}. */
+  setUIMessageStreamOptions,
   /** Check if the current turn was stopped by the user. See {@link isStopped}. */
   isStopped,
   /** Clean up aborted parts from a UIMessage. See {@link cleanupAbortedParts}. */