From 50c86cdc1ed713ac2dcdbb697ba21be3eb172a5f Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Sun, 16 Aug 2026 20:55:35 -0700 Subject: [PATCH 1/4] feat(chat): validate custom agent client data --- .changeset/quiet-chats-validate.md | 5 + docs/ai-chat/client-protocol.mdx | 2 +- docs/ai-chat/custom-agents.mdx | 120 +++-- docs/ai-chat/reference.mdx | 4 +- docs/ai-chat/types.mdx | 4 +- packages/trigger-sdk/src/v3/ai.ts | 414 +++++++++++++--- ...ustom-agent-client-data-validation.test.ts | 460 ++++++++++++++++++ 7 files changed, 903 insertions(+), 106 deletions(-) create mode 100644 .changeset/quiet-chats-validate.md create mode 100644 packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts diff --git a/.changeset/quiet-chats-validate.md b/.changeset/quiet-chats-validate.md new file mode 100644 index 0000000000..9eba1990ca --- /dev/null +++ b/.changeset/quiet-chats-validate.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Custom chat agents now validate and parse client data declared with `chat.withClientData({ schema })` before passing it to agent code. diff --git a/docs/ai-chat/client-protocol.mdx b/docs/ai-chat/client-protocol.mdx index d039b39366..3aed294e3a 100644 --- a/docs/ai-chat/client-protocol.mdx +++ b/docs/ai-chat/client-protocol.mdx @@ -771,7 +771,7 @@ type ChatTaskWirePayload - **`metadata` is the wire envelope for `clientData`.** The agent's `clientData` (typed via `chat.withClientData({ schema })`) is read from this field at run boot. If the agent declares e.g. `{ userId: string, model?: string }`, then every `kind: "message"` payload — and the `triggerConfig.basePayload` you sent at session create — must carry a matching `metadata.userId`. The agent rejects messages whose metadata fails schema validation. + **`metadata` is the wire envelope for `clientData`.** The agent's `clientData` (typed via `chat.withClientData({ schema })`) is read from this field at run boot. If the agent declares e.g. `{ userId: string, model?: string }`, then the `triggerConfig.basePayload` you sent at session create and every non-close `kind: "message"` payload must carry a matching `metadata.userId`. Invalid metadata is not passed to agent code. Async reads produce an error chunk followed by `turn-complete`; raw `chat.messages.on()` subscriptions use `onClientDataValidationError` and the task log so they do not end an active response. ### Sending a message diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 197bff6b5e..e9f149457e 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -19,61 +19,91 @@ Inside the wrapper, pick one of two loop styles: - **[Managed loop](#managed-loop-chatcreatesession)** — `chat.createSession()` yields turns; the SDK handles stop signals, accumulation, idle suspend/resume, and turn-complete signaling. You write the turn body. - **[Hand-rolled loop](#hand-rolled-loop-with-primitives)** — you write the loop itself with `chat.messages`, `MessageAccumulator`, `pipeAndCapture`, and `writeTurnComplete`. The right choice when you need complete control over `.toUIMessageStream()` (e.g. `onFinish`, `originalMessages`) beyond what `chat.setUIMessageStreamOptions()` provides, or you're implementing a custom protocol. +### Validating client data + +Use `chat.withClientData({ schema })` to validate `payload.metadata`. Custom agents parse the initial payload and later message and action frames before passing them to `run`, `chat.messages`, or `chat.createSession`. Schema defaults and transforms are included in the value your code receives. Close frames are not validated. + +If validation fails in `run`, `chat.createSession()`, or an async read such as `wait()`, the SDK writes an error chunk followed by `turn-complete` and skips the invalid frame. The invalid value never reaches your run handler or turn loop. If the initial payload is invalid, the task waits for the next valid frame instead of starting `run` with bad data. Without a schema, metadata is passed through unchanged. + +`chat.messages.on()` is different because a subscribed frame can arrive while the current response is still streaming. Ending the turn at that point would cut off the response. The SDK skips the frame, logs the validation error, and calls `onClientDataValidationError` if you set it: + +```ts +import { chat } from "@trigger.dev/sdk/ai"; +import { z } from "zod"; + +export const myChat = chat + .withClientData({ schema: z.object({ userId: z.string() }) }) + .customAgent({ + id: "my-chat", + onClientDataValidationError: ({ error, payload }) => { + console.warn("Invalid client data", { error, trigger: payload.trigger }); + }, + run: async (payload) => { + // ... + }, + }); +``` + +`chat.messages.peek()` validates synchronously and throws validation errors to the caller. If your schema only supports asynchronous parsing, use `once()`, `wait()`, or `waitWithIdleTimeout()` instead. + ## Managed loop: chat.createSession() `chat.createSession()` gives you an async iterator of `ChatTurn` objects. Each turn arrives with the accumulated history, a combined stop+cancel signal, and helpers to finish the turn: ```ts trigger/my-chat.ts -import { chat, type ChatTaskWirePayload } from "@trigger.dev/sdk/ai"; +import { chat } from "@trigger.dev/sdk/ai"; import { streamText, stepCountIs } from "ai"; import { anthropic } from "@ai-sdk/anthropic"; - -export const myChat = chat.customAgent({ - id: "my-chat", - run: async (payload: ChatTaskWirePayload, { signal }) => { - // One-time initialization — plain code, no hooks. Upsert, not create: - // continuation runs boot with the row already in place. - const clientData = payload.metadata as { userId: string }; - await db.chat.upsert({ - where: { id: payload.chatId }, - create: { id: payload.chatId, userId: clientData.userId }, - update: {}, - }); - - const session = chat.createSession(payload, { - signal, - idleTimeoutInSeconds: 60, - timeout: "1h", - }); - - for await (const turn of session) { - // Persist the incoming user message BEFORE streaming — this is your - // onTurnStart equivalent. Without it, a page reload mid-stream - // restores the assistant text (replayed from the session) but loses - // the user message that prompted it. - await db.chat.update({ - where: { id: turn.chatId }, - data: { messages: turn.uiMessages }, +import { z } from "zod"; + +export const myChat = chat + .withClientData({ schema: z.object({ userId: z.string() }) }) + .customAgent({ + id: "my-chat", + run: async (payload, { signal }) => { + // One-time initialization — plain code, no hooks. Upsert, not create: + // continuation runs boot with the row already in place. + const clientData = payload.metadata!; + await db.chat.upsert({ + where: { id: payload.chatId }, + create: { id: payload.chatId, userId: clientData.userId }, + update: {}, }); - const result = streamText({ - model: anthropic("claude-sonnet-4-5"), - messages: turn.messages, - abortSignal: turn.signal, - stopWhen: stepCountIs(15), + const session = chat.createSession(payload, { + signal, + idleTimeoutInSeconds: 60, + timeout: "1h", }); - // Pipe, capture, accumulate, and signal turn-complete — all in one call - await turn.complete(result); - - // Persist the full exchange after the turn — your onTurnComplete equivalent - await db.chat.update({ - where: { id: turn.chatId }, - data: { messages: turn.uiMessages }, - }); - } - }, -}); + for await (const turn of session) { + // Persist the incoming user message BEFORE streaming — this is your + // onTurnStart equivalent. Without it, a page reload mid-stream + // restores the assistant text (replayed from the session) but loses + // the user message that prompted it. + await db.chat.update({ + where: { id: turn.chatId }, + data: { messages: turn.uiMessages }, + }); + + const result = streamText({ + model: anthropic("claude-sonnet-4-5"), + messages: turn.messages, + abortSignal: turn.signal, + stopWhen: stepCountIs(15), + }); + + // Pipe, capture, accumulate, and signal turn-complete — all in one call + await turn.complete(result); + + // Persist the full exchange after the turn — your onTurnComplete equivalent + await db.chat.update({ + where: { id: turn.chatId }, + data: { messages: turn.uiMessages }, + }); + } + }, + }); ``` @@ -102,7 +132,7 @@ Each turn yielded by the iterator provides: | `number` | `number` | Turn number (0-indexed) | | `chatId` | `string` | Chat session ID | | `trigger` | `string` | What triggered this turn | -| `clientData` | `unknown` | Client data from the transport | +| `clientData` | Schema output or `unknown` | Parsed client data when `withClientData` is configured | | `messages` | `ModelMessage[]` | Full accumulated model messages — pass to `streamText` | | `uiMessages` | `UIMessage[]` | Full accumulated UI messages — use for persistence | | `signal` | `AbortSignal` | Combined stop+cancel signal (fresh each turn) | diff --git a/docs/ai-chat/reference.mdx b/docs/ai-chat/reference.mdx index da08f9a047..6960e8161d 100644 --- a/docs/ai-chat/reference.mdx +++ b/docs/ai-chat/reference.mdx @@ -546,7 +546,7 @@ Use this when you need [`InferChatUIMessage`](#inferchatuimessage) / typed `data ## `chat.withClientData` -Returns a [`ChatBuilder`](/ai-chat/types#chatbuilder) with a fixed client data schema. All hooks and `run` get typed `clientData` without passing `clientDataSchema` in `.agent()` options. +Returns a [`ChatBuilder`](/ai-chat/types#chatbuilder) with a fixed client data schema. Managed-agent hooks and `run` get typed `clientData` without passing `clientDataSchema` in `.agent()` options. Custom agents parse `payload.metadata` on the initial payload and later input frames before passing it to user code. ```ts chat.withClientData({ schema: TSchema }): ChatBuilder; @@ -556,6 +556,8 @@ chat.withClientData({ schema: TSchema }): ChatBuilder(fn: (writer: ChatWriter) => Promise | T): Pr return result; } +type ChatCustomAgentClientDataParser = { + parse: (value: unknown) => Promise | unknown; + parseSync: (value: unknown) => unknown; +}; + +type ChatCustomAgentClientDataErrorHandler = (event: { + error: unknown; + payload: ChatTaskWirePayload; +}) => Promise | void; + +const chatCustomAgentClientDataParserKey = locals.create( + "chat.customAgentClientDataParser" +); +const chatCustomAgentClientDataErrorHandlerKey = + locals.create("chat.customAgentClientDataErrorHandler"); + +function shouldValidateChatCustomAgentPayload(payload: ChatTaskWirePayload): boolean { + return ( + payload.trigger !== "close" && locals.get(chatCustomAgentClientDataParserKey) !== undefined + ); +} + +function getChatCustomAgentSyncSchemaParseFn(schema: TaskSchema): (value: unknown) => unknown { + const parser = schema as any; + + if (typeof parser === "function" && typeof parser.assert === "function") { + return parser.assert.bind(parser); + } + + if (typeof parser === "function") { + return (value) => { + const result = parser(value); + if (result && typeof result.then === "function") { + void Promise.resolve(result).catch(() => {}); + throw new Error( + "chat.messages.peek() cannot validate clientData with an asynchronous schema. " + + "Use chat.messages.once(), chat.messages.wait(), or chat.messages.waitWithIdleTimeout()." + ); + } + return result; + }; + } + + if (typeof parser.parse === "function") { + return parser.parse.bind(parser); + } + + if (typeof parser.validateSync === "function") { + return parser.validateSync.bind(parser); + } + + if (typeof parser.create === "function") { + return parser.create.bind(parser); + } + + if (typeof parser.assert === "function") { + return (value) => { + parser.assert(value); + return value; + }; + } + + return () => { + throw new Error( + "chat.messages.peek() cannot validate clientData with this schema. " + + "Use chat.messages.once(), chat.messages.wait(), or chat.messages.waitWithIdleTimeout()." + ); + }; +} + +async function reportChatCustomAgentClientDataError( + payload: ChatTaskWirePayload, + error: unknown, + writeToStream: boolean +): Promise { + const errorText = error instanceof Error ? error.message : "An unexpected error occurred"; + logger.warn("chat.customAgent: clientData validation failed", { + chatId: payload.chatId, + trigger: payload.trigger, + error: errorText, + }); + + const errorHandler = locals.get(chatCustomAgentClientDataErrorHandlerKey); + if (errorHandler) { + try { + await errorHandler({ error, payload }); + } catch (handlerError) { + logger.warn("chat.customAgent: clientData validation error handler failed", { + chatId: payload.chatId, + trigger: payload.trigger, + error: handlerError instanceof Error ? handlerError.message : String(handlerError), + }); + } + } + + if (!writeToStream) { + return; + } + + try { + await withChatWriter((writer) => { + writer.write({ type: "error", errorText } as any); + }); + await chatWriteTurnComplete(); + } catch (signalError) { + logger.warn("chat.customAgent: failed to report clientData validation error", { + chatId: payload.chatId, + trigger: payload.trigger, + error: signalError instanceof Error ? signalError.message : String(signalError), + }); + } +} + +type ChatCustomAgentPayloadValidationResult = + | { ok: true; payload: TPayload } + | { ok: false }; + +async function validateChatCustomAgentPayload( + payload: TPayload, + options: { writeErrorToStream?: boolean } = {} +): Promise> { + const parser = locals.get(chatCustomAgentClientDataParserKey); + if (!parser || payload.trigger === "close") { + return { ok: true, payload }; + } + + try { + const metadata = await parser.parse(payload.metadata); + return { ok: true, payload: { ...payload, metadata } }; + } catch (error) { + await reportChatCustomAgentClientDataError(payload, error, options.writeErrorToStream ?? true); + return { ok: false }; + } +} + +function validateChatCustomAgentPayloadSync( + payload: TPayload +): TPayload { + const parser = locals.get(chatCustomAgentClientDataParserKey); + if (!parser || payload.trigger === "close") { + return payload; + } + + try { + return { ...payload, metadata: parser.parseSync(payload.metadata) }; + } catch (error) { + logger.warn("chat.customAgent: clientData validation failed in chat.messages.peek()", { + chatId: payload.chatId, + trigger: payload.trigger, + error: error instanceof Error ? error.message : String(error), + }); + throw error; + } +} + // `ChatTaskWirePayload` and `ChatInputChunk` live in `./ai-shared.ts` so // browser bundles (which import them via `chat-client.ts` / `chat.ts`) // can pull the types without dragging `ai.ts` into the client graph. @@ -1543,20 +1698,41 @@ export type ChatTaskRunPayload< // keep their original shape. Each accessor resolves the session handle // lazily via `getChatSession()` so the module-level references stay // compatible with the pre-migration wiring. +function subscribeToRawChatMessages(handler: (payload: ChatTaskWirePayload) => unknown) { + return getChatSession().in.on((chunk) => { + if (chunk.kind === "message") { + // Returning `true` marks the record CONSUMED at the manager level: + // it is neither buffered for a later `once()` nor re-delivered by + // the buffer drain when the next turn re-attaches its handler. + void Promise.resolve(handler(chunk.payload)).catch(() => {}); + return true; + } + return undefined; + }); +} + const messagesInput: RealtimeDefinedInputStream = { id: "chat-messages", on(handler) { - return getChatSession().in.on((chunk) => { - if (chunk.kind === "message") { - // Returning `true` marks the record CONSUMED at the manager level: - // it is neither buffered for a later `once()` nor re-delivered by - // the buffer drain when the next turn re-attaches its handler. - // Without this, a message arriving mid-stream was delivered twice - // and ran a duplicate turn. - void Promise.resolve(handler(chunk.payload)).catch(() => {}); - return true; - } - return undefined; + if (!locals.get(chatCustomAgentClientDataParserKey)) { + return subscribeToRawChatMessages(handler); + } + + let delivery = Promise.resolve(); + return subscribeToRawChatMessages((payload) => { + delivery = delivery + .then(async () => { + const result = await validateChatCustomAgentPayload(payload, { + // A subscription may receive a frame while the current turn is + // still streaming. Completing that turn here would close the + // active response, so on() reports through the callback and log. + writeErrorToStream: false, + }); + if (result.ok) { + await handler(result.payload); + } + }) + .catch(() => {}); }); }, once(options) { @@ -1575,8 +1751,15 @@ const messagesInput: RealtimeDefinedInputStream = { return; } if (result.output.kind === "message") { - resolve({ ok: true, output: result.output.payload }); - return; + if (!shouldValidateChatCustomAgentPayload(result.output.payload)) { + resolve({ ok: true, output: result.output.payload }); + return; + } + const validated = await validateChatCustomAgentPayload(result.output.payload); + if (validated.ok) { + resolve({ ok: true, output: validated.payload }); + return; + } } // Non-message chunks (stops) are handled by the stopInput // facade's persistent listener; loop and wait for the next. @@ -1604,7 +1787,9 @@ const messagesInput: RealtimeDefinedInputStream = { }, peek() { const chunk = getChatSession().in.peek(); - if (chunk && chunk.kind === "message") return chunk.payload; + if (chunk && chunk.kind === "message") { + return validateChatCustomAgentPayloadSync(chunk.payload); + } return undefined; }, wait(options) { @@ -1617,8 +1802,15 @@ const messagesInput: RealtimeDefinedInputStream = { return; } if (result.output.kind === "message") { - resolve({ ok: true, output: result.output.payload }); - return; + if (!shouldValidateChatCustomAgentPayload(result.output.payload)) { + resolve({ ok: true, output: result.output.payload }); + return; + } + const validated = await validateChatCustomAgentPayload(result.output.payload); + if (validated.ok) { + resolve({ ok: true, output: validated.payload }); + return; + } } // Stop chunks are handled by the stopInput facade's persistent // listener; loop back into the suspending wait. @@ -1633,7 +1825,13 @@ const messagesInput: RealtimeDefinedInputStream = { const result = await getChatSession().in.waitWithIdleTimeout(options); if (!result.ok) return result; if (result.output.kind === "message") { - return { ok: true, output: result.output.payload }; + if (!shouldValidateChatCustomAgentPayload(result.output.payload)) { + return { ok: true, output: result.output.payload }; + } + const validated = await validateChatCustomAgentPayload(result.output.payload); + if (validated.ok) { + return { ok: true, output: validated.payload }; + } } // Swallow stop-kind chunks — persistent stop listener already handled // the abort; we just loop for the next message. @@ -5323,9 +5521,34 @@ type ChatCustomAgentOptions< ChatTaskWirePayload>, unknown >, - "triggerSource" | "agentConfig" + "triggerSource" | "agentConfig" | "run" > & { + /** + * Schema for validating `metadata` from the frontend. + * + * The initial payload and later `chat.messages` frames are parsed before + * user code receives them. Invalid input is skipped. Async reads write an + * error chunk followed by `turn-complete`; subscriptions use + * `onClientDataValidationError` because a turn may still be streaming. + */ clientDataSchema?: TClientDataSchema; + /** + * Called when a custom-agent input fails `clientDataSchema` validation. + * + * Async reads also write an error chunk followed by `turn-complete`. + * `chat.messages.on()` cannot safely complete a turn that may still be + * streaming, so subscribed frames are reported through this callback and + * the task log instead. + */ + onClientDataValidationError?: (event: { + error: unknown; + payload: ChatTaskWirePayload>; + }) => Promise | void; + run: TaskOptions< + TIdentifier, + ChatTaskWirePayload>, + unknown + >["run"]; }; function chatCustomAgent< @@ -5335,7 +5558,11 @@ function chatCustomAgent< >( options: ChatCustomAgentOptions ): Task>, unknown> { - const { clientDataSchema, run: userRun, ...restOptions } = options; + const { clientDataSchema, onClientDataValidationError, run: userRun, ...restOptions } = options; + const parseClientData = clientDataSchema ? getSchemaParseFn(clientDataSchema) : undefined; + const parseClientDataSync = clientDataSchema + ? getChatCustomAgentSyncSchemaParseFn(clientDataSchema) + : undefined; const task = createTask< TIdentifier, @@ -5362,6 +5589,18 @@ function chatCustomAgent< locals.set(chatSessionHandleKey, sessions.open(payload.chatId)); locals.set(chatExternalIdKey, payload.chatId); locals.set(chatAgentRunContextKey, runOptions.ctx); + if (parseClientData && parseClientDataSync) { + locals.set(chatCustomAgentClientDataParserKey, { + parse: parseClientData, + parseSync: parseClientDataSync, + }); + } + if (onClientDataValidationError) { + locals.set( + chatCustomAgentClientDataErrorHandlerKey, + onClientDataValidationError as ChatCustomAgentClientDataErrorHandler + ); + } // Initialize the turn-complete trim slot so `chat.writeTurnComplete` // trims `session.out` back to the previous turn boundary. Without // this the slot is undefined and the trim never runs, so `.out` @@ -5374,7 +5613,48 @@ function chatCustomAgent< // listener — otherwise a continuation boot replays already-answered // messages into the loop's first wait. await seedSessionInResumeCursorForCustomLoop(payload); - return userRun(payload, runOptions); + + // Keep the schema-free path identical to the original custom-agent + // wrapper, including when userRun starts executing. + if (!parseClientData) { + return userRun(payload, runOptions); + } + + const validated = await validateChatCustomAgentPayload(payload); + if (validated.ok) { + return userRun( + validated.payload as ChatTaskWirePayload>, + runOptions + ); + } + + // The Session base payload is sticky across continuation runs. If it is + // invalid, returning here would boot the same bad metadata again on the + // next message. Stay attached and wait for a valid wire frame instead. + const next = await messagesInput.waitWithIdleTimeout({ + idleTimeoutInSeconds: payload.idleTimeoutInSeconds ?? 30, + timeout: "1h", + spanName: "waiting for valid clientData", + }); + if (!next.ok || next.output.trigger === "close") { + return; + } + + // Normal input frames omit run-level boot context. Carry it forward so + // a continuation still tells the custom loop to restore prior state. + const recoveredPayload = { + ...next.output, + continuation: next.output.continuation ?? payload.continuation, + previousRunId: next.output.previousRunId ?? payload.previousRunId, + sessionId: next.output.sessionId ?? payload.sessionId, + idleTimeoutInSeconds: next.output.idleTimeoutInSeconds ?? payload.idleTimeoutInSeconds, + headStartMessages: next.output.headStartMessages ?? payload.headStartMessages, + }; + + return userRun( + recoveredPayload as ChatTaskWirePayload>, + runOptions + ); }, }); @@ -9441,7 +9721,7 @@ export type ChatSessionOptions = { pendingMessages?: PendingMessagesOptions; }; -export type ChatTurn = { +export type ChatTurn = { /** Turn number (0-indexed). */ number: number; /** Chat session ID. */ @@ -9449,7 +9729,7 @@ export type ChatTurn = { /** What triggered this turn. */ trigger: string; /** Client data from the transport (`metadata` field on the wire payload). */ - clientData: unknown; + clientData: TClientData; /** Full accumulated model messages — pass directly to `streamText`. */ readonly messages: ModelMessage[]; /** Full accumulated UI messages — use for persistence. */ @@ -9548,10 +9828,10 @@ export type ChatTurn = { * }); * ``` */ -function createChatSession( - payload: ChatTaskWirePayload, +function createChatSession( + payload: ChatTaskWirePayload, options: ChatSessionOptions -): AsyncIterable { +): AsyncIterable> { const { signal: runSignal, idleTimeoutInSeconds: sessionIdleTimeoutOpt, @@ -9585,7 +9865,7 @@ function createChatSession( let activeMsgSub: { off: () => void } | undefined; return { - async next(): Promise> { + async next(): Promise>> { activeMsgSub?.off(); activeMsgSub = undefined; if (!booted) { @@ -9647,7 +9927,7 @@ function createChatSession( return { done: true, value: undefined }; } const continuationBoot = isMessagelessContinuationBoot; - currentPayload = result.output; + currentPayload = result.output as ChatTaskWirePayload; // Preserve the continuation flag — the wire payload of the next // message doesn't carry it, and `turn.continuation` is how the // user knows to seed history (e.g. `turn.setMessages(stored)`). @@ -9659,8 +9939,24 @@ function createChatSession( // Subsequent turns: drain buffered mid-turn messages first (they // were consumed and won't be re-delivered), then wait. if (turn > 0) { - if (sessionPendingWire.length > 0) { - currentPayload = sessionPendingWire.shift()!; + let bufferedPayload: ChatTaskWirePayload | undefined; + while (sessionPendingWire.length > 0) { + const candidate = sessionPendingWire.shift()!; + if (!locals.get(chatCustomAgentClientDataParserKey)) { + // Avoid adding an async boundary when no schema is configured. + bufferedPayload = candidate as ChatTaskWirePayload; + break; + } + + const validated = await validateChatCustomAgentPayload(candidate); + if (validated.ok) { + bufferedPayload = validated.payload as ChatTaskWirePayload; + break; + } + } + + if (bufferedPayload) { + currentPayload = bufferedPayload; } else { // chat.requestUpgrade() / chat.endRun() — exit before waiting if (locals.get(chatUpgradeRequestedKey) || locals.get(chatEndRunRequestedKey)) { @@ -9677,7 +9973,7 @@ function createChatSession( stop.cleanup(); return { done: true, value: undefined }; } - currentPayload = next.output; + currentPayload = next.output as ChatTaskWirePayload; } } @@ -9707,37 +10003,39 @@ function createChatSession( }); // Listen for messages during streaming (steering + next-turn buffer) - const sessionMsgSub = messagesInput.on(async (msg) => { - if (sessionPendingMessages) { - // Steering route — the frontend re-sends non-injected - // messages on turn complete, so don't also buffer the wire. - // Slim wire: at most one delta message per record. Read - // `msg.message` directly — no array slicing needed. - const lastUIMessage = msg.message; - if (lastUIMessage) { - if (sessionPendingMessages.onReceived) { + const sessionMsgSub = sessionPendingMessages + ? messagesInput.on(async (msg) => { + // Steering route — the frontend re-sends non-injected + // messages on turn complete, so don't also buffer the wire. + // Slim wire: at most one delta message per record. Read + // `msg.message` directly — no array slicing needed. + const lastUIMessage = msg.message; + if (lastUIMessage) { + if (sessionPendingMessages.onReceived) { + try { + await sessionPendingMessages.onReceived({ + message: lastUIMessage, + chatId: currentPayload.chatId, + turn, + }); + } catch { + /* non-fatal */ + } + } try { - await sessionPendingMessages.onReceived({ - message: lastUIMessage, - chatId: currentPayload.chatId, - turn, - }); + const modelMsgs = await toModelMessages([lastUIMessage]); + turnSteeringQueue.push({ uiMessage: lastUIMessage, modelMessages: modelMsgs }); } catch { /* non-fatal */ } } - try { - const modelMsgs = await toModelMessages([lastUIMessage]); - turnSteeringQueue.push({ uiMessage: lastUIMessage, modelMessages: modelMsgs }); - } catch { - /* non-fatal */ - } - } - return; - } - - sessionPendingWire.push(msg); - }); + }) + : subscribeToRawChatMessages((msg) => { + // Buffer synchronously in wire order. Validation happens when + // the frame becomes the next turn, after the active response + // has completed and it is safe to write an error boundary. + sessionPendingWire.push(msg); + }); activeMsgSub = sessionMsgSub; // Accumulate messages. Slim wire: pass the single delta message as @@ -9773,11 +10071,11 @@ function createChatSession( const combinedSignal = AbortSignal.any([runSignal, stop.signal]); - const turnObj: ChatTurn = { + const turnObj: ChatTurn = { number: turn, chatId: currentPayload.chatId, trigger: currentPayload.trigger, - clientData: currentPayload.metadata, + clientData: currentPayload.metadata as TClientData, get messages() { return accumulator.modelMessages; }, diff --git a/packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts b/packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts new file mode 100644 index 0000000000..b432ca55f9 --- /dev/null +++ b/packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts @@ -0,0 +1,460 @@ +// Import the test harness first so chat tasks register in its resource catalog. +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { describe, expect, expectTypeOf, it } from "vitest"; +import { z } from "zod"; +import { chat } from "../src/v3/ai.js"; + +function userMessage(text: string, id: string) { + return { + id, + role: "user" as const, + parts: [{ type: "text" as const, text }], + }; +} + +function deferred() { + let resolve!: () => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} + +async function waitFor(check: () => boolean, timeoutMs = 5_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((resolve) => setTimeout(resolve, 10)); + } + throw new Error("waitFor timed out"); +} + +describe("chat.customAgent clientData validation", () => { + it("passes parsed clientData to run and createSession turns", async () => { + const clientData = { userId: "user_123", attempt: "42" }; + let initialClientData: unknown; + let turnClientData: unknown; + + const agent = chat + .withClientData({ + schema: z.object({ + userId: z.string(), + attempt: z.coerce.number().int(), + }), + }) + .customAgent({ + id: "custom-agent-client-data-valid", + run: async (payload, { signal }) => { + expectTypeOf(payload.metadata).toEqualTypeOf< + { userId: string; attempt: number } | undefined + >(); + initialClientData = payload.metadata; + + const session = chat.createSession(payload, { + signal, + idleTimeoutInSeconds: 1, + }); + for await (const turn of session) { + expectTypeOf(turn.clientData).toEqualTypeOf<{ + userId: string; + attempt: number; + }>(); + turnClientData = turn.clientData; + await turn.done(); + break; + } + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-valid-chat", + clientData, + }); + + try { + await waitFor(() => initialClientData !== undefined); + await harness.sendMessage(userMessage("hello", "message-1")); + + expect(initialClientData).toEqual({ userId: "user_123", attempt: 42 }); + expect(turnClientData).toEqual({ userId: "user_123", attempt: 42 }); + } finally { + await harness.close(); + } + }); + + it("reports an invalid frame without passing it to the turn loop", async () => { + const clientData: { userId: string; attempt: unknown } = { + userId: "user_123", + attempt: "1", + }; + let started = false; + const receivedClientData: unknown[] = []; + + const agent = chat + .withClientData({ + schema: z.object({ + userId: z.string(), + attempt: z.coerce.number().int(), + }), + }) + .customAgent({ + id: "custom-agent-client-data-invalid-frame", + run: async (payload, { signal }) => { + started = true; + const session = chat.createSession(payload, { + signal, + idleTimeoutInSeconds: 1, + }); + for await (const turn of session) { + receivedClientData.push(turn.clientData); + await turn.done(); + } + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-invalid-frame-chat", + clientData, + }); + + try { + await waitFor(() => started); + clientData.attempt = "not-a-number"; + + const invalidTurn = await harness.sendMessage(userMessage("invalid", "message-1")); + + expect(receivedClientData).toHaveLength(0); + expect(invalidTurn.chunks).toEqual([ + expect.objectContaining({ type: "error", errorText: expect.any(String) }), + ]); + expect(invalidTurn.rawChunks).toContainEqual( + expect.objectContaining({ type: "trigger:turn-complete" }) + ); + + clientData.attempt = "2"; + await harness.sendMessage(userMessage("valid", "message-2")); + await waitFor(() => receivedClientData.length === 1); + + expect(receivedClientData).toEqual([{ userId: "user_123", attempt: 2 }]); + } finally { + await harness.close(); + } + }); + + it("waits for valid clientData when the initial payload is invalid", async () => { + let runCalls = 0; + let receivedClientData: unknown; + let receivedContinuation: boolean | undefined; + let receivedPreviousRunId: string | undefined; + const clientData: { userId: unknown } = { userId: 123 }; + + const agent = chat + .withClientData({ + schema: z.object({ userId: z.string() }), + }) + .customAgent({ + id: "custom-agent-client-data-invalid-initial", + run: async (payload) => { + runCalls++; + receivedClientData = payload.metadata; + receivedContinuation = payload.continuation; + receivedPreviousRunId = payload.previousRunId; + await chat.writeTurnComplete(); + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-invalid-initial-chat", + clientData, + continuation: true, + previousRunId: "run_previous", + }); + + try { + await waitFor(() => + harness.allRawChunks.some( + (chunk) => + typeof chunk === "object" && + chunk !== null && + (chunk as { type?: string }).type === "trigger:turn-complete" + ) + ); + + expect(runCalls).toBe(0); + expect(harness.allChunks).toEqual([ + expect.objectContaining({ type: "error", errorText: expect.any(String) }), + ]); + + clientData.userId = "user_123"; + await harness.sendMessage(userMessage("retry", "message-1")); + + expect(runCalls).toBe(1); + expect(receivedClientData).toEqual({ userId: "user_123" }); + expect(receivedContinuation).toBe(true); + expect(receivedPreviousRunId).toBe("run_previous"); + } finally { + await harness.close(); + } + }); + + it("keeps async chat.messages.on deliveries in wire order", async () => { + const clientData = { sequence: 0 }; + const parserStarts: number[] = []; + const received: number[] = []; + let started = false; + const finished = deferred(); + + const agent = chat + .withClientData({ + schema: async (value: unknown) => { + const sequence = (value as { sequence: number }).sequence; + parserStarts.push(sequence); + if (sequence === 1) { + await new Promise((resolve) => setTimeout(resolve, 50)); + } + return { sequence }; + }, + }) + .customAgent({ + id: "custom-agent-client-data-async-order", + run: async () => { + started = true; + const subscription = chat.messages.on(async (payload) => { + received.push((payload.metadata as { sequence: number }).sequence); + await chat.writeTurnComplete(); + if (received.length === 2) { + finished.resolve(); + } + }); + await finished.promise; + subscription.off(); + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-async-order-chat", + clientData, + }); + + try { + await waitFor(() => started); + + clientData.sequence = 1; + const first = harness.sendMessage(userMessage("first", "message-1")); + await waitFor(() => parserStarts.includes(1)); + + clientData.sequence = 2; + const second = harness.sendMessage(userMessage("second", "message-2")); + + await Promise.all([first, second]); + await waitFor(() => received.length === 2); + + expect(received).toEqual([1, 2]); + } finally { + finished.resolve(); + await harness.close(); + } + }); + + it("delivers frames that arrived before chat.messages.on is removed", async () => { + const clientData = { blocked: false }; + const parserStarted = deferred(); + const releaseParser = deferred(); + const delivered = deferred(); + let removeSubscription: (() => void) | undefined; + let handlerCalls = 0; + let started = false; + + const agent = chat + .withClientData({ + schema: async (value: unknown) => { + const blocked = (value as { blocked: boolean }).blocked; + if (blocked) { + parserStarted.resolve(); + await releaseParser.promise; + } + return { blocked }; + }, + }) + .customAgent({ + id: "custom-agent-client-data-off-after-arrival", + run: async (_payload, { signal }) => { + started = true; + const subscription = chat.messages.on(async () => { + handlerCalls++; + await chat.writeTurnComplete(); + delivered.resolve(); + }); + removeSubscription = () => subscription.off(); + await Promise.race([ + delivered.promise, + new Promise((resolve) => { + signal.addEventListener("abort", () => resolve(), { once: true }); + }), + ]); + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-off-after-arrival-chat", + clientData, + }); + + try { + await waitFor(() => started); + clientData.blocked = true; + const send = harness.sendMessage(userMessage("hello", "message-1")); + await parserStarted.promise; + + removeSubscription!(); + releaseParser.resolve(); + + await send; + await delivered.promise; + expect(handlerCalls).toBe(1); + } finally { + releaseParser.resolve(); + await harness.close(); + } + }); + + it("does not complete an active turn when a buffered frame is invalid", async () => { + const clientData: { attempt: unknown } = { attempt: "1" }; + const firstTurnStarted = deferred(); + const releaseFirstTurn = deferred(); + const validationErrors: unknown[] = []; + const receivedClientData: unknown[] = []; + let started = false; + + const agent = chat + .withClientData({ schema: z.object({ attempt: z.coerce.number().int() }) }) + .customAgent({ + id: "custom-agent-client-data-buffered-invalid", + onClientDataValidationError: ({ error }) => { + validationErrors.push(error); + }, + run: async (payload, { signal }) => { + started = true; + const session = chat.createSession(payload, { + signal, + idleTimeoutInSeconds: 1, + }); + for await (const turn of session) { + receivedClientData.push(turn.clientData); + firstTurnStarted.resolve(); + await releaseFirstTurn.promise; + await turn.done(); + } + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-buffered-invalid-chat", + clientData, + }); + + try { + await waitFor(() => started); + const first = harness.sendMessage(userMessage("first", "message-1")); + await firstTurnStarted.promise; + + clientData.attempt = "not-a-number"; + const invalid = harness.sendMessage(userMessage("invalid", "message-2")); + await new Promise((resolve) => setTimeout(resolve, 75)); + + expect(validationErrors).toHaveLength(0); + expect(harness.allRawChunks).toHaveLength(0); + + releaseFirstTurn.resolve(); + await Promise.all([first, invalid]); + await waitFor(() => validationErrors.length === 1); + + expect(receivedClientData).toEqual([{ attempt: 1 }]); + expect(harness.allChunks).toContainEqual( + expect.objectContaining({ type: "error", errorText: expect.any(String) }) + ); + } finally { + releaseFirstTurn.resolve(); + await harness.close(); + } + }); + + it("reports invalid chat.messages.on frames without calling the subscriber", async () => { + const clientData: { userId: unknown } = { userId: "user_123" }; + const validationErrors: unknown[] = []; + let handlerCalls = 0; + let started = false; + + const agent = chat.withClientData({ schema: z.object({ userId: z.string() }) }).customAgent({ + id: "custom-agent-client-data-on-invalid", + onClientDataValidationError: ({ error }) => { + validationErrors.push(error); + }, + run: async (_payload, { signal }) => { + started = true; + const subscription = chat.messages.on(() => { + handlerCalls++; + }); + await new Promise((resolve) => { + signal.addEventListener("abort", () => resolve(), { once: true }); + }); + subscription.off(); + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-on-invalid-chat", + clientData, + }); + + try { + await waitFor(() => started); + clientData.userId = 123; + void harness.sendMessage(userMessage("invalid", "message-1")); + await waitFor(() => validationErrors.length === 1); + + expect(handlerCalls).toBe(0); + expect(harness.allRawChunks).toHaveLength(0); + } finally { + await harness.close(); + } + }); + + it("passes clientData through unchanged when no schema is configured", async () => { + const clientData = { userId: "user_123", nested: { enabled: true } }; + let initialClientData: unknown; + let turnClientData: unknown; + + const agent = chat.customAgent({ + id: "custom-agent-client-data-no-schema", + run: async (payload, { signal }) => { + initialClientData = payload.metadata; + const session = chat.createSession(payload, { + signal, + idleTimeoutInSeconds: 1, + }); + for await (const turn of session) { + turnClientData = turn.clientData; + await turn.done(); + break; + } + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-no-schema-chat", + clientData, + }); + + try { + await waitFor(() => initialClientData !== undefined); + await harness.sendMessage(userMessage("hello", "message-1")); + + expect(initialClientData).toBe(clientData); + expect(turnClientData).toBe(clientData); + } finally { + await harness.close(); + } + }); +}); From 80790a7258f53b030f097ce0569dfb7cd8532b82 Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Mon, 17 Aug 2026 01:05:24 -0700 Subject: [PATCH 2/4] fix(chat): harden custom agent validation recovery --- docs/ai-chat/custom-agents.mdx | 2 +- packages/trigger-sdk/src/v3/ai.ts | 32 ++++++- ...ustom-agent-client-data-validation.test.ts | 92 +++++++++++++++++++ 3 files changed, 123 insertions(+), 3 deletions(-) diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index e9f149457e..71e2f624ad 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -23,7 +23,7 @@ Inside the wrapper, pick one of two loop styles: Use `chat.withClientData({ schema })` to validate `payload.metadata`. Custom agents parse the initial payload and later message and action frames before passing them to `run`, `chat.messages`, or `chat.createSession`. Schema defaults and transforms are included in the value your code receives. Close frames are not validated. -If validation fails in `run`, `chat.createSession()`, or an async read such as `wait()`, the SDK writes an error chunk followed by `turn-complete` and skips the invalid frame. The invalid value never reaches your run handler or turn loop. If the initial payload is invalid, the task waits for the next valid frame instead of starting `run` with bad data. Without a schema, metadata is passed through unchanged. +If validation fails in `run`, `chat.createSession()`, or an async read such as `wait()`, the SDK writes an error chunk followed by `turn-complete` and skips the invalid frame. The invalid value never reaches your run handler or turn loop. If the initial payload is invalid, the task waits for the next valid frame instead of starting `run` with bad data. On a [head-start handover](/ai-chat/fast-starts#handover-with-custom-agents) boot with invalid client data, the SDK drains the warm handover signal first: a skip ends the run, and a real handover partial is discarded with a logged warning. Without a schema, metadata is passed through unchanged. `chat.messages.on()` is different because a subscribed frame can arrive while the current response is still streaming. Ending the turn at that point would cut off the response. The SDK skips the frame, logs the validation error, and calls `onClientDataValidationError` if you set it: diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 9c88c4ab35..03405f1a4e 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -5539,10 +5539,13 @@ type ChatCustomAgentOptions< * `chat.messages.on()` cannot safely complete a turn that may still be * streaming, so subscribed frames are reported through this callback and * the task log instead. + * + * `payload.metadata` is typed `unknown`: this callback only fires when + * the metadata failed to parse, so it can be any shape the client sent. */ onClientDataValidationError?: (event: { error: unknown; - payload: ChatTaskWirePayload>; + payload: ChatTaskWirePayload; }) => Promise | void; run: TaskOptions< TIdentifier, @@ -5628,6 +5631,28 @@ function chatCustomAgent< ); } + // A handover-prepare boot parks the warm handler's signal on + // `session.in`. Drain it with the handover facade BEFORE the message + // wait below — that facade consumes-and-discards non-message chunks + // and would swallow the signal (see `waitForHandover`). The warm + // partial cannot be spliced without valid clientData: mirror the + // normal flow for skip/crash (exit without a turn) and drop a real + // partial — the error chunk above already reported the failure. + if (payload.trigger === "handover-prepare") { + const signal = await waitForHandover({ + payload, + timeout: "1h", + spanName: "waiting for handover signal (invalid clientData)", + }); + if (!signal || signal.kind === "handover-skip") { + return; + } + logger.warn( + "chat.customAgent: dropping head-start handover partial — clientData failed validation", + { chatId: payload.chatId, isFinal: signal.isFinal } + ); + } + // The Session base payload is sticky across continuation runs. If it is // invalid, returning here would boot the same bad metadata again on the // next message. Stay attached and wait for a valid wire frame instead. @@ -8561,7 +8586,10 @@ export interface ChatBuilder< options: ChatCustomAgentOptions ) => Task, unknown> : ( - options: ChatCustomAgentOptions + options: Omit< + ChatCustomAgentOptions, + "clientDataSchema" + > ) => Task>, unknown>; } diff --git a/packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts b/packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts index b432ca55f9..249ab8a854 100644 --- a/packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts +++ b/packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts @@ -421,6 +421,98 @@ describe("chat.customAgent clientData validation", () => { } }); + it("exits without a turn when a handover-prepare boot has invalid clientData and the warm handler skips", async () => { + const clientData: { userId: unknown } = { userId: 123 }; + let runCalls = 0; + + const agent = chat.withClientData({ schema: z.object({ userId: z.string() }) }).customAgent({ + id: "custom-agent-client-data-handover-skip", + run: async () => { + runCalls++; + await chat.writeTurnComplete(); + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-handover-skip-chat", + mode: "handover-prepare", + clientData, + }); + + try { + await waitFor(() => + harness.allChunks.some((chunk) => (chunk as { type?: string }).type === "error") + ); + expect(runCalls).toBe(0); + + // The recovery path must drain the skip via the handover facade and + // end the run, mirroring the normal handover-skip exit. + await harness.sendHandoverSkip(); + + // The run has exited — a valid frame must NOT boot the loop. (Without + // the drain, the run would still be sitting in the message wait and + // would process it.) Fire-and-forget: no turn-complete will arrive. + clientData.userId = "user_123"; + void harness.sendMessage(userMessage("late", "message-1")).catch(() => {}); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(runCalls).toBe(0); + } finally { + await harness.close(); + } + }); + + it("drops the head-start partial and recovers on the next valid frame when a handover-prepare boot has invalid clientData", async () => { + const clientData: { userId: unknown } = { userId: 123 }; + let runCalls = 0; + let receivedTrigger: string | undefined; + let receivedClientData: unknown; + + const agent = chat.withClientData({ schema: z.object({ userId: z.string() }) }).customAgent({ + id: "custom-agent-client-data-handover-drop", + run: async (payload) => { + runCalls++; + receivedTrigger = payload.trigger; + receivedClientData = payload.metadata; + await chat.writeTurnComplete(); + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-handover-drop-chat", + mode: "handover-prepare", + clientData, + }); + + try { + await waitFor(() => + harness.allChunks.some((chunk) => (chunk as { type?: string }).type === "error") + ); + expect(runCalls).toBe(0); + + // Resolves on the next turn-complete — the recovered message turn below. + const handover = harness.sendHandover({ + partialAssistantMessage: [ + { role: "assistant", content: [{ type: "text", text: "warm partial" }] }, + ], + }); + // Let the recovery drain consume the handover signal before the + // message frame goes out — a frame arriving mid-drain would be + // discarded by the handover facade (same as the pre-existing turn-0 + // handover wait in chat.createSession). + await new Promise((resolve) => setTimeout(resolve, 50)); + + clientData.userId = "user_123"; + await harness.sendMessage(userMessage("retry", "message-1")); + await handover; + + expect(runCalls).toBe(1); + expect(receivedTrigger).toBe("submit-message"); + expect(receivedClientData).toEqual({ userId: "user_123" }); + } finally { + await harness.close(); + } + }); + it("passes clientData through unchanged when no schema is configured", async () => { const clientData = { userId: "user_123", nested: { enabled: true } }; let initialClientData: unknown; From cb35a3a2832d7636466b85d9c60e8e2aa3b2b4e9 Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Mon, 17 Aug 2026 10:36:54 -0700 Subject: [PATCH 3/4] fix(chat): tighten custom agent validation lifecycle --- docs/ai-chat/custom-agents.mdx | 6 +- packages/trigger-sdk/src/v3/ai.ts | 332 +++++++++++------ ...ustom-agent-client-data-validation.test.ts | 343 +++++++++++++++--- 3 files changed, 527 insertions(+), 154 deletions(-) diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 71e2f624ad..6f962c4b55 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -23,9 +23,11 @@ Inside the wrapper, pick one of two loop styles: Use `chat.withClientData({ schema })` to validate `payload.metadata`. Custom agents parse the initial payload and later message and action frames before passing them to `run`, `chat.messages`, or `chat.createSession`. Schema defaults and transforms are included in the value your code receives. Close frames are not validated. -If validation fails in `run`, `chat.createSession()`, or an async read such as `wait()`, the SDK writes an error chunk followed by `turn-complete` and skips the invalid frame. The invalid value never reaches your run handler or turn loop. If the initial payload is invalid, the task waits for the next valid frame instead of starting `run` with bad data. On a [head-start handover](/ai-chat/fast-starts#handover-with-custom-agents) boot with invalid client data, the SDK drains the warm handover signal first: a skip ends the run, and a real handover partial is discarded with a logged warning. Without a schema, metadata is passed through unchanged. +If validation fails for a submitted turn or an async read such as `wait()`, the SDK writes an error chunk followed by `turn-complete` and skips the invalid frame. The invalid value never reaches your run handler or turn loop. The task then waits for the next valid frame. A messageless preload or continuation boot has no submitted turn to complete, so the SDK reports the error through the task log and `onClientDataValidationError` while it waits. -`chat.messages.on()` is different because a subscribed frame can arrive while the current response is still streaming. Ending the turn at that point would cut off the response. The SDK skips the frame, logs the validation error, and calls `onClientDataValidationError` if you set it: +An invalid [head-start handover](/ai-chat/fast-starts#handover-with-custom-agents) boot fails closed. The SDK waits for the warm handler to finish so stream ordering stays intact. A handover skip ends the run. A real handover writes the validation error and `turn-complete` after the warm output, then ends the run. Without a schema, metadata is passed through unchanged. + +`chat.messages.on()` is different because a subscribed frame can arrive while the current response is still streaming. Ending the turn at that point would cut off the response. While the subscription is active, the SDK skips the frame, logs the validation error, and calls `onClientDataValidationError` if you set it. This also applies to the steering subscription created by `chat.createSession({ pendingMessages })`. Calling `off()` prevents queued validation from invoking your handler or error callback. ```ts import { chat } from "@trigger.dev/sdk/ai"; diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 03405f1a4e..1cbfaca247 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1453,6 +1453,17 @@ function shouldValidateChatCustomAgentPayload(payload: ChatTaskWirePayload): boo ); } +function assertChatCustomAgentSyncParseResult(result: unknown): unknown { + if (result && typeof (result as { then?: unknown }).then === "function") { + void Promise.resolve(result).catch(() => {}); + throw new Error( + "chat.messages.peek() cannot validate clientData with an asynchronous schema. " + + "Use chat.messages.once(), chat.messages.wait(), or chat.messages.waitWithIdleTimeout()." + ); + } + return result; +} + function getChatCustomAgentSyncSchemaParseFn(schema: TaskSchema): (value: unknown) => unknown { const parser = schema as any; @@ -1461,21 +1472,11 @@ function getChatCustomAgentSyncSchemaParseFn(schema: TaskSchema): (value: unknow } if (typeof parser === "function") { - return (value) => { - const result = parser(value); - if (result && typeof result.then === "function") { - void Promise.resolve(result).catch(() => {}); - throw new Error( - "chat.messages.peek() cannot validate clientData with an asynchronous schema. " + - "Use chat.messages.once(), chat.messages.wait(), or chat.messages.waitWithIdleTimeout()." - ); - } - return result; - }; + return (value) => assertChatCustomAgentSyncParseResult(parser(value)); } if (typeof parser.parse === "function") { - return parser.parse.bind(parser); + return (value) => assertChatCustomAgentSyncParseResult(parser.parse(value)); } if (typeof parser.validateSync === "function") { @@ -1501,10 +1502,29 @@ function getChatCustomAgentSyncSchemaParseFn(schema: TaskSchema): (value: unknow }; } +async function writeChatCustomAgentClientDataErrorToStream( + payload: ChatTaskWirePayload, + error: unknown +): Promise { + const errorText = error instanceof Error ? error.message : "An unexpected error occurred"; + try { + await withChatWriter((writer) => { + writer.write({ type: "error", errorText } as any); + }); + await chatWriteTurnComplete(); + } catch (signalError) { + logger.warn("chat.customAgent: failed to report clientData validation error", { + chatId: payload.chatId, + trigger: payload.trigger, + error: signalError instanceof Error ? signalError.message : String(signalError), + }); + } +} + async function reportChatCustomAgentClientDataError( payload: ChatTaskWirePayload, error: unknown, - writeToStream: boolean + options: { writeToStream: boolean; callHandler?: boolean } ): Promise { const errorText = error instanceof Error ? error.message : "An unexpected error occurred"; logger.warn("chat.customAgent: clientData validation failed", { @@ -1513,7 +1533,10 @@ async function reportChatCustomAgentClientDataError( error: errorText, }); - const errorHandler = locals.get(chatCustomAgentClientDataErrorHandlerKey); + const errorHandler = + options.callHandler === false + ? undefined + : locals.get(chatCustomAgentClientDataErrorHandlerKey); if (errorHandler) { try { await errorHandler({ error, payload }); @@ -1526,31 +1549,18 @@ async function reportChatCustomAgentClientDataError( } } - if (!writeToStream) { + if (!options.writeToStream) { return; } - - try { - await withChatWriter((writer) => { - writer.write({ type: "error", errorText } as any); - }); - await chatWriteTurnComplete(); - } catch (signalError) { - logger.warn("chat.customAgent: failed to report clientData validation error", { - chatId: payload.chatId, - trigger: payload.trigger, - error: signalError instanceof Error ? signalError.message : String(signalError), - }); - } + await writeChatCustomAgentClientDataErrorToStream(payload, error); } type ChatCustomAgentPayloadValidationResult = | { ok: true; payload: TPayload } - | { ok: false }; + | { ok: false; error: unknown }; -async function validateChatCustomAgentPayload( - payload: TPayload, - options: { writeErrorToStream?: boolean } = {} +async function parseChatCustomAgentPayload( + payload: TPayload ): Promise> { const parser = locals.get(chatCustomAgentClientDataParserKey); if (!parser || payload.trigger === "close") { @@ -1561,11 +1571,23 @@ async function validateChatCustomAgentPayload( + payload: TPayload, + options: { writeErrorToStream?: boolean } = {} +): Promise> { + const result = await parseChatCustomAgentPayload(payload); + if (!result.ok) { + await reportChatCustomAgentClientDataError(payload, result.error, { + writeToStream: options.writeErrorToStream ?? true, + }); + } + return result; +} + function validateChatCustomAgentPayloadSync( payload: TPayload ): TPayload { @@ -1698,7 +1720,14 @@ export type ChatTaskRunPayload< // keep their original shape. Each accessor resolves the session handle // lazily via `getChatSession()` so the module-level references stay // compatible with the pre-migration wiring. -function subscribeToRawChatMessages(handler: (payload: ChatTaskWirePayload) => unknown) { +type ChatMessageSubscription = { + off: () => void; + drain?: () => Promise; +}; + +function subscribeToRawChatMessages( + handler: (payload: ChatTaskWirePayload) => unknown +): ChatMessageSubscription { return getChatSession().in.on((chunk) => { if (chunk.kind === "message") { // Returning `true` marks the record CONSUMED at the manager level: @@ -1711,6 +1740,55 @@ function subscribeToRawChatMessages(handler: (payload: ChatTaskWirePayload) => u }); } +function subscribeToValidatedChatMessages( + handler: (payload: ChatTaskWirePayload, isActive: () => boolean) => unknown, + options: { + onAfterOff?: (payload: ChatTaskWirePayload) => unknown; + onInvalidAfterOff?: (payload: ChatTaskWirePayload, error: unknown) => unknown; + } = {} +): ChatMessageSubscription { + let active = true; + let delivery = Promise.resolve(); + const subscription = subscribeToRawChatMessages((payload) => { + delivery = delivery + .then(async () => { + const result = await parseChatCustomAgentPayload(payload); + if (!result.ok) { + if (active) { + // Completing the turn here could close an active response. + await reportChatCustomAgentClientDataError(payload, result.error, { + writeToStream: false, + }); + } else if (options.onInvalidAfterOff) { + await options.onInvalidAfterOff(payload, result.error); + } else { + // The subscription was removed while parsing. Keep the failure + // observable without invoking a user callback after off(). + await reportChatCustomAgentClientDataError(payload, result.error, { + writeToStream: false, + callHandler: false, + }); + } + return; + } + if (active) { + await handler(result.payload, () => active); + } else { + await options.onAfterOff?.(result.payload); + } + }) + .catch(() => {}); + }); + + return { + off() { + active = false; + subscription.off(); + }, + drain: () => delivery, + }; +} + const messagesInput: RealtimeDefinedInputStream = { id: "chat-messages", on(handler) { @@ -1718,22 +1796,7 @@ const messagesInput: RealtimeDefinedInputStream = { return subscribeToRawChatMessages(handler); } - let delivery = Promise.resolve(); - return subscribeToRawChatMessages((payload) => { - delivery = delivery - .then(async () => { - const result = await validateChatCustomAgentPayload(payload, { - // A subscription may receive a frame while the current turn is - // still streaming. Completing that turn here would close the - // active response, so on() reports through the callback and log. - writeErrorToStream: false, - }); - if (result.ok) { - await handler(result.payload); - } - }) - .catch(() => {}); - }); + return subscribeToValidatedChatMessages((payload) => handler(payload)); }, once(options) { const ctx = taskContext.ctx; @@ -5527,18 +5590,18 @@ type ChatCustomAgentOptions< * Schema for validating `metadata` from the frontend. * * The initial payload and later `chat.messages` frames are parsed before - * user code receives them. Invalid input is skipped. Async reads write an - * error chunk followed by `turn-complete`; subscriptions use - * `onClientDataValidationError` because a turn may still be streaming. + * user code receives them. Invalid submitted turns and async reads write an + * error chunk followed by `turn-complete`. Messageless boots and active + * subscriptions use `onClientDataValidationError` and the task log because + * there is no submitted turn to complete or a response may still be streaming. */ clientDataSchema?: TClientDataSchema; /** * Called when a custom-agent input fails `clientDataSchema` validation. * - * Async reads also write an error chunk followed by `turn-complete`. - * `chat.messages.on()` cannot safely complete a turn that may still be - * streaming, so subscribed frames are reported through this callback and - * the task log instead. + * Submitted turns and async reads also write an error chunk followed by + * `turn-complete`. Messageless boots and active `chat.messages.on()` + * subscriptions are reported through this callback and the task log only. * * `payload.metadata` is typed `unknown`: this callback only fires when * the metadata failed to parse, so it can be any shape the client sent. @@ -5623,7 +5686,20 @@ function chatCustomAgent< return userRun(payload, runOptions); } - const validated = await validateChatCustomAgentPayload(payload); + const isHandoverBoot = payload.trigger === "handover-prepare"; + const isMessagelessBoot = + payload.trigger === "preload" || + (payload.continuation === true && + payload.message === undefined && + payload.trigger !== "action" && + payload.trigger !== "regenerate-message" && + !isHandoverBoot); + const validated = await validateChatCustomAgentPayload(payload, { + // Preload and continuation boots do not represent a submitted turn, + // so there is no sender waiting for a terminal frame. Handover errors + // must be written after the warm response flushes and signals below. + writeErrorToStream: !isMessagelessBoot && !isHandoverBoot, + }); if (validated.ok) { return userRun( validated.payload as ChatTaskWirePayload>, @@ -5631,14 +5707,7 @@ function chatCustomAgent< ); } - // A handover-prepare boot parks the warm handler's signal on - // `session.in`. Drain it with the handover facade BEFORE the message - // wait below — that facade consumes-and-discards non-message chunks - // and would swallow the signal (see `waitForHandover`). The warm - // partial cannot be spliced without valid clientData: mirror the - // normal flow for skip/crash (exit without a turn) and drop a real - // partial — the error chunk above already reported the failure. - if (payload.trigger === "handover-prepare") { + if (isHandoverBoot) { const signal = await waitForHandover({ payload, timeout: "1h", @@ -5647,10 +5716,11 @@ function chatCustomAgent< if (!signal || signal.kind === "handover-skip") { return; } - logger.warn( - "chat.customAgent: dropping head-start handover partial — clientData failed validation", - { chatId: payload.chatId, isFinal: signal.isFinal } - ); + + // The head-start writer flushes before sending this signal. Writing + // the terminal error now preserves stream order and closes the stitch. + await writeChatCustomAgentClientDataErrorToStream(payload, validated.error); + return; } // The Session base payload is sticky across continuation runs. If it is @@ -5673,7 +5743,6 @@ function chatCustomAgent< previousRunId: next.output.previousRunId ?? payload.previousRunId, sessionId: next.output.sessionId ?? payload.sessionId, idleTimeoutInSeconds: next.output.idleTimeoutInSeconds ?? payload.idleTimeoutInSeconds, - headStartMessages: next.output.headStartMessages ?? payload.headStartMessages, }; return userRun( @@ -9887,14 +9956,23 @@ function createChatSession( // Messages consumed mid-turn, dispatched one per next(). Iterator-level // for the same reason as the agent loop's `pendingWireMessages`: // consumed records never replay, so a turn-local buffer loses them. - const sessionPendingWire: ChatTaskWirePayload[] = []; + const sessionPendingWire: Array< + | { payload: ChatTaskWirePayload; validation: "unvalidated" | "valid" } + | { payload: ChatTaskWirePayload; validation: "invalid"; error: unknown } + > = []; // The current turn's message subscription — detached defensively at the // top of next() in case user code threw without complete()/done(). - let activeMsgSub: { off: () => void } | undefined; + let activeMsgSub: ChatMessageSubscription | undefined; return { async next(): Promise>> { - activeMsgSub?.off(); + if (activeMsgSub?.drain) { + activeMsgSub.off(); + await activeMsgSub.drain(); + } else { + // Keep the schema-free path free of a new async boundary. + activeMsgSub?.off(); + } activeMsgSub = undefined; if (!booted) { booted = true; @@ -9970,13 +10048,22 @@ function createChatSession( let bufferedPayload: ChatTaskWirePayload | undefined; while (sessionPendingWire.length > 0) { const candidate = sessionPendingWire.shift()!; - if (!locals.get(chatCustomAgentClientDataParserKey)) { + if (candidate.validation === "invalid") { + await reportChatCustomAgentClientDataError(candidate.payload, candidate.error, { + writeToStream: true, + }); + continue; + } + if ( + candidate.validation === "valid" || + !locals.get(chatCustomAgentClientDataParserKey) + ) { // Avoid adding an async boundary when no schema is configured. - bufferedPayload = candidate as ChatTaskWirePayload; + bufferedPayload = candidate.payload as ChatTaskWirePayload; break; } - const validated = await validateChatCustomAgentPayload(candidate); + const validated = await validateChatCustomAgentPayload(candidate.payload); if (validated.ok) { bufferedPayload = validated.payload as ChatTaskWirePayload; break; @@ -10031,38 +10118,72 @@ function createChatSession( }); // Listen for messages during streaming (steering + next-turn buffer) - const sessionMsgSub = sessionPendingMessages - ? messagesInput.on(async (msg) => { - // Steering route — the frontend re-sends non-injected - // messages on turn complete, so don't also buffer the wire. - // Slim wire: at most one delta message per record. Read - // `msg.message` directly — no array slicing needed. - const lastUIMessage = msg.message; - if (lastUIMessage) { - if (sessionPendingMessages.onReceived) { - try { - await sessionPendingMessages.onReceived({ - message: lastUIMessage, - chatId: currentPayload.chatId, - turn, - }); - } catch { - /* non-fatal */ - } - } - try { - const modelMsgs = await toModelMessages([lastUIMessage]); - turnSteeringQueue.push({ uiMessage: lastUIMessage, modelMessages: modelMsgs }); - } catch { - /* non-fatal */ - } + const handleSteeringMessage = async ( + msg: ChatTaskWirePayload, + isActive: () => boolean = () => true + ) => { + const bufferForNextTurn = () => { + sessionPendingWire.push({ payload: msg, validation: "valid" }); + }; + if (!isActive()) { + bufferForNextTurn(); + return; + } + + // Steering route — the frontend re-sends non-injected + // messages on turn complete, so don't also buffer the wire. + // Slim wire: at most one delta message per record. Read + // `msg.message` directly — no array slicing needed. + const lastUIMessage = msg.message; + if (lastUIMessage) { + if (sessionPendingMessages?.onReceived) { + try { + await sessionPendingMessages.onReceived({ + message: lastUIMessage, + chatId: currentPayload.chatId, + turn, + }); + } catch { + /* non-fatal */ } - }) + } + if (!isActive()) { + bufferForNextTurn(); + return; + } + try { + const modelMsgs = await toModelMessages([lastUIMessage]); + if (!isActive()) { + bufferForNextTurn(); + return; + } + turnSteeringQueue.push({ uiMessage: lastUIMessage, modelMessages: modelMsgs }); + } catch { + /* non-fatal */ + } + } + }; + + const sessionMsgSub: ChatMessageSubscription = sessionPendingMessages + ? locals.get(chatCustomAgentClientDataParserKey) + ? subscribeToValidatedChatMessages(handleSteeringMessage, { + onAfterOff: (msg) => { + sessionPendingWire.push({ payload: msg, validation: "valid" }); + }, + onInvalidAfterOff: (msg, error) => { + sessionPendingWire.push({ payload: msg, validation: "invalid", error }); + }, + }) + : messagesInput.on(async (msg) => { + // Steering route — the frontend re-sends non-injected + // messages on turn complete, so don't also buffer the wire. + await handleSteeringMessage(msg); + }) : subscribeToRawChatMessages((msg) => { // Buffer synchronously in wire order. Validation happens when // the frame becomes the next turn, after the active response // has completed and it is safe to write an error boundary. - sessionPendingWire.push(msg); + sessionPendingWire.push({ payload: msg, validation: "unvalidated" }); }); activeMsgSub = sessionMsgSub; @@ -10280,7 +10401,6 @@ function createChatSession( } } - sessionMsgSub.off(); await chatWriteTurnComplete(); return response; }, diff --git a/packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts b/packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts index 249ab8a854..492be61be8 100644 --- a/packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts +++ b/packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts @@ -142,11 +142,12 @@ describe("chat.customAgent clientData validation", () => { } }); - it("waits for valid clientData when the initial payload is invalid", async () => { + it("waits without completing a turn when a messageless continuation boot is invalid", async () => { let runCalls = 0; let receivedClientData: unknown; let receivedContinuation: boolean | undefined; let receivedPreviousRunId: string | undefined; + const validationErrors: unknown[] = []; const clientData: { userId: unknown } = { userId: 123 }; const agent = chat @@ -155,6 +156,9 @@ describe("chat.customAgent clientData validation", () => { }) .customAgent({ id: "custom-agent-client-data-invalid-initial", + onClientDataValidationError: ({ error }) => { + validationErrors.push(error); + }, run: async (payload) => { runCalls++; receivedClientData = payload.metadata; @@ -171,6 +175,48 @@ describe("chat.customAgent clientData validation", () => { previousRunId: "run_previous", }); + try { + await waitFor(() => validationErrors.length === 1); + + expect(runCalls).toBe(0); + expect(harness.allRawChunks).toHaveLength(0); + + clientData.userId = "user_123"; + const recovered = await harness.sendMessage(userMessage("retry", "message-1")); + + expect(runCalls).toBe(1); + expect(receivedClientData).toEqual({ userId: "user_123" }); + expect(receivedContinuation).toBe(true); + expect(receivedPreviousRunId).toBe("run_previous"); + expect(recovered.chunks).toHaveLength(0); + expect(recovered.rawChunks).toEqual([ + expect.objectContaining({ type: "trigger:turn-complete" }), + ]); + } finally { + await harness.close(); + } + }); + + it("completes an invalid submitted boot before waiting for valid clientData", async () => { + const clientData: { userId: unknown } = { userId: 123 }; + let runCalls = 0; + let receivedClientData: unknown; + + const agent = chat.withClientData({ schema: z.object({ userId: z.string() }) }).customAgent({ + id: "custom-agent-client-data-invalid-submitted-boot", + run: async (payload) => { + runCalls++; + receivedClientData = payload.metadata; + await chat.writeTurnComplete(); + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-invalid-submitted-boot-chat", + mode: "submit-message", + clientData, + }); + try { await waitFor(() => harness.allRawChunks.some( @@ -191,8 +237,6 @@ describe("chat.customAgent clientData validation", () => { expect(runCalls).toBe(1); expect(receivedClientData).toEqual({ userId: "user_123" }); - expect(receivedContinuation).toBe(true); - expect(receivedPreviousRunId).toBe("run_previous"); } finally { await harness.close(); } @@ -257,13 +301,14 @@ describe("chat.customAgent clientData validation", () => { } }); - it("delivers frames that arrived before chat.messages.on is removed", async () => { + it("does not deliver frames or validation callbacks after chat.messages.on is removed", async () => { const clientData = { blocked: false }; const parserStarted = deferred(); const releaseParser = deferred(); - const delivered = deferred(); + const parserFinished = deferred(); let removeSubscription: (() => void) | undefined; let handlerCalls = 0; + let validationErrorCalls = 0; let started = false; const agent = chat @@ -273,26 +318,26 @@ describe("chat.customAgent clientData validation", () => { if (blocked) { parserStarted.resolve(); await releaseParser.promise; + parserFinished.resolve(); + throw new Error("invalid after unsubscribe"); } return { blocked }; }, }) .customAgent({ id: "custom-agent-client-data-off-after-arrival", + onClientDataValidationError: () => { + validationErrorCalls++; + }, run: async (_payload, { signal }) => { started = true; const subscription = chat.messages.on(async () => { handlerCalls++; - await chat.writeTurnComplete(); - delivered.resolve(); }); removeSubscription = () => subscription.off(); - await Promise.race([ - delivered.promise, - new Promise((resolve) => { - signal.addEventListener("abort", () => resolve(), { once: true }); - }), - ]); + await new Promise((resolve) => { + signal.addEventListener("abort", () => resolve(), { once: true }); + }); }, }); @@ -304,21 +349,68 @@ describe("chat.customAgent clientData validation", () => { try { await waitFor(() => started); clientData.blocked = true; - const send = harness.sendMessage(userMessage("hello", "message-1")); + void harness.sendMessage(userMessage("hello", "message-1")); await parserStarted.promise; removeSubscription!(); releaseParser.resolve(); - await send; - await delivered.promise; - expect(handlerCalls).toBe(1); + await parserFinished.promise; + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(handlerCalls).toBe(0); + expect(validationErrorCalls).toBe(0); } finally { releaseParser.resolve(); await harness.close(); } }); + it("throws from chat.messages.peek when an object parser returns a promise", async () => { + const clientData = { userId: "user_123" }; + let started = false; + let peekError: unknown; + + const agent = chat + .withClientData({ + schema: { + parse: async (value: unknown) => value as { userId: string }, + } as any, + }) + .customAgent({ + id: "custom-agent-client-data-async-object-peek", + run: async (_payload, { signal }) => { + started = true; + while (!signal.aborted) { + try { + chat.messages.peek(); + } catch (error) { + peekError = error; + await chat.writeTurnComplete(); + return; + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-async-object-peek-chat", + clientData, + }); + + try { + await waitFor(() => started); + const send = harness.sendMessage(userMessage("hello", "message-1")); + await waitFor(() => peekError !== undefined); + await send; + + expect(peekError).toBeInstanceOf(Error); + expect((peekError as Error).message).toContain("asynchronous schema"); + } finally { + await harness.close(); + } + }); + it("does not complete an active turn when a buffered frame is invalid", async () => { const clientData: { attempt: unknown } = { attempt: "1" }; const firstTurnStarted = deferred(); @@ -380,6 +472,169 @@ describe("chat.customAgent clientData validation", () => { } }); + it("buffers a steering frame whose validation finishes after the turn closes", async () => { + const clientData = { sequence: 0 }; + const parserStarted = deferred(); + const releaseParser = deferred(); + const firstTurnStarted = deferred(); + const releaseFirstTurn = deferred(); + const firstDoneStarted = deferred(); + const secondTurnFinished = deferred(); + const receivedSequences: number[] = []; + const receivedMessageIds: string[][] = []; + let started = false; + + const agent = chat + .withClientData({ + schema: async (value: unknown) => { + const sequence = (value as { sequence: number }).sequence; + if (sequence === 2) { + parserStarted.resolve(); + await releaseParser.promise; + } + return { sequence }; + }, + }) + .customAgent({ + id: "custom-agent-client-data-late-steering-validation", + run: async (payload, { signal }) => { + started = true; + const session = chat.createSession(payload, { + signal, + idleTimeoutInSeconds: 1, + pendingMessages: {}, + }); + for await (const turn of session) { + receivedSequences.push(turn.clientData.sequence); + receivedMessageIds.push(turn.uiMessages.map((message) => message.id)); + if (turn.number === 0) { + firstTurnStarted.resolve(); + await releaseFirstTurn.promise; + firstDoneStarted.resolve(); + await turn.done(); + continue; + } + await turn.done(); + secondTurnFinished.resolve(); + break; + } + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-late-steering-validation-chat", + clientData, + }); + + try { + await waitFor(() => started); + clientData.sequence = 1; + const first = harness.sendMessage(userMessage("first", "message-1")); + await firstTurnStarted.promise; + + clientData.sequence = 2; + void harness.sendMessage(userMessage("second", "message-2")); + await parserStarted.promise; + + releaseFirstTurn.resolve(); + await firstDoneStarted.promise; + await Promise.resolve(); + releaseParser.resolve(); + + await first; + await secondTurnFinished.promise; + expect(receivedSequences).toEqual([1, 2]); + expect(receivedMessageIds).toEqual([["message-1"], ["message-1", "message-2"]]); + } finally { + releaseFirstTurn.resolve(); + releaseParser.resolve(); + await harness.close(); + } + }); + + it("does not reparse an invalid steering frame after the turn closes", async () => { + const clientData = { sequence: 0 }; + const parserStarted = deferred(); + const releaseParser = deferred(); + const firstTurnStarted = deferred(); + const releaseFirstTurn = deferred(); + const firstDoneStarted = deferred(); + const validationErrors: unknown[] = []; + const receivedSequences: number[] = []; + let lateFrameParseCalls = 0; + let started = false; + + const agent = chat + .withClientData({ + schema: async (value: unknown) => { + const sequence = (value as { sequence: number }).sequence; + if (sequence === 2) { + lateFrameParseCalls++; + parserStarted.resolve(); + await releaseParser.promise; + if (lateFrameParseCalls === 1) { + throw new Error("invalid late frame"); + } + } + return { sequence }; + }, + }) + .customAgent({ + id: "custom-agent-client-data-late-invalid-steering", + onClientDataValidationError: ({ error }) => { + validationErrors.push(error); + }, + run: async (payload, { signal }) => { + started = true; + const session = chat.createSession(payload, { + signal, + idleTimeoutInSeconds: 1, + pendingMessages: {}, + }); + for await (const turn of session) { + receivedSequences.push(turn.clientData.sequence); + firstTurnStarted.resolve(); + await releaseFirstTurn.promise; + firstDoneStarted.resolve(); + await turn.done(); + } + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-late-invalid-steering-chat", + clientData, + }); + + try { + await waitFor(() => started); + clientData.sequence = 1; + const first = harness.sendMessage(userMessage("first", "message-1")); + await firstTurnStarted.promise; + + clientData.sequence = 2; + void harness.sendMessage(userMessage("second", "message-2")); + await parserStarted.promise; + + releaseFirstTurn.resolve(); + await firstDoneStarted.promise; + await Promise.resolve(); + releaseParser.resolve(); + + await first; + await waitFor(() => validationErrors.length === 1); + expect(lateFrameParseCalls).toBe(1); + expect(receivedSequences).toEqual([1]); + expect(harness.allChunks).toContainEqual( + expect.objectContaining({ type: "error", errorText: "invalid late frame" }) + ); + } finally { + releaseFirstTurn.resolve(); + releaseParser.resolve(); + await harness.close(); + } + }); + it("reports invalid chat.messages.on frames without calling the subscriber", async () => { const clientData: { userId: unknown } = { userId: "user_123" }; const validationErrors: unknown[] = []; @@ -423,10 +678,14 @@ describe("chat.customAgent clientData validation", () => { it("exits without a turn when a handover-prepare boot has invalid clientData and the warm handler skips", async () => { const clientData: { userId: unknown } = { userId: 123 }; + const validationErrors: unknown[] = []; let runCalls = 0; const agent = chat.withClientData({ schema: z.object({ userId: z.string() }) }).customAgent({ id: "custom-agent-client-data-handover-skip", + onClientDataValidationError: ({ error }) => { + validationErrors.push(error); + }, run: async () => { runCalls++; await chat.writeTurnComplete(); @@ -440,12 +699,11 @@ describe("chat.customAgent clientData validation", () => { }); try { - await waitFor(() => - harness.allChunks.some((chunk) => (chunk as { type?: string }).type === "error") - ); + await waitFor(() => validationErrors.length === 1); expect(runCalls).toBe(0); + expect(harness.allRawChunks).toHaveLength(0); - // The recovery path must drain the skip via the handover facade and + // The validation path must drain the skip via the handover facade and // end the run, mirroring the normal handover-skip exit. await harness.sendHandoverSkip(); @@ -461,53 +719,46 @@ describe("chat.customAgent clientData validation", () => { } }); - it("drops the head-start partial and recovers on the next valid frame when a handover-prepare boot has invalid clientData", async () => { + it("fails an invalid handover boot after the warm handler signals", async () => { const clientData: { userId: unknown } = { userId: 123 }; + const validationErrors: unknown[] = []; let runCalls = 0; - let receivedTrigger: string | undefined; - let receivedClientData: unknown; const agent = chat.withClientData({ schema: z.object({ userId: z.string() }) }).customAgent({ - id: "custom-agent-client-data-handover-drop", - run: async (payload) => { + id: "custom-agent-client-data-handover-invalid", + onClientDataValidationError: ({ error }) => { + validationErrors.push(error); + }, + run: async () => { runCalls++; - receivedTrigger = payload.trigger; - receivedClientData = payload.metadata; await chat.writeTurnComplete(); }, }); const harness = mockChatAgent(agent, { - chatId: "custom-agent-client-data-handover-drop-chat", + chatId: "custom-agent-client-data-handover-invalid-chat", mode: "handover-prepare", clientData, }); try { - await waitFor(() => - harness.allChunks.some((chunk) => (chunk as { type?: string }).type === "error") - ); + await waitFor(() => validationErrors.length === 1); expect(runCalls).toBe(0); + expect(harness.allRawChunks).toHaveLength(0); - // Resolves on the next turn-complete — the recovered message turn below. - const handover = harness.sendHandover({ + const handover = await harness.sendHandover({ partialAssistantMessage: [ { role: "assistant", content: [{ type: "text", text: "warm partial" }] }, ], }); - // Let the recovery drain consume the handover signal before the - // message frame goes out — a frame arriving mid-drain would be - // discarded by the handover facade (same as the pre-existing turn-0 - // handover wait in chat.createSession). - await new Promise((resolve) => setTimeout(resolve, 50)); - clientData.userId = "user_123"; - await harness.sendMessage(userMessage("retry", "message-1")); - await handover; - - expect(runCalls).toBe(1); - expect(receivedTrigger).toBe("submit-message"); - expect(receivedClientData).toEqual({ userId: "user_123" }); + expect(runCalls).toBe(0); + expect(handover.chunks).toEqual([ + expect.objectContaining({ type: "error", errorText: expect.any(String) }), + ]); + expect(handover.rawChunks).toContainEqual( + expect.objectContaining({ type: "trigger:turn-complete" }) + ); } finally { await harness.close(); } From 881cd8e4899a7f77dd822febdc25c401f336f43e Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Mon, 17 Aug 2026 14:28:12 -0700 Subject: [PATCH 4/4] fix(chat): narrow custom agent validation contract --- docs/ai-chat/client-protocol.mdx | 4 +- docs/ai-chat/custom-agents.mdx | 12 ++- docs/ai-chat/reference.mdx | 2 +- packages/trigger-sdk/src/v3/ai.ts | 19 +++-- packages/trigger-sdk/src/v3/chat.ts | 5 +- ...ustom-agent-client-data-validation.test.ts | 80 +++++++++++++++++-- 6 files changed, 103 insertions(+), 19 deletions(-) diff --git a/docs/ai-chat/client-protocol.mdx b/docs/ai-chat/client-protocol.mdx index 3aed294e3a..981559d7a5 100644 --- a/docs/ai-chat/client-protocol.mdx +++ b/docs/ai-chat/client-protocol.mdx @@ -832,7 +832,9 @@ Custom actions (undo, rollback, edit) ride on the same `.in` channel using `kind } ``` -Actions wake the agent from suspension (same as messages) and fire the `onAction` hook — they are not turns, so `run()` and turn lifecycle hooks do not fire. If `onAction` returns a `StreamTextResult`, the response is auto-piped to the frontend (but still no `run()` or `onTurnComplete`). The `action` payload is validated against the agent's `actionSchema`. If the agent didn't register an `actionSchema` (or your `action` payload doesn't match it), validation fails the same way `metadata` does — `.in/append` returns `200 OK`, but the run trace shows `chat turn N [ERROR]` and the wire emits a `turn-complete` control record with no other chunks. See [Actions](/ai-chat/actions) for the agent-side schema setup. +For managed `chat.agent()` tasks, actions wake the agent from suspension (same as messages) and fire the `onAction` hook — they are not turns, so `run()` and turn lifecycle hooks do not fire. If `onAction` returns a `StreamTextResult`, the response is auto-piped to the frontend (but still no `run()` or `onTurnComplete`). The `action` payload is validated against the agent's `actionSchema`. If the agent didn't register an `actionSchema` (or your `action` payload doesn't match it), validation fails the same way `metadata` does — `.in/append` returns `200 OK`, but the run trace shows `chat turn N [ERROR]` and the wire emits a `turn-complete` control record with no other chunks. See [Actions](/ai-chat/actions) for the agent-side schema setup. + +Raw `chat.customAgent()` tasks receive `action` as `unknown` and must validate it in their own loop. ### Regenerating the last response diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 6f962c4b55..3bea5788a9 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -21,13 +21,19 @@ Inside the wrapper, pick one of two loop styles: ### Validating client data -Use `chat.withClientData({ schema })` to validate `payload.metadata`. Custom agents parse the initial payload and later message and action frames before passing them to `run`, `chat.messages`, or `chat.createSession`. Schema defaults and transforms are included in the value your code receives. Close frames are not validated. +Use `chat.withClientData({ schema })` to validate `payload.metadata`. Custom agents parse the metadata on the initial payload and every later non-close input frame before passing it to `run`, `chat.messages`, or `chat.createSession`. Schema defaults and transforms are included in the value your code receives. -If validation fails for a submitted turn or an async read such as `wait()`, the SDK writes an error chunk followed by `turn-complete` and skips the invalid frame. The invalid value never reaches your run handler or turn loop. The task then waits for the next valid frame. A messageless preload or continuation boot has no submitted turn to complete, so the SDK reports the error through the task log and `onClientDataValidationError` while it waits. +This only validates `metadata`. A raw custom agent does not expose an action schema, so `payload.action` remains `unknown`. Validate the full frame or action payload in your own loop when you need that boundary. + +If validation fails for a submitted turn or an async read such as `wait()`, the SDK consumes and skips the invalid frame, writes an `Invalid client data` error followed by `turn-complete`, then waits for the next valid frame. The invalid value is not returned to the raw caller. The detailed validator error is available in the task log and `onClientDataValidationError`, but it is not sent to the client. + +This convenience path settles the invalid input before the read returns. If your raw loop needs to coordinate validation with persistence or settlement, omit `withClientData({ schema })` and validate the full wire frame in the loop instead. A messageless preload or continuation boot has no submitted turn to complete, so the SDK reports the error through the task log and callback while it waits. An invalid [head-start handover](/ai-chat/fast-starts#handover-with-custom-agents) boot fails closed. The SDK waits for the warm handler to finish so stream ordering stays intact. A handover skip ends the run. A real handover writes the validation error and `turn-complete` after the warm output, then ends the run. Without a schema, metadata is passed through unchanged. -`chat.messages.on()` is different because a subscribed frame can arrive while the current response is still streaming. Ending the turn at that point would cut off the response. While the subscription is active, the SDK skips the frame, logs the validation error, and calls `onClientDataValidationError` if you set it. This also applies to the steering subscription created by `chat.createSession({ pendingMessages })`. Calling `off()` prevents queued validation from invoking your handler or error callback. +`chat.messages.on()` is different because a subscribed frame can arrive while the current response is still streaming. Ending the turn at that point would cut off the response. While the subscription is active, the SDK skips an invalid frame, logs the validation error, and calls `onClientDataValidationError` if you set it. This also applies to the steering subscription created by `chat.createSession({ pendingMessages })`. + +Calling `off()` stops the subscription from accepting new frames. A valid frame accepted before `off()` still finishes validation and is delivered to the handler. An invalid frame that finishes validation after `off()` is logged without calling the handler or error callback. ```ts import { chat } from "@trigger.dev/sdk/ai"; diff --git a/docs/ai-chat/reference.mdx b/docs/ai-chat/reference.mdx index 6960e8161d..a3afd7d971 100644 --- a/docs/ai-chat/reference.mdx +++ b/docs/ai-chat/reference.mdx @@ -786,7 +786,7 @@ Send a custom action to the agent. Actions wake the agent from suspension and fi transport.sendAction(chatId: string, action: unknown): Promise> ``` -The action payload is validated against the agent's `actionSchema` on the backend. +For managed `chat.agent()` tasks, the action payload is validated against the agent's `actionSchema` on the backend. Raw `chat.customAgent()` tasks receive it as `unknown` and must validate it themselves. ```tsx // Undo button diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 1cbfaca247..09afc3f7ea 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1441,6 +1441,8 @@ type ChatCustomAgentClientDataErrorHandler = (event: { payload: ChatTaskWirePayload; }) => Promise | void; +const CHAT_CUSTOM_AGENT_CLIENT_DATA_ERROR_TEXT = "Invalid client data"; + const chatCustomAgentClientDataParserKey = locals.create( "chat.customAgentClientDataParser" ); @@ -1503,13 +1505,14 @@ function getChatCustomAgentSyncSchemaParseFn(schema: TaskSchema): (value: unknow } async function writeChatCustomAgentClientDataErrorToStream( - payload: ChatTaskWirePayload, - error: unknown + payload: ChatTaskWirePayload ): Promise { - const errorText = error instanceof Error ? error.message : "An unexpected error occurred"; try { await withChatWriter((writer) => { - writer.write({ type: "error", errorText } as any); + writer.write({ + type: "error", + errorText: CHAT_CUSTOM_AGENT_CLIENT_DATA_ERROR_TEXT, + } as any); }); await chatWriteTurnComplete(); } catch (signalError) { @@ -1552,7 +1555,7 @@ async function reportChatCustomAgentClientDataError( if (!options.writeToStream) { return; } - await writeChatCustomAgentClientDataErrorToStream(payload, error); + await writeChatCustomAgentClientDataErrorToStream(payload); } type ChatCustomAgentPayloadValidationResult = @@ -1796,7 +1799,8 @@ const messagesInput: RealtimeDefinedInputStream = { return subscribeToRawChatMessages(handler); } - return subscribeToValidatedChatMessages((payload) => handler(payload)); + const deliver = (payload: ChatTaskWirePayload) => handler(payload); + return subscribeToValidatedChatMessages(deliver, { onAfterOff: deliver }); }, once(options) { const ctx = taskContext.ctx; @@ -5594,6 +5598,7 @@ type ChatCustomAgentOptions< * error chunk followed by `turn-complete`. Messageless boots and active * subscriptions use `onClientDataValidationError` and the task log because * there is no submitted turn to complete or a response may still be streaming. + * This validates `metadata` only; raw `action` payloads remain `unknown`. */ clientDataSchema?: TClientDataSchema; /** @@ -5719,7 +5724,7 @@ function chatCustomAgent< // The head-start writer flushes before sending this signal. Writing // the terminal error now preserves stream order and closes the stitch. - await writeChatCustomAgentClientDataErrorToStream(payload, validated.error); + await writeChatCustomAgentClientDataErrorToStream(payload); return; } diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index a7c7125575..39c91eeb7f 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -95,7 +95,10 @@ export type ChatTaskWirePayload { }; let started = false; const receivedClientData: unknown[] = []; + const validationErrors: unknown[] = []; const agent = chat .withClientData({ @@ -100,6 +101,9 @@ describe("chat.customAgent clientData validation", () => { }) .customAgent({ id: "custom-agent-client-data-invalid-frame", + onClientDataValidationError: ({ error }) => { + validationErrors.push(error); + }, run: async (payload, { signal }) => { started = true; const session = chat.createSession(payload, { @@ -126,8 +130,10 @@ describe("chat.customAgent clientData validation", () => { expect(receivedClientData).toHaveLength(0); expect(invalidTurn.chunks).toEqual([ - expect.objectContaining({ type: "error", errorText: expect.any(String) }), + expect.objectContaining({ type: "error", errorText: "Invalid client data" }), ]); + expect(validationErrors).toHaveLength(1); + expect(validationErrors[0]).toBeInstanceOf(z.ZodError); expect(invalidTurn.rawChunks).toContainEqual( expect.objectContaining({ type: "trigger:turn-complete" }) ); @@ -229,7 +235,7 @@ describe("chat.customAgent clientData validation", () => { expect(runCalls).toBe(0); expect(harness.allChunks).toEqual([ - expect.objectContaining({ type: "error", errorText: expect.any(String) }), + expect.objectContaining({ type: "error", errorText: "Invalid client data" }), ]); clientData.userId = "user_123"; @@ -301,7 +307,7 @@ describe("chat.customAgent clientData validation", () => { } }); - it("does not deliver frames or validation callbacks after chat.messages.on is removed", async () => { + it("does not report an invalid frame whose validation finishes after chat.messages.on is removed", async () => { const clientData = { blocked: false }; const parserStarted = deferred(); const releaseParser = deferred(); @@ -365,6 +371,68 @@ describe("chat.customAgent clientData validation", () => { } }); + it("delivers a valid frame accepted before chat.messages.on is removed", async () => { + const clientData = { blocked: false }; + const parserStarted = deferred(); + const releaseParser = deferred(); + const delivered = deferred(); + let removeSubscription: (() => void) | undefined; + let receivedMetadata: unknown; + let handlerCalls = 0; + let started = false; + + const agent = chat + .withClientData({ + schema: async (value: unknown) => { + const blocked = (value as { blocked: boolean }).blocked; + if (blocked) { + parserStarted.resolve(); + await releaseParser.promise; + } + return { blocked, parsed: true as const }; + }, + }) + .customAgent({ + id: "custom-agent-client-data-deliver-pending-after-off", + run: async (_payload, { signal }) => { + started = true; + const subscription = chat.messages.on(async (payload) => { + handlerCalls++; + receivedMetadata = payload.metadata; + await chat.writeTurnComplete(); + delivered.resolve(); + }); + removeSubscription = () => subscription.off(); + await new Promise((resolve) => { + signal.addEventListener("abort", () => resolve(), { once: true }); + }); + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-deliver-pending-after-off-chat", + clientData, + }); + + try { + await waitFor(() => started); + clientData.blocked = true; + const send = harness.sendMessage(userMessage("hello", "message-1")); + await parserStarted.promise; + + removeSubscription!(); + releaseParser.resolve(); + + await send; + await delivered.promise; + expect(handlerCalls).toBe(1); + expect(receivedMetadata).toEqual({ blocked: true, parsed: true }); + } finally { + releaseParser.resolve(); + await harness.close(); + } + }); + it("throws from chat.messages.peek when an object parser returns a promise", async () => { const clientData = { userId: "user_123" }; let started = false; @@ -464,7 +532,7 @@ describe("chat.customAgent clientData validation", () => { expect(receivedClientData).toEqual([{ attempt: 1 }]); expect(harness.allChunks).toContainEqual( - expect.objectContaining({ type: "error", errorText: expect.any(String) }) + expect.objectContaining({ type: "error", errorText: "Invalid client data" }) ); } finally { releaseFirstTurn.resolve(); @@ -626,7 +694,7 @@ describe("chat.customAgent clientData validation", () => { expect(lateFrameParseCalls).toBe(1); expect(receivedSequences).toEqual([1]); expect(harness.allChunks).toContainEqual( - expect.objectContaining({ type: "error", errorText: "invalid late frame" }) + expect.objectContaining({ type: "error", errorText: "Invalid client data" }) ); } finally { releaseFirstTurn.resolve(); @@ -754,7 +822,7 @@ describe("chat.customAgent clientData validation", () => { expect(runCalls).toBe(0); expect(handover.chunks).toEqual([ - expect.objectContaining({ type: "error", errorText: expect.any(String) }), + expect.objectContaining({ type: "error", errorText: "Invalid client data" }), ]); expect(handover.rawChunks).toContainEqual( expect.objectContaining({ type: "trigger:turn-complete" })