diff --git a/.changeset/ai-client-store-snapshot.md b/.changeset/ai-client-store-snapshot.md new file mode 100644 index 0000000000..9e3146b980 --- /dev/null +++ b/.changeset/ai-client-store-snapshot.md @@ -0,0 +1,25 @@ +--- +'@tanstack/ai-client': minor +'@tanstack/ai-react': minor +'@tanstack/ai-solid': minor +'@tanstack/ai-vue': minor +'@tanstack/ai-svelte': minor +'@tanstack/ai-angular': minor +'@tanstack/ai-preact': minor +'@tanstack/ai-octane': minor +'@tanstack/ai-remix': minor +--- + +Hold client UI state in a TanStack Store atom. Framework hooks read `getSnapshot()` instead of copying fields through change callbacks. React `useChat()` first commits `initialMessages` and `initialResumeSnapshot`, then applies browser persistence after attach. This timing applies to client-only rendering and keeps the server snapshot stable during SSR hydration. Until the stored transcript is applied, `isHydrating` is `true`. + +The snapshot `messages` and `queue` arrays are frozen in every adapter. An in-place change such as `messages.push()` throws a `TypeError`. Use `setMessages` to change messages. The hook return types now show this: `messages` and `queue` are `ReadonlyArray`, and so is `RealtimeClientState.messages`. If you pass `messages` to your own component or function, type that parameter as `ReadonlyArray`. + +Snapshot message `parts` are frozen at runtime too, but their TypeScript type stays mutable, because a readonly part list is not assignable to `UIMessage`. Tool `input`, `output`, and subagent transcripts on a snapshot are copies. + +A message or part that did not change keeps the same object from one snapshot to the next. Memoized message rows do not render again when a different message streams. + +`ChatClient` adds `subscribeSnapshot` / `getSnapshot`. The live connection method stays `subscribe()`. Generation, video, realtime, and BYOK clients expose `subscribe` / `getSnapshot`. `AudioRecorder.subscribe` still passes the new state value. + +`RealtimeClient` behavior changes: when `connect()` connects, it clears the messages and the pending user and assistant transcripts from the previous session. A transcript that arrives while the microphone starts is kept. A failed connection keeps the previous transcript. `interrupt()` now clears `pendingAssistantTranscript`. + +`setMessages` accepts the hook's readonly `messages`, including nested subagent cards. A listener that throws does not skip the other listeners in any client, and it does not leave `isLoading` stuck. The error goes to `reportError` in the browser and to `console.error` in other runtimes. diff --git a/.changeset/ai-harness-provider-keys.md b/.changeset/ai-harness-provider-keys.md new file mode 100644 index 0000000000..574005d03f --- /dev/null +++ b/.changeset/ai-harness-provider-keys.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-harness': minor +--- + +Add the `providerKeys({ providers })` plugin to `@tanstack/ai-harness/plugins`. Users of a shipped harness connect their own model provider in the app, with no `.env` file. It adds `/connect:`, `/disconnect:`, and `/keys`. `/connect:` runs the provider's `signIn` when it has one, else it opens `keyUrl` when set and asks the user to paste the key as a secret question. The key is saved in the credential store of the user, and a `keyedAdapter(...)` for that provider uses it. `/keys` and the plugin state show where each key comes from: a saved key (the last 4 characters only), the env var, or no key. diff --git a/.changeset/evaluate-usage-on-failure.md b/.changeset/evaluate-usage-on-failure.md new file mode 100644 index 0000000000..19ad310db8 --- /dev/null +++ b/.changeset/evaluate-usage-on-failure.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': patch +--- + +`decide()` now reports token usage when the provider answered but the answers cannot be mapped, for example when an answer is missing, has the wrong type, or has no probability for its value. Before, `decide()` threw without reporting the usage that the provider already billed. The `onUsage` middleware hook and the `evaluate:usage` event now get that usage once, before `onError` runs. `onFinish` still does not run for a failed call, and `decide()` throws the mapping error. If an `onUsage` hook throws, that error is thrown instead, as it is for a successful call. The mapping error is still logged. diff --git a/.changeset/json-schema-tool-input.md b/.changeset/json-schema-tool-input.md new file mode 100644 index 0000000000..7b930dbbca --- /dev/null +++ b/.changeset/json-schema-tool-input.md @@ -0,0 +1,17 @@ +--- +'@tanstack/ai': minor +--- + +Check a raw JSON Schema tool input before the tool runs. This covers plain JSON Schema tools and every MCP tool. + +Behavior change: + +- A tool with a JSON Schema `inputSchema` now gets a tool error (state `output-error`) when the input does not fit the schema. The tool does not run. Before, any input ran. +- The check runs at the same place as the Standard Schema check: after approval edits and middleware `onBeforeToolCall`. +- The draft comes from `$schema` (draft-04, draft-07, 2019-09, 2020-12). With no `$schema`, the check uses 2020-12. `format` is not checked. +- OpenAPI 3.0 schemas work: `nullable: true` lets the field be `null`, and a boolean `exclusiveMinimum` or `exclusiveMaximum` works in every draft. +- The error names each failed location (up to 10) and shows the received arguments: `Input validation failed for tool : #/count: ...`. +- A schema that the check cannot use, such as one with a remote `$ref`, gives a tool error on each call. +- The check uses `@cfworker/json-schema`. It does not compile code, so it runs on Cloudflare Workers. + +Migration: if a JSON Schema tool must accept input that its schema rejects, widen the schema. diff --git a/.changeset/native-compact-retry.md b/.changeset/native-compact-retry.md new file mode 100644 index 0000000000..de48c7d7aa --- /dev/null +++ b/.changeset/native-compact-retry.md @@ -0,0 +1,6 @@ +--- +'@tanstack/ai-compaction': patch +'@tanstack/openai-base': patch +--- + +A failed native `compact` is not saved as the checkpoint or the durable compaction record. The next process calls `compact` again. A compact result keeps image files, documents, and audio, and sends those parts back on the next request. diff --git a/.changeset/opencode-teardown-bounded.md b/.changeset/opencode-teardown-bounded.md new file mode 100644 index 0000000000..9f4b78fda9 --- /dev/null +++ b/.changeset/opencode-teardown-bounded.md @@ -0,0 +1,6 @@ +--- +'@tanstack/ai-opencode': patch +'@tanstack/ai-sandbox': patch +--- + +`RUN_FINISHED` now arrives when a teardown step does not settle. The OpenCode adapter gives each teardown step (session dispose, server dispose, tool bridge close) one second and logs a step that fails. The host tool bridge's `close()` now closes open connections, so an agent that holds a request open on the bridge no longer blocks the end of a Claude Code, Codex, or OpenCode run. diff --git a/.changeset/persistence-resume-tool-results.md b/.changeset/persistence-resume-tool-results.md new file mode 100644 index 0000000000..032f553bfd --- /dev/null +++ b/.changeset/persistence-resume-tool-results.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-persistence': patch +--- + +Store a tool result once when a run resumes from a client-tool interrupt. The client posts earlier tool results under interrupt snapshot ids the store never saw, so `withPersistence` appended a second copy and the next provider request failed. diff --git a/.changeset/raw-tool-arguments.md b/.changeset/raw-tool-arguments.md new file mode 100644 index 0000000000..9f97815a79 --- /dev/null +++ b/.changeset/raw-tool-arguments.md @@ -0,0 +1,19 @@ +--- +'@tanstack/ai': minor +'@tanstack/openai-base': patch +'@tanstack/ai-openrouter': patch +'@tanstack/ai-bedrock': patch +'@tanstack/ai-mistral': patch +'@tanstack/ai-gemini': patch +'@tanstack/ai-ollama': patch +'@tanstack/ai-anthropic': patch +--- + +Keep the provider's raw tool argument string for replay. + +- Adapters now send the raw argument string as `args` on `TOOL_CALL_END`. +- The tool call in the history keeps that exact string when it is the same JSON value as the parsed input, so the next request sends the provider's own bytes back (for example `1.50`, or an id past `Number.MAX_SAFE_INTEGER`). When the parsed input differs, only the changed values are written again. +- The `TOOL_CALL_END` that reaches the client carries the same string in `metadata.tanstack.args`, so the client history matches the server history. It no longer also carries `metadata.tanstack.input`: in-process `chat()` and `StreamProcessor` parse `input` from `args`. Run the client and the server on the same release: an older client ignores `args`, so it builds the tool input from the streamed deltas only, and a tool call with no deltas (for example an Anthropic server tool) gets no input. +- History converters for Bedrock and the Gemini Interactions API send malformed tool arguments as `{}` instead of failing the request. +- Gemini Interactions: the initial `step.arguments` snapshot goes out as a `TOOL_CALL_ARGS` delta only when no argument deltas follow, so the deltas always add up to the final arguments. +- Malformed tool arguments are no longer turned into `{}` by an adapter. The tool gets a "Failed to parse tool arguments as JSON" error that the model can repair. Mistral still ends the run with an error on malformed arguments, as before. diff --git a/.changeset/reasoning-start-metadata.md b/.changeset/reasoning-start-metadata.md new file mode 100644 index 0000000000..80a0920ec1 --- /dev/null +++ b/.changeset/reasoning-start-metadata.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': patch +--- + +Keep custom metadata from `REASONING_MESSAGE_START` on the assistant `UIMessage`. The thinking text was already kept. Metadata that arrived only on the start event was dropped. diff --git a/.changeset/responses-cache-write-tokens.md b/.changeset/responses-cache-write-tokens.md new file mode 100644 index 0000000000..19c4667afc --- /dev/null +++ b/.changeset/responses-cache-write-tokens.md @@ -0,0 +1,5 @@ +--- +'@tanstack/openai-base': patch +--- + +Responses API usage now reports `promptTokensDetails.cacheWriteTokens` from `input_tokens_details.cache_write_tokens`, the same as the Chat Completions path. diff --git a/.changeset/stream-message-ownership.md b/.changeset/stream-message-ownership.md new file mode 100644 index 0000000000..51afe5d9f2 --- /dev/null +++ b/.changeset/stream-message-ownership.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': patch +--- + +File each model call on its own assistant message. Reasoning that arrives before the call's text id, a tool call whose parent id is new, and a signature that names an earlier step stay on the message that owns them. diff --git a/.changeset/structured-output-json-object.md b/.changeset/structured-output-json-object.md new file mode 100644 index 0000000000..e5b3d88868 --- /dev/null +++ b/.changeset/structured-output-json-object.md @@ -0,0 +1,12 @@ +--- +'@tanstack/ai-utils': minor +'@tanstack/openai-base': minor +'@tanstack/ai-openrouter': minor +'@tanstack/ai-groq': minor +--- + +Structured output sends `response_format: { type: "json_object" }` for OpenRouter and Groq models that declare a response format without strict structured outputs. These models reject strict `json_schema`. In JSON mode the adapter adds the schema as a system prompt. `await chat({ outputSchema })` checks the reply when the schema is a Standard Schema, such as Zod. Other models keep strict `json_schema`. + +`modelOptions` can set `json_object` for any model. On the OpenAI-base adapters (OpenAI, Grok, Groq and others), a `json_object` in `modelOptions` was replaced by strict `json_schema`. It is now sent as given, with the schema in a system prompt. In JSON mode, `supportsCombinedToolsAndSchema()` returns `false`. Tools and the schema then go out in separate requests, so the tool-calling turns do not get the schema prompt. + +OpenRouter catalog ids with a `:variant`, such as `google/gemma-4-26b-a4b-it:free`, now match their own catalog entry. Before, only the id without the variant was checked. diff --git a/.changeset/subagent-budget-edges.md b/.changeset/subagent-budget-edges.md new file mode 100644 index 0000000000..a762791500 --- /dev/null +++ b/.changeset/subagent-budget-edges.md @@ -0,0 +1,11 @@ +--- +'@tanstack/ai': patch +'@tanstack/ai-persistence': patch +--- + +Fix subagent limit and continue edges. + +- A background `start` that throws gives the `maxCalls` slot back. `maxConcurrent` counts that child until `done` settles. `start` now returns `done`. +- A schema result that is not an object does not drop `sessionId` or `background`. A top-level `anyOf` schema still shows those fields. +- A suspended child with no stored `base` keeps its transcript. +- When the parent deadline is already past, the model gets `subagent limit reached (no time left)`. diff --git a/.changeset/tool-input-coercion.md b/.changeset/tool-input-coercion.md new file mode 100644 index 0000000000..3e85c9f7c5 --- /dev/null +++ b/.changeset/tool-input-coercion.md @@ -0,0 +1,12 @@ +--- +'@tanstack/ai': minor +--- + +Repair common model mistakes in a tool input before the final input check. + +- When a JSON Schema tool's input (and every MCP tool's input) fails the check, it is repaired and checked again: `"5"` as `5` (decimal text only), `"true"` as `true`, `5` as `"5"` for a string, and `null` as `""`, `0`, or `false` for a required field. A `null` for an optional field that does not allow `null` is removed. The repair follows unions, `$ref`, `if`/`then`/`else`, `prefixItems`, `patternProperties`, and `additionalProperties`. +- A Standard Schema runs first as written. When it fails and the schema can export JSON Schema (Zod, ArkType, Valibot), the input is repaired and the schema runs again. A failure reports the issues of the first run. +- The received arguments do not change. The tool gets the repaired copy. +- New `chat({ repairToolInput })` option. `false` turns the repair off. A function replaces the built-in repair: it gets `{ tool, toolCall, input, error }` and returns the input to check again, or `undefined` to send the validation error to the model. + +Behavior change: a call that failed the check only because of these mistakes now runs with the repaired input. diff --git a/.changeset/tool-input-final-check.md b/.changeset/tool-input-final-check.md new file mode 100644 index 0000000000..39fb39220b --- /dev/null +++ b/.changeset/tool-input-final-check.md @@ -0,0 +1,22 @@ +--- +'@tanstack/ai': minor +--- + +Check a tool input at the last moment, just before the tool runs or the client gets the call. + +Behavior change: + +- The tool's Standard Schema `inputSchema` now checks the input after approval `editedArgs` and after middleware `onBeforeToolCall`. Before, a `transformArgs` value reached `execute` or the client without a check. An input that does not fit the schema now becomes a tool error (state `output-error`), and the tool does not run. An approval edit is also still checked when the run resumes. +- `onBeforeToolCall` gets the parsed input before the schema check. It has no schema defaults or transforms yet, and it can be invalid. +- Async Standard Schema input validators are awaited. Before, every call to such a tool failed with an input validation error. +- `onBeforeToolCall` now also runs for a client tool before the server sends the call to the client. This includes a client request at an interrupt boundary. It does not run while a call waits for approval or for a denied call. When the client result comes back, the result is kept, with no second hook run and no second check. `otelMiddleware` records no `execute_tool` span for a client tool, because the server does not run it. +- A validation error after the hook runs `onAfterToolCall` with `ok: false`. +- The tool error that the model sees now also shows the arguments that the model sent. When an approval edit or middleware changed them, the error says so. +- While a call waits for approval, an approved client tool in the same batch is now sent to the client. Before, it waited for the next round trip. +- When the only call that waits for approval fails its input check, the other calls of the batch now run. Before, they got no result. + +Migration: + +- If a middleware `transformArgs` or an approval edit sends a value that the schema rejects, make the value fit the schema, or widen the schema. +- If `onBeforeToolCall` reads `hookCtx.args`, check its shape first. Do not expect schema defaults there. +- If a middleware counts or blocks calls in `onBeforeToolCall` and must see only server tools, return early when `hookCtx.tool?.execute` is missing. diff --git a/docs/adapters/groq.md b/docs/adapters/groq.md index d230f5d547..e450ddd895 100644 --- a/docs/adapters/groq.md +++ b/docs/adapters/groq.md @@ -291,6 +291,10 @@ Creates a Groq summarization adapter with an explicit API key. Creates a Groq transcription (speech-to-text) adapter. The short form reads `GROQ_API_KEY` from the environment; the `create*` form takes an explicit API key. Supported models: `whisper-large-v3-turbo`, `whisper-large-v3`. +## Structured output + +Some Groq models, for example `llama-3.3-70b-versatile`, reject strict JSON schema. For these models, `chat({ outputSchema })` sends `json_object` and adds your schema to the system prompt. Read [When strict JSON schema is rejected](../structured-outputs/overview.md#when-strict-json-schema-is-rejected). + ## Limitations - **Text-to-Speech**: Groq does not currently expose a TTS adapter. Use OpenAI, Gemini, ElevenLabs, or fal for speech generation. diff --git a/docs/adapters/openrouter.md b/docs/adapters/openrouter.md index f180126cd3..1d1a4491bd 100644 --- a/docs/adapters/openrouter.md +++ b/docs/adapters/openrouter.md @@ -190,6 +190,27 @@ export async function POST(request: Request) { } ``` +## JSON mode + +Some OpenRouter models reject strict JSON schema. For these models, `chat({ outputSchema })` sends `json_object` and adds your schema to the system prompt. + +To send JSON mode to any other model, set `modelOptions.responseFormat`: + +```typescript +import { chat } from "@tanstack/ai"; +import { openRouterText } from "@tanstack/ai-openrouter"; +import { z } from "zod"; + +const person = await chat({ + adapter: openRouterText("openai/gpt-5.6-sol"), + messages: [{ role: "user", content: "Return Ada Lovelace, age 36." }], + modelOptions: { responseFormat: { type: "json_object" } }, + outputSchema: z.object({ name: z.string(), age: z.number() }), +}); +``` + +Read [When strict JSON schema is rejected](../structured-outputs/overview.md#when-strict-json-schema-is-rejected). + ## Tools and structured output together You can pass both `tools` and `outputSchema` on one `chat()` call. For some diff --git a/docs/advanced/middleware.md b/docs/advanced/middleware.md index e0800eb5e5..c52d5e90ea 100644 --- a/docs/advanced/middleware.md +++ b/docs/advanced/middleware.md @@ -668,7 +668,26 @@ system prompt, is in [Apply Answers](../interrupts/apply-answers). ### onBeforeToolCall -Called before each tool executes. The first middleware that returns a non-void decision short-circuits — remaining middleware are skipped for that tool call. +Called before a server tool runs, and before a client tool goes to the client. The first middleware that returns a non-void decision skips the remaining middleware for that call. + +Each tool call goes through these steps: + +1. The engine parses the JSON arguments from the model. +2. `onBeforeToolCall` gets the parsed input. +3. The tool's `inputSchema` checks the final input. An async validator is awaited. A schema transform applies once, here. +4. The server tool runs, or the client gets the call, with the checked input. + +The schema checks the input after your hook: + +- `hookCtx.args` has no schema defaults or transforms yet, and it can be invalid. Check its shape before you read it. +- A `transformArgs` value or an in-place edit must fit the schema. If it does not fit, the call gets a tool error, `onAfterToolCall` runs with `ok: false`, and the tool does not run. + +When the hook runs: + +- A client tool runs the hook before the server sends the call to the client. The hook does not run when the client result comes back. +- A client tool gets `onAfterToolCall` only when the hook skips the call or the schema check fails. The result from the client does not run it. If your middleware pairs the two hooks, return early when `hookCtx.tool?.execute` is missing. +- A call that waits for approval does not run the hook. A denied call does not run it either. +- After approval, the hook gets the approved input, including any `editedArgs`. See [Tool Approval](../interrupts/tool-approval). ```typescript import { type ChatMiddleware } from "@tanstack/ai"; diff --git a/docs/api/ai-angular.md b/docs/api/ai-angular.md index a5888bfb20..b14b8b9b03 100644 --- a/docs/api/ai-angular.md +++ b/docs/api/ai-angular.md @@ -151,7 +151,7 @@ import type { type TSchema = any; interface InjectChatResult { - messages: Signal; + messages: Signal>; sendMessage: ( content: string | MultimodalContent, options?: SendMessageOptions, @@ -171,7 +171,7 @@ interface InjectChatResult { reload: () => Promise; stop: () => void; clear: () => void; - setMessages: (messages: UIMessage[]) => void; + setMessages: (messages: ReadonlyArray) => void; isLoading: Signal; error: Signal; status: Signal; diff --git a/docs/api/ai-client.md b/docs/api/ai-client.md index e7118ec7a7..97f191379e 100644 --- a/docs/api/ai-client.md +++ b/docs/api/ai-client.md @@ -274,6 +274,56 @@ Stop tailing and drop the connection. Keeps messages, the run pointer and the ru id, so a later `attach()` continues where it left off. See [Lifecycle](#lifecycle-attach-and-detach). +#### `getSnapshot()` + +Returns the current `ChatClientSnapshot`. The field names are +`messages`, `status`, `isLoading`, `error`, `queue`, `runId`, `isSubscribed`, +`connectionStatus`, `sessionGenerating`, `hasOlderMessages`, `isHydrating`, +and `interruptState`. The snapshot object, `messages`, `queue`, and their +entries are frozen copies. `interruptState` is not copied. In-place +`messages.push` throws. Use `setMessages`. + +`subscribeSnapshot(listener)` returns an unsubscribe function. It does not +fire with the current value. Read `getSnapshot()` first. This is not +`subscribe()`, which starts the live connection loop. + +```typescript +import { ChatClient, fetchServerSentEvents } from "@tanstack/ai-client"; + +const client = new ChatClient({ + connection: fetchServerSentEvents("/api/chat"), +}); + +function logSnapshot() { + const snapshot = client.getSnapshot(); + console.log(snapshot.status, snapshot.messages.length); +} + +logSnapshot(); + +const unsubscribe = client.subscribeSnapshot(() => { + logSnapshot(); +}); + +client.attach(); +client.sendMessage("Hello"); +``` + +`GenerationClient`, `VideoGenerationClient`, `RealtimeClient`, and +`ByokClient` use `subscribe(listener)` / `getSnapshot()`. The listener takes +no argument. Call `getSnapshot()` inside it. A listener that throws does not +skip the other listeners. The error goes to `reportError` in the browser and +to `console.error` in other runtimes. + +When `RealtimeClient.connect()` connects, it clears the previous transcript. +A transcript that arrives while the microphone starts is kept. A failed +connection keeps the previous transcript. Snapshot `parts` are frozen at +runtime, but their TypeScript type stays mutable. Tool `input`, `output`, and +subagent transcripts on a snapshot are copies. + +`AudioRecorder.subscribe` still passes the new state: +`subscribe((state) => ...)`. `getSnapshot()` returns that same state value. + #### `stop()` Stops the current response generation. It aborts the in-flight request. @@ -296,7 +346,7 @@ import { client } from "./client"; client.clear(); ``` -#### `setMessagesManually(messages: UIMessage[])` +#### `setMessagesManually(messages: ReadonlyArray)` Manually sets the messages array. diff --git a/docs/api/ai-octane.md b/docs/api/ai-octane.md index 040e4e960b..6f99cdb537 100644 --- a/docs/api/ai-octane.md +++ b/docs/api/ai-octane.md @@ -170,7 +170,7 @@ import type { } from '@tanstack/ai-client' interface UseChatReturn { - messages: Array + messages: ReadonlyArray sendMessage: ( content: string | MultimodalContent, options?: SendMessageOptions, @@ -196,9 +196,9 @@ interface UseChatReturn { connectionStatus: ConnectionStatus sessionGenerating: boolean isHydrating: boolean - setMessages: (messages: Array) => void + setMessages: (messages: ReadonlyArray) => void clear: () => void - queue: Array + queue: ReadonlyArray cancelQueued: (id: string) => void runId: string | null } diff --git a/docs/api/ai-preact.md b/docs/api/ai-preact.md index f556403749..de8ca0dd3d 100644 --- a/docs/api/ai-preact.md +++ b/docs/api/ai-preact.md @@ -147,7 +147,7 @@ import type { } from "@tanstack/ai-client"; interface UseChatReturn { - messages: UIMessage[]; + messages: ReadonlyArray; sendMessage: ( content: string | MultimodalContent, options?: SendMessageOptions, @@ -169,7 +169,7 @@ interface UseChatReturn { isLoading: boolean; isHydrating: boolean; error: Error | undefined; - setMessages: (messages: UIMessage[]) => void; + setMessages: (messages: ReadonlyArray) => void; clear: () => void; } ``` diff --git a/docs/api/ai-react.md b/docs/api/ai-react.md index f0d88a651d..891c67ad38 100644 --- a/docs/api/ai-react.md +++ b/docs/api/ai-react.md @@ -180,7 +180,7 @@ import type { } from "@tanstack/ai-client"; interface UseChatReturn { - messages: UIMessage[]; + messages: ReadonlyArray; sendMessage: ( content: string | MultimodalContent, options?: SendMessageOptions, @@ -202,7 +202,7 @@ interface UseChatReturn { isLoading: boolean; isHydrating: boolean; error: Error | undefined; - setMessages: (messages: UIMessage[]) => void; + setMessages: (messages: ReadonlyArray) => void; clear: () => void; } ``` diff --git a/docs/api/ai-remix.md b/docs/api/ai-remix.md index 04b3977541..c93848aca3 100644 --- a/docs/api/ai-remix.md +++ b/docs/api/ai-remix.md @@ -216,7 +216,7 @@ import type { } from '@tanstack/ai-client' interface CreateChatReturn { - messages: Array + messages: ReadonlyArray sendMessage: ( content: string | MultimodalContent, options?: SendMessageOptions, @@ -240,9 +240,9 @@ interface CreateChatReturn { connectionStatus: ConnectionStatus sessionGenerating: boolean isHydrating: boolean - setMessages: (messages: Array) => void + setMessages: (messages: ReadonlyArray) => void clear: () => void - queue: Array + queue: ReadonlyArray cancelQueued: (id: string) => void runId: string | null } diff --git a/docs/api/ai-solid.md b/docs/api/ai-solid.md index f431bcc06b..54f5be2304 100644 --- a/docs/api/ai-solid.md +++ b/docs/api/ai-solid.md @@ -177,7 +177,7 @@ import type { } from "@tanstack/ai-client"; interface UseChatReturn { - messages: Accessor; + messages: Accessor>; sendMessage: ( content: string | MultimodalContent, options?: SendMessageOptions, @@ -199,7 +199,7 @@ interface UseChatReturn { isLoading: Accessor; isHydrating: Accessor; error: Accessor; - setMessages: (messages: UIMessage[]) => void; + setMessages: (messages: ReadonlyArray) => void; clear: () => void; } ``` diff --git a/docs/api/ai-svelte.md b/docs/api/ai-svelte.md index 912a9807b5..c323d88851 100644 --- a/docs/api/ai-svelte.md +++ b/docs/api/ai-svelte.md @@ -170,7 +170,7 @@ import type { import type { ModelMessage } from "@tanstack/ai"; interface CreateChatReturn { - readonly messages: UIMessage[]; + readonly messages: ReadonlyArray; sendMessage: ( content: string | MultimodalContent, options?: SendMessageOptions, @@ -196,7 +196,7 @@ interface CreateChatReturn { readonly connectionStatus: ConnectionStatus; readonly sessionGenerating: boolean; readonly isHydrating: boolean; - setMessages: (messages: UIMessage[]) => void; + setMessages: (messages: ReadonlyArray) => void; clear: () => void; /** @deprecated Use `updateForwardedProps` instead. */ updateBody: (body: Record) => void; diff --git a/docs/api/ai-vue.md b/docs/api/ai-vue.md index 8d352d9a94..1aed44c089 100644 --- a/docs/api/ai-vue.md +++ b/docs/api/ai-vue.md @@ -168,7 +168,7 @@ import type { } from "@tanstack/ai-client"; interface UseChatReturn { - messages: DeepReadonly>; + messages: Readonly>>; sendMessage: ( content: string | MultimodalContent, options?: SendMessageOptions, @@ -194,7 +194,7 @@ interface UseChatReturn { connectionStatus: DeepReadonly>; sessionGenerating: DeepReadonly>; isHydrating: DeepReadonly>; - setMessages: (messages: UIMessage[]) => void; + setMessages: (messages: ReadonlyArray) => void; clear: () => void; } ``` diff --git a/docs/chat/stream-events.md b/docs/chat/stream-events.md index a88110295d..0c6fd3ff54 100644 --- a/docs/chat/stream-events.md +++ b/docs/chat/stream-events.md @@ -238,7 +238,7 @@ The media generation hooks take a `threadId` too. There it names a slot, not a c ## Tool input and output -SSE and HTTP `TOOL_CALL_END` does not carry parsed `input`. In-process `chat()` still has `input`. Tool input and output also live on `UIMessage` parts. +SSE and HTTP `TOOL_CALL_END` does not carry parsed `input`. `metadata.tanstack.args` carries the tool argument string as the provider sent it. In-process `chat()` still has `input`, parsed from that string. Tool input and output also live on `UIMessage` parts. On the server, feed chunks into `StreamProcessor`. On the client, read `useChat` `messages`. diff --git a/docs/chat/subagents.md b/docs/chat/subagents.md index 21e49c317a..15f9c24634 100644 --- a/docs/chat/subagents.md +++ b/docs/chat/subagents.md @@ -330,7 +330,7 @@ Every agent tool takes the agent's own fields, plus two more: - The fields of `inputSchema`, for an agent with one. The tool checks them against the schema. - `prompt`: the task as text, for an agent without `inputSchema`. The child gets it as the last user message. - `sessionId`: the `subagentRunId` of an earlier result. The call continues that child. The tool has this field only when `withPersistence` is on the parent `chat()`. -- `background`: `true` starts the child and returns its `subagentRunId` at once. The tool has this field only when the host sets `subagents.binding.start`. +- `background`: `true` starts the child and returns its `subagentRunId` at once. The tool has this field only when the host sets `subagents.binding.start`. The host returns `done`. `maxConcurrent` counts that child until `done` settles. When the run uses `sessionId` or `background`, that name is reserved. If an `inputSchema` has a field with that name, `chat()` throws before the run starts. @@ -423,6 +423,8 @@ The whole tree shares one budget. A child that calls `ctx.chat({ subagents })` p - Set `limits` on the first `chat()`. A nested chat in a child uses that budget, so its own `limits` do not apply. - A child that a router picks does not count against the limits. The children it starts do count, one level deeper. - A child that runs past `timeoutMs` stops, and the model gets `subagent timed out after 120000 ms`. +- When the parent has no time left, the model gets `subagent limit reached (no time left)`. +- A limit of `timeoutMs: 0` gets `subagent limit reached (timeoutMs 0)`. ## Keys for agents diff --git a/docs/config.json b/docs/config.json index 6c49b8298a..fc89d9c9dd 100644 --- a/docs/config.json +++ b/docs/config.json @@ -160,7 +160,7 @@ "label": "Subagents", "to": "chat/subagents", "addedAt": "2026-09-21", - "updatedAt": "2026-10-13" + "updatedAt": "2026-10-14" }, { "label": "Agentic Cycle", @@ -247,7 +247,7 @@ "label": "Overview", "to": "structured-outputs/overview", "addedAt": "2026-05-19", - "updatedAt": "2026-08-19" + "updatedAt": "2026-10-15" }, { "label": "One-Shot Extraction", @@ -294,7 +294,7 @@ "label": "Server Tools", "to": "tools/server-tools", "addedAt": "2026-04-15", - "updatedAt": "2026-10-13" + "updatedAt": "2026-10-14" }, { "label": "Client Tools", @@ -353,7 +353,7 @@ "label": "Middleware", "to": "advanced/middleware", "addedAt": "2026-04-15", - "updatedAt": "2026-10-13" + "updatedAt": "2026-10-14" }, { "label": "Built-in Middleware", @@ -880,7 +880,7 @@ "label": "Evaluate", "to": "evaluate/evaluate", "addedAt": "2026-09-18", - "updatedAt": "2026-10-08" + "updatedAt": "2026-10-13" } ], "tab": "guides" @@ -1036,6 +1036,12 @@ "addedAt": "2026-10-13", "updatedAt": "2026-10-14" }, + { + "label": "Connect model providers", + "to": "harness/provider-keys", + "addedAt": "2026-10-13", + "updatedAt": "2026-10-14" + }, { "label": "Build your own UI", "to": "harness/custom-ui", @@ -1407,7 +1413,7 @@ "label": "Groq", "to": "adapters/groq", "addedAt": "2026-04-15", - "updatedAt": "2026-10-09" + "updatedAt": "2026-10-15" }, { "label": "Mistral", @@ -1459,7 +1465,7 @@ "label": "OpenRouter Adapter", "to": "adapters/openrouter", "addedAt": "2026-04-15", - "updatedAt": "2026-10-13" + "updatedAt": "2026-10-15" }, { "label": "Perplexity Search", @@ -1647,7 +1653,7 @@ "label": "@tanstack/ai-client", "to": "api/ai-client", "addedAt": "2026-04-15", - "updatedAt": "2026-10-08" + "updatedAt": "2026-10-14" }, { "label": "@tanstack/ai-react", diff --git a/docs/evaluate/evaluate.md b/docs/evaluate/evaluate.md index d71187a051..3a27b976d2 100644 --- a/docs/evaluate/evaluate.md +++ b/docs/evaluate/evaluate.md @@ -443,6 +443,10 @@ try { If `questions` is empty, `decide()` throws before any request. If a question key is `meta`, `decide()` throws. +If the provider answers but `decide()` cannot map an answer, `decide()` throws. +For example, an answer is missing, has the wrong type, or has no probability +for its value. `onUsage` and the `evaluate:usage` event still get the token +usage of that call, before `onError` runs. `onFinish` does not run. ## Runnable Example diff --git a/docs/harness/overview.md b/docs/harness/overview.md index 91830b3ef4..b84963fc68 100644 --- a/docs/harness/overview.md +++ b/docs/harness/overview.md @@ -58,6 +58,8 @@ export const assistant = defineHarness({ `agents` are agents you run from code. Put an agent in `subagents` instead when the model must call it as a tool, the same as in `chat()`. +To ship the harness to users with no `.env` file, let each user connect their own key. See [Connect model providers](./provider-keys). + ## 2. Open a session A host runs sessions. Give it the stores from `@tanstack/ai-persistence`, then open a session for a conversation id. diff --git a/docs/harness/provider-keys.md b/docs/harness/provider-keys.md new file mode 100644 index 0000000000..12fd9800b0 --- /dev/null +++ b/docs/harness/provider-keys.md @@ -0,0 +1,201 @@ +--- +title: Connect model providers +id: harness-provider-keys +order: 11 +description: "Let each user connect their own OpenAI, Anthropic, xAI, fal, or OpenRouter key inside your app with /connect. A harness that you ship needs no .env file." +keywords: + - tanstack ai + - harness + - byok + - api key + - providerKeys + - keyedAdapter +--- + +You ship your harness to users: a CLI that they install, a desktop app, or a hosted app with accounts. `openaiText('gpt-6.1-sol')` reads `OPENAI_API_KEY` from the environment. Your users do not have your `.env` file, and they do not start your app from a terminal with a key set. + +With `providerKeys`, each user runs `/connect:openai` in your app and pastes their own key. The harness saves the key for that user and uses it for every turn. + +## 1. Define keyed adapters + +A keyed adapter waits for a key. `keyedAdapter` takes the provider and a function that makes the adapter from the key: + +```ts group=harness-provider-keys +import { keyedAdapter } from '@tanstack/ai' +import { createAnthropicChat } from '@tanstack/ai-anthropic' +import { anthropicByok } from '@tanstack/ai-anthropic/byok' +import { createOpenaiChat } from '@tanstack/ai-openai' +import { openaiByok } from '@tanstack/ai-openai/byok' + +const gpt = keyedAdapter(openaiByok, (key) => createOpenaiChat('gpt-6.1-sol', key)) +const claude = keyedAdapter(anthropicByok, (key) => + createAnthropicChat('claude-sonnet-5-5', key), +) +``` + +- `openaiByok` describes the provider: the id `openai`, the label `OpenAI`, and the env var `OPENAI_API_KEY`. +- Each provider package exports one from its `/byok` subpath, for example `anthropicByok`, `grokByok`, `falByok`, and `openrouterByok`. +- `keyedAdapter` works for every adapter kind: text, image, speech, audio, and video. + +Before each turn, the session calls your function with the key of the user of that session. + +## 2. Add the providerKeys plugin + +Give the harness a keyed adapter as its `adapter`. Then list in `providerKeys` every provider that users can connect: + +```ts group=harness-provider-keys +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { modelPicker, providerKeys } from '@tanstack/ai-harness/plugins' +import { falByok } from '@tanstack/ai-fal/byok' +import { grokByok } from '@tanstack/ai-grok/byok' +import { openrouterByok } from '@tanstack/ai-openrouter/byok' +import { openrouterSignIn } from '@tanstack/ai-openrouter/pkce' +import { memoryPersistence } from '@tanstack/ai-persistence' + +export const assistant = defineHarness({ + name: 'acme/assistant', + adapter: gpt, + plugins: () => [ + providerKeys({ + providers: [ + { ...openaiByok, keyUrl: 'https://platform.openai.com/api-keys' }, + anthropicByok, + grokByok, + falByok, + { ...openrouterByok, signIn: openrouterSignIn() }, + ], + }), + modelPicker({ choices: { gpt, claude }, default: 'gpt' }), + ], +}) + +const host = createHarnessHost({ persistence: memoryPersistence() }) +const session = await host.open(assistant, { + threadId: 'ada-chat', + principal: { id: 'ada' }, +}) +``` + +- For each provider, `providerKeys` adds `/connect:` and `/disconnect:`. It also adds `/keys`. +- `modelPicker` and `compact` also take keyed adapters. After `/model claude`, the next turn uses the Anthropic key of the user. +- `signIn` is optional. OpenRouter has one. See [Sign in with OpenRouter](#sign-in-with-openrouter). +- `keyUrl` is optional. It is the page where the user makes a key. `/connect:` opens it in the browser, then asks for the key. + +## 3. Connect a provider + +`/connect:openai` asks the user for the key. The question has `secret: true`, so show a password field for it: + +```ts group=harness-provider-keys +const stopAsking = session.onEvent('CUSTOM', async ({ event }) => { + if (event.name !== 'harness.question') return + const [question] = session.snapshot().pendingQuestions + // question.secret is true. Read the key without showing it. + if (question) await session.answer({ questionId: question.questionId, value: 'sk-proj-...' }) +}) + +console.log(await session.command('connect:openai')) // Connected to OpenAI (key ...abcd). +stopAsking() +``` + +- With `keyUrl`, the page to make a key opens first. The session sends a `harness.auth_required` event with the URL, and your app opens it. +- `/disconnect:openai` deletes the key. +- `/keys` shows where each key comes from: + +```text +OpenAI: connected (key ...abcd) +Anthropic: from the ANTHROPIC_API_KEY env var +xAI Grok: missing. Run /connect grok. +fal.ai: missing. Run /connect fal. +OpenRouter: missing. Run /connect openrouter. +``` + +In your own UI, read the same list from `session.snapshot().plugins['tanstack/provider-keys'].providers`. Each entry has an `id`, a `label`, and a `state`: `connected`, `env`, or `missing`. + +### Sign in with OpenRouter + +OpenRouter has a browser sign-in, so users do not copy a key. `openrouterSignIn()` in the providers list turns it on: + +1. Run `/connect:openrouter`. +2. The session sends a `harness.auth_required` event with the OpenRouter sign-in page. Your app opens it in the browser. Sign in and approve. +3. The browser shows `You are signed in to OpenRouter. You can close this tab.` +4. The command returns `Connected to OpenRouter (key ...wxyz).` + +The sign-in uses PKCE. It listens on `127.0.0.1` on a random port for one callback, and it checks a random state. After 10 minutes it stops. To change the time, pass `openrouterSignIn({ timeoutMs })`. + +The listener runs on the computer of the harness. So the sign-in works when the harness and the browser are on the same computer: a CLI or a desktop app. In a hosted app, remove `signIn`, and users paste their OpenRouter key. + +OpenAI and Anthropic have no sign-in for other apps, so users paste a key. + +## When a key is missing + +If a turn needs a key that the user did not connect, the turn stops before the model call: + +- Clients get a `harness.auth_required` event with the provider id. +- The turn fails with `Sign in to openai first. Run /connect openai.` + +After `/connect:openai`, send the message again. + +## Where the keys are stored + +`/connect:` saves the key in the credential store of the host (`stores.credentials`), as an `api_key` credential with the provider id. The store keeps it for the principal of the session: + +- CLI or desktop app: a file in the home folder of the user. +- Hosted app: your database. Each signed-in user is a principal, so each user has their own keys. Write the store with `defineCredentialStore` from `@tanstack/ai-persistence`. +- `memoryPersistence()`: the keys stay in memory until the process stops. + +Encrypt the keys at rest in a real store. + +## Use an Anthropic token + +The key that a user pastes can be a Claude OAuth token (`sk-ant-oat...`). `createAnthropicChat` picks the auth mode from the token, so the keyed adapter above works for it too. See [Bearer and OAuth tokens](../adapters/anthropic#bearer-and-oauth-tokens). + +## The env var still works + +If the user has no saved key, the session reads the env var of the provider, for example `OPENAI_API_KEY`. A saved key wins over the env var. So you can develop with a `.env` file, and your users connect in the app. + +## Use the key of the user in an agent + +Your agent makes images with fal, so it needs the fal key of the user. Every agent gets `ctx.keys`. In a harness, `ctx.keys` reads the same keys as the turn: + +```ts group=harness-provider-keys +import { defineAgent } from '@tanstack/ai' +import { falImage } from '@tanstack/ai-fal' +import { z } from 'zod' + +const flux = keyedAdapter(falByok, (key) => falImage('fal-ai/flux/dev', { apiKey: key })) + +export const illustrator = defineAgent({ + name: 'illustrator', + description: 'Draws an image from a prompt', + inputSchema: z.object({ prompt: z.string() }), + run: async (ctx) => + ctx.generateImage({ + adapter: await ctx.keys.adapter(flux), + prompt: ctx.input.prompt, + }), +}) +``` + +Put the agent in `subagents.agents` of the harness, and the model calls it as a tool. + +- `ctx.keys.adapter(adapter)`: makes a keyed adapter with the key of the user. A plain adapter comes back unchanged. +- `ctx.keys.get(provider)`: the key, or `null`. +- `ctx.keys.require(provider)`: the key. If it is missing, the run stops with `harness.auth_required`, the same as a turn. +- The children of the agent get the same keys. +- Plugins get the same keys as `ctx.keys` in `setup`. + +Outside a harness, `ctx.keys` reads the env var. See [Keys inside agents](../advanced/byok#keys-inside-agents). + +## Keep keys safe + +- A key shows masked only: `...` and the last 4 characters. +- The plugin never puts a key in an event, the transcript, or a command result. +- The answer to the key question does not go into the inbox of the session. +- Do not log a key. For error messages, use `maskKey` from `@tanstack/ai/byok`. + +## What you have now + +- `/connect:` for OpenAI, Anthropic, xAI, fal, and OpenRouter in the harness that you ship, with no `.env` file. +- One key for each user, in your credential store. +- The env var as a fallback while you develop. +- Agents and plugins that use the key of the user through `ctx.keys`. diff --git a/docs/persistence/client-persistence.md b/docs/persistence/client-persistence.md index eb1ed4e231..7989d62003 100644 --- a/docs/persistence/client-persistence.md +++ b/docs/persistence/client-persistence.md @@ -54,10 +54,15 @@ to pick one and how it differs from the `runId` the hook reports. The client stores one record per `threadId`, the transcript plus a small resume pointer. On the next load `useChat` reads it and: -- **Repaints the transcript** from storage with no network. Sync adapters - (`localStorage` / `sessionStorage`) hydrate during construction. IndexedDB - hydrates asynchronously after the database opens. Until then, `isHydrating` - is `true`, so you can show a loading state instead of an empty chat. +- **Repaints the transcript** from storage with no network. A direct + `ChatClient` and most framework wrappers hydrate sync adapters + (`localStorage` / `sessionStorage`) during construction. React is + different: the first committed render uses `initialMessages` and + `initialResumeSnapshot`, and browser persistence is applied after the hook + attaches. This keeps the server snapshot stable during SSR hydration. + IndexedDB always hydrates asynchronously after the database opens. Until the + stored transcript is applied, `isHydrating` is `true`, so you can show a + loading state instead of an empty chat. - **Rehydrates a pending interrupt**, so an approval prompt comes back exactly as it was. - **Rejoins an in-flight run**, if a reply was still streaming when the page diff --git a/docs/structured-outputs/overview.md b/docs/structured-outputs/overview.md index db260f9778..f06ef2e538 100644 --- a/docs/structured-outputs/overview.md +++ b/docs/structured-outputs/overview.md @@ -62,6 +62,39 @@ Every adapter handles structured output through its provider's native API: The provider-specific details are handled for you. The same `chat({ outputSchema })` call works across all of them. For a coding agent in a sandbox, see [Harness Agents](./harnesses). +### When strict JSON schema is rejected + +Some models reject strict JSON schema. They accept JSON mode (`response_format` type `json_object`) only. A strict request to one of these models fails before the model runs. + +The OpenRouter and Groq adapters find these models in their catalogs. For these models, `chat({ outputSchema })` does two things: + +- It sends JSON mode. +- It adds your schema to the system prompt, because JSON mode does not send the schema. + +When your schema is a Standard Schema, such as Zod, `await chat({ outputSchema })` still checks the reply. JSON of the wrong shape throws. TanStack AI does not check a plain JSON Schema or a streamed reply. Check those in your own code. + +```typescript +import { chat } from "@tanstack/ai"; +import { openRouterText } from "@tanstack/ai-openrouter"; +import { z } from "zod"; + +const Person = z.object({ name: z.string(), age: z.number() }); + +const person = await chat({ + adapter: openRouterText("qwen/qwen3.7-flash"), + messages: [{ role: "user", content: "Return Ada Lovelace, age 36." }], + outputSchema: Person, +}); + +person.name; +person.age; +``` + +To send JSON mode to a model that is not in the catalog, set it in `modelOptions`: + +- OpenRouter: `modelOptions: { responseFormat: { type: "json_object" } }` +- Groq: `modelOptions: { response_format: { type: "json_object" } }` + ### Anthropic schema complexity limits Anthropic compiles a structured-output schema into a grammar and rejects schemas it considers too large or too complex with a 400 error — typically `Schema is too complex for compilation` or `output_config.format.schema: Invalid schema: The compiled grammar is too large`. This affects Claude models directly and `anthropic/*` models routed through OpenRouter, even when every other provider accepts the same schema. diff --git a/docs/tools/server-tools.md b/docs/tools/server-tools.md index 12aaf118a7..9106cb0e00 100644 --- a/docs/tools/server-tools.md +++ b/docs/tools/server-tools.md @@ -355,7 +355,96 @@ const getUserData = getUserDataDef.server(async ({ userId }) => { **Throwing vs. returning an error:** if your `.server()` function throws, the SDK catches it and surfaces it as a tool-result *error* (the model sees the failure but you lose control over the message). Returning a structured `{ error }` shape keeps the model in control of how to recover and is usually preferable. Either way, when an `outputSchema` is defined the returned value is validated against it (Zod) before being added to the conversation — so include the `error` field in your `outputSchema` if you return it. -Malformed JSON arguments and Standard Schema input validation failures are also surfaced as tool-result errors. The tool implementation is not called, and the agent loop can return the error to the model so it can repair the tool call. This also applies to Chat Completions tools with no required input fields. +Malformed JSON arguments and input validation failures (Standard Schema or JSON Schema) are also surfaced as tool-result errors. The tool implementation is not called, and the agent loop can return the error to the model so it can repair the tool call. This also applies to Chat Completions tools with no required input fields. + +The `inputSchema` checks the input just before the tool runs: + +- The check runs after approval edits and after middleware `onBeforeToolCall`. An approval edit that does not fit the schema stops the resume with an `invalid-edited-args` error. A middleware change that does not fit becomes a tool error. +- An async validator is awaited. +- The tool gets the schema's output, so a transform such as `z.string().transform(Number)` applies once. + +Models often send a value with the wrong type, such as `"5"` for a number. The check repairs these against the schema before it fails the call: + +| The model sends | The schema wants | The tool gets | +| --- | --- | --- | +| `"5"` | `number` or `integer` | `5` | +| `"true"` | `boolean` | `true` | +| `5` | `string` | `"5"` | +| `null` for an optional field | a type with no `null` | no field | +| `null` for a required field | `string`, `number`, or `boolean` | `""`, `0`, or `false` | + +- A JSON Schema tool, and an MCP tool, gets the repair when its input fails the check. A valid input runs as it is. +- A Standard Schema that can export JSON Schema (Zod, ArkType, Valibot) runs first as written. If it fails, the input is repaired and the schema runs again. If it still fails, the error shows the issues of the first run. +- An input that cannot be repaired, such as `"many"` for a number, becomes a tool error. + +### Turn off or replace the repair + +Sometimes the built-in repair is the wrong move. Maybe your tool moves money, and a `null` must never turn into `0`. Or maybe your model makes the same odd mistake every time, and you know the fix. Pass `repairToolInput` to `chat()` and take control. + +Turn the repair off, so a bad input always goes back to the model as an error: + +```typescript +import { chat, toServerSentEventsResponse } from "@tanstack/ai"; +import { openaiText } from "@tanstack/ai-openai"; +import { transferFunds } from "./tools"; + +export async function POST(request: Request) { + const { messages } = await request.json(); + + const stream = chat({ + adapter: openaiText("gpt-6.1-sol"), + messages, + tools: [transferFunds], + repairToolInput: false, + }); + + return toServerSentEventsResponse(stream); +} +``` + +Or write your own fix. Say the model keeps sending `"limit": "all"` to `search_products`. Map it to a number: + +```typescript +import { chat, toServerSentEventsResponse } from "@tanstack/ai"; +import { openaiText } from "@tanstack/ai-openai"; +import { searchProducts } from "./tools"; + +export async function POST(request: Request) { + const { messages } = await request.json(); + + const stream = chat({ + adapter: openaiText("gpt-6.1-sol"), + messages, + tools: [searchProducts], + repairToolInput: ({ tool, input }) => { + if ( + tool.name === "search_products" && + typeof input === "object" && + input !== null && + "limit" in input && + input.limit === "all" + ) { + return { ...input, limit: 100 }; + } + return undefined; + }, + }); + + return toServerSentEventsResponse(stream); +} +``` + +How your function works: + +- It runs only when an input fails the check. A valid input never gets to it. +- It replaces the built-in repair. A `"5"` stays a string unless your function changes it. +- It gets `tool`, `toolCall`, `input`, and `error` (the validation error). +- Return the new input, and the check runs again. If that check fails, the model gets the original error. +- Return `undefined` to send the original error to the model. If your function throws, the model gets the thrown message. + +The option works the same for Standard Schema tools, JSON Schema tools, and MCP tools. + +A client tool, or a tool that waits for approval, is checked twice: once before it goes out, and once when it comes back. So your function can run twice for one call. ## Using JSON Schema @@ -403,14 +492,24 @@ const getUserData = getUserDataDef.server(async (args) => { }); ``` -> **Note:** JSON Schema tools skip runtime validation. Zod schemas are recommended for full type safety and validation. +The JSON Schema checks each call before the tool runs. MCP tools get the same check. If the input does not fit, the call gets a tool error, and the error names each failed field: + +```text +Input validation failed for tool get_user_data: #: Instance does not have required property "userId". +``` + +- The check uses the draft that `$schema` names: draft-04, draft-07, 2019-09, or 2020-12. With no `$schema`, it uses 2020-12. +- Got an OpenAPI 3.0 schema? `nullable: true` works, so that field can be `null`. A boolean `exclusiveMinimum` or `exclusiveMaximum` works too. +- The check does not test `format` (such as `uri` or `email`). A `format` only describes the value. +- The check does not compile code, so it also runs on Cloudflare Workers. +- `args` is still typed as `unknown`. A Zod schema gives you the type too. > **Tip:** When you pass `.client()` tools to `useChat`, a check on `part.name` narrows `part.input` and `part.output`. See [Type-safe tool call events](../chat/stream-events#type-safe-tool-call-events). ## Best Practices 1. **Keep tools focused** - Each tool should do one thing well -2. **Validate inputs** - Use Zod schemas to ensure type safety (JSON Schema skips validation) +2. **Validate inputs** - The `inputSchema` checks every call. A Zod schema also gives you types 3. **Handle errors** - Return meaningful error messages 4. **Use descriptions** - Clear descriptions help the model use tools correctly 5. **Secure sensitive operations** - Never expose API keys or secrets to the client diff --git a/examples/react/basic-chat/src/routes/index.tsx b/examples/react/basic-chat/src/routes/index.tsx index 5090278290..5cba00f8fe 100644 --- a/examples/react/basic-chat/src/routes/index.tsx +++ b/examples/react/basic-chat/src/routes/index.tsx @@ -8,7 +8,7 @@ import { import { OpenRouterKeyForm } from '@/components/open-router-key-form' import { byok } from '@/lib/byok' -function Messages({ messages }: { messages: Array }) { +function Messages({ messages }: { messages: ReadonlyArray }) { const messagesContainerRef = useRef(null) const visibleMessages = messages.filter((message) => message.parts.some((part) => part.type === 'text' && part.content.trim()), diff --git a/examples/sandbox-cloudflare/src/routes/index.tsx b/examples/sandbox-cloudflare/src/routes/index.tsx index ecdc069bc2..afb07f2812 100644 --- a/examples/sandbox-cloudflare/src/routes/index.tsx +++ b/examples/sandbox-cloudflare/src/routes/index.tsx @@ -216,7 +216,7 @@ function SandboxWaiting({ kind }: { kind: SandboxWaitKind }) { function sandboxWaitKind( isLoading: boolean, - messages: Array, + messages: ReadonlyArray, ): SandboxWaitKind | false { if ( !isLoading || @@ -232,7 +232,7 @@ function Messages({ messages, waiting, }: { - messages: Array + messages: ReadonlyArray waiting: SandboxWaitKind | false }) { const containerRef = useRef(null) diff --git a/examples/sandbox-web/src/routes/index.tsx b/examples/sandbox-web/src/routes/index.tsx index 15f64826e9..d146caf42d 100644 --- a/examples/sandbox-web/src/routes/index.tsx +++ b/examples/sandbox-web/src/routes/index.tsx @@ -160,7 +160,7 @@ function SandboxWaiting({ kind }: { kind: SandboxWaitKind }) { /** First message in a thread vs follow-up while waiting for the first chunk. */ function sandboxWaitKind( isLoading: boolean, - messages: Array, + messages: ReadonlyArray, ): SandboxWaitKind | false { if ( !isLoading || @@ -176,7 +176,7 @@ function Messages({ messages, waiting, }: { - messages: Array + messages: ReadonlyArray waiting: SandboxWaitKind | false }) { const containerRef = useRef(null) diff --git a/examples/ts-code-mode-web/src/components/MessageSizeOverlay.tsx b/examples/ts-code-mode-web/src/components/MessageSizeOverlay.tsx index c9e91185b2..c34c66e24a 100644 --- a/examples/ts-code-mode-web/src/components/MessageSizeOverlay.tsx +++ b/examples/ts-code-mode-web/src/components/MessageSizeOverlay.tsx @@ -4,7 +4,7 @@ import type { UIMessage } from '@tanstack/ai-react' import type { VMEvent } from './JavaScriptVM' interface MessageSizeOverlayProps { - messages: Array + messages: ReadonlyArray toolCallEvents: Map> } @@ -20,7 +20,7 @@ function formatBytes(bytes: number): string { /** * Calculate the actual byte size of the messages array */ -function calculateActualSize(messages: Array): number { +function calculateActualSize(messages: ReadonlyArray): number { return new TextEncoder().encode(JSON.stringify(messages)).length } @@ -29,7 +29,7 @@ function calculateActualSize(messages: Array): number { * as individual tool call parts in the messages array. */ function calculateTheoreticalSize( - messages: Array, + messages: ReadonlyArray, toolCallEvents: Map>, ): number { // Start with the actual messages size diff --git a/examples/ts-code-mode-web/src/routes/_banking-demo/banking-demo.tsx b/examples/ts-code-mode-web/src/routes/_banking-demo/banking-demo.tsx index 6dd38976e6..94bb0601f5 100644 --- a/examples/ts-code-mode-web/src/routes/_banking-demo/banking-demo.tsx +++ b/examples/ts-code-mode-web/src/routes/_banking-demo/banking-demo.tsx @@ -136,7 +136,7 @@ function ToolCallDisplay({ } // Messages component for chat transcript -function Messages({ messages }: { messages: Array }) { +function Messages({ messages }: { messages: ReadonlyArray }) { const containerRef = useRef(null) useEffect(() => { diff --git a/examples/ts-code-mode-web/src/routes/_database-demo/database-demo.tsx b/examples/ts-code-mode-web/src/routes/_database-demo/database-demo.tsx index d4d9c966e2..e55e250188 100644 --- a/examples/ts-code-mode-web/src/routes/_database-demo/database-demo.tsx +++ b/examples/ts-code-mode-web/src/routes/_database-demo/database-demo.tsx @@ -355,7 +355,7 @@ function Messages({ messages, toolCallEvents, }: { - messages: Array + messages: ReadonlyArray toolCallEvents: Map> }) { const containerRef = useRef(null) @@ -641,7 +641,7 @@ const GOLD_PROMPT_MAP: Record = { 'For each city, show the total revenue, number of unique customers, and the most purchased product category': 4, } -function extractReportText(messages: Array): string { +function extractReportText(messages: ReadonlyArray): string { return messages .filter((m) => m.role === 'assistant') .flatMap((m) => diff --git a/examples/ts-code-mode-web/src/routes/_npm-github-chat/npm-github-chat.tsx b/examples/ts-code-mode-web/src/routes/_npm-github-chat/npm-github-chat.tsx index e9da7004b6..31fb8c8b34 100644 --- a/examples/ts-code-mode-web/src/routes/_npm-github-chat/npm-github-chat.tsx +++ b/examples/ts-code-mode-web/src/routes/_npm-github-chat/npm-github-chat.tsx @@ -215,7 +215,7 @@ function Messages({ messages, toolCallEvents, }: { - messages: Array + messages: ReadonlyArray toolCallEvents: Map> }) { const messagesContainerRef = useRef(null) @@ -454,7 +454,7 @@ function Messages({ } // Calculate context sizes for the chat metrics -function calculateActualSize(messages: Array): number { +function calculateActualSize(messages: ReadonlyArray): number { return new TextEncoder().encode(JSON.stringify(messages)).length } @@ -496,7 +496,7 @@ function CodeModePage() { const npmDataIdCounter = useRef(0) // Ref to hold current messages for the client tool to access - const messagesRef = useRef>([]) + const messagesRef = useRef>([]) const body = useMemo( () => ({ diff --git a/examples/ts-code-mode-web/src/routes/_reporting/reporting-agent.tsx b/examples/ts-code-mode-web/src/routes/_reporting/reporting-agent.tsx index 09d5999a3f..a9467bccec 100644 --- a/examples/ts-code-mode-web/src/routes/_reporting/reporting-agent.tsx +++ b/examples/ts-code-mode-web/src/routes/_reporting/reporting-agent.tsx @@ -223,7 +223,7 @@ function Messages({ messages, toolCallEvents, }: { - messages: Array + messages: ReadonlyArray toolCallEvents: Map> }) { const messagesContainerRef = useRef(null) @@ -437,12 +437,12 @@ function EmptyState() { } // Calculate context sizes -function calculateActualSize(messages: Array): number { +function calculateActualSize(messages: ReadonlyArray): number { return new TextEncoder().encode(JSON.stringify(messages)).length } function calculateTheoreticalSize( - messages: Array, + messages: ReadonlyArray, toolCallEvents: Map>, ): number { let theoreticalSize = calculateActualSize(messages) diff --git a/examples/ts-react-chat/package.json b/examples/ts-react-chat/package.json index 3ce3a7f02d..dace1d669a 100644 --- a/examples/ts-react-chat/package.json +++ b/examples/ts-react-chat/package.json @@ -54,7 +54,7 @@ "@tanstack/react-router-devtools": "^1.158.4", "@tanstack/react-router-ssr-query": "^1.158.4", "@tanstack/react-start": "^1.168.60", - "@tanstack/react-store": "^0.8.0", + "@tanstack/react-store": "^0.11.1", "@tanstack/router-plugin": "^1.168.42", "@tanstack/store": "^0.11.1", "dockerode": "^4.0.2", diff --git a/examples/ts-react-chat/src/routes/app-studio.tsx b/examples/ts-react-chat/src/routes/app-studio.tsx index 3381273cb4..ee4878f794 100644 --- a/examples/ts-react-chat/src/routes/app-studio.tsx +++ b/examples/ts-react-chat/src/routes/app-studio.tsx @@ -137,7 +137,7 @@ function ThreadNav({ ) } -function collectPreviewUrls(messages: Array): Set { +function collectPreviewUrls(messages: ReadonlyArray): Set { const urls = new Set() for (const message of messages) { for (const part of message.parts) { @@ -155,7 +155,7 @@ function collectPreviewUrls(messages: Array): Set { } function latestPreview( - messages: Array, + messages: ReadonlyArray, skip: ReadonlySet = new Set(), ): string | null { for (let index = messages.length - 1; index >= 0; index--) { @@ -175,7 +175,10 @@ function latestPreview( return null } -function hasUserText(messages: Array, text: string): boolean { +function hasUserText( + messages: ReadonlyArray, + text: string, +): boolean { return messages.some( (message) => message.role === 'user' && @@ -683,7 +686,7 @@ function MessageList({ messages, isLoading, }: { - messages: Array + messages: ReadonlyArray isLoading: boolean }) { const [node, setNode] = useState(null) diff --git a/examples/ts-react-chat/src/routes/capability-demo.tsx b/examples/ts-react-chat/src/routes/capability-demo.tsx index 7625084e61..7275a865d3 100644 --- a/examples/ts-react-chat/src/routes/capability-demo.tsx +++ b/examples/ts-react-chat/src/routes/capability-demo.tsx @@ -11,7 +11,7 @@ interface SessionConfig { seq: number } -function Messages({ messages }: { messages: Array }) { +function Messages({ messages }: { messages: ReadonlyArray }) { const visible = messages.filter((m) => m.parts.some((p) => p.type === 'text' && p.content.trim()), ) diff --git a/examples/ts-react-chat/src/routes/example.runtime-context.tsx b/examples/ts-react-chat/src/routes/example.runtime-context.tsx index 5a0366fa7d..d48651f0c0 100644 --- a/examples/ts-react-chat/src/routes/example.runtime-context.tsx +++ b/examples/ts-react-chat/src/routes/example.runtime-context.tsx @@ -123,7 +123,7 @@ function RuntimeContextToolResult({ function RuntimeMessages({ messages, }: { - messages: Array + messages: ReadonlyArray }) { const messagesContainerRef = useRef(null) const visibleMessages = messages.filter((message) => diff --git a/examples/ts-react-chat/src/routes/index.tsx b/examples/ts-react-chat/src/routes/index.tsx index 83e2469bdf..307b0c79a6 100644 --- a/examples/ts-react-chat/src/routes/index.tsx +++ b/examples/ts-react-chat/src/routes/index.tsx @@ -109,7 +109,7 @@ function Messages({ messages, interrupts, }: { - messages: Array + messages: ReadonlyArray interrupts: BoundInterrupts }) { const messagesContainerRef = useRef(null) diff --git a/examples/ts-react-chat/src/routes/mcp-apps.tsx b/examples/ts-react-chat/src/routes/mcp-apps.tsx index dc1041c899..23c5915c34 100644 --- a/examples/ts-react-chat/src/routes/mcp-apps.tsx +++ b/examples/ts-react-chat/src/routes/mcp-apps.tsx @@ -82,7 +82,7 @@ function Messages({ messages, bridge, }: { - messages: Array + messages: ReadonlyArray bridge: McpAppBridge }) { const visible = messages.filter((m) => diff --git a/examples/ts-react-chat/src/routes/mcp-demo.tsx b/examples/ts-react-chat/src/routes/mcp-demo.tsx index f3d4e1f8bf..71f0d05058 100644 --- a/examples/ts-react-chat/src/routes/mcp-demo.tsx +++ b/examples/ts-react-chat/src/routes/mcp-demo.tsx @@ -135,7 +135,7 @@ function ToolResultView({ part }: { part: ToolResultPart }) { ) } -function Messages({ messages }: { messages: Array }) { +function Messages({ messages }: { messages: ReadonlyArray }) { const messagesContainerRef = useRef(null) const visibleMessages = messages.filter((message) => diff --git a/examples/ts-react-chat/src/routes/queueing.tsx b/examples/ts-react-chat/src/routes/queueing.tsx index 351dcda80d..a6ffbe8110 100644 --- a/examples/ts-react-chat/src/routes/queueing.tsx +++ b/examples/ts-react-chat/src/routes/queueing.tsx @@ -54,9 +54,9 @@ function textOf(message: UIMessage): string { } interface PanelChat { - messages: Array + messages: ReadonlyArray isLoading: boolean - queue: Array + queue: ReadonlyArray cancelQueued: (id: string) => void stop: () => void error: Error | undefined diff --git a/examples/ts-react-chat/src/routes/repo-report.tsx b/examples/ts-react-chat/src/routes/repo-report.tsx index a88e227ccd..aee0310bbc 100644 --- a/examples/ts-react-chat/src/routes/repo-report.tsx +++ b/examples/ts-react-chat/src/routes/repo-report.tsx @@ -29,7 +29,7 @@ export const Route = createFileRoute('/repo-report')({ }) function reportFromMessages( - messages: Array, + messages: ReadonlyArray, ): RepoReportCard | undefined { for (const message of [...messages].reverse()) { for (const part of [...message.parts].reverse()) { diff --git a/examples/ts-react-chat/src/routes/sandboxes-durable.tsx b/examples/ts-react-chat/src/routes/sandboxes-durable.tsx index 340162f778..69bb1c2d18 100644 --- a/examples/ts-react-chat/src/routes/sandboxes-durable.tsx +++ b/examples/ts-react-chat/src/routes/sandboxes-durable.tsx @@ -380,7 +380,7 @@ function SandboxWaiting({ kind }: { kind: SandboxWaitKind }) { export function sandboxWaitKind( isLoading: boolean, hasActiveRun: boolean, - messages: Array, + messages: ReadonlyArray, ): SandboxWaitKind | false { // "Nothing has come back for the current turn yet." An EMPTY transcript counts, // and that is the case this used to get wrong: chat persistence saves the pending @@ -408,7 +408,7 @@ function Messages({ messages, waiting, }: { - messages: Array + messages: ReadonlyArray waiting: SandboxWaitKind | false }) { const containerRef = useRef(null) diff --git a/examples/ts-react-chat/src/routes/sandboxes.tsx b/examples/ts-react-chat/src/routes/sandboxes.tsx index a9a76237cb..625828a4b6 100644 --- a/examples/ts-react-chat/src/routes/sandboxes.tsx +++ b/examples/ts-react-chat/src/routes/sandboxes.tsx @@ -221,7 +221,7 @@ function SandboxWaiting({ kind }: { kind: SandboxWaitKind }) { function sandboxWaitKind( isLoading: boolean, - messages: Array, + messages: ReadonlyArray, ): SandboxWaitKind | false { if ( !isLoading || @@ -243,7 +243,7 @@ function Messages({ messages, waiting, }: { - messages: Array + messages: ReadonlyArray waiting: SandboxWaitKind | false }) { const containerRef = useRef(null) diff --git a/examples/ts-react-native-chat/src/App.tsx b/examples/ts-react-native-chat/src/App.tsx index 973bf42b26..0e8b8284b9 100644 --- a/examples/ts-react-native-chat/src/App.tsx +++ b/examples/ts-react-native-chat/src/App.tsx @@ -191,7 +191,7 @@ function getRecipeSource(part?: StructuredRecipePart): RecipeValue | undefined { } function findLatestRecipePart( - messages: Array, + messages: ReadonlyArray, ): StructuredRecipePart | undefined { for (const message of [...messages].reverse()) { for (const part of [...message.parts].reverse()) { @@ -417,7 +417,7 @@ function UserRequest({ message }: { message: UIMessage }) { ) } -function RequestHistory({ messages }: { messages: Array }) { +function RequestHistory({ messages }: { messages: ReadonlyArray }) { const userMessages = messages.filter((message) => message.role === 'user') if (userMessages.length === 0) return null diff --git a/examples/ts-solid-chat/package.json b/examples/ts-solid-chat/package.json index 5c031cf2f3..ebdeb83aac 100644 --- a/examples/ts-solid-chat/package.json +++ b/examples/ts-solid-chat/package.json @@ -27,7 +27,7 @@ "@tanstack/solid-router-devtools": "^1.139.10", "@tanstack/solid-router-ssr-query": "^1.139.10", "@tanstack/solid-start": "^1.168.57", - "@tanstack/solid-store": "^0.8.0", + "@tanstack/solid-store": "^0.11.1", "@tanstack/store": "^0.11.1", "highlight.js": "^11.11.1", "lucide-solid": "^0.554.0", diff --git a/examples/ts-solid-chat/src/routes/index.tsx b/examples/ts-solid-chat/src/routes/index.tsx index 05a2b1b749..7e215f42f0 100644 --- a/examples/ts-solid-chat/src/routes/index.tsx +++ b/examples/ts-solid-chat/src/routes/index.tsx @@ -30,7 +30,7 @@ function ChatInputArea(props: { children: JSXElement }) { } function Messages(props: { - messages: ChatMessages + messages: Readonly addToolApprovalResponse: (response: { id: string approved: boolean @@ -175,7 +175,7 @@ function Messages(props: { } function DebugPanel(props: { - messages: Array + messages: ReadonlyArray chunks: Array onClearChunks: () => void }) { diff --git a/examples/ts-svelte-chat/src/lib/components/Messages.svelte b/examples/ts-svelte-chat/src/lib/components/Messages.svelte index f45228e6f3..1eca32c504 100644 --- a/examples/ts-svelte-chat/src/lib/components/Messages.svelte +++ b/examples/ts-svelte-chat/src/lib/components/Messages.svelte @@ -8,7 +8,7 @@ import ThinkingPart from './ThinkingPart.svelte' interface Props { - messages: Array + messages: ReadonlyArray addToolApprovalResponse: (response: { id: string approved: boolean diff --git a/packages/ai-angular/src/inject-chat.ts b/packages/ai-angular/src/inject-chat.ts index 8065f36fe3..1514f1f22b 100644 --- a/packages/ai-angular/src/inject-chat.ts +++ b/packages/ai-angular/src/inject-chat.ts @@ -21,13 +21,9 @@ import type { StreamChunk, } from '@tanstack/ai' import type { - ChatClientState, ResolvableChatInterrupt, - ChatInterruptState, ChatResumeState, - ConnectionStatus, InferredClientContext, - QueuedMessage, SendMessageOptions, StructuredOutputPart, } from '@tanstack/ai-client' @@ -39,9 +35,6 @@ import type { UIMessage, } from './types' -const EMPTY_INTERRUPTS = Object.freeze([]) -const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) - export function injectChat< const TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, @@ -65,26 +58,6 @@ export function injectChat< const destroyRef = inject(DestroyRef) const injector = inject(Injector) - const messages = signal>>( - options.initialMessages || [], - ) - const isLoading = signal(false) - const hasOlderMessages = signal(false) - const error = signal(undefined) - const status = signal('ready') - const isSubscribed = signal(false) - const connectionStatus = signal('disconnected') - const sessionGenerating = signal(false) - const isHydrating = signal(false) - const queue = signal>([]) - const runId = signal(null) - const interruptState = signal>({ - interrupts: EMPTY_INTERRUPTS, - pendingInterrupts: EMPTY_INTERRUPTS, - interruptErrors: EMPTY_INTERRUPT_ERRORS, - resuming: false, - }) - // Reactive option sources. Plain values become constant computeds. const bodySource = options.body !== undefined ? toReactive(options.body) : undefined @@ -147,13 +120,7 @@ export function injectChat< onChunk: (chunk: StreamChunk) => options.onChunk?.(chunk), onFinish: (message) => options.onFinish?.(message), onError: (err) => options.onError?.(err), - onRunIdChange: (nextRunId) => runId.set(nextRunId), - // No `onResumeStateChange`: the run identity is surfaced as the `runId` - // signal (via `onRunIdChange`) and pending interrupts arrive through - // `onInterruptStateChange`, so there is nothing left for it to do — and it - // is not a public option here, matching the other framework packages. onInterruptStateChange: (nextInterruptState, context) => { - interruptState.set(nextInterruptState) options.onInterruptStateChange?.(nextInterruptState, context) }, tools: toolsSource?.(), @@ -165,25 +132,15 @@ export function injectChat< ...(options.streamProcessor !== undefined && { streamProcessor: options.streamProcessor, }), - onMessagesChange: (m: Array>) => { - messages.set(m) - hasOlderMessages.set(client.getHasOlderMessages()) - }, - onLoadingChange: (v: boolean) => isLoading.set(v), - onStatusChange: (v: ChatClientState) => status.set(v), - onErrorChange: (v: Error | undefined) => error.set(v), - onSubscriptionChange: (v: boolean) => isSubscribed.set(v), - onConnectionStatusChange: (v: ConnectionStatus) => connectionStatus.set(v), - onSessionGeneratingChange: (v: boolean) => sessionGenerating.set(v), - onHydratingChange: (v: boolean) => isHydrating.set(v), ...(options.queue !== undefined && { queue: options.queue }), - onQueueChange: (nextQueue: Array) => queue.set(nextQueue), }) - messages.set(client.getMessages()) - interruptState.set(client.getInterruptState()) - hasOlderMessages.set(client.getHasOlderMessages()) - isHydrating.set(client.getIsHydrating()) + const snapshot = signal(client.getSnapshot()) + const messages = computed(() => snapshot().messages) + const interruptState = computed(() => snapshot().interruptState) + const unsubscribeSnapshot = client.subscribeSnapshot(() => + snapshot.set(client.getSnapshot()), + ) // START TAILING HERE, not in the constructor. A client is idle until something // attaches it, so a client that gets built and thrown away never opens a @@ -242,6 +199,7 @@ export function injectChat< ) destroyRef.onDestroy(() => { + unsubscribeSnapshot() // Release the connection first: the counterpart of the `attach` above. client.detach() if (liveSource?.()) { @@ -303,11 +261,10 @@ export function injectChat< } const loadOlderMessages = async () => { await client.loadOlderMessages() - hasOlderMessages.set(client.getHasOlderMessages()) } const stop = () => client.stop() const clear = () => client.clear() - const setMessages = (m: Array>) => + const setMessages = (m: ReadonlyArray>) => client.setMessagesManually(m) const addToolResult = async (result: { toolCallId: string @@ -350,27 +307,27 @@ export function injectChat< // oxlint-disable-next-line eslint-js/no-restricted-syntax -- return shape diverges from conditional InjectChatResult; TS can't structurally narrow the TSchema-gated partial/final signals return { - messages: messages.asReadonly(), + messages, sendMessage, - queue: queue.asReadonly(), + queue: computed(() => snapshot().queue), cancelQueued, append, reload, stop, - isLoading: isLoading.asReadonly(), - hasOlderMessages: hasOlderMessages.asReadonly(), + isLoading: computed(() => snapshot().isLoading), + hasOlderMessages: computed(() => snapshot().hasOlderMessages), loadOlderMessages, - error: error.asReadonly(), - status: status.asReadonly(), - isSubscribed: isSubscribed.asReadonly(), - connectionStatus: connectionStatus.asReadonly(), - sessionGenerating: sessionGenerating.asReadonly(), - isHydrating: isHydrating.asReadonly(), + error: computed(() => snapshot().error), + status: computed(() => snapshot().status), + isSubscribed: computed(() => snapshot().isSubscribed), + connectionStatus: computed(() => snapshot().connectionStatus), + sessionGenerating: computed(() => snapshot().sessionGenerating), + isHydrating: computed(() => snapshot().isHydrating), setMessages, clear, addToolResult, addToolApprovalResponse, - runId: runId.asReadonly(), + runId: computed(() => snapshot().runId), interrupts, pendingInterrupts, interruptErrors, diff --git a/packages/ai-angular/src/inject-generate-video.ts b/packages/ai-angular/src/inject-generate-video.ts index 3287fbea29..43f19d7afb 100644 --- a/packages/ai-angular/src/inject-generate-video.ts +++ b/packages/ai-angular/src/inject-generate-video.ts @@ -181,27 +181,6 @@ export function injectGenerateVideo( onStatusUpdate: (s: VideoStatusInfo) => { if (!disposed) options.onStatusUpdate?.(s) }, - onResultChange: (r: TOutput | null) => { - if (!disposed) result.set(r) - }, - onLoadingChange: (l: boolean) => { - if (!disposed) isLoading.set(l) - }, - onErrorChange: (e: Error | undefined) => { - if (!disposed) error.set(e) - }, - onStatusChange: (s: GenerationClientState) => { - if (!disposed) status.set(s) - }, - onJobIdChange: (id: string | null) => { - if (!disposed) jobId.set(id) - }, - onVideoStatusChange: (s: VideoStatusInfo | null) => { - if (!disposed) videoStatus.set(s) - }, - onResumeStateChange: (rs: { runId: string } | null) => { - if (!disposed) runId.set(rs?.runId ?? null) - }, } let client: VideoGenerationClient @@ -221,6 +200,19 @@ export function injectGenerateVideo( ) } + const applySnapshot = () => { + const next = client.getSnapshot() + result.set(next.result) + isLoading.set(next.isLoading) + error.set(next.error) + status.set(next.status) + runId.set(next.runId) + jobId.set(next.jobId) + videoStatus.set(next.videoStatus) + } + applySnapshot() + const unsubscribeSnapshot = client.subscribe(applySnapshot) + if (bodySource) { effect( () => { @@ -242,6 +234,7 @@ export function injectGenerateVideo( ) destroyRef.onDestroy(() => { disposed = true + unsubscribeSnapshot() client.dispose() }) diff --git a/packages/ai-angular/src/inject-generation.ts b/packages/ai-angular/src/inject-generation.ts index dfd095e079..a103f38744 100644 --- a/packages/ai-angular/src/inject-generation.ts +++ b/packages/ai-angular/src/inject-generation.ts @@ -205,21 +205,6 @@ export function injectGeneration< onChunk: (c: StreamChunk) => { if (!disposed) options.onChunk?.(c) }, - onResultChange: (r: TOutput | null) => { - if (!disposed) result.set(r) - }, - onLoadingChange: (l: boolean) => { - if (!disposed) isLoading.set(l) - }, - onErrorChange: (e: Error | undefined) => { - if (!disposed) error.set(e) - }, - onStatusChange: (s: GenerationClientState) => { - if (!disposed) status.set(s) - }, - onResumeStateChange: (rs) => { - if (!disposed) runId.set(rs?.runId ?? null) - }, } const persistenceProps = @@ -253,6 +238,17 @@ export function injectGeneration< ) } + const applySnapshot = () => { + const next = client.getSnapshot() + result.set(next.result) + isLoading.set(next.isLoading) + error.set(next.error) + status.set(next.status) + runId.set(next.runId) + } + applySnapshot() + const unsubscribeSnapshot = client.subscribe(applySnapshot) + if (bodySource) { effect( () => { @@ -274,6 +270,7 @@ export function injectGeneration< ) destroyRef.onDestroy(() => { disposed = true + unsubscribeSnapshot() client.dispose() }) diff --git a/packages/ai-angular/src/types.ts b/packages/ai-angular/src/types.ts index 6a6da218a3..fc044ae599 100644 --- a/packages/ai-angular/src/types.ts +++ b/packages/ai-angular/src/types.ts @@ -136,7 +136,7 @@ interface BaseInjectChatResult< readonly [], > { /** Current messages in the conversation. */ - messages: Signal>> + messages: Signal>> /** * Send a message (string or multimodal content). * Pass `{ whenBusy }` to override the queue policy for a single send, or @@ -218,7 +218,7 @@ interface BaseInjectChatResult< /** Current error, if any. */ error: Signal /** Set messages manually. */ - setMessages: (messages: Array>) => void + setMessages: (messages: ReadonlyArray>) => void /** Clear all messages. */ clear: () => void /** Current generation status. */ diff --git a/packages/ai-angular/tests/inject-chat.test.ts b/packages/ai-angular/tests/inject-chat.test.ts index ee23b0b103..606a8bf759 100644 --- a/packages/ai-angular/tests/inject-chat.test.ts +++ b/packages/ai-angular/tests/inject-chat.test.ts @@ -391,3 +391,51 @@ describe('injectChat — structured output', () => { expect(result.final()).toBeNull() }) }) + +describe('injectChat snapshot', () => { + it('reads one frozen client snapshot and drops it on destroy', async () => { + let resolveRead!: (value: { + messages: Array<{ + id: string + role: 'user' + parts: Array<{ type: 'text'; content: string }> + }> + }) => void + const read = new Promise[0]>((resolve) => { + resolveRead = resolve + }) + const { result, flush, destroy } = renderInjectChat({ + connection: createMockConnectionAdapter({ + chunks: createTextChunks('Hi'), + }), + threadId: 'angular-snapshot', + persistence: { + getItem: () => read, + setItem: async () => {}, + removeItem: async () => {}, + }, + }) + flush() + + expect(result.isHydrating()).toBe(true) + resolveRead({ + messages: [ + { id: 'm1', role: 'user', parts: [{ type: 'text', content: 'hi' }] }, + ], + }) + await vi.waitFor(() => expect(result.isHydrating()).toBe(false)) + expect(result.messages()).toHaveLength(1) + + await result.sendMessage('Hello') + expect(result.isLoading()).toBe(false) + expect(result.status()).toBe('ready') + const messages = result.messages() + expect(messages).toHaveLength(3) + expect(Object.isFrozen(messages)).toBe(true) + expect(() => (messages as Array).push({})).toThrow(TypeError) + + destroy() + result.setMessages([]) + expect(result.messages()).toBe(messages) + }) +}) diff --git a/packages/ai-anthropic/src/adapters/text.ts b/packages/ai-anthropic/src/adapters/text.ts index d850bbe1f3..d02313231b 100644 --- a/packages/ai-anthropic/src/adapters/text.ts +++ b/packages/ai-anthropic/src/adapters/text.ts @@ -1512,7 +1512,13 @@ export class AnthropicTextAdapter< let accumulatedSignature = '' const toolCallsMap = new Map< number, - { id: string; name: string; input: string; started: boolean } + { + id: string + name: string + input: string + started: boolean + hasInputDelta: boolean + } >() let currentToolIndex = -1 // Server-side tools share the `input_json_delta` wire format with client @@ -1577,8 +1583,11 @@ export class AnthropicTextAdapter< toolCallsMap.set(currentToolIndex, { id: event.content_block.id, name: event.content_block.name, - input: '', + // The initial input snapshot. The first `input_json_delta` + // replaces it. + input: JSON.stringify(event.content_block.input) ?? '', started: false, + hasInputDelta: false, }) } else if (event.content_block.type === 'server_tool_use') { currentServerTool = { @@ -1851,6 +1860,10 @@ export class AnthropicTextAdapter< } } + if (!existing.hasInputDelta) { + existing.input = '' + existing.hasInputDelta = true + } existing.input += event.delta.partial_json yield { @@ -1932,13 +1945,14 @@ export class AnthropicTextAdapter< } } - // Emit TOOL_CALL_END - let parsedInput: unknown = {} + // Emit TOOL_CALL_END. Malformed JSON sends no `input`, so core + // reports the error from the raw `args`. + let parsedInput: unknown try { const parsed = existing.input ? JSON.parse(existing.input) : {} parsedInput = parsed && typeof parsed === 'object' ? parsed : {} } catch { - parsedInput = {} + // Keep `parsedInput` undefined for the tool error. } yield { @@ -1948,7 +1962,8 @@ export class AnthropicTextAdapter< toolName: existing.name, model, timestamp: Date.now(), - input: parsedInput, + args: existing.input, + ...(parsedInput !== undefined && { input: parsedInput }), } // Reset so a new TEXT_MESSAGE_START is emitted if text follows tool calls diff --git a/packages/ai-anthropic/tests/anthropic-adapter.test.ts b/packages/ai-anthropic/tests/anthropic-adapter.test.ts index 86e950aa5a..4c51a2db89 100644 --- a/packages/ai-anthropic/tests/anthropic-adapter.test.ts +++ b/packages/ai-anthropic/tests/anthropic-adapter.test.ts @@ -12,6 +12,7 @@ import { AnthropicTextAdapter } from '../src/adapters/text' import type { AnthropicTextProviderOptions } from '../src/adapters/text' import type { AnthropicDocumentMetadata } from '../src/message-types' import { ANTHROPIC_MAX_NONSTREAMING_TOKENS } from '../src/model-meta' +import { createSilentLogger } from './utils/logger' import { z } from 'zod' const mocks = vi.hoisted(() => { @@ -2520,6 +2521,98 @@ describe('Anthropic stream processing', () => { 'text B', ]) }) + + /** The TOOL_CALL_END of one streamed `tool_use` block. */ + async function toolCallEnd( + initialInput: Record | undefined, + deltas: Array, + ) { + mocks.betaMessagesCreate.mockResolvedValueOnce( + (async function* () { + yield { + type: 'content_block_start', + index: 0, + content_block: { + type: 'tool_use', + id: 'toolu_raw', + name: 'inspect', + ...(initialInput !== undefined && { input: initialInput }), + }, + } + for (const partial_json of deltas) { + yield { + type: 'content_block_delta', + index: 0, + delta: { type: 'input_json_delta', partial_json }, + } + } + yield { type: 'content_block_stop', index: 0 } + yield { + type: 'message_delta', + delta: { stop_reason: 'tool_use' }, + usage: { output_tokens: 1 }, + } + yield { type: 'message_stop' } + })(), + ) + + const chunks: Array = [] + for await (const chunk of createAdapter('claude-opus-4-1').chatStream({ + model: 'claude-opus-4-1', + messages: [{ role: 'user', content: 'Go' }], + logger: createSilentLogger(), + })) { + chunks.push(chunk) + } + return chunks.find((chunk) => chunk.type === 'TOOL_CALL_END') + } + + it.each([ + { + name: 'the streamed input', + initial: {}, + deltas: ['{"count":2}'], + args: '{"count":2}', + input: { count: 2 }, + }, + { + name: 'the streamed bytes as sent', + initial: {}, + deltas: ['{"n": ', '1.50}'], + args: '{"n": 1.50}', + input: { n: 1.5 }, + }, + { + name: 'the initial input snapshot when no delta arrives', + initial: { value: 'start' }, + deltas: [], + args: '{"value":"start"}', + input: { value: 'start' }, + }, + { + name: 'the streamed input over the initial snapshot', + initial: { value: 'start' }, + deltas: ['{"count":2}'], + args: '{"count":2}', + input: { count: 2 }, + }, + ])( + 'sends $name as the raw TOOL_CALL_END args', + async ({ initial, deltas, args, input }) => { + expect(await toolCallEnd(initial, deltas)).toMatchObject({ args, input }) + }, + ) + + it('sends malformed tool input raw on TOOL_CALL_END with no input', async () => { + const end = await toolCallEnd({}, ['{"a":']) + expect(end).toMatchObject({ args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + + it('sends empty input as {} with empty raw args', async () => { + const end = await toolCallEnd(undefined, []) + expect(end).toMatchObject({ input: {}, args: '' }) + }) }) describe('Anthropic adapter error handling', () => { diff --git a/packages/ai-bedrock/src/converse/message-converter.ts b/packages/ai-bedrock/src/converse/message-converter.ts index 602a0adf15..7900075848 100644 --- a/packages/ai-bedrock/src/converse/message-converter.ts +++ b/packages/ai-bedrock/src/converse/message-converter.ts @@ -276,11 +276,10 @@ function messageToBlocks( } // null → no text blocks - // Append toolUse blocks for assistant tool calls. Malformed or non-object - // arguments come from a prior assistant turn the engine already accepted, so - // they signal a real upstream problem — throw rather than silently coercing to - // `{}` and forwarding a corrupted tool call (the adapter's catch surfaces it - // as a RUN_ERROR). + // Append toolUse blocks for assistant tool calls. Malformed JSON stays in + // the history next to its tool error, so it goes out as `{}`. Non-object + // arguments signal a real upstream problem and throw (the adapter's catch + // surfaces it as a RUN_ERROR). if (msg.role === 'assistant' && msg.toolCalls) { for (const call of msg.toolCalls) blocks.push(toolUseBlock(call)) } @@ -292,13 +291,11 @@ function toolUseBlock( call: NonNullable[number], ): ContentBlock { const rawArguments = call.function.arguments || '{}' - let parsed: unknown + let parsed: unknown = {} try { parsed = JSON.parse(rawArguments) - } catch (error: unknown) { - throw new Error( - `Bedrock Converse: tool call "${call.function.name}" has malformed JSON arguments (${String(error)}). Raw: ${rawArguments}`, - ) + } catch { + // The tool result already told the model the arguments were malformed. } if (parsed === null || typeof parsed !== 'object' || Array.isArray(parsed)) { throw new Error( diff --git a/packages/ai-bedrock/src/converse/stream-processor.ts b/packages/ai-bedrock/src/converse/stream-processor.ts index 911d90d7e3..2be2caa5cb 100644 --- a/packages/ai-bedrock/src/converse/stream-processor.ts +++ b/packages/ai-bedrock/src/converse/stream-processor.ts @@ -72,6 +72,7 @@ export async function* processConverseStream( id: string name: string started: boolean + arguments: string } >() @@ -167,6 +168,7 @@ export async function* processConverseStream( id, name, started: true, + arguments: '', }) yield { type: EventType.TOOL_CALL_START, @@ -189,6 +191,7 @@ export async function* processConverseStream( if (delta && 'toolUse' in delta && delta.toolUse?.input !== undefined) { const toolCall = toolCallsByIndex.get(index) if (toolCall?.started) { + toolCall.arguments += delta.toolUse.input yield { type: EventType.TOOL_CALL_ARGS, toolCallId: toolCall.id, @@ -295,6 +298,7 @@ export async function* processConverseStream( yield { type: EventType.TOOL_CALL_END, toolCallId: toolCall.id, + args: toolCall.arguments, toolCallName: toolCall.name, toolName: toolCall.name, timestamp: Date.now(), @@ -337,6 +341,7 @@ export async function* processConverseStream( yield { type: EventType.TOOL_CALL_END, toolCallId: toolCall.id, + args: toolCall.arguments, toolCallName: toolCall.name, toolName: toolCall.name, timestamp: Date.now(), diff --git a/packages/ai-bedrock/tests/converse/message-converter.test.ts b/packages/ai-bedrock/tests/converse/message-converter.test.ts index 0001e8b107..95409cfa18 100644 --- a/packages/ai-bedrock/tests/converse/message-converter.test.ts +++ b/packages/ai-bedrock/tests/converse/message-converter.test.ts @@ -117,22 +117,25 @@ describe('toConverseMessages', () => { }) }) - it('throws on malformed JSON tool-call arguments instead of forwarding {}', () => { - expect(() => - toConverseMessages([ - { - role: 'assistant', - content: '', - toolCalls: [ - { - id: 't1', - type: 'function', - function: { name: 'getX', arguments: '{"a":' }, - }, - ], - }, - ]), - ).toThrow(/malformed JSON/) + it('sends malformed JSON tool-call arguments as {}', () => { + const { + messages: [message], + } = toConverseMessages([ + { + role: 'assistant', + content: '', + toolCalls: [ + { + id: 't1', + type: 'function', + function: { name: 'getX', arguments: '{"a":' }, + }, + ], + }, + ]) + expect(message?.content).toEqual([ + { toolUse: { toolUseId: 't1', name: 'getX', input: {} } }, + ]) }) it('throws when tool-call arguments parse to a non-object', () => { diff --git a/packages/ai-bedrock/tests/converse/stream-processor.test.ts b/packages/ai-bedrock/tests/converse/stream-processor.test.ts index 620b304d45..9d546c403c 100644 --- a/packages/ai-bedrock/tests/converse/stream-processor.test.ts +++ b/packages/ai-bedrock/tests/converse/stream-processor.test.ts @@ -377,4 +377,66 @@ describe('processConverseStream', () => { events.filter((e) => e.type === EventType.TOOL_CALL_END), ).toHaveLength(2) }) + + it('puts the raw streamed argument string on TOOL_CALL_END', async () => { + const events = await collect( + { + contentBlockStart: { + start: { toolUse: { toolUseId: 't1', name: 'getX' } }, + contentBlockIndex: 0, + }, + }, + { + contentBlockDelta: { + delta: { toolUse: { input: '{"n": ' } }, + contentBlockIndex: 0, + }, + }, + { + contentBlockDelta: { + delta: { toolUse: { input: '1.50}' } }, + contentBlockIndex: 0, + }, + }, + { contentBlockStop: { contentBlockIndex: 0 } }, + { messageStop: { stopReason: 'tool_use' } }, + ) + const end = events.find((e) => e.type === EventType.TOOL_CALL_END) + expect(end).toMatchObject({ toolCallId: 't1', args: '{"n": 1.50}' }) + }) + + it('puts malformed raw arguments on TOOL_CALL_END without an input', async () => { + const events = await collect( + { + contentBlockStart: { + start: { toolUse: { toolUseId: 't1', name: 'getX' } }, + contentBlockIndex: 0, + }, + }, + { + contentBlockDelta: { + delta: { toolUse: { input: '{"a":' } }, + contentBlockIndex: 0, + }, + }, + // Truncated: the drained END still carries the raw string. + ) + const end = events.find((e) => e.type === EventType.TOOL_CALL_END) + expect(end).toMatchObject({ toolCallId: 't1', args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + + it('sends empty args on TOOL_CALL_END when no argument fragment arrived', async () => { + const events = await collect( + { + contentBlockStart: { + start: { toolUse: { toolUseId: 't1', name: 'getX' } }, + contentBlockIndex: 0, + }, + }, + { contentBlockStop: { contentBlockIndex: 0 } }, + ) + const end = events.find((e) => e.type === EventType.TOOL_CALL_END) + expect(end).toMatchObject({ toolCallId: 't1', args: '' }) + }) }) diff --git a/packages/ai-client/package.json b/packages/ai-client/package.json index 9692ffe9d8..bc589d2ede 100644 --- a/packages/ai-client/package.json +++ b/packages/ai-client/package.json @@ -71,7 +71,8 @@ "dependencies": { "@tanstack/ai": "workspace:^", "@tanstack/ai-event-client": "workspace:^", - "@tanstack/ai-utils": "workspace:^" + "@tanstack/ai-utils": "workspace:^", + "@tanstack/store": "^0.11.1" }, "devDependencies": { "@standard-schema/spec": "^1.1.0", diff --git a/packages/ai-client/src/audio-recorder.ts b/packages/ai-client/src/audio-recorder.ts index 0f2b79737f..d719dada87 100644 --- a/packages/ai-client/src/audio-recorder.ts +++ b/packages/ai-client/src/audio-recorder.ts @@ -1,5 +1,6 @@ import { arrayBufferToBase64 } from '@tanstack/ai-utils' import type { AudioPart } from '@tanstack/ai/client' +import { notifyListeners } from './snapshot-atom' /** Lifecycle state of an {@link AudioRecorder}. */ export type AudioRecorderState = 'idle' | 'recording' | 'stopping' @@ -63,12 +64,12 @@ export class AudioRecorder { private chunks: Array = [] private startedAt = 0 private _state: AudioRecorderState = 'idle' + private readonly listeners = new Set<(state: AudioRecorderState) => void>() // True while start() is awaiting getUserMedia (state is still 'idle' then). private starting = false // Set by cancel()/teardown during that window so start() releases the // freshly acquired stream instead of beginning a leaked recording. private pendingCancel = false - private readonly listeners = new Set<(state: AudioRecorderState) => void>() private stopResolve: ((recording: AudioRecording) => void) | null = null private stopReject: ((error: Error) => void) | null = null @@ -92,20 +93,19 @@ export class AudioRecorder { subscribe(cb: (state: AudioRecorderState) => void): () => void { this.listeners.add(cb) - return () => { - this.listeners.delete(cb) - } + return () => this.listeners.delete(cb) } + /** Current recorder lifecycle state. */ + getSnapshot = (): AudioRecorderState => this._state + private setState(state: AudioRecorderState): void { this._state = state - for (const cb of this.listeners) { - cb(state) - } + notifyListeners(this.listeners, state) } async start(): Promise { - if (this._state !== 'idle' || this.starting) { + if (this.state !== 'idle' || this.starting) { return } this.starting = true @@ -168,7 +168,7 @@ export class AudioRecorder { } stop(): Promise { - if (this._state !== 'recording' || !this.recorder) { + if (this.state !== 'recording' || !this.recorder) { return Promise.reject( new Error('AudioRecorder.stop() called while not recording'), ) @@ -179,7 +179,7 @@ export class AudioRecorder { // Some browsers/codecs never fire onstop; this watchdog unwedges the // recorder instead of leaking this promise forever. const watchdog = setTimeout(() => { - if (this._state !== 'stopping') { + if (this.state !== 'stopping') { return } // Detach handlers so a late onstop/onerror from the stalled recorder @@ -215,7 +215,7 @@ export class AudioRecorder { this.pendingCancel = true return } - if (this._state === 'idle') { + if (this.state === 'idle') { return } const recorder = this.recorder @@ -244,8 +244,9 @@ export class AudioRecorder { const reject = this.stopReject this.stopResolve = null this.stopReject = null - this.setState('idle') + // Settle first. The pending stop() must not depend on listener code. reject?.(new Error('Recording cancelled')) + this.setState('idle') } private async finalize(): Promise { @@ -270,8 +271,8 @@ export class AudioRecorder { const resolve = this.stopResolve this.stopResolve = null this.stopReject = null - this.setState('idle') resolve?.(recording) + this.setState('idle') } catch (err) { this.handleError( err instanceof Error ? err : new Error('Failed to finalize recording'), @@ -286,11 +287,11 @@ export class AudioRecorder { const reject = this.stopReject this.stopResolve = null this.stopReject = null - this.setState('idle') - // Settle the pending stop() promise before invoking the user callback so a - // throwing onError can't strand the awaiter (the two error channels are - // independent — see start()/stop() docs). + // Settle the pending stop() promise before listeners and the user callback + // run, so a throw there can't strand the awaiter (the two error channels + // are independent — see start()/stop() docs). reject?.(error) + this.setState('idle') this.notifyError(error) } diff --git a/packages/ai-client/src/byok/client.ts b/packages/ai-client/src/byok/client.ts index a61867397e..1514172e18 100644 --- a/packages/ai-client/src/byok/client.ts +++ b/packages/ai-client/src/byok/client.ts @@ -4,6 +4,7 @@ import { isProviderId, maskKey, } from '@tanstack/ai/byok' +import { notifyListeners } from '../snapshot-atom' import { memoryStorage } from './storage' import type { ByokProvider, ProviderId } from '@tanstack/ai/byok' import type { Keyring, KeyringStorage } from './storage' @@ -108,9 +109,9 @@ export class ByokClient { #coverageAll = false #coverage: Record = {} readonly #companions: Record> = {} - readonly #listeners = new Set<() => void>() - #snapshot: ByokSnapshot #storageError: string | null = null + #snapshot: ByokSnapshot + readonly #listeners = new Set<() => void>() readonly #ready: Promise constructor(options: DefineByokOptions = {}) { @@ -136,20 +137,25 @@ export class ByokClient { subscribe = (listener: () => void): (() => void) => { this.#listeners.add(listener) - return () => { - this.#listeners.delete(listener) - } + return () => this.#listeners.delete(listener) } getSnapshot = (): ByokSnapshot => this.#snapshot #buildSnapshot(): ByokSnapshot { - return { - status: { ...this.#statuses }, + return Object.freeze({ + status: Object.freeze( + Object.fromEntries( + Object.entries(this.#statuses).map(([key, value]) => [ + key, + value ? Object.freeze({ ...value }) : value, + ]), + ), + ), locked: this.#locked, - prompt: this.#prompt, + prompt: this.#prompt ? Object.freeze({ ...this.#prompt }) : null, storageError: this.#storageError, - } + }) } keys(): Keyring { @@ -393,7 +399,7 @@ export class ByokClient { #emit(): void { this.#snapshot = this.#buildSnapshot() - for (const listener of this.#listeners) listener() + notifyListeners(this.#listeners) } } diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index 4a3c19c544..44138c6a0f 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -57,9 +57,18 @@ import type { ChatDevtoolsBridge, ChatDevtoolsBridgeOptions, } from './devtools' +import { + createAtom, + freezeSnapshotMessages, + freezeSnapshotPart, + patchAtom, + subscribeAtom, +} from './snapshot-atom' +import type { Atom } from './snapshot-atom' import type { BoundInterrupts, ChatClientOptions, + ChatClientSnapshot, ChatClientState, ChatFetcher, ChatInterruptState, @@ -374,6 +383,29 @@ function rebuildsAssistantMessage(chunk: StreamChunk): boolean { type SubagentCard = Extract +/** + * Snapshot messages are frozen, and a host can pass them back in + * (`setMessagesManually(snapshot.messages)`, `initialMessages`). The client + * writes subagent handles into message parts, nested cards included, so copy + * each frozen part down through subagent transcripts. + */ +function thawMessages(messages: ReadonlyArray): Array { + return messages.map((message) => ({ + ...message, + parts: message.parts.map((part) => { + if (!Object.isFrozen(part)) return part + if (part.type !== 'subagent') return { ...part } + return { + ...part, + subagent: { + ...part.subagent, + messages: thawMessages(part.subagent.messages), + }, + } + }), + })) +} + /** Every subagent card in the messages, nested cards included. */ function collectSubagentParts( messages: ReadonlyArray, @@ -434,7 +466,54 @@ export class ChatClient< // without a storage adapter. private readonly persistor?: ChatPersistor private readonly clearedStreamTracker = new ClearedStreamTracker() - private currentRunId: string | null = null + private readonly snapshotAtom: Atom> = + createAtom>({ + messages: [], + status: 'ready', + isLoading: false, + error: undefined, + isSubscribed: false, + connectionStatus: 'disconnected', + sessionGenerating: false, + queue: [], + runId: null, + hasOlderMessages: false, + isHydrating: false, + interruptState: { + interrupts: Object.freeze([]), + pendingInterrupts: Object.freeze([]), + interruptErrors: Object.freeze([]), + resuming: false, + } as ChatInterruptState, + }) + // The snapshot atom is the only copy of this state. Setters patch it. + private get status() { + return this.snapshotAtom.get().status + } + private get isLoading() { + return this.snapshotAtom.get().isLoading + } + private get error() { + return this.snapshotAtom.get().error + } + private get isSubscribed() { + return this.snapshotAtom.get().isSubscribed + } + private get connectionStatus() { + return this.snapshotAtom.get().connectionStatus + } + private get sessionGenerating() { + return this.snapshotAtom.get().sessionGenerating + } + private get hasOlderMessages() { + return this.snapshotAtom.get().hasOlderMessages + } + private get isHydrating() { + return this.snapshotAtom.get().isHydrating + } + private get currentRunId() { + return this.snapshotAtom.get().runId + } // Interrupt-resume tracking: the run/thread of the most recent interrupted // run, so approvals/client-tool results can be sent back. Cleared when the // run terminates. This is STATE (interrupt) resume, not delivery/cursor. @@ -494,11 +573,6 @@ export class ChatClient< * deliverers cannot both append a user message before only one stream runs. */ private deliverClaim = false - private isLoading = false - private isSubscribed = false - private error: Error | undefined = undefined - private status: ChatClientState = 'ready' - private connectionStatus: ConnectionStatus = 'disconnected' private abortController: AbortController | null = null private readonly clientToolsRef: { current: Map } private readonly devtoolsBridge: ChatDevtoolsBridge @@ -538,7 +612,6 @@ export class ChatClient< // continuationPending was true (chained approval scenario) private continuationSkipped = false private draining = false - private sessionGenerating = false private readonly activeRunIds = new Set() /** Latched by `dispose()`; stops any late async callback starting new work. */ private disposed = false @@ -550,8 +623,6 @@ export class ChatClient< private hydrationInFlight: number | undefined /** An async storage adapter is still reading the persisted transcript. */ private storeHydrating = false - /** See {@link getIsHydrating}. */ - private isHydrating = false /** Constructor inputs `attach()` needs on every re-attach, not just the first. */ private readonly rejoinRunId: string | null | undefined private readonly cachesMessages: boolean @@ -560,7 +631,6 @@ export class ChatClient< * `persistence === true`. */ private readonly historyPageSize: number | undefined - private hasOlderMessages = false private olderMessagesCursor: string | undefined private readonly knownServerMessageIds = new Set() private loadOlderMessagesInFlight = false @@ -721,7 +791,7 @@ export class ChatClient< // hydrates from the server on mount, keyed by threadId.) const initialMessages = syncPersistedState ? syncPersistedState.messages - : options.initialMessages + : options.initialMessages && thawMessages(options.initialMessages) // A durable snapshot read synchronously from storage wins over the // in-memory `initialResumeSnapshot` fallback applied above. A snapshot with // pending interrupts rehydrates the interrupt UI; a bare in-flight run is @@ -767,6 +837,12 @@ export class ChatClient< // its live handle before anyone reads the messages. this.syncSubagentHandles() this.persistor?.notifyMessagesChanged(messages) + patchAtom(this.snapshotAtom, { + messages: freezeSnapshotMessages( + this.frozenMessages, + messages as Array>, + ), + }) this.callbacksRef.current.onMessagesChange(messages) // The bridge only snapshots on status changes, and the panel draws a // child card from the snapshot, so without this the card would update @@ -1014,8 +1090,11 @@ export class ChatClient< this.cachesMessages = cachesMessages // Server-authoritative mode loads the thread on the first `attach()`. // Start as hydrating so the first render shows loading, not an empty thread. - this.isHydrating = - this.storeHydrating || (!cachesMessages && !!this.connection.hydrate) + this.snapshotAtom.set( + this.readSnapshot( + this.storeHydrating || (!cachesMessages && !!this.connection.hydrate), + ), + ) // NO TAILING HERE, deliberately. Constructing a client must not open a // connection. // @@ -1575,7 +1654,7 @@ export class ChatClient< private setCurrentRunId(runId: string | null): void { if (this.currentRunId === runId) return - this.currentRunId = runId + patchAtom(this.snapshotAtom, { runId }) this.callbacksRef.current.onRunIdChange(runId) } @@ -1767,8 +1846,62 @@ export class ChatClient< return `${prefix}-${Date.now()}-${Math.random().toString(36).substring(7)}` } + /** Frozen copies of unchanged messages and parts, keyed by the live object. */ + private readonly frozenMessages = new WeakMap() + + private freezeSnapshotQueue( + queue: Array, + ): ReadonlyArray { + return Object.freeze( + queue.map((item) => { + if (typeof item.content !== 'object') return Object.freeze({ ...item }) + const content = Array.isArray(item.content.content) + ? Object.freeze( + item.content.content.map((part) => freezeSnapshotPart(part)), + ) + : item.content.content + return Object.freeze({ + ...item, + content: Object.freeze({ ...item.content, content }), + }) + }), + ) as ReadonlyArray + } + + /** The atom with messages, queue, and interrupt state read fresh. */ + private readSnapshot( + isHydrating: boolean, + ): ChatClientSnapshot { + return Object.freeze({ + ...this.snapshotAtom.get(), + messages: freezeSnapshotMessages( + this.frozenMessages, + this.processor.getMessages() as Array>, + ), + queue: this.freezeSnapshotQueue(this.getQueue()), + isHydrating, + interruptState: this.interruptManager.getState(), + }) + } + + /** + * Subscribe to UI snapshot changes. Does not fire with the current value; + * read {@link getSnapshot} first. + * + * Named separately from {@link subscribe} (the live connection loop). + */ + subscribeSnapshot = (listener: () => void): (() => void) => + subscribeAtom(this.snapshotAtom, listener) + + /** + * Current UI snapshot. The snapshot, `messages`, `queue`, and their entries + * are frozen copies. `interruptState` is the interrupt manager's own state. + */ + getSnapshot = (): ChatClientSnapshot => + this.snapshotAtom.get() + private setIsLoading(isLoading: boolean): void { - this.isLoading = isLoading + patchAtom(this.snapshotAtom, { isLoading }) this.callbacksRef.current.onLoadingChange(isLoading) this.events.loadingChanged(isLoading) if (isLoading) this.supersedeHydration() @@ -1786,19 +1919,19 @@ export class ChatClient< } private setStatus(status: ChatClientState): void { - this.status = status + patchAtom(this.snapshotAtom, { status }) this.callbacksRef.current.onStatusChange(status) this.devtoolsBridge.emitSnapshot() } private setIsSubscribed(isSubscribed: boolean): void { - this.isSubscribed = isSubscribed + patchAtom(this.snapshotAtom, { isSubscribed }) this.callbacksRef.current.onSubscriptionChange(isSubscribed) this.devtoolsBridge.emitSnapshot() } private setConnectionStatus(status: ConnectionStatus): void { - this.connectionStatus = status + patchAtom(this.snapshotAtom, { connectionStatus: status }) this.callbacksRef.current.onConnectionStatusChange(status) this.devtoolsBridge.emitSnapshot() } @@ -1807,13 +1940,13 @@ export class ChatClient< const isHydrating = this.hydrationInFlight !== undefined || this.storeHydrating if (this.isHydrating === isHydrating) return - this.isHydrating = isHydrating + patchAtom(this.snapshotAtom, { isHydrating }) this.callbacksRef.current.onHydratingChange(isHydrating) } private setSessionGenerating(isGenerating: boolean): void { if (this.sessionGenerating === isGenerating) return - this.sessionGenerating = isGenerating + patchAtom(this.snapshotAtom, { sessionGenerating: isGenerating }) this.callbacksRef.current.onSessionGeneratingChange(isGenerating) this.devtoolsBridge.emitSnapshot() } @@ -1823,6 +1956,7 @@ export class ChatClient< // Capture state before invoking callbacks so a synchronous nested change // cannot pair this publication's source with a later manager snapshot. const interruptState = this.interruptManager.getState() + patchAtom(this.snapshotAtom, { interruptState }) // Persist (or clear) the durable resume snapshot so a full page reload can // rehydrate pending interrupts and rejoin the run. Folded into the same // persistence adapter that stores messages (one record per chat). @@ -1864,7 +1998,7 @@ export class ChatClient< } private setError(error: Error | undefined): void { - this.error = error + patchAtom(this.snapshotAtom, { error }) this.callbacksRef.current.onErrorChange(error) this.events.errorChanged(error?.message || null) } @@ -3355,8 +3489,12 @@ export class ChatClient< } } + private setHasOlderMessages(hasOlderMessages: boolean) { + patchAtom(this.snapshotAtom, { hasOlderMessages }) + } + private resetHistoryPaging() { - this.hasOlderMessages = false + this.setHasOlderMessages(false) this.olderMessagesCursor = undefined this.knownServerMessageIds.clear() this.historyGeneration++ @@ -3364,11 +3502,11 @@ export class ChatClient< private applyHydrationPage(page: ChatHydrationResult['page']) { if (page?.truncated === true) { - this.hasOlderMessages = true + this.setHasOlderMessages(true) this.olderMessagesCursor = page.cursor return } - this.hasOlderMessages = false + this.setHasOlderMessages(false) this.olderMessagesCursor = undefined } @@ -3510,7 +3648,9 @@ export class ChatClient< } private emitQueueChange(): void { - this.callbacksRef.current.onQueueChange(this.getQueue()) + const queue = this.getQueue() + patchAtom(this.snapshotAtom, { queue: this.freezeSnapshotQueue(queue) }) + this.callbacksRef.current.onQueueChange(queue) this.devtoolsBridge.emitSnapshot() } @@ -3593,8 +3733,8 @@ export class ChatClient< /** * Manually set messages */ - setMessagesManually(messages: Array>): void { - this.processor.setMessages(messages) + setMessagesManually(messages: ReadonlyArray>): void { + this.processor.setMessages(thawMessages(messages)) this.devtoolsBridge.emitSnapshot() } diff --git a/packages/ai-client/src/generation-client.ts b/packages/ai-client/src/generation-client.ts index ffc43decad..7512508ddc 100644 --- a/packages/ai-client/src/generation-client.ts +++ b/packages/ai-client/src/generation-client.ts @@ -30,8 +30,17 @@ import type { GenerationDevtoolsBridge, GenerationDevtoolsBridgeOptions, } from './devtools' +import { + cloneSnapshotValue, + createAtom, + patchAtom, + subscribeAtom, +} from './snapshot-atom' +import type { Atom } from './snapshot-atom' +import { finishStoppedRun, owns, ownsRun, whileOwned } from './run-ownership' import type { GenerationClientOptions, + GenerationClientSnapshot, GenerationClientState, GenerationFetcher, GenerationPersistenceOptions, @@ -135,6 +144,16 @@ export class GenerationClient< private isLoading = false private error: Error | undefined = undefined private status: GenerationClientState = 'idle' + private readonly snapshotAtom: Atom> = + createAtom>( + Object.freeze({ + result: null, + isLoading: false, + error: undefined, + status: 'idle', + runId: null, + }), + ) private resumeSnapshot: GenerationResumeSnapshot | undefined private lastEmittedResumeState: string | undefined private abortController: AbortController | null = null @@ -266,16 +285,23 @@ export class GenerationClient< this.input = input this.progress = null - const runId = this.devtoolsBridge.beginRun(input) - this.setIsLoading(true) - this.setStatus('generating') - this.setError(undefined) - const abortController = new AbortController() this.abortController = abortController const { signal } = abortController + const runId = this.devtoolsBridge.beginRun(input) + patchAtom(this.snapshotAtom, { runId }) try { + // Inside the try: a throwing callback fails this run, and `finally` + // still clears `isLoading`. + if ( + !ownsRun(this.abortController, abortController) || + !this.setIsLoading(true, abortController) || + !this.setStatus('generating', abortController) || + !this.setError(undefined, abortController) + ) + return + let headers: Record | undefined if (this.byok) { const provider = resolveByokProviderId( @@ -284,6 +310,7 @@ export class GenerationClient< byokFallbackProviderId(this.byok), ) headers = await prepareResolvedByokHeaders(this.byok, provider) + if (!ownsRun(this.abortController, abortController)) return } if (this.fetcher) { @@ -292,19 +319,22 @@ export class GenerationClient< input, headers === undefined ? { signal } : { signal, headers }, ) - if (signal.aborted) return + if (!ownsRun(this.abortController, abortController)) return if (result instanceof Response) { // Server function returned SSE Response — parse stream await this.processStream( parseSSEResponse(result, signal), runId, - signal, + abortController, ) } else { this.devtoolsBridge.ensureRunStarted(runId) - this.setResult(result) - this.setStatus('success') - this.completePlainFetcherResumeSnapshot(result) + if ( + !this.setResult(result, abortController) || + !this.setStatus('success', abortController) + ) + return + this.completePlainFetcherResumeSnapshot(result, abortController) } } else if (this.connection) { // Streaming adapter path @@ -315,13 +345,16 @@ export class GenerationClient< signal, this.createRunContext(runId, headers), ) - await this.processStream(stream, runId, signal) + await this.processStream(stream, runId, abortController) } else { throw new Error( 'GenerationClient requires either a connection or fetcher option', ) } - if (!signal.aborted && this.status === 'success') { + if ( + ownsRun(this.abortController, abortController) && + this.status === 'success' + ) { // Bump progress to 100 on successful completion so devtools // snapshots reflect the final state. The bridge mirrors this in // the run's recorded progress, but the snapshot reads `progress` @@ -334,7 +367,7 @@ export class GenerationClient< ) } } catch (err: unknown) { - if (signal.aborted) return + if (!ownsRun(this.abortController, abortController)) return const error = err instanceof Error ? err : new Error(String(err)) if (error instanceof ByokMissingError) { this.byok?.request(error.provider, 'missing') @@ -342,18 +375,31 @@ export class GenerationClient< if (error instanceof ByokBlockedError && error.reason === 'locked') { this.byok?.request(error.provider, 'locked') } - this.setError(error) - this.setStatus('error') - this.recordResumeSnapshotError(error) + if (!ownsRun(this.abortController, abortController)) return + // Finish the devtools run first. A callback below can call `stop()`, + // and `stop()` would record this failed run as cancelled. this.devtoolsBridge.finishRun( this.devtoolsBridge.getActiveRunId() ?? runId, 'run:errored', 'errored', error.message, ) - this.callbacksRef.onError?.(error) + try { + if ( + this.setError(error, abortController) && + this.setStatus('error', abortController) + ) { + this.recordResumeSnapshotError(error, abortController) + } + } finally { + // The run failed even if a callback above stopped or replaced it. + this.callbacksRef.onError?.(error) + } } finally { if (this.abortController === abortController) { + patchAtom(this.snapshotAtom, { runId: null }) + // Release the run before user code runs, so a throwing + // `onLoadingChange` cannot keep a finished run as the current one. this.abortController = null this.setIsLoading(false) } @@ -375,8 +421,9 @@ export class GenerationClient< private async processStream( source: AsyncIterable, fallbackRunId: string, - signal: AbortSignal, + abortController: AbortController, ): Promise { + const { signal } = abortController let streamRunId: string | undefined let sawTerminalChunk = false @@ -385,7 +432,11 @@ export class GenerationClient< const chunk = restoreInboundChunk(raw) this.callbacksRef.onChunk?.(chunk) - this.observeResumeSnapshot(chunk) + if ( + !ownsRun(this.abortController, abortController) || + !this.observeResumeSnapshot(chunk, abortController) + ) + return const chunkRunId = 'runId' in chunk && typeof chunk.runId === 'string' ? chunk.runId @@ -395,19 +446,21 @@ export class GenerationClient< switch (chunk.type) { case 'RUN_STARTED': { streamRunId = chunk.runId + patchAtom(this.snapshotAtom, { runId: chunk.runId }) + if (!ownsRun(this.abortController, abortController)) return this.devtoolsBridge.ensureRunStarted(chunk.runId) break } case 'CUSTOM': { this.devtoolsBridge.ensureRunStarted(streamRunId ?? fallbackRunId) if (chunk.name === GENERATION_EVENTS.RESULT) { - this.setResult(chunk.value as TResult) + if (!this.setResult(chunk.value as TResult, abortController)) return } else if (chunk.name === GENERATION_EVENTS.PROGRESS) { const { progress, message } = chunk.value as { progress: number message?: string } - this.setProgress(progress, message) + if (!this.setProgress(progress, message, abortController)) return } break } @@ -415,7 +468,7 @@ export class GenerationClient< streamRunId = chunk.runId sawTerminalChunk = true this.devtoolsBridge.ensureRunStarted(chunk.runId) - this.setStatus('success') + if (!this.setStatus('success', abortController)) return break } case 'RUN_ERROR': { @@ -443,27 +496,38 @@ export class GenerationClient< */ stop(): void { const runId = this.devtoolsBridge.getActiveRunId() + const stoppedStatus = this.status if (this.abortController) { this.abortController.abort() this.abortController = null } - this.setIsLoading(false) - if (this.status === 'generating') { - this.setStatus('idle') - if (runId) { - this.devtoolsBridge.finishRun(runId, 'run:cancelled', 'cancelled') - } - } + finishStoppedRun(this.devtoolsBridge, runId, stoppedStatus, this.error) // A stopped run is no longer resumable. Without this the in-memory // snapshot stays `running`, and a remount's `maybeResumeInFlight` would // rejoin a run the user just cancelled. - if (this.resumeSnapshot && this.resumeSnapshot.status === 'running') { + const resumeStopped = this.resumeSnapshot?.status === 'running' + if (this.resumeSnapshot && resumeStopped) { this.resumeSnapshot = { ...this.resumeSnapshot, resumeState: null, status: 'idle', } - this.notifyResumeSnapshotChanged() + } + // Each step below can run user code. That code can start a new run, so + // clear the flag `generate()` checks now and publish it last. If a new + // run starts, it owns the client, so stop here and keep its state. + this.isLoading = false + patchAtom(this.snapshotAtom, { runId: null }) + if (this.abortController) return + try { + if (stoppedStatus === 'generating') { + this.setStatus('idle') + if (this.abortController) return + } + if (resumeStopped) this.notifyResumeSnapshotChanged() + } finally { + // A throwing callback above must not leave the snapshot loading. + if (!this.abortController) this.setIsLoading(false) } } @@ -475,6 +539,7 @@ export class GenerationClient< */ reset(): void { this.stop() + if (this.abortController || this.isLoading) return this.setResult(null) this.input = null this.progress = null @@ -539,7 +604,9 @@ export class GenerationClient< this.abortController.abort() this.abortController = null } - this.setIsLoading(false) + if (this.resumeSnapshot?.status !== 'running') { + patchAtom(this.snapshotAtom, { runId: null }) + } this.devtoolsBridge.dispose() this.devtoolsMounted = false // Re-arm mount hydration + rejoin so a remount resumes from the (preserved) @@ -547,12 +614,26 @@ export class GenerationClient< // `maybeResumeInFlight`, both individually guarded. this.serverHydrationStarted = false this.rejoinedRunId = undefined + this.setIsLoading(false) } // =========================== // Getters // =========================== + /** + * Subscribe to UI snapshot changes. Does not fire with the current value; + * read {@link getSnapshot} first. + */ + subscribe = (listener: () => void): (() => void) => + subscribeAtom(this.snapshotAtom, listener) + + /** + * Current UI snapshot. Frozen envelope. A plain-object or array `result` is + * a shallow copy. Any other `result` object is the value itself. + */ + getSnapshot = (): GenerationClientSnapshot => this.snapshotAtom.get() + getResult(): TOutput | null { return this.result } @@ -600,27 +681,29 @@ export class GenerationClient< // Private state setters // =========================== - private setResult(rawResult: TResult | null): void { + // Each setter returns false once `controller` no longer owns the client. + + private setResult( + rawResult: TResult | null, + controller?: AbortController, + ): boolean { if (rawResult === null) { this.result = null - this.callbacksRef.onResultChange?.(null) - this.devtoolsBridge.recordResultChange() - return + return this.publishResult(controller) } if (this.callbacksRef.onResult) { const transformed = this.callbacksRef.onResult(rawResult) + if (!owns(this.abortController, controller)) return false if (transformed === null) { // null return → keep previous result unchanged, just re-emit this.devtoolsBridge.emitState() - return + return true } if (transformed !== undefined) { // Non-null, non-undefined → use transformed value this.result = transformed - this.callbacksRef.onResultChange?.(this.result) - this.devtoolsBridge.recordResultChange() - return + return this.publishResult(controller) } } @@ -629,29 +712,67 @@ export class GenerationClient< // `TOutput` defaults to `TResult`, so the runtime cast is sound. // oxlint-disable-next-line eslint-js/no-restricted-syntax -- TOutput defaults to TResult when no onResult transform is supplied this.result = rawResult as unknown as TOutput - this.callbacksRef.onResultChange?.(this.result) - this.devtoolsBridge.recordResultChange() + return this.publishResult(controller) + } + + private publishResult(controller?: AbortController): boolean { + const result = this.result + patchAtom(this.snapshotAtom, { result: cloneSnapshotValue(result) }) + return whileOwned( + () => this.abortController, + controller, + () => this.callbacksRef.onResultChange?.(result), + () => this.devtoolsBridge.recordResultChange(), + ) } - private setIsLoading(isLoading: boolean): void { + private setIsLoading( + isLoading: boolean, + controller?: AbortController, + ): boolean { this.isLoading = isLoading - this.callbacksRef.onLoadingChange?.(isLoading) - this.devtoolsBridge.recordLoadingChange() + patchAtom(this.snapshotAtom, { isLoading }) + return whileOwned( + () => this.abortController, + controller, + () => this.callbacksRef.onLoadingChange?.(isLoading), + () => this.devtoolsBridge.recordLoadingChange(), + ) } - private setError(error: Error | undefined): void { + private setError( + error: Error | undefined, + controller?: AbortController, + ): boolean { this.error = error - this.callbacksRef.onErrorChange?.(error) - this.devtoolsBridge.recordErrorChange(error) + patchAtom(this.snapshotAtom, { error }) + return whileOwned( + () => this.abortController, + controller, + () => this.callbacksRef.onErrorChange?.(error), + () => this.devtoolsBridge.recordErrorChange(error), + ) } - private setStatus(status: GenerationClientState): void { + private setStatus( + status: GenerationClientState, + controller?: AbortController, + ): boolean { this.status = status - this.callbacksRef.onStatusChange?.(status) - this.devtoolsBridge.recordStatusChange(status) + patchAtom(this.snapshotAtom, { status }) + return whileOwned( + () => this.abortController, + controller, + () => this.callbacksRef.onStatusChange?.(status), + () => this.devtoolsBridge.recordStatusChange(status), + ) } - private setProgress(value: number, message?: string): void { + private setProgress( + value: number, + message?: string, + controller?: AbortController, + ): boolean { this.progress = { value, ...(message ? { message } : {}), @@ -661,7 +782,11 @@ export class GenerationClient< } else { this.callbacksRef.onProgress?.(value, message) } - this.devtoolsBridge.recordProgressChange() + return whileOwned( + () => this.abortController, + controller, + () => this.devtoolsBridge.recordProgressChange(), + ) } private createDevtoolsMetadata( @@ -698,12 +823,15 @@ export class GenerationClient< } } - private observeResumeSnapshot(chunk: StreamChunk): void { + private observeResumeSnapshot( + chunk: StreamChunk, + controller: AbortController, + ): boolean { this.resumeSnapshot = updateGenerationResumeSnapshot( this.resumeSnapshot, chunk, ) - this.notifyResumeSnapshotChanged() + return this.notifyResumeSnapshotChanged(controller) } /** @@ -711,9 +839,13 @@ export class GenerationClient< * The snapshot stays internal (persistence + devtools); the hook consumes * `resumeState`, mirroring the chat client. */ - private notifyResumeSnapshotChanged(): void { - this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot) - this.emitResumeState() + private notifyResumeSnapshotChanged(controller?: AbortController): boolean { + return whileOwned( + () => this.abortController, + controller, + () => this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot), + () => this.emitResumeState(controller), + ) } /** @@ -727,7 +859,7 @@ export class GenerationClient< * boundaries and when artifacts land, so skip the notification unless it * materially changed — same gate the persistence writes use. */ - private emitResumeState(): void { + private emitResumeState(controller?: AbortController): void { const snapshot = this.resumeSnapshot const state = snapshot?.resumeState const resumeState: GenerationResumeState | null = state @@ -743,6 +875,12 @@ export class GenerationClient< return } this.lastEmittedResumeState = signature + if (resumeState?.runId) { + patchAtom(this.snapshotAtom, { runId: resumeState.runId }) + } else if (!this.isLoading) { + patchAtom(this.snapshotAtom, { runId: null }) + } + if (!owns(this.abortController, controller)) return this.callbacksRef.onResumeStateChange?.(resumeState) } @@ -871,7 +1009,10 @@ export class GenerationClient< * stale `error` from a previous run is intentionally dropped — this run * succeeded. */ - private completePlainFetcherResumeSnapshot(rawResult: unknown): void { + private completePlainFetcherResumeSnapshot( + rawResult: unknown, + controller: AbortController, + ): void { const previous = this.resumeSnapshot const result = createGenerationResultSnapshot(rawResult) this.resumeSnapshot = { @@ -888,7 +1029,7 @@ export class GenerationClient< ? { result: { ...previous.result } } : {}), } - this.notifyResumeSnapshotChanged() + this.notifyResumeSnapshotChanged(controller) } /** @@ -897,17 +1038,22 @@ export class GenerationClient< * mark the snapshot `error`, leaving a persisted record that claims the * run is still in flight. */ - private recordResumeSnapshotError(error: Error): void { + private recordResumeSnapshotError( + error: Error, + controller?: AbortController, + ): boolean { // Surface the failure on the OBSERVABLE fields FIRST: a rejoin (or live // stream) that emits RUN_ERROR has already flipped the snapshot to `error` // via `observeResumeSnapshot`, so the early-return below would otherwise // skip this and leave `status` stuck on `generating` — the run would look // like it is still going forever. The guard avoids a duplicate `error` // emission on the live `generate()` path, which sets the status itself. - if (this.status !== 'error') this.setStatus('error') - this.setError(error) - if (this.resumeSnapshot?.status === 'error') return - if (!this.resumeSnapshot && !this.serverDriven) return + if (this.status !== 'error' && !this.setStatus('error', controller)) { + return false + } + if (!this.setError(error, controller)) return false + if (this.resumeSnapshot?.status === 'error') return true + if (!this.resumeSnapshot && !this.serverDriven) return true const previous = this.resumeSnapshot this.resumeSnapshot = { schemaVersion: 1, @@ -920,7 +1066,7 @@ export class GenerationClient< ...(previous?.result ? { result: { ...previous.result } } : {}), error: { message: error.message }, } - this.notifyResumeSnapshotChanged() + return this.notifyResumeSnapshotChanged(controller) } /** @@ -1044,23 +1190,26 @@ export class GenerationClient< this.rejoinedRunId = runId const controller = new AbortController() this.abortController = controller - this.setIsLoading(true) - this.setStatus('generating') void (async () => { try { + if ( + !this.setIsLoading(true, controller) || + !this.setStatus('generating', controller) + ) + return await this.processStream( joinRun(runId, controller.signal), runId, - controller.signal, + controller, ) } catch (error) { - if (!controller.signal.aborted) { + if (ownsRun(this.abortController, controller)) { const failure = error instanceof Error ? error : new Error(String(error)) // Settles `status`/`error` AND rewrites the snapshot to a terminal // `error` with a null `resumeState`, so the next mount does not // rejoin this run again. - this.recordResumeSnapshotError(failure) + if (!this.recordResumeSnapshotError(failure, controller)) return this.callbacksRef.onError?.(failure) } } finally { @@ -1068,6 +1217,8 @@ export class GenerationClient< // fresh `generate()` may have replaced the controller while the tail // was settling, and that live run owns `isLoading` now. if (this.abortController === controller) { + // `emitResumeState` skips this while `isLoading` is still true. + patchAtom(this.snapshotAtom, { runId: null }) this.abortController = null this.setIsLoading(false) } diff --git a/packages/ai-client/src/generation-types.ts b/packages/ai-client/src/generation-types.ts index bcc029aaeb..fd4f2399ee 100644 --- a/packages/ai-client/src/generation-types.ts +++ b/packages/ai-client/src/generation-types.ts @@ -68,6 +68,29 @@ export type InferGenerationOutput = TFn extends ( */ export type GenerationClientState = 'idle' | 'generating' | 'success' | 'error' +/** + * UI snapshot published by {@link GenerationClient.subscribe} / + * {@link GenerationClient.getSnapshot}. + */ +export interface GenerationClientSnapshot { + readonly result: TOutput | null + readonly isLoading: boolean + readonly error: Error | undefined + readonly status: GenerationClientState + readonly runId: string | null +} + +/** + * UI snapshot published by {@link VideoGenerationClient.subscribe} / + * {@link VideoGenerationClient.getSnapshot}. + */ +export interface VideoGenerationClientSnapshot< + TOutput = unknown, +> extends GenerationClientSnapshot { + readonly jobId: string | null + readonly videoStatus: VideoStatusInfo | null +} + /** * Status of a persisted/restored generation run. * diff --git a/packages/ai-client/src/index.ts b/packages/ai-client/src/index.ts index 444e87a3b3..15376535c2 100644 --- a/packages/ai-client/src/index.ts +++ b/packages/ai-client/src/index.ts @@ -72,6 +72,7 @@ export type { ChatRequestBody, InferChatMessages, InferredClientContext, + ChatClientSnapshot, ChatClientState, ConnectionStatus, ChatFetcher, @@ -92,7 +93,9 @@ export type { export type { InferGenerationOutput, InferGenerationOutputFromReturn, + GenerationClientSnapshot, GenerationClientState, + VideoGenerationClientSnapshot, GenerationPersistenceOptions, GenerationClientOptions, GenerationFetcher, diff --git a/packages/ai-client/src/realtime-client.ts b/packages/ai-client/src/realtime-client.ts index 7808818745..298fb7133d 100644 --- a/packages/ai-client/src/realtime-client.ts +++ b/packages/ai-client/src/realtime-client.ts @@ -1,4 +1,11 @@ import { convertSchemaToJsonSchema } from '@tanstack/ai/client' +import { + createAtom, + freezeSnapshotMessages, + notifyListeners, + subscribeAtom, +} from './snapshot-atom' +import type { Atom } from './snapshot-atom' import type { AnyClientTool, AudioVisualization, @@ -48,6 +55,8 @@ export class RealtimeClient { private readonly stateChangeCallbacks: Set = new Set() private unsubscribers: Array<() => void> = [] + // Declared before `snapshotAtom`, whose initializer reads it. + private readonly frozenMessages = new WeakMap() private state: RealtimeClientState = { status: 'idle', @@ -57,6 +66,9 @@ export class RealtimeClient { pendingAssistantTranscript: null, error: null, } + private readonly snapshotAtom: Atom = createAtom( + this.buildSnapshot(), + ) constructor(options: RealtimeClientOptions) { this.options = { @@ -88,7 +100,10 @@ export class RealtimeClient { return } - this.updateState({ status: 'connecting', error: null }) + this.updateState({ + status: 'connecting', + error: null, + }) try { // Fetch token from server @@ -108,6 +123,15 @@ export class RealtimeClient { toolsList, ) + // A new session starts here. Clear the last transcript before events + // flow, so a greeting that arrives while audio capture starts is kept. + // A connect that fails before this point keeps the last transcript. + this.updateState({ + messages: [], + pendingUserTranscript: null, + pendingAssistantTranscript: null, + }) + // Subscribe to connection events this.subscribeToConnectionEvents() @@ -119,7 +143,10 @@ export class RealtimeClient { await this.connection.startAudioCapture() } - this.updateState({ status: 'connected', mode: 'listening' }) + this.updateState({ + status: 'connected', + mode: 'listening', + }) this.options.onConnect?.() } catch (error) { const err = error instanceof Error ? error : new Error(String(error)) @@ -195,6 +222,7 @@ export class RealtimeClient { return } this.connection.interrupt() + this.updateState({ pendingAssistantTranscript: null }) } // ============================================================================ @@ -260,7 +288,7 @@ export class RealtimeClient { } /** Get conversation messages */ - get messages(): Array { + get messages(): ReadonlyArray { return this.state.messages } @@ -331,6 +359,16 @@ export class RealtimeClient { } } + /** + * Subscribe to UI snapshot changes. Does not fire with the current value; + * read {@link getSnapshot} first. + */ + subscribe = (listener: () => void): (() => void) => + subscribeAtom(this.snapshotAtom, listener) + + /** Current UI snapshot. */ + getSnapshot = (): RealtimeClientState => this.snapshotAtom.get() + // ============================================================================ // Cleanup // ============================================================================ @@ -349,12 +387,25 @@ export class RealtimeClient { // ============================================================================ private updateState(updates: Partial): void { + const changed = Object.keys(updates).some((key) => { + const stateKey = key as keyof RealtimeClientState + return !Object.is(this.state[stateKey], updates[stateKey]) + }) + if (!changed) return + const messagesChanged = + 'messages' in updates && updates.messages !== this.state.messages this.state = { ...this.state, ...updates } + // A transcript or mode update keeps the same `messages` array, so UI + // memo keyed on it does not run again for every transcript token. + const snapshot = messagesChanged + ? this.buildSnapshot() + : Object.freeze({ + ...this.state, + messages: this.snapshotAtom.get().messages, + }) + this.snapshotAtom.set(snapshot) - // Notify callbacks - for (const callback of this.stateChangeCallbacks) { - callback(this.state) - } + notifyListeners(this.stateChangeCallbacks, snapshot) // Notify specific callbacks if ('status' in updates && updates.status !== undefined) { @@ -365,6 +416,16 @@ export class RealtimeClient { } } + private buildSnapshot(): RealtimeClientState { + return Object.freeze({ + ...this.state, + messages: freezeSnapshotMessages( + this.frozenMessages, + this.state.messages, + ), + }) + } + private addMessage(message: RealtimeMessage): void { this.updateState({ messages: [...this.state.messages, message], diff --git a/packages/ai-client/src/realtime-types.ts b/packages/ai-client/src/realtime-types.ts index c6f2470a97..235190dc53 100644 --- a/packages/ai-client/src/realtime-types.ts +++ b/packages/ai-client/src/realtime-types.ts @@ -108,12 +108,12 @@ export interface RealtimeClientOptions { * Internal state of the RealtimeClient */ export interface RealtimeClientState { - status: RealtimeStatus - mode: RealtimeMode - messages: Array - pendingUserTranscript: string | null - pendingAssistantTranscript: string | null - error: Error | null + readonly status: RealtimeStatus + readonly mode: RealtimeMode + readonly messages: ReadonlyArray + readonly pendingUserTranscript: string | null + readonly pendingAssistantTranscript: string | null + readonly error: Error | null } /** diff --git a/packages/ai-client/src/run-ownership.ts b/packages/ai-client/src/run-ownership.ts new file mode 100644 index 0000000000..c5c269bbe5 --- /dev/null +++ b/packages/ai-client/src/run-ownership.ts @@ -0,0 +1,52 @@ +import type { GenerationDevtoolsBridge } from './devtools' +import type { GenerationClientState } from './generation-types' + +/** True while `controller` is the client's current, unaborted run. */ +export function ownsRun( + current: AbortController | null, + controller: AbortController, +): boolean { + return !controller.signal.aborted && current === controller +} + +/** True without a controller: that caller does not own a run. */ +export function owns( + current: AbortController | null, + controller: AbortController | undefined, +): boolean { + return !controller || ownsRun(current, controller) +} + +/** + * Run `steps` in order while `controller` owns the client. A step can run + * user code, and that code can call `stop()` or start a new run. Returns + * whether `controller` still owns the client. + */ +export function whileOwned( + getCurrent: () => AbortController | null, + controller: AbortController | undefined, + ...steps: Array<() => unknown> +): boolean { + for (const step of steps) { + if (!owns(getCurrent(), controller)) return false + step() + } + return owns(getCurrent(), controller) +} + +/** Close a stopped run in devtools with the status it had when stopped. */ +export function finishStoppedRun( + bridge: Pick, 'finishRun'>, + runId: string | null, + status: GenerationClientState, + error: Error | undefined, +): void { + if (!runId) return + if (status === 'success') { + bridge.finishRun(runId, 'run:completed', 'completed') + } else if (status === 'error') { + bridge.finishRun(runId, 'run:errored', 'errored', error?.message) + } else { + bridge.finishRun(runId, 'run:cancelled', 'cancelled') + } +} diff --git a/packages/ai-client/src/snapshot-atom.ts b/packages/ai-client/src/snapshot-atom.ts new file mode 100644 index 0000000000..e691397392 --- /dev/null +++ b/packages/ai-client/src/snapshot-atom.ts @@ -0,0 +1,200 @@ +import type { Atom } from '@tanstack/store' + +export { createAtom } from '@tanstack/store' +export type { Atom } + +function isPlainObject(value: unknown): value is Record { + if (value === null || typeof value !== 'object' || Array.isArray(value)) { + return false + } + const proto: unknown = Object.getPrototypeOf(value) + return proto === Object.prototype || proto === null +} + +/** + * Shallow copy, then freeze. Use for published snapshot fields that must not + * alias client internals. Only plain objects and arrays are copied. Any other + * object (Blob, Map, Date, class instance) is returned as it is, because a + * spread copy loses its prototype. + */ +export function cloneSnapshotValue(value: T): T { + if (Array.isArray(value)) return Object.freeze([...value]) as T + return isPlainObject(value) ? (Object.freeze({ ...value }) as T) : value +} + +/** + * Report an error thrown by a listener without stopping the caller. The + * browser's `reportError` reaches `window.onerror` and error trackers. Other + * runtimes log it. + */ +export function reportListenerError(error: unknown): void { + if (typeof globalThis.reportError === 'function') { + globalThis.reportError(error) + } else { + console.error(error) + } +} + +/** Call every listener. One listener's throw does not skip the others. */ +export function notifyListeners>( + listeners: Iterable<(...args: TArgs) => void>, + ...args: TArgs +): void { + for (const listener of listeners) { + try { + listener(...args) + } catch (error) { + reportListenerError(error) + } + } +} + +/** Deep-freeze plain JSON. Class instances stay as they are. */ +function freezePlainData(value: unknown): unknown { + if (Array.isArray(value)) { + return Object.freeze(value.map(freezePlainData)) + } + if (!isPlainObject(value)) return value + const copy: Record = {} + for (const key of Object.keys(value)) { + copy[key] = freezePlainData(value[key]) + } + return Object.freeze(copy) +} + +/** + * Freeze a copy of one message part. Tool `input` / `output`, `source`, + * `approval`, tool-result `content`, and subagent transcripts are copies. + * A subagent `stop` function stays the live function. + */ +export function freezeSnapshotPart( + part: TPart, + cache: WeakMap = new WeakMap(), +): TPart { + const record = part as Record + const copy: Record = { ...record } + if (isPlainObject(record.source)) { + copy.source = Object.freeze({ ...record.source }) + } + if (record.type === 'tool-call' && isPlainObject(record.approval)) { + copy.approval = Object.freeze({ ...record.approval }) + } + if (record.type === 'tool-result' && Array.isArray(record.content)) { + copy.content = Object.freeze( + record.content.map((contentPart: unknown) => + typeof contentPart === 'object' && contentPart !== null + ? freezeSnapshotPart(contentPart, cache) + : contentPart, + ), + ) + } + if ('input' in record) copy.input = freezePlainData(record.input) + if ('output' in record) copy.output = freezePlainData(record.output) + if (record.type === 'subagent' && isPlainObject(record.subagent)) { + const sub = record.subagent + copy.subagent = Object.freeze( + Array.isArray(sub.messages) + ? { + ...sub, + messages: freezeSnapshotMessages( + cache, + sub.messages as Array<{ parts: ReadonlyArray }>, + ), + } + : { ...sub }, + ) + } + return Object.freeze(copy) as TPart +} + +function cached( + cache: WeakMap, + value: T, + freeze: (value: T) => T, +): T { + let frozen = cache.get(value) as T | undefined + if (!frozen) { + frozen = freeze(value) + cache.set(value, frozen) + } + return frozen +} + +/** + * Freeze a copy of each message and its parts. The client replaces a message + * or a part object when it changes, so `cache` gives an unchanged message or + * part the same frozen copy. UI memo keys on message identity, and a streamed + * token copies only the part that changed. + */ +export function freezeSnapshotMessages< + TMessage extends { parts: ReadonlyArray }, +>( + cache: WeakMap, + messages: ReadonlyArray, +): ReadonlyArray { + return Object.freeze( + messages.map((message) => + cached(cache, message, () => + Object.freeze({ + ...message, + parts: Object.freeze( + message.parts.map((part) => + cached(cache, part, () => freezeSnapshotPart(part, cache)), + ), + ), + }), + ), + ), + ) +} + +/** + * Merge a partial update into an object atom and freeze the result. Skips + * notify when every provided field is `Object.is`-equal to the current value. + */ +export function patchAtom( + atom: Atom, + patch: Partial, +): void { + atom.set((prev) => { + let changed = false + const next = { ...prev } + for (const key of Object.keys(patch) as Array) { + if (Object.is(prev[key], patch[key])) continue + next[key] = patch[key] as T[keyof T] + changed = true + } + return changed ? Object.freeze(next) : prev + }) +} + +const MAX_LISTENER_REENTRY = 8 + +/** + * Subscribe to an atom with the `() => void` unsubscribe that framework + * stores expect. + */ +export function subscribeAtom( + atom: Pick, 'get' | 'subscribe'>, + listener: () => void, +): () => void { + const { unsubscribe } = atom.subscribe(() => { + // Store calls this inside an effect. A change made while `listener` + // runs (a Solid effect or Vue watcher that calls `stop()`) does not + // notify this subscriber again, so call it again until the value stays. + // One listener's throw must not skip the other subscribers, and must + // not escape into the setter that published this snapshot. + let seen: T + let spins = 0 + do { + seen = atom.get() + try { + listener() + } catch (error) { + reportListenerError(error) + } + spins++ + } while (atom.get() !== seen && spins < MAX_LISTENER_REENTRY) + }) + return unsubscribe +} diff --git a/packages/ai-client/src/types.ts b/packages/ai-client/src/types.ts index 5cbca7c5f3..0030286806 100644 --- a/packages/ai-client/src/types.ts +++ b/packages/ai-client/src/types.ts @@ -368,6 +368,31 @@ export type ToolResultState = */ export type ChatClientState = 'ready' | 'submitted' | 'streaming' | 'error' +/** + * UI snapshot published by {@link ChatClient.subscribeSnapshot} / + * {@link ChatClient.getSnapshot}. + */ +export interface ChatClientSnapshot< + TTools extends ReadonlyArray = any, + TInterrupts extends ReadonlyArray> = + readonly [], +> { + readonly messages: ReadonlyArray> + readonly status: ChatClientState + readonly isLoading: boolean + readonly error: Error | undefined + readonly isSubscribed: boolean + readonly connectionStatus: ConnectionStatus + readonly sessionGenerating: boolean + readonly queue: ReadonlyArray + readonly runId: string | null + /** True when the last hydrate or older-page response said more messages exist. */ + readonly hasOlderMessages: boolean + /** True while a store or server load may still repaint the transcript. */ + readonly isHydrating: boolean + readonly interruptState: ChatInterruptState +} + /** * Connection lifecycle state for the subscription loop. */ @@ -1076,7 +1101,8 @@ export interface ChatClientBaseOptions< onError?: (error: Error) => void /** - * Callback when messages change + * Callback when messages change. Receives the client's live array, not the + * frozen snapshot copy. Do not change it in place; use `setMessagesManually`. */ onMessagesChange?: (messages: Array>) => void @@ -1134,7 +1160,7 @@ export interface ChatClientBaseOptions< /** * Callback when the pending send queue changes (enqueue, cancel, drain, - * or flush). + * or flush). Receives a live array, not the frozen snapshot copy. */ onQueueChange?: (queue: Array) => void diff --git a/packages/ai-client/src/video-generation-client.ts b/packages/ai-client/src/video-generation-client.ts index 3f1502444f..24e497993e 100644 --- a/packages/ai-client/src/video-generation-client.ts +++ b/packages/ai-client/src/video-generation-client.ts @@ -30,6 +30,14 @@ import type { VideoDevtoolsBridge, VideoDevtoolsBridgeOptions, } from './devtools' +import { + cloneSnapshotValue, + createAtom, + patchAtom, + subscribeAtom, +} from './snapshot-atom' +import type { Atom } from './snapshot-atom' +import { finishStoppedRun, owns, ownsRun, whileOwned } from './run-ownership' import type { GenerationClientState, GenerationFetcher, @@ -40,6 +48,7 @@ import type { VideoGenerateInput, VideoGenerateResult, VideoGenerationClientOptions, + VideoGenerationClientSnapshot, VideoStatusInfo, } from './generation-types' @@ -140,6 +149,18 @@ export class VideoGenerationClient { private isLoading = false private error: Error | undefined = undefined private status: GenerationClientState = 'idle' + private readonly snapshotAtom: Atom> = + createAtom>( + Object.freeze({ + result: null, + isLoading: false, + error: undefined, + status: 'idle', + runId: null, + jobId: null, + videoStatus: null, + }), + ) private resumeSnapshot: GenerationResumeSnapshot | undefined private abortController: AbortController | null = null private rejoinedRunId: string | undefined @@ -257,18 +278,25 @@ export class VideoGenerationClient { this.input = input this.progress = null - const runId = this.devtoolsBridge.beginRun(input) - this.setIsLoading(true) - this.setStatus('generating') - this.setError(undefined) - this.setJobId(null) - this.setVideoStatus(null) - const abortController = new AbortController() this.abortController = abortController const { signal } = abortController + const runId = this.devtoolsBridge.beginRun(input) + patchAtom(this.snapshotAtom, { runId }) try { + // Inside the try: a throwing callback fails this run, and `finally` + // still clears `isLoading`. + if ( + !ownsRun(this.abortController, abortController) || + !this.setIsLoading(true, abortController) || + !this.setStatus('generating', abortController) || + !this.setError(undefined, abortController) || + !this.setJobId(null, abortController) || + !this.setVideoStatus(null, abortController) + ) + return + let headers: Record | undefined if (this.byok) { const provider = resolveByokProviderId( @@ -277,10 +305,11 @@ export class VideoGenerationClient { byokFallbackProviderId(this.byok), ) headers = await prepareResolvedByokHeaders(this.byok, provider) + if (!ownsRun(this.abortController, abortController)) return } if (this.fetcher) { - await this.generateWithFetcher(input, signal, runId, headers) + await this.generateWithFetcher(input, abortController, runId, headers) } else if (this.connection) { const mergedData = { ...this.body, ...input } const stream = this.connection.connect( @@ -289,13 +318,16 @@ export class VideoGenerationClient { signal, this.createRunContext(runId, headers), ) - await this.processStream(stream, runId, signal) + await this.processStream(stream, runId, abortController) } else { throw new Error( 'VideoGenerationClient requires either a connection or fetcher option', ) } - if (!signal.aborted && this.status === 'success') { + if ( + ownsRun(this.abortController, abortController) && + this.status === 'success' + ) { this.devtoolsBridge.finishRun( this.devtoolsBridge.getActiveRunId() ?? runId, 'run:completed', @@ -303,7 +335,7 @@ export class VideoGenerationClient { ) } } catch (err: unknown) { - if (signal.aborted) return + if (!ownsRun(this.abortController, abortController)) return const error = err instanceof Error ? err : new Error(String(err)) if (error instanceof ByokMissingError) { this.byok?.request(error.provider, 'missing') @@ -311,18 +343,31 @@ export class VideoGenerationClient { if (error instanceof ByokBlockedError && error.reason === 'locked') { this.byok?.request(error.provider, 'locked') } - this.setError(error) - this.setStatus('error') - this.recordResumeSnapshotError(error) + if (!ownsRun(this.abortController, abortController)) return + // Finish the devtools run first. A callback below can call `stop()`, + // and `stop()` would record this failed run as cancelled. this.devtoolsBridge.finishRun( this.devtoolsBridge.getActiveRunId() ?? runId, 'run:errored', 'errored', error.message, ) - this.callbacksRef.onError?.(error) + try { + if ( + this.setError(error, abortController) && + this.setStatus('error', abortController) + ) { + this.recordResumeSnapshotError(error, abortController) + } + } finally { + // The run failed even if a callback above stopped or replaced it. + this.callbacksRef.onError?.(error) + } } finally { if (this.abortController === abortController) { + patchAtom(this.snapshotAtom, { runId: null }) + // Release the run before user code runs, so a throwing + // `onLoadingChange` cannot keep a finished run as the current one. this.abortController = null this.setIsLoading(false) } @@ -334,27 +379,35 @@ export class VideoGenerationClient { */ private async generateWithFetcher( input: VideoGenerateInput, - signal: AbortSignal, + abortController: AbortController, runId: string, headers?: Record, ): Promise { if (!this.fetcher) return + const { signal } = abortController // Fetcher returns a completed result directly, or a Response with SSE body const result = await this.fetcher( input, headers === undefined ? { signal } : { signal, headers }, ) - if (signal.aborted) return + if (!ownsRun(this.abortController, abortController)) return if (result instanceof Response) { // Server function returned SSE Response — parse stream - await this.processStream(parseSSEResponse(result, signal), runId, signal) + await this.processStream( + parseSSEResponse(result, signal), + runId, + abortController, + ) } else { this.devtoolsBridge.ensureRunStarted(runId) - this.setResult(result) - this.setStatus('success') - this.completePlainFetcherResumeSnapshot(result) + if ( + !this.setResult(result, abortController) || + !this.setStatus('success', abortController) + ) + return + this.completePlainFetcherResumeSnapshot(result, abortController) } } @@ -370,8 +423,9 @@ export class VideoGenerationClient { private async processStream( source: AsyncIterable, fallbackRunId: string, - signal: AbortSignal, + abortController: AbortController, ): Promise { + const { signal } = abortController let streamRunId: string | undefined let sawTerminalChunk = false @@ -380,7 +434,11 @@ export class VideoGenerationClient { const chunk = restoreInboundChunk(raw) this.callbacksRef.onChunk?.(chunk) - this.observeResumeSnapshot(chunk) + if ( + !ownsRun(this.abortController, abortController) || + !this.observeResumeSnapshot(chunk, abortController) + ) + return const chunkRunId = 'runId' in chunk && typeof chunk.runId === 'string' ? chunk.runId @@ -390,6 +448,8 @@ export class VideoGenerationClient { switch (chunk.type) { case 'RUN_STARTED': { streamRunId = chunk.runId + patchAtom(this.snapshotAtom, { runId: chunk.runId }) + if (!ownsRun(this.abortController, abortController)) return this.devtoolsBridge.ensureRunStarted(chunk.runId) break } @@ -397,23 +457,45 @@ export class VideoGenerationClient { this.devtoolsBridge.ensureRunStarted(streamRunId ?? fallbackRunId) if (chunk.name === GENERATION_EVENTS.VIDEO_JOB_CREATED) { const { jobId } = chunk.value as { jobId: string } - this.setJobId(jobId) - this.callbacksRef.onJobCreated?.(jobId) + if ( + !this.setJobId(jobId, abortController) || + !whileOwned( + () => this.abortController, + abortController, + () => this.callbacksRef.onJobCreated?.(jobId), + ) + ) + return } else if (chunk.name === GENERATION_EVENTS.VIDEO_STATUS) { const statusInfo = chunk.value as VideoStatusInfo - this.setVideoStatus(statusInfo) - this.callbacksRef.onStatusUpdate?.(statusInfo) - if (statusInfo.progress !== undefined) { - this.setProgress(statusInfo.progress) - } + if ( + !this.setVideoStatus(statusInfo, abortController) || + !whileOwned( + () => this.abortController, + abortController, + () => this.callbacksRef.onStatusUpdate?.(statusInfo), + ) + ) + return + if ( + statusInfo.progress !== undefined && + !this.setProgress(statusInfo.progress, undefined, abortController) + ) + return } else if (chunk.name === GENERATION_EVENTS.RESULT) { - this.setResult(chunk.value as VideoGenerateResult) + if ( + !this.setResult( + chunk.value as VideoGenerateResult, + abortController, + ) + ) + return } else if (chunk.name === GENERATION_EVENTS.PROGRESS) { const { progress, message } = chunk.value as { progress: number message?: string } - this.setProgress(progress, message) + if (!this.setProgress(progress, message, abortController)) return } break } @@ -421,7 +503,7 @@ export class VideoGenerationClient { streamRunId = chunk.runId sawTerminalChunk = true this.devtoolsBridge.ensureRunStarted(chunk.runId) - this.setStatus('success') + if (!this.setStatus('success', abortController)) return break } case 'RUN_ERROR': { @@ -449,27 +531,38 @@ export class VideoGenerationClient { */ stop(): void { const runId = this.devtoolsBridge.getActiveRunId() + const stoppedStatus = this.status if (this.abortController) { this.abortController.abort() this.abortController = null } - this.setIsLoading(false) - if (this.status === 'generating') { - this.setStatus('idle') - if (runId) { - this.devtoolsBridge.finishRun(runId, 'run:cancelled', 'cancelled') - } - } + finishStoppedRun(this.devtoolsBridge, runId, stoppedStatus, this.error) // A stopped run is no longer resumable. Without this the in-memory // snapshot stays `running`, and a remount's `maybeResumeInFlight` would // rejoin a run the user just cancelled. - if (this.resumeSnapshot && this.resumeSnapshot.status === 'running') { + const resumeStopped = this.resumeSnapshot?.status === 'running' + if (this.resumeSnapshot && resumeStopped) { this.resumeSnapshot = { ...this.resumeSnapshot, resumeState: null, status: 'idle', } - this.notifyResumeSnapshotChanged() + } + // Each step below can run user code. That code can start a new run, so + // clear the flag `generate()` checks now and publish it last. If a new + // run starts, it owns the client, so stop here and keep its state. + this.isLoading = false + patchAtom(this.snapshotAtom, { runId: null }) + if (this.abortController) return + try { + if (stoppedStatus === 'generating') { + this.setStatus('idle') + if (this.abortController) return + } + if (resumeStopped) this.notifyResumeSnapshotChanged() + } finally { + // A throwing callback above must not leave the snapshot loading. + if (!this.abortController) this.setIsLoading(false) } } @@ -481,6 +574,7 @@ export class VideoGenerationClient { */ reset(): void { this.stop() + if (this.abortController || this.isLoading) return this.setResult(null) this.input = null this.progress = null @@ -549,17 +643,34 @@ export class VideoGenerationClient { this.abortController.abort() this.abortController = null } - this.setIsLoading(false) + if (this.resumeSnapshot?.status !== 'running') { + patchAtom(this.snapshotAtom, { runId: null }) + } this.devtoolsBridge.dispose() this.devtoolsMounted = false this.serverHydrationStarted = false this.rejoinedRunId = undefined + this.setIsLoading(false) } // =========================== // Getters // =========================== + /** + * Subscribe to UI snapshot changes. Does not fire with the current value; + * read {@link getSnapshot} first. + */ + subscribe = (listener: () => void): (() => void) => + subscribeAtom(this.snapshotAtom, listener) + + /** + * Current UI snapshot. Frozen envelope. A plain-object or array `result` is + * a shallow copy. Any other `result` object is the value itself. + */ + getSnapshot = (): VideoGenerationClientSnapshot => + this.snapshotAtom.get() + getResult(): TOutput | null { return this.result } @@ -615,34 +726,43 @@ export class VideoGenerationClient { // Private state setters // =========================== - private setResult(rawResult: VideoGenerateResult | null): void { + // Each setter returns false once `controller` no longer owns the client. + + private setResult( + rawResult: VideoGenerateResult | null, + controller?: AbortController, + ): boolean { if (rawResult === null) { this.result = null - this.callbacksRef.onResultChange?.(null) - this.devtoolsBridge.recordResultChange() - return + return this.publishResult(controller) } const completedStatus = this.createCompletedVideoStatus(rawResult) - if (this.progress?.value !== 100) { - this.setProgress(100, this.progress?.message) + if ( + this.progress?.value !== 100 && + !this.setProgress(100, this.progress?.message, controller) + ) { + return false + } + if ( + !this.setJobId(rawResult.jobId, controller) || + !this.setVideoStatus(completedStatus, controller) + ) { + return false } - this.setJobId(rawResult.jobId) - this.setVideoStatus(completedStatus) if (this.callbacksRef.onResult) { const transformed = this.callbacksRef.onResult(rawResult) + if (!owns(this.abortController, controller)) return false if (transformed === null) { // null return → keep previous result unchanged, just re-emit this.devtoolsBridge.emitState() - return + return true } if (transformed !== undefined) { // Non-null, non-undefined → use transformed value this.result = transformed - this.callbacksRef.onResultChange?.(this.result) - this.devtoolsBridge.recordResultChange() - return + return this.publishResult(controller) } } @@ -652,41 +772,98 @@ export class VideoGenerationClient { // sound. // oxlint-disable-next-line eslint-js/no-restricted-syntax -- TOutput defaults to VideoGenerateResult when no onResult transform is supplied this.result = rawResult as unknown as TOutput - this.callbacksRef.onResultChange?.(this.result) - this.devtoolsBridge.recordResultChange() + return this.publishResult(controller) + } + + private publishResult(controller?: AbortController): boolean { + const result = this.result + patchAtom(this.snapshotAtom, { result: cloneSnapshotValue(result) }) + return whileOwned( + () => this.abortController, + controller, + () => this.callbacksRef.onResultChange?.(result), + () => this.devtoolsBridge.recordResultChange(), + ) } - private setJobId(jobId: string | null): void { + private setJobId( + jobId: string | null, + controller?: AbortController, + ): boolean { this.jobId = jobId - this.callbacksRef.onJobIdChange?.(jobId) - this.devtoolsBridge.recordJobIdChange() + patchAtom(this.snapshotAtom, { jobId }) + return whileOwned( + () => this.abortController, + controller, + () => this.callbacksRef.onJobIdChange?.(jobId), + () => this.devtoolsBridge.recordJobIdChange(), + ) } - private setVideoStatus(status: VideoStatusInfo | null): void { - this.videoStatus = status - this.callbacksRef.onVideoStatusChange?.(status) - this.devtoolsBridge.recordVideoStatusChange() + private setVideoStatus( + status: VideoStatusInfo | null, + controller?: AbortController, + ): boolean { + const videoStatus = status ? { ...status } : null + this.videoStatus = videoStatus + patchAtom(this.snapshotAtom, { + videoStatus: videoStatus ? Object.freeze({ ...videoStatus }) : null, + }) + return whileOwned( + () => this.abortController, + controller, + () => this.callbacksRef.onVideoStatusChange?.(videoStatus), + () => this.devtoolsBridge.recordVideoStatusChange(), + ) } - private setIsLoading(isLoading: boolean): void { + private setIsLoading( + isLoading: boolean, + controller?: AbortController, + ): boolean { this.isLoading = isLoading - this.callbacksRef.onLoadingChange?.(isLoading) - this.devtoolsBridge.recordLoadingChange() + patchAtom(this.snapshotAtom, { isLoading }) + return whileOwned( + () => this.abortController, + controller, + () => this.callbacksRef.onLoadingChange?.(isLoading), + () => this.devtoolsBridge.recordLoadingChange(), + ) } - private setError(error: Error | undefined): void { + private setError( + error: Error | undefined, + controller?: AbortController, + ): boolean { this.error = error - this.callbacksRef.onErrorChange?.(error) - this.devtoolsBridge.recordErrorChange(error) + patchAtom(this.snapshotAtom, { error }) + return whileOwned( + () => this.abortController, + controller, + () => this.callbacksRef.onErrorChange?.(error), + () => this.devtoolsBridge.recordErrorChange(error), + ) } - private setStatus(status: GenerationClientState): void { + private setStatus( + status: GenerationClientState, + controller?: AbortController, + ): boolean { this.status = status - this.callbacksRef.onStatusChange?.(status) - this.devtoolsBridge.recordStatusChange(status) + patchAtom(this.snapshotAtom, { status }) + return whileOwned( + () => this.abortController, + controller, + () => this.callbacksRef.onStatusChange?.(status), + () => this.devtoolsBridge.recordStatusChange(status), + ) } - private setProgress(value: number, message?: string): void { + private setProgress( + value: number, + message?: string, + controller?: AbortController, + ): boolean { this.progress = { value, ...(message ? { message } : {}), @@ -696,7 +873,11 @@ export class VideoGenerationClient { } else { this.callbacksRef.onProgress?.(value, message) } - this.devtoolsBridge.recordProgressChange() + return whileOwned( + () => this.abortController, + controller, + () => this.devtoolsBridge.recordProgressChange(), + ) } private createCompletedVideoStatus( @@ -744,22 +925,29 @@ export class VideoGenerationClient { } } - private observeResumeSnapshot(chunk: StreamChunk): void { + private observeResumeSnapshot( + chunk: StreamChunk, + controller: AbortController, + ): boolean { this.resumeSnapshot = updateGenerationResumeSnapshot( this.resumeSnapshot, chunk, ) - this.notifyResumeSnapshotChanged() + return this.notifyResumeSnapshotChanged(controller) } /** Notify the internal snapshot listener AND emit the public resume state. */ - private notifyResumeSnapshotChanged(): void { - this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot) - this.emitResumeState() + private notifyResumeSnapshotChanged(controller?: AbortController): boolean { + return whileOwned( + () => this.abortController, + controller, + () => this.callbacksRef.onResumeSnapshotChange?.(this.resumeSnapshot), + () => this.emitResumeState(controller), + ) } /** Derive the public `resumeState` from the internal snapshot. */ - private emitResumeState(): void { + private emitResumeState(controller?: AbortController): void { const snapshot = this.resumeSnapshot const state = snapshot?.resumeState const resumeState: GenerationResumeState | null = state @@ -770,6 +958,12 @@ export class VideoGenerationClient { : {}), } : null + if (resumeState?.runId) { + patchAtom(this.snapshotAtom, { runId: resumeState.runId }) + } else if (!this.isLoading) { + patchAtom(this.snapshotAtom, { runId: null }) + } + if (!owns(this.abortController, controller)) return this.callbacksRef.onResumeStateChange?.(resumeState) } @@ -883,7 +1077,10 @@ export class VideoGenerationClient { * stale `error` from a previous run is intentionally dropped — this run * succeeded. */ - private completePlainFetcherResumeSnapshot(rawResult: unknown): void { + private completePlainFetcherResumeSnapshot( + rawResult: unknown, + controller: AbortController, + ): void { const previous = this.resumeSnapshot const result = createGenerationResultSnapshot(rawResult) this.resumeSnapshot = { @@ -900,7 +1097,7 @@ export class VideoGenerationClient { ? { result: { ...previous.result } } : {}), } - this.notifyResumeSnapshotChanged() + this.notifyResumeSnapshotChanged(controller) } /** @@ -909,16 +1106,21 @@ export class VideoGenerationClient { * mark the snapshot `error`, leaving a persisted record that claims the * run is still in flight. */ - private recordResumeSnapshotError(error: Error): void { + private recordResumeSnapshotError( + error: Error, + controller?: AbortController, + ): boolean { // Surface the failure on the observable fields FIRST, unconditionally (see // the note in GenerationClient.recordResumeSnapshotError): a RUN_ERROR // already flipped the snapshot to `error`, so the early-return would else // skip this and leave `status` stuck on `generating`. The guard avoids a // duplicate `error` emission on the live `generate()` path. - if (this.status !== 'error') this.setStatus('error') - this.setError(error) - if (this.resumeSnapshot?.status === 'error') return - if (!this.resumeSnapshot && !this.serverDriven) return + if (this.status !== 'error' && !this.setStatus('error', controller)) { + return false + } + if (!this.setError(error, controller)) return false + if (this.resumeSnapshot?.status === 'error') return true + if (!this.resumeSnapshot && !this.serverDriven) return true const previous = this.resumeSnapshot this.resumeSnapshot = { schemaVersion: 1, @@ -931,7 +1133,7 @@ export class VideoGenerationClient { ...(previous?.result ? { result: { ...previous.result } } : {}), error: { message: error.message }, } - this.notifyResumeSnapshotChanged() + return this.notifyResumeSnapshotChanged(controller) } /** @@ -1048,23 +1250,26 @@ export class VideoGenerationClient { this.rejoinedRunId = runId const controller = new AbortController() this.abortController = controller - this.setIsLoading(true) - this.setStatus('generating') void (async () => { try { + if ( + !this.setIsLoading(true, controller) || + !this.setStatus('generating', controller) + ) + return await this.processStream( joinRun(runId, controller.signal), runId, - controller.signal, + controller, ) } catch (error) { - if (!controller.signal.aborted) { + if (ownsRun(this.abortController, controller)) { const failure = error instanceof Error ? error : new Error(String(error)) // Settles `status`/`error` AND rewrites the snapshot to a terminal // `error` with a null `resumeState`, so the next mount does not // rejoin this run again. - this.recordResumeSnapshotError(failure) + if (!this.recordResumeSnapshotError(failure, controller)) return this.callbacksRef.onError?.(failure) } } finally { @@ -1072,6 +1277,8 @@ export class VideoGenerationClient { // fresh `generate()` may have replaced the controller while the tail // was settling, and that live run owns `isLoading` now. if (this.abortController === controller) { + // `emitResumeState` skips this while `isLoading` is still true. + patchAtom(this.snapshotAtom, { runId: null }) this.abortController = null this.setIsLoading(false) } diff --git a/packages/ai-client/tests/audio-recorder.test.ts b/packages/ai-client/tests/audio-recorder.test.ts index 49c773b0cd..13ba5625f8 100644 --- a/packages/ai-client/tests/audio-recorder.test.ts +++ b/packages/ai-client/tests/audio-recorder.test.ts @@ -75,6 +75,21 @@ afterEach(() => { }) describe('AudioRecorder', () => { + it('uses Set subscription behavior and reports repeated states', async () => { + const recorder = new AudioRecorder() + const listener = vi.fn() + const first = recorder.subscribe(listener) + recorder.subscribe(listener) + + getUserMedia.mockRejectedValueOnce(new Error('denied')) + await expect(recorder.start()).rejects.toThrow('denied') + expect(listener).toHaveBeenCalledTimes(1) + first() + getUserMedia.mockRejectedValueOnce(new Error('denied')) + await expect(recorder.start()).rejects.toThrow('denied') + expect(listener).toHaveBeenCalledTimes(1) + }) + it('isSupported() is true when media APIs exist', () => { expect(AudioRecorder.isSupported()).toBe(true) }) @@ -106,6 +121,25 @@ describe('AudioRecorder', () => { expect(states).toContain('idle') }) + it('a throwing listener does not skip others or strand stop()', async () => { + const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}) + const recorder = new AudioRecorder() + const states: Array = [] + recorder.subscribe(() => { + throw new Error('listener failed') + }) + recorder.subscribe((s) => states.push(s)) + + await recorder.start() + const recording = await recorder.stop() + + expect(recording.base64).toBe('AQID') + expect(states).toEqual(['recording', 'stopping', 'idle']) + expect(recorder.state).toBe('idle') + expect(consoleError).toHaveBeenCalledWith(new Error('listener failed')) + consoleError.mockRestore() + }) + it('stops microphone tracks on stop', async () => { const stream = makeStream() getUserMedia.mockResolvedValueOnce(stream) diff --git a/packages/ai-client/tests/byok-client.test.ts b/packages/ai-client/tests/byok-client.test.ts index 6754240720..467ecd0d9f 100644 --- a/packages/ai-client/tests/byok-client.test.ts +++ b/packages/ai-client/tests/byok-client.test.ts @@ -8,6 +8,51 @@ describe('defineByok memory', () => { vi.unstubAllGlobals() }) + it('isolates and freezes nested snapshot state', async () => { + const byok = defineByok() + await byok.update('openai', 'sk-abcdefghij') + const snapshot = byok.getSnapshot() + + expect(Object.isFrozen(snapshot)).toBe(true) + expect(Object.isFrozen(snapshot.status)).toBe(true) + expect(Object.isFrozen(snapshot.status.openai)).toBe(true) + Reflect.set(snapshot.status.openai!, 'masked', 'leak') + expect(byok.getSnapshot().status.openai).toEqual({ + state: 'set', + masked: 'ghij', + }) + }) + + it('uses Set subscription behavior', () => { + const byok = defineByok() + const listener = vi.fn() + const first = byok.subscribe(listener) + byok.subscribe(listener) + + byok.request('openai', 'missing') + expect(listener).toHaveBeenCalledTimes(1) + first() + byok.request('openai', 'missing') + expect(listener).toHaveBeenCalledTimes(1) + }) + + it('a throwing listener does not skip the next listener', () => { + const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}) + const byok = defineByok() + const listener = vi.fn() + byok.subscribe(() => { + throw new Error('listener failed') + }) + byok.subscribe(listener) + + byok.request('openai', 'missing') + + expect(listener).toHaveBeenCalledTimes(1) + expect(byok.getSnapshot().prompt?.provider).toBe('openai') + expect(consoleError).toHaveBeenCalledWith(new Error('listener failed')) + consoleError.mockRestore() + }) + it('updates and snapshots without exposing the raw key', async () => { const byok = defineByok({ storage: memoryStorage() }) await byok.update('openai', 'sk-abcdefghij') diff --git a/packages/ai-client/tests/byok-generation-client.test.ts b/packages/ai-client/tests/byok-generation-client.test.ts index f11d2d79e1..7338305c0f 100644 --- a/packages/ai-client/tests/byok-generation-client.test.ts +++ b/packages/ai-client/tests/byok-generation-client.test.ts @@ -87,6 +87,64 @@ describe('GenerationClient byok', () => { expect(client.getError()).toBeInstanceOf(ByokMissingError) }) + it('stops the failed run when the BYOK request subscriber stops', async () => { + const byok = defineByok({ storage: memoryStorage() }) + await byok.update('elevenlabs', ELEVENLABS_KEY) + const onErrorChange = vi.fn() + let client!: GenerationClient<{ prompt: string }, { id: string }> + client = new GenerationClient({ + fetcher: async (): Promise<{ id: string }> => { + throw new ByokMissingError('elevenlabs') + }, + byok, + byokProvider: () => 'elevenlabs', + onErrorChange, + }) + byok.subscribe(() => { + if (byok.getSnapshot().prompt) client.stop() + }) + + await client.generate({ prompt: 'first' }) + + expect(client.getStatus()).toBe('idle') + expect(client.getError()).toBeUndefined() + expect(onErrorChange).not.toHaveBeenCalledWith(expect.any(Error)) + }) + + it('does not let a failed BYOK run overwrite a replacement', async () => { + const byok = defineByok({ storage: memoryStorage() }) + await byok.update('elevenlabs', ELEVENLABS_KEY) + let releaseSecond!: (value: { id: string }) => void + const secondResult = new Promise<{ id: string }>((resolve) => { + releaseSecond = resolve + }) + let secondGenerate: Promise | undefined + let replaced = false + const client = new GenerationClient({ + fetcher: async (input) => { + if (input.prompt === 'first') { + throw new ByokMissingError('elevenlabs') + } + return secondResult + }, + byok, + byokProvider: () => 'elevenlabs', + }) + byok.subscribe(() => { + if (!byok.getSnapshot().prompt || replaced) return + replaced = true + client.stop() + secondGenerate = client.generate({ prompt: 'second' }) + }) + + await client.generate({ prompt: 'first' }) + + expect(client.getStatus()).toBe('generating') + expect(client.getError()).toBeUndefined() + releaseSecond({ id: 'second' }) + await secondGenerate + }) + it('requests a missing key when the fetcher returns a byokMissing Response', async () => { const byok = defineByok({ storage: memoryStorage() }) await byok.update('elevenlabs', ELEVENLABS_KEY) diff --git a/packages/ai-client/tests/byok-video-generation-client.test.ts b/packages/ai-client/tests/byok-video-generation-client.test.ts index 166b3eb4d4..2b7633c385 100644 --- a/packages/ai-client/tests/byok-video-generation-client.test.ts +++ b/packages/ai-client/tests/byok-video-generation-client.test.ts @@ -92,6 +92,64 @@ describe('VideoGenerationClient byok', () => { expect(client.getError()).toBeInstanceOf(ByokMissingError) }) + it('stops the failed run when the BYOK request subscriber stops', async () => { + const byok = defineByok({ storage: memoryStorage() }) + await byok.update('elevenlabs', ELEVENLABS_KEY) + const onErrorChange = vi.fn() + let client!: VideoGenerationClient + client = new VideoGenerationClient({ + fetcher: async () => { + throw new ByokMissingError('elevenlabs') + }, + byok, + byokProvider: () => 'elevenlabs', + onErrorChange, + }) + byok.subscribe(() => { + if (byok.getSnapshot().prompt) client.stop() + }) + + await client.generate({ prompt: 'first' }) + + expect(client.getStatus()).toBe('idle') + expect(client.getError()).toBeUndefined() + expect(onErrorChange).not.toHaveBeenCalledWith(expect.any(Error)) + }) + + it('does not let a failed BYOK run overwrite a replacement', async () => { + const byok = defineByok({ storage: memoryStorage() }) + await byok.update('elevenlabs', ELEVENLABS_KEY) + let releaseSecond!: (value: typeof VIDEO_RESULT) => void + const secondResult = new Promise((resolve) => { + releaseSecond = resolve + }) + let secondGenerate: Promise | undefined + let replaced = false + const client = new VideoGenerationClient({ + fetcher: async (input) => { + if (input.prompt === 'first') { + throw new ByokMissingError('elevenlabs') + } + return secondResult + }, + byok, + byokProvider: () => 'elevenlabs', + }) + byok.subscribe(() => { + if (!byok.getSnapshot().prompt || replaced) return + replaced = true + client.stop() + secondGenerate = client.generate({ prompt: 'second' }) + }) + + await client.generate({ prompt: 'first' }) + + expect(client.getStatus()).toBe('generating') + expect(client.getError()).toBeUndefined() + releaseSecond(VIDEO_RESULT) + await secondGenerate + }) + it('requests a missing key when the fetcher returns a byokMissing Response', async () => { const byok = defineByok({ storage: memoryStorage() }) await byok.update('elevenlabs', ELEVENLABS_KEY) diff --git a/packages/ai-client/tests/chat-client-history.test.ts b/packages/ai-client/tests/chat-client-history.test.ts index b558bea6e6..1125eda5a4 100644 --- a/packages/ai-client/tests/chat-client-history.test.ts +++ b/packages/ai-client/tests/chat-client-history.test.ts @@ -139,6 +139,7 @@ describe('ChatClient history paging', () => { 'm3', ]) }) + expect(client.getSnapshot().hasOlderMessages).toBe(true) await client.loadOlderMessages() @@ -148,6 +149,12 @@ describe('ChatClient history paging', () => { 'm3', ]) expect(client.getHasOlderMessages()).toBe(false) + expect(client.getSnapshot().hasOlderMessages).toBe(false) + expect(client.getSnapshot().messages.map((message) => message.id)).toEqual([ + 'm1', + 'm2', + 'm3', + ]) }) it('send with history.pageSize posts only the new turn', async () => { @@ -538,5 +545,6 @@ describe('ChatClient history paging', () => { expect(client.getMessages()).toEqual([]) expect(client.getHasOlderMessages()).toBe(false) + expect(client.getSnapshot().hasOlderMessages).toBe(false) }) }) diff --git a/packages/ai-client/tests/chat-client-subagents.test.ts b/packages/ai-client/tests/chat-client-subagents.test.ts index a4c3a0a175..ff484d327d 100644 --- a/packages/ai-client/tests/chat-client-subagents.test.ts +++ b/packages/ai-client/tests/chat-client-subagents.test.ts @@ -297,6 +297,98 @@ describe('ChatClient subagent handles for restored messages', () => { expect(client.getSubagents()).toEqual([]) }) + it('accepts frozen snapshot messages back without throwing', () => { + const client = new ChatClient({ + connection: createMockConnectionAdapter({ chunks: [] }), + initialMessages: [ + { + id: 'a1', + role: 'assistant', + parts: [ + { + type: 'subagent', + subagent: { + id: 'sub-restored', + name: 'researcher', + status: 'finished', + messages: [], + }, + }, + ], + }, + ], + }) + const saved = client.getSnapshot().messages + expect(Object.isFrozen(saved[0]?.parts[0])).toBe(true) + + client.clear() + expect(() => client.setMessagesManually([...saved])).not.toThrow() + expect(typeof client.getSubagents()[0]?.stop).toBe('function') + + const fork = new ChatClient({ + connection: createMockConnectionAdapter({ chunks: [] }), + initialMessages: [...saved], + }) + expect(typeof fork.getSubagents()[0]?.stop).toBe('function') + // Each client owns its handle, so `stop()` reaches the right client. + expect(fork.getSubagents()[0]).not.toBe(client.getSubagents()[0]) + }) + + it('accepts a frozen nested subagent card back without throwing', () => { + const client = new ChatClient({ + connection: createMockConnectionAdapter({ chunks: [] }), + initialMessages: [ + { + id: 'a1', + role: 'assistant', + parts: [ + { + type: 'subagent', + subagent: { + id: 'sub-outer', + name: 'planner', + status: 'finished', + messages: [ + { + id: 'child-1', + role: 'assistant', + parts: [ + { + type: 'subagent', + subagent: { + id: 'sub-inner', + name: 'researcher', + status: 'finished', + messages: [], + }, + }, + ], + }, + ], + }, + }, + ], + }, + ], + }) + const saved = client.getSnapshot().messages + const outer = saved[0]?.parts[0] + expect( + outer?.type === 'subagent' && + Object.isFrozen(outer.subagent.messages[0]?.parts[0]), + ).toBe(true) + + client.clear() + expect(() => client.setMessagesManually(saved)).not.toThrow() + expect(client.getSubagents().map((handle) => handle.id)).toEqual([ + 'sub-outer', + 'sub-inner', + ]) + for (const handle of client.getSubagents()) { + expect(typeof handle.stop).toBe('function') + } + }) + it('gives a nested card a handle with stop', () => { const client = new ChatClient({ connection: createMockConnectionAdapter({ chunks: [] }), diff --git a/packages/ai-client/tests/chat-client.test.ts b/packages/ai-client/tests/chat-client.test.ts index 84e4f09bfa..384b0eef25 100644 --- a/packages/ai-client/tests/chat-client.test.ts +++ b/packages/ai-client/tests/chat-client.test.ts @@ -68,6 +68,226 @@ describe('ChatClient', () => { expect(client.getError()).toBeUndefined() }) + it('exposes a UI snapshot that subscribeSnapshot notifies', () => { + const adapter = createMockConnectionAdapter() + const initialMessages: Array = [ + { + id: 'msg-1', + role: 'user', + parts: [{ type: 'text', content: 'Hello' }], + createdAt: new Date(), + }, + ] + const client = new ChatClient({ + connection: adapter, + initialMessages, + }) + + const snapshot = client.getSnapshot() + expect(snapshot.messages).toEqual(initialMessages) + expect(snapshot.status).toBe('ready') + expect(snapshot.isLoading).toBe(false) + expect(client.getSnapshot()).toBe(snapshot) + + const listener = vi.fn() + const stop = client.subscribeSnapshot(listener) + client.setMessagesManually([]) + expect(listener).toHaveBeenCalledTimes(1) + expect(client.getSnapshot().messages).toEqual([]) + expect(client.getSnapshot()).not.toBe(snapshot) + stop() + client.setMessagesManually([ + { + id: 'msg-2', + role: 'user', + parts: [{ type: 'text', content: 'After' }], + }, + ]) + expect(listener).toHaveBeenCalledTimes(1) + }) + + it('isolates nested message parts from callers and old snapshots', () => { + const adapter = createMockConnectionAdapter() + const source = { + type: 'url' as const, + value: 'https://example.com/a.png', + } + const client = new ChatClient({ + connection: adapter, + initialMessages: [ + { + id: 'msg-1', + role: 'user', + parts: [{ type: 'image', source }], + createdAt: new Date(), + }, + ], + }) + const oldSnapshot = client.getSnapshot() + + source.value = 'https://example.com/changed.png' + expect(oldSnapshot.messages[0]?.parts[0]).toMatchObject({ + source: { value: 'https://example.com/a.png' }, + }) + const publishedPart = oldSnapshot.messages[0]?.parts[0] + expect(Object.isFrozen(publishedPart)).toBe(true) + if (!publishedPart || !('source' in publishedPart)) { + throw new Error('Expected a message part with a source') + } + expect(Object.isFrozen(publishedPart.source)).toBe(true) + }) + + it('copies tool input, approval, and subagent transcripts', () => { + const input = { city: 'Paris', nested: { n: 1 } } + const approval = { id: 'ap-1', needsApproval: true } + const output = { ok: true } + const child = { + id: 'child-1', + role: 'assistant' as const, + parts: [{ type: 'text' as const, content: 'from child' }], + } + const client = new ChatClient({ + connection: createMockConnectionAdapter(), + initialMessages: [ + { + id: 'msg-1', + role: 'assistant', + parts: [ + { + type: 'tool-call', + id: 'call-1', + name: 'lookup', + arguments: '{"city":"Paris"}', + state: 'input-complete', + input, + output, + approval, + }, + { + type: 'tool-result', + toolCallId: 'call-1', + content: [{ type: 'text', content: 'done' }], + state: 'complete', + }, + { + type: 'subagent', + subagent: { + id: 'sub-1', + name: 'researcher', + status: 'finished', + messages: [child], + }, + }, + ], + }, + ], + }) + const parts = client.getSnapshot().messages[0]?.parts + const toolCall = parts?.[0] + const toolResult = parts?.[1] + const subagentPart = parts?.[2] + if ( + toolCall?.type !== 'tool-call' || + toolResult?.type !== 'tool-result' || + subagentPart?.type !== 'subagent' || + !Array.isArray(toolResult.content) + ) { + throw new Error('Expected tool-call, tool-result, and subagent parts') + } + + expect(Object.isFrozen(toolCall.input)).toBe(true) + expect(Object.isFrozen(toolCall.approval)).toBe(true) + expect(Object.isFrozen(toolCall.output)).toBe(true) + expect(Reflect.set(toolCall.input, 'city', 'London')).toBe(false) + if (!toolCall.approval) { + throw new Error('Expected tool-call approval') + } + expect(Reflect.set(toolCall.approval, 'approved', true)).toBe(false) + expect(input.city).toBe('Paris') + expect(approval).not.toHaveProperty('approved') + expect(Object.isFrozen(toolResult.content)).toBe(true) + expect(Object.isFrozen(toolResult.content[0])).toBe(true) + + const publishedChild = subagentPart.subagent.messages[0] + expect(Object.isFrozen(subagentPart.subagent.messages)).toBe(true) + expect(publishedChild).not.toBe(child) + expect(Reflect.set(publishedChild?.parts ?? [], 'length', 0)).toBe(false) + expect(child.parts).toHaveLength(1) + const handle = client.getSubagents()[0] + expect(subagentPart.subagent).not.toBe(handle) + expect(subagentPart.subagent.stop).toBe(handle?.stop) + }) + + it('isolates messages published by normal updates', () => { + const client = new ChatClient({ + connection: createMockConnectionAdapter(), + }) + const oldSnapshot = client.getSnapshot() + const source = { + type: 'url' as const, + value: 'https://example.com/updated.png', + } + + client.setMessagesManually([ + { + id: 'updated-1', + role: 'user', + parts: [{ type: 'image', source }], + }, + ]) + + const snapshot = client.getSnapshot() + const part = snapshot.messages[0]?.parts[0] + expect(snapshot).not.toBe(oldSnapshot) + expect(oldSnapshot.messages).toEqual([]) + expect(Object.isFrozen(snapshot.messages)).toBe(true) + expect(Object.isFrozen(snapshot.messages[0])).toBe(true) + expect(Object.isFrozen(snapshot.messages[0]?.parts)).toBe(true) + expect(Object.isFrozen(part)).toBe(true) + expect(part).toMatchObject({ source }) + + if (!part || !('source' in part)) { + throw new Error('Expected a message part with a source') + } + const publishedSource = part.source + expect(Object.isFrozen(publishedSource)).toBe(true) + expect(Reflect.set(publishedSource, 'value', 'changed')).toBe(false) + expect(client.getMessages()[0]?.parts[0]).toMatchObject({ source }) + expect(client.getSnapshot()).toBe(snapshot) + }) + + it('keeps the identity of unchanged messages across snapshots', async () => { + const client = new ChatClient({ + connection: createMockConnectionAdapter({ + chunks: createTextChunks('Hi there'), + }), + initialMessages: [ + { + id: 'earlier-1', + role: 'user', + parts: [{ type: 'text', content: 'Earlier' }], + }, + ], + }) + const first = client.getSnapshot().messages[0] + const seen: Array = [] + const stop = client.subscribeSnapshot(() => { + seen.push(client.getSnapshot().messages[0]) + }) + + await client.sendMessage('Hello') + stop() + + const messages = client.getSnapshot().messages + expect(messages.length).toBeGreaterThan(2) + expect(seen.length).toBeGreaterThan(1) + for (const message of seen) { + expect(message).toBe(first) + } + expect(messages[0]).toBe(first) + expect(Object.isFrozen(first)).toBe(true) + }) + it('should initialize with provided messages', () => { const adapter = createMockConnectionAdapter() const initialMessages: Array = [ @@ -2270,6 +2490,61 @@ describe('ChatClient', () => { { type: 'text', content: 'Second' }, ]) }) + + it('isolates queue entries published by normal updates', async () => { + const client = new ChatClient({ + connection: createMockConnectionAdapter({ + chunks: createTextChunks('Response'), + chunkDelay: 100, + }), + }) + const firstSend = client.sendMessage('First') + const oldSnapshot = client.getSnapshot() + + await client.sendMessage({ + content: [ + { + type: 'image', + source: { + type: 'url', + value: 'https://example.com/queued.png', + }, + }, + ], + }) + + const snapshot = client.getSnapshot() + const queuedContent = snapshot.queue[0]?.content + expect(snapshot).not.toBe(oldSnapshot) + expect(oldSnapshot.queue).toEqual([]) + expect(Object.isFrozen(snapshot.queue)).toBe(true) + expect(Object.isFrozen(snapshot.queue[0])).toBe(true) + expect(typeof queuedContent).toBe('object') + if (typeof queuedContent === 'object') { + expect(Object.isFrozen(queuedContent)).toBe(true) + expect(Object.isFrozen(queuedContent.content)).toBe(true) + const queuedPart = Array.isArray(queuedContent.content) + ? queuedContent.content[0] + : undefined + expect(Object.isFrozen(queuedPart)).toBe(true) + if (!queuedPart || !('source' in queuedPart)) { + throw new Error('Expected a queued content part with a source') + } + const publishedSource = queuedPart.source + expect(Object.isFrozen(publishedSource)).toBe(true) + expect(Reflect.set(publishedSource, 'value', 'changed')).toBe(false) + } + expect(client.getQueue()[0]?.content).toMatchObject({ + content: [ + { + source: { value: 'https://example.com/queued.png' }, + }, + ], + }) + expect(client.getSnapshot()).toBe(snapshot) + + await firstSend + }) }) describe('append', () => { diff --git a/packages/ai-client/tests/client-snapshots.test.ts b/packages/ai-client/tests/client-snapshots.test.ts new file mode 100644 index 0000000000..3450436671 --- /dev/null +++ b/packages/ai-client/tests/client-snapshots.test.ts @@ -0,0 +1,393 @@ +import { describe, expect, it, vi } from 'vitest' +import { createMockConnectionAdapter, createTextChunks } from './test-utils' +import { + AudioRecorder, + ChatClient, + GenerationClient, + RealtimeClient, + VideoGenerationClient, +} from '../src' +import type { RealtimeAdapter, RealtimeConnection } from '../src' + +describe('client UI snapshots', () => { + it('publishes and clears the active generation run id', async () => { + let resolve!: (value: { id: string }) => void + const pending = new Promise<{ id: string }>((done) => { + resolve = done + }) + const client = new GenerationClient({ fetcher: async () => pending }) + + const generating = client.generate({ prompt: 'test' }) + expect(client.getSnapshot().runId).toEqual(expect.any(String)) + expect(client.getSnapshot().runId).not.toBe('') + + resolve({ id: '1' }) + await generating + expect(client.getSnapshot().runId).toBeNull() + }) + + it('GenerationClient notifies subscribe and keeps getSnapshot identity until a change', async () => { + const client = new GenerationClient({ + fetcher: async () => ({ id: '1' }), + }) + const snapshot = client.getSnapshot() + expect(snapshot).toEqual({ + result: null, + isLoading: false, + error: undefined, + status: 'idle', + runId: null, + }) + expect(client.getSnapshot()).toBe(snapshot) + + const listener = vi.fn() + const stop = client.subscribe(listener) + await client.generate({ prompt: 'test' }) + + expect(listener).toHaveBeenCalled() + expect(client.getSnapshot().status).toBe('success') + expect(client.getSnapshot().result).toEqual({ id: '1' }) + expect(client.getSnapshot()).not.toBe(snapshot) + const published = client.getSnapshot().result + expect(published).toEqual({ id: '1' }) + if (published) { + Reflect.set(published, 'id', 'mutated') + } + expect(client.getResult()).toEqual({ id: '1' }) + expect(Object.isFrozen(published)).toBe(true) + expect(client.getSnapshot().result).toEqual({ id: '1' }) + const calls = listener.mock.calls.length + stop() + await client.generate({ prompt: 'again' }) + expect(listener).toHaveBeenCalledTimes(calls) + }) + + it('a throwing generation subscriber does not stick isLoading', async () => { + const client = new GenerationClient({ + fetcher: async () => ({ id: '1' }), + }) + const error = vi.spyOn(console, 'error').mockImplementation(() => {}) + const stop = client.subscribe(() => { + throw new Error('boom') + }) + + await client.generate({ prompt: 'test' }) + + expect(client.getSnapshot().isLoading).toBe(false) + expect(client.getSnapshot().status).toBe('success') + stop() + error.mockRestore() + await client.generate({ prompt: 'next' }) + expect(client.getSnapshot().isLoading).toBe(false) + expect(client.getSnapshot().status).toBe('success') + }) + + it('a throwing chat snapshot subscriber does not stick isLoading', async () => { + const client = new ChatClient({ + connection: createMockConnectionAdapter({ + chunks: createTextChunks('Hello'), + }), + }) + const error = vi.spyOn(console, 'error').mockImplementation(() => {}) + const second = vi.fn() + const stopFirst = client.subscribeSnapshot(() => { + throw new Error('listener failed') + }) + const stopSecond = client.subscribeSnapshot(second) + + await client.sendMessage('Hi') + + expect(second).toHaveBeenCalled() + expect(client.getSnapshot().isLoading).toBe(false) + expect(client.getSnapshot().status).toBe('ready') + expect(client.getSnapshot().messages.at(-1)?.role).toBe('assistant') + stopFirst() + stopSecond() + error.mockRestore() + }) + + it('VideoGenerationClient snapshot includes job fields', async () => { + const result = { + jobId: 'job-1', + status: 'completed' as const, + url: 'https://example.com/video.mp4', + } + const client = new VideoGenerationClient({ + fetcher: async () => result, + }) + const snapshot = client.getSnapshot() + expect(snapshot.jobId).toBeNull() + expect(snapshot.videoStatus).toBeNull() + expect(client.getSnapshot()).toBe(snapshot) + + const listener = vi.fn() + const stop = client.subscribe(listener) + await client.generate({ prompt: 'test video' }) + + expect(listener).toHaveBeenCalled() + expect(client.getSnapshot().status).toBe('success') + expect(client.getSnapshot().result).toEqual(result) + expect(client.getSnapshot().jobId).toBe('job-1') + expect(client.getSnapshot().videoStatus?.status).toBe('completed') + expect(client.getSnapshot()).not.toBe(snapshot) + const completed = client.getSnapshot() + expect(Object.isFrozen(completed.videoStatus)).toBe(true) + Reflect.set(completed.videoStatus!, 'status', 'failed') + expect(client.getVideoStatus()?.status).toBe('completed') + const calls = listener.mock.calls.length + stop() + await client.generate({ prompt: 'again' }) + expect(listener).toHaveBeenCalledTimes(calls) + }) + + it('AudioRecorder getSnapshot matches state', () => { + const recorder = new AudioRecorder() + expect(recorder.getSnapshot()).toBe('idle') + expect(recorder.getSnapshot()).toBe(recorder.state) + }) + + it('RealtimeClient subscribe fires when connect fails', async () => { + const adapter: RealtimeAdapter = { + provider: 'test', + connect: async () => { + throw new Error('unused') + }, + } + const client = new RealtimeClient({ + getToken: async () => { + throw new Error('no token') + }, + adapter, + }) + const snapshot = client.getSnapshot() + expect(snapshot.status).toBe('idle') + expect(client.getSnapshot()).toBe(snapshot) + + const listener = vi.fn() + const stop = client.subscribe(listener) + await expect(client.connect()).rejects.toThrow('no token') + + expect(listener).toHaveBeenCalled() + expect(client.getSnapshot().status).toBe('error') + expect(client.getSnapshot()).not.toBe(snapshot) + const calls = listener.mock.calls.length + stop() + await expect(client.connect()).rejects.toThrow('no token') + expect(listener).toHaveBeenCalledTimes(calls) + client.destroy() + }) + + it('RealtimeClient does not read connection control from its public snapshot', async () => { + const getToken = vi.fn(async () => { + throw new Error('token attempted') + }) + const client = new RealtimeClient({ + getToken, + adapter: { provider: 'test', connect: vi.fn() }, + }) + + Reflect.set(client.getSnapshot(), 'status', 'connected') + await expect(client.connect()).rejects.toThrow('token attempted') + expect(getToken).toHaveBeenCalledOnce() + }) + + it('RealtimeClient isolates connection control from state change callbacks', async () => { + const getToken = vi.fn(async () => ({ + provider: 'test', + token: 'token', + expiresAt: Date.now() + 120_000, + config: {}, + })) + const connect = vi.fn(async () => { + throw new Error('connect attempted') + }) + const client = new RealtimeClient({ + getToken, + adapter: { provider: 'test', connect }, + }) + client.onStateChange((state) => { + Reflect.set(state, 'status', 'connected') + }) + + await expect(client.connect()).rejects.toThrow('connect attempted') + await expect(client.connect()).rejects.toThrow('connect attempted') + + expect(getToken).toHaveBeenCalledTimes(2) + expect(connect).toHaveBeenCalledTimes(2) + }) + + it('keeps the transcript when connect fails and clears it after a later success', async () => { + let connects = 0 + const connection: RealtimeConnection = { + disconnect: async () => {}, + startAudioCapture: async () => {}, + stopAudioCapture: () => {}, + sendText: () => {}, + sendImage: () => {}, + sendToolResult: () => {}, + updateSession: () => {}, + interrupt: () => {}, + on: () => () => {}, + getAudioVisualization: () => ({ + inputLevel: 0, + outputLevel: 0, + getInputFrequencyData: () => new Uint8Array(), + getOutputFrequencyData: () => new Uint8Array(), + getInputTimeDomainData: () => new Uint8Array(), + getOutputTimeDomainData: () => new Uint8Array(), + inputSampleRate: 0, + outputSampleRate: 0, + }), + } + const client = new RealtimeClient({ + getToken: async () => { + connects++ + if (connects === 2) throw new Error('no token') + return { + provider: 'test', + token: 'token', + expiresAt: Date.now() + 120_000, + config: {}, + } + }, + adapter: { provider: 'test', connect: async () => connection }, + autoCapture: false, + }) + + await client.connect() + client.sendText('keep me') + expect(client.getSnapshot().messages).toHaveLength(1) + await client.disconnect() + expect(client.getSnapshot().messages).toHaveLength(1) + + await expect(client.connect()).rejects.toThrow('no token') + expect(client.getSnapshot().status).toBe('error') + expect(client.getSnapshot().messages).toHaveLength(1) + + await client.connect() + expect(client.getSnapshot().status).toBe('connected') + expect(client.getSnapshot().messages).toEqual([]) + client.destroy() + }) + + it('RealtimeClient keeps a transcript that arrives while audio capture starts', async () => { + const handlers = new Map void>() + const connection: RealtimeConnection = { + disconnect: async () => {}, + // The provider greets the user while the mic permission prompt is open. + startAudioCapture: async () => { + handlers.get('transcript')?.({ + role: 'assistant', + transcript: 'Hello', + isFinal: false, + }) + }, + stopAudioCapture: () => {}, + sendText: () => {}, + sendImage: () => {}, + sendToolResult: () => {}, + updateSession: () => {}, + interrupt: () => {}, + on: (event: string, handler: (payload: any) => void) => { + handlers.set(event, handler) + return () => handlers.delete(event) + }, + getAudioVisualization: () => ({ + inputLevel: 0, + outputLevel: 0, + getInputFrequencyData: () => new Uint8Array(), + getOutputFrequencyData: () => new Uint8Array(), + getInputTimeDomainData: () => new Uint8Array(), + getOutputTimeDomainData: () => new Uint8Array(), + inputSampleRate: 0, + outputSampleRate: 0, + }), + } as RealtimeConnection + const client = new RealtimeClient({ + getToken: async () => ({ + provider: 'test', + token: 'token', + expiresAt: Date.now() + 120_000, + config: {}, + }), + adapter: { provider: 'test', connect: async () => connection }, + }) + + await client.connect() + + expect(client.getSnapshot().status).toBe('connected') + expect(client.getSnapshot().pendingAssistantTranscript).toBe('Hello') + client.destroy() + }) + + it('RealtimeClient state change callbacks are isolated from a throwing one', async () => { + const client = new RealtimeClient({ + getToken: async () => { + throw new Error('no token') + }, + adapter: { provider: 'test', connect: vi.fn() }, + }) + const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}) + const seen: Array = [] + client.onStateChange(() => { + throw new Error('listener failed') + }) + client.onStateChange((state) => seen.push(state.status)) + + await expect(client.connect()).rejects.toThrow('no token') + + expect(seen).toEqual(['connecting', 'error']) + expect(consoleError).toHaveBeenCalledWith(new Error('listener failed')) + consoleError.mockRestore() + }) + + it('RealtimeClient keeps the identity of unchanged messages across snapshots', async () => { + const connection: RealtimeConnection = { + disconnect: async () => {}, + startAudioCapture: async () => {}, + stopAudioCapture: () => {}, + sendText: () => {}, + sendImage: () => {}, + sendToolResult: () => {}, + updateSession: () => {}, + interrupt: () => {}, + on: () => () => {}, + getAudioVisualization: () => ({ + inputLevel: 0, + outputLevel: 0, + getInputFrequencyData: () => new Uint8Array(), + getOutputFrequencyData: () => new Uint8Array(), + getInputTimeDomainData: () => new Uint8Array(), + getOutputTimeDomainData: () => new Uint8Array(), + inputSampleRate: 0, + outputSampleRate: 0, + }), + } + const client = new RealtimeClient({ + getToken: async () => ({ + provider: 'test', + token: 'token', + expiresAt: Date.now() + 120_000, + config: {}, + }), + adapter: { provider: 'test', connect: async () => connection }, + autoCapture: false, + }) + await client.connect() + + client.sendText('first') + const first = client.getSnapshot().messages[0] + client.sendText('second') + + const messages = client.getSnapshot().messages + expect(messages).toHaveLength(2) + expect(messages[0]).toBe(first) + expect(Object.isFrozen(first)).toBe(true) + + // A state change that is not a message change keeps the same array. + client.stopListening() + expect(client.getSnapshot().mode).toBe('idle') + expect(client.getSnapshot().messages).toBe(messages) + client.destroy() + }) +}) diff --git a/packages/ai-client/tests/generation-client.test.ts b/packages/ai-client/tests/generation-client.test.ts index 8b941689dd..7a20ecb6e2 100644 --- a/packages/ai-client/tests/generation-client.test.ts +++ b/packages/ai-client/tests/generation-client.test.ts @@ -115,6 +115,7 @@ describe('GenerationClient', () => { expect(onResultChange).toHaveBeenCalledWith(mockResult) expect(client.getResult()).toEqual(mockResult) expect(client.getStatus()).toBe('success') + expect(client.getSnapshot().runId).toBeNull() expect(client.getIsLoading()).toBe(false) }) @@ -245,6 +246,57 @@ describe('GenerationClient', () => { expect(client.getStatus()).toBe('success') }) + it('keeps the local run id when a stream chunk has no resume pointer', async () => { + const gate = createDeferred() + const chunks: Array = [ + { + type: EventType.CUSTOM, + name: 'generation:progress', + value: { percent: 1 }, + timestamp: Date.now(), + }, + { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: 'generation:result', + value: { id: '1' }, + timestamp: Date.now(), + }, + { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + metadata: { tanstack: { finishReason: 'stop' as const } }, + timestamp: Date.now(), + }, + ] + const connection = { + async *connect() { + yield chunks[0]! + await gate.promise + for (const chunk of chunks.slice(1)) { + yield chunk + } + }, + } + + const client = new GenerationClient({ connection }) + const pending = client.generate({ prompt: 'test' }) + await waitForCondition(() => { + expect(client.getSnapshot().isLoading).toBe(true) + }) + expect(client.getSnapshot().runId).toEqual(expect.any(String)) + expect(client.getSnapshot().runId).not.toBe('') + gate.resolve() + await pending + expect(client.getSnapshot().runId).toBeNull() + }) + it('should handle RUN_ERROR from stream', async () => { const onError = vi.fn() @@ -274,6 +326,7 @@ describe('GenerationClient', () => { expect(onError).toHaveBeenCalledWith(expect.any(Error)) expect(client.getStatus()).toBe('error') expect(client.getError()?.message).toBe('Generation failed') + expect(client.getSnapshot().runId).toBeNull() }) it('should report progress from CUSTOM progress events', async () => { @@ -408,10 +461,12 @@ describe('GenerationClient', () => { const generatePromise = client.generate({ prompt: 'test' }) expect(client.getIsLoading()).toBe(true) + expect(client.getSnapshot().runId).toEqual(expect.any(String)) client.stop() expect(client.getIsLoading()).toBe(false) expect(client.getStatus()).toBe('idle') + expect(client.getSnapshot().runId).toBeNull() resolvePromise!({ id: '1' }) await generatePromise @@ -432,6 +487,7 @@ describe('GenerationClient', () => { expect(client.getResult()).toBeNull() expect(client.getError()).toBeUndefined() expect(client.getStatus()).toBe('idle') + expect(client.getSnapshot().runId).toBeNull() }) }) @@ -621,12 +677,14 @@ describe('GenerationClient', () => { await waitForCondition(() => { expect(signals).toHaveLength(2) expect(client.getIsLoading()).toBe(true) + expect(client.getSnapshot().runId).toBe('run-2') }) firstCanFinish.resolve(undefined) await firstGenerate expect(client.getIsLoading()).toBe(true) + expect(client.getSnapshot().runId).toBe('run-2') client.stop() expect(signals[1]?.aborted).toBe(true) @@ -758,12 +816,14 @@ describe('GenerationClient', () => { await waitForCondition(() => { expect(signals).toHaveLength(2) expect(client.getIsLoading()).toBe(true) + expect(client.getSnapshot().runId).toBe('run-2') }) firstCanFinish.resolve(undefined) await firstGenerate expect(client.getIsLoading()).toBe(true) + expect(client.getSnapshot().runId).toBe('run-2') client.stop() expect(signals[1]?.aborted).toBe(true) @@ -1504,6 +1564,7 @@ describe('GenerationClient', () => { }) expect(joinRun).toHaveBeenCalledWith(runId, expect.anything()) expect(client.getIsLoading()).toBe(false) + expect(client.getSnapshot().runId).toBeNull() }) it('surfaces an error (never stays stuck generating) when the rejoin throws', async () => { @@ -2199,3 +2260,613 @@ describe('GenerationClient', () => { }) }) }) + +// One harness over both clients, so each run-ownership race is pinned once. +interface RunClient { + generate(input: { prompt: string }): Promise + stop(): void + reset(): void + dispose(): void + mountDevtools(): void + subscribe(listener: () => void): () => void + getSnapshot(): { runId: string | null; isLoading: boolean } + getIsLoading(): boolean + getStatus(): string + getError(): Error | undefined + getResumeSnapshot(): unknown +} + +interface RunHarness { + client: RunClient + resultId: () => string | undefined +} + +interface RunHarnessOptions { + fetcher?: (prompt: string) => Promise + connection?: ConnectConnectionAdapter + onChunk?: (chunk: StreamChunk) => void + onLoadingChange?: (isLoading: boolean) => void + onStatusChange?: (status: string) => void + onErrorChange?: (error: Error | undefined) => void + onError?: (error: Error) => void + onResultChange?: (resultId: string | undefined) => void +} + +const promptOf = (input: { prompt: unknown }) => + typeof input.prompt === 'string' ? input.prompt : '' + +const runClientKinds: Array<{ + name: string + create: (options: RunHarnessOptions) => RunHarness +}> = [ + { + name: 'GenerationClient', + create: (options) => { + const { fetcher, connection, onResultChange, ...callbacks } = options + const client = new GenerationClient<{ prompt: string }, { id: string }>({ + ...callbacks, + ...(connection + ? { connection } + : { + fetcher: async (input: { prompt: string }) => ({ + id: await (fetcher ?? (async (prompt) => prompt))(input.prompt), + }), + }), + onResultChange: (result) => onResultChange?.(result?.id), + }) + return { client, resultId: () => client.getResult()?.id } + }, + }, + { + name: 'VideoGenerationClient', + create: (options) => { + const { fetcher, connection, onResultChange, ...callbacks } = options + const client = new VideoGenerationClient({ + ...callbacks, + ...(connection + ? { connection } + : { + fetcher: async (input: { prompt: unknown }) => { + const jobId = await (fetcher ?? (async (prompt) => prompt))( + promptOf(input), + ) + return { + jobId, + status: 'completed' as const, + url: `https://example.com/${jobId}.mp4`, + } + }, + }), + onResultChange: (result) => onResultChange?.(result?.jobId), + }) + return { client, resultId: () => client.getResult()?.jobId } + }, + }, +] + +describe.each(runClientKinds)('$name run ownership', ({ create }) => { + it('should not overwrite a run started by the stopped loading callback', async () => { + const first = createDeferred() + const second = createDeferred() + let secondGenerate: Promise | undefined + let secondRunId: string | null = null + let replaced = false + let h!: RunHarness + h = create({ + fetcher: (prompt) => + prompt === 'first' ? first.promise : second.promise, + onLoadingChange: (isLoading) => { + if (isLoading || replaced) return + replaced = true + secondGenerate = h.client.generate({ prompt: 'second' }) + secondRunId = h.client.getSnapshot().runId + }, + }) + + const firstGenerate = h.client.generate({ prompt: 'first' }) + h.client.stop() + + expect(secondRunId).not.toBeNull() + expect(h.client.getSnapshot().runId).toBe(secondRunId) + expect(h.client.getIsLoading()).toBe(true) + expect(h.client.getStatus()).toBe('generating') + first.resolve('first') + second.resolve('second') + await Promise.all([firstGenerate, secondGenerate]) + }) + + it('should not overwrite a run started by the stopped status callback', async () => { + const first = createDeferred() + const second = createDeferred() + let secondGenerate: Promise | undefined + let h!: RunHarness + h = create({ + fetcher: (prompt) => + prompt === 'first' ? first.promise : second.promise, + onStatusChange: (status) => { + if (status !== 'idle' || secondGenerate) return + secondGenerate = h.client.generate({ prompt: 'second' }) + }, + }) + + const firstGenerate = h.client.generate({ prompt: 'first' }) + h.client.stop() + + expect(secondGenerate).toBeDefined() + expect(h.client.getIsLoading()).toBe(true) + expect(h.client.getSnapshot().isLoading).toBe(true) + expect(h.client.getStatus()).toBe('generating') + first.resolve('first') + second.resolve('second') + await Promise.all([firstGenerate, secondGenerate]) + }) + + it('should publish isLoading false when a stop callback throws', async () => { + const first = createDeferred() + let h!: RunHarness + h = create({ + fetcher: () => first.promise, + onStatusChange: (status) => { + if (status === 'idle') throw new Error('status failed') + }, + }) + + const run = h.client.generate({ prompt: 'first' }) + expect(() => h.client.stop()).toThrow('status failed') + + expect(h.client.getIsLoading()).toBe(false) + expect(h.client.getSnapshot().isLoading).toBe(false) + first.resolve('first') + await run + }) + + it('should not lock generate() when the loading callback throws', async () => { + let failNext = true + const h = create({ + onLoadingChange: (isLoading) => { + if (!isLoading || !failNext) return + failNext = false + throw new Error('loading failed') + }, + }) + + await h.client.generate({ prompt: 'first' }) + + expect(h.client.getStatus()).toBe('error') + expect(h.client.getError()?.message).toBe('loading failed') + expect(h.client.getSnapshot().isLoading).toBe(false) + + await h.client.generate({ prompt: 'second' }) + + expect(h.client.getStatus()).toBe('success') + expect(h.resultId()).toBe('second') + }) + + it('should stop before transport work when loading callback stops', async () => { + const fetcher = vi.fn(async (prompt: string) => prompt) + let h!: RunHarness + h = create({ + fetcher, + onLoadingChange: (isLoading) => { + if (isLoading) h.client.stop() + }, + }) + + await h.client.generate({ prompt: 'test' }) + + expect(fetcher).not.toHaveBeenCalled() + expect(h.client.getStatus()).toBe('idle') + expect(h.client.getSnapshot().runId).toBeNull() + }) + + it('should not complete resume state after success callback stops', async () => { + let h!: RunHarness + h = create({ + onStatusChange: (status) => { + if (status === 'success') h.client.stop() + }, + }) + + await h.client.generate({ prompt: 'test' }) + + expect(h.client.getStatus()).toBe('success') + expect(h.client.getSnapshot().runId).toBeNull() + expect(h.client.getResumeSnapshot()).toBeUndefined() + }) + + it('should not let initialization continue after loading callback starts a replacement', async () => { + const second = createDeferred() + const prompts: Array = [] + let replaced = false + let secondGenerate: Promise | undefined + let h!: RunHarness + h = create({ + fetcher: async (prompt) => { + prompts.push(prompt) + return prompt === 'second' ? second.promise : 'first' + }, + onLoadingChange: (isLoading) => { + if (!isLoading || replaced) return + replaced = true + h.client.stop() + secondGenerate = h.client.generate({ prompt: 'second' }) + }, + }) + + await h.client.generate({ prompt: 'first' }) + + expect(prompts).toEqual(['second']) + expect(h.client.getStatus()).toBe('generating') + second.resolve('second') + await secondGenerate + }) + + it('should not let an old result callback complete a replacement run', async () => { + const second = createDeferred() + let secondGenerate: Promise | undefined + let secondRunId: string | null = null + let h!: RunHarness + h = create({ + fetcher: async (prompt) => + prompt === 'first' ? 'first' : second.promise, + onResultChange: (resultId) => { + if (resultId !== 'first') return + h.client.stop() + secondGenerate = h.client.generate({ prompt: 'second' }) + secondRunId = h.client.getSnapshot().runId + }, + }) + + await h.client.generate({ prompt: 'first' }) + + expect(secondRunId).not.toBeNull() + expect(h.client.getSnapshot().runId).toBe(secondRunId) + expect(h.client.getStatus()).toBe('generating') + expect(h.client.getResumeSnapshot()).toBeUndefined() + second.resolve('second') + await secondGenerate + }) + + it('should not let an old error callback overwrite a replacement run', async () => { + const second = createDeferred() + let secondGenerate: Promise | undefined + let secondRunId: string | null = null + let h!: RunHarness + h = create({ + fetcher: async (prompt) => { + if (prompt === 'first') throw new Error('first failed') + return second.promise + }, + onErrorChange: (error) => { + if (error?.message !== 'first failed') return + h.client.stop() + secondGenerate = h.client.generate({ prompt: 'second' }) + secondRunId = h.client.getSnapshot().runId + }, + }) + + await h.client.generate({ prompt: 'first' }) + + expect(secondRunId).not.toBeNull() + expect(h.client.getSnapshot().runId).toBe(secondRunId) + expect(h.client.getStatus()).toBe('generating') + second.resolve('second') + await secondGenerate + }) + + it('should remain stopped when the error callback stops, and still report the error once', async () => { + const onError = vi.fn() + let h!: RunHarness + h = create({ + fetcher: async () => { + throw new Error('failed') + }, + onErrorChange: (error) => { + if (error) h.client.stop() + }, + onError, + }) + + await h.client.generate({ prompt: 'test' }) + + expect(h.client.getStatus()).toBe('idle') + expect(h.client.getSnapshot().runId).toBeNull() + expect(onError).toHaveBeenCalledTimes(1) + expect(onError).toHaveBeenCalledWith( + expect.objectContaining({ message: 'failed' }), + ) + }) + + it('should not let an old error status callback overwrite a replacement run', async () => { + const second = createDeferred() + let secondGenerate: Promise | undefined + let replaced = false + let h!: RunHarness + h = create({ + fetcher: async (prompt) => { + if (prompt === 'first') throw new Error('first failed') + return second.promise + }, + onStatusChange: (status) => { + if (status !== 'error' || replaced) return + replaced = true + h.client.stop() + secondGenerate = h.client.generate({ prompt: 'second' }) + }, + }) + + await h.client.generate({ prompt: 'first' }) + + expect(h.client.getStatus()).toBe('generating') + expect(h.client.getError()).toBeUndefined() + second.resolve('second') + await secondGenerate + }) + + it('should stop before transport when the run id subscriber stops', async () => { + const fetcher = vi.fn(async (prompt: string) => prompt) + const onLoadingChange = vi.fn() + const h = create({ fetcher, onLoadingChange }) + h.client.subscribe(() => { + if (h.client.getSnapshot().runId) h.client.stop() + }) + + await h.client.generate({ prompt: 'test' }) + + expect(fetcher).not.toHaveBeenCalled() + expect(onLoadingChange).not.toHaveBeenCalledWith(true) + expect(h.client.getSnapshot().runId).toBeNull() + }) + + it('should not continue an old run after the run id subscriber starts a replacement', async () => { + const second = createDeferred() + const prompts: Array = [] + let secondGenerate: Promise | undefined + let replaced = false + const h = create({ + fetcher: async (prompt) => { + prompts.push(prompt) + return prompt === 'second' ? second.promise : 'first' + }, + }) + h.client.subscribe(() => { + if (!h.client.getSnapshot().runId || replaced) return + replaced = true + h.client.stop() + secondGenerate = h.client.generate({ prompt: 'second' }) + }) + + await h.client.generate({ prompt: 'first' }) + + expect(prompts).toEqual(['second']) + expect(h.client.getStatus()).toBe('generating') + second.resolve('second') + await secondGenerate + }) + + it('should not call the loading callback after its snapshot subscriber stops', async () => { + const fetcher = vi.fn(async (prompt: string) => prompt) + const onLoadingChange = vi.fn() + const h = create({ fetcher, onLoadingChange }) + h.client.subscribe(() => { + if (h.client.getSnapshot().isLoading) h.client.stop() + }) + + await h.client.generate({ prompt: 'test' }) + + expect(fetcher).not.toHaveBeenCalled() + expect(onLoadingChange).not.toHaveBeenCalledWith(true) + }) + + it('should not call the old status callback after its subscriber starts a replacement', async () => { + const second = createDeferred() + const generatingStatuses: Array = [] + let secondGenerate: Promise | undefined + let replaced = false + const h = create({ + fetcher: async (prompt) => + prompt === 'second' ? second.promise : 'first', + onStatusChange: (status) => { + if (status === 'generating') generatingStatuses.push(status) + }, + }) + h.client.subscribe(() => { + if (h.client.getStatus() !== 'generating' || replaced) return + replaced = true + h.client.stop() + secondGenerate = h.client.generate({ prompt: 'second' }) + }) + + await h.client.generate({ prompt: 'first' }) + + expect(generatingStatuses).toEqual(['generating']) + expect(h.client.getStatus()).toBe('generating') + second.resolve('second') + await secondGenerate + }) + + it('should not call the old result callback after its subscriber starts a replacement', async () => { + const second = createDeferred() + const onResultChange = vi.fn() + let secondGenerate: Promise | undefined + let replaced = false + const h = create({ + fetcher: async (prompt) => + prompt === 'second' ? second.promise : 'first', + onResultChange, + }) + h.client.subscribe(() => { + if (h.resultId() !== 'first' || replaced) return + replaced = true + h.client.stop() + secondGenerate = h.client.generate({ prompt: 'second' }) + }) + + await h.client.generate({ prompt: 'first' }) + + expect(onResultChange).not.toHaveBeenCalledWith('first') + expect(h.client.getStatus()).toBe('generating') + second.resolve('second') + await secondGenerate + }) + + it('should not call the old error callback after its subscriber starts a replacement', async () => { + const second = createDeferred() + const onErrorChange = vi.fn() + let secondGenerate: Promise | undefined + let replaced = false + const h = create({ + fetcher: async (prompt) => { + if (prompt === 'first') throw new Error('first failed') + return second.promise + }, + onErrorChange, + }) + h.client.subscribe(() => { + if (h.client.getError()?.message !== 'first failed' || replaced) return + replaced = true + h.client.stop() + secondGenerate = h.client.generate({ prompt: 'second' }) + }) + + await h.client.generate({ prompt: 'first' }) + + expect(onErrorChange).not.toHaveBeenCalledWith(expect.any(Error)) + expect(h.client.getStatus()).toBe('generating') + second.resolve('second') + await secondGenerate + }) + + it('should not overwrite a run started by the reset loading callback', async () => { + const first = createDeferred() + const second = createDeferred() + let replaceOnStop = false + let secondGenerate: Promise | undefined + let secondRunId: string | null = null + let h!: RunHarness + h = create({ + fetcher: async (prompt) => { + if (prompt === 'seed') return 'seed' + return prompt === 'first' ? first.promise : second.promise + }, + onLoadingChange: (isLoading) => { + if (isLoading || !replaceOnStop) return + replaceOnStop = false + secondGenerate = h.client.generate({ prompt: 'second' }) + secondRunId = h.client.getSnapshot().runId + }, + }) + await h.client.generate({ prompt: 'seed' }) + const firstGenerate = h.client.generate({ prompt: 'first' }) + replaceOnStop = true + + h.client.reset() + + expect(secondRunId).not.toBeNull() + expect(h.client.getSnapshot().runId).toBe(secondRunId) + expect(h.client.getStatus()).toBe('generating') + expect(h.client.getIsLoading()).toBe(true) + expect(h.resultId()).toBe('seed') + first.resolve('first') + second.resolve('second') + await Promise.all([firstGenerate, secondGenerate]) + }) + + it('should not overwrite a run started by the dispose loading callback', async () => { + const first = createDeferred() + const second = createDeferred() + let replaceOnDispose = false + let secondGenerate: Promise | undefined + let secondRunId: string | null = null + let h!: RunHarness + h = create({ + fetcher: (prompt) => + prompt === 'first' ? first.promise : second.promise, + onLoadingChange: (isLoading) => { + if (isLoading || !replaceOnDispose) return + replaceOnDispose = false + h.client.mountDevtools() + secondGenerate = h.client.generate({ prompt: 'second' }) + secondRunId = h.client.getSnapshot().runId + }, + }) + const firstGenerate = h.client.generate({ prompt: 'first' }) + replaceOnDispose = true + + h.client.dispose() + + expect(secondRunId).not.toBeNull() + expect(h.client.getSnapshot().runId).toBe(secondRunId) + expect(h.client.getStatus()).toBe('generating') + expect(h.client.getIsLoading()).toBe(true) + first.resolve('first') + second.resolve('second') + await Promise.all([firstGenerate, secondGenerate]) + }) + + it('should not publish a run id after onChunk stops the run', async () => { + const connection: ConnectConnectionAdapter = { + async *connect() { + yield { + type: EventType.RUN_STARTED as const, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + } + }, + } + let h!: RunHarness + h = create({ + connection, + onChunk: (chunk) => { + if (chunk.type === EventType.RUN_STARTED) h.client.stop() + }, + }) + + await h.client.generate({ prompt: 'test' }) + + expect(h.client.getSnapshot().runId).toBeNull() + expect(h.client.getStatus()).toBe('idle') + }) + + it('should not let an old onChunk continuation overwrite a newer run', async () => { + let connectionCount = 0 + const secondBlocked = createDeferred() + const connection: ConnectConnectionAdapter = { + async *connect() { + connectionCount++ + if (connectionCount === 1) { + yield { + type: EventType.RUN_STARTED as const, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + } + return + } + await secondBlocked.promise + }, + } + let secondGenerate: Promise | undefined + let secondLocalRunId: string | null = null + let h!: RunHarness + h = create({ + connection, + onChunk: (chunk) => { + if (chunk.type !== EventType.RUN_STARTED) return + h.client.stop() + secondGenerate = h.client.generate({ prompt: 'second' }) + secondLocalRunId = h.client.getSnapshot().runId + }, + }) + + await h.client.generate({ prompt: 'first' }) + + expect(secondLocalRunId).not.toBeNull() + expect(h.client.getSnapshot().runId).toBe(secondLocalRunId) + expect(h.client.getIsLoading()).toBe(true) + h.client.stop() + secondBlocked.resolve() + await secondGenerate + }) +}) diff --git a/packages/ai-client/tests/generation-devtools.test.ts b/packages/ai-client/tests/generation-devtools.test.ts index 133519a42d..85d35f4863 100644 --- a/packages/ai-client/tests/generation-devtools.test.ts +++ b/packages/ai-client/tests/generation-devtools.test.ts @@ -7,6 +7,15 @@ import { VideoGenerationClient } from '../src/video-generation-client' import type { StreamChunk } from '@tanstack/ai' import type { ConnectConnectionAdapter } from '../src/connection-adapters' +interface StopRunOptions { + threadId: string + devtools: { hookName: string } + fetchText: () => Promise + onLoadingChange?: (isLoading: boolean) => void + onStatusChange?: (status: string) => void + onErrorChange?: (error: Error | undefined) => void +} + interface DevtoolsEvent { type: string payload: TPayload @@ -289,6 +298,121 @@ describe('generation client devtools bridge', () => { client.dispose() }) + describe.each([ + { + name: 'generation', + create: (options: StopRunOptions) => + new GenerationClient<{ prompt: string }, { text: string }>({ + ...options, + fetcher: async () => ({ text: await options.fetchText() }), + }), + successResult: { text: 'done' }, + }, + { + name: 'video', + create: (options: StopRunOptions) => + new VideoGenerationClient({ + ...options, + fetcher: async () => ({ + jobId: await options.fetchText(), + status: 'completed' as const, + url: 'https://example.com/video-1.mp4', + }), + }), + successResult: expect.objectContaining({ jobId: 'done' }), + }, + ])( + '$name devtools run stopped from a callback', + ({ name, create, successResult }) => { + it('finishes the run when a runId subscriber stops it', async () => { + const fetchText = vi.fn(async () => 'done') + const client = create({ + threadId: `${name}-run-id-stop`, + devtools: { hookName: 'useRunIdStop' }, + fetchText, + }) + const unsubscribe = client.subscribe(() => { + if (client.getSnapshot().runId) client.stop() + }) + + await client.generate({ prompt: 'stop early' }) + + expect(fetchText).not.toHaveBeenCalled() + expect(latestSnapshotState().activeRunId).toBeNull() + unsubscribe() + client.dispose() + }) + + it('finishes the run when loading callback stops it', async () => { + const fetchText = vi.fn(async () => 'done') + let client!: ReturnType + client = create({ + threadId: `${name}-loading-stop`, + devtools: { hookName: 'useLoadingStop' }, + fetchText, + onLoadingChange: (isLoading) => { + if (isLoading) client.stop() + }, + }) + + await client.generate({ prompt: 'stop while loading' }) + + expect(fetchText).not.toHaveBeenCalled() + expect(latestSnapshotState().activeRunId).toBeNull() + client.dispose() + }) + + it('finishes a completed run when success callback stops it', async () => { + let client!: ReturnType + client = create({ + threadId: `${name}-success-stop`, + devtools: { hookName: 'useSuccessStop' }, + fetchText: async () => 'done', + onStatusChange: (status) => { + if (status === 'success') client.stop() + }, + }) + + await client.generate({ prompt: 'finish then stop' }) + + expect(latestSnapshotState().activeRunId).toBeNull() + expect(latestGenerationRuns().at(-1)).toEqual( + expect.objectContaining({ status: 'success', result: successResult }), + ) + client.dispose() + }) + + it.each(['onStatusChange', 'onErrorChange'] as const)( + 'records an errored run when %s stops it', + async (callback) => { + let client!: ReturnType + const stopOnError = () => { + if (client.getError()) client.stop() + } + client = create({ + threadId: `${name}-error-stop-${callback}`, + devtools: { hookName: 'useErrorStop' }, + fetchText: async () => { + throw new Error(`${name} failed`) + }, + [callback]: stopOnError, + }) + + await client.generate({ prompt: 'fail then stop' }) + + expect(latestSnapshotState().activeRunId).toBeNull() + expect(latestGenerationRuns().at(-1)).toEqual( + expect.objectContaining({ + status: 'error', + error: `${name} failed`, + }), + ) + client.dispose() + }, + ) + }, + ) + it('includes input, progress, and renderable previews in generation snapshots', async () => { const client = new GenerationClient({ threadId: 'image-hook', diff --git a/packages/ai-client/tests/snapshot-atom.test.ts b/packages/ai-client/tests/snapshot-atom.test.ts new file mode 100644 index 0000000000..b915b7dd77 --- /dev/null +++ b/packages/ai-client/tests/snapshot-atom.test.ts @@ -0,0 +1,133 @@ +import { describe, expect, it, vi } from 'vitest' +import { + cloneSnapshotValue, + createAtom, + freezeSnapshotMessages, + patchAtom, + subscribeAtom, +} from '../src/snapshot-atom' + +describe('snapshot-atom', () => { + it('patchAtom freezes the patched snapshot and keeps identity without a change', () => { + const atom = createAtom({ count: 1 }) + patchAtom(atom, { count: 2 }) + const snapshot = atom.get() + + expect(Object.isFrozen(snapshot)).toBe(true) + patchAtom(atom, { count: 2 }) + expect(atom.get()).toBe(snapshot) + }) + + it('patchAtom skips notify when every field is unchanged', () => { + const atom = createAtom({ count: 1, label: 'a' }) + const listener = vi.fn() + const stop = subscribeAtom(atom, listener) + + patchAtom(atom, { count: 1, label: 'a' }) + expect(listener).not.toHaveBeenCalled() + expect(atom.get()).toEqual({ count: 1, label: 'a' }) + + patchAtom(atom, { count: 2 }) + expect(listener).toHaveBeenCalledTimes(1) + expect(atom.get()).toEqual({ count: 2, label: 'a' }) + stop() + }) + + it('subscribeAtom adapts { unsubscribe } to a plain cleanup', () => { + const atom = createAtom(0) + const listener = vi.fn() + const stop = subscribeAtom(atom, listener) + + atom.set(1) + expect(listener).toHaveBeenCalledTimes(1) + stop() + atom.set(2) + expect(listener).toHaveBeenCalledTimes(1) + }) + + it('isolates a throwing listener from the next subscriber', () => { + const atom = createAtom(0) + const second = vi.fn() + const stopFirst = subscribeAtom(atom, () => { + throw new Error('listener failed') + }) + const stopSecond = subscribeAtom(atom, second) + const error = vi.spyOn(console, 'error').mockImplementation(() => {}) + + atom.set(1) + + expect(second).toHaveBeenCalledTimes(1) + expect(error).toHaveBeenCalled() + stopFirst() + stopSecond() + error.mockRestore() + }) + + it('subscribeAtom calls the listener again for a change made inside it', () => { + const atom = createAtom({ isLoading: true, status: 'generating' }) + const seen: Array = [] + const stop = subscribeAtom(atom, () => { + seen.push(atom.get().isLoading) + // A Solid effect or Vue watcher that calls `stop()` from here. + if (atom.get().status === 'idle' && atom.get().isLoading) { + patchAtom(atom, { isLoading: false }) + } + }) + + patchAtom(atom, { status: 'idle' }) + expect(seen).toEqual([true, false]) + stop() + }) + + it('stops calling a listener that changes the atom on every call', () => { + const atom = createAtom({ count: 0 }) + let calls = 0 + const stop = subscribeAtom(atom, () => { + calls++ + patchAtom(atom, { count: atom.get().count + 1 }) + }) + + patchAtom(atom, { count: 1 }) + + expect(calls).toBe(8) + stop() + }) + + it('reuses the frozen copy of an unchanged message and part', () => { + const cache = new WeakMap() + const text = { type: 'text', content: 'hi' } + const tool = { type: 'tool-call', input: { q: 'a' } } + const first = { id: 'm1', parts: [text, tool] } + + const [frozenFirst] = freezeSnapshotMessages(cache, [first]) + // The processor replaces the message and the changed part only. + const second = { ...first, parts: [{ ...text, content: 'hi!' }, tool] } + const [frozenSecond] = freezeSnapshotMessages(cache, [second]) + + expect(frozenSecond).not.toBe(frozenFirst) + expect(frozenSecond?.parts[1]).toBe(frozenFirst?.parts[1]) + expect(frozenSecond?.parts[0]).not.toBe(frozenFirst?.parts[0]) + expect(freezeSnapshotMessages(cache, [second])[0]).toBe(frozenSecond) + }) + + it('cloneSnapshotValue copies plain values and keeps other objects', () => { + class Result { + read() { + return 'ok' + } + } + const blob = new Blob(['x']) + const map = new Map([['a', 1]]) + const instance = new Result() + + expect(cloneSnapshotValue(blob)).toBe(blob) + expect(cloneSnapshotValue(map)).toBe(map) + expect(cloneSnapshotValue(instance).read()).toBe('ok') + + const plain = { url: 'a' } + const copy = cloneSnapshotValue(plain) + expect(copy).toEqual(plain) + expect(copy).not.toBe(plain) + expect(Object.isFrozen(copy)).toBe(true) + }) +}) diff --git a/packages/ai-client/tests/video-generation-client.test.ts b/packages/ai-client/tests/video-generation-client.test.ts index be788e4d57..bdb9a356ce 100644 --- a/packages/ai-client/tests/video-generation-client.test.ts +++ b/packages/ai-client/tests/video-generation-client.test.ts @@ -224,6 +224,7 @@ describe('VideoGenerationClient', () => { ) expect(client.getStatus()).toBe('success') expect(client.getJobId()).toBe('job-123') + expect(client.getSnapshot().runId).toBeNull() }) it('should track video status updates', async () => { @@ -314,6 +315,7 @@ describe('VideoGenerationClient', () => { expect(onError).toHaveBeenCalledWith(expect.any(Error)) expect(client.getStatus()).toBe('error') expect(client.getError()?.message).toBe('Video generation failed') + expect(client.getSnapshot().runId).toBeNull() }) it('should report progress from video:status events', async () => { @@ -483,10 +485,12 @@ describe('VideoGenerationClient', () => { const generatePromise = client.generate({ prompt: 'test' }) expect(client.getIsLoading()).toBe(true) + expect(client.getSnapshot().runId).toEqual(expect.any(String)) client.stop() expect(client.getIsLoading()).toBe(false) expect(client.getStatus()).toBe('idle') + expect(client.getSnapshot().runId).toBeNull() resolvePromise!({ jobId: 'job-1', @@ -561,6 +565,7 @@ describe('VideoGenerationClient', () => { expect(client.getVideoStatus()).toBeNull() expect(client.getError()).toBeUndefined() expect(client.getStatus()).toBe('idle') + expect(client.getSnapshot().runId).toBeNull() }) }) diff --git a/packages/ai-compaction/src/index.test.ts b/packages/ai-compaction/src/index.test.ts index 3c103fb763..f4cd633858 100644 --- a/packages/ai-compaction/src/index.test.ts +++ b/packages/ai-compaction/src/index.test.ts @@ -3243,6 +3243,40 @@ describe('withCompaction native', () => { expect(other?.providerMessages?.[0]?.content).toBe('gpt-5.4') }) + it('does not keep a failed native compact as the checkpoint', async () => { + const store = memoryStore() + const strategy = vi.fn(() => [text('user', 'cut')]) + const messages = [big('user'), big('assistant'), big('user')] + const context = () => { + const ctx = checkpointContext(store) + ctx.model = 'gpt-5.5' + return ctx + } + const failed = await runOnConfig( + withCompaction({ + maxTokens: 50, + strategy, + native: { compact: () => Promise.reject(new Error('404 Not Found')) }, + }), + messages, + context(), + ) + + expect(failed).toStrictEqual({ providerMessages: [text('user', 'cut')] }) + + const compact = vi.fn(() => Promise.resolve([text('user', 'sealed')])) + const retried = await runOnConfig( + withCompaction({ maxTokens: 50, strategy, native: { compact } }), + messages, + context(), + ) + + expect(compact).toHaveBeenCalledTimes(1) + expect(retried).toStrictEqual({ + providerMessages: [text('user', 'sealed')], + }) + }) + it('passes the system prompts to the adapter compact', async () => { const compact = vi.fn((_options: TextCompactOptions) => Promise.resolve([text('user', 'kept')]), diff --git a/packages/ai-compaction/src/index.ts b/packages/ai-compaction/src/index.ts index 7633839581..690b69fa7f 100644 --- a/packages/ai-compaction/src/index.ts +++ b/packages/ai-compaction/src/index.ts @@ -1153,7 +1153,9 @@ export function withCompaction( await store?.delete(USAGE_NAMESPACE, ctx.threadId) } - if (logRecords) { + // A failed compact must not be stored. This call still uses the + // fallback. The next process tries compact again. + if (!error && logRecords) { // The model gets `next` now. The record makes the log fold the same. if (!ctx.signal?.aborted) { await logRecords.append([ @@ -1167,7 +1169,7 @@ export function withCompaction( }), ]) } - } else if (metadata && inputMessages === messages) { + } else if (!error && metadata && inputMessages === messages) { const checkpoint: CompactionCheckpoint = { schemaVersion: 1, sourceMessageCount: messages.length, diff --git a/packages/ai-gemini/src/adapters/text.ts b/packages/ai-gemini/src/adapters/text.ts index 8286eb1624..16499b1845 100644 --- a/packages/ai-gemini/src/adapters/text.ts +++ b/packages/ai-gemini/src/adapters/text.ts @@ -305,7 +305,12 @@ function buildInteractionsInput( const signature = metadata?.thoughtSignature if (typeof signature === 'string' && signature) steps.push({ type: 'thought', signature, summary: [] }) - const args: unknown = JSON.parse(toolCall.function.arguments || '{}') + // Malformed JSON stays in the history next to its tool error, so it + // goes out as `{}`. + let args: unknown = {} + try { + args = JSON.parse(toolCall.function.arguments || '{}') + } catch {} if (typeof args !== 'object' || args === null || Array.isArray(args)) throw new Error( `Gemini Interactions function ${toolCall.function.name} requires object arguments`, @@ -1158,12 +1163,14 @@ export class GeminiTextAdapter< // this chunk (including UNEXPECTED_TOOL_CALL finishes) were already // registered and started by the per-part loop above. for (const [toolCallId, toolCallData] of toolCallMap.entries()) { - let parsedInput: unknown = {} + // Malformed JSON sends no `input`, so core reports the error from + // the raw `args`. + let parsedInput: unknown try { const parsed = JSON.parse(toolCallData.args) parsedInput = parsed && typeof parsed === 'object' ? parsed : {} } catch { - parsedInput = {} + // Keep `parsedInput` undefined for the tool error. } yield { @@ -1173,7 +1180,8 @@ export class GeminiTextAdapter< toolName: toolCallData.name, model, timestamp: Date.now(), - input: parsedInput, + args: toolCallData.args, + ...(parsedInput !== undefined && { input: parsedInput }), } } diff --git a/packages/ai-gemini/src/experimental/text-interactions/adapter.ts b/packages/ai-gemini/src/experimental/text-interactions/adapter.ts index 40dc9a26ce..4da75e0593 100644 --- a/packages/ai-gemini/src/experimental/text-interactions/adapter.ts +++ b/packages/ai-gemini/src/experimental/text-interactions/adapter.ts @@ -920,6 +920,22 @@ function parsePartialToolArguments( } } +// The raw `args` and parsed `input` of a TOOL_CALL_END. The streamed argument +// text wins over the initial `step.arguments` snapshot. Malformed JSON sends +// no `input`, so core reports the error from the raw `args`. +function toolCallEndArguments( + streamed: string | undefined, + state: ToolCallState, +): { args: string; input?: unknown } { + const args = streamed?.trim() ? streamed : JSON.stringify(state.args) + try { + const parsed: unknown = JSON.parse(args) + return { args, input: parsed && typeof parsed === 'object' ? parsed : {} } + } catch { + return { args } + } +} + // `satisfies` pins these arrays to the SDK's narrow mime-type unions: if // Google removes a format the build breaks, and if they add one ours keeps // working (we just won't accept the new one until added here). @@ -1235,6 +1251,37 @@ async function* translateInteractionEvents( } } + // Ends a tool call. With no streamed argument text, the initial + // `step.arguments` snapshot goes out as the only TOOL_CALL_ARGS delta, so + // the deltas always add up to the final arguments. + const endToolCall = function* ( + toolCallId: string, + state: ToolCallState, + ): Generator { + const end = toolCallEndArguments( + argStringByToolCallId.get(toolCallId), + state, + ) + if (!argStringByToolCallId.get(toolCallId)?.trim() && end.args !== '{}') { + yield { + type: EventType.TOOL_CALL_ARGS, + toolCallId, + model, + timestamp, + delta: end.args, + args: end.args, + } + } + yield { + type: EventType.TOOL_CALL_END, + toolCallId, + toolName: state.name, + model, + timestamp, + ...end, + } + } + // Seals any in-flight messages and tool calls. Called both on the // normal terminal path (`interaction.complete`) and on the error path // (`error` SSE event + premature EOF) so the StreamProcessor never @@ -1245,14 +1292,7 @@ async function* translateInteractionEvents( for (const [toolCallId, state] of toolCalls) { if (state.ended) continue state.ended = true - yield { - type: EventType.TOOL_CALL_END, - toolCallId, - toolName: state.name, - model, - timestamp, - input: state.args, - } + yield* endToolCall(toolCallId, state) } if (hasEmittedTextMessageStart) { hasEmittedTextMessageStart = false @@ -1307,8 +1347,9 @@ async function* translateInteractionEvents( indexToToolCallId.set(index, toolCallId) // `step.arguments` is required on FunctionCallStep but may // be an empty `{}` placeholder when streaming, where the - // real args arrive as `arguments_delta` events. Treat both - // uniformly: stash whatever we got, stringify once. + // real args arrive as `arguments_delta` events. Keep this + // initial snapshot apart from the streamed text: it is sent + // only when no argument delta arrives (see `endToolCall`). const initialArgs = step.arguments const state: ToolCallState = { name: step.name, @@ -1318,12 +1359,6 @@ async function* translateInteractionEvents( ended: false, } toolCalls.set(toolCallId, state) - argStringByToolCallId.set( - toolCallId, - Object.keys(initialArgs).length > 0 - ? JSON.stringify(initialArgs) - : '', - ) yield { type: EventType.TOOL_CALL_START, toolCallId, @@ -1337,17 +1372,6 @@ async function* translateInteractionEvents( timestamp, index: state.index, } - if (Object.keys(initialArgs).length > 0) { - const argsJson = JSON.stringify(initialArgs) - yield { - type: EventType.TOOL_CALL_ARGS, - toolCallId, - model, - timestamp, - delta: argsJson, - args: argsJson, - } - } break } case 'thought': { @@ -1579,8 +1603,9 @@ async function* translateInteractionEvents( break } const state = toolCalls.get(toolCallId) - if (!state) break - const fragment = delta.arguments ?? '' + // A delta with no arguments keeps the initial snapshot. + if (!state || typeof delta.arguments !== 'string') break + const fragment = delta.arguments const buffer = (argStringByToolCallId.get(toolCallId) ?? '') + fragment argStringByToolCallId.set(toolCallId, buffer) @@ -1693,14 +1718,7 @@ async function* translateInteractionEvents( const state = toolCalls.get(toolCallId) if (state && !state.ended) { state.ended = true - yield { - type: EventType.TOOL_CALL_END, - toolCallId, - toolName: state.name, - model, - timestamp, - input: state.args, - } + yield* endToolCall(toolCallId, state) } indexToToolCallId.delete(event.index) } diff --git a/packages/ai-gemini/tests/gemini-adapter.test.ts b/packages/ai-gemini/tests/gemini-adapter.test.ts index 86fee86930..f3c4dc586a 100644 --- a/packages/ai-gemini/tests/gemini-adapter.test.ts +++ b/packages/ai-gemini/tests/gemini-adapter.test.ts @@ -570,6 +570,83 @@ describe('GeminiAdapter through AI', () => { expect(endIds.sort()).toEqual(startIds.sort()) }) + it.each([ + { args: { count: 2 }, raw: '{"count":2}', input: { count: 2 } }, + { args: '{"n": 1.50}', raw: '{"n": 1.50}', input: { n: 1.5 } }, + ])( + 'sends the raw function arguments $raw on TOOL_CALL_END', + async ({ args, raw, input }) => { + mocks.generateContentStreamSpy.mockResolvedValue( + createStream([ + { + candidates: [ + { + content: { + parts: [ + { functionCall: { id: 'call_raw', name: 'inspect', args } }, + ], + }, + finishReason: 'STOP', + }, + ], + }, + ]), + ) + + const chunks: Array = [] + for await (const chunk of createTextAdapter().chatStream({ + model: 'gemini-2.5-pro', + logger: testLogger, + messages: [{ role: 'user', content: 'Go' }], + })) { + chunks.push(chunk) + } + + expect(chunks.find((c) => c.type === 'TOOL_CALL_END')).toMatchObject({ + args: raw, + input, + }) + }, + ) + + it('sends malformed function arguments raw on TOOL_CALL_END with no input', async () => { + mocks.generateContentStreamSpy.mockResolvedValue( + createStream([ + { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call_bad', + name: 'inspect', + args: '{"a":', + }, + }, + ], + }, + finishReason: 'STOP', + }, + ], + }, + ]), + ) + + const chunks: Array = [] + for await (const chunk of createTextAdapter().chatStream({ + model: 'gemini-2.5-pro', + logger: testLogger, + messages: [{ role: 'user', content: 'Go' }], + })) { + chunks.push(chunk) + } + + const end = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(end).toMatchObject({ args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + it('emits grounded web search sources and omits the provider call on the next turn', async () => { const groundingMetadata = { webSearchQueries: ['latest release'], diff --git a/packages/ai-gemini/tests/replay-parity.test.ts b/packages/ai-gemini/tests/replay-parity.test.ts index f4ba64fe63..201e2ade36 100644 --- a/packages/ai-gemini/tests/replay-parity.test.ts +++ b/packages/ai-gemini/tests/replay-parity.test.ts @@ -723,6 +723,45 @@ describe('Gemini replay parity', () => { ).toBe('interaction-real') }) + it('replays malformed tool arguments to the agentic API as {}', async () => { + mocks.interactions.mockResolvedValue({ id: 'i', output_text: 'ok' }) + await collect( + chat({ + adapter: new GeminiTextAdapter({ apiKey: 'key' }, 'gemini-3.7-flash'), + messages: [ + { + role: 'assistant', + content: null, + toolCalls: [ + { + id: 'call-1', + type: 'function', + function: { name: 'check', arguments: '{"a":' }, + }, + ], + }, + { role: 'tool', toolCallId: 'call-1', content: 'bad arguments' }, + { + role: 'user', + content: [ + { + type: 'video', + source: { type: 'url', value: 'u', mimeType: 'video/mp4' }, + metadata: { processing: 'agentic' }, + }, + ], + }, + ], + }), + ) + expect(mocks.interactions.mock.calls[0]?.[0].input).toContainEqual({ + type: 'function_call', + id: 'call-1', + name: 'check', + arguments: {}, + }) + }) + it('keeps API source independent across concurrent calls on one adapter', async () => { mocks.interactions.mockResolvedValue({ id: 'interaction-concurrent', diff --git a/packages/ai-gemini/tests/text-interactions-adapter.test.ts b/packages/ai-gemini/tests/text-interactions-adapter.test.ts index 30dffa518c..0bff21caed 100644 --- a/packages/ai-gemini/tests/text-interactions-adapter.test.ts +++ b/packages/ai-gemini/tests/text-interactions-adapter.test.ts @@ -563,11 +563,9 @@ describe('GeminiTextInteractionsAdapter', () => { expect(endEvent.input).toEqual({ location: 'Berlin', unit: 'celsius' }) }) - it('preserves the last good args when the arguments stream truncates mid-fragment', async () => { - // If the stream ends after an incomplete fragment, the tool input must - // retain whatever was parsed so far rather than being reset to `{}` — a - // strict JSON.parse would throw on the partial buffer and lose the - // already-merged keys. + it('keeps truncated arguments raw on TOOL_CALL_END with no input', async () => { + // A partial value can appear while the arguments stream. It is not the + // final input: core reports the malformed raw arguments as a tool error. mocks.interactionsCreateSpy.mockResolvedValue( mkStream([ { @@ -602,21 +600,104 @@ describe('GeminiTextInteractionsAdapter', () => { ]), ) - const adapter = createAdapter() const chunks = await collectChunks( - chat({ - adapter, + createAdapter().chatStream({ + model: 'gemini-2.5-flash', + logger: testLogger, messages: [{ role: 'user', content: 'Weather in London?' }], tools: [{ name: 'lookup_weather', description: 'Return the weather' }], }), ) - // No parse-failure RUN_ERROR, and the completed key survives truncation. expect(chunks.find((c) => c.type === 'RUN_ERROR')).toBeUndefined() - const endEvent = chunks.find((c) => c.type === 'TOOL_CALL_END') as any - expect(endEvent.input).toEqual({ city: 'London' }) + const endEvent = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(endEvent).toMatchObject({ args: '{"city":"London","unit":' }) + expect(endEvent).not.toHaveProperty('input') }) + const rawArgumentCases: Array<{ + name: string + deltas: Array + args: string + input: Record + }> = [ + { + name: 'the initial snapshot when no argument delta arrives', + deltas: [], + args: '{"initial":1}', + input: { initial: 1 }, + }, + { + name: 'the streamed text over the initial snapshot', + deltas: ['{"n": ', '1.50}'], + args: '{"n": 1.50}', + input: { n: 1.5 }, + }, + { + name: 'the initial snapshot when a delta has no arguments', + deltas: [undefined], + args: '{"initial":1}', + input: { initial: 1 }, + }, + ] + + it.each(rawArgumentCases)( + 'sends $name as the raw TOOL_CALL_END args', + async ({ deltas, args, input }) => { + mocks.interactionsCreateSpy.mockResolvedValue( + mkStream([ + { + event_type: 'interaction.created', + interaction: { id: 'int_raw', status: 'in_progress' }, + }, + { + event_type: 'step.start', + index: 0, + step: { + type: 'function_call', + id: 'call_raw', + name: 'inspect', + arguments: { initial: 1 }, + }, + }, + ...deltas.map((fragment) => ({ + event_type: 'step.delta', + index: 0, + delta: { + type: 'arguments_delta', + ...(fragment !== undefined && { arguments: fragment }), + }, + })), + { event_type: 'step.stop', index: 0 }, + { + event_type: 'interaction.completed', + interaction: { id: 'int_raw', status: 'completed' }, + }, + ]), + ) + + const chunks = await collectChunks( + createAdapter().chatStream({ + model: 'gemini-2.5-flash', + logger: testLogger, + messages: [{ role: 'user', content: 'Go' }], + tools: [{ name: 'inspect', description: 'Inspect' }], + }), + ) + + expect(chunks.find((c) => c.type === 'TOOL_CALL_END')).toMatchObject({ + args, + input, + }) + // A client that only reads the deltas gets the same arguments. + const streamed = chunks + .filter((c) => c.type === 'TOOL_CALL_ARGS') + .map((c) => ('delta' in c ? c.delta : '')) + .join('') + expect(streamed).toBe(args) + }, + ) + it('translates thought_summary deltas into REASONING events', async () => { mocks.interactionsCreateSpy.mockResolvedValue( mkStream([ diff --git a/packages/ai-groq/src/adapters/text.ts b/packages/ai-groq/src/adapters/text.ts index 0b7f778a52..5427ff752f 100644 --- a/packages/ai-groq/src/adapters/text.ts +++ b/packages/ai-groq/src/adapters/text.ts @@ -3,6 +3,7 @@ import { OpenAIBaseChatCompletionsTextAdapter } from '@tanstack/openai-base' import { GROQ_MODEL_INPUT_MODALITIES, GROQ_MODEL_REASONING, + groqModelPrefersJsonObject, } from '../model-meta' import { getGroqApiKeyFromEnv, withGroqDefaults } from '../utils/client' import { makeGroqStructuredOutputCompatibleWithMap } from '../utils/schema-converter' @@ -196,6 +197,10 @@ export class GroqTextAdapter< override supportsCombinedToolsAndSchema(): boolean { return false } + + protected override modelPrefersJsonObject(model: string): boolean { + return groqModelPrefersJsonObject(model) + } } function isRecord(value: unknown): value is Record { diff --git a/packages/ai-groq/src/model-meta.ts b/packages/ai-groq/src/model-meta.ts index 0660304b8a..88f9d5da95 100644 --- a/packages/ai-groq/src/model-meta.ts +++ b/packages/ai-groq/src/model-meta.ts @@ -336,6 +336,32 @@ const QWEN3_32B = { reasoning: EFFORT_NONE_DEFAULT, } as const satisfies ModelMeta +const GROQ_CHAT_MODEL_META = [ + LLAMA_3_1_8B_INSTANT, + LLAMA_3_3_70B_VERSATILE, + LLAMA_4_MAVERICK_17B_128E_INSTRUCT, + LLAMA_4_SCOUT_17B_16E_INSTRUCT, + LLAMA_GUARD_4_12B, + LLAMA_PROMPT_GUARD_2_86M, + LLAMA_PROMPT_GUARD_2_22M, + GPT_OSS_20B, + GPT_OSS_120B, + GPT_OSS_SAFEGUARD_20B, + KIMI_K2_INSTRUCT_0905, + QWEN3_32B, +] as const + +/** + * `true` when structured output must send JSON mode (`json_object`). + * The model declares `json_object` and does not declare `json_schema`. + */ +export function groqModelPrefersJsonObject(model: string): boolean { + const found = GROQ_CHAT_MODEL_META.find((item) => item.name === model) + if (!found) return false + const features: ReadonlyArray = found.supports.features + return features.includes('json_object') && !features.includes('json_schema') +} + /** * All supported Groq chat model identifiers. */ diff --git a/packages/ai-groq/tests/groq-adapter.test.ts b/packages/ai-groq/tests/groq-adapter.test.ts index ee7cb9737d..506a25c44d 100644 --- a/packages/ai-groq/tests/groq-adapter.test.ts +++ b/packages/ai-groq/tests/groq-adapter.test.ts @@ -143,6 +143,64 @@ describe('Groq adapters', () => { expect(adapter).toBeDefined() }) + it('sends json_object on a json_object-only model', async () => { + const mockCreate = setupMockSdkClient([], { + choices: [{ message: { content: '{"name":"Ada"}' } }], + }) + const adapter = createGroqText('llama-3.3-70b-versatile', 'test-api-key') + + await adapter.structuredOutput({ + chatOptions: { + model: 'llama-3.3-70b-versatile', + messages: [{ role: 'user', content: 'Return a person as json' }], + logger: testLogger, + }, + outputSchema: { + type: 'object', + properties: { name: { type: 'string' } }, + required: ['name'], + }, + }) + + expect(mockCreate.mock.calls[0]?.[0]).toMatchObject({ + response_format: { type: 'json_object' }, + messages: [ + { + role: 'system', + content: expect.stringContaining('"required":["name"]'), + }, + { role: 'user', content: 'Return a person as json' }, + ], + }) + }) + + it('keeps json_schema on a model that declares json_schema', async () => { + const mockCreate = setupMockSdkClient([], { + choices: [{ message: { content: '{"name":"Ada"}' } }], + }) + const adapter = createGroqText( + 'moonshotai/kimi-k2-instruct-0905', + 'test-api-key', + ) + + await adapter.structuredOutput({ + chatOptions: { + model: 'moonshotai/kimi-k2-instruct-0905', + messages: [{ role: 'user', content: 'Return a person as json' }], + logger: testLogger, + }, + outputSchema: { + type: 'object', + properties: { name: { type: 'string' } }, + required: ['name'], + }, + }) + + expect(mockCreate.mock.calls[0]?.[0]).toMatchObject({ + response_format: { type: 'json_schema' }, + }) + }) + it('opts out of native combined tools+schema mode (#605) — Groq API rejects the combination', () => { const adapter = createGroqText('llama-3.3-70b-versatile', 'test-api-key') // The OpenAI Chat Completions base defaults to `true`; Groq must diff --git a/packages/ai-harness/src/first-party/index.ts b/packages/ai-harness/src/first-party/index.ts index e0a09d0ada..0370211b18 100644 --- a/packages/ai-harness/src/first-party/index.ts +++ b/packages/ai-harness/src/first-party/index.ts @@ -24,6 +24,12 @@ export { formatTodos, todos } from './todos' export type { Todo } from './todos' export { fileCommands, projectInstructions } from './files' export { compact, usage } from './session-tools' +export { providerKeys } from './provider-keys' +export type { + KeyProvider, + ProviderKeyStatus, + ProviderKeysState, +} from './provider-keys' export { GoalMet, goal, selectGoal } from './goal' export type { Goal } from './goal' export { agents, builtInAgents } from './agents' diff --git a/packages/ai-harness/src/first-party/provider-keys.ts b/packages/ai-harness/src/first-party/provider-keys.ts new file mode 100644 index 0000000000..3dae90d270 --- /dev/null +++ b/packages/ai-harness/src/first-party/provider-keys.ts @@ -0,0 +1,193 @@ +import { maskKey } from '@tanstack/ai/byok' +import { envKeyOf } from '../auth' +import { defineCommand } from '../commands' +import { definePlugin } from '../plugins' +import type { ByokProvider } from '@tanstack/ai/byok' +import type { CredentialsAccess } from '../auth' +import type { AnyCommand, CommandContext } from '../commands' + +/** Signs the user in and resolves to an API key, for example `openrouterSignIn()`. */ +type ProviderSignIn = (ctx: { + open: (url: string) => void + signal?: AbortSignal +}) => Promise + +/** + * A provider the user can connect: its BYOK descriptor, with a sign-in when + * it has one. `keyUrl` is the provider's page to make a key: `/connect` opens + * it in the browser before it asks for the key. + */ +export type KeyProvider = ByokProvider & { + signIn?: ProviderSignIn + keyUrl?: string +} + +export interface ProviderKeyStatus { + id: string + label: string + /** `connected`: a saved key. `env`: the env var. `missing`: neither. */ + state: 'connected' | 'env' | 'missing' +} + +/** What UIs read from `state.plugins['tanstack/provider-keys']`. */ +export interface ProviderKeysState { + providers: Array +} + +/** + * `promise`, or a rejection once `signal` aborts. A cancelled `/connect` + * stops waiting for the answer, and a late answer saves nothing. + */ +function untilAborted(promise: Promise, signal: AbortSignal) { + signal.throwIfAborted() + return Promise.race([ + promise, + new Promise((_, reject) => + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }), + ), + ]) +} + +/** The state of one provider, and its `/keys` line. The key shows masked only. */ +async function describeKey( + credentials: CredentialsAccess, + provider: ByokProvider, +) { + const { id, label } = provider + const status = (state: ProviderKeyStatus['state']) => ({ id, label, state }) + const saved = await credentials.get(id) + if (saved?.type === 'api_key') { + return { + status: status('connected'), + line: `${label}: connected (key ...${maskKey(saved.value)})`, + } + } + const env = envKeyOf(provider) + if (env) { + return { + status: status('env'), + line: `${label}: from the ${env.name} env var`, + } + } + return { + status: status('missing'), + line: `${label}: missing. Run /connect ${id}.`, + } +} + +/** + * Let each user connect their own model providers inside the app, so a + * shipped harness needs no `.env` file. For each provider it adds + * `/connect ` and `/disconnect `, and `/keys` lists where each key + * comes from. `/connect` runs the provider's `signIn` when it has one (the + * host opens the browser), else it asks the user to paste the key. The key + * is saved in the credential store of the session's principal. A + * `keyedAdapter(...)` for that provider then uses it. The env var still works + * when no key is saved. Keys show masked only (the last 4 characters). + * + * @param options.providers - BYOK descriptors, for example `openaiByok`. Add + * `signIn` for a provider with a browser sign-in, for example + * `{ ...openrouterByok, signIn: openrouterSignIn() }`, or `keyUrl` to open + * the provider's key page before the question, for example + * `{ ...openaiByok, keyUrl: 'https://platform.openai.com/api-keys' }`. + * + * @example + * ```ts + * plugins: () => [ + * providerKeys({ + * providers: [ + * openaiByok, + * anthropicByok, + * { ...openrouterByok, signIn: openrouterSignIn() }, + * ], + * }), + * ] + * ``` + */ +export function providerKeys(options: { + providers: ReadonlyArray +}) { + const { providers } = options + return definePlugin({ + name: 'tanstack/provider-keys', + setup: async (ctx) => { + const state = ctx.state({ providers: [] }) + /** + * Check every provider, keep the result for UIs, and return the `/keys` + * lines. A command passes the credentials of the user who runs it. + */ + const refresh = async (credentials = ctx.credentials) => { + const described = await Promise.all( + providers.map((provider) => describeKey(credentials, provider)), + ) + await state.update(() => ({ + providers: described.map((entry) => entry.status), + })) + return described.map((entry) => entry.line).join('\n') + } + await refresh() + + /** + * The key from the provider's sign-in, else the one the user pastes. + * It asks the user who runs the command. + */ + const readKey = async ( + provider: KeyProvider, + { signal, session }: CommandContext, + ) => { + const open = (url: string) => + session.authRequired({ connector: provider.id, url }) + if (provider.signIn) return provider.signIn({ open, signal }) + // No sign-in for other apps: open the page to make a key, then ask. + if (provider.keyUrl) open(provider.keyUrl) + return session.ask({ + message: provider.keyUrl + ? `Paste your ${provider.label} API key. The page to make one is open in your browser.` + : `Paste your ${provider.label} API key`, + secret: true, + }) + } + + const commands: Record = { + keys: defineCommand({ + description: + 'Show the model providers and where their keys come from', + run: (_input, { credentials }) => refresh(credentials), + }), + } + for (const provider of providers) { + const { id, label } = provider + commands[`connect:${id}`] = defineCommand({ + description: `Connect ${label} with your own key`, + // The key belongs to the user who runs the command. + run: async (_input, command) => { + const { signal, credentials } = command + const key = ( + await untilAborted(readKey(provider, command), signal) + ).trim() + if (key === '') { + throw new Error(`No ${label} key was given. Nothing was saved.`) + } + await credentials.set({ + id, + credential: { type: 'api_key', value: key }, + }) + await refresh(credentials) + return `Connected to ${label} (key ...${maskKey(key)}).` + }, + }) + commands[`disconnect:${id}`] = defineCommand({ + description: `Remove your ${label} key`, + run: async (_input, { credentials }) => { + await credentials.delete(id) + await refresh(credentials) + return `Disconnected from ${label}.` + }, + }) + } + return { commands } + }, + }) +} diff --git a/packages/ai-harness/tests/provider-keys.test.ts b/packages/ai-harness/tests/provider-keys.test.ts new file mode 100644 index 0000000000..2809228fce --- /dev/null +++ b/packages/ai-harness/tests/provider-keys.test.ts @@ -0,0 +1,466 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { EventType, defineAgent } from '@tanstack/ai' +import { defineByokProvider, keyedAdapter } from '@tanstack/ai/byok' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { HARNESS_EVENTS, createHarnessHost, defineHarness } from '../src' +import { compact, modelPicker, providerKeys } from '../src/first-party' +import { mockAdapter, text, toolCall } from './helpers' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { + AnyAgent, + HarnessPlugin, + HarnessSession, + SessionEvent, +} from '../src' + +const acme = defineByokProvider({ + id: 'acme', + label: 'Acme', + env: 'ACME_TEST_KEY', +}) +const beta = defineByokProvider({ + id: 'beta', + label: 'Beta', + env: 'BETA_TEST_KEY', +}) +const gamma = defineByokProvider({ id: 'gamma', label: 'Gamma' }) + +/** The credential scope of the test user. */ +const USER = { threadId: 't', userId: 'user-1' } +const SAVED_KEY = 'sk-acme-saved-1234' +const PASTED_KEY = 'sk-acme-pasted-5678' + +afterEach(() => { + vi.unstubAllEnvs() +}) + +/** A keyed adapter for `acme` whose turns are scripted. Records each key it got. */ +function keyedMock(replies: Parameters[0]) { + const mock = mockAdapter(replies) + const keys: Array = [] + const adapter = keyedAdapter(acme, (key) => { + keys.push(key) + return mock.adapter + }) + return { adapter, keys, calls: mock.calls } +} + +async function open(options: { + adapter: AnyTextAdapter | ReturnType['adapter'] + plugins?: Array + agents?: Array + subagents?: Array +}) { + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/provider-keys', + adapter: options.adapter, + agents: options.agents ?? [], + subagents: { agents: options.subagents ?? [] }, + plugins: () => options.plugins ?? [], + }), + { threadId: 't', principal: { id: 'user-1' } }, + ) + const events: Array = [] + const reader = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ signal: reader.signal })) + events.push(entry) + })() + const saveKey = (value: string) => + persistence.stores.credentials.set({ + scope: USER, + id: 'acme', + credential: { + type: 'api_key', + value, + }, + }) + const close = async () => { + reader.abort() + await reading + await host.close() + } + return { session, persistence, events, saveKey, close } +} + +/** The connectors named by the `harness.auth_required` events. */ +function authRequired(events: ReadonlyArray) { + return events.flatMap(({ event }) => + event.type === EventType.CUSTOM && + event.name === HARNESS_EVENTS.authRequired + ? [event.value] + : [], + ) +} + +/** Start `/connect acme`, and answer its question with `answer`. */ +async function connectWith(session: HarnessSession, answer: string) { + const result = session.command('connect:acme') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [question] = session.snapshot().pendingQuestions + const receipt = await session.answer({ + questionId: question?.questionId ?? '', + value: answer, + }) + return { question, receipt, result: await result } +} + +describe('a keyed main adapter', () => { + it('uses the key saved in the credential store before the env var', async () => { + vi.stubEnv('ACME_TEST_KEY', 'sk-acme-env-0000') + const keyed = keyedMock([() => text('hello from acme')]) + const { session, saveKey, close } = await open({ adapter: keyed.adapter }) + await saveKey(SAVED_KEY) + expect(await session.prompt('hi')).toEqual({ + text: 'hello from acme', + media: [], + }) + expect(keyed.keys).toEqual([SAVED_KEY]) + await close() + }) + + it('reads the env var when no key is saved', async () => { + vi.stubEnv('ACME_TEST_KEY', 'sk-acme-env-0000') + const keyed = keyedMock([() => text('hello')]) + const { session, close } = await open({ adapter: keyed.adapter }) + await session.prompt('hi') + expect(keyed.keys).toEqual(['sk-acme-env-0000']) + await close() + }) + + it('stops the turn and asks for a sign-in when there is no key', async () => { + vi.stubEnv('ACME_TEST_KEY', '') + const keyed = keyedMock([() => text('never')]) + const { session, events, close } = await open({ adapter: keyed.adapter }) + await expect(session.prompt('hi')).rejects.toThrow( + 'Sign in to acme first. Run /connect acme.', + ) + expect(keyed.calls).toHaveLength(0) + expect(authRequired(events)).toEqual([{ connector: 'acme' }]) + await close() + }) + + it('builds a keyed model picker choice after /model', async () => { + const plain = mockAdapter([() => text('plain answer')]) + const keyed = keyedMock([() => text('keyed answer')]) + const { session, saveKey, close } = await open({ + adapter: plain.adapter, + plugins: [ + modelPicker({ + choices: { plain: plain.adapter, acme: keyed.adapter }, + default: 'plain', + }), + ], + }) + await saveKey(SAVED_KEY) + await session.command('model', 'acme') + expect(await session.prompt('hi')).toEqual({ + text: 'keyed answer', + media: [], + }) + expect(keyed.keys).toEqual([SAVED_KEY]) + await close() + }) +}) + +describe('ctx.keys in agents', () => { + /** An agent that builds a keyed adapter with `ctx.keys` and answers with it. */ + function writerWith(adapter: ReturnType['adapter']) { + return defineAgent({ + name: 'writer', + description: 'Writes with the Acme model', + run: async (ctx) => + ctx.chat({ adapter: await ctx.keys.adapter(adapter), stream: false }), + }) + } + + it('gives a background agent the session keys', async () => { + const keyed = keyedMock([() => text('agent answer')]) + const { session, saveKey, close } = await open({ + adapter: mockAdapter([]).adapter, + agents: [writerWith(keyed.adapter)], + }) + await saveKey(SAVED_KEY) + expect(await session.agent('writer')?.run()).toBe('agent answer') + expect(keyed.keys).toEqual([SAVED_KEY]) + await close() + }) + + it('gives a subagent the model calls the session keys', async () => { + const keyed = keyedMock([() => text('draft')]) + const lead = mockAdapter([ + () => toolCall('writer', {}), + () => text('The writer is done.'), + ]) + const { session, saveKey, close } = await open({ + adapter: lead.adapter, + subagents: [writerWith(keyed.adapter)], + }) + await saveKey(SAVED_KEY) + expect(await session.prompt('write it')).toEqual({ + text: 'The writer is done.', + media: [], + }) + expect(keyed.keys).toEqual([SAVED_KEY]) + await close() + }) +}) + +describe('compact with a keyed adapter', () => { + it('compacts with the saved key', async () => { + const summarizer = keyedMock([() => text('We talked.')]) + const main = mockAdapter([() => text('one'), () => text('two')]) + const { session, saveKey, close } = await open({ + adapter: main.adapter, + plugins: [compact({ adapter: summarizer.adapter })], + }) + await saveKey(SAVED_KEY) + await session.prompt('first') + await session.prompt('second') + expect(await session.command('compact')).toBe( + 'Compacted 4 messages into a summary.', + ) + expect(summarizer.keys).toEqual([SAVED_KEY]) + await close() + }) +}) + +describe('providerKeys', () => { + function plugin() { + return providerKeys({ providers: [acme, beta, gamma] }) + } + + it('saves a pasted key and shows it masked', async () => { + const { session, persistence, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [plugin()], + }) + const { question, result } = await connectWith(session, ` ${PASTED_KEY} `) + expect(question).toMatchObject({ + message: 'Paste your Acme API key', + secret: true, + }) + expect(result).toBe('Connected to Acme (key ...5678).') + expect( + await persistence.stores.credentials.get({ scope: USER, id: 'acme' }), + ).toEqual({ + type: 'api_key', + value: PASTED_KEY, + }) + expect(await session.command('keys')).toContain( + 'Acme: connected (key ...5678)', + ) + await close() + }) + + it('scrubs a key that /connect saved from the errors of the session', async () => { + const leaky = defineAgent({ + name: 'leaky', + description: 'Fails with the key in its error', + run: async () => { + throw new Error(`fetch failed: Bearer ${PASTED_KEY}`) + }, + }) + const { session, close } = await open({ + adapter: mockAdapter([]).adapter, + agents: [leaky], + plugins: [plugin()], + }) + await connectWith(session, PASTED_KEY) + const errors: Array = [] + session.onEvent(EventType.RUN_ERROR, ({ event }) => { + errors.push(event.message) + }) + await expect(session.agent('leaky')?.run()).rejects.toThrow('[redacted]') + expect(errors).toEqual(['fetch failed: Bearer [redacted]']) + await close() + }) + + it('refuses an empty key', async () => { + const { session, persistence, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [plugin()], + }) + await expect(connectWith(session, ' ')).rejects.toThrow( + 'No Acme key was given. Nothing was saved.', + ) + expect( + await persistence.stores.credentials.get({ scope: USER, id: 'acme' }), + ).toBeNull() + await close() + }) + + it('stops a cancelled /connect, and a late answer saves nothing', async () => { + const { session, persistence, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [plugin()], + }) + const connecting = session.command('connect:acme') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [question] = session.snapshot().pendingQuestions + await connecting.cancel() + await expect(connecting).rejects.toBeDefined() + if (question) + await session.answer({ + questionId: question.questionId, + value: PASTED_KEY, + }) + expect( + await persistence.stores.credentials.get({ scope: USER, id: 'acme' }), + ).toBeNull() + await close() + }) + + it('deletes the key on disconnect', async () => { + const { session, persistence, saveKey, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [plugin()], + }) + await saveKey(SAVED_KEY) + expect(await session.command('disconnect:acme')).toBe( + 'Disconnected from Acme.', + ) + expect( + await persistence.stores.credentials.get({ scope: USER, id: 'acme' }), + ).toBeNull() + expect(await session.command('keys')).toContain( + 'Acme: missing. Run /connect acme.', + ) + await close() + }) + + it('opens the key page of a provider with keyUrl before it asks', async () => { + const { session, persistence, events, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [ + providerKeys({ + providers: [{ ...acme, keyUrl: 'https://acme.example/keys' }], + }), + ], + }) + const { question, result } = await connectWith(session, PASTED_KEY) + expect(authRequired(events)).toEqual([ + { connector: 'acme', url: 'https://acme.example/keys' }, + ]) + expect(question).toMatchObject({ + message: + 'Paste your Acme API key. The page to make one is open in your browser.', + secret: true, + }) + expect(result).toBe('Connected to Acme (key ...5678).') + expect( + await persistence.stores.credentials.get({ scope: USER, id: 'acme' }), + ).toEqual({ + type: 'api_key', + value: PASTED_KEY, + }) + await close() + }) + + it('runs signIn with open through authRequired and saves its key', async () => { + const signIn = async (ctx: { open: (url: string) => void }) => { + ctx.open('https://acme.example/sign-in') + return 'sk-acme-signed-9876' + } + const { session, persistence, events, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [providerKeys({ providers: [{ ...acme, signIn }] })], + }) + expect(await session.command('connect:acme')).toBe( + 'Connected to Acme (key ...9876).', + ) + expect(authRequired(events)).toEqual([ + { connector: 'acme', url: 'https://acme.example/sign-in' }, + ]) + expect( + await persistence.stores.credentials.get({ scope: USER, id: 'acme' }), + ).toEqual({ + type: 'api_key', + value: 'sk-acme-signed-9876', + }) + await close() + }) + + it('lists every provider with its state in the plugin state', async () => { + vi.stubEnv('BETA_TEST_KEY', 'sk-beta-env-4321') + const { session, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [plugin()], + }) + expect(session.snapshot().plugins['tanstack/provider-keys']).toEqual({ + providers: [ + { id: 'acme', label: 'Acme', state: 'missing' }, + { id: 'beta', label: 'Beta', state: 'env' }, + { id: 'gamma', label: 'Gamma', state: 'missing' }, + ], + }) + expect(await session.command('keys')).toBe( + [ + 'Acme: missing. Run /connect acme.', + 'Beta: from the BETA_TEST_KEY env var', + 'Gamma: missing. Run /connect gamma.', + ].join('\n'), + ) + await connectWith(session, PASTED_KEY) + expect(session.snapshot().plugins['tanstack/provider-keys']).toEqual({ + providers: [ + { id: 'acme', label: 'Acme', state: 'connected' }, + { id: 'beta', label: 'Beta', state: 'env' }, + { id: 'gamma', label: 'Gamma', state: 'missing' }, + ], + }) + await close() + }) + + it('marks the key question as secret in the snapshot', async () => { + const { session, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [plugin()], + }) + const connecting = session.command('connect:acme') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toMatchObject([ + { message: 'Paste your Acme API key', secret: true }, + ]), + ) + const [question] = session.snapshot().pendingQuestions + if (question) + await session.answer({ + questionId: question.questionId, + value: PASTED_KEY, + }) + await connecting + await close() + }) + + it('keeps the key out of events, results, the inbox, and the transcript', async () => { + const keyed = keyedMock([() => text('hello')]) + const { session, persistence, events, close } = await open({ + adapter: keyed.adapter, + plugins: [plugin()], + }) + const { receipt, result } = await connectWith(session, PASTED_KEY) + const results = [ + result, + await session.command('keys'), + await session.prompt('hi'), + ] + expect(keyed.keys).toEqual([PASTED_KEY]) + const inbox = await persistence.stores.inbox.get({ + threadId: 't', + inputId: receipt.inputId, + }) + const transcript = await persistence.stores.messages.loadThread('t') + const everything = JSON.stringify({ events, results, inbox, transcript }) + expect(everything).not.toContain(PASTED_KEY) + expect(everything).toContain('...5678') + await close() + }) +}) diff --git a/packages/ai-mistral/src/adapters/text.ts b/packages/ai-mistral/src/adapters/text.ts index 40845ff0b5..b141e7f72e 100644 --- a/packages/ai-mistral/src/adapters/text.ts +++ b/packages/ai-mistral/src/adapters/text.ts @@ -750,6 +750,7 @@ export class MistralTextAdapter< toolName: toolCall.name, model: chunkModel, timestamp, + args: toolCall.arguments, input: parsedInput, }) } @@ -834,6 +835,9 @@ export class MistralTextAdapter< } for (const [, toolCall] of toolCallsInProgress) { if (toolCall.started && !toolCall.ended) { + // Parse before marking ended, so a parse error still lets the + // catch below send TOOL_CALL_END with the raw args. + const parsedInput = parseToolCallInput(toolCall, normalizeToolInput) toolCall.ended = true hasEmittedToolCall = true yield asChunk({ @@ -843,7 +847,8 @@ export class MistralTextAdapter< toolName: toolCall.name, model: lastChunkModel, timestamp, - input: parseToolCallInput(toolCall, normalizeToolInput), + args: toolCall.arguments, + input: parsedInput, }) } } @@ -900,9 +905,9 @@ export class MistralTextAdapter< for (const [, toolCall] of toolCallsInProgress) { if (toolCall.started && !toolCall.ended) { toolCall.ended = true - // Best-effort parse for the partial args; if invalid, surface - // empty input rather than throwing inside the cleanup path. - let partialInput: unknown = {} + // Best-effort parse for the partial args. If invalid, leave input + // out and send the raw args, so core reports the parse error. + let partialInput: unknown try { partialInput = toolCall.arguments ? normalizeToolInput( @@ -911,7 +916,7 @@ export class MistralTextAdapter< ) : {} } catch { - partialInput = {} + partialInput = undefined } yield asChunk({ type: 'TOOL_CALL_END', @@ -920,7 +925,8 @@ export class MistralTextAdapter< toolName: toolCall.name, model: lastChunkModel, timestamp, - input: partialInput, + args: toolCall.arguments, + ...(partialInput !== undefined && { input: partialInput }), }) } } diff --git a/packages/ai-mistral/tests/mistral-adapter.test.ts b/packages/ai-mistral/tests/mistral-adapter.test.ts index 5212814b9a..0d0fea0483 100644 --- a/packages/ai-mistral/tests/mistral-adapter.test.ts +++ b/packages/ai-mistral/tests/mistral-adapter.test.ts @@ -1087,6 +1087,109 @@ describe('Mistral AG-UI event emission', () => { expect(caught?.message).toMatch(/lookup_weather/) }) + it('puts the raw argument string on TOOL_CALL_END', async () => { + setupMockStream([ + { + id: 'cmpl-raw', + model: 'mistral-large-latest', + choices: [ + { + index: 0, + delta: { + toolCalls: [ + { + index: 0, + id: 'tc_raw', + function: { + name: 'lookup_weather', + arguments: '{"location": "Berlin", ', + }, + }, + ], + }, + finishReason: null, + }, + ], + }, + { + id: 'cmpl-raw', + model: 'mistral-large-latest', + choices: [ + { + index: 0, + delta: { + toolCalls: [{ index: 0, function: { arguments: '"n": 1.50}' } }], + }, + finishReason: 'tool_calls', + }, + ], + }, + ]) + + const adapter = createMistralText('mistral-large-latest', 'test-api-key') + const chunks: Array = [] + for await (const chunk of adapter.chatStream( + chatOpts({ + model: 'mistral-large-latest', + messages: [{ role: 'user', content: 'Weather in Berlin?' }], + tools: [weatherTool], + }), + )) { + chunks.push(chunk) + } + + const end = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(end).toMatchObject({ + toolCallId: 'tc_raw', + args: '{"location": "Berlin", "n": 1.50}', + input: { location: 'Berlin' }, + }) + }) + + it('sends malformed raw arguments on TOOL_CALL_END without an input', async () => { + setupMockStream([ + { + id: 'cmpl-bad', + model: 'mistral-large-latest', + choices: [ + { + index: 0, + delta: { + toolCalls: [ + { + index: 0, + id: 'tc_bad', + function: { name: 'lookup_weather', arguments: '{"a":' }, + }, + ], + }, + finishReason: 'tool_calls', + }, + ], + }, + ]) + + const adapter = createMistralText('mistral-large-latest', 'test-api-key') + const chunks: Array = [] + await expect( + (async () => { + for await (const chunk of adapter.chatStream( + chatOpts({ + model: 'mistral-large-latest', + messages: [{ role: 'user', content: 'x' }], + tools: [weatherTool], + }), + )) { + chunks.push(chunk) + } + })(), + ).rejects.toThrow(/Failed to parse tool call arguments/) + + const end = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(end).toMatchObject({ toolCallId: 'tc_bad', args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + it('sends stream_options.include_usage so Mistral returns usage on streaming', async () => { let capturedBody: unknown vi.stubGlobal( diff --git a/packages/ai-octane/src/realtime-types.ts b/packages/ai-octane/src/realtime-types.ts index 4bf2e45a94..24e06016b2 100644 --- a/packages/ai-octane/src/realtime-types.ts +++ b/packages/ai-octane/src/realtime-types.ts @@ -104,7 +104,7 @@ export interface UseRealtimeChatReturn { /** Current mode (idle, listening, thinking, speaking) */ mode: RealtimeMode /** Conversation messages */ - messages: Array + messages: ReadonlyArray /** User transcript while speaking (before finalized) */ pendingUserTranscript: string | null /** Assistant transcript while speaking (before finalized) */ diff --git a/packages/ai-octane/src/types.ts b/packages/ai-octane/src/types.ts index 255843d7be..61cc46e294 100644 --- a/packages/ai-octane/src/types.ts +++ b/packages/ai-octane/src/types.ts @@ -162,7 +162,7 @@ interface BaseUseChatReturn< * `messages[i].parts.find(p => p.type === 'structured-output')` is typed * with the schema's inferred shape — `data: T`, `partial: DeepPartial`. */ - messages: Array> + messages: ReadonlyArray> /** * Send a message and get a response. @@ -181,7 +181,7 @@ interface BaseUseChatReturn< * Pending messages queued while the client is busy (streaming, claiming a * send, or draining). Separate from `messages` until they drain. */ - queue: Array + queue: ReadonlyArray /** * Cancel a queued message before it drains. No-op if already sent. @@ -313,7 +313,7 @@ interface BaseUseChatReturn< /** * Set messages manually */ - setMessages: (messages: Array>) => void + setMessages: (messages: ReadonlyArray>) => void /** * Clear all messages diff --git a/packages/ai-octane/src/use-chat.tsrx b/packages/ai-octane/src/use-chat.tsrx index e1e4724eaa..10ae644e9a 100644 --- a/packages/ai-octane/src/use-chat.tsrx +++ b/packages/ai-octane/src/use-chat.tsrx @@ -11,13 +11,9 @@ import type { StreamChunk, } from '@tanstack/ai/client' import type { - ChatClientState, ResolvableChatInterrupt, - ChatInterruptState, ChatResumeState, - ConnectionStatus, InferredClientContext, - QueuedMessage, SendMessageOptions, StructuredOutputPart, } from '@tanstack/ai-client' @@ -30,9 +26,6 @@ import type { UseChatReturn, } from './types' -const EMPTY_INTERRUPTS = Object.freeze([]) -const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) - export function useChat< const TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, @@ -49,54 +42,13 @@ export function useChat< const hookId = useId() const clientId = options.threadId ?? hookId - const [messages, setMessages] = useState>>( - options.initialMessages || [], - ) - const [isLoading, setIsLoading] = useState(false) - const [hasOlderMessages, setHasOlderMessages] = useState(false) - const [error, setError] = useState(undefined) - const [status, setStatus] = useState('ready') - const [isSubscribed, setIsSubscribed] = useState(false) - const [connectionStatus, setConnectionStatus] = - useState('disconnected') - const [sessionGenerating, setSessionGenerating] = useState(false) - const [queue, setQueue] = useState>([]) - const [runId, setRunId] = useState(null) - const [interruptState, setInterruptState] = useState< - ChatInterruptState - >(() => ({ - interrupts: EMPTY_INTERRUPTS, - pendingInterrupts: EMPTY_INTERRUPTS, - interruptErrors: EMPTY_INTERRUPT_ERRORS, - resuming: false, - })) - - type Partial = DeepPartial>> - type Final = InferSchemaType> - - // Track the rendered snapshot so a replacement client only synchronizes - // state when its messages actually differ. - const messagesRef = useRef>>( - options.initialMessages || [], - ) const subscribedRef = useRef(false) const activeClientRef = useRef(null) const seenClientRef = useRef(null) - - // Update ref synchronously during render so it's always current when useMemo runs. - messagesRef.current = messages - - // Track current options in a ref to avoid recreating client when options change const optionsRef = useRef>(options) optionsRef.current = options - const syncResumeState = useCallback((target: ChatClient | null) => { - if (!target) return - setRunId(target.getCurrentRunId()) - setInterruptState(target.getInterruptState()) - }, []) - // Create ChatClient instance with callbacks to sync state const { client, initialization } = useMemo(() => { const messagesToUse = options.initialMessages || [] @@ -213,77 +165,11 @@ export function useChat< ...(options.streamProcessor !== undefined && { streamProcessor: options.streamProcessor, }), - onMessagesChange: (newMessages: Array>) => { - runOrQueueForActiveInstance(() => { - setMessages(newMessages) - const currentInstance = getActiveInstance() - if (!currentInstance) return - setHasOlderMessages(currentInstance.getHasOlderMessages()) - }) - }, - onLoadingChange: (newIsLoading: boolean) => { - runOrQueueForActiveInstance(() => { - const currentInstance = getActiveInstance() - if (!currentInstance) return - setIsLoading(newIsLoading) - syncResumeState(currentInstance) - }) - }, - onErrorChange: (newError: Error | undefined) => { - runOrQueueForActiveInstance(() => { - setError(newError) - }) - }, - onStatusChange: (status: ChatClientState) => { - runOrQueueForActiveInstance(() => { - setStatus(status) - }) - }, - onSubscriptionChange: (nextIsSubscribed: boolean) => { - runOrQueueForActiveInstance(() => { - setIsSubscribed(nextIsSubscribed) - }) - }, - onConnectionStatusChange: (nextStatus: ConnectionStatus) => { - runOrQueueForActiveInstance(() => { - setConnectionStatus(nextStatus) - }) - }, - onSessionGeneratingChange: (isGenerating: boolean) => { - runOrQueueForActiveInstance(() => { - setSessionGenerating(isGenerating) - }) - }, - onHydratingChange: (nextIsHydrating: boolean) => { - runOrQueueForActiveInstance(() => { - setHydration({ client: instance, value: nextIsHydrating }) - }) - }, ...(optionsRef.current.queue !== undefined && { queue: optionsRef.current.queue, }), - onQueueChange: (nextQueue: Array) => { - runOrQueueForActiveInstance(() => { - setQueue(nextQueue) - }) - }, - onRunIdChange: (nextRunId) => { - runOrQueueForActiveInstance(() => { - setRunId(nextRunId) - }) - }, - onResumeStateChange: (_nextResumeState, nextPendingInterrupts) => { - runOrQueueForActiveInstance(() => { - setInterruptState((current) => ({ - ...current, - interrupts: nextPendingInterrupts, - pendingInterrupts: nextPendingInterrupts, - })) - }) - }, onInterruptStateChange: (nextInterruptState, context) => { runOrQueueForActiveInstance(() => { - setInterruptState(nextInterruptState) optionsRef.current.onInterruptStateChange?.( nextInterruptState, context, @@ -293,18 +179,21 @@ export function useChat< }) instanceHolder.current = instance return { client: instance, initialization: initializationState } - }, [clientId, syncResumeState]) + }, [clientId]) + + type Partial = DeepPartial>> + type Final = InferSchemaType> - // Seeded from the client so the first render already shows hydrating, not an - // empty chat. The constructor sets it without firing a callback. Tagged with - // its client: after a thread switch, the new client's value shows on the - // first render, before the effects sync it. - const [hydration, setHydration] = useState(() => ({ + // Constructor hydration (sync store / resume snapshot) must be on the first + // paint. `useEffect` runs after paint, so seed from the snapshot here. + // Tagged with its client: after a thread switch, the new client's snapshot + // shows on the first render, before the effect below subscribes. + const [held, setHeld] = useState(() => ({ client, - value: client.getIsHydrating(), + snapshot: client.getSnapshot(), })) - const isHydrating = - hydration.client === client ? hydration.value : client.getIsHydrating() + const snapshot = + held.client === client ? held.snapshot : client.getSnapshot() useEffect(() => { activeClientRef.current = client @@ -325,12 +214,16 @@ export function useChat< }, [client, initialization]) useEffect(() => { - const clientMessages = client.getMessages() - if (clientMessages !== messagesRef.current) { - setMessages(clientMessages) + const applySnapshot = () => { + const next = client.getSnapshot() + setHeld((prev) => + prev.client === client && prev.snapshot === next + ? prev + : { client, snapshot: next }, + ) } - setHasOlderMessages(client.getHasOlderMessages()) - setHydration({ client, value: client.getIsHydrating() }) + applySnapshot() + return client.subscribeSnapshot(applySnapshot) }, [client]) // OCTANE DIVERGENCE: upstream captures the initial transport. ChatClient @@ -400,7 +293,6 @@ export function useChat< useEffect(() => { client.attach() client.mountDevtools() - syncResumeState(client) return () => { client.detach() if (optionsRef.current.live) { @@ -413,20 +305,16 @@ export function useChat< activeClientRef.current = null } } - }, [client, syncResumeState]) + }, [client]) const sendMessage = useCallback( async ( content: string | MultimodalContent, sendOptions?: SendMessageOptions, ) => { - try { - await client.sendMessage(content, undefined, sendOptions) - } finally { - syncResumeState(client) - } + await client.sendMessage(content, undefined, sendOptions) }, - [client, syncResumeState], + [client], ) const cancelQueued = useCallback( @@ -438,26 +326,17 @@ export function useChat< const append = useCallback( async (message: ModelMessage | UIMessage) => { - try { - await client.append(message) - } finally { - syncResumeState(client) - } + await client.append(message) }, - [client, syncResumeState], + [client], ) const reload = useCallback(async () => { - try { - await client.reload() - } finally { - syncResumeState(client) - } - }, [client, syncResumeState]) + await client.reload() + }, [client]) const loadOlderMessages = useCallback(async () => { await client.loadOlderMessages() - setHasOlderMessages(client.getHasOlderMessages()) }, [client]) const stop = useCallback(() => { @@ -466,11 +345,10 @@ export function useChat< const clear = useCallback(() => { client.clear() - syncResumeState(client) - }, [client, syncResumeState]) + }, [client]) const setMessagesManually = useCallback( - (newMessages: Array>) => { + (newMessages: ReadonlyArray>) => { client.setMessagesManually(newMessages) }, [client], @@ -492,21 +370,15 @@ export function useChat< const addToolApprovalResponse = useCallback( async (response: { id: string; approved: boolean }) => { await client.addToolApprovalResponse(response) - syncResumeState(client) }, - [client, syncResumeState], + [client], ) const resumeInterrupts = useCallback( async (resumeItems: Array, state?: ChatResumeState) => { - const result = await client.resumeInterrupts( - resumeItems, - state ?? undefined, - ) - syncResumeState(client) - return result + return await client.resumeInterrupts(resumeItems, state ?? undefined) }, - [client, syncResumeState], + [client], ) const resolveInterrupts = useCallback( @@ -550,7 +422,7 @@ export function useChat< // a stale assistant turn or a system prompt) we deliberately return null // rather than scanning historical assistants — otherwise a `final` from a // previous session would leak into the hook value on first render. - const renderedMessages = client.getMessages() + const renderedMessages = snapshot.messages const activeStructuredPart = useMemo(() => { let lastUserIndex = -1 @@ -595,26 +467,26 @@ export function useChat< append, reload, stop, - isLoading, - hasOlderMessages, + isLoading: snapshot.isLoading, + hasOlderMessages: snapshot.hasOlderMessages, loadOlderMessages, - error, - status, - isSubscribed, - connectionStatus, - sessionGenerating, - isHydrating, + error: snapshot.error, + status: snapshot.status, + isSubscribed: snapshot.isSubscribed, + connectionStatus: snapshot.connectionStatus, + sessionGenerating: snapshot.sessionGenerating, + isHydrating: snapshot.isHydrating, setMessages: setMessagesManually, clear, addToolResult, addToolApprovalResponse, - queue, + queue: snapshot.queue, cancelQueued, - runId, - interrupts: interruptState.interrupts, - pendingInterrupts: interruptState.pendingInterrupts, - interruptErrors: interruptState.interruptErrors, - resuming: interruptState.resuming, + runId: snapshot.runId, + interrupts: snapshot.interruptState.interrupts, + pendingInterrupts: snapshot.interruptState.pendingInterrupts, + interruptErrors: snapshot.interruptState.interruptErrors, + resuming: snapshot.interruptState.resuming, resolveInterrupts, cancelInterrupts, retryInterrupts, diff --git a/packages/ai-octane/src/use-generate-video.tsrx b/packages/ai-octane/src/use-generate-video.tsrx index 5772a97326..c18dddc7bf 100644 --- a/packages/ai-octane/src/use-generate-video.tsrx +++ b/packages/ai-octane/src/use-generate-video.tsrx @@ -172,12 +172,6 @@ export function useGenerateVideo( onStatusUpdate: (s: VideoStatusInfo) => { optionsRef.current.onStatusUpdate?.(s) }, - onResultChange: setResult, - onLoadingChange: setIsLoading, - onErrorChange: setError, - onStatusChange: setStatus, - onJobIdChange: setJobId, - onVideoStatusChange: setVideoStatus, } if (opts.connection) { @@ -199,6 +193,20 @@ export function useGenerateVideo( ) }, [clientId]) + useEffect(() => { + const applySnapshot = () => { + const next = client.getSnapshot() + setResult(next.result) + setIsLoading(next.isLoading) + setError(next.error) + setStatus(next.status) + setJobId(next.jobId) + setVideoStatus(next.videoStatus) + } + applySnapshot() + return client.subscribe(applySnapshot) + }, [client]) + // Sync body changes without recreating client useEffect(() => { // Conditional spread: target uses strict-optional `body?: T`. diff --git a/packages/ai-octane/src/use-generation.tsrx b/packages/ai-octane/src/use-generation.tsrx index e3ca19c084..23eb3b0489 100644 --- a/packages/ai-octane/src/use-generation.tsrx +++ b/packages/ai-octane/src/use-generation.tsrx @@ -160,10 +160,6 @@ export function useGeneration< onChunk: (c: StreamChunk) => { optionsRef.current.onChunk?.(c) }, - onResultChange: setResult, - onLoadingChange: setIsLoading, - onErrorChange: setError, - onStatusChange: setStatus, } if (opts.connection) { @@ -185,6 +181,18 @@ export function useGeneration< ) }, [clientId]) + useEffect(() => { + const applySnapshot = () => { + const next = client.getSnapshot() + setResult(next.result) + setIsLoading(next.isLoading) + setError(next.error) + setStatus(next.status) + } + applySnapshot() + return client.subscribe(applySnapshot) + }, [client]) + // Sync body changes without recreating client useEffect(() => { // Conditional spread: target uses strict-optional `body?: T`. diff --git a/packages/ai-octane/src/use-realtime-chat.tsrx b/packages/ai-octane/src/use-realtime-chat.tsrx index b486b8cc12..03f4730f5d 100644 --- a/packages/ai-octane/src/use-realtime-chat.tsrx +++ b/packages/ai-octane/src/use-realtime-chat.tsrx @@ -61,7 +61,7 @@ export function useRealtimeChat( // State const [status, setStatus] = useState('idle') const [mode, setMode] = useState('idle') - const [messages, setMessages] = useState>([]) + const [messages, setMessages] = useState>([]) const [pendingUserTranscript, setPendingUserTranscript] = useState< string | null >(null) @@ -124,17 +124,14 @@ export function useRealtimeChat( semanticEagerness: initial.semanticEagerness, }), onStatusChange: (newStatus) => { - setStatus(newStatus) // OCTANE DIVERGENCE: upstream declares this callback but does not // forward status changes beyond its own hook state. optionsRef.current.onStatusChange?.(newStatus) }, onModeChange: (newMode) => { - setMode(newMode) optionsRef.current.onModeChange?.(newMode) }, onMessage: (message) => { - setMessages((prev) => [...prev, message]) optionsRef.current.onMessage?.(message) }, onUsage: (usage) => { @@ -144,31 +141,36 @@ export function useRealtimeChat( optionsRef.current.onGoAway?.(timeLeft) }, onError: (err) => { - setError(err) optionsRef.current.onError?.(err) }, onConnect: () => { - setError(null) optionsRef.current.onConnect?.() }, onDisconnect: () => { optionsRef.current.onDisconnect?.() }, onInterrupted: () => { - setPendingAssistantTranscript(null) optionsRef.current.onInterrupted?.() }, }) - - // Subscribe to state changes for transcripts - clientRef.current.onStateChange((state) => { - setPendingUserTranscript(state.pendingUserTranscript) - setPendingAssistantTranscript(state.pendingAssistantTranscript) - }) } const client = clientRef.current + useEffect(() => { + const applySnapshot = () => { + const next = client.getSnapshot() + setStatus(next.status) + setMode(next.mode) + setMessages(next.messages) + setPendingUserTranscript(next.pendingUserTranscript) + setPendingAssistantTranscript(next.pendingAssistantTranscript) + setError(next.error) + } + applySnapshot() + return client.subscribe(applySnapshot) + }, [client]) + // Audio level animation loop useEffect(() => { function updateLevels() { @@ -200,10 +202,6 @@ export function useRealtimeChat( // Connection methods const connect = useCallback(async () => { - setError(null) - setMessages([]) - setPendingUserTranscript(null) - setPendingAssistantTranscript(null) await client.connect() }, [client]) diff --git a/packages/ai-octane/tests/conformance/use-chat.test.ts b/packages/ai-octane/tests/conformance/use-chat.test.ts index ccc5147141..7fbd8f2df4 100644 --- a/packages/ai-octane/tests/conformance/use-chat.test.ts +++ b/packages/ai-octane/tests/conformance/use-chat.test.ts @@ -215,15 +215,19 @@ describe('useChat', () => { removeItem: vi.fn(), } - const { result } = renderUseChat({ - connection: adapter, - threadId: 'persisted-chat', - persistence, + let firstMessages: ReadonlyArray | undefined + const { result } = renderHook(() => { + const chat = useChat({ + connection: adapter, + threadId: 'persisted-chat', + persistence, + }) + firstMessages ??= chat.messages + return chat }) - await waitFor(() => { - expect(result.current.messages).toEqual(persistedMessages) - }) + expect(firstMessages).toEqual(persistedMessages) + expect(result.current.messages).toEqual(persistedMessages) expect(persistence.getItem).toHaveBeenCalledWith('persisted-chat') }) @@ -2449,4 +2453,87 @@ describe('useChat', () => { ) }) }) + + describe('snapshot', () => { + function message(id: string): UIMessage { + return { + id, + role: 'user', + parts: [{ type: 'text', content: id }], + createdAt: new Date(), + } + } + + it('returns the frozen snapshot array', async () => { + const { result } = renderUseChat({ + connection: createMockConnectionAdapter({ + chunks: createTextChunks('Hi'), + }), + }) + + await act(async () => { + await result.current.sendMessage('Hello') + }) + + const messages = result.current.messages + expect(messages).toHaveLength(2) + expect(Object.isFrozen(messages)).toBe(true) + expect(() => (messages as Array).push({})).toThrow(TypeError) + }) + + it("shows the new client's snapshot on the first render after a thread switch", () => { + const stored: Record> = { + a: [message('a')], + b: [message('b')], + } + const persistence = { + getItem: vi.fn((id: string) => stored[id] ?? []), + setItem: vi.fn(), + removeItem: vi.fn(), + } + const renders: Array<[string, string | undefined]> = [] + + const { result } = renderHook(() => { + const [threadId, setId] = useState('a') + const chat = useChat({ + connection: createMockConnectionAdapter(), + threadId, + persistence, + }) + renders.push([threadId, chat.messages[0]?.id]) + return { ...chat, setId } + }) + + act(() => { + result.current.setId('b') + }) + + expect(result.current.messages).toEqual(stored.b) + for (const [threadId, firstId] of renders) { + expect(firstId).toBe(threadId) + } + }) + + it('stops reading the client snapshot after unmount', async () => { + let renders = 0 + const { result, unmount } = renderHook(() => { + renders++ + return useChat({ connection: createMockConnectionAdapter() }) + }) + const { setMessages } = result.current + + act(() => { + setMessages([message('live')]) + }) + expect(result.current.messages).toHaveLength(1) + + unmount() + const rendersAtUnmount = renders + act(() => { + setMessages([message('after-unmount')]) + }) + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(renders).toBe(rendersAtUnmount) + }) + }) }) diff --git a/packages/ai-octane/typetests/use-chat.test-d.tsx b/packages/ai-octane/typetests/use-chat.test-d.tsx index 03b84a5639..475e16c6cd 100644 --- a/packages/ai-octane/typetests/use-chat.test-d.tsx +++ b/packages/ai-octane/typetests/use-chat.test-d.tsx @@ -51,7 +51,7 @@ describe('useChat() return type', () => { type R = UseChatReturn expectTypeOf().toBeFunction() expectTypeOf().toBeBoolean() - expectTypeOf().toBeArray() + expectTypeOf().toExtend>() }) it('options accept outputSchema with the schema type', () => { @@ -81,7 +81,7 @@ describe('useChat() return type', () => { it('exposes queue, runId, and interrupt controls', () => { type R = UseChatReturn - expectTypeOf().toEqualTypeOf>() + expectTypeOf().toEqualTypeOf>() expectTypeOf().toEqualTypeOf<(id: string) => void>() expectTypeOf().toEqualTypeOf() expectTypeOf().toBeBoolean() diff --git a/packages/ai-ollama/src/adapters/text.ts b/packages/ai-ollama/src/adapters/text.ts index bc87d6cfee..ca1b59df65 100644 --- a/packages/ai-ollama/src/adapters/text.ts +++ b/packages/ai-ollama/src/adapters/text.ts @@ -355,7 +355,7 @@ export class OllamaTextAdapter extends BaseTextAdapter< } // Serialize arguments to a string for the TOOL_CALL_ARGS event - let parsedInput: unknown = {} + let parsedInput: unknown const argsStr = typeof actualToolCall.function.arguments === 'string' ? actualToolCall.function.arguments @@ -364,7 +364,8 @@ export class OllamaTextAdapter extends BaseTextAdapter< const parsed = JSON.parse(argsStr) parsedInput = parsed && typeof parsed === 'object' ? parsed : {} } catch { - parsedInput = actualToolCall.function.arguments + // Leave input out. Core reports the malformed raw args. + parsedInput = undefined } // Emit TOOL_CALL_ARGS with full args (Ollama doesn't stream args incrementally) @@ -385,7 +386,8 @@ export class OllamaTextAdapter extends BaseTextAdapter< toolName: actualToolCall.function.name || '', model: chunk.model, timestamp: Date.now(), - input: parsedInput, + args: argsStr, + ...(parsedInput !== undefined && { input: parsedInput }), }) return events diff --git a/packages/ai-ollama/tests/text-adapter.test.ts b/packages/ai-ollama/tests/text-adapter.test.ts index dd759c34b7..14dde3a044 100644 --- a/packages/ai-ollama/tests/text-adapter.test.ts +++ b/packages/ai-ollama/tests/text-adapter.test.ts @@ -276,6 +276,53 @@ describe('OllamaTextAdapter.chatStream (tool calls)', () => { expect(startChunk!.toolCallId).toMatch(/^search_\d+/) }) + async function toolCallEnd(args: unknown) { + chatMock.mockResolvedValueOnce( + asyncIterable([ + { + message: { + role: 'assistant', + content: '', + tool_calls: [ + { id: 'tc-raw', function: { name: 'search', arguments: args } }, + ], + }, + done: true, + done_reason: 'stop', + }, + ]), + ) + const chunks = await collectStream( + createOllamaChat('llama3.2').chatStream({ + logger: testLogger, + model: 'llama3.2', + messages: [{ role: 'user', content: 'find cats' }], + tools: [searchTool], + }), + ) + return chunks.find((c) => c.type === 'TOOL_CALL_END') + } + + it('puts the raw argument string on TOOL_CALL_END', async () => { + const end = await toolCallEnd('{"q": "cats", "n": 1.50}') + expect(end).toMatchObject({ + toolCallId: 'tc-raw', + args: '{"q": "cats", "n": 1.50}', + input: { q: 'cats', n: 1.5 }, + }) + }) + + it('puts the serialized object arguments on TOOL_CALL_END', async () => { + const end = await toolCallEnd({ q: 'cats' }) + expect(end).toMatchObject({ args: '{"q":"cats"}', input: { q: 'cats' } }) + }) + + it('sends malformed raw arguments on TOOL_CALL_END without an input', async () => { + const end = await toolCallEnd('{"a":') + expect(end).toMatchObject({ toolCallId: 'tc-raw', args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + it('forwards tools to the ollama client in provider format', async () => { chatMock.mockResolvedValueOnce( asyncIterable([ diff --git a/packages/ai-opencode/src/adapters/text.ts b/packages/ai-opencode/src/adapters/text.ts index 0d3636fdcc..bb0d39cb12 100644 --- a/packages/ai-opencode/src/adapters/text.ts +++ b/packages/ai-opencode/src/adapters/text.ts @@ -54,6 +54,7 @@ import type { OpencodeTextProviderOptions } from '../provider-options' const DEFAULT_WORKDIR = '/workspace' const DEFAULT_PORT = 4096 +const TEARDOWN_STEP_MS = 1000 export interface OpencodeTextConfig { /** Working directory inside the sandbox. Defaults to `/workspace`. */ @@ -544,9 +545,26 @@ export class OpencodeTextAdapter< externalSignal.removeEventListener('abort', onAbort) } channel.close() - await handle?.dispose() - await server?.dispose() - await bridge?.close() + // Each step gets a time limit: a step that never settles (a sandbox + // `kill()` that never returns, a bridge request held open) must not + // hold back RUN_FINISHED. + const bounded = (step: () => Promise | undefined) => + Promise.race([ + Promise.resolve() + .then(step) + .catch((error: unknown) => + logger.errors('opencode teardown failed', { + error, + source: 'opencode.chatStream', + }), + ), + new Promise((resolve) => + setTimeout(resolve, TEARDOWN_STEP_MS).unref?.(), + ), + ]) + await bounded(() => handle?.dispose()) + await bounded(() => server?.dispose()) + await bounded(() => bridge?.close()) } } diff --git a/packages/ai-opencode/tests/teardown.test.ts b/packages/ai-opencode/tests/teardown.test.ts new file mode 100644 index 0000000000..546f47db67 --- /dev/null +++ b/packages/ai-opencode/tests/teardown.test.ts @@ -0,0 +1,58 @@ +/** + * A teardown step that never settles must not hold back the end of the run + * (issue #1638). The server's `dispose()` and the bridge's `close()` both + * hang here; the stream must still end. + */ +import { describe, expect, it, vi } from 'vitest' +import { + provideSandbox, + provideToolBridgeProvisioner, +} from '@tanstack/ai-sandbox' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { opencodeText } from '../src/index' +import type { CapabilityContext } from '@tanstack/ai' +import type { SandboxHandle } from '@tanstack/ai-sandbox' + +const never = () => new Promise(() => {}) + +vi.mock('../src/process/sandbox-server', () => ({ + startOpencodeServerInSandbox: () => + Promise.resolve({ baseUrl: 'http://127.0.0.1:1', dispose: never }), +})) +vi.mock('../src/process/server', () => ({ + startOpencodeSession: () => Promise.reject(new Error('no session')), +})) + +const MODEL = 'anthropic/claude-sonnet-4-5' + +describe('opencode teardown', () => { + it('ends the stream when dispose() and close() never settle', async () => { + const capabilities = { + capabilities: { markProvided: () => {}, has: () => true }, + } as unknown as CapabilityContext + provideSandbox(capabilities, { + id: 'sbx', + provider: 'mock', + } as unknown as SandboxHandle) + provideToolBridgeProvisioner(capabilities, { + provision: () => + Promise.resolve({ + name: 'tanstack', + url: 'http://127.0.0.1:1/mcp', + token: 'token', + close: never, + }), + }) + const types: Array = [] + for await (const chunk of opencodeText(MODEL).chatStream({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + logger: resolveDebugOption(false), + capabilities, + tools: [{ name: 'getWeather', description: 'weather' }], + })) { + types.push(chunk.type) + } + expect(types).toContain('RUN_ERROR') + }) +}) diff --git a/packages/ai-openrouter/src/adapters/responses-text.ts b/packages/ai-openrouter/src/adapters/responses-text.ts index 89a9356964..901fcd4dc6 100644 --- a/packages/ai-openrouter/src/adapters/responses-text.ts +++ b/packages/ai-openrouter/src/adapters/responses-text.ts @@ -19,6 +19,10 @@ import { generateId } from '@tanstack/ai-utils' import { extractRequestOptions } from '../internal/request-options' import { clientForCall } from '../internal/wrap-fetch' import { openRouterSupportsCombinedToolsAndSchema } from '../internal/combined-tools-and-schema' +import { + openRouterUsesJsonObject, + withJsonModePrompt, +} from '../internal/json-object-mode' import { makeStructuredOutputCompatible } from '../internal/schema-converter' import { convertFunctionToolToResponsesFormat, @@ -280,12 +284,14 @@ export class OpenRouterResponsesTextAdapter< options: StructuredOutputOptions, ): Promise> { const { chatOptions, outputSchema } = options - const responsesRequest = this.mapOptionsToRequest(chatOptions) - - const jsonSchema = this.makeStructuredOutputCompatible( - outputSchema, - outputSchema.required, + const jsonMode = openRouterUsesJsonObject( + chatOptions.model || this.model, + chatOptions.modelOptions, ) + const responsesRequest = this.mapOptionsToRequest( + jsonMode ? withJsonModePrompt(chatOptions, outputSchema) : chatOptions, + ) + const textFormat = this.structuredTextFormat(jsonMode, outputSchema) try { chatOptions.logger.request( @@ -303,12 +309,7 @@ export class OpenRouterResponsesTextAdapter< ...responsesRequest, stream: false, text: { - format: { - type: 'json_schema', - name: 'structured_output', - schema: jsonSchema, - strict: true, - }, + format: textFormat, }, }, }, @@ -411,12 +412,14 @@ export class OpenRouterResponsesTextAdapter< options: StructuredOutputOptions, ): AsyncIterable { const { chatOptions, outputSchema } = options - const responsesRequest = this.mapOptionsToRequest(chatOptions) - - const jsonSchema = this.makeStructuredOutputCompatible( - outputSchema, - outputSchema.required, + const jsonMode = openRouterUsesJsonObject( + chatOptions.model || this.model, + chatOptions.modelOptions, + ) + const responsesRequest = this.mapOptionsToRequest( + jsonMode ? withJsonModePrompt(chatOptions, outputSchema) : chatOptions, ) + const textFormat = this.structuredTextFormat(jsonMode, outputSchema) const aguiState = { runId: generateId(this.name), @@ -522,12 +525,7 @@ export class OpenRouterResponsesTextAdapter< ...responsesRequest, stream: true, text: { - format: { - type: 'json_schema', - name: 'structured_output', - schema: jsonSchema, - strict: true, - }, + format: textFormat, }, }, }, @@ -1471,7 +1469,7 @@ export class OpenRouterResponsesTextAdapter< rawArguments: chunk.arguments, }, ) - parsedInput = {} + parsedInput = undefined } } @@ -1482,7 +1480,8 @@ export class OpenRouterResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: chunk.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } } @@ -1552,7 +1551,7 @@ export class OpenRouterResponsesTextAdapter< rawArguments: rawArgs, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -1562,7 +1561,8 @@ export class OpenRouterResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: rawArgs, + ...(parsedInput !== undefined && { input: parsedInput }), } metadata.ended = true metadata.pendingArguments = undefined @@ -1679,7 +1679,7 @@ export class OpenRouterResponsesTextAdapter< rawArguments: rawArgs, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -1689,7 +1689,8 @@ export class OpenRouterResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: rawArgs, + ...(parsedInput !== undefined && { input: parsedInput }), } metadata.ended = true metadata.pendingArguments = undefined @@ -1863,8 +1864,9 @@ export class OpenRouterResponsesTextAdapter< ) : undefined - // Attach text.format json_schema only when outputSchema is set and every - // routed model is in the combined-capable set. + // Attach text.format json_schema only when outputSchema is set and the + // combined gate passes (every routed model is combined-capable, and not + // JSON mode). const combinedOutputSchema: JSONSchema | undefined = options.outputSchema const combinedSchema = combinedOutputSchema && @@ -1957,13 +1959,43 @@ export class OpenRouterResponsesTextAdapter< /** * Combined mode is safe only when this model and every `modelOptions.models` - * fallback are in `OPENROUTER_COMBINED_TOOLS_AND_SCHEMA_MODELS`. - * `:variant` suffixes are routing directives and do not change the gate. + * fallback are in `OPENROUTER_COMBINED_TOOLS_AND_SCHEMA_MODELS`. JSON mode + * is never combined: the schema goes in a system prompt on a separate + * structured-output request, not on the tool-calling turns. */ supportsCombinedToolsAndSchema( modelOptions?: OpenRouterResponsesTextProviderOptions, ): boolean { - return openRouterSupportsCombinedToolsAndSchema(this.model, modelOptions) + return ( + !openRouterUsesJsonObject(this.model, modelOptions) && + openRouterSupportsCombinedToolsAndSchema(this.model, modelOptions) + ) + } + + /** `json_object` in JSON mode, else strict `json_schema` from `outputSchema`. */ + private structuredTextFormat( + jsonMode: boolean, + outputSchema: JSONSchema, + ): + | { type: 'json_object' } + | { + type: 'json_schema' + name: 'structured_output' + schema: Record + strict: true + } { + if (jsonMode) { + return { type: 'json_object' } + } + return { + type: 'json_schema', + name: 'structured_output', + schema: this.makeStructuredOutputCompatible( + outputSchema, + outputSchema.required, + ), + strict: true, + } } /** diff --git a/packages/ai-openrouter/src/adapters/text.ts b/packages/ai-openrouter/src/adapters/text.ts index 677d8dc636..027913b697 100644 --- a/packages/ai-openrouter/src/adapters/text.ts +++ b/packages/ai-openrouter/src/adapters/text.ts @@ -19,6 +19,10 @@ import { extractRequestOptions } from '../internal/request-options' import { clientForCall } from '../internal/wrap-fetch' import { makeStructuredOutputCompatible } from '../internal/schema-converter' import { openRouterSupportsCombinedToolsAndSchema } from '../internal/combined-tools-and-schema' +import { + openRouterUsesJsonObject, + withJsonModePrompt, +} from '../internal/json-object-mode' import { OPENROUTER_MODEL_INPUT_MODALITIES, OPENROUTER_MODEL_REASONING, @@ -289,16 +293,23 @@ export class OpenRouterTextAdapter< /** * Generate structured output via OpenRouter's `responseFormat`. Uses * `stream: false`. Default is strict `json_schema` from `outputSchema`. - * Callers can opt into JSON mode with - * `modelOptions.responseFormat: { type: 'json_object' }`. + * JSON mode is `modelOptions.responseFormat: { type: 'json_object' }`, or + * a catalog model that has a response format and no strict structured + * outputs. JSON mode adds the schema as a system prompt. */ async structuredOutput( options: StructuredOutputOptions>, ): Promise> { const { chatOptions, outputSchema } = options - const chatRequest = this.mapOptionsToRequest(chatOptions) + const jsonMode = openRouterUsesJsonObject( + chatOptions.model || this.model, + chatOptions.modelOptions, + ) + const chatRequest = this.mapOptionsToRequest( + jsonMode ? withJsonModePrompt(chatOptions, outputSchema) : chatOptions, + ) const responseFormat = this.resolveStructuredResponseFormat( - chatRequest.responseFormat, + jsonMode, outputSchema, ) @@ -438,9 +449,15 @@ export class OpenRouterTextAdapter< options: StructuredOutputOptions>, ): AsyncIterable { const { chatOptions, outputSchema } = options - const chatRequest = this.mapOptionsToRequest(chatOptions) + const jsonMode = openRouterUsesJsonObject( + chatOptions.model || this.model, + chatOptions.modelOptions, + ) + const chatRequest = this.mapOptionsToRequest( + jsonMode ? withJsonModePrompt(chatOptions, outputSchema) : chatOptions, + ) const responseFormat = this.resolveStructuredResponseFormat( - chatRequest.responseFormat, + jsonMode, outputSchema, ) @@ -788,15 +805,14 @@ export class OpenRouterTextAdapter< /** * Resolve the provider request format for a schema-bearing call. * - * Explicit `modelOptions.responseFormat: { type: 'json_object' }` is - * forwarded. Every other value is replaced with strict `json_schema` - * generated from `outputSchema`. + * `json_object` in JSON mode. Every other call gets strict + * `json_schema` from `outputSchema`. */ protected resolveStructuredResponseFormat( - requested: ChatRequest['responseFormat'], + jsonMode: boolean, outputSchema: JSONSchema, ): NonNullable { - if (requested?.type === 'json_object') { + if (jsonMode) { return { type: 'json_object' } } @@ -1167,7 +1183,7 @@ export class OpenRouterTextAdapter< rawArguments: toolCall.arguments, }, ) - parsedInput = {} + parsedInput = undefined } } @@ -1179,7 +1195,8 @@ export class OpenRouterTextAdapter< toolName: toolCall.name, model: chunk.model || options.model, timestamp: Date.now(), - input: parsedInput, + args: toolCall.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } emittedAnyToolCallEnd = true } @@ -1230,7 +1247,7 @@ export class OpenRouterTextAdapter< rawArguments: toolCall.arguments, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -1240,7 +1257,8 @@ export class OpenRouterTextAdapter< toolName: toolCall.name, model: lastModel || options.model, timestamp: Date.now(), - input: parsedInput, + args: toolCall.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } emittedAnyToolCallEnd = true } @@ -1481,16 +1499,11 @@ export class OpenRouterTextAdapter< ? convertToolsToProviderFormat(options.tools) : undefined - // Attach json_schema only when outputSchema is set, every routed model - // is in the combined set, and the caller did not opt into JSON mode. + // Attach json_schema only when outputSchema is set and the combined gate + // passes (every routed model is in the combined set, and not JSON mode). const combinedOutputSchema: JSONSchema | undefined = options.outputSchema - const requestedResponseFormat = - options.modelOptions != null && 'responseFormat' in options.modelOptions - ? options.modelOptions.responseFormat - : undefined const combinedSchema = combinedOutputSchema && - requestedResponseFormat?.type !== 'json_object' && this.supportsCombinedToolsAndSchema(options.modelOptions) ? this.makeStructuredOutputCompatible( combinedOutputSchema, @@ -1536,13 +1549,17 @@ export class OpenRouterTextAdapter< /** * Combined mode is safe only when this model and every `modelOptions.models` - * fallback are in `OPENROUTER_COMBINED_TOOLS_AND_SCHEMA_MODELS`. - * `:variant` suffixes are routing directives and do not change the gate. + * fallback are in `OPENROUTER_COMBINED_TOOLS_AND_SCHEMA_MODELS`. JSON mode + * is never combined: the schema goes in a system prompt on a separate + * structured-output request, not on the tool-calling turns. */ supportsCombinedToolsAndSchema( modelOptions?: ResolveProviderOptions, ): boolean { - return openRouterSupportsCombinedToolsAndSchema(this.model, modelOptions) + return ( + !openRouterUsesJsonObject(this.model, modelOptions) && + openRouterSupportsCombinedToolsAndSchema(this.model, modelOptions) + ) } /** diff --git a/packages/ai-openrouter/src/internal/combined-tools-and-schema.ts b/packages/ai-openrouter/src/internal/combined-tools-and-schema.ts index 9892a1ff22..c72a432b01 100644 --- a/packages/ai-openrouter/src/internal/combined-tools-and-schema.ts +++ b/packages/ai-openrouter/src/internal/combined-tools-and-schema.ts @@ -1,23 +1,11 @@ import { OPENROUTER_COMBINED_TOOLS_AND_SCHEMA_MODELS } from '../model-meta' - -type OpenRouterCombinedModelOptions = { - models?: ReadonlyArray | undefined -} - -function stripOpenRouterModelVariant(model: string): string { - const variantIndex = model.indexOf(':') - return variantIndex === -1 ? model : model.slice(0, variantIndex) -} +import { openRouterModelIn, openRouterRoutedModels } from './model-ids' export function openRouterSupportsCombinedToolsAndSchema( model: string, - modelOptions?: OpenRouterCombinedModelOptions | undefined, + modelOptions?: unknown, ): boolean { - const candidates = [model, ...(modelOptions?.models ?? [])].map( - stripOpenRouterModelVariant, - ) - - return candidates.every((candidate) => - OPENROUTER_COMBINED_TOOLS_AND_SCHEMA_MODELS.has(candidate), + return openRouterRoutedModels(model, modelOptions).every((id) => + openRouterModelIn(OPENROUTER_COMBINED_TOOLS_AND_SCHEMA_MODELS, id), ) } diff --git a/packages/ai-openrouter/src/internal/json-object-mode.ts b/packages/ai-openrouter/src/internal/json-object-mode.ts new file mode 100644 index 0000000000..2a1b472ce9 --- /dev/null +++ b/packages/ai-openrouter/src/internal/json-object-mode.ts @@ -0,0 +1,45 @@ +import { jsonModeSystemPrompt } from '@tanstack/ai-utils' +import { OPENROUTER_JSON_OBJECT_ONLY_MODELS } from '../model-meta' +import { openRouterModelIn, openRouterRoutedModels } from './model-ids' +import type { TextOptions } from '@tanstack/ai' + +/** + * True when this OpenRouter call must send `json_object`. The caller set it + * in `modelOptions` (`responseFormat` or `text.format`), or a routed model + * (`model` with its `variant`, or `modelOptions.models`) has + * `response_format` and not `structured_outputs` in the catalog. + */ +export function openRouterUsesJsonObject( + model: string, + modelOptions: unknown, +): boolean { + const record = (modelOptions ?? {}) as { + responseFormat?: { type?: unknown } + text?: { format?: { type?: unknown } } + } + if ( + record.responseFormat?.type === 'json_object' || + record.text?.format?.type === 'json_object' + ) { + return true + } + return openRouterRoutedModels(model, modelOptions).some((id) => + openRouterModelIn(OPENROUTER_JSON_OBJECT_ONLY_MODELS, id), + ) +} + +/** Add the schema as a system prompt, because JSON mode does not send it. */ +export function withJsonModePrompt< + TProviderOptions extends Record, +>( + options: TextOptions, + schema: unknown, +): TextOptions { + return { + ...options, + systemPrompts: [ + ...(options.systemPrompts ?? []), + jsonModeSystemPrompt(schema), + ], + } +} diff --git a/packages/ai-openrouter/src/internal/model-ids.ts b/packages/ai-openrouter/src/internal/model-ids.ts new file mode 100644 index 0000000000..9d8497e740 --- /dev/null +++ b/packages/ai-openrouter/src/internal/model-ids.ts @@ -0,0 +1,31 @@ +/** + * The ids a call routes to: `model` with any `modelOptions.variant` suffix, + * then each `modelOptions.models` fallback. + */ +export function openRouterRoutedModels( + model: string, + modelOptions: unknown, +): Array { + const record = (modelOptions ?? {}) as { variant?: unknown; models?: unknown } + const primary = + typeof record.variant === 'string' && record.variant + ? `${model}:${record.variant}` + : model + const fallbacks = Array.isArray(record.models) + ? record.models.filter((id): id is string => typeof id === 'string') + : [] + return [primary, ...fallbacks] +} + +/** + * The catalog lists some `:variant` ids (for example `:free`) as their own + * models. Check the full id first, then the id without its `:variant`. + */ +export function openRouterModelIn( + set: ReadonlySet, + id: string, +): boolean { + if (set.has(id)) return true + const variantIndex = id.indexOf(':') + return variantIndex !== -1 && set.has(id.slice(0, variantIndex)) +} diff --git a/packages/ai-openrouter/src/model-meta.ts b/packages/ai-openrouter/src/model-meta.ts index 198c95bfd7..de2270f0b6 100644 --- a/packages/ai-openrouter/src/model-meta.ts +++ b/packages/ai-openrouter/src/model-meta.ts @@ -23338,3 +23338,45 @@ export const OPENROUTER_COMBINED_TOOLS_AND_SCHEMA_MODELS: ReadonlySet = Z_AI_GLM_5_3_FLASH_BATCH.id, Z_AI_GLM_5_3_BATCH.id, ]) + +/** + * Generated: chat models whose catalog `supported_parameters` include + * `response_format` and do not include `structured_outputs`. + * Structured output sends JSON mode for these ids. + */ +export const OPENROUTER_JSON_OBJECT_ONLY_MODELS: ReadonlySet = + new Set([ + AION_LABS_AION_2_0.id, + AION_LABS_AION_3_0.id, + AION_LABS_AION_3_0_MINI.id, + AION_LABS_AION_3_5.id, + AION_LABS_AION_3_5_MINI.id, + COGNITIVECOMPUTATIONS_DOLPHIN_MISTRAL_24B_VENICE_EDITION.id, + DEEPSEEK_DEEPSEEK_R1.id, + GOOGLE_GEMINI_3_1_FLASH_LITE_IMAGE.id, + GOOGLE_GEMMA_4_26B_A4B_IT_FREE.id, + GOOGLE_GEMMA_4_31B_IT_FREE.id, + IBM_GRANITE_GRANITE_4_0_H_MICRO.id, + INCLUSIONAI_LING_3_0_FLASH.id, + MANCER_WEAVER.id, + MICROSOFT_WIZARDLM_2_8X22B.id, + MINIMAX_MINIMAX_M2_1.id, + MINIMAX_MINIMAX_M2_7.id, + NOUSRESEARCH_HERMES_4_405B.id, + QWEN_QWEN_2_5_CODER_32B_INSTRUCT.id, + QWEN_QWEN3_235B_A22B.id, + QWEN_QWEN3_235B_A22B_THINKING_2507.id, + QWEN_QWEN3_30B_A3B_THINKING_2507.id, + QWEN_QWEN3_8B.id, + QWEN_QWEN3_CODER_FLASH.id, + QWEN_QWEN3_7_FLASH.id, + RELACE_RELACE_SEARCH.id, + UNBIASED_PARETO.id, + Z_AI_GLM_4_5.id, + Z_AI_GLM_4_5V.id, + Z_AI_GLM_4_6V.id, + Z_AI_GLM_5_TURBO.id, + Z_AI_GLM_5_3_FLASHX.id, + Z_AI_GLM_5_3_PRIME.id, + Z_AI_GLM_5V_TURBO.id, + ]) diff --git a/packages/ai-openrouter/tests/openrouter-adapter.test.ts b/packages/ai-openrouter/tests/openrouter-adapter.test.ts index 00b6235e78..fc70af778c 100644 --- a/packages/ai-openrouter/tests/openrouter-adapter.test.ts +++ b/packages/ai-openrouter/tests/openrouter-adapter.test.ts @@ -522,6 +522,54 @@ describe('OpenRouter adapter option mapping', () => { } }) + // finishReason=null closes the call in the post-stream drain. + it.each([ + { raw: '{"n": 1.50}', finishReason: 'tool_calls' }, + { raw: '{"n": 1.50}', finishReason: null }, + { raw: '{"a":', finishReason: 'tool_calls' }, + { raw: '{"a":', finishReason: null }, + ])( + 'sends the raw argument string $raw as args on TOOL_CALL_END (finishReason=$finishReason)', + async ({ raw, finishReason }) => { + setupMockSdkClient([ + { + id: 'chatcmpl-raw', + model: 'openai/gpt-4o-mini', + choices: [ + { + delta: { + toolCalls: [ + { + index: 0, + id: 'call_raw', + type: 'function', + function: { name: 'lookup_weather', arguments: raw }, + }, + ], + }, + finishReason, + }, + ], + }, + ]) + + const chunks: Array = [] + for await (const chunk of createAdapter().chatStream({ + model: 'openai/gpt-4o-mini', + messages: [{ role: 'user', content: 'Weather?' }], + tools: [weatherTool], + logger: testLogger, + })) { + chunks.push(chunk) + } + + const toolEnd = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(toolEnd).toMatchObject({ toolCallId: 'call_raw', args: raw }) + if (raw === '{"a":') expect(toolEnd).not.toHaveProperty('input') + else expect(toolEnd).toMatchObject({ input: { n: 1.5 } }) + }, + ) + it('emits parentMessageId on tool-first tool calls matching the assistant message id', async () => { // Tool call streams before any text content. parentMessageId must bind the // tool call to the same assistant message id the eventual TEXT_MESSAGE_START @@ -1724,6 +1772,79 @@ describe('OpenRouter structured output', () => { expect(rawParams.chatRequest.stream).toBe(true) }) + it('sends json_object and the schema prompt for a response-format-only model', async () => { + setupMockSdkClient([ + { + id: 'c-auto', + model: 'qwen/qwen3.7-flash', + choices: [ + { + delta: { content: '{"name":"Alice","age":30}' }, + finishReason: 'stop', + }, + ], + }, + ]) + const adapter = createOpenRouterText('qwen/qwen3.7-flash', 'test-key') + + const result = await chat({ + adapter, + messages: [{ role: 'user', content: 'Give me a person as json' }], + outputSchema: { + type: 'object', + properties: { + name: { type: 'string' }, + age: { type: 'number' }, + }, + required: ['name', 'age'], + }, + }) + + expect(result).toEqual({ name: 'Alice', age: 30 }) + const structuredCall = mockSend.mock.calls.find( + ([args]: Array) => args.chatRequest.responseFormat, + ) + expect(structuredCall?.[0].chatRequest.responseFormat).toEqual({ + type: 'json_object', + }) + expect(structuredCall?.[0].chatRequest.messages[0]).toEqual({ + role: 'system', + content: expect.stringContaining('"required":["name","age"]'), + }) + }) + + it('keeps json_schema on a strict structured-output model', async () => { + setupMockSdkClient([ + { + id: 'c-strict', + model: 'openai/gpt-4o-mini', + choices: [ + { + delta: { content: '{"name":"Alice","age":30}' }, + finishReason: 'stop', + }, + ], + }, + ]) + const adapter = createAdapter() + + await chat({ + adapter, + messages: [{ role: 'user', content: 'Give me a person as json' }], + outputSchema: { + type: 'object', + properties: { + name: { type: 'string' }, + age: { type: 'number' }, + }, + required: ['name', 'age'], + }, + }) + + const [rawParams] = mockSend.mock.calls[0]! + expect(rawParams.chatRequest.responseFormat.type).toBe('json_schema') + }) + it('respects includeUsage false for structured output streams', async () => { setupMockSdkClient([ { diff --git a/packages/ai-openrouter/tests/openrouter-combined-structured-output.test.ts b/packages/ai-openrouter/tests/openrouter-combined-structured-output.test.ts index 4386a9dc22..390c1996b3 100644 --- a/packages/ai-openrouter/tests/openrouter-combined-structured-output.test.ts +++ b/packages/ai-openrouter/tests/openrouter-combined-structured-output.test.ts @@ -7,6 +7,7 @@ import { OPENROUTER_COMBINED_TOOLS_AND_SCHEMA_MODELS, } from '../src/model-meta' import { openRouterSupportsCombinedToolsAndSchema } from '../src/internal/combined-tools-and-schema' +import { openRouterUsesJsonObject } from '../src/internal/json-object-mode' import { createOpenRouterResponsesText } from '../src/adapters/responses-text' import { createOpenRouterText } from '../src/adapters/text' import type { AdapterYieldChunk, Tool } from '@tanstack/ai' @@ -430,6 +431,110 @@ describe('OpenRouter combined tools + outputSchema', () => { }) }) + describe('JSON mode', () => { + it('finds JSON-only models by full id, variant, and fallback', () => { + // The catalog lists this `:free` id on its own. + expect( + openRouterUsesJsonObject('google/gemma-4-26b-a4b-it:free', undefined), + ).toBe(true) + expect( + openRouterUsesJsonObject('google/gemma-4-26b-a4b-it', { + variant: 'free', + }), + ).toBe(true) + expect(openRouterUsesJsonObject('google/gemma-4-26b-a4b-it', {})).toBe( + false, + ) + expect(openRouterUsesJsonObject('qwen/qwen3.7-flash:nitro', {})).toBe( + true, + ) + expect( + openRouterUsesJsonObject('openai/gpt-4o', { + models: ['qwen/qwen3.7-flash'], + }), + ).toBe(true) + expect(openRouterUsesJsonObject('openai/gpt-4o', {})).toBe(false) + }) + + it('turns combined mode off when the caller sets json_object', () => { + expect( + createOpenRouterText( + 'openai/gpt-4o', + 'k', + ).supportsCombinedToolsAndSchema({ + responseFormat: { type: 'json_object' }, + }), + ).toBe(false) + expect( + createOpenRouterResponsesText( + 'openai/gpt-4o', + 'k', + ).supportsCombinedToolsAndSchema({ + text: { format: { type: 'json_object' } }, + }), + ).toBe(false) + }) + + it('keeps the schema prompt off tool turns and sends json_object on the final call', async () => { + setupChatStream(jsonStopChunks) + const adapter = createOpenRouterText('qwen/qwen3.7-flash', 'k') + const chunks: Array = [] + + for await (const chunk of chat({ + adapter, + messages: [{ role: 'user', content: 'hi' }], + tools, + outputSchema: AnswerSchema, + stream: true, + })) { + chunks.push(chunk) + } + + expect(mockChatSend).toHaveBeenCalledTimes(2) + const [toolTurn, final] = mockChatSend.mock.calls.map( + ([params]) => params.chatRequest, + ) + expect(toolTurn.responseFormat).toBeUndefined() + expect(JSON.stringify(toolTurn.messages)).not.toContain('JSON Schema') + expect(final.responseFormat).toEqual({ type: 'json_object' }) + expect( + final.messages.filter((m: { role: string }) => m.role === 'system'), + ).toEqual([ + { + role: 'system', + content: expect.stringContaining('"required":["answer"]'), + }, + ]) + expect(readCompleteObject(chunks)).toEqual({ answer: 'ok' }) + }) + + it('sends text.format json_object and the schema as instructions on the Responses adapter', async () => { + setupResponsesStream(responsesJsonChunks) + const adapter = createOpenRouterResponsesText('qwen/qwen3.7-flash', 'k') + const chunks: Array = [] + + for await (const chunk of chat({ + adapter, + messages: [{ role: 'user', content: 'hi' }], + tools, + outputSchema: AnswerSchema, + stream: true, + })) { + chunks.push(chunk) + } + + expect(mockResponsesSend).toHaveBeenCalledTimes(2) + const [toolTurn, final] = mockResponsesSend.mock.calls.map( + ([params]) => params.responsesRequest, + ) + expect(toolTurn.text?.format).toBeUndefined() + expect(toolTurn.instructions).toBeUndefined() + expect(final.text?.format).toEqual({ type: 'json_object' }) + expect(final.instructions).toContain('"required":["answer"]') + expect(readCompleteObject(chunks)).toEqual({ answer: 'ok' }) + }) + }) + describe('set integrity', () => { const catalog = new Set(OPENROUTER_CHAT_MODELS) diff --git a/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts b/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts index e02554a996..caeaf5bc9d 100644 --- a/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts +++ b/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts @@ -1353,6 +1353,108 @@ describe('OpenRouter responses adapter — stream event bridge', () => { }) }) + describe('raw tool arguments', () => { + type EndSite = 'arguments.done' | 'output_item.done' | 'response.completed' + + // Each site emits TOOL_CALL_END from a different event. Without a name on + // output_item.added, args.done is deferred to a later backfill. + function rawArgsStream( + site: EndSite, + raw: string, + ): Array> { + const named = { + type: 'function_call', + id: 'item_raw', + callId: 'call_raw', + name: 'lookup_weather', + } + return [ + { + type: 'response.output_item.added', + sequenceNumber: 0, + outputIndex: 0, + item: + site === 'arguments.done' + ? { ...named, arguments: '' } + : { + type: 'function_call', + id: 'item_raw', + callId: 'call_raw', + arguments: '', + }, + }, + { + type: 'response.function_call_arguments.done', + sequenceNumber: 1, + itemId: 'item_raw', + outputIndex: 0, + arguments: raw, + }, + ...(site === 'output_item.done' + ? [ + { + type: 'response.output_item.done', + sequenceNumber: 2, + outputIndex: 0, + item: { ...named, arguments: raw }, + }, + ] + : []), + { + type: 'response.completed', + sequenceNumber: 3, + response: { + model: 'm', + output: [{ ...named, arguments: raw }], + usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 }, + }, + }, + ] + } + + async function toolEnds(site: EndSite, raw: string) { + setupMockSdkClient(rawArgsStream(site, raw)) + const chunks: Array = [] + for await (const chunk of createAdapter().chatStream({ + model: 'openai/gpt-4o-mini', + messages: [{ role: 'user', content: 'hi' }], + tools: [weatherTool], + logger: testLogger, + })) { + chunks.push(chunk) + } + return chunks.filter((chunk) => chunk.type === 'TOOL_CALL_END') + } + + it.each([ + 'arguments.done', + 'output_item.done', + 'response.completed', + ])( + 'sends the raw argument string as args on TOOL_CALL_END (%s)', + async (site) => { + const ends = await toolEnds(site, '{"n": 1.50}') + expect(ends).toHaveLength(1) + expect(ends[0]).toMatchObject({ + toolCallId: 'call_raw', + args: '{"n": 1.50}', + input: { n: 1.5 }, + }) + }, + ) + + it.each([ + 'arguments.done', + 'output_item.done', + 'response.completed', + ])('sends malformed arguments as args with no input (%s)', async (site) => { + const ends = await toolEnds(site, '{"a":') + expect(ends).toHaveLength(1) + expect(ends[0]).toMatchObject({ toolCallId: 'call_raw', args: '{"a":' }) + expect(ends[0]).not.toHaveProperty('input') + }) + }) + it('round-trips distinct Responses item and call IDs through a server tool', async () => { const firstTurn = [ { diff --git a/packages/ai-persistence/src/merge-stored.ts b/packages/ai-persistence/src/merge-stored.ts index 4c8ba03c81..6b29219b98 100644 --- a/packages/ai-persistence/src/merge-stored.ts +++ b/packages/ai-persistence/src/merge-stored.ts @@ -1,10 +1,5 @@ import type { ModelMessage } from '@tanstack/ai' -// Empty incoming keeps stored. Non-empty: the last incoming id that already -// exists in stored is a cutoff (reload drops the old assistant after that -// user). Same id is replaced in place, but a stored mid-conversation record -// stays when the incoming copy has none. New ids and messages with no id are -// appended. /** How many stored messages the merge keeps. Less than `stored.length` on a reload. */ export function storedCutoff( stored: ReadonlyArray, @@ -19,6 +14,12 @@ export function storedCutoff( return stored.length } +// Empty incoming keeps stored. Non-empty: the last incoming id that already +// exists in stored is a cutoff (reload drops the old assistant after that +// user). Same id, or a tool result with the same toolCallId, is replaced in +// place, but a stored mid-conversation record stays when the incoming copy has +// none. An id-less message that matches the prefix row at its index is +// dropped. Other messages are appended. export function mergeStoredMessages( stored: ReadonlyArray, incoming: ReadonlyArray, @@ -55,6 +56,19 @@ export function mergeStoredMessages( if (!message) continue const id = message.id if (id && storedIds.has(id)) continue + // A resumed run posts stored tool results under the interrupt snapshot's + // ids, which the store never saw. Match tool results by toolCallId. + if (message.role === 'tool' && message.toolCallId !== undefined) { + const storedIndex = prefix.findIndex( + (stored) => + stored.role === 'tool' && stored.toolCallId === message.toolCallId, + ) + const stored = prefix[storedIndex] + if (stored) { + merged[storedIndex] = keepStoredRecord(message, stored) + continue + } + } // Attach/reload can post the stored transcript again with no ids. Keep the // prefix row instead of appending a second copy of the same turn. if (!id) { diff --git a/packages/ai-persistence/src/subagent-runs.ts b/packages/ai-persistence/src/subagent-runs.ts index 6842c47040..b05b190183 100644 --- a/packages/ai-persistence/src/subagent-runs.ts +++ b/packages/ai-persistence/src/subagent-runs.ts @@ -505,7 +505,10 @@ export function createSubagentRunRecorder(stores: { initialMessages: modelMessagesToUIMessages(initial), }), status: 'running', - base: resumes ? (marker.base ?? 0) : stored.length, + // No stored base means the cut point is unknown. Keep every message + // that is already stored. `0` would drop them on the next continue. + base: + resumes && marker.base !== undefined ? marker.base : stored.length, // A nested child belongs to the thread of the child that started it. parentThreadId: ownerThread({ threadId: input.threadId, diff --git a/packages/ai-persistence/tests/merge-stored.test.ts b/packages/ai-persistence/tests/merge-stored.test.ts new file mode 100644 index 0000000000..0e4ca54c0a --- /dev/null +++ b/packages/ai-persistence/tests/merge-stored.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, it } from 'vitest' +import { mergeStoredMessages } from '../src/merge-stored' +import type { ModelMessage } from '@tanstack/ai' + +const user: ModelMessage = { id: 'u1', role: 'user', content: 'hi' } +const assistant: ModelMessage = { + id: 'a1', + role: 'assistant', + content: null, + toolCalls: [ + { + id: 'call_1', + type: 'function', + function: { name: 'lookup', arguments: '{}' }, + }, + ], +} +const result = (content: string, id?: string): ModelMessage => ({ + ...(id && { id }), + role: 'tool', + toolCallId: 'call_1', + content, +}) + +describe('mergeStoredMessages: tool results', () => { + it('replaces a stored tool result posted under a new id in place', () => { + const merged = mergeStoredMessages( + [user, assistant, result('old')], + [user, assistant, result('new', 'snapshot_r1_2')], + ) + expect(merged).toEqual([user, assistant, result('new', 'snapshot_r1_2')]) + }) + + it('appends tool results that have no toolCallId', () => { + const stored: ModelMessage = { role: 'tool', content: 'a' } + const incoming: ModelMessage = { id: 'x', role: 'tool', content: 'b' } + expect(mergeStoredMessages([stored], [incoming])).toEqual([ + stored, + incoming, + ]) + }) +}) diff --git a/packages/ai-persistence/tests/message-ids-across-calls.test.ts b/packages/ai-persistence/tests/message-ids-across-calls.test.ts index 259f5e5668..6907a6db25 100644 --- a/packages/ai-persistence/tests/message-ids-across-calls.test.ts +++ b/packages/ai-persistence/tests/message-ids-across-calls.test.ts @@ -119,6 +119,22 @@ function model(withThinking: boolean) { ] } +const toolCall = ( + toolCallId: string, + toolCallName: string, + parentMessageId: string, +): Array => [ + { + type: EventType.TOOL_CALL_START, + toolCallId, + toolCallName, + parentMessageId, + timestamp: t, + }, + { type: EventType.TOOL_CALL_ARGS, toolCallId, delta: '{}', timestamp: t }, + { type: EventType.TOOL_CALL_END, toolCallId, timestamp: t }, +] + const call1Count = (messages: ReadonlyArray) => messages.flatMap((m) => m.toolCalls ?? []).filter((c) => c.id === 'call_1') .length @@ -196,4 +212,69 @@ describe('useChat + withPersistence: message ids across model calls', () => { expect(call1Count(stored)).toBe(1) }) } + + it('stores each tool result once when a run resumes from a client tool', async () => { + const persistence = memoryPersistence() + // Server tool, then a client tool that interrupts the run, then text. + const { adapter, prompts } = mockAdapter([ + run('r1', toolCall('call_1', 'lookup', 'm1'), 'tool_calls'), + run('r1', toolCall('call_2', 'ask', 'm2'), 'tool_calls'), + run('r1', text('m3', 'done'), 'stop'), + ]) + const tools = [ + { name: 'lookup', description: 'Lookup', execute: () => ({ ok: true }) }, + { name: 'ask', description: 'Ask the user' }, + ] + const middleware = [withPersistence(persistence)] + + const processor = new StreamProcessor() + const user1: UIMessage = { + id: 'u1', + role: 'user', + parts: [{ type: 'text', content: 'hi' }], + } + processor.setMessages([user1]) + processor.prepareAssistantMessage() + for await (const chunk of chat({ + adapter, + messages: await serverMessages([user1], 'r1'), + tools, + runId: 'r1', + threadId: 't1', + middleware, + }) as AsyncIterable) { + processor.processChunk(chunk) + } + processor.finalizeStream() + + // The interrupt snapshot gave the stored server tool result a snapshot id. + processor.addToolResult('call_2', { answer: 'yes' }) + for await (const _ of chat({ + adapter, + messages: await serverMessages(processor.getMessages(), 'r1'), + tools, + runId: 'r1', + threadId: 't1', + resume: [ + { + interruptId: 'client_tool_call_2', + status: 'resolved', + payload: { answer: 'yes' }, + }, + ], + middleware, + }) as AsyncIterable) { + // drain + } + + const stored = threadMessages( + await persistence.stores.messages!.loadThread('t1'), + ) + expect( + stored.filter((m) => m.role === 'tool').map((m) => m.toolCallId), + ).toEqual(['call_1', 'call_2']) + expect( + prompts[2]!.filter((m) => m.role === 'tool').map((m) => m.toolCallId), + ).toEqual(['call_1', 'call_2']) + }) }) diff --git a/packages/ai-persistence/tests/subagent-persistence-cards.test.ts b/packages/ai-persistence/tests/subagent-persistence-cards.test.ts index 750edced4f..c00f64814d 100644 --- a/packages/ai-persistence/tests/subagent-persistence-cards.test.ts +++ b/packages/ai-persistence/tests/subagent-persistence-cards.test.ts @@ -823,6 +823,43 @@ describe('subagent run recorder', () => { }) }) + it('keeps stored messages when a suspended child has no base', async () => { + const { stores } = memoryPersistence() + await stores.messages.saveThread('subagent:child', [ + { + id: 'm1', + role: 'assistant', + content: 'Old notes', + metadata: { + tanstack: { subagent: { name: 'writer', status: 'suspended' } }, + 'tanstack:subagentChild': { + parentThreadId: 'desk', + prompt: 'Take notes', + }, + }, + }, + ]) + const recorder = createSubagentRunRecorder({ + messages: stores.messages, + intervalMs: 0, + }) + await feed(recorder, 'run-2', [ + childStarted('child', { parentToolCallId: 'call_2' }), + ...textMessage('m2', 'More notes', { subagentRunId: 'child' }), + { + type: EventType.SUBAGENT_FINISHED, + subagentRunId: 'child', + outcome: { type: 'suspended', interruptIds: ['approval'] }, + timestamp: t, + }, + ]) + + const loaded = await recorder.loadChild('child', { threadId: 'desk' }) + expect(loaded?.messages.map((message) => message.content)).toContain( + 'Old notes', + ) + }) + it('loads nothing for an unknown child', async () => { const { stores } = memoryPersistence() const recorder = createSubagentRunRecorder({ messages: stores.messages }) diff --git a/packages/ai-preact/src/types.ts b/packages/ai-preact/src/types.ts index 3171c45478..6af42c8597 100644 --- a/packages/ai-preact/src/types.ts +++ b/packages/ai-preact/src/types.ts @@ -98,7 +98,7 @@ export interface UseChatReturn< /** * Current messages in the conversation */ - messages: Array> + messages: ReadonlyArray> /** * Send a message and get a response. @@ -114,7 +114,7 @@ export interface UseChatReturn< /** * Pending messages queued while a stream is in flight. */ - queue: Array + queue: ReadonlyArray /** * Cancel a queued message before it drains. No-op if already sent. @@ -215,7 +215,7 @@ export interface UseChatReturn< /** * Set messages manually */ - setMessages: (messages: Array>) => void + setMessages: (messages: ReadonlyArray>) => void /** * Clear all messages diff --git a/packages/ai-preact/src/use-chat.ts b/packages/ai-preact/src/use-chat.ts index 4b235626ae..309f4fb478 100644 --- a/packages/ai-preact/src/use-chat.ts +++ b/packages/ai-preact/src/use-chat.ts @@ -1,21 +1,11 @@ import { ChatClient } from '@tanstack/ai-client' import { createChatDevtoolsBridge } from '@tanstack/ai-client/devtools' -import { - useCallback, - useEffect, - useId, - useMemo, - useRef, - useState, -} from 'preact/hooks' +import { useSyncExternalStore } from 'preact/compat' +import { useCallback, useEffect, useId, useMemo, useRef } from 'preact/hooks' import type { - ChatClientState, ResolvableChatInterrupt, - ChatInterruptState, ChatResumeState, - ConnectionStatus, InferredClientContext, - QueuedMessage, SendMessageOptions, } from '@tanstack/ai-client' import type { @@ -32,9 +22,6 @@ import type { UseChatReturn, } from './types' -const EMPTY_INTERRUPTS = Object.freeze([]) -const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) - export function useChat< const TTools extends ReadonlyArray = any, TContext = InferredClientContext, @@ -50,33 +37,6 @@ export function useChat< const hookId = useId() const clientId = options.threadId ?? hookId - const [messages, setMessages] = useState>>( - options.initialMessages || [], - ) - const [isLoading, setIsLoading] = useState(false) - const [hasOlderMessages, setHasOlderMessages] = useState(false) - const [error, setError] = useState(undefined) - const [status, setStatus] = useState('ready') - const [isSubscribed, setIsSubscribed] = useState(false) - const [connectionStatus, setConnectionStatus] = - useState('disconnected') - const [sessionGenerating, setSessionGenerating] = useState(false) - const [queue, setQueue] = useState>([]) - const [runId, setRunId] = useState(null) - const [interruptState, setInterruptState] = useState< - ChatInterruptState - >(() => ({ - interrupts: EMPTY_INTERRUPTS, - pendingInterrupts: EMPTY_INTERRUPTS, - interruptErrors: EMPTY_INTERRUPT_ERRORS, - resuming: false, - })) - - // Track current messages in a ref to preserve them when client is recreated - const messagesRef = useRef>>( - options.initialMessages || [], - ) - const isFirstMountRef = useRef(true) const activeClientRef = useRef(null) const cleanupInvalidationRef = useRef | null>( null, @@ -90,19 +50,8 @@ export function useChat< optionsRef.current = options - const syncResumeState = useCallback((target: ChatClient | null) => { - if (!target) return - setRunId(target.getCurrentRunId()) - setInterruptState(target.getInterruptState()) - }, []) - - useEffect(() => { - messagesRef.current = messages - }, [messages]) - const { client, initialization } = useMemo(() => { const messagesToUse = options.initialMessages || [] - isFirstMountRef.current = false // Build options with conditional spreads for fields whose source // type is `T | undefined` but the ChatClient target uses a strict @@ -219,77 +168,11 @@ export function useChat< ...(options.streamProcessor !== undefined && { streamProcessor: options.streamProcessor, }), - onMessagesChange: (newMessages: Array>) => { - runOrQueueForActiveInstance(() => { - setMessages(newMessages) - const currentInstance = getActiveInstance() - if (!currentInstance) return - setHasOlderMessages(currentInstance.getHasOlderMessages()) - }) - }, - onLoadingChange: (newIsLoading: boolean) => { - runOrQueueForActiveInstance(() => { - const currentInstance = getActiveInstance() - if (!currentInstance) return - setIsLoading(newIsLoading) - syncResumeState(currentInstance) - }) - }, - onStatusChange: (newStatus: ChatClientState) => { - runOrQueueForActiveInstance(() => { - setStatus(newStatus) - }) - }, - onErrorChange: (newError: Error | undefined) => { - runOrQueueForActiveInstance(() => { - setError(newError) - }) - }, - onSubscriptionChange: (nextIsSubscribed: boolean) => { - runOrQueueForActiveInstance(() => { - setIsSubscribed(nextIsSubscribed) - }) - }, - onConnectionStatusChange: (nextStatus: ConnectionStatus) => { - runOrQueueForActiveInstance(() => { - setConnectionStatus(nextStatus) - }) - }, - onSessionGeneratingChange: (isGenerating: boolean) => { - runOrQueueForActiveInstance(() => { - setSessionGenerating(isGenerating) - }) - }, - onHydratingChange: (nextIsHydrating: boolean) => { - runOrQueueForActiveInstance(() => { - setHydration({ client: instance, value: nextIsHydrating }) - }) - }, ...(optionsRef.current.queue !== undefined && { queue: optionsRef.current.queue, }), - onQueueChange: (nextQueue: Array) => { - runOrQueueForActiveInstance(() => { - setQueue(nextQueue) - }) - }, - onRunIdChange: (nextRunId) => { - runOrQueueForActiveInstance(() => { - setRunId(nextRunId) - }) - }, - onResumeStateChange: (_nextResumeState, nextPendingInterrupts) => { - runOrQueueForActiveInstance(() => { - setInterruptState((current) => ({ - ...current, - interrupts: nextPendingInterrupts, - pendingInterrupts: nextPendingInterrupts, - })) - }) - }, onInterruptStateChange: (nextInterruptState, context) => { runOrQueueForActiveInstance(() => { - setInterruptState(nextInterruptState) optionsRef.current.onInterruptStateChange?.( nextInterruptState, context, @@ -299,18 +182,7 @@ export function useChat< }) instanceHolder.current = instance return { client: instance, initialization: initializationState } - }, [clientId, syncResumeState]) - - // Seeded from the client so the first render already shows hydrating, not an - // empty chat. The constructor sets it without firing a callback. Tagged with - // its client: after a thread switch, the new client's value shows on the - // first render, before the effects sync it. - const [hydration, setHydration] = useState(() => ({ - client, - value: client.getIsHydrating(), - })) - const isHydrating = - hydration.client === client ? hydration.value : client.getIsHydrating() + }, [clientId]) useEffect(() => { activeClientRef.current = client @@ -330,14 +202,10 @@ export function useChat< } }, [client, initialization]) - useEffect(() => { - const clientMessages = client.getMessages() - if (clientMessages !== messagesRef.current) { - setMessages(clientMessages) - } - setHasOlderMessages(client.getHasOlderMessages()) - setHydration({ client, value: client.getIsHydrating() }) - }, [client]) + const snapshot = useSyncExternalStore( + client.subscribeSnapshot, + client.getSnapshot, + ) // Sync body / forwardedProps changes to the client. // Both populate the same wire payload; `forwardedProps` is preferred @@ -405,10 +273,6 @@ export function useChat< cleanupInvalidationRef.current = null } client.mountDevtools() - // Delivery-durability resume is transparent: the resumable SSE connection - // adapter reattaches via the browser's native Last-Event-ID on reconnect. - // We only seed interrupt (state) resume from the client here. - syncResumeState(client) return () => { cleanupInvalidationRef.current = setTimeout(() => { @@ -437,7 +301,7 @@ export function useChat< } cleanupDisposalRef.current = disposal } - }, [client, syncResumeState]) + }, [client]) // All callback options are read through optionsRef at call time, so fresh // closures from each render are picked up without recreating the client. @@ -446,44 +310,31 @@ export function useChat< content: string | MultimodalContent, sendOptions?: SendMessageOptions, ) => { - try { - await client.sendMessage(content, undefined, sendOptions) - } finally { - syncResumeState(client) - } + await client.sendMessage(content, undefined, sendOptions) }, - [client, syncResumeState], + [client], ) const cancelQueued = useCallback( (id: string) => { client.cancelQueued(id) }, - [client, syncResumeState], + [client], ) const append = useCallback( async (message: ModelMessage | UIMessage) => { - try { - await client.append(message) - } finally { - syncResumeState(client) - } + await client.append(message) }, - [client, syncResumeState], + [client], ) const reload = useCallback(async () => { - try { - await client.reload() - } finally { - syncResumeState(client) - } - }, [client, syncResumeState]) + await client.reload() + }, [client]) const loadOlderMessages = useCallback(async () => { await client.loadOlderMessages() - setHasOlderMessages(client.getHasOlderMessages()) }, [client]) const stop = useCallback(() => { @@ -492,11 +343,10 @@ export function useChat< const clear = useCallback(() => { client.clear() - syncResumeState(client) - }, [client, syncResumeState]) + }, [client]) const setMessagesManually = useCallback( - (newMessages: Array>) => { + (newMessages: ReadonlyArray>) => { client.setMessagesManually(newMessages) }, [client], @@ -518,18 +368,15 @@ export function useChat< const addToolApprovalResponse = useCallback( async (response: { id: string; approved: boolean }) => { await client.addToolApprovalResponse(response) - syncResumeState(client) }, - [client, syncResumeState], + [client], ) const resumeInterrupts = useCallback( async (resumeItems: Array, state?: ChatResumeState) => { - const result = await client.resumeInterrupts(resumeItems, state) - syncResumeState(client) - return result + return await client.resumeInterrupts(resumeItems, state) }, - [client, syncResumeState], + [client], ) const resolveInterrupts = useCallback( @@ -563,7 +410,7 @@ export function useChat< [client], ) - const renderedMessages = client.getMessages() + const renderedMessages = snapshot.messages return { messages: renderedMessages, @@ -571,26 +418,26 @@ export function useChat< append, reload, stop, - isLoading, - hasOlderMessages, + isLoading: snapshot.isLoading, + hasOlderMessages: snapshot.hasOlderMessages, loadOlderMessages, - error, - status, - isSubscribed, - connectionStatus, - sessionGenerating, - isHydrating, + error: snapshot.error, + status: snapshot.status, + isSubscribed: snapshot.isSubscribed, + connectionStatus: snapshot.connectionStatus, + sessionGenerating: snapshot.sessionGenerating, + isHydrating: snapshot.isHydrating, setMessages: setMessagesManually, clear, addToolResult, addToolApprovalResponse, - queue, + queue: snapshot.queue, cancelQueued, - runId, - interrupts: interruptState.interrupts, - pendingInterrupts: interruptState.pendingInterrupts, - interruptErrors: interruptState.interruptErrors, - resuming: interruptState.resuming, + runId: snapshot.runId, + interrupts: snapshot.interruptState.interrupts, + pendingInterrupts: snapshot.interruptState.pendingInterrupts, + interruptErrors: snapshot.interruptState.interruptErrors, + resuming: snapshot.interruptState.resuming, resolveInterrupts, cancelInterrupts, retryInterrupts, diff --git a/packages/ai-preact/tests/chat-ui/create-ui.test.tsx b/packages/ai-preact/tests/chat-ui/create-ui.test.tsx index 96249626d6..1443072611 100644 --- a/packages/ai-preact/tests/chat-ui/create-ui.test.tsx +++ b/packages/ai-preact/tests/chat-ui/create-ui.test.tsx @@ -120,9 +120,9 @@ describe('createChatHook', () => { function Screen() { const chat = useAppChat({ - initialMessages: [messageWithToolResults] as ChatUIHost< - typeof chatOptions - >['messages'], + initialMessages: [messageWithToolResults] as Array< + ChatUIHost['messages'][number] + >, }) return } diff --git a/packages/ai-preact/tests/use-chat.test.ts b/packages/ai-preact/tests/use-chat.test.ts index 625fd667e7..1babc4229d 100644 --- a/packages/ai-preact/tests/use-chat.test.ts +++ b/packages/ai-preact/tests/use-chat.test.ts @@ -213,6 +213,43 @@ describe('useChat', () => { expect(result.current.messages).toEqual(initialMessages) }) + it('throws on in-place messages.push and accepts setMessages', () => { + const { result } = renderUseChat({ + connection: createMockConnectionAdapter(), + initialMessages: [ + { + id: 'msg-1', + role: 'user', + parts: [{ type: 'text', content: 'Hello' }], + createdAt: new Date(), + }, + ], + }) + + expect(Object.isFrozen(result.current.messages)).toBe(true) + expect(() => { + Array.prototype.push.call(result.current.messages, { + id: 'x', + role: 'user', + parts: [{ type: 'text', content: 'nope' }], + createdAt: new Date(), + }) + }).toThrow(TypeError) + + act(() => { + result.current.setMessages([ + ...result.current.messages, + { + id: 'ok', + role: 'user', + parts: [{ type: 'text', content: 'ok' }], + createdAt: new Date(), + }, + ]) + }) + expect(result.current.messages).toHaveLength(2) + }) + it('should initialize with persisted messages', async () => { const adapter = createMockConnectionAdapter() const persistedMessages: Array = [ diff --git a/packages/ai-react/src/realtime-types.ts b/packages/ai-react/src/realtime-types.ts index e581502370..675dfd6835 100644 --- a/packages/ai-react/src/realtime-types.ts +++ b/packages/ai-react/src/realtime-types.ts @@ -104,7 +104,7 @@ export interface UseRealtimeChatReturn { /** Current mode (idle, listening, thinking, speaking) */ mode: RealtimeMode /** Conversation messages */ - messages: Array + messages: ReadonlyArray /** User transcript while speaking (before finalized) */ pendingUserTranscript: string | null /** Assistant transcript while speaking (before finalized) */ diff --git a/packages/ai-react/src/types.ts b/packages/ai-react/src/types.ts index 99b63abc5b..1b95c84063 100644 --- a/packages/ai-react/src/types.ts +++ b/packages/ai-react/src/types.ts @@ -177,7 +177,7 @@ interface BaseUseChatReturn< * When `subagents` is supplied, a `type: 'subagent'` part narrows on * `subagent.name`, and that child's `messages` use the agent's tools. */ - messages: Array> + messages: ReadonlyArray> /** * Live child-agent invocations, nested cards included. Each entry is the @@ -203,7 +203,7 @@ interface BaseUseChatReturn< * Pending messages queued while the client is busy (streaming, claiming a * send, or draining). Separate from `messages` until they drain. */ - queue: Array + queue: ReadonlyArray /** * Cancel a queued message before it drains. No-op if already sent. @@ -335,7 +335,7 @@ interface BaseUseChatReturn< /** * Set messages manually */ - setMessages: (messages: Array>) => void + setMessages: (messages: ReadonlyArray>) => void /** * Clear all messages diff --git a/packages/ai-react/src/use-chat.ts b/packages/ai-react/src/use-chat.ts index 7cd04c5120..7d10da9b30 100644 --- a/packages/ai-react/src/use-chat.ts +++ b/packages/ai-react/src/use-chat.ts @@ -11,13 +11,10 @@ import type { StreamChunk, } from '@tanstack/ai/client' import type { - ChatClientState, + ChatClientPersistence, ResolvableChatInterrupt, - ChatInterruptState, ChatResumeState, - ConnectionStatus, InferredClientContext, - QueuedMessage, SendMessageOptions, StructuredOutputPart, } from '@tanstack/ai-client' @@ -30,9 +27,6 @@ import type { UseChatReturn, } from './types' -const EMPTY_INTERRUPTS = Object.freeze([]) -const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) - export function useChat< const TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, @@ -51,36 +45,9 @@ export function useChat< const hookId = useId() const clientId = options.threadId ?? hookId - const [messages, setMessages] = useState>>( - options.initialMessages || [], - ) - const [isLoading, setIsLoading] = useState(false) - const [hasOlderMessages, setHasOlderMessages] = useState(false) - const [error, setError] = useState(undefined) - const [status, setStatus] = useState('ready') - const [isSubscribed, setIsSubscribed] = useState(false) - const [connectionStatus, setConnectionStatus] = - useState('disconnected') - const [sessionGenerating, setSessionGenerating] = useState(false) - const [queue, setQueue] = useState>([]) - const [runId, setRunId] = useState(null) - const [interruptState, setInterruptState] = useState< - ChatInterruptState - >(() => ({ - interrupts: EMPTY_INTERRUPTS, - pendingInterrupts: EMPTY_INTERRUPTS, - interruptErrors: EMPTY_INTERRUPT_ERRORS, - resuming: false, - })) - type Partial = DeepPartial>> type Final = InferSchemaType> - // Track current messages in a ref to preserve them when client is recreated - const messagesRef = useRef>>( - options.initialMessages || [], - ) - const isFirstMountRef = useRef(true) const subscribedRef = useRef(false) const activeClientRef = useRef(null) const cleanupInvalidationRef = useRef | null>( @@ -91,9 +58,6 @@ export function useChat< timeout: ReturnType } | null>(null) - // Update ref synchronously during render so it's always current when useMemo runs. - messagesRef.current = messages - // Track current options in a ref to avoid recreating client when options change const optionsRef = useRef>( @@ -101,222 +65,141 @@ export function useChat< ) optionsRef.current = options - const syncResumeState = useCallback((target: ChatClient | null) => { - if (!target) return - setRunId(target.getCurrentRunId()) - setInterruptState(target.getInterruptState()) - }, []) - - // Create ChatClient instance with callbacks to sync state - const { client, initialization } = useMemo(() => { - const messagesToUse = options.initialMessages || [] - isFirstMountRef.current = false - - // Build options with conditional spreads for fields whose source - // type is `T | undefined` but the ChatClient target uses a strict - // optional (`field?: T`) — `exactOptionalPropertyTypes` rejects - // assigning `undefined` to those, so we omit the key when absent. - const initialOptions = optionsRef.current - const transport = initialOptions.connection - ? { connection: initialOptions.connection } - : { fetcher: initialOptions.fetcher } - - const instanceHolder: { - current: ChatClient | undefined - } = { current: undefined } - const getActiveInstance = () => { - const currentInstance = instanceHolder.current - if (!currentInstance || activeClientRef.current !== currentInstance) { - return undefined + // Create ChatClient once per clientId. User callbacks read latest options. + const { client, initialization, serverSnapshot, startPersistenceHydration } = + useMemo(() => { + const messagesToUse = options.initialMessages || [] + + const initialOptions = optionsRef.current + const { persistence, release: releasePersistenceHydration } = + gatePersistence(initialOptions.persistence) + const transport = initialOptions.connection + ? { connection: initialOptions.connection } + : { fetcher: initialOptions.fetcher } + + const instanceHolder: { + current: ChatClient | undefined + } = { current: undefined } + const getActiveInstance = () => { + const currentInstance = instanceHolder.current + if (!currentInstance || activeClientRef.current !== currentInstance) { + return undefined + } + return currentInstance } - return currentInstance - } - // ChatClient may publish while its constructor is running or while async - // persistence resolves before commit. Preserve those exact notifications - // until this render commits; invoking them here would run state setters and - // user callbacks for a client React may abandon. - const initializationState = { - ready: false, - callbacks: [] as Array<() => void>, - } - const runOrQueueForActiveInstance = (callback: () => void) => { - if (!initializationState.ready) { - initializationState.callbacks.push(callback) - return + // Queue user callbacks until this instance is the committed client. + // Invoking them here would run onChunk/onFinish/onError for a client + // React may abandon. Browser persistence is gated until attach so the + // first paint stays on initialMessages / initialResumeSnapshot. + const initializationState = { + ready: false, + callbacks: [] as Array<() => void>, } - const currentInstance = instanceHolder.current - if (!currentInstance || activeClientRef.current !== currentInstance) - return - callback() - } - const instance = new ChatClient({ - devtoolsBridgeFactory: createChatDevtoolsBridge, - ...transport, - initialMessages: messagesToUse, - ...(typeof initialOptions.threadId === 'string' && - initialOptions.persistence === true - ? { - persistence: true, - threadId: initialOptions.threadId, - ...(initialOptions.history !== undefined && { - history: initialOptions.history, - }), - } - : typeof initialOptions.threadId === 'string' && - initialOptions.persistence + const runOrQueueForActiveInstance = (callback: () => void) => { + if (!initializationState.ready) { + initializationState.callbacks.push(callback) + return + } + const currentInstance = instanceHolder.current + if (!currentInstance || activeClientRef.current !== currentInstance) + return + callback() + } + const instance = new ChatClient({ + devtoolsBridgeFactory: createChatDevtoolsBridge, + ...transport, + initialMessages: messagesToUse, + ...(typeof initialOptions.threadId === 'string' && persistence === true ? { - persistence: initialOptions.persistence, + persistence: true, threadId: initialOptions.threadId, + ...(initialOptions.history !== undefined && { + history: initialOptions.history, + }), } - : { - ...(initialOptions.threadId !== undefined && { + : typeof initialOptions.threadId === 'string' && persistence + ? { + persistence, threadId: initialOptions.threadId, + } + : { + ...(initialOptions.threadId !== undefined && { + threadId: initialOptions.threadId, + }), }), - }), - ...(initialOptions.body !== undefined && { body: initialOptions.body }), - ...(initialOptions.forwardedProps !== undefined && { - forwardedProps: initialOptions.forwardedProps, - }), - ...(initialOptions.byok !== undefined && { byok: initialOptions.byok }), - byokProvider: () => optionsRef.current.byokProvider?.(), - ...(initialOptions.initialResumeSnapshot !== undefined && { - initialResumeSnapshot: initialOptions.initialResumeSnapshot, - }), - ...(initialOptions.context !== undefined && { - context: initialOptions.context, - }), - devtools: { - ...initialOptions.devtools, - framework: 'react', - hookName: 'useChat', - outputKind: initialOptions.outputSchema ? 'structured' : 'chat', - }, - onResponse: (response) => { - // ChatClient awaits this return value. Queuing would drop the promise. - if (!getActiveInstance()) return - return optionsRef.current.onResponse?.(response) - }, - onChunk: (chunk: StreamChunk) => { - runOrQueueForActiveInstance(() => { - optionsRef.current.onChunk?.(chunk) - }) - }, - onFinish: (message: UIMessage) => { - runOrQueueForActiveInstance(() => { - optionsRef.current.onFinish?.(message) - }) - }, - onError: (error: Error) => { - runOrQueueForActiveInstance(() => { - optionsRef.current.onError?.(error) - }) - }, - ...(initialOptions.tools !== undefined && { - tools: initialOptions.tools, - }), - ...(initialOptions.interrupts !== undefined && { - interrupts: initialOptions.interrupts, - }), - onCustomEvent: (eventType, data, context) => { - runOrQueueForActiveInstance(() => { - optionsRef.current.onCustomEvent?.(eventType, data, context) - }) - }, - ...(options.streamProcessor !== undefined && { - streamProcessor: options.streamProcessor, - }), - onMessagesChange: (newMessages: Array>) => { - runOrQueueForActiveInstance(() => { - setMessages(newMessages) - const currentInstance = getActiveInstance() - if (!currentInstance) return - setHasOlderMessages(currentInstance.getHasOlderMessages()) - }) - }, - onLoadingChange: (newIsLoading: boolean) => { - runOrQueueForActiveInstance(() => { - const currentInstance = getActiveInstance() - if (!currentInstance) return - setIsLoading(newIsLoading) - syncResumeState(currentInstance) - }) - }, - onErrorChange: (newError: Error | undefined) => { - runOrQueueForActiveInstance(() => { - setError(newError) - }) - }, - onStatusChange: (status: ChatClientState) => { - runOrQueueForActiveInstance(() => { - setStatus(status) - }) - }, - onSubscriptionChange: (nextIsSubscribed: boolean) => { - runOrQueueForActiveInstance(() => { - setIsSubscribed(nextIsSubscribed) - }) - }, - onConnectionStatusChange: (nextStatus: ConnectionStatus) => { - runOrQueueForActiveInstance(() => { - setConnectionStatus(nextStatus) - }) - }, - onSessionGeneratingChange: (isGenerating: boolean) => { - runOrQueueForActiveInstance(() => { - setSessionGenerating(isGenerating) - }) - }, - onHydratingChange: (nextIsHydrating: boolean) => { - runOrQueueForActiveInstance(() => { - setHydration({ client: instance, value: nextIsHydrating }) - }) - }, - ...(optionsRef.current.queue !== undefined && { - queue: optionsRef.current.queue, - }), - onQueueChange: (nextQueue: Array) => { - runOrQueueForActiveInstance(() => { - setQueue(nextQueue) - }) - }, - onRunIdChange: (nextRunId) => { - runOrQueueForActiveInstance(() => { - setRunId(nextRunId) - }) - }, - onResumeStateChange: (_nextResumeState, nextPendingInterrupts) => { - runOrQueueForActiveInstance(() => { - setInterruptState((current) => ({ - ...current, - interrupts: nextPendingInterrupts, - pendingInterrupts: nextPendingInterrupts, - })) - }) - }, - onInterruptStateChange: (nextInterruptState, context) => { - runOrQueueForActiveInstance(() => { - setInterruptState(nextInterruptState) - optionsRef.current.onInterruptStateChange?.( - nextInterruptState, - context, - ) - }) - }, - }) - instanceHolder.current = instance - return { client: instance, initialization: initializationState } - }, [clientId, syncResumeState]) - - // Seeded from the client so the first render already shows hydrating, not an - // empty chat. The constructor sets it without firing a callback. Tagged with - // its client: after a thread switch, the new client's value shows on the - // first render, before the effects sync it. - const [hydration, setHydration] = useState(() => ({ - client, - value: client.getIsHydrating(), - })) - const isHydrating = - hydration.client === client ? hydration.value : client.getIsHydrating() + ...(initialOptions.body !== undefined && { body: initialOptions.body }), + ...(initialOptions.forwardedProps !== undefined && { + forwardedProps: initialOptions.forwardedProps, + }), + ...(initialOptions.byok !== undefined && { byok: initialOptions.byok }), + byokProvider: () => optionsRef.current.byokProvider?.(), + ...(initialOptions.initialResumeSnapshot !== undefined && { + initialResumeSnapshot: initialOptions.initialResumeSnapshot, + }), + ...(initialOptions.context !== undefined && { + context: initialOptions.context, + }), + devtools: { + ...initialOptions.devtools, + framework: 'react', + hookName: 'useChat', + outputKind: initialOptions.outputSchema ? 'structured' : 'chat', + }, + onResponse: (response) => { + // ChatClient awaits this return value. Queuing would drop the promise. + if (!getActiveInstance()) return + return optionsRef.current.onResponse?.(response) + }, + onChunk: (chunk: StreamChunk) => { + runOrQueueForActiveInstance(() => { + optionsRef.current.onChunk?.(chunk) + }) + }, + onFinish: (message: UIMessage) => { + runOrQueueForActiveInstance(() => { + optionsRef.current.onFinish?.(message) + }) + }, + onError: (error: Error) => { + runOrQueueForActiveInstance(() => { + optionsRef.current.onError?.(error) + }) + }, + ...(initialOptions.tools !== undefined && { + tools: initialOptions.tools, + }), + ...(initialOptions.interrupts !== undefined && { + interrupts: initialOptions.interrupts, + }), + onCustomEvent: (eventType, data, context) => { + runOrQueueForActiveInstance(() => { + optionsRef.current.onCustomEvent?.(eventType, data, context) + }) + }, + ...(options.streamProcessor !== undefined && { + streamProcessor: options.streamProcessor, + }), + ...(optionsRef.current.queue !== undefined && { + queue: optionsRef.current.queue, + }), + onInterruptStateChange: (nextInterruptState, context) => { + runOrQueueForActiveInstance(() => { + optionsRef.current.onInterruptStateChange?.( + nextInterruptState, + context, + ) + }) + }, + }) + instanceHolder.current = instance + const capturedServerSnapshot = instance.getSnapshot() + return { + client: instance, + initialization: initializationState, + serverSnapshot: capturedServerSnapshot, + startPersistenceHydration: releasePersistenceHydration, + } + }, [clientId]) useEffect(() => { activeClientRef.current = client @@ -336,14 +219,23 @@ export function useChat< } }, [client, initialization]) + // Plain state, not useSyncExternalStore: that renders synchronously on every + // change, so a buffered stream rendered once per chunk. setState batches the + // changes into one render per task. + const [held, setHeld] = useState(() => ({ client, snapshot: serverSnapshot })) useEffect(() => { - const clientMessages = client.getMessages() - if (clientMessages !== messagesRef.current) { - setMessages(clientMessages) + const sync = () => { + const next = client.getSnapshot() + setHeld((prev) => + prev.client === client && prev.snapshot === next + ? prev + : { client, snapshot: next }, + ) } - setHasOlderMessages(client.getHasOlderMessages()) - setHydration({ client, value: client.getIsHydrating() }) + sync() + return client.subscribeSnapshot(sync) }, [client]) + const snapshot = held.client === client ? held.snapshot : serverSnapshot // Sync each wire-payload slot in its own effect so an unrelated option // changing doesn't re-run the others. `updateOptions` declares strict-optional @@ -407,10 +299,11 @@ export function useChat< // run is picked straight back up from the durable log. useEffect(() => { client.attach() + startPersistenceHydration() return () => { client.detach() } - }, [client]) + }, [client, startPersistenceHydration]) useEffect(() => { if (cleanupDisposalRef.current?.client === client) { @@ -422,10 +315,6 @@ export function useChat< cleanupInvalidationRef.current = null } client.mountDevtools() - // Delivery-durability resume is transparent: the resumable SSE connection - // adapter re-attaches via the browser's native Last-Event-ID on reconnect. - // We only seed interrupt (state) resume from the client here. - syncResumeState(client) return () => { cleanupInvalidationRef.current = setTimeout(() => { @@ -458,51 +347,38 @@ export function useChat< } cleanupDisposalRef.current = disposal } - }, [client, syncResumeState]) + }, [client]) const sendMessage = useCallback( async ( content: string | MultimodalContent, sendOptions?: SendMessageOptions, ) => { - try { - await client.sendMessage(content, undefined, sendOptions) - } finally { - syncResumeState(client) - } + await client.sendMessage(content, undefined, sendOptions) }, - [client, syncResumeState], + [client], ) const cancelQueued = useCallback( (id: string) => { client.cancelQueued(id) }, - [client, syncResumeState], + [client], ) const append = useCallback( async (message: ModelMessage | UIMessage) => { - try { - await client.append(message) - } finally { - syncResumeState(client) - } + await client.append(message) }, - [client, syncResumeState], + [client], ) const reload = useCallback(async () => { - try { - await client.reload() - } finally { - syncResumeState(client) - } - }, [client, syncResumeState]) + await client.reload() + }, [client]) const loadOlderMessages = useCallback(async () => { await client.loadOlderMessages() - setHasOlderMessages(client.getHasOlderMessages()) }, [client]) const stop = useCallback(() => { @@ -511,11 +387,10 @@ export function useChat< const clear = useCallback(() => { client.clear() - syncResumeState(client) - }, [client, syncResumeState]) + }, [client]) const setMessagesManually = useCallback( - (newMessages: Array>) => { + (newMessages: ReadonlyArray>) => { client.setMessagesManually(newMessages) }, [client], @@ -537,21 +412,15 @@ export function useChat< const addToolApprovalResponse = useCallback( async (response: { id: string; approved: boolean }) => { await client.addToolApprovalResponse(response) - syncResumeState(client) }, - [client, syncResumeState], + [client], ) const resumeInterrupts = useCallback( async (resumeItems: Array, state?: ChatResumeState) => { - const result = await client.resumeInterrupts( - resumeItems, - state ?? undefined, - ) - syncResumeState(client) - return result + return await client.resumeInterrupts(resumeItems, state ?? undefined) }, - [client, syncResumeState], + [client], ) const resolveInterrupts = useCallback( @@ -595,7 +464,7 @@ export function useChat< // a stale assistant turn or a system prompt) we deliberately return null // rather than scanning historical assistants — otherwise a `final` from a // previous session would leak into the hook value on first render. - const renderedMessages = client.getMessages() + const renderedMessages = snapshot.messages const activeStructuredPart = useMemo(() => { let lastUserIndex = -1 @@ -643,26 +512,26 @@ export function useChat< append, reload, stop, - isLoading, - hasOlderMessages, + isLoading: snapshot.isLoading, + hasOlderMessages: snapshot.hasOlderMessages, loadOlderMessages, - error, - status, - isSubscribed, - connectionStatus, - sessionGenerating, - isHydrating, + error: snapshot.error, + status: snapshot.status, + isSubscribed: snapshot.isSubscribed, + connectionStatus: snapshot.connectionStatus, + sessionGenerating: snapshot.sessionGenerating, + isHydrating: snapshot.isHydrating, setMessages: setMessagesManually, clear, addToolResult, addToolApprovalResponse, - queue, + queue: snapshot.queue, cancelQueued, - runId, - interrupts: interruptState.interrupts, - pendingInterrupts: interruptState.pendingInterrupts, - interruptErrors: interruptState.interruptErrors, - resuming: interruptState.resuming, + runId: snapshot.runId, + interrupts: snapshot.interruptState.interrupts, + pendingInterrupts: snapshot.interruptState.pendingInterrupts, + interruptErrors: snapshot.interruptState.interruptErrors, + resuming: snapshot.interruptState.resuming, resolveInterrupts, cancelInterrupts, retryInterrupts, @@ -672,3 +541,47 @@ export function useChat< final, } as unknown as UseChatReturn } + +/** + * Hold a browser store's first read until `release()`. React's first commit + * then matches the server snapshot (`initialMessages`), and the stored + * transcript applies after the hook attaches. + */ +function gatePersistence>( + source: boolean | ChatClientPersistence | undefined, +): { + persistence: boolean | ChatClientPersistence | undefined + release: () => void +} { + if (!source || source === true) { + return { persistence: source, release: () => {} } + } + const store = source + let released = false + let startRead: (() => void) | undefined + let read: ReturnType['getItem']> | undefined + return { + persistence: { + getItem(id) { + read ??= new Promise((resolve, reject) => { + startRead = () => { + startRead = undefined + try { + resolve(store.getItem(id)) + } catch (error) { + reject(error) + } + } + }) + if (released) startRead?.() + return read + }, + setItem: (id, state) => store.setItem(id, state), + removeItem: (id) => store.removeItem(id), + }, + release: () => { + released = true + startRead?.() + }, + } +} diff --git a/packages/ai-react/src/use-generate-video.ts b/packages/ai-react/src/use-generate-video.ts index 708281fee2..393c17f1ab 100644 --- a/packages/ai-react/src/use-generate-video.ts +++ b/packages/ai-react/src/use-generate-video.ts @@ -1,6 +1,13 @@ import { VideoGenerationClient } from '@tanstack/ai-client' import { createVideoDevtoolsBridge } from '@tanstack/ai-client/devtools' -import { useCallback, useEffect, useId, useMemo, useRef, useState } from 'react' +import { + useCallback, + useEffect, + useId, + useMemo, + useRef, + useSyncExternalStore, +} from 'react' import type { StreamChunk } from '@tanstack/ai' import type { AIDevtoolsDisplayOptions, @@ -175,14 +182,6 @@ export function useGenerateVideo( // The hook identity is `threadId`. `hookId` is only a React recreation key. const clientIdentity = options.threadId ?? hookId - const [result, setResult] = useState(null) - const [jobId, setJobId] = useState(null) - const [videoStatus, setVideoStatus] = useState(null) - const [isLoading, setIsLoading] = useState(false) - const [error, setError] = useState(undefined) - const [status, setStatus] = useState('idle') - const [runId, setRunId] = useState(null) - const optionsRef = useRef(options) optionsRef.current = options const disposedRef = useRef(false) @@ -235,27 +234,6 @@ export function useGenerateVideo( onStatusUpdate: (s: VideoStatusInfo) => { if (!disposedRef.current) optionsRef.current.onStatusUpdate?.(s) }, - onResultChange: (r: TOutput | null) => { - if (!disposedRef.current) setResult(r) - }, - onLoadingChange: (l: boolean) => { - if (!disposedRef.current) setIsLoading(l) - }, - onErrorChange: (e: Error | undefined) => { - if (!disposedRef.current) setError(e) - }, - onStatusChange: (s: GenerationClientState) => { - if (!disposedRef.current) setStatus(s) - }, - onJobIdChange: (id: string | null) => { - if (!disposedRef.current) setJobId(id) - }, - onVideoStatusChange: (s: VideoStatusInfo | null) => { - if (!disposedRef.current) setVideoStatus(s) - }, - onResumeStateChange: (rs: { runId: string } | null) => { - if (!disposedRef.current) setRunId(rs?.runId ?? null) - }, } const persistenceProps = @@ -289,6 +267,12 @@ export function useGenerateVideo( ) }, [clientIdentity, hookId]) + const snapshot = useSyncExternalStore( + client.subscribe, + client.getSnapshot, + client.getSnapshot, + ) + // Sync body changes without recreating client useEffect(() => { // Conditional spread: target uses strict-optional `body?: T`. @@ -327,14 +311,14 @@ export function useGenerateVideo( return { generate, - result, - jobId, - videoStatus, - isLoading, - error, - status, + result: snapshot.result, + jobId: snapshot.jobId, + videoStatus: snapshot.videoStatus, + isLoading: snapshot.isLoading, + error: snapshot.error, + status: snapshot.status, stop, reset, - runId, + runId: snapshot.runId, } } diff --git a/packages/ai-react/src/use-generation.ts b/packages/ai-react/src/use-generation.ts index cc077bd8d5..cdcc84d55a 100644 --- a/packages/ai-react/src/use-generation.ts +++ b/packages/ai-react/src/use-generation.ts @@ -1,6 +1,13 @@ import { GenerationClient } from '@tanstack/ai-client' import { createGenerationDevtoolsBridge } from '@tanstack/ai-client/devtools' -import { useCallback, useEffect, useId, useMemo, useRef, useState } from 'react' +import { + useCallback, + useEffect, + useId, + useMemo, + useRef, + useSyncExternalStore, +} from 'react' import type { StreamChunk } from '@tanstack/ai' import type { AIDevtoolsClientMetadata, @@ -203,12 +210,6 @@ export function useGenerationWithDevtoolsIdentity< // The hook identity is `threadId`. `hookId` is only a React recreation key. const clientIdentity = options.threadId ?? hookId - const [result, setResult] = useState(null) - const [isLoading, setIsLoading] = useState(false) - const [error, setError] = useState(undefined) - const [status, setStatus] = useState('idle') - const [runId, setRunId] = useState(null) - const optionsRef = useRef(options) optionsRef.current = options const disposedRef = useRef(false) @@ -258,21 +259,6 @@ export function useGenerationWithDevtoolsIdentity< onChunk: (c: StreamChunk) => { if (!disposedRef.current) optionsRef.current.onChunk?.(c) }, - onResultChange: (r) => { - if (!disposedRef.current) setResult(r) - }, - onLoadingChange: (l) => { - if (!disposedRef.current) setIsLoading(l) - }, - onErrorChange: (e) => { - if (!disposedRef.current) setError(e) - }, - onStatusChange: (s) => { - if (!disposedRef.current) setStatus(s) - }, - onResumeStateChange: (rs) => { - if (!disposedRef.current) setRunId(rs?.runId ?? null) - }, } const persistenceProps = @@ -306,6 +292,12 @@ export function useGenerationWithDevtoolsIdentity< ) }, [clientIdentity, hookId]) + const snapshot = useSyncExternalStore( + client.subscribe, + client.getSnapshot, + client.getSnapshot, + ) + // Sync body changes without recreating client useEffect(() => { // Conditional spread: target uses strict-optional `body?: T`. @@ -344,12 +336,12 @@ export function useGenerationWithDevtoolsIdentity< return { generate, - result, - isLoading, - error, - status, + result: snapshot.result, + isLoading: snapshot.isLoading, + error: snapshot.error, + status: snapshot.status, stop, reset, - runId, + runId: snapshot.runId, } } diff --git a/packages/ai-react/src/use-realtime-chat.ts b/packages/ai-react/src/use-realtime-chat.ts index bddef502c6..9646765285 100644 --- a/packages/ai-react/src/use-realtime-chat.ts +++ b/packages/ai-react/src/use-realtime-chat.ts @@ -1,11 +1,12 @@ -import { useCallback, useEffect, useRef, useState } from 'react' +import { + useCallback, + useEffect, + useRef, + useState, + useSyncExternalStore, +} from 'react' import { RealtimeClient } from '@tanstack/ai-client' -import type { - RealtimeMessage, - RealtimeMode, - RealtimeSessionConfig, - RealtimeStatus, -} from '@tanstack/ai' +import type { RealtimeSessionConfig } from '@tanstack/ai' import type { UseRealtimeChatOptions, UseRealtimeChatReturn, @@ -58,17 +59,6 @@ const emptyTimeDomainData = new Uint8Array(128).fill(128) export function useRealtimeChat( options: UseRealtimeChatOptions, ): UseRealtimeChatReturn { - // State - const [status, setStatus] = useState('idle') - const [mode, setMode] = useState('idle') - const [messages, setMessages] = useState>([]) - const [pendingUserTranscript, setPendingUserTranscript] = useState< - string | null - >(null) - const [pendingAssistantTranscript, setPendingAssistantTranscript] = useState< - string | null - >(null) - const [error, setError] = useState(null) const [inputLevel, setInputLevel] = useState(0) const [outputLevel, setOutputLevel] = useState(0) @@ -113,15 +103,10 @@ export function useRealtimeChat( ...(initial.semanticEagerness !== undefined && { semanticEagerness: initial.semanticEagerness, }), - onStatusChange: (newStatus) => { - setStatus(newStatus) - }, onModeChange: (newMode) => { - setMode(newMode) optionsRef.current.onModeChange?.(newMode) }, onMessage: (message) => { - setMessages((prev) => [...prev, message]) optionsRef.current.onMessage?.(message) }, onUsage(usage) { @@ -131,30 +116,26 @@ export function useRealtimeChat( optionsRef.current.onGoAway?.(timeLeft) }, onError: (err) => { - setError(err) optionsRef.current.onError?.(err) }, onConnect: () => { - setError(null) optionsRef.current.onConnect?.() }, onDisconnect: () => { optionsRef.current.onDisconnect?.() }, onInterrupted: () => { - setPendingAssistantTranscript(null) optionsRef.current.onInterrupted?.() }, }) - - // Subscribe to state changes for transcripts - clientRef.current.onStateChange((state) => { - setPendingUserTranscript(state.pendingUserTranscript) - setPendingAssistantTranscript(state.pendingAssistantTranscript) - }) } const client = clientRef.current + const snapshot = useSyncExternalStore( + client.subscribe, + client.getSnapshot, + client.getSnapshot, + ) // Audio level animation loop useEffect(() => { @@ -166,7 +147,7 @@ export function useRealtimeChat( animationFrameRef.current = requestAnimationFrame(updateLevels) } - if (status === 'connected') { + if (snapshot.status === 'connected') { updateLevels() } @@ -176,7 +157,7 @@ export function useRealtimeChat( animationFrameRef.current = null } } - }, [status]) + }, [snapshot.status]) // Cleanup on unmount useEffect(() => { @@ -187,10 +168,6 @@ export function useRealtimeChat( // Connection methods const connect = useCallback(async () => { - setError(null) - setMessages([]) - setPendingUserTranscript(null) - setPendingAssistantTranscript(null) await client.connect() }, [client]) @@ -258,16 +235,16 @@ export function useRealtimeChat( return { // Connection state - status, - error, + status: snapshot.status, + error: snapshot.error, connect, disconnect, // Conversation state - mode, - messages, - pendingUserTranscript, - pendingAssistantTranscript, + mode: snapshot.mode, + messages: snapshot.messages, + pendingUserTranscript: snapshot.pendingUserTranscript, + pendingAssistantTranscript: snapshot.pendingAssistantTranscript, // Voice control startListening, diff --git a/packages/ai-react/tests/chat-ui/create-ui-types.test.tsx b/packages/ai-react/tests/chat-ui/create-ui-types.test.tsx index 32fa04dccc..1aea6497e5 100644 --- a/packages/ai-react/tests/chat-ui/create-ui-types.test.tsx +++ b/packages/ai-react/tests/chat-ui/create-ui-types.test.tsx @@ -128,7 +128,9 @@ it('types tool and interrupt component props from chatOptions', () => { ChatUIHost >() expectTypeOf(UI.useChatContext().sendMessage).toBeFunction() - expectTypeOf(UI.useChatContext().queue).toBeArray() + expectTypeOf(UI.useChatContext().queue).toExtend< + ReadonlyArray + >() expectTypeOf['Queue']>().not.toBeAny() return }, diff --git a/packages/ai-react/tests/chat-ui/create-ui.test.tsx b/packages/ai-react/tests/chat-ui/create-ui.test.tsx index 01c451190b..82de8a989f 100644 --- a/packages/ai-react/tests/chat-ui/create-ui.test.tsx +++ b/packages/ai-react/tests/chat-ui/create-ui.test.tsx @@ -119,9 +119,9 @@ describe('createChatHook', () => { function Screen() { const chat = useAppChat({ - initialMessages: [messageWithToolResults] as ChatUIHost< - typeof chatOptions - >['messages'], + initialMessages: [messageWithToolResults] as Array< + ChatUIHost['messages'][number] + >, }) return } diff --git a/packages/ai-react/tests/use-chat-types.test.ts b/packages/ai-react/tests/use-chat-types.test.ts index 17521be061..a7eafa6561 100644 --- a/packages/ai-react/tests/use-chat-types.test.ts +++ b/packages/ai-react/tests/use-chat-types.test.ts @@ -48,7 +48,7 @@ describe('useChat() return type', () => { type R = UseChatReturn expectTypeOf().toBeFunction() expectTypeOf().toBeBoolean() - expectTypeOf().toBeArray() + expectTypeOf().toExtend>() }) it('options accept outputSchema with the schema type', () => { diff --git a/packages/ai-react/tests/use-chat.test.ts b/packages/ai-react/tests/use-chat.test.ts index 5678552bbd..7fcd4699f3 100644 --- a/packages/ai-react/tests/use-chat.test.ts +++ b/packages/ai-react/tests/use-chat.test.ts @@ -1,7 +1,13 @@ import { EventType } from '@tanstack/ai' import { ChatClient } from '@tanstack/ai-client' import { act, render, renderHook, waitFor } from '@testing-library/react' -import { StrictMode, Suspense, createElement, useState } from 'react' +import { + StrictMode, + Suspense, + createElement, + useLayoutEffect, + useState, +} from 'react' import { afterEach, describe, expect, it, vi } from 'vitest' import { useChat } from '../src/use-chat' import { @@ -204,6 +210,43 @@ describe('useChat', () => { expect(result.current.messages).toEqual(initialMessages) }) + it('throws on in-place messages.push and accepts setMessages', () => { + const { result } = renderUseChat({ + connection: createMockConnectionAdapter(), + initialMessages: [ + { + id: 'msg-1', + role: 'user', + parts: [{ type: 'text', content: 'Hello' }], + createdAt: new Date(), + }, + ], + }) + + expect(Object.isFrozen(result.current.messages)).toBe(true) + expect(() => { + Array.prototype.push.call(result.current.messages, { + id: 'x', + role: 'user', + parts: [{ type: 'text', content: 'nope' }], + createdAt: new Date(), + }) + }).toThrow(TypeError) + + act(() => { + result.current.setMessages([ + ...result.current.messages, + { + id: 'ok', + role: 'user', + parts: [{ type: 'text', content: 'ok' }], + createdAt: new Date(), + }, + ]) + }) + expect(result.current.messages).toHaveLength(2) + }) + it('should initialize with persisted messages', async () => { const adapter = createMockConnectionAdapter() const persistedMessages: Array = [ @@ -232,7 +275,7 @@ describe('useChat', () => { expect(persistence.getItem).toHaveBeenCalledWith('persisted-chat') }) - it('should forward synchronous persisted interrupt hydration', () => { + it('forwards persisted interrupt hydration from a synchronous store', async () => { const onInterruptStateChange = vi.fn() const persistence = { getItem: vi.fn(() => ({ @@ -254,7 +297,9 @@ describe('useChat', () => { { wrapper: StrictMode }, ) - expect(result.current.interrupts).toHaveLength(2) + await waitFor(() => { + expect(result.current.interrupts).toHaveLength(2) + }) expect(onInterruptStateChange).toHaveBeenCalledOnce() expect(onInterruptStateChange).toHaveBeenCalledWith( expect.objectContaining({ interrupts: result.current.interrupts }), @@ -262,6 +307,360 @@ describe('useChat', () => { ) }) + it('keeps the server message count for the first hydration render', async () => { + const initialMessages: Array = [ + { + id: 'server-message', + role: 'user', + parts: [{ type: 'text', content: 'Server message' }], + createdAt: new Date(), + }, + ] + const persistedMessages: Array = [ + ...initialMessages, + { + id: 'persisted-message', + role: 'assistant', + parts: [{ type: 'text', content: 'Persisted message' }], + createdAt: new Date(), + }, + ] + const persistence = { + getItem: vi.fn(() => persistedMessages), + setItem: vi.fn(), + removeItem: vi.fn(), + } + const renderCounts: Array = [] + const onRecoverableError = vi.fn() + const container = document.createElement('div') + container.innerHTML = '
1
' + document.body.append(container) + + function HydratedChat() { + const chat = useChat({ + connection: createMockConnectionAdapter(), + threadId: 'hydrated-message-chat', + initialMessages, + persistence, + }) + renderCounts.push(chat.messages.length) + return createElement( + 'div', + { 'data-testid': 'message-count' }, + chat.messages.length, + ) + } + + const view = render( + createElement(StrictMode, null, createElement(HydratedChat)), + { container, hydrate: true, onRecoverableError }, + ) + + expect(renderCounts[0]).toBe(1) + await waitFor(() => { + expect(view.getByTestId('message-count').textContent).toBe('2') + }) + expect(persistence.getItem).toHaveBeenCalledOnce() + expect(onRecoverableError).not.toHaveBeenCalled() + }) + + it('keeps server interrupts empty until persistence hydration starts after commit', async () => { + const persistence = { + getItem: vi.fn(() => ({ + messages: [], + resume: createInterruptResumeSnapshot(), + })), + setItem: vi.fn(), + removeItem: vi.fn(), + } + const renderCounts: Array = [] + const hydrationCommitStates: Array = [] + const onRecoverableError = vi.fn() + const container = document.createElement('div') + container.innerHTML = '
0
' + document.body.append(container) + + function HydratedInterrupts() { + let committed = false + useLayoutEffect(() => { + committed = true + }, []) + const chat = useChat({ + connection: createMockConnectionAdapter(), + threadId: 'hydrated-interrupt-chat', + persistence, + onInterruptStateChange: (_state, context) => { + if (context.source === 'hydrate') { + hydrationCommitStates.push(committed) + } + }, + }) + renderCounts.push(chat.interrupts.length) + return createElement( + 'div', + { 'data-testid': 'interrupt-count' }, + chat.interrupts.length, + ) + } + + const view = render( + createElement(StrictMode, null, createElement(HydratedInterrupts)), + { container, hydrate: true, onRecoverableError }, + ) + + expect(renderCounts[0]).toBe(0) + await waitFor(() => { + expect(view.getByTestId('interrupt-count').textContent).toBe('2') + }) + expect(hydrationCommitStates).toEqual([true]) + expect(persistence.getItem).toHaveBeenCalledOnce() + expect(onRecoverableError).not.toHaveBeenCalled() + }) + + it('keeps the initial resume snapshot through hydration before rejoining the stored run', async () => { + const initialResumeSnapshot = createInterruptResumeSnapshot() + const persistence = { + getItem: vi.fn(() => ({ + messages: [], + resume: { + schemaVersion: 2, + resumeState: { + threadId: 'hydrated-resume-chat', + runId: 'stored-run', + }, + }, + })), + setItem: vi.fn(), + removeItem: vi.fn(), + } + const joinRun = vi.fn(async function* (_runId: string) {}) + const connection: ResumableConnectConnectionAdapter = { + connect: async function* () {}, + joinRun, + } + const renderCounts: Array = [] + const onRecoverableError = vi.fn() + const container = document.createElement('div') + container.innerHTML = '
2
' + document.body.append(container) + + function HydratedResume() { + const chat = useChat({ + connection, + threadId: 'hydrated-resume-chat', + initialResumeSnapshot, + persistence, + }) + renderCounts.push(chat.interrupts.length) + return createElement( + 'div', + { 'data-testid': 'interrupt-count' }, + chat.interrupts.length, + ) + } + + const view = render(createElement(HydratedResume), { + container, + hydrate: true, + onRecoverableError, + }) + + expect(renderCounts[0]).toBe(2) + expect(joinRun).not.toHaveBeenCalled() + await waitFor(() => { + expect(view.getByTestId('interrupt-count').textContent).toBe('0') + expect(joinRun).toHaveBeenCalledOnce() + }) + expect(joinRun).toHaveBeenCalledWith('stored-run', expect.anything()) + expect(persistence.getItem).toHaveBeenCalledOnce() + expect(onRecoverableError).not.toHaveBeenCalled() + }) + + it('preserves the persistence adapter receiver for reads, writes, and removal', async () => { + class ReceiverSensitivePersistence implements ChatClientPersistence { + calls: Array = [] + + getItem(id: string) { + this.calls.push(`get:${id}`) + return [] + } + + setItem(id: string) { + this.calls.push(`set:${id}`) + } + + removeItem(id: string) { + this.calls.push(`remove:${id}`) + } + } + + const persistence = new ReceiverSensitivePersistence() + const { result } = renderUseChat({ + connection: createMockConnectionAdapter({ + chunks: createTextChunks('Persisted reply'), + }), + threadId: 'receiver-chat', + persistence, + }) + + await waitFor(() => { + expect(persistence.calls).toContain('get:receiver-chat') + }) + await act(async () => { + await result.current.sendMessage('Persist me') + }) + await waitFor(() => { + expect(persistence.calls).toContain('set:receiver-chat') + }) + const removalsBeforeClear = persistence.calls.filter( + (call) => call === 'remove:receiver-chat', + ).length + act(() => result.current.clear()) + await waitFor(() => { + const removalsAfterClear = persistence.calls.filter( + (call) => call === 'remove:receiver-chat', + ).length + expect(removalsAfterClear).toBeGreaterThan(removalsBeforeClear) + }) + }) + + it('keeps initial state when a synchronous persistence read throws during hydration', async () => { + const initialMessages: Array = [ + { + id: 'initial-message', + role: 'user', + parts: [{ type: 'text', content: 'Initial' }], + createdAt: new Date(), + }, + ] + const persistence = { + getItem: vi.fn(() => { + throw new Error('storage unavailable') + }), + setItem: vi.fn(), + removeItem: vi.fn(), + } + const onInterruptStateChange = vi.fn() + const onRecoverableError = vi.fn() + const container = document.createElement('div') + container.innerHTML = '
1
' + document.body.append(container) + + function HydratedChat() { + const chat = useChat({ + connection: createMockConnectionAdapter(), + threadId: 'throwing-sync-chat', + initialMessages, + persistence, + onInterruptStateChange, + }) + return createElement( + 'div', + { 'data-testid': 'message-count' }, + chat.messages.length, + ) + } + + const view = render(createElement(HydratedChat), { + container, + hydrate: true, + onRecoverableError, + }) + + await waitFor(() => { + expect(persistence.getItem).toHaveBeenCalledOnce() + }) + expect(view.getByTestId('message-count').textContent).toBe('1') + expect(onInterruptStateChange).not.toHaveBeenCalled() + expect(onRecoverableError).not.toHaveBeenCalled() + }) + + it('keeps initial state when an asynchronous persistence read rejects', async () => { + const initialMessages: Array = [ + { + id: 'initial-message', + role: 'user', + parts: [{ type: 'text', content: 'Initial' }], + createdAt: new Date(), + }, + ] + const persistence = { + getItem: vi.fn(async () => { + throw new Error('storage unavailable') + }), + setItem: vi.fn(), + removeItem: vi.fn(), + } + const onInterruptStateChange = vi.fn() + const { result } = renderUseChat({ + connection: createMockConnectionAdapter(), + threadId: 'throwing-async-chat', + initialMessages, + persistence, + onInterruptStateChange, + }) + + await waitFor(() => { + expect(persistence.getItem).toHaveBeenCalledOnce() + }) + expect(result.current.messages).toEqual(initialMessages) + expect(onInterruptStateChange).not.toHaveBeenCalled() + }) + + it('applies asynchronous persistence only after attach and commit', async () => { + const hydration = createDeferred<{ + messages: Array + resume: ReturnType + }>() + let committed = false + const readCommitStates: Array = [] + const hydrationCommitStates: Array = [] + const persistence = { + getItem: vi.fn(() => { + readCommitStates.push(committed) + return hydration.promise + }), + setItem: vi.fn(), + removeItem: vi.fn(), + } + + function useHydratedChat() { + useLayoutEffect(() => { + committed = true + }, []) + const chat = useChat({ + connection: createMockConnectionAdapter(), + threadId: 'async-hydration-chat', + persistence, + onInterruptStateChange: (_state, context) => { + if (context.source === 'hydrate') { + hydrationCommitStates.push(committed) + } + }, + }) + return chat.interrupts.length + } + + const { result } = renderHook(() => useHydratedChat()) + + expect(result.current).toBe(0) + await waitFor(() => { + expect(persistence.getItem).toHaveBeenCalledOnce() + }) + expect(readCommitStates).toEqual([true]) + await act(async () => { + hydration.resolve({ + messages: [], + resume: createInterruptResumeSnapshot(), + }) + await hydration.promise + }) + await waitFor(() => { + expect(result.current).toBe(2) + }) + expect(hydrationCommitStates).toEqual([true]) + }) + it('does not publish async hydration from an abandoned render', async () => { const hydration = createDeferred<{ messages: Array @@ -302,6 +701,7 @@ describe('useChat', () => { await hydration.promise }) + expect(persistence.getItem).not.toHaveBeenCalled() expect(onInterruptStateChange).not.toHaveBeenCalled() }) @@ -330,8 +730,8 @@ describe('useChat', () => { await waitFor(() => { expect(persistence.getItem).toHaveBeenCalledWith('persisted-empty-chat') + expect(result.current.messages).toEqual([]) }) - expect(result.current.messages).toEqual([]) }) it('should ignore async persisted messages from a previous id', async () => { diff --git a/packages/ai-remix/src/create-chat.ts b/packages/ai-remix/src/create-chat.ts index 590af4557f..ef3b2d6e1f 100644 --- a/packages/ai-remix/src/create-chat.ts +++ b/packages/ai-remix/src/create-chat.ts @@ -2,13 +2,9 @@ import { ChatClient } from '@tanstack/ai-client' import { createChatDevtoolsBridge } from '@tanstack/ai-client/devtools' import type { Handle } from 'remix/ui' import type { - ChatClientState, ResolvableChatInterrupt, - ChatInterruptState, ChatResumeState, - ConnectionStatus, InferredClientContext, - QueuedMessage, SendMessageOptions, StructuredOutputPart, } from '@tanstack/ai-client' @@ -29,16 +25,12 @@ import type { UIMessage, } from './types.ts' -const EMPTY_INTERRUPTS = Object.freeze([]) -const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) - /** * Create a chat helper for a Remix component. * * Call this in setup with the component Handle from `remix/ui`. The helper - * wraps ChatClient and stores chat state in local variables. ChatClient state - * callbacks write those variables and then call `handle.update()`, so render - * reads the latest snapshot through getters. + * wraps ChatClient and calls `handle.update()` when the client snapshot + * changes, so render reads the latest frozen snapshot through getters. * * The default thread id is `options.threadId ?? handle.id`. Cleanup runs when * `handle.signal` aborts. Do not pass identification in `options.devtools`; @@ -81,23 +73,6 @@ export function createChat< handle: Pick, options: CreateChatOptions, ) { - let messages = options.initialMessages || [] - let isLoading = false - let hasOlderMessages = false - let error: Error | undefined - let status: ChatClientState = 'ready' - let isSubscribed = false - let connectionStatus: ConnectionStatus = 'disconnected' - let sessionGenerating = false - let isHydrating = false - let queue: Array = [] - let runId: string | null = null - let interruptState: ChatInterruptState = { - interrupts: EMPTY_INTERRUPTS, - pendingInterrupts: EMPTY_INTERRUPTS, - interruptErrors: EMPTY_INTERRUPT_ERRORS, - resuming: false, - } let closed = false type Partial = DeepPartial>> @@ -168,69 +143,19 @@ export function createChat< ...(options.streamProcessor !== undefined && { streamProcessor: options.streamProcessor, }), - onMessagesChange: (newMessages: Array>) => { - messages = newMessages - hasOlderMessages = client.getHasOlderMessages() - commit() - }, - onLoadingChange: (newIsLoading: boolean) => { - isLoading = newIsLoading - syncResumeState() - commit() - }, - onStatusChange: (newStatus: ChatClientState) => { - status = newStatus - commit() - }, - onErrorChange: (newError: Error | undefined) => { - error = newError - commit() - }, - onSubscriptionChange: (nextIsSubscribed: boolean) => { - isSubscribed = nextIsSubscribed - commit() - }, - onConnectionStatusChange: (nextStatus: ConnectionStatus) => { - connectionStatus = nextStatus - commit() - }, - onSessionGeneratingChange: (isGenerating: boolean) => { - sessionGenerating = isGenerating - commit() - }, - onHydratingChange: (nextIsHydrating: boolean) => { - isHydrating = nextIsHydrating - commit() - }, ...(options.queue !== undefined && { queue: options.queue }), - onQueueChange: (nextQueue: Array) => { - queue = nextQueue - commit() - }, - onRunIdChange: (nextRunId) => { - runId = nextRunId - commit() - }, onInterruptStateChange: (nextInterruptState, context) => { - interruptState = nextInterruptState options.onInterruptStateChange?.(nextInterruptState, context) - commit() }, }) - function syncResumeState() { - runId = client.getCurrentRunId() - interruptState = client.getInterruptState() - } - - messages = client.getMessages() - isHydrating = client.getIsHydrating() - interruptState = client.getInterruptState() - runId = client.getCurrentRunId() + // Render reads the frozen client snapshot through the getters below. + const unsubscribeSnapshot = client.subscribeSnapshot(commit) function close() { if (closed) return closed = true + unsubscribeSnapshot() client.detach() if (options.live) { client.unsubscribe() @@ -265,37 +190,23 @@ export function createChat< sendOptions?: SendMessageOptions, ) => { syncTools() - try { - await client.sendMessage(content, undefined, sendOptions) - } finally { - syncResumeState() - } + await client.sendMessage(content, undefined, sendOptions) } const cancelQueued = (id: string) => client.cancelQueued(id) const append = async (message: ModelMessage | UIMessage) => { syncTools() - try { - await client.append(message) - } finally { - syncResumeState() - } + await client.append(message) } const reload = async () => { syncTools() - try { - await client.reload() - } finally { - syncResumeState() - } + await client.reload() } const loadOlderMessages = async () => { await client.loadOlderMessages() - hasOlderMessages = client.getHasOlderMessages() - commit() } const stop = () => { @@ -304,10 +215,9 @@ export function createChat< const clear = () => { client.clear() - syncResumeState() } - const setMessages = (newMessages: Array>) => { + const setMessages = (newMessages: ReadonlyArray>) => { client.setMessagesManually(newMessages) } @@ -329,7 +239,6 @@ export function createChat< }) => { syncTools() await client.addToolApprovalResponse(response) - syncResumeState() } const resumeInterrupts = async ( @@ -337,9 +246,7 @@ export function createChat< state?: ChatResumeState, ) => { syncTools() - const result = await client.resumeInterrupts(resumeItems, state) - syncResumeState() - return result + return await client.resumeInterrupts(resumeItems, state) } const resolveInterrupts = ( @@ -375,6 +282,7 @@ export function createChat< } function activeStructuredPart(): StructuredOutputPart | null { + const { messages } = client.getSnapshot() let lastUserIndex = -1 for (let i = messages.length - 1; i >= 0; i--) { if (messages[i]?.role === 'user') { @@ -396,49 +304,49 @@ export function createChat< return { get messages() { - return messages + return client.getSnapshot().messages }, get isLoading() { - return isLoading + return client.getSnapshot().isLoading }, get hasOlderMessages() { - return hasOlderMessages + return client.getSnapshot().hasOlderMessages }, get error() { - return error + return client.getSnapshot().error }, get status() { - return status + return client.getSnapshot().status }, get isSubscribed() { - return isSubscribed + return client.getSnapshot().isSubscribed }, get connectionStatus() { - return connectionStatus + return client.getSnapshot().connectionStatus }, get sessionGenerating() { - return sessionGenerating + return client.getSnapshot().sessionGenerating }, get isHydrating() { - return isHydrating + return client.getSnapshot().isHydrating }, get queue() { - return queue + return client.getSnapshot().queue }, get runId() { - return runId + return client.getSnapshot().runId }, get interrupts() { - return interruptState.interrupts + return client.getSnapshot().interruptState.interrupts }, get pendingInterrupts() { - return interruptState.interrupts + return client.getSnapshot().interruptState.interrupts }, get interruptErrors() { - return interruptState.interruptErrors + return client.getSnapshot().interruptState.interruptErrors }, get resuming() { - return interruptState.resuming + return client.getSnapshot().interruptState.resuming }, get partial() { const part = activeStructuredPart() diff --git a/packages/ai-remix/src/types.ts b/packages/ai-remix/src/types.ts index 33964c582b..57afe14197 100644 --- a/packages/ai-remix/src/types.ts +++ b/packages/ai-remix/src/types.ts @@ -168,7 +168,7 @@ interface BaseCreateChatReturn< * `messages[i].parts.find(p => p.type === 'structured-output')` is typed * with the schema's inferred shape: `data: T`, `partial: DeepPartial`. */ - messages: Array> + messages: ReadonlyArray> /** * Send a message and get a response. @@ -187,7 +187,7 @@ interface BaseCreateChatReturn< * Pending messages queued while the client is busy (streaming, claiming a * send, or draining). Separate from `messages` until they drain. */ - queue: Array + queue: ReadonlyArray /** * Cancel a queued message before it drains. No-op if already sent. @@ -320,7 +320,7 @@ interface BaseCreateChatReturn< /** * Set messages manually */ - setMessages: (messages: Array>) => void + setMessages: (messages: ReadonlyArray>) => void /** * Clear all messages diff --git a/packages/ai-remix/tests/create-chat.test.ts b/packages/ai-remix/tests/create-chat.test.ts index f90e6d2910..ebb1ef1053 100644 --- a/packages/ai-remix/tests/create-chat.test.ts +++ b/packages/ai-remix/tests/create-chat.test.ts @@ -121,4 +121,26 @@ describe('createChat', () => { expect(handle.update).not.toHaveBeenCalled() }) + + it('reads the frozen client snapshot and updates on every snapshot change until abort', async () => { + const { handle, abort } = createHandle() + const chat = createChat(handle, { + connection: createConnection([textChunk('Hi'), runFinished()]), + }) + + await chat.sendMessage('Hello') + const messages = chat.messages + expect(Object.isFrozen(messages)).toBe(true) + expect(() => (messages as Array).push({})).toThrow(TypeError) + + vi.mocked(handle.update).mockClear() + chat.setMessages(messages.slice(0, 1)) + expect(chat.messages).toHaveLength(1) + expect(handle.update).toHaveBeenCalledTimes(1) + + abort.abort() + vi.mocked(handle.update).mockClear() + chat.setMessages([]) + expect(handle.update).not.toHaveBeenCalled() + }) }) diff --git a/packages/ai-sandbox/src/tool-bridge.ts b/packages/ai-sandbox/src/tool-bridge.ts index 5d7a117a31..6d370716ca 100644 --- a/packages/ai-sandbox/src/tool-bridge.ts +++ b/packages/ai-sandbox/src/tool-bridge.ts @@ -363,8 +363,15 @@ export async function startHostToolBridge( name: BRIDGED_MCP_SERVER_NAME, url, token, - close: () => - new Promise((resolve) => httpServer.close(() => resolve())), + close: () => { + const closed = new Promise((resolve) => + httpServer.close(() => resolve()), + ) + // `close()` alone waits for every open connection. An agent that keeps + // a request or a keep-alive socket open would hold the run's teardown. + httpServer.closeAllConnections() + return closed + }, } } diff --git a/packages/ai-sandbox/tests/tool-bridge.test.ts b/packages/ai-sandbox/tests/tool-bridge.test.ts index da8d9dab65..827691812d 100644 --- a/packages/ai-sandbox/tests/tool-bridge.test.ts +++ b/packages/ai-sandbox/tests/tool-bridge.test.ts @@ -3,6 +3,7 @@ * client (no `claude` needed): list tools + call a tool, asserting the host * `execute()` runs and the result comes back, plus bearer-token enforcement. */ +import { request } from 'node:http' import { afterEach, describe, expect, it } from 'vitest' import { Client } from '@modelcontextprotocol/sdk/client/index.js' import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js' @@ -96,6 +97,19 @@ describe('startHostToolBridge', () => { await expect(client.connect(transport)).rejects.toThrow() await client.close().catch(() => {}) }) + + it('closes while a request is still open', async () => { + const open = await startHostToolBridge([], { hostForSandbox: '127.0.0.1' }) + // Send headers and part of a body, then never finish the request. + const req = request(open.url, { + method: 'POST', + headers: { authorization: `Bearer ${open.token}` }, + }) + req.on('error', () => {}) + req.write('{') + await new Promise((resolve) => setTimeout(resolve, 50)) + await open.close() + }) }) describe('hostForSandbox', () => { diff --git a/packages/ai-solid/src/types.ts b/packages/ai-solid/src/types.ts index fff7873e6a..e800fff851 100644 --- a/packages/ai-solid/src/types.ts +++ b/packages/ai-solid/src/types.ts @@ -148,13 +148,13 @@ interface BaseUseChatReturn< * `messages()[i].parts.find(p => p.type === 'structured-output')` is typed * by the schema — `data: T`, `partial: DeepPartial`. */ - messages: Accessor>> + messages: Accessor>> /** * Pending messages queued while the client is busy (streaming, claiming a * send, or draining). Separate from `messages` until they drain. */ - queue: Accessor> + queue: Accessor> /** * Cancel a queued message before it drains. No-op if already sent. @@ -270,7 +270,7 @@ interface BaseUseChatReturn< /** * Set messages manually */ - setMessages: (messages: Array>) => void + setMessages: (messages: ReadonlyArray>) => void /** * Clear all messages diff --git a/packages/ai-solid/src/use-chat.ts b/packages/ai-solid/src/use-chat.ts index 21f759034e..4d46d999bb 100644 --- a/packages/ai-solid/src/use-chat.ts +++ b/packages/ai-solid/src/use-chat.ts @@ -11,13 +11,9 @@ import { import { ChatClient } from '@tanstack/ai-client' import { createChatDevtoolsBridge } from '@tanstack/ai-client/devtools' import type { - ChatClientState, ResolvableChatInterrupt, - ChatInterruptState, ChatResumeState, - ConnectionStatus, InferredClientContext, - QueuedMessage, SendMessageOptions, StructuredOutputPart, } from '@tanstack/ai-client' @@ -38,9 +34,6 @@ import type { UseChatReturn, } from './types' -const EMPTY_INTERRUPTS = Object.freeze([]) -const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) - export function useChat< const TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, @@ -59,34 +52,6 @@ export function useChat< // The hook's identity is its `threadId`. Reload with the same `threadId` // restores the same conversation. - const [messages, setMessages] = createSignal>>( - options.initialMessages || [], - ) - const [isLoading, setIsLoading] = createSignal(false) - const [hasOlderMessages, setHasOlderMessages] = createSignal(false) - const [error, setError] = createSignal(undefined) - const [status, setStatus] = createSignal('ready') - const [isSubscribed, setIsSubscribed] = createSignal(false) - const [connectionStatus, setConnectionStatus] = - createSignal('disconnected') - const [sessionGenerating, setSessionGenerating] = createSignal(false) - const [isHydrating, setIsHydrating] = createSignal(false) - const [queue, setQueue] = createSignal>([]) - const [runId, setRunId] = createSignal(null) - const [interruptState, setInterruptState] = createSignal< - ChatInterruptState - >({ - interrupts: EMPTY_INTERRUPTS, - pendingInterrupts: EMPTY_INTERRUPTS, - interruptErrors: EMPTY_INTERRUPT_ERRORS, - resuming: false, - }) - - const syncResumeState = () => { - setRunId(client().getCurrentRunId()) - setInterruptState(client().getInterruptState()) - } - // Structured-output `partial` / `final` are derived from `messages` — // specifically from the structured-output part on the latest assistant // message (the one after the most recent user message). Per-turn parts @@ -94,7 +59,7 @@ export function useChat< type Partial = DeepPartial>> type Final = InferSchemaType> - // Create ChatClient instance with callbacks to sync state. + // Build the client once per threadId. User callbacks read latest options. // Every user-provided callback is wrapped so the LATEST `options.xxx` value // is read at call time. Direct assignment would freeze the callback to the // reference we saw at creation; the wrapper lets reactive `options` or @@ -169,48 +134,8 @@ export function useChat< ...(options.streamProcessor !== undefined && { streamProcessor: options.streamProcessor, }), - onMessagesChange: (newMessages: Array>) => { - setMessages(newMessages) - setHasOlderMessages(instance.getHasOlderMessages()) - }, - onLoadingChange: (newIsLoading: boolean) => { - setIsLoading(newIsLoading) - syncResumeState() - }, - onStatusChange: (newStatus: ChatClientState) => { - setStatus(newStatus) - }, - onErrorChange: (newError: Error | undefined) => { - setError(newError) - }, - onSubscriptionChange: (nextIsSubscribed: boolean) => { - setIsSubscribed(nextIsSubscribed) - }, - onConnectionStatusChange: (nextStatus: ConnectionStatus) => { - setConnectionStatus(nextStatus) - }, - onSessionGeneratingChange: (isGenerating: boolean) => { - setSessionGenerating(isGenerating) - }, - onHydratingChange: (nextIsHydrating: boolean) => { - setIsHydrating(nextIsHydrating) - }, ...(options.queue !== undefined && { queue: options.queue }), - onQueueChange: (nextQueue: Array) => { - setQueue(nextQueue) - }, - onRunIdChange: (nextRunId) => { - setRunId(nextRunId) - }, - onResumeStateChange: (_nextResumeState, nextPendingInterrupts) => { - setInterruptState((current) => ({ - ...current, - interrupts: nextPendingInterrupts, - pendingInterrupts: nextPendingInterrupts, - })) - }, onInterruptStateChange: (nextInterruptState, context) => { - setInterruptState(nextInterruptState) options.onInterruptStateChange?.(nextInterruptState, context) }, }) @@ -221,10 +146,18 @@ export function useChat< // new client and drop the transcript. The effects below sync those values. const client = createMemo(on(() => options.threadId, createClient)) - setMessages(client().getMessages()) - setHasOlderMessages(client().getHasOlderMessages()) - setIsHydrating(client().getIsHydrating()) - syncResumeState() + // One signal: the snapshot is frozen and replaced on every change. Seeded + // here so constructor hydration (a sync store) is on the first paint. + const [snapshot, setSnapshot] = createSignal(client().getSnapshot()) + createEffect(() => { + const target = client() + setSnapshot(target.getSnapshot()) + onCleanup(target.subscribeSnapshot(() => setSnapshot(target.getSnapshot()))) + }) + const pick = >(key: TKey) => + createMemo(() => snapshot()[key]) + const messages = pick('messages') + const interruptState = pick('interruptState') // Sync body / forwardedProps changes to the client. // Both populate the same wire payload; `forwardedProps` is preferred @@ -256,6 +189,11 @@ export function useChat< client().unsubscribe() } + // First paint must see constructor hydration (sync store / resume snapshot) + // and the live subscribe above. `createEffect` runs after paint, so the + // subscribe effect cannot be the first apply. + setSnapshot(client().getSnapshot()) + createEffect(() => { if (options.live) { client().subscribe() @@ -271,10 +209,6 @@ export function useChat< // connections per origin until the page reloaded. client().attach() client().mountDevtools() - // Delivery-durability resume is transparent: the resumable SSE connection - // adapter reattaches via the browser's native Last-Event-ID on reconnect. - // We only seed interrupt (state) resume from the client here. - syncResumeState() }) // Cleanup on unmount: stop any in-flight requests. @@ -298,32 +232,19 @@ export function useChat< content: string | MultimodalContent, sendOptions?: SendMessageOptions, ) => { - try { - await client().sendMessage(content, undefined, sendOptions) - } finally { - syncResumeState() - } + await client().sendMessage(content, undefined, sendOptions) } const append = async (message: ModelMessage | UIMessage) => { - try { - await client().append(message) - } finally { - syncResumeState() - } + await client().append(message) } const reload = async () => { - try { - await client().reload() - } finally { - syncResumeState() - } + await client().reload() } const loadOlderMessages = async () => { await client().loadOlderMessages() - setHasOlderMessages(client().getHasOlderMessages()) } const stop = () => { @@ -332,10 +253,11 @@ export function useChat< const clear = () => { client().clear() - syncResumeState() } - const setMessagesManually = (newMessages: Array>) => { + const setMessagesManually = ( + newMessages: ReadonlyArray>, + ) => { client().setMessagesManually(newMessages) } @@ -354,7 +276,6 @@ export function useChat< approved: boolean }) => { await client().addToolApprovalResponse(response) - syncResumeState() } const cancelQueued = (id: string) => client().cancelQueued(id) @@ -363,9 +284,7 @@ export function useChat< resumeItems: Array, state?: ChatResumeState, ) => { - const result = await client().resumeInterrupts(resumeItems, state) - syncResumeState() - return result + return await client().resumeInterrupts(resumeItems, state) } const resolveInterrupts = ( @@ -443,26 +362,26 @@ export function useChat< // oxlint-disable-next-line eslint-js/no-restricted-syntax -- primitive return shape diverges from generic UseChatReturn; TS can't structurally narrow the conditional partial/final fields return { messages, - queue, + queue: pick('queue'), sendMessage, append, reload, stop, - isLoading, - hasOlderMessages, + isLoading: pick('isLoading'), + hasOlderMessages: pick('hasOlderMessages'), loadOlderMessages, - error, - status, - isSubscribed, - connectionStatus, - sessionGenerating, - isHydrating, + error: pick('error'), + status: pick('status'), + isSubscribed: pick('isSubscribed'), + connectionStatus: pick('connectionStatus'), + sessionGenerating: pick('sessionGenerating'), + isHydrating: pick('isHydrating'), setMessages: setMessagesManually, clear, addToolResult, addToolApprovalResponse, cancelQueued, - runId, + runId: pick('runId'), interrupts, pendingInterrupts, interruptErrors, diff --git a/packages/ai-solid/src/use-generate-video.ts b/packages/ai-solid/src/use-generate-video.ts index 42acd27e53..721401311c 100644 --- a/packages/ai-solid/src/use-generate-video.ts +++ b/packages/ai-solid/src/use-generate-video.ts @@ -1,6 +1,7 @@ import { VideoGenerationClient } from '@tanstack/ai-client' import { createVideoDevtoolsBridge } from '@tanstack/ai-client/devtools' import { + batch, createEffect, createSignal, onCleanup, @@ -244,27 +245,6 @@ export function useGenerateVideo( onStatusUpdate: (s: VideoStatusInfo) => { if (!disposed) options.onStatusUpdate?.(s) }, - onResultChange: (r: TOutput | null) => { - if (!disposed) setResult(() => r) - }, - onLoadingChange: (l: boolean) => { - if (!disposed) setIsLoading(l) - }, - onErrorChange: (e: Error | undefined) => { - if (!disposed) setError(e) - }, - onStatusChange: (s: GenerationClientState) => { - if (!disposed) setStatus(s) - }, - onJobIdChange: (id: string | null) => { - if (!disposed) setJobId(id) - }, - onVideoStatusChange: (s: VideoStatusInfo | null) => { - if (!disposed) setVideoStatus(s) - }, - onResumeStateChange: (rs: { runId: string } | null) => { - if (!disposed) setRunId(rs?.runId ?? null) - }, } if (options.connection) { @@ -286,6 +266,22 @@ export function useGenerateVideo( ) }) + const applySnapshot = () => { + const next = client.getSnapshot() + // One batch, so effects never see half of a snapshot. + batch(() => { + setResult(() => next.result) + setIsLoading(next.isLoading) + setError(next.error) + setStatus(next.status) + setRunId(next.runId) + setJobId(next.jobId) + setVideoStatus(next.videoStatus) + }) + } + applySnapshot() + onCleanup(client.subscribe(applySnapshot)) + // Sync body changes without recreating client createEffect(() => { const currentBody = options.body diff --git a/packages/ai-solid/src/use-generation.ts b/packages/ai-solid/src/use-generation.ts index eedc8eb5c2..3a89d4c4c6 100644 --- a/packages/ai-solid/src/use-generation.ts +++ b/packages/ai-solid/src/use-generation.ts @@ -1,6 +1,7 @@ import { GenerationClient } from '@tanstack/ai-client' import { createGenerationDevtoolsBridge } from '@tanstack/ai-client/devtools' import { + batch, createEffect, createSignal, onCleanup, @@ -230,21 +231,6 @@ export function useGeneration< onChunk: (c: StreamChunk) => { if (!disposed) options.onChunk?.(c) }, - onResultChange: (r) => { - if (!disposed) setResult(() => r) - }, - onLoadingChange: (l) => { - if (!disposed) setIsLoading(l) - }, - onErrorChange: (e) => { - if (!disposed) setError(e) - }, - onStatusChange: (s) => { - if (!disposed) setStatus(s) - }, - onResumeStateChange: (rs) => { - if (!disposed) setRunId(rs?.runId ?? null) - }, } const persistenceProps = @@ -280,6 +266,20 @@ export function useGeneration< ) }) + const applySnapshot = () => { + const next = client.getSnapshot() + // One batch, so effects never see half of a snapshot. + batch(() => { + setResult(() => next.result) + setIsLoading(next.isLoading) + setError(next.error) + setStatus(next.status) + setRunId(next.runId) + }) + } + applySnapshot() + onCleanup(client.subscribe(applySnapshot)) + // Sync body changes without recreating client createEffect(() => { const currentBody = options.body diff --git a/packages/ai-solid/tests/use-chat.test.ts b/packages/ai-solid/tests/use-chat.test.ts index e4927b8fc1..04f2a39857 100644 --- a/packages/ai-solid/tests/use-chat.test.ts +++ b/packages/ai-solid/tests/use-chat.test.ts @@ -154,14 +154,20 @@ describe('useChat', () => { removeItem: vi.fn(), } - const { result } = renderUseChat({ - connection: adapter, - threadId: 'persisted-chat', - persistence: persistence, + let firstMessages: ReadonlyArray | undefined + const rendered = renderHook(() => { + const chat = useChat({ + connection: adapter, + threadId: 'persisted-chat', + persistence, + }) + firstMessages ??= chat.messages() + return chat }) + expect(firstMessages).toEqual(persistedMessages) await waitFor(() => { - expect(result.current.messages).toEqual(persistedMessages) + expect(rendered.result.messages()).toEqual(persistedMessages) }) expect(persistence.getItem).toHaveBeenCalledWith('persisted-chat') }) diff --git a/packages/ai-svelte/src/create-chat.svelte.ts b/packages/ai-svelte/src/create-chat.svelte.ts index ea0176406c..b82cb69276 100644 --- a/packages/ai-svelte/src/create-chat.svelte.ts +++ b/packages/ai-svelte/src/create-chat.svelte.ts @@ -2,13 +2,9 @@ import { ChatClient } from '@tanstack/ai-client' import { createChatDevtoolsBridge } from '@tanstack/ai-client/devtools' import { onMount } from 'svelte' import type { - ChatClientState, ResolvableChatInterrupt, - ChatInterruptState, ChatResumeState, - ConnectionStatus, InferredClientContext, - QueuedMessage, SendMessageOptions, StructuredOutputPart, } from '@tanstack/ai-client' @@ -29,9 +25,6 @@ import type { UIMessage, } from './types' -const EMPTY_INTERRUPTS = Object.freeze([]) -const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) - /** * Creates a reactive chat instance for Svelte 5. * @@ -72,25 +65,6 @@ export function createChat< >( options: CreateChatOptions, ): CreateChatReturn { - // Create reactive state using Svelte 5 runes - let messages = $state>>(options.initialMessages || []) - let isLoading = $state(false) - let hasOlderMessages = $state(false) - let error = $state(undefined) - let status = $state('ready') - let isSubscribed = $state(false) - let connectionStatus = $state('disconnected') - let sessionGenerating = $state(false) - let isHydrating = $state(false) - let queue = $state>([]) - let runId = $state(null) - let interruptState = $state.raw>({ - interrupts: EMPTY_INTERRUPTS, - pendingInterrupts: EMPTY_INTERRUPTS, - interruptErrors: EMPTY_INTERRUPT_ERRORS, - resuming: false, - }) - // Structured-output `partial` / `final` are derived from `messages` — // specifically from the structured-output part on the latest assistant // message (the one after the most recent user message). Per-turn parts @@ -175,53 +149,17 @@ export function createChat< ...(options.streamProcessor !== undefined && { streamProcessor: options.streamProcessor, }), - onMessagesChange: (newMessages: Array>) => { - messages = newMessages - hasOlderMessages = client.getHasOlderMessages() - }, - onLoadingChange: (newIsLoading: boolean) => { - isLoading = newIsLoading - syncResumeState() - }, - onStatusChange: (newStatus: ChatClientState) => { - status = newStatus - }, - onErrorChange: (newError: Error | undefined) => { - error = newError - }, - onSubscriptionChange: (nextIsSubscribed: boolean) => { - isSubscribed = nextIsSubscribed - }, - onConnectionStatusChange: (nextStatus: ConnectionStatus) => { - connectionStatus = nextStatus - }, - onSessionGeneratingChange: (isGenerating: boolean) => { - sessionGenerating = isGenerating - }, - onHydratingChange: (nextIsHydrating: boolean) => { - isHydrating = nextIsHydrating - }, ...(options.queue !== undefined && { queue: options.queue }), - onQueueChange: (nextQueue: Array) => { - queue = nextQueue - }, - onRunIdChange: (nextRunId) => { - runId = nextRunId - }, onInterruptStateChange: (nextInterruptState, context) => { - interruptState = nextInterruptState options.onInterruptStateChange?.(nextInterruptState, context) }, }) - function syncResumeState() { - runId = client.getCurrentRunId() - interruptState = client.getInterruptState() - } - - messages = client.getMessages() - isHydrating = client.getIsHydrating() - interruptState = client.getInterruptState() + // One raw rune: the snapshot is frozen and replaced on every change. + let snapshot = $state.raw(client.getSnapshot()) + const unsubscribeSnapshot = client.subscribeSnapshot(() => { + snapshot = client.getSnapshot() + }) if (options.live) { client.subscribe() @@ -245,7 +183,6 @@ export function createChat< // Delivery-durability resume is transparent: the resumable SSE // connection adapter reattaches via the browser's native // Last-Event-ID on reconnect. We only seed interrupt (state) resume. - syncResumeState() client.attach() // ONLY THE VIEW ON SCREEN HOLDS A STREAM. `onMount`'s returned function // runs when the component is destroyed, which is the one automatic @@ -257,6 +194,7 @@ export function createChat< // `detach` keeps the transcript and the resume pointer, so re-entering // the view picks the run back up from the durable log. return () => { + unsubscribeSnapshot() client.detach() } }) @@ -274,34 +212,21 @@ export function createChat< content: string | MultimodalContent, sendOptions?: SendMessageOptions, ) => { - try { - await client.sendMessage(content, undefined, sendOptions) - } finally { - syncResumeState() - } + await client.sendMessage(content, undefined, sendOptions) } const cancelQueued = (id: string) => client.cancelQueued(id) const append = async (message: ModelMessage | UIMessage) => { - try { - await client.append(message) - } finally { - syncResumeState() - } + await client.append(message) } const reload = async () => { - try { - await client.reload() - } finally { - syncResumeState() - } + await client.reload() } const loadOlderMessages = async () => { await client.loadOlderMessages() - hasOlderMessages = client.getHasOlderMessages() } const stop = () => { @@ -309,15 +234,15 @@ export function createChat< } const dispose = () => { + unsubscribeSnapshot() client.dispose() } const clear = () => { client.clear() - syncResumeState() } - const setMessages = (newMessages: Array>) => { + const setMessages = (newMessages: ReadonlyArray>) => { client.setMessagesManually(newMessages) } @@ -336,16 +261,13 @@ export function createChat< approved: boolean }) => { await client.addToolApprovalResponse(response) - syncResumeState() } const resumeInterrupts = async ( resumeItems: Array, state?: ChatResumeState, ) => { - const result = await client.resumeInterrupts(resumeItems, state) - syncResumeState() - return result + return await client.resumeInterrupts(resumeItems, state) } const resolveInterrupts = ( @@ -399,6 +321,7 @@ export function createChat< // multi-line scan re-runs whenever `messages` changes. const activeStructuredPart: StructuredOutputPart | null = $derived.by(() => { let lastUserIndex = -1 + const messages = snapshot.messages for (let i = messages.length - 1; i >= 0; i--) { if (messages[i]?.role === 'user') { lastUserIndex = i @@ -434,49 +357,49 @@ export function createChat< // oxlint-disable-next-line eslint-js/no-restricted-syntax -- rune return shape diverges from generic CreateChatReturn due to TSchema conditional partial/final fields; TS can't structurally narrow. return { get messages() { - return messages + return snapshot.messages }, get isLoading() { - return isLoading + return snapshot.isLoading }, get hasOlderMessages() { - return hasOlderMessages + return snapshot.hasOlderMessages }, get error() { - return error + return snapshot.error }, get status() { - return status + return snapshot.status }, get isSubscribed() { - return isSubscribed + return snapshot.isSubscribed }, get connectionStatus() { - return connectionStatus + return snapshot.connectionStatus }, get sessionGenerating() { - return sessionGenerating + return snapshot.sessionGenerating }, get isHydrating() { - return isHydrating + return snapshot.isHydrating }, get queue() { - return queue + return snapshot.queue }, get runId() { - return runId + return snapshot.runId }, get interrupts() { - return interruptState.interrupts + return snapshot.interruptState.interrupts }, get pendingInterrupts() { - return interruptState.interrupts + return snapshot.interruptState.interrupts }, get interruptErrors() { - return interruptState.interruptErrors + return snapshot.interruptState.interruptErrors }, get resuming() { - return interruptState.resuming + return snapshot.interruptState.resuming }, get partial() { return partial diff --git a/packages/ai-svelte/src/create-generate-video.svelte.ts b/packages/ai-svelte/src/create-generate-video.svelte.ts index d8b8615bcf..4448245f8e 100644 --- a/packages/ai-svelte/src/create-generate-video.svelte.ts +++ b/packages/ai-svelte/src/create-generate-video.svelte.ts @@ -235,34 +235,6 @@ export function createGenerateVideo( onStatusUpdate: (s: VideoStatusInfo) => { if (!disposed) options.onStatusUpdate?.(s) }, - onResultChange: (r: TOutput | null) => { - if (disposed) return - result = r - }, - onLoadingChange: (l: boolean) => { - if (disposed) return - isLoading = l - }, - onErrorChange: (e: Error | undefined) => { - if (disposed) return - error = e - }, - onStatusChange: (s: GenerationClientState) => { - if (disposed) return - status = s - }, - onJobIdChange: (id: string | null) => { - if (disposed) return - jobId = id - }, - onVideoStatusChange: (s: VideoStatusInfo | null) => { - if (disposed) return - videoStatus = s - }, - onResumeStateChange: (rs: { runId: string } | null) => { - if (disposed) return - runId = rs?.runId ?? null - }, } let client: VideoGenerationClient @@ -283,6 +255,19 @@ export function createGenerateVideo( ) } + const applySnapshot = () => { + const next = client.getSnapshot() + result = next.result + isLoading = next.isLoading + error = next.error + status = next.status + runId = next.runId + jobId = next.jobId + videoStatus = next.videoStatus + } + applySnapshot() + let unsubscribeSnapshot = client.subscribe(applySnapshot) + // Mount devtools only. Generation runs are never auto-started on setup — // persisted state is read-only for display. client.mountDevtools() @@ -297,7 +282,11 @@ export function createGenerateVideo( // frameworks revive via mountDevtools() in their mount effects), so an // explicit generate() after dispose() is the Svelte revive path: bring // the client and the reactive bindings back together. - disposed = false + if (disposed) { + disposed = false + applySnapshot() + unsubscribeSnapshot = client.subscribe(applySnapshot) + } client.mountDevtools() await client.generate(input) } @@ -312,6 +301,8 @@ export function createGenerateVideo( const dispose = () => { disposed = true + unsubscribeSnapshot() + unsubscribeSnapshot = () => {} client.dispose() } diff --git a/packages/ai-svelte/src/create-generation.svelte.ts b/packages/ai-svelte/src/create-generation.svelte.ts index bf99c9dd3b..1c659dd2f8 100644 --- a/packages/ai-svelte/src/create-generation.svelte.ts +++ b/packages/ai-svelte/src/create-generation.svelte.ts @@ -233,26 +233,6 @@ export function createGeneration< onChunk: (c: StreamChunk) => { if (!disposed) options.onChunk?.(c) }, - onResultChange: (r: TOutput | null) => { - if (disposed) return - result = r - }, - onLoadingChange: (l: boolean) => { - if (disposed) return - isLoading = l - }, - onErrorChange: (e: Error | undefined) => { - if (disposed) return - error = e - }, - onStatusChange: (s: GenerationClientState) => { - if (disposed) return - status = s - }, - onResumeStateChange: (rs) => { - if (disposed) return - runId = rs?.runId ?? null - }, } const persistenceProps = @@ -287,6 +267,17 @@ export function createGeneration< ) } + const applySnapshot = () => { + const next = client.getSnapshot() + result = next.result + isLoading = next.isLoading + error = next.error + status = next.status + runId = next.runId + } + applySnapshot() + let unsubscribeSnapshot = client.subscribe(applySnapshot) + // Mount devtools only. Generation runs are never auto-started on setup — // persisted state is read-only for display. client.mountDevtools() @@ -301,7 +292,11 @@ export function createGeneration< // frameworks revive via mountDevtools() in their mount effects), so an // explicit generate() after dispose() is the Svelte revive path: bring // the client and the reactive bindings back together. - disposed = false + if (disposed) { + disposed = false + applySnapshot() + unsubscribeSnapshot = client.subscribe(applySnapshot) + } client.mountDevtools() await client.generate(input) } @@ -316,6 +311,8 @@ export function createGeneration< const dispose = () => { disposed = true + unsubscribeSnapshot() + unsubscribeSnapshot = () => {} client.dispose() } diff --git a/packages/ai-svelte/src/types.ts b/packages/ai-svelte/src/types.ts index 894e10267b..868a41cdac 100644 --- a/packages/ai-svelte/src/types.ts +++ b/packages/ai-svelte/src/types.ts @@ -153,7 +153,7 @@ interface BaseCreateChatReturn< * 'structured-output')` is typed by the schema — `data: T`, * `partial: DeepPartial`. */ - readonly messages: Array> + readonly messages: ReadonlyArray> /** * Send a message and get a response. @@ -169,7 +169,7 @@ interface BaseCreateChatReturn< /** * Pending messages queued while a stream is in flight. */ - readonly queue: Array + readonly queue: ReadonlyArray /** * Cancel a queued message before it drains. No-op if already sent. @@ -279,7 +279,7 @@ interface BaseCreateChatReturn< /** * Set messages manually */ - setMessages: (messages: Array>) => void + setMessages: (messages: ReadonlyArray>) => void /** * Clear all messages diff --git a/packages/ai-svelte/tests/create-generation-dispose.test.ts b/packages/ai-svelte/tests/create-generation-dispose.test.ts new file mode 100644 index 0000000000..31920ea009 --- /dev/null +++ b/packages/ai-svelte/tests/create-generation-dispose.test.ts @@ -0,0 +1,41 @@ +import { describe, expect, it } from 'vitest' +import { createGeneration } from '../src/create-generation.svelte' +import { createGenerateVideo } from '../src/create-generate-video.svelte' + +describe('Svelte generate() after dispose() rebinds the snapshot', () => { + it('createGeneration updates result after dispose then generate', async () => { + const gen = createGeneration<{ prompt: string }, { id: string }>({ + fetcher: async (input) => ({ id: input.prompt }), + }) + + await gen.generate({ prompt: 'a' }) + expect(gen.result).toEqual({ id: 'a' }) + + gen.dispose() + await gen.generate({ prompt: 'b' }) + expect(gen.result).toEqual({ id: 'b' }) + }) + + it('createGenerateVideo updates result after dispose then generate', async () => { + const video = createGenerateVideo({ + fetcher: async (input) => { + if (typeof input.prompt !== 'string') { + throw new Error('Expected a text prompt') + } + + return { + jobId: input.prompt, + status: 'completed' as const, + url: `https://example.com/${input.prompt}.mp4`, + } + }, + }) + + await video.generate({ prompt: 'one' }) + expect(video.result?.jobId).toBe('one') + + video.dispose() + await video.generate({ prompt: 'two' }) + expect(video.result?.jobId).toBe('two') + }) +}) diff --git a/packages/ai-svelte/tests/use-chat.test.ts b/packages/ai-svelte/tests/use-chat.test.ts index 8a0bd123ff..c9219c6fed 100644 --- a/packages/ai-svelte/tests/use-chat.test.ts +++ b/packages/ai-svelte/tests/use-chat.test.ts @@ -419,4 +419,28 @@ describe('createChat', () => { await promise.catch(() => {}) }) }) + + describe('snapshot', () => { + it('publishes frozen messages and stops updating after dispose', async () => { + const chat = createChat({ + connection: createMockConnectionAdapter({ + chunks: createTextChunks('Hi'), + }), + }) + + await chat.sendMessage('Hello') + const messages = chat.messages + expect(messages).toHaveLength(2) + expect(Object.isFrozen(messages)).toBe(true) + expect(() => (messages as Array).push({})).toThrow(TypeError) + + // The frozen list goes back in as it is. + chat.setMessages(messages) + expect(chat.messages).toHaveLength(2) + + chat.dispose() + chat.setMessages([]) + expect(chat.messages).toHaveLength(2) + }) + }) }) diff --git a/packages/ai-utils/src/index.ts b/packages/ai-utils/src/index.ts index 619338eee6..abedd7a63d 100644 --- a/packages/ai-utils/src/index.ts +++ b/packages/ai-utils/src/index.ts @@ -7,3 +7,4 @@ export { base64ToArrayBuffer, base64ToUint8Array, } from './base64' +export { jsonModeSystemPrompt } from './json-mode' diff --git a/packages/ai-utils/src/json-mode.ts b/packages/ai-utils/src/json-mode.ts new file mode 100644 index 0000000000..34a77997a1 --- /dev/null +++ b/packages/ai-utils/src/json-mode.ts @@ -0,0 +1,8 @@ +/** + * System prompt for JSON mode (`json_object`). JSON mode does not send the + * schema, so the prompt carries it. The prompt also has the word "JSON", + * which some providers require in JSON mode. + */ +export function jsonModeSystemPrompt(schema: unknown): string { + return `Respond with only a JSON object that matches this JSON Schema:\n${JSON.stringify(schema)}` +} diff --git a/packages/ai-vue/src/types.ts b/packages/ai-vue/src/types.ts index dcb68e85c1..cde837e3ef 100644 --- a/packages/ai-vue/src/types.ts +++ b/packages/ai-vue/src/types.ts @@ -157,7 +157,7 @@ interface BaseUseChatReturn< * `messages[i].parts.find(p => p.type === 'structured-output')` is typed * by the schema — `data: T`, `partial: DeepPartial`. */ - messages: Readonly>>> + messages: Readonly>>> /** * Send a message and get a response. @@ -173,7 +173,7 @@ interface BaseUseChatReturn< /** * Pending messages queued while a stream is in flight. */ - queue: Readonly>> + queue: Readonly>> /** * Cancel a queued message before it drains. No-op if already sent. @@ -276,7 +276,7 @@ interface BaseUseChatReturn< /** * Set messages manually */ - setMessages: (messages: Array>) => void + setMessages: (messages: ReadonlyArray>) => void /** * Clear all messages diff --git a/packages/ai-vue/src/use-chat.ts b/packages/ai-vue/src/use-chat.ts index 08fe88e942..ec1186b6de 100644 --- a/packages/ai-vue/src/use-chat.ts +++ b/packages/ai-vue/src/use-chat.ts @@ -19,13 +19,9 @@ import type { StreamChunk, } from '@tanstack/ai' import type { - ChatClientState, ResolvableChatInterrupt, - ChatInterruptState, ChatResumeState, - ConnectionStatus, InferredClientContext, - QueuedMessage, SendMessageOptions, StructuredOutputPart, } from '@tanstack/ai-client' @@ -37,9 +33,6 @@ import type { UseChatReturn, } from './types' -const EMPTY_INTERRUPTS = Object.freeze([]) -const EMPTY_INTERRUPT_ERRORS = Object.freeze([]) - export function useChat< const TTools extends ReadonlyArray = any, TSchema extends SchemaInput | undefined = undefined, @@ -55,26 +48,6 @@ export function useChat< TInterrupts > = {} as UseChatOptions, ): UseChatReturn { - const messages = shallowRef>>( - options.initialMessages || [], - ) - const isLoading = shallowRef(false) - const hasOlderMessages = shallowRef(false) - const error = shallowRef(undefined) - const status = shallowRef('ready') - const isSubscribed = shallowRef(false) - const connectionStatus = shallowRef('disconnected') - const sessionGenerating = shallowRef(false) - const isHydrating = shallowRef(false) - const queue = shallowRef>([]) - const runId = shallowRef(null) - const interruptState = shallowRef>({ - interrupts: EMPTY_INTERRUPTS, - pendingInterrupts: EMPTY_INTERRUPTS, - interruptErrors: EMPTY_INTERRUPT_ERRORS, - resuming: false, - }) - // Structured-output `partial` / `final` are derived from `messages` — // specifically from the structured-output part on the latest assistant // message (the one after the most recent user message). This keeps @@ -162,53 +135,21 @@ export function useChat< ...(options.streamProcessor !== undefined && { streamProcessor: options.streamProcessor, }), - onMessagesChange: (newMessages: Array>) => { - messages.value = newMessages - hasOlderMessages.value = client.getHasOlderMessages() - }, - onLoadingChange: (newIsLoading: boolean) => { - isLoading.value = newIsLoading - syncResumeState() - }, - onStatusChange: (newStatus: ChatClientState) => { - status.value = newStatus - }, - onErrorChange: (newError: Error | undefined) => { - error.value = newError - }, - onSubscriptionChange: (nextIsSubscribed: boolean) => { - isSubscribed.value = nextIsSubscribed - }, - onConnectionStatusChange: (nextStatus: ConnectionStatus) => { - connectionStatus.value = nextStatus - }, - onSessionGeneratingChange: (isGenerating: boolean) => { - sessionGenerating.value = isGenerating - }, - onHydratingChange: (nextIsHydrating: boolean) => { - isHydrating.value = nextIsHydrating - }, ...(options.queue !== undefined && { queue: options.queue }), - onQueueChange: (nextQueue: Array) => { - queue.value = nextQueue - }, - onRunIdChange: (nextRunId) => { - runId.value = nextRunId - }, onInterruptStateChange: (nextInterruptState, context) => { - interruptState.value = nextInterruptState options.onInterruptStateChange?.(nextInterruptState, context) }, }) - function syncResumeState() { - runId.value = client.getCurrentRunId() - interruptState.value = client.getInterruptState() - } - - messages.value = client.getMessages() - isHydrating.value = client.getIsHydrating() - interruptState.value = client.getInterruptState() + // One shallow ref: the snapshot is frozen and replaced on every change. + const snapshot = shallowRef(client.getSnapshot()) + onScopeDispose( + client.subscribeSnapshot(() => { + snapshot.value = client.getSnapshot() + }), + ) + const pick = (key: TKey) => + readonly(computed(() => snapshot.value[key])) // Sync body / forwardedProps changes to the client. // Both populate the same wire payload; `forwardedProps` is preferred @@ -261,10 +202,6 @@ export function useChat< // connections per origin until the page reloaded. client.attach() client.mountDevtools() - // Delivery-durability resume is transparent: the resumable SSE connection - // adapter reattaches via the browser's native Last-Event-ID on reconnect. - // We only seed interrupt (state) resume from the client here. - syncResumeState() }) // Cleanup on unmount: stop any in-flight requests @@ -288,34 +225,21 @@ export function useChat< content: string | MultimodalContent, sendOptions?: SendMessageOptions, ) => { - try { - await client.sendMessage(content, undefined, sendOptions) - } finally { - syncResumeState() - } + await client.sendMessage(content, undefined, sendOptions) } const cancelQueued = (id: string) => client.cancelQueued(id) const append = async (message: ModelMessage | UIMessage) => { - try { - await client.append(message) - } finally { - syncResumeState() - } + await client.append(message) } const reload = async () => { - try { - await client.reload() - } finally { - syncResumeState() - } + await client.reload() } const loadOlderMessages = async () => { await client.loadOlderMessages() - hasOlderMessages.value = client.getHasOlderMessages() } const stop = () => { @@ -324,10 +248,11 @@ export function useChat< const clear = () => { client.clear() - syncResumeState() } - const setMessagesManually = (newMessages: Array>) => { + const setMessagesManually = ( + newMessages: ReadonlyArray>, + ) => { client.setMessagesManually(newMessages) } @@ -346,22 +271,23 @@ export function useChat< approved: boolean }) => { await client.addToolApprovalResponse(response) - syncResumeState() } const resumeInterrupts = async ( resumeItems: Array, state?: ChatResumeState, ) => { - const result = await client.resumeInterrupts(resumeItems, state) - syncResumeState() - return result + return await client.resumeInterrupts(resumeItems, state) } - const interrupts = computed(() => interruptState.value.interrupts) - const pendingInterrupts = computed(() => interruptState.value.interrupts) - const interruptErrors = computed(() => interruptState.value.interruptErrors) - const resuming = computed(() => interruptState.value.resuming) + const interrupts = computed(() => snapshot.value.interruptState.interrupts) + const pendingInterrupts = computed( + () => snapshot.value.interruptState.interrupts, + ) + const interruptErrors = computed( + () => snapshot.value.interruptState.interruptErrors, + ) + const resuming = computed(() => snapshot.value.interruptState.resuming) const resolveInterrupts = ( resolution: @@ -398,7 +324,7 @@ export function useChat< // otherwise a `final` from a previous session would leak into the value // on first render. const activeStructuredPart = computed(() => { - const list = messages.value + const list = snapshot.value.messages let lastUserIndex = -1 for (let i = list.length - 1; i >= 0; i--) { if (list[i]?.role === 'user') { @@ -436,27 +362,27 @@ export function useChat< // didn't supply `outputSchema`. // oxlint-disable-next-line eslint-js/no-restricted-syntax -- composable return shape diverges from conditional UseChatReturn; TS can't structurally narrow the TSchema-gated partial/final refs return { - messages: readonly(messages), + messages: pick('messages'), sendMessage, - queue: readonly(queue), + queue: pick('queue'), cancelQueued, append, reload, stop, - isLoading: readonly(isLoading), - hasOlderMessages: readonly(hasOlderMessages), + isLoading: pick('isLoading'), + hasOlderMessages: pick('hasOlderMessages'), loadOlderMessages, - error: readonly(error), - status: readonly(status), - isSubscribed: readonly(isSubscribed), - connectionStatus: readonly(connectionStatus), - sessionGenerating: readonly(sessionGenerating), - isHydrating: readonly(isHydrating), + error: pick('error'), + status: pick('status'), + isSubscribed: pick('isSubscribed'), + connectionStatus: pick('connectionStatus'), + sessionGenerating: pick('sessionGenerating'), + isHydrating: pick('isHydrating'), setMessages: setMessagesManually, clear, addToolResult, addToolApprovalResponse, - runId: readonly(runId), + runId: pick('runId'), interrupts: readonly(interrupts), pendingInterrupts: readonly(pendingInterrupts), interruptErrors: readonly(interruptErrors), diff --git a/packages/ai-vue/src/use-generate-video.ts b/packages/ai-vue/src/use-generate-video.ts index fb2ac7a2a7..66ae22cd0a 100644 --- a/packages/ai-vue/src/use-generate-video.ts +++ b/packages/ai-vue/src/use-generate-video.ts @@ -231,34 +231,6 @@ export function useGenerateVideo( onStatusUpdate: (s: VideoStatusInfo) => { if (!disposed) options.onStatusUpdate?.(s) }, - onResultChange: (r: TOutput | null) => { - if (disposed) return - result.value = r - }, - onLoadingChange: (l: boolean) => { - if (disposed) return - isLoading.value = l - }, - onErrorChange: (e: Error | undefined) => { - if (disposed) return - error.value = e - }, - onStatusChange: (s: GenerationClientState) => { - if (disposed) return - status.value = s - }, - onJobIdChange: (id: string | null) => { - if (disposed) return - jobId.value = id - }, - onVideoStatusChange: (s: VideoStatusInfo | null) => { - if (disposed) return - videoStatus.value = s - }, - onResumeStateChange: (rs: { runId: string } | null) => { - if (disposed) return - runId.value = rs?.runId ?? null - }, } let client: VideoGenerationClient @@ -279,6 +251,19 @@ export function useGenerateVideo( ) } + const applySnapshot = () => { + const next = client.getSnapshot() + result.value = next.result + isLoading.value = next.isLoading + error.value = next.error + status.value = next.status + runId.value = next.runId + jobId.value = next.jobId + videoStatus.value = next.videoStatus + } + applySnapshot() + onScopeDispose(client.subscribe(applySnapshot)) + // Sync body changes to the client. // Conditional spread: `updateOptions` declares `body?: Record` // (strict optional) and rejects explicit `undefined` under EOPT. diff --git a/packages/ai-vue/src/use-generation.ts b/packages/ai-vue/src/use-generation.ts index cbec8f435f..2095da2d28 100644 --- a/packages/ai-vue/src/use-generation.ts +++ b/packages/ai-vue/src/use-generation.ts @@ -220,26 +220,6 @@ export function useGeneration< onChunk: (c: StreamChunk) => { if (!disposed) options.onChunk?.(c) }, - onResultChange: (r: TOutput | null) => { - if (disposed) return - result.value = r - }, - onLoadingChange: (l: boolean) => { - if (disposed) return - isLoading.value = l - }, - onErrorChange: (e: Error | undefined) => { - if (disposed) return - error.value = e - }, - onStatusChange: (s: GenerationClientState) => { - if (disposed) return - status.value = s - }, - onResumeStateChange: (rs) => { - if (disposed) return - runId.value = rs?.runId ?? null - }, } const persistenceProps = @@ -274,6 +254,17 @@ export function useGeneration< ) } + const applySnapshot = () => { + const next = client.getSnapshot() + result.value = next.result + isLoading.value = next.isLoading + error.value = next.error + status.value = next.status + runId.value = next.runId + } + applySnapshot() + onScopeDispose(client.subscribe(applySnapshot)) + // Sync body changes to the client. // Conditional spread: `updateOptions` declares `body?: Record` // (strict optional) and rejects explicit `undefined` under EOPT. diff --git a/packages/ai-vue/tests/use-chat.test.ts b/packages/ai-vue/tests/use-chat.test.ts index 501e9f4e01..63a6533e0b 100644 --- a/packages/ai-vue/tests/use-chat.test.ts +++ b/packages/ai-vue/tests/use-chat.test.ts @@ -1,5 +1,6 @@ import { ChatClient } from '@tanstack/ai-client' -import { flushPromises } from '@vue/test-utils' +import { flushPromises, mount } from '@vue/test-utils' +import { defineComponent } from 'vue' import { afterEach, describe, expect, it, vi } from 'vitest' import { createInterruptResumeSnapshot, @@ -8,6 +9,7 @@ import { createToolCallChunks, renderUseChat, } from './test-utils' +import { useChat } from '../src/use-chat' import type { UIMessage } from '../src/types' import type { ModelMessage } from '@tanstack/ai' @@ -1292,4 +1294,55 @@ describe('useChat', () => { }) }) }) + + describe('snapshot', () => { + it('reads one frozen client snapshot and drops it on scope dispose', async () => { + let resolveRead!: (value: { messages: Array }) => void + const read = new Promise<{ messages: Array }>((resolve) => { + resolveRead = resolve + }) + let chat!: ReturnType + const wrapper = mount( + defineComponent({ + setup() { + chat = useChat({ + connection: createMockConnectionAdapter({ + chunks: createTextChunks('Hi'), + }), + threadId: 'vue-snapshot', + persistence: { + getItem: () => read, + setItem: async () => {}, + removeItem: async () => {}, + }, + }) + return {} + }, + template: '
', + }), + ) + + expect(chat.isHydrating.value).toBe(true) + resolveRead({ + messages: [ + { id: 'm1', role: 'user', parts: [{ type: 'text', content: 'hi' }] }, + ], + }) + await vi.waitFor(() => expect(chat.isHydrating.value).toBe(false)) + expect(chat.messages.value).toHaveLength(1) + + await chat.sendMessage('Hello') + await flushPromises() + expect(chat.isLoading.value).toBe(false) + expect(chat.status.value).toBe('ready') + const messages = chat.messages.value + expect(messages).toHaveLength(3) + expect(Object.isFrozen(messages)).toBe(true) + expect(() => (messages as Array).push({})).toThrow(TypeError) + + wrapper.unmount() + chat.setMessages([]) + expect(chat.messages.value).toBe(messages) + }) + }) }) diff --git a/packages/ai/package.json b/packages/ai/package.json index 9f2eca7f77..81fd734828 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -102,6 +102,7 @@ ], "dependencies": { "@ag-ui/core": "1.0.0", + "@cfworker/json-schema": "^4.1.1", "@standard-schema/spec": "^1.1.0", "@tanstack/ai-event-client": "workspace:^", "@tanstack/ai-utils": "workspace:^", diff --git a/packages/ai/skills/ai-core/adapter-configuration/SKILL.md b/packages/ai/skills/ai-core/adapter-configuration/SKILL.md index a5803d164a..adceb6ced1 100644 --- a/packages/ai/skills/ai-core/adapter-configuration/SKILL.md +++ b/packages/ai/skills/ai-core/adapter-configuration/SKILL.md @@ -373,6 +373,20 @@ adapter for a custom model that doesn't support the combination, return > Note `seed-2-0-lite-260428` (the obvious default) is one of them; use > `seed-2-0-lite-260228` or `dola-seed-2-1-turbo-260628` for typed output. +> **JSON mode (`json_object`) is never combined.** Some OpenRouter and Groq +> models reject strict `json_schema` and accept only `json_object` +> (`OPENROUTER_JSON_OBJECT_ONLY_MODELS`, Groq models whose `features` list +> `json_object` without `json_schema`). For these, structured output sends +> `json_object` and adds the schema as a system prompt. A caller can also +> set JSON mode in `modelOptions` (`responseFormat` / `text.format` on +> OpenRouter, `response_format` / `text.format` on OpenAI-base adapters). +> In JSON mode `supportsCombinedToolsAndSchema()` returns `false`, so the +> schema prompt goes only on the separate finalization request. A subclass +> of the OpenAI-base adapters marks its own JSON-only models by overriding +> `protected modelPrefersJsonObject(model: string): boolean`. With a plain +> JSON Schema `outputSchema`, TanStack AI does not validate the reply; use +> a Standard Schema (Zod, Valibot, ArkType) to get it validated. + ### 6. OpenAI-Compatible Providers Any provider that implements the OpenAI **Chat Completions** API (DeepSeek, diff --git a/packages/ai/skills/ai-core/structured-outputs/SKILL.md b/packages/ai/skills/ai-core/structured-outputs/SKILL.md index 3bea8cbaea..3ab402722c 100644 --- a/packages/ai/skills/ai-core/structured-outputs/SKILL.md +++ b/packages/ai/skills/ai-core/structured-outputs/SKILL.md @@ -182,23 +182,25 @@ The terminal event is a `CUSTOM` chunk: `{ type: 'CUSTOM', name: 'structured-out **Adapter coverage for streaming:** -| Adapter | `outputSchema` + `stream: true` | -| ----------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------- | -| `@tanstack/ai-openai` (Responses + Chat Completions) | **Native combined mode (#605)** — schema wired into the regular `chatStream` call alongside `tools`; engine harvests JSON, no finalization round-trip | -| `@tanstack/ai-anthropic` (Claude 4.5+ only) | **Native combined mode (#605)** — `output_config.format` + `tools` in one beta Messages call. Older Claude models fall back | -| `@tanstack/ai-gemini` (Gemini 3.x only) | **Native combined mode (#605)** — `responseSchema` + `tools` in one `generateContentStream`. Gemini 2.x falls back | -| `@tanstack/ai-grok` | **Native combined mode (#605)** — OpenAI Responses `text.format` + `tools` for grok-4.6, grok-4.5, grok-4.3, and grok-build-0.1 | -| `@tanstack/ai-openrouter` | Native single-request stream (legacy `structuredOutputStream` path; per-call combined-mode lookup is a follow-up) | -| `@tanstack/ai-groq` | Legacy `structuredOutputStream` only (no tools — Groq's API rejects schema + tools + stream) | -| `@tanstack/ai-bedrock` | Separate native `structuredOutputStream` finalization through Converse or an OpenAI-compatible API | -| `@tanstack/ai-byteplus` | Native combined mode on supported models; unsupported models emit `RUN_ERROR` | -| `@tanstack/ai-cloudflare` | Native `structuredOutputStream` without tools; with tools, a separate finalization call (Workers AI models answer the tool turn in prose) | -| `@tanstack/ai-claude-code` | Combined + event source — `--json-schema` on the same harness turn. Read `useChat().final`. See Pattern 6. | -| `@tanstack/ai-codex` | Combined + event source — `--output-schema` on the same harness turn. Read `useChat().final`. See Pattern 6. | -| `@tanstack/ai-opencode` | Combined + event source — prompt-and-parse. Read `useChat().final`. See Pattern 6. | -| `@tanstack/ai-grok-build` | Combined + event source — prompt-and-parse (ACP and streaming-json). Read `useChat().final` or the `structured-output` part. See Pattern 6. | -| `@tanstack/ai-acp` (`acpCompatible`) | Combined + event source — prompt-and-parse. Read `useChat().final` or the `structured-output` part. See Pattern 6. | -| All other adapters (ollama, older Claude, Gemini 2.x) | Fallback: runs non-streaming `structuredOutput`, emits one `structured-output.complete` event | +| Adapter | `outputSchema` + `stream: true` | +| ---------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------- | +| `@tanstack/ai-openai` (Responses + Chat Completions) | **Native combined mode (#605)** — schema wired into the regular `chatStream` call alongside `tools`; engine harvests JSON, no finalization round-trip | +| `@tanstack/ai-anthropic` (Claude 4.5+ only) | **Native combined mode (#605)** — `output_config.format` + `tools` in one beta Messages call. Older Claude models fall back | +| `@tanstack/ai-gemini` (Gemini 3.x only) | **Native combined mode (#605)** — `responseSchema` + `tools` in one `generateContentStream`. Gemini 2.x falls back | +| `@tanstack/ai-grok` | **Native combined mode (#605)** — OpenAI Responses `text.format` + `tools` for grok-4.6, grok-4.5, grok-4.3, and grok-build-0.1 | +| `@tanstack/ai-openrouter` | Native single-request stream (legacy `structuredOutputStream` path; per-call combined-mode lookup is a follow-up) | +| `@tanstack/ai-groq` | Legacy `structuredOutputStream` only (no tools — Groq's API rejects schema + tools + stream) | +| `@tanstack/ai-bedrock` | Separate native `structuredOutputStream` finalization through Converse or an OpenAI-compatible API | +| `@tanstack/ai-byteplus` | Native combined mode on supported models; unsupported models emit `RUN_ERROR` | +| `@tanstack/ai-cloudflare` | Native `structuredOutputStream` without tools; with tools, a separate finalization call (Workers AI models answer the tool turn in prose) | +| `@tanstack/ai-claude-code` | Combined + event source — `--json-schema` on the same harness turn. Read `useChat().final`. See Pattern 6. | +| `@tanstack/ai-codex` | Combined + event source — `--output-schema` on the same harness turn. Read `useChat().final`. See Pattern 6. | +| `@tanstack/ai-opencode` | Combined + event source — prompt-and-parse. Read `useChat().final`. See Pattern 6. | +| `@tanstack/ai-grok-build` | Combined + event source — prompt-and-parse (ACP and streaming-json). Read `useChat().final` or the `structured-output` part. See Pattern 6. | +| `@tanstack/ai-acp` (`acpCompatible`) | Combined + event source — prompt-and-parse. Read `useChat().final` or the `structured-output` part. See Pattern 6. | + +**JSON mode (`json_object`).** Some OpenRouter and Groq models reject strict `json_schema`. For those models the adapter sends `json_object` and adds the schema as a system prompt. The model is asked for the shape but is not held to it. A caller can force JSON mode with `modelOptions` (`responseFormat: { type: 'json_object' }` on OpenRouter, `response_format` on Groq and other OpenAI-base adapters). JSON mode is never combined with tools: the schema prompt goes only on the finalization call. Only `await chat({ outputSchema })` with a Standard Schema validates the reply, so prefer Zod/Valibot/ArkType over a plain JSON Schema on these models. +| All other adapters (ollama, older Claude, Gemini 2.x) | Fallback: runs non-streaming `structuredOutput`, emits one `structured-output.complete` event | **Native-combined output vs separate finalization** is signaled by the adapter's optional `supportsCombinedToolsAndSchema(modelOptions)` method. When diff --git a/packages/ai/src/activities/chat/agents/bound.ts b/packages/ai/src/activities/chat/agents/bound.ts index 0b4008797c..6b4efb5e90 100644 --- a/packages/ai/src/activities/chat/agents/bound.ts +++ b/packages/ai/src/activities/chat/agents/bound.ts @@ -83,7 +83,11 @@ export interface SubagentBinding { parentToolCallId?: string }, options: { wake: true }, - ) => Promise<{ subagentRunId: string }> + ) => Promise<{ + subagentRunId: string + /** Settles when the child ends. `maxConcurrent` counts the child until then. */ + done: Promise + }> } /** diff --git a/packages/ai/src/activities/chat/agents/limits.ts b/packages/ai/src/activities/chat/agents/limits.ts index 7210880ec0..72ea9ea5ff 100644 --- a/packages/ai/src/activities/chat/agents/limits.ts +++ b/packages/ai/src/activities/chat/agents/limits.ts @@ -54,6 +54,11 @@ export class SubagentBudget { return undefined } + /** Undo one `reserve`. A background start that throws never began. */ + release(): void { + this.shared.calls = Math.max(0, this.shared.calls - 1) + } + /** Milliseconds a child started now may run, or `undefined` for no limit. */ childTimeout(now = Date.now()): number | undefined { const own = this.limits.timeoutMs diff --git a/packages/ai/src/activities/chat/agents/spawn.ts b/packages/ai/src/activities/chat/agents/spawn.ts index f8d5bd3753..ec54627d7a 100644 --- a/packages/ai/src/activities/chat/agents/spawn.ts +++ b/packages/ai/src/activities/chat/agents/spawn.ts @@ -1047,34 +1047,33 @@ function agentToolSchema( type: 'object', properties: { prompt: { type: 'string' } }, } - // ponytail: a schema with no top-level properties (a top-level anyOf) gets - // no sessionId or background. Wrap it in allOf if one needs them. - if (schema.properties === undefined || names.length === 0) return schema - return { - ...schema, - properties: { - ...schema.properties, - ...(names.includes('sessionId') && { - sessionId: { - type: 'string', - description: - 'The subagentRunId of an earlier result. The call continues that child.', - }, - }), - ...(names.includes('background') && { - background: { - type: 'boolean', - description: - 'true starts the child and returns its subagentRunId at once.', - }, - }), - }, + if (names.length === 0) return schema + const extra: JSONSchema['properties'] = { + ...(names.includes('sessionId') && { + sessionId: { + type: 'string', + description: + 'The subagentRunId of an earlier result. The call continues that child.', + }, + }), + ...(names.includes('background') && { + background: { + type: 'boolean', + description: + 'true starts the child and returns its subagentRunId at once.', + }, + }), } + // A top-level anyOf has no `properties`. Put the reserved fields beside it. + if (schema.properties === undefined) { + return { allOf: [schema, { type: 'object', properties: extra }] } + } + return { ...schema, properties: { ...schema.properties, ...extra } } } /** - * The input schema of an agent tool. The tool loop checks a Standard Schema - * before middleware runs, as for any tool. With reserved fields, a wrapper + * The input schema of an agent tool. The tool loop checks it after + * middleware, just before the tool runs, as for any tool. With reserved fields, a wrapper * checks the agent's fields and passes the reserved ones through. */ function agentToolInput( @@ -1092,21 +1091,34 @@ function agentToolInput( vendor: 'tanstack-ai', validate: (value) => { const { fields, reserved } = splitReserved(value, names) - // ponytail: a schema that turns the input into a non-object (a - // top-level transform) loses the reserved fields. - const add = (result: StandardSchemaV1.Result) => - result.issues !== undefined || - Object.keys(reserved).length === 0 || - typeof result.value !== 'object' || - result.value === null || - Array.isArray(result.value) - ? result - : { - value: { - ...(result.value as Record), - ...reserved, + const add = (result: StandardSchemaV1.Result) => { + if ( + result.issues !== undefined || + Object.keys(reserved).length === 0 + ) { + return result + } + if ( + typeof result.value !== 'object' || + result.value === null || + Array.isArray(result.value) + ) { + return { + issues: [ + { + message: + 'Agent input must be an object when the call uses sessionId or background.', }, - } + ], + } + } + return { + value: { + ...(result.value as Record), + ...reserved, + }, + } + } const result = inner['~standard'].validate(fields) return result instanceof Promise ? result.then(add) : add(result) }, @@ -1215,7 +1227,9 @@ export function createSyntheticSubagentTools( const refusal = suspended ? undefined : timeout === 0 - ? `subagent limit reached (timeoutMs ${budget.limits.timeoutMs ?? 0})` + ? budget.limits.timeoutMs === 0 + ? 'subagent limit reached (timeoutMs 0)' + : 'subagent limit reached (no time left)' : budget.reserve(active) if (refusal !== undefined) { return { @@ -1361,16 +1375,28 @@ export function createSyntheticSubagentTools( // A background child counts against the tree budget too. const refusal = budget.reserve(active) if (refusal !== undefined) throw new Error(refusal) - const { subagentRunId } = await start( - { - agent: agent.name, - ...fields, - ...(context?.toolCallId !== undefined && { - parentToolCallId: context.toolCallId, - }), - }, - { wake: true }, - ) + active += 1 + let started: { subagentRunId: string; done: Promise } + try { + started = await start( + { + agent: agent.name, + ...fields, + ...(context?.toolCallId !== undefined && { + parentToolCallId: context.toolCallId, + }), + }, + { wake: true }, + ) + } catch (error) { + active -= 1 + budget.release() + throw error + } + void Promise.resolve(started.done).finally(() => { + active -= 1 + }) + const { subagentRunId } = started return { subagentRunId, text: '', diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 31b483eb63..5fead436a2 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -85,6 +85,7 @@ import { SUBAGENT_PLAN_KEY, readSubagentTurn } from './agents/turn' import { MiddlewareAbortError, ToolCallManager, + approvalResolution, executeToolCalls, } from './tools/tool-calls' import { @@ -136,8 +137,11 @@ import type { import type { ApprovalRequest, ClientToolRequest, + ExecuteToolCallsResult, McpInputRequest, + ToolExecutionMiddlewareHooks, ToolResult, + ToolResumeExecutionState, } from './tools/tool-calls' import type { ApprovalSchemaConfig, @@ -615,6 +619,11 @@ export interface TextActivityOptions< agentLoopStrategy?: TextOptions['agentLoopStrategy'] /** How the server tools of one model turn run. Default `'parallel'`. */ toolExecution?: TextOptions['toolExecution'] + /** + * How a tool input that fails the final check is repaired. Default: the + * built-in repair. `false`: no repair. A function replaces it. + */ + repairToolInput?: TextOptions['repairToolInput'] /** * Retry a rate-limited model call after the wait the provider asks for. * @@ -2352,8 +2361,19 @@ class TextEngine< } private handleToolCallEndEvent(chunk: ToolCallEndEvent): void { - this.toolCallManager.completeToolCall(chunk) + const final = this.toolCallManager.completeToolCall(chunk) const end = chunk as AdapterYieldChunk + if (final) { + // The client stores the same argument string as the server history. + // The wire carries only `args`; `input` is parsed back from it. + end.args = final.args + end.input = final.input + const metadata = tanstackMetadata(end) + if (metadata) { + delete metadata.args + delete metadata.input + } + } const state = end.state ?? tanstackMetadata(end)?.state if (state !== 'output-error' || end.result === undefined) return this.handleToolCallResultEvent({ @@ -2607,44 +2627,12 @@ class TextEngine< clientToolResults, (eventName, data, options) => this.createCustomEventChunk(eventName, data, options), - { - onBeforeToolCall: async (toolCall, tool, args) => { - this.logger.tools(`phase=before name=${toolCall.function.name}`, { - name: toolCall.function.name, - args, - }) - const hookCtx = { - toolCall, - tool, - args, - toolName: toolCall.function.name, - toolCallId: toolCall.id, - } - return this.middlewareRunner.runOnBeforeToolCall( - this.middlewareCtx, - hookCtx, - ) - }, - onAfterToolCall: async (info) => { - this.logger.tools(`phase=after name=${info.toolName}`, { - name: info.toolName, - result: info.result, - }) - return this.middlewareRunner.runOnAfterToolCall( - this.middlewareCtx, - info, - ) - }, - }, + this.toolMiddlewareHooks(), this.middlewareCtx.context, this.toolAbortSignal, - { - clientToolErrors: this.resumeClientToolErrors, - deniedToolResults: this.resumeDeniedToolResults, - cancelledToolCallIds: this.resumeCancelledToolCallIds, - inputResponses: this.resumeInputResponses, - }, + this.toolResumeState(), this.params.toolExecution, + this.params.repairToolInput, ) // Consume the async generator, yielding custom events and collecting the return value @@ -2873,44 +2861,12 @@ class TextEngine< clientToolResults, (eventName, data, options) => this.createCustomEventChunk(eventName, data, options), - { - onBeforeToolCall: async (toolCall, tool, args) => { - this.logger.tools(`phase=before name=${toolCall.function.name}`, { - name: toolCall.function.name, - args, - }) - const hookCtx = { - toolCall, - tool, - args, - toolName: toolCall.function.name, - toolCallId: toolCall.id, - } - return this.middlewareRunner.runOnBeforeToolCall( - this.middlewareCtx, - hookCtx, - ) - }, - onAfterToolCall: async (info) => { - this.logger.tools(`phase=after name=${info.toolName}`, { - name: info.toolName, - result: info.result, - }) - return this.middlewareRunner.runOnAfterToolCall( - this.middlewareCtx, - info, - ) - }, - }, + this.toolMiddlewareHooks(), this.middlewareCtx.context, this.toolAbortSignal, - { - clientToolErrors: this.resumeClientToolErrors, - deniedToolResults: this.resumeDeniedToolResults, - cancelledToolCallIds: this.resumeCancelledToolCallIds, - inputResponses: this.resumeInputResponses, - }, + this.toolResumeState(), this.params.toolExecution, + this.params.repairToolInput, ) // Consume the async generator, yielding custom events and collecting the return value @@ -2957,7 +2913,8 @@ class TextEngine< finishEvent, toolCalls, afterToolRequests, - executionResult.inputRequired, + // Its results already went out with the chunks above. + { ...executionResult, results: [] }, ) this.setToolPhase('wait') return @@ -3907,7 +3864,8 @@ class TextEngine< requests?: ReadonlyArray< GenericInterruptRequest> >, - inputRequired: ReadonlyArray = [], + // The tool phase result, so an afterTools boundary runs no hook twice. + executed?: ExecuteToolCallsResult, ): AsyncGenerator { this.middlewareCtx.phase = phase const boundaryRequests = @@ -3937,62 +3895,107 @@ class TextEngine< this.addTerminalAssistantMessages() } } - const actionable = this.getBoundaryActionableToolRequests(toolCalls) + const actionable = + executed ?? + (yield* this.drainToolCallGenerator(this.boundaryToolRequests(toolCalls))) + for (const chunk of this.buildToolResultChunks( + actionable.results, + finishEvent, + )) { + yield* this.pipeThroughMiddleware(chunk) + } yield* this.emitActionableInterruptBoundary( finishEvent, - actionable.approvals, - actionable.clientRequests, + actionable.needsApproval, + actionable.needsClientExecution, boundaryRequests, - inputRequired, + actionable.inputRequired, + actionable.subagentInterrupts, ) return true } - private getBoundaryActionableToolRequests( + /** + * The approval previews and client requests a boundary pauses with. They get + * the same middleware hook and final input check as the tool phase, and a + * skip or an invalid input becomes a tool result. No server tool runs here: + * a server tool is in the list only while it waits for approval. + */ + private boundaryToolRequests( toolCalls: ReadonlyArray, - ): { - approvals: Array - clientRequests: Array - } { + ): ReturnType { const { approvals, clientToolResults } = this.collectClientState() - const approvalRequests: Array = [] - const clientRequests: Array = [] - for (const toolCall of toolCalls) { + const tools = this.resolveExecutableTools(toolCalls) + const pendingCalls = toolCalls.filter((toolCall) => { // Provider-executed calls are complete; never surface them as client work. - if (isProviderExecutedToolCall(toolCall)) continue - const tool = this.resolveExecutableTools([toolCall]).find( + if (isProviderExecutedToolCall(toolCall)) return false + const tool = tools.find( (candidate) => candidate.name === toolCall.function.name, - ) as RuntimeToolWithApproval | undefined - if (!tool) continue - let input: unknown = {} - try { - const parsed = JSON.parse(toolCall.function.arguments.trim() || '{}') - input = parsed && typeof parsed === 'object' ? parsed : {} - } catch { - input = {} + ) + if (!tool) return false + if ( + tool.needsApproval && + approvalResolution(approvals, toolCall.id) === undefined + ) { + return true } - const approvalId = `approval_${toolCall.id}` - if (tool.needsApproval && !approvals.has(approvalId)) { - approvalRequests.push({ - toolCallId: toolCall.id, - toolName: toolCall.function.name, - input, - approvalId, - }) - } else if ( + return ( !tool.execute && !clientToolResults.has(toolCall.id) && !this.resumeClientToolErrors.has(toolCall.id) && !this.resumeCancelledToolCallIds.has(toolCall.id) - ) { - clientRequests.push({ - toolCallId: toolCall.id, + ) + }) + return executeToolCalls( + pendingCalls, + tools, + approvals, + clientToolResults, + undefined, + this.toolMiddlewareHooks(), + this.middlewareCtx.context, + this.toolAbortSignal, + this.toolResumeState(), + undefined, + this.params.repairToolInput, + ) + } + + private toolMiddlewareHooks(): ToolExecutionMiddlewareHooks { + return { + onBeforeToolCall: async (toolCall, tool, args) => { + this.logger.tools(`phase=before name=${toolCall.function.name}`, { + name: toolCall.function.name, + args, + }) + return this.middlewareRunner.runOnBeforeToolCall(this.middlewareCtx, { + toolCall, + tool, + args, toolName: toolCall.function.name, - input, + toolCallId: toolCall.id, }) - } + }, + onAfterToolCall: async (info) => { + this.logger.tools(`phase=after name=${info.toolName}`, { + name: info.toolName, + result: info.result, + }) + return this.middlewareRunner.runOnAfterToolCall( + this.middlewareCtx, + info, + ) + }, + } + } + + private toolResumeState(): ToolResumeExecutionState { + return { + clientToolErrors: this.resumeClientToolErrors, + deniedToolResults: this.resumeDeniedToolResults, + cancelledToolCallIds: this.resumeCancelledToolCallIds, + inputResponses: this.resumeInputResponses, } - return { approvals: approvalRequests, clientRequests } } private completeEphemeralInterruptBindings(chunk: StreamChunk): StreamChunk { diff --git a/packages/ai/src/activities/chat/middleware/types.ts b/packages/ai/src/activities/chat/middleware/types.ts index 2acb594465..101b920a6f 100644 --- a/packages/ai/src/activities/chat/middleware/types.ts +++ b/packages/ai/src/activities/chat/middleware/types.ts @@ -439,7 +439,10 @@ export interface ToolCallHookContext { toolCall: ToolCall /** The resolved tool definition, if found */ tool: Tool | undefined - /** Parsed arguments for the tool call */ + /** + * Parsed arguments for the tool call, before the `inputSchema` check: no + * schema defaults or transforms yet, and possibly invalid. + */ args: unknown /** Name of the tool */ toolName: string @@ -845,8 +848,10 @@ export interface ChatMiddleware< | Promise | null> /** - * Called before a tool is executed. + * Called before a server tool runs and before a client tool is sent to the + * client. Not called while the call waits for approval. * Can observe, transform args, skip execution, or abort the run. + * The tool's `inputSchema` checks the args after this hook. */ onBeforeToolCall?: ( ctx: ChatMiddlewareContext, diff --git a/packages/ai/src/activities/chat/stream/processor.ts b/packages/ai/src/activities/chat/stream/processor.ts index 12c960686d..34bb5f607e 100644 --- a/packages/ai/src/activities/chat/stream/processor.ts +++ b/packages/ai/src/activities/chat/stream/processor.ts @@ -35,6 +35,10 @@ import { withTanstackMetadata, } from '../../../utilities/merge-metadata' import { getChunkRunId } from '../../../utilities/chunk-ids' +import { + finishToolCallArguments, + parseToolArguments, +} from '../../../utilities/tool-call-arguments' import type { AdapterYieldChunk } from '../../../utilities/adapter-yield-chunk' import { isContentPartArray, @@ -214,38 +218,6 @@ function interruptBatchHasGeneric(interrupts: Array): boolean { }) } -/** - * The canonical arguments string for a `TOOL_CALL_END.input`, or `undefined` - * when JSON cannot carry it: `JSON.stringify` returns `undefined` (despite - * its declared type) for a top-level function or symbol and throws on BigInt - * and circular references. `null` is not tool arguments either, so it also - * keeps the streamed value instead of writing `arguments = "null"`. - */ -function serializeToolInput(input: unknown): string | undefined { - if (input === undefined || input === null) return undefined - try { - const serialized: unknown = JSON.stringify(input) - return typeof serialized === 'string' ? serialized : undefined - } catch { - return undefined - } -} - -/** - * The input of a tool call from its raw arguments: `{}` for an empty string, - * `undefined` when the arguments are not complete JSON (for example cut off) - * or are `null`. - */ -function parseToolArguments(args: string): unknown { - if (args.trim() === '') return {} - try { - const parsed: unknown = JSON.parse(args) - return parsed ?? undefined - } catch { - return undefined - } -} - /** * StreamProcessor - State machine for processing AI response streams * @@ -298,6 +270,16 @@ export class StreamProcessor { // only them. Messages written with no known run wait in the unassigned set. private readonly callMessageIds = new Map>() private readonly unassignedCallMessageIds = new Set() + // Messages of a finished tool-call turn. The next model call must not + // file its reasoning or tools onto them. A `stop` finish does not seal, + // so a trailing reasoning delta still joins that answer. + private readonly sealedMessageIds = new Set() + // Step id and reasoning message id to the assistant message that owns + // that thinking step. A later signature names one of those ids. + private readonly thinkingOwners = new Map< + string, + { messageId: string; stepId: string } + >() // Tool calls started since a RUN_STARTED that found no other run active, in // order. A success RUN_FINISHED hands the unanswered ones to the client. private runToolCallIds: Array = [] @@ -644,6 +626,12 @@ export class StreamProcessor { for (const id of this.unassignedCallMessageIds) { if (!keptIds.has(id)) this.unassignedCallMessageIds.delete(id) } + for (const id of this.sealedMessageIds) { + if (!keptIds.has(id)) this.sealedMessageIds.delete(id) + } + for (const [key, owner] of this.thinkingOwners) { + if (!keptIds.has(owner.messageId)) this.thinkingOwners.delete(key) + } this.messages = this.messages.slice(0, index + 1) this.emitMessagesChange() } @@ -663,6 +651,8 @@ export class StreamProcessor { this.structuredOutputUpdateBatches.clear() this.pendingManualMessageId = null this.legacyManualMessage = false + this.sealedMessageIds.clear() + this.thinkingOwners.clear() this.openChunk = null this.childProcessors.clear() this.childToolCalls.clear() @@ -816,12 +806,28 @@ export class StreamProcessor { break case 'REASONING_START': - case 'REASONING_MESSAGE_START': case 'REASONING_MESSAGE_END': case 'REASONING_END': // No special handling needed break + case 'REASONING_MESSAGE_START': { + // Same destination as content events: the assistant UIMessage. + // No metadata: do not create an empty assistant row. + const start = chunk as Extract< + StreamChunk, + { type: 'REASONING_MESSAGE_START' } + > + if (start.metadata == null) break + const target = this.unscopedAssistant(start) + const { messageId } = this.ensureAssistantMessage( + target.id ?? undefined, + !target.create, + ) + this.mergeMessageMetadata(messageId, start.metadata) + break + } + case 'REASONING_MESSAGE_CONTENT': this.handleReasoningMessageContentEvent( chunk as Extract, @@ -1149,6 +1155,7 @@ export class StreamProcessor { // Set iteration is insertion-order; reverse-iterate to search from the end const ids = Array.from(this.activeMessageIds).reverse() for (const id of ids) { + if (this.sealedMessageIds.has(id)) continue const state = this.messageStates.get(id) if (state && state.role === 'assistant' && belongsToCurrentCall(id)) { return id @@ -1157,8 +1164,10 @@ export class StreamProcessor { // finalizeStream() clears activeMessageIds but keeps messageStates. // Leftover reasoning after an early RUN_FINISHED must resume that // assistant. A new user turn calls prepareAssistantMessage(), which - // clears messageStates first. + // clears messageStates first. A sealed tool-call turn stays closed + // so the next model call opens its own message. for (const [id, state] of [...this.messageStates].reverse()) { + if (this.sealedMessageIds.has(id)) continue if (state.role === 'assistant' && belongsToCurrentCall(id)) { return id } @@ -1166,6 +1175,89 @@ export class StreamProcessor { return null } + /** The run that currently owns `messageId`, if one was recorded. */ + private runOfMessage(messageId: string): string | undefined { + for (const [runId, ids] of this.callMessageIds) { + if (ids.has(messageId)) return runId + } + return undefined + } + + /** Latest unsealed assistant message written by `runId`. */ + private openMessageOfRun(runId: string): string | null { + const ids = this.callMessageIds.get(runId) + if (!ids) return null + for (const id of [...ids].reverse()) { + if (this.sealedMessageIds.has(id)) continue + const state = this.messageStates.get(id) + if (state?.role === 'assistant') return id + } + return null + } + + /** + * Where an event with no assistant message id belongs. + * `create` means open a new message. Do not borrow another run's. + */ + private unscopedAssistant(chunk: StreamChunk): { + id: string | null + create: boolean + } { + const explicit = getChunkRunId(chunk) + const onlyRun = + this.activeRuns.size === 1 ? [...this.activeRuns][0] : undefined + const runId = explicit ?? onlyRun + if (runId !== undefined) { + const open = this.openMessageOfRun(runId) + if (open) return { id: open, create: false } + if (explicit !== undefined || this.activeRuns.has(runId)) { + return { id: null, create: true } + } + } + const active = this.getActiveAssistantMessageId() + if (active) return { id: active, create: false } + return { id: null, create: true } + } + + private rememberThinking( + messageId: string, + stepId: string, + reasoningMessageId?: string, + ): void { + const owner = { messageId, stepId } + this.thinkingOwners.set(stepId, owner) + if (reasoningMessageId !== undefined && reasoningMessageId !== stepId) { + this.thinkingOwners.set(reasoningMessageId, owner) + } + } + + /** Put `signature` on the thinking step `namedId` already belongs to. */ + private applyKnownSignature( + namedId: string | undefined, + signature: string, + ): boolean { + if (!namedId) return false + const owner = this.thinkingOwners.get(namedId) + if (!owner) return false + const state = this.getMessageState(owner.messageId) + if (!state) return false + state.thinkingStepSignatures.set(owner.stepId, signature) + const content = state.thinkingSteps.get(owner.stepId) ?? '' + if (!state.thinkingSteps.has(owner.stepId)) { + state.thinkingSteps.set(owner.stepId, content) + state.thinkingStepOrder.push(owner.stepId) + } + this.messages = updateThinkingPart( + this.messages, + owner.messageId, + owner.stepId, + content, + signature, + ) + this.emitMessagesChange() + return true + } + private resumeAssistantState(id: string, state: MessageStreamState): void { this.activeMessageIds.add(id) if (state.isComplete || this.isDone) { @@ -1627,6 +1719,12 @@ export class StreamProcessor { } if (this.unassignedCallMessageIds.delete(fromId)) this.unassignedCallMessageIds.add(toId) + if (this.sealedMessageIds.delete(fromId)) this.sealedMessageIds.add(toId) + for (const [key, owner] of this.thinkingOwners) { + if (owner.messageId === fromId) { + this.thinkingOwners.set(key, { messageId: toId, stepId: owner.stepId }) + } + } } /** @@ -1651,9 +1749,17 @@ export class StreamProcessor { const uiRole: 'system' | 'user' | 'assistant' = role === 'user' || role === 'system' ? role : 'assistant' - // Case 1: A manual message was created via startAssistantMessage() - if (this.pendingManualMessageId) { - const pendingId = this.pendingManualMessageId + // Case 1: A placeholder (reasoning before the server names the message, + // or startAssistantMessage()) takes this id. A placeholder of another + // run does not. That run keeps its own message. + const pendingId = this.pendingManualMessageId + const pendingRun = pendingId ? this.runOfMessage(pendingId) : undefined + const targetRun = this.runOfMessage(messageId) + const separateRuns = + pendingRun !== undefined && + targetRun !== undefined && + pendingRun !== targetRun + if (pendingId && !separateRuns) { this.pendingManualMessageId = null this.legacyManualMessage = false @@ -2323,10 +2429,16 @@ export class StreamProcessor { if (pendingId && parentId && parentId !== pendingId) { const parent = this.messages.find((m) => m.id === parentId) const placeholder = this.messages.find((m) => m.id === pendingId) + const pendingRun = this.runOfMessage(pendingId) + const parentRun = parent ? this.runOfMessage(parent.id) : undefined + const separateRuns = + pendingRun !== undefined && + parentRun !== undefined && + pendingRun !== parentRun if (!parent) { this.pendingManualMessageId = null this.renameMessage(pendingId, parentId) - } else if (parent.role === 'assistant' && placeholder) { + } else if (!separateRuns && parent.role === 'assistant' && placeholder) { // The parent already exists. Fold the placeholder into it, so the // thinking stays in the same message as the tool call. this.pendingManualMessageId = null @@ -2480,9 +2592,9 @@ export class StreamProcessor { * Handle TOOL_CALL_END event — arguments are finalized (input-complete). * Tool output arrives on TOOL_CALL_RESULT, not on this event. * - * If TOOL_CALL_END carries parsed `input`, it is the canonical arguments: - * write it into the accumulated string and override the rendered part's - * `input` with it. That covers adapters that deliver the whole input on END + * The END's raw `args` and parsed `input` replace the streamed arguments + * (see `finishToolCallArguments`), and the part's `input` is parsed from + * the result. That covers adapters that deliver the whole input on END * (e.g. Anthropic server_tool_use / web_search — issue #839) and adapters * that stream the wire arguments and then normalize them (OpenAI strict-mode * null widening, undone by the adapter after #939), so `arguments` and @@ -2499,47 +2611,19 @@ export class StreamProcessor { const msgState = this.getMessageState(messageId) if (!msgState) return - // The parsed input can ride on the spec `input` field or, for adapters - // that only stamp it into TanStack metadata, on `metadata.tanstack.input`. - const input = - chunk.input !== undefined - ? chunk.input - : (tanstackMetadata(chunk)?.input as unknown) - // Transition the tool call to input-complete (the authoritative completion signal) const existingToolCall = msgState.toolCalls.get(chunk.toolCallId) if (existingToolCall && existingToolCall.state !== 'input-complete') { - // The parsed input replaces the accumulated arguments string, so - // completeToolCall's strict parse surfaces the canonical value on the - // ToolCallPart even when the streamed deltas carried a provider-side - // reshaping of it. An input JSON cannot carry is not canonical: both - // `arguments` and `input` then stay with the streamed value, so the - // two never disagree. - const serializedInput = serializeToolInput(input) - if (serializedInput !== undefined) { - existingToolCall.arguments = serializedInput - } + // The END's arguments replace the streamed ones, so completeToolCall's + // strict parse surfaces the canonical value on the ToolCallPart even + // when the streamed deltas carried a provider-side reshaping of it. + existingToolCall.arguments = finishToolCallArguments( + existingToolCall.arguments, + chunk, + ).args const index = msgState.toolCallOrder.indexOf(chunk.toolCallId) this.completeToolCall(messageId, index, existingToolCall) - - // Canonicalize on the parsed input: overrides the accumulated-args parse - // that completeToolCall wrote (adapters may coerce values differently - // between streamed args and the final structured input). - if (serializedInput !== undefined) { - existingToolCall.parsedArguments = input - this.messages = updateToolCallPart(this.messages, messageId, { - id: existingToolCall.id, - name: existingToolCall.name, - arguments: existingToolCall.arguments, - state: 'input-complete', - input, - ...(existingToolCall.metadata !== undefined && { - metadata: existingToolCall.metadata, - }), - }) - this.emitMessagesChange() - } } } @@ -2688,6 +2772,9 @@ export class StreamProcessor { this.finishReason === 'tool_calls' && chunk.outcome?.type !== 'interrupt' if (isIntermediateToolTurn) { + // The next model call is a new assistant message. Leave this + // call's messages closed so its reasoning cannot land here. + for (const id of ids) this.sealedMessageIds.add(id) return } if (!chunk.outcome || chunk.outcome.type === 'success') { @@ -2939,7 +3026,8 @@ export class StreamProcessor { chunk: Extract, ): void { const stepId = chunk.stepName || generateMessageId() - const activeId = this.getActiveAssistantMessageId() + const target = this.unscopedAssistant(chunk) + const activeId = target.create ? null : target.id if (activeId) { const state = this.getMessageState(activeId) if (state) { @@ -2948,6 +3036,7 @@ export class StreamProcessor { state.thinkingSteps.set(stepId, '') state.thinkingStepOrder.push(stepId) } + this.rememberThinking(activeId, stepId) // Clear any pending stepId from a prior STEP_STARTED that fired // before the assistant message existed. Now that we're tracking // the step directly on message state, the pending value is stale @@ -2998,14 +3087,23 @@ export class StreamProcessor { chunk: Extract, ): void { const extra = chunk as AdapterYieldChunk - const signature = extra.signature + const metadataSignature = tanstackMetadata(chunk)?.signature + const signature = + extra.signature ?? + (typeof metadataSignature === 'string' ? metadataSignature : undefined) if (!signature) return + const namedStep = extra.stepId ?? chunk.stepName + if (this.applyKnownSignature(namedStep, signature)) return + + const target = this.unscopedAssistant(chunk) const { messageId, state } = this.ensureAssistantMessage( - this.getActiveAssistantMessageId() ?? undefined, + target.id ?? undefined, + !target.create, ) - const stepId = state.currentThinkingStepId ?? extra.stepId + const stepId = state.currentThinkingStepId ?? namedStep if (!stepId) return + this.rememberThinking(messageId, stepId) const thinking = state.thinkingSteps.get(stepId) if (thinking === undefined) return @@ -3029,8 +3127,10 @@ export class StreamProcessor { private handleReasoningMessageContentEvent( chunk: Extract, ): void { + const target = this.unscopedAssistant(chunk) const { messageId, state } = this.ensureAssistantMessage( - this.getActiveAssistantMessageId() ?? undefined, + target.id ?? undefined, + !target.create, ) this.mergeMessageMetadata(messageId, chunk.metadata) @@ -3045,6 +3145,7 @@ export class StreamProcessor { state.thinkingStepOrder.push(stepId) state.currentThinkingStepId = stepId } + this.rememberThinking(messageId, stepId, chunk.messageId) const nextThinking = (state.thinkingSteps.get(stepId) ?? '') + delta state.thinkingSteps.set(stepId, nextThinking) @@ -3077,10 +3178,15 @@ export class StreamProcessor { return } + if (this.applyKnownSignature(chunk.entityId, encryptedValue)) return + + const target = this.unscopedAssistant(chunk) const { messageId, state } = this.ensureAssistantMessage( - this.getActiveAssistantMessageId() ?? undefined, + target.id ?? undefined, + !target.create, ) const stepId = state.currentThinkingStepId ?? chunk.entityId + this.rememberThinking(messageId, stepId, chunk.entityId) this.mergeMessageMetadata(messageId, chunk.metadata) state.thinkingStepSignatures.set(stepId, encryptedValue) const content = state.thinkingSteps.get(stepId) ?? '' @@ -3720,6 +3826,8 @@ export class StreamProcessor { this.pendingManualMessageId = null this.legacyManualMessage = false this.pendingThinkingStepId = null + this.sealedMessageIds.clear() + this.thinkingOwners.clear() this.openChunk = null this.finishReason = null this.hasError = false diff --git a/packages/ai/src/activities/chat/tools/input-validation.ts b/packages/ai/src/activities/chat/tools/input-validation.ts new file mode 100644 index 0000000000..b84d083b55 --- /dev/null +++ b/packages/ai/src/activities/chat/tools/input-validation.ts @@ -0,0 +1,557 @@ +import { dereference, initialBaseURI, validate } from '@cfworker/json-schema' +import { + StandardSchemaValidationError, + convertSchemaToJsonSchema, + isStandardJSONSchema, + isStandardSchema, + validateWithStandardSchema, +} from './schema-converter' +import type { StandardSchemaV1 } from '@standard-schema/spec' +import type { Schema, SchemaDraft } from '@cfworker/json-schema' +import type { JSONSchema, SchemaInput } from '../../../types' + +/** A JSON Schema ready for `validate`: the copy, its draft, its `$ref` lookup. */ +interface CompiledSchema { + schema: Schema | boolean + draft: SchemaDraft + lookup: Record +} + +const compiled = new WeakMap() + +/** Keywords whose value is data, not a schema. */ +const DATA_KEYWORDS = new Set([ + 'const', + 'enum', + 'default', + 'examples', + 'example', +]) + +/** Keywords whose value maps names to schemas. */ +const SCHEMA_MAP_KEYWORDS = new Set([ + 'properties', + 'patternProperties', + '$defs', + 'definitions', + 'dependentSchemas', + 'dependencies', +]) + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value) +} + +/** Draft-04 and OpenAPI 3.0 write these bounds as booleans. */ +const BOOLEAN_BOUNDS = [ + ['exclusiveMinimum', 'minimum'], + ['exclusiveMaximum', 'maximum'], +] as const + +/** Keywords that can reject `null` even when `type` allows it. */ +const NULL_REJECTING_KEYWORDS = ['$ref', 'allOf', 'anyOf', 'oneOf', 'not'] + +/** Keywords that name a schema. They stay on a wrapped schema. */ +const NAMING_KEYWORDS = new Set(['$id', 'id', '$schema', '$anchor']) + +/** A key as a JSON Pointer token, URI-encoded as cfworker stores paths. */ +function pointerToken(key: string) { + return encodeURI(key.replaceAll('~', '~0').replaceAll('/', '~1')) +} + +/** Calls `visit` on each schema in `node`, inner schemas first. */ +function eachSchema( + node: unknown, + path: string, + visit: (schema: Record, path: string) => void, +): void { + if (Array.isArray(node)) { + node.forEach((item, index) => eachSchema(item, `${path}/${index}`, visit)) + return + } + if (!isRecord(node)) return + for (const [key, value] of Object.entries(node)) { + if (DATA_KEYWORDS.has(key)) continue + const keyPath = `${path}/${pointerToken(key)}` + if (SCHEMA_MAP_KEYWORDS.has(key) && isRecord(value)) { + for (const [name, schema] of Object.entries(value)) { + eachSchema(schema, `${keyPath}/${pointerToken(name)}`, visit) + } + } else { + eachSchema(value, keyPath, visit) + } + } + visit(node, path) +} + +/** + * Make one schema of the JSON copy ready for the validator, in place: + * - `format` goes away. It is an annotation, as in 2019-09 and later: a tool + * that declares `format: 'uri'` still gets a plain path. + * - A boolean bound (`exclusiveMinimum: true`, draft-04 and OpenAPI 3.0) + * becomes the number form, so it works in every draft. + * - OpenAPI 3.0 `nullable: true` lets the value be `null`: `null` joins + * `type` and `enum`. A schema with a `$ref`, `allOf`, `anyOf`, `oneOf`, or + * `not` becomes `{ if: { type: 'null' }, else: schema }`, and its path goes + * in `wrapped`. A failed `else` is an `if` error, which the error text + * leaves out, so the model never reads "Expected null". + */ +function prepareSchema( + schema: Record, + path: string, + wrapped: Array, +) { + if (typeof schema.format === 'string') delete schema.format + for (const [bound, limit] of BOOLEAN_BOUNDS) { + if (typeof schema[bound] !== 'boolean') continue + if (schema[bound] && typeof schema[limit] === 'number') { + schema[bound] = schema[limit] + delete schema[limit] + } else { + delete schema[bound] + } + } + if (schema.nullable !== true) return + if (schema.type !== undefined) { + schema.type = [...new Set([schema.type, 'null'].flat())] + } + if (Array.isArray(schema.enum) && !schema.enum.includes(null)) { + schema.enum.push(null) + } + if (!NULL_REJECTING_KEYWORDS.some((key) => key in schema)) return + const inner: Record = {} + for (const [key, value] of Object.entries(schema)) { + if (NAMING_KEYWORDS.has(key)) continue + inner[key] = value + delete schema[key] + } + schema.if = { type: 'null' } + schema.else = inner + wrapped.push(path) +} + +/** Any base works. A ref resolves as a URL, as cfworker reads it. */ +const REF_BASE = 'https://ref.invalid/' + +/** A `$ref` into a wrapped schema, moved to where that schema went. */ +function movedRef(ref: string, wrappedPath: string) { + if (!ref.startsWith('#')) return ref + const target = new URL(ref, REF_BASE).href + const from = new URL(wrappedPath, REF_BASE).href + return target.startsWith(`${from}/`) + ? `${from}/else${target.slice(from.length)}`.slice(REF_BASE.length) + : ref +} + +/** The draft named by `$schema`, else 2020-12. */ +function draftOf(schema: JSONSchema): SchemaDraft { + const uri = typeof schema.$schema === 'string' ? schema.$schema : '' + // Boolean bounds are numbers by now, and cfworker reads numeric bounds as + // booleans in its draft-04 mode. Draft 7 checks the rest the same way. + if (/draft-0[467]/.test(uri)) return '7' + if (uri.includes('2019-09')) return '2019-09' + return '2020-12' +} + +function compile(schema: JSONSchema | boolean): CompiledSchema { + if (typeof schema === 'boolean') { + return { schema, draft: '2020-12', lookup: dereference(schema) } + } + const cached = compiled.get(schema) + if (cached) return cached + // `dereference` writes lookup data onto the schema it gets. A JSON copy + // keeps the tool's own schema, which also goes to the provider, unchanged. + const copy: Schema = JSON.parse(JSON.stringify(schema)) + const wrapped: Array = [] + eachSchema(copy, '#', (node, path) => prepareSchema(node, path, wrapped)) + // Inner schemas were wrapped first, so their moves apply first. + // ponytail: refs are read as paths from the root. An absolute ref, and a + // ref inside a nullable schema with its own `$id`, are not moved. + if (wrapped.length > 0) { + eachSchema(copy, '#', (node) => { + if (typeof node.$ref === 'string') { + node.$ref = wrapped.reduce(movedRef, node.$ref) + } + }) + } + const result = { + schema: copy, + draft: draftOf(schema), + lookup: dereference(copy), + } + compiled.set(schema, result) + return result +} + +const DECIMAL = /^-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?$/ + +function convertPrimitive(value: unknown, type: string) { + switch (type) { + case 'number': + case 'integer': { + if (value === null) return 0 + if (typeof value === 'boolean') return value ? 1 : 0 + // Decimal text only: no hex, and no integer past the safe range. + if (typeof value !== 'string' || !DECIMAL.test(value)) return value + const number = Number(value) + const valid = + type === 'integer' + ? Number.isSafeInteger(number) + : Number.isFinite(number) + return valid ? number : value + } + case 'boolean': + if (value === null || value === 'false' || value === 0) return false + if (value === 'true' || value === 1) return true + return value + case 'string': + if (value === null) return '' + return typeof value === 'number' || typeof value === 'boolean' + ? String(value) + : value + case 'null': + return value === '' || value === 0 || value === false ? null : value + default: + return value + } +} + +function matchesType(value: unknown, type: string) { + switch (type) { + case 'null': + return value === null + case 'array': + return Array.isArray(value) + case 'object': + return ( + typeof value === 'object' && value !== null && !Array.isArray(value) + ) + case 'integer': + return typeof value === 'number' && Number.isInteger(value) + default: + return typeof value === type + } +} + +/** + * Repair common model mistakes against the schema before the check: `"5"` + * for 5, `"true"` for true, `null` for an optional field. It walks unions, + * `$ref`, `if`/`then`/`else`, tuples, and pattern and additional properties. + * A value that cannot be repaired stays as it is, so the check reports it. + */ +function coerceArguments(value: unknown, root: CompiledSchema) { + const { draft, lookup } = root + + function check(schema: Schema | boolean, candidate: unknown) { + if (typeof schema === 'boolean') return schema + return validate(candidate, schema, draft, lookup).valid + } + + function refTarget(schema: Schema) { + const uri = schema.__absolute_ref__ + const target = uri === undefined ? undefined : lookup[uri] + return typeof target === 'object' ? target : undefined + } + + function coerceUnion( + candidate: unknown, + schemas: Array, + seen: Set, + required: Set, + ) { + if (schemas.some((schema) => check(schema, candidate))) return candidate + for (const schema of schemas) { + const converted = coerce( + structuredClone(candidate), + schema, + seen, + required, + ) + if (check(schema, converted)) return converted + } + return candidate + } + + function requiredKeys( + schema: Schema | boolean, + candidate: unknown, + seen = new Set(), + ): Set { + if (typeof schema === 'boolean' || seen.has(schema)) return new Set() + const nextSeen = new Set(seen).add(schema) + const required = new Set( + Array.isArray(schema.required) ? schema.required : [], + ) + const children: Array = [...(schema.allOf ?? [])] + const target = refTarget(schema) + if (target) children.push(target) + if (schema.if !== undefined) { + const branch = check(schema.if, candidate) ? schema.then : schema.else + if (branch) children.push(branch) + } + for (const child of children) { + for (const key of requiredKeys(child, candidate, nextSeen)) { + required.add(key) + } + } + return required + } + + function coerce( + candidate: unknown, + schema: Schema | boolean, + seen = new Set(), + inheritedRequired = new Set(), + ): unknown { + if (typeof schema === 'boolean') return candidate + // A reference cycle at the same value cannot add another conversion. + if (seen.has(schema)) return candidate + const nextSeen = new Set(seen).add(schema) + const required = requiredKeys(schema, candidate) + for (const key of inheritedRequired) required.add(key) + let next = candidate + const target = refTarget(schema) + if (target) next = coerce(next, target, nextSeen, required) + for (const nested of schema.allOf ?? []) { + next = coerce(next, nested, nextSeen, required) + } + if (schema.anyOf) next = coerceUnion(next, schema.anyOf, nextSeen, required) + if (schema.oneOf) next = coerceUnion(next, schema.oneOf, nextSeen, required) + + const types = + typeof schema.type === 'string' ? [schema.type] : (schema.type ?? []) + if (!types.some((type) => matchesType(next, type))) { + // `null` last: `0` for `['null', 'string']` is `"0"`, not `null`. + const order: Array = types.filter((type) => type !== 'null') + if (types.includes('null')) order.push('null') + for (const type of order) { + const converted = convertPrimitive(next, type) + if (converted !== next) { + next = converted + break + } + } + } + + if (typeof next === 'object' && next !== null && !Array.isArray(next)) { + // Own keys only, `__proto__` included. + const object: Record = Object.fromEntries( + Object.entries(next), + ) + const properties = Object.entries(schema.properties ?? {}) + const knownKeys = new Set(properties.map(([key]) => key)) + for (const [key, property] of properties) { + if (!Object.hasOwn(object, key)) continue + if ( + object[key] === null && + !required.has(key) && + !check(property, null) + ) { + // A strict-mode provider sends `null` for an optional field. + delete object[key] + } else { + object[key] = coerce(object[key], property) + } + } + for (const [pattern, property] of Object.entries( + schema.patternProperties ?? {}, + )) { + const expression = new RegExp(pattern, 'u') + for (const key of Object.keys(object)) { + if (!expression.test(key)) continue + knownKeys.add(key) + object[key] = coerce(object[key], property) + } + } + const additional = schema.additionalProperties + if (typeof additional === 'object') { + for (const [key, entry] of Object.entries(object)) { + if (!knownKeys.has(key)) object[key] = coerce(entry, additional) + } + } + next = object + } + if (Array.isArray(next)) { + const tuple = + schema.prefixItems ?? + (Array.isArray(schema.items) ? schema.items : undefined) + const rest = Array.isArray(schema.items) ? undefined : schema.items + next = next.map((item: unknown, index) => { + const itemSchema = tuple?.[index] ?? rest + return itemSchema === undefined ? item : coerce(item, itemSchema) + }) + } + if (schema.if !== undefined) { + const branch = check(schema.if, next) ? schema.then : schema.else + if (branch) next = coerce(next, branch, nextSeen, required) + } + return next + } + + return coerce(value, root.schema) +} + +/** Errors that only say that a nested check failed. The nested error says why. */ +const WRAPPER_KEYWORDS = new Set([ + 'properties', + 'patternProperties', + 'items', + 'prefixItems', + 'additionalItems', + 'unevaluatedItems', + '$ref', + '$recursiveRef', + 'allOf', + // "Instance does not match "then" schema." (or "else") + 'if', +]) + +/** More errors than this are counted, not listed. */ +const MAX_ERRORS = 10 + +function validateJsonInput( + schema: JSONSchema | boolean, + received: unknown, + toolName: string, + repair: boolean, +) { + let errors: string + try { + // Check the JSON form: a key set to `undefined` is the same as no key, + // and an input with no JSON form (`undefined`) is `null`. + const json = JSON.stringify(received) ?? 'null' + const value: unknown = JSON.parse(json) + const root = compile(schema) + const check = (candidate: unknown) => + validate(candidate, root.schema, root.draft, root.lookup, false) + const result = check(value) + if (result.valid) return received + // Repair only an input that fails. The error reports the input as sent. + if (repair) { + let converted = value + try { + converted = coerceArguments(value, root) + } catch { + // A schema that cannot be used: the first check reports it. + } + if (JSON.stringify(converted) !== json && check(converted).valid) { + return converted + } + } + const leaves = result.errors.filter( + (error) => !WRAPPER_KEYWORDS.has(error.keyword), + ) + const shown = leaves.length > 0 ? leaves : result.errors + errors = shown + .slice(0, MAX_ERRORS) + .map((error) => `${error.instanceLocation}: ${error.error}`) + .join(', ') + if (shown.length > MAX_ERRORS) { + errors += `, and ${shown.length - MAX_ERRORS} more` + } + } catch (error) { + // A schema the validator cannot use, such as an unresolved `$ref`, or a + // value with no JSON form. The first line is enough for the model. + // Drop the validator base (in a browser, the page URL and its query). + const message = error instanceof Error ? error.message : String(error) + errors = message + .split('\n', 1) + .join('') + .replaceAll(initialBaseURI.href, '') + .replace(/ +Absolute URI "[^"]*"\.?/, '') + .trim() + } + throw new Error(`Input validation failed for tool ${toolName}: ${errors}`) +} + +function standardInputError( + issues: ReadonlyArray, + toolName: string, +) { + const error = new StandardSchemaValidationError(issues) + error.message = `Input validation failed for tool ${toolName}: ${error.message}` + return error +} + +/** + * How a failed tool input is repaired. `undefined`: the built-in repair. + * `false`: no repair. A function: called instead of the built-in repair. It + * returns the input to check again, or `undefined` to keep the error. + */ +export type ToolInputRepair = + | false + | ((input: unknown, error: Error) => unknown | Promise) + +/** + * The final check of a tool input, just before the tool runs or the client + * gets the call. A Standard Schema runs here (an async validator too), and + * its output is the input, so transforms apply once. A JSON Schema (an MCP + * tool, a plain JSON Schema tool) is checked with `@cfworker/json-schema`. + * A failed input is repaired as `repair` says, then checked again. The error + * of the first check goes back to the model as the tool result. + */ +export async function validateToolInput( + schema: SchemaInput | boolean | undefined, + received: unknown, + toolName: string, + repair?: ToolInputRepair, +): Promise { + try { + return await checkToolInput( + schema, + received, + toolName, + repair === undefined, + ) + } catch (error) { + if (typeof repair !== 'function') throw error + const repaired = await repair( + received, + error instanceof Error ? error : new Error(String(error)), + ) + if (repaired === undefined) throw error + try { + return await checkToolInput(schema, repaired, toolName, false) + } catch { + throw error + } + } +} + +async function checkToolInput( + schema: SchemaInput | boolean | undefined, + received: unknown, + toolName: string, + repair: boolean, +): Promise { + if (schema === undefined) return received + if (typeof schema === 'boolean') { + return validateJsonInput(schema, received, toolName, repair) + } + if (!isStandardSchema(schema)) { + const jsonSchema = convertSchemaToJsonSchema(schema, { io: 'input' }) + return jsonSchema + ? validateJsonInput(jsonSchema, received, toolName, repair) + : received + } + const first = await validateWithStandardSchema(schema, received) + if (first.success) return first.data + // The authored schema runs first. When it fails and it can export JSON + // Schema, repair the input and run the authored schema again. A failure + // reports the issues of the first run. + const original = standardInputError(first.issues, toolName) + if (!repair || !isStandardJSONSchema(schema)) throw original + let converted: unknown + try { + const jsonSchema = convertSchemaToJsonSchema(schema, { io: 'input' }) + if (!jsonSchema) throw original + converted = validateJsonInput(jsonSchema, received, toolName, true) + } catch { + throw original + } + // Nothing was repaired: a second run gives the same issues. + if (converted === received) throw original + const retried = await validateWithStandardSchema(schema, converted) + if (!retried.success) throw original + return retried.data +} diff --git a/packages/ai/src/activities/chat/tools/tool-calls.ts b/packages/ai/src/activities/chat/tools/tool-calls.ts index 2b889b2934..bd697e8edf 100644 --- a/packages/ai/src/activities/chat/tools/tool-calls.ts +++ b/packages/ai/src/activities/chat/tools/tool-calls.ts @@ -1,7 +1,9 @@ +import { finishToolCallArguments } from '../../../utilities/tool-call-arguments' import { normalizeToolResult } from '../../../utilities/tool-result' import { tanstackMetadata } from '../../../utilities/merge-metadata' import { isProviderExecutedToolCall } from '../../../utilities/provider-executed' import { mergeStreams } from '../../../utilities/merge-streams' +import { validateToolInput } from './input-validation' import type { AdapterYieldChunk } from '../../../utilities/adapter-yield-chunk' import { StandardSchemaValidationError, @@ -351,15 +353,16 @@ export class ToolCallManager< * Complete a tool call with its final input * Called when TOOL_CALL_END is received */ - completeToolCall(event: ToolCallEndEvent): void { + completeToolCall( + event: ToolCallEndEvent, + ): { args: string; input: unknown } | undefined { for (const toolCall of this.toolCallsMap.values()) { if (toolCall.id !== event.toolCallId) continue - if (event.input === undefined) return - const normalized = - event.input && typeof event.input === 'object' ? event.input : {} - toolCall.function.arguments = JSON.stringify(normalized) - return + const final = finishToolCallArguments(toolCall.function.arguments, event) + toolCall.function.arguments = final.args + return final } + return undefined } /** @@ -415,20 +418,7 @@ export class ToolCallManager< ) } - // Validate input against inputSchema (for Standard Schema compliant schemas) - if (tool.inputSchema && isStandardSchema(tool.inputSchema)) { - try { - args = parseWithStandardSchema(tool.inputSchema, args) - } catch (validationError: unknown) { - const message = - validationError instanceof Error - ? validationError.message - : 'Validation failed' - throw new Error( - `Input validation failed for tool ${tool.name}: ${message}`, - ) - } - } + args = await validateToolInput(tool.inputSchema, args, tool.name) // Execute the tool const executionContext = { @@ -600,7 +590,7 @@ export interface ToolResumeExecutionState { inputResponses?: ReadonlyMap } -function approvalResolution( +export function approvalResolution( approvals: ReadonlyMap, toolCallId: string, ): ToolApprovalResolution | undefined { @@ -625,7 +615,7 @@ function deniedApprovalResult(resolution: ToolApprovalResolution): unknown { : { error: 'User declined tool execution' } } -interface ExecuteToolCallsResult { +export interface ExecuteToolCallsResult { /** Tool results ready to send to LLM */ results: Array /** Tools that need user approval before execution */ @@ -989,6 +979,8 @@ async function buildClientToolResult( * @param toolExecution - `'parallel'` (default) prepares every call in call * order, then starts the server tools together. `'sequential'` runs one * call at a time. Results come back in call order either way. + * @param repairToolInput - How an input that fails the final check is + * repaired: the built-in repair (default), none (`false`), or a function. */ export async function* executeToolCalls( toolCalls: Array, @@ -1005,6 +997,7 @@ export async function* executeToolCalls( abortSignal?: AbortSignal, resumeState?: ToolResumeExecutionState, toolExecution: NonNullable = 'parallel', + repairToolInput?: TextOptions['repairToolInput'], ): AsyncGenerator { const results: Array = [] const needsApproval: Array = [] @@ -1091,6 +1084,8 @@ export async function* executeToolCalls( !resumeState?.cancelledToolCallIds?.has(tc.id) ) }) + // The calls that the gating holds back. + const held: Array = [] for (const toolCall of toolCalls) { // Provider-executed tools (Anthropic web_search / web_fetch) already ran @@ -1112,13 +1107,17 @@ export async function* executeToolCalls( continue } - // Skip non-pending tools while approvals are outstanding + // Skip non-pending tools while approvals are outstanding. A client tool + // with no approval, or an approved one, is still sent to the client. if (hasPendingApprovals) { - const isPendingApproval = - tool.needsApproval && - approvalResolution(approvals, toolCall.id) === undefined - const isPlainClientRequest = !tool.needsApproval && !tool.execute - if (!isPendingApproval && !isPlainClientRequest) { + const resolution = approvalResolution(approvals, toolCall.id) + const isPendingApproval = tool.needsApproval && resolution === undefined + const isClientRequest = + !tool.execute && + (!tool.needsApproval || + (resolution !== undefined && isApproved(resolution))) + if (!isPendingApproval && !isClientRequest) { + held.push(toolCall) continue } } @@ -1155,28 +1154,135 @@ export async function* executeToolCalls( continue } } + // The input that the model sent, to tell later edits apart from it. + const sentInput = JSON.stringify(input) + + const resolution = tool.needsApproval + ? approvalResolution(approvals, toolCall.id) + : undefined + // A call that waits for approval shows its checked input. Middleware + // runs only once the call is approved. + const needsPreview = tool.needsApproval && resolution === undefined + if (resolution !== undefined && !isApproved(resolution)) { + // User declined + results.push({ + toolCallId: toolCall.id, + toolName, + result: + resumeState?.deniedToolResults?.get(toolCall.id) ?? + deniedApprovalResult(resolution), + input, + state: 'output-error', + outcome: 'denied', + }) + continue + } + if (resolution !== undefined) { + input = editedApprovalArgs(resolution) ?? input + } + // A client tool whose result came back already ran with the input that + // the hook and the check gave it before dispatch. Checking the raw model + // input again would drop a real result. + const clientError = resumeState?.clientToolErrors?.get(toolCall.id) + if ( + !tool.execute && + (clientResults.has(toolCall.id) || clientError !== undefined) + ) { + results.push( + await buildClientToolResult( + toolCall.id, + toolName, + tool, + clientResults.get(toolCall.id), + input, + clientError, + ), + ) + continue + } - // Validate input against inputSchema (for Standard Schema compliant schemas) - if (tool.inputSchema && isStandardSchema(tool.inputSchema)) { + if (!needsPreview && middlewareHooks) { + let decision: Awaited> try { - input = parseWithStandardSchema(tool.inputSchema, input) - } catch (validationError: unknown) { - const message = - validationError instanceof Error - ? validationError.message - : 'Validation failed' - results.push({ + decision = await applyBeforeToolCallDecision( + toolCall, + tool, + input, + toolName, + middlewareHooks, + results, + ) + } catch (error) { + // The calls prepared so far still run, as they would one at a time. + failures.push(error) + break + } + if (!decision.proceed) continue + input = decision.input + } + + // The final check, after approval edits and middleware: invalid input + // becomes a tool error, and the checked value is what runs. + try { + input = await validateToolInput( + tool.inputSchema, + input, + toolName, + typeof repairToolInput === 'function' + ? (failed, error) => + repairToolInput({ tool, toolCall, input: failed, error }) + : repairToolInput, + ) + } catch (error) { + // The model sees its own arguments, never values that middleware put + // in. They stay out of the error, so tracing does not record them. + // When an edit or middleware changed the input, the model is told so, + // so it does not retry the same valid call. + let changed: boolean + try { + changed = JSON.stringify(input) !== sentInput + } catch { + changed = true + } + const message = `${ + error instanceof Error ? error.message : 'Validation failed' + }\nReceived arguments: ${argsStr}${ + changed + ? '\nAn approval edit or middleware changed these arguments before the check. The changed input failed the check.' + : '' + }` + await pushResultAndRunAfterHook( + results, + { toolCallId: toolCall.id, toolName, - result: { - error: `Input validation failed for tool ${tool.name}: ${message}`, - }, - // raw parse may have failed validation — still attach best-effort input + result: { error: message }, input, state: 'output-error', - }) - continue - } + duration: 0, + }, + { + toolCall, + tool, + toolName, + toolCallId: toolCall.id, + ok: false, + duration: 0, + error, + }, + // A preview call runs no hooks. + needsPreview ? undefined : middlewareHooks, + ) + continue + } + if (needsPreview) { + needsApproval.push({ + toolCallId: toolCall.id, + toolName, + input, + approvalId: `approval_${toolCall.id}`, + }) + continue } // Create a ToolExecutionContext for this tool call with event emission @@ -1207,142 +1313,13 @@ export async function* executeToolCalls( }, } as ToolExecutionContext - // CASE 1: Client-side tool (no execute function) + // Client-side tool (no execute function), approved or with no approval if (!tool.execute) { - // Check if tool needs approval - if (tool.needsApproval) { - const approvalId = `approval_${toolCall.id}` - const resolution = approvalResolution(approvals, toolCall.id) - - // Check if approval decision exists - if (resolution !== undefined) { - const approved = isApproved(resolution) - - if (approved) { - input = editedApprovalArgs(resolution) ?? input - // Approved - check if client has executed - const clientError = resumeState?.clientToolErrors?.get(toolCall.id) - if (clientResults.has(toolCall.id) || clientError !== undefined) { - results.push( - await buildClientToolResult( - toolCall.id, - toolName, - tool, - clientResults.get(toolCall.id), - input, - clientError, - ), - ) - } else { - // Approved but not executed yet - request client execution - needsClientExecution.push({ - toolCallId: toolCall.id, - toolName, - input, - }) - } - } else { - // User declined - results.push({ - toolCallId: toolCall.id, - toolName, - result: - resumeState?.deniedToolResults?.get(toolCall.id) ?? - deniedApprovalResult(resolution), - input, - state: 'output-error', - outcome: 'denied', - }) - } - } else { - // Need approval first - needsApproval.push({ - toolCallId: toolCall.id, - toolName: toolCall.function.name, - input, - approvalId, - }) - } - } else { - // No approval needed - check if client has executed - const clientError = resumeState?.clientToolErrors?.get(toolCall.id) - if (clientResults.has(toolCall.id) || clientError !== undefined) { - results.push( - await buildClientToolResult( - toolCall.id, - toolName, - tool, - clientResults.get(toolCall.id), - input, - clientError, - ), - ) - } else { - // Request client execution - needsClientExecution.push({ - toolCallId: toolCall.id, - toolName, - input, - }) - } - } + needsClientExecution.push({ toolCallId: toolCall.id, toolName, input }) continue } - // CASE 2: Server tool with approval required - if (tool.needsApproval) { - const approvalId = `approval_${toolCall.id}` - const resolution = approvalResolution(approvals, toolCall.id) - - if (resolution === undefined) { - // Need approval - needsApproval.push({ - toolCallId: toolCall.id, - toolName, - input, - approvalId, - }) - continue - } - if (!isApproved(resolution)) { - // User declined - results.push({ - toolCallId: toolCall.id, - toolName, - result: - resumeState?.deniedToolResults?.get(toolCall.id) ?? - deniedApprovalResult(resolution), - input, - state: 'output-error', - outcome: 'denied', - }) - continue - } - // Approved: run it like any other server tool below. - input = editedApprovalArgs(resolution) ?? input - } - - // CASE 3: Server tool, approved or with no approval - if (middlewareHooks) { - let decision: Awaited> - try { - decision = await applyBeforeToolCallDecision( - toolCall, - tool, - input, - toolName, - middlewareHooks, - results, - ) - } catch (error) { - // The calls prepared so far still run, as they would one at a time. - failures.push(error) - break - } - if (!decision.proceed) continue - input = decision.input - } - + // Server tool, approved or with no approval const run = runServerTool( toolCall, tool, @@ -1362,12 +1339,36 @@ export async function* executeToolCalls( yield* mergeStreams(runs) if (failures.length > 0) throw failures[0] + // No approval waits after all: each held-back call's preview failed its + // check or its arguments. The held calls run now, so each call has a result. + let heldInterrupts: Array = [] + if (held.length > 0 && needsApproval.length === 0) { + const rest = yield* executeToolCalls( + held, + tools, + approvals, + clientResults, + createCustomEventChunk, + middlewareHooks, + userContext, + abortSignal, + resumeState, + toolExecution, + repairToolInput, + ) + results.push(...rest.results) + needsClientExecution.push(...rest.needsClientExecution) + inputRequired.push(...rest.inputRequired) + heldInterrupts = rest.subagentInterrupts + } + // Parallel tools finish in any order. The model gets the results, input // requests, and interrupts in the order it made the calls. const callOrder = new Map(toolCalls.map((tc, index) => [tc.id, index])) const byCallOrder = (a: { toolCallId: string }, b: { toolCallId: string }) => (callOrder.get(a.toolCallId) ?? 0) - (callOrder.get(b.toolCallId) ?? 0) results.sort(byCallOrder) + needsClientExecution.sort(byCallOrder) inputRequired.sort(byCallOrder) return { @@ -1375,8 +1376,9 @@ export async function* executeToolCalls( needsApproval, needsClientExecution, inputRequired, - subagentInterrupts: toolCalls.flatMap( - (tc) => interruptsByCall.get(tc.id) ?? [], - ), + subagentInterrupts: [ + ...toolCalls.flatMap((tc) => interruptsByCall.get(tc.id) ?? []), + ...heldInterrupts, + ], } } diff --git a/packages/ai/src/activities/evaluate/index.ts b/packages/ai/src/activities/evaluate/index.ts index 2fa3f1f0de..1c44c42c7c 100644 --- a/packages/ai/src/activities/evaluate/index.ts +++ b/packages/ai/src/activities/evaluate/index.ts @@ -548,7 +548,24 @@ export async function decide< logger, }) - const answers = mapAnswers(questions, result.answers) + let answers + try { + answers = mapAnswers(questions, result.answers) + } catch (error) { + // Log first: a throwing onUsage hook replaces this error below. + logger.errors('evaluate answer mapping failed', { + error, + source: 'evaluate', + }) + aiEventClient.emit('evaluate:usage', { + requestId, + model: result.model, + usage: result.usage, + timestamp: Date.now(), + }) + await runGenerationUsage(middleware, mwCtx, result.usage) + throw error + } const duration = Date.now() - startTime aiEventClient.emit('evaluate:request:completed', { diff --git a/packages/ai/src/middlewares/otel.ts b/packages/ai/src/middlewares/otel.ts index cef80e2e8d..00994530f7 100644 --- a/packages/ai/src/middlewares/otel.ts +++ b/packages/ai/src/middlewares/otel.ts @@ -899,6 +899,9 @@ export function otelMiddleware( onBeforeToolCall(ctx, hookCtx) { safeCall('otel.onBeforeToolCall', () => { + // A client tool runs in the client, not here, so it gets no + // execute_tool span. Middleware sees its input before dispatch. + if (hookCtx.tool && !hookCtx.tool.execute) return const state = stateByCtx.get(ctx) if (!state) return const parent = state.currentIterationSpan ?? state.rootSpan diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 36a4345902..336fcbe2a1 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -773,6 +773,8 @@ export interface TanStackRunMetadata { state?: ToolOutputState /** Parsed `TOOL_CALL_END` input. Spec `TOOL_CALL_END` has no top-level `input`. */ input?: unknown + /** The raw tool argument string on `TOOL_CALL_END`, as the provider sent it. */ + args?: string } /** @@ -1277,6 +1279,20 @@ export interface ResolvedPromptCache { */ export type FetchWrapper = (next: typeof fetch) => typeof fetch +/** + * What a `repairToolInput` function gets: a tool input that failed the check. + */ +export interface ToolInputRepairContext { + /** The tool the model called. */ + tool: AnyTool + /** The call, with its id and the raw arguments text. */ + toolCall: ToolCall + /** The input that failed, after approval edits and middleware. */ + input: unknown + /** The validation error. A Standard Schema error also has `issues`. */ + error: Error +} + /** * Options passed into the SDK and further piped to the AI provider. */ @@ -1314,6 +1330,16 @@ export interface TextOptions< * starts them together, and `'sequential'` runs them one at a time. */ toolExecution?: 'parallel' | 'sequential' + /** + * How a tool input that fails the final check is repaired. By default, + * common model mistakes such as `"5"` for a number are repaired against the + * schema. `false` turns the repair off. A function replaces the built-in + * repair: it returns the input to check again, or `undefined` to send the + * validation error to the model. + */ + repairToolInput?: + | false + | ((context: ToolInputRepairContext) => unknown | Promise) /** * Retry a model call when the provider error has a wait time (usually a * 429). `chat()` waits for the `retryAfterMs` time of the `RUN_ERROR`, then diff --git a/packages/ai/src/utilities/normalize-stream-chunk.test.ts b/packages/ai/src/utilities/normalize-stream-chunk.test.ts index 7954dc147a..27e42095cd 100644 --- a/packages/ai/src/utilities/normalize-stream-chunk.test.ts +++ b/packages/ai/src/utilities/normalize-stream-chunk.test.ts @@ -343,6 +343,37 @@ describe('normalizeStreamChunk', () => { expect(out).not.toHaveProperty('toolCallName') }) + it('sends TOOL_CALL_END args without a second copy as input', () => { + const out = normalizeOne({ + type: EventType.TOOL_CALL_END, + toolCallId: 'tc1', + args: '{"q": 1}', + input: { q: 1 }, + } as AdapterYieldChunk) + + expect(out).toEqual({ + type: EventType.TOOL_CALL_END, + toolCallId: 'tc1', + metadata: { tanstack: { args: '{"q": 1}' } }, + }) + }) + + it('rewrites TOOL_CALL_END args to match an input that differs', () => { + const out = normalizeOne({ + type: EventType.TOOL_CALL_END, + toolCallId: 'tc1', + args: '{"a": 1, "b": null}', + input: { a: 1 }, + metadata: { tanstack: { input: { a: 1 } } }, + } as AdapterYieldChunk) + + expect(out).toEqual({ + type: EventType.TOOL_CALL_END, + toolCallId: 'tc1', + metadata: { tanstack: { args: '{"a":1}' } }, + }) + }) + it('stringifies an array TOOL_CALL_END result onto RESULT content', () => { const parts = [{ type: 'text', content: 'hello' }] const out = normalizeAll({ diff --git a/packages/ai/src/utilities/normalize-stream-chunk.ts b/packages/ai/src/utilities/normalize-stream-chunk.ts index a706252269..6d4cbcbe78 100644 --- a/packages/ai/src/utilities/normalize-stream-chunk.ts +++ b/packages/ai/src/utilities/normalize-stream-chunk.ts @@ -3,7 +3,8 @@ import type { StreamChunk, ToolCallResultEvent } from '../types' import type { AdapterYieldChunk } from './adapter-yield-chunk' import { isTanstackUsage, toSpecTokenUsage } from './ag-ui-usage' import type { MetadataRecord } from './merge-metadata' -import { withTanstackMetadata } from './merge-metadata' +import { tanstackMetadata, withTanstackMetadata } from './merge-metadata' +import { finishToolCallArguments } from './tool-call-arguments' import { reasoningEncryptedValue } from './reasoning-encrypted-value' import { specKeysFor } from './spec-event-keys' @@ -136,6 +137,17 @@ export function normalizeStreamChunk( if (chunk.type === EventType.TOOL_CALL_START) { skipLeftover.add('toolName') } + // A TOOL_CALL_END sends its arguments once, as `args` made to match + // `input`. `restorePublicUsage` parses `input` back from it. + let sendsArgs = false + if ( + chunk.type === EventType.TOOL_CALL_END && + typeof (chunk.args ?? tanstackMetadata(chunk)?.args) === 'string' + ) { + sendsArgs = true + tanstack.args = finishToolCallArguments('', chunk).args + skipLeftover.add('input') + } for (const key of Object.keys(chunk)) { if (specKeys.has(key) || skipLeftover.has(key)) continue if (tanstack[key] !== undefined) continue @@ -164,6 +176,8 @@ export function normalizeStreamChunk( ? specChunk : withTanstackMetadata(specChunk, tanstack) + if (sendsArgs) delete tanstackMetadata(normalized)?.input + const extras = encryptedValueExtras(chunk) const main: Array = [normalized as StreamChunk] diff --git a/packages/ai/src/utilities/restore-inbound-chunk.test.ts b/packages/ai/src/utilities/restore-inbound-chunk.test.ts index 26d8a451ba..01f632f4e8 100644 --- a/packages/ai/src/utilities/restore-inbound-chunk.test.ts +++ b/packages/ai/src/utilities/restore-inbound-chunk.test.ts @@ -80,6 +80,21 @@ describe('restorePublicUsage', () => { } expect(chunk.input).toEqual({ q: 'sf' }) }) + + it('restores TOOL_CALL_END input from metadata.tanstack.args', () => { + const chunk: StreamChunk = { + type: EventType.TOOL_CALL_END, + toolCallId: 'tc1', + metadata: { tanstack: { args: '{"q": "sf"}' } }, + } + + restorePublicUsage(chunk) + + if (chunk.type !== EventType.TOOL_CALL_END) { + throw new Error('expected TOOL_CALL_END') + } + expect(chunk.input).toEqual({ q: 'sf' }) + }) }) describe('restoreInboundChunk', () => { diff --git a/packages/ai/src/utilities/restore-inbound-chunk.ts b/packages/ai/src/utilities/restore-inbound-chunk.ts index 460ee8137a..ccf24f3849 100644 --- a/packages/ai/src/utilities/restore-inbound-chunk.ts +++ b/packages/ai/src/utilities/restore-inbound-chunk.ts @@ -3,6 +3,7 @@ import type { StreamChunk } from '../types' import type { AdapterYieldChunk } from './adapter-yield-chunk' import { isTanstackUsage, rebuildTokenUsage } from './ag-ui-usage' import { tanstackMetadata } from './merge-metadata' +import { parseToolArguments } from './tool-call-arguments' /** * Rebuild TanStack `TokenUsage` (`promptTokens`) from spec `usage[]` plus @@ -34,7 +35,12 @@ export function restorePublicUsage(chunk: StreamChunk): StreamChunk { } if (chunk.type === EventType.TOOL_CALL_END && chunk.input === undefined) { - const input = tanstackMetadata(chunk)?.input + const tanstack = tanstackMetadata(chunk) + const input = + tanstack?.input ?? + (typeof tanstack?.args === 'string' + ? parseToolArguments(tanstack.args) + : undefined) if (input !== undefined) { chunk.input = input } diff --git a/packages/ai/src/utilities/tool-call-arguments.ts b/packages/ai/src/utilities/tool-call-arguments.ts new file mode 100644 index 0000000000..d90f8d8bb9 --- /dev/null +++ b/packages/ai/src/utilities/tool-call-arguments.ts @@ -0,0 +1,244 @@ +import { tanstackMetadata } from './merge-metadata' +import type { ToolCallEndEvent } from '../types' +import type { AdapterYieldChunk } from './adapter-yield-chunk' + +type JsonNode = { + start: number + end: number + value: unknown + keyToken?: string + properties?: Map + items?: Array +} + +function sameJsonValue(left: unknown, right: unknown): boolean { + if (Object.is(left, right)) return true + if (Array.isArray(left) && Array.isArray(right)) + return ( + left.length === right.length && + left.every((value, index) => sameJsonValue(value, right[index])) + ) + if ( + left === null || + right === null || + typeof left !== 'object' || + typeof right !== 'object' || + Array.isArray(left) || + Array.isArray(right) + ) + return false + const leftKeys = Object.keys(left) + const rightKeys = Object.keys(right) + return ( + leftKeys.length === rightKeys.length && + leftKeys.every( + (key) => + Object.prototype.hasOwnProperty.call(right, key) && + sameJsonValue(Reflect.get(left, key), Reflect.get(right, key)), + ) + ) +} + +function scanJson(raw: string, parsed: unknown): JsonNode { + let cursor = 0 + const whitespace = () => { + while ( + raw[cursor] === ' ' || + raw[cursor] === '\t' || + raw[cursor] === '\r' || + raw[cursor] === '\n' + ) + cursor++ + } + const stringToken = () => { + const start = cursor++ + while (cursor < raw.length) { + if (raw[cursor] === '\\') cursor += 2 + else if (raw[cursor++] === '"') break + } + return raw.slice(start, cursor) + } + const value = (parsedValue: unknown): JsonNode => { + whitespace() + const start = cursor + if (raw[cursor] === '{') { + cursor++ + whitespace() + const properties = new Map() + while (raw[cursor] !== '}') { + const keyToken = stringToken() + const key: unknown = JSON.parse(keyToken) + whitespace() + cursor++ + const child = + parsedValue !== null && + typeof parsedValue === 'object' && + typeof key === 'string' + ? Reflect.get(parsedValue, key) + : undefined + const childNode = value(child) + if (typeof key === 'string') + properties.set(key, { ...childNode, keyToken }) + whitespace() + if (raw[cursor] !== ',') break + cursor++ + whitespace() + } + cursor++ + return { start, end: cursor, value: parsedValue, properties } + } + if (raw[cursor] === '[') { + cursor++ + whitespace() + const items: Array = [] + while (raw[cursor] !== ']') { + items.push( + value( + Array.isArray(parsedValue) ? parsedValue[items.length] : undefined, + ), + ) + whitespace() + if (raw[cursor] !== ',') break + cursor++ + } + cursor++ + return { start, end: cursor, value: parsedValue, items } + } + if (raw[cursor] === '"') stringToken() + else + while (cursor < raw.length && !/[\s,}\]]/.test(raw[cursor] ?? '')) + cursor++ + return { start, end: cursor, value: parsedValue } + } + return value(parsed) +} + +function serializeValue( + value: unknown, + raw: string, + node: JsonNode | undefined, + seen: Set, +): string | undefined { + if (node && sameJsonValue(node.value, value)) + return raw.slice(node.start, node.end) + if (typeof value === 'number') { + if (Object.is(value, -0)) return '-0' + if (value === Infinity) return '1e999' + if (value === -Infinity) return '-1e999' + return JSON.stringify(value) + } + if (value === null || typeof value !== 'object') return JSON.stringify(value) + if (seen.has(value)) + throw new TypeError('Cannot serialize circular tool input') + seen.add(value) + try { + if (Array.isArray(value)) + return ( + '[' + + Array.from( + value, + (item, index) => + serializeValue(item, raw, node?.items?.[index], seen) ?? 'null', + ).join(',') + + ']' + ) + const properties: Array = [] + for (const key of Object.keys(value)) { + const serialized = serializeValue( + Reflect.get(value, key), + raw, + node?.properties?.get(key), + seen, + ) + if (serialized !== undefined) + properties.push( + (node?.properties?.get(key)?.keyToken ?? JSON.stringify(key)) + + ':' + + serialized, + ) + } + return '{' + properties.join(',') + '}' + } finally { + seen.delete(value) + } +} + +/** + * Keep `raw` when it is the same JSON value as `input`. When they differ, + * write `input` again and keep the raw tokens of the values that did not + * change. Empty or missing raw arguments use the input. + */ +export function reconcileToolCallArguments( + raw: string | undefined, + input: unknown, +): string { + if (input === undefined) return raw ?? '' + if (raw === undefined) + return serializeValue(input, '', undefined, new Set()) ?? '' + let parsed: unknown + try { + parsed = JSON.parse(raw) + } catch { + // Empty raw arguments are no arguments: the input is the value to keep. + if (raw.trim() === '') + return serializeValue(input, '', undefined, new Set()) ?? raw + return raw + } + if (sameJsonValue(parsed, input)) return raw + return serializeValue(input, raw, scanJson(raw, parsed), new Set()) ?? raw +} + +/** + * The input of a tool call from its raw arguments: `{}` for an empty string, + * `undefined` when the arguments are not complete JSON (for example cut off) + * or are `null`. + */ +export function parseToolArguments(args: string): unknown { + if (args.trim() === '') return {} + try { + const parsed: unknown = JSON.parse(args) + return parsed ?? undefined + } catch { + return undefined + } +} + +/** + * The final arguments and input of a tool call at TOOL_CALL_END. `streamed` + * is the text of its TOOL_CALL_ARGS deltas. The END's raw `args` (the + * provider's own string) win over it, and stay as they are when they are the + * same JSON value as the END's `input`. Malformed arguments also stay, so the + * tool reports the error and the model can repair them. `input` is always + * the parse of `args`. + */ +export function finishToolCallArguments( + streamed: string, + end: ToolCallEndEvent, +): { args: string; input: unknown } { + const metadata = tanstackMetadata(end) + // Wire data: only a string is raw arguments. + const endArgs: unknown = (end as AdapterYieldChunk).args ?? metadata?.args + const rawArgs = typeof endArgs === 'string' ? endArgs : undefined + const endInput = end.input ?? metadata?.input + // Only an object or array is tool input. Anything else (`null`, a + // function) keeps the raw or streamed arguments. + const input = + endInput !== null && typeof endInput === 'object' ? endInput : undefined + let raw: string | undefined = rawArgs ?? streamed + // Streamed text that is not JSON gives way to the input. + if (rawArgs === undefined && input !== undefined) { + try { + JSON.parse(raw) + } catch { + raw = undefined + } + } + let args: string + try { + args = reconcileToolCallArguments(raw, input) + } catch { + // An input with no JSON form (a cycle, a bigint): keep the raw arguments. + args = raw ?? streamed + } + return { args, input: parseToolArguments(args) } +} diff --git a/packages/ai/tests/chat-client-tool-wait.test.ts b/packages/ai/tests/chat-client-tool-wait.test.ts index a49217ae20..24ee54158d 100644 --- a/packages/ai/tests/chat-client-tool-wait.test.ts +++ b/packages/ai/tests/chat-client-tool-wait.test.ts @@ -214,7 +214,11 @@ describe('client-tool wait lifecycle', () => { }), ) - expect(fakeTracer.spans).toHaveLength(2) + // The client runs the tool, so the server records no execute_tool span. + expect(fakeTracer.spans.map((span) => span.name)).toEqual([ + 'chat test-model', + 'chat test-model #0', + ]) expect(fakeTracer.spans.every((span) => span.ended)).toBe(true) expect( fakeMeter.records.filter( diff --git a/packages/ai/tests/evaluate.test.ts b/packages/ai/tests/evaluate.test.ts index 392742ae2d..aa7145b1c0 100644 --- a/packages/ai/tests/evaluate.test.ts +++ b/packages/ai/tests/evaluate.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it } from 'vitest' +import { aiEventClient } from '@tanstack/ai-event-client' import { boolean, choice, decide, score } from '../src/index' import type { EvaluateAdapter } from '../src/activities/evaluate/adapter' import type { @@ -85,6 +86,60 @@ function recordingMiddleware() { return { middleware, events } } +function recordUsageEvents() { + const usage: Array<{ requestId: string; model: string; usage: TokenUsage }> = + [] + const unsubscribe = aiEventClient.on( + 'evaluate:usage', + (event) => { + const { requestId, model } = event.payload + usage.push({ requestId, model, usage: event.payload.usage }) + }, + { withEventTarget: true }, + ) + return { usage, unsubscribe } +} + +function recordOrder(onUsage?: () => void) { + const order: Array = [] + const unsubscribes = [ + aiEventClient.on( + 'evaluate:request:completed', + () => { + order.push('event:completed') + }, + { withEventTarget: true }, + ), + aiEventClient.on( + 'evaluate:usage', + () => { + order.push('event:usage') + }, + { withEventTarget: true }, + ), + ] + const middleware: GenerationMiddleware = { + name: 'order', + onUsage: () => { + order.push('onUsage') + onUsage?.() + }, + onFinish: () => { + order.push('onFinish') + }, + onAbort: () => { + order.push('onAbort') + }, + onError: () => { + order.push('onError') + }, + } + const unsubscribe = () => { + for (const off of unsubscribes) off() + } + return { order, middleware, unsubscribe } +} + const mixedQuestions = { queue: choice({ instructions: 'Which team should handle this ticket?', @@ -502,4 +557,212 @@ describe('decide()', () => { expect(events.error).toHaveLength(1) expect(events.abort).toHaveLength(0) }) + + it('reports usage once on success', async () => { + const usage = { promptTokens: 3, completionTokens: 1, totalTokens: 4 } + const { middleware, events } = recordingMiddleware() + const recorded = recordUsageEvents() + const adapter = mockEvaluateAdapter(async () => + adapterResult(mixedWireAnswers, usage), + ) + + try { + await decide({ + adapter, + state: 'ticket', + questions: mixedQuestions, + middleware: [middleware], + }) + } finally { + recorded.unsubscribe() + } + + expect(events.usage).toEqual([usage]) + expect(recorded.usage).toEqual([ + { requestId: events.start[0]!.requestId, model: 'jev-1.13.0', usage }, + ]) + }) + + it('reports usage once, then fires onError, when the answers cannot be mapped', async () => { + const usage = { promptTokens: 3, completionTokens: 1, totalTokens: 4 } + const order: Array = [] + const { middleware, events } = recordingMiddleware() + const recorded = recordUsageEvents() + const adapter = mockEvaluateAdapter(async () => + adapterResult({ queue: mixedWireAnswers.queue! }, usage), + ) + + let thrown: unknown + try { + await decide({ + adapter, + state: 'ticket', + questions: mixedQuestions, + middleware: [ + middleware, + { + onUsage: () => { + order.push('usage') + }, + onError: () => { + order.push('error') + }, + }, + ], + debug: false, + }) + } catch (error) { + thrown = error + } finally { + recorded.unsubscribe() + } + + expect(thrown).toBeInstanceOf(Error) + expect((thrown as Error).message).toBe( + 'decide(): missing answer for question "urgency"', + ) + expect(events.error).toHaveLength(1) + expect(events.error[0]!.error).toBe(thrown) + expect(events.usage).toEqual([usage]) + expect(recorded.usage).toEqual([ + { requestId: events.start[0]!.requestId, model: 'jev-1.13.0', usage }, + ]) + expect(order).toEqual(['usage', 'error']) + expect(events.finish).toHaveLength(0) + }) + + it('reports no usage when the adapter throws', async () => { + const { middleware, events } = recordingMiddleware() + const recorded = recordUsageEvents() + const adapter = mockEvaluateAdapter(async () => { + throw new Error('evaluate boom') + }) + + try { + await expect( + decide({ + adapter, + state: 'ticket', + questions: mixedQuestions, + middleware: [middleware], + debug: false, + }), + ).rejects.toThrow('evaluate boom') + } finally { + recorded.unsubscribe() + } + + expect(events.usage).toHaveLength(0) + expect(recorded.usage).toHaveLength(0) + expect(events.error).toHaveLength(1) + }) + + it('keeps the success order: completed event, usage event, onUsage, onFinish', async () => { + const recorded = recordOrder() + const adapter = mockEvaluateAdapter(async () => + adapterResult(mixedWireAnswers), + ) + + try { + await decide({ + adapter, + state: 'ticket', + questions: mixedQuestions, + middleware: [recorded.middleware], + }) + } finally { + recorded.unsubscribe() + } + + expect(recorded.order).toEqual([ + 'event:completed', + 'event:usage', + 'onUsage', + 'onFinish', + ]) + }) + + it('still emits completed and runs onError when onUsage throws on success', async () => { + const hookError = new Error('onUsage boom') + const recorded = recordOrder(() => { + throw hookError + }) + const adapter = mockEvaluateAdapter(async () => + adapterResult(mixedWireAnswers), + ) + + let thrown: unknown + try { + await decide({ + adapter, + state: 'ticket', + questions: mixedQuestions, + middleware: [recorded.middleware], + debug: false, + }) + } catch (error) { + thrown = error + } finally { + recorded.unsubscribe() + } + + expect(thrown).toBe(hookError) + expect(recorded.order).toEqual([ + 'event:completed', + 'event:usage', + 'onUsage', + 'onError', + ]) + }) + + it('emits no completed event when the answers cannot be mapped', async () => { + const recorded = recordOrder() + const adapter = mockEvaluateAdapter(async () => + adapterResult({ queue: mixedWireAnswers.queue! }), + ) + + try { + await expect( + decide({ + adapter, + state: 'ticket', + questions: mixedQuestions, + middleware: [recorded.middleware], + debug: false, + }), + ).rejects.toThrow('decide(): missing answer for question "urgency"') + } finally { + recorded.unsubscribe() + } + + expect(recorded.order).toEqual(['event:usage', 'onUsage', 'onError']) + }) + + it('still runs onError once when onUsage throws after the answers cannot be mapped', async () => { + const hookError = new Error('onUsage boom') + const recorded = recordOrder(() => { + throw hookError + }) + const adapter = mockEvaluateAdapter(async () => + adapterResult({ queue: mixedWireAnswers.queue! }), + ) + + let thrown: unknown + try { + await decide({ + adapter, + state: 'ticket', + questions: mixedQuestions, + middleware: [recorded.middleware], + debug: false, + }) + } catch (error) { + thrown = error + } finally { + recorded.unsubscribe() + } + + expect(thrown).toBe(hookError) + expect(recorded.order).toEqual(['event:usage', 'onUsage', 'onError']) + }) }) diff --git a/packages/ai/tests/reasoning-start-metadata.test.ts b/packages/ai/tests/reasoning-start-metadata.test.ts new file mode 100644 index 0000000000..69e0187810 --- /dev/null +++ b/packages/ai/tests/reasoning-start-metadata.test.ts @@ -0,0 +1,103 @@ +import { describe, expect, it } from 'vitest' +import { StreamProcessor } from '../src/activities/chat/stream/processor' +import { EventType } from '../src/types' + +describe('REASONING_MESSAGE_START metadata', () => { + it('does not create an assistant message when the start event has no metadata', () => { + const processor = new StreamProcessor() + processor.processChunk({ + type: EventType.REASONING_MESSAGE_START, + messageId: 'reasoning-1', + role: 'reasoning', + timestamp: 1, + }) + expect(processor.getMessages()).toEqual([]) + }) + + it('keeps custom metadata from the start event on the assistant UIMessage', () => { + const processor = new StreamProcessor() + processor.processChunk({ + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: 1, + }) + processor.processChunk({ + type: EventType.REASONING_MESSAGE_START, + messageId: 'reasoning-1', + role: 'reasoning', + timestamp: 2, + metadata: { source: 'reasoning', traceId: 'trace-1' }, + }) + processor.processChunk({ + type: EventType.REASONING_MESSAGE_CONTENT, + messageId: 'reasoning-1', + delta: 'Let me think', + timestamp: 3, + }) + + const message = processor.getMessages()[0] + const thinking = message?.parts.find((part) => part.type === 'thinking') + expect(thinking && 'content' in thinking ? thinking.content : '').toBe( + 'Let me think', + ) + expect(message?.metadata).toEqual({ + source: 'reasoning', + traceId: 'trace-1', + }) + }) + + it('keeps start metadata on the second open run', () => { + const processor = new StreamProcessor() + processor.processChunk({ + type: EventType.RUN_STARTED, + runId: 'run-A', + threadId: 'thread-1', + timestamp: 1, + }) + processor.processChunk({ + type: EventType.TEXT_MESSAGE_START, + messageId: 'message-A', + role: 'assistant', + timestamp: 2, + }) + processor.processChunk({ + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'message-A', + delta: 'A output', + timestamp: 3, + }) + processor.processChunk({ + type: EventType.RUN_STARTED, + runId: 'run-B', + threadId: 'thread-1', + timestamp: 4, + }) + processor.processChunk({ + type: EventType.REASONING_MESSAGE_START, + messageId: 'reasoning-B', + role: 'reasoning', + timestamp: 5, + metadata: { + source: 'reasoning', + traceId: 'trace-b', + tanstack: { runId: 'run-B' }, + }, + }) + processor.processChunk({ + type: EventType.REASONING_MESSAGE_CONTENT, + messageId: 'reasoning-B', + delta: 'B reasoning', + timestamp: 6, + }) + + const messages = processor.getMessages() + expect(messages[0]?.id).toBe('message-A') + expect(messages[0]?.metadata).toBeUndefined() + expect(messages[1]?.metadata).toEqual({ + source: 'reasoning', + traceId: 'trace-b', + tanstack: { runId: 'run-B' }, + }) + }) +}) diff --git a/packages/ai/tests/stream-processor.test.ts b/packages/ai/tests/stream-processor.test.ts index cca84350b0..1a718be7b1 100644 --- a/packages/ai/tests/stream-processor.test.ts +++ b/packages/ai/tests/stream-processor.test.ts @@ -5347,6 +5347,155 @@ describe('StreamProcessor', () => { ) }) + it('files each model call on its own assistant message (#1533)', () => { + // Public agent-loop order: one RUN_STARTED, then a tool_calls + // RUN_FINISHED between calls. Later calls do not send RUN_STARTED. + const processor = new StreamProcessor() + processor.processChunk(ev.runStarted('run-1')) + processor.processChunk(ev.stepStarted('step-a')) + processor.processChunk(ev.reasoningContent('think A', 'rs_A')) + processor.processChunk(ev.textStart('msg-1')) + processor.processChunk(ev.textContent('hello', 'msg-1')) + processor.processChunk( + chunk(EventType.TOOL_CALL_START, { + toolCallId: 'call_S', + toolCallName: 'search', + parentMessageId: 'msg-1', + }), + ) + processor.processChunk(ev.toolArgs('call_S', '{}')) + processor.processChunk(ev.toolEnd('call_S')) + processor.processChunk(ev.runFinished('tool_calls', 'run-1')) + processor.processChunk(ev.toolResult('call_S', '{"ok":true}')) + + processor.processChunk( + chunk(EventType.TOOL_CALL_START, { + toolCallId: 'call_G', + toolCallName: 'get', + parentMessageId: 'msg-2', + }), + ) + processor.processChunk(ev.toolArgs('call_G', '{}')) + processor.processChunk(ev.toolEnd('call_G')) + processor.processChunk(ev.runFinished('tool_calls', 'run-2')) + processor.processChunk(ev.toolResult('call_G', '{"ok":true}')) + + processor.processChunk(ev.stepStarted('step-b')) + processor.processChunk(ev.reasoningContent('think B', 'rs_B')) + processor.processChunk(ev.textStart('msg-3')) + processor.processChunk(ev.textContent('bye', 'msg-3')) + processor.processChunk(ev.runFinished('stop', 'run-3')) + + const messages = processor.getMessages() + expect(messages.map((message) => message.id)).toEqual([ + 'msg-1', + 'msg-2', + 'msg-3', + ]) + expect(messages[0]?.parts.map((part) => part.type)).toEqual([ + 'thinking', + 'text', + 'tool-call', + 'tool-result', + ]) + expect(messages[1]?.parts.map((part) => part.type)).toEqual([ + 'tool-call', + 'tool-result', + ]) + expect(messages[2]?.parts.map((part) => part.type)).toEqual([ + 'thinking', + 'text', + ]) + const thinking = (index: number) => { + const part = messages[index]?.parts.find( + (item) => item.type === 'thinking', + ) + if (part?.type !== 'thinking') throw new Error('expected thinking') + return part + } + expect(thinking(0).content).toBe('think A') + expect(thinking(0).stepId).toBe('step-a') + expect(thinking(2).content).toBe('think B') + expect(thinking(2).stepId).toBe('step-b') + }) + + it('puts a signature on the step it names, not the active step (#1533)', () => { + const processor = new StreamProcessor() + processor.processChunk(ev.runStarted()) + processor.processChunk(ev.stepStarted('step-1')) + processor.processChunk(ev.reasoningContent('first', 'rs_A')) + processor.processChunk(ev.stepStarted('step-2')) + processor.processChunk(ev.reasoningContent('second', 'rs_B')) + processor.processChunk( + chunk(EventType.REASONING_ENCRYPTED_VALUE, { + subtype: 'message', + entityId: 'rs_A', + encryptedValue: 'sig-a', + }), + ) + processor.processChunk(ev.runFinished('tool_calls')) + processor.processChunk(ev.reasoningContent('next', 'rs_C')) + processor.processChunk( + chunk(EventType.REASONING_ENCRYPTED_VALUE, { + subtype: 'message', + entityId: 'rs_A', + encryptedValue: 'sig-a-late', + }), + ) + + const messages = processor.getMessages() + expect(messages).toHaveLength(2) + const first = messages[0]?.parts.filter( + (part) => part.type === 'thinking', + ) + const second = messages[1]?.parts.filter( + (part) => part.type === 'thinking', + ) + expect(first).toEqual([ + { + type: 'thinking', + content: 'first', + stepId: 'step-1', + signature: 'sig-a-late', + }, + { type: 'thinking', content: 'second', stepId: 'step-2' }, + ]) + expect(second).toEqual([ + { type: 'thinking', content: 'next', stepId: 'rs_C' }, + ]) + }) + + it('keeps reasoning from a second open run off the first message (#1533)', () => { + const processor = new StreamProcessor() + processor.processChunk(ev.runStarted('run-A')) + processor.processChunk(ev.textStart('message-A')) + processor.processChunk(ev.textContent('A output', 'message-A')) + processor.processChunk(ev.runStarted('run-B')) + processor.processChunk( + chunk(EventType.REASONING_MESSAGE_CONTENT, { + messageId: 'reasoning-B', + delta: 'B reasoning', + metadata: { tanstack: { runId: 'run-B' } }, + }), + ) + + const messages = processor.getMessages() + expect(messages).toHaveLength(2) + expect(messages[0]?.id).toBe('message-A') + expect(messages[0]?.metadata).toBeUndefined() + expect(messages[0]?.parts).toEqual([ + { type: 'text', content: 'A output' }, + ]) + const thinking = messages[1]?.parts[0] + expect(thinking).toMatchObject({ + type: 'thinking', + content: 'B reasoning', + }) + expect(messages[1]?.metadata).toMatchObject({ + tanstack: { runId: 'run-B' }, + }) + }) + it('flushes leftover text after a stop RUN_FINISHED and fires onStreamEnd once', () => { const events = spyEvents() const processor = new StreamProcessor({ events }) diff --git a/packages/ai/tests/subagent-interrupts.test.ts b/packages/ai/tests/subagent-interrupts.test.ts index e7c1aabee6..28aa6a5a59 100644 --- a/packages/ai/tests/subagent-interrupts.test.ts +++ b/packages/ai/tests/subagent-interrupts.test.ts @@ -2,6 +2,8 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { chat } from '../src/activities/chat' import { defineAgent } from '../src/activities/chat/agents/define-agent' +import { defineChatMiddleware } from '../src/activities/chat/middleware/define' +import { defineInterrupt } from '../src/interrupt-definition' import { StreamProcessor } from '../src/activities/chat/stream/processor' import { uiMessagesToWire } from '../src/utilities/ag-ui-wire' import { fromSpecTokenUsage } from '../src/utilities/ag-ui-usage' @@ -493,6 +495,50 @@ describe('subagent interrupts', () => { expect(second.at(-1)).not.toHaveProperty('outcome.type', 'interrupt') }) + it('keeps the child interrupt next to an afterTools boundary interrupt', async () => { + const { agent } = cleanerAgent() + const review = defineInterrupt({ + id: 'after-review', + responseSchema: z.boolean(), + }) + const chunks = await collectChunks( + chat({ + adapter: toolParent('{}'), + threadId: 't', + runId: 'run-1', + messages: [user], + subagents: { agents: [agent] }, + interrupts: [review], + middleware: [ + defineChatMiddleware({ + name: 'after-review', + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterTools') return undefined + return { + interrupts: [ + review.interrupt({ + key: 'review', + reason: 'review', + message: 'Review', + }), + ], + } + }, + }), + ], + }) as AsyncIterable, + ) + expect(chunks.at(-1)).toMatchObject({ + type: EventType.RUN_FINISHED, + outcome: { + type: 'interrupt', + interrupts: expect.arrayContaining([ + expect.objectContaining({ id: 'approval_call_c' }), + ]), + }, + }) + }) + it('gives a resumed child the same input the model wrote', async () => { const { run, execute } = cleanerChild() const inputs: Array<{ folder: string }> = [] diff --git a/packages/ai/tests/subagent-limits.test.ts b/packages/ai/tests/subagent-limits.test.ts index 9e74ad74c4..eed0f45c59 100644 --- a/packages/ai/tests/subagent-limits.test.ts +++ b/packages/ai/tests/subagent-limits.test.ts @@ -123,6 +123,61 @@ describe('subagent limits', () => { expect(results(chunks).c1.error).toBe('subagent timed out after 20 ms') }) + it('says no time left when the parent deadline is already past', async () => { + vi.useFakeTimers({ toFake: ['Date'] }) + vi.setSystemTime(0) + const grandchild = defineAgent({ + name: 'grandchild', + description: 'Deep', + run: async () => 'deep', + }) + const child = defineAgent({ + name: 'child', + description: 'Calls a grandchild', + run: (ctx) => { + vi.setSystemTime(5_000) + return ctx.chat({ + adapter: callsInOneTurn('grandchild', ['g1']).adapter, + subagents: { agents: [grandchild] }, + }) + }, + }) + try { + const { adapter } = callsInOneTurn('child', ['c1']) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + subagents: { agents: [child], limits: { timeoutMs: 5_000 } }, + }) as AsyncIterable, + ) + expect(JSON.stringify(chunks)).toContain( + 'subagent limit reached (no time left)', + ) + } finally { + vi.useRealTimers() + } + }) + + it('names timeoutMs 0 when that limit is zero', async () => { + const worker = defineAgent({ + name: 'worker', + description: 'Works', + run: async () => 'done', + }) + const { adapter } = callsInOneTurn('worker', ['c1']) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + subagents: { agents: [worker], limits: { timeoutMs: 0 } }, + }) as AsyncIterable, + ) + expect(results(chunks).c1.error).toBe( + 'subagent limit reached (timeoutMs 0)', + ) + }) + it('refuses children past maxConcurrent in one turn', async () => { let release!: () => void const gate = new Promise((resolve) => { diff --git a/packages/ai/tests/subagent-tool-input.test.ts b/packages/ai/tests/subagent-tool-input.test.ts index ac2f30bacb..ef7be89a80 100644 --- a/packages/ai/tests/subagent-tool-input.test.ts +++ b/packages/ai/tests/subagent-tool-input.test.ts @@ -220,7 +220,7 @@ describe('agent tool input', () => { }, }) - it('checks the agent fields before middleware, with sessionId in the args', async () => { + it('checks the agent fields after middleware, with sessionId in the args', async () => { const deep = defineAgent({ name: 'deep', description: 'Takes a depth', @@ -250,11 +250,34 @@ describe('agent tool input', () => { }) as AsyncIterable, ) - // The schema default is in place when middleware runs. - expect(seen).toEqual([{ topic: 'tides', depth: 2 }]) + // Middleware sees the model's input. The final check adds the default. + expect(seen).toEqual([{ topic: 'tides' }]) expect(runs[0]?.input).toEqual({ topic: 'tides', depth: 2 }) }) + it('checks a middleware rewrite of the agent fields before the child starts', async () => { + const rewrite: ChatMiddleware = { + name: 'rewrite', + onBeforeToolCall: () => ({ + type: 'transformArgs', + args: { topic: { text: 'tides' } }, + }), + } + const { chunks, result } = await callSubagent( + 'researcher', + { topic: 'tides' }, + { middleware: [rewrite] }, + ) + + expect(result.error).toContain( + 'Input validation failed for tool researcher', + ) + expect(runs).toEqual([]) + expect( + chunks.some((chunk) => chunk.type === EventType.SUBAGENT_STARTED), + ).toBe(false) + }) + it('throws when an inputSchema uses a reserved field that the run uses', () => { const { adapter } = parentCalling('taken', {}) @@ -290,7 +313,12 @@ describe('agent tool input', () => { {}, { middleware: [storedChildren(oldChild)], - binding: { start: async () => ({ subagentRunId: 'unused' }) }, + binding: { + start: async () => ({ + subagentRunId: 'unused', + done: Promise.resolve(), + }), + }, }, ) @@ -423,7 +451,10 @@ describe('agent tool input', () => { }) it('starts a background child through the start hook', async () => { - const start = vi.fn(async () => ({ subagentRunId: 'subagent-bg' })) + const start = vi.fn(async () => ({ + subagentRunId: 'subagent-bg', + done: Promise.resolve(), + })) const { result } = await callSubagent( 'writer', { prompt: 'Draft', background: true }, @@ -441,8 +472,143 @@ describe('agent tool input', () => { expect(runs).toEqual([]) }) + it('rejects sessionId when the agent schema result is not an object', async () => { + const flat = defineAgent({ + name: 'flat', + description: 'Returns text', + inputSchema: { + '~standard': { + version: 1 as const, + vendor: 'test', + validate: () => ({ value: 'plain' }), + jsonSchema: { + input: () => ({ type: 'string' }), + output: () => ({ type: 'string' }), + }, + }, + }, + run: async (ctx) => { + runs.push({ agent: 'flat', input: ctx.input, messages: ctx.messages }) + return 'ran' + }, + }) + const { adapter } = parentCalling('flat', { sessionId: 'subagent-old' }) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + middleware: [storedChildren(oldChild)], + subagents: { agents: [flat] }, + }) as AsyncIterable, + ) + + expect(toolResult(chunks).error).toContain( + 'Agent input must be an object when the call uses sessionId or background.', + ) + expect(runs).toEqual([]) + }) + + it('advertises sessionId beside a top-level anyOf schema', async () => { + const either = defineAgent({ + name: 'either', + description: 'One of two shapes', + inputSchema: z.union([ + z.object({ topic: z.string() }), + z.object({ note: z.string() }), + ]), + run: async () => 'ok', + }) + const { adapter, calls } = parentCalling('either', { topic: 'tides' }) + await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + middleware: [storedChildren(oldChild)], + subagents: { agents: [either] }, + }) as AsyncIterable, + ) + + expect(JSON.stringify(toolSchemas(calls).either)).toContain('sessionId') + }) + + it('holds maxConcurrent until a background child ends', async () => { + let release!: () => void + const done = new Promise((resolve) => { + release = resolve + }) + const start = vi.fn(async () => ({ subagentRunId: 'subagent-bg', done })) + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('c1', 'writer'), + ev.toolArgs('c1', JSON.stringify({ prompt: 'A', background: true })), + ev.toolStart('c2', 'writer'), + ev.toolArgs('c2', JSON.stringify({ prompt: 'B', background: true })), + ev.runFinished('tool_calls'), + ], + [ev.runStarted(), ev.runFinished('stop')], + ], + }) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + subagents: { + agents: [writer], + limits: { maxConcurrent: 1 }, + binding: { start }, + }, + }) as AsyncIterable, + ) + release() + + expect(start).toHaveBeenCalledTimes(1) + expect(JSON.stringify(chunks)).toContain( + 'subagent limit reached (maxConcurrent 1)', + ) + }) + + it('frees a maxCalls slot when a background start throws', async () => { + let n = 0 + const start = vi.fn(async () => { + n += 1 + if (n === 1) throw new Error('host down') + return { subagentRunId: 'subagent-bg', done: Promise.resolve() } + }) + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('c1', 'writer'), + ev.toolArgs('c1', JSON.stringify({ prompt: 'A', background: true })), + ev.toolStart('c2', 'writer'), + ev.toolArgs('c2', JSON.stringify({ prompt: 'B', background: true })), + ev.runFinished('tool_calls'), + ], + [ev.runStarted(), ev.runFinished('stop')], + ], + }) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Go' }], + toolExecution: 'sequential', + subagents: { + agents: [writer], + limits: { maxCalls: 1 }, + binding: { start }, + }, + }) as AsyncIterable, + ) + + expect(start).toHaveBeenCalledTimes(2) + expect(JSON.stringify(chunks)).toContain('host down') + expect(JSON.stringify(chunks)).toContain('subagent-bg') + }) + const host: SubagentBinding = { - start: async () => ({ subagentRunId: 'unused' }), + start: async () => ({ subagentRunId: 'unused', done: Promise.resolve() }), } const toolErrors: Array< [ @@ -463,9 +629,10 @@ describe('agent tool input', () => { [ 'a field that fails the agent schema', 'researcher', - { topic: 3 }, + // A number would be repaired into a string. A list cannot be. + { topic: [] }, {}, - 'Input validation failed for tool researcher: Validation failed: Invalid input: expected string, received number', + 'Input validation failed for tool researcher: Validation failed: Invalid input: expected string, received array', ], [ 'an unknown sessionId', @@ -484,9 +651,11 @@ describe('agent tool input', () => { [ 'a background that is not a boolean', 'writer', - { background: 'true' }, + // 'true' would be repaired into true. 'yes' cannot be. + { background: 'yes' }, { binding: host }, - 'prompt and sessionId must be strings, and background a boolean.', + // The JSON Schema of the tool checks the reserved fields too. + 'Input validation failed for tool writer: #/background: Instance type "string" is invalid. Expected "boolean".', ], [ 'background with a sessionId', @@ -501,7 +670,9 @@ describe('agent tool input', () => { async (_label, agent, args, options, error) => { const { chunks, result } = await callSubagent(agent, args, options) - expect(result).toEqual({ error }) + // A schema error also shows the received arguments after the message. + expect(Object.keys(result)).toEqual(['error']) + expect(result.error.startsWith(error)).toBe(true) expect(runs).toEqual([]) // The parent run goes on after the tool error. expect(chunks.some((chunk) => chunk.type === EventType.RUN_ERROR)).toBe( diff --git a/packages/ai/tests/tool-call-arguments.test.ts b/packages/ai/tests/tool-call-arguments.test.ts new file mode 100644 index 0000000000..f0ed922b14 --- /dev/null +++ b/packages/ai/tests/tool-call-arguments.test.ts @@ -0,0 +1,223 @@ +import { describe, expect, it } from 'vitest' +import { chat } from '../src/activities/chat/index' +import { EventType } from '../src/types' +import type { ToolCallEndEvent } from '../src/types' +import { tanstackMetadata } from '../src/utilities/merge-metadata' +import { + finishToolCallArguments, + reconcileToolCallArguments, +} from '../src/utilities/tool-call-arguments' +import { collectChunks, createMockAdapter, ev } from './test-utils' + +describe('lexical tool argument reconciliation', () => { + it.each([ + ' 1e999 ', + '-0', + '9007199254740993', + '1.20e+3', + '"\\u0061"', + '{"a":1e999,"b":[-0,9007199254740993,"\\u0061"]}', + ])('keeps unchanged raw JSON exactly: %s', (raw) => { + expect(reconcileToolCallArguments(raw, JSON.parse(raw))).toBe(raw) + }) + it('reuses unchanged nested values after removing an optional null', () => { + const raw = + '{"removed":null,"nested":{"number":1e999,"text":"\\u0061"},"array":[-0,9007199254740993]}' + expect( + reconcileToolCallArguments(raw, { + nested: { number: Infinity, text: 'a' }, + array: [-0, 9007199254740992], + }), + ).toBe( + '{"nested":{"number":1e999,"text":"\\u0061"},"array":[-0,9007199254740993]}', + ) + }) + it('retains unchanged array positions while correcting another position', () => { + expect(reconcileToolCallArguments('[1e999,2,-0]', [Infinity, 3, -0])).toBe( + '[1e999,3,-0]', + ) + }) + it('uses valid JSON spelling for corrected numbers', () => { + expect( + reconcileToolCallArguments('{}', { + positive: Infinity, + negative: -Infinity, + zero: -0, + }), + ).toBe('{"positive":1e999,"negative":-1e999,"zero":-0}') + }) + it('does not repair malformed raw JSON with a parsed projection', () => { + expect(reconcileToolCallArguments('{"broken":', {})).toBe('{"broken":') + }) + it('serializes END-only input when no raw arguments exist', () => { + expect(reconcileToolCallArguments(undefined, null)).toBe('null') + expect(reconcileToolCallArguments(undefined, [1, true])).toBe('[1,true]') + }) + it('leaves raw arguments alone when input is undefined', () => { + expect(reconcileToolCallArguments('1e999', undefined)).toBe('1e999') + }) +}) + +describe('raw JSON key spelling', () => { + it('keeps an escaped key when removing an unrelated null', () => { + expect( + reconcileToolCallArguments('{"\\u0061":9007199254740993,"drop":null}', { + a: 9007199254740992, + }), + ).toBe('{"\\u0061":9007199254740993}') + }) + it('uses the last duplicate key value and its original spelling', () => { + expect( + reconcileToolCallArguments('{"a":1,"\\u0061":1e999,"drop":null}', { + a: Infinity, + }), + ).toBe('{"\\u0061":1e999}') + }) + it('serializes sparse custom array positions as null slots', () => { + const values: Array = new Array(2) + values[1] = -0 + expect(reconcileToolCallArguments(undefined, values)).toBe('[null,-0]') + }) +}) + +describe('finishToolCallArguments', () => { + const end = ( + fields: Partial & { args?: string }, + ): ToolCallEndEvent => ({ + type: EventType.TOOL_CALL_END, + toolCallId: 'call-1', + timestamp: 0, + ...fields, + }) + + it('keeps raw args when the input is null', () => { + expect( + finishToolCallArguments('', end({ args: '{"q":"x"}', input: null })), + ).toEqual({ args: '{"q":"x"}', input: { q: 'x' } }) + }) + + it('keeps malformed raw args with no input', () => { + expect(finishToolCallArguments('', end({ args: '{"q":' }))).toEqual({ + args: '{"q":', + input: undefined, + }) + }) + + it('keeps the streamed bytes when they match the input', () => { + expect( + finishToolCallArguments('{"n": 1.50}', end({ input: { n: 1.5 } })), + ).toEqual({ args: '{"n": 1.50}', input: { n: 1.5 } }) + }) + + it('ignores args that are not a string', () => { + expect( + finishToolCallArguments( + '{"a":1}', + end({ metadata: { tanstack: { args: { a: 1 } } } } as never), + ), + ).toEqual({ args: '{"a":1}', input: { a: 1 } }) + }) + + it('reads args and input from metadata.tanstack', () => { + expect( + finishToolCallArguments( + '', + end({ + metadata: { + tanstack: { args: '{"a": 1, "b": null}', input: { a: 1 } }, + }, + }), + ), + ).toEqual({ args: '{"a":1}', input: { a: 1 } }) + }) +}) + +describe('raw tool arguments through chat()', () => { + async function run(raw: string, input?: unknown) { + const executed: Array = [] + const { adapter, calls } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('call-1', 'check'), + ev.toolArgs('call-1', raw), + { + type: EventType.TOOL_CALL_END, + toolCallId: 'call-1', + timestamp: Date.now(), + args: raw, + ...(input !== undefined ? { input } : {}), + }, + ev.runFinished('tool_calls'), + ], + [ev.runStarted('run-2'), ev.runFinished('stop', 'run-2')], + ], + }) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'check' }], + tools: [ + { + name: 'check', + description: 'Check', + execute: (value: unknown) => { + executed.push(value) + return 'ok' + }, + }, + ], + }), + ) + return { chunks, calls, executed } + } + + function replayedArguments( + calls: ReturnType['calls'], + ) { + return calls[1]?.messages + .flatMap((message) => message.toolCalls ?? []) + .map((call) => call.function.arguments) + } + + it('replays the exact provider bytes', async () => { + const raw = '{"count": 1.50, "id": 9007199254740993}' + const { calls, executed } = await run(raw, JSON.parse(raw)) + expect(replayedArguments(calls)).toEqual([raw]) + expect(executed).toEqual([{ count: 1.5, id: 9007199254740992 }]) + }) + + it('sends the raw arguments to the client in TOOL_CALL_END metadata', async () => { + const raw = '{"count": 1.50}' + const { chunks } = await run(raw, JSON.parse(raw)) + const end = chunks.find((chunk) => chunk.type === EventType.TOOL_CALL_END) + const metadata = end ? tanstackMetadata(end) : undefined + expect(metadata && 'args' in metadata ? metadata.args : undefined).toBe(raw) + // `input` is not sent twice; in process it is parsed back from `args`. + expect(metadata && 'input' in metadata).toBe(false) + expect(end && 'input' in end ? end.input : undefined).toEqual({ + count: 1.5, + }) + }) + + it('keeps an in-process input that JSON spells as -0 or 1e999', async () => { + const raw = '{"zero": -0, "big": 1e999}' + const { chunks } = await run(raw, JSON.parse(raw)) + const end = chunks.find((chunk) => chunk.type === EventType.TOOL_CALL_END) + expect(end && 'input' in end ? end.input : undefined).toEqual({ + zero: -0, + big: Infinity, + }) + }) + + it('reports malformed raw arguments instead of running with {}', async () => { + const { calls, executed } = await run('{"count":') + expect(executed).toEqual([]) + const toolMessage = calls[1]?.messages.find( + (message) => message.role === 'tool', + ) + expect(JSON.stringify(toolMessage)).toContain( + 'Failed to parse tool arguments as JSON', + ) + }) +}) diff --git a/packages/ai/tests/tool-input-validation.test.ts b/packages/ai/tests/tool-input-validation.test.ts new file mode 100644 index 0000000000..14b5dec60a --- /dev/null +++ b/packages/ai/tests/tool-input-validation.test.ts @@ -0,0 +1,1527 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { chat } from '../src/activities/chat/index' +import { defineChatMiddleware } from '../src/activities/chat/middleware/define' +import { defineInterrupt } from '../src/interrupt-definition' +import { executeToolCalls } from '../src/activities/chat/tools/tool-calls' +import { validateToolInput } from '../src/activities/chat/tools/input-validation' +import { otelMiddleware } from '../src/middlewares/otel' +import { EventType } from '../src/types' +import { collectChunks, createMockAdapter, ev } from './test-utils' +import { createFakeMeter, createFakeTracer } from './middlewares/fake-otel' +import type { StandardSchemaV1 } from '@standard-schema/spec' +import type { ChatMiddleware } from '../src/activities/chat/middleware/types' +import type { ToolExecutionMiddlewareHooks } from '../src/activities/chat/tools/tool-calls' +import type { ToolApprovalResolution } from '../src/interrupts' +import type { + AnyTool, + JSONSchema, + ModelMessage, + SchemaInput, + StreamChunk, + TextOptions, + Tool, + ToolCall, + ToolInputRepairContext, +} from '../src/types' + +const schema = z.object({ count: z.number().int().min(1) }) + +/** Checks the schema alone: the built-in repair would fix some inputs. */ +const noRepair = { repairToolInput: false } as const + +async function runTool( + rawArguments: string, + inputSchema: SchemaInput | boolean, + options: { + client?: boolean + needsApproval?: boolean + approvals?: Map + middleware?: ToolExecutionMiddlewareHooks + clientResults?: Map + repairToolInput?: TextOptions['repairToolInput'] + } = {}, +) { + const executed: Array = [] + const tool: AnyTool = { + name: 'check', + description: 'Check input', + inputSchema, + needsApproval: options.needsApproval, + ...(options.client + ? {} + : { + execute: (input: unknown) => { + executed.push(input) + return input + }, + }), + } + const call: ToolCall = { + id: 'call-1', + type: 'function', + function: { name: 'check', arguments: rawArguments }, + } + const generator = executeToolCalls( + [call], + [tool], + options.approvals, + options.clientResults, + undefined, + options.middleware, + undefined, + undefined, + undefined, + undefined, + options.repairToolInput, + ) + while (true) { + const next = await generator.next() + if (next.done) return { ...next.value, executed } + } +} + +describe('final tool input validation', () => { + it('awaits an async Standard Schema validator', async () => { + const asyncSchema: StandardSchemaV1 = { + '~standard': { + version: 1, + vendor: 'test', + validate: async (value) => + typeof value === 'object' && + value !== null && + 'count' in value && + typeof value.count === 'number' + ? { value: { count: value.count * 10 } } + : { issues: [{ message: 'count must be a number' }] }, + }, + } + const ok = await runTool('{"count":2}', asyncSchema) + expect(ok.executed).toEqual([{ count: 20 }]) + const bad = await runTool('{"count":"2"}', asyncSchema) + expect(bad.executed).toEqual([]) + expect(bad.results[0]?.state).toBe('output-error') + expect(bad.results[0]?.result).toEqual({ + error: expect.stringContaining('count must be a number'), + }) + }) + + it('includes the received arguments in the error', async () => { + const result = await runTool('{"count":0}', schema) + expect(result.results[0]?.result).toEqual({ + error: expect.stringMatching( + /^Input validation failed for tool check: [\s\S]*\nReceived arguments: \{"count":0\}$/, + ), + }) + expect(result.executed).toEqual([]) + }) + + it('echoes only the model arguments, not values middleware added', async () => { + const result = await runTool('{"count":0}', schema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 0, token: 'secret-token' }, + }), + }, + }) + expect(result.results[0]?.state).toBe('output-error') + expect(JSON.stringify(result.results[0]?.result)).not.toContain( + 'secret-token', + ) + // The model learns that its own arguments were changed. + expect(JSON.stringify(result.results[0]?.result)).toContain( + 'An approval edit or middleware changed these arguments', + ) + }) + + it('keeps empty and null arguments as {}', async () => { + for (const raw of ['', ' ', 'null']) { + const result = await runTool(raw, z.object({})) + expect(result.executed).toEqual([{}]) + } + }) + + it.each([false, true])( + 'rejects a middleware replacement before server or client execution: client=%s', + async (client) => { + const result = await runTool('{"count":2}', schema, { + client, + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 'wrong' }, + }), + }, + }) + expect(result.results[0]?.state).toBe('output-error') + expect(result.executed).toEqual([]) + expect(result.needsClientExecution).toEqual([]) + }, + ) + + it('rejects a middleware edit made in place', async () => { + const result = await runTool('{"count":2}', schema, { + middleware: { + onBeforeToolCall: async (_call, _tool, input) => { + if (typeof input === 'object' && input !== null) + Reflect.set(input, 'count', 'wrong') + }, + }, + }) + expect(result.results[0]?.state).toBe('output-error') + expect(result.executed).toEqual([]) + }) + + it('runs onAfterToolCall with the validation error', async () => { + const after: Array = [] + await runTool('{"count":2}', schema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 'wrong' }, + }), + onAfterToolCall: async (info) => { + after.push(info.ok) + }, + }, + }) + expect(after).toEqual([false]) + }) + + it('lets middleware repair invalid input', async () => { + const result = await runTool('{"count":"wrong"}', schema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 4 }, + }), + }, + }) + expect(result.executed).toEqual([{ count: 4 }]) + }) + + it('shows middleware the raw input and runs a transform once', async () => { + const seen: Array = [] + let transforms = 0 + const inputSchema = z.object({ + count: z.string().transform((value) => { + transforms++ + return Number(value) + }), + }) + const result = await runTool('{"count":"2"}', inputSchema, { + middleware: { + onBeforeToolCall: async (_call, _tool, input) => { + seen.push(structuredClone(input)) + }, + }, + }) + expect(seen).toEqual([{ count: '2' }]) + expect(result.executed).toEqual([{ count: 2 }]) + expect(transforms).toBe(1) + }) + + it('sends the checked middleware input to the client', async () => { + const result = await runTool('{"count":2}', schema, { + client: true, + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 5 }, + }), + }, + }) + expect(result.needsClientExecution).toEqual([ + { toolCallId: 'call-1', toolName: 'check', input: { count: 5 } }, + ]) + }) + + it('checks an approval edit for a client tool', async () => { + const result = await runTool('{"count":2}', schema, { + client: true, + needsApproval: true, + approvals: new Map([ + ['call-1', { approved: true, editedArgs: { count: 'wrong' } }], + ]), + }) + expect(result.results[0]?.state).toBe('output-error') + expect(result.needsClientExecution).toEqual([]) + }) + + it('keeps a returned client result when middleware repaired the input', async () => { + let hooks = 0 + const result = await runTool('{"count":"wrong"}', schema, { + client: true, + clientResults: new Map([['call-1', 'client result']]), + middleware: { + onBeforeToolCall: async () => { + hooks++ + return { type: 'transformArgs', args: { count: 4 } } + }, + }, + }) + expect(hooks).toBe(0) + expect(result.results[0]?.state).not.toBe('output-error') + expect(result.results[0]?.result).toBe('client result') + }) + + it('rejects an invalid approval edit for a server tool', async () => { + const result = await runTool('{"count":2}', schema, { + needsApproval: true, + approvals: new Map([ + ['call-1', { approved: true, editedArgs: { count: 0 } }], + ]), + }) + expect(result.results[0]?.state).toBe('output-error') + expect(result.executed).toEqual([]) + }) + + it.each([false, true])( + 'keeps the transformed approval edit: client=%s', + async (client) => { + const result = await runTool('"2"', z.string().transform(Number), { + client, + needsApproval: true, + approvals: new Map([['call-1', { approved: true, editedArgs: '4' }]]), + }) + if (client) { + expect(result.needsClientExecution).toEqual([ + { toolCallId: 'call-1', toolName: 'check', input: 4 }, + ]) + } else { + expect(result.executed).toEqual([4]) + } + }, + ) + + it.each([false, true])( + 'runs no middleware and dispatches nothing while approval is pending: client=%s', + async (client) => { + let middlewareCalls = 0 + const result = await runTool('{"count":2}', schema, { + client, + needsApproval: true, + middleware: { + onBeforeToolCall: async () => { + middlewareCalls++ + }, + }, + }) + expect(result.needsApproval).toEqual([ + { + toolCallId: 'call-1', + toolName: 'check', + input: { count: 2 }, + approvalId: 'approval_call-1', + }, + ]) + expect(result.needsClientExecution).toEqual([]) + expect(middlewareCalls).toBe(0) + expect(result.executed).toEqual([]) + }, + ) + + it.each([false, true])( + 'runs no middleware for a denied tool: client=%s', + async (client) => { + let middlewareCalls = 0 + const result = await runTool('{"count":2}', schema, { + client, + needsApproval: true, + approvals: new Map([['call-1', false]]), + middleware: { + onBeforeToolCall: async () => { + middlewareCalls++ + }, + }, + }) + expect(result.results[0]?.outcome).toBe('denied') + expect(result.needsClientExecution).toEqual([]) + expect(middlewareCalls).toBe(0) + expect(result.executed).toEqual([]) + }, + ) +}) + +describe('a batch with an approval', () => { + const call = (id: string, name: string, args: string): ToolCall => ({ + id, + type: 'function', + function: { name, arguments: args }, + }) + + async function runBatch( + calls: Array, + tools: Array, + approvals?: Map, + ) { + const generator = executeToolCalls(calls, tools, approvals) + while (true) { + const next = await generator.next() + if (next.done) return next.value + } + } + + it('runs the held calls when the only approval preview fails its check', async () => { + const execute = vi.fn(() => 'ran') + const result = await runBatch( + [call('a', 'gated', '{"count":0}'), call('b', 'plain', '{}')], + [ + { + name: 'gated', + description: 'Gated', + inputSchema: schema, + needsApproval: true, + execute: vi.fn(), + }, + { name: 'plain', description: 'Plain', execute }, + ], + ) + expect(result.needsApproval).toEqual([]) + expect(execute).toHaveBeenCalledTimes(1) + expect(result.results.map((r) => [r.toolCallId, r.state])).toEqual([ + ['a', 'output-error'], + ['b', undefined], + ]) + }) + + it('sends an approved client call while another call waits for approval', async () => { + const result = await runBatch( + [call('a', 'clientGated', '{}'), call('b', 'serverGated', '{}')], + [ + { name: 'clientGated', description: 'Client', needsApproval: true }, + { + name: 'serverGated', + description: 'Server', + needsApproval: true, + execute: vi.fn(), + }, + ], + new Map([['a', true]]), + ) + expect(result.needsClientExecution.map((r) => r.toolCallId)).toEqual(['a']) + expect(result.needsApproval.map((r) => r.toolCallId)).toEqual(['b']) + }) +}) + +describe('JSON Schema tool input', () => { + const jsonSchema: JSONSchema = { + type: 'object', + properties: { count: { type: 'integer', minimum: 1 } }, + required: ['count'], + additionalProperties: false, + } + + it('runs valid input', async () => { + const result = await runTool('{"count":2}', jsonSchema) + expect(result.executed).toEqual([{ count: 2 }]) + }) + + it.each([ + '{}', + '{"count":"wrong"}', + '{"count":0}', + '{"count":2,"extra":true}', + ])('rejects invalid input without executing: %s', async (raw) => { + const result = await runTool(raw, jsonSchema) + expect(result.executed).toEqual([]) + expect(result.results[0]?.state).toBe('output-error') + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Input validation failed for tool check'), + }) + }) + + it('names the failed location and shows the received arguments', async () => { + const result = await runTool('{"count":0}', jsonSchema) + expect(result.results[0]?.result).toEqual({ + error: expect.stringMatching(/#\/count: .*\n.*Received arguments:/s), + }) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Received arguments: {"count":0}'), + }) + }) + + it('checks local references and conditional fields', async () => { + const schema: JSONSchema = { + type: 'object', + $defs: { count: { type: 'integer' } }, + properties: { + mode: { type: 'string' }, + count: { $ref: '#/$defs/count' }, + }, + if: { properties: { mode: { const: 'count' } }, required: ['mode'] }, + then: { required: ['count'] }, + } + expect( + (await runTool('{"mode":"count","count":3}', schema)).executed, + ).toEqual([{ mode: 'count', count: 3 }]) + const missing = await runTool('{"mode":"count"}', schema) + expect(missing.executed).toEqual([]) + // The "then" wrapper error is left out. The nested error says why. + expect(missing.results[0]?.result).toEqual({ + error: expect.stringContaining( + 'check: #: Instance does not have required property "count".\n', + ), + }) + expect((await runTool('{"count":"x"}', schema)).executed).toEqual([]) + }) + + it('uses the draft named by $schema', async () => { + // Draft 4 spells an exclusive bound as a boolean next to `maximum`. + const schema: JSONSchema = JSON.parse( + JSON.stringify({ + $schema: 'http://json-schema.org/draft-04/schema#', + type: 'object', + properties: { + count: { type: 'number', maximum: 5, exclusiveMaximum: true }, + }, + }), + ) + expect((await runTool('{"count":4}', schema)).executed).toEqual([ + { count: 4 }, + ]) + expect((await runTool('{"count":5}', schema)).executed).toEqual([]) + }) + + it('runs empty and null arguments as {} when nothing is required', async () => { + for (const raw of ['', 'null']) { + const result = await runTool(raw, { type: 'object', properties: {} }) + expect(result.executed).toEqual([{}]) + } + }) + + it('rejects empty arguments when a field is required', async () => { + const result = await runTool('', jsonSchema) + expect(result.executed).toEqual([]) + expect(result.results[0]?.state).toBe('output-error') + }) + + it('checks a middleware rewrite and an approval edit', async () => { + const rewritten = await runTool('{"count":2}', jsonSchema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 'wrong' }, + }), + }, + }) + expect(rewritten.executed).toEqual([]) + const edited = await runTool('{"count":2}', jsonSchema, { + needsApproval: true, + approvals: new Map([ + ['call-1', { approved: true, editedArgs: { count: -1 } }], + ]), + }) + expect(edited.executed).toEqual([]) + expect(edited.results[0]?.state).toBe('output-error') + }) + + it('turns an unusable schema into a tool error', async () => { + const result = await runTool('{"count":2}', { + type: 'object', + properties: { count: { $ref: 'https://example.com/count.json' } }, + }) + expect(result.executed).toEqual([]) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Unresolved $ref'), + }) + expect(JSON.stringify(result.results[0]?.result)).not.toContain( + 'Known schemas', + ) + }) + + it('omits the validator base URI from a schema error', async () => { + const missing = await runTool('{"count":2}', { + type: 'object', + properties: { count: { $ref: '#/$defs/missing' } }, + }) + expect(missing.results[0]?.result).toEqual({ + error: expect.stringContaining('Unresolved $ref "#/$defs/missing"'), + }) + const missingError = JSON.stringify(missing.results[0]?.result) + expect(missingError).not.toContain('Absolute URI') + expect(missingError).not.toContain('github.com') + expect(missingError).not.toContain('Known schemas') + + const duplicate = await runTool('{"count":2}', { + type: 'object', + properties: { count: { $id: 'https://github.com/cfworker' } }, + }) + const duplicateError = JSON.stringify(duplicate.results[0]?.result) + expect(duplicate.executed).toEqual([]) + expect(duplicateError).toContain('Duplicate schema URI') + expect(duplicateError).not.toContain('github.com') + }) + + it.each([false, true])( + 'a boolean schema %s runs only when the schema is true', + async (inputSchema) => { + const result = await runTool('{"count":2}', inputSchema) + expect(result.executed).toEqual(inputSchema ? [{ count: 2 }] : []) + if (!inputSchema) expect(result.results[0]?.state).toBe('output-error') + }, + ) + + it.each([ + undefined, + 'http://json-schema.org/schema#', + 'http://json-schema.org/draft-07/schema#', + 'https://json-schema.org/draft/2019-09/schema', + ])('reads a boolean exclusive bound with $schema %s', async ($schema) => { + const schema: JSONSchema = JSON.parse( + JSON.stringify({ + $schema, + type: 'object', + properties: { + count: { type: 'number', maximum: 5, exclusiveMaximum: true }, + }, + }), + ) + expect((await runTool('{"count":4}', schema)).executed).toEqual([ + { count: 4 }, + ]) + expect((await runTool('{"count":5}', schema)).executed).toEqual([]) + }) + + it('keeps a number bound next to a boolean bound', async () => { + const schema: JSONSchema = JSON.parse( + JSON.stringify({ + $schema: 'https://json-schema.org/draft/2020-12/schema', + type: 'object', + properties: { + low: { type: 'number', minimum: 0, exclusiveMinimum: false }, + high: { type: 'number', exclusiveMinimum: 5 }, + }, + }), + ) + expect((await runTool('{"low":0,"high":6}', schema)).executed).toEqual([ + { low: 0, high: 6 }, + ]) + expect((await runTool('{"high":5}', schema)).executed).toEqual([]) + }) + + it('lets an OpenAPI 3.0 nullable reference be null', async () => { + const schema: JSONSchema = { + type: 'object', + definitions: { User: { type: 'object', required: ['id'] } }, + properties: { + owner: { nullable: true, allOf: [{ $ref: '#/definitions/User' }] }, + editor: { nullable: true, oneOf: [{ $ref: '#/definitions/User' }] }, + }, + } + expect( + (await runTool('{"owner":null,"editor":null}', schema)).executed, + ).toEqual([{ owner: null, editor: null }]) + expect( + (await runTool('{"owner":{"id":1},"editor":{"id":2}}', schema)).executed, + ).toEqual([{ owner: { id: 1 }, editor: { id: 2 } }]) + const wrong = await runTool('{"owner":{}}', schema) + expect(wrong.executed).toEqual([]) + expect(wrong.results[0]?.result).toEqual({ + error: expect.stringContaining( + 'check: #/owner: Instance does not have required property "id".\n', + ), + }) + }) + + it('moves a percent-encoded $ref and a ref into a nullable root', async () => { + const schema: JSONSchema = { + $id: 'https://example.com/tool', + nullable: true, + allOf: [{ $ref: '#/definitions/Input' }], + definitions: { + Input: { + type: 'object', + properties: { + 'my field': { + nullable: true, + anyOf: [{ type: 'integer' }, { type: 'string' }], + }, + 'a/b': { + $ref: '#/definitions/Input/properties/my%20field/anyOf/0', + }, + '50%': { + nullable: true, + allOf: [{ type: 'object', required: ['id'] }], + }, + copy: { $ref: '#/definitions/Input/properties/50%25/allOf/0' }, + }, + }, + }, + } + expect( + (await runTool('{"my field":null,"a/b":3}', schema)).executed, + ).toEqual([{ 'my field': null, 'a/b': 3 }]) + expect((await runTool('{"a/b":null}', schema, noRepair)).executed).toEqual( + [], + ) + expect((await runTool('{"a/b":"3"}', schema, noRepair)).executed).toEqual( + [], + ) + expect( + (await runTool('{"50%":null,"copy":{"id":1}}', schema)).executed, + ).toEqual([{ '50%': null, copy: { id: 1 } }]) + expect((await runTool('{"copy":null}', schema, noRepair)).executed).toEqual( + [], + ) + }) + + it('keeps a $ref that points into a nullable schema', async () => { + // zod-to-json-schema with `target: 'openApi3'` reuses a schema by path. + const address = { + type: 'object', + properties: { city: { type: 'string' } }, + required: ['city'], + } + const schema: JSONSchema = { + type: 'object', + properties: { + home: { anyOf: [address, { type: 'string' }], nullable: true }, + work: { $ref: '#/properties/home/anyOf/0' }, + office: { nullable: true, $ref: '#/properties/home/anyOf/0' }, + }, + } + expect( + (await runTool('{"home":null,"work":{"city":"x"},"office":null}', schema)) + .executed, + ).toEqual([{ home: null, work: { city: 'x' }, office: null }]) + expect((await runTool('{"work":null}', schema, noRepair)).executed).toEqual( + [], + ) + expect((await runTool('{"office":{}}', schema, noRepair)).executed).toEqual( + [], + ) + expect((await runTool('{"home":1}', schema, noRepair)).executed).toEqual([]) + }) + + it('moves a $ref through nested nullable schemas', async () => { + const schema: JSONSchema = { + type: 'object', + properties: { + outer: { + nullable: true, + allOf: [ + { + type: 'object', + properties: { + inner: { + nullable: true, + oneOf: [{ type: 'integer' }, { type: 'string' }], + }, + }, + }, + ], + }, + copy: { $ref: '#/properties/outer/allOf/0/properties/inner/oneOf/0' }, + }, + } + expect( + (await runTool('{"outer":{"inner":null},"copy":3}', schema)).executed, + ).toEqual([{ outer: { inner: null }, copy: 3 }]) + expect((await runTool('{"copy":"3"}', schema, noRepair)).executed).toEqual( + [], + ) + expect( + (await runTool('{"outer":{"inner":true}}', schema, noRepair)).executed, + ).toEqual([]) + }) + + it('keeps draft-07 $ref siblings ignored in a nullable schema', async () => { + const schema: JSONSchema = { + $schema: 'http://json-schema.org/draft-07/schema#', + type: 'object', + definitions: { P: { type: 'object' } }, + properties: { + p: { $ref: '#/definitions/P', nullable: true, required: ['name'] }, + }, + } + expect((await runTool('{"p":{"id":1}}', schema)).executed).toEqual([ + { p: { id: 1 } }, + ]) + expect((await runTool('{"p":null}', schema)).executed).toEqual([ + { p: null }, + ]) + }) + + it('lets an OpenAPI 3.0 nullable field be null', async () => { + const schema: JSONSchema = { + type: 'object', + properties: { + name: { type: 'string', nullable: true }, + kind: { type: 'string', enum: ['a', 'b'], nullable: true }, + size: { type: 'integer' }, + }, + } + expect( + (await runTool('{"name":null,"kind":null}', schema)).executed, + ).toEqual([{ name: null, kind: null }]) + expect((await runTool('{"kind":"c"}', schema)).executed).toEqual([]) + expect((await runTool('{"size":null}', schema, noRepair)).executed).toEqual( + [], + ) + }) + + it('repairs a value inside an OpenAPI 3.0 nullable schema', async () => { + const schema: JSONSchema = { + type: 'object', + properties: { + count: { nullable: true, allOf: [{ type: 'integer' }] }, + size: { type: 'integer', nullable: true }, + }, + } + expect( + (await runTool('{"count":"3","size":"4"}', schema)).executed, + ).toEqual([{ count: 3, size: 4 }]) + expect( + (await runTool('{"count":null,"size":null}', schema)).executed, + ).toEqual([{ count: null, size: null }]) + }) + + it('treats format as an annotation', async () => { + const result = await runTool('{"path":"/tmp/report.txt"}', { + type: 'object', + properties: { path: { type: 'string', format: 'uri' } }, + required: ['path'], + }) + expect(result.executed).toEqual([{ path: '/tmp/report.txt' }]) + }) + + it('keeps a format key inside const and enum values', async () => { + const input = { + file: { format: 'pdf' }, + kind: { format: 'csv' }, + default: '/tmp/report.txt', + } + const result = await runTool(JSON.stringify(input), { + type: 'object', + properties: { + file: { const: { format: 'pdf' } }, + kind: { enum: [{ format: 'csv' }] }, + // A property named like a data keyword is still a schema. + default: { type: 'string', format: 'uri' }, + }, + }) + expect(result.executed).toEqual([input]) + }) + + it('lists at most 10 errors', async () => { + const result = await runTool( + JSON.stringify({ xs: Array.from({ length: 12 }, () => 'x') }), + { + type: 'object', + properties: { xs: { type: 'array', items: { type: 'integer' } } }, + }, + ) + const error = JSON.stringify(result.results[0]?.result) + expect(error).toContain('#/xs/9: ') + expect(error).not.toContain('#/xs/10: ') + expect(error).toContain(', and 2 more') + }) + + it('checks a middleware rewrite to undefined as null', async () => { + const result = await runTool('{"count":2}', jsonSchema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: undefined, + }), + }, + }) + expect(result.executed).toEqual([]) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining( + 'check: #: Instance type "null" is invalid. Expected "object".', + ), + }) + }) + + it('treats a key set to undefined as a missing key', async () => { + const result = await runTool('{"count":2}', jsonSchema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 2, note: undefined }, + }), + }, + }) + expect(result.executed).toEqual([{ count: 2, note: undefined }]) + }) + + it('does not change the tool schema', async () => { + const schema: JSONSchema = { + type: 'object', + $defs: { count: { type: 'integer' } }, + properties: { + count: { $ref: '#/$defs/count' }, + name: { type: 'string', format: 'uri', nullable: true }, + kind: { enum: ['a'], nullable: true }, + owner: { nullable: true, allOf: [{ $ref: '#/$defs/count' }] }, + }, + } + const before = JSON.stringify(schema) + await runTool('{"count":2}', schema) + expect(JSON.stringify(schema)).toBe(before) + expect(Object.getOwnPropertyNames(schema.properties?.count)).toEqual([ + '$ref', + ]) + }) + + it('works when runtime code generation is blocked', async () => { + vi.stubGlobal('Function', function blockedCodeGeneration() { + throw new Error('Code generation is blocked') + }) + try { + const schema: JSONSchema = { + type: 'object', + properties: { count: { type: 'integer' } }, + } + expect((await runTool('{"count":2}', schema)).executed).toEqual([ + { count: 2 }, + ]) + expect((await runTool('{"count":"2"}', schema)).executed).toEqual([ + { count: 2 }, + ]) + expect((await runTool('{"count":"x"}', schema)).executed).toEqual([]) + } finally { + vi.unstubAllGlobals() + } + }) +}) + +describe('tool input repair', () => { + const schema: JSONSchema = { + type: 'object', + properties: { count: { type: 'integer', minimum: 1 } }, + required: ['count'], + additionalProperties: false, + } + + it('coerces a numeric field before execution', async () => { + const result = await runTool('{"count":"2"}', schema) + expect(result.executed).toEqual([{ count: 2 }]) + }) + + it('does not change the received arguments', async () => { + const received = { count: '2' } + expect(await validateToolInput(schema, received, 'check')).toEqual({ + count: 2, + }) + expect(received).toEqual({ count: '2' }) + }) + + it('shows the original arguments when the repaired value still fails', async () => { + const result = await runTool('{"count":"0"}', schema) + expect(result.executed).toEqual([]) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Received arguments: {"count":"0"}'), + }) + }) + + it('follows local references and drops a null optional field', async () => { + const result = await runTool('{"count":"3","label":null,"omitted":null}', { + type: 'object', + $defs: { count: { type: 'integer', minimum: 1 } }, + properties: { + count: { $ref: '#/$defs/count' }, + label: { type: ['string', 'null'] }, + omitted: { type: 'string' }, + }, + required: ['count'], + }) + expect(result.executed).toEqual([{ count: 3, label: null }]) + }) + + it.each<[unknown, JSONSchema, unknown]>([ + ['12', { anyOf: [{ type: 'integer' }, { type: 'string' }] }, '12'], + ['12', { type: ['integer', 'string'] }, '12'], + ['3', { type: 'integer' }, 3], + [['1', '2'], { type: 'array', items: { type: 'integer' } }, [1, 2]], + [ + ['1', 'x'], + { prefixItems: [{ type: 'integer' }, { type: 'string' }] }, + [1, 'x'], + ], + ['true', { type: 'boolean' }, true], + [null, { type: 'string' }, ''], + ])('repairs %j against %j', async (received, inputSchema, expected) => { + expect(await validateToolInput(inputSchema, received, 'check')).toEqual( + expected, + ) + }) + + it('rejects a primitive that cannot be repaired', async () => { + await expect( + validateToolInput({ type: 'integer' }, 'wrong', 'check'), + ).rejects.toThrow('Input validation failed for tool check') + }) + + it('keeps a valid referenced union branch', async () => { + const result = await runTool('{"count":"12"}', { + type: 'object', + $defs: { numeric: { type: 'integer' }, text: { type: 'string' } }, + properties: { + count: { + anyOf: [{ $ref: '#/$defs/numeric' }, { $ref: '#/$defs/text' }], + }, + }, + required: ['count'], + }) + expect(result.executed).toEqual([{ count: '12' }]) + }) + + it('coerces conditional branch fields', async () => { + const result = await runTool('{"mode":"count","count":"4"}', { + type: 'object', + if: { properties: { mode: { const: 'count' } }, required: ['mode'] }, + then: { + properties: { count: { type: 'integer' } }, + required: ['count'], + }, + }) + expect(result.executed).toEqual([{ mode: 'count', count: 4 }]) + }) + + it('keeps a null field that a conditional branch requires', async () => { + const result = await runTool('{"mode":"count","count":null}', { + type: 'object', + properties: { mode: { type: 'string' }, count: { type: 'integer' } }, + if: { properties: { mode: { const: 'count' } }, required: ['mode'] }, + then: { required: ['count'] }, + }) + expect(result.executed).toEqual([{ mode: 'count', count: 0 }]) + }) + + it('keeps a null field that an intersection requires', async () => { + const result = await runTool('{"count":null}', { + type: 'object', + allOf: [ + { properties: { count: { type: 'integer' } } }, + { required: ['count'] }, + ], + }) + expect(result.executed).toEqual([{ count: 0 }]) + }) + + it('keeps a null field required outside a union branch', async () => { + const result = await runTool('{"count":null}', { + type: 'object', + required: ['count'], + anyOf: [{ properties: { count: { type: 'integer' } } }], + }) + expect(result.executed).toEqual([{ count: 0 }]) + }) + + it('coerces pattern and additional properties', async () => { + const result = await runTool('{"n_1":"4","extra":"5"}', { + type: 'object', + patternProperties: { '^n_': { type: 'integer' } }, + additionalProperties: { type: 'number' }, + }) + expect(result.executed).toEqual([{ n_1: 4, extra: 5 }]) + }) + + it('coerces an approval edit before execution', async () => { + const result = await runTool('{"count":2}', schema, { + needsApproval: true, + approvals: new Map([ + ['call-1', { approved: true, editedArgs: { count: '4' } }], + ]), + }) + expect(result.executed).toEqual([{ count: 4 }]) + }) + + it('coerces a middleware rewrite before a client request', async () => { + const result = await runTool('{"count":2}', schema, { + client: true, + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: '5' }, + }), + }, + }) + expect(result.needsClientExecution).toEqual([ + { toolCallId: 'call-1', toolName: 'check', input: { count: 5 } }, + ]) + }) + + it('repairs input for a Standard Schema that exports JSON Schema', async () => { + const result = await runTool( + '{"count":"2","label":null}', + z.object({ count: z.number(), label: z.string().optional() }), + ) + expect(result.executed).toEqual([{ count: 2 }]) + }) + + it('runs the authored transform on the repaired input', async () => { + const result = await runTool( + '{"count":"2"}', + z.object({ count: z.number().transform((value) => value * 10) }), + ) + expect(result.executed).toEqual([{ count: 20 }]) + }) + + it('reports the issues of the first run when repair does not help', async () => { + const result = await runTool( + '{"count":"wrong"}', + z.object({ count: z.number() }), + ) + expect(result.executed).toEqual([]) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Received arguments: {"count":"wrong"}'), + }) + }) +}) + +const clientCheck = ( + inputSchema: SchemaInput, + needsApproval = false, +): Tool => ({ + name: 'check', + description: 'Check input', + inputSchema, + needsApproval, +}) + +const history = (input: unknown): Array => [ + { + role: 'assistant', + content: null, + toolCalls: [ + { + id: 'call', + type: 'function', + function: { name: 'check', arguments: JSON.stringify(input) }, + }, + ], + }, +] + +function descriptors(chunks: Array) { + return chunks.flatMap((chunk) => + chunk.type === EventType.RUN_FINISHED && chunk.outcome?.type === 'interrupt' + ? chunk.outcome.interrupts + : [], + ) +} + +function clientInput(chunks: Array) { + return descriptors(chunks).find((entry) => entry.id === 'client_tool_call') + ?.metadata?.input +} + +const review = defineInterrupt({ + id: 'boundary-review', + responseSchema: z.boolean(), +}) + +// Pauses at a tool boundary, so client requests come from that boundary. +const reviewAt = (phase: 'beforeTools' | 'afterTools') => + defineChatMiddleware({ + name: 'boundary', + onInterruptBoundary(ctx) { + if (ctx.phase !== phase) return undefined + return { + interrupts: [ + review.interrupt({ + key: 'review', + reason: 'review', + message: 'Review', + }), + ], + } + }, + }) + +async function initial( + input: unknown, + tool: Tool, + middleware: Array = [], + boundary?: 'beforeTools' | 'afterTools', + repairToolInput?: TextOptions['repairToolInput'], +) { + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('call', 'check'), + ev.toolArgs('call', JSON.stringify(input)), + ev.runFinished('tool_calls'), + ], + ], + }) + const messages: Array = [{ role: 'user', content: 'check' }] + if (!boundary) { + return collectChunks( + chat({ adapter, messages, tools: [tool], middleware, repairToolInput }), + ) + } + return collectChunks( + chat({ + adapter, + messages, + tools: [tool], + interrupts: [review], + middleware: [reviewAt(boundary), ...middleware], + repairToolInput, + }), + ) +} + +describe('client tool input in chat()', () => { + it.each([false, true])( + 'shows middleware the raw input, then sends the checked input: boundary=%s', + async (boundary) => { + let transforms = 0 + const seen: Array = [] + const inputSchema = z.object({ + count: z.string().transform((value) => { + transforms++ + return Number(value) + }), + }) + const chunks = await initial( + { count: '2' }, + clientCheck(inputSchema), + [ + { + onBeforeToolCall(_ctx, hook) { + seen.push(structuredClone(hook.args)) + }, + }, + ], + boundary ? 'beforeTools' : undefined, + ) + expect(seen).toEqual([{ count: '2' }]) + expect(transforms).toBe(1) + expect(clientInput(chunks)).toEqual({ count: 2 }) + }, + ) + + it.each([false, true])( + 'sends no client request when middleware makes the input invalid: boundary=%s', + async (boundary) => { + const chunks = await initial( + { count: 2 }, + clientCheck(schema), + [ + { + onBeforeToolCall: () => ({ + type: 'transformArgs', + args: { count: 'wrong' }, + }), + }, + ], + boundary ? 'beforeTools' : undefined, + ) + expect(clientInput(chunks)).toBeUndefined() + expect( + chunks.some( + (chunk) => + chunk.type === EventType.TOOL_CALL_RESULT && + typeof chunk.content === 'string' && + chunk.content.includes('Input validation failed'), + ), + ).toBe(true) + }, + ) + + it('runs no middleware for the pending approval preview', async () => { + let hooks = 0 + const chunks = await initial( + { count: '2' }, + clientCheck(z.object({ count: z.string().transform(Number) }), true), + [ + { + onBeforeToolCall() { + hooks++ + }, + }, + ], + ) + expect(hooks).toBe(0) + expect( + descriptors(chunks).find((entry) => entry.id === 'approval_call') + ?.metadata?.input, + ).toEqual({ count: 2 }) + expect(clientInput(chunks)).toBeUndefined() + }) + + it('checks an approval edit after middleware on resume', async () => { + const { adapter, calls } = createMockAdapter({ iterations: [] }) + const seen: Array = [] + const inputSchema = z.object({ + count: z.number().transform((value) => value + 1), + }) + const chunks = await collectChunks( + chat({ + adapter, + parentRunId: 'old', + messages: history({ count: 1 }), + tools: [clientCheck(inputSchema, true)], + resume: [ + { + interruptId: 'approval_call', + status: 'resolved', + payload: { approved: true, editedArgs: { count: 2 } }, + }, + ], + middleware: [ + { + onBeforeToolCall(_ctx, hook) { + seen.push(structuredClone(hook.args)) + }, + }, + ], + }), + ) + expect(calls).toHaveLength(0) + expect(seen).toEqual([{ count: 2 }]) + expect(clientInput(chunks)).toEqual({ count: 3 }) + }) + + it('keeps the client result, with no second hook run, when it comes back', async () => { + const { adapter, calls } = createMockAdapter({ + iterations: [[ev.runStarted(), ev.runFinished()]], + }) + let hooks = 0 + await collectChunks( + chat({ + adapter, + parentRunId: 'old', + messages: history({ count: 'wrong' }), + tools: [clientCheck(schema)], + resume: [ + { + interruptId: 'client_tool_call', + status: 'resolved', + payload: { ok: true }, + }, + ], + middleware: [ + { + onBeforeToolCall() { + hooks++ + }, + }, + ], + }), + ) + expect(hooks).toBe(0) + const results = calls[0]?.messages.filter( + (message) => message.role === 'tool', + ) + expect(results).toHaveLength(1) + expect(JSON.stringify(results?.[0])).not.toContain( + 'Input validation failed', + ) + }) + + it.each([false, true])( + 'runs the hook once for a client tool at an afterTools boundary: invalid=%s', + async (invalid) => { + let hooks = 0 + const chunks = await initial( + { count: 2 }, + clientCheck(schema), + [ + { + onBeforeToolCall() { + hooks++ + return invalid + ? { type: 'transformArgs', args: { count: 'wrong' } } + : undefined + }, + }, + ], + 'afterTools', + ) + expect(hooks).toBe(1) + const toolResults = chunks.filter( + (chunk) => chunk.type === EventType.TOOL_CALL_RESULT, + ) + expect(toolResults).toHaveLength(invalid ? 1 : 0) + if (invalid) { + expect(clientInput(chunks)).toBeUndefined() + } else { + expect(clientInput(chunks)).toEqual({ count: 2 }) + } + }, + ) + + it('records no execute_tool span for a client tool', async () => { + const fakeTracer = createFakeTracer() + await initial({ count: 2 }, clientCheck(schema), [ + otelMiddleware({ + tracer: fakeTracer.tracer, + meter: createFakeMeter().meter, + }), + ]) + expect( + fakeTracer.spans.some((span) => span.name.startsWith('execute_tool')), + ).toBe(false) + }) +}) + +describe('repairToolInput', () => { + const countSchema: JSONSchema = { + type: 'object', + properties: { count: { type: 'integer' } }, + required: ['count'], + } + + /** One model turn that calls `check` with each input, then a turn that answers. */ + async function run( + inputs: Array, + inputSchema: SchemaInput, + repairToolInput?: TextOptions['repairToolInput'], + ) { + const executed: Array = [] + const tool: Tool = { + name: 'check', + description: 'Check', + inputSchema, + execute: (input) => { + executed.push(input) + return 'ok' + }, + } + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ...inputs.flatMap((input, index) => [ + ev.toolStart(`call-${index}`, 'check'), + ev.toolArgs(`call-${index}`, JSON.stringify(input)), + ev.toolEnd(`call-${index}`), + ]), + ev.runFinished('tool_calls'), + ], + [ev.runStarted(), ev.textContent('Done'), ev.runFinished('stop')], + ], + }) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'check' }], + tools: [tool], + repairToolInput, + }), + ) + const results = chunks.flatMap((chunk) => + chunk.type === EventType.TOOL_CALL_RESULT && + typeof chunk.content === 'string' + ? [chunk.content] + : [], + ) + const errors = results.filter((content) => + content.includes('Input validation failed'), + ) + return { executed, results, errors } + } + + it('turns the built-in repair off with false', async () => { + const result = await run([{ count: '3' }], countSchema, false) + expect(result.executed).toEqual([]) + expect(result.errors).toHaveLength(1) + }) + + it('turns the repair off with false for a Standard Schema', async () => { + const result = await run( + [{ count: '3' }], + z.object({ count: z.number() }), + false, + ) + expect(result.executed).toEqual([]) + expect(result.errors).toHaveLength(1) + }) + + it.each([false, true])( + 'turns the repair off with false for a client tool: boundary=%s', + async (boundary) => { + const chunks = await initial( + { count: '3' }, + clientCheck(schema), + [], + boundary ? 'beforeTools' : undefined, + false, + ) + expect(clientInput(chunks)).toBeUndefined() + expect( + chunks.some( + (chunk) => + chunk.type === EventType.TOOL_CALL_RESULT && + typeof chunk.content === 'string' && + chunk.content.includes('Input validation failed'), + ), + ).toBe(true) + }, + ) + + it('replaces the built-in repair with a function', async () => { + const repair = vi.fn(({ input }: ToolInputRepairContext) => + typeof input === 'object' && + input !== null && + 'count' in input && + input.count === 'many' + ? { count: 7 } + : undefined, + ) + const result = await run( + [{ count: 'many' }, { count: '3' }], + countSchema, + repair, + ) + // `'3'` is not repaired: the function said `undefined`. + expect(result.executed).toEqual([{ count: 7 }]) + expect(result.errors).toHaveLength(1) + expect(result.errors[0]).toContain( + String.raw`Received arguments: {\"count\":\"3\"}`, + ) + const context = repair.mock.calls[0]?.[0] + expect(context?.tool.name).toBe('check') + expect(context?.toolCall.id).toBe('call-0') + expect(context?.input).toEqual({ count: 'many' }) + expect(context?.error.message).toContain('Input validation failed') + }) + + it('does not call the function for a valid input', async () => { + const repair = vi.fn(() => ({ count: 9 })) + const result = await run([{ count: 2 }], countSchema, repair) + expect(repair).not.toHaveBeenCalled() + expect(result.executed).toEqual([{ count: 2 }]) + }) + + it('reports the original error when the repaired input still fails', async () => { + const result = await run([{ count: 'many' }], countSchema, () => ({ + count: 'still wrong', + })) + expect(result.executed).toEqual([]) + expect(result.errors[0]).toContain( + String.raw`Received arguments: {\"count\":\"many\"}`, + ) + }) + + it('runs a Standard Schema, with its transform, on the repaired input', async () => { + const result = await run( + [{ count: 'two' }], + z.object({ count: z.number().transform((value) => value * 10) }), + () => ({ count: 2 }), + ) + expect(result.executed).toEqual([{ count: 20 }]) + }) + + it('makes the thrown error the tool error', async () => { + const result = await run([{ count: 'many' }], countSchema, () => { + throw new Error('cannot repair') + }) + expect(result.executed).toEqual([]) + // The model's own arguments follow the thrown message. + expect(result.results).toEqual([ + JSON.stringify({ + error: 'cannot repair\nReceived arguments: {"count":"many"}', + }), + ]) + }) +}) diff --git a/packages/openai-base/src/adapters/chat-completions-text.ts b/packages/openai-base/src/adapters/chat-completions-text.ts index ae0e6172f8..6c4698e476 100644 --- a/packages/openai-base/src/adapters/chat-completions-text.ts +++ b/packages/openai-base/src/adapters/chat-completions-text.ts @@ -27,6 +27,10 @@ import type { StructuredOutputCompatibility, } from '../utils/schema-converter' import { buildChatCompletionsUsage } from '../usage' +import { + explicitJsonObject, + withJsonModePrompt, +} from '../utils/json-object-mode' import { toChatCompletionsToolChoice } from '../tools/tool-choice' import { convertToolsToChatCompletionsFormat } from './chat-completions-tool-converter' import type OpenAI from 'openai' @@ -220,6 +224,7 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< toolCallId, toolCallName: rejectedToolCall.toolName, toolName: rejectedToolCall.toolName, + args: rejectedToolCall.arguments, ...(rejectedToolCall.input !== undefined && { input: rejectedToolCall.input, }), @@ -299,11 +304,13 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< options: StructuredOutputOptions, ): Promise> { const { chatOptions, outputSchema } = options - const requestParams = this.mapOptionsToRequest(chatOptions) - - const jsonSchema = this.makeStructuredOutputCompatible( + const jsonMode = this.usesJsonObject(chatOptions) + const requestParams = this.mapOptionsToRequest( + jsonMode ? withJsonModePrompt(chatOptions, outputSchema) : chatOptions, + ) + const responseFormat = this.chatCompletionsResponseFormat( + jsonMode, outputSchema, - outputSchema.required, ) try { @@ -325,14 +332,7 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< { ...cleanParams, stream: false, - response_format: { - type: 'json_schema', - json_schema: { - name: 'structured_output', - schema: jsonSchema, - strict: true, - }, - }, + response_format: responseFormat, }, extractRequestOptions(chatOptions.request), ) @@ -408,11 +408,13 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< options: StructuredOutputOptions, ): AsyncIterable { const { chatOptions, outputSchema } = options - const requestParams = this.mapOptionsToRequest(chatOptions) - - const jsonSchema = this.makeStructuredOutputCompatible( + const jsonMode = this.usesJsonObject(chatOptions) + const requestParams = this.mapOptionsToRequest( + jsonMode ? withJsonModePrompt(chatOptions, outputSchema) : chatOptions, + ) + const responseFormat = this.chatCompletionsResponseFormat( + jsonMode, outputSchema, - outputSchema.required, ) const aguiState = { @@ -493,14 +495,7 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< ...cleanParams, stream: true, stream_options: { include_usage: true }, - response_format: { - type: 'json_schema', - json_schema: { - name: 'structured_output', - schema: jsonSchema, - strict: true, - }, - }, + response_format: responseFormat, }, extractRequestOptions(chatOptions.request), ) @@ -818,6 +813,56 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< .schema } + /** + * `true` when this model takes `json_object` and rejects strict + * `json_schema`. Subclasses answer from their model catalog. Default `false`. + */ + protected modelPrefersJsonObject(_model: string): boolean { + return false + } + + /** JSON mode when `modelOptions` sets it or the model catalog needs it. */ + protected usesJsonObject(options: { + model: string + modelOptions?: unknown + }): boolean { + return ( + explicitJsonObject(options.modelOptions) || + this.modelPrefersJsonObject(options.model) + ) + } + + protected chatCompletionsResponseFormat( + jsonMode: boolean, + schema: Record, + ): + | { type: 'json_object' } + | { + type: 'json_schema' + json_schema: { + name: string + schema: Record + strict: true + } + } { + if (jsonMode) { + return { type: 'json_object' } + } + return { + type: 'json_schema', + json_schema: { + name: 'structured_output', + schema: this.makeStructuredOutputCompatible( + schema, + Array.isArray(schema.required) + ? (schema.required as Array) + : undefined, + ), + strict: true, + }, + } + } + /** * Extract reasoning content from a stream chunk. Default returns * `undefined` because the OpenAI Chat Completions chunk shape doesn't @@ -1166,7 +1211,8 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< toolName: toolCall.name, model: chunk.model || options.model, timestamp: Date.now(), - input: parsedInput, + args: toolCall.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } emittedAnyToolCallEnd = true } @@ -1236,7 +1282,8 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< toolName: toolCall.name, model: lastModel || options.model, timestamp: Date.now(), - input: parsedInput, + args: toolCall.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } pendingToolCount += 1 emittedAnyToolCallEnd = true @@ -1468,24 +1515,18 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< // `response_format` alongside any `tools`. Modern OpenAI-compatible // Chat Completions accepts both together and emits the schema- // constrained text on the natural final turn. - const combinedSchema = options.outputSchema as - | Record - | undefined + // JSON mode is never combined (see `supportsCombinedToolsAndSchema`). + // A caller's own `json_object` reaches the request through the + // `modelOptions` spread. + const combinedSchema = this.usesJsonObject(options) + ? undefined + : (options.outputSchema as Record | undefined) const responseFormat = combinedSchema ? { - response_format: { - type: 'json_schema' as const, - json_schema: { - name: 'structured_output', - schema: this.makeStructuredOutputCompatible( - combinedSchema, - Array.isArray(combinedSchema.required) - ? (combinedSchema.required as Array) - : undefined, - ), - strict: true, - }, - }, + response_format: this.chatCompletionsResponseFormat( + false, + combinedSchema, + ), } : undefined @@ -1534,8 +1575,10 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< * return `false` because its API rejects schema + tools + stream with a * 400. */ - supportsCombinedToolsAndSchema(): boolean { - return true + supportsCombinedToolsAndSchema(modelOptions?: unknown): boolean { + // JSON mode sends the schema as a system prompt on a separate + // structured-output request, not on the tool-calling turns. + return !this.usesJsonObject({ model: this.model, modelOptions }) } /** diff --git a/packages/openai-base/src/adapters/responses-text.ts b/packages/openai-base/src/adapters/responses-text.ts index cb8c05b562..8803b2fb0a 100644 --- a/packages/openai-base/src/adapters/responses-text.ts +++ b/packages/openai-base/src/adapters/responses-text.ts @@ -30,6 +30,10 @@ import type { StructuredOutputCompatibility, } from '../utils/schema-converter' import { buildResponsesUsage } from '../usage' +import { + explicitJsonObject, + withJsonModePrompt, +} from '../utils/json-object-mode' import { toResponsesToolChoice } from '../tools/tool-choice' import { getOpenAIProviderToolKind } from '../tools/openai-provider-tool' import { convertToolsToResponsesFormat } from './responses-tool-converter' @@ -322,17 +326,56 @@ function fromCompactedItem( content: item.content.map((part): ContentPart => { if (isRecord(part) && typeof part.text === 'string') return { type: 'text', content: part.text } - if ( - isRecord(part) && - part.type === 'input_image' && - typeof part.image_url === 'string' - ) - return { - type: 'image', - source: { type: 'url', value: part.image_url }, - } - // ponytail: text and image urls only. Map files here when a - // compacted history holds them. + if (isRecord(part) && part.type === 'input_image') { + if (typeof part.image_url === 'string') + return { + type: 'image', + source: { type: 'url', value: part.image_url }, + } + if (typeof part.file_id === 'string') + return { + type: 'image', + source: { type: 'file', value: part.file_id }, + } + } + if (isRecord(part) && part.type === 'input_file') { + if (typeof part.file_id === 'string') + return { + type: 'document', + source: { type: 'file', value: part.file_id }, + } + if (typeof part.file_url === 'string') + return { + type: 'document', + source: { type: 'url', value: part.file_url }, + } + if ( + typeof part.file_data === 'string' && + typeof part.filename === 'string' + ) + return { + type: 'document', + source: { + type: 'data', + value: part.file_data, + mimeType: 'application/pdf', + }, + metadata: { filename: part.filename }, + } + // Audio leaves as input_file with file_data and no filename. + if (typeof part.file_data === 'string') + return { + type: 'audio', + source: { + type: 'data', + value: part.file_data, + mimeType: 'application/octet-stream', + }, + } + } + // ponytail: other item types still throw. A throw with no HTTP + // status does not latch. Pass an item through when a real compact + // result holds one. throw new Error( `Cannot map a compacted ${isRecord(part) ? String(part.type) : 'part'} part.`, ) @@ -589,13 +632,11 @@ export abstract class OpenAIBaseResponsesTextAdapter< options: StructuredOutputOptions, ): Promise> { const { chatOptions, outputSchema } = options - const requestParams = this.mapOptionsToRequest(chatOptions) - - // Apply provider-specific transformations for structured output compatibility - const jsonSchema = this.makeStructuredOutputCompatible( - outputSchema, - outputSchema.required, + const jsonMode = this.usesJsonObject(chatOptions) + const requestParams = this.mapOptionsToRequest( + jsonMode ? withJsonModePrompt(chatOptions, outputSchema) : chatOptions, ) + const textFormat = this.responsesStructuredFormat(jsonMode, outputSchema) try { // Strip streaming-only fields a subclass override of mapOptionsToRequest @@ -621,14 +662,8 @@ export abstract class OpenAIBaseResponsesTextAdapter< { ...(cleanParams as Omit), stream: false, - // Configure structured output via text.format text: { - format: { - type: 'json_schema', - name: 'structured_output', - schema: jsonSchema, - strict: true, - }, + format: textFormat, }, }, extractRequestOptions(chatOptions.request), @@ -703,12 +738,11 @@ export abstract class OpenAIBaseResponsesTextAdapter< options: StructuredOutputOptions, ): AsyncIterable { const { chatOptions, outputSchema } = options - const requestParams = this.mapOptionsToRequest(chatOptions) - - const jsonSchema = this.makeStructuredOutputCompatible( - outputSchema, - outputSchema.required, + const jsonMode = this.usesJsonObject(chatOptions) + const requestParams = this.mapOptionsToRequest( + jsonMode ? withJsonModePrompt(chatOptions, outputSchema) : chatOptions, ) + const textFormat = this.responsesStructuredFormat(jsonMode, outputSchema) const aguiState = { runId: generateId(this.name), @@ -808,12 +842,7 @@ export abstract class OpenAIBaseResponsesTextAdapter< ...cleanParams, stream: true, text: { - format: { - type: 'json_schema', - name: 'structured_output', - schema: jsonSchema, - strict: true, - }, + format: textFormat, }, }, extractRequestOptions(chatOptions.request), @@ -1204,6 +1233,52 @@ export abstract class OpenAIBaseResponsesTextAdapter< .schema } + /** + * `true` when this model takes `json_object` and rejects strict + * `json_schema`. Subclasses answer from their model catalog. Default `false`. + */ + protected modelPrefersJsonObject(_model: string): boolean { + return false + } + + /** JSON mode when `modelOptions` sets it or the model catalog needs it. */ + protected usesJsonObject(options: { + model: string + modelOptions?: unknown + }): boolean { + return ( + explicitJsonObject(options.modelOptions) || + this.modelPrefersJsonObject(options.model) + ) + } + + protected responsesStructuredFormat( + jsonMode: boolean, + schema: Record, + ): + | { type: 'json_object' } + | { + type: 'json_schema' + name: string + schema: Record + strict: true + } { + if (jsonMode) { + return { type: 'json_object' } + } + return { + type: 'json_schema', + name: 'structured_output', + schema: this.makeStructuredOutputCompatible( + schema, + Array.isArray(schema.required) + ? (schema.required as Array) + : undefined, + ), + strict: true, + } + } + /** * Final shaping pass applied to parsed structured-output JSON before it is * returned to the caller. Default is a passthrough. @@ -2074,7 +2149,7 @@ export abstract class OpenAIBaseResponsesTextAdapter< rawArguments: chunk.arguments, }, ) - parsedInput = {} + parsedInput = undefined } } @@ -2085,7 +2160,8 @@ export abstract class OpenAIBaseResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: chunk.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } } @@ -2171,7 +2247,7 @@ export abstract class OpenAIBaseResponsesTextAdapter< rawArguments: rawArgs, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -2181,7 +2257,8 @@ export abstract class OpenAIBaseResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: rawArgs, + ...(parsedInput !== undefined && { input: parsedInput }), } metadata.ended = true metadata.pendingArguments = undefined @@ -2342,7 +2419,7 @@ export abstract class OpenAIBaseResponsesTextAdapter< rawArguments: rawArgs, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -2352,7 +2429,8 @@ export abstract class OpenAIBaseResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: rawArgs, + ...(parsedInput !== undefined && { input: parsedInput }), } metadata.ended = true metadata.pendingArguments = undefined @@ -2576,23 +2654,16 @@ export abstract class OpenAIBaseResponsesTextAdapter< // (pre-converted at the activity boundary). Wire it into `text.format` // alongside any `tools` — the Responses API supports both together and // emits the schema-constrained text on the natural final turn. - const combinedSchema = options.outputSchema as - | Record - | undefined + // JSON mode is never combined (see `supportsCombinedToolsAndSchema`). + // A caller's own `json_object` reaches the request through the + // `modelOptions` spread. + const combinedSchema = this.usesJsonObject(options) + ? undefined + : (options.outputSchema as Record | undefined) const textFormat = combinedSchema ? { text: { - format: { - type: 'json_schema' as const, - name: 'structured_output', - schema: this.makeStructuredOutputCompatible( - combinedSchema, - Array.isArray(combinedSchema.required) - ? (combinedSchema.required as Array) - : undefined, - ), - strict: true, - }, + format: this.responsesStructuredFormat(false, combinedSchema), }, } : undefined @@ -2719,8 +2790,10 @@ export abstract class OpenAIBaseResponsesTextAdapter< * together in a single streaming request (per issue #605). Subclasses * that route to providers without this capability should override. */ - supportsCombinedToolsAndSchema(): boolean { - return true + supportsCombinedToolsAndSchema(modelOptions?: unknown): boolean { + // JSON mode sends the schema as a system prompt on a separate + // structured-output request, not on the tool-calling turns. + return !this.usesJsonObject({ model: this.model, modelOptions }) } /** diff --git a/packages/openai-base/src/usage.ts b/packages/openai-base/src/usage.ts index 2555494bf0..e0a8387faf 100644 --- a/packages/openai-base/src/usage.ts +++ b/packages/openai-base/src/usage.ts @@ -90,8 +90,8 @@ export function buildChatCompletionsUsage( * `ResponseUsage` object. * * Shared by every provider that routes through - * {@link OpenAIBaseResponsesTextAdapter}. Surfaces cached prompt tokens and - * reasoning detail tokens when present. Returns `undefined` when the provider + * {@link OpenAIBaseResponsesTextAdapter}. Surfaces cache read, cache write, + * and reasoning detail tokens when present. Returns `undefined` when the provider * reported no usage object, so callers omit the field rather than fabricating * zeroed totals. */ @@ -116,6 +116,15 @@ export function buildResponsesUsage( } } + // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition + const cacheWriteTokens = usage.input_tokens_details?.cache_write_tokens + if (cacheWriteTokens && cacheWriteTokens > 0) { + result.promptTokensDetails = { + ...result.promptTokensDetails, + cacheWriteTokens, + } + } + // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition const reasoningTokens = usage.output_tokens_details?.reasoning_tokens if (reasoningTokens && reasoningTokens > 0) { diff --git a/packages/openai-base/src/utils/json-object-mode.ts b/packages/openai-base/src/utils/json-object-mode.ts new file mode 100644 index 0000000000..bcc72bf550 --- /dev/null +++ b/packages/openai-base/src/utils/json-object-mode.ts @@ -0,0 +1,34 @@ +import { jsonModeSystemPrompt } from '@tanstack/ai-utils' +import type { TextOptions } from '@tanstack/ai' + +/** Caller set JSON mode on a Chat Completions or Responses body. */ +export function explicitJsonObject(modelOptions: unknown): boolean { + if (modelOptions == null || typeof modelOptions !== 'object') return false + if ('response_format' in modelOptions) { + const format = (modelOptions as { response_format?: { type?: unknown } }) + .response_format + if (format?.type === 'json_object') return true + } + if ('text' in modelOptions) { + const text = (modelOptions as { text?: { format?: { type?: unknown } } }) + .text + if (text?.format?.type === 'json_object') return true + } + return false +} + +/** Add the schema as a system prompt, because JSON mode does not send it. */ +export function withJsonModePrompt< + TProviderOptions extends Record, +>( + options: TextOptions, + schema: unknown, +): TextOptions { + return { + ...options, + systemPrompts: [ + ...(options.systemPrompts ?? []), + jsonModeSystemPrompt(schema), + ], + } +} diff --git a/packages/openai-base/tests/chat-completions-structured-output-stream.test.ts b/packages/openai-base/tests/chat-completions-structured-output-stream.test.ts index be2ee84657..46a30e3d78 100644 --- a/packages/openai-base/tests/chat-completions-structured-output-stream.test.ts +++ b/packages/openai-base/tests/chat-completions-structured-output-stream.test.ts @@ -284,6 +284,37 @@ describe('OpenAIBaseChatCompletionsTextAdapter.structuredOutputStream', () => { expect(request.tools).toBeUndefined() }) + it('sends json_object and the schema as one system prompt when modelOptions sets json_object', async () => { + setupStreamingMock([deltaChunk('{"name":"X","age":1}'), finishChunk()]) + const adapter = new TestAdapter() + + await collect( + adapter.structuredOutputStream!({ + chatOptions: { + model: 'test-model', + messages: [{ role: 'user', content: 'extract' }], + logger: testLogger, + modelOptions: { response_format: { type: 'json_object' } }, + // A direct caller can also set it here; the prompt must not repeat. + outputSchema: personSchema, + }, + outputSchema: personSchema, + }), + ) + + const request = mockCreate.mock.calls[0]![0] as { + response_format?: { type: string; json_schema?: unknown } + messages: Array<{ role: string; content: unknown }> + } + expect(request.response_format).toEqual({ type: 'json_object' }) + expect(request.messages.filter((m) => m.role === 'system')).toEqual([ + { + role: 'system', + content: expect.stringContaining(JSON.stringify(personSchema)), + }, + ]) + }) + it('accumulates JSON text from deltas and emits exactly one structured-output.complete', async () => { setupStreamingMock([ deltaChunk('{"name":"A"'), diff --git a/packages/openai-base/tests/chat-completions-text.test.ts b/packages/openai-base/tests/chat-completions-text.test.ts index 1fe050309a..d43d5ab770 100644 --- a/packages/openai-base/tests/chat-completions-text.test.ts +++ b/packages/openai-base/tests/chat-completions-text.test.ts @@ -1240,7 +1240,8 @@ describe('OpenAIBaseChatCompletionsTextAdapter', () => { const toolEnd = chunks.find((chunk) => chunk.type === 'TOOL_CALL_END') expect(toolEnd).toBeDefined() - expect(toolEnd?.input).toBeUndefined() + expect(toolEnd).toMatchObject({ args: '{"location":' }) + expect(toolEnd).not.toHaveProperty('input') expect( chunks.find((chunk) => chunk.type === 'TOOL_CALL_ARGS'), ).toMatchObject({ @@ -1263,6 +1264,64 @@ describe('OpenAIBaseChatCompletionsTextAdapter', () => { ) }) + describe('raw tool arguments', () => { + it.each([true, false])( + 'sends the accumulated argument string as args on TOOL_CALL_END (finish_reason=%s)', + async (withFinishReason) => { + setupMockSdkClient([ + { + id: 'chatcmpl-raw', + model: 'test-model', + choices: [ + { + delta: { + tool_calls: [ + { + index: 0, + id: 'call_raw', + type: 'function', + function: { name: 'lookup_weather', arguments: '{"n": ' }, + }, + ], + }, + finish_reason: null, + }, + ], + }, + { + id: 'chatcmpl-raw', + model: 'test-model', + choices: [ + { + delta: { + tool_calls: [{ index: 0, function: { arguments: '1.50}' } }], + }, + finish_reason: withFinishReason ? 'tool_calls' : null, + }, + ], + }, + ]) + const adapter = new TestChatCompletionsAdapter(testConfig, 'test-model') + const chunks: Array = [] + for await (const chunk of adapter.chatStream({ + logger: testLogger, + model: 'test-model', + messages: [{ role: 'user', content: 'Weather?' }], + tools: [weatherTool], + })) { + chunks.push(chunk) + } + + const toolEnd = chunks.find((chunk) => chunk.type === 'TOOL_CALL_END') + expect(toolEnd).toMatchObject({ + toolCallId: 'call_raw', + args: '{"n": 1.50}', + input: { n: 1.5 }, + }) + }, + ) + }) + describe('subclassing', () => { it('allows subclassing with custom name', () => { class MyProviderAdapter extends OpenAIBaseChatCompletionsTextAdapter { @@ -1434,6 +1493,43 @@ describe('OpenAIBaseChatCompletionsTextAdapter', () => { ) }) + it('is not combined in JSON mode, and keeps the caller json_object without a schema prompt', async () => { + const streamChunks = [ + { + id: 'chatcmpl-1', + model: 'test-model', + choices: [{ delta: { content: 'ok' }, finish_reason: 'stop' }], + }, + ] + + setupMockSdkClient(streamChunks) + const adapter = new TestChatCompletionsAdapter(testConfig, 'test-model') + const modelOptions = { response_format: { type: 'json_object' } } + expect(adapter.supportsCombinedToolsAndSchema(modelOptions)).toBe(false) + + for await (const _ of adapter.chatStream({ + logger: testLogger, + model: 'test-model', + messages: [{ role: 'user', content: 'Hello' }], + tools: [weatherTool], + modelOptions, + outputSchema: { + type: 'object', + properties: { city: { type: 'string' } }, + required: ['city'], + }, + })) { + // drain + } + + const request = mockCreate.mock.calls[0]![0] as { + response_format?: unknown + messages: Array<{ role: string }> + } + expect(request.response_format).toEqual({ type: 'json_object' }) + expect(request.messages.some((m) => m.role === 'system')).toBe(false) + }) + it('omits response_format when outputSchema is not set', async () => { const streamChunks = [ { diff --git a/packages/openai-base/tests/responses-compact-wire.test.ts b/packages/openai-base/tests/responses-compact-wire.test.ts index 6d45c86ce6..ea5427f3e2 100644 --- a/packages/openai-base/tests/responses-compact-wire.test.ts +++ b/packages/openai-base/tests/responses-compact-wire.test.ts @@ -7,7 +7,9 @@ import type { ClientOptions } from 'openai' const logger = resolveDebugOption(false) const model = 'gpt-5.5' -class Responses extends OpenAIBaseResponsesTextAdapter {} +class Responses extends OpenAIBaseResponsesTextAdapter { + override readonly supportsFileSources = true +} const compacted = { id: 'resp_compact_1', @@ -191,4 +193,77 @@ describe('Responses adapter compact', () => { ], }) }) + + it('sends file, audio, and document parts back on the next call', async () => { + const pdfData = `data:${'application/pdf'};base64,${Buffer.from('%PDF-').toString('base64')}` + const mediaContent = [ + { type: 'input_text', text: 'See the file' }, + { type: 'input_image', file_id: 'file_img' }, + { type: 'input_file', file_id: 'file_pdf' }, + { type: 'input_file', file_url: 'https://example.com/note.pdf' }, + { type: 'input_file', filename: 'note.pdf', file_data: pdfData }, + { + type: 'input_file', + file_data: 'data:application/octet-stream;base64,QQ==', + }, + ] + const withMedia = { + ...compacted, + output: [ + { + type: 'message', + id: 'msg_media', + role: 'user', + status: 'completed', + content: mediaContent, + }, + { type: 'compaction', id: 'cmp_1', encrypted_content: 'sealed' }, + ], + } + const { client, sent } = recordingClient( + () => + new Response(`data: ${JSON.stringify(done)}\n\n`, { + headers: { 'content-type': 'text/event-stream' }, + }), + ) + const adapter = new Responses(model, 'openai', client) + const history = await new Responses( + model, + 'openai', + recordingClient(() => Response.json(withMedia)).client, + ).compact({ model, messages: [{ role: 'user', content: 'See the file' }] }) + + for await (const _ of adapter.chatStream({ + logger, + model, + messages: history, + })) { + // drain + } + + expect(sent[0]?.body).toMatchObject({ + input: [ + { + type: 'message', + role: 'user', + content: [ + { type: 'input_text', text: 'See the file' }, + { type: 'input_image', file_id: 'file_img', detail: 'auto' }, + { type: 'input_file', file_id: 'file_pdf' }, + { type: 'input_file', file_url: 'https://example.com/note.pdf' }, + { + type: 'input_file', + filename: 'note.pdf', + file_data: pdfData, + }, + { + type: 'input_file', + file_data: 'data:application/octet-stream;base64,QQ==', + }, + ], + }, + { type: 'compaction', id: 'cmp_1', encrypted_content: 'sealed' }, + ], + }) + }) }) diff --git a/packages/openai-base/tests/responses-text.test.ts b/packages/openai-base/tests/responses-text.test.ts index 404d41239a..4ebcfe8005 100644 --- a/packages/openai-base/tests/responses-text.test.ts +++ b/packages/openai-base/tests/responses-text.test.ts @@ -1869,6 +1869,108 @@ describe('OpenAIBaseResponsesTextAdapter', () => { }) }) + describe('raw tool arguments', () => { + type EndSite = 'arguments.done' | 'output_item.done' | 'response.completed' + + // Each site emits TOOL_CALL_END from a different event. Without a name on + // output_item.added, args.done is deferred to a later backfill. + function rawArgsStream( + site: EndSite, + raw: string, + ): Array> { + const named = { + type: 'function_call', + id: 'fc_raw', + name: 'lookup_weather', + } + return [ + { + type: 'response.created', + response: { + id: 'resp-raw', + model: 'test-model', + status: 'in_progress', + }, + }, + { + type: 'response.output_item.added', + output_index: 0, + item: + site === 'arguments.done' + ? named + : { type: 'function_call', id: 'fc_raw' }, + }, + { + type: 'response.function_call_arguments.done', + item_id: 'fc_raw', + arguments: raw, + }, + ...(site === 'output_item.done' + ? [ + { + type: 'response.output_item.done', + output_index: 0, + item: { ...named, arguments: raw }, + }, + ] + : []), + { + type: 'response.completed', + response: { + id: 'resp-raw', + model: 'test-model', + status: 'completed', + output: [{ ...named, arguments: raw }], + usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 }, + }, + }, + ] + } + + async function toolEnds(site: EndSite, raw: string) { + setupMockResponsesClient(rawArgsStream(site, raw)) + const adapter = new TestResponsesAdapter(testConfig, 'test-model') + const chunks: Array = [] + for await (const chunk of adapter.chatStream({ + logger: testLogger, + model: 'test-model', + messages: [{ role: 'user', content: 'q' }], + tools: [weatherTool], + })) { + chunks.push(chunk) + } + return chunks.filter((chunk) => chunk.type === 'TOOL_CALL_END') + } + + it.each([ + 'arguments.done', + 'output_item.done', + 'response.completed', + ])( + 'sends the raw argument string as args on TOOL_CALL_END (%s)', + async (site) => { + const ends = await toolEnds(site, '{"n": 1.50}') + expect(ends).toHaveLength(1) + expect(ends[0]).toMatchObject({ + toolCallId: 'fc_raw', + args: '{"n": 1.50}', + input: { n: 1.5 }, + }) + }, + ) + + it.each([ + 'arguments.done', + 'output_item.done', + 'response.completed', + ])('sends malformed arguments as args with no input (%s)', async (site) => { + const ends = await toolEnds(site, '{"a":') + expect(ends).toHaveLength(1) + expect(ends[0]).toMatchObject({ toolCallId: 'fc_raw', args: '{"a":' }) + expect(ends[0]).not.toHaveProperty('input') + }) + }) + describe('content_part events', () => { it('emits TEXT_MESSAGE_START on content_part.added with output_text', async () => { const streamChunks = [ @@ -2471,6 +2573,50 @@ describe('OpenAIBaseResponsesTextAdapter', () => { ) }) + it('sends text.format json_object and the schema as instructions when modelOptions sets json_object', async () => { + setupMockResponsesClient([], { + output: [ + { + type: 'message', + content: [ + { + type: 'output_text', + text: '{"name":"Alice","age":30}', + }, + ], + }, + ], + }) + + const adapter = new TestResponsesAdapter(testConfig, 'test-model') + + await adapter.structuredOutput({ + chatOptions: { + logger: testLogger, + model: 'test-model', + messages: [{ role: 'user', content: 'Give me a person' }], + modelOptions: { text: { format: { type: 'json_object' } } }, + }, + outputSchema: { + type: 'object', + properties: { + name: { type: 'string' }, + age: { type: 'number' }, + }, + required: ['name', 'age'], + }, + }) + + expect(mockResponsesCreate).toHaveBeenCalledWith( + expect.objectContaining({ + stream: false, + text: { format: { type: 'json_object' } }, + instructions: expect.stringContaining('"required":["name","age"]'), + }), + expect.anything(), + ) + }) + it('forwards response.usage on structuredOutput', async () => { setupMockResponsesClient([], { output: [ diff --git a/packages/openai-base/tests/responses-usage.test.ts b/packages/openai-base/tests/responses-usage.test.ts new file mode 100644 index 0000000000..ad41790bc9 --- /dev/null +++ b/packages/openai-base/tests/responses-usage.test.ts @@ -0,0 +1,22 @@ +import { describe, expect, it } from 'vitest' +import { buildResponsesUsage } from '../src/usage' + +// OpenAI Responses usage for a call that writes the prompt cache (issue #1734). +describe('buildResponsesUsage', () => { + it('maps cache reads and cache writes', () => { + expect( + buildResponsesUsage({ + input_tokens: 1941, + input_tokens_details: { cached_tokens: 3, cache_write_tokens: 1938 }, + output_tokens: 5, + output_tokens_details: { reasoning_tokens: 0 }, + total_tokens: 1946, + }), + ).toEqual({ + promptTokens: 1941, + completionTokens: 5, + totalTokens: 1946, + promptTokensDetails: { cachedTokens: 3, cacheWriteTokens: 1938 }, + }) + }) +}) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 3bf31b5d36..2ca0f8f472 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1072,10 +1072,10 @@ importers: version: 2.68.0 '@mcp-ui/client': specifier: ^7.1.1 - version: 7.1.1(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + version: 7.1.1(@cfworker/json-schema@4.1.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) '@modelcontextprotocol/sdk': specifier: ^1.29.0 - version: 1.29.0(supports-color@7.2.0)(zod@4.2.1) + version: 1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1) '@ngrok/ngrok': specifier: ^1.7.0 version: 1.7.0 @@ -1191,8 +1191,8 @@ importers: specifier: ^1.168.60 version: 1.168.60(crossws@0.4.6(srvx@0.11.22))(esbuild@0.28.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(rolldown@1.2.9)(rollup@4.60.1)(vite-plugin-solid@2.11.14(solid-js@1.9.10)(vite@8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0)))(vite@8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0)) '@tanstack/react-store': - specifier: ^0.8.0 - version: 0.8.0(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + specifier: ^0.11.1 + version: 0.11.2(react-dom@19.2.3(react@19.2.3))(react@19.2.3) '@tanstack/router-plugin': specifier: ^1.168.42 version: 1.168.42(@tanstack/react-router@1.170.41(react-dom@19.2.3(react@19.2.3))(react@19.2.3))(esbuild@0.28.1)(rolldown@1.2.9)(rollup@4.60.1)(vite-plugin-solid@2.11.14(solid-js@1.9.10)(vite@8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0)))(vite@8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0)) @@ -1792,8 +1792,8 @@ importers: specifier: ^1.168.57 version: 1.168.57(@tanstack/react-router@1.170.41(react-dom@19.2.3(react@19.2.3))(react@19.2.3))(crossws@0.4.6(srvx@0.11.22))(esbuild@0.28.1)(rolldown@1.2.9)(rollup@4.60.1)(solid-js@1.9.10)(vite-plugin-solid@2.11.14(solid-js@1.9.10)(vite@8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0)))(vite@8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0)) '@tanstack/solid-store': - specifier: ^0.8.0 - version: 0.8.0(solid-js@1.9.10) + specifier: ^0.11.1 + version: 0.11.2(solid-js@1.9.10) '@tanstack/store': specifier: ^0.11.1 version: 0.11.2 @@ -2011,6 +2011,9 @@ importers: '@ag-ui/core': specifier: 1.0.0 version: 1.0.0(zod@4.2.1) + '@cfworker/json-schema': + specifier: ^4.1.1 + version: 4.1.1 '@standard-schema/spec': specifier: ^1.1.0 version: 1.1.0 @@ -2227,6 +2230,9 @@ importers: '@tanstack/ai-utils': specifier: workspace:^ version: link:../ai-utils + '@tanstack/store': + specifier: ^0.11.1 + version: 0.11.2 devDependencies: '@standard-schema/spec': specifier: ^1.1.0 @@ -2509,7 +2515,7 @@ importers: dependencies: '@google/genai': specifier: ^2.10.0 - version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.3.6)) + version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6)) '@tanstack/ai-utils': specifier: workspace:^ version: link:../ai-utils @@ -3013,7 +3019,7 @@ importers: devDependencies: '@mcp-ui/client': specifier: ^7.1.1 - version: 7.1.1(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + version: 7.1.1(@cfworker/json-schema@4.1.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) '@tanstack/ai': specifier: workspace:* version: link:../ai @@ -3044,7 +3050,7 @@ importers: devDependencies: '@mcp-ui/client': specifier: ^7.1.1 - version: 7.1.1(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + version: 7.1.1(@cfworker/json-schema@4.1.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) '@standard-schema/spec': specifier: ^1.1.0 version: 1.1.0 @@ -3137,7 +3143,7 @@ importers: dependencies: '@modelcontextprotocol/sdk': specifier: ^1.29.0 - version: 1.29.0(zod@4.3.6) + version: 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6) '@tanstack/ai-skills': specifier: workspace:^ version: link:../ai-skills @@ -3162,7 +3168,7 @@ importers: dependencies: '@blaxel/core': specifier: ^0.3.10 - version: 0.3.11(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2)) + version: 0.3.11(@cfworker/json-schema@4.1.1)(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2)) devDependencies: '@tanstack/ai-sandbox': specifier: workspace:* @@ -3490,7 +3496,7 @@ importers: devDependencies: '@google/genai': specifier: ^2.10.0 - version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.3.6)) + version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6)) '@tanstack/ai': specifier: workspace:* version: link:../ai @@ -3681,7 +3687,7 @@ importers: version: 1.34.0(vitest@4.1.11(@opentelemetry/api@1.9.1)(@types/node@24.10.3)(@vitest/coverage-v8@4.1.10(vitest@4.1.10))(happy-dom@20.11.2)(jsdom@27.4.0(@noble/hashes@2.3.0))(vite@8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0))) '@modelcontextprotocol/sdk': specifier: ^1.29.0 - version: 1.29.0(zod@4.3.6) + version: 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6) '@openrouter/sdk': specifier: 0.13.20 version: 0.13.20 @@ -5152,6 +5158,9 @@ packages: '@bufbuild/protobuf@2.15.0': resolution: {integrity: sha512-DAheWUkVr/SJTWCc+lg9dhY0eN4SaWlf4+bG1KzHeXbnqt0AfB/NX0Z+VunGlM1ki1B4zVvye27MpKh/svySUA==} + '@cfworker/json-schema@4.1.1': + resolution: {integrity: sha512-gAmrUZSGtKc3AiBL71iNWxDsyUC5uMaKKGdvzYsBoTW/xi42JQHl7eKV2OYzCUqvc+D2RCcf7EXY2iCyFIk6og==} + '@changesets/apply-release-plan@7.1.0': resolution: {integrity: sha512-yq8ML3YS7koKQ/9bk1PqO0HMzApIFNwjlwCnwFEXMzNe8NpzeeYYKCmnhWJGkN8g7E51MnWaSbqRcTcdIxUgnQ==} @@ -5915,12 +5924,6 @@ packages: cpu: [x64] os: [win32] - '@eslint-community/eslint-utils@4.10.1': - resolution: {integrity: sha512-cuadcxVFE8sDK6iWJbs8Sn0av2Nrh2QSGQhVlBW9AaAHqHwjWsZHT8LJ4hFGPh7ASBV2deFdM7H/DPjulmh8rg==} - engines: {node: ^12.22.0 || ^14.17.0 || >=16.0.0} - peerDependencies: - eslint: ^6.0.0 || ^7.0.0 || >=8.0.0 - '@eslint-community/eslint-utils@4.9.1': resolution: {integrity: sha512-phrYmNiYppR7znFEdqgfWHXR6NCkZEK7hwWDHZUjit/2/U0r6XvkDl0SYnoM51Hq7FhCGdLDT6zxCCOY1hexsQ==} engines: {node: ^12.22.0 || ^14.17.0 || >=16.0.0} @@ -10868,12 +10871,6 @@ packages: react: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0 react-dom: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0 - '@tanstack/react-store@0.8.0': - resolution: {integrity: sha512-1vG9beLIuB7q69skxK9r5xiLN3ztzIPfSQSs0GfeqWGO2tGIyInZx0x1COhpx97RKaONSoAb8C3dxacWksm1ow==} - peerDependencies: - react: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0 - react-dom: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0 - '@tanstack/router-core@1.171.34': resolution: {integrity: sha512-X+19pGuuLVPvvAvPYw9v/+ktRk5cCll8AtfpBpOhDRmVgUMrfqze0467xtq8YRVVs/mmo5rVtWn/7Tp32RDBKw==} engines: {node: '>=20.19'} @@ -11004,8 +11001,8 @@ packages: vite: optional: true - '@tanstack/solid-store@0.8.0': - resolution: {integrity: sha512-JwqTedbxyOGw7mfmdGkB0RGgefRCw/tNauc8tlMcaS1mV5wTFT8c1KIB3LgttuHaanMJEBeqQJ7bc/R0WTP1fA==} + '@tanstack/solid-store@0.11.2': + resolution: {integrity: sha512-lscar3xeRKfs/RQBcZjR9NG55sfMP5Pb8S/URohUqcnSB8Yu19JLE/hPjfSd7pKmCpjYq1ZbjqtpIc5bDE1lAQ==} peerDependencies: solid-js: ^1.6.0 @@ -11040,9 +11037,6 @@ packages: '@tanstack/store@0.11.2': resolution: {integrity: sha512-sJ4mjol8uQsHV0gOJzzjwXfh2Fwm+Sz0+8deqiTm4jGbMdjzNSW+xZCFm0kUa870uhd8yi+DpKnZb5Kc4apa0Q==} - '@tanstack/store@0.8.0': - resolution: {integrity: sha512-Om+BO0YfMZe//X2z0uLF2j+75nQga6TpTJgLJQBiq85aOyZNIhkCgleNcud2KQg4k4v9Y9l+Uhru3qWMPGTOzQ==} - '@tanstack/typedoc-config@0.3.4': resolution: {integrity: sha512-WXJavpiNqumNwdBQTM3wBqbHaULgoZ3cIri7xJinKqX/mOGXgc6ZghfAardD2PydP8tOcPf+xhrd1Qk/am9DMg==} engines: {node: '>=18'} @@ -14353,10 +14347,6 @@ packages: resolution: {integrity: sha512-Hs59xBNfUIunMFgWAbGX5cq6893IbWg4KnrjbYwX3tx0ztorVgTDA6B2sxf8ejHJ4wz8BqGUMYlnzNBer5NvGg==} engines: {node: '>= 4'} - ignore@7.0.9: - resolution: {integrity: sha512-brTTsvFRt5C1gGHtPst/281UjPD5t9fBqbgoMPlVWy11ZLTPfu7HxK4ZYqO9H7o/yC9rSTCI85EaQ4OoY12qYw==} - engines: {node: '>= 4'} - image-size@0.5.5: resolution: {integrity: sha512-6TDAlDPZxUFCv+fuOkIoXT/V/f3Qbq8e37p+YOiYrUv3v9cc3/6x78VdfPgFVaB9dZYeLUfKgHRebpkm/oP2VQ==} engines: {node: '>=0.10.0'} @@ -21034,10 +21024,10 @@ snapshots: '@bcoe/v8-coverage@1.0.2': {} - '@blaxel/core@0.3.11(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2))': + '@blaxel/core@0.3.11(@cfworker/json-schema@4.1.1)(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2))': dependencies: '@hey-api/client-fetch': 0.10.2(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2)) - '@modelcontextprotocol/sdk': 1.29.0(zod@4.3.6) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6) archiver: 7.0.1 axios: 1.18.1(debug@4.4.3(supports-color@7.2.0))(supports-color@7.2.0) dockerfile-ast: 0.7.1 @@ -21068,6 +21058,8 @@ snapshots: '@bufbuild/protobuf@2.15.0': {} + '@cfworker/json-schema@4.1.1': {} + '@changesets/apply-release-plan@7.1.0': dependencies: '@changesets/config': 3.1.3 @@ -21719,11 +21711,6 @@ snapshots: '@esbuild/win32-x64@0.28.1': optional: true - '@eslint-community/eslint-utils@4.10.1(eslint@9.39.4(jiti@2.7.0))': - dependencies: - eslint: 9.39.4(jiti@2.7.0) - eslint-visitor-keys: 3.4.3 - '@eslint-community/eslint-utils@4.9.1(eslint@9.39.4(jiti@2.7.0))': dependencies: eslint: 9.39.4(jiti@2.7.0) @@ -22269,27 +22256,27 @@ snapshots: '@shikijs/types': 3.23.0 '@shikijs/vscode-textmate': 10.0.2 - '@google/genai@2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.3.6))': + '@google/genai@2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6))': dependencies: google-auth-library: 10.5.0 p-retry: 4.6.2 protobufjs: 7.6.4 ws: 8.21.0 optionalDependencies: - '@modelcontextprotocol/sdk': 1.29.0(zod@4.3.6) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6) transitivePeerDependencies: - bufferutil - supports-color - utf-8-validate - '@google/genai@2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.6.5))': + '@google/genai@2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5))': dependencies: google-auth-library: 10.5.0 p-retry: 4.6.2 protobufjs: 7.6.4 ws: 8.21.0 optionalDependencies: - '@modelcontextprotocol/sdk': 1.29.0(zod@4.6.5) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5) transitivePeerDependencies: - bufferutil - supports-color @@ -22683,10 +22670,10 @@ snapshots: - encoding - supports-color - '@mcp-ui/client@7.1.1(react-dom@19.2.3(react@19.2.3))(react@19.2.3)': + '@mcp-ui/client@7.1.1(@cfworker/json-schema@4.1.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)': dependencies: - '@modelcontextprotocol/ext-apps': 1.7.4(@modelcontextprotocol/sdk@1.29.0(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76) - '@modelcontextprotocol/sdk': 1.29.0(zod@3.25.76) + '@modelcontextprotocol/ext-apps': 1.7.4(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(zod@3.25.76) react: 19.2.3 react-dom: 19.2.3(react@19.2.3) zod: 3.25.76 @@ -22784,16 +22771,16 @@ snapshots: dependencies: zod: 4.3.6 - '@modelcontextprotocol/ext-apps@1.7.4(@modelcontextprotocol/sdk@1.29.0(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76)': + '@modelcontextprotocol/ext-apps@1.7.4(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76)': dependencies: - '@modelcontextprotocol/sdk': 1.29.0(supports-color@7.2.0)(zod@4.2.1) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1) '@standard-schema/spec': 1.1.0 zod: 3.25.76 optionalDependencies: react: 19.2.3 react-dom: 19.2.3(react@19.2.3) - '@modelcontextprotocol/sdk@1.29.0(supports-color@7.2.0)(zod@4.2.1)': + '@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1)': dependencies: '@hono/node-server': 1.19.17(hono@4.13.8) ajv: 8.20.0 @@ -22812,10 +22799,12 @@ snapshots: raw-body: 3.0.2 zod: 4.2.1 zod-to-json-schema: 3.25.2(zod@4.2.1) + optionalDependencies: + '@cfworker/json-schema': 4.1.1 transitivePeerDependencies: - supports-color - '@modelcontextprotocol/sdk@1.29.0(zod@3.25.76)': + '@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@3.25.76)': dependencies: '@hono/node-server': 1.19.17(hono@4.13.8) ajv: 8.20.0 @@ -22834,10 +22823,12 @@ snapshots: raw-body: 3.0.2 zod: 3.25.76 zod-to-json-schema: 3.25.2(zod@3.25.76) + optionalDependencies: + '@cfworker/json-schema': 4.1.1 transitivePeerDependencies: - supports-color - '@modelcontextprotocol/sdk@1.29.0(zod@4.3.6)': + '@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6)': dependencies: '@hono/node-server': 1.19.17(hono@4.13.8) ajv: 8.20.0 @@ -22856,10 +22847,12 @@ snapshots: raw-body: 3.0.2 zod: 4.3.6 zod-to-json-schema: 3.25.2(zod@4.3.6) + optionalDependencies: + '@cfworker/json-schema': 4.1.1 transitivePeerDependencies: - supports-color - '@modelcontextprotocol/sdk@1.29.0(zod@4.6.5)': + '@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5)': dependencies: '@hono/node-server': 1.19.17(hono@4.13.8) ajv: 8.20.0 @@ -22878,6 +22871,8 @@ snapshots: raw-body: 3.0.2 zod: 4.6.5 zod-to-json-schema: 3.25.2(zod@4.6.5) + optionalDependencies: + '@cfworker/json-schema': 4.1.1 transitivePeerDependencies: - supports-color optional: true @@ -27624,13 +27619,6 @@ snapshots: react-dom: 19.2.3(react@19.2.3) use-sync-external-store: 1.6.0(react@19.2.3) - '@tanstack/react-store@0.8.0(react-dom@19.2.3(react@19.2.3))(react@19.2.3)': - dependencies: - '@tanstack/store': 0.8.0 - react: 19.2.3 - react-dom: 19.2.3(react@19.2.3) - use-sync-external-store: 1.6.0(react@19.2.3) - '@tanstack/router-core@1.171.34': dependencies: '@tanstack/history': 1.162.4 @@ -27832,9 +27820,9 @@ snapshots: - vite-plugin-solid - webpack - '@tanstack/solid-store@0.8.0(solid-js@1.9.10)': + '@tanstack/solid-store@0.11.2(solid-js@1.9.10)': dependencies: - '@tanstack/store': 0.8.0 + '@tanstack/store': 0.11.2 solid-js: 1.9.10 '@tanstack/start-client-core@1.170.34': @@ -27993,8 +27981,6 @@ snapshots: '@tanstack/store@0.11.2': {} - '@tanstack/store@0.8.0': {} - '@tanstack/typedoc-config@0.3.4': dependencies: typedoc: 0.28.20(@typescript/typescript6@6.0.2) @@ -28402,7 +28388,7 @@ snapshots: '@typescript-eslint/utils': 8.59.4(eslint@9.39.4(jiti@2.7.0))(typescript@7.0.2) '@typescript-eslint/visitor-keys': 8.59.4 eslint: 9.39.4(jiti@2.7.0) - ignore: 7.0.9 + ignore: 7.0.5 natural-compare: 1.4.0 ts-api-utils: 2.5.0(typescript@7.0.2) typescript: 7.0.2 @@ -28523,7 +28509,7 @@ snapshots: '@typescript-eslint/utils@8.59.4(eslint@9.39.4(jiti@2.7.0))(typescript@7.0.2)': dependencies: - '@eslint-community/eslint-utils': 4.10.1(eslint@9.39.4(jiti@2.7.0)) + '@eslint-community/eslint-utils': 4.9.1(eslint@9.39.4(jiti@2.7.0)) '@typescript-eslint/scope-manager': 8.59.4 '@typescript-eslint/types': 8.59.4 '@typescript-eslint/typescript-estree': 8.59.4(typescript@7.0.2) @@ -30861,7 +30847,7 @@ snapshots: eslint@9.39.4(jiti@2.7.0): dependencies: - '@eslint-community/eslint-utils': 4.10.1(eslint@9.39.4(jiti@2.7.0)) + '@eslint-community/eslint-utils': 4.9.1(eslint@9.39.4(jiti@2.7.0)) '@eslint-community/regexpp': 4.12.2 '@eslint/config-array': 0.21.2 '@eslint/config-helpers': 0.4.2 @@ -31964,9 +31950,6 @@ snapshots: ignore@7.0.5: {} - ignore@7.0.9: - optional: true - image-size@0.5.5: optional: true diff --git a/scripts/convert-openrouter-models.ts b/scripts/convert-openrouter-models.ts index 4d392f550b..5af4cc99e2 100644 --- a/scripts/convert-openrouter-models.ts +++ b/scripts/convert-openrouter-models.ts @@ -68,6 +68,7 @@ const perModelInputModalities: Record = {} const imageModels = new Set([]) const chatModels = new Set([]) const combinedToolsAndSchemaModels = new Set([]) +const jsonObjectOnlyModels = new Set([]) const videoModels = new Set([]) const reasoningModels = new Set([]) const reasoningShapes: Array = [] @@ -89,6 +90,22 @@ export const OPENROUTER_COMBINED_TOOLS_AND_SCHEMA_MODELS: ReadonlySet = .join('\n')}\n ])` } +/** + * Catalog ids whose `supported_parameters` include `response_format` and not + * `structured_outputs`. Structured output sends JSON mode for these. + */ +function generateJsonObjectOnlyModelsSet(): string { + return `/** + * Generated: chat models whose catalog \`supported_parameters\` include + * \`response_format\` and do not include \`structured_outputs\`. + * Structured output sends JSON mode for these ids. + */ +export const OPENROUTER_JSON_OBJECT_ONLY_MODELS: ReadonlySet = + new Set([\n${Array.from(jsonObjectOnlyModels) + .map((id) => ` ${id},`) + .join('\n')}\n ])` +} + function generateChatModelsArray(): string { const modelIds = Array.from(chatModels) if (modelIds.length === 0) { @@ -331,6 +348,9 @@ function generateModelMetaString(model: OpenRouterModel): string { ) { combinedToolsAndSchemaModels.add(`${constName}.id`) } + if (params.has('response_format') && !params.has('structured_outputs')) { + jsonObjectOnlyModels.add(`${constName}.id`) + } } if (outputModalities.includes('image')) { imageModels.add(`${constName}.id`) @@ -499,6 +519,8 @@ ${generateVideoModelMeta()} ${generateImageModelsArray()} ${generateCombinedToolsAndSchemaModelsSet()} + +${generateJsonObjectOnlyModelsSet()} ` console.log(file) writeFile('packages/ai-openrouter/src/model-meta.ts', file).then(() => { diff --git a/testing/e2e/fixtures/harness-first-party/basic.json b/testing/e2e/fixtures/harness-first-party/basic.json index 1c1b3e99ce..198db10502 100644 --- a/testing/e2e/fixtures/harness-first-party/basic.json +++ b/testing/e2e/fixtures/harness-first-party/basic.json @@ -121,6 +121,14 @@ "content": "The report is done." } }, + { + "match": { + "userMessage": "[first-party-keys] hello" + }, + "response": { + "content": "Hello with your key." + } + }, { "match": { "userMessage": "[first-party-code-mode] which tools" diff --git a/testing/e2e/fixtures/middleware-test/raw-tool-args.json b/testing/e2e/fixtures/middleware-test/raw-tool-args.json new file mode 100644 index 0000000000..97e159b49d --- /dev/null +++ b/testing/e2e/fixtures/middleware-test/raw-tool-args.json @@ -0,0 +1,27 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[raw-tool-args] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "get_weather", + "arguments": "{\"city\": \"NYC\", \"days\": 1.50}" + } + ] + } + }, + { + "match": { + "userMessage": "[raw-tool-args] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "The weather is sunny." + } + } + ] +} diff --git a/testing/e2e/global-setup.ts b/testing/e2e/global-setup.ts index f3a2cba472..4f06314e99 100644 --- a/testing/e2e/global-setup.ts +++ b/testing/e2e/global-setup.ts @@ -157,6 +157,9 @@ export default async function globalSetup() { // Provider-executed web search responses need native Responses/Gemini wire // items that aimock does not synthesize. These mounts feed the adapter E2E // route deterministic search calls and grounding metadata. + // OpenAI Responses usage for a call that writes the prompt cache (#1734). + // aimock's native Responses helper cannot emit `input_tokens_details`. + mock.mount('/openai-responses-usage', openaiResponsesUsageMount()) mock.mount('/provider-search-openai', openaiProviderSearchMount()) mock.mount('/provider-search-gemini', geminiProviderSearchMount()) @@ -177,6 +180,12 @@ export default async function globalSetup() { // `usage` survives onto `RUN_FINISHED.usage` on the fallback path. mock.mount('/anthropic-structured-usage', anthropicStructuredUsageMount()) + // decide() usage on a failed answer mapping. aimock has no OpenAI Decisions + // endpoint, so this mount hand-crafts a `/v1/decisions` response with valid + // usage and a choice with no entry in `probabilities`. The companion spec + // asserts `decide()` rejects and still reports that usage. + mock.mount('/evaluate-usage-on-failure', evaluateUsageOnFailureMount()) + // BytePlus. Ark (chat, Seedream image, Seedance video) serves its whole data // plane under /api/v3, and Seed Speech (TTS/ASR) mounts its own endpoints at // the same prefix on a different host — which collapses onto one prefix here @@ -1958,6 +1967,68 @@ function openaiUsageDetailsMount( } } +function openaiResponsesUsageMount(): Mountable { + return { + async handleRequest( + req: http.IncomingMessage, + res: http.ServerResponse, + pathname: string, + ): Promise { + if (req.method !== 'POST' || pathname !== '/v1/responses') { + return false + } + await drainBody(req) + + const message = { + id: 'msg_e2e_usage', + type: 'message', + role: 'assistant', + status: 'completed', + content: [{ type: 'output_text', text: 'Hi', annotations: [] }], + } + const events = [ + { + type: 'response.output_text.delta', + item_id: 'msg_e2e_usage', + output_index: 0, + content_index: 0, + delta: 'Hi', + }, + { + type: 'response.completed', + response: { + id: 'resp_e2e_usage', + object: 'response', + status: 'completed', + model: 'gpt-4o', + output: [message], + usage: { + input_tokens: 1941, + input_tokens_details: { + cached_tokens: 0, + cache_write_tokens: 1938, + }, + output_tokens: 5, + output_tokens_details: { reasoning_tokens: 0 }, + total_tokens: 1946, + }, + }, + }, + ] + + res.statusCode = 200 + res.setHeader('Content-Type', 'text/event-stream') + res.setHeader('Cache-Control', 'no-cache') + for (const event of events) { + res.write(`data: ${JSON.stringify(event)}\n\n`) + } + res.write('data: [DONE]\n\n') + res.end() + return true + }, + } +} + function openaiProviderSearchMount(): Mountable { return { async handleRequest( @@ -2187,6 +2258,52 @@ function anthropicStructuredUsageMount(): Mountable { } } +/** + * Mounts the OpenAI Decisions `/v1/decisions` response for the + * `evaluate-usage-on-failure` spec: valid usage, but a `queue` answer whose + * choice (`sales`) has no entry in `probabilities`, so `decide()` throws + * "missing probability for choice". + */ +function evaluateUsageOnFailureMount(): Mountable { + return { + async handleRequest( + req: http.IncomingMessage, + res: http.ServerResponse, + pathname: string, + ): Promise { + if (req.method !== 'POST' || pathname !== '/v1/decisions') { + return false + } + await drainBody(req) + res.statusCode = 200 + res.setHeader('Content-Type', 'application/json') + res.end( + JSON.stringify({ + model: 'gpt-6-luna', + answers: [ + { + type: 'choice', + name: 'queue', + choice: 'sales', + confidence: 0.9, + probabilities: [ + { value: 'billing', probability: 0.6 }, + { value: 'tech', probability: 0.4 }, + ], + }, + ], + usage: { + input_tokens: 42, + output_tokens: 7, + total_tokens: 49, + }, + }), + ) + return true + }, + } +} + function buildToolPlusServerToolEvents(): Array> { const messageId = 'msg_bug_604' const model = 'claude-sonnet-4-5' diff --git a/testing/e2e/src/components/ChatUI.tsx b/testing/e2e/src/components/ChatUI.tsx index 2971dc7177..5e021de485 100644 --- a/testing/e2e/src/components/ChatUI.tsx +++ b/testing/e2e/src/components/ChatUI.tsx @@ -15,7 +15,7 @@ import { ApprovalPrompt } from '@/components/ApprovalPrompt' interface ChatUIProps< TTools extends ReadonlyArray = ReadonlyArray, > { - messages: Array + messages: ReadonlyArray isLoading: boolean onSendMessage: (text: string) => void /** Sends the typed prompt plus an attached image as an image content part. */ @@ -48,7 +48,7 @@ interface ChatUIProps< * `useChat` and auto-sent FIFO once the run settles. Rendered in a * region separate from `messages` so e2e tests can assert queued state * distinctly from the delivered conversation. */ - queue?: Array + queue?: ReadonlyArray /** Remove a queued message before it drains. */ cancelQueued?: (id: string) => void /** Block new input while pending interrupts await resolution. */ diff --git a/testing/e2e/src/components/ImageGenUI.tsx b/testing/e2e/src/components/ImageGenUI.tsx index 93cec8b7ce..bd59c7d159 100644 --- a/testing/e2e/src/components/ImageGenUI.tsx +++ b/testing/e2e/src/components/ImageGenUI.tsx @@ -1,4 +1,4 @@ -import { useState } from 'react' +import { useRef, useState } from 'react' import { useGenerateImage, fetchServerSentEvents, @@ -67,8 +67,12 @@ export function ImageGenUI({ } } - const { generate, result, isLoading, error, status } = + const { generate, result, isLoading, error, status, runId } = useGenerateImage(connectionOptions()) + const activeTuple = useRef({ runId: '', status: 'idle', isLoading: false }) + if (runId && status === 'generating' && isLoading) { + activeTuple.current = { runId, status, isLoading } + } const handleGenerate = async () => { if (!imageFile) { @@ -89,6 +93,12 @@ export function ImageGenUI({ return (
+ {activeTuple.current.runId} + {activeTuple.current.status} + + {String(activeTuple.current.isLoading)} + + {runId ?? ''}
{ if (!imageFile) { @@ -92,6 +96,12 @@ export function VideoGenUI({ return (
+ {activeTuple.current.runId} + {activeTuple.current.status} + + {String(activeTuple.current.isLoading)} + + {runId ?? ''}
rootRouteImport, } as any) +const CallOwnershipRoute = CallOwnershipRouteImport.update({ + id: '/call-ownership', + path: '/call-ownership', + getParentRoute: () => rootRouteImport, +} as any) const ChatClientDefaultBridgeRoute = ChatClientDefaultBridgeRouteImport.update({ id: '/chat-client-default-bridge', path: '/chat-client-default-bridge', @@ -509,6 +520,11 @@ const ApiByteplusSeedance1080pWireRoute = path: '/api/byteplus-seedance-1080p-wire', getParentRoute: () => rootRouteImport, } as any) +const ApiCallOwnershipWireRoute = ApiCallOwnershipWireRouteImport.update({ + id: '/api/call-ownership-wire', + path: '/api/call-ownership-wire', + getParentRoute: () => rootRouteImport, +} as any) const ApiChatRoute = ApiChatRouteImport.update({ id: '/api/chat', path: '/api/chat', @@ -520,6 +536,12 @@ const ApiChatCompletionsIncompleteStreamRoute = path: '/api/chat-completions-incomplete-stream', getParentRoute: () => rootRouteImport, } as any) +const ApiClientToolResumeStoreRoute = + ApiClientToolResumeStoreRouteImport.update({ + id: '/api/client-tool-resume-store', + path: '/api/client-tool-resume-store', + getParentRoute: () => rootRouteImport, + } as any) const ApiCloudflareBindingWireRoute = ApiCloudflareBindingWireRouteImport.update({ id: '/api/cloudflare-binding-wire', @@ -571,6 +593,12 @@ const ApiEmbeddingRoute = ApiEmbeddingRouteImport.update({ path: '/api/embedding', getParentRoute: () => rootRouteImport, } as any) +const ApiEvaluateUsageOnFailureRoute = + ApiEvaluateUsageOnFailureRouteImport.update({ + id: '/api/evaluate-usage-on-failure', + path: '/api/evaluate-usage-on-failure', + getParentRoute: () => rootRouteImport, + } as any) const ApiFakeTextRoute = ApiFakeTextRouteImport.update({ id: '/api/fake-text', path: '/api/fake-text', @@ -654,6 +682,11 @@ const ApiJoinRunClientToolRoute = ApiJoinRunClientToolRouteImport.update({ path: '/api/join-run-client-tool', getParentRoute: () => rootRouteImport, } as any) +const ApiJsonSchemaToolInputRoute = ApiJsonSchemaToolInputRouteImport.update({ + id: '/api/json-schema-tool-input', + path: '/api/json-schema-tool-input', + getParentRoute: () => rootRouteImport, +} as any) const ApiJsonTransportRoute = ApiJsonTransportRouteImport.update({ id: '/api/json-transport', path: '/api/json-transport', @@ -824,6 +857,11 @@ const ApiOpenaiMalformedToolArgumentsRoute = path: '/api/openai-malformed-tool-arguments', getParentRoute: () => rootRouteImport, } as any) +const ApiOpenaiResponsesUsageRoute = ApiOpenaiResponsesUsageRouteImport.update({ + id: '/api/openai-responses-usage', + path: '/api/openai-responses-usage', + getParentRoute: () => rootRouteImport, +} as any) const ApiOpenaiShellSkillsWireRoute = ApiOpenaiShellSkillsWireRouteImport.update({ id: '/api/openai-shell-skills-wire', @@ -1100,6 +1138,7 @@ export interface FileRoutesByFullPath { '/': typeof IndexRoute '/activity-test': typeof ActivityTestRoute '/byok': typeof ByokRoute + '/call-ownership': typeof CallOwnershipRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute '/chat-client-stream-processing': typeof ChatClientStreamProcessingRoute '/chat-client-watcher-status': typeof ChatClientWatcherStatusRoute @@ -1158,8 +1197,10 @@ export interface FileRoutesByFullPath { '/api/byok-chat': typeof ApiByokChatRoute '/api/byteplus-encrypted-reasoning-wire': typeof ApiByteplusEncryptedReasoningWireRoute '/api/byteplus-seedance-1080p-wire': typeof ApiByteplusSeedance1080pWireRoute + '/api/call-ownership-wire': typeof ApiCallOwnershipWireRoute '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute + '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute '/api/compaction-usage-wire': typeof ApiCompactionUsageWireRoute '/api/compaction-wire': typeof ApiCompactionWireRoute @@ -1170,6 +1211,7 @@ export interface FileRoutesByFullPath { '/api/durable-delivery': typeof ApiDurableDeliveryRoute '/api/durable-takeover': typeof ApiDurableTakeoverRoute '/api/embedding': typeof ApiEmbeddingRoute + '/api/evaluate-usage-on-failure': typeof ApiEvaluateUsageOnFailureRoute '/api/fake-text': typeof ApiFakeTextRoute '/api/file-source-wire': typeof ApiFileSourceWireRoute '/api/foreign-chunk-events': typeof ApiForeignChunkEventsRoute @@ -1186,6 +1228,7 @@ export interface FileRoutesByFullPath { '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/join-run-client-tool': typeof ApiJoinRunClientToolRoute + '/api/json-schema-tool-input': typeof ApiJsonSchemaToolInputRoute '/api/json-transport': typeof ApiJsonTransportRoute '/api/keyed-adapter': typeof ApiKeyedAdapterRoute '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -1219,6 +1262,7 @@ export interface FileRoutesByFullPath { '/api/openai-completed-response-text': typeof ApiOpenaiCompletedResponseTextRoute '/api/openai-image-2-5-models': typeof ApiOpenaiImage25ModelsRoute '/api/openai-malformed-tool-arguments': typeof ApiOpenaiMalformedToolArgumentsRoute + '/api/openai-responses-usage': typeof ApiOpenaiResponsesUsageRoute '/api/openai-shell-skills-wire': typeof ApiOpenaiShellSkillsWireRoute '/api/openai-strict-tool-null-wire': typeof ApiOpenaiStrictToolNullWireRoute '/api/openai-usage-details': typeof ApiOpenaiUsageDetailsRoute @@ -1276,6 +1320,7 @@ export interface FileRoutesByTo { '/': typeof IndexRoute '/activity-test': typeof ActivityTestRoute '/byok': typeof ByokRoute + '/call-ownership': typeof CallOwnershipRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute '/chat-client-stream-processing': typeof ChatClientStreamProcessingRoute '/chat-client-watcher-status': typeof ChatClientWatcherStatusRoute @@ -1334,8 +1379,10 @@ export interface FileRoutesByTo { '/api/byok-chat': typeof ApiByokChatRoute '/api/byteplus-encrypted-reasoning-wire': typeof ApiByteplusEncryptedReasoningWireRoute '/api/byteplus-seedance-1080p-wire': typeof ApiByteplusSeedance1080pWireRoute + '/api/call-ownership-wire': typeof ApiCallOwnershipWireRoute '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute + '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute '/api/compaction-usage-wire': typeof ApiCompactionUsageWireRoute '/api/compaction-wire': typeof ApiCompactionWireRoute @@ -1346,6 +1393,7 @@ export interface FileRoutesByTo { '/api/durable-delivery': typeof ApiDurableDeliveryRoute '/api/durable-takeover': typeof ApiDurableTakeoverRoute '/api/embedding': typeof ApiEmbeddingRoute + '/api/evaluate-usage-on-failure': typeof ApiEvaluateUsageOnFailureRoute '/api/fake-text': typeof ApiFakeTextRoute '/api/file-source-wire': typeof ApiFileSourceWireRoute '/api/foreign-chunk-events': typeof ApiForeignChunkEventsRoute @@ -1362,6 +1410,7 @@ export interface FileRoutesByTo { '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/join-run-client-tool': typeof ApiJoinRunClientToolRoute + '/api/json-schema-tool-input': typeof ApiJsonSchemaToolInputRoute '/api/json-transport': typeof ApiJsonTransportRoute '/api/keyed-adapter': typeof ApiKeyedAdapterRoute '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -1395,6 +1444,7 @@ export interface FileRoutesByTo { '/api/openai-completed-response-text': typeof ApiOpenaiCompletedResponseTextRoute '/api/openai-image-2-5-models': typeof ApiOpenaiImage25ModelsRoute '/api/openai-malformed-tool-arguments': typeof ApiOpenaiMalformedToolArgumentsRoute + '/api/openai-responses-usage': typeof ApiOpenaiResponsesUsageRoute '/api/openai-shell-skills-wire': typeof ApiOpenaiShellSkillsWireRoute '/api/openai-strict-tool-null-wire': typeof ApiOpenaiStrictToolNullWireRoute '/api/openai-usage-details': typeof ApiOpenaiUsageDetailsRoute @@ -1453,6 +1503,7 @@ export interface FileRoutesById { '/': typeof IndexRoute '/activity-test': typeof ActivityTestRoute '/byok': typeof ByokRoute + '/call-ownership': typeof CallOwnershipRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute '/chat-client-stream-processing': typeof ChatClientStreamProcessingRoute '/chat-client-watcher-status': typeof ChatClientWatcherStatusRoute @@ -1511,8 +1562,10 @@ export interface FileRoutesById { '/api/byok-chat': typeof ApiByokChatRoute '/api/byteplus-encrypted-reasoning-wire': typeof ApiByteplusEncryptedReasoningWireRoute '/api/byteplus-seedance-1080p-wire': typeof ApiByteplusSeedance1080pWireRoute + '/api/call-ownership-wire': typeof ApiCallOwnershipWireRoute '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute + '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute '/api/compaction-usage-wire': typeof ApiCompactionUsageWireRoute '/api/compaction-wire': typeof ApiCompactionWireRoute @@ -1523,6 +1576,7 @@ export interface FileRoutesById { '/api/durable-delivery': typeof ApiDurableDeliveryRoute '/api/durable-takeover': typeof ApiDurableTakeoverRoute '/api/embedding': typeof ApiEmbeddingRoute + '/api/evaluate-usage-on-failure': typeof ApiEvaluateUsageOnFailureRoute '/api/fake-text': typeof ApiFakeTextRoute '/api/file-source-wire': typeof ApiFileSourceWireRoute '/api/foreign-chunk-events': typeof ApiForeignChunkEventsRoute @@ -1539,6 +1593,7 @@ export interface FileRoutesById { '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/join-run-client-tool': typeof ApiJoinRunClientToolRoute + '/api/json-schema-tool-input': typeof ApiJsonSchemaToolInputRoute '/api/json-transport': typeof ApiJsonTransportRoute '/api/keyed-adapter': typeof ApiKeyedAdapterRoute '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -1572,6 +1627,7 @@ export interface FileRoutesById { '/api/openai-completed-response-text': typeof ApiOpenaiCompletedResponseTextRoute '/api/openai-image-2-5-models': typeof ApiOpenaiImage25ModelsRoute '/api/openai-malformed-tool-arguments': typeof ApiOpenaiMalformedToolArgumentsRoute + '/api/openai-responses-usage': typeof ApiOpenaiResponsesUsageRoute '/api/openai-shell-skills-wire': typeof ApiOpenaiShellSkillsWireRoute '/api/openai-strict-tool-null-wire': typeof ApiOpenaiStrictToolNullWireRoute '/api/openai-usage-details': typeof ApiOpenaiUsageDetailsRoute @@ -1631,6 +1687,7 @@ export interface FileRouteTypes { | '/' | '/activity-test' | '/byok' + | '/call-ownership' | '/chat-client-default-bridge' | '/chat-client-stream-processing' | '/chat-client-watcher-status' @@ -1689,8 +1746,10 @@ export interface FileRouteTypes { | '/api/byok-chat' | '/api/byteplus-encrypted-reasoning-wire' | '/api/byteplus-seedance-1080p-wire' + | '/api/call-ownership-wire' | '/api/chat' | '/api/chat-completions-incomplete-stream' + | '/api/client-tool-resume-store' | '/api/cloudflare-binding-wire' | '/api/compaction-usage-wire' | '/api/compaction-wire' @@ -1701,6 +1760,7 @@ export interface FileRouteTypes { | '/api/durable-delivery' | '/api/durable-takeover' | '/api/embedding' + | '/api/evaluate-usage-on-failure' | '/api/fake-text' | '/api/file-source-wire' | '/api/foreign-chunk-events' @@ -1717,6 +1777,7 @@ export interface FileRouteTypes { | '/api/image' | '/api/interrupts-test' | '/api/join-run-client-tool' + | '/api/json-schema-tool-input' | '/api/json-transport' | '/api/keyed-adapter' | '/api/lazy-tools-wire' @@ -1750,6 +1811,7 @@ export interface FileRouteTypes { | '/api/openai-completed-response-text' | '/api/openai-image-2-5-models' | '/api/openai-malformed-tool-arguments' + | '/api/openai-responses-usage' | '/api/openai-shell-skills-wire' | '/api/openai-strict-tool-null-wire' | '/api/openai-usage-details' @@ -1807,6 +1869,7 @@ export interface FileRouteTypes { | '/' | '/activity-test' | '/byok' + | '/call-ownership' | '/chat-client-default-bridge' | '/chat-client-stream-processing' | '/chat-client-watcher-status' @@ -1865,8 +1928,10 @@ export interface FileRouteTypes { | '/api/byok-chat' | '/api/byteplus-encrypted-reasoning-wire' | '/api/byteplus-seedance-1080p-wire' + | '/api/call-ownership-wire' | '/api/chat' | '/api/chat-completions-incomplete-stream' + | '/api/client-tool-resume-store' | '/api/cloudflare-binding-wire' | '/api/compaction-usage-wire' | '/api/compaction-wire' @@ -1877,6 +1942,7 @@ export interface FileRouteTypes { | '/api/durable-delivery' | '/api/durable-takeover' | '/api/embedding' + | '/api/evaluate-usage-on-failure' | '/api/fake-text' | '/api/file-source-wire' | '/api/foreign-chunk-events' @@ -1893,6 +1959,7 @@ export interface FileRouteTypes { | '/api/image' | '/api/interrupts-test' | '/api/join-run-client-tool' + | '/api/json-schema-tool-input' | '/api/json-transport' | '/api/keyed-adapter' | '/api/lazy-tools-wire' @@ -1926,6 +1993,7 @@ export interface FileRouteTypes { | '/api/openai-completed-response-text' | '/api/openai-image-2-5-models' | '/api/openai-malformed-tool-arguments' + | '/api/openai-responses-usage' | '/api/openai-shell-skills-wire' | '/api/openai-strict-tool-null-wire' | '/api/openai-usage-details' @@ -1983,6 +2051,7 @@ export interface FileRouteTypes { | '/' | '/activity-test' | '/byok' + | '/call-ownership' | '/chat-client-default-bridge' | '/chat-client-stream-processing' | '/chat-client-watcher-status' @@ -2041,8 +2110,10 @@ export interface FileRouteTypes { | '/api/byok-chat' | '/api/byteplus-encrypted-reasoning-wire' | '/api/byteplus-seedance-1080p-wire' + | '/api/call-ownership-wire' | '/api/chat' | '/api/chat-completions-incomplete-stream' + | '/api/client-tool-resume-store' | '/api/cloudflare-binding-wire' | '/api/compaction-usage-wire' | '/api/compaction-wire' @@ -2053,6 +2124,7 @@ export interface FileRouteTypes { | '/api/durable-delivery' | '/api/durable-takeover' | '/api/embedding' + | '/api/evaluate-usage-on-failure' | '/api/fake-text' | '/api/file-source-wire' | '/api/foreign-chunk-events' @@ -2069,6 +2141,7 @@ export interface FileRouteTypes { | '/api/image' | '/api/interrupts-test' | '/api/join-run-client-tool' + | '/api/json-schema-tool-input' | '/api/json-transport' | '/api/keyed-adapter' | '/api/lazy-tools-wire' @@ -2102,6 +2175,7 @@ export interface FileRouteTypes { | '/api/openai-completed-response-text' | '/api/openai-image-2-5-models' | '/api/openai-malformed-tool-arguments' + | '/api/openai-responses-usage' | '/api/openai-shell-skills-wire' | '/api/openai-strict-tool-null-wire' | '/api/openai-usage-details' @@ -2160,6 +2234,7 @@ export interface RootRouteChildren { IndexRoute: typeof IndexRoute ActivityTestRoute: typeof ActivityTestRoute ByokRoute: typeof ByokRoute + CallOwnershipRoute: typeof CallOwnershipRoute ChatClientDefaultBridgeRoute: typeof ChatClientDefaultBridgeRoute ChatClientStreamProcessingRoute: typeof ChatClientStreamProcessingRoute ChatClientWatcherStatusRoute: typeof ChatClientWatcherStatusRoute @@ -2218,8 +2293,10 @@ export interface RootRouteChildren { ApiByokChatRoute: typeof ApiByokChatRoute ApiByteplusEncryptedReasoningWireRoute: typeof ApiByteplusEncryptedReasoningWireRoute ApiByteplusSeedance1080pWireRoute: typeof ApiByteplusSeedance1080pWireRoute + ApiCallOwnershipWireRoute: typeof ApiCallOwnershipWireRoute ApiChatRoute: typeof ApiChatRoute ApiChatCompletionsIncompleteStreamRoute: typeof ApiChatCompletionsIncompleteStreamRoute + ApiClientToolResumeStoreRoute: typeof ApiClientToolResumeStoreRoute ApiCloudflareBindingWireRoute: typeof ApiCloudflareBindingWireRoute ApiCompactionUsageWireRoute: typeof ApiCompactionUsageWireRoute ApiCompactionWireRoute: typeof ApiCompactionWireRoute @@ -2230,6 +2307,7 @@ export interface RootRouteChildren { ApiDurableDeliveryRoute: typeof ApiDurableDeliveryRoute ApiDurableTakeoverRoute: typeof ApiDurableTakeoverRoute ApiEmbeddingRoute: typeof ApiEmbeddingRoute + ApiEvaluateUsageOnFailureRoute: typeof ApiEvaluateUsageOnFailureRoute ApiFakeTextRoute: typeof ApiFakeTextRoute ApiFileSourceWireRoute: typeof ApiFileSourceWireRoute ApiForeignChunkEventsRoute: typeof ApiForeignChunkEventsRoute @@ -2246,6 +2324,7 @@ export interface RootRouteChildren { ApiImageRoute: typeof ApiImageRouteWithChildren ApiInterruptsTestRoute: typeof ApiInterruptsTestRoute ApiJoinRunClientToolRoute: typeof ApiJoinRunClientToolRoute + ApiJsonSchemaToolInputRoute: typeof ApiJsonSchemaToolInputRoute ApiJsonTransportRoute: typeof ApiJsonTransportRoute ApiKeyedAdapterRoute: typeof ApiKeyedAdapterRoute ApiLazyToolsWireRoute: typeof ApiLazyToolsWireRoute @@ -2279,6 +2358,7 @@ export interface RootRouteChildren { ApiOpenaiCompletedResponseTextRoute: typeof ApiOpenaiCompletedResponseTextRoute ApiOpenaiImage25ModelsRoute: typeof ApiOpenaiImage25ModelsRoute ApiOpenaiMalformedToolArgumentsRoute: typeof ApiOpenaiMalformedToolArgumentsRoute + ApiOpenaiResponsesUsageRoute: typeof ApiOpenaiResponsesUsageRoute ApiOpenaiShellSkillsWireRoute: typeof ApiOpenaiShellSkillsWireRoute ApiOpenaiStrictToolNullWireRoute: typeof ApiOpenaiStrictToolNullWireRoute ApiOpenaiUsageDetailsRoute: typeof ApiOpenaiUsageDetailsRoute @@ -2351,6 +2431,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ByokRouteImport parentRoute: typeof rootRouteImport } + '/call-ownership': { + id: '/call-ownership' + path: '/call-ownership' + fullPath: '/call-ownership' + preLoaderRoute: typeof CallOwnershipRouteImport + parentRoute: typeof rootRouteImport + } '/chat-client-default-bridge': { id: '/chat-client-default-bridge' path: '/chat-client-default-bridge' @@ -2764,6 +2851,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiByteplusSeedance1080pWireRouteImport parentRoute: typeof rootRouteImport } + '/api/call-ownership-wire': { + id: '/api/call-ownership-wire' + path: '/api/call-ownership-wire' + fullPath: '/api/call-ownership-wire' + preLoaderRoute: typeof ApiCallOwnershipWireRouteImport + parentRoute: typeof rootRouteImport + } '/api/chat': { id: '/api/chat' path: '/api/chat' @@ -2778,6 +2872,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiChatCompletionsIncompleteStreamRouteImport parentRoute: typeof rootRouteImport } + '/api/client-tool-resume-store': { + id: '/api/client-tool-resume-store' + path: '/api/client-tool-resume-store' + fullPath: '/api/client-tool-resume-store' + preLoaderRoute: typeof ApiClientToolResumeStoreRouteImport + parentRoute: typeof rootRouteImport + } '/api/cloudflare-binding-wire': { id: '/api/cloudflare-binding-wire' path: '/api/cloudflare-binding-wire' @@ -2848,6 +2949,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiEmbeddingRouteImport parentRoute: typeof rootRouteImport } + '/api/evaluate-usage-on-failure': { + id: '/api/evaluate-usage-on-failure' + path: '/api/evaluate-usage-on-failure' + fullPath: '/api/evaluate-usage-on-failure' + preLoaderRoute: typeof ApiEvaluateUsageOnFailureRouteImport + parentRoute: typeof rootRouteImport + } '/api/fake-text': { id: '/api/fake-text' path: '/api/fake-text' @@ -2960,6 +3068,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiJoinRunClientToolRouteImport parentRoute: typeof rootRouteImport } + '/api/json-schema-tool-input': { + id: '/api/json-schema-tool-input' + path: '/api/json-schema-tool-input' + fullPath: '/api/json-schema-tool-input' + preLoaderRoute: typeof ApiJsonSchemaToolInputRouteImport + parentRoute: typeof rootRouteImport + } '/api/json-transport': { id: '/api/json-transport' path: '/api/json-transport' @@ -3191,6 +3306,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiOpenaiMalformedToolArgumentsRouteImport parentRoute: typeof rootRouteImport } + '/api/openai-responses-usage': { + id: '/api/openai-responses-usage' + path: '/api/openai-responses-usage' + fullPath: '/api/openai-responses-usage' + preLoaderRoute: typeof ApiOpenaiResponsesUsageRouteImport + parentRoute: typeof rootRouteImport + } '/api/openai-shell-skills-wire': { id: '/api/openai-shell-skills-wire' path: '/api/openai-shell-skills-wire' @@ -3613,6 +3735,7 @@ const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, ActivityTestRoute: ActivityTestRoute, ByokRoute: ByokRoute, + CallOwnershipRoute: CallOwnershipRoute, ChatClientDefaultBridgeRoute: ChatClientDefaultBridgeRoute, ChatClientStreamProcessingRoute: ChatClientStreamProcessingRoute, ChatClientWatcherStatusRoute: ChatClientWatcherStatusRoute, @@ -3675,9 +3798,11 @@ const rootRouteChildren: RootRouteChildren = { ApiByteplusEncryptedReasoningWireRoute: ApiByteplusEncryptedReasoningWireRoute, ApiByteplusSeedance1080pWireRoute: ApiByteplusSeedance1080pWireRoute, + ApiCallOwnershipWireRoute: ApiCallOwnershipWireRoute, ApiChatRoute: ApiChatRoute, ApiChatCompletionsIncompleteStreamRoute: ApiChatCompletionsIncompleteStreamRoute, + ApiClientToolResumeStoreRoute: ApiClientToolResumeStoreRoute, ApiCloudflareBindingWireRoute: ApiCloudflareBindingWireRoute, ApiCompactionUsageWireRoute: ApiCompactionUsageWireRoute, ApiCompactionWireRoute: ApiCompactionWireRoute, @@ -3688,6 +3813,7 @@ const rootRouteChildren: RootRouteChildren = { ApiDurableDeliveryRoute: ApiDurableDeliveryRoute, ApiDurableTakeoverRoute: ApiDurableTakeoverRoute, ApiEmbeddingRoute: ApiEmbeddingRoute, + ApiEvaluateUsageOnFailureRoute: ApiEvaluateUsageOnFailureRoute, ApiFakeTextRoute: ApiFakeTextRoute, ApiFileSourceWireRoute: ApiFileSourceWireRoute, ApiForeignChunkEventsRoute: ApiForeignChunkEventsRoute, @@ -3704,6 +3830,7 @@ const rootRouteChildren: RootRouteChildren = { ApiImageRoute: ApiImageRouteWithChildren, ApiInterruptsTestRoute: ApiInterruptsTestRoute, ApiJoinRunClientToolRoute: ApiJoinRunClientToolRoute, + ApiJsonSchemaToolInputRoute: ApiJsonSchemaToolInputRoute, ApiJsonTransportRoute: ApiJsonTransportRoute, ApiKeyedAdapterRoute: ApiKeyedAdapterRoute, ApiLazyToolsWireRoute: ApiLazyToolsWireRoute, @@ -3737,6 +3864,7 @@ const rootRouteChildren: RootRouteChildren = { ApiOpenaiCompletedResponseTextRoute: ApiOpenaiCompletedResponseTextRoute, ApiOpenaiImage25ModelsRoute: ApiOpenaiImage25ModelsRoute, ApiOpenaiMalformedToolArgumentsRoute: ApiOpenaiMalformedToolArgumentsRoute, + ApiOpenaiResponsesUsageRoute: ApiOpenaiResponsesUsageRoute, ApiOpenaiShellSkillsWireRoute: ApiOpenaiShellSkillsWireRoute, ApiOpenaiStrictToolNullWireRoute: ApiOpenaiStrictToolNullWireRoute, ApiOpenaiUsageDetailsRoute: ApiOpenaiUsageDetailsRoute, diff --git a/testing/e2e/src/routes/$provider/$feature.tsx b/testing/e2e/src/routes/$provider/$feature.tsx index 0385423107..23e6d856ce 100644 --- a/testing/e2e/src/routes/$provider/$feature.tsx +++ b/testing/e2e/src/routes/$provider/$feature.tsx @@ -400,6 +400,7 @@ function ChatFeature({ interrupts, stop, clear, + setMessages, queue, cancelQueued, } = useChat({ @@ -475,6 +476,13 @@ function ChatFeature({ +
)} { + const now = Date.now() + return (async function* () { + yield { type: 'RUN_STARTED', threadId, runId, timestamp: now } + yield { type: 'STEP_STARTED', stepName: 'step-a', timestamp: now } + yield { + type: 'REASONING_MESSAGE_CONTENT', + messageId: 'rs_A', + delta: 'think A', + timestamp: now, + } + yield { + type: 'TEXT_MESSAGE_START', + messageId: 'msg-1', + role: 'assistant', + timestamp: now, + } + yield { + type: 'TEXT_MESSAGE_CONTENT', + messageId: 'msg-1', + delta: 'hello', + timestamp: now, + } + yield { + type: 'TOOL_CALL_START', + toolCallId: 'call_S', + toolCallName: 'search', + parentMessageId: 'msg-1', + timestamp: now, + } + yield { + type: 'TOOL_CALL_ARGS', + toolCallId: 'call_S', + delta: '{}', + timestamp: now, + } + yield { type: 'TOOL_CALL_END', toolCallId: 'call_S', timestamp: now } + yield { + type: 'RUN_FINISHED', + threadId, + runId, + finishReason: 'tool_calls', + timestamp: now, + } + yield { + type: 'TOOL_CALL_RESULT', + toolCallId: 'call_S', + content: '{"ok":true}', + messageId: 'result-call_S', + timestamp: now, + } + yield { + type: 'TOOL_CALL_START', + toolCallId: 'call_G', + toolCallName: 'get', + parentMessageId: 'msg-2', + timestamp: now, + } + yield { + type: 'TOOL_CALL_ARGS', + toolCallId: 'call_G', + delta: '{}', + timestamp: now, + } + yield { type: 'TOOL_CALL_END', toolCallId: 'call_G', timestamp: now } + yield { + type: 'RUN_FINISHED', + threadId, + runId: 'run-2', + finishReason: 'tool_calls', + timestamp: now, + } + yield { + type: 'TOOL_CALL_RESULT', + toolCallId: 'call_G', + content: '{"ok":true}', + messageId: 'result-call_G', + timestamp: now, + } + yield { type: 'STEP_STARTED', stepName: 'step-b', timestamp: now } + yield { + type: 'REASONING_MESSAGE_CONTENT', + messageId: 'rs_B', + delta: 'think B', + timestamp: now, + } + yield { + type: 'TEXT_MESSAGE_START', + messageId: 'msg-3', + role: 'assistant', + timestamp: now, + } + yield { + type: 'TEXT_MESSAGE_CONTENT', + messageId: 'msg-3', + delta: 'bye', + timestamp: now, + } + yield { + type: 'RUN_FINISHED', + threadId, + runId: 'run-3', + finishReason: 'stop', + timestamp: now, + } + })() as AsyncIterable +} + +export const Route = createFileRoute('/api/call-ownership-wire')({ + server: { + handlers: { + POST: async ({ request }) => { + const body: unknown = await request.json() + const threadId = + typeof body === 'object' && body !== null && 'threadId' in body + ? String((body as Record).threadId) + : 'thread-1' + return toServerSentEventsResponse( + callOwnershipRun(threadId, `run-${threadId}`), + ) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/api.client-tool-resume-store.ts b/testing/e2e/src/routes/api.client-tool-resume-store.ts new file mode 100644 index 0000000000..1c39c0f4bc --- /dev/null +++ b/testing/e2e/src/routes/api.client-tool-resume-store.ts @@ -0,0 +1,230 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + chat, + chatParamsFromRequestBody, + createChatOptions, +} from '@tanstack/ai' +import { StreamProcessor, uiMessagesToWire } from '@tanstack/ai/client' +import { createAnthropicChat } from '@tanstack/ai-anthropic' +import { memoryPersistence, withPersistence } from '@tanstack/ai-persistence' + +const DUMMY_KEY = 'sk-ant-e2e-test-dummy-key' + +type SseEvent = Record & { type: string } + +function messageStart(id: string): SseEvent { + return { + type: 'message_start', + message: { + id, + type: 'message', + role: 'assistant', + content: [], + model: 'claude-sonnet-4-5', + stop_reason: null, + stop_sequence: null, + usage: { input_tokens: 5, output_tokens: 0 }, + }, + } +} + +function text(index: number, value: string) { + return [ + { + type: 'content_block_start', + index, + content_block: { type: 'text', text: '' }, + }, + { + type: 'content_block_delta', + index, + delta: { type: 'text_delta', text: value }, + }, + { type: 'content_block_stop', index }, + ] +} + +function toolUse(index: number, id: string, name: string) { + return [ + { + type: 'content_block_start', + index, + content_block: { type: 'tool_use', id, name, input: {} }, + }, + { + type: 'content_block_delta', + index, + delta: { type: 'input_json_delta', partial_json: '{}' }, + }, + { type: 'content_block_stop', index }, + ] +} + +function end(stopReason: string) { + return [ + { + type: 'message_delta', + delta: { stop_reason: stopReason, stop_sequence: null }, + usage: { output_tokens: 20 }, + }, + { type: 'message_stop' }, + ] +} + +function sse(events: Array): Response { + return new Response( + events + .map( + (event) => `event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`, + ) + .join(''), + { status: 200, headers: { 'Content-Type': 'text/event-stream' } }, + ) +} + +/** Answers each call with the next scripted stream and records its body. */ +function scriptedFetch(streams: Array>) { + const bodies: Array<{ + messages: Array<{ role: string; content: unknown }> + }> = [] + const fetchImpl: typeof fetch = async (input, init) => { + const req = input instanceof Request ? input : new Request(input, init) + bodies.push(JSON.parse(await req.text())) + return sse(streams[bodies.length - 1] ?? []) + } + return { bodies, fetchImpl } +} + +/** Every `tool_result` id the request sends back to Claude. */ +function toolResultIds(body: { messages: Array<{ content: unknown }> }) { + return body.messages.flatMap((message) => + Array.isArray(message.content) + ? message.content + .filter((block: { type: string }) => block.type === 'tool_result') + .map((block: { tool_use_id: string }) => block.tool_use_id) + : [], + ) +} + +/** + * Claude calls a server tool, then a client tool, and the run stops on the + * client-tool interrupt. The browser takes the interrupt `MESSAGES_SNAPSHOT`, + * which gives the stored server tool result a snapshot id, adds the client + * tool's output, and resumes the run. `withPersistence` merges what it posts + * with the stored transcript. + * + * Before the fix, the store had never seen the snapshot id, so the merge + * appended the server tool result a second time and the resumed call sent + * Claude a second `tool_result` for `toolu_lookup`, with no matching `tool_use` + * in the assistant turn before it. + * + * `StreamProcessor` → `uiMessagesToWire` → `chatParamsFromRequestBody` → + * `chat()`, as a browser client and `/api/chat` do, answered by a capturing + * `fetch` (no aimock fixture). + */ +export const Route = createFileRoute('/api/client-tool-resume-store')({ + server: { + handlers: { + POST: async () => { + try { + const claude = scriptedFetch([ + [ + messageStart('msg_call_1'), + ...toolUse(0, 'toolu_lookup', 'lookup'), + ...end('tool_use'), + ], + [ + messageStart('msg_call_2'), + ...toolUse(0, 'toolu_ask', 'ask'), + ...end('tool_use'), + ], + [ + messageStart('msg_call_3'), + ...text(0, 'It is 42.'), + ...end('end_turn'), + ], + ]) + const adapter = createAnthropicChat('claude-sonnet-4-5', DUMMY_KEY, { + fetch: claude.fetchImpl, + }) + const persistence = memoryPersistence() + const tools = [ + { + name: 'lookup', + description: 'Look up the answer.', + execute: () => ({ answer: 42 }), + }, + { name: 'ask', description: 'Ask the user.' }, + ] + const threadId = 'client-tool-resume-store' + const runId = 'client-tool-resume-store-1' + const client = new StreamProcessor({}) + client.addUserMessage('What is the answer?') + + const send = async ( + resume?: Array<{ + interruptId: string + status: 'resolved' + payload: unknown + }>, + ) => { + const params = await chatParamsFromRequestBody({ + threadId, + runId, + // `JSON.parse(JSON.stringify(...))` stands in for the HTTP hop. + messages: JSON.parse( + JSON.stringify(uiMessagesToWire(client.getMessages())), + ), + tools: [], + context: [], + ...(resume && { resume }), + }) + client.prepareAssistantMessage() + for await (const chunk of chat({ + ...createChatOptions({ adapter }), + messages: params.messages, + ...(params.resume && { resume: params.resume }), + tools, + threadId, + runId, + middleware: [withPersistence(persistence)], + stream: true, + })) { + client.processChunk(chunk) + } + client.finalizeStream() + } + + await send() + client.addToolResult('toolu_ask', { confirmed: true }) + await send([ + { + interruptId: 'client_tool_toolu_ask', + status: 'resolved', + payload: { confirmed: true }, + }, + ]) + + const stored = + (await persistence.stores.messages?.loadThread(threadId)) ?? [] + return Response.json({ + ok: true, + resumedToolResultIds: claude.bodies[2] + ? toolResultIds(claude.bodies[2]) + : [], + storedToolResultIds: stored.flatMap((message) => + message.role === 'tool' && message.toolCallId + ? [message.toolCallId] + : [], + ), + }) + } catch (error) { + return Response.json({ + ok: false, + error: error instanceof Error ? error.message : String(error), + }) + } + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/api.evaluate-usage-on-failure.ts b/testing/e2e/src/routes/api.evaluate-usage-on-failure.ts new file mode 100644 index 0000000000..360b9b114c --- /dev/null +++ b/testing/e2e/src/routes/api.evaluate-usage-on-failure.ts @@ -0,0 +1,56 @@ +import { createFileRoute } from '@tanstack/react-router' +import { choice, decide } from '@tanstack/ai' +import { createOpenaiDecider } from '@tanstack/ai-openai' + +const LLMOCK_DEFAULT_BASE = process.env.LLMOCK_URL || 'http://127.0.0.1:4010' +const DUMMY_KEY = 'sk-e2e-test-dummy-key' + +/** + * Drives `decide()` through the OpenAI evaluate adapter against the + * `/evaluate-usage-on-failure` aimock mount. The mount answers with valid + * usage but picks a choice with no entry in `probabilities`, so `decide()` + * fails while mapping the answer. The route reports the rejection together + * with the usage seen by the `onUsage` middleware hook. + */ +export const Route = createFileRoute('/api/evaluate-usage-on-failure')({ + server: { + handlers: { + POST: async () => { + const adapter = createOpenaiDecider('gpt-6-luna', DUMMY_KEY, { + baseURL: `${LLMOCK_DEFAULT_BASE}/evaluate-usage-on-failure/v1`, + }) + + const onUsageReports: Array = [] + + let error: { name: string; message: string } | undefined + try { + await decide({ + adapter, + state: { ticket: 'I was charged twice this month.' }, + questions: { + queue: choice({ + instructions: 'Which team should handle this ticket?', + options: { billing: null, tech: null }, + }), + }, + middleware: [ + { + name: 'usage-observer', + onUsage: (_ctx, usage) => { + onUsageReports.push(usage) + }, + }, + ], + }) + } catch (caught) { + error = + caught instanceof Error + ? { name: caught.name, message: caught.message } + : { name: 'Unknown', message: String(caught) } + } + + return Response.json({ error, onUsageReports }) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/api.harness-first-party.ts b/testing/e2e/src/routes/api.harness-first-party.ts index be06b724de..ed86376f8f 100644 --- a/testing/e2e/src/routes/api.harness-first-party.ts +++ b/testing/e2e/src/routes/api.harness-first-party.ts @@ -5,7 +5,7 @@ import { join } from 'node:path' import { promisify } from 'node:util' import { createFileRoute } from '@tanstack/react-router' import { codeMode } from '@tanstack/ai-code-mode/harness' -import { defineAgent, toolDefinition } from '@tanstack/ai' +import { defineAgent, keyedAdapter, toolDefinition } from '@tanstack/ai' import { createHarnessHost, defineHarness, @@ -17,14 +17,17 @@ import { goal, permissions, projectInstructions, + providerKeys, selectGoal, title, } from '@tanstack/ai-harness/plugins' +import { createOpenaiChat } from '@tanstack/ai-openai' +import { openaiByok } from '@tanstack/ai-openai/byok' import { memoryPersistence } from '@tanstack/ai-persistence' import { z } from 'zod' import type { AnyTextAdapter } from '@tanstack/ai' import type { HarnessHost } from '@tanstack/ai-harness' -import { createTextAdapter } from '@/lib/providers' +import { LLMOCK_DEFAULT_BASE, createTextAdapter } from '@/lib/providers' const run = promisify(execFile) @@ -197,6 +200,55 @@ async function sessionTitle( return { text: turn.text, title: entry?.title, failures } } +/** The key that the user pastes in the `provider-keys` scenario. */ +const USER_KEY = 'sk-e2e-user-key-7f3a' + +/** + * `providerKeys()`: `/connect:openai` asks for the key as a secret question, + * saves the answer, and the next turn builds the model adapter with that key. + */ +async function connectedKey(host: HarnessHost, testId: string | undefined) { + // The keys the adapter is built with. aimock's journal redacts the + // authorization header, so the spec reads them here. + const keysUsed: Array = [] + const gpt = keyedAdapter(openaiByok, (key) => { + keysUsed.push(key) + return createOpenaiChat('gpt-6.1-sol', key, { + baseURL: `${LLMOCK_DEFAULT_BASE}/v1`, + ...(testId ? { defaultHeaders: { 'X-Test-Id': testId } } : {}), + }) + }) + const session = await host.open( + defineHarness({ + name: 'e2e/harness-provider-keys', + adapter: gpt, + plugins: () => [providerKeys({ providers: [openaiByok] })], + }), + { threadId: 'e2e-keys', principal: { id: 'ada' } }, + ) + const connecting = session.command('connect:openai') + const deadline = Date.now() + 5_000 + while ( + session.snapshot().pendingQuestions.length === 0 && + Date.now() < deadline + ) { + await new Promise((resolve) => setTimeout(resolve, 20)) + } + const [question] = session.snapshot().pendingQuestions + // Without a question, `connecting` never settles: fail now. + if (!question) throw new Error('/connect:openai asked no question.') + await session.answer({ questionId: question.questionId, value: USER_KEY }) + const connected = await connecting + const turn = await session.prompt('[first-party-keys] hello') + return { + secret: question.secret ?? false, + connected, + text: turn.text, + key: USER_KEY, + usedUserKey: keysUsed.includes(USER_KEY), + } +} + /** * `goal()`: `/goal` starts a turn, the judge says that the goal is met, and * no next turn starts. @@ -262,7 +314,7 @@ async function codeModeTools(host: HarnessHost, model: () => AnyTextAdapter) { /** * First-party harness plugins, one scenario per request: - * `?scenario=agents|plan-subagent|env|title|goal|code-mode&testId=...`. The models are + * `?scenario=agents|plan-subagent|env|title|goal|provider-keys|code-mode&testId=...`. The models are * OpenAI adapters against aimock. */ export const Route = createFileRoute('/api/harness-first-party')({ @@ -294,6 +346,9 @@ export const Route = createFileRoute('/api/harness-first-party')({ ).adapter return Response.json(await sessionTitle(host, model, titleModel)) } + if (scenario === 'provider-keys') { + return Response.json(await connectedKey(host, testId)) + } if (scenario === 'goal') { const judgeModel = createTextAdapter( 'openai', diff --git a/testing/e2e/src/routes/api.json-schema-tool-input.ts b/testing/e2e/src/routes/api.json-schema-tool-input.ts new file mode 100644 index 0000000000..5cba53fc80 --- /dev/null +++ b/testing/e2e/src/routes/api.json-schema-tool-input.ts @@ -0,0 +1,87 @@ +import { createFileRoute } from '@tanstack/react-router' +import { chat, toolDefinition } from '@tanstack/ai' +import { fakeText } from '@tanstack/ai/testing' + +// No network: `fakeText` answers from a script. The first answer calls a +// JSON Schema tool twice: once with a value that cannot be repaired, once +// with a number sent as a string. It also calls an OpenAPI 3.0 tool with +// `null` for a `nullable` field. The second answer repeats the tool results. +// MCP tools have the same raw JSON Schema input. With `{ repair: 'custom' }`, +// a `repairToolInput` function replaces the built-in repair: it maps +// `'many'` to 2 and repairs nothing else. +export const Route = createFileRoute('/api/json-schema-tool-input')({ + server: { + handlers: { + POST: async ({ request }) => { + const body: unknown = await request.json().catch(() => ({})) + const custom = + typeof body === 'object' && + body !== null && + 'repair' in body && + body.repair === 'custom' + + const fake = fakeText() + fake.setResponses([ + { + toolCalls: [ + { name: 'count_items', input: { count: 'many' } }, + { name: 'count_items', input: { count: '3' } }, + { name: 'label_item', input: { label: null } }, + ], + }, + ({ request }) => ({ + text: request.messages + .filter((message) => message.role === 'tool') + .map((message) => `Tool said: ${String(message.content)}`) + .join('\n'), + }), + ]) + + const runs: Array = [] + const countItems = toolDefinition({ + name: 'count_items', + description: 'Count the items', + inputSchema: { + $schema: 'http://json-schema.org/draft-07/schema#', + type: 'object', + properties: { count: { type: 'integer' } }, + required: ['count'], + }, + }).server((input) => { + runs.push(input) + return 'counted' + }) + const labelItem = toolDefinition({ + name: 'label_item', + description: 'Label the item', + inputSchema: { + type: 'object', + properties: { label: { type: 'string', nullable: true } }, + required: ['label'], + }, + }).server((input) => { + runs.push(input) + return 'labeled' + }) + + const { text } = await chat({ + adapter: fake, + messages: [{ role: 'user', content: 'Count the items' }], + tools: [countItems, labelItem], + stream: false, + ...(custom && { + repairToolInput: ({ input }) => + typeof input === 'object' && + input !== null && + 'count' in input && + input.count === 'many' + ? { count: 2 } + : undefined, + }), + }) + + return Response.json({ text, runs }) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/api.middleware-test.ts b/testing/e2e/src/routes/api.middleware-test.ts index 8c7213fe5a..f5d7cb672a 100644 --- a/testing/e2e/src/routes/api.middleware-test.ts +++ b/testing/e2e/src/routes/api.middleware-test.ts @@ -177,6 +177,16 @@ const toolReplaceMiddleware: ChatMiddleware = { }, } +// Rewrites the args to a value the tool's inputSchema rejects. The final +// input check runs after this hook, so the tool must not run. +const toolBadArgsMiddleware: ChatMiddleware = { + name: 'tool-bad-args', + onBeforeToolCall(_ctx, hookCtx) { + if (hookCtx.toolName !== 'get_weather') return undefined + return { type: 'transformArgs' as const, args: { city: { name: 'NYC' } } } + }, +} + /** * Capability provide/consume flow (`capability` mode). * @@ -579,6 +589,8 @@ export const Route = createFileRoute('/api/middleware-test')({ middleware.push(toolSkipMiddleware) if (middlewareMode === 'tool-replace') middleware.push(toolReplaceMiddleware) + if (middlewareMode === 'tool-bad-args') + middleware.push(toolBadArgsMiddleware) if (middlewareMode === 'capability') { // Order matters: the provider's setup() must run before the // consumer reads the capability. Array order is preserved. @@ -663,7 +675,7 @@ export const Route = createFileRoute('/api/middleware-test')({ const tools = genericScenario && testId ? genericTools(testId, genericScenario) - : scenario === 'with-tool' + : scenario === 'with-tool' || scenario === 'raw-tool-args' ? [weatherTool] : scenario === 'structured-client-tool-wait' ? [clientContextToolDefinition] diff --git a/testing/e2e/src/routes/api.openai-responses-usage.ts b/testing/e2e/src/routes/api.openai-responses-usage.ts new file mode 100644 index 0000000000..c2a2d6c82b --- /dev/null +++ b/testing/e2e/src/routes/api.openai-responses-usage.ts @@ -0,0 +1,49 @@ +import { createFileRoute } from '@tanstack/react-router' +import { chat, createChatOptions } from '@tanstack/ai' +import { createOpenaiChat } from '@tanstack/ai-openai' + +const LLMOCK_DEFAULT_BASE = process.env.LLMOCK_URL || 'http://127.0.0.1:4010' +const DUMMY_KEY = 'sk-e2e-test-dummy-key' + +/** + * Drives the OpenAI Responses adapter against a hand-crafted aimock mount + * (`/openai-responses-usage`) whose `response.completed` usage reports + * `input_tokens_details.cache_write_tokens`. The companion spec asserts it + * reaches `RUN_FINISHED.usage` as `promptTokensDetails.cacheWriteTokens`. + */ +export const Route = createFileRoute('/api/openai-responses-usage')({ + server: { + handlers: { + POST: async () => { + const adapter = createOpenaiChat('gpt-4o', DUMMY_KEY, { + baseURL: `${LLMOCK_DEFAULT_BASE}/openai-responses-usage/v1`, + }) + + let usage: Record | undefined + try { + for await (const chunk of chat({ + ...createChatOptions({ adapter }), + messages: [{ role: 'user', content: 'hi' }], + })) { + if (chunk.type === 'RUN_FINISHED') { + usage = chunk.usage as Record | undefined + } + } + } catch (error) { + return new Response( + JSON.stringify({ + ok: false, + error: error instanceof Error ? error.message : String(error), + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ) + } + + return new Response(JSON.stringify({ ok: true, usage }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/api.openrouter-json-object-wire.ts b/testing/e2e/src/routes/api.openrouter-json-object-wire.ts index 0c5a820415..2a29b788ce 100644 --- a/testing/e2e/src/routes/api.openrouter-json-object-wire.ts +++ b/testing/e2e/src/routes/api.openrouter-json-object-wire.ts @@ -18,6 +18,9 @@ export const Route = createFileRoute('/api/openrouter-json-object-wire')({ POST: async ({ request }) => { const url = new URL(request.url) const testId = url.searchParams.get('testId') ?? undefined + // `via=catalog`: a json_object-only catalog model, no modelOptions. + const catalog = url.searchParams.get('via') === 'catalog' + const model = catalog ? 'qwen/qwen3.7-flash' : 'openai/gpt-4o' const httpClient = new HTTPClient() if (testId) { @@ -28,7 +31,7 @@ export const Route = createFileRoute('/api/openrouter-json-object-wire')({ }) } - const adapter = createOpenRouterText('openai/gpt-4o', DUMMY_KEY, { + const adapter = createOpenRouterText(model, DUMMY_KEY, { serverURL: `${LLMOCK_DEFAULT_BASE}/v1`, httpClient, }) @@ -42,9 +45,9 @@ export const Route = createFileRoute('/api/openrouter-json-object-wire')({ content: '[json-object-wire] return a person as json', }, ], - modelOptions: { - responseFormat: { type: 'json_object' }, - }, + ...(catalog + ? {} + : { modelOptions: { responseFormat: { type: 'json_object' } } }), outputSchema: { type: 'object', properties: { diff --git a/testing/e2e/src/routes/api.tools-test.ts b/testing/e2e/src/routes/api.tools-test.ts index 3e40bc8957..43f5815264 100644 --- a/testing/e2e/src/routes/api.tools-test.ts +++ b/testing/e2e/src/routes/api.tools-test.ts @@ -65,9 +65,10 @@ function createProviderFreeAdapter(scenario: string): AnyTextAdapter { : scenario === 'client-tool-input-error' || scenario === 'invalid-client-tool-retry' ? { - arguments: '{"message":42,"type":"info"}', + // An object cannot be repaired into the string `message`. + arguments: '{"message":{"text":"hi"},"type":"info"}', initialText: 'Showing a notification.', - input: { message: 42, type: 'info' }, + input: { message: { text: 'hi' }, type: 'info' }, name: scenario === 'invalid-client-tool-retry' ? 'invalid-client-tool-retry-test' diff --git a/testing/e2e/src/routes/call-ownership.tsx b/testing/e2e/src/routes/call-ownership.tsx new file mode 100644 index 0000000000..e477470aca --- /dev/null +++ b/testing/e2e/src/routes/call-ownership.tsx @@ -0,0 +1,50 @@ +import { useEffect } from 'react' +import { createFileRoute } from '@tanstack/react-router' +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' + +/** + * Harness for issue #1533. `/api/call-ownership-wire` streams three model + * calls the way `chat()` publishes them: one `RUN_STARTED`, then a + * `tool_calls` `RUN_FINISHED` between calls, with reasoning before each + * call's text id. The page prints one assistant message per call. + */ +function partLabel(part: { type: string; content?: unknown }): string { + if ( + (part.type === 'thinking' || part.type === 'text') && + typeof part.content === 'string' + ) { + return `${part.type}:${part.content}` + } + return part.type +} + +function CallOwnershipPage() { + const { messages, sendMessage } = useChat({ + threadId: 'call-ownership-1', + connection: fetchServerSentEvents('/api/call-ownership-wire'), + }) + + const summary = messages + .filter((message) => message.role === 'assistant') + .map((message) => { + const parts = message.parts.map((part) => partLabel(part)).join(',') + return `${message.id} ${parts}` + }) + .join(';') + + useEffect(() => { + void sendMessage('go') + // Fire the single run once on mount; the harness route ignores the content. + // eslint-disable-next-line react-hooks/exhaustive-deps + }, []) + + return ( +
+
{summary}
+
+ ) +} + +export const Route = createFileRoute('/call-ownership')({ + component: CallOwnershipPage, +}) diff --git a/testing/e2e/src/routes/middleware-test.tsx b/testing/e2e/src/routes/middleware-test.tsx index 6bde171f32..fc326cc99b 100644 --- a/testing/e2e/src/routes/middleware-test.tsx +++ b/testing/e2e/src/routes/middleware-test.tsx @@ -19,6 +19,7 @@ const MIDDLEWARE_MODES = [ { id: 'chunk-transform', label: 'Chunk Transform (prefix text)' }, { id: 'tool-skip', label: 'Tool Skip (skip with custom result)' }, { id: 'tool-replace', label: 'Tool Replace (replace result after call)' }, + { id: 'tool-bad-args', label: 'Tool Bad Args (invalid args rewrite)' }, { id: 'capability', label: 'Capability (provide/consume prefix)' }, { id: 'phase-recorder', label: 'Phase Recorder (capture phase + chunks)' }, { id: 'otel', label: 'OpenTelemetry (capture spans/metrics)' }, @@ -253,6 +254,7 @@ function MiddlewareTestPage() { +