feat(chat): runtime clientData validation for custom agents (#4646)
## Summary
`chat.withClientData({ schema }).customAgent()` now parses
`payload.metadata` before passing it to `run`, `chat.messages`, or
`chat.createSession`. Schema defaults and transforms are preserved.
Custom agents without a schema keep the existing pass-through behavior.
This does not change `chat.agent()`. Raw custom agents do not expose an
action schema, so `payload.action` remains `unknown`.
## Validation failures
Invalid client data is logged and never passed to user code. The client
receives a fixed `Invalid client data` error; validator details stay in
the task log and `onClientDataValidationError`.
- Submitted turns and async reads write the error followed by
`turn-complete`, then wait for the next valid frame. This settles the
invalid input before the raw read returns. Callers that need to
coordinate validation with their own persistence or settlement should
omit the schema and validate the full frame in their loop.
- Messageless preload and continuation boots call
`onClientDataValidationError` and wait without writing a terminal frame.
- Active `chat.messages.on()` subscriptions skip invalid frames and call
`onClientDataValidationError` without ending the response. `off()` stops
new frames. A valid frame accepted before `off()` finishes validation
and is delivered; an invalid pending frame is logged without invoking
user callbacks.
- `chat.messages.peek()` throws synchronously.
- Invalid head-start handovers fail closed. A skip ends the run. A real
handover writes the validation error after the warm output, writes
`turn-complete`, and ends the run.
Validation is automatic when a schema is declared. We can make it opt-in
or return a typed failure if maintainers prefer that contract.
## Testing
- `pnpm --filter @trigger.dev/sdk run test -- --run`
- `pnpm --filter @trigger.dev/sdk run typecheck`
- `pnpm run build --filter @trigger.dev/sdk`
- `pnpm run lint`
- Formatting checks pass
## ✅ Checklist
- [x] I followed the contributing guide
- [x] The PR title follows the convention
- [x] I ran and tested the change
## Changelog
Custom chat agents now validate and parse client data declared with
`chat.withClientData({ schema })` before passing it to agent code.
## Screenshots
Not applicable.
---------
Co-authored-by: Eric Allam <eallam@icloud.com>
This commit is contained in:
Graham Tremper
@@ -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.
|
||||
@@ -771,7 +771,7 @@ type ChatTaskWirePayload<TMessage extends UIMessage = UIMessage, TMetadata = unk
|
||||
```
|
||||
|
||||
<Note>
|
||||
**`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.
|
||||
</Note>
|
||||
|
||||
### 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
|
||||
|
||||
|
||||
@@ -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";
|
||||
import { z } from "zod";
|
||||
|
||||
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 },
|
||||
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);
|
||||
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 },
|
||||
});
|
||||
|
||||
// Persist the full exchange after the turn — your onTurnComplete equivalent
|
||||
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 },
|
||||
});
|
||||
}
|
||||
},
|
||||
});
|
||||
```
|
||||
|
||||
<Warning>
|
||||
@@ -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) |
|
||||
|
||||
@@ -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<TSchema>({ schema: TSchema }): ChatBuilder<UIMessage, TSchema>;
|
||||
chat.withClientData<TSchema extends TaskSchema>(config: {
|
||||
schema: TSchema;
|
||||
reportErrorAt?: "turn-end" | "arrival";
|
||||
onValidationError?: (event: {
|
||||
error: unknown;
|
||||
payload: ChatTaskWirePayload;
|
||||
}) => Promise<void> | void;
|
||||
}): ChatBuilder<UIMessage, TSchema>;
|
||||
```
|
||||
|
||||
| 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<ReadableStream<UIMessageChunk>>
|
||||
```
|
||||
|
||||
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
|
||||
|
||||
@@ -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()`.
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -95,7 +95,10 @@ export type ChatTaskWirePayload<TMessage extends UIMessage = UIMessage, TMetadat
|
||||
| "handover-prepare";
|
||||
messageId?: string;
|
||||
metadata?: TMetadata;
|
||||
/** Custom action payload when `trigger` is `"action"`. Validated against `actionSchema` on the backend. */
|
||||
/**
|
||||
* Custom action payload when `trigger` is `"action"`. Managed agents validate
|
||||
* this against `actionSchema`; raw custom agents must validate it themselves.
|
||||
*/
|
||||
action?: unknown;
|
||||
/** Whether this run is continuing an existing chat whose previous run ended. */
|
||||
continuation?: boolean;
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
// Import the test harness first so chat tasks register in its resource catalog.
|
||||
import { mockChatAgent } from "../src/v3/test/index.js";
|
||||
|
||||
import { describe, expect, it } from "vitest";
|
||||
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<void>((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);
|
||||
});
|
||||
});
|
||||
@@ -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<void>((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<void>((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<void>((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<void>((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();
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -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<void>((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<typeof userMessage>) => Promise<unknown> },
|
||||
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();
|
||||
}
|
||||
}
|
||||
);
|
||||
});
|
||||
Reference in New Issue
Block a user