From 610067f29a731185fcdb01468b7c40c8c59dda47 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Fri, 6 Mar 2026 16:19:25 +0000 Subject: [PATCH] Make clientData typesafe and pass to all chat.task hooks --- packages/core/src/v3/index.ts | 1 + packages/trigger-sdk/src/v3/ai.ts | 108 ++++++++++++++---- packages/trigger-sdk/src/v3/chat-react.ts | 3 +- packages/trigger-sdk/src/v3/chat.ts | 13 ++- .../ai-chat/src/components/chat-app.tsx | 4 +- references/ai-chat/src/trigger/chat.ts | 13 ++- 6 files changed, 104 insertions(+), 38 deletions(-) diff --git a/packages/core/src/v3/index.ts b/packages/core/src/v3/index.ts index 2757363f4..883da2885 100644 --- a/packages/core/src/v3/index.ts +++ b/packages/core/src/v3/index.ts @@ -80,6 +80,7 @@ export { getSchemaParseFn, type AnySchemaParseFn, type SchemaParseFn, + type inferSchemaOut, isSchemaZodEsque, isSchemaValibotEsque, isSchemaArkTypeEsque, diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 3140402d4..3fad1565d 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -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 = { +type ChatTaskWirePayload = { messages: TMessage[]; chatId: string; trigger: "submit-message" | "regenerate-message"; messageId?: string; - metadata?: unknown; + metadata?: TMetadata; }; /** @@ -196,7 +198,7 @@ type ChatTaskWirePayload = { * 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 = { /** 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 = ChatTaskPayload & ChatTaskSignals; // Input streams for bidirectional chat communication const messagesInput = streams.input({ id: CHAT_MESSAGES_STREAM_ID }); @@ -384,13 +386,13 @@ async function pipeChat( /** * Event passed to the `onChatStart` callback. */ -export type ChatStartEvent = { +export type ChatStartEvent = { /** 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 = { /** 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 = { /** 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 = Omit< - TaskOptions, - "run" -> & { +export type ChatTaskOptions< + TIdentifier extends string, + TClientDataSchema extends TaskSchema | undefined = undefined, +> = Omit, "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 = Omit< * **Auto-piping:** If this function returns a value with `.toUIMessageStream()`, * the stream is automatically piped to the frontend. */ - run: (payload: ChatTaskRunPayload) => Promise; + run: (payload: ChatTaskRunPayload>) => Promise; /** * Called on the first turn (turn 0) of a new run, before the `run` function executes. @@ -477,7 +503,7 @@ export type ChatTaskOptions = Omit< * } * ``` */ - onChatStart?: (event: ChatStartEvent) => Promise | void; + onChatStart?: (event: ChatStartEvent>) => Promise | void; /** * Called at the start of every turn, after message accumulation and `onChatStart` (turn 0), @@ -493,7 +519,7 @@ export type ChatTaskOptions = Omit< * } * ``` */ - onTurnStart?: (event: TurnStartEvent) => Promise | void; + onTurnStart?: (event: TurnStartEvent>) => Promise | void; /** * Called after each turn completes (after the response is captured, before waiting @@ -508,7 +534,7 @@ export type ChatTaskOptions = Omit< * } * ``` */ - onTurnComplete?: (event: TurnCompleteEvent) => Promise | void; + onTurnComplete?: (event: TurnCompleteEvent>) => Promise | void; /** * Maximum number of conversational turns (message round-trips) a single run @@ -578,11 +604,15 @@ export type ChatTaskOptions = Omit< * }); * ``` */ -function chatTask( - options: ChatTaskOptions -): Task { +function chatTask< + TIdentifier extends string, + TClientDataSchema extends TaskSchema | undefined = undefined, +>( + options: ChatTaskOptions +): Task>, unknown> { const { run: userRun, + clientDataSchema, onChatStart, onTurnStart, onTurnComplete, @@ -593,7 +623,11 @@ function chatTask( ...restOptions } = options; - return createTask({ + const parseClientData = clientDataSchema + ? getSchemaParseFn(clientDataSchema) + : undefined; + + return createTask>, 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( 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; const lastUserMessage = extractLastUserMessageText(uiMessages); const turnAttributes: Attributes = { @@ -738,7 +775,7 @@ function chatTask( await onChatStart({ chatId: currentWirePayload.chatId, messages: accumulatedMessages, - clientData: wireMetadata, + clientData, runId: currentRunId, chatAccessToken: turnAccessToken, }); @@ -765,6 +802,7 @@ function chatTask( turn, runId: currentRunId, chatAccessToken: turnAccessToken, + clientData, }); }, { @@ -783,11 +821,11 @@ function chatTask( 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( 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; + * // { model?: string; userId: string } + * ``` + */ +export type InferChatClientData = TTask extends Task< + string, + ChatTaskWirePayload, + any +> + ? TMetadata + : unknown; + export const chat = { /** Create a chat task. See {@link chatTask}. */ task: chatTask, diff --git a/packages/trigger-sdk/src/v3/chat-react.ts b/packages/trigger-sdk/src/v3/chat-react.ts index 1ee48a4b2..612f0c184 100644 --- a/packages/trigger-sdk/src/v3/chat-react.ts +++ b/packages/trigger-sdk/src/v3/chat-react.ts @@ -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 = Omit< - TriggerChatTransportOptions, + TriggerChatTransportOptions>, "task" > & { /** The task ID. Strongly typed when a task type parameter is provided. */ diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 6c7d39424..366ecaf52 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -35,7 +35,7 @@ const DEFAULT_STREAM_TIMEOUT_SECONDS = 120; /** * Options for creating a TriggerChatTransport. */ -export type TriggerChatTransportOptions = { +export type TriggerChatTransportOptions = { /** * 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; + clientData?: TClientData extends Record ? TClientData : Record; + /** * Restore active chat sessions from external storage (e.g. localStorage). @@ -254,7 +255,7 @@ export class TriggerChatTransport implements ChatTransport { 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; diff --git a/references/ai-chat/src/components/chat-app.tsx b/references/ai-chat/src/components/chat-app.tsx index 8ffc7b41b..a00695ec4 100644 --- a/references/ai-chat/src/components/chat-app.tsx +++ b/references/ai-chat/src/components/chat-app.tsx @@ -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({ 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"], }, diff --git a/references/ai-chat/src/trigger/chat.ts b/references/ai-chat/src/trigger/chat.ts index 61c455341..68c65d750 100644 --- a/references/ai-chat/src/trigger/chat.ts +++ b/references/ai-chat/src/trigger/chat.ts @@ -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, - } + }, }); }, });