diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 89b758e34..a209c12ea 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -317,6 +317,97 @@ export { CHAT_MESSAGES_STREAM_ID, CHAT_STOP_STREAM_ID }; */ const chatStream = streams.define({ id: _CHAT_STREAM_KEY }); +// --------------------------------------------------------------------------- +// ChatWriter — stream writer for callbacks +// --------------------------------------------------------------------------- + +/** + * A stream writer passed to chat lifecycle callbacks (`onPreload`, `onChatStart`, + * `onTurnStart`, `onTurnComplete`, `onCompacted`). + * + * Write custom `UIMessageChunk` parts (e.g. `data-*` parts) directly to the chat + * stream without the ceremony of `chat.stream.writer({ execute })`. + * + * The writer is lazy — no stream overhead if you don't call `write()` or `merge()`. + * + * @example + * ```ts + * onTurnStart: async ({ writer }) => { + * writer.write({ type: "data-status", data: { loading: true } }); + * }, + * onTurnComplete: async ({ writer, uiMessages }) => { + * writer.write({ type: "data-analytics", data: { messageCount: uiMessages.length } }); + * }, + * ``` + */ +export type ChatWriter = { + /** Write a single UIMessageChunk to the chat stream. */ + write(part: UIMessageChunk): void; + /** Merge another stream's chunks into the chat stream. */ + merge(stream: ReadableStream): void; +}; + +/** + * Creates a lazy ChatWriter that only opens a realtime stream on first use. + * Call `flush()` after the callback returns to await stream completion. + * @internal + */ +function createLazyChatWriter(): { writer: ChatWriter; flush: () => Promise } { + let writeImpl: ((part: UIMessageChunk) => void) | null = null; + let mergeImpl: ((stream: ReadableStream) => void) | null = null; + let waitPromise: (() => Promise) | null = null; + let resolveExecute: (() => void) | null = null; + + function ensureInitialized() { + if (writeImpl) return; + + const executePromise = new Promise((resolve) => { + resolveExecute = resolve; + }); + + const { waitUntilComplete } = chatStream.writer({ + collapsed: true, + spanName: "callback writer", + execute: ({ write, merge }) => { + writeImpl = write; + mergeImpl = merge; + return executePromise; // Keep execute alive until flush() + }, + }); + waitPromise = waitUntilComplete; + } + + return { + writer: { + write(part: UIMessageChunk) { + ensureInitialized(); + writeImpl!(part); + }, + merge(stream: ReadableStream) { + ensureInitialized(); + mergeImpl!(stream); + }, + }, + async flush() { + if (resolveExecute) { + resolveExecute(); // Signal execute to complete + await waitPromise!(); // Wait for stream to finish piping + } + }, + }; +} + +/** + * Runs a callback with a lazy ChatWriter, flushing the stream after completion. + * @internal + */ +async function withChatWriter(fn: (writer: ChatWriter) => Promise | T): Promise { + const { writer, flush } = createLazyChatWriter(); + const result = await fn(writer); + await flush(); + return result; +} + /** * The wire payload shape sent by `TriggerChatTransport`. * Uses `metadata` to match the AI SDK's `ChatRequestOptions` field name. @@ -748,6 +839,8 @@ export type CompactedEvent = { chatId?: string; /** The current turn number (if running inside a chat.task). */ turn?: number; + /** Stream writer — write custom `UIMessageChunk` parts to the chat stream. Lazy: no overhead if unused. */ + writer: ChatWriter; }; /** @@ -988,7 +1081,7 @@ async function chatCompact( const { waitUntilComplete } = streams.writer(CHAT_STREAM_KEY, { spanName: "stream compaction chunks", collapsed: true, - execute: async ({ write }) => { + execute: async ({ write, merge }) => { write({ type: "step-start" }); write({ type: "data-compaction", @@ -1012,7 +1105,8 @@ async function chatCompact( { role: "assistant" as const, content: [{ type: "text" as const, text: `[Conversation summary]\n\n${summary}` }] }, ]); - // Fire onCompacted hook + // Fire onCompacted hook — pass the existing writer so the callback + // can write custom chunks without creating a separate stream. const onCompactedHook = locals.get(chatOnCompactedKey); if (onCompactedHook) { await onCompactedHook({ @@ -1026,6 +1120,7 @@ async function chatCompact( stepNumber, chatId: turnCtx?.chatId, turn: turnCtx?.turn, + writer: { write, merge }, }); } @@ -1576,6 +1671,8 @@ export type PreloadEvent = { chatAccessToken: string; /** Custom data from the frontend. */ clientData?: TClientData; + /** Stream writer — write custom `UIMessageChunk` parts to the chat stream. Lazy: no overhead if unused. */ + writer: ChatWriter; }; /** @@ -1598,6 +1695,8 @@ export type ChatStartEvent = { previousRunId?: string; /** Whether this run was preloaded before the first message. */ preloaded: boolean; + /** Stream writer — write custom `UIMessageChunk` parts to the chat stream. Lazy: no overhead if unused. */ + writer: ChatWriter; }; /** @@ -1628,6 +1727,8 @@ export type TurnStartEvent = { previousTurnUsage?: LanguageModelUsage; /** Cumulative token usage across all completed turns so far. */ totalUsage: LanguageModelUsage; + /** Stream writer — write custom `UIMessageChunk` parts to the chat stream. Lazy: no overhead if unused. */ + writer: ChatWriter; }; /** @@ -1685,6 +1786,15 @@ export type TurnCompleteEvent = { totalUsage: LanguageModelUsage; }; +/** + * Event passed to the `onBeforeTurnComplete` callback. + * Same as `TurnCompleteEvent` but includes a `writer` since the stream is still open. + */ +export type BeforeTurnCompleteEvent = TurnCompleteEvent & { + /** Stream writer — write custom `UIMessageChunk` parts to the chat stream. Lazy: no overhead if unused. */ + writer: ChatWriter; +}; + export type ChatTaskOptions< TIdentifier extends string, TClientDataSchema extends TaskSchema | undefined = undefined, @@ -1774,17 +1884,17 @@ export type ChatTaskOptions< * * @example * ```ts - * onBeforeTurnComplete: async ({ messages, uiMessages, usage }) => { + * onBeforeTurnComplete: async ({ writer, usage }) => { * if (usage?.inputTokens && usage.inputTokens > 5000) { - * await chat.stream.append({ type: "data-compaction", id: generateId(), data: { status: "compacting" } }); + * writer.write({ type: "data-compaction", id: generateId(), data: { status: "compacting" } }); * // ... compact messages ... * chat.setMessages(compactedMessages); - * await chat.stream.append({ type: "data-compaction", id: generateId(), data: { status: "complete" } }); + * writer.write({ type: "data-compaction", id: generateId(), data: { status: "complete" } }); * } * } * ``` */ - onBeforeTurnComplete?: (event: TurnCompleteEvent>) => Promise | void; + onBeforeTurnComplete?: (event: BeforeTurnCompleteEvent>) => Promise | void; /** * Called when conversation compaction occurs (via `chat.compact()` or @@ -2120,11 +2230,14 @@ function chatTask< await tracer.startActiveSpan( "onPreload()", async () => { - await onPreload({ - chatId: payload.chatId, - runId: currentRunId, - chatAccessToken: preloadAccessToken, - clientData: preloadClientData, + await withChatWriter(async (writer) => { + await onPreload({ + chatId: payload.chatId, + runId: currentRunId, + chatAccessToken: preloadAccessToken, + clientData: preloadClientData, + writer, + }); }); }, { @@ -2333,15 +2446,18 @@ function chatTask< await tracer.startActiveSpan( "onChatStart()", async () => { - await onChatStart({ - chatId: currentWirePayload.chatId, - messages: accumulatedMessages, - clientData, - runId: currentRunId, - chatAccessToken: turnAccessToken, - continuation, - previousRunId, - preloaded, + await withChatWriter(async (writer) => { + await onChatStart({ + chatId: currentWirePayload.chatId, + messages: accumulatedMessages, + clientData, + runId: currentRunId, + chatAccessToken: turnAccessToken, + continuation, + previousRunId, + preloaded, + writer, + }); }); }, { @@ -2364,19 +2480,22 @@ function chatTask< await tracer.startActiveSpan( "onTurnStart()", async () => { - await onTurnStart({ - chatId: currentWirePayload.chatId, - messages: accumulatedMessages, - uiMessages: accumulatedUIMessages, - turn, - runId: currentRunId, - chatAccessToken: turnAccessToken, - clientData, - continuation, - previousRunId, - preloaded, - previousTurnUsage, - totalUsage: cumulativeUsage, + await withChatWriter(async (writer) => { + await onTurnStart({ + chatId: currentWirePayload.chatId, + messages: accumulatedMessages, + uiMessages: accumulatedUIMessages, + turn, + runId: currentRunId, + chatAccessToken: turnAccessToken, + clientData, + continuation, + previousRunId, + preloaded, + previousTurnUsage, + totalUsage: cumulativeUsage, + writer, + }); }); // Check if onTurnStart replaced messages (compaction) @@ -2617,7 +2736,7 @@ function chatTask< const { waitUntilComplete } = streams.writer(CHAT_STREAM_KEY, { spanName: "stream compaction chunks", collapsed: true, - execute: async ({ write }) => { + execute: async ({ write, merge }) => { write({ type: "data-compaction", id: compactionId, @@ -2670,6 +2789,7 @@ function chatTask< stepNumber: -1, // outer loop, not a step chatId: currentWirePayload.chatId, turn, + writer: { write, merge }, }); } @@ -2733,7 +2853,9 @@ function chatTask< await tracer.startActiveSpan( "onBeforeTurnComplete()", async () => { - await onBeforeTurnComplete(turnCompleteEvent); + await withChatWriter(async (writer) => { + await onBeforeTurnComplete({ ...turnCompleteEvent, writer }); + }); // Check if the hook replaced messages (compaction) const override = locals.get(chatOverrideMessagesKey); diff --git a/references/ai-chat/src/trigger/chat.ts b/references/ai-chat/src/trigger/chat.ts index be385c6be..898802c1f 100644 --- a/references/ai-chat/src/trigger/chat.ts +++ b/references/ai-chat/src/trigger/chat.ts @@ -424,7 +424,9 @@ export const aiChat = chat.task({ summaryLength: summary.length, }); }, - onTurnStart: async ({ chatId, uiMessages }) => { + onTurnStart: async ({ chatId, uiMessages, writer }) => { + writer.write({ type: "data-turn-status", data: { status: "preparing" } }); + // Persist messages so mid-stream refresh still shows the user message. // Deferred — runs in parallel with streaming, awaited before onTurnComplete. chat.defer(prisma.chat.update({ where: { id: chatId }, data: { messages: uiMessages as any } }));