From 50c86cdc1ed713ac2dcdbb697ba21be3eb172a5f Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Sun, 16 Aug 2026 20:55:35 -0700 Subject: [PATCH 01/11] 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 00000000000..9eba1990caa --- /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 d039b39366a..3aed294e3a2 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 197bff6b5e1..e9f149457e1 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 da08f9a0473..6960e8161d6 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 00000000000..b432ca55f9a --- /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 02/11] 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 e9f149457e1..71e2f624ad8 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 9c88c4ab35f..03405f1a4ee 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 b432ca55f9a..249ab8a8542 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 03/11] 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 71e2f624ad8..6f962c4b555 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 03405f1a4ee..1cbfaca2471 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 249ab8a8542..492be61be89 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 04/11] 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 3aed294e3a2..981559d7a5a 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 6f962c4b555..3bea5788a90 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 6960e8161d6..a3afd7d971e 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 1cbfaca2471..09afc3f7eae 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 a7c7125575e..39c91eeb7f7 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" }) From 19585dc450191399b1fe794e3be1b80e0b5ab5ac Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 27 Aug 2026 16:02:28 +0100 Subject: [PATCH 05/11] refactor(chat): scope the clientData validation options to withClientData The error-timing knob sat on the customAgent config as `clientDataValidationErrorTiming`, which needed the prefix to say what it applied to. Moved inside `withClientData`, where the schema already establishes the subject, it is just `reportErrorAt`. The validation callback gets the same treatment as `onValidationError`, composing with the task-level hook the way the builder's other hooks already do. The customAgent config keeps `clientDataSchema` and `onClientDataValidationError` for callers that configure it directly rather than through the builder, and the internal field is renamed to match. --- docs/ai-chat/custom-agents.mdx | 12 ++++- packages/trigger-sdk/src/v3/ai.ts | 53 ++++++++++++++++--- ...tom-agent-client-data-error-timing.test.ts | 8 +-- 3 files changed, 60 insertions(+), 13 deletions(-) diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 9e7d9351788..161cb2b04e4 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -33,7 +33,17 @@ An invalid [head-start handover](/ai-chat/fast-starts#handover-with-custom-agent `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 })`. -By default the client-visible `Invalid client data` error for such a frame is held until the turn ends, so a bad send cannot truncate an answer the user is already reading. Set `clientDataValidationErrorTiming: "arrival"` if you would rather surface it as soon as validation fails, accepting that it ends the response in progress. `onClientDataValidationError` and the task log fire on arrival in both modes, and the frame is never delivered as a turn either way. +By default the client-visible `Invalid client data` error for such a frame is held until the turn ends, so a bad send cannot truncate an answer the user is already reading. Pass `reportErrorAt: "arrival"` to `withClientData` if you would rather surface it as soon as validation fails, accepting that it ends the response in progress: + +```ts +chat.withClientData({ + schema: z.object({ userId: z.string() }), + reportErrorAt: "arrival", + onValidationError: ({ error, payload }) => logger.warn("bad client data", { error }), +}); +``` + +The validation callback and the task log fire on arrival in both modes, and the frame is never delivered as a turn either way. 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. diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 23196afe865..40be90b85b0 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -6116,7 +6116,7 @@ type ChatCustomAgentOptions< * this only governs the stream-visible error. The frame is never delivered as * a turn in either mode. */ - clientDataValidationErrorTiming?: "turn-end" | "arrival"; + clientDataReportErrorAt?: "turn-end" | "arrival"; run: TaskOptions< TIdentifier, ChatTaskWirePayload>, @@ -6134,7 +6134,7 @@ function chatCustomAgent< const { clientDataSchema, onClientDataValidationError, - clientDataValidationErrorTiming, + clientDataReportErrorAt, run: userRun, ...restOptions } = options; @@ -6181,10 +6181,7 @@ function chatCustomAgent< onClientDataValidationError as ChatCustomAgentClientDataErrorHandler ); } - locals.set( - chatCustomAgentClientDataErrorTimingKey, - clientDataValidationErrorTiming ?? "turn-end" - ); + locals.set(chatCustomAgentClientDataErrorTimingKey, clientDataReportErrorAt ?? "turn-end"); // 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` @@ -9014,9 +9011,29 @@ export interface ChatBuilder< config?: ChatWithUIMessageConfig ): ChatBuilder; - /** Fix the client data schema. Returns a new builder preserving all accumulated state. */ + /** + * Fix the client data schema, and how validation failures are handled. + * Returns a new builder preserving all accumulated state. + */ withClientData(config: { schema: TSchema; + /** + * When a frame that arrived mid-turn fails validation, decides when the + * client-visible error is written. + * + * `"turn-end"` (default) waits for the turn to close, so a bad send cannot + * truncate an answer already being read. `"arrival"` writes it as soon as + * validation fails, ending the response in progress. + * + * `onValidationError` and the task log fire on arrival either way, and the + * frame is never delivered as a turn. + */ + reportErrorAt?: "turn-end" | "arrival"; + /** Called when an input fails validation. Composes with the task-level hook. */ + onValidationError?: (event: { + error: unknown; + payload: ChatTaskWirePayload; + }) => Promise | void; }): ChatBuilder; /** Register a builder-level `onBoot` hook. Runs before the task-level hook if both are set. */ @@ -9140,6 +9157,8 @@ type ChatBuilderHooks = { type ChatBuilderConfig = { uiStreamOptions?: ChatUIMessageStreamOptions; clientDataSchema?: TaskSchema; + clientDataReportErrorAt?: "turn-end" | "arrival"; + clientDataOnValidationError?: ChatCustomAgentClientDataErrorHandler; hooks: ChatBuilderHooks; }; @@ -9167,10 +9186,17 @@ function createChatBuilder< }); }, - withClientData(cdConfig: { schema: TSchema }) { + withClientData(cdConfig: { + schema: TSchema; + reportErrorAt?: "turn-end" | "arrival"; + onValidationError?: ChatCustomAgentClientDataErrorHandler; + }) { return createChatBuilder({ ...config, clientDataSchema: cdConfig.schema, + clientDataReportErrorAt: cdConfig.reportErrorAt ?? config.clientDataReportErrorAt, + clientDataOnValidationError: + cdConfig.onValidationError ?? config.clientDataOnValidationError, }); }, @@ -9283,6 +9309,13 @@ function createChatBuilder< return chatCustomAgent({ ...options, ...(config.clientDataSchema ? { clientDataSchema: config.clientDataSchema } : {}), + ...(config.clientDataReportErrorAt + ? { clientDataReportErrorAt: config.clientDataReportErrorAt } + : {}), + onClientDataValidationError: composeHooks( + config.clientDataOnValidationError, + options.onClientDataValidationError + ), }); }, } as unknown as ChatBuilder; @@ -9338,9 +9371,13 @@ function withUIMessage( */ function withClientData(config: { schema: TSchema; + reportErrorAt?: "turn-end" | "arrival"; + onValidationError?: ChatCustomAgentClientDataErrorHandler; }): ChatBuilder { return createChatBuilder({ clientDataSchema: config.schema, + clientDataReportErrorAt: config.reportErrorAt, + clientDataOnValidationError: config.onValidationError, hooks: {}, }); } diff --git a/packages/trigger-sdk/test/custom-agent-client-data-error-timing.test.ts b/packages/trigger-sdk/test/custom-agent-client-data-error-timing.test.ts index 4f51b56e32f..8ae66d51a7a 100644 --- a/packages/trigger-sdk/test/custom-agent-client-data-error-timing.test.ts +++ b/packages/trigger-sdk/test/custom-agent-client-data-error-timing.test.ts @@ -56,13 +56,13 @@ describe("chat.customAgent clientDataValidationErrorTiming", () => { } return { sequence }; }, + ...(timing ? { reportErrorAt: timing } : {}), + onValidationError: ({ error }) => { + validationErrors.push(error); + }, }) .customAgent({ id: `custom-agent-client-data-timing-${timing ?? "default"}`, - ...(timing ? { clientDataValidationErrorTiming: timing } : {}), - onClientDataValidationError: ({ error }) => { - validationErrors.push(error); - }, run: async (payload, { signal }) => { started = true; const session = chat.createSession(payload, { From d40b8583bc81cceb5452bfd11a8b8d0d6700c7b0 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 27 Aug 2026 17:40:09 +0100 Subject: [PATCH 06/11] refactor(sdk): drop the unused drain hook from chat message subscriptions The drain() exposed on ChatMessageSubscription had no call sites. off() detaches from the router but does not cancel work already chained, so the parse chain runs to completion on its own and an invalid frame is still reported. Nothing has to wait on it. Documents the one edge that is genuinely uncovered instead: a parse still in flight when the task itself returns, where teardown can cut the report short. That wants a bounded wait if it is ever added, since the chain awaits a user-supplied schema. --- packages/trigger-sdk/src/v3/ai.ts | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 40be90b85b0..4943507d3dd 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1796,7 +1796,6 @@ async function waitOnChatRoute( type ChatMessageSubscription = { off: () => void; - drain?: () => Promise; }; /** @@ -1826,8 +1825,16 @@ function subscribeToRawChatMessages( * delivery of every later message, and only when a schema is declared. Raw * delivery has always been fire-and-forget, so awaiting here would also make * handler concurrency differ between the validated and unvalidated paths for no - * stated reason. `drain()` still covers the parse chain, which is what callers - * need before ending a run. + * stated reason. + * + * `off()` detaches from the router but does not cancel work already chained: + * the parse chain is a live promise chain and runs to completion on its own, so + * a frame that arrives just before a turn closes is still parsed and a failure + * is still reported (through the `!active` branch). Nothing therefore has to + * wait on it, which is why no `drain()` hook is exposed. The one uncovered edge + * is a parse still in flight when the task itself returns, where teardown can + * cut the report short; give this a bounded wait at the run-end boundary rather + * than an unbounded one, since the chain awaits a user-supplied schema. */ function subscribeToValidatedChatMessages( handler: (payload: ChatTaskWirePayload, isActive: () => boolean) => unknown, @@ -1874,7 +1881,6 @@ function subscribeToValidatedChatMessages( active = false; subscription.off(); }, - drain: () => delivery, }; } @@ -1958,7 +1964,6 @@ function observeValidatedChatMessages( active = false; subscription.off(); }, - drain: () => delivery, }; } From 7b0251017ffca7c5959c7fef5af97b34a0bdc09f Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 27 Aug 2026 17:40:09 +0100 Subject: [PATCH 07/11] docs(ai-chat): document the clientData validation error timing options The chat.withClientData reference still carried the older signature with only schema, so reportErrorAt and onValidationError were missing from the page that defines the function. clientDataReportErrorAt, the flat option on chat.customAgent, was absent from the docs entirely while its sibling callback was documented in six places. Adds both to the reference with types and defaults, states what reportErrorAt does and does not govern, names the flat equivalents for callers who pass options directly, and cross-links from the types page. --- docs/ai-chat/reference.mdx | 18 ++++++++++++++---- docs/ai-chat/types.mdx | 2 ++ 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/docs/ai-chat/reference.mdx b/docs/ai-chat/reference.mdx index 84791b5274e..eac6ecaefd6 100644 --- a/docs/ai-chat/reference.mdx +++ b/docs/ai-chat/reference.mdx @@ -550,15 +550,25 @@ Use this when you need [`InferChatUIMessage`](#inferchatuimessage) / typed `data 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; +chat.withClientData({ + schema: TSchema; + reportErrorAt?: "turn-end" | "arrival"; + onValidationError?: (event: { error: unknown; payload: ChatTaskWirePayload }) => Promise | void; +}): ChatBuilder; ``` -| Parameter | Type | Description | -| --------- | ------------ | -------------------------------------------------- | -| `schema` | `TaskSchema` | Zod, ArkType, Valibot, or any supported schema lib | +| Parameter | Type | Default | Description | +| ------------------- | --------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------ | +| `schema` | `TaskSchema` | required | Zod, ArkType, Valibot, or any supported schema lib | +| `reportErrorAt` | `"turn-end" \| "arrival"` | `"turn-end"` | When the client-visible `Invalid client data` error is written for a frame that failed validation mid-turn | +| `onValidationError` | `(event) => void` | — | Called when an input fails validation. Composes with the task-level `onClientDataValidationError` rather than replacing it | + +`reportErrorAt` governs only the stream-visible error. `"turn-end"` holds it until the turn closes, so a bad send cannot truncate an answer the user is already reading; `"arrival"` writes it as soon as validation fails, ending the response in progress. `onValidationError` and the task log fire on arrival in both modes, and the frame is never delivered as a turn either way. For `chat.customAgent()`, invalid client data is skipped. Async reads emit an error chunk followed by `turn-complete`. A `chat.messages.on()` subscription uses the task's `onClientDataValidationError` callback and task log instead, so an active response is not ended early. Without a schema, metadata is passed through unchanged. +Passing options directly to `chat.customAgent()` instead of through the builder uses the flat equivalents: `clientDataSchema`, `clientDataReportErrorAt`, and `onClientDataValidationError`. + Full guide: [Typed client data](/ai-chat/types#typed-client-data-with-chatwithclientdata). ## `ChatWithUIMessageConfig` diff --git a/docs/ai-chat/types.mdx b/docs/ai-chat/types.mdx index 328a2c5e303..e0b6dcc873f 100644 --- a/docs/ai-chat/types.mdx +++ b/docs/ai-chat/types.mdx @@ -169,6 +169,8 @@ export const myChat = chat The schema runs at runtime for both `.agent()` and `.customAgent()`. Custom agents validate the initial payload and later `chat.messages` frames. Invalid frames are not passed to user code. Async reads emit an error chunk followed by `turn-complete`; `chat.messages.on()` reports through `onClientDataValidationError` and the task log so it does not end an active response. Without a schema, metadata is passed through unchanged. +`withClientData` also takes `reportErrorAt` and `onValidationError` alongside `schema`. See [chat.withClientData](/ai-chat/reference#chatwithclientdata) for both, and [Validating client data](/ai-chat/custom-agents#validating-client-data) for the custom-agent walkthrough. + ## ChatBuilder Both `chat.withUIMessage()` and `chat.withClientData()` return a **ChatBuilder** — a chainable object that accumulates configuration before creating the agent with `.agent()`. From c5e0f9f09543d67c77d3621b3c3a4baff1c9751d Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 27 Aug 2026 17:58:53 +0100 Subject: [PATCH 08/11] fix(sdk): correct three clientData validation faults found in review Steering under an async schema queued messages against the wrong record. The validating observer parses before handing a frame to the steering queue, so the handler runs after an await, and the sequence was tracked in one slot shared by every frame. With two frames in flight that slot already held the newer sequence when the older frame's parse resolved, so the queue entry named the wrong record. The drain takes records by sequence, so the first entry consumed the second frame's record, the second entry's take found nothing and was skipped, and the injected message's own record stayed on the channel to be answered again as a later turn. One message was lost and another answered twice. The sequence now travels with the record. chat.messages.next() returned the payload unparsed. Every sibling read validates, so a hand-rolled loop got raw metadata with no schema defaults or transforms, and because it consumes it could also take a record the observer still owned and parse it a second time. It now follows the same claim-and-validate path as the other reads. The default "turn-end" error timing never wrote its error. Suppressing the write while a turn was open dropped it rather than holding it, so a client whose frame was rejected mid-turn got no signal at all, since the callback and task log are server-side only. Failures are now held and written just before the turn-complete chunk, matching the error-then-turn-complete order the other paths use. Adds a regression test for the sequence fault, verified red against the shared-slot version, and asserts the deferred error is actually written rather than only that it is absent while the turn is open. --- packages/trigger-sdk/src/v3/ai.ts | 98 ++++++-- ...tom-agent-client-data-error-timing.test.ts | 5 + .../test/custom-agent-steering-seqnum.test.ts | 231 ++++++++++++++++++ 3 files changed, 311 insertions(+), 23 deletions(-) create mode 100644 packages/trigger-sdk/test/custom-agent-steering-seqnum.test.ts diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 4943507d3dd..d7f96a212f9 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1894,10 +1894,14 @@ function subscribeToValidatedChatMessages( * buffer of its own, because the record was never removed from the channel. */ function observeValidatedChatMessages( - handler: (payload: ChatTaskWirePayload, isActive: () => boolean) => unknown, - options: { - onSeqNum?: (seqNum: number) => void; - } = {} + /** + * Receives the observed record's own `seqNum`. It is passed per record rather + * than tracked outside, because the handler runs after an await: with two + * frames in flight a shared slot already holds the newer sequence by the time + * the older frame's parse resolves, and the steering queue would then take + * the wrong record off the channel. + */ + handler: (payload: ChatTaskWirePayload, seqNum: number, isActive: () => boolean) => unknown ): ChatMessageSubscription { let active = true; let delivery = Promise.resolve(); @@ -1908,7 +1912,7 @@ function observeValidatedChatMessages( const subscription = chatInputRouter().observe(CHAT_ROUTE_MESSAGES, (record) => { const payload = (record.data as Extract).payload; - options.onSeqNum?.(record.seqNum); + const seqNum = record.seqNum; /** * Claimed synchronously, before any await, so a read cannot validate the * same record. The promise lets a read wait for the outcome rather than @@ -1937,12 +1941,15 @@ function observeValidatedChatMessages( * a later one, so removing it loses nothing that was still owed. */ const timing = locals.get(chatCustomAgentClientDataErrorTimingKey) ?? "turn-end"; + const writeNow = timing === "arrival" || !active; await reportChatCustomAgentClientDataError(payload, result.error, { - // A terminal error written into a live response can close it, so by - // default it waits for the turn to end. `"arrival"` opts into the - // earlier, more disruptive report. - writeToStream: timing === "arrival" || !active, + writeToStream: writeNow, }); + if (!writeNow) { + const deferred = locals.get(chatCustomAgentDeferredClientDataErrorsKey) ?? []; + deferred.push(payload); + locals.set(chatCustomAgentDeferredClientDataErrorsKey, deferred); + } // Drop the record: an invalid frame is never answered, by this run or // a later one. Released so a waiting read stops waiting, finds it // gone, and goes back to waiting for the next message. @@ -1954,7 +1961,7 @@ function observeValidatedChatMessages( // Valid, so release. The record is still queued, and whichever comes // first, an injection or a later turn's read, now owns it. releaseClaim(); - void Promise.resolve(handler(result.payload, () => active)).catch(() => {}); + void Promise.resolve(handler(result.payload, seqNum, () => active)).catch(() => {}); }) .catch(() => {}); }); @@ -2045,13 +2052,31 @@ const messagesInput: ChatMessages = { ); } - const record = await chatInputRouter().next(CHAT_ROUTE_MESSAGES, { - timeoutMs: timeoutInSeconds === undefined ? undefined : timeoutInSeconds * 1000, - }); - if (!record) return undefined; + // Consuming read, so it takes the same claim-and-validate path as the other + // reads: a record the observer still owns is put back and awaited, and an + // invalid payload is reported and skipped rather than surfaced raw. + while (true) { + const record = await chatInputRouter().next(CHAT_ROUTE_MESSAGES, { + timeoutMs: timeoutInSeconds === undefined ? undefined : timeoutInSeconds * 1000, + }); + if (!record) return undefined; - const chunk = record.data as Extract; - return { id: record.id, seqNum: record.seqNum, payload: chunk.payload }; + const claim = observerClaimFor(record); + if (claim) { + chatInputRouter().untake(CHAT_ROUTE_MESSAGES, record); + await claim; + continue; + } + + const chunk = record.data as Extract; + if (!shouldValidateChatCustomAgentPayload(chunk.payload)) { + return { id: record.id, seqNum: record.seqNum, payload: chunk.payload }; + } + const validated = await validateChatCustomAgentPayload(chunk.payload); + if (validated.ok) { + return { id: record.id, seqNum: record.seqNum, payload: validated.payload }; + } + } }, wait(options) { return new ManualWaitpointPromise(async (resolve, reject) => { @@ -3212,6 +3237,17 @@ const chatAgentRunContextKey = locals.create("chat.agentRunConte const chatCustomAgentClientDataErrorTimingKey = locals.create<"turn-end" | "arrival">( "chat.customAgent.clientDataErrorTiming" ); +/** + * @internal Client-data errors held back until the open turn closes. + * + * The default timing exists so a bad send cannot truncate an answer already + * being read, which means the write has to happen later rather than not at all: + * the callback and the task log are server-side, so dropping the stream write + * would leave the client with no signal that its frame was rejected. + */ +const chatCustomAgentDeferredClientDataErrorsKey = locals.create( + "chat.customAgent.deferredClientDataErrors" +); /** * @internal Sequences the validating observer has claimed. * @@ -9888,6 +9924,7 @@ function createStopSignal(): { async function chatWriteTurnComplete(options?: { publicAccessToken?: string; }): Promise<{ lastEventId?: string; sessionInEventId?: string }> { + await flushDeferredChatCustomAgentClientDataErrors(); const result = await writeTurnCompleteChunk(undefined, options?.publicAccessToken); // Same cursor written to the `session-in-event-id` header inside // `writeTurnCompleteChunk`; surfaced here so the caller can persist it. @@ -9898,6 +9935,25 @@ async function chatWriteTurnComplete(options?: { }; } +/** + * Writes the client-data errors held back by the default `"turn-end"` timing. + * + * Ordered before the turn-complete chunk, matching the error-then-turn-complete + * shape the submitted-turn and async-read paths already write. + */ +async function flushDeferredChatCustomAgentClientDataErrors(): Promise { + const deferred = locals.get(chatCustomAgentDeferredClientDataErrorsKey); + if (!deferred || deferred.length === 0) return; + locals.set(chatCustomAgentDeferredClientDataErrorsKey, []); + for (const payload of deferred) { + try { + await writeChatCustomAgentClientDataErrorToStream(payload); + } catch { + /* non-fatal */ + } + } +} + /** * The outcome of a turn's stream, reported by {@link pipeChatAndCapture}. * @@ -10784,13 +10840,9 @@ function createChatSession( const sessionMsgSub: ChatMessageSubscription | undefined = sessionPendingMessages ? locals.get(chatCustomAgentClientDataParserKey) - ? (() => { - let lastSeqNum: number | undefined; - return observeValidatedChatMessages( - (msg, isActive) => handleSteeringMessage(msg, lastSeqNum, isActive), - { onSeqNum: (seqNum) => (lastSeqNum = seqNum) } - ); - })() + ? observeValidatedChatMessages((msg, seqNum, isActive) => + handleSteeringMessage(msg, seqNum, isActive) + ) : chatInputRouter().observe(CHAT_ROUTE_MESSAGES, (record) => { const msg = (record.data as Extract).payload; void handleSteeringMessage(msg, record.seqNum); diff --git a/packages/trigger-sdk/test/custom-agent-client-data-error-timing.test.ts b/packages/trigger-sdk/test/custom-agent-client-data-error-timing.test.ts index 8ae66d51a7a..3f0d09d9f61 100644 --- a/packages/trigger-sdk/test/custom-agent-client-data-error-timing.test.ts +++ b/packages/trigger-sdk/test/custom-agent-client-data-error-timing.test.ts @@ -100,6 +100,10 @@ describe("chat.customAgent clientDataValidationErrorTiming", () => { releaseFirstTurn.resolve(); await first; + // The default holds the write until the turn closes, so the assertion has + // to wait for it: without this, "no chunk while open" would also pass if + // the error were never written at all. + await waitFor(() => errorChunks(harness).length > 0, "deferred error written"); return { chunksWhileOpen, chunksAfter: errorChunks(harness).length, validationErrors }; } finally { releaseFirstTurn.resolve(); @@ -112,6 +116,7 @@ describe("chat.customAgent clientDataValidationErrorTiming", () => { // The handler always fires on arrival; only the stream write is held back. expect(result.validationErrors).toHaveLength(1); expect(result.chunksWhileOpen).toBe(0); + expect(result.chunksAfter).toBeGreaterThan(0); }); it('writes the terminal error immediately with "arrival"', { timeout: 30_000 }, async () => { diff --git a/packages/trigger-sdk/test/custom-agent-steering-seqnum.test.ts b/packages/trigger-sdk/test/custom-agent-steering-seqnum.test.ts new file mode 100644 index 00000000000..6c52c7ed3a8 --- /dev/null +++ b/packages/trigger-sdk/test/custom-agent-steering-seqnum.test.ts @@ -0,0 +1,231 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { sessionStreams } from "@trigger.dev/core/v3"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { simulateReadableStream, stepCountIs, streamText, tool } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { chat } from "../src/v3/ai.js"; + +/** + * The validating steering observer parses a frame before handing it to the + * steering queue, so the handler runs after an await. With two frames in + * flight, anything that tracks "the current sequence" outside the handler + * already holds the newer frame's number by the time the older frame's parse + * resolves, and the queue entry is built against the wrong record. + * + * That matters because `drainSteeringQueue` takes the record by `seqNum`. An + * entry carrying the wrong one removes a message nobody answered, and leaves + * the injected message's own record on the channel where a later turn answers + * it a second time. + * + * The scenario below therefore needs two frames observed before either parse + * resolves, which the gated async schema guarantees without depending on + * timing: mid-turn frames block on `parseGate`, while the boot frame + * (`sequence` 0) passes straight through so the run can start. Turn 1 takes two + * steps so it has a `prepareStep` boundary to inject at, held open by a tool + * gate; later turns answer in one step. + * + * The whole batch is injected, so the fix shows up twice over: with a shared + * sequence both entries carry the newer one, the first `take` consumes the + * second frame's record and the second `take` finds nothing, so only one + * message is injected, the other is lost outright, and the injected message's + * own record survives to be answered again as a second turn. + */ + +const USAGE = { + inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 1, text: 1, reasoning: undefined }, +}; + +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, label = "condition", timeoutMs = 8_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 10)); + } + throw new Error(`waitFor timed out: ${label}`); +} + +function lastUserText(prompt: { role: string; content: unknown }[]): string { + const users = prompt.filter((m) => m.role === "user"); + const last = users[users.length - 1]; + return Array.isArray(last?.content) + ? (last.content as { type: string; text?: string }[]) + .filter((p) => p.type === "text") + .map((p) => p.text ?? "") + .join("") + : ""; +} + +function textChunks(text: string): LanguageModelV3StreamPart[] { + return [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: text }, + { type: "text-end", id: "t1" }, + { type: "finish", finishReason: { unified: "stop", raw: "stop" }, usage: USAGE }, + ]; +} + +function toolCallChunks(callId: string): LanguageModelV3StreamPart[] { + return [ + { type: "tool-input-start", id: callId, toolName: "gate" }, + { type: "tool-input-delta", id: callId, delta: JSON.stringify({ q: "x" }) }, + { type: "tool-input-end", id: callId }, + { type: "tool-call", toolCallId: callId, toolName: "gate", input: JSON.stringify({ q: "x" }) }, + { type: "finish", finishReason: { unified: "tool-calls", raw: "tool_calls" }, usage: USAGE }, + ]; +} + +function gatedTwoStepModel(answers: string[]) { + let call = 0; + return new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + const isFirstTurnToolStep = call++ === 0; + const text = `ANSWER(${lastUserText(prompt)})`; + if (!isFirstTurnToolStep) answers.push(text); + return { + stream: simulateReadableStream({ + chunks: isFirstTurnToolStep ? toolCallChunks("tc-1") : textChunks(text), + initialDelayInMs: 10, + chunkDelayInMs: 2, + }), + }; + }, + }); +} + +type SeqReader = { lastSeqNum(sessionId: string, io: "in" | "out"): number | undefined }; + +/** Appends a message and resolves once the channel has actually taken it. */ +async function sendAndLand( + harness: { sendMessage: (m: ReturnType) => Promise }, + chatId: string, + text: string, + id: string +) { + const seqs = sessionStreams as unknown as SeqReader; + const before = seqs.lastSeqNum(chatId, "in") ?? -1; + void harness.sendMessage(userMessage(text, id)); + await waitFor(() => (seqs.lastSeqNum(chatId, "in") ?? -1) > before, `append ${id}`); +} + +describe("chat.customAgent steering under async client-data validation", () => { + it( + "takes the record it injected, so the declined message still gets its own turn", + { timeout: 30_000 }, + async () => { + const chatId = "steering-seqnum-chat"; + const clientData = { sequence: 0 }; + const parseGate = deferred(); + const toolGate = deferred(); + const answers: string[] = []; + const injected: string[] = []; + const received: string[] = []; + let toolEntered = false; + let turnCount = 0; + + const gateTool = tool({ + description: "blocks until the test opens it", + inputSchema: z.object({ q: z.string() }), + execute: async () => { + toolEntered = true; + await toolGate.promise; + return "ok"; + }, + }); + + const model = gatedTwoStepModel(answers); + + const agent = chat + .withClientData({ + schema: async (value: unknown) => { + const sequence = (value as { sequence: number }).sequence; + if (sequence > 0) await parseGate.promise; + return { sequence }; + }, + }) + .customAgent({ + id: "custom-agent-steering-seqnum", + run: async (payload, { signal }) => { + const session = chat.createSession(payload, { + signal, + idleTimeoutInSeconds: 1, + pendingMessages: { + shouldInject: () => true, + onReceived: ({ message }) => { + received.push(message.id); + }, + onInjected: ({ messages }) => { + injected.push(...messages.map((m) => m.id)); + }, + }, + }); + + for await (const turn of session) { + turnCount++; + await turn.complete( + streamText({ + model, + messages: turn.messages, + abortSignal: turn.signal, + prepareStep: turn.prepareStep(), + tools: { gate: gateTool }, + stopWhen: stepCountIs(5), + }) + ); + } + }, + }); + + const harness = mockChatAgent(agent, { chatId, clientData }); + + try { + const opening = harness.sendMessage(userMessage("opening", "m-0")); + void opening.catch(() => {}); + await waitFor(() => turnCount === 1, "turn 1 started"); + await waitFor(() => toolEntered, "tool entered, boundary pending"); + + clientData.sequence = 1; + await sendAndLand(harness, chatId, "first", "m-a"); + clientData.sequence = 2; + await sendAndLand(harness, chatId, "second", "m-b"); + + parseGate.resolve(); + await waitFor( + () => received.includes("m-a") && received.includes("m-b"), + "both frames validated and queued" + ); + + toolGate.resolve(); + await waitFor( + () => injected.includes("m-a") && injected.includes("m-b"), + "both frames injected" + ); + + await opening; + + expect(injected).toEqual(["m-a", "m-b"]); + expect(turnCount).toBe(1); + expect(answers).toHaveLength(1); + } finally { + parseGate.resolve(); + toolGate.resolve(); + await harness.close(); + } + } + ); +}); From b058d33465a006380258656e350231c81fec678e Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 27 Aug 2026 18:00:40 +0100 Subject: [PATCH 09/11] docs(ai-chat): use the declared signature for chat.withClientData The block read as a call expression with a type-literal body, which is not valid TypeScript and did not match the declaration form the neighbouring chat.withUIMessage section uses. --- docs/ai-chat/reference.mdx | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/docs/ai-chat/reference.mdx b/docs/ai-chat/reference.mdx index eac6ecaefd6..db53d964eb3 100644 --- a/docs/ai-chat/reference.mdx +++ b/docs/ai-chat/reference.mdx @@ -550,10 +550,13 @@ Use this when you need [`InferChatUIMessage`](#inferchatuimessage) / typed `data 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({ +chat.withClientData(config: { schema: TSchema; reportErrorAt?: "turn-end" | "arrival"; - onValidationError?: (event: { error: unknown; payload: ChatTaskWirePayload }) => Promise | void; + onValidationError?: (event: { + error: unknown; + payload: ChatTaskWirePayload; + }) => Promise | void; }): ChatBuilder; ``` From b70a880e74a07314396d1e04c33dcd22f17ffe4f Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 27 Aug 2026 18:18:24 +0100 Subject: [PATCH 10/11] docs(ai-chat): scope reportErrorAt to steering frames The paragraph introducing reportErrorAt followed the one about chat.messages.on(), so it read as governing that path too. It does not: the timing is read only by the session's steering observer, and a raw subscription always reports through the callback and the task log without writing to the stream. Says which path the option applies to, and why the two differ: a raw subscription has no turn boundary the SDK can key a deferred write to, while the session owns one. --- docs/ai-chat/custom-agents.mdx | 6 ++++-- docs/ai-chat/reference.mdx | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 161cb2b04e4..842835bf736 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -31,9 +31,11 @@ This convenience path settles the invalid input before the read returns. If your 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 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 })`. +`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. A raw subscription has no turn boundary the SDK can key a later write to, so it reports through the callback and the task log only and never writes to the stream. -By default the client-visible `Invalid client data` error for such a frame is held until the turn ends, so a bad send cannot truncate an answer the user is already reading. Pass `reportErrorAt: "arrival"` to `withClientData` if you would rather surface it as soon as validation fails, accepting that it ends the response in progress: +The steering subscription created by `chat.createSession({ pendingMessages })` skips an invalid frame the same way, but the session does own the turn boundary, so it can write the client-visible error once the turn has closed. `reportErrorAt` governs that write and applies to steering frames only, not to `chat.messages.on()`. + +By default the `Invalid client data` error for a steering frame is held until the turn ends, so a bad send cannot truncate an answer the user is already reading. Pass `reportErrorAt: "arrival"` to `withClientData` if you would rather surface it as soon as validation fails, accepting that it ends the response in progress: ```ts chat.withClientData({ diff --git a/docs/ai-chat/reference.mdx b/docs/ai-chat/reference.mdx index db53d964eb3..18ee0c58190 100644 --- a/docs/ai-chat/reference.mdx +++ b/docs/ai-chat/reference.mdx @@ -563,10 +563,10 @@ chat.withClientData(config: { | Parameter | Type | Default | Description | | ------------------- | --------------------------- | ------------ | ------------------------------------------------------------------------------------------------------------------------ | | `schema` | `TaskSchema` | required | Zod, ArkType, Valibot, or any supported schema lib | -| `reportErrorAt` | `"turn-end" \| "arrival"` | `"turn-end"` | When the client-visible `Invalid client data` error is written for a frame that failed validation mid-turn | +| `reportErrorAt` | `"turn-end" \| "arrival"` | `"turn-end"` | When the client-visible `Invalid client data` error is written for a steering frame that failed validation mid-turn | | `onValidationError` | `(event) => void` | — | Called when an input fails validation. Composes with the task-level `onClientDataValidationError` rather than replacing it | -`reportErrorAt` governs only the stream-visible error. `"turn-end"` holds it until the turn closes, so a bad send cannot truncate an answer the user is already reading; `"arrival"` writes it as soon as validation fails, ending the response in progress. `onValidationError` and the task log fire on arrival in both modes, and the frame is never delivered as a turn either way. +`reportErrorAt` governs only the stream-visible error, and only for frames arriving on the steering subscription created by `chat.createSession({ pendingMessages })`. `"turn-end"` holds it until the turn closes, so a bad send cannot truncate an answer the user is already reading; `"arrival"` writes it as soon as validation fails, ending the response in progress. `onValidationError` and the task log fire on arrival in both modes, and the frame is never delivered as a turn either way. A raw `chat.messages.on()` subscription has no turn boundary to defer to, so it always reports through the callback and the task log without writing to the stream. For `chat.customAgent()`, invalid client data is skipped. Async reads emit an error chunk followed by `turn-complete`. A `chat.messages.on()` subscription uses the task's `onClientDataValidationError` callback and task log instead, so an active response is not ended early. Without a schema, metadata is passed through unchanged. From cc49c01e183c4e88e6e993ca1fbd1acdbfae3dd2 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Thu, 27 Aug 2026 18:30:17 +0100 Subject: [PATCH 11/11] fix(sdk): treat the read timeout as a total budget across skipped frames chat.messages.next() and once() re-issued the wait with the caller's full timeout on every skipped invalid frame, so the timeout was per attempt rather than a total. A client sending invalid frames faster than the timeout kept the read blocked indefinitely and the documented return on timeout never fired. Both now compute a deadline once and pass the remaining time to each attempt, which also preserves the existing zero-timeout behaviour of draining only what is already buffered. --- packages/trigger-sdk/src/v3/ai.ts | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index d7f96a212f9..ead72cf1aaa 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1996,11 +1996,18 @@ const messagesInput: ChatMessages = { }, once(options) { return new InputStreamOncePromise((resolve, reject) => { - // Same skip-and-wait rule as `waitWithIdleTimeout`: a payload that fails - // validation is reported and not surfaced. + /** + * Same skip-and-wait rule as `waitWithIdleTimeout`: a payload that fails + * validation is reported and not surfaced. The timeout is a total budget + * across retries, so skipping a frame cannot extend the wait forever. + */ + const deadline = + options?.timeoutMs === undefined ? undefined : Date.now() + options.timeoutMs; const take = () => { chatInputRouter() - .next(CHAT_ROUTE_MESSAGES, { timeoutMs: options?.timeoutMs }) + .next(CHAT_ROUTE_MESSAGES, { + timeoutMs: deadline === undefined ? undefined : Math.max(0, deadline - Date.now()), + }) .then(async (record) => { if (!record) { resolve({ @@ -2055,9 +2062,18 @@ const messagesInput: ChatMessages = { // Consuming read, so it takes the same claim-and-validate path as the other // reads: a record the observer still owns is put back and awaited, and an // invalid payload is reported and skipped rather than surfaced raw. + const totalMs = timeoutInSeconds === undefined ? undefined : timeoutInSeconds * 1000; + /** + * The caller's timeout is a total budget, not a per-attempt one. Skipping an + * invalid frame must not buy another full wait, or a client sending invalid + * frames faster than the timeout would keep the read blocked indefinitely + * and it would never return. + */ + const deadline = totalMs === undefined ? undefined : Date.now() + totalMs; + while (true) { const record = await chatInputRouter().next(CHAT_ROUTE_MESSAGES, { - timeoutMs: timeoutInSeconds === undefined ? undefined : timeoutInSeconds * 1000, + timeoutMs: deadline === undefined ? undefined : Math.max(0, deadline - Date.now()), }); if (!record) return undefined;