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

Add a writer to easily write chunks in callbacks

Browse Source
This commit is contained in:
Eric Allam
2026-03-26 08:21:12 +00:00
parent 843b85e356
commit a503e99a57
2 changed files with 160 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/trigger-sdk/src/v3/ai.ts
+157 -35
View File
@@ -317,6 +317,97 @@ export { CHAT_MESSAGES_STREAM_ID, CHAT_STOP_STREAM_ID };
*/
const chatStream = streams.define<UIMessageChunk>({ 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<UIMessageChunk>): 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<void> } {
let writeImpl: ((part: UIMessageChunk) => void) | null = null;
let mergeImpl: ((stream: ReadableStream<UIMessageChunk>) => void) | null = null;
let waitPromise: (() => Promise<unknown>) | null = null;
let resolveExecute: (() => void) | null = null;
function ensureInitialized() {
if (writeImpl) return;
const executePromise = new Promise<void>((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<UIMessageChunk>) {
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<T>(fn: (writer: ChatWriter) => Promise<T> | T): Promise<T> {
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<TClientData = unknown> = {
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<TClientData = unknown> = {
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<TClientData = unknown> = {
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<TClientData = unknown> = {
totalUsage: LanguageModelUsage;
};
/**
* Event passed to the `onBeforeTurnComplete` callback.
* Same as `TurnCompleteEvent` but includes a `writer` since the stream is still open.
*/
export type BeforeTurnCompleteEvent<TClientData = unknown> = TurnCompleteEvent<TClientData> & {
/** 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<inferSchemaOut<TClientDataSchema>>) => Promise<void> | void;
onBeforeTurnComplete?: (event: BeforeTurnCompleteEvent<inferSchemaOut<TClientDataSchema>>) => Promise<void> | 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);
references/ai-chat/src/trigger/chat.ts
+3 -1
View File
@@ -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 } }));