github-featured/triggerdotdev--trigger.dev
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Add support for toUIMessageStream() options

Browse Source
This commit is contained in:
Eric Allam
2026-03-10 14:54:34 +00:00
parent a458120c59
commit 627db05c56
3 changed files with 122 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/trigger-sdk/src/v3/ai.ts
+103 -1
View File
@@ -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}. */
pnpm-lock.yaml
Generated
+6 -6
View File
@@ -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
references/ai-chat/src/trigger/chat.ts
+13 -1
View File
@@ -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