Make clientData typesafe and pass to all chat.task hooks

This commit is contained in:
Eric Allam
2026-03-06 16:19:25 +00:00
parent 4a1a0f1701
commit 610067f29a
6 changed files with 104 additions and 38 deletions
+1
View File
@@ -80,6 +80,7 @@ export {
getSchemaParseFn,
type AnySchemaParseFn,
type SchemaParseFn,
type inferSchemaOut,
isSchemaZodEsque,
isSchemaValibotEsque,
isSchemaArkTypeEsque,
+84 -24
View File
@@ -1,11 +1,13 @@
import {
accessoryAttributes,
AnyTask,
getSchemaParseFn,
isSchemaZodEsque,
SemanticInternalAttributes,
Task,
taskContext,
type inferSchemaIn,
type inferSchemaOut,
type PipeStreamOptions,
type TaskIdentifier,
type TaskOptions,
@@ -178,12 +180,12 @@ export { CHAT_MESSAGES_STREAM_ID, CHAT_STOP_STREAM_ID };
* Uses `metadata` to match the AI SDK's `ChatRequestOptions` field name.
* @internal
*/
type ChatTaskWirePayload<TMessage extends UIMessage = UIMessage> = {
type ChatTaskWirePayload<TMessage extends UIMessage = UIMessage, TMetadata = unknown> = {
messages: TMessage[];
chatId: string;
trigger: "submit-message" | "regenerate-message";
messageId?: string;
metadata?: unknown;
metadata?: TMetadata;
};
/**
@@ -196,7 +198,7 @@ type ChatTaskWirePayload<TMessage extends UIMessage = UIMessage> = {
* The backend accumulates the full conversation history across turns, so the frontend
* only needs to send new messages after the first turn.
*/
export type ChatTaskPayload = {
export type ChatTaskPayload<TClientData = unknown> = {
/** Model-ready messages — pass directly to `streamText({ messages })`. */
messages: ModelMessage[];
@@ -214,7 +216,7 @@ export type ChatTaskPayload = {
messageId?: string;
/** Custom data from the frontend (passed via `metadata` on `sendMessage()` or the transport). */
clientData?: unknown;
clientData?: TClientData;
};
/**
@@ -233,7 +235,7 @@ export type ChatTaskSignals = {
* The full payload passed to a `chatTask` run function.
* Extends `ChatTaskPayload` (the wire payload) with abort signals.
*/
export type ChatTaskRunPayload = ChatTaskPayload & ChatTaskSignals;
export type ChatTaskRunPayload<TClientData = unknown> = ChatTaskPayload<TClientData> & ChatTaskSignals;
// Input streams for bidirectional chat communication
const messagesInput = streams.input<ChatTaskWirePayload>({ id: CHAT_MESSAGES_STREAM_ID });
@@ -384,13 +386,13 @@ async function pipeChat(
/**
* Event passed to the `onChatStart` callback.
*/
export type ChatStartEvent = {
export type ChatStartEvent<TClientData = unknown> = {
/** The unique identifier for the chat session. */
chatId: string;
/** The initial model-ready messages for this conversation. */
messages: ModelMessage[];
/** Custom data from the frontend (passed via `metadata` on `sendMessage()` or the transport). */
clientData: unknown;
clientData: TClientData;
/** The Trigger.dev run ID for this conversation. */
runId: string;
/** A scoped access token for this chat run. Persist this for frontend reconnection. */
@@ -400,7 +402,7 @@ export type ChatStartEvent = {
/**
* Event passed to the `onTurnStart` callback.
*/
export type TurnStartEvent = {
export type TurnStartEvent<TClientData = unknown> = {
/** The unique identifier for the chat session. */
chatId: string;
/** The accumulated model-ready messages (all turns so far, including new user message). */
@@ -413,12 +415,14 @@ export type TurnStartEvent = {
runId: string;
/** A scoped access token for this chat run. */
chatAccessToken: string;
/** Custom data from the frontend. */
clientData?: TClientData;
};
/**
* Event passed to the `onTurnComplete` callback.
*/
export type TurnCompleteEvent = {
export type TurnCompleteEvent<TClientData = unknown> = {
/** The unique identifier for the chat session. */
chatId: string;
/** The full accumulated conversation in model format (all turns so far). */
@@ -448,12 +452,34 @@ export type TurnCompleteEvent = {
chatAccessToken: string;
/** The last event ID from the stream writer. Use this with `resume: true` to avoid replaying events after refresh. */
lastEventId?: string;
/** Custom data from the frontend. */
clientData?: TClientData;
};
export type ChatTaskOptions<TIdentifier extends string> = Omit<
TaskOptions<TIdentifier, ChatTaskWirePayload, unknown>,
"run"
> & {
export type ChatTaskOptions<
TIdentifier extends string,
TClientDataSchema extends TaskSchema | undefined = undefined,
> = Omit<TaskOptions<TIdentifier, ChatTaskWirePayload, unknown>, "run"> & {
/**
* Schema for validating `clientData` from the frontend.
* Accepts Zod, ArkType, Valibot, or any supported schema library.
* When provided, `clientData` is parsed and typed in all hooks and `run`.
*
* @example
* ```ts
* import { z } from "zod";
*
* chat.task({
* id: "my-chat",
* clientDataSchema: z.object({ model: z.string().optional(), userId: z.string() }),
* run: async ({ messages, clientData, signal }) => {
* // clientData is typed as { model?: string; userId: string }
* },
* });
* ```
*/
clientDataSchema?: TClientDataSchema;
/**
* The run function for the chat task.
*
@@ -463,7 +489,7 @@ export type ChatTaskOptions<TIdentifier extends string> = Omit<
* **Auto-piping:** If this function returns a value with `.toUIMessageStream()`,
* the stream is automatically piped to the frontend.
*/
run: (payload: ChatTaskRunPayload) => Promise<unknown>;
run: (payload: ChatTaskRunPayload<inferSchemaOut<TClientDataSchema>>) => Promise<unknown>;
/**
* Called on the first turn (turn 0) of a new run, before the `run` function executes.
@@ -477,7 +503,7 @@ export type ChatTaskOptions<TIdentifier extends string> = Omit<
* }
* ```
*/
onChatStart?: (event: ChatStartEvent) => Promise<void> | void;
onChatStart?: (event: ChatStartEvent<inferSchemaOut<TClientDataSchema>>) => Promise<void> | void;
/**
* Called at the start of every turn, after message accumulation and `onChatStart` (turn 0),
@@ -493,7 +519,7 @@ export type ChatTaskOptions<TIdentifier extends string> = Omit<
* }
* ```
*/
onTurnStart?: (event: TurnStartEvent) => Promise<void> | void;
onTurnStart?: (event: TurnStartEvent<inferSchemaOut<TClientDataSchema>>) => Promise<void> | void;
/**
* Called after each turn completes (after the response is captured, before waiting
@@ -508,7 +534,7 @@ export type ChatTaskOptions<TIdentifier extends string> = Omit<
* }
* ```
*/
onTurnComplete?: (event: TurnCompleteEvent) => Promise<void> | void;
onTurnComplete?: (event: TurnCompleteEvent<inferSchemaOut<TClientDataSchema>>) => Promise<void> | void;
/**
* Maximum number of conversational turns (message round-trips) a single run
@@ -578,11 +604,15 @@ export type ChatTaskOptions<TIdentifier extends string> = Omit<
* });
* ```
*/
function chatTask<TIdentifier extends string>(
options: ChatTaskOptions<TIdentifier>
): Task<TIdentifier, ChatTaskWirePayload, unknown> {
function chatTask<
TIdentifier extends string,
TClientDataSchema extends TaskSchema | undefined = undefined,
>(
options: ChatTaskOptions<TIdentifier, TClientDataSchema>
): Task<TIdentifier, ChatTaskWirePayload<UIMessage, inferSchemaIn<TClientDataSchema>>, unknown> {
const {
run: userRun,
clientDataSchema,
onChatStart,
onTurnStart,
onTurnComplete,
@@ -593,7 +623,11 @@ function chatTask<TIdentifier extends string>(
...restOptions
} = options;
return createTask<TIdentifier, ChatTaskWirePayload, unknown>({
const parseClientData = clientDataSchema
? getSchemaParseFn(clientDataSchema)
: undefined;
return createTask<TIdentifier, ChatTaskWirePayload<UIMessage, inferSchemaIn<TClientDataSchema>>, unknown>({
...restOptions,
run: async (payload: ChatTaskWirePayload, { signal: runSignal }) => {
// Set gen_ai.conversation.id on the run-level span for dashboard context
@@ -626,6 +660,9 @@ function chatTask<TIdentifier extends string>(
for (let turn = 0; turn < maxTurns; turn++) {
// Extract turn-level context before entering the span
const { metadata: wireMetadata, messages: uiMessages, ...restWire } = currentWirePayload;
const clientData = (parseClientData
? await parseClientData(wireMetadata)
: wireMetadata) as inferSchemaOut<TClientDataSchema>;
const lastUserMessage = extractLastUserMessageText(uiMessages);
const turnAttributes: Attributes = {
@@ -738,7 +775,7 @@ function chatTask<TIdentifier extends string>(
await onChatStart({
chatId: currentWirePayload.chatId,
messages: accumulatedMessages,
clientData: wireMetadata,
clientData,
runId: currentRunId,
chatAccessToken: turnAccessToken,
});
@@ -765,6 +802,7 @@ function chatTask<TIdentifier extends string>(
turn,
runId: currentRunId,
chatAccessToken: turnAccessToken,
clientData,
});
},
{
@@ -783,11 +821,11 @@ function chatTask<TIdentifier extends string>(
const result = await userRun({
...restWire,
messages: accumulatedMessages,
clientData: wireMetadata,
clientData,
signal: combinedSignal,
cancelSignal,
stopSignal,
});
} as any);
// Auto-pipe if the run function returned a StreamTextResult or similar,
// but only if pipeChat() wasn't already called manually during this turn.
@@ -866,6 +904,7 @@ function chatTask<TIdentifier extends string>(
runId: currentRunId,
chatAccessToken: turnAccessToken,
lastEventId: turnCompleteResult.lastEventId,
clientData,
});
},
{
@@ -1023,6 +1062,27 @@ function setWarmTimeoutInSeconds(seconds: number): void {
metadata.set(WARM_TIMEOUT_METADATA_KEY, seconds);
}
/**
* Extracts the client data (metadata) type from a chat task.
* Use this to type the `metadata` option on the transport.
*
* @example
* ```ts
* import type { InferChatClientData } from "@trigger.dev/sdk/ai";
* import type { myChat } from "@/trigger/chat";
*
* type MyClientData = InferChatClientData<typeof myChat>;
* // { model?: string; userId: string }
* ```
*/
export type InferChatClientData<TTask extends AnyTask> = TTask extends Task<
string,
ChatTaskWirePayload<any, infer TMetadata>,
any
>
? TMetadata
: unknown;
export const chat = {
/** Create a chat task. See {@link chatTask}. */
task: chatTask,
+2 -1
View File
@@ -29,6 +29,7 @@ import {
type TriggerChatTransportOptions,
} from "./chat.js";
import type { AnyTask, TaskIdentifier } from "@trigger.dev/core/v3";
import type { InferChatClientData } from "./ai.js";
/**
* Options for `useTriggerChatTransport`, with a type-safe `task` field.
@@ -39,7 +40,7 @@ import type { AnyTask, TaskIdentifier } from "@trigger.dev/core/v3";
* ```
*/
export type UseTriggerChatTransportOptions<TTask extends AnyTask = AnyTask> = Omit<
TriggerChatTransportOptions,
TriggerChatTransportOptions<InferChatClientData<TTask>>,
"task"
> & {
/** The task ID. Strongly typed when a task type parameter is provided. */
+7 -6
View File
@@ -35,7 +35,7 @@ const DEFAULT_STREAM_TIMEOUT_SECONDS = 120;
/**
* Options for creating a TriggerChatTransport.
*/
export type TriggerChatTransportOptions = {
export type TriggerChatTransportOptions<TClientData = unknown> = {
/**
* The Trigger.dev task ID to trigger for chat completions.
* This task should be defined using `chatTask()` from `@trigger.dev/sdk/ai`,
@@ -84,22 +84,23 @@ export type TriggerChatTransportOptions = {
streamTimeoutSeconds?: number;
/**
* Default metadata included in every request payload.
* Default client data included in every request payload.
* Merged with per-call `metadata` from `sendMessage()` — per-call values
* take precedence over transport-level defaults.
*
* Useful for data that should accompany every message, like a user ID.
* When the task uses `clientDataSchema`, this is typed to match the schema.
*
* @example
* ```ts
* new TriggerChatTransport({
* task: "my-chat",
* accessToken,
* metadata: { userId: currentUser.id },
* clientData: { userId: currentUser.id },
* });
* ```
*/
metadata?: Record<string, unknown>;
clientData?: TClientData extends Record<string, unknown> ? TClientData : Record<string, unknown>;
/**
* Restore active chat sessions from external storage (e.g. localStorage).
@@ -254,7 +255,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
this.streamKey = options.streamKey ?? DEFAULT_STREAM_KEY;
this.extraHeaders = options.headers ?? {};
this.streamTimeoutSeconds = options.streamTimeoutSeconds ?? DEFAULT_STREAM_TIMEOUT_SECONDS;
this.defaultMetadata = options.metadata;
this.defaultMetadata = options.clientData;
this.triggerOptions = options.triggerOptions;
this._onSessionChange = options.onSessionChange;
@@ -3,6 +3,7 @@
import type { UIMessage } from "ai";
import { generateId } from "ai";
import { useTriggerChatTransport } from "@trigger.dev/sdk/chat/react";
import type { aiChat } from "@/trigger/chat";
import { useCallback, useEffect, useState } from "react";
import { Chat } from "@/components/chat";
import { ChatSidebar } from "@/components/chat-sidebar";
@@ -56,12 +57,13 @@ export function ChatApp({
[]
);
const transport = useTriggerChatTransport({
const transport = useTriggerChatTransport<typeof aiChat>({
task: "ai-chat",
accessToken: getChatToken,
baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL,
sessions: initialSessions,
onSessionChange: handleSessionChange,
clientData: { userId: "user_123" },
triggerOptions: {
tags: ["user:user_123"],
},
+7 -6
View File
@@ -86,6 +86,7 @@ declare const Deno: unknown;
export const aiChat = chat.task({
id: "ai-chat",
clientDataSchema: z.object({ model: z.string().optional(), userId: z.string() }),
warmTimeoutInSeconds: 60,
chatAccessTokenTTL: "2h",
onChatStart: async ({ chatId, runId, chatAccessToken }) => {
@@ -125,20 +126,20 @@ export const aiChat = chat.task({
});
},
run: async ({ messages, clientData, stopSignal }) => {
const { model: modelId } = z
.object({ model: z.string().optional() })
.parse(clientData ?? {});
return streamText({
model: getModel(modelId),
model: getModel(clientData?.model),
system: "You are a helpful assistant. Be concise and friendly.",
messages,
tools: { inspectEnvironment },
stopWhen: stepCountIs(10),
abortSignal: stopSignal,
providerOptions: {
openai: { user: clientData?.userId },
anthropic: { metadata: { user_id: clientData?.userId } },
},
experimental_telemetry: {
isEnabled: true,
}
},
});
},
});