From 7487a603041a5fa0b7e588f8e8a65ea72e4c5d1d Mon Sep 17 00:00:00 2001 From: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Date: Fri, 23 Oct 2026 16:24:10 +1100 Subject: [PATCH 1/2] feat(ai-ui-message-stream): client tools, structured output and reconnect for AI SDK routes TanStack useChat gets more from an AI SDK route: - fromChatTransport sends the client's tool declarations in the request body. clientToolsFromAISDKBody(body, serverTools) turns them into AI SDK tools without execute on the route, so client and WebMCP tools need no copy there. Server tools win a name clash. - With outputSchema, an Output.object answer becomes a typed structured-output part with partial objects and a checked final value. - A dropped stream reconnects with createReconnectTracker, the backoff of the TanStack HTTP connections, set with the reconnect option. @tanstack/ai-client now exports createReconnectTracker. Tests prove client persistence, the message queue, devtools and loading saved AI SDK chats over the bridge. --- .changeset/ai-ui-message-stream-client.md | 12 + docs/migration/ai-sdk-compatibility.md | 114 ++++---- .../tanstack-client-ai-sdk-server.md | 109 +++++++- packages/ai-client/src/index.ts | 2 + packages/ai-ui-message-stream/README.md | 19 +- packages/ai-ui-message-stream/package.json | 5 - .../ai-ui-message-stream/src/client-tools.ts | 60 ++++ packages/ai-ui-message-stream/src/index.ts | 91 +++++-- packages/ai-ui-message-stream/src/stream.ts | 68 ++++- .../tests/client-extras.test.ts | 257 ++++++++++++++++++ .../tests/client-tools.test.ts | 135 +++++++++ .../tests/reconnect.test.ts | 67 +++++ skills/tanstack-ai-migration/SKILL.md | 4 +- 13 files changed, 842 insertions(+), 101 deletions(-) create mode 100644 .changeset/ai-ui-message-stream-client.md create mode 100644 packages/ai-ui-message-stream/src/client-tools.ts create mode 100644 packages/ai-ui-message-stream/tests/client-extras.test.ts create mode 100644 packages/ai-ui-message-stream/tests/client-tools.test.ts create mode 100644 packages/ai-ui-message-stream/tests/reconnect.test.ts diff --git a/.changeset/ai-ui-message-stream-client.md b/.changeset/ai-ui-message-stream-client.md new file mode 100644 index 0000000000..a12909be45 --- /dev/null +++ b/.changeset/ai-ui-message-stream-client.md @@ -0,0 +1,12 @@ +--- +'@tanstack/ai-ui-message-stream': minor +'@tanstack/ai-client': patch +--- + +TanStack `useChat` gets more from an AI SDK route through `fromChatTransport`: + +- Client tools go to the route in the request body. Read them on the route with `clientToolsFromAISDKBody(body, serverTools)`, so the route needs no copy. +- With `outputSchema`, an `Output.object` answer is a typed `structured-output` part with partial objects. +- A dropped stream reconnects with the backoff of the TanStack HTTP connections. Set it with the `reconnect` option. + +`@tanstack/ai-client` exports `createReconnectTracker` and its `ReconnectTracker` type. diff --git a/docs/migration/ai-sdk-compatibility.md b/docs/migration/ai-sdk-compatibility.md index 6964deeedc..f4d689fa95 100644 --- a/docs/migration/ai-sdk-compatibility.md +++ b/docs/migration/ai-sdk-compatibility.md @@ -29,29 +29,29 @@ Each row has one of these results: ## 📋 Features -| Feature | AI SDK client, TanStack server | TanStack client, AI SDK server | -| -------------------------------------------------------------------- | --------------------------------------------------------- | ------------------------------------------------------------------------------------------ | -| Text | ✅ Same | ✅ Same | -| Reasoning and its signature | ✅ Same | ✅ Same | -| Tool-call provider data (Gemini `thoughtSignature`, OpenAI item ids) | 🔄 Mapped to `callProviderMetadata` and back | 🔄 Mapped to tool-call `metadata` and back | -| Server tools | ✅ Same | ✅ Same | -| Provider-executed tools (web search) | 🔄 Mapped to `providerExecuted` | 🔄 Mapped to the TanStack provider-executed flag | -| Client tools | ✅ Same | ✅ Same. Pass the tools to `useChat`. | -| Tool approvals | 🔄 Mapped to `tool-approval-request` and back to `resume` | 🔄 Mapped to a `tool-approval` interrupt. The approval `signature` goes back to the route. | -| Tool errors and denials | ✅ Same | ✅ Same | -| Steps of a tool loop | 🔄 Mapped to `start-step` / `finish-step` | 🔄 Mapped to `STEP_*` events, and to `step-start` parts in the history | -| Files and images from the user | 🔄 Mapped to content parts | 🔄 Mapped to file parts | -| Web search sources | 🔄 Mapped to `source-url` chunks | 🔄 Mapped to custom events | -| Model, response ids and token usage | 🔄 Mapped to message metadata under `tanstack` | ❌ Dropped (the AI SDK stream has none) | -| Message metadata of the app | ✅ Same | 🔄 Mapped to `metadata.aiSdk` and back | -| Activity | 🔄 Mapped to a `data-` part | 🔄 Mapped from a `data-` part with an `id` | -| Subagents | 🧩 Grouped into one `data-subagent` card per child | ❌ Dropped (the AI SDK has no subagents) | -| Shared state | ➡️ One way, to a `data-state` part | ❌ Dropped (the AI SDK has no state) | -| Structured output | ✅ Same (the JSON is text) | 🔄 Mapped to a text part in the history | -| Custom events | 🔄 Mapped to a transient `data-` chunk | 🔄 Mapped from a transient `data-` chunk | -| Stream resume | ✅ Same, with a durability log | ✅ Same. A reload needs `threadId` on `fromChatTransport`. | -| System messages | ❌ Dropped. The server owns its system prompts. | ❌ Dropped. `streamText` rejects them. | -| Generated files | ❌ Dropped (`chat()` sends none) | 🔄 Mapped to custom events | +| Feature | AI SDK client, TanStack server | TanStack client, AI SDK server | +| -------------------------------------------------------------------- | --------------------------------------------------------- | -------------------------------------------------------------------------------------------------- | +| Text | ✅ Same | ✅ Same | +| Reasoning and its signature | ✅ Same | ✅ Same | +| Tool-call provider data (Gemini `thoughtSignature`, OpenAI item ids) | 🔄 Mapped to `callProviderMetadata` and back | 🔄 Mapped to tool-call `metadata` and back | +| Server tools | ✅ Same | ✅ Same | +| Provider-executed tools (web search) | 🔄 Mapped to `providerExecuted` | 🔄 Mapped to the TanStack provider-executed flag | +| Client tools | ✅ Same | 🔄 Mapped. Their schemas go to the route in the body. Read them with `clientToolsFromAISDKBody`. | +| Tool approvals | 🔄 Mapped to `tool-approval-request` and back to `resume` | 🔄 Mapped to a `tool-approval` interrupt. The approval `signature` goes back to the route. | +| Tool errors and denials | ✅ Same | ✅ Same | +| Steps of a tool loop | 🔄 Mapped to `start-step` / `finish-step` | 🔄 Mapped to `STEP_*` events, and to `step-start` parts in the history | +| Files and images from the user | 🔄 Mapped to content parts | 🔄 Mapped to file parts | +| Web search sources | 🔄 Mapped to `source-url` chunks | 🔄 Mapped to custom events | +| Model, response ids and token usage | 🔄 Mapped to message metadata under `tanstack` | ❌ Dropped (the AI SDK stream has none) | +| Message metadata of the app | ✅ Same | 🔄 Mapped to `metadata.aiSdk` and back | +| Activity | 🔄 Mapped to a `data-` part | 🔄 Mapped from a `data-` part with an `id` | +| Subagents | 🧩 Grouped into one `data-subagent` card per child | ❌ Dropped (the AI SDK has no subagents) | +| Shared state | ➡️ One way, to a `data-state` part | ❌ Dropped (the AI SDK has no state) | +| Structured output | ✅ Same (the JSON is text) | 🔄 Mapped to a typed `structured-output` part with `outputSchema`, and back to text in the history | +| Custom events | 🔄 Mapped to a transient `data-` chunk | 🔄 Mapped from a transient `data-` chunk | +| Stream resume | ✅ Same, with a durability log | ✅ Same. A reload needs `threadId` on `fromChatTransport`. | +| System messages | ❌ Dropped. The server owns its system prompts. | ❌ Dropped. `streamText` rejects them. | +| Generated files | ❌ Dropped (`chat()` sends none) | 🔄 Mapped to custom events | ## ▲ AI SDK client, TanStack server @@ -109,39 +109,40 @@ Each row has one of these results: ### ⬇️ Reply: AI SDK chunks to TanStack events -| AI SDK chunk | AG-UI event | Result | -| -------------------------------------------------------------------------------------- | --------------------------------------------------------------- | -------------------------------------------------------------------------------------- | -| `start` | `RUN_STARTED` | 🔄 Mapped. The message id is the AG-UI message id. | -| `messageMetadata` on `start`, `message-metadata`, `finish` | `metadata.aiSdk` on the message, at the end of the run | 🔄 Mapped | -| `start-step` / `finish-step` | `STEP_STARTED` / `STEP_FINISHED` | ✅ Same | -| `reset-step` | | ❌ Dropped | -| `text-start` / `text-delta` / `text-end` | `TEXT_MESSAGE_START` / `CONTENT` / `END` | 🧩 Grouped. All text parts go into one assistant message. | -| `providerMetadata` on text chunks | | ❌ Dropped. TanStack keeps no metadata per text part. | -| `text-end.citations` | `CUSTOM` named `citations` | 🔄 Mapped | -| `reasoning-start` / `reasoning-delta` / `reasoning-end` | `REASONING_MESSAGE_*` | ✅ Same. Each reasoning part gets its own id. | -| `providerMetadata` on `reasoning-end` | `REASONING_ENCRYPTED_VALUE` | 🔄 Mapped. It goes back to the route unchanged. | -| `tool-input-start` / `tool-input-delta` | `TOOL_CALL_START` / `TOOL_CALL_ARGS` | ✅ Same | -| `providerMetadata` on tool input chunks | tool-call `metadata.aiSdk.callProviderMetadata` | 🔄 Mapped | -| `providerExecuted` | tool-call `metadata.providerExecuted` | 🔄 Mapped. The client does not run the call. | -| `inputEncoding: 'text'` | One `TOOL_CALL_ARGS` with the final input | 🔄 Mapped | -| `tool-input-available` | `TOOL_CALL_END`, after any approval or output for the same call | ✅ Same | -| `tool-input-error` | `TOOL_CALL_END` plus an error `TOOL_CALL_RESULT` | 🔄 Mapped | -| `tool-output-available` | `TOOL_CALL_RESULT` | ✅ Same | -| `providerMetadata` on tool output chunks | tool-call `metadata.aiSdk.resultProviderMetadata` | 🔄 Mapped when it comes right after the call. Later, it is dropped. | -| `tool-output-available` with `preliminary: true` | | ❌ Dropped. Only the final output is sent. | -| `tool-output-error` / `tool-output-denied` | Error `TOOL_CALL_RESULT` | 🔄 Mapped | -| `tool-approval-request` | Approval interrupt on `RUN_FINISHED` | 🔄 Mapped. Without `tools`, it is a generic interrupt: resolve it with `{ approved }`. | -| `signature`, `approvalDescriptor`, `inputSchemaInput`, `reason` of an approval request | tool-call `metadata.aiSdk.approval` | 🔄 Mapped | -| `tool-approval-request` with `isAutomatic` | | ❌ Dropped. The route already decided, and the output or denial follows. | -| `tool-approval-response` | | ❌ Dropped | -| `finish` | `RUN_FINISHED` with the finish reason | 🔄 Mapped | -| `finish` with `tool-calls` | `RUN_FINISHED` with no finish reason | 🔄 Mapped. In TanStack, `tool_calls` means more steps follow in the same stream. | -| `error` | `RUN_ERROR` | 🔄 Mapped | -| `abort` | `RUN_FINISHED` with a `cancelled` outcome | 🔄 Mapped | -| `data-` with an `id` and object data | `ACTIVITY_SNAPSHOT` of type `` | 🔄 Mapped. The same id replaces the activity. | -| Other `data-` (transient, no `id`, or not an object) | `CUSTOM` named `data-` | 🔄 Mapped | -| `source-url`, `source-document`, `file`, `reasoning-file` | `CUSTOM` with the same name | 🔄 Mapped | -| `custom` | | ❌ Dropped | +| AI SDK chunk | AG-UI event | Result | +| -------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------- | +| `start` | `RUN_STARTED` | 🔄 Mapped. The message id is the AG-UI message id. | +| `messageMetadata` on `start`, `message-metadata`, `finish` | `metadata.aiSdk` on the message, at the end of the run | 🔄 Mapped | +| `start-step` / `finish-step` | `STEP_STARTED` / `STEP_FINISHED` | ✅ Same | +| `reset-step` | | ❌ Dropped | +| `text-start` / `text-delta` / `text-end` | `TEXT_MESSAGE_START` / `CONTENT` / `END` | 🧩 Grouped. All text parts go into one assistant message. | +| Text of an `Output.object` answer, with `outputSchema` | `structured-output.start`, the text, then `structured-output.complete` with the checked object | 🔄 Mapped. A schema failure ends the run with `RUN_ERROR`. | +| `providerMetadata` on text chunks | | ❌ Dropped. TanStack keeps no metadata per text part. | +| `text-end.citations` | `CUSTOM` named `citations` | 🔄 Mapped | +| `reasoning-start` / `reasoning-delta` / `reasoning-end` | `REASONING_MESSAGE_*` | ✅ Same. Each reasoning part gets its own id. | +| `providerMetadata` on `reasoning-end` | `REASONING_ENCRYPTED_VALUE` | 🔄 Mapped. It goes back to the route unchanged. | +| `tool-input-start` / `tool-input-delta` | `TOOL_CALL_START` / `TOOL_CALL_ARGS` | ✅ Same | +| `providerMetadata` on tool input chunks | tool-call `metadata.aiSdk.callProviderMetadata` | 🔄 Mapped | +| `providerExecuted` | tool-call `metadata.providerExecuted` | 🔄 Mapped. The client does not run the call. | +| `inputEncoding: 'text'` | One `TOOL_CALL_ARGS` with the final input | 🔄 Mapped | +| `tool-input-available` | `TOOL_CALL_END`, after any approval or output for the same call | ✅ Same | +| `tool-input-error` | `TOOL_CALL_END` plus an error `TOOL_CALL_RESULT` | 🔄 Mapped | +| `tool-output-available` | `TOOL_CALL_RESULT` | ✅ Same | +| `providerMetadata` on tool output chunks | tool-call `metadata.aiSdk.resultProviderMetadata` | 🔄 Mapped when it comes right after the call. Later, it is dropped. | +| `tool-output-available` with `preliminary: true` | | ❌ Dropped. Only the final output is sent. | +| `tool-output-error` / `tool-output-denied` | Error `TOOL_CALL_RESULT` | 🔄 Mapped | +| `tool-approval-request` | Approval interrupt on `RUN_FINISHED` | 🔄 Mapped. Without `tools`, it is a generic interrupt: resolve it with `{ approved }`. | +| `signature`, `approvalDescriptor`, `inputSchemaInput`, `reason` of an approval request | tool-call `metadata.aiSdk.approval` | 🔄 Mapped | +| `tool-approval-request` with `isAutomatic` | | ❌ Dropped. The route already decided, and the output or denial follows. | +| `tool-approval-response` | | ❌ Dropped | +| `finish` | `RUN_FINISHED` with the finish reason | 🔄 Mapped | +| `finish` with `tool-calls` | `RUN_FINISHED` with no finish reason | 🔄 Mapped. In TanStack, `tool_calls` means more steps follow in the same stream. | +| `error` | `RUN_ERROR` | 🔄 Mapped | +| `abort` | `RUN_FINISHED` with a `cancelled` outcome | 🔄 Mapped | +| `data-` with an `id` and object data | `ACTIVITY_SNAPSHOT` of type `` | 🔄 Mapped. The same id replaces the activity. | +| Other `data-` (transient, no `id`, or not an object) | `CUSTOM` named `data-` | 🔄 Mapped | +| `source-url`, `source-document`, `file`, `reasoning-file` | `CUSTOM` with the same name | 🔄 Mapped | +| `custom` | | ❌ Dropped | ### ⬆️ Request: TanStack messages to AI SDK message parts @@ -155,6 +156,7 @@ Each row has one of these results: | Provider-executed `tool-call` | `dynamic-tool` with `providerExecuted: true` | 🔄 Mapped | | `tool-call` with an approval answer | `dynamic-tool` with `approval-responded`, and the stored `signature` | 🔄 Mapped | | `tool-call` that waits for approval | `dynamic-tool` with `approval-requested` | 🔄 Mapped | +| Client tool declarations of the run | `clientTools` in the request body | 🔄 Mapped | | A tool result followed by more of the reply | `step-start` before the next step | 🔄 Mapped | | `structured-output`, complete | `text` with the JSON | 🔄 Mapped | | Message `metadata.aiSdk` | Message `metadata` | 🔄 Mapped | diff --git a/docs/migration/tanstack-client-ai-sdk-server.md b/docs/migration/tanstack-client-ai-sdk-server.md index 897b926524..4930d5e8af 100644 --- a/docs/migration/tanstack-client-ai-sdk-server.md +++ b/docs/migration/tanstack-client-ai-sdk-server.md @@ -19,6 +19,14 @@ Point TanStack AI `useChat` at a Vercel AI SDK route. Wrap an AI SDK `ChatTransp Use this page to move the client first. To move the server first, see [AI SDK useChat with a TanStack AI server](./ai-sdk-client-tanstack-server). +Against your AI SDK route, TanStack `useChat` gives you: + +- Client persistence and resume after a reload, with no server code. +- A [message queue](../chat/queueing) while a reply streams. +- Client tools that the route reads from the request, with no copy on the route. +- Typed structured output, with partial objects while it streams. +- Typed approvals, TanStack devtools, and hooks for React, Solid, Vue, Svelte, Preact, Angular and more. + > [!WARNING] > This bridge is a stop gap while you port. Some features change shape or are dropped on the way (see [AI SDK bridge compatibility](./ai-sdk-compatibility)). When the client works, port the route to `chat()` with the [migration guide](./migration-from-vercel-ai), then remove the bridge. @@ -96,7 +104,38 @@ export async function POST(request: Request) { ## 4. Run client tools and approve tool calls -Define the tools of the route on the client too, with the same names. Pass them to `useChat` and to `fromChatTransport`: +Define client tools on the client only. The connection sends their names, descriptions and JSON Schemas to the route in the request body. Read them on the route with `clientToolsFromAISDKBody`, together with the server tools: + +```typescript +import { openai } from '@ai-sdk/openai' +import { clientToolsFromAISDKBody } from '@tanstack/ai-ui-message-stream' +import { convertToModelMessages, streamText, tool } from 'ai' +import { z } from 'zod' + +const getWeather = tool({ + description: 'Get the weather for a city', + inputSchema: z.object({ city: z.string() }), + execute: async ({ city }) => ({ city, temperature: 72 }), +}) + +export async function POST(request: Request) { + const body = await request.json() + + const result = streamText({ + model: openai('gpt-6.1-sol'), + messages: await convertToModelMessages(body.messages), + tools: clientToolsFromAISDKBody(body, { getWeather }), + }) + + return result.toUIMessageStreamResponse() +} +``` + +- A server tool wins when a client tool has the same name. +- A client tool has no `execute` on the route, so `streamText` stops at the call. +- Browser tools from [WebMCP](../tools/webmcp) reach the route the same way. + +Pass the tools to `useChat`, and to `fromChatTransport` for approvals: - **Client tool:** give it a `.client()` implementation. When the route stops at the call, the client runs the tool and sends the output back. - **Tool approval:** a tool that needs approval on the route shows up in `interrupts`. Resolve it with `interrupt.resolveInterrupt(true)` or `false`. The answer goes back to the route, and `streamText` runs or skips the tool. @@ -148,13 +187,75 @@ export function Chat() { } ``` -## 5. Show data parts as activity +## 5. Get typed structured output + +When the route streams an object with `Output.object`, pass the same schema to `useChat` and to `fromChatTransport`. The answer is a `structured-output` part, with partial objects while it streams and the checked `data` at the end: + +```tsx +import { useChat } from '@tanstack/ai-react' +import { fromChatTransport } from '@tanstack/ai-ui-message-stream' +import { DefaultChatTransport } from 'ai' +import { z } from 'zod' + +const recipe = z.object({ name: z.string(), steps: z.array(z.string()) }) + +const connection = fromChatTransport( + new DefaultChatTransport({ api: '/api/recipe' }), + { outputSchema: recipe }, +) + +export function Recipe() { + const { messages } = useChat({ + connection, + threadId: 'recipe', + outputSchema: recipe, + }) + + return messages.flatMap((message) => + message.parts.map((part, i) => + part.type === 'structured-output' ? ( +

