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..981559d7a5a 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 @@ -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 197bff6b5e1..3bea5788a90 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -19,61 +19,99 @@ 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 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. + +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 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"; +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 +140,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..a3afd7d971e 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> ``` -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/docs/ai-chat/types.mdx b/docs/ai-chat/types.mdx index 4cbf12d16c0..328a2c5e303 100644 --- a/docs/ai-chat/types.mdx +++ b/docs/ai-chat/types.mdx @@ -140,7 +140,7 @@ You can also import `InferChatUIMessage` from `@trigger.dev/sdk/ai` in non-React ## Typed client data with `chat.withClientData` -`chat.withClientData({ schema })` returns a [ChatBuilder](#chatbuilder) that fixes the client data schema. All hooks and `run` receive typed `clientData` without needing `clientDataSchema` in `.agent()` options. +`chat.withClientData({ schema })` returns a [ChatBuilder](#chatbuilder) that fixes the client data schema. Managed-agent hooks and `run` receive typed `clientData` without needing `clientDataSchema` in `.agent()` options. A `.customAgent()` run receives the parsed schema output in `payload.metadata`, and `chat.createSession()` yields it as `turn.clientData`. ```ts import { chat } from "@trigger.dev/sdk/ai"; @@ -167,6 +167,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. + ## ChatBuilder Both `chat.withUIMessage()` and `chat.withClientData()` return a **ChatBuilder** — a chainable object that accumulates configuration before creating the agent with `.agent()`. diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 844d506079b..09afc3f7eae 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1431,6 +1431,186 @@ async function withChatWriter(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 CHAT_CUSTOM_AGENT_CLIENT_DATA_ERROR_TEXT = "Invalid client data"; + +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 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; + + if (typeof parser === "function" && typeof parser.assert === "function") { + return parser.assert.bind(parser); + } + + if (typeof parser === "function") { + return (value) => assertChatCustomAgentSyncParseResult(parser(value)); + } + + if (typeof parser.parse === "function") { + return (value) => assertChatCustomAgentSyncParseResult(parser.parse(value)); + } + + 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 writeChatCustomAgentClientDataErrorToStream( + payload: ChatTaskWirePayload +): Promise { + try { + await withChatWriter((writer) => { + writer.write({ + type: "error", + errorText: CHAT_CUSTOM_AGENT_CLIENT_DATA_ERROR_TEXT, + } 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, + options: { writeToStream: boolean; callHandler?: 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 = + options.callHandler === false + ? undefined + : 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 (!options.writeToStream) { + return; + } + await writeChatCustomAgentClientDataErrorToStream(payload); +} + +type ChatCustomAgentPayloadValidationResult = + | { ok: true; payload: TPayload } + | { ok: false; error: unknown }; + +async function parseChatCustomAgentPayload( + payload: TPayload +): 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) { + return { ok: false, error }; + } +} + +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 { + 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,21 +1723,84 @@ 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. +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: + // 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; + }); +} + +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) { - 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); + } + + const deliver = (payload: ChatTaskWirePayload) => handler(payload); + return subscribeToValidatedChatMessages(deliver, { onAfterOff: deliver }); }, once(options) { const ctx = taskContext.ctx; @@ -1575,8 +1818,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 +1854,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 +1869,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 +1892,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 +5588,38 @@ 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 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. + * This validates `metadata` only; raw `action` payloads remain `unknown`. + */ clientDataSchema?: TClientDataSchema; + /** + * Called when a custom-agent input fails `clientDataSchema` validation. + * + * 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. + */ + onClientDataValidationError?: (event: { + error: unknown; + payload: ChatTaskWirePayload; + }) => Promise | void; + run: TaskOptions< + TIdentifier, + ChatTaskWirePayload>, + unknown + >["run"]; }; function chatCustomAgent< @@ -5335,7 +5629,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 +5660,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 +5684,76 @@ 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 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>, + runOptions + ); + } + + if (isHandoverBoot) { + const signal = await waitForHandover({ + payload, + timeout: "1h", + spanName: "waiting for handover signal (invalid clientData)", + }); + if (!signal || signal.kind === "handover-skip") { + return; + } + + // The head-start writer flushes before sending this signal. Writing + // the terminal error now preserves stream order and closes the stitch. + await writeChatCustomAgentClientDataErrorToStream(payload); + return; + } + + // 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, + }; + + return userRun( + recoveredPayload as ChatTaskWirePayload>, + runOptions + ); }, }); @@ -8281,7 +8660,10 @@ export interface ChatBuilder< options: ChatCustomAgentOptions ) => Task, unknown> : ( - options: ChatCustomAgentOptions + options: Omit< + ChatCustomAgentOptions, + "clientDataSchema" + > ) => Task>, unknown>; } @@ -9441,7 +9823,7 @@ export type ChatSessionOptions = { pendingMessages?: PendingMessagesOptions; }; -export type ChatTurn = { +export type ChatTurn = { /** Turn number (0-indexed). */ number: number; /** Chat session ID. */ @@ -9449,7 +9831,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 +9930,10 @@ export type ChatTurn = { * }); * ``` */ -function createChatSession( - payload: ChatTaskWirePayload, +function createChatSession( + payload: ChatTaskWirePayload, options: ChatSessionOptions -): AsyncIterable { +): AsyncIterable> { const { signal: runSignal, idleTimeoutInSeconds: sessionIdleTimeoutOpt, @@ -9579,14 +9961,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(); + async next(): Promise>> { + 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; @@ -9647,7 +10038,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 +10050,33 @@ 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 (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.payload as ChatTaskWirePayload; + break; + } + + const validated = await validateChatCustomAgentPayload(candidate.payload); + 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 +10093,7 @@ function createChatSession( stop.cleanup(); return { done: true, value: undefined }; } - currentPayload = next.output; + currentPayload = next.output as ChatTaskWirePayload; } } @@ -9707,37 +10123,73 @@ 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) { - try { - await sessionPendingMessages.onReceived({ - message: lastUIMessage, - chatId: currentPayload.chatId, - turn, - }); - } 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 { - const modelMsgs = await toModelMessages([lastUIMessage]); - turnSteeringQueue.push({ uiMessage: lastUIMessage, modelMessages: modelMsgs }); + await sessionPendingMessages.onReceived({ + message: lastUIMessage, + chatId: currentPayload.chatId, + turn, + }); } catch { /* non-fatal */ } } - return; + if (!isActive()) { + bufferForNextTurn(); + return; + } + try { + const modelMsgs = await toModelMessages([lastUIMessage]); + if (!isActive()) { + bufferForNextTurn(); + return; + } + turnSteeringQueue.push({ uiMessage: lastUIMessage, modelMessages: modelMsgs }); + } catch { + /* non-fatal */ + } } + }; - sessionPendingWire.push(msg); - }); + 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({ payload: msg, validation: "unvalidated" }); + }); activeMsgSub = sessionMsgSub; // Accumulate messages. Slim wire: pass the single delta message as @@ -9773,11 +10225,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; }, @@ -9954,7 +10406,6 @@ function createChatSession( } } - sessionMsgSub.off(); await chatWriteTurnComplete(); return response; }, 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 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 validationErrors: unknown[] = []; + + const agent = chat + .withClientData({ + schema: z.object({ + userId: z.string(), + attempt: z.coerce.number().int(), + }), + }) + .customAgent({ + id: "custom-agent-client-data-invalid-frame", + 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); + 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: "Invalid client data" }), + ]); + expect(validationErrors).toHaveLength(1); + expect(validationErrors[0]).toBeInstanceOf(z.ZodError); + 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 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 + .withClientData({ + schema: z.object({ userId: z.string() }), + }) + .customAgent({ + id: "custom-agent-client-data-invalid-initial", + onClientDataValidationError: ({ error }) => { + validationErrors.push(error); + }, + 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(() => 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( + (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: "Invalid client data" }), + ]); + + clientData.userId = "user_123"; + await harness.sendMessage(userMessage("retry", "message-1")); + + expect(runCalls).toBe(1); + expect(receivedClientData).toEqual({ userId: "user_123" }); + } 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("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(); + const parserFinished = deferred(); + let removeSubscription: (() => void) | undefined; + let handlerCalls = 0; + let validationErrorCalls = 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; + 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++; + }); + removeSubscription = () => subscription.off(); + await 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; + void harness.sendMessage(userMessage("hello", "message-1")); + await parserStarted.promise; + + removeSubscription!(); + releaseParser.resolve(); + + 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("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; + 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(); + 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: "Invalid client data" }) + ); + } finally { + releaseFirstTurn.resolve(); + await harness.close(); + } + }); + + 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 client data" }) + ); + } 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[] = []; + 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("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(); + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-handover-skip-chat", + mode: "handover-prepare", + clientData, + }); + + try { + await waitFor(() => validationErrors.length === 1); + expect(runCalls).toBe(0); + expect(harness.allRawChunks).toHaveLength(0); + + // The validation 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("fails an invalid handover boot after the warm handler signals", 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-invalid", + onClientDataValidationError: ({ error }) => { + validationErrors.push(error); + }, + run: async () => { + runCalls++; + await chat.writeTurnComplete(); + }, + }); + + const harness = mockChatAgent(agent, { + chatId: "custom-agent-client-data-handover-invalid-chat", + mode: "handover-prepare", + clientData, + }); + + try { + await waitFor(() => validationErrors.length === 1); + expect(runCalls).toBe(0); + expect(harness.allRawChunks).toHaveLength(0); + + const handover = await harness.sendHandover({ + partialAssistantMessage: [ + { role: "assistant", content: [{ type: "text", text: "warm partial" }] }, + ], + }); + + expect(runCalls).toBe(0); + expect(handover.chunks).toEqual([ + expect.objectContaining({ type: "error", errorText: "Invalid client data" }), + ]); + expect(handover.rawChunks).toContainEqual( + expect.objectContaining({ type: "trigger:turn-complete" }) + ); + } 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(); + } + }); +});