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 a9acb1614f0..8bdea37f1fe 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 8c66919fac6..842835bf736 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -19,61 +19,113 @@ 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. 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. + +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({ + 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. + +```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 +154,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 983bdca5321..18ee0c58190 100644 --- a/docs/ai-chat/reference.mdx +++ b/docs/ai-chat/reference.mdx @@ -547,15 +547,30 @@ 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; +chat.withClientData(config: { + 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 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, 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. + +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). @@ -785,7 +800,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/docs/ai-chat/types.mdx b/docs/ai-chat/types.mdx index 4cbf12d16c0..e0b6dcc873f 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,10 @@ 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()`. diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 7b101c0e62a..ead72cf1aaa 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -1389,6 +1389,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. @@ -1540,7 +1720,7 @@ async function waitOnChatRoute( onSuspend?: () => Promise | void; onResume?: () => Promise | void; } -): Promise<{ ok: true; output: T } | { ok: false; error?: Error }> { +): Promise<{ ok: true; output: T; record: SessionStreamRecord } | { ok: false; error?: Error }> { const router = chatInputRouter(); const session = getChatSession(); @@ -1552,13 +1732,13 @@ async function waitOnChatRoute( const warm = await router.next(route, { timeoutMs: idleMs }); if (warm) { span.setAttribute("wait.resolved", "idle"); - return { ok: true as const, output: warm.data as T }; + return { ok: true as const, output: warm.data as T, record: warm }; } } else { const buffered = await router.next(route, { timeoutMs: 0 }); if (buffered) { span.setAttribute("wait.resolved", "buffered"); - return { ok: true as const, output: buffered.data as T }; + return { ok: true as const, output: buffered.data as T, record: buffered }; } } @@ -1596,7 +1776,7 @@ async function waitOnChatRoute( if (!record) continue; if (options.onResume) await options.onResume(); - return { ok: true as const, output: record.data as T }; + return { ok: true as const, output: record.data as T, record }; } }, { @@ -1614,35 +1794,256 @@ async function waitOnChatRoute( ); } +type ChatMessageSubscription = { + off: () => void; +}; + +/** + * Raw `.in` message delivery, with no client-data validation. + * + * Sits directly on the router's `messages` route, so a record handed to the + * handler is consumed and the resume floor is free to advance past it. + */ +function subscribeToRawChatMessages( + handler: (payload: ChatTaskWirePayload) => unknown +): ChatMessageSubscription { + return chatInputRouter().on(CHAT_ROUTE_MESSAGES, (record) => { + const chunk = record.data as Extract; + void Promise.resolve(handler(chunk.payload)).catch(() => {}); + }); +} + +/** + * Message delivery with client-data validation in front of it. + * + * Parses are chained so payloads are validated in wire order, which is the + * point of the feature: a later message must not be validated against state a + * earlier one has not established yet. + * + * The handler is dispatched but **not** awaited by that chain. Awaiting it + * would serialise user code, so one slow or never-resolving handler would stall + * 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. + * + * `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, + 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) { + void Promise.resolve(handler(result.payload, () => active)).catch(() => {}); + } else { + void Promise.resolve(options.onAfterOff?.(result.payload)).catch(() => {}); + } + }) + .catch(() => {}); + }); + + return { + off() { + active = false; + subscription.off(); + }, + }; +} + +/** + * Validated delivery that does **not** consume. + * + * The steering path must observe rather than take: a record it does not inject + * has to stay queued so a later turn answers it. Validation still runs in wire + * order in front of the handler, so `clientData` is parsed before steering code + * sees it, and a frame whose parse finishes after the turn closed needs no + * buffer of its own, because the record was never removed from the channel. + */ +function observeValidatedChatMessages( + /** + * 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(); + const claims = + locals.get(chatObserverClaimedSeqNumsKey) ?? + new Map; release: () => void }>(); + locals.set(chatObserverClaimedSeqNumsKey, claims); + + const subscription = chatInputRouter().observe(CHAT_ROUTE_MESSAGES, (record) => { + const payload = (record.data as Extract).payload; + 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 + * spin on a record it is not allowed to take. + */ + let release!: () => void; + const released = new Promise((resolve) => { + release = resolve; + }); + claims.set(record.seqNum, { released, release }); + const releaseClaim = () => { + claims.delete(record.seqNum); + release(); + }; + delivery = delivery + .then(async () => { + const result = await parseChatCustomAgentPayload(payload); + if (!result.ok) { + /** + * Reported once, then taken off the channel. + * + * Observing does not consume, so without this the record would be + * read again by the next turn and parsed a second time, firing a + * user-supplied schema and reporting the same failure twice. A frame + * that fails validation is never going to be answered by this run or + * 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, { + 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. + chatInputRouter().take(CHAT_ROUTE_MESSAGES, record.seqNum); + releaseClaim(); + return; + } + + // 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, seqNum, () => active)).catch(() => {}); + }) + .catch(() => {}); + }); + + return { + off() { + active = false; + subscription.off(); + }, + }; +} + +/** + * Whether a record is currently owned by the validating observer. + * + * A read that pulls one puts it back and keeps waiting, so the observer's parse + * stays the only one and the record is not consumed out from under it. + */ +function observerClaimFor(record: SessionStreamRecord): Promise | undefined { + return locals.get(chatObserverClaimedSeqNumsKey)?.get(record.seqNum)?.released; +} + const messagesInput: ChatMessages = { id: "chat-messages", on(handler) { - return chatInputRouter().on(CHAT_ROUTE_MESSAGES, (record) => { - const chunk = record.data as Extract; - void Promise.resolve(handler(chunk.payload)).catch(() => {}); - }); + if (!locals.get(chatCustomAgentClientDataParserKey)) { + return subscribeToRawChatMessages(handler); + } + + const deliver = (payload: ChatTaskWirePayload) => handler(payload); + return subscribeToValidatedChatMessages(deliver, { onAfterOff: deliver }); }, once(options) { return new InputStreamOncePromise((resolve, reject) => { - chatInputRouter() - .next(CHAT_ROUTE_MESSAGES, { timeoutMs: options?.timeoutMs }) - .then((record) => { - if (!record) { - resolve({ - ok: false, - error: new InputStreamTimeoutError("chat-messages", options?.timeoutMs ?? 0), - }); - return; - } - const chunk = record.data as Extract; - resolve({ ok: true, output: chunk.payload }); - }, reject); + /** + * 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: deadline === undefined ? undefined : Math.max(0, deadline - Date.now()), + }) + .then(async (record) => { + if (!record) { + resolve({ + ok: false, + error: new InputStreamTimeoutError("chat-messages", options?.timeoutMs ?? 0), + }); + return; + } + const claim = observerClaimFor(record); + if (claim) { + chatInputRouter().untake(CHAT_ROUTE_MESSAGES, record); + await claim; + take(); + return; + } + const payload = (record.data as Extract).payload; + if (!shouldValidateChatCustomAgentPayload(payload)) { + resolve({ ok: true, output: payload }); + return; + } + const validated = await validateChatCustomAgentPayload(payload); + if (validated.ok) { + resolve({ ok: true, output: validated.payload }); + return; + } + take(); + }, reject); + }; + take(); }); }, peek() { const record = chatInputRouter().peek(CHAT_ROUTE_MESSAGES); if (!record) return undefined; - return (record.data as Extract).payload; + const payload = (record.data as Extract).payload; + return validateChatCustomAgentPayloadSync(payload); }, async hasPending() { return chatInputRouter().hasPending(CHAT_ROUTE_MESSAGES); @@ -1658,39 +2059,106 @@ 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. + 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; - const chunk = record.data as Extract; - return { id: record.id, seqNum: record.seqNum, payload: chunk.payload }; + while (true) { + const record = await chatInputRouter().next(CHAT_ROUTE_MESSAGES, { + timeoutMs: deadline === undefined ? undefined : Math.max(0, deadline - Date.now()), + }); + if (!record) return undefined; + + 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) => { try { - const result = await waitOnChatRoute>( - CHAT_ROUTE_MESSAGES, - { timeout: options?.timeout, spanName: options?.spanName } - ); - resolve( - result.ok - ? { ok: true, output: result.output.payload } - : { ok: false, error: result.error ?? new Error("Timed out") } - ); + // Same skip-and-wait rule as `waitWithIdleTimeout`, looping back into + // the suspending wait rather than surfacing an unvalidated payload. + while (true) { + const result = await waitOnChatRoute>( + CHAT_ROUTE_MESSAGES, + { timeout: options?.timeout, spanName: options?.spanName } + ); + if (!result.ok) { + resolve({ ok: false, error: result.error ?? new Error("Timed out") }); + return; + } + const claim = observerClaimFor(result.record); + if (claim) { + chatInputRouter().untake(CHAT_ROUTE_MESSAGES, result.record); + await claim; + continue; + } + const payload = result.output.payload; + if (!shouldValidateChatCustomAgentPayload(payload)) { + resolve({ ok: true, output: payload }); + return; + } + const validated = await validateChatCustomAgentPayload(payload); + if (validated.ok) { + resolve({ ok: true, output: validated.payload }); + return; + } + } } catch (error) { reject(error); } }); }, async waitWithIdleTimeout(options) { - const result = await waitOnChatRoute>( - CHAT_ROUTE_MESSAGES, - options - ); - return result.ok - ? { ok: true, output: result.output.payload } - : { ok: false, error: result.error }; + /** + * A payload that fails client-data validation is reported and skipped + * rather than surfaced, so the caller never sees an unvalidated + * `clientData`. The loop waits for the next message instead of returning an + * error, which keeps a single bad client send from ending the turn. + */ + while (true) { + const result = await waitOnChatRoute>( + CHAT_ROUTE_MESSAGES, + options + ); + if (!result.ok) return { ok: false, error: result.error }; + + const claim = observerClaimFor(result.record); + if (claim) { + chatInputRouter().untake(CHAT_ROUTE_MESSAGES, result.record); + await claim; + continue; + } + const payload = result.output.payload; + if (!shouldValidateChatCustomAgentPayload(payload)) { + return { ok: true, output: payload }; + } + const validated = await validateChatCustomAgentPayload(payload); + if (validated.ok) { + return { ok: true, output: validated.payload }; + } + } }, async send(_runId, data, options) { await getChatSession().in.send( @@ -2781,6 +3249,32 @@ const chatOnCompactedKey = locals.create<(event: CompactedEvent) => Promise | void>("chat.onCompacted"); /** @internal Full task `ctx` for the active `chat.agent` run (for hooks invoked from nested compaction). */ const chatAgentRunContextKey = locals.create("chat.agentRunContext"); +/** @internal When a mid-turn validation failure writes its terminal error. */ +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. + * + * The observer watches without consuming, so a record it is still validating is + * also visible to a turn's read. Claiming makes ownership single: a read that + * pulls a claimed record puts it back and keeps waiting, so exactly one of them + * validates it and the schema runs once per frame. + */ +const chatObserverClaimedSeqNumsKey = locals.create< + Map; release: () => void }> +>("chat.customAgent.observerClaims"); /** @internal Marks the root run created by `chat.customAgent()`. */ const chatCustomAgentRunKey = locals.create("chat.customAgentRun"); /** @internal Number of active `chat.createSession()` iterators in this run. */ @@ -5638,9 +6132,53 @@ 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; + /** + * When a frame that arrived mid-turn fails `clientDataSchema` validation, + * decides when the terminal error reaches the client. + * + * - `"turn-end"` (default) waits until the turn has closed. A terminal error + * written into a live response can close it, so this keeps a bad client + * send from truncating an answer the user is already reading. + * - `"arrival"` writes it as soon as validation fails, which surfaces the + * problem sooner at the cost of ending the response in progress. + * + * `onClientDataValidationError` and the task log fire on arrival either way; + * this only governs the stream-visible error. The frame is never delivered as + * a turn in either mode. + */ + clientDataReportErrorAt?: "turn-end" | "arrival"; + run: TaskOptions< + TIdentifier, + ChatTaskWirePayload>, + unknown + >["run"]; }; function chatCustomAgent< @@ -5650,7 +6188,17 @@ function chatCustomAgent< >( options: ChatCustomAgentOptions ): Task>, unknown> { - const { clientDataSchema, run: userRun, ...restOptions } = options; + const { + clientDataSchema, + onClientDataValidationError, + clientDataReportErrorAt, + run: userRun, + ...restOptions + } = options; + const parseClientData = clientDataSchema ? getSchemaParseFn(clientDataSchema) : undefined; + const parseClientDataSync = clientDataSchema + ? getChatCustomAgentSyncSchemaParseFn(clientDataSchema) + : undefined; const task = createTask< TIdentifier, @@ -5678,6 +6226,19 @@ function chatCustomAgent< locals.set(chatExternalIdKey, payload.chatId); locals.set(chatAgentRunContextKey, runOptions.ctx); locals.set(chatCustomAgentRunKey, true); + if (parseClientData && parseClientDataSync) { + locals.set(chatCustomAgentClientDataParserKey, { + parse: parseClientData, + parseSync: parseClientDataSync, + }); + } + if (onClientDataValidationError) { + locals.set( + chatCustomAgentClientDataErrorHandlerKey, + onClientDataValidationError as ChatCustomAgentClientDataErrorHandler + ); + } + 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` @@ -5689,7 +6250,76 @@ function chatCustomAgent< await installChatInputRouter(payload.chatId, { resuming: Boolean(payload.continuation), }); - 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 + ); }, }); @@ -8438,9 +9068,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. */ @@ -8540,7 +9190,10 @@ export interface ChatBuilder< options: ChatCustomAgentOptions ) => Task, unknown> : ( - options: ChatCustomAgentOptions + options: Omit< + ChatCustomAgentOptions, + "clientDataSchema" + > ) => Task>, unknown>; } @@ -8561,6 +9214,8 @@ type ChatBuilderHooks = { type ChatBuilderConfig = { uiStreamOptions?: ChatUIMessageStreamOptions; clientDataSchema?: TaskSchema; + clientDataReportErrorAt?: "turn-end" | "arrival"; + clientDataOnValidationError?: ChatCustomAgentClientDataErrorHandler; hooks: ChatBuilderHooks; }; @@ -8588,10 +9243,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, }); }, @@ -8704,6 +9366,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; @@ -8759,9 +9428,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: {}, }); } @@ -9267,6 +9940,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. @@ -9277,6 +9951,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}. * @@ -9765,7 +10458,7 @@ export type ChatSessionOptions = { pendingMessages?: PendingMessagesOptions; }; -export type ChatTurn = { +export type ChatTurn = { /** Turn number (0-indexed). */ number: number; /** Chat session ID. */ @@ -9773,7 +10466,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. */ @@ -9837,14 +10530,14 @@ export type ChatTurn = { | undefined; }; -function trackActiveChatSessionIterator( - iterator: AsyncIterator -): AsyncIterator { +function trackActiveChatSessionIterator( + iterator: AsyncIterator +): AsyncIterator { locals.set(chatActiveSessionIteratorsKey, (locals.get(chatActiveSessionIteratorsKey) ?? 0) + 1); let active = true; let closing = false; let activeNextCalls = 0; - let closePromise: Promise> | undefined; + let closePromise: Promise> | undefined; const nextSettledWaiters = new Set<() => void>(); function finish() { @@ -9869,7 +10562,7 @@ function trackActiveChatSessionIterator( return new Promise((resolve) => nextSettledWaiters.add(resolve)); } - function closeIterator(): Promise> { + function closeIterator(): Promise> { closing = true; if (!closePromise) { closePromise = (async () => { @@ -9895,7 +10588,7 @@ function trackActiveChatSessionIterator( } activeNextCalls++; - let result: IteratorResult; + let result: IteratorResult; try { result = await iterator.next(); } catch (error) { @@ -9964,10 +10657,10 @@ function trackActiveChatSessionIterator( * }); * ``` */ -function createChatSession( - payload: ChatTaskWirePayload, +function createChatSession( + payload: ChatTaskWirePayload, options: ChatSessionOptions -): AsyncIterable { +): AsyncIterable> { const { signal: runSignal, idleTimeoutInSeconds: sessionIdleTimeoutOpt, @@ -9994,10 +10687,10 @@ function createChatSession( let cumulativeUsage: LanguageModelUsage = emptyUsage(); // 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; - const iterator: AsyncIterator = { - async next(): Promise> { + const iterator: AsyncIterator> = { + async next(): Promise>> { activeMsgSub?.off(); activeMsgSub = undefined; if (!booted) { @@ -10061,7 +10754,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)`). @@ -10091,7 +10784,12 @@ function createChatSession( stop.cleanup(); return { done: true, value: undefined }; } - currentPayload = next.output; + /** + * The facade validated `clientData` against the configured schema + * before returning, so the parsed shape is what arrives here. The + * cast carries that across the untyped wire payload boundary. + */ + currentPayload = next.output as typeof currentPayload; } // Check limits @@ -10124,34 +10822,47 @@ function createChatSession( * mid-turn message stays queued on the router, which is what keeps the * resume floor behind it. */ - const sessionMsgSub = sessionPendingMessages - ? chatInputRouter().observe(CHAT_ROUTE_MESSAGES, async (record) => { - const msg = (record.data as Extract).payload; - 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, - seqNum: record.seqNum, - }); - } catch { - /* non-fatal */ - } + const handleSteeringMessage = async ( + msg: ChatTaskWirePayload, + seqNum: number | undefined, + isActive: () => boolean = () => true + ) => { + const lastUIMessage = msg.message; + if (lastUIMessage) { + if (sessionPendingMessages?.onReceived) { + try { + await sessionPendingMessages.onReceived({ + message: lastUIMessage, + chatId: currentPayload.chatId, + turn, + }); + } catch { + /* non-fatal */ } - }) + } + if (!isActive()) return; + try { + const modelMsgs = await toModelMessages([lastUIMessage]); + turnSteeringQueue.push({ + uiMessage: lastUIMessage, + modelMessages: modelMsgs, + seqNum, + }); + } catch { + /* non-fatal */ + } + } + }; + + const sessionMsgSub: ChatMessageSubscription | undefined = sessionPendingMessages + ? locals.get(chatCustomAgentClientDataParserKey) + ? 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); + }) : undefined; activeMsgSub = sessionMsgSub; @@ -10188,11 +10899,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/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 2f9391aa482..531bc36b3d2 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, label = "condition", timeoutMs = 8_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: ${label}`); +} + +function errorChunks(harness: { allChunks: unknown[] }) { + return (harness.allChunks as { type?: string; errorText?: string }[]).filter( + (c) => c.type === "error" + ); +} + +/** + * A terminal error written into a live response can close it, so by default a + * mid-turn validation failure is not surfaced to the stream until the turn ends. + * `clientDataValidationErrorTiming: "arrival"` opts into the earlier report. + * + * Both cases run the identical scenario, with a frame that fails validation + * while the first turn is still open, so the only difference is the setting. + */ +describe("chat.customAgent clientDataValidationErrorTiming", () => { + async function run(timing: "turn-end" | "arrival" | undefined) { + const clientData = { sequence: 0 }; + const firstTurnStarted = deferred(); + const releaseFirstTurn = deferred(); + const validationErrors: unknown[] = []; + let started = false; + + const agent = chat + .withClientData({ + schema: async (value: unknown) => { + const sequence = (value as { sequence: number }).sequence; + if (sequence === 2) { + throw new Error("invalid mid-turn frame"); + } + return { sequence }; + }, + ...(timing ? { reportErrorAt: timing } : {}), + onValidationError: ({ error }) => { + validationErrors.push(error); + }, + }) + .customAgent({ + id: `custom-agent-client-data-timing-${timing ?? "default"}`, + run: async (payload, { signal }) => { + started = true; + const session = chat.createSession(payload, { + signal, + idleTimeoutInSeconds: 1, + pendingMessages: {}, + }); + for await (const turn of session) { + firstTurnStarted.resolve(); + await releaseFirstTurn.promise; + await turn.done(); + break; + } + }, + }); + + const harness = mockChatAgent(agent, { + chatId: `custom-agent-client-data-timing-${timing ?? "default"}-chat`, + clientData, + }); + + try { + await waitFor(() => started, "run started"); + clientData.sequence = 1; + const first = harness.sendMessage(userMessage("first", "message-1")); + await firstTurnStarted.promise; + + // Lands while turn 1 is still open, and fails validation. + clientData.sequence = 2; + void harness.sendMessage(userMessage("second", "message-2")); + await waitFor(() => validationErrors.length === 1, "handler fired"); + + // Observed while the turn is still open, before it is released. + const chunksWhileOpen = errorChunks(harness).length; + + 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(); + await harness.close(); + } + } + + it("defers the terminal error past an open turn by default", { timeout: 30_000 }, async () => { + const result = await run(undefined); + // 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 () => { + const result = await run("arrival"); + expect(result.validationErrors).toHaveLength(1); + expect(result.chunksWhileOpen).toBeGreaterThan(0); + }); +}); 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..7923f5efeb3 --- /dev/null +++ b/packages/trigger-sdk/test/custom-agent-client-data-validation.test.ts @@ -0,0 +1,871 @@ +// 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 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(); + } + }); +}); 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(); + } + } + ); +});