{part.partial?.name}

+ ) : null, + ), + ) +} +``` + +If the text is not valid JSON for the schema, the run ends with an error. + +## 6. Load a saved chat + +If your app saves AI SDK messages, convert them with `fromAISDKMessages` and pass them as `initialMessages`: + +```tsx +import { useChat } from '@tanstack/ai-react' +import { + fromAISDKMessages, + fromChatTransport, +} from '@tanstack/ai-ui-message-stream' +import { DefaultChatTransport } from 'ai' +import type { UIMessage } from 'ai' + +const connection = fromChatTransport( + new DefaultChatTransport({ api: '/api/chat' }), +) + +export function SavedChat({ saved }: { saved: Array }) { + return useChat({ + connection, + threadId: 'support-chat', + initialMessages: fromAISDKMessages(saved), + }) +} +``` + +## 7. Show data parts as activity A data part with an `id` (`writer.write({ type: 'data-weather', id, data })` on the route) arrives as an activity message of type `weather`. A new value for the same id replaces it. A transient data part, or one with no `id`, arrives as a custom event named `data-`. -## 6. Resume a stream +## 8. Resume a stream -A dropped connection needs nothing. The connection calls `reconnectToStream` on the transport, and the reply continues. +A dropped connection needs nothing. The connection calls `reconnectToStream` on the transport, and the reply continues. It waits longer between tries that get no new data, as the TanStack HTTP connections do. To change the waits and the limit, pass `reconnect` to `fromChatTransport`, for example `{ maxAttempts: 10 }`. To continue a reply after a page reload, add client persistence, and pass the `threadId` to `fromChatTransport`. The route must support the AI SDK resume `GET` (`/api/chat//stream`): diff --git a/packages/ai-client/src/index.ts b/packages/ai-client/src/index.ts index c4561ca23e..aa41eb719c 100644 --- a/packages/ai-client/src/index.ts +++ b/packages/ai-client/src/index.ts @@ -175,6 +175,8 @@ export { StreamTruncatedError, DurableStreamIncompleteError, StreamReconnectLimitError, + createReconnectTracker, + type ReconnectTracker, type ConnectConnectionAdapter, type ChatHydrateOptions, type ChatHydrationResult, diff --git a/packages/ai-ui-message-stream/README.md b/packages/ai-ui-message-stream/README.md index 13d3b49ea3..e496c4d4b6 100644 --- a/packages/ai-ui-message-stream/README.md +++ b/packages/ai-ui-message-stream/README.md @@ -78,15 +78,16 @@ export function Chat() { ## API -| Export | Direction | -| -------------------------------------------------------- | ---------------------------------------------------------------------- | -| `chatParamsFromAISDKRequest(request)` | AI SDK `useChat` request to `chat()` params, approval answers included | -| `toUIMessageStream(stream)` | `chat()` events to an AI SDK UI message stream | -| `toResumableUIMessageStreamResponse(stream, durability)` | A `chat()` run that AI SDK `useChat({ resume: true })` can rejoin | -| `resumeUIMessageStreamResponse(durability)` | The AI SDK resume `GET` | -| `fromChatTransport(transport, { threadId, tools })` | AI SDK `ChatTransport` to a TanStack `useChat` connection | -| `fromAISDKMessages` / `toAISDKMessages` | Message history in each direction | -| `fromUIMessageStream(stream, ids)` | AI SDK UI message stream to TanStack events | +| Export | Direction | +| ---------------------------------------------------------------------------- | -------------------------------------------------------------------------- | +| `chatParamsFromAISDKRequest(request)` | AI SDK `useChat` request to `chat()` params, approval answers included | +| `toUIMessageStream(stream)` | `chat()` events to an AI SDK UI message stream | +| `toResumableUIMessageStreamResponse(stream, durability)` | A `chat()` run that AI SDK `useChat({ resume: true })` can rejoin | +| `resumeUIMessageStreamResponse(durability)` | The AI SDK resume `GET` | +| `fromChatTransport(transport, { threadId, tools, outputSchema, reconnect })` | AI SDK `ChatTransport` to a TanStack `useChat` connection | +| `clientToolsFromAISDKBody(body, serverTools)` | On an AI SDK route: the client tools of the request, plus the server tools | +| `fromAISDKMessages` / `toAISDKMessages` | Message history in each direction | +| `fromUIMessageStream(stream, ids)` | AI SDK UI message stream to TanStack events | Text, reasoning, tools, tool approvals, activity, subagents and stream resume cross in both directions where the other side has them. Shared state goes from a TanStack server to an AI SDK client only, because the AI SDK client does not send state. diff --git a/packages/ai-ui-message-stream/package.json b/packages/ai-ui-message-stream/package.json index 8f377b7261..c25632ba95 100644 --- a/packages/ai-ui-message-stream/package.json +++ b/packages/ai-ui-message-stream/package.json @@ -49,11 +49,6 @@ "@tanstack/ai-client": "workspace:^", "ai": "^7.0.0" }, - "peerDependenciesMeta": { - "@tanstack/ai-client": { - "optional": true - } - }, "devDependencies": { "@ag-ui/core": "1.0.0", "@tanstack/ai": "workspace:*", diff --git a/packages/ai-ui-message-stream/src/client-tools.ts b/packages/ai-ui-message-stream/src/client-tools.ts new file mode 100644 index 0000000000..bcddace4d3 --- /dev/null +++ b/packages/ai-ui-message-stream/src/client-tools.ts @@ -0,0 +1,60 @@ +import { jsonSchema, tool } from 'ai' +import type { ToolSet } from 'ai' + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value) +} + +/** + * Add the tools that a TanStack `useChat` client declares (`clientTools` in + * the request body, sent by `fromChatTransport`) to the tools of an AI SDK + * route. A client tool has no `execute`, so `streamText` stops at its call + * and the client runs it. A server tool wins on a name clash. + * + * ```ts + * const body = await request.json() + * const result = streamText({ + * model, + * messages: await convertToModelMessages(body.messages), + * tools: clientToolsFromAISDKBody(body, { weather }), + * }) + * ``` + * + * @throws An error when `clientTools` is present but is not a list of + * `{ name, description, parameters }` with a JSON Schema object. + */ +export function clientToolsFromAISDKBody( + body: unknown, + serverTools: TServerTools, +): TServerTools & ToolSet { + const declared = isRecord(body) ? body.clientTools : undefined + if (declared === undefined) return serverTools + if (!Array.isArray(declared)) { + throw new Error('`clientTools` must be an array.') + } + // `fromEntries` defines own properties, so a `__proto__` name is inert. + const clientTools: ToolSet = Object.fromEntries( + declared.map((entry: unknown) => { + if ( + !isRecord(entry) || + typeof entry.name !== 'string' || + entry.name === '' || + typeof entry.description !== 'string' || + !isRecord(entry.parameters) + ) { + throw new Error( + 'Each client tool must have a `name`, a `description` and JSON Schema `parameters`.', + ) + } + // No `execute`: `streamText` stops at the call and the client runs it. + return [ + entry.name, + tool({ + description: entry.description, + inputSchema: jsonSchema(entry.parameters), + }), + ] + }), + ) + return { ...clientTools, ...serverTools } +} diff --git a/packages/ai-ui-message-stream/src/index.ts b/packages/ai-ui-message-stream/src/index.ts index 536024ffb2..15c5e1ca8e 100644 --- a/packages/ai-ui-message-stream/src/index.ts +++ b/packages/ai-ui-message-stream/src/index.ts @@ -1,6 +1,10 @@ import { modelMessagesToUIMessages } from '@tanstack/ai' import type { ModelMessage, UIMessage } from '@tanstack/ai' -import type { ConnectConnectionAdapter } from '@tanstack/ai-client' +import { createReconnectTracker } from '@tanstack/ai-client' +import type { + ConnectConnectionAdapter, + ReconnectOptions, +} from '@tanstack/ai-client' import type { ChatTransport, UIMessage as AISDKUIMessage, @@ -22,6 +26,7 @@ export { chatParamsFromAISDKRequestBody, } from './request' export type { AISDKChatParams } from './request' +export { clientToolsFromAISDKBody } from './client-tools' export type { SubagentData } from './data-parts' function asUIMessages( @@ -35,25 +40,26 @@ function asUIMessages( ) } -type ReconnectOptions = Parameters< +type ReconnectRequest = Parameters< ChatTransport['reconnectToStream'] >[0] /** * Read `stream`, and when the connection drops, read the rest through - * `transport.reconnectToStream`. A replayed prefix is skipped. + * `transport.reconnectToStream`, with the backoff and limits of the TanStack + * HTTP connections. A replayed prefix is skipped. */ function reconnecting( transport: ChatTransport, stream: ReadableStream, - options: ReconnectOptions, + request: ReconnectRequest, + reconnect: ReconnectOptions | undefined, ): ReadableStream { + const tracker = createReconnectTracker(reconnect) let reader = stream.getReader() let delivered = 0 let skip = 0 - // ponytail: a reconnect that delivers nothing new ends the stream with the - // error. Add attempts and backoff if one retry without progress is too few. - let deliveredAtReconnect: number | undefined + let deliveredAtLastDrop = 0 return new ReadableStream({ async pull(controller) { for (;;) { @@ -69,10 +75,20 @@ function reconnecting( delivered++ return controller.enqueue(value) } catch (error) { - if (options.abortSignal?.aborted) throw error - if (deliveredAtReconnect === delivered) throw error - deliveredAtReconnect = delivered - const resumed = await transport.reconnectToStream(options) + if (request.abortSignal?.aborted) throw error + // Throws StreamReconnectLimitError when the run is stuck. + await tracker.waitBeforeReconnect( + delivered > deliveredAtLastDrop, + request.abortSignal, + ) + deliveredAtLastDrop = delivered + let resumed: ReadableStream | null + try { + resumed = await transport.reconnectToStream(request) + } catch { + // The old reader fails again, so the next pass waits and retries. + continue + } if (!resumed) throw error reader = resumed.getReader() skip = transport.resumeStreamIsReplay === true ? delivered : 0 @@ -93,23 +109,39 @@ function reconnecting( * request. * * When the stream drops, the adapter reconnects with - * `transport.reconnectToStream` and continues the run. Set `threadId` to also + * `transport.reconnectToStream` and continues the run, with the same backoff + * as the TanStack HTTP connections (`reconnect` sets it). Set `threadId` to also * rejoin a run after a page reload (`joinRun`), with a chat `persistence` * adapter on the client. For that, the transport must replay the stream from * its start (`resumeStreamIsReplay`, true for `DefaultChatTransport`). * * Pass the same `tools` as `useChat({ tools })`, so a tool approval is a * `tool-approval` interrupt that `resolveInterrupt(true)` answers. + * + * Pass the same `outputSchema` as `useChat({ outputSchema })` when the route + * uses `Output.object({ schema })`, so the answer is a typed + * `structured-output` part. + * + * The client's tools (JSON Schema) go in the body as `clientTools`. Read them + * on the route with `clientToolsFromAISDKBody`, so the route needs no copy. */ export function fromChatTransport( transport: ChatTransport, { threadId, tools, - }: Pick & { threadId?: string } = {}, + outputSchema, + reconnect, + }: Pick & { + threadId?: string + /** Backoff and limits for reconnects after a dropped stream. */ + reconnect?: ReconnectOptions + } = {}, ): ConnectConnectionAdapter { return { async *connect(messages, data, abortSignal, runContext) { + const { clientTools } = runContext + const body = clientTools?.length ? { ...data, clientTools } : data const stream = await transport.sendMessages({ trigger: 'submit-message', chatId: runContext.threadId, @@ -117,15 +149,20 @@ export function fromChatTransport( messages: toAISDKMessages(asUIMessages(messages)), abortSignal, ...(runContext.headers && { headers: runContext.headers }), - ...(data && { body: data }), + ...(body && { body }), }) yield* fromUIMessageStream( - reconnecting(transport, stream, { - chatId: runContext.threadId, - abortSignal, - ...(runContext.headers && { headers: runContext.headers }), - }), - { ...runContext, tools }, + reconnecting( + transport, + stream, + { + chatId: runContext.threadId, + abortSignal, + ...(runContext.headers && { headers: runContext.headers }), + }, + reconnect, + ), + { ...runContext, tools, outputSchema }, ) }, // A rejoin gets only the run id, and the AI SDK reconnects by chat id. @@ -136,11 +173,15 @@ export function fromChatTransport( const stream = await transport.reconnectToStream(options) // Thrown before the first event, so the client drops its resume pointer. if (!stream) throw new Error(`No active stream for chat ${threadId}`) - yield* fromUIMessageStream(reconnecting(transport, stream, options), { - threadId, - runId, - tools, - }) + yield* fromUIMessageStream( + reconnecting(transport, stream, options, reconnect), + { + threadId, + runId, + tools, + outputSchema, + }, + ) }, }), } diff --git a/packages/ai-ui-message-stream/src/stream.ts b/packages/ai-ui-message-stream/src/stream.ts index e9f8d41b28..7ff6ca5e4e 100644 --- a/packages/ai-ui-message-stream/src/stream.ts +++ b/packages/ai-ui-message-stream/src/stream.ts @@ -6,12 +6,14 @@ import { hashSchemaInput, normalizeApprovalSchema, readInterruptBinding, + validateWithStandardSchema, withInterruptBinding, } from '@tanstack/ai' import type { AnyClientTool, Interrupt, RunFinishedEvent, + SchemaInput, StreamChunk, ToolCallStartEvent, ToolResultOutcome, @@ -427,6 +429,14 @@ export interface FromUIMessageStreamOptions { * one of them with `needsApproval: true` is a `tool-approval` interrupt. */ tools?: ReadonlyArray + /** + * The schema of the route's `Output.object({ schema })`, as passed to + * `useChat({ outputSchema })`. The answer text becomes a `structured-output` + * part with partial objects while it streams. At the end the text is + * checked against the schema: a valid answer is the part's `data`, and an + * invalid one ends the run with a RUN_ERROR. + */ + outputSchema?: SchemaInput } /** @@ -454,7 +464,12 @@ export interface FromUIMessageStreamOptions { */ export async function* fromUIMessageStream( stream: ReadableStream, - { threadId, runId, tools: clientTools }: FromUIMessageStreamOptions, + { + threadId, + runId, + tools: clientTools, + outputSchema, + }: FromUIMessageStreamOptions, ): AsyncGenerator { let messageId = `${runId}-assistant` let ended = false @@ -478,6 +493,40 @@ export async function* fromUIMessageStream( const approvals: Array = [] // ponytail: shallow merge, where the AI SDK merges nested objects too. let messageMetadata: unknown + // The answer text, once text starts with `outputSchema` set. + // ponytail: all text of the stream is the answer, so text in a tool-call + // step joins it. Start the answer per step if a route sends such text. + let structuredRaw: string | undefined + + /** + * `structured-output.complete` with the checked answer, or a RUN_ERROR + * when the text is not JSON or fails the schema. + */ + async function structuredEnd(): Promise { + if (structuredRaw === undefined) return undefined + const object = parseJson(structuredRaw) + const result = + object === undefined + ? undefined + : await validateWithStandardSchema(outputSchema, object) + if (result?.success) { + return { + type: EventType.CUSTOM, + name: 'structured-output.complete', + value: { object: result.data, raw: structuredRaw, messageId }, + } + } + return { + type: EventType.RUN_ERROR, + code: 'structured-output-validation-failed', + message: + result === undefined + ? 'The structured output is not valid JSON' + : `The structured output does not match the schema: ${result.issues + .map((issue) => issue.message) + .join(', ')}`, + } + } function mergeMessageMetadata(value: unknown) { if (value === undefined) return @@ -649,6 +698,16 @@ export async function* fromUIMessageStream( // Every AI SDK text part belongs to the one assistant message. The // TanStack client starts a new text segment on each START. case 'text-start': + // The TanStack client puts the text after this event in a + // `structured-output` part, as for `chat({ outputSchema })`. + if (outputSchema !== undefined && structuredRaw === undefined) { + structuredRaw = '' + yield { + type: EventType.CUSTOM, + name: 'structured-output.start', + value: { messageId }, + } + } yield { type: EventType.TEXT_MESSAGE_START, messageId, @@ -656,6 +715,7 @@ export async function* fromUIMessageStream( } break case 'text-delta': + if (structuredRaw !== undefined) structuredRaw += chunk.delta yield { type: EventType.TEXT_MESSAGE_CONTENT, messageId, @@ -873,6 +933,9 @@ export async function* fromUIMessageStream( chunk.finishReason === undefined ? undefined : FROM_AI_SDK_FINISH[chunk.finishReason] + const end = await structuredEnd() + if (end) yield end + if (end?.type === EventType.RUN_ERROR) return yield* finished( runMetadata(finishReason == null ? undefined : { finishReason }), ) @@ -887,6 +950,9 @@ export async function* fromUIMessageStream( } yield* endTools() + const end = await structuredEnd() + if (end) yield end + if (end?.type === EventType.RUN_ERROR) return yield* finished(runMetadata()) } diff --git a/packages/ai-ui-message-stream/tests/client-extras.test.ts b/packages/ai-ui-message-stream/tests/client-extras.test.ts new file mode 100644 index 0000000000..ec26bb2704 --- /dev/null +++ b/packages/ai-ui-message-stream/tests/client-extras.test.ts @@ -0,0 +1,257 @@ +import { describe, expect, it, vi } from 'vitest' +import { EventSchemas } from '@ag-ui/core/schemas' +import { EventType } from '@tanstack/ai' +import { ChatClient } from '@tanstack/ai-client' +import { createChatDevtoolsBridge } from '@tanstack/ai-client/devtools' +import { convertToModelMessages } from 'ai' +import { z } from 'zod' +import { + fromAISDKMessages, + fromChatTransport, + fromUIMessageStream, +} from '../src' +import type { StreamChunk } from '@tanstack/ai' +import type { ChatDevtoolsBridge } from '@tanstack/ai-client/devtools' +import type { ChatTransport, UIMessage, UIMessageChunk } from 'ai' + +function chunkStream( + list: Array, +): ReadableStream { + return new ReadableStream({ + start(controller) { + for (const chunk of list) controller.enqueue(chunk) + controller.close() + }, + }) +} + +/** A transport that answers request N with `turns[N]`, after `gate` if set. */ +function fakeTransport( + turns: Array>, + gate?: Promise, +) { + const requests: Array> = [] + const transport: ChatTransport = { + async sendMessages({ messages }) { + requests.push(messages) + const turn = turns[requests.length - 1] ?? [] + if (requests.length === 1) await gate + return chunkStream(turn) + }, + reconnectToStream: () => Promise.resolve(null), + } + return { requests, transport } +} + +/** One AI SDK reply with message id `messageId`, streaming `deltas`. */ +function textTurn( + messageId: string, + ...deltas: Array +): Array { + return [ + { type: 'start', messageId }, + { type: 'text-start', id: '0' }, + ...deltas.map( + (delta): UIMessageChunk => ({ + type: 'text-delta', + id: '0', + delta, + }), + ), + { type: 'text-end', id: '0' }, + { type: 'finish', finishReason: 'stop' }, + ] +} + +const recipe = z.object({ name: z.string(), steps: z.array(z.string()) }) + +async function collect( + list: Array, + outputSchema?: typeof recipe, +) { + const out: Array = [] + for await (const event of fromUIMessageStream(chunkStream(list), { + threadId: 't', + runId: 'r', + outputSchema, + })) + out.push(event) + for (const event of out) EventSchemas.parse(event) + return out +} + +describe('AI SDK server, TanStack client: client features', () => { + it('shows an Output.object answer as a typed structured-output part', async () => { + const { transport } = fakeTransport([ + textTurn('a-1', '{"name":"Pan', 'cakes","steps":["Mix",', '"Fry"]}'), + ]) + const partials: Array = [] + const client = new ChatClient({ + threadId: 'thread-1', + connection: fromChatTransport(transport, { outputSchema: recipe }), + onMessagesChange: (messages) => { + const part = messages + .at(-1) + ?.parts.find((p) => p.type === 'structured-output') + if (part?.status === 'streaming') partials.push(part.partial) + }, + }) + + await client.sendMessage('A recipe') + expect(partials).toContainEqual({ name: 'Pan' }) + expect(partials).toContainEqual({ name: 'Pancakes', steps: ['Mix'] }) + const assistant = client.getMessages().at(-1) + expect(assistant?.parts).toEqual([ + { + type: 'structured-output', + status: 'complete', + raw: '{"name":"Pancakes","steps":["Mix","Fry"]}', + data: { name: 'Pancakes', steps: ['Mix', 'Fry'] }, + partial: { name: 'Pancakes', steps: ['Mix', 'Fry'] }, + }, + ]) + }) + + it('wraps the answer text in structured-output events', async () => { + const out = await collect( + textTurn('a-1', '{"name":"A","steps":[]}'), + recipe, + ) + expect( + out.map((e) => (e.type === EventType.CUSTOM ? e.name : e.type)), + ).toEqual([ + EventType.RUN_STARTED, + 'structured-output.start', + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + EventType.TEXT_MESSAGE_END, + 'structured-output.complete', + EventType.RUN_FINISHED, + ]) + expect( + out.find( + (e) => + e.type === EventType.CUSTOM && + e.name === 'structured-output.complete', + ), + ).toMatchObject({ + value: { object: { name: 'A', steps: [] }, messageId: 'a-1' }, + }) + }) + + it('ends with a RUN_ERROR when the answer fails the schema', async () => { + for (const text of ['{"name":1}', 'not json']) { + const out = await collect(textTurn('a-1', text), recipe) + expect(out.at(-1)).toMatchObject({ + type: EventType.RUN_ERROR, + code: 'structured-output-validation-failed', + }) + expect(out.some((e) => e.type === EventType.RUN_FINISHED)).toBe(false) + } + }) + + it('keeps plain text without outputSchema', async () => { + const out = await collect(textTurn('a-1', '{"name":"A"}')) + expect(out.some((e) => e.type === EventType.CUSTOM)).toBe(false) + }) + + it('queues a message sent while a reply streams and sends it after', async () => { + let release!: () => void + const gate = new Promise((resolve) => (release = resolve)) + const { requests, transport } = fakeTransport( + [textTurn('a-1', 'First.'), textTurn('a-2', 'Second.')], + gate, + ) + const client = new ChatClient({ + threadId: 'thread-1', + connection: fromChatTransport(transport), + }) + + const first = client.sendMessage('one') + await vi.waitFor(() => expect(client.getIsLoading()).toBe(true)) + await client.sendMessage('two') + expect(client.getQueue().map((q) => q.content)).toEqual(['two']) + expect(requests).toHaveLength(1) + + release() + await first + await vi.waitFor(() => expect(requests).toHaveLength(2)) + await vi.waitFor(() => expect(client.getIsLoading()).toBe(false)) + expect(client.getQueue()).toEqual([]) + expect( + requests[1]!.map((m) => + m.parts.flatMap((p) => (p.type === 'text' ? [p.text] : [])).join(''), + ), + ).toEqual(['one', 'First.', 'two']) + }) + + it('emits devtools events for a bridged run', async () => { + const { transport } = fakeTransport([textTurn('a-1', 'Hi.')]) + // The bridge useChat installs, as the framework hooks pass it. + let bridge: ChatDevtoolsBridge | undefined + const client = new ChatClient({ + threadId: 'thread-1', + connection: fromChatTransport(transport), + devtoolsBridgeFactory: (options) => + (bridge = createChatDevtoolsBridge(options)), + }) + if (!bridge) throw new Error('expected a devtools bridge') + const messageSent = vi.spyOn(bridge.events, 'messageSent') + const textUpdated = vi.spyOn(bridge.events, 'textUpdated') + const loadingChanged = vi.spyOn(bridge.events, 'loadingChanged') + + await client.sendMessage('Hello') + expect(messageSent).toHaveBeenCalledWith(expect.any(String), 'Hello') + expect(textUpdated).toHaveBeenCalledWith(expect.any(String), 'a-1', 'Hi.') + expect(loadingChanged.mock.calls.map(([loading]) => loading)).toEqual([ + true, + false, + ]) + }) + + it('loads saved AI SDK history and sends it back unchanged', async () => { + const saved: Array = [ + { id: 'u-1', role: 'user', parts: [{ type: 'text', text: 'Weather?' }] }, + { + id: 'a-1', + role: 'assistant', + // As AI SDK useChat saves a two-step answer. + parts: [ + { type: 'step-start' }, + { + type: 'dynamic-tool', + toolName: 'weather', + toolCallId: 'c-1', + input: { city: 'Paris' }, + state: 'output-available', + output: { temp: 20 }, + }, + { type: 'step-start' }, + { type: 'text', text: 'It is 20 degrees.' }, + ], + }, + ] + const { requests, transport } = fakeTransport([textTurn('a-3', 'Yes.')]) + const client = new ChatClient({ + threadId: 'thread-1', + connection: fromChatTransport(transport), + initialMessages: fromAISDKMessages(saved), + }) + expect(client.getMessages().map((m) => m.id)).toEqual(['u-1', 'a-1']) + + await client.sendMessage('Warm?') + // The leading step-start is dropped. The route's model messages match. + expect(await convertToModelMessages(requests[0]!.slice(0, 2))).toEqual( + await convertToModelMessages(saved), + ) + expect(requests[0]![1]!.parts.map((p) => p.type)).toEqual([ + 'dynamic-tool', + 'step-start', + 'text', + ]) + expect(requests[0]![2]).toMatchObject({ + role: 'user', + parts: [{ type: 'text', text: 'Warm?' }], + }) + }) +}) diff --git a/packages/ai-ui-message-stream/tests/client-tools.test.ts b/packages/ai-ui-message-stream/tests/client-tools.test.ts new file mode 100644 index 0000000000..1cda39a063 --- /dev/null +++ b/packages/ai-ui-message-stream/tests/client-tools.test.ts @@ -0,0 +1,135 @@ +import { describe, expect, it, vi } from 'vitest' +import { toolDefinition } from '@tanstack/ai/client' +import { ChatClient } from '@tanstack/ai-client' +import { convertToModelMessages, streamText, tool } from 'ai' +import { convertArrayToReadableStream, MockLanguageModelV4 } from 'ai/test' +import { z } from 'zod' +import { clientToolsFromAISDKBody, fromChatTransport } from '../src' +import type { ChatTransport, UIMessage } from 'ai' + +const usage = { + inputTokens: { + total: 1, + noCache: undefined, + cacheRead: undefined, + cacheWrite: undefined, + }, + outputTokens: { total: 1, text: undefined, reasoning: undefined }, +} + +describe('client tools over fromChatTransport', () => { + it('declares the client tools to an AI SDK route, which stops at the call', async () => { + const model = new MockLanguageModelV4({ + doStream: [ + { + stream: convertArrayToReadableStream([ + { type: 'stream-start', warnings: [] }, + { + type: 'tool-call', + toolCallId: 'call-1', + toolName: 'getLocation', + input: '{"precise":true}', + }, + { + type: 'finish', + finishReason: { unified: 'tool-calls', raw: undefined }, + usage, + }, + ]), + }, + { + stream: convertArrayToReadableStream([ + { type: 'stream-start', warnings: [] }, + { type: 'text-start', id: 't' }, + { type: 'text-delta', id: 't', delta: 'You are in Sydney.' }, + { type: 'text-end', id: 't' }, + { + type: 'finish', + finishReason: { unified: 'stop', raw: undefined }, + usage, + }, + ]), + }, + ], + }) + const bodies: Array = [] + // The AI SDK route, called in process. + const transport: ChatTransport = { + async sendMessages({ messages, body }) { + bodies.push(body) + const result = streamText({ + model, + messages: await convertToModelMessages(messages), + tools: clientToolsFromAISDKBody(body, {}), + }) + return result.toUIMessageStream() + }, + reconnectToStream: () => Promise.resolve(null), + } + const getLocation = vi.fn(() => ({ city: 'Sydney' })) + const client = new ChatClient({ + threadId: 'thread-1', + connection: fromChatTransport(transport), + tools: [ + toolDefinition({ + name: 'getLocation', + description: 'Get the user location', + inputSchema: z.object({ precise: z.boolean() }), + }).client(getLocation), + ], + }) + + await client.sendMessage('Where am I?') + await vi.waitFor(() => + expect(JSON.stringify(client.getMessages())).toContain('Sydney.'), + ) + + expect(bodies[0]).toMatchObject({ + clientTools: [ + { + name: 'getLocation', + description: 'Get the user location', + parameters: { type: 'object' }, + }, + ], + }) + // The route offered the tool to the model and did not run it. + expect(model.doStreamCalls[0]?.tools).toMatchObject([ + { type: 'function', name: 'getLocation' }, + ]) + expect(getLocation).toHaveBeenCalledWith( + { precise: true }, + expect.anything(), + ) + // The second request carried the client's output to the model. + expect(JSON.stringify(model.doStreamCalls[1]?.prompt)).toContain('Sydney') + }) + + it('lets a server tool win a name clash and rejects a bad declaration', () => { + const server = tool({ + description: 'Server weather', + inputSchema: z.object({}), + execute: () => 'sunny', + }) + const declared = (name: string) => ({ + name, + description: 'Client', + parameters: { type: 'object' }, + }) + const tools = clientToolsFromAISDKBody( + { clientTools: [declared('weather'), declared('__proto__')] }, + { weather: server }, + ) + expect(tools.weather).toBe(server) + expect(Object.keys(tools).sort()).toEqual(['__proto__', 'weather']) + expect(Object.getPrototypeOf(tools)).toBe(Object.prototype) + + expect(clientToolsFromAISDKBody({ id: 'x' }, { weather: server })).toEqual({ + weather: server, + }) + expect(() => clientToolsFromAISDKBody({ clientTools: 'all' }, {})).toThrow() + expect(() => + clientToolsFromAISDKBody({ clientTools: [{ name: 'x' }] }, {}), + ).toThrow() + }) +}) diff --git a/packages/ai-ui-message-stream/tests/reconnect.test.ts b/packages/ai-ui-message-stream/tests/reconnect.test.ts new file mode 100644 index 0000000000..39b0ca6291 --- /dev/null +++ b/packages/ai-ui-message-stream/tests/reconnect.test.ts @@ -0,0 +1,67 @@ +import { describe, expect, it } from 'vitest' +import { StreamReconnectLimitError } from '@tanstack/ai-client' +import { fromChatTransport } from '../src' +import type { ChatTransport, UIMessage, UIMessageChunk } from 'ai' + +const REPLY: Array = [ + { type: 'start', messageId: 'a-1' }, + { type: 'text-start', id: 't' }, + { type: 'text-delta', id: 't', delta: 'One ' }, + { type: 'text-delta', id: 't', delta: 'two ' }, + { type: 'text-delta', id: 't', delta: 'three.' }, + { type: 'text-end', id: 't' }, + { type: 'finish', finishReason: 'stop' }, +] + +/** The reply from the start, failing after `dropAfter` chunks. */ +function replay(dropAfter = Infinity): ReadableStream { + let index = 0 + return new ReadableStream({ + pull(controller) { + if (index === dropAfter) controller.error(new TypeError('network error')) + else if (index === REPLY.length) controller.close() + else controller.enqueue(REPLY[index++]!) + }, + }) +} + +async function run( + drops: Array, + reconnect: { maxAttempts?: number } = {}, +): Promise<{ text: string; reconnects: number }> { + let reconnects = 0 + const transport: ChatTransport = { + resumeStreamIsReplay: true, + sendMessages: () => Promise.resolve(replay(drops[0])), + reconnectToStream: () => { + reconnects++ + return Promise.resolve(replay(drops[reconnects])) + }, + } + const connection = fromChatTransport(transport, { + reconnect: { initialDelayMs: 1, ...reconnect }, + }) + let text = '' + for await (const event of connection.connect([], undefined, undefined, { + threadId: 'chat-1', + runId: 'run-1', + })) { + if (event.type === 'TEXT_MESSAGE_CONTENT') text += event.delta + } + return { text, reconnects } +} + +describe('fromChatTransport reconnect', () => { + it('reconnects with backoff and skips the replayed prefix', async () => { + // Drops after 3 chunks, then the first replay drops after 4. + const result = await run([3, 4]) + expect(result).toEqual({ text: 'One two three.', reconnects: 2 }) + }) + + it('stops at the limit when reconnects make no progress', async () => { + // Every replay drops before it passes the chunks already delivered. + await expect(run([3, 2, 2, 2], { maxAttempts: 2 })).rejects.toBeInstanceOf( + StreamReconnectLimitError, + ) + }) +}) diff --git a/skills/tanstack-ai-migration/SKILL.md b/skills/tanstack-ai-migration/SKILL.md index 839df34e61..12528d384a 100644 --- a/skills/tanstack-ai-migration/SKILL.md +++ b/skills/tanstack-ai-migration/SKILL.md @@ -46,7 +46,9 @@ place to stay: finish the port of the other side, then remove it. https://tanstack.com/ai/latest/docs/migration/ai-sdk-client-tanstack-server - Client first: keep the `streamText` route and pass `fromChatTransport(new DefaultChatTransport({ api }), { threadId, tools })` - as the `connection` of TanStack `useChat`. + as the `connection` of TanStack `useChat`. On the route, + `clientToolsFromAISDKBody(body, serverTools)` reads the client's tools, so + they need no copy there. https://tanstack.com/ai/latest/docs/migration/tanstack-client-ai-sdk-server Tool approvals, client tools, activity, subagents and stream resume cross the From a5764368075859af771c3721adc85eedd5b0c229 Mon Sep 17 00:00:00 2001 From: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Date: Fri, 23 Oct 2026 16:25:17 +1100 Subject: [PATCH 2/2] docs(ai-ui-message-stream): note that text in a tool step breaks structured output --- docs/migration/tanstack-client-ai-sdk-server.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/migration/tanstack-client-ai-sdk-server.md b/docs/migration/tanstack-client-ai-sdk-server.md index 4930d5e8af..edd0cae07a 100644 --- a/docs/migration/tanstack-client-ai-sdk-server.md +++ b/docs/migration/tanstack-client-ai-sdk-server.md @@ -221,7 +221,7 @@ export function Recipe() { } ``` -If the text is not valid JSON for the schema, the run ends with an error. +If the text is not valid JSON for the schema, the run ends with an error. All text of the reply counts as the answer, so a route that also writes text in a step that calls tools fails the check. ## 6. Load a saved chat