From 627db05c567ed2511b8aec18445f1d9b7d408271 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Tue, 10 Mar 2026 14:54:34 +0000 Subject: [PATCH] Add support for toUIMessageStream() options --- packages/trigger-sdk/src/v3/ai.ts | 104 ++++++++++++++++++++++++- pnpm-lock.yaml | 12 +-- references/ai-chat/src/trigger/chat.ts | 14 +++- 3 files changed, 122 insertions(+), 8 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 46237d9cd..a54eaa73e 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/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>>("chat.defer"); */ const chatPipeCountKey = locals.create("chat.pipeCount"); const chatStopControllerKey = locals.create("chat.stopController"); +/** Static (task-level) UIMessageStream options, set once during chatTask setup. @internal */ +const chatUIStreamStaticKey = locals.create("chat.uiMessageStreamOptions.static"); +/** Per-turn UIMessageStream options, set via chat.setUIMessageStreamOptions(). @internal */ +const chatUIStreamPerTurnKey = locals.create("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, + "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((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}. */ diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 6897b2fe8..e3c5fb7a9 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1138,7 +1138,7 @@ importers: version: 18.3.1 react-email: specifier: ^2.1.1 - version: 2.1.2(@opentelemetry/api@1.9.0)(@swc/helpers@0.5.15)(eslint@8.31.0) + version: 2.1.2(@opentelemetry/api@1.9.0)(@swc/helpers@0.5.15)(bufferutil@4.0.9)(eslint@8.31.0) resend: specifier: ^3.2.0 version: 3.2.0 @@ -39919,7 +39919,7 @@ snapshots: react: 18.2.0 react-dom: 18.2.0(react@18.2.0) - react-email@2.1.2(@opentelemetry/api@1.9.0)(@swc/helpers@0.5.15)(eslint@8.31.0): + react-email@2.1.2(@opentelemetry/api@1.9.0)(@swc/helpers@0.5.15)(bufferutil@4.0.9)(eslint@8.31.0): dependencies: '@babel/parser': 7.24.1 '@radix-ui/colors': 1.0.1 @@ -39956,8 +39956,8 @@ snapshots: react: 18.3.1 react-dom: 18.2.0(react@18.3.1) shelljs: 0.8.5 - socket.io: 4.7.3 - socket.io-client: 4.7.3 + socket.io: 4.7.3(bufferutil@4.0.9) + socket.io-client: 4.7.3(bufferutil@4.0.9) sonner: 1.3.1(react-dom@18.2.0(react@18.3.1))(react@18.3.1) source-map-js: 1.0.2 stacktrace-parser: 0.1.10 @@ -41207,7 +41207,7 @@ snapshots: - supports-color - utf-8-validate - socket.io-client@4.7.3: + socket.io-client@4.7.3(bufferutil@4.0.9): dependencies: '@socket.io/component-emitter': 3.1.0 debug: 4.3.7(supports-color@10.0.0) @@ -41236,7 +41236,7 @@ snapshots: transitivePeerDependencies: - supports-color - socket.io@4.7.3: + socket.io@4.7.3(bufferutil@4.0.9): dependencies: accepts: 1.3.8 base64id: 2.0.0 diff --git a/references/ai-chat/src/trigger/chat.ts b/references/ai-chat/src/trigger/chat.ts index 3cb32a16a..ed9304282 100644 --- a/references/ai-chat/src/trigger/chat.ts +++ b/references/ai-chat/src/trigger/chat.ts @@ -1,5 +1,5 @@ import { chat, ai, type ChatTaskWirePayload } from "@trigger.dev/sdk/ai"; -import { schemaTask, task } from "@trigger.dev/sdk"; +import { logger, schemaTask, task } from "@trigger.dev/sdk"; import { streamText, tool, dynamicTool, stepCountIs, generateId } from "ai"; import type { LanguageModel, Tool as AITool, UIMessage } from "ai"; import { openai } from "@ai-sdk/openai"; @@ -231,6 +231,18 @@ export const aiChat = chat.task({ clientDataSchema: z.object({ model: z.string().optional(), userId: z.string() }), warmTimeoutInSeconds: 60, chatAccessTokenTTL: "2h", + uiMessageStreamOptions: { + sendReasoning: true, + onError: (error) => { + // Log the full error server-side for debugging + logger.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."; + }, + }, onPreload: async ({ chatId, runId, chatAccessToken, clientData }) => { if (!clientData) return; // Eagerly initialize before the user's first message arrives