Make clientData typesafe and pass to all chat.task hooks
This commit is contained in:
@@ -80,6 +80,7 @@ export {
|
||||
getSchemaParseFn,
|
||||
type AnySchemaParseFn,
|
||||
type SchemaParseFn,
|
||||
type inferSchemaOut,
|
||||
isSchemaZodEsque,
|
||||
isSchemaValibotEsque,
|
||||
isSchemaArkTypeEsque,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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. */
|
||||
|
||||
@@ -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"],
|
||||
},
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
},
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user