From 114416ce1eeddfa795325f330edbb4ba90955f5c Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 14 Jul 2026 14:53:48 +0200 Subject: [PATCH 01/13] feat: multi-capability assistants (defineAssistant + useAssistant) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Declare an assistant's capabilities once on the server and consume them from one typed client hook, over a single endpoint. Server (`@tanstack/ai/assistant`) - `defineAssistant(config)`: an inert registry of per-capability callbacks (chat, image, audio, speech, video, transcription, summarize), each invoking an existing activity. Exposes one `handler(request)` that parses the AG-UI request, routes by a `capability` discriminator, and streams the result back over the existing SSE wire. Nothing is constructed until a request arrives. - Every callback — chat and one-shots alike — receives the same validated AG-UI envelope: `chat` gets `AssistantChatRequest`; one-shots get their input plus `threadId`/`runId`/`parentRunId`/`state`/`aguiContext`/ `forwardedProps` and the raw `request`, all parsed through the same request validator (`chatParamsFromRequestBody`). Untrusted input fields (e.g. `req.prompt`) stay `unknown` so callbacks narrow before use. Client core (`@tanstack/ai-client/assistant`) - `AssistantClient` composes one `ChatClient` + one `GenerationClient` per declared capability behind a single connection. `AssistantSystem` maps declared capabilities to fully-typed surfaces, inferred from the definition — no call-site generics. - The imperative methods resolve to their results, so capabilities chain with a plain `await` instead of reading reactive state (which is stale in an async closure under React): `assistant..generate(input)` resolves to the generation result (`TResult | null`), and `assistant.chat.sendMessage(...)` resolves to the validated structured `final` when the chat callback declared an `outputSchema`, otherwise to the messages array. Reactive `.result`/`.messages`/`.partial`/`.final` remain for rendering. Framework hooks - `useAssistant` for React / Solid / Vue and `createAssistant` for Svelte (`@tanstack/ai-/assistant`). Each capability entry is the corresponding primitive hook's own return (`useChat` / `useGeneration`). Chat type inference - `chat()`'s return carries an optional type-only phantom (`ChatResultMeta`), so `assistant.chat` infers typed `partial`/`final` from an `outputSchema` and types message tool-call parts from the callback's tools — no client re-declaration. Runtime unchanged. Client-executed tools still pass through `chat: { tools }` for their browser implementations (their code can't cross the wire). Docs (`docs/assistant`): a 3-page section — Overview (what an assistant is, when to reach for one vs. single hooks, and typed chat/tools/structured output), Scenarios (sequential pipeline + CMS content workflow, chained via the resolved `await` returns), and Multiple Assistants (specialized per-surface assistants across routes, plus composing two assistants together on one page). Plus an `ai-core/assistant` agent skill, E2E coverage (chat verified end-to-end; image one-shot skipped pending an aimock media fixture), and a changeset. --- .changeset/assistant-multi-capability.md | 10 + docs/assistant/multiple-assistants.md | 280 +++++++++++++ docs/assistant/overview.md | 298 ++++++++++++++ docs/assistant/scenarios.md | 252 ++++++++++++ docs/config.json | 23 ++ packages/ai-client/package.json | 4 + packages/ai-client/src/assistant-client.ts | 79 ++++ .../ai-client/src/assistant-structured.ts | 46 +++ packages/ai-client/src/assistant-types.ts | 356 +++++++++++++++++ packages/ai-client/src/assistant.ts | 9 + .../ai-client/tests/assistant-client.test.ts | 165 ++++++++ .../tests/assistant-structured.test.ts | 112 ++++++ packages/ai-client/vite.config.ts | 2 +- packages/ai-react/package.json | 4 + packages/ai-react/src/assistant.ts | 1 + packages/ai-react/src/use-assistant.ts | 251 ++++++++++++ .../ai-react/tests/use-assistant.test.tsx | 114 ++++++ packages/ai-react/vite.config.ts | 2 +- packages/ai-solid/package.json | 4 + packages/ai-solid/src/assistant.ts | 1 + packages/ai-solid/src/use-assistant.ts | 335 ++++++++++++++++ .../ai-solid/tests/use-assistant.test.tsx | 180 +++++++++ packages/ai-solid/tsdown.config.ts | 2 +- packages/ai-svelte/package.json | 5 + packages/ai-svelte/src/assistant.ts | 1 + .../ai-svelte/src/create-assistant.svelte.ts | 306 +++++++++++++++ .../tests/create-assistant.svelte.test.ts | 102 +++++ packages/ai-vue/package.json | 4 + packages/ai-vue/src/assistant.ts | 1 + packages/ai-vue/src/use-assistant.ts | 259 ++++++++++++ packages/ai-vue/tests/use-assistant.test.ts | 152 ++++++++ packages/ai-vue/tsdown.config.ts | 2 +- packages/ai/package.json | 4 + packages/ai/skills/ai-core/assistant/SKILL.md | 367 ++++++++++++++++++ .../skills/ai-core/chat-experience/SKILL.md | 4 + .../custom-backend-integration/SKILL.md | 4 + .../skills/ai-core/media-generation/SKILL.md | 4 + packages/ai/src/activities/assistant/index.ts | 131 +++++++ packages/ai/src/activities/assistant/types.ts | 118 ++++++ packages/ai/src/activities/chat/index.ts | 33 +- packages/ai/src/assistant.ts | 13 + packages/ai/src/index.ts | 10 + packages/ai/src/types.ts | 34 ++ .../tests/activities/assistant/index.test.ts | 140 +++++++ .../activities/chat/chat-result-meta.test.ts | 97 +++++ packages/ai/vite.config.ts | 1 + testing/e2e/fixtures/assistant/basic.json | 10 + testing/e2e/src/lib/feature-support.ts | 1 + testing/e2e/src/lib/features.ts | 4 + testing/e2e/src/lib/types.ts | 2 + testing/e2e/src/routeTree.gen.ts | 42 ++ testing/e2e/src/routes/api.assistant.ts | 52 +++ testing/e2e/src/routes/assistant.tsx | 128 ++++++ testing/e2e/tests/assistant.spec.ts | 62 +++ 54 files changed, 4614 insertions(+), 9 deletions(-) create mode 100644 .changeset/assistant-multi-capability.md create mode 100644 docs/assistant/multiple-assistants.md create mode 100644 docs/assistant/overview.md create mode 100644 docs/assistant/scenarios.md create mode 100644 packages/ai-client/src/assistant-client.ts create mode 100644 packages/ai-client/src/assistant-structured.ts create mode 100644 packages/ai-client/src/assistant-types.ts create mode 100644 packages/ai-client/src/assistant.ts create mode 100644 packages/ai-client/tests/assistant-client.test.ts create mode 100644 packages/ai-client/tests/assistant-structured.test.ts create mode 100644 packages/ai-react/src/assistant.ts create mode 100644 packages/ai-react/src/use-assistant.ts create mode 100644 packages/ai-react/tests/use-assistant.test.tsx create mode 100644 packages/ai-solid/src/assistant.ts create mode 100644 packages/ai-solid/src/use-assistant.ts create mode 100644 packages/ai-solid/tests/use-assistant.test.tsx create mode 100644 packages/ai-svelte/src/assistant.ts create mode 100644 packages/ai-svelte/src/create-assistant.svelte.ts create mode 100644 packages/ai-svelte/tests/create-assistant.svelte.test.ts create mode 100644 packages/ai-vue/src/assistant.ts create mode 100644 packages/ai-vue/src/use-assistant.ts create mode 100644 packages/ai-vue/tests/use-assistant.test.ts create mode 100644 packages/ai/skills/ai-core/assistant/SKILL.md create mode 100644 packages/ai/src/activities/assistant/index.ts create mode 100644 packages/ai/src/activities/assistant/types.ts create mode 100644 packages/ai/src/assistant.ts create mode 100644 packages/ai/tests/activities/assistant/index.test.ts create mode 100644 packages/ai/tests/activities/chat/chat-result-meta.test.ts create mode 100644 testing/e2e/fixtures/assistant/basic.json create mode 100644 testing/e2e/src/routes/api.assistant.ts create mode 100644 testing/e2e/src/routes/assistant.tsx create mode 100644 testing/e2e/tests/assistant.spec.ts diff --git a/.changeset/assistant-multi-capability.md b/.changeset/assistant-multi-capability.md new file mode 100644 index 000000000..c02ea564b --- /dev/null +++ b/.changeset/assistant-multi-capability.md @@ -0,0 +1,10 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-client': minor +'@tanstack/ai-react': minor +'@tanstack/ai-solid': minor +'@tanstack/ai-vue': minor +'@tanstack/ai-svelte': minor +--- + +Add defineAssistant (server) and useAssistant (client) for multi-capability assistants. diff --git a/docs/assistant/multiple-assistants.md b/docs/assistant/multiple-assistants.md new file mode 100644 index 000000000..ae865ba10 --- /dev/null +++ b/docs/assistant/multiple-assistants.md @@ -0,0 +1,280 @@ +--- +title: Multiple Assistants +id: multiple-assistants +order: 3 +description: "Give each product surface its own assistant — a support page with chat only, a content studio with chat + image + speech — each with its own route, its own useAssistant, and its own system prompt. Learn when to split into several assistants versus one broad one." +keywords: + - tanstack ai + - assistant + - defineAssistant + - useAssistant + - multiple assistants + - system prompt + - routes +--- + +**One `defineAssistant` per product surface.** An assistant bundles a *fixed* set of capabilities behind one endpoint with one system prompt. When different pages or features want different capability sets or different instructions, don't stretch a single assistant to cover them all — define one per surface. Each gets its own route and its own `useAssistant` on its own page. + +A typical app ends up with a few: + +| Surface | Capabilities | Why its own assistant | +|---|---|---| +| Customer **support** page | `chat` only | Tight, support-flavored system prompt; no image/speech to expose | +| **Content studio** page | `chat` + `image` + `speech` | A creative workflow that chains capabilities | +| Internal **dashboard** | `chat` + `summarize` | Different auth boundary; different instructions | + +## Two assistants, two routes + +Give each assistant its own module and its own handler. A small helper keeps the shared adapter config — the model choice and any provider options — in one place, so both assistants stay in sync without duplicating it: + +```ts +// api/assistants.ts +import { chat, generateImage, generateSpeech } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiImage, openaiSpeech, openaiText } from '@tanstack/ai-openai' + +// One place to change the model / provider options for every assistant. +const textModel = () => openaiText('gpt-5.5') + +// Support: chat only, with a support-flavored system prompt. +export const supportAssistant = defineAssistant({ + chat: (req) => + chat({ + adapter: textModel(), + messages: req.messages, + systemPrompts: [ + 'You are a customer-support agent for Acme. Be concise and friendly.', + ], + }), +}) + +// Studio: chat + image + speech, for a creative workflow. +export const studioAssistant = defineAssistant({ + chat: (req) => + chat({ + adapter: textModel(), + messages: req.messages, + systemPrompts: ['You are a creative copywriter.'], + }), + image: (req) => { + if (typeof req.prompt !== 'string') { + throw new Error('image prompt must be a string') + } + return generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: req.prompt, + }) + }, + speech: (req) => + generateSpeech({ adapter: openaiSpeech('tts-1'), text: req.text }), +}) + +// Two routes — one per assistant. +export const supportPOST = (request: Request) => supportAssistant.handler(request) +export const studioPOST = (request: Request) => studioAssistant.handler(request) +``` + +In a real app these are two route files — `/api/support` and `/api/studio` — each exporting its own `POST`. The `supportPOST` / `studioPOST` names above just let both live in one snippet. + +## Two pages, two clients + +Each page calls `useAssistant` against *its own* assistant and *its own* connection. The support page only ever sees `assistant.chat`, because that's all `supportAssistant` declared: + +```tsx +// pages/SupportPage.tsx +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { chat } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiText } from '@tanstack/ai-openai' + +const textModel = () => openaiText('gpt-5.5') + +// The same object your /api/support route exports. +const supportAssistant = defineAssistant({ + chat: (req) => + chat({ + adapter: textModel(), + messages: req.messages, + systemPrompts: [ + 'You are a customer-support agent for Acme. Be concise and friendly.', + ], + }), +}) + +function SupportPage() { + const assistant = useAssistant(supportAssistant, { + connection: fetchServerSentEvents('/api/support'), + }) + + return ( +
+ + {assistant.chat.messages.map((message) => ( +

+ {message.parts.find((part) => part.type === 'text')?.content} +

+ ))} +
+ ) +} +``` + +The studio page points at `/api/studio` and gets `assistant.chat`, `assistant.image`, and `assistant.speech`: + +```tsx +// pages/StudioPage.tsx +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { chat, generateImage, generateSpeech } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiImage, openaiSpeech, openaiText } from '@tanstack/ai-openai' + +const textModel = () => openaiText('gpt-5.5') + +// The same object your /api/studio route exports. +const studioAssistant = defineAssistant({ + chat: (req) => + chat({ + adapter: textModel(), + messages: req.messages, + systemPrompts: ['You are a creative copywriter.'], + }), + image: (req) => { + if (typeof req.prompt !== 'string') { + throw new Error('image prompt must be a string') + } + return generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: req.prompt, + }) + }, + speech: (req) => + generateSpeech({ adapter: openaiSpeech('tts-1'), text: req.text }), +}) + +function StudioPage() { + const assistant = useAssistant(studioAssistant, { + connection: fetchServerSentEvents('/api/studio'), + }) + + return ( +
+ + + {assistant.image.result?.images[0]?.url && ( + + )} +
+ ) +} +``` + +Each page's `assistant` object is typed to exactly its assistant's capabilities — `SupportPage` has no `assistant.image` to misuse, and there's no way to accidentally call the support route with an image request. + +## Two assistants on one page + +Multiple assistants aren't only for *separate* pages. You can co-locate two **specialized** assistants in the same component — each keeping its own route, connection, auth boundary, and type surface — and compose them in one UI. The glue is the resolved `await` returns: every method hands its result straight back, so one assistant's output flows into the other's input without any reactive-state juggling. + +Here a campaign builder pairs a **structured-output copy** assistant (its own `/api/copy` route) with a **media** assistant (its own `/api/media` route): + +```tsx +// pages/CampaignBuilder.tsx +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { chat, generateImage, generateSpeech } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiImage, openaiSpeech, openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +const TaglineSchema = z.object({ headline: z.string(), subhead: z.string() }) + +// Copy: structured-output chat, its own /api/copy route. +const copyAssistant = defineAssistant({ + chat: (req) => + chat({ + adapter: openaiText('gpt-5.5'), + messages: req.messages, + outputSchema: TaglineSchema, + stream: true, + }), +}) + +// Media: image + speech, its own /api/media route. +const mediaAssistant = defineAssistant({ + image: (req) => { + if (typeof req.prompt !== 'string') { + throw new Error('image prompt must be a string') + } + return generateImage({ adapter: openaiImage('gpt-image-2'), prompt: req.prompt }) + }, + speech: (req) => + generateSpeech({ adapter: openaiSpeech('tts-1'), text: req.text }), +}) + +function CampaignBuilder() { + const copy = useAssistant(copyAssistant, { + connection: fetchServerSentEvents('/api/copy'), + }) + const media = useAssistant(mediaAssistant, { + connection: fetchServerSentEvents('/api/media'), + }) + + async function build(brief: string) { + // One assistant drafts the copy... + const tagline = await copy.chat.sendMessage(`Write a tagline for: ${brief}`) + if (!tagline) return + // ...the other turns it into art and audio. Two routes, one workflow. + await media.image.generate({ prompt: tagline.headline }) + await media.speech.generate({ text: tagline.subhead }) + } + + return ( +
+ +

{copy.chat.partial.headline ?? 'Drafting…'}

+ {media.image.result?.images[0]?.url && ( + + )} +
+ ) +} +``` + +Each assistant keeps its own connection, route, auth boundary, and type surface: `copy` has no `.image` and `media` has no `.chat`, so neither can misuse the other's capabilities or route. Yet they compose cleanly in one component, because each method resolves to its result — `copy.chat.sendMessage` hands back the validated `{ headline, subhead }` (its callback declared `outputSchema`), which threads directly into `media.image.generate` and `media.speech.generate`. The reactive fields (`copy.chat.partial`, `media.image.result`) are still there for rendering. + +Contrast this with a single broad assistant that declared all three capabilities behind **one** route: that's the right choice when the capabilities truly share a connection and boundary — see [Scenarios](./scenarios) for that one-assistant chaining shape. Reach for two assistants on one page when the capabilities need to appear together but have **different boundaries** — separate routes, auth, or models — that you don't want to collapse into one endpoint. + +## Split, or one broad assistant? + +Both are valid. Choose by how the capabilities actually relate: + +**Split into multiple assistants when:** + +- **Different product surfaces.** A support widget and a content studio are different features with different users — separate assistants keep their prompts, capabilities, and routes independent. +- **Different capability sets.** If support never needs image generation, don't expose it there. A narrower assistant is a narrower attack surface and a simpler client type. +- **Different auth or rate-limit boundaries.** Separate routes let you guard `/api/support` and `/api/studio` independently. + +**Use one broad assistant when:** + +- **The capabilities are always used together** on the same page or in the same workflow — that's exactly the [pipeline / content-workflow](./scenarios) shape, where one capability feeds the next. +- **They share a system prompt and auth boundary.** If splitting would just duplicate the same configuration twice, keep it as one. + +**Co-locate multiple assistants on one page when:** the capabilities need to appear together in the same UI but have **different boundaries** — separate routes, auth, or models. This is the middle ground between the two options above: separate assistants (so each keeps its own boundary and type surface), composed in one component via their resolved `await` returns (see [Two assistants on one page](#two-assistants-on-one-page)). + +The rule of thumb: **split by surface, not by capability count.** A single page that happens to use three capabilities is one assistant; three pages that each use chat are three assistants. And when one page needs two *different* boundaries side by side, that's two assistants on one page. + +## Next + +- [Assistant Overview](./overview) — capabilities, typing, and when to reach for an assistant at all. +- [Scenarios](./scenarios) — chaining capabilities within a single assistant. diff --git a/docs/assistant/overview.md b/docs/assistant/overview.md new file mode 100644 index 000000000..dd532c668 --- /dev/null +++ b/docs/assistant/overview.md @@ -0,0 +1,298 @@ +--- +title: Assistant Overview +id: assistant-overview +order: 1 +description: "An assistant bundles several AI capabilities — chat, image, speech, and more — behind ONE endpoint and hands you ONE typed client. Learn when to reach for an assistant instead of the single-capability hooks, and how to define and drive one." +keywords: + - tanstack ai + - assistant + - defineAssistant + - useAssistant + - multi-capability + - chat + - image generation + - text-to-speech +--- + +**An assistant bundles multiple AI capabilities — `chat`, `image`, `speech`, and more — behind a single request handler on the server, and gives you back a single typed client on the browser.** You declare the capabilities once with `defineAssistant()`; `useAssistant()` returns one `assistant` object keyed by exactly the capabilities you declared, all sharing one connection. + +`defineAssistant()` lives on the `/assistant` subpath of `@tanstack/ai`, and `useAssistant()` on the `/assistant` subpath of `@tanstack/ai-react` (Solid, Vue, and Svelte follow the same pattern) — not the package root. + +```ts +// api/assistant.ts +import { chat, generateImage } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiImage, openaiText } from '@tanstack/ai-openai' + +export const blogAssistant = defineAssistant({ + chat: (req) => + chat({ adapter: openaiText('gpt-5.5'), messages: req.messages }), + image: (req) => { + // `req.prompt` arrives as `unknown` — narrow it before use. + if (typeof req.prompt !== 'string') { + throw new Error('image prompt must be a string') + } + return generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: req.prompt, + }) + }, +}) + +export const POST = (request: Request) => blogAssistant.handler(request) +``` + +That one route now serves both chat and image generation. On the client, one hook drives both: + +```tsx +// components/Assistant.tsx +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { chat, generateImage } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiImage, openaiText } from '@tanstack/ai-openai' + +// The same object your server route exports — share it from one module in a +// real app; repeated here so this snippet type-checks on its own. +const blogAssistant = defineAssistant({ + chat: (req) => + chat({ adapter: openaiText('gpt-5.5'), messages: req.messages }), + image: (req) => { + if (typeof req.prompt !== 'string') { + throw new Error('image prompt must be a string') + } + return generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: req.prompt, + }) + }, +}) + +function Assistant() { + const assistant = useAssistant(blogAssistant, { + connection: fetchServerSentEvents('/api/assistant'), + }) + + return ( +
+ + {assistant.chat.messages.map((message) => ( +

+ {message.parts.find((part) => part.type === 'text')?.content} +

+ ))} + + + {assistant.image.result?.images[0]?.url && ( + + )} +
+ ) +} +``` + +`assistant.chat` behaves exactly like [`useChat`](../chat/streaming) — streaming, tool calls, and message parts all work the same way. Each one-shot capability (`assistant.image`, `assistant.speech`, ...) behaves like its matching single hook ([`useGenerateImage`](../media/image-generation#hook-api), and so on): a `.generate(input)` method plus reactive `.result`, `.isLoading`, and `.error`. + +Both `assistant..generate(input)` and `assistant.chat.sendMessage(...)` **resolve to their result**, so you can chain them with a single `await` — `generate` resolves to the generation result (or `null`), and `sendMessage` resolves to the structured `final` (when the chat callback set `outputSchema`) or the messages array otherwise (`const result = await assistant.image.generate(...)`). The reactive `.result` / `.messages` / `.partial` / `.final` remain for rendering. → Chaining these awaited returns is exactly where an assistant pays off; see [Scenarios](./scenarios). + +## When should you reach for an assistant? + +An assistant is not "a better `useChat`." It earns its place only when you have **two or more capabilities that live on the same surface** — a page, a workflow, a pipeline. If you need exactly one capability, the single hook is simpler and there is nothing to gain from the extra layer. + +| Your situation | Reach for | Why | +|---|---|---| +| One page needs chat **and** image **and** speech | An assistant | One connection, one client object, one route to deploy and secure | +| One capability feeds the next (generate an image → write copy about it → narrate it) | An assistant | Typed handoffs between capabilities on the same `assistant` object | +| A multi-step content pipeline behind one endpoint | An assistant | The whole workflow is one deployable unit with one auth boundary | +| A page that only sends chat messages | [`useChat`](../chat/streaming) | No second capability to share — the assistant layer adds nothing | +| A page that only generates images | [`useGenerateImage`](../media/generation-hooks) | Same — a single hook is the whole story | + +**The honest tradeoff:** an assistant centralizes routing and typing across capabilities, but it also couples them into one definition and one endpoint. If two features never appear together and want different auth or system prompts, don't force them into one assistant — give each its own (see [Multiple Assistants](./multiple-assistants)). Reach for an assistant when capabilities *share* something (a connection, a page, or each other's output); keep them separate when they don't. + +## Defining capabilities (server) + +Each capability is a plain callback: `(req) => `. The `chat` capability returns the stream `chat()` produces; one-shot capabilities (`image`, `speech`, `audio`, `video`, `transcription`, `summarize`) return either a `Promise` of their result or — with `stream: true` — an `AsyncIterable`. Every key is optional; the client only ever sees the capabilities you declare. + +`defineAssistant()` itself is inert: none of these callbacks run, and no adapter is constructed, until a request actually reaches `blogAssistant.handler`. The handler discriminates each incoming request by capability — routed there by the client — parses it into the matching request shape (`AssistantChatRequest`, `AssistantImageRequest`, `AssistantSpeechRequest`, ...), and streams whatever the callback returns back over Server-Sent Events. + +Because one-shot request fields such as `req.prompt` arrive as `unknown` from the wire, narrow them with a `typeof` check (as above) before handing them to the adapter — never assume the shape. + +## Typed chat: tools and structured output + +The `chat` capability inherits `useChat`'s full type inference, driven entirely by what the server callback passes to `chat()` — you don't re-declare anything on the client. + +**Callback tools auto-type `assistant.chat.messages`.** Pass `tools` to `chat()` in the server callback and the tool-call/result parts on `assistant.chat.messages` are typed from those tools, with nothing to register on the client: + +```tsx +// components/WeatherChat.tsx +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { chat, toolDefinition } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +const getWeatherDef = toolDefinition({ + name: 'get_weather', + description: 'Get the current weather for a city', + inputSchema: z.object({ city: z.string() }), + outputSchema: z.object({ tempF: z.number(), conditions: z.string() }), +}) + +const getWeather = getWeatherDef.server(async ({ city }) => { + return { tempF: 72, conditions: `Sunny in ${city}` } +}) + +// The same object your server route exports — share it from one module in a +// real app; repeated here so this snippet type-checks on its own. +const blogAssistant = defineAssistant({ + chat: (req) => + chat({ + adapter: openaiText('gpt-5.5'), + messages: req.messages, + tools: [getWeather], + }), +}) + +function WeatherChat() { + const assistant = useAssistant(blogAssistant, { + connection: fetchServerSentEvents('/api/assistant'), + }) + + const weatherCall = assistant.chat.messages + .at(-1) + ?.parts.find( + (part) => part.type === 'tool-call' && part.name === 'get_weather', + ) + + return ( +
+ + {weatherCall?.type === 'tool-call' && weatherCall.output && ( +

{weatherCall.output.conditions}

+ )} +
+ ) +} +``` + +No `chat: { tools }` option was passed to `useAssistant` — `weatherCall.output` is still narrowed to `{ tempF: number; conditions: string }`, inferred entirely from `tools: [getWeather]` on the server callback. + +**`chat: { tools }` is optional — you need it only for client-executed tools.** A [client tool](../tools/client-tools)'s `.client()` implementation runs in the browser, so its code can't cross the wire. The server callback sees only the tool's *definition* (for typing and to tell the model it exists); the client registers the runtime implementation via `chat: { tools }`. Types still come from the server callback either way: + +```tsx +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { chat, toolDefinition } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +const showToastDef = toolDefinition({ + name: 'show_toast', + description: 'Show a browser notification', + inputSchema: z.object({ message: z.string() }), +}) + +const showToast = showToastDef.client((input) => { + console.log(input.message) + return { ok: true } +}) + +// The same object your server route exports — share it from one module in a +// real app; repeated here so this snippet type-checks on its own. +const blogAssistant = defineAssistant({ + chat: (req) => + chat({ + adapter: openaiText('gpt-5.5'), + messages: req.messages, + tools: [showToastDef], // definition only — the client executes it + }), +}) + +function useAssistantWithClientTool() { + return useAssistant(blogAssistant, { + connection: fetchServerSentEvents('/api/assistant'), + chat: { tools: [showToast] }, + }) +} +``` + +**`outputSchema` → typed `assistant.chat.partial` / `.final`.** If the `chat` callback passes an `outputSchema` to `chat()`, `assistant.chat` picks up a progressively-parsed `partial` (`DeepPartial`) and a validated terminal `final` — the same inference [`useChat({ outputSchema })`](../structured-outputs/streaming) gives you: + +```tsx +// components/BlogOutlineForm.tsx +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { chat } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +const BlogOutlineSchema = z.object({ + title: z.string(), + sections: z.array(z.string()), +}) + +// The same object your server route exports — share it from one module in a +// real app; repeated here so this snippet type-checks on its own. +const blogAssistant = defineAssistant({ + chat: (req) => + chat({ + adapter: openaiText('gpt-5.5'), + messages: req.messages, + outputSchema: BlogOutlineSchema, + stream: true, + }), +}) + +function BlogOutlineForm() { + const assistant = useAssistant(blogAssistant, { + connection: fetchServerSentEvents('/api/assistant'), + }) + + return ( +
+ +

Title: {assistant.chat.partial.title ?? '…'}

+
    + {assistant.chat.partial.sections?.map((section, i) => ( +
  • {section}
    • + ))} +
+ {assistant.chat.final && ( + 
{JSON.stringify(assistant.chat.final, null, 2)}
+ )} +
+ ) +} +``` + +`partial` and `final` only exist on the type when the `chat` callback declares `outputSchema` — omit it and `assistant.chat` has no such fields, exactly like `useChat` without `outputSchema`. See [Structured Outputs](../structured-outputs/overview) for the full story. + +## Which page do I read? + +| You want to… | Read | +|---|---| +| Chain capabilities end-to-end — run one after another and thread each result into the next (a pipeline, a content workflow) | [Scenarios](./scenarios) | +| Give different pages or features their own capability sets, routes, and system prompts | [Multiple Assistants](./multiple-assistants) | +| Understand the underlying chat surface (streaming, tools, message parts) | [Chat: Streaming](../chat/streaming) | +| Drive a single capability without an assistant | [Media & Generation Hooks](../media/generation-hooks) | diff --git a/docs/assistant/scenarios.md b/docs/assistant/scenarios.md new file mode 100644 index 000000000..28ac3171f --- /dev/null +++ b/docs/assistant/scenarios.md @@ -0,0 +1,252 @@ +--- +title: Assistant Scenarios +id: assistant-scenarios +order: 2 +description: "Journey-style recipes for chaining assistant capabilities: a sequential pipeline (image → chat → speech) and a one-endpoint content workflow (draft → hero image → audio version). See why an assistant beats wiring separate hooks when steps feed each other." +keywords: + - tanstack ai + - assistant + - useAssistant + - pipeline + - content workflow + - chaining capabilities + - multimodal +--- + +**These are point A → point B recipes.** Each one takes a concrete goal and walks the full flow across an assistant's capabilities. The thread running through both: because every capability hangs off the same `assistant` object, the output of one step is right there — typed — for the next step to consume. Wiring the same flow with separate `useChat` / `useGenerateImage` / `useGenerateSpeech` hooks means three connections, three client objects, and manually shuttling values between them. + +If you only need one capability, you don't need any of this — reach for the single hook. See [When should you reach for an assistant?](./overview#when-should-you-reach-for-an-assistant) on the overview. + +## Recipe 1 — Sequential pipeline + +**Goal:** generate an image, have the model write copy *about* that image, then narrate the copy aloud — each step awaiting the last and threading its result forward. + +Declare all three capabilities in one assistant on the server: + +```ts +// api/assistant.ts +import { chat, generateImage, generateSpeech } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiImage, openaiSpeech, openaiText } from '@tanstack/ai-openai' + +export const studioAssistant = defineAssistant({ + chat: (req) => + chat({ adapter: openaiText('gpt-5.5'), messages: req.messages }), + image: (req) => { + if (typeof req.prompt !== 'string') { + throw new Error('image prompt must be a string') + } + return generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: req.prompt, + }) + }, + speech: (req) => + generateSpeech({ + adapter: openaiSpeech('tts-1'), + text: req.text, + voice: req.voice, + }), +}) + +export const POST = (request: Request) => studioAssistant.handler(request) +``` + +On the client, the whole pipeline is one function. Each `await` resolves when that step finishes, and its result is on the same `assistant` object for the next step to read: + +```tsx +// components/Pipeline.tsx +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { chat, generateImage, generateSpeech } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiImage, openaiSpeech, openaiText } from '@tanstack/ai-openai' + +// The same object your server route exports — share it from one module in a +// real app; repeated here so this snippet type-checks on its own. +const studioAssistant = defineAssistant({ + chat: (req) => + chat({ adapter: openaiText('gpt-5.5'), messages: req.messages }), + image: (req) => { + if (typeof req.prompt !== 'string') { + throw new Error('image prompt must be a string') + } + return generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: req.prompt, + }) + }, + speech: (req) => + generateSpeech({ + adapter: openaiSpeech('tts-1'), + text: req.text, + voice: req.voice, + }), +}) + +function Pipeline() { + const assistant = useAssistant(studioAssistant, { + connection: fetchServerSentEvents('/api/assistant'), + }) + + async function run(subject: string) { + // 1. Generate an image — the result comes straight back from the await. + const image = await assistant.image.generate({ prompt: subject }) + const imageUrl = image?.images[0]?.url + if (!imageUrl) return + + // 2. Hand the image to chat and ask for copy about it. + // sendMessage resolves to the messages (no outputSchema here). + const messages = await assistant.chat.sendMessage({ + content: [ + { type: 'image', source: { type: 'url', value: imageUrl } }, + { type: 'text', content: 'Write a short, punchy caption for this image.' }, + ], + }) + + // 3. Pull the reply text out of the returned messages and narrate it. + const caption = messages + .at(-1) + ?.parts.find((part) => part.type === 'text')?.content + if (!caption) return + + await assistant.speech.generate({ text: caption }) + } + + return ( +
+ + {assistant.image.result?.images[0]?.url && ( + + )} +
+ ) +} +``` + +**Why an assistant here:** each step's `await` hands back its own result — `const image = await assistant.image.generate(...)`, `const messages = await assistant.chat.sendMessage(...)` — and you thread that returned value straight into the next call. No reactive-state reads between steps: the image URL feeds chat as [multimodal input](../media/image-generation#image-conditioned-generation), and the chat reply (pulled from the returned `messages`) feeds speech — all typed, all on one connection. With three separate hooks you'd hand-carry each value across three independent clients and manage three connections. The reactive `assistant.image.result` / `assistant.chat.messages` are still there for *rendering* (see the JSX below); the pipeline just doesn't need them. See [Text-to-Speech](../media/text-to-speech#playing-audio-in-the-browser) for turning `assistant.speech.result?.audio` into playable audio. + +## Recipe 2 — Content workflow (CMS) + +**Goal:** a "content" assistant powering a mini CMS. From one prompt it (a) drafts a blog post as a **typed object** (`title` + `body`) via structured-output chat, (b) generates a hero image from the drafted title, and (c) produces an audio version of the body via speech — all behind **one endpoint**. + +The server declares the structured-output `chat`, plus `image` and `speech`: + +```ts +// api/content.ts +import { chat, generateImage, generateSpeech } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiImage, openaiSpeech, openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +const BlogPostSchema = z.object({ + title: z.string(), + body: z.string(), +}) + +export const contentAssistant = defineAssistant({ + chat: (req) => + chat({ + adapter: openaiText('gpt-5.5'), + messages: req.messages, + outputSchema: BlogPostSchema, + stream: true, + }), + image: (req) => { + if (typeof req.prompt !== 'string') { + throw new Error('image prompt must be a string') + } + return generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: req.prompt, + }) + }, + speech: (req) => + generateSpeech({ adapter: openaiSpeech('tts-1'), text: req.text }), +}) + +export const POST = (request: Request) => contentAssistant.handler(request) +``` + +The client orchestrates the workflow off the typed `final` object — `assistant.chat.final` is narrowed to `{ title: string; body: string } | null` because the `chat` callback declared `outputSchema`: + +```tsx +// components/ContentStudio.tsx +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { chat, generateImage, generateSpeech } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiImage, openaiSpeech, openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +const BlogPostSchema = z.object({ + title: z.string(), + body: z.string(), +}) + +// The same object your server route exports — share it from one module in a +// real app; repeated here so this snippet type-checks on its own. +const contentAssistant = defineAssistant({ + chat: (req) => + chat({ + adapter: openaiText('gpt-5.5'), + messages: req.messages, + outputSchema: BlogPostSchema, + stream: true, + }), + image: (req) => { + if (typeof req.prompt !== 'string') { + throw new Error('image prompt must be a string') + } + return generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: req.prompt, + }) + }, + speech: (req) => + generateSpeech({ adapter: openaiSpeech('tts-1'), text: req.text }), +}) + +function ContentStudio() { + const assistant = useAssistant(contentAssistant, { + connection: fetchServerSentEvents('/api/content'), + }) + + async function publish(topic: string) { + // (a) Draft the post as a typed object — sendMessage resolves to the + // validated `final` because the chat callback declared outputSchema. + const draft = await assistant.chat.sendMessage( + `Write a short blog post about ${topic}.`, + ) // { title: string; body: string } | null + if (!draft) return + + // (b) Generate a hero image from the drafted title. + await assistant.image.generate({ prompt: `Hero image for: ${draft.title}` }) + + // (c) Produce an audio version of the body. + await assistant.speech.generate({ text: draft.body }) + } + + return ( +
+ + {/* Live progress while the draft streams in */} +

{assistant.chat.partial.title ?? 'Drafting…'}

+

{assistant.chat.partial.body}

+ {assistant.image.result?.images[0]?.url && ( + {assistant.chat.final?.title + )} +
+ ) +} +``` + +**Why an assistant here:** `sendMessage` resolves to the validated structured `final` — `const draft = await assistant.chat.sendMessage(...)` hands you the typed `{ title, body }` object directly, which you thread into the next calls (`draft.title` becomes the image prompt, `draft.body` becomes the speech text). These are ordinary typed property reads on the awaited return, not values you serialized out of one hook and back into another, and not a reactive-state read after the await. The entire workflow ships as a single endpoint (`/api/content`) with one auth boundary and one connection, which is exactly the shape a content-management feature wants. Because the `chat` callback declares `outputSchema`, you also get `assistant.chat.partial` for a live preview while the draft streams — see [Structured Outputs: Streaming UIs](../structured-outputs/streaming). + +## Next + +- [Multiple Assistants](./multiple-assistants) — when different pages or features each need their own assistant. +- [Assistant Overview](./overview) — the full capability + typing reference. +- [Client Tools](../tools/client-tools) — run tool implementations in the browser inside `assistant.chat`. diff --git a/docs/config.json b/docs/config.json index 7debc1ba4..65dc9e66d 100644 --- a/docs/config.json +++ b/docs/config.json @@ -308,6 +308,29 @@ } ] }, + { + "label": "Assistant", + "children": [ + { + "label": "Overview", + "to": "assistant/overview", + "addedAt": "2026-07-13", + "updatedAt": "2026-07-14" + }, + { + "label": "Scenarios", + "to": "assistant/scenarios", + "addedAt": "2026-07-14", + "updatedAt": "2026-07-14" + }, + { + "label": "Multiple Assistants", + "to": "assistant/multiple-assistants", + "addedAt": "2026-07-14", + "updatedAt": "2026-07-14" + } + ] + }, { "label": "Middleware", "children": [ diff --git a/packages/ai-client/package.json b/packages/ai-client/package.json index 37dfbd1b2..55d348ad7 100644 --- a/packages/ai-client/package.json +++ b/packages/ai-client/package.json @@ -42,6 +42,10 @@ "./devtools": { "types": "./dist/esm/devtools.d.ts", "import": "./dist/esm/devtools.js" + }, + "./assistant": { + "types": "./dist/esm/assistant.d.ts", + "import": "./dist/esm/assistant.js" } }, "files": [ diff --git a/packages/ai-client/src/assistant-client.ts b/packages/ai-client/src/assistant-client.ts new file mode 100644 index 000000000..2f235fd73 --- /dev/null +++ b/packages/ai-client/src/assistant-client.ts @@ -0,0 +1,79 @@ +import { ChatClient } from './chat-client.js' +import { GenerationClient } from './generation-client.js' +import type { AssistantDefinition } from '@tanstack/ai/assistant' +import type { AnyClientTool } from '@tanstack/ai/client' +import type { + AssistantClientOptions, + OneShotCapabilityName, +} from './assistant-types.js' + +/** + * Composes one `ChatClient` and/or one `GenerationClient` per one-shot + * capability declared on an `AssistantDefinition`, sharing the single + * `connection` adapter passed in `options`. + * + * No connection wrapping happens here: each sub-client tags its own + * requests with the capability name — `ChatClient` via + * `forwardedProps: { capability: 'chat' }`, `GenerationClient` via + * `body: { capability: }` — so a single server endpoint can route + * on that field. + */ +export class AssistantClient< + TDef extends AssistantDefinition = AssistantDefinition, + TChatTools extends ReadonlyArray = [], +> { + /** The chat sub-client, present only if the definition declares `chat`. */ + chat?: ChatClient + private readonly oneShots = new Map>() + readonly capabilities: ReadonlyArray + + constructor(options: AssistantClientOptions) { + const { assistant, connection, id, threadId, chat, callbacks } = options + this.capabilities = assistant.capabilities + + for (const capability of assistant.capabilities) { + if (capability === 'chat') { + this.chat = new ChatClient({ + connection, + id: id ? `${id}:chat` : undefined, + threadId, + tools: chat?.tools, + forwardedProps: { ...chat?.forwardedProps, capability: 'chat' }, + ...callbacks?.chat, + }) + continue + } + + const oneShotCapability = capability as OneShotCapabilityName + this.oneShots.set( + oneShotCapability, + new GenerationClient({ + connection, + id: id ? `${id}:${oneShotCapability}` : undefined, + body: { capability: oneShotCapability }, + ...callbacks?.oneShot?.(oneShotCapability), + }), + ) + } + } + + /** Whether the assistant declares the given capability. */ + has(capability: string): boolean { + return this.capabilities.includes(capability) + } + + /** The one-shot `GenerationClient` for a declared capability, if any. */ + get( + capability: OneShotCapabilityName, + ): GenerationClient | undefined { + return this.oneShots.get(capability) + } + + /** Tears down the chat client and every one-shot client. */ + dispose(): void { + this.chat?.dispose() + for (const oneShot of this.oneShots.values()) { + oneShot.dispose() + } + } +} diff --git a/packages/ai-client/src/assistant-structured.ts b/packages/ai-client/src/assistant-structured.ts new file mode 100644 index 000000000..03b4b9d78 --- /dev/null +++ b/packages/ai-client/src/assistant-structured.ts @@ -0,0 +1,46 @@ +import type { UIMessage } from './types.js' + +/** + * Computes the "active" structured-output `partial`/`final` values from a + * `UIMessage` history — the framework-agnostic core of `useChat`'s + * `partial`/`final` fields, shared by the assistant hooks. + * + * The "active" structured-output part is the one on the assistant message + * that follows the latest user message. No such message exists between + * `sendMessage()` and the first chunk, so `partial`/`final` naturally read + * as cleared. Historical parts on earlier assistant messages remain + * available via `messages` directly. + * + * When there is NO user message yet (e.g. `initialMessages` contains only a + * stale assistant turn or a system prompt) this deliberately returns + * `{ partial: {}, final: null }` rather than scanning historical + * assistants — otherwise a `final` from a previous session would leak into + * the hook value on first render. + */ +export function computeStructuredParts(messages: ReadonlyArray): { + partial: unknown + final: unknown +} { + let lastUserIndex = -1 + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i]?.role === 'user') { + lastUserIndex = i + break + } + } + if (lastUserIndex === -1) return { partial: {}, final: null } + + for (let i = messages.length - 1; i > lastUserIndex; i--) { + const m = messages[i] + if (m?.role !== 'assistant') continue + const part = m.parts.find((p) => p.type === 'structured-output') + if (part) { + return { + partial: part.partial ?? part.data ?? {}, + final: part.status === 'complete' ? part.data : null, + } + } + } + + return { partial: {}, final: null } +} diff --git a/packages/ai-client/src/assistant-types.ts b/packages/ai-client/src/assistant-types.ts new file mode 100644 index 000000000..e2594b136 --- /dev/null +++ b/packages/ai-client/src/assistant-types.ts @@ -0,0 +1,356 @@ +// packages/ai-client/src/assistant-types.ts +import type { + AudioGenerationResult, + ImageGenerationResult, + InferChatSchema, + InferChatTools, + SummarizationResult, + TTSResult, + TranscriptionResult, + VideoJobResult, +} from '@tanstack/ai' +import type { AssistantDefinition } from '@tanstack/ai/assistant' +import type { + AnyClientTool, + InferSchemaType, + ModelMessage, + SchemaInput, +} from '@tanstack/ai/client' +import type { ConnectConnectionAdapter } from './connection-adapters.js' +import type { + ChatClientOptions, + ChatClientState, + ConnectionStatus, + MultimodalContent, + UIMessage, +} from './types.js' +import type { + AudioGenerateInput, + GenerationClientOptions, + GenerationClientState, + ImageGenerateInput, + SpeechGenerateInput, + SummarizeGenerateInput, + TranscriptionGenerateInput, + VideoGenerateInput, +} from './generation-types.js' + +/** + * Reactive state callbacks forwarded into the underlying sub-clients. + * + * `ChatClient` and `GenerationClient` only accept these callbacks via their + * constructors (`updateOptions` does not accept them), so `AssistantClient` + * must thread them through at construction time rather than attaching them + * post-construction. + */ +export interface AssistantClientCallbacks { + /** Reactive state callbacks forwarded to the `ChatClient` constructor. */ + chat?: Pick< + ChatClientOptions, + | 'onMessagesChange' + | 'onLoadingChange' + | 'onErrorChange' + | 'onStatusChange' + | 'onSubscriptionChange' + | 'onConnectionStatusChange' + | 'onSessionGeneratingChange' + > + /** + * Reactive state callbacks forwarded to a one-shot capability's + * `GenerationClient` constructor. Invoked once per declared one-shot + * capability so each sub-client can be wired independently. + */ + oneShot?: ( + capability: OneShotCapabilityName, + ) => Pick< + GenerationClientOptions, + 'onResultChange' | 'onLoadingChange' | 'onErrorChange' | 'onStatusChange' + > +} + +/** Options for AssistantClient (framework-agnostic core). */ +export interface AssistantClientOptions< + TDef extends AssistantDefinition, + /** + * @deprecated Chat tool typing is now inferred from the assistant + * definition's chat callback (`TDef`), not from this generic. Retained only + * so the framework hooks (`ai-react`/`ai-solid`/`ai-vue`/`ai-svelte`), which + * still thread it through, keep compiling until they are updated. It no + * longer influences the chat surface's message/tool types. + */ + TChatTools extends ReadonlyArray = [], +> { + assistant: TDef + connection: ConnectConnectionAdapter + id?: string + threadId?: string + /** Chat-only options mirrored onto the underlying ChatClient. */ + chat?: { + tools?: TChatTools + forwardedProps?: Record + } + /** + * Reactive state callbacks forwarded into the sub-clients' constructors. + * Set by framework hooks (not users) to wire up reactive state. + */ + callbacks?: AssistantClientCallbacks +} + +/** Maps declared capability names to their client generate-input type. */ +export interface GenerateInputByCapability { + image: ImageGenerateInput + audio: AudioGenerateInput + speech: SpeechGenerateInput + video: VideoGenerateInput + transcription: TranscriptionGenerateInput + summarize: SummarizeGenerateInput +} + +/** Maps declared capability names to their result type. */ +export interface ResultByCapability { + image: ImageGenerationResult + audio: AudioGenerationResult + speech: TTSResult + video: VideoJobResult + transcription: TranscriptionResult + summarize: SummarizationResult +} + +export type OneShotCapabilityName = keyof GenerateInputByCapability + +/** + * Recursive partial — every property and every nested array element is + * optional. Types the in-flight `partial` value the chat surface exposes while + * a structured-output stream is still arriving (the JSON has shape but is + * incomplete). Mirrors the `DeepPartial` used by the framework `useChat` + * hooks. + */ +export type DeepPartial = + T extends ReadonlyArray + ? Array> + : T extends object + ? { [K in keyof T]?: DeepPartial } + : T + +/** + * Map a captured chat-callback tool (server / definition / client) to a + * `definition`-shaped object so it satisfies the `AnyClientTool` constraint on + * `UIMessage`/`ToolCallPart` **without** relaxing the core message types. + * + * The chat callback's `tools` are typically server (`.server()`) or bare + * definition tools, whose `__toolSide` is `'server'` / `'definition'` — + * `AnyClientTool` rejects `'server'`. We rebuild each tool as a + * `'definition'`-side object, preserving the exact fields the `Infer*` helpers + * and `ToolCallPartForTool` read (`name`, `inputSchema`, `outputSchema`, + * `needsApproval`). `description` is required by `AnyClientTool`'s underlying + * `Tool` shape, so it is kept as `string` (its value never feeds message + * typing). + */ +type ToClientToolShape = TTool extends { + name: infer TName extends string +} + ? { + __toolSide: 'definition' + name: TName + description: string + inputSchema?: TTool extends { inputSchema?: infer TIn } ? TIn : undefined + outputSchema?: TTool extends { outputSchema?: infer TOut } + ? TOut + : undefined + needsApproval?: TTool extends { needsApproval?: infer TNa } ? TNa : false + } + : never + +/** Map a captured tool tuple to a client-tool-shaped tuple (see above). */ +type ToClientTools = TTools extends readonly [ + infer THead, + ...infer TRest, +] + ? [ToClientToolShape, ...ToClientTools] + : [] + +/** Recover the chat callback's return type from an assistant definition. */ +type ChatCallbackReturn = + TDef extends AssistantDefinition + ? TCaps extends { chat: (...args: Array) => infer TRet } + ? TRet + : never + : never + +/** + * Client-tool-shaped tuple recovered from the chat callback's captured tools. + * Drives the `messages` tool-call/result part typing on the chat surface. + */ +export type ChatToolsOf = ToClientTools< + NonNullable>> +> + +/** Structured-output schema recovered from the chat callback (or `undefined`). */ +export type ChatSchemaOf = InferChatSchema> + +/** The base chat capability surface — mirrors the frameworks' useChat return. */ +export interface AssistantChatSurfaceBase< + TTools extends ReadonlyArray = [], +> { + /** Current messages in the conversation. */ + messages: Array> + + /** + * Append a message to the conversation. + */ + append: (message: ModelMessage | UIMessage) => Promise + + /** + * Reload the last assistant message. + */ + reload: () => Promise + + /** + * Stop the current response generation. + */ + stop: () => void + + /** + * Clear all messages. + */ + clear: () => void + + /** + * Set messages manually. + */ + setMessages: (messages: Array>) => void + + /** + * Add the result of a client-side tool execution. + */ + addToolResult: (result: { + toolCallId: string + tool: string + output: any + state?: 'output-available' | 'output-error' + errorText?: string + }) => Promise + + /** + * Respond to a tool approval request. + */ + addToolApprovalResponse: (response: { + id: string // approval.id, not toolCallId + approved: boolean + }) => Promise + + /** + * Whether a response is currently being generated. + */ + isLoading: boolean + + /** + * Current error, if any. + */ + error: Error | undefined + + /** + * Current status of the chat client. + */ + status: ChatClientState + + /** + * Whether the subscription loop is currently active. + */ + isSubscribed: boolean + + /** + * Current connection lifecycle status. + */ + connectionStatus: ConnectionStatus + + /** + * Whether the shared session is actively generating. + */ + sessionGenerating: boolean +} + +/** + * The chat capability surface. Extends {@link AssistantChatSurfaceBase} and, + * when the chat callback declared an `outputSchema` (`TSchema extends + * SchemaInput`), adds the typed structured-output fields `partial` / `final` + * — mirroring the framework `useChat` hooks' conditional return. When no + * schema was declared, neither field is present. + */ +export type AssistantChatSurface< + TTools extends ReadonlyArray = [], + TSchema = undefined, +> = AssistantChatSurfaceBase & { + /** + * Send a message and get a response. Can be a simple string or multimodal + * content with images, audio, etc. + * + * Resolves once the assistant turn completes: to the schema-validated final + * structured output (`InferSchemaType | null`) when the chat + * callback declared an `outputSchema`, otherwise to the resulting messages + * array (`Array>`). + */ + sendMessage: ( + content: string | MultimodalContent, + ) => Promise< + TSchema extends SchemaInput + ? InferSchemaType | null + : Array> + > +} & (TSchema extends SchemaInput + ? { + /** + * Live, progressively-parsed structured output while the stream is + * still arriving. Resets on every new run. + */ + partial: DeepPartial> + /** + * Final, schema-validated structured output. `null` until the terminal + * structured-output event arrives. Resets on every new run. + */ + final: InferSchemaType | null + } + : Record) + +/** The one-shot capability surface — mirrors useGeneration return. */ +export interface AssistantGenerationSurface { + generate: (input: TInput) => Promise + result: TResult | null + isLoading: boolean + error: Error | undefined + status: GenerationClientState + stop: () => void + reset: () => void +} + +/** + * The full typed system returned by useAssistant. + * + * The `chat` surface is fully inferred from the assistant definition's chat + * callback — its captured tools type `chat.messages`' tool-call/result parts, + * and its `outputSchema` (if any) adds typed `chat.partial` / `chat.final`. + */ +export type AssistantSystem< + TDef extends AssistantDefinition, + /** + * @deprecated No longer used for chat tool typing — chat tools are inferred + * from `TDef`'s chat callback. Retained so the framework hooks that still + * pass a second generic keep compiling; it has no effect on the surface. + */ + TChatTools extends ReadonlyArray = [], +> = + // `TChatTools` is deprecated/unused for typing; the constraint always holds, + // so this reads the parameter (keeping it for hook back-compat) without + // affecting the resulting surface. + [TChatTools] extends [ReadonlyArray] + ? { + [K in keyof TDef['~caps'] & string]: K extends 'chat' + ? AssistantChatSurface, ChatSchemaOf> + : K extends OneShotCapabilityName + ? AssistantGenerationSurface< + GenerateInputByCapability[K], + ResultByCapability[K] + > + : never + } + : never diff --git a/packages/ai-client/src/assistant.ts b/packages/ai-client/src/assistant.ts new file mode 100644 index 000000000..40e7f3c07 --- /dev/null +++ b/packages/ai-client/src/assistant.ts @@ -0,0 +1,9 @@ +export { AssistantClient } from './assistant-client' +export { computeStructuredParts } from './assistant-structured' +export type { + AssistantClientOptions, + AssistantSystem, + AssistantChatSurface, + AssistantGenerationSurface, + OneShotCapabilityName, +} from './assistant-types' diff --git a/packages/ai-client/tests/assistant-client.test.ts b/packages/ai-client/tests/assistant-client.test.ts new file mode 100644 index 000000000..f4f9a7846 --- /dev/null +++ b/packages/ai-client/tests/assistant-client.test.ts @@ -0,0 +1,165 @@ +import { describe, expect, expectTypeOf, it } from 'vitest' +import { chat, toolDefinition } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { z } from 'zod' +import { AssistantClient } from '../src/assistant-client.js' +import { fetchServerSentEvents } from '../src/connection-adapters.js' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { AssistantSystem } from '../src/assistant-types.js' + +// Declared, never executed — used only inside chat callbacks that the type +// tests never invoke (`defineAssistant` only runs `Object.keys`). Ambient, so +// it emits no runtime binding and is never read. +declare const adapter: AnyTextAdapter + +describe('AssistantClient', () => { + it('creates one sub-client per declared capability', () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({ id: '', model: '', images: [] }) as any, + }) + const client = new AssistantClient({ + assistant, + connection: fetchServerSentEvents('/api/assistant'), + }) + + expect(client.has('chat')).toBe(true) + expect(client.has('image')).toBe(true) + expect(client.has('speech')).toBe(false) + expect(client.chat).toBeDefined() + expect(client.get('image')).toBeDefined() + expect(client.capabilities).toEqual(['chat', 'image']) + }) + + it('tags each sub-client with its capability', () => { + const assistant = defineAssistant({ + image: async () => ({ id: '', model: '', images: [] }) as any, + }) + const client = new AssistantClient({ + assistant, + connection: fetchServerSentEvents('/api/assistant'), + }) + + expect(client.chat).toBeUndefined() + expect(client.has('chat')).toBe(false) + expect(client.get('image')).toBeDefined() + }) + + it('dispose() tears down the chat client and every one-shot client', () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({ id: '', model: '', images: [] }) as any, + speech: async () => ({}) as any, + }) + const client = new AssistantClient({ + assistant, + connection: fetchServerSentEvents('/api/assistant'), + }) + + const chatDisposeCalls: Array = [] + const imageDisposeCalls: Array = [] + const speechDisposeCalls: Array = [] + client.chat!.dispose = () => { + chatDisposeCalls.push(true) + } + client.get('image')!.dispose = () => { + imageDisposeCalls.push(true) + } + client.get('speech')!.dispose = () => { + speechDisposeCalls.push(true) + } + + client.dispose() + + expect(chatDisposeCalls).toHaveLength(1) + expect(imageDisposeCalls).toHaveLength(1) + expect(speechDisposeCalls).toHaveLength(1) + }) + + it('AssistantSystem exposes only declared capabilities, typed', () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({ id: '', model: '', images: [] }) as any, + }) + type Sys = AssistantSystem + expectTypeOf().toHaveProperty('chat') + expectTypeOf().toHaveProperty('image') + // @ts-expect-error speech was not declared + expectTypeOf().toHaveProperty('speech') + expectTypeOf().toMatchTypeOf<{ + id: string + model: string + images: Array + } | null>() + }) +}) + +/** + * Type-only tests: the chat surface is inferred from the chat callback's + * return — its captured tools type the message tool-call parts, and its + * `outputSchema` (if any) adds typed `partial` / `final`. `defineAssistant` is + * called for real (only `Object.keys` runs); the `chat(...)` callbacks are + * never invoked, so a declared `AnyTextAdapter` never runs. + */ +describe('AssistantSystem chat surface inference', () => { + it('callback tools narrow the chat messages tool-call parts', () => { + const weatherDef = toolDefinition({ + name: 'get_weather', + description: 'Get the weather for a city', + inputSchema: z.object({ city: z.string() }), + outputSchema: z.object({ tempC: z.number() }), + }) + + const assistant = defineAssistant({ + chat: (req) => + chat({ adapter, messages: req.messages, tools: [weatherDef] }), + }) + + type Sys = AssistantSystem + type Part = Sys['chat']['messages'][number]['parts'][number] + type WeatherCall = Extract + + // The tool-call part named `get_weather` exists and its input/output are + // narrowed to the tool's schema-inferred types. + expectTypeOf().toEqualTypeOf<'get_weather'>() + expectTypeOf().toEqualTypeOf< + { city: string } | undefined + >() + expectTypeOf().toEqualTypeOf< + { tempC: number } | undefined + >() + }) + + it('outputSchema adds typed partial/final to the chat surface', () => { + const outputSchema = z.object({ answer: z.string(), score: z.number() }) + + const assistant = defineAssistant({ + chat: (req) => + chat({ adapter, messages: req.messages, outputSchema, stream: true }), + }) + + type ChatSurface = AssistantSystem['chat'] + + expectTypeOf().toEqualTypeOf<{ + answer: string + score: number + } | null>() + expectTypeOf().toEqualTypeOf<{ + answer?: string + score?: number + }>() + }) + + it('a plain chat (no schema) exposes no partial/final', () => { + const assistant = defineAssistant({ + chat: (req) => chat({ adapter, messages: req.messages }), + }) + + type ChatSurface = AssistantSystem['chat'] + + // @ts-expect-error no outputSchema → no `partial` on the chat surface + expectTypeOf().toHaveProperty('partial') + // @ts-expect-error no outputSchema → no `final` on the chat surface + expectTypeOf().toHaveProperty('final') + }) +}) diff --git a/packages/ai-client/tests/assistant-structured.test.ts b/packages/ai-client/tests/assistant-structured.test.ts new file mode 100644 index 000000000..438afdad3 --- /dev/null +++ b/packages/ai-client/tests/assistant-structured.test.ts @@ -0,0 +1,112 @@ +import { describe, expect, it } from 'vitest' +import { computeStructuredParts } from '../src/assistant-structured.js' +import type { UIMessage } from '../src/types.js' + +function userMessage(id: string): UIMessage { + return { + id, + role: 'user', + parts: [{ type: 'text', content: 'hi' }], + } +} + +function assistantWithStructuredOutput( + id: string, + part: { + status: 'streaming' | 'complete' | 'error' + partial?: unknown + data?: unknown + raw?: string + }, +): UIMessage { + return { + id, + role: 'assistant', + parts: [ + { + type: 'structured-output', + status: part.status, + partial: part.partial, + data: part.data, + raw: part.raw ?? '', + }, + ], + } +} + +describe('computeStructuredParts', () => { + it('returns { partial: {}, final: null } when there is no user message', () => { + const messages: Array = [ + assistantWithStructuredOutput('a1', { + status: 'complete', + data: { foo: 'bar' }, + }), + ] + + expect(computeStructuredParts(messages)).toEqual({ + partial: {}, + final: null, + }) + }) + + it('returns partial from a streaming structured-output part after the latest user message', () => { + const messages: Array = [ + userMessage('u1'), + assistantWithStructuredOutput('a1', { + status: 'streaming', + partial: { name: 'Al' }, + }), + ] + + expect(computeStructuredParts(messages)).toEqual({ + partial: { name: 'Al' }, + final: null, + }) + }) + + it('falls back to data for partial when a streaming part has no partial field', () => { + const messages: Array = [ + userMessage('u1'), + assistantWithStructuredOutput('a1', { + status: 'streaming', + data: { name: 'Al' }, + }), + ] + + expect(computeStructuredParts(messages)).toEqual({ + partial: { name: 'Al' }, + final: null, + }) + }) + + it('returns final from a complete structured-output part after the latest user message', () => { + const messages: Array = [ + userMessage('u1'), + assistantWithStructuredOutput('a1', { + status: 'complete', + partial: { name: 'Alem' }, + data: { name: 'Alem' }, + }), + ] + + expect(computeStructuredParts(messages)).toEqual({ + partial: { name: 'Alem' }, + final: { name: 'Alem' }, + }) + }) + + it('does not leak a final from a structured-output part before the latest user message', () => { + const messages: Array = [ + assistantWithStructuredOutput('a1', { + status: 'complete', + data: { name: 'stale' }, + }), + userMessage('u1'), + ] + + expect(computeStructuredParts(messages)).toEqual({ + partial: {}, + final: null, + }) + }) +}) diff --git a/packages/ai-client/vite.config.ts b/packages/ai-client/vite.config.ts index b36468a9a..b60ca079b 100644 --- a/packages/ai-client/vite.config.ts +++ b/packages/ai-client/vite.config.ts @@ -40,7 +40,7 @@ export default mergeConfig( // implementations; declare it as its own entry so the build emits // it independently and the main entry can stay free of the bridge // classes (they're imported only via `import type` from clients). - entry: ['./src/index.ts', './src/devtools.ts'], + entry: ['./src/index.ts', './src/devtools.ts', './src/assistant.ts'], srcDir: './src', cjs: false, }), diff --git a/packages/ai-react/package.json b/packages/ai-react/package.json index a4ffca2b9..eb82e4bca 100644 --- a/packages/ai-react/package.json +++ b/packages/ai-react/package.json @@ -28,6 +28,10 @@ "./mcp-apps": { "types": "./dist/esm/mcp-apps.d.ts", "import": "./dist/esm/mcp-apps.js" + }, + "./assistant": { + "types": "./dist/esm/assistant.d.ts", + "import": "./dist/esm/assistant.js" } }, "files": [ diff --git a/packages/ai-react/src/assistant.ts b/packages/ai-react/src/assistant.ts new file mode 100644 index 000000000..6a8666b6e --- /dev/null +++ b/packages/ai-react/src/assistant.ts @@ -0,0 +1 @@ +export { useAssistant } from './use-assistant' diff --git a/packages/ai-react/src/use-assistant.ts b/packages/ai-react/src/use-assistant.ts new file mode 100644 index 000000000..f803cd3a5 --- /dev/null +++ b/packages/ai-react/src/use-assistant.ts @@ -0,0 +1,251 @@ +import { + AssistantClient, + computeStructuredParts, +} from '@tanstack/ai-client/assistant' +import { useEffect, useId, useMemo, useRef, useState } from 'react' +import type { AssistantDefinition } from '@tanstack/ai/assistant' +import type { AnyClientTool } from '@tanstack/ai/client' +import type { + AssistantClientOptions, + AssistantSystem, + OneShotCapabilityName, +} from '@tanstack/ai-client/assistant' +import type { + ChatClientState, + ConnectionStatus, + GenerationClientState, +} from '@tanstack/ai-client' + +/** Reactive chat-capability state mirrored from the underlying ChatClient. */ +interface ChatState { + messages: Array + isLoading: boolean + error: Error | undefined + status: ChatClientState + isSubscribed: boolean + connectionStatus: ConnectionStatus + sessionGenerating: boolean +} + +/** Reactive one-shot-capability state mirrored from a GenerationClient. */ +interface OneShotState { + result: any + isLoading: boolean + error: Error | undefined + status: GenerationClientState +} + +const initialChatState: ChatState = { + messages: [], + isLoading: false, + error: undefined, + status: 'ready', + isSubscribed: false, + connectionStatus: 'disconnected', + sessionGenerating: false, +} + +const initialOneShotState: OneShotState = { + result: null, + isLoading: false, + error: undefined, + status: 'idle', +} + +/** + * React hook wrapping `AssistantClient`, composing the existing chat + + * generation clients behind one endpoint into a single typed system keyed by + * the assistant's declared capabilities. + * + * @example + * ```tsx + * const system = useAssistant(assistant, { + * connection: fetchServerSentEvents('/api/assistant'), + * }) + * + * await system.chat.sendMessage('hi') + * await system.image.generate({ prompt: 'a fox' }) + * ``` + */ +export function useAssistant< + TDef extends AssistantDefinition, + TChatTools extends ReadonlyArray = [], +>( + assistant: TDef, + options: Omit< + AssistantClientOptions, + 'assistant' | 'callbacks' + >, +): AssistantSystem { + const hookId = useId() + const clientId = options.id ?? hookId + + const optionsRef = useRef(options) + optionsRef.current = options + const activeClientRef = useRef | null>(null) + + const [chatState, setChatState] = useState(initialChatState) + const [oneShotState, setOneShotState] = useState< + Record + >({}) + + const client = useMemo(() => { + const initial = optionsRef.current + + const instance = new AssistantClient({ + assistant, + connection: initial.connection, + id: clientId, + threadId: initial.threadId, + chat: initial.chat, + callbacks: { + chat: { + onMessagesChange: (m) => { + if (activeClientRef.current !== instance) return + setChatState((s) => ({ ...s, messages: m })) + }, + onLoadingChange: (v) => { + if (activeClientRef.current !== instance) return + setChatState((s) => ({ ...s, isLoading: v })) + }, + onErrorChange: (v) => { + if (activeClientRef.current !== instance) return + setChatState((s) => ({ ...s, error: v })) + }, + onStatusChange: (v) => { + if (activeClientRef.current !== instance) return + setChatState((s) => ({ ...s, status: v })) + }, + onSubscriptionChange: (v) => { + if (activeClientRef.current !== instance) return + setChatState((s) => ({ ...s, isSubscribed: v })) + }, + onConnectionStatusChange: (v) => { + if (activeClientRef.current !== instance) return + setChatState((s) => ({ ...s, connectionStatus: v })) + }, + onSessionGeneratingChange: (v) => { + if (activeClientRef.current !== instance) return + setChatState((s) => ({ ...s, sessionGenerating: v })) + }, + }, + oneShot: (cap) => ({ + onResultChange: (r) => { + if (activeClientRef.current !== instance) return + setOneShotState((s) => ({ + ...s, + [cap]: { ...(s[cap] ?? initialOneShotState), result: r }, + })) + }, + onLoadingChange: (v) => { + if (activeClientRef.current !== instance) return + setOneShotState((s) => ({ + ...s, + [cap]: { ...(s[cap] ?? initialOneShotState), isLoading: v }, + })) + }, + onErrorChange: (v) => { + if (activeClientRef.current !== instance) return + setOneShotState((s) => ({ + ...s, + [cap]: { ...(s[cap] ?? initialOneShotState), error: v }, + })) + }, + onStatusChange: (v) => { + if (activeClientRef.current !== instance) return + setOneShotState((s) => ({ + ...s, + [cap]: { ...(s[cap] ?? initialOneShotState), status: v }, + })) + }, + }), + }, + }) + activeClientRef.current = instance + return instance + // Client is intentionally keyed only on clientId, mirroring useChat. + }, [clientId]) + + useEffect(() => { + activeClientRef.current = client + return () => { + if (activeClientRef.current === client) { + activeClientRef.current = null + } + client.dispose() + } + }, [client]) + + const { partial, final } = useMemo( + () => computeStructuredParts(chatState.messages), + [chatState.messages], + ) + + const system = useMemo(() => { + const out: Record = {} + + for (const cap of client.capabilities) { + if (cap === 'chat') { + const c = client.chat + if (!c) continue + out.chat = { + messages: chatState.messages, + sendMessage: async (content: any) => { + await c.sendMessage(content) + const msgs = c.getMessages() + const hasStructured = + msgs + .at(-1) + ?.parts.some((p) => p.type === 'structured-output') ?? false + // Cast: TS can't structurally narrow the conditional return of + // AssistantChatSurface['sendMessage'] against this runtime branch, + // mirroring the assembled-`system` cast below. + return ( + hasStructured ? computeStructuredParts(msgs).final : msgs + ) as any + }, + append: (message: any) => c.append(message), + reload: () => c.reload(), + stop: () => c.stop(), + clear: () => c.clear(), + setMessages: (messages: any) => c.setMessagesManually(messages), + addToolResult: (result: any) => c.addToolResult(result), + addToolApprovalResponse: (response: any) => + c.addToolApprovalResponse(response), + isLoading: chatState.isLoading, + error: chatState.error, + status: chatState.status, + isSubscribed: chatState.isSubscribed, + connectionStatus: chatState.connectionStatus, + sessionGenerating: chatState.sessionGenerating, + // Runtime shape unconditionally exposes partial/final; the public + // AssistantSystem type hides them when the chat capability's + // outputSchema is absent, matching useChat's behavior. + partial, + final, + } + continue + } + + const g = client.get(cap as OneShotCapabilityName) + if (!g) continue + const slice = oneShotState[cap] ?? initialOneShotState + out[cap] = { + generate: async (input: any) => { + await g.generate(input) + return g.getResult() + }, + result: slice.result, + isLoading: slice.isLoading, + error: slice.error, + status: slice.status, + stop: () => g.stop(), + reset: () => g.reset(), + } + } + + return out as AssistantSystem + }, [client, chatState, oneShotState, partial, final]) + + return system +} diff --git a/packages/ai-react/tests/use-assistant.test.tsx b/packages/ai-react/tests/use-assistant.test.tsx new file mode 100644 index 000000000..153952945 --- /dev/null +++ b/packages/ai-react/tests/use-assistant.test.tsx @@ -0,0 +1,114 @@ +import { describe, expect, it } from 'vitest' +import { act, renderHook } from '@testing-library/react' +import { defineAssistant } from '@tanstack/ai/assistant' +import { stream } from '@tanstack/ai-client' +import { useAssistant } from '../src/use-assistant.js' + +// A connection adapter that replays canned chunks for both capabilities, +// branching on the `capability` discriminator forwarded by AssistantClient. +function fakeConnection() { + return stream(async function* (_messages, data) { + const capability = (data as Record | undefined)?.capability + + if (capability === 'chat') { + yield { type: 'RUN_STARTED', threadId: 't', runId: 'r' } as any + yield { + type: 'TEXT_MESSAGE_START', + messageId: 'm1', + role: 'assistant', + } as any + yield { + type: 'TEXT_MESSAGE_CONTENT', + messageId: 'm1', + delta: 'hello', + } as any + yield { type: 'TEXT_MESSAGE_END', messageId: 'm1' } as any + yield { type: 'RUN_FINISHED', threadId: 't', runId: 'r' } as any + } else { + yield { type: 'RUN_STARTED', threadId: 't', runId: 'r' } as any + yield { + type: 'CUSTOM', + name: 'generation:result', + value: { id: 'i', model: 'gpt-image-1', images: [{ url: 'u' }] }, + } as any + yield { type: 'RUN_FINISHED', threadId: 't', runId: 'r' } as any + } + }) +} + +describe('useAssistant', () => { + it('exposes only the declared capabilities', () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({}) as any, + }) + const { result } = renderHook(() => + useAssistant(assistant, { connection: fakeConnection() }), + ) + expect(result.current.chat).toBeDefined() + expect(result.current.image).toBeDefined() + expect((result.current as any).speech).toBeUndefined() + }) + + it('generate() on a one-shot capability populates its result', async () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({}) as any, + }) + const { result } = renderHook(() => + useAssistant(assistant, { connection: fakeConnection() }), + ) + + // Holder object (not a bare `let`) so TS keeps the awaited return type at + // the read site — a `let` assigned only inside the async `act` callback + // gets narrowed away. The assignment still type-checks generate()'s return. + const captured: { value: typeof result.current.image.result } = { + value: null, + } + await act(async () => { + captured.value = await result.current.image.generate({ prompt: 'a fox' }) + }) + + // generate() resolves to the fresh result... + expect(captured.value?.images[0]?.url).toBe('u') + // ...and the reactive result state is still populated. + expect(result.current.image.result?.images[0]?.url).toBe('u') + }) + + it('sendMessage() on the chat capability populates messages', async () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({}) as any, + }) + const { result } = renderHook(() => + useAssistant(assistant, { connection: fakeConnection() }), + ) + + // Holder object (see note in the generate() test) keeps the awaited + // sendMessage() return type at the read site. + const captured: { value: typeof result.current.chat.messages } = { + value: [], + } + await act(async () => { + captured.value = await result.current.chat.sendMessage('hi') + }) + + // With no outputSchema, sendMessage() resolves to the messages array... + expect(captured.value.length).toBeGreaterThan(0) + // ...and the reactive messages state is still populated. + expect(result.current.chat.messages.length).toBeGreaterThan(0) + }) + + it('exposes chat.partial/final with cleared defaults on first render', () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({}) as any, + }) + const { result } = renderHook(() => + useAssistant(assistant, { connection: fakeConnection() }), + ) + + expect((result.current.chat as any).partial).toEqual({}) + expect((result.current.chat as any).final).toBeNull() + }) +}) diff --git a/packages/ai-react/vite.config.ts b/packages/ai-react/vite.config.ts index a697ee62f..891058a8d 100644 --- a/packages/ai-react/vite.config.ts +++ b/packages/ai-react/vite.config.ts @@ -29,7 +29,7 @@ const config = defineConfig({ export default mergeConfig( config, tanstackViteConfig({ - entry: ['./src/index.ts', './src/mcp-apps.ts'], + entry: ['./src/index.ts', './src/mcp-apps.ts', './src/assistant.ts'], srcDir: './src', cjs: false, }), diff --git a/packages/ai-solid/package.json b/packages/ai-solid/package.json index 3d83b1856..554b8a403 100644 --- a/packages/ai-solid/package.json +++ b/packages/ai-solid/package.json @@ -24,6 +24,10 @@ ".": { "types": "./dist/index.d.ts", "import": "./dist/index.js" + }, + "./assistant": { + "types": "./dist/assistant.d.ts", + "import": "./dist/assistant.js" } }, "files": [ diff --git a/packages/ai-solid/src/assistant.ts b/packages/ai-solid/src/assistant.ts new file mode 100644 index 000000000..6a8666b6e --- /dev/null +++ b/packages/ai-solid/src/assistant.ts @@ -0,0 +1 @@ +export { useAssistant } from './use-assistant' diff --git a/packages/ai-solid/src/use-assistant.ts b/packages/ai-solid/src/use-assistant.ts new file mode 100644 index 000000000..bf0d64d57 --- /dev/null +++ b/packages/ai-solid/src/use-assistant.ts @@ -0,0 +1,335 @@ +import { + createMemo, + createSignal, + createUniqueId, + onCleanup, + onMount, +} from 'solid-js' + +import { + AssistantClient, + computeStructuredParts, +} from '@tanstack/ai-client/assistant' +import type { + AnyClientTool, + ChatClientState, + ConnectionStatus, + GenerationClientState, + MultimodalContent, + UIMessage, +} from '@tanstack/ai-client' +import type { + AssistantClientOptions, + AssistantSystem, + OneShotCapabilityName, +} from '@tanstack/ai-client/assistant' +import type { ModelMessage } from '@tanstack/ai' +import type { AssistantDefinition } from '@tanstack/ai/assistant' + +/** Reactive chat sub-state, mirrored from the `ChatClient` callbacks. */ +interface ChatState> { + messages: Array> + isLoading: boolean + error: Error | undefined + status: ChatClientState + isSubscribed: boolean + connectionStatus: ConnectionStatus + sessionGenerating: boolean +} + +/** Reactive one-shot sub-state, mirrored from a `GenerationClient`'s callbacks. */ +interface OneShotState { + result: unknown + isLoading: boolean + error: Error | undefined + status: GenerationClientState +} + +const defaultOneShotState: OneShotState = { + result: null, + isLoading: false, + error: undefined, + status: 'idle', +} + +/** + * Solid hook for a multi-capability `AssistantDefinition` — composes one + * reactive surface per declared capability (`chat`, and/or one or more + * one-shot generation capabilities) backed by a single `AssistantClient`. + * + * Mirrors the `useChat` idiom: state lives in `createSignal`s, the client is + * built once in a `createMemo` (`clientId` is a stable per-hook id), and every reactive + * callback is threaded into the `AssistantClient` constructor via its + * `callbacks` option — `ChatClient` and `GenerationClient` only accept these + * callbacks at construction time, not through `updateOptions`. + * + * @example + * ```tsx + * const assistant = useAssistant(myAssistant, { + * connection: fetchServerSentEvents('/api/assistant'), + * }) + * + * await assistant.chat.sendMessage('Hello') + * await assistant.image.generate({ prompt: 'A sunset' }) + * ``` + */ +export function useAssistant< + TDef extends AssistantDefinition, + TChatTools extends ReadonlyArray = [], +>( + assistant: TDef, + options: Omit< + AssistantClientOptions, + 'assistant' | 'callbacks' + >, +): AssistantSystem { + const hookId = createUniqueId() + const clientId = options.id || hookId + + const [chatState, setChatState] = createSignal>({ + messages: [], + isLoading: false, + error: undefined, + status: 'ready', + isSubscribed: false, + connectionStatus: 'disconnected', + sessionGenerating: false, + }) + + const initialOneShotState: Record = {} + for (const capability of assistant.capabilities) { + if (capability !== 'chat') { + initialOneShotState[capability] = { ...defaultOneShotState } + } + } + const [oneShotState, setOneShotState] = + createSignal>(initialOneShotState) + + // Build the AssistantClient with every reactive callback wired through the + // constructor's `callbacks` option (VP2: ChatClient/GenerationClient only + // accept these at construction time). + const client = createMemo(() => { + return new AssistantClient({ + assistant, + connection: options.connection, + id: clientId, + ...(options.threadId !== undefined && { threadId: options.threadId }), + ...(options.chat !== undefined && { chat: options.chat }), + callbacks: { + chat: { + onMessagesChange: (messages) => + setChatState((s) => ({ ...s, messages })), + onLoadingChange: (isLoading) => + setChatState((s) => ({ ...s, isLoading })), + onErrorChange: (error) => setChatState((s) => ({ ...s, error })), + onStatusChange: (status) => setChatState((s) => ({ ...s, status })), + onSubscriptionChange: (isSubscribed) => + setChatState((s) => ({ ...s, isSubscribed })), + onConnectionStatusChange: (connectionStatus) => + setChatState((s) => ({ ...s, connectionStatus })), + onSessionGeneratingChange: (sessionGenerating) => + setChatState((s) => ({ ...s, sessionGenerating })), + }, + oneShot: (capability) => ({ + onResultChange: (result) => + setOneShotState((s) => ({ + ...s, + [capability]: { + ...(s[capability] ?? defaultOneShotState), + result, + }, + })), + onLoadingChange: (isLoading) => + setOneShotState((s) => ({ + ...s, + [capability]: { + ...(s[capability] ?? defaultOneShotState), + isLoading, + }, + })), + onErrorChange: (error) => + setOneShotState((s) => ({ + ...s, + [capability]: { + ...(s[capability] ?? defaultOneShotState), + error, + }, + })), + onStatusChange: (status) => + setOneShotState((s) => ({ + ...s, + [capability]: { + ...(s[capability] ?? defaultOneShotState), + status, + }, + })), + }), + }, + }) + }) + + // Sync initial chat state now that the client (and its `chat` sub-client, + // if declared) exists — mirrors useChat's `setMessages(client().getMessages())`. + const initialChatClient = client().chat + if (initialChatClient) { + setChatState({ + messages: initialChatClient.getMessages(), + isLoading: initialChatClient.getIsLoading(), + error: initialChatClient.getError(), + status: initialChatClient.getStatus(), + isSubscribed: initialChatClient.getIsSubscribed(), + connectionStatus: initialChatClient.getConnectionStatus(), + sessionGenerating: initialChatClient.getSessionGenerating(), + }) + } + + onMount(() => { + client().chat?.mountDevtools() + }) + + // Cleanup on unmount: tear down the chat client and every one-shot client. + onCleanup(() => { + client().dispose() + }) + + // Narrowing helpers: capability presence is guaranteed by construction (a + // surface for `capability` is only built below when the assistant actually + // declares it), but the sub-client getters are typed as `T | undefined`, so + // reads go through an explicit check rather than a non-null assertion. + const requireChatClient = () => { + const chatClient = client().chat + if (chatClient === undefined) { + throw new Error( + 'useAssistant: "chat" capability was not declared on this assistant', + ) + } + return chatClient + } + + const requireOneShotClient = (capability: OneShotCapabilityName) => { + const oneShotClient = client().get(capability) + if (oneShotClient === undefined) { + throw new Error( + `useAssistant: "${capability}" capability was not declared on this assistant`, + ) + } + return oneShotClient + } + + // Built dynamically per declared capability below; the object shape can't + // be statically checked against the mapped `AssistantSystem` type, so it's + // assembled as `Record` and cast once at the end. + const system: Record = {} + + for (const capability of assistant.capabilities) { + if (capability === 'chat') { + system.chat = { + get messages() { + return chatState().messages + }, + get isLoading() { + return chatState().isLoading + }, + get error() { + return chatState().error + }, + get status() { + return chatState().status + }, + get isSubscribed() { + return chatState().isSubscribed + }, + get connectionStatus() { + return chatState().connectionStatus + }, + get sessionGenerating() { + return chatState().sessionGenerating + }, + // Runtime shape unconditionally exposes partial/final; the public + // AssistantSystem type hides them when the chat capability's + // outputSchema is absent, matching useChat's behavior. + get partial() { + return computeStructuredParts(chatState().messages).partial + }, + get final() { + return computeStructuredParts(chatState().messages).final + }, + sendMessage: async (content: string | MultimodalContent) => { + const chatClient = requireChatClient() + await chatClient.sendMessage(content) + const msgs = chatClient.getMessages() + const hasStructured = + msgs.at(-1)?.parts.some((p) => p.type === 'structured-output') ?? + false + // Cast: TS can't structurally narrow the conditional return of + // AssistantChatSurface['sendMessage'] against this runtime branch, + // mirroring the assembled-`system` cast at the end of this hook. + return ( + hasStructured ? computeStructuredParts(msgs).final : msgs + ) as any + }, + append: async (message: ModelMessage | UIMessage) => { + await requireChatClient().append(message) + }, + reload: async () => { + await requireChatClient().reload() + }, + stop: () => { + requireChatClient().stop() + }, + clear: () => { + requireChatClient().clear() + }, + setMessages: (messages: Array>) => { + requireChatClient().setMessagesManually(messages) + }, + addToolResult: async (result: { + toolCallId: string + tool: string + output: any + state?: 'output-available' | 'output-error' + errorText?: string + }) => { + await requireChatClient().addToolResult(result) + }, + addToolApprovalResponse: async (response: { + id: string + approved: boolean + }) => { + await requireChatClient().addToolApprovalResponse(response) + }, + } + continue + } + + const oneShotCapability = capability as OneShotCapabilityName + system[capability] = { + get result() { + return oneShotState()[oneShotCapability]?.result ?? null + }, + get isLoading() { + return oneShotState()[oneShotCapability]?.isLoading ?? false + }, + get error() { + return oneShotState()[oneShotCapability]?.error + }, + get status() { + return oneShotState()[oneShotCapability]?.status ?? 'idle' + }, + generate: async (input: Record) => { + const oneShotClient = requireOneShotClient(oneShotCapability) + await oneShotClient.generate(input) + return oneShotClient.getResult() + }, + stop: () => { + requireOneShotClient(oneShotCapability).stop() + }, + reset: () => { + requireOneShotClient(oneShotCapability).reset() + }, + } + } + + // eslint-disable-next-line no-restricted-syntax -- built dynamically per declared capability; TS can't structurally verify the assembled object against the mapped AssistantSystem type + return system as unknown as AssistantSystem +} diff --git a/packages/ai-solid/tests/use-assistant.test.tsx b/packages/ai-solid/tests/use-assistant.test.tsx new file mode 100644 index 000000000..bc0c97cb8 --- /dev/null +++ b/packages/ai-solid/tests/use-assistant.test.tsx @@ -0,0 +1,180 @@ +import { renderHook, waitFor } from '@solidjs/testing-library' +import { describe, expect, it } from 'vitest' +import { defineAssistant } from '@tanstack/ai/assistant' +import { stream } from '@tanstack/ai-client' +import type { StreamChunk } from '@tanstack/ai' +import { useAssistant } from '../src/use-assistant.js' +import { createTextChunks } from './test-utils' + +/** + * A canned connection shared by chat and one-shot sub-clients (mirroring how + * `AssistantClient` composes them). Routes on `data.capability`, which + * `AssistantClient` tags onto every request: `forwardedProps.capability` for + * chat, `body.capability` for one-shot generation. + */ +function createAssistantConnection() { + return stream(async function* (_messages, data): AsyncGenerator { + const capability = (data as Record | undefined)?.capability + + if (capability === 'chat') { + yield* createTextChunks('Hello from chat') + return + } + + yield { + type: 'RUN_STARTED', + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + } as StreamChunk + yield { + type: 'CUSTOM', + name: 'generation:result', + value: { id: 'img-1', model: 'test-model', images: [] }, + timestamp: Date.now(), + } as StreamChunk + yield { + type: 'RUN_FINISHED', + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + } as StreamChunk + }) +} + +describe('useAssistant', () => { + it('exposes only the declared capabilities', () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({ id: '', model: '', images: [] }) as any, + }) + + const { result } = renderHook(() => + useAssistant(assistant, { connection: createAssistantConnection() }), + ) + + expect(result.chat).toBeDefined() + expect(result.image).toBeDefined() + expect((result as any).speech).toBeUndefined() + }) + + it('populates the one-shot result after generate()', async () => { + const assistant = defineAssistant({ + image: async () => ({ id: '', model: '', images: [] }) as any, + }) + + const { result } = renderHook(() => + useAssistant(assistant, { connection: createAssistantConnection() }), + ) + + expect(result.image.result).toBeNull() + expect(result.image.isLoading).toBe(false) + expect(result.image.status).toBe('idle') + + // generate() resolves to the fresh result... + const generated = await result.image.generate({ prompt: 'A sunset' }) + expect(generated).toEqual({ + id: 'img-1', + model: 'test-model', + images: [], + }) + + // ...and the reactive result state is still populated. + expect(result.image.result).toEqual({ + id: 'img-1', + model: 'test-model', + images: [], + }) + expect(result.image.status).toBe('success') + expect(result.image.isLoading).toBe(false) + }) + + it('populates chat messages after sendMessage()', async () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + }) + + const { result } = renderHook(() => + useAssistant(assistant, { connection: createAssistantConnection() }), + ) + + expect(result.chat.messages).toEqual([]) + + // With no outputSchema, sendMessage() resolves to the messages array. + const returned = await result.chat.sendMessage('Hi there') + expect(returned.length).toBeGreaterThanOrEqual(2) + + await waitFor(() => { + expect(result.chat.messages.length).toBeGreaterThanOrEqual(2) + }) + + const userMessage = result.chat.messages.find((m) => m.role === 'user') + expect(userMessage).toBeDefined() + if (userMessage) { + expect(userMessage.parts[0]).toEqual({ + type: 'text', + content: 'Hi there', + }) + } + + const assistantMessage = result.chat.messages.find( + (m) => m.role === 'assistant', + ) + expect(assistantMessage).toBeDefined() + const textPart = assistantMessage?.parts.find((p) => p.type === 'text') + expect(textPart).toBeDefined() + if (textPart && textPart.type === 'text') { + expect(textPart.content).toBe('Hello from chat') + } + }) + + it('supports both chat and one-shot capabilities on the same assistant', async () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({ id: '', model: '', images: [] }) as any, + }) + + const { result } = renderHook(() => + useAssistant(assistant, { connection: createAssistantConnection() }), + ) + + await result.chat.sendMessage('Hi there') + await waitFor(() => { + expect(result.chat.messages.length).toBeGreaterThanOrEqual(2) + }) + + await result.image.generate({ prompt: 'A sunset' }) + expect(result.image.result).toEqual({ + id: 'img-1', + model: 'test-model', + images: [], + }) + }) + + it('exposes structured partial/final on the chat surface with no structured part present', () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + }) + + const { result } = renderHook(() => + useAssistant(assistant, { connection: createAssistantConnection() }), + ) + + expect((result.chat as any).partial).toEqual({}) + expect((result.chat as any).final).toBeNull() + }) + + it('disposes both the chat and one-shot sub-clients on cleanup', () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({ id: '', model: '', images: [] }) as any, + }) + + const { result, cleanup } = renderHook(() => + useAssistant(assistant, { connection: createAssistantConnection() }), + ) + + expect(result.chat).toBeDefined() + expect(() => cleanup()).not.toThrow() + }) +}) diff --git a/packages/ai-solid/tsdown.config.ts b/packages/ai-solid/tsdown.config.ts index 3c1187801..de631c408 100644 --- a/packages/ai-solid/tsdown.config.ts +++ b/packages/ai-solid/tsdown.config.ts @@ -1,7 +1,7 @@ import { defineConfig } from 'tsdown' export default defineConfig({ - entry: ['./src/index.ts'], + entry: ['./src/index.ts', './src/assistant.ts'], format: ['esm'], unbundle: true, dts: true, diff --git a/packages/ai-svelte/package.json b/packages/ai-svelte/package.json index 808307237..fa06b7a56 100644 --- a/packages/ai-svelte/package.json +++ b/packages/ai-svelte/package.json @@ -26,6 +26,11 @@ "types": "./dist/index.d.ts", "svelte": "./dist/index.js", "import": "./dist/index.js" + }, + "./assistant": { + "types": "./dist/assistant.d.ts", + "svelte": "./dist/assistant.js", + "import": "./dist/assistant.js" } }, "files": [ diff --git a/packages/ai-svelte/src/assistant.ts b/packages/ai-svelte/src/assistant.ts new file mode 100644 index 000000000..7d2764959 --- /dev/null +++ b/packages/ai-svelte/src/assistant.ts @@ -0,0 +1 @@ +export { createAssistant } from './create-assistant.svelte' diff --git a/packages/ai-svelte/src/create-assistant.svelte.ts b/packages/ai-svelte/src/create-assistant.svelte.ts new file mode 100644 index 000000000..ea3213dd0 --- /dev/null +++ b/packages/ai-svelte/src/create-assistant.svelte.ts @@ -0,0 +1,306 @@ +import { + AssistantClient, + computeStructuredParts, +} from '@tanstack/ai-client/assistant' +import type { + AssistantClientOptions, + AssistantSystem, + OneShotCapabilityName, +} from '@tanstack/ai-client/assistant' +import type { + AnyClientTool, + ChatClientState, + ConnectionStatus, + GenerationClientState, +} from '@tanstack/ai-client' +import type { AssistantDefinition } from '@tanstack/ai/assistant' +import type { ModelMessage } from '@tanstack/ai' +import type { MultimodalContent, UIMessage } from './types' + +/** Reactive state for a single one-shot (non-chat) capability. */ +interface OneShotState { + result: unknown + isLoading: boolean + error: Error | undefined + status: GenerationClientState +} + +/** + * Creates a reactive assistant instance for Svelte 5. + * + * This function wraps the `AssistantClient` from `@tanstack/ai-client` and + * exposes reactive state using Svelte 5 runes, one surface per declared + * capability (`chat` plus any one-shot generation capabilities). The + * returned object exposes reactive getters that automatically update when + * state changes. + * + * @example + * ```svelte + * + * + *
+ * {#each assistant.chat.messages as message} + *
{message.role}: {message.parts[0].content}
+ * {/each} + * + * + * + *
+ * ``` + */ +export function createAssistant< + TDef extends AssistantDefinition, + TChatTools extends ReadonlyArray = [], +>( + assistant: TDef, + options: Omit< + AssistantClientOptions, + 'assistant' | 'callbacks' + >, +): AssistantSystem & { dispose: () => void } { + // Reactive state for the chat capability, if declared. + let chatMessages = $state>>([]) + let chatIsLoading = $state(false) + let chatError = $state(undefined) + let chatStatus = $state('ready') + let chatIsSubscribed = $state(false) + let chatConnectionStatus = $state('disconnected') + let chatSessionGenerating = $state(false) + + // Derived structured-output `partial`/`final`, recomputed whenever + // `chatMessages` changes (mirrors `useChat`/`useAssistant`'s `useMemo`). + const structuredParts = $derived(computeStructuredParts(chatMessages)) + + // Reactive state per one-shot capability, keyed by capability name. + // Initialized eagerly for every declared one-shot capability, since + // `assistant.capabilities` is fixed at creation time (mirrors + // `AssistantClient`'s own constructor, which iterates the same array). + const oneShotStates = $state>({}) + for (const capability of assistant.capabilities) { + if (capability === 'chat') continue + oneShotStates[capability] = { + result: null, + isLoading: false, + error: undefined, + status: 'idle', + } + } + + // Returns (and lazily creates) the reactive state slot for a one-shot + // capability, avoiding a non-null assertion at each call site below. + const ensureOneShotState = (capability: string): OneShotState => { + let state = oneShotStates[capability] + if (!state) { + state = { + result: null, + isLoading: false, + error: undefined, + status: 'idle', + } + oneShotStates[capability] = state + } + return state + } + + // Create the AssistantClient eagerly, once. Reactive callbacks are passed + // into the constructor (the only place ChatClient/GenerationClient accept + // them) rather than via `updateOptions`, mirroring `createChat`. + const client = new AssistantClient({ + assistant, + connection: options.connection, + ...(options.id !== undefined && { id: options.id }), + ...(options.threadId !== undefined && { threadId: options.threadId }), + ...(options.chat !== undefined && { chat: options.chat }), + callbacks: { + chat: { + onMessagesChange: (newMessages) => { + chatMessages = newMessages + }, + onLoadingChange: (newIsLoading) => { + chatIsLoading = newIsLoading + }, + onErrorChange: (newError) => { + chatError = newError + }, + onStatusChange: (newStatus) => { + chatStatus = newStatus + }, + onSubscriptionChange: (nextIsSubscribed) => { + chatIsSubscribed = nextIsSubscribed + }, + onConnectionStatusChange: (nextStatus) => { + chatConnectionStatus = nextStatus + }, + onSessionGeneratingChange: (isGenerating) => { + chatSessionGenerating = isGenerating + }, + }, + oneShot: (capability) => ({ + onResultChange: (result) => { + ensureOneShotState(capability).result = result + }, + onLoadingChange: (isLoading) => { + ensureOneShotState(capability).isLoading = isLoading + }, + onErrorChange: (error) => { + ensureOneShotState(capability).error = error + }, + onStatusChange: (status) => { + ensureOneShotState(capability).status = status + }, + }), + }, + }) + + chatMessages = client.chat?.getMessages() ?? [] + + // Note: No auto-cleanup in Svelte — call `dispose()` in your component's + // cleanup if needed. Unlike React/Vue/Solid, Svelte 5 runes like $effect + // can only be used during component initialization. + const dispose = () => { + client.dispose() + } + + // Chat capability methods. + const sendMessage = async (content: string | MultimodalContent) => { + await client.chat?.sendMessage(content) + const msgs = client.chat?.getMessages() ?? [] + const hasStructured = + msgs.at(-1)?.parts.some((p) => p.type === 'structured-output') ?? false + // Cast: TS can't structurally narrow the conditional return of + // AssistantChatSurface['sendMessage'] against this runtime branch, + // mirroring the assembled-`system` cast at the end of this function. + return (hasStructured ? computeStructuredParts(msgs).final : msgs) as any + } + + const append = async (message: ModelMessage | UIMessage) => { + await client.chat?.append(message) + } + + const reload = async () => { + await client.chat?.reload() + } + + const chatStop = () => { + client.chat?.stop() + } + + const clear = () => { + client.chat?.clear() + } + + const setMessages = (newMessages: Array>) => { + client.chat?.setMessagesManually(newMessages) + } + + const addToolResult = async (result: { + toolCallId: string + tool: string + output: any + state?: 'output-available' | 'output-error' + errorText?: string + }) => { + await client.chat?.addToolResult(result) + } + + const addToolApprovalResponse = async (response: { + id: string + approved: boolean + }) => { + await client.chat?.addToolApprovalResponse(response) + } + + // Build the returned system: one entry per declared capability, plus a + // top-level `dispose`. Uses getters so Svelte tracks the underlying + // `$state` without a `$` prefix. + const system: Record = {} + + for (const capability of assistant.capabilities) { + if (capability === 'chat') { + system.chat = { + get messages() { + return chatMessages + }, + get isLoading() { + return chatIsLoading + }, + get error() { + return chatError + }, + get status() { + return chatStatus + }, + get isSubscribed() { + return chatIsSubscribed + }, + get connectionStatus() { + return chatConnectionStatus + }, + get sessionGenerating() { + return chatSessionGenerating + }, + // Runtime shape unconditionally exposes partial/final; the public + // AssistantSystem type hides them when the chat capability's + // outputSchema is absent, matching useChat's behavior. + get partial() { + return structuredParts.partial + }, + get final() { + return structuredParts.final + }, + sendMessage, + append, + reload, + stop: chatStop, + clear, + setMessages, + addToolResult, + addToolApprovalResponse, + } + continue + } + + const capabilityName = capability as OneShotCapabilityName + system[capabilityName] = { + get result() { + return oneShotStates[capabilityName]?.result ?? null + }, + get isLoading() { + return oneShotStates[capabilityName]?.isLoading ?? false + }, + get error() { + return oneShotStates[capabilityName]?.error + }, + get status() { + return oneShotStates[capabilityName]?.status ?? 'idle' + }, + generate: async (input: any) => { + const oneShotClient = client.get(capabilityName) + await oneShotClient?.generate(input) + return oneShotClient?.getResult() ?? null + }, + stop: () => { + client.get(capabilityName)?.stop() + }, + reset: () => { + client.get(capabilityName)?.reset() + }, + } + } + + system.dispose = dispose + + // eslint-disable-next-line no-restricted-syntax -- built dynamically from a runtime `assistant.capabilities` array; the static AssistantSystem shape can't be verified structurally here, plus the added `dispose` field. + return system as unknown as AssistantSystem & { + dispose: () => void + } +} diff --git a/packages/ai-svelte/tests/create-assistant.svelte.test.ts b/packages/ai-svelte/tests/create-assistant.svelte.test.ts new file mode 100644 index 000000000..423dcf398 --- /dev/null +++ b/packages/ai-svelte/tests/create-assistant.svelte.test.ts @@ -0,0 +1,102 @@ +import { describe, expect, it } from 'vitest' +import { defineAssistant } from '@tanstack/ai/assistant' +import { stream } from '@tanstack/ai-client' +import { createAssistant } from '../src/create-assistant.svelte.js' + +// A connection adapter that replays canned chunks for both capabilities, +// branching on the `capability` discriminator forwarded by AssistantClient. +function fakeConnection() { + return stream(async function* (_messages, data) { + const capability = (data as Record | undefined)?.capability + + if (capability === 'chat') { + yield { type: 'RUN_STARTED', threadId: 't', runId: 'r' } as any + yield { + type: 'TEXT_MESSAGE_START', + messageId: 'm1', + role: 'assistant', + } as any + yield { + type: 'TEXT_MESSAGE_CONTENT', + messageId: 'm1', + delta: 'hello', + } as any + yield { type: 'TEXT_MESSAGE_END', messageId: 'm1' } as any + yield { type: 'RUN_FINISHED', threadId: 't', runId: 'r' } as any + } else { + yield { type: 'RUN_STARTED', threadId: 't', runId: 'r' } as any + yield { + type: 'CUSTOM', + name: 'generation:result', + value: { id: 'i', model: 'gpt-image-1', images: [{ url: 'u' }] }, + } as any + yield { type: 'RUN_FINISHED', threadId: 't', runId: 'r' } as any + } + }) +} + +describe('createAssistant', () => { + it('exposes only the declared capabilities', () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({}) as any, + }) + + const system = createAssistant(assistant, { connection: fakeConnection() }) + + expect(system.chat).toBeDefined() + expect(system.image).toBeDefined() + expect((system as any).speech).toBeUndefined() + + system.dispose() + }) + + it('generate() on a one-shot capability populates its result', async () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({}) as any, + }) + + const system = createAssistant(assistant, { connection: fakeConnection() }) + + // generate() resolves to the fresh result... + const generated = await system.image.generate({ prompt: 'a fox' }) + expect(generated?.images[0]?.url).toBe('u') + + // ...and the reactive result state is still populated. + expect(system.image.result?.images[0]?.url).toBe('u') + + system.dispose() + }) + + it('sendMessage() on the chat capability populates messages', async () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({}) as any, + }) + + const system = createAssistant(assistant, { connection: fakeConnection() }) + + // With no outputSchema, sendMessage() resolves to the messages array. + const returned = await system.chat.sendMessage('hi') + expect(returned.length).toBeGreaterThan(0) + + expect(system.chat.messages.length).toBeGreaterThan(0) + + system.dispose() + }) + + it('exposes structured partial/final on the chat surface with no structured part', () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({}) as any, + }) + + const system = createAssistant(assistant, { connection: fakeConnection() }) + + expect((system.chat as any).partial).toEqual({}) + expect((system.chat as any).final).toBeNull() + + system.dispose() + }) +}) diff --git a/packages/ai-vue/package.json b/packages/ai-vue/package.json index 86add75f4..357a6717c 100644 --- a/packages/ai-vue/package.json +++ b/packages/ai-vue/package.json @@ -24,6 +24,10 @@ ".": { "types": "./dist/index.d.ts", "import": "./dist/index.js" + }, + "./assistant": { + "types": "./dist/assistant.d.ts", + "import": "./dist/assistant.js" } }, "files": [ diff --git a/packages/ai-vue/src/assistant.ts b/packages/ai-vue/src/assistant.ts new file mode 100644 index 000000000..6a8666b6e --- /dev/null +++ b/packages/ai-vue/src/assistant.ts @@ -0,0 +1 @@ +export { useAssistant } from './use-assistant' diff --git a/packages/ai-vue/src/use-assistant.ts b/packages/ai-vue/src/use-assistant.ts new file mode 100644 index 000000000..7b312dfce --- /dev/null +++ b/packages/ai-vue/src/use-assistant.ts @@ -0,0 +1,259 @@ +import { + AssistantClient, + computeStructuredParts, +} from '@tanstack/ai-client/assistant' +import { computed, onMounted, onScopeDispose, readonly, shallowRef } from 'vue' +import type { ModelMessage } from '@tanstack/ai' +import type { AssistantDefinition } from '@tanstack/ai/assistant' +import type { + AssistantClientOptions, + AssistantSystem, + OneShotCapabilityName, +} from '@tanstack/ai-client/assistant' +import type { + AnyClientTool, + ChatClientState, + ConnectionStatus, + GenerationClientState, + MultimodalContent, + UIMessage, +} from '@tanstack/ai-client' + +/** Per-capability one-shot generation state, tracked in a single record ref. */ +interface OneShotState { + result: unknown + isLoading: boolean + error: Error | undefined + status: GenerationClientState +} + +const DEFAULT_ONE_SHOT_STATE: OneShotState = { + result: null, + isLoading: false, + error: undefined, + status: 'idle', +} + +/** + * Vue composable for `AssistantDefinition`-based assistants: one composed + * system exposing a typed surface per declared capability (`chat` plus any + * one-shot capabilities like `image`, `speech`, `transcription`, etc.), + * sharing a single connection. + * + * Mirrors the `useChat` idiom: the `AssistantClient` is built eagerly, once, + * with reactive state callbacks wired into its constructor — `ChatClient` + * and `GenerationClient` (the sub-clients `AssistantClient` composes) only + * accept these callbacks via construction, not `updateOptions`, so they must + * be threaded through up front rather than attached after the fact. + * + * @example + * ```vue + * + * + * + * ``` + */ +export function useAssistant< + TDef extends AssistantDefinition, + TChatTools extends ReadonlyArray = [], +>( + assistant: TDef, + options: Omit< + AssistantClientOptions, + 'assistant' | 'callbacks' + >, +): AssistantSystem { + // Chat sub-client state. + const chatMessages = shallowRef>>([]) + const chatIsLoading = shallowRef(false) + const chatError = shallowRef(undefined) + const chatStatus = shallowRef('ready') + const chatIsSubscribed = shallowRef(false) + const chatConnectionStatus = shallowRef('disconnected') + const chatSessionGenerating = shallowRef(false) + + // One-shot capability state, keyed by capability name. A single record ref + // (rather than one ref per capability) keeps the update path uniform + // regardless of how many one-shot capabilities the assistant declares. + const oneShotState = shallowRef>({}) + + const updateOneShot = (capability: string, patch: Partial) => { + const previous = oneShotState.value[capability] ?? DEFAULT_ONE_SHOT_STATE + oneShotState.value = { + ...oneShotState.value, + [capability]: { ...previous, ...patch }, + } + } + + // Build the AssistantClient eagerly, once (no memo). Reactive callbacks are + // passed into the constructor's `callbacks` option — not wired up via + // `updateOptions` afterwards — because the sub-clients (`ChatClient`, + // `GenerationClient`) only accept these callbacks at construction time. + const client = new AssistantClient({ + ...options, + assistant, + callbacks: { + chat: { + onMessagesChange: (messages) => { + chatMessages.value = messages + }, + onLoadingChange: (isLoading) => { + chatIsLoading.value = isLoading + }, + onErrorChange: (error) => { + chatError.value = error + }, + onStatusChange: (status) => { + chatStatus.value = status + }, + onSubscriptionChange: (isSubscribed) => { + chatIsSubscribed.value = isSubscribed + }, + onConnectionStatusChange: (connectionStatus) => { + chatConnectionStatus.value = connectionStatus + }, + onSessionGeneratingChange: (sessionGenerating) => { + chatSessionGenerating.value = sessionGenerating + }, + }, + oneShot: (capability) => ({ + onResultChange: (result) => updateOneShot(capability, { result }), + onLoadingChange: (isLoading) => + updateOneShot(capability, { isLoading }), + onErrorChange: (error) => updateOneShot(capability, { error }), + onStatusChange: (status) => updateOneShot(capability, { status }), + }), + }, + }) + + onMounted(() => { + client.chat?.mountDevtools() + }) + + // Cleanup on unmount: tears down the chat sub-client and every one-shot + // sub-client (stops in-flight requests, unregisters devtools). + onScopeDispose(() => { + client.dispose() + }) + + const structuredParts = computed(() => + computeStructuredParts(chatMessages.value), + ) + + const system: Record = {} + + for (const capability of client.capabilities) { + if (capability === 'chat') { + const chatClient = client.chat + // `client.capabilities` only contains 'chat' when the AssistantClient + // constructor actually built the chat sub-client, so this is always + // defined here — but the field type stays optional, so guard rather + // than assert. + if (!chatClient) continue + + system.chat = { + messages: readonly(chatMessages), + sendMessage: async (content: string | MultimodalContent) => { + await chatClient.sendMessage(content) + const msgs = chatClient.getMessages() + const hasStructured = + msgs.at(-1)?.parts.some((p) => p.type === 'structured-output') ?? + false + // Cast: TS can't structurally narrow the conditional return of + // AssistantChatSurface['sendMessage'] against this runtime branch, + // mirroring the assembled-`system` cast at the end of this composable. + return ( + hasStructured ? computeStructuredParts(msgs).final : msgs + ) as any + }, + append: async (message: ModelMessage | UIMessage) => { + await chatClient.append(message) + }, + reload: async () => { + await chatClient.reload() + }, + stop: () => { + chatClient.stop() + }, + clear: () => { + chatClient.clear() + }, + setMessages: (messages: Array>) => { + chatClient.setMessagesManually(messages) + }, + addToolResult: async (result: { + toolCallId: string + tool: string + output: any + state?: 'output-available' | 'output-error' + errorText?: string + }) => { + await chatClient.addToolResult(result) + }, + addToolApprovalResponse: async (response: { + id: string + approved: boolean + }) => { + await chatClient.addToolApprovalResponse(response) + }, + isLoading: readonly(chatIsLoading), + error: readonly(chatError), + status: readonly(chatStatus), + isSubscribed: readonly(chatIsSubscribed), + connectionStatus: readonly(chatConnectionStatus), + sessionGenerating: readonly(chatSessionGenerating), + // Runtime shape unconditionally exposes partial/final; the public + // AssistantSystem type hides them when the chat capability's + // outputSchema is absent, matching useChat's behavior. + partial: readonly(computed(() => structuredParts.value.partial)), + final: readonly(computed(() => structuredParts.value.final)), + } + continue + } + + const oneShotCapability = capability as OneShotCapabilityName + const oneShotClient = client.get(oneShotCapability) + // Same reasoning as `chatClient` above: declared in `capabilities` + // implies the sub-client was constructed, but the map lookup is + // typed as possibly `undefined`. + if (!oneShotClient) continue + + system[capability] = { + generate: async (input: Record) => { + await oneShotClient.generate(input) + return oneShotClient.getResult() + }, + result: computed(() => oneShotState.value[capability]?.result ?? null), + isLoading: computed( + () => oneShotState.value[capability]?.isLoading ?? false, + ), + error: computed(() => oneShotState.value[capability]?.error), + status: computed(() => oneShotState.value[capability]?.status ?? 'idle'), + stop: () => { + oneShotClient.stop() + }, + reset: () => { + oneShotClient.reset() + }, + } + } + + // The runtime shape (refs nested per capability) diverges from the + // declared `AssistantSystem` type (plain values) — the same divergence + // `useChat` accepts for its return, since consumers unwrap refs via + // `.value` (script) or Vue's template auto-unwrapping. + // eslint-disable-next-line no-restricted-syntax -- composable return shape (nested refs per capability) diverges from the framework-agnostic AssistantSystem type (plain values); TS can't structurally relate the two + return system as unknown as AssistantSystem +} diff --git a/packages/ai-vue/tests/use-assistant.test.ts b/packages/ai-vue/tests/use-assistant.test.ts new file mode 100644 index 000000000..8b1f40f60 --- /dev/null +++ b/packages/ai-vue/tests/use-assistant.test.ts @@ -0,0 +1,152 @@ +import { flushPromises, mount } from '@vue/test-utils' +import { defineComponent } from 'vue' +import { describe, expect, it } from 'vitest' +import { defineAssistant } from '@tanstack/ai/assistant' +import { stream } from '@tanstack/ai-client' +import { useAssistant } from '../src/use-assistant.js' +import { createTextChunks } from './test-utils' +import type { StreamChunk } from '@tanstack/ai' +import type { AssistantDefinition } from '@tanstack/ai/assistant' +import type { ConnectConnectionAdapter } from '@tanstack/ai-client' + +// Helper mirroring the `generation:result` CUSTOM chunk used across the +// ai-vue generation tests (see use-generation.test.ts). +function createGenerationChunks(result: unknown): Array { + return [ + { type: 'RUN_STARTED', runId: 'run-1', timestamp: Date.now() }, + { + type: 'CUSTOM', + name: 'generation:result', + value: result, + timestamp: Date.now(), + }, + { + type: 'RUN_FINISHED', + runId: 'run-1', + finishReason: 'stop', + timestamp: Date.now(), + }, + ] as unknown as Array +} + +/** + * A single canned connection shared by chat and one-shot sub-clients. + * `AssistantClient` tags each sub-client's requests with `capability` (via + * `forwardedProps` for chat, `body` for one-shot), so a single `stream()` + * adapter can route on `data.capability` the same way a real server handler + * (`defineAssistant(...).handler`) would. + */ +function createAssistantConnection() { + return stream(async function* (_messages, data) { + if (data?.capability === 'chat') { + yield* createTextChunks('Hello from assistant') + return + } + if (data?.capability === 'image') { + yield* createGenerationChunks({ + id: 'img-1', + model: 'test', + images: ['a.png'], + }) + return + } + }) +} + +function renderUseAssistant( + assistant: AssistantDefinition, + options: { connection: ConnectConnectionAdapter }, +) { + const TestComponent = defineComponent({ + setup() { + return { system: useAssistant(assistant, options) } + }, + template: '
', + }) + + const wrapper = mount(TestComponent) + return { wrapper, system: wrapper.vm.system as any } +} + +describe('useAssistant', () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({ id: '', model: '', images: [] }) as any, + }) + + it('exposes only the declared capabilities', () => { + const { system } = renderUseAssistant(assistant, { + connection: createAssistantConnection(), + }) + + expect(system.chat).toBeDefined() + expect(system.image).toBeDefined() + expect(system.speech).toBeUndefined() + }) + + it('exposes structured partial/final with no structured part', () => { + const { system } = renderUseAssistant(assistant, { + connection: createAssistantConnection(), + }) + + expect(system.chat.partial.value).toEqual({}) + expect(system.chat.final.value).toBeNull() + }) + + it('sendMessage populates chat messages', async () => { + const { system } = renderUseAssistant(assistant, { + connection: createAssistantConnection(), + }) + + // With no outputSchema, sendMessage() resolves to the messages array. + const returned = await system.chat.sendMessage('Hi') + expect(returned.length).toBeGreaterThan(0) + await flushPromises() + + expect(system.chat.messages.value.length).toBeGreaterThan(0) + const assistantMessage = system.chat.messages.value.find( + (m: any) => m.role === 'assistant', + ) + expect(assistantMessage).toBeDefined() + const textPart = assistantMessage?.parts.find((p: any) => p.type === 'text') + expect(textPart?.content).toBe('Hello from assistant') + }) + + it('generate populates the one-shot result', async () => { + const { system } = renderUseAssistant(assistant, { + connection: createAssistantConnection(), + }) + + expect(system.image.result.value).toBeNull() + expect(system.image.isLoading.value).toBe(false) + + // generate() resolves to the fresh result... + const generated = await system.image.generate({ prompt: 'a cat' }) + expect(generated).toEqual({ + id: 'img-1', + model: 'test', + images: ['a.png'], + }) + await flushPromises() + + // ...and the reactive result ref is still populated. + expect(system.image.result.value).toEqual({ + id: 'img-1', + model: 'test', + images: ['a.png'], + }) + expect(system.image.isLoading.value).toBe(false) + expect(system.image.status.value).toBe('success') + }) + + it('disposes the underlying client on unmount without throwing', async () => { + const { wrapper, system } = renderUseAssistant(assistant, { + connection: createAssistantConnection(), + }) + + await system.chat.sendMessage('Hi') + await flushPromises() + + expect(() => wrapper.unmount()).not.toThrow() + }) +}) diff --git a/packages/ai-vue/tsdown.config.ts b/packages/ai-vue/tsdown.config.ts index 3c1187801..de631c408 100644 --- a/packages/ai-vue/tsdown.config.ts +++ b/packages/ai-vue/tsdown.config.ts @@ -1,7 +1,7 @@ import { defineConfig } from 'tsdown' export default defineConfig({ - entry: ['./src/index.ts'], + entry: ['./src/index.ts', './src/assistant.ts'], format: ['esm'], unbundle: true, dts: true, diff --git a/packages/ai/package.json b/packages/ai/package.json index a1dcf9a97..07b04624a 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -25,6 +25,10 @@ "types": "./dist/esm/index.d.ts", "import": "./dist/esm/index.js" }, + "./assistant": { + "types": "./dist/esm/assistant.d.ts", + "import": "./dist/esm/assistant.js" + }, "./client": { "types": "./dist/esm/client.d.ts", "import": "./dist/esm/client.js" diff --git a/packages/ai/skills/ai-core/assistant/SKILL.md b/packages/ai/skills/ai-core/assistant/SKILL.md new file mode 100644 index 000000000..e6ccbed9a --- /dev/null +++ b/packages/ai/skills/ai-core/assistant/SKILL.md @@ -0,0 +1,367 @@ +--- +name: ai-core/assistant +description: > + Multi-capability assistant composition: defineAssistant() registers + per-capability callbacks (chat/image/audio/speech/video/transcription/ + summarize) on the server behind one handler; useAssistant() consumes all + declared capabilities from a single client hook with no generics, types + inferred from the server definition. Each capability entry mirrors the + matching primitive hook (useChat, useGenerateImage, etc.) exactly. +type: sub-skill +library: tanstack-ai +library_version: '0.10.0' +sources: + - 'TanStack/ai:packages/ai/src/activities/assistant/index.ts' + - 'TanStack/ai:packages/ai/src/activities/assistant/types.ts' + - 'TanStack/ai:packages/ai-client/src/assistant-client.ts' + - 'TanStack/ai:packages/ai-client/src/assistant-types.ts' + - 'TanStack/ai:packages/ai-react/src/use-assistant.ts' +--- + +# Assistant + +This skill builds on ai-core, ai-core/chat-experience, and +ai-core/media-generation. Read them first. + +`defineAssistant` + `useAssistant` are a **composition layer**, not a new +activity or wire format. They wire together activities you already know +(`chat()`, `generateImage()`, `generateSpeech()`, …) behind one server +endpoint and one client hook. There is no new client state machine: each +declared capability on the client is exactly the same surface as the +matching primitive hook (`useChat`, `useGenerateImage`, `useGenerateAudio`, +`useGenerateSpeech`, `useGenerateVideo`, `useTranscription`, `useSummarize`). + +## Setup — Chat + Image Assistant End-to-End + +### Server: `defineAssistant` + a single `handler` + +Each capability key maps to a callback `(req) => `. The +definition is **inert** — `defineAssistant` only stores the callbacks and a +static list of declared capability names. It constructs nothing (no +adapters, no connections) until a request actually reaches `handler`, so +it's safe to import into an isomorphic module shared with the client. + +```typescript +// src/lib/assistant.ts — shared/isomorphic module +import { chat, generateImage, generateSpeech } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { + openaiText, + openaiImage, + openaiSpeech, +} from '@tanstack/ai-openai/adapters' +import { getWeather } from './tools' + +export const blogAssistant = defineAssistant({ + chat: (req) => + chat({ + adapter: openaiText('gpt-5.5'), + messages: req.messages, + threadId: req.threadId, + runId: req.runId, + tools: [getWeather], + }), + + image: (req) => + generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: req.prompt, + size: req.size, + numberOfImages: req.numberOfImages, + }), + + speech: (req) => + generateSpeech({ + adapter: openaiSpeech('tts-1'), + text: req.text, + voice: req.voice ?? 'alloy', + }), +}) +``` + +```typescript +// src/routes/api.assistant.ts — server route, single handler +import { createFileRoute } from '@tanstack/react-router' +import { assistant } from '../lib/assistant' + +export const Route = createFileRoute('/api/assistant')({ + server: { + handlers: { + POST: (request) => blogAssistant.handler(request), + }, + }, +}) +``` + +`handler` is the **only** thing the route needs to call. Internally it +routes by a `capability` discriminator carried on the request body: `'chat'` +dispatches to the AG-UI `RunAgentInput` parsing path (same as a standalone +chat route), and every other declared key is a one-shot generation request +parsed into that capability's input shape. Unknown or undeclared +capabilities get a `400` before any callback runs. + +### Client: `useAssistant` + +```tsx +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { assistant } from '../lib/assistant' + +function AssistantPanel() { + const assistant = useAssistant(blogAssistant, { + connection: fetchServerSentEvents('/api/assistant'), + }) + + return ( +
+
+ {assistant.chat.messages.map((message) => ( +
{message.role}
+ ))} +
+ + + + {assistant.image.result?.images.map((img, i) => ( + + ))} +
+ ) +} +``` + +`assistant` is typed from the `assistant` value passed in — **no generics** at +the call site. Only capabilities declared in `defineAssistant` appear on +`assistant`; referencing an undeclared key is a compile error. `assistant.chat` +is the full `useChat` return (`messages`, `sendMessage`, `isLoading`, +`error`, `status`, `stop`, `clear`, `addToolResult`, …); `assistant.image` / +`assistant.audio` / `assistant.speech` / `assistant.video` / +`assistant.transcription` / `assistant.summarize` are each the full +`useGeneration`-style return (`generate`, `result`, `isLoading`, `error`, +`status`, `stop`, `reset`). + +Vue/Solid/Svelte have identical patterns with different hook imports +(e.g. `import { useAssistant } from '@tanstack/ai-solid/assistant'`). + +## Core Patterns + +### 1. Chat tools auto-type from the server callback; `chat: { tools }` is only for client-executed runtime + +Tools passed to `chat({ tools: [...] })` inside the server callback +automatically type `assistant.chat.messages`' tool-call/result parts — +narrowed by tool name, input, and output — with **no** client-side +re-declaration needed for typing: + +```typescript +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { blogAssistant } from '../lib/assistant' // chat callback passed tools: [getWeather] + +const assistant = useAssistant(blogAssistant, { + connection: fetchServerSentEvents('/api/assistant'), +}) + +// assistant.chat.messages parts are already narrowed by tool name — inferred +// from the server callback's `tools: [getWeather]`, not from anything passed +// here. +``` + +`chat: { tools }` on `useAssistant` still exists, but only for one reason: +a client-**executed** tool's `.client()` implementation runs in the browser, +so its code can't cross the wire — the server callback only ever sees the +tool's _definition_ (for the model and for typing). Pass the client +implementation there to register its runtime: + +```typescript +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { blogAssistant } from '../lib/assistant' // chat callback passed tools: [showToastDef] +import { showToast } from '../lib/tools' // showToastDef.client((input) => ...) + +const assistant = useAssistant(blogAssistant, { + connection: fetchServerSentEvents('/api/assistant'), + chat: { tools: [showToast] }, // runtime only — types already came from the callback +}) +``` + +`chat.forwardedProps` is also available on the same option, merged into +every chat request alongside the reserved `capability` field. + +### 2. Structured output via `outputSchema` in the chat callback + +If the `chat` callback passes `outputSchema` to `chat()`, `assistant.chat` +picks up typed `partial` (progressive `DeepPartial`) and `final` (validated +terminal object) fields — the same conditional shape `useChat({ +outputSchema })` returns. Omit `outputSchema` and neither field is present on +the type. + +```typescript +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { blogAssistant } from '../lib/assistant' // chat callback passed outputSchema: BlogOutlineSchema + +const assistant = useAssistant(blogAssistant, { + connection: fetchServerSentEvents('/api/assistant'), +}) + +assistant.chat.partial.title // string | undefined — fills in as JSON streams +assistant.chat.final // full schema type | null — set once the run completes +``` + +### 3. Manual chaining across capabilities + +There is no auto-chaining or shared "artifact workspace" — `useAssistant` +is a pure composition layer. To use one capability's result as input to +another, read the result value and pass it explicitly: + +```typescript +await assistant.image.generate({ prompt: 'a lighthouse at dusk' }) + +// after assistant.image.result is populated: +if (assistant.image.result) { + assistant.chat.sendMessage({ + content: [ + { + type: 'image', + source: { type: 'url', value: assistant.image.result.images[0].url }, + }, + { type: 'text', content: 'Write a short caption for this image.' }, + ], + }) +} +``` + +### 4. Only declared capabilities are constructed + +`defineAssistant({ chat, image })` produces a client `assistant` with exactly +`assistant.chat` and `assistant.image` — no `assistant.audio`, `assistant.video`, etc. +Add or remove capability keys on the server definition to change what the +client can call; there's nothing else to keep in sync. + +### 5. Sharing one connection across capabilities + +All capabilities declared in one `defineAssistant` call share a single +`connection` passed to `useAssistant` — one endpoint, one adapter. Each +underlying sub-client (a `ChatClient` for `chat`, a `GenerationClient` per +one-shot capability) tags its own requests with the capability name, so the +single `handler` on the server can route correctly. This mirrors the shared +`connection` pattern already used by `useChat` / `useGenerateImage` +individually — see ai-core/chat-experience and ai-core/media-generation. + +## Common Mistakes + +### a. HIGH: Constructing adapters outside the capability callback + +```typescript +// WRONG — adapter constructed eagerly at module load, defeats "inert" guarantee +const textAdapter = openaiText('gpt-5.5') +const blogAssistant = defineAssistant({ + chat: (req) => chat({ adapter: textAdapter, messages: req.messages }), +}) + +// CORRECT — construct inside the callback, per request +const blogAssistant = defineAssistant({ + chat: (req) => + chat({ adapter: openaiText('gpt-5.5'), messages: req.messages }), +}) +``` + +Adapter construction is cheap (it wraps config; the provider HTTP client is +created lazily), but constructing outside the callback breaks the +"`defineAssistant` is inert, safe to import into the client bundle" +guarantee — it now runs provider setup at import time. + +### b. HIGH: Writing a custom handler that branches on capability manually + +```typescript +// WRONG — reimplementing what `handler` already does +export const POST = async (request: Request) => { + const body = await request.json() + if (body.capability === 'chat') { + return toServerSentEventsResponse( + chat({ adapter, messages: body.messages }), + ) + } + // ...manual branching for every capability +} + +// CORRECT — defineAssistant's handler already parses, routes, and serializes +export const POST = (request: Request) => blogAssistant.handler(request) +``` + +### c. HIGH: Passing `model` or top-level generation options to `useAssistant` + +There is no `model`/`tools`-for-generation option on `useAssistant` itself. +Model choice and per-capability options belong inside the server callback +(`openaiImage('gpt-image-2')`, `openaiSpeech('tts-1')`, …); the +only client-side option `useAssistant` accepts besides `connection` is +`chat: { tools, forwardedProps }` (`tools` here registers **client-executed** +tools' runtime implementations only — typing already comes from the server +callback) and `threadId`/`id`. + +```typescript +// WRONG — no such options on useAssistant +useAssistant(blogAssistant, { connection, model: 'gpt-5.5', prompt: '...' }) + +// CORRECT — model/prompt options live in the server callback; useAssistant +// only takes connection, threadId/id, and chat.{tools,forwardedProps} +useAssistant(blogAssistant, { + connection: fetchServerSentEvents('/api/assistant'), +}) +``` + +### d. MEDIUM: Expecting `defineAssistant` to auto-chain results + +```typescript +// WRONG — assuming the assistant remembers assistant.image.result automatically +assistant.chat.sendMessage('use the image I just generated') + +// CORRECT — thread the result value through explicitly (see Core Pattern 2) +assistant.chat.sendMessage({ + content: [ + { + type: 'image', + source: { type: 'url', value: assistant.image.result.images[0].url }, + }, + { type: 'text', content: 'Use this image.' }, + ], +}) +``` + +Chaining is manual by design (v1 non-goal: no shared artifact workspace or +auto-resolution of prior outputs). + +### e. MEDIUM: Declaring a capability on the server but never checking it client-side + +Every capability declared in `defineAssistant` is unconditionally present on +`assistant` — there's no need to guard with `assistant.image?.generate`. If a +capability is optional per-deployment, omit the key from the +`defineAssistant` config entirely rather than declaring it and ignoring it +on the client. + +## Cross-References + +- See also: **ai-core/chat-experience/SKILL.md** -- `assistant.chat` is the + same `useChat` surface; streaming, tool rendering, and multimodal + messages all apply unchanged. +- See also: **ai-core/media-generation/SKILL.md** -- `assistant.image` / + `assistant.audio` / `assistant.speech` / `assistant.video` / + `assistant.transcription` / `assistant.summarize` are each the same + `useGeneration`-style surface documented there. +- See also: **ai-core/tool-calling/SKILL.md** -- Tools passed to the `chat` + capability callback follow the same server/client tool patterns as a + standalone `chat()` route. +- See also: **ai-core/custom-backend-integration/SKILL.md** -- The shared + `connection` passed to `useAssistant` is the same `ConnectConnectionAdapter` + used by `useChat` / `useGenerate*`. diff --git a/packages/ai/skills/ai-core/chat-experience/SKILL.md b/packages/ai/skills/ai-core/chat-experience/SKILL.md index 3ae4408dc..2aa0606a5 100644 --- a/packages/ai/skills/ai-core/chat-experience/SKILL.md +++ b/packages/ai/skills/ai-core/chat-experience/SKILL.md @@ -620,3 +620,7 @@ If not handled, the UI appears to hang with no feedback. - See also: **ai-core/tool-calling/SKILL.md** -- Most chats include tools - See also: **ai-core/adapter-configuration/SKILL.md** -- Adapter choice affects available features - See also: **ai-core/middleware/SKILL.md** -- Use middleware for analytics and lifecycle events +- See also: **ai-core/assistant/SKILL.md** -- To compose chat with other + capabilities (image/audio/speech/video/transcription/summarize) behind a + single `defineAssistant` handler + `useAssistant` hook instead of wiring up + `chat()` and `useChat` alone diff --git a/packages/ai/skills/ai-core/custom-backend-integration/SKILL.md b/packages/ai/skills/ai-core/custom-backend-integration/SKILL.md index 37275b6f7..3cb3fb1a3 100644 --- a/packages/ai/skills/ai-core/custom-backend-integration/SKILL.md +++ b/packages/ai/skills/ai-core/custom-backend-integration/SKILL.md @@ -461,3 +461,7 @@ Source: `docs/protocol/http-stream-protocol.md` - See also: **ai-core/ag-ui-protocol/SKILL.md** -- Understanding the AG-UI protocol helps build compatible custom servers - See also: **ai-core/chat-experience/SKILL.md** -- Full chat setup patterns including server-side `chat()` and `toServerSentEventsResponse()` - See also: **ai-core/middleware/SKILL.md** -- Use middleware for analytics and lifecycle events on the server side +- See also: **ai-core/assistant/SKILL.md** -- If the backend needs to compose + chat with other capabilities (image/audio/speech/video/transcription/ + summarize) behind one endpoint, `defineAssistant` + `useAssistant` may fit + better than a custom connection adapter wrapping `useChat` alone diff --git a/packages/ai/skills/ai-core/media-generation/SKILL.md b/packages/ai/skills/ai-core/media-generation/SKILL.md index 268a85f74..b0e6aaefc 100644 --- a/packages/ai/skills/ai-core/media-generation/SKILL.md +++ b/packages/ai/skills/ai-core/media-generation/SKILL.md @@ -876,3 +876,7 @@ and piping into a custom logger. returns unexpected output or fails mid-stream, toggle `debug: true` on any `generate*()` call to see request metadata, raw provider chunks, and errors. Covers per-category toggling and piping into pino/winston. +- See also: **ai-core/assistant/SKILL.md** -- To combine multiple media + activities (image/audio/speech/video/transcription) with chat behind one + `defineAssistant` handler + `useAssistant` hook instead of wiring up each + `generate*()` activity and its own hook separately. diff --git a/packages/ai/src/activities/assistant/index.ts b/packages/ai/src/activities/assistant/index.ts new file mode 100644 index 000000000..bb2422ac9 --- /dev/null +++ b/packages/ai/src/activities/assistant/index.ts @@ -0,0 +1,131 @@ +import { chatParamsFromRequestBody } from '../../utilities/chat-params.js' +import { toServerSentEventsResponse } from '../../stream-to-response.js' +import { streamGenerationResult } from '../stream-generation-result.js' +import type { + AssistantChatRequest, + AssistantConfig, + AssistantDefinition, +} from './types.js' + +export type * from './types.js' + +function isAsyncIterable(value: unknown): value is AsyncIterable { + return ( + value != null && typeof value === 'object' && Symbol.asyncIterator in value + ) +} + +/** + * Define a server-side assistant: a registry of per-capability callbacks + * exposing a single request `handler`. Inert — no adapters, connections, + * or other resources are constructed until a request actually arrives. + */ +export function defineAssistant( + config: T, +): AssistantDefinition { + const capabilities = Object.keys(config) as Array + + const handler = async (request: Request): Promise => { + let body: unknown + try { + body = await request.json() + } catch { + return new Response('Invalid JSON body', { status: 400 }) + } + + const bodyRecord = body as Record + const forwardedProps: Record = + (bodyRecord.forwardedProps as Record | undefined) ?? + (bodyRecord.data as Record | undefined) ?? + {} + const capability = forwardedProps.capability + // Gate on OWN, DECLARED capabilities only — `capabilities` is + // `Object.keys(config)`, so this excludes inherited `Object.prototype` + // members (`toString`, `valueOf`, `hasOwnProperty`, `constructor`, …) + // that `capability in config` / `Reflect.get(config, capability)` would + // otherwise treat as valid callbacks. + const isDeclaredCapability = + typeof capability === 'string' && + (capabilities as Array).includes(capability) + + // `Reflect.get` returns `any`, so this is a single narrowing cast (not + // `as unknown as`) — a per-key cast from `AssistantConfig` directly + // fails structurally because callback params are contravariant per key. + const callback = isDeclaredCapability + ? (Reflect.get(config, capability) as ((req: any) => unknown) | undefined) + : undefined + + if (!isDeclaredCapability || !callback) { + return new Response( + `Unknown assistant capability: ${String(capability)}`, + { status: 400 }, + ) + } + + if (capability === 'chat') { + // chatParamsFromRequestBody rejects with an AGUIError (not a thrown + // Response) when the body doesn't conform to AG-UI's RunAgentInput + // shape. Surface that as a 400, mirroring the e2e `api.chat.ts` route. + let params + try { + params = await chatParamsFromRequestBody(body) + } catch (error) { + if (error instanceof Response) return error + return new Response( + error instanceof Error ? error.message : 'Invalid request body', + { status: 400 }, + ) + } + const chatReq: AssistantChatRequest = { ...params, request } + const result = callback(chatReq) + if (!isAsyncIterable(result)) { + return new Response( + 'assistant chat capability must return a streaming ChatStream', + { status: 400 }, + ) + } + return toServerSentEventsResponse(result as AsyncIterable) + } + + // One-shot capabilities: validate the AG-UI envelope (same as `chat`), + // then pull the generation input out of forwardedProps. + let params + try { + params = await chatParamsFromRequestBody(body) + } catch (error) { + if (error instanceof Response) return error + return new Response( + error instanceof Error ? error.message : 'Invalid request body', + { status: 400 }, + ) + } + const { capability: _omit, ...input } = params.forwardedProps + const oneShotReq = { + ...input, + threadId: params.threadId, + runId: params.runId, + ...(params.parentRunId !== undefined && { + parentRunId: params.parentRunId, + }), + state: params.state, + aguiContext: params.aguiContext, + forwardedProps: params.forwardedProps, + request, + } + const result = callback(oneShotReq) + + if (isAsyncIterable(result)) { + // User opted into stream:true; the activity already emits generation:result. + return toServerSentEventsResponse(result as AsyncIterable) + } + // Non-streaming activity promise → wrap so the client's GenerationClient sees it. + return toServerSentEventsResponse( + streamGenerationResult(() => Promise.resolve(result), { + threadId: params.threadId, + runId: params.runId, + }), + ) + } + + return { capabilities, handler, '~caps': config } +} diff --git a/packages/ai/src/activities/assistant/types.ts b/packages/ai/src/activities/assistant/types.ts new file mode 100644 index 000000000..3f3efb05a --- /dev/null +++ b/packages/ai/src/activities/assistant/types.ts @@ -0,0 +1,118 @@ +// packages/ai/src/activities/assistant/types.ts +import type { Context as AGUIContext } from '@ag-ui/core' +import type { + AudioGenerationResult, + ChatStream, + ImageGenerationResult, + JSONSchema, + ModelMessage, + StreamChunk, + StructuredOutputStream, + SummarizationResult, + TTSResult, + TranscriptionResult, + UIMessage, + VideoJobResult, +} from '../../types.js' + +/** The parsed request handed to the `chat` capability callback. */ +export interface AssistantChatRequest { + messages: Array + threadId: string + runId: string + parentRunId?: string + /** Client-declared tools (name/description/JSON-schema) from the AG-UI body. */ + tools: Array<{ name: string; description: string; parameters: JSONSchema }> + forwardedProps: Record + state: unknown + aguiContext: Array + /** Raw request, for escape hatches (headers, auth). */ + request: Request +} + +/** Base for one-shot capability requests: parsed generation input + raw request. */ +interface AssistantOneShotBase { + threadId: string + runId: string + parentRunId?: string + state: unknown + aguiContext: Array + forwardedProps: Record + request: Request +} + +export interface AssistantImageRequest extends AssistantOneShotBase { + prompt: unknown + numberOfImages?: number + size?: unknown + modelOptions?: Record +} +export interface AssistantAudioRequest extends AssistantOneShotBase { + prompt: unknown + duration?: number + modelOptions?: Record +} +export interface AssistantSpeechRequest extends AssistantOneShotBase { + text: string + voice?: string + format?: 'mp3' | 'opus' | 'aac' | 'flac' | 'wav' | 'pcm' + speed?: number + modelOptions?: Record +} +export interface AssistantVideoRequest extends AssistantOneShotBase { + prompt: unknown + size?: unknown + duration?: unknown + modelOptions?: Record +} +export interface AssistantTranscriptionRequest extends AssistantOneShotBase { + audio: string | File | Blob | ArrayBuffer + language?: string + prompt?: string + responseFormat?: 'json' | 'text' | 'srt' | 'verbose_json' | 'vtt' + modelOptions?: Record +} +export interface AssistantSummarizeRequest extends AssistantOneShotBase { + text: string + maxLength?: number + style?: 'bullet-points' | 'paragraph' | 'concise' + focus?: Array + modelOptions?: Record +} + +/** A capability callback returns either a stream or a promise of the activity result. */ +type MaybeStream = AsyncIterable | Promise + +/** The shape a user passes to `defineAssistant`. Every key optional. */ +export interface AssistantConfig { + chat?: ( + req: AssistantChatRequest, + ) => + | ChatStream + | StructuredOutputStream + | Promise + | Promise + image?: (req: AssistantImageRequest) => MaybeStream + audio?: (req: AssistantAudioRequest) => MaybeStream + speech?: (req: AssistantSpeechRequest) => MaybeStream + video?: (req: AssistantVideoRequest) => MaybeStream + transcription?: ( + req: AssistantTranscriptionRequest, + ) => MaybeStream + summarize?: ( + req: AssistantSummarizeRequest, + ) => MaybeStream +} + +export type AssistantCapabilityName = keyof AssistantConfig + +export interface AssistantDefinition< + T extends AssistantConfig = AssistantConfig, +> { + /** The declared capability names (for the client to enumerate). */ + readonly capabilities: ReadonlyArray + /** Single request handler; routes by the `capability` discriminator. */ + handler: (request: Request) => Promise + /** Type-only carrier of the config for client inference. Never read at runtime. */ + readonly '~caps': T +} diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 93e7a6481..e61dc5f21 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -45,6 +45,7 @@ import type { import type { AgentLoopStrategy, AnyTool, + ChatResultMeta, ChatStream, ConstrainedModelMessage, CustomEvent, @@ -2712,7 +2713,8 @@ export function chat< TTools, TMiddleware >, -): TextActivityResult { +): TextActivityResult & + ChatResultMeta { validateCapabilities(options.middleware ?? [], options.adapter) const { outputSchema, stream } = options @@ -2721,36 +2723,57 @@ export function chat< // output. Without an explicit `stream: true`, schema-bearing calls run the // agent loop and resolve to a typed Promise>. if (outputSchema && stream === true) { + // phantom-only cast: ~chatMeta never exists at runtime; carries + // TTools/TSchema/TStream for downstream inference. The `unknown` + // hop is required because intersecting the deferred conditional + // `TextActivityResult` with the (all-optional) + // `ChatResultMeta` makes TS eagerly compare the concrete runtime + // return type against every branch of the conditional instead of + // deferring, which trips the overlap check. + // eslint-disable-next-line no-restricted-syntax -- phantom-only cast, see comment above return runStreamingStructuredOutput({ ...options, outputSchema, stream, - }) as TextActivityResult + }) as unknown as TextActivityResult & + ChatResultMeta } // If outputSchema is provided, run agentic structured output (Promise) if (outputSchema) { + // phantom-only cast: ~chatMeta never exists at runtime; carries + // TTools/TSchema/TStream for downstream inference (see note above). + // eslint-disable-next-line no-restricted-syntax -- phantom-only cast, see comment above return runAgenticStructuredOutput({ ...options, outputSchema, - }) as TextActivityResult + }) as unknown as TextActivityResult & + ChatResultMeta } // If stream is explicitly false, run non-streaming text if (stream === false) { + // phantom-only cast: ~chatMeta never exists at runtime; carries + // TTools/TSchema/TStream for downstream inference (see note above). + // eslint-disable-next-line no-restricted-syntax -- phantom-only cast, see comment above return runNonStreamingText({ ...options, outputSchema: undefined, stream, - }) as TextActivityResult + }) as unknown as TextActivityResult & + ChatResultMeta } // Otherwise, run streaming text (default) + // phantom-only cast: ~chatMeta never exists at runtime; carries + // TTools/TSchema/TStream for downstream inference (see note above). + // eslint-disable-next-line no-restricted-syntax -- phantom-only cast, see comment above return runStreamingText({ ...options, outputSchema: undefined, stream, - }) as TextActivityResult + }) as unknown as TextActivityResult & + ChatResultMeta } /** diff --git a/packages/ai/src/assistant.ts b/packages/ai/src/assistant.ts new file mode 100644 index 000000000..377478178 --- /dev/null +++ b/packages/ai/src/assistant.ts @@ -0,0 +1,13 @@ +export { defineAssistant } from './activities/assistant/index' +export type { + AssistantConfig, + AssistantDefinition, + AssistantCapabilityName, + AssistantChatRequest, + AssistantImageRequest, + AssistantAudioRequest, + AssistantSpeechRequest, + AssistantVideoRequest, + AssistantTranscriptionRequest, + AssistantSummarizeRequest, +} from './activities/assistant/index' diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 102987c02..f043b53cc 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -159,6 +159,16 @@ export type { // All types export * from './types' +// chat() result phantom (tool/schema/stream inference for downstream +// consumers like useAssistant) — re-exported explicitly for discoverability +// alongside the other activity-specific type blocks above. +export type { + ChatResultMeta, + InferChatSchema, + InferChatStream, + InferChatTools, +} from './types' + export { firstSentence, renderLazyCatalogEntry, diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 983e220e8..892024ac6 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -1527,6 +1527,40 @@ export type StructuredOutputStream = AsyncIterable< | ToolInputAvailableEvent > +/** Type-only carrier attached to chat()'s return so downstream consumers + * (useAssistant) can recover the tool set + structured-output schema without + * the info being erased. Never present at runtime — optional phantom. */ +export interface ChatResultMeta { + readonly '~chatMeta'?: { + tools: TTools + schema: TSchema + stream: TStream + } +} + +/** Recovers the tool set passed to `chat({ tools })` from its return type. */ +export type InferChatTools = T extends { + '~chatMeta'?: { tools: infer TTools } | undefined +} + ? TTools + : never + +/** Recovers the `outputSchema` passed to `chat({ outputSchema })` from its + * return type. `undefined` when no schema was provided. */ +export type InferChatSchema = T extends { + '~chatMeta'?: { schema: infer TSchema } | undefined +} + ? TSchema + : undefined + +/** Recovers the `stream` option passed to `chat({ stream })` from its return + * type. */ +export type InferChatStream = T extends { + '~chatMeta'?: { stream: infer TStream } | undefined +} + ? TStream + : boolean + // ============================================================================ // AG-UI Reasoning Event Interfaces // ============================================================================ diff --git a/packages/ai/tests/activities/assistant/index.test.ts b/packages/ai/tests/activities/assistant/index.test.ts new file mode 100644 index 000000000..2b1bf4539 --- /dev/null +++ b/packages/ai/tests/activities/assistant/index.test.ts @@ -0,0 +1,140 @@ +import { describe, expect, it, vi } from 'vitest' +import { defineAssistant } from '../../../src/activities/assistant/index.js' +import type { AssistantConfig } from '../../../src/activities/assistant/types.js' +import type { ChatStream } from '../../../src/types.js' + +// Compile-only: a streaming chat callback (the shape `chat()` returns by +// default) must be assignable to `AssistantConfig['chat']` without a cast. +// Regression test for the `any`-collapse bug where `TextActivityResult` silently dropped the default `ChatStream` branch. +const _chatCbAssignable: AssistantConfig['chat'] = (_req) => + undefined as unknown as ChatStream +void _chatCbAssignable + +describe('defineAssistant', () => { + it('is inert: does not invoke any capability callback at define time', () => { + const chatCb = vi.fn() + const imageCb = vi.fn() + const assistant = defineAssistant({ chat: chatCb, image: imageCb }) + expect(chatCb).not.toHaveBeenCalled() + expect(imageCb).not.toHaveBeenCalled() + expect(assistant.capabilities.slice().sort()).toEqual(['chat', 'image']) + expect(typeof assistant.handler).toBe('function') + }) + + it('is exported from the assistant subpath entry', async () => { + const mod = await import('../../../src/assistant.js') + expect(typeof mod.defineAssistant).toBe('function') + }, 30000) +}) + +function runAgentBody(capability: string, extra: Record = {}) { + return { + threadId: 't1', + runId: 'r1', + state: {}, + messages: [], + tools: [], + context: [], + forwardedProps: { capability, ...extra }, + data: { capability, ...extra }, + } +} + +function req(body: unknown) { + return new Request('http://x/api/assistant', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify(body), + }) +} + +async function readSse(res: Response): Promise> { + const text = await res.text() + return text + .split('\n\n') + .map((l) => l.replace(/^data: /, '').trim()) + .filter(Boolean) + .map((l) => JSON.parse(l)) +} + +describe('assistant.handler', () => { + it('400s on unknown capability', async () => { + const assistant = defineAssistant({ chat: async function* () {} as any }) + const res = await assistant.handler(req(runAgentBody('image'))) + expect(res.status).toBe(400) + }) + + it('400s on an inherited Object.prototype key used as capability', async () => { + const assistant = defineAssistant({ chat: async function* () {} as any }) + const res = await assistant.handler(req(runAgentBody('toString'))) + expect(res.status).toBe(400) + }) + + it('routes chat and streams the callback iterable', async () => { + const assistant = defineAssistant({ + chat: async function* (r: any) { + yield { + type: 'RUN_STARTED', + threadId: r.threadId, + runId: r.runId, + } as any + yield { + type: 'RUN_FINISHED', + threadId: r.threadId, + runId: r.runId, + } as any + }, + }) + const res = await assistant.handler(req(runAgentBody('chat'))) + expect(res.headers.get('content-type')).toContain('text/event-stream') + const chunks = await readSse(res) + expect(chunks[0].type).toBe('RUN_STARTED') + }) + + it('wraps a one-shot promise as a generation:result CUSTOM event', async () => { + const assistant = defineAssistant({ + image: async (r: any) => ({ + id: 'img1', + model: 'gpt-image-1', + images: [{ url: `generated:${String(r.prompt)}` }], + }), + }) + const res = await assistant.handler( + req(runAgentBody('image', { prompt: 'a fox' })), + ) + const chunks = await readSse(res) + const custom = chunks.find((c) => c.type === 'CUSTOM') + expect(custom.name).toBe('generation:result') + expect(custom.value.images[0].url).toBe('generated:a fox') + }) + + it('gives a one-shot capability callback the validated AG-UI envelope', async () => { + const assistant = defineAssistant({ + image: async (r: any) => ({ + id: 'img1', + model: 'gpt-image-1', + // Echo the envelope fields into the result so we can assert on + // them via the streamed generation:result CUSTOM event. + images: [ + { + url: `${String(r.prompt)}|${r.threadId}|${r.runId}|${JSON.stringify(r.state)}|${JSON.stringify(r.aguiContext)}|${r.forwardedProps.capability}`, + }, + ], + }), + }) + const body = runAgentBody('image', { prompt: 'a fox' }) + const res = await assistant.handler(req(body)) + const chunks = await readSse(res) + const custom = chunks.find((c) => c.type === 'CUSTOM') + expect(custom.name).toBe('generation:result') + const [prompt, threadId, runId, state, aguiContext, capability] = + custom.value.images[0].url.split('|') + expect(prompt).toBe('a fox') + expect(threadId).toBe(body.threadId) + expect(runId).toBe(body.runId) + expect(state).toBe(JSON.stringify(body.state)) + expect(aguiContext).toBe(JSON.stringify(body.context)) + expect(capability).toBe('image') + }) +}) diff --git a/packages/ai/tests/activities/chat/chat-result-meta.test.ts b/packages/ai/tests/activities/chat/chat-result-meta.test.ts new file mode 100644 index 000000000..bc8e96ff7 --- /dev/null +++ b/packages/ai/tests/activities/chat/chat-result-meta.test.ts @@ -0,0 +1,97 @@ +/** + * Type-only tests for the `chat()` result phantom (`ChatResultMeta`). + * + * `chat()` returns `TextActivityResult & + * ChatResultMeta` — an optional, runtime-absent + * carrier that lets downstream consumers (e.g. `useAssistant`) recover the + * tool set + structured-output schema that would otherwise be erased by the + * time the caller only has `ReturnType`. + * + * The phantom is optional (`'~chatMeta'?:`), so it must not change + * assignability of `chat()`'s return to its pre-existing consumers (see + * `chat-result-types.test.ts` and `activities/assistant/index.test.ts` for + * the pinned pre-phantom shapes). + */ +import { describe, expectTypeOf, it } from 'vitest' +import { z } from 'zod' +import { chat } from '../../../src/activities/chat/index.js' +import { toolDefinition } from '../../../src/activities/chat/tools/tool-definition.js' +import { createMockAdapter } from '../../test-utils.js' +import type { + ChatStream, + InferChatSchema, + InferChatTools, + StreamChunk, +} from '../../../src/types.js' + +describe('chat() result phantom (~chatMeta)', () => { + const { adapter } = createMockAdapter({ iterations: [[]] }) + + it('InferChatTools recovers the tool tuple passed to `tools`', () => { + const weatherTool = toolDefinition({ + name: 'get_weather', + description: 'Get the weather for a city', + inputSchema: z.object({ city: z.string() }), + }).server(() => ({ tempC: 20 })) + + const result = chat({ + adapter, + messages: [{ role: 'user', content: 'Hi' }], + tools: [weatherTool], + }) + + type R = typeof result + type Tools = InferChatTools + + // The captured tool set includes the tool passed to `tools`. + expectTypeOf().toEqualTypeOf() + expectTypeOf().toEqualTypeOf<'get_weather'>() + }) + + it('InferChatSchema recovers the `outputSchema` passed to chat()', () => { + const schema = z.object({ greeting: z.string() }) + + const result = chat({ + adapter, + messages: [{ role: 'user', content: 'Hi' }], + outputSchema: schema, + stream: true, + }) + + type R = typeof result + type Schema = InferChatSchema + + expectTypeOf().toEqualTypeOf() + }) + + it('InferChatSchema is `undefined` when no outputSchema is passed', () => { + const result = chat({ + adapter, + messages: [{ role: 'user', content: 'Hi' }], + }) + + type R = typeof result + type Schema = InferChatSchema + + expectTypeOf().toEqualTypeOf() + }) + + it('the phantom does not break assignability to ChatStream / AsyncIterable', () => { + const result = chat({ + adapter, + messages: [{ role: 'user', content: 'Hi' }], + }) + + type R = typeof result + + expectTypeOf().toMatchTypeOf() + expectTypeOf().toMatchTypeOf>() + + // Regression guard: existing consumers that assign `chat()`'s return + // directly to `AsyncIterable` (e.g. + // `toServerSentEventsResponse`) must still compile with the phantom + // present. + const _assignable: AsyncIterable = result + void _assignable + }) +}) diff --git a/packages/ai/vite.config.ts b/packages/ai/vite.config.ts index 5189bd6eb..066420b9d 100644 --- a/packages/ai/vite.config.ts +++ b/packages/ai/vite.config.ts @@ -31,6 +31,7 @@ export default mergeConfig( tanstackViteConfig({ entry: [ './src/index.ts', + './src/assistant.ts', './src/client.ts', './src/activities/index.ts', './src/middlewares/index.ts', diff --git a/testing/e2e/fixtures/assistant/basic.json b/testing/e2e/fixtures/assistant/basic.json new file mode 100644 index 000000000..81e0fe911 --- /dev/null +++ b/testing/e2e/fixtures/assistant/basic.json @@ -0,0 +1,10 @@ +{ + "fixtures": [ + { + "match": { "userMessage": "[assistant] recommend a guitar" }, + "response": { + "content": "I'd recommend the Fender Stratocaster — versatile and comfortable." + } + } + ] +} diff --git a/testing/e2e/src/lib/feature-support.ts b/testing/e2e/src/lib/feature-support.ts index cd3e2203d..38862e26a 100644 --- a/testing/e2e/src/lib/feature-support.ts +++ b/testing/e2e/src/lib/feature-support.ts @@ -254,6 +254,7 @@ export const matrix: Record> = { // Only Gemini currently surfaces a first-class stateful conversation API via // the adapter (geminiTextInteractions, behind @tanstack/ai-gemini/experimental). 'stateful-interactions': new Set(['gemini']), + assistant: new Set(['openai']), } export function isSupported(provider: Provider, feature: Feature): boolean { diff --git a/testing/e2e/src/lib/features.ts b/testing/e2e/src/lib/features.ts index f6a0b57d1..947e0f72c 100644 --- a/testing/e2e/src/lib/features.ts +++ b/testing/e2e/src/lib/features.ts @@ -144,4 +144,8 @@ export const featureConfigs: Record = { tools: [], modelOptions: {}, }, + assistant: { + tools: [], + modelOptions: {}, + }, } diff --git a/testing/e2e/src/lib/types.ts b/testing/e2e/src/lib/types.ts index 0e362b5b1..f148d3b74 100644 --- a/testing/e2e/src/lib/types.ts +++ b/testing/e2e/src/lib/types.ts @@ -44,6 +44,7 @@ export type Feature = | 'image-to-video' | 'interactions-video' | 'stateful-interactions' + | 'assistant' export const ALL_PROVIDERS: Provider[] = [ 'openai', @@ -90,4 +91,5 @@ export const ALL_FEATURES: Feature[] = [ 'image-to-video', 'interactions-video', 'stateful-interactions', + 'assistant', ] diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 14ffbf330..15d13759e 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -19,6 +19,7 @@ import { Route as DevtoolsRouteARouteImport } from './routes/devtools-route-a' import { Route as DevtoolsGenerationHooksRouteImport } from './routes/devtools-generation-hooks' import { Route as DevtoolsChatRouteImport } from './routes/devtools-chat' import { Route as ChatClientDefaultBridgeRouteImport } from './routes/chat-client-default-bridge' +import { Route as AssistantRouteImport } from './routes/assistant' import { Route as IndexRouteImport } from './routes/index' import { Route as ProviderIndexRouteImport } from './routes/$provider/index' import { Route as ApiVideoRouteImport } from './routes/api.video' @@ -47,6 +48,7 @@ import { Route as ApiLazyToolsWireRouteImport } from './routes/api.lazy-tools-wi import { Route as ApiImageRouteImport } from './routes/api.image' import { Route as ApiChatRouteImport } from './routes/api.chat' import { Route as ApiAudioRouteImport } from './routes/api.audio' +import { Route as ApiAssistantRouteImport } from './routes/api.assistant' import { Route as ApiArktypeToolWireRouteImport } from './routes/api.arktype-tool-wire' import { Route as ApiAnthropicStructuredUsageRouteImport } from './routes/api.anthropic-structured-usage' import { Route as ApiAnthropicSkillsWireRouteImport } from './routes/api.anthropic-skills-wire' @@ -108,6 +110,11 @@ const ChatClientDefaultBridgeRoute = ChatClientDefaultBridgeRouteImport.update({ path: '/chat-client-default-bridge', getParentRoute: () => rootRouteImport, } as any) +const AssistantRoute = AssistantRouteImport.update({ + id: '/assistant', + path: '/assistant', + getParentRoute: () => rootRouteImport, +} as any) const IndexRoute = IndexRouteImport.update({ id: '/', path: '/', @@ -252,6 +259,11 @@ const ApiAudioRoute = ApiAudioRouteImport.update({ path: '/api/audio', getParentRoute: () => rootRouteImport, } as any) +const ApiAssistantRoute = ApiAssistantRouteImport.update({ + id: '/api/assistant', + path: '/api/assistant', + getParentRoute: () => rootRouteImport, +} as any) const ApiArktypeToolWireRoute = ApiArktypeToolWireRouteImport.update({ id: '/api/arktype-tool-wire', path: '/api/arktype-tool-wire', @@ -306,6 +318,7 @@ const ApiAudioStreamRoute = ApiAudioStreamRouteImport.update({ export interface FileRoutesByFullPath { '/': typeof IndexRoute + '/assistant': typeof AssistantRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute '/devtools-chat': typeof DevtoolsChatRoute '/devtools-generation-hooks': typeof DevtoolsGenerationHooksRoute @@ -321,6 +334,7 @@ export interface FileRoutesByFullPath { '/api/anthropic-skills-wire': typeof ApiAnthropicSkillsWireRoute '/api/anthropic-structured-usage': typeof ApiAnthropicStructuredUsageRoute '/api/arktype-tool-wire': typeof ApiArktypeToolWireRoute + '/api/assistant': typeof ApiAssistantRoute '/api/audio': typeof ApiAudioRouteWithChildren '/api/chat': typeof ApiChatRoute '/api/image': typeof ApiImageRouteWithChildren @@ -356,6 +370,7 @@ export interface FileRoutesByFullPath { } export interface FileRoutesByTo { '/': typeof IndexRoute + '/assistant': typeof AssistantRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute '/devtools-chat': typeof DevtoolsChatRoute '/devtools-generation-hooks': typeof DevtoolsGenerationHooksRoute @@ -371,6 +386,7 @@ export interface FileRoutesByTo { '/api/anthropic-skills-wire': typeof ApiAnthropicSkillsWireRoute '/api/anthropic-structured-usage': typeof ApiAnthropicStructuredUsageRoute '/api/arktype-tool-wire': typeof ApiArktypeToolWireRoute + '/api/assistant': typeof ApiAssistantRoute '/api/audio': typeof ApiAudioRouteWithChildren '/api/chat': typeof ApiChatRoute '/api/image': typeof ApiImageRouteWithChildren @@ -407,6 +423,7 @@ export interface FileRoutesByTo { export interface FileRoutesById { __root__: typeof rootRouteImport '/': typeof IndexRoute + '/assistant': typeof AssistantRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute '/devtools-chat': typeof DevtoolsChatRoute '/devtools-generation-hooks': typeof DevtoolsGenerationHooksRoute @@ -422,6 +439,7 @@ export interface FileRoutesById { '/api/anthropic-skills-wire': typeof ApiAnthropicSkillsWireRoute '/api/anthropic-structured-usage': typeof ApiAnthropicStructuredUsageRoute '/api/arktype-tool-wire': typeof ApiArktypeToolWireRoute + '/api/assistant': typeof ApiAssistantRoute '/api/audio': typeof ApiAudioRouteWithChildren '/api/chat': typeof ApiChatRoute '/api/image': typeof ApiImageRouteWithChildren @@ -459,6 +477,7 @@ export interface FileRouteTypes { fileRoutesByFullPath: FileRoutesByFullPath fullPaths: | '/' + | '/assistant' | '/chat-client-default-bridge' | '/devtools-chat' | '/devtools-generation-hooks' @@ -474,6 +493,7 @@ export interface FileRouteTypes { | '/api/anthropic-skills-wire' | '/api/anthropic-structured-usage' | '/api/arktype-tool-wire' + | '/api/assistant' | '/api/audio' | '/api/chat' | '/api/image' @@ -509,6 +529,7 @@ export interface FileRouteTypes { fileRoutesByTo: FileRoutesByTo to: | '/' + | '/assistant' | '/chat-client-default-bridge' | '/devtools-chat' | '/devtools-generation-hooks' @@ -524,6 +545,7 @@ export interface FileRouteTypes { | '/api/anthropic-skills-wire' | '/api/anthropic-structured-usage' | '/api/arktype-tool-wire' + | '/api/assistant' | '/api/audio' | '/api/chat' | '/api/image' @@ -559,6 +581,7 @@ export interface FileRouteTypes { id: | '__root__' | '/' + | '/assistant' | '/chat-client-default-bridge' | '/devtools-chat' | '/devtools-generation-hooks' @@ -574,6 +597,7 @@ export interface FileRouteTypes { | '/api/anthropic-skills-wire' | '/api/anthropic-structured-usage' | '/api/arktype-tool-wire' + | '/api/assistant' | '/api/audio' | '/api/chat' | '/api/image' @@ -610,6 +634,7 @@ export interface FileRouteTypes { } export interface RootRouteChildren { IndexRoute: typeof IndexRoute + AssistantRoute: typeof AssistantRoute ChatClientDefaultBridgeRoute: typeof ChatClientDefaultBridgeRoute DevtoolsChatRoute: typeof DevtoolsChatRoute DevtoolsGenerationHooksRoute: typeof DevtoolsGenerationHooksRoute @@ -625,6 +650,7 @@ export interface RootRouteChildren { ApiAnthropicSkillsWireRoute: typeof ApiAnthropicSkillsWireRoute ApiAnthropicStructuredUsageRoute: typeof ApiAnthropicStructuredUsageRoute ApiArktypeToolWireRoute: typeof ApiArktypeToolWireRoute + ApiAssistantRoute: typeof ApiAssistantRoute ApiAudioRoute: typeof ApiAudioRouteWithChildren ApiChatRoute: typeof ApiChatRoute ApiImageRoute: typeof ApiImageRouteWithChildren @@ -726,6 +752,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ChatClientDefaultBridgeRouteImport parentRoute: typeof rootRouteImport } + '/assistant': { + id: '/assistant' + path: '/assistant' + fullPath: '/assistant' + preLoaderRoute: typeof AssistantRouteImport + parentRoute: typeof rootRouteImport + } '/': { id: '/' path: '/' @@ -922,6 +955,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiAudioRouteImport parentRoute: typeof rootRouteImport } + '/api/assistant': { + id: '/api/assistant' + path: '/api/assistant' + fullPath: '/api/assistant' + preLoaderRoute: typeof ApiAssistantRouteImport + parentRoute: typeof rootRouteImport + } '/api/arktype-tool-wire': { id: '/api/arktype-tool-wire' path: '/api/arktype-tool-wire' @@ -1055,6 +1095,7 @@ const ApiVideoRouteWithChildren = ApiVideoRoute._addFileChildren( const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, + AssistantRoute: AssistantRoute, ChatClientDefaultBridgeRoute: ChatClientDefaultBridgeRoute, DevtoolsChatRoute: DevtoolsChatRoute, DevtoolsGenerationHooksRoute: DevtoolsGenerationHooksRoute, @@ -1070,6 +1111,7 @@ const rootRouteChildren: RootRouteChildren = { ApiAnthropicSkillsWireRoute: ApiAnthropicSkillsWireRoute, ApiAnthropicStructuredUsageRoute: ApiAnthropicStructuredUsageRoute, ApiArktypeToolWireRoute: ApiArktypeToolWireRoute, + ApiAssistantRoute: ApiAssistantRoute, ApiAudioRoute: ApiAudioRouteWithChildren, ApiChatRoute: ApiChatRoute, ApiImageRoute: ApiImageRouteWithChildren, diff --git a/testing/e2e/src/routes/api.assistant.ts b/testing/e2e/src/routes/api.assistant.ts new file mode 100644 index 000000000..02bf054cd --- /dev/null +++ b/testing/e2e/src/routes/api.assistant.ts @@ -0,0 +1,52 @@ +import { createFileRoute } from '@tanstack/react-router' +import { chat, generateImage, maxIterations } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import type { Provider } from '@/lib/types' +import { createTextAdapter } from '@/lib/providers' +import { createImageAdapter } from '@/lib/media-providers' + +export const Route = createFileRoute('/api/assistant')({ + server: { + handlers: { + POST: async ({ request }) => { + await import('@/lib/llmock-server').then((m) => m.ensureLLMock()) + + // `assistant.handler(request)` below consumes the body itself via + // `request.json()`, so peek at a clone here to pull the + // provider/testId/aimockPort needed to construct adapters, and pass + // the original (unconsumed) request through to the handler. + const peekBody = await request.clone().json() + const peekData = peekBody.forwardedProps ?? peekBody.data ?? peekBody + const { provider, testId, aimockPort } = peekData as { + provider: Provider + testId?: string + aimockPort?: number + } + + const assistant = defineAssistant({ + chat: (req) => + chat({ + ...createTextAdapter( + provider, + undefined, + aimockPort, + testId, + 'assistant', + ), + messages: req.messages, + agentLoopStrategy: maxIterations(5), + threadId: req.threadId, + runId: req.runId, + }), + image: (req) => + generateImage({ + adapter: createImageAdapter(provider, aimockPort, testId), + prompt: req.prompt as string, + }), + }) + + return assistant.handler(request) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/assistant.tsx b/testing/e2e/src/routes/assistant.tsx new file mode 100644 index 000000000..963b0a5e9 --- /dev/null +++ b/testing/e2e/src/routes/assistant.tsx @@ -0,0 +1,128 @@ +import { createFileRoute } from '@tanstack/react-router' +import { chat, generateImage, maxIterations } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { openaiImage, openaiText } from '@tanstack/ai-openai' +import { ChatUI } from '@/components/ChatUI' +import type { Provider } from '@/lib/types' + +export interface AssistantRouteSearch { + provider: Provider + testId?: string + aimockPort?: number +} + +const DEFAULT_PROVIDER: Provider = 'openai' + +function parseAssistantRouteSearch( + search: Record, +): AssistantRouteSearch { + const aimockPort = + typeof search.aimockPort === 'string' + ? Number.parseInt(search.aimockPort, 10) + : undefined + const provider = + typeof search.provider === 'string' + ? (search.provider as Provider) + : DEFAULT_PROVIDER + + return { + provider, + ...(typeof search.testId === 'string' ? { testId: search.testId } : {}), + ...(aimockPort !== undefined && !Number.isNaN(aimockPort) + ? { aimockPort } + : {}), + } +} + +export const Route = createFileRoute('/assistant')({ + component: AssistantRoute, + validateSearch: parseAssistantRouteSearch, +}) + +// Client-side assistant definition. `defineAssistant` is inert — none of +// these callbacks ever run in the browser. `useAssistant` only reads the +// declared capability names (and their types) off this value to build the +// typed client system; the actual chat/image requests are always sent to +// the server route `/api/assistant`, which defines its own (real) callbacks. +// Mirroring the server route's two capabilities (chat + image) here just +// keeps the client types in sync with what the server actually supports. +// +// Uses the single-provider openai adapters directly (not the multi-provider +// `@/lib/providers` / `@/lib/media-providers` factories) so the browser +// bundle doesn't pull in every provider SDK (e.g. ollama's `node:fs`) — the +// callbacks are inert on the client, so only their return TYPES matter. +const assistant = defineAssistant({ + chat: (req) => + chat({ + adapter: openaiText('gpt-5.5'), + messages: req.messages, + agentLoopStrategy: maxIterations(5), + threadId: req.threadId, + runId: req.runId, + }), + image: (req) => + generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: req.prompt as string, + }), +}) + +function AssistantRoute() { + const { provider, testId, aimockPort } = Route.useSearch() + + const system = useAssistant(assistant, { + connection: fetchServerSentEvents('/api/assistant'), + chat: { forwardedProps: { provider, testId, aimockPort } }, + }) + + // The image one-shot's `GenerationClient` (unlike the chat sub-client) has + // no `forwardedProps` option on `useAssistant` — its request body is fixed + // to `{ capability: 'image' }` by `AssistantClient`. So the routing + // metadata (provider/testId/aimockPort) the server needs rides along on + // the `generate()` call's input instead: `GenerationClient.generate` merges + // its input directly into the wire `forwardedProps`, and the server route + // reads `provider`/`testId`/`aimockPort` back out of `forwardedProps`. + // Declared as a variable (not an inline literal) so these extra fields + // don't trip `ImageGenerateInput`'s excess-property check. + const imageGenerateInput = { + prompt: '[assistant] a fox', + provider, + testId, + aimockPort, + } + + return ( +
+
+ + + {system.image.result?.images[0]?.url ?? ''} + +
+
+ { + system.chat.sendMessage(text) + }} + onStop={system.chat.stop} + /> +
+
+ ) +} diff --git a/testing/e2e/tests/assistant.spec.ts b/testing/e2e/tests/assistant.spec.ts new file mode 100644 index 000000000..46797116c --- /dev/null +++ b/testing/e2e/tests/assistant.spec.ts @@ -0,0 +1,62 @@ +import { test, expect } from './fixtures' +import { + getLastAssistantMessage, + sendMessage, + waitForResponse, +} from './helpers' +import { providersFor } from './test-matrix' + +// The assistant feature has its own dedicated page (`/assistant`) that drives +// `useAssistant`, which posts to the `/api/assistant` handler — NOT the generic +// matrix page (`/$provider/$feature` → `/api/chat`). So we navigate directly +// rather than via `featureUrl`, exercising the real defineAssistant/useAssistant +// path end to end. +function assistantUrl(provider: string, testId: string, aimockPort: number) { + const params = new URLSearchParams({ + provider, + testId, + aimockPort: String(aimockPort), + }) + return `/assistant?${params.toString()}` +} + +for (const provider of providersFor('assistant')) { + test.describe(`${provider} — assistant`, () => { + test('chat capability streams through useAssistant + the assistant handler', async ({ + page, + testId, + aimockPort, + }) => { + await page.goto(assistantUrl(provider, testId, aimockPort)) + + await sendMessage(page, '[assistant] recommend a guitar') + await waitForResponse(page) + + const response = await getLastAssistantMessage(page) + expect(response).toContain('Fender Stratocaster') + }) + + // TODO(assistant image e2e): the click + one-shot dispatch through + // `/api/assistant` (capability=image) fires correctly, but aimock returns + // no image body — image generation is mocked programmatically via + // `registerMediaFixtures` (`match.endpoint`), not the userMessage-keyed + // fixtures used for chat, so the assistant image request has no registered + // response. The one-shot path is already covered by unit tests + // (`@tanstack/ai`: handler emits a `generation:result` CUSTOM event; + // `@tanstack/ai-react`: `system.image.generate` populates `result`). + // Un-skip once an assistant image media fixture is registered in global-setup. + test.skip('image one-shot leg runs through the assistant handler', async ({ + page, + testId, + aimockPort, + }) => { + await page.goto(assistantUrl(provider, testId, aimockPort)) + + await page.getByTestId('assistant-generate-image').click() + + await expect(page.getByTestId('assistant-image-result')).not.toBeEmpty({ + timeout: 15_000, + }) + }) + }) +} From 1d127c74bfb1251b5c684ab45d6af4c025e6acb2 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Tue, 14 Jul 2026 13:08:22 +0000 Subject: [PATCH 02/13] ci: apply automated fixes --- packages/ai-react/src/use-assistant.ts | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/packages/ai-react/src/use-assistant.ts b/packages/ai-react/src/use-assistant.ts index f803cd3a5..b9e53a4cd 100644 --- a/packages/ai-react/src/use-assistant.ts +++ b/packages/ai-react/src/use-assistant.ts @@ -194,9 +194,8 @@ export function useAssistant< await c.sendMessage(content) const msgs = c.getMessages() const hasStructured = - msgs - .at(-1) - ?.parts.some((p) => p.type === 'structured-output') ?? false + msgs.at(-1)?.parts.some((p) => p.type === 'structured-output') ?? + false // Cast: TS can't structurally narrow the conditional return of // AssistantChatSurface['sendMessage'] against this runtime branch, // mirroring the assembled-`system` cast below. From 2a3e333b9d3c84f44acb1e54708d149f31bd5f78 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 14 Jul 2026 17:12:53 +0200 Subject: [PATCH 03/13] docs(examples): add Blog Studio assistant example to ts-react-chat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A topic → finished blog post demo of the multi-capability assistant. One `defineAssistant` route hosts structured-output chat, image, and speech; the page drafts the post, then illustrates and narrates it in parallel via the resolved `await` returns, and renders the result as an editorial article (hero image, Markdown body, voice-over player) in a two-column studio layout. Failed image/speech steps degrade gracefully with an honest per-step status; narration text is Markdown-stripped and length-capped for the TTS limit. --- .../ts-react-chat/src/components/Header.tsx | 14 + examples/ts-react-chat/src/routeTree.gen.ts | 42 ++ .../src/routes/api.blog-studio.ts | 82 ++++ .../ts-react-chat/src/routes/blog-studio.tsx | 416 ++++++++++++++++++ examples/ts-react-chat/src/routes/index.tsx | 8 + 5 files changed, 562 insertions(+) create mode 100644 examples/ts-react-chat/src/routes/api.blog-studio.ts create mode 100644 examples/ts-react-chat/src/routes/blog-studio.tsx diff --git a/examples/ts-react-chat/src/components/Header.tsx b/examples/ts-react-chat/src/components/Header.tsx index 061f74d78..36023287f 100644 --- a/examples/ts-react-chat/src/components/Header.tsx +++ b/examples/ts-react-chat/src/components/Header.tsx @@ -16,6 +16,7 @@ import { Mic, MessageSquare, Music, + Newspaper, Plug, Server, Sparkles, @@ -206,6 +207,19 @@ export default function Header() { Examples

+ setIsOpen(false)} + className="flex items-center gap-3 p-3 rounded-lg hover:bg-gray-800 transition-colors mb-1" + activeProps={{ + className: + 'flex items-center gap-3 p-3 rounded-lg bg-cyan-600 hover:bg-cyan-700 transition-colors mb-1', + }} + > + + Blog Studio + + setIsOpen(false)} diff --git a/examples/ts-react-chat/src/routeTree.gen.ts b/examples/ts-react-chat/src/routeTree.gen.ts index b223bce09..06b43bd5b 100644 --- a/examples/ts-react-chat/src/routeTree.gen.ts +++ b/examples/ts-react-chat/src/routeTree.gen.ts @@ -21,6 +21,7 @@ import { Route as ImageToolReproRouteImport } from './routes/image-tool-repro' import { Route as ImageGenRouteImport } from './routes/image-gen' import { Route as GenerationHooksRouteImport } from './routes/generation-hooks' import { Route as CapabilityDemoRouteImport } from './routes/capability-demo' +import { Route as BlogStudioRouteImport } from './routes/blog-studio' import { Route as IndexRouteImport } from './routes/index' import { Route as GenerationsVideoRouteImport } from './routes/generations.video' import { Route as GenerationsTranscriptionRouteImport } from './routes/generations.transcription' @@ -48,6 +49,7 @@ import { Route as ApiMcpAppsCallRouteImport } from './routes/api.mcp-apps-call' import { Route as ApiImageToolReproRouteImport } from './routes/api.image-tool-repro' import { Route as ApiImageGenRouteImport } from './routes/api.image-gen' import { Route as ApiCapabilityDemoRouteImport } from './routes/api.capability-demo' +import { Route as ApiBlogStudioRouteImport } from './routes/api.blog-studio' import { Route as ExampleGuitarsIndexRouteImport } from './routes/example.guitars/index' import { Route as ExampleGuitarsGuitarIdRouteImport } from './routes/example.guitars/$guitarId' import { Route as ApiGenerateVideoRouteImport } from './routes/api.generate.video' @@ -115,6 +117,11 @@ const CapabilityDemoRoute = CapabilityDemoRouteImport.update({ path: '/capability-demo', getParentRoute: () => rootRouteImport, } as any) +const BlogStudioRoute = BlogStudioRouteImport.update({ + id: '/blog-studio', + path: '/blog-studio', + getParentRoute: () => rootRouteImport, +} as any) const IndexRoute = IndexRouteImport.update({ id: '/', path: '/', @@ -253,6 +260,11 @@ const ApiCapabilityDemoRoute = ApiCapabilityDemoRouteImport.update({ path: '/api/capability-demo', getParentRoute: () => rootRouteImport, } as any) +const ApiBlogStudioRoute = ApiBlogStudioRouteImport.update({ + id: '/api/blog-studio', + path: '/api/blog-studio', + getParentRoute: () => rootRouteImport, +} as any) const ExampleGuitarsIndexRoute = ExampleGuitarsIndexRouteImport.update({ id: '/example/guitars/', path: '/example/guitars/', @@ -286,6 +298,7 @@ const ApiGenerateAudioRoute = ApiGenerateAudioRouteImport.update({ export interface FileRoutesByFullPath { '/': typeof IndexRoute + '/blog-studio': typeof BlogStudioRoute '/capability-demo': typeof CapabilityDemoRoute '/generation-hooks': typeof GenerationHooksRoute '/image-gen': typeof ImageGenRoute @@ -298,6 +311,7 @@ export interface FileRoutesByFullPath { '/server-fn-chat': typeof ServerFnChatRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute + '/api/blog-studio': typeof ApiBlogStudioRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute '/api/image-gen': typeof ApiImageGenRoute '/api/image-tool-repro': typeof ApiImageToolReproRoute @@ -333,6 +347,7 @@ export interface FileRoutesByFullPath { } export interface FileRoutesByTo { '/': typeof IndexRoute + '/blog-studio': typeof BlogStudioRoute '/capability-demo': typeof CapabilityDemoRoute '/generation-hooks': typeof GenerationHooksRoute '/image-gen': typeof ImageGenRoute @@ -345,6 +360,7 @@ export interface FileRoutesByTo { '/server-fn-chat': typeof ServerFnChatRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute + '/api/blog-studio': typeof ApiBlogStudioRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute '/api/image-gen': typeof ApiImageGenRoute '/api/image-tool-repro': typeof ApiImageToolReproRoute @@ -381,6 +397,7 @@ export interface FileRoutesByTo { export interface FileRoutesById { __root__: typeof rootRouteImport '/': typeof IndexRoute + '/blog-studio': typeof BlogStudioRoute '/capability-demo': typeof CapabilityDemoRoute '/generation-hooks': typeof GenerationHooksRoute '/image-gen': typeof ImageGenRoute @@ -393,6 +410,7 @@ export interface FileRoutesById { '/server-fn-chat': typeof ServerFnChatRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute + '/api/blog-studio': typeof ApiBlogStudioRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute '/api/image-gen': typeof ApiImageGenRoute '/api/image-tool-repro': typeof ApiImageToolReproRoute @@ -430,6 +448,7 @@ export interface FileRouteTypes { fileRoutesByFullPath: FileRoutesByFullPath fullPaths: | '/' + | '/blog-studio' | '/capability-demo' | '/generation-hooks' | '/image-gen' @@ -442,6 +461,7 @@ export interface FileRouteTypes { | '/server-fn-chat' | '/threads' | '/typesafe-tools' + | '/api/blog-studio' | '/api/capability-demo' | '/api/image-gen' | '/api/image-tool-repro' @@ -477,6 +497,7 @@ export interface FileRouteTypes { fileRoutesByTo: FileRoutesByTo to: | '/' + | '/blog-studio' | '/capability-demo' | '/generation-hooks' | '/image-gen' @@ -489,6 +510,7 @@ export interface FileRouteTypes { | '/server-fn-chat' | '/threads' | '/typesafe-tools' + | '/api/blog-studio' | '/api/capability-demo' | '/api/image-gen' | '/api/image-tool-repro' @@ -524,6 +546,7 @@ export interface FileRouteTypes { id: | '__root__' | '/' + | '/blog-studio' | '/capability-demo' | '/generation-hooks' | '/image-gen' @@ -536,6 +559,7 @@ export interface FileRouteTypes { | '/server-fn-chat' | '/threads' | '/typesafe-tools' + | '/api/blog-studio' | '/api/capability-demo' | '/api/image-gen' | '/api/image-tool-repro' @@ -572,6 +596,7 @@ export interface FileRouteTypes { } export interface RootRouteChildren { IndexRoute: typeof IndexRoute + BlogStudioRoute: typeof BlogStudioRoute CapabilityDemoRoute: typeof CapabilityDemoRoute GenerationHooksRoute: typeof GenerationHooksRoute ImageGenRoute: typeof ImageGenRoute @@ -584,6 +609,7 @@ export interface RootRouteChildren { ServerFnChatRoute: typeof ServerFnChatRoute ThreadsRoute: typeof ThreadsRoute TypesafeToolsRoute: typeof TypesafeToolsRoute + ApiBlogStudioRoute: typeof ApiBlogStudioRoute ApiCapabilityDemoRoute: typeof ApiCapabilityDemoRoute ApiImageGenRoute: typeof ApiImageGenRoute ApiImageToolReproRoute: typeof ApiImageToolReproRoute @@ -704,6 +730,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof CapabilityDemoRouteImport parentRoute: typeof rootRouteImport } + '/blog-studio': { + id: '/blog-studio' + path: '/blog-studio' + fullPath: '/blog-studio' + preLoaderRoute: typeof BlogStudioRouteImport + parentRoute: typeof rootRouteImport + } '/': { id: '/' path: '/' @@ -893,6 +926,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiCapabilityDemoRouteImport parentRoute: typeof rootRouteImport } + '/api/blog-studio': { + id: '/api/blog-studio' + path: '/api/blog-studio' + fullPath: '/api/blog-studio' + preLoaderRoute: typeof ApiBlogStudioRouteImport + parentRoute: typeof rootRouteImport + } '/example/guitars/': { id: '/example/guitars/' path: '/example/guitars' @@ -940,6 +980,7 @@ declare module '@tanstack/react-router' { const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, + BlogStudioRoute: BlogStudioRoute, CapabilityDemoRoute: CapabilityDemoRoute, GenerationHooksRoute: GenerationHooksRoute, ImageGenRoute: ImageGenRoute, @@ -952,6 +993,7 @@ const rootRouteChildren: RootRouteChildren = { ServerFnChatRoute: ServerFnChatRoute, ThreadsRoute: ThreadsRoute, TypesafeToolsRoute: TypesafeToolsRoute, + ApiBlogStudioRoute: ApiBlogStudioRoute, ApiCapabilityDemoRoute: ApiCapabilityDemoRoute, ApiImageGenRoute: ApiImageGenRoute, ApiImageToolReproRoute: ApiImageToolReproRoute, diff --git a/examples/ts-react-chat/src/routes/api.blog-studio.ts b/examples/ts-react-chat/src/routes/api.blog-studio.ts new file mode 100644 index 000000000..9db94c5e4 --- /dev/null +++ b/examples/ts-react-chat/src/routes/api.blog-studio.ts @@ -0,0 +1,82 @@ +import { createFileRoute } from '@tanstack/react-router' +import { chat, generateImage, generateSpeech } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiImage, openaiSpeech, openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +/** + * The shape a blog post is drafted into. The `chat` capability declares this + * as its `outputSchema`, so `assistant.chat.sendMessage(...)` resolves to a + * validated `{ title, subtitle, body } | null` on the client — no manual + * parsing, and `assistant.chat.partial` streams a live draft. + * + * Exported so the client page (`blog-studio.tsx`) can import the same schema + * and stay in type-sync. Only this schema lives at module scope; the actual + * `defineAssistant` (with its server adapters) is built inside the POST + * handler so importing the schema never pulls provider SDKs into the browser + * bundle. + */ +export const BlogPostSchema = z.object({ + title: z.string().describe('A punchy, editorial blog post title'), + subtitle: z.string().describe('A one-sentence standfirst / subtitle'), + body: z + .string() + .describe( + 'The full blog post body as GitHub-flavored Markdown: use ## / ### ' + + 'headings, short paragraphs, and the occasional list. ~400-600 words.', + ), +}) + +export type BlogPost = z.infer + +const SYSTEM_PROMPT = + 'You are a seasoned staff writer. Given a topic, write one engaging, ' + + 'well-structured blog post. Return a title, a short subtitle, and the ' + + 'body as GitHub-flavored Markdown with section headings and tight ' + + 'paragraphs. Be vivid and concrete; avoid filler and clichés.' + +export const Route = createFileRoute('/api/blog-studio')({ + server: { + handlers: { + // One endpoint, three capabilities. `defineAssistant` is inert until + // `handler(request)` runs, at which point it parses the AG-UI request, + // routes by the `capability` discriminator the client sends, and streams + // the result back over SSE. + POST: async ({ request }) => { + const assistant = defineAssistant({ + // Write the post as a typed object (structured output + streaming). + chat: (req) => + chat({ + adapter: openaiText('gpt-5.5'), + messages: req.messages, + systemPrompts: [SYSTEM_PROMPT], + outputSchema: BlogPostSchema, + stream: true, + threadId: req.threadId, + runId: req.runId, + }), + // Generate a landscape hero / OG image from the drafted title. + image: (req) => { + if (typeof req.prompt !== 'string') { + throw new Error('image prompt must be a string') + } + return generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: req.prompt, + size: '1536x1024', + }) + }, + // Narrate the post body. + speech: (req) => + generateSpeech({ + adapter: openaiSpeech('tts-1'), + text: req.text, + voice: 'alloy', + }), + }) + + return assistant.handler(request) + }, + }, + }, +}) diff --git a/examples/ts-react-chat/src/routes/blog-studio.tsx b/examples/ts-react-chat/src/routes/blog-studio.tsx new file mode 100644 index 000000000..314aced5b --- /dev/null +++ b/examples/ts-react-chat/src/routes/blog-studio.tsx @@ -0,0 +1,416 @@ +import { useEffect, useRef, useState } from 'react' +import type { ReactNode } from 'react' +import { createFileRoute } from '@tanstack/react-router' +import { chat, generateImage, generateSpeech } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { openaiImage, openaiSpeech, openaiText } from '@tanstack/ai-openai' +import ReactMarkdown from 'react-markdown' +import remarkGfm from 'remark-gfm' +import { + AlertTriangle, + Check, + Image as ImageIcon, + Loader2, + Newspaper, + PenLine, + Sparkles, + Volume2, + Wand2, +} from 'lucide-react' +import { BlogPostSchema } from './api.blog-studio' + +// Client-side assistant definition. `defineAssistant` is INERT in the browser +// — these callbacks never run; `useAssistant` only reads the declared +// capability names and their return types off this value to build the fully +// typed client. Every request is sent to the `/api/blog-studio` route, which +// owns the real callbacks. Mirroring the three capabilities (and the chat +// `outputSchema`) here keeps the client types in sync with the server. +// +// Single-provider openai adapters are used directly (not the multi-provider +// factories) so the browser bundle doesn't pull in every provider SDK. +const blogAssistant = defineAssistant({ + chat: (req) => + chat({ + adapter: openaiText('gpt-5.5'), + messages: req.messages, + outputSchema: BlogPostSchema, + stream: true, + }), + image: (req) => + generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: typeof req.prompt === 'string' ? req.prompt : '', + }), + speech: (req) => + generateSpeech({ adapter: openaiSpeech('tts-1'), text: req.text }), +}) + +export const Route = createFileRoute('/blog-studio')({ + component: BlogStudio, +}) + +// The chat step runs first; once the draft is ready, the image and voice-over +// are produced in parallel. Each one-shot step tracks its own outcome because +// `generate()` resolves to `null` on failure rather than rejecting. +type Phase = 'idle' | 'writing' | 'producing' | 'done' +type StepOutcome = 'pending' | 'active' | 'ok' | 'failed' +type StepState = 'pending' | 'active' | 'done' | 'failed' + +function outcomeToState(outcome: StepOutcome): StepState { + return outcome === 'ok' + ? 'done' + : outcome === 'active' + ? 'active' + : outcome === 'failed' + ? 'failed' + : 'pending' +} + +function base64ToObjectUrl(base64: string, contentType: string): string { + const binary = atob(base64) + const bytes = new Uint8Array(binary.length) + for (let i = 0; i < binary.length; i++) bytes[i] = binary.charCodeAt(i) + return URL.createObjectURL(new Blob([bytes], { type: contentType })) +} + +// Prepare the post body for narration: strip Markdown so TTS doesn't read the +// syntax aloud, and cap the length at a sentence boundary (OpenAI TTS rejects +// input over 4096 characters, and long posts easily exceed it). +function forNarration(markdown: string, max = 4000): string { + const plain = markdown + .replace(/^#{1,6}\s+/gm, '') // headings + .replace(/^\s*[-*+]\s+/gm, '') // list bullets + .replace(/^\s*>\s?/gm, '') // blockquotes + .replace(/\*\*(.*?)\*\*/g, '$1') // bold + .replace(/\*(.*?)\*/g, '$1') // italic + .replace(/`([^`]+)`/g, '$1') // inline code + .replace(/\[([^\]]+)\]\([^)]*\)/g, '$1') // links → link text + .replace(/\n{3,}/g, '\n\n') + .trim() + if (plain.length <= max) return plain + const clipped = plain.slice(0, max) + const boundary = Math.max( + clipped.lastIndexOf('. '), + clipped.lastIndexOf('\n'), + ) + return (boundary > max / 2 ? clipped.slice(0, boundary + 1) : clipped).trim() +} + +function BlogStudio() { + const assistant = useAssistant(blogAssistant, { + connection: fetchServerSentEvents('/api/blog-studio'), + }) + + const [topic, setTopic] = useState('') + const [phase, setPhase] = useState('idle') + const [imageUrl, setImageUrl] = useState(null) + const [audioUrl, setAudioUrl] = useState(null) + const [imageOutcome, setImageOutcome] = useState('pending') + const [audioOutcome, setAudioOutcome] = useState('pending') + const [error, setError] = useState(null) + const audioObjectUrl = useRef(null) + + // Revoke the last audio object URL when the component unmounts. + useEffect( + () => () => { + if (audioObjectUrl.current) URL.revokeObjectURL(audioObjectUrl.current) + }, + [], + ) + + const isRunning = phase === 'writing' || phase === 'producing' + + async function run() { + const t = topic.trim() + if (!t || isRunning) return + setError(null) + setImageUrl(null) + setAudioUrl(null) + // Revoke the previous run's audio URL now, so it's freed regardless of + // whether this run's narration succeeds. + if (audioObjectUrl.current) { + URL.revokeObjectURL(audioObjectUrl.current) + audioObjectUrl.current = null + } + setImageOutcome('pending') + setAudioOutcome('pending') + setPhase('writing') + + try { + // 1. Draft the post. `sendMessage` resolves to the structured result + // because the chat callback declared an `outputSchema`. We re-validate + // it with the same schema before chaining: if the run errored (e.g. a + // missing server API key) there's no completed structured output, so + // guarding on the shape — not just a null check — keeps the pipeline + // from proceeding with a half-finished draft. + const drafted = await assistant.chat.sendMessage( + `Write a blog post about: ${t}`, + ) + const parsed = BlogPostSchema.safeParse(drafted) + if (!parsed.success) { + setError( + 'Could not draft the post. Make sure OPENAI_API_KEY is set on the ' + + 'server, then try again.', + ) + setPhase('idle') + return + } + const post = parsed.data + + // 2. Illustrate and narrate in parallel — both derive only from the + // finished draft, so there's no reason to wait for one before the + // other. Each resolves to `null` on failure (generate never rejects), + // so each records its own outcome and the article ships best-effort. + setPhase('producing') + setImageOutcome('active') + setAudioOutcome('active') + + const illustrate = (async () => { + const image = await assistant.image.generate({ + prompt: + `A striking editorial hero image for a blog post titled ` + + `"${post.title}". ${post.subtitle}. Modern, clean, cinematic, ` + + `high quality, no text.`, + }) + const shot = image?.images[0] + const url = + shot?.url ?? + (shot?.b64Json ? `data:image/png;base64,${shot.b64Json}` : null) + setImageUrl(url) + setImageOutcome(url ? 'ok' : 'failed') + })() + + const narrate = (async () => { + const speech = await assistant.speech.generate({ + text: forNarration(post.body), + }) + if (speech?.audio) { + const audio = base64ToObjectUrl( + speech.audio, + speech.contentType ?? 'audio/mpeg', + ) + audioObjectUrl.current = audio + setAudioUrl(audio) + setAudioOutcome('ok') + } else { + setAudioOutcome('failed') + } + })() + + await Promise.all([illustrate, narrate]) + setPhase('done') + } catch (err) { + setError(err instanceof Error ? err.message : 'Something went wrong.') + setPhase('idle') + } + } + + const draft = assistant.chat.partial + const post = assistant.chat.final + const title = post?.title ?? draft.title + const subtitle = post?.subtitle ?? draft.subtitle + const body = post?.body ?? draft.body ?? '' + const showArticle = Boolean(post) || (phase === 'writing' && (title || body)) + + return ( +
+ {/* Left: the studio controls */} + + + {/* Right: the post */} +
+ {showArticle ? ( +
+ {/* Hero image */} +
+ {imageUrl ? ( + {title + ) : ( +
+ {imageOutcome === 'active' ? ( +
+ + Illustrating… +
+ ) : imageOutcome === 'failed' ? ( +
+ + + Couldn't generate a hero image + +
+ ) : ( + + )} +
+ )} +
+ +
+

+ {title ?? 'Untitled'} +

+ {subtitle && ( +

{subtitle}

+ )} + + {/* Byline + voice-over */} +
+
+ +
+ Written & narrated by TanStack AI + {audioUrl ? ( +
+ +
+ ) : audioOutcome === 'active' ? ( + + Recording + voice-over… + + ) : audioOutcome === 'failed' ? ( + + + Voice-over unavailable + + ) : null} +
+ + {/* Body */} +
+ + {body} + +
+
+
+ ) : ( +
+
+ +

+ Your finished post — hero image, article, and voice-over — will + appear here. +

+
+
+ )} +
+
+ ) +} + +function StepRow({ + label, + icon, + state, +}: { + label: string + icon: ReactNode + state: StepState +}) { + const cls = + state === 'active' + ? 'text-amber-700 font-medium' + : state === 'done' + ? 'text-stone-500' + : state === 'failed' + ? 'text-amber-600' + : 'text-stone-300' + return ( + + {state === 'active' ? ( + + ) : state === 'done' ? ( + + ) : state === 'failed' ? ( + + ) : ( + icon + )} + {label} + + ) +} diff --git a/examples/ts-react-chat/src/routes/index.tsx b/examples/ts-react-chat/src/routes/index.tsx index b3228dd67..99c3a10d9 100644 --- a/examples/ts-react-chat/src/routes/index.tsx +++ b/examples/ts-react-chat/src/routes/index.tsx @@ -10,6 +10,7 @@ import { ImagePlus, Mic, Music, + Newspaper, Send, Square, Video, @@ -144,6 +145,13 @@ function Messages({

+ + + Blog Studio + Date: Tue, 14 Jul 2026 18:30:43 +0200 Subject: [PATCH 04/13] feat(assistant): per-capability onResult transform on useAssistant MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `useAssistant` now accepts a per-one-shot-capability `{ onResult, forwardedProps }` option (keyed by capability name), mirroring the standalone generation hooks. `onResult` runs on the raw backend result and its return type becomes `assistant..result` (via `InferGenerationOutput`); the parameter is typed as the raw result. Full inference both directions — locked by an `expectTypeOf` test. Threaded through `AssistantClient` into each `GenerationClient`, across all four framework hooks. Also refactors the Blog Studio example to derive 100% of its UI state from the assistant surface — zero `useState`/`useEffect`/`useRef`. Chat `isLoading`/`partial`/`final` and each one-shot's `result`/`status`/`error` drive the stepper, hero, and article; the voice-over uses `speech.onResult` to expose a ready `{ src }` data URL (no object-URL, no cleanup). Docs (assistant overview) and the ai-core/assistant skill updated for the new option. --- docs/assistant/overview.md | 57 ++++ .../ts-react-chat/src/routes/blog-studio.tsx | 277 +++++++----------- packages/ai-client/src/assistant-client.ts | 13 +- packages/ai-client/src/assistant-types.ts | 46 ++- packages/ai-react/src/use-assistant.ts | 24 +- .../ai-react/tests/use-assistant.test.tsx | 39 ++- packages/ai-solid/src/use-assistant.ts | 22 +- .../ai-svelte/src/create-assistant.svelte.ts | 23 +- packages/ai-vue/src/use-assistant.ts | 18 +- packages/ai/skills/ai-core/assistant/SKILL.md | 38 ++- 10 files changed, 338 insertions(+), 219 deletions(-) diff --git a/docs/assistant/overview.md b/docs/assistant/overview.md index dd532c668..4663c843a 100644 --- a/docs/assistant/overview.md +++ b/docs/assistant/overview.md @@ -288,6 +288,63 @@ function BlogOutlineForm() { `partial` and `final` only exist on the type when the `chat` callback declares `outputSchema` — omit it and `assistant.chat` has no such fields, exactly like `useChat` without `outputSchema`. See [Structured Outputs](../structured-outputs/overview) for the full story. +## Transforming one-shot results + +Each one-shot capability accepts an optional `onResult` transform, keyed by the +capability name on `useAssistant`'s options — the same shape as the standalone +generation hooks. It runs on the raw backend result before it lands on +`assistant..result`, and **its return type becomes that +capability's `result` type**. Use it to reshape a result once, at the hook, +instead of deriving it in every component that reads it: + +```tsx +// components/CoverArt.tsx +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { chat, generateImage } from '@tanstack/ai' +import { defineAssistant } from '@tanstack/ai/assistant' +import { openaiImage, openaiText } from '@tanstack/ai-openai' + +const blogAssistant = defineAssistant({ + chat: (req) => + chat({ adapter: openaiText('gpt-5.5'), messages: req.messages }), + image: (req) => { + if (typeof req.prompt !== 'string') { + throw new Error('image prompt must be a string') + } + return generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: req.prompt, + }) + }, +}) + +function CoverArt() { + const assistant = useAssistant(blogAssistant, { + connection: fetchServerSentEvents('/api/assistant'), + // `result` is the raw `ImageGenerationResult`; the return type flows + // through to `assistant.image.result`. + image: { onResult: (result) => result.images[0]?.url ?? null }, + }) + + return ( +
+ + {/* `assistant.image.result` is now `string | null`, not the raw result */} + {assistant.image.result && } +
+ ) +} +``` + +Return `undefined` (or nothing) from `onResult` to keep the raw result. Each +capability's option also takes `forwardedProps` to merge extra fields into that +capability's request body. + ## Which page do I read? | You want to… | Read | diff --git a/examples/ts-react-chat/src/routes/blog-studio.tsx b/examples/ts-react-chat/src/routes/blog-studio.tsx index 314aced5b..2519c7047 100644 --- a/examples/ts-react-chat/src/routes/blog-studio.tsx +++ b/examples/ts-react-chat/src/routes/blog-studio.tsx @@ -1,5 +1,4 @@ -import { useEffect, useRef, useState } from 'react' -import type { ReactNode } from 'react' +import type { FormEvent, ReactNode } from 'react' import { createFileRoute } from '@tanstack/react-router' import { chat, generateImage, generateSpeech } from '@tanstack/ai' import { defineAssistant } from '@tanstack/ai/assistant' @@ -51,30 +50,21 @@ export const Route = createFileRoute('/blog-studio')({ component: BlogStudio, }) -// The chat step runs first; once the draft is ready, the image and voice-over -// are produced in parallel. Each one-shot step tracks its own outcome because -// `generate()` resolves to `null` on failure rather than rejecting. -type Phase = 'idle' | 'writing' | 'producing' | 'done' -type StepOutcome = 'pending' | 'active' | 'ok' | 'failed' type StepState = 'pending' | 'active' | 'done' | 'failed' -function outcomeToState(outcome: StepOutcome): StepState { - return outcome === 'ok' - ? 'done' - : outcome === 'active' - ? 'active' - : outcome === 'failed' +// The one-shot generation status maps directly onto a step state. +function statusToStep( + status: 'idle' | 'generating' | 'success' | 'error', +): StepState { + return status === 'generating' + ? 'active' + : status === 'success' + ? 'done' + : status === 'error' ? 'failed' : 'pending' } -function base64ToObjectUrl(base64: string, contentType: string): string { - const binary = atob(base64) - const bytes = new Uint8Array(binary.length) - for (let i = 0; i < binary.length; i++) bytes[i] = binary.charCodeAt(i) - return URL.createObjectURL(new Blob([bytes], { type: contentType })) -} - // Prepare the post body for narration: strip Markdown so TTS doesn't read the // syntax aloud, and cap the length at a sentence boundary (OpenAI TTS rejects // input over 4096 characters, and long posts easily exceed it). @@ -101,119 +91,76 @@ function forNarration(markdown: string, max = 4000): string { function BlogStudio() { const assistant = useAssistant(blogAssistant, { connection: fetchServerSentEvents('/api/blog-studio'), - }) - - const [topic, setTopic] = useState('') - const [phase, setPhase] = useState('idle') - const [imageUrl, setImageUrl] = useState(null) - const [audioUrl, setAudioUrl] = useState(null) - const [imageOutcome, setImageOutcome] = useState('pending') - const [audioOutcome, setAudioOutcome] = useState('pending') - const [error, setError] = useState(null) - const audioObjectUrl = useRef(null) - - // Revoke the last audio object URL when the component unmounts. - useEffect( - () => () => { - if (audioObjectUrl.current) URL.revokeObjectURL(audioObjectUrl.current) + // Transform the raw TTS result into a ready-to-play data URL at the hook. + // `result` is typed as the raw `TTSResult`; `assistant.speech.result` + // becomes `{ src: string } | null` — no component state, no cleanup. + speech: { + onResult: (result) => ({ + src: `data:${result.contentType ?? 'audio/mpeg'};base64,${result.audio}`, + }), }, - [], - ) + }) - const isRunning = phase === 'writing' || phase === 'producing' + // Everything below is derived from the assistant's reactive state — no + // useState / useEffect. The chat carries loading/partial/final; each one-shot + // carries result/status/error. + const { chat: post, image, speech } = assistant + const draft = post.partial + const finished = post.final + const title = finished?.title ?? draft.title + const subtitle = finished?.subtitle ?? draft.subtitle + const body = finished?.body ?? draft.body ?? '' - async function run() { - const t = topic.trim() - if (!t || isRunning) return - setError(null) - setImageUrl(null) - setAudioUrl(null) - // Revoke the previous run's audio URL now, so it's freed regardless of - // whether this run's narration succeeds. - if (audioObjectUrl.current) { - URL.revokeObjectURL(audioObjectUrl.current) - audioObjectUrl.current = null - } - setImageOutcome('pending') - setAudioOutcome('pending') - setPhase('writing') + const heroImage = image.result?.images[0] + const imageUrl = + heroImage?.url ?? + (heroImage?.b64Json ? `data:image/png;base64,${heroImage.b64Json}` : null) + // `speech.result` is the `onResult` transform's output — `{ src } | null`. + const audioSrc = speech.result?.src ?? null - try { - // 1. Draft the post. `sendMessage` resolves to the structured result - // because the chat callback declared an `outputSchema`. We re-validate - // it with the same schema before chaining: if the run errored (e.g. a - // missing server API key) there's no completed structured output, so - // guarding on the shape — not just a null check — keeps the pipeline - // from proceeding with a half-finished draft. - const drafted = await assistant.chat.sendMessage( - `Write a blog post about: ${t}`, - ) - const parsed = BlogPostSchema.safeParse(drafted) - if (!parsed.success) { - setError( - 'Could not draft the post. Make sure OPENAI_API_KEY is set on the ' + - 'server, then try again.', - ) - setPhase('idle') - return - } - const post = parsed.data + const isRunning = post.isLoading || image.isLoading || speech.isLoading + const hasRun = post.messages.length > 0 || isRunning + const writingStep: StepState = post.error + ? 'failed' + : post.isLoading + ? 'active' + : hasRun + ? 'done' + : 'pending' + const showArticle = Boolean(finished) || Boolean(title) || Boolean(body) - // 2. Illustrate and narrate in parallel — both derive only from the - // finished draft, so there's no reason to wait for one before the - // other. Each resolves to `null` on failure (generate never rejects), - // so each records its own outcome and the article ships best-effort. - setPhase('producing') - setImageOutcome('active') - setAudioOutcome('active') + async function run(e: FormEvent) { + e.preventDefault() + const topic = String( + new FormData(e.currentTarget).get('topic') ?? '', + ).trim() + if (!topic || isRunning) return - const illustrate = (async () => { - const image = await assistant.image.generate({ - prompt: - `A striking editorial hero image for a blog post titled ` + - `"${post.title}". ${post.subtitle}. Modern, clean, cinematic, ` + - `high quality, no text.`, - }) - const shot = image?.images[0] - const url = - shot?.url ?? - (shot?.b64Json ? `data:image/png;base64,${shot.b64Json}` : null) - setImageUrl(url) - setImageOutcome(url ? 'ok' : 'failed') - })() + // Reset the surfaces so a re-run starts clean (clears prior draft/results). + post.clear() + image.reset() + speech.reset() - const narrate = (async () => { - const speech = await assistant.speech.generate({ - text: forNarration(post.body), - }) - if (speech?.audio) { - const audio = base64ToObjectUrl( - speech.audio, - speech.contentType ?? 'audio/mpeg', - ) - audioObjectUrl.current = audio - setAudioUrl(audio) - setAudioOutcome('ok') - } else { - setAudioOutcome('failed') - } - })() + // 1. Draft the post. `sendMessage` resolves to the schema-validated result; + // re-validate before chaining so a failed run doesn't proceed with a + // half-finished draft. + const parsed = BlogPostSchema.safeParse( + await post.sendMessage(`Write a blog post about: ${topic}`), + ) + if (!parsed.success) return + const draftedPost = parsed.data - await Promise.all([illustrate, narrate]) - setPhase('done') - } catch (err) { - setError(err instanceof Error ? err.message : 'Something went wrong.') - setPhase('idle') - } + // 2. Illustrate and narrate in parallel — fire and forget; the surfaces' + // reactive result/status/error drive the UI. `generate()` never rejects. + void image.generate({ + prompt: + `A striking editorial hero image for a blog post titled ` + + `"${draftedPost.title}". ${draftedPost.subtitle}. Modern, clean, ` + + `cinematic, high quality, no text.`, + }) + void speech.generate({ text: forNarration(draftedPost.body) }) } - const draft = assistant.chat.partial - const post = assistant.chat.final - const title = post?.title ?? draft.title - const subtitle = post?.subtitle ?? draft.subtitle - const body = post?.body ?? draft.body ?? '' - const showArticle = Boolean(post) || (phase === 'writing' && (title || body)) - return (
{/* Left: the studio controls */} @@ -233,62 +180,58 @@ function BlogStudio() { endpoint.

- - setTopic(e.target.value)} - onKeyDown={(e) => { - if (e.key === 'Enter') void run() - }} - placeholder="e.g. The quiet comeback of urban foxes" - disabled={isRunning} - className="mb-3 w-full rounded-lg border border-stone-300 bg-white px-4 py-3 text-stone-800 shadow-sm placeholder:text-stone-400 focus:border-amber-500 focus:outline-none focus:ring-2 focus:ring-amber-500/30 disabled:opacity-60" - /> - +
+ + + +
- {phase !== 'idle' && ( + {hasRun && (
} - // The stepper only renders once phase !== 'idle', so writing is - // either in progress or already finished. - state={phase === 'writing' ? 'active' : 'done'} + state={writingStep} /> } - state={outcomeToState(imageOutcome)} + state={statusToStep(image.status)} /> } - state={outcomeToState(audioOutcome)} + state={statusToStep(speech.status)} />
)} - {error && ( + {post.error && (
- {error} + {post.error.message}
)} @@ -307,12 +250,12 @@ function BlogStudio() { /> ) : (
- {imageOutcome === 'active' ? ( + {image.status === 'generating' ? (
Illustrating…
- ) : imageOutcome === 'failed' ? ( + ) : image.status === 'error' ? (
@@ -340,17 +283,17 @@ function BlogStudio() {
Written & narrated by TanStack AI - {audioUrl ? ( + {audioSrc ? (
-
- ) : audioOutcome === 'active' ? ( + ) : speech.status === 'generating' ? ( Recording voice-over… - ) : audioOutcome === 'failed' ? ( + ) : speech.status === 'error' ? ( Voice-over unavailable diff --git a/packages/ai-client/src/assistant-client.ts b/packages/ai-client/src/assistant-client.ts index 2f235fd73..284ba8f39 100644 --- a/packages/ai-client/src/assistant-client.ts +++ b/packages/ai-client/src/assistant-client.ts @@ -5,6 +5,7 @@ import type { AnyClientTool } from '@tanstack/ai/client' import type { AssistantClientOptions, OneShotCapabilityName, + OneShotCapabilityOptions, } from './assistant-types.js' /** @@ -45,12 +46,22 @@ export class AssistantClient< } const oneShotCapability = capability as OneShotCapabilityName + const capOptions: OneShotCapabilityOptions | undefined = + options[oneShotCapability] this.oneShots.set( oneShotCapability, new GenerationClient({ connection, id: id ? `${id}:${oneShotCapability}` : undefined, - body: { capability: oneShotCapability }, + // Merge per-capability forwardedProps under the routing discriminator. + body: { + ...capOptions?.forwardedProps, + capability: oneShotCapability, + }, + // Forward the result transform, if the caller supplied one. + ...(capOptions?.onResult !== undefined && { + onResult: capOptions.onResult, + }), ...callbacks?.oneShot?.(oneShotCapability), }), ) diff --git a/packages/ai-client/src/assistant-types.ts b/packages/ai-client/src/assistant-types.ts index e2594b136..ae9def5e9 100644 --- a/packages/ai-client/src/assistant-types.ts +++ b/packages/ai-client/src/assistant-types.ts @@ -29,6 +29,7 @@ import type { GenerationClientOptions, GenerationClientState, ImageGenerateInput, + InferGenerationOutput, SpeechGenerateInput, SummarizeGenerateInput, TranscriptionGenerateInput, @@ -68,6 +69,22 @@ export interface AssistantClientCallbacks { > } +/** + * Per-one-shot-capability options. Mirrors the standalone generation hooks: + * an optional `onResult` transform (its return type becomes the capability's + * `result` type) plus extra `forwardedProps` merged into the request body. + */ +export interface OneShotCapabilityOptions { + /** + * Transform the raw backend result before it is stored on the surface. + * `assistant..result` becomes this function's (non-nullish) + * return type. Return `undefined`/`null`/`void` to keep the raw result. + */ + onResult?: (result: TResult) => unknown + /** Extra fields merged into this capability's request body. */ + forwardedProps?: Record +} + /** Options for AssistantClient (framework-agnostic core). */ export interface AssistantClientOptions< TDef extends AssistantDefinition, @@ -89,6 +106,13 @@ export interface AssistantClientOptions< tools?: TChatTools forwardedProps?: Record } + /** Per-capability one-shot options (`onResult` transform, forwardedProps). */ + image?: OneShotCapabilityOptions + audio?: OneShotCapabilityOptions + speech?: OneShotCapabilityOptions + video?: OneShotCapabilityOptions + transcription?: OneShotCapabilityOptions + summarize?: OneShotCapabilityOptions /** * Reactive state callbacks forwarded into the sub-clients' constructors. * Set by framework hooks (not users) to wire up reactive state. @@ -330,6 +354,20 @@ export interface AssistantGenerationSurface { * callback — its captured tools type `chat.messages`' tool-call/result parts, * and its `outputSchema` (if any) adds typed `chat.partial` / `chat.final`. */ +/** + * The stored `result` type for a one-shot capability `K`: the `onResult` + * transform's return type when the options declare one (via + * {@link InferGenerationOutput}), otherwise the raw backend result. + */ +export type OneShotResultType< + TOptions, + TCap extends OneShotCapabilityName, +> = TCap extends keyof TOptions + ? TOptions[TCap] extends { onResult?: infer TFn } + ? InferGenerationOutput + : ResultByCapability[TCap] + : ResultByCapability[TCap] + export type AssistantSystem< TDef extends AssistantDefinition, /** @@ -338,6 +376,12 @@ export type AssistantSystem< * pass a second generic keep compiling; it has no effect on the surface. */ TChatTools extends ReadonlyArray = [], + /** + * The user's options object. Its per-capability `onResult` transforms drive + * each one-shot capability's `result` type. Defaults to `unknown` (no + * transforms → raw backend result types). + */ + TOptions = unknown, > = // `TChatTools` is deprecated/unused for typing; the constraint always holds, // so this reads the parameter (keeping it for hook back-compat) without @@ -349,7 +393,7 @@ export type AssistantSystem< : K extends OneShotCapabilityName ? AssistantGenerationSurface< GenerateInputByCapability[K], - ResultByCapability[K] + OneShotResultType > : never } diff --git a/packages/ai-react/src/use-assistant.ts b/packages/ai-react/src/use-assistant.ts index b9e53a4cd..5d5af2539 100644 --- a/packages/ai-react/src/use-assistant.ts +++ b/packages/ai-react/src/use-assistant.ts @@ -4,7 +4,6 @@ import { } from '@tanstack/ai-client/assistant' import { useEffect, useId, useMemo, useRef, useState } from 'react' import type { AssistantDefinition } from '@tanstack/ai/assistant' -import type { AnyClientTool } from '@tanstack/ai/client' import type { AssistantClientOptions, AssistantSystem, @@ -69,20 +68,21 @@ const initialOneShotState: OneShotState = { */ export function useAssistant< TDef extends AssistantDefinition, - TChatTools extends ReadonlyArray = [], ->( - assistant: TDef, - options: Omit< - AssistantClientOptions, + // Capture the options object so per-capability `onResult` transforms flow + // into each one-shot capability's `result` type. `any` tools keep the + // client-executed `chat.tools` option unconstrained (tool typing on the chat + // surface comes from the assistant definition, not this generic). + TOptions extends Omit< + AssistantClientOptions, 'assistant' | 'callbacks' >, -): AssistantSystem { +>(assistant: TDef, options: TOptions): AssistantSystem { const hookId = useId() const clientId = options.id ?? hookId const optionsRef = useRef(options) optionsRef.current = options - const activeClientRef = useRef | null>(null) + const activeClientRef = useRef | null>(null) const [chatState, setChatState] = useState(initialChatState) const [oneShotState, setOneShotState] = useState< @@ -92,12 +92,10 @@ export function useAssistant< const client = useMemo(() => { const initial = optionsRef.current - const instance = new AssistantClient({ + const instance = new AssistantClient({ + ...initial, assistant, - connection: initial.connection, id: clientId, - threadId: initial.threadId, - chat: initial.chat, callbacks: { chat: { onMessagesChange: (m) => { @@ -243,7 +241,7 @@ export function useAssistant< } } - return out as AssistantSystem + return out as AssistantSystem }, [client, chatState, oneShotState, partial, final]) return system diff --git a/packages/ai-react/tests/use-assistant.test.tsx b/packages/ai-react/tests/use-assistant.test.tsx index 153952945..49c22fa6c 100644 --- a/packages/ai-react/tests/use-assistant.test.tsx +++ b/packages/ai-react/tests/use-assistant.test.tsx @@ -1,7 +1,8 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, expectTypeOf, it } from 'vitest' import { act, renderHook } from '@testing-library/react' import { defineAssistant } from '@tanstack/ai/assistant' import { stream } from '@tanstack/ai-client' +import type { ImageGenerationResult } from '@tanstack/ai' import { useAssistant } from '../src/use-assistant.js' // A connection adapter that replays canned chunks for both capabilities, @@ -111,4 +112,40 @@ describe('useAssistant', () => { expect((result.current.chat as any).partial).toEqual({}) expect((result.current.chat as any).final).toBeNull() }) + + it('applies a one-shot onResult transform and infers its type', async () => { + const assistant = defineAssistant({ + chat: async function* () {} as any, + image: async () => ({}) as any, + }) + const { result } = renderHook(() => + useAssistant(assistant, { + connection: fakeConnection(), + image: { + onResult: (raw) => { + // The transform's input is the raw backend result, fully typed + // (`toEqualTypeOf` fails if `raw` were `any`). + expectTypeOf(raw).toEqualTypeOf() + return raw.images[0]?.url ?? null + }, + }, + }), + ) + + // The surface `result` is the transform's (non-nullish) return type — not + // the raw result, and not `any`. + expectTypeOf(result.current.image.result).toEqualTypeOf() + + const captured: { value: typeof result.current.image.result } = { + value: null, + } + await act(async () => { + captured.value = await result.current.image.generate({ prompt: 'a fox' }) + }) + + // The transformed value flows to both the resolved return and the + // reactive result state. + expect(captured.value).toBe('u') + expect(result.current.image.result).toBe('u') + }) }) diff --git a/packages/ai-solid/src/use-assistant.ts b/packages/ai-solid/src/use-assistant.ts index bf0d64d57..e749add1e 100644 --- a/packages/ai-solid/src/use-assistant.ts +++ b/packages/ai-solid/src/use-assistant.ts @@ -76,13 +76,17 @@ const defaultOneShotState: OneShotState = { export function useAssistant< TDef extends AssistantDefinition, TChatTools extends ReadonlyArray = [], + // Capture the options so per-capability `onResult` transforms flow into each + // one-shot capability's `result` type (`any` tools keep `chat.tools` + // unconstrained; chat tool typing comes from the definition). + TOptions extends Omit< + AssistantClientOptions, + 'assistant' | 'callbacks' + > = Omit, 'assistant' | 'callbacks'>, >( assistant: TDef, - options: Omit< - AssistantClientOptions, - 'assistant' | 'callbacks' - >, -): AssistantSystem { + options: TOptions, +): AssistantSystem { const hookId = createUniqueId() const clientId = options.id || hookId @@ -109,12 +113,10 @@ export function useAssistant< // constructor's `callbacks` option (VP2: ChatClient/GenerationClient only // accept these at construction time). const client = createMemo(() => { - return new AssistantClient({ + return new AssistantClient({ + ...options, assistant, - connection: options.connection, id: clientId, - ...(options.threadId !== undefined && { threadId: options.threadId }), - ...(options.chat !== undefined && { chat: options.chat }), callbacks: { chat: { onMessagesChange: (messages) => @@ -331,5 +333,5 @@ export function useAssistant< } // eslint-disable-next-line no-restricted-syntax -- built dynamically per declared capability; TS can't structurally verify the assembled object against the mapped AssistantSystem type - return system as unknown as AssistantSystem + return system as unknown as AssistantSystem } diff --git a/packages/ai-svelte/src/create-assistant.svelte.ts b/packages/ai-svelte/src/create-assistant.svelte.ts index ea3213dd0..0e011a1e1 100644 --- a/packages/ai-svelte/src/create-assistant.svelte.ts +++ b/packages/ai-svelte/src/create-assistant.svelte.ts @@ -60,13 +60,17 @@ interface OneShotState { export function createAssistant< TDef extends AssistantDefinition, TChatTools extends ReadonlyArray = [], + // Capture the options so per-capability `onResult` transforms flow into each + // one-shot capability's `result` type (`any` tools keep `chat.tools` + // unconstrained; chat tool typing comes from the definition). + TOptions extends Omit< + AssistantClientOptions, + 'assistant' | 'callbacks' + > = Omit, 'assistant' | 'callbacks'>, >( assistant: TDef, - options: Omit< - AssistantClientOptions, - 'assistant' | 'callbacks' - >, -): AssistantSystem & { dispose: () => void } { + options: TOptions, +): AssistantSystem & { dispose: () => void } { // Reactive state for the chat capability, if declared. let chatMessages = $state>>([]) let chatIsLoading = $state(false) @@ -114,12 +118,9 @@ export function createAssistant< // Create the AssistantClient eagerly, once. Reactive callbacks are passed // into the constructor (the only place ChatClient/GenerationClient accept // them) rather than via `updateOptions`, mirroring `createChat`. - const client = new AssistantClient({ + const client = new AssistantClient({ + ...options, assistant, - connection: options.connection, - ...(options.id !== undefined && { id: options.id }), - ...(options.threadId !== undefined && { threadId: options.threadId }), - ...(options.chat !== undefined && { chat: options.chat }), callbacks: { chat: { onMessagesChange: (newMessages) => { @@ -300,7 +301,7 @@ export function createAssistant< system.dispose = dispose // eslint-disable-next-line no-restricted-syntax -- built dynamically from a runtime `assistant.capabilities` array; the static AssistantSystem shape can't be verified structurally here, plus the added `dispose` field. - return system as unknown as AssistantSystem & { + return system as unknown as AssistantSystem & { dispose: () => void } } diff --git a/packages/ai-vue/src/use-assistant.ts b/packages/ai-vue/src/use-assistant.ts index 7b312dfce..c82b4068c 100644 --- a/packages/ai-vue/src/use-assistant.ts +++ b/packages/ai-vue/src/use-assistant.ts @@ -68,13 +68,17 @@ const DEFAULT_ONE_SHOT_STATE: OneShotState = { export function useAssistant< TDef extends AssistantDefinition, TChatTools extends ReadonlyArray = [], + // Capture the options so per-capability `onResult` transforms flow into each + // one-shot capability's `result` type (`any` tools keep `chat.tools` + // unconstrained; chat tool typing comes from the definition). + TOptions extends Omit< + AssistantClientOptions, + 'assistant' | 'callbacks' + > = Omit, 'assistant' | 'callbacks'>, >( assistant: TDef, - options: Omit< - AssistantClientOptions, - 'assistant' | 'callbacks' - >, -): AssistantSystem { + options: TOptions, +): AssistantSystem { // Chat sub-client state. const chatMessages = shallowRef>>([]) const chatIsLoading = shallowRef(false) @@ -101,7 +105,7 @@ export function useAssistant< // passed into the constructor's `callbacks` option — not wired up via // `updateOptions` afterwards — because the sub-clients (`ChatClient`, // `GenerationClient`) only accept these callbacks at construction time. - const client = new AssistantClient({ + const client = new AssistantClient({ ...options, assistant, callbacks: { @@ -255,5 +259,5 @@ export function useAssistant< // `useChat` accepts for its return, since consumers unwrap refs via // `.value` (script) or Vue's template auto-unwrapping. // eslint-disable-next-line no-restricted-syntax -- composable return shape (nested refs per capability) diverges from the framework-agnostic AssistantSystem type (plain values); TS can't structurally relate the two - return system as unknown as AssistantSystem + return system as unknown as AssistantSystem } diff --git a/packages/ai/skills/ai-core/assistant/SKILL.md b/packages/ai/skills/ai-core/assistant/SKILL.md index e6ccbed9a..96b348194 100644 --- a/packages/ai/skills/ai-core/assistant/SKILL.md +++ b/packages/ai/skills/ai-core/assistant/SKILL.md @@ -198,6 +198,26 @@ const assistant = useAssistant(blogAssistant, { `chat.forwardedProps` is also available on the same option, merged into every chat request alongside the reserved `capability` field. +Each **one-shot** capability accepts an optional transform too, keyed by the +capability name (`image`, `speech`, `audio`, `video`, `transcription`, +`summarize`): `image: { onResult, forwardedProps }`. `onResult` runs on the +raw backend result and its return type becomes `assistant..result` +— mirroring the standalone generation hooks. Return nothing to keep the raw +result. + +```typescript +import { fetchServerSentEvents } from '@tanstack/ai-react' +import { useAssistant } from '@tanstack/ai-react/assistant' +import { blogAssistant } from '../lib/assistant' + +const assistant = useAssistant(blogAssistant, { + connection: fetchServerSentEvents('/api/assistant'), + image: { onResult: (result) => result.images[0]?.url ?? null }, +}) + +assistant.image.result // string | null — the transform's return type +``` + ### 2. Structured output via `outputSchema` in the chat callback If the `chat` callback passes `outputSchema` to `chat()`, `assistant.chat` @@ -302,20 +322,22 @@ export const POST = (request: Request) => blogAssistant.handler(request) ### c. HIGH: Passing `model` or top-level generation options to `useAssistant` -There is no `model`/`tools`-for-generation option on `useAssistant` itself. -Model choice and per-capability options belong inside the server callback -(`openaiImage('gpt-image-2')`, `openaiSpeech('tts-1')`, …); the -only client-side option `useAssistant` accepts besides `connection` is +There is no `model`/`prompt` option on `useAssistant` itself. Model choice and +generation parameters belong inside the server callback +(`openaiImage('gpt-image-2')`, `openaiSpeech('tts-1')`, …). The client-side +options `useAssistant` accepts besides `connection` are `threadId`/`id`, `chat: { tools, forwardedProps }` (`tools` here registers **client-executed** tools' runtime implementations only — typing already comes from the server -callback) and `threadId`/`id`. +callback), and a per-one-shot-capability `{ onResult, forwardedProps }` (a +result transform + extra request-body fields — not model/prompt). ```typescript -// WRONG — no such options on useAssistant +// WRONG — no model/prompt options on useAssistant useAssistant(blogAssistant, { connection, model: 'gpt-5.5', prompt: '...' }) -// CORRECT — model/prompt options live in the server callback; useAssistant -// only takes connection, threadId/id, and chat.{tools,forwardedProps} +// CORRECT — model/prompt live in the server callback; useAssistant takes +// connection, threadId/id, chat.{tools,forwardedProps}, and per-capability +// { onResult, forwardedProps } useAssistant(blogAssistant, { connection: fetchServerSentEvents('/api/assistant'), }) From a075915acce8d6d80955cde5b39ad852b54d11b3 Mon Sep 17 00:00:00 2001 From: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Date: Thu, 16 Jul 2026 12:06:33 +1000 Subject: [PATCH 05/13] examples: add hooks and server Blog Studio pipeline variants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire client-hooks and server-fn studio demos next to the assistant version so the same draft → (hero ∥ narration) flow can be compared by where the pipeline is composed. Share schema/helpers, register routes and nav, and stream server progress over SSE. --- .../ts-react-chat/src/components/Header.tsx | 30 +- .../src/lib/blog-studio-server-fns.ts | 381 +++++++++++ examples/ts-react-chat/src/lib/blog-studio.ts | 67 ++ examples/ts-react-chat/src/routeTree.gen.ts | 63 ++ .../src/routes/api.blog-draft.ts | 65 ++ .../src/routes/api.blog-studio.ts | 36 +- .../src/routes/blog-studio-hooks.tsx | 451 +++++++++++++ .../src/routes/blog-studio-server.tsx | 599 ++++++++++++++++++ .../ts-react-chat/src/routes/blog-studio.tsx | 48 +- examples/ts-react-chat/src/routes/index.tsx | 14 + 10 files changed, 1687 insertions(+), 67 deletions(-) create mode 100644 examples/ts-react-chat/src/lib/blog-studio-server-fns.ts create mode 100644 examples/ts-react-chat/src/lib/blog-studio.ts create mode 100644 examples/ts-react-chat/src/routes/api.blog-draft.ts create mode 100644 examples/ts-react-chat/src/routes/blog-studio-hooks.tsx create mode 100644 examples/ts-react-chat/src/routes/blog-studio-server.tsx diff --git a/examples/ts-react-chat/src/components/Header.tsx b/examples/ts-react-chat/src/components/Header.tsx index 36023287f..813fc6c9d 100644 --- a/examples/ts-react-chat/src/components/Header.tsx +++ b/examples/ts-react-chat/src/components/Header.tsx @@ -11,10 +11,10 @@ import { Guitar, Home, Image, - Menu, LayoutGrid, - Mic, + Menu, MessageSquare, + Mic, Music, Newspaper, Plug, @@ -220,6 +220,32 @@ export default function Header() { Blog Studio + setIsOpen(false)} + className="flex items-center gap-3 p-3 rounded-lg hover:bg-gray-800 transition-colors mb-1" + activeProps={{ + className: + 'flex items-center gap-3 p-3 rounded-lg bg-cyan-600 hover:bg-cyan-700 transition-colors mb-1', + }} + > + + Blog Studio (hooks) + + + setIsOpen(false)} + className="flex items-center gap-3 p-3 rounded-lg hover:bg-gray-800 transition-colors mb-1" + activeProps={{ + className: + 'flex items-center gap-3 p-3 rounded-lg bg-cyan-600 hover:bg-cyan-700 transition-colors mb-1', + }} + > + + Blog Studio (server) + + setIsOpen(false)} diff --git a/examples/ts-react-chat/src/lib/blog-studio-server-fns.ts b/examples/ts-react-chat/src/lib/blog-studio-server-fns.ts new file mode 100644 index 000000000..2172bb6be --- /dev/null +++ b/examples/ts-react-chat/src/lib/blog-studio-server-fns.ts @@ -0,0 +1,381 @@ +import { createServerFn } from '@tanstack/react-start' +import { getRequest } from '@tanstack/react-start/server' +import { z } from 'zod' +import { + EventType, + chat, + generateImage, + generateSpeech, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { openaiImage, openaiSpeech, openaiText } from '@tanstack/ai-openai' +import { + BLOG_STUDIO_SYSTEM_PROMPT, + BlogPostSchema, + forNarration, + heroPromptFor, +} from './blog-studio' +import type { + ImageGenerationResult, + StreamChunk, + TTSResult, +} from '@tanstack/ai' +import type { BlogPost } from './blog-studio' + +// ============================================================================= +// Blog Studio (server) — server-composed pipeline with SSE progress +// ============================================================================= + +export type BlogStudioStep = 'drafting' | 'heroImage' | 'narration' + +/** CUSTOM event names emitted by `createBlogPostStreamFn`. */ +export const BLOG_STUDIO_SERVER_EVENTS = { + /** `{ step, status: 'started' | 'done' | 'error', result?, error? }` */ + STEP: 'pipeline:step', + /** Final `{ post, hero, audio }` once every step has finished. */ + RESULT: 'pipeline:result', +} as const + +export type BlogStudioStepEvent = + | { + step: BlogStudioStep + status: 'started' + } + | { + step: 'drafting' + status: 'done' + result: BlogPost + } + | { + step: 'heroImage' + status: 'done' + result: ImageGenerationResult + } + | { + step: 'narration' + status: 'done' + result: TTSResult + } + | { + step: BlogStudioStep + status: 'error' + error: string + } + +export type BlogStudioResultEvent = { + post: BlogPost + hero: ImageGenerationResult + audio: TTSResult +} + +function pipelineCustom(name: string, value: unknown): StreamChunk { + return { + type: EventType.CUSTOM, + name, + value, + timestamp: Date.now(), + } +} + +function stepEvent(value: BlogStudioStepEvent): StreamChunk { + return pipelineCustom(BLOG_STUDIO_SERVER_EVENTS.STEP, value) +} + +/** + * Unbounded push channel so parallel steps can interleave progress events + * into a single async iterable. + */ +function createEventChannel(): { + push: (chunk: StreamChunk) => void + close: () => void + [Symbol.asyncIterator]: () => AsyncIterator +} { + const queue: Array = [] + let notify: (() => void) | null = null + let closed = false + + return { + push(chunk) { + if (closed) return + queue.push(chunk) + notify?.() + }, + close() { + closed = true + notify?.() + }, + [Symbol.asyncIterator]() { + return { + async next(): Promise> { + for (;;) { + const chunk = queue.shift() + if (chunk !== undefined) return { value: chunk, done: false } + if (closed) return { value: undefined, done: true } + await new Promise((resolve) => { + notify = resolve + }) + notify = null + } + }, + } + }, + } +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null +} + +function throwIfAborted(signal: AbortSignal): void { + if (signal.aborted) { + throw new DOMException('Aborted', 'AbortError') + } +} + +function isAbortError(err: unknown): boolean { + return ( + (err instanceof DOMException && err.name === 'AbortError') || + (err instanceof Error && err.name === 'AbortError') + ) +} + +/** + * One controller shared by the pipeline, provider calls, and the SSE + * response body. Chains Start's request signal (client Stop / disconnect) + * so abort propagates end to end. + */ +function linkAbortController(requestSignal: AbortSignal): AbortController { + const abortController = new AbortController() + if (requestSignal.aborted) { + abortController.abort(requestSignal.reason) + } else { + requestSignal.addEventListener( + 'abort', + () => abortController.abort(requestSignal.reason), + { once: true }, + ) + } + return abortController +} + +/** + * Server-side composition: draft → (hero ∥ narration). + * + * Streams: + * - live structured-output chunks while drafting + * - `pipeline:step` started/done/error for each step + * - `pipeline:result` with the finished artifact + * + * Abort checks run only after `await` points — `signal.aborted` is live and + * can flip while work is in flight; a check on a brand-new controller with + * no await in between is always false. + */ +async function* createBlogPostPipeline( + topic: string, + abortController: AbortController, +): AsyncGenerator { + const { signal } = abortController + const runId = `blog-server-${Date.now()}` + const threadId = 'blog-studio-server' + + yield { + type: EventType.RUN_STARTED, + runId, + threadId, + timestamp: Date.now(), + } + + try { + // ── 1. Draft (streamed structured output) ───────────────────────────── + yield stepEvent({ step: 'drafting', status: 'started' }) + + const draftStream = chat({ + adapter: openaiText('gpt-5.5'), + messages: [ + { role: 'user', content: `Write a blog post about: ${topic}` }, + ], + systemPrompts: [BLOG_STUDIO_SYSTEM_PROMPT], + outputSchema: BlogPostSchema, + stream: true, + abortController, + }) + + let post: BlogPost | undefined + for await (const chunk of draftStream) { + // After each streamed chunk (async gap) — abort may have landed. + throwIfAborted(signal) + // Forward live draft chunks; skip nested run lifecycle so our outer + // RUN_* stay authoritative for this server-fn response. + if ( + chunk.type !== EventType.RUN_STARTED && + chunk.type !== EventType.RUN_FINISHED && + chunk.type !== EventType.RUN_ERROR + ) { + yield chunk + } + if ( + chunk.type === EventType.CUSTOM && + chunk.name === 'structured-output.complete' + ) { + const value = chunk.value + // Prefer the validated object from the orchestrator; re-parse so a + // malformed payload can't proceed as a half-typed draft. + if (isRecord(value)) { + const parsed = BlogPostSchema.safeParse(value.object) + if (parsed.success) post = parsed.data + } + } + } + + throwIfAborted(signal) + + if (!post) { + const message = 'Drafting did not produce a structured blog post' + yield stepEvent({ step: 'drafting', status: 'error', error: message }) + throw new Error(message) + } + + yield stepEvent({ step: 'drafting', status: 'done', result: post }) + + // ── 2. Hero + narration in parallel ─────────────────────────────────── + // Image/speech activities don't take AbortSignal; on abort we stop + // consuming the channel and throw. Close the channel so `for await` + // unblocks even if providers are still finishing. + const channel = createEventChannel() + const closeChannelOnAbort = () => channel.close() + signal.addEventListener('abort', closeChannelOnAbort, { once: true }) + + channel.push(stepEvent({ step: 'heroImage', status: 'started' })) + channel.push(stepEvent({ step: 'narration', status: 'started' })) + + const work = Promise.all([ + generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: heroPromptFor(post), + size: '1536x1024', + }).then( + (hero) => { + throwIfAborted(signal) + channel.push( + stepEvent({ step: 'heroImage', status: 'done', result: hero }), + ) + return hero + }, + (err: unknown) => { + if (isAbortError(err) || signal.aborted) throw err + const message = + err instanceof Error ? err.message : 'Image generation failed' + channel.push( + stepEvent({ step: 'heroImage', status: 'error', error: message }), + ) + throw err + }, + ), + generateSpeech({ + adapter: openaiSpeech('tts-1'), + text: forNarration(post.body), + voice: 'alloy', + }).then( + (audio) => { + throwIfAborted(signal) + channel.push( + stepEvent({ step: 'narration', status: 'done', result: audio }), + ) + return audio + }, + (err: unknown) => { + if (isAbortError(err) || signal.aborted) throw err + const message = + err instanceof Error ? err.message : 'Speech generation failed' + channel.push( + stepEvent({ step: 'narration', status: 'error', error: message }), + ) + throw err + }, + ), + ]).finally(() => { + signal.removeEventListener('abort', closeChannelOnAbort) + channel.close() + }) + + for await (const chunk of channel) { + throwIfAborted(signal) + yield chunk + } + + const [hero, audio] = await work + throwIfAborted(signal) + + yield pipelineCustom(BLOG_STUDIO_SERVER_EVENTS.RESULT, { + post, + hero, + audio, + } satisfies BlogStudioResultEvent) + + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + finishReason: 'stop', + timestamp: Date.now(), + } + } catch (err) { + if (signal.aborted || isAbortError(err)) { + yield { + type: EventType.RUN_ERROR, + message: 'Aborted', + timestamp: Date.now(), + } + return + } + const message = err instanceof Error ? err.message : 'Pipeline failed' + yield { + type: EventType.RUN_ERROR, + message, + timestamp: Date.now(), + } + } +} + +/** + * One server function: composes draft → (hero ∥ narration) and streams step + * progress over SSE. Pair with a client that reads the Response body as SSE. + * + * Abort sources (merged into one controller): + * - the HTTP request signal (`getRequest().signal` — client Stop / disconnect) + * - SSE body cancel via `toServerSentEventsResponse({ abortController })` + * + * Note: Start 1.159 does not put `signal` on the handler ctx; client `signal` + * is honored by aborting the underlying fetch/request, which trips + * `getRequest().signal`. + */ +export const createBlogPostStreamFn = createServerFn({ method: 'POST' }) + .inputValidator(z.object({ topic: z.string().min(1) })) + .handler(({ data }) => { + const abortController = linkAbortController(getRequest().signal) + return toServerSentEventsResponse( + createBlogPostPipeline(data.topic, abortController), + { abortController }, + ) + }) + +/** One-shot regenerate — same models as the pipeline, no streaming. */ +export const regenerateBlogHeroFn = createServerFn({ method: 'POST' }) + .inputValidator(z.object({ prompt: z.string().min(1) })) + .handler(({ data }) => + generateImage({ + adapter: openaiImage('gpt-image-2'), + prompt: data.prompt, + size: '1536x1024', + }), + ) + +export const regenerateBlogNarrationFn = createServerFn({ method: 'POST' }) + .inputValidator(z.object({ text: z.string().min(1) })) + .handler(({ data }) => + generateSpeech({ + adapter: openaiSpeech('tts-1'), + text: data.text, + voice: 'alloy', + }), + ) diff --git a/examples/ts-react-chat/src/lib/blog-studio.ts b/examples/ts-react-chat/src/lib/blog-studio.ts new file mode 100644 index 000000000..64c37bf62 --- /dev/null +++ b/examples/ts-react-chat/src/lib/blog-studio.ts @@ -0,0 +1,67 @@ +import { z } from 'zod' + +/** + * Shared Blog Studio helpers used by the server-fn pipeline and the + * hooks-based client pipeline. Schema + prompt builders live here so both + * paths draft the same shape and produce matching hero / narration inputs. + * + * The assistant version (`/blog-studio`) mirrors this schema on its API route + * for the typed `useAssistant` client. + */ + +export const BlogPostSchema = z.object({ + title: z.string().describe('A punchy, editorial blog post title'), + subtitle: z.string().describe('A one-sentence standfirst / subtitle'), + body: z + .string() + .describe( + 'The full blog post body as GitHub-flavored Markdown: use ## / ### ' + + 'headings, short paragraphs, and the occasional list. ~400-600 words.', + ), +}) + +export type BlogPost = z.infer + +export const BLOG_STUDIO_SYSTEM_PROMPT = + 'You are a seasoned staff writer. Given a topic, write one engaging, ' + + 'well-structured blog post. Return a title, a short subtitle, and the ' + + 'body as GitHub-flavored Markdown with section headings and tight ' + + 'paragraphs. Be vivid and concrete; avoid filler and clichés.' + +/** + * Build the hero-image prompt from a drafted post. Used by the server-side + * server pipeline and by the client's "Regenerate hero image" button so both + * produce the same style of illustration. + */ +export function heroPromptFor(post: BlogPost): string { + return ( + `A striking editorial hero image for a blog post titled ` + + `"${post.title}". ${post.subtitle}. Modern, clean, cinematic, ` + + `high quality, no text.` + ) +} + +/** + * Prepare the post body for narration: strip Markdown so TTS doesn't read the + * syntax aloud, and cap the length at a sentence boundary (OpenAI TTS rejects + * input over 4096 characters, and long posts easily exceed it). + */ +export function forNarration(markdown: string, max = 4000): string { + const plain = markdown + .replace(/^#{1,6}\s+/gm, '') // headings + .replace(/^\s*[-*+]\s+/gm, '') // list bullets + .replace(/^\s*>\s?/gm, '') // blockquotes + .replace(/\*\*(.*?)\*\*/g, '$1') // bold + .replace(/\*(.*?)\*/g, '$1') // italic + .replace(/`([^`]+)`/g, '$1') // inline code + .replace(/\[([^\]]+)\]\([^)]*\)/g, '$1') // links → link text + .replace(/\n{3,}/g, '\n\n') + .trim() + if (plain.length <= max) return plain + const clipped = plain.slice(0, max) + const boundary = Math.max( + clipped.lastIndexOf('. '), + clipped.lastIndexOf('\n'), + ) + return (boundary > max / 2 ? clipped.slice(0, boundary + 1) : clipped).trim() +} diff --git a/examples/ts-react-chat/src/routeTree.gen.ts b/examples/ts-react-chat/src/routeTree.gen.ts index 06b43bd5b..dd6bb728e 100644 --- a/examples/ts-react-chat/src/routeTree.gen.ts +++ b/examples/ts-react-chat/src/routeTree.gen.ts @@ -21,6 +21,8 @@ import { Route as ImageToolReproRouteImport } from './routes/image-tool-repro' import { Route as ImageGenRouteImport } from './routes/image-gen' import { Route as GenerationHooksRouteImport } from './routes/generation-hooks' import { Route as CapabilityDemoRouteImport } from './routes/capability-demo' +import { Route as BlogStudioServerRouteImport } from './routes/blog-studio-server' +import { Route as BlogStudioHooksRouteImport } from './routes/blog-studio-hooks' import { Route as BlogStudioRouteImport } from './routes/blog-studio' import { Route as IndexRouteImport } from './routes/index' import { Route as GenerationsVideoRouteImport } from './routes/generations.video' @@ -50,6 +52,7 @@ import { Route as ApiImageToolReproRouteImport } from './routes/api.image-tool-r import { Route as ApiImageGenRouteImport } from './routes/api.image-gen' import { Route as ApiCapabilityDemoRouteImport } from './routes/api.capability-demo' import { Route as ApiBlogStudioRouteImport } from './routes/api.blog-studio' +import { Route as ApiBlogDraftRouteImport } from './routes/api.blog-draft' import { Route as ExampleGuitarsIndexRouteImport } from './routes/example.guitars/index' import { Route as ExampleGuitarsGuitarIdRouteImport } from './routes/example.guitars/$guitarId' import { Route as ApiGenerateVideoRouteImport } from './routes/api.generate.video' @@ -117,6 +120,16 @@ const CapabilityDemoRoute = CapabilityDemoRouteImport.update({ path: '/capability-demo', getParentRoute: () => rootRouteImport, } as any) +const BlogStudioServerRoute = BlogStudioServerRouteImport.update({ + id: '/blog-studio-server', + path: '/blog-studio-server', + getParentRoute: () => rootRouteImport, +} as any) +const BlogStudioHooksRoute = BlogStudioHooksRouteImport.update({ + id: '/blog-studio-hooks', + path: '/blog-studio-hooks', + getParentRoute: () => rootRouteImport, +} as any) const BlogStudioRoute = BlogStudioRouteImport.update({ id: '/blog-studio', path: '/blog-studio', @@ -265,6 +278,11 @@ const ApiBlogStudioRoute = ApiBlogStudioRouteImport.update({ path: '/api/blog-studio', getParentRoute: () => rootRouteImport, } as any) +const ApiBlogDraftRoute = ApiBlogDraftRouteImport.update({ + id: '/api/blog-draft', + path: '/api/blog-draft', + getParentRoute: () => rootRouteImport, +} as any) const ExampleGuitarsIndexRoute = ExampleGuitarsIndexRouteImport.update({ id: '/example/guitars/', path: '/example/guitars/', @@ -299,6 +317,8 @@ const ApiGenerateAudioRoute = ApiGenerateAudioRouteImport.update({ export interface FileRoutesByFullPath { '/': typeof IndexRoute '/blog-studio': typeof BlogStudioRoute + '/blog-studio-hooks': typeof BlogStudioHooksRoute + '/blog-studio-server': typeof BlogStudioServerRoute '/capability-demo': typeof CapabilityDemoRoute '/generation-hooks': typeof GenerationHooksRoute '/image-gen': typeof ImageGenRoute @@ -311,6 +331,7 @@ export interface FileRoutesByFullPath { '/server-fn-chat': typeof ServerFnChatRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute + '/api/blog-draft': typeof ApiBlogDraftRoute '/api/blog-studio': typeof ApiBlogStudioRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute '/api/image-gen': typeof ApiImageGenRoute @@ -348,6 +369,8 @@ export interface FileRoutesByFullPath { export interface FileRoutesByTo { '/': typeof IndexRoute '/blog-studio': typeof BlogStudioRoute + '/blog-studio-hooks': typeof BlogStudioHooksRoute + '/blog-studio-server': typeof BlogStudioServerRoute '/capability-demo': typeof CapabilityDemoRoute '/generation-hooks': typeof GenerationHooksRoute '/image-gen': typeof ImageGenRoute @@ -360,6 +383,7 @@ export interface FileRoutesByTo { '/server-fn-chat': typeof ServerFnChatRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute + '/api/blog-draft': typeof ApiBlogDraftRoute '/api/blog-studio': typeof ApiBlogStudioRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute '/api/image-gen': typeof ApiImageGenRoute @@ -398,6 +422,8 @@ export interface FileRoutesById { __root__: typeof rootRouteImport '/': typeof IndexRoute '/blog-studio': typeof BlogStudioRoute + '/blog-studio-hooks': typeof BlogStudioHooksRoute + '/blog-studio-server': typeof BlogStudioServerRoute '/capability-demo': typeof CapabilityDemoRoute '/generation-hooks': typeof GenerationHooksRoute '/image-gen': typeof ImageGenRoute @@ -410,6 +436,7 @@ export interface FileRoutesById { '/server-fn-chat': typeof ServerFnChatRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute + '/api/blog-draft': typeof ApiBlogDraftRoute '/api/blog-studio': typeof ApiBlogStudioRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute '/api/image-gen': typeof ApiImageGenRoute @@ -449,6 +476,8 @@ export interface FileRouteTypes { fullPaths: | '/' | '/blog-studio' + | '/blog-studio-hooks' + | '/blog-studio-server' | '/capability-demo' | '/generation-hooks' | '/image-gen' @@ -461,6 +490,7 @@ export interface FileRouteTypes { | '/server-fn-chat' | '/threads' | '/typesafe-tools' + | '/api/blog-draft' | '/api/blog-studio' | '/api/capability-demo' | '/api/image-gen' @@ -498,6 +528,8 @@ export interface FileRouteTypes { to: | '/' | '/blog-studio' + | '/blog-studio-hooks' + | '/blog-studio-server' | '/capability-demo' | '/generation-hooks' | '/image-gen' @@ -510,6 +542,7 @@ export interface FileRouteTypes { | '/server-fn-chat' | '/threads' | '/typesafe-tools' + | '/api/blog-draft' | '/api/blog-studio' | '/api/capability-demo' | '/api/image-gen' @@ -547,6 +580,8 @@ export interface FileRouteTypes { | '__root__' | '/' | '/blog-studio' + | '/blog-studio-hooks' + | '/blog-studio-server' | '/capability-demo' | '/generation-hooks' | '/image-gen' @@ -559,6 +594,7 @@ export interface FileRouteTypes { | '/server-fn-chat' | '/threads' | '/typesafe-tools' + | '/api/blog-draft' | '/api/blog-studio' | '/api/capability-demo' | '/api/image-gen' @@ -597,6 +633,8 @@ export interface FileRouteTypes { export interface RootRouteChildren { IndexRoute: typeof IndexRoute BlogStudioRoute: typeof BlogStudioRoute + BlogStudioHooksRoute: typeof BlogStudioHooksRoute + BlogStudioServerRoute: typeof BlogStudioServerRoute CapabilityDemoRoute: typeof CapabilityDemoRoute GenerationHooksRoute: typeof GenerationHooksRoute ImageGenRoute: typeof ImageGenRoute @@ -609,6 +647,7 @@ export interface RootRouteChildren { ServerFnChatRoute: typeof ServerFnChatRoute ThreadsRoute: typeof ThreadsRoute TypesafeToolsRoute: typeof TypesafeToolsRoute + ApiBlogDraftRoute: typeof ApiBlogDraftRoute ApiBlogStudioRoute: typeof ApiBlogStudioRoute ApiCapabilityDemoRoute: typeof ApiCapabilityDemoRoute ApiImageGenRoute: typeof ApiImageGenRoute @@ -730,6 +769,20 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof CapabilityDemoRouteImport parentRoute: typeof rootRouteImport } + '/blog-studio-server': { + id: '/blog-studio-server' + path: '/blog-studio-server' + fullPath: '/blog-studio-server' + preLoaderRoute: typeof BlogStudioServerRouteImport + parentRoute: typeof rootRouteImport + } + '/blog-studio-hooks': { + id: '/blog-studio-hooks' + path: '/blog-studio-hooks' + fullPath: '/blog-studio-hooks' + preLoaderRoute: typeof BlogStudioHooksRouteImport + parentRoute: typeof rootRouteImport + } '/blog-studio': { id: '/blog-studio' path: '/blog-studio' @@ -933,6 +986,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiBlogStudioRouteImport parentRoute: typeof rootRouteImport } + '/api/blog-draft': { + id: '/api/blog-draft' + path: '/api/blog-draft' + fullPath: '/api/blog-draft' + preLoaderRoute: typeof ApiBlogDraftRouteImport + parentRoute: typeof rootRouteImport + } '/example/guitars/': { id: '/example/guitars/' path: '/example/guitars' @@ -981,6 +1041,8 @@ declare module '@tanstack/react-router' { const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, BlogStudioRoute: BlogStudioRoute, + BlogStudioHooksRoute: BlogStudioHooksRoute, + BlogStudioServerRoute: BlogStudioServerRoute, CapabilityDemoRoute: CapabilityDemoRoute, GenerationHooksRoute: GenerationHooksRoute, ImageGenRoute: ImageGenRoute, @@ -993,6 +1055,7 @@ const rootRouteChildren: RootRouteChildren = { ServerFnChatRoute: ServerFnChatRoute, ThreadsRoute: ThreadsRoute, TypesafeToolsRoute: TypesafeToolsRoute, + ApiBlogDraftRoute: ApiBlogDraftRoute, ApiBlogStudioRoute: ApiBlogStudioRoute, ApiCapabilityDemoRoute: ApiCapabilityDemoRoute, ApiImageGenRoute: ApiImageGenRoute, diff --git a/examples/ts-react-chat/src/routes/api.blog-draft.ts b/examples/ts-react-chat/src/routes/api.blog-draft.ts new file mode 100644 index 000000000..0ee04b793 --- /dev/null +++ b/examples/ts-react-chat/src/routes/api.blog-draft.ts @@ -0,0 +1,65 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + chat, + chatParamsFromRequestBody, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' +import { BLOG_STUDIO_SYSTEM_PROMPT, BlogPostSchema } from '../lib/blog-studio' +import type { StreamChunk } from '@tanstack/ai' + +/** + * Draft-only endpoint for the hooks-based Blog Studio. Returns a streaming + * structured-output chat so `useChat({ outputSchema })` can drive the client + * draft step before image/speech fan-out. + */ +export const Route = createFileRoute('/api/blog-draft')({ + server: { + handlers: { + POST: async ({ request }) => { + // Attach the listener first, then re-check aborted so a race between + // an early check and addEventListener can't drop the abort. + const abortController = new AbortController() + const onAbort = () => abortController.abort() + request.signal.addEventListener('abort', onAbort, { once: true }) + if (request.signal.aborted) { + onAbort() + return new Response(null, { status: 499 }) + } + + let params + try { + params = await chatParamsFromRequestBody(await request.json()) + } catch (error) { + return new Response( + error instanceof Error ? error.message : 'Bad request', + { status: 400 }, + ) + } + + try { + const stream = chat({ + adapter: openaiText('gpt-5.5'), + messages: params.messages, + systemPrompts: [BLOG_STUDIO_SYSTEM_PROMPT], + outputSchema: BlogPostSchema, + stream: true, + threadId: params.threadId, + runId: params.runId, + abortController, + }) as AsyncIterable + return toServerSentEventsResponse(stream, { abortController }) + } catch (error) { + console.error('[api/blog-draft] Error:', error) + return new Response( + JSON.stringify({ error: 'Internal server error' }), + { + status: 500, + headers: { 'Content-Type': 'application/json' }, + }, + ) + } + }, + }, + }, +}) diff --git a/examples/ts-react-chat/src/routes/api.blog-studio.ts b/examples/ts-react-chat/src/routes/api.blog-studio.ts index 9db94c5e4..c5748491d 100644 --- a/examples/ts-react-chat/src/routes/api.blog-studio.ts +++ b/examples/ts-react-chat/src/routes/api.blog-studio.ts @@ -2,38 +2,10 @@ import { createFileRoute } from '@tanstack/react-router' import { chat, generateImage, generateSpeech } from '@tanstack/ai' import { defineAssistant } from '@tanstack/ai/assistant' import { openaiImage, openaiSpeech, openaiText } from '@tanstack/ai-openai' -import { z } from 'zod' +import { BLOG_STUDIO_SYSTEM_PROMPT, BlogPostSchema } from '../lib/blog-studio' -/** - * The shape a blog post is drafted into. The `chat` capability declares this - * as its `outputSchema`, so `assistant.chat.sendMessage(...)` resolves to a - * validated `{ title, subtitle, body } | null` on the client — no manual - * parsing, and `assistant.chat.partial` streams a live draft. - * - * Exported so the client page (`blog-studio.tsx`) can import the same schema - * and stay in type-sync. Only this schema lives at module scope; the actual - * `defineAssistant` (with its server adapters) is built inside the POST - * handler so importing the schema never pulls provider SDKs into the browser - * bundle. - */ -export const BlogPostSchema = z.object({ - title: z.string().describe('A punchy, editorial blog post title'), - subtitle: z.string().describe('A one-sentence standfirst / subtitle'), - body: z - .string() - .describe( - 'The full blog post body as GitHub-flavored Markdown: use ## / ### ' + - 'headings, short paragraphs, and the occasional list. ~400-600 words.', - ), -}) - -export type BlogPost = z.infer - -const SYSTEM_PROMPT = - 'You are a seasoned staff writer. Given a topic, write one engaging, ' + - 'well-structured blog post. Return a title, a short subtitle, and the ' + - 'body as GitHub-flavored Markdown with section headings and tight ' + - 'paragraphs. Be vivid and concrete; avoid filler and clichés.' +// Re-export so the assistant page can import schema + types from this route +// module (keeps client types in sync without pulling server adapters). export const Route = createFileRoute('/api/blog-studio')({ server: { @@ -49,7 +21,7 @@ export const Route = createFileRoute('/api/blog-studio')({ chat({ adapter: openaiText('gpt-5.5'), messages: req.messages, - systemPrompts: [SYSTEM_PROMPT], + systemPrompts: [BLOG_STUDIO_SYSTEM_PROMPT], outputSchema: BlogPostSchema, stream: true, threadId: req.threadId, diff --git a/examples/ts-react-chat/src/routes/blog-studio-hooks.tsx b/examples/ts-react-chat/src/routes/blog-studio-hooks.tsx new file mode 100644 index 000000000..219e220ec --- /dev/null +++ b/examples/ts-react-chat/src/routes/blog-studio-hooks.tsx @@ -0,0 +1,451 @@ +import { Link, createFileRoute } from '@tanstack/react-router' +import { useEffect } from 'react' +import { + fetchServerSentEvents, + useChat, + useGenerateImage, + useGenerateSpeech, +} from '@tanstack/ai-react' +import ReactMarkdown from 'react-markdown' +import remarkGfm from 'remark-gfm' +import { + AlertTriangle, + Check, + Image as ImageIcon, + Loader2, + Newspaper, + PenLine, + Sparkles, + Square, + Volume2, + Wand2, +} from 'lucide-react' +import { BlogPostSchema, forNarration, heroPromptFor } from '../lib/blog-studio' +import type { GenerationClientState } from '@tanstack/ai-client' +import type { ImageGenerationResult, TTSResult } from '@tanstack/ai' +import type { FormEvent, ReactNode } from 'react' + +export const Route = createFileRoute('/blog-studio-hooks')({ + component: BlogStudioHooks, +}) + +type StepState = 'pending' | 'active' | 'done' | 'failed' + +function generationToStep(status: GenerationClientState): StepState { + return status === 'generating' + ? 'active' + : status === 'success' + ? 'done' + : status === 'error' + ? 'failed' + : 'pending' +} + +function imageUrlOf(result: ImageGenerationResult | null): string | null { + const image = result?.images[0] + if (!image) return null + return ( + image.url ?? + (image.b64Json ? `data:image/png;base64,${image.b64Json}` : null) + ) +} + +function audioSrcOf(result: TTSResult | null): string | null { + if (!result) return null + return `data:${result.contentType ?? 'audio/mpeg'};base64,${result.audio}` +} + +function BlogStudioHooks() { + const hero = useGenerateImage({ + connection: fetchServerSentEvents('/api/generate/image'), + body: { model: 'gpt-image-2' }, + }) + + const narration = useGenerateSpeech({ + connection: fetchServerSentEvents('/api/generate/speech'), + body: { provider: 'openai' }, + }) + + const { + sendMessage, + isLoading: isDrafting, + error: draftError, + stop: stopDraft, + partial, + final, + clear, + } = useChat({ + id: 'blog-studio-hooks-draft', + outputSchema: BlogPostSchema, + connection: fetchServerSentEvents('/api/blog-draft'), + }) + + // `final` only becomes non-null once the structured draft is complete. + // `generate()` is a no-op if that hook is already in flight, so a + // duplicate effect run (e.g. Strict Mode) is harmless. + useEffect(() => { + if (!final) return + + void Promise.all([ + hero.generate({ + prompt: heroPromptFor(final), + size: '1536x1024', + }), + narration.generate({ + text: forNarration(final.body), + voice: 'alloy', + }), + ]) + }, [final, hero.generate, narration.generate]) + + const post = final + const imageUrl = imageUrlOf(hero.result) + const audioSrc = audioSrcOf(narration.result) + + const isRunning = isDrafting || hero.isLoading || narration.isLoading + + const hasRun = Boolean(isDrafting || final || draftError) + + const liveDraft = partial + const draftedChars = liveDraft.body?.length ?? 0 + const writingStep: StepState = isDrafting + ? 'active' + : post + ? 'done' + : draftError + ? 'failed' + : 'pending' + + const shownTitle = post?.title ?? liveDraft.title + const shownSubtitle = post?.subtitle ?? liveDraft.subtitle + const shownBody = post?.body ?? liveDraft.body + const showArticle = Boolean(shownTitle || shownBody) + + const heroStep = generationToStep(hero.status) + const narrationStep = generationToStep(narration.status) + + function stopAll() { + stopDraft() + hero.stop() + narration.stop() + } + + function run(e: FormEvent) { + e.preventDefault() + const topic = String( + new FormData(e.currentTarget).get('topic') ?? '', + ).trim() + if (!topic || isRunning) return + + clear() + hero.reset() + narration.reset() + + void sendMessage(`Write a blog post about: ${topic}`) + } + + const pipelineError = + draftError?.message ?? hero.error?.message ?? narration.error?.message + + return ( +
+ + +
+ {showArticle ? ( +
+
+ {imageUrl && !hero.isLoading ? ( + {shownTitle + ) : ( +
+ {hero.isLoading ? ( +
+ + Illustrating… +
+ ) : heroStep === 'failed' ? ( +
+ + + Couldn't generate a hero image + +
+ ) : ( + + )} +
+ )} +
+ +
+ {shownTitle ? ( +

+ {shownTitle} +

+ ) : ( +
+ )} + {shownSubtitle && ( +

{shownSubtitle}

+ )} + +
+
+ +
+ Written & narrated by TanStack AI + {narration.isLoading ? ( + + Recording + voice-over… + + ) : audioSrc ? ( +
+ +
+ ) : narrationStep === 'failed' ? ( + + + Voice-over unavailable + + ) : null} +
+ +
+ + {shownBody ?? ''} + + {writingStep === 'active' && ( + + )} +
+
+
+ ) : isRunning ? ( +
+
+ +

+ {draftedChars > 0 + ? `Drafting the article… ${draftedChars} characters so far.` + : 'Starting the draft…'} +

+
+
+ ) : ( +
+
+ +

+ Your finished post — hero image, article, and voice-over — will + appear here. +

+
+
+ )} +
+
+ ) +} + +function StepRow({ + label, + icon, + state, + detail, +}: { + label: string + icon: ReactNode + state: StepState + detail?: string +}) { + const cls = + state === 'active' + ? 'text-violet-700 font-medium' + : state === 'done' + ? 'text-stone-500' + : state === 'failed' + ? 'text-amber-600' + : 'text-stone-300' + return ( + + {state === 'active' ? ( + + ) : state === 'done' ? ( + + ) : state === 'failed' ? ( + + ) : ( + icon + )} + {label} + {detail && ( + + {detail} + + )} + + ) +} diff --git a/examples/ts-react-chat/src/routes/blog-studio-server.tsx b/examples/ts-react-chat/src/routes/blog-studio-server.tsx new file mode 100644 index 000000000..d886e6abd --- /dev/null +++ b/examples/ts-react-chat/src/routes/blog-studio-server.tsx @@ -0,0 +1,599 @@ +import { Link, createFileRoute } from '@tanstack/react-router' +import { useRef, useState } from 'react' +import ReactMarkdown from 'react-markdown' +import remarkGfm from 'remark-gfm' +import { + AlertTriangle, + Check, + Image as ImageIcon, + Loader2, + Newspaper, + PenLine, + Sparkles, + Square, + Volume2, + Wand2, +} from 'lucide-react' +import { EventType } from '@tanstack/ai' +import { forNarration, heroPromptFor } from '../lib/blog-studio' +import { + BLOG_STUDIO_SERVER_EVENTS, + createBlogPostStreamFn, + regenerateBlogHeroFn, + regenerateBlogNarrationFn, +} from '../lib/blog-studio-server-fns' +import type { ImageGenerationResult, TTSResult } from '@tanstack/ai' +import type { BlogPost } from '../lib/blog-studio' +import type { BlogStudioStep } from '../lib/blog-studio-server-fns' +import type { FormEvent, ReactNode } from 'react' + +export const Route = createFileRoute('/blog-studio-server')({ + component: BlogStudioServer, +}) + +type StepState = 'pending' | 'active' | 'done' | 'failed' + +function imageUrlOf(result: ImageGenerationResult | null): string | null { + const image = result?.images[0] + if (!image) return null + return ( + image.url ?? + (image.b64Json ? `data:image/png;base64,${image.b64Json}` : null) + ) +} + +function audioSrcOf(result: TTSResult | null): string | null { + if (!result) return null + return `data:${result.contentType ?? 'audio/mpeg'};base64,${result.audio}` +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null +} + +function isBlogPost(value: unknown): value is BlogPost { + return ( + isRecord(value) && + typeof value.title === 'string' && + typeof value.subtitle === 'string' && + typeof value.body === 'string' + ) +} + +function isImageResult(value: unknown): value is ImageGenerationResult { + return isRecord(value) && Array.isArray(value.images) +} + +function isTtsResult(value: unknown): value is TTSResult { + return isRecord(value) && typeof value.audio === 'string' +} + +/** Read `data: {…}\\n\\n` SSE from a Response (same format as toServerSentEventsResponse). */ +async function* readSseJson( + response: Response, + signal: AbortSignal, +): AsyncGenerator { + if (!response.ok) { + throw new Error( + `HTTP error! status: ${response.status} ${response.statusText}`, + ) + } + const body = response.body + if (!body) throw new Error('Response has no body') + + const reader = body.getReader() + const decoder = new TextDecoder() + let buffer = '' + + try { + while (!signal.aborted) { + const { done, value } = await reader.read() + if (done) break + buffer += decoder.decode(value, { stream: true }) + const lines = buffer.split('\n') + buffer = lines.pop() ?? '' + for (const line of lines) { + const trimmed = line.trim() + if (!trimmed.startsWith('data:')) continue + const data = trimmed.slice(5).trim() + if (!data || data === '[DONE]') continue + try { + yield JSON.parse(data) as unknown + } catch { + // skip malformed + } + } + } + } finally { + reader.releaseLock() + } +} + +function BlogStudioServer() { + const abortRef = useRef(null) + + const [topic, setTopic] = useState('') + const [hasRun, setHasRun] = useState(false) + const [isRunning, setIsRunning] = useState(false) + const [error, setError] = useState(null) + + const [writingStep, setWritingStep] = useState('pending') + const [heroStep, setHeroStep] = useState('pending') + const [narrationStep, setNarrationStep] = useState('pending') + + const [post, setPost] = useState(null) + const [hero, setHero] = useState(null) + const [audio, setAudio] = useState(null) + const [draftChars, setDraftChars] = useState(0) + + const [heroBusy, setHeroBusy] = useState(false) + const [narrationBusy, setNarrationBusy] = useState(false) + + const imageUrl = imageUrlOf(hero) + const audioSrc = audioSrcOf(audio) + const showArticle = Boolean(post) + + function stop() { + abortRef.current?.abort() + abortRef.current = null + setIsRunning(false) + setHeroBusy(false) + setNarrationBusy(false) + } + + function resetPipeline() { + setError(null) + setPost(null) + setHero(null) + setAudio(null) + setDraftChars(0) + setWritingStep('pending') + setHeroStep('pending') + setNarrationStep('pending') + } + + function isBlogStudioStep(step: string): step is BlogStudioStep { + return step === 'drafting' || step === 'heroImage' || step === 'narration' + } + + function applyStepEvent(value: unknown) { + if (!isRecord(value) || typeof value.step !== 'string') return + if (!isBlogStudioStep(value.step)) return + const step = value.step + const status = value.status + const setStep = + step === 'drafting' + ? setWritingStep + : step === 'heroImage' + ? setHeroStep + : setNarrationStep + + if (status === 'started') { + setStep('active') + return + } + if (status === 'error') { + setStep('failed') + if (typeof value.error === 'string') setError(value.error) + return + } + if (status === 'done') { + setStep('done') + if (step === 'drafting' && isBlogPost(value.result)) { + setPost(value.result) + } else if (step === 'heroImage' && isImageResult(value.result)) { + setHero(value.result) + } else if (step === 'narration' && isTtsResult(value.result)) { + setAudio(value.result) + } + } + } + + async function run(e: FormEvent) { + e.preventDefault() + const nextTopic = topic.trim() + if (!nextTopic || isRunning) return + + stop() + resetPipeline() + setHasRun(true) + setIsRunning(true) + + const abort = new AbortController() + abortRef.current = abort + + try { + // Server function returns an SSE Response. One request: server chains + // draft → (hero ∥ narration) and streams pipeline:step / pipeline:result + // plus live structured-output deltas while drafting. + const response = await createBlogPostStreamFn({ + data: { topic: nextTopic }, + signal: abort.signal, + }) + + if (!(response instanceof Response)) { + throw new Error('Expected an SSE Response from createBlogPostStreamFn') + } + + for await (const chunk of readSseJson(response, abort.signal)) { + if (!isRecord(chunk) || typeof chunk.type !== 'string') continue + + if (chunk.type === EventType.RUN_ERROR) { + const message = + typeof chunk.message === 'string' + ? chunk.message + : 'Pipeline failed' + if (message !== 'Aborted') setError(message) + setWritingStep((s) => (s === 'done' ? s : 'failed')) + setHeroStep((s) => (s === 'active' ? 'failed' : s)) + setNarrationStep((s) => (s === 'active' ? 'failed' : s)) + break + } + + // Live draft JSON length while structured output streams. + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) { + const delta = chunk.delta + if (typeof delta === 'string') { + setDraftChars((n) => n + delta.length) + } + } + + if (chunk.type === EventType.CUSTOM) { + if (chunk.name === BLOG_STUDIO_SERVER_EVENTS.STEP) { + applyStepEvent(chunk.value) + } else if ( + chunk.name === BLOG_STUDIO_SERVER_EVENTS.RESULT && + isRecord(chunk.value) + ) { + if (isBlogPost(chunk.value.post)) setPost(chunk.value.post) + if (isImageResult(chunk.value.hero)) setHero(chunk.value.hero) + if (isTtsResult(chunk.value.audio)) setAudio(chunk.value.audio) + setWritingStep('done') + setHeroStep((s) => (s === 'failed' ? s : 'done')) + setNarrationStep((s) => (s === 'failed' ? s : 'done')) + } + } + } + } catch (err) { + if (abort.signal.aborted) return + setError(err instanceof Error ? err.message : 'Pipeline failed') + setWritingStep((s) => (s === 'done' ? s : 'failed')) + } finally { + if (abortRef.current === abort) { + abortRef.current = null + setIsRunning(false) + } + } + } + + async function regenerateHero() { + if (!post || isRunning || heroBusy) return + setHeroBusy(true) + setHeroStep('active') + setError(null) + try { + const result = await regenerateBlogHeroFn({ + data: { prompt: heroPromptFor(post) }, + }) + setHero(result) + setHeroStep('done') + } catch (err) { + setHeroStep('failed') + setError(err instanceof Error ? err.message : 'Hero regenerate failed') + } finally { + setHeroBusy(false) + } + } + + async function regenerateNarration() { + if (!post || isRunning || narrationBusy) return + setNarrationBusy(true) + setNarrationStep('active') + setError(null) + try { + const result = await regenerateBlogNarrationFn({ + data: { text: forNarration(post.body) }, + }) + setAudio(result) + setNarrationStep('done') + } catch (err) { + setNarrationStep('failed') + setError( + err instanceof Error ? err.message : 'Narration regenerate failed', + ) + } finally { + setNarrationBusy(false) + } + } + + return ( +
+ + +
+ {showArticle && post ? ( +
+
+ {imageUrl && !heroBusy && heroStep !== 'active' ? ( + {post.title} + ) : ( +
+ {heroBusy || heroStep === 'active' ? ( +
+ + Illustrating… +
+ ) : heroStep === 'failed' ? ( +
+ + + Couldn't generate a hero image + +
+ ) : ( + + )} +
+ )} +
+ +
+

+ {post.title} +

+ {post.subtitle && ( +

{post.subtitle}

+ )} + +
+
+ +
+ Written & narrated by TanStack AI + {narrationBusy || narrationStep === 'active' ? ( + + Recording + voice-over… + + ) : audioSrc ? ( +
+ +
+ ) : narrationStep === 'failed' ? ( + + + Voice-over unavailable + + ) : null} +
+ +
+ + {post.body} + +
+
+
+ ) : isRunning ? ( +
+
+ +

+ {writingStep === 'active' + ? draftChars > 0 + ? `Drafting… ${draftChars} characters so far.` + : 'Drafting the article…' + : 'Illustrating and recording voice-over…'} +

+
+
+ ) : ( +
+
+ +

+ Your finished post streams in step-by-step from a single server + function. +

+
+
+ )} +
+
+ ) +} + +function StepRow({ + label, + icon, + state, + detail, +}: { + label: string + icon: ReactNode + state: StepState + detail?: string +}) { + const cls = + state === 'active' + ? 'text-sky-700 font-medium' + : state === 'done' + ? 'text-stone-500' + : state === 'failed' + ? 'text-amber-600' + : 'text-stone-300' + return ( + + {state === 'active' ? ( + + ) : state === 'done' ? ( + + ) : state === 'failed' ? ( + + ) : ( + icon + )} + {label} + {detail && ( + + {detail} + + )} + + ) +} diff --git a/examples/ts-react-chat/src/routes/blog-studio.tsx b/examples/ts-react-chat/src/routes/blog-studio.tsx index 2519c7047..1669ecc6d 100644 --- a/examples/ts-react-chat/src/routes/blog-studio.tsx +++ b/examples/ts-react-chat/src/routes/blog-studio.tsx @@ -1,5 +1,4 @@ -import type { FormEvent, ReactNode } from 'react' -import { createFileRoute } from '@tanstack/react-router' +import { Link, createFileRoute } from '@tanstack/react-router' import { chat, generateImage, generateSpeech } from '@tanstack/ai' import { defineAssistant } from '@tanstack/ai/assistant' import { fetchServerSentEvents } from '@tanstack/ai-react' @@ -18,7 +17,8 @@ import { Volume2, Wand2, } from 'lucide-react' -import { BlogPostSchema } from './api.blog-studio' +import { BlogPostSchema, forNarration, heroPromptFor } from '../lib/blog-studio' +import type { FormEvent, ReactNode } from 'react' // Client-side assistant definition. `defineAssistant` is INERT in the browser // — these callbacks never run; `useAssistant` only reads the declared @@ -65,29 +65,6 @@ function statusToStep( : 'pending' } -// Prepare the post body for narration: strip Markdown so TTS doesn't read the -// syntax aloud, and cap the length at a sentence boundary (OpenAI TTS rejects -// input over 4096 characters, and long posts easily exceed it). -function forNarration(markdown: string, max = 4000): string { - const plain = markdown - .replace(/^#{1,6}\s+/gm, '') // headings - .replace(/^\s*[-*+]\s+/gm, '') // list bullets - .replace(/^\s*>\s?/gm, '') // blockquotes - .replace(/\*\*(.*?)\*\*/g, '$1') // bold - .replace(/\*(.*?)\*/g, '$1') // italic - .replace(/`([^`]+)`/g, '$1') // inline code - .replace(/\[([^\]]+)\]\([^)]*\)/g, '$1') // links → link text - .replace(/\n{3,}/g, '\n\n') - .trim() - if (plain.length <= max) return plain - const clipped = plain.slice(0, max) - const boundary = Math.max( - clipped.lastIndexOf('. '), - clipped.lastIndexOf('\n'), - ) - return (boundary > max / 2 ? clipped.slice(0, boundary + 1) : clipped).trim() -} - function BlogStudio() { const assistant = useAssistant(blogAssistant, { connection: fetchServerSentEvents('/api/blog-studio'), @@ -152,12 +129,7 @@ function BlogStudio() { // 2. Illustrate and narrate in parallel — fire and forget; the surfaces' // reactive result/status/error drive the UI. `generate()` never rejects. - void image.generate({ - prompt: - `A striking editorial hero image for a blog post titled ` + - `"${draftedPost.title}". ${draftedPost.subtitle}. Modern, clean, ` + - `cinematic, high quality, no text.`, - }) + void image.generate({ prompt: heroPromptFor(draftedPost) }) void speech.generate({ text: forNarration(draftedPost.body) }) } @@ -174,11 +146,21 @@ function BlogStudio() {

Turn a topic into a finished post

-

+

One assistant writes the article, then illustrates it and records a voice-over in parallel — chained from a single prompt over one endpoint.

+

+ Prefer the server to own the whole pipeline? See the{' '} + + server version + + . +

diff --git a/packages/ai-client/src/chain-client.ts b/packages/ai-client/src/chain-client.ts new file mode 100644 index 000000000..39de8a70f --- /dev/null +++ b/packages/ai-client/src/chain-client.ts @@ -0,0 +1,378 @@ +import { CHAIN_EVENTS } from '@tanstack/ai' +import { GENERATION_EVENTS } from './generation-types' +import { parseSSEResponse } from './sse-parser' +import { chainStepKey } from './chain-types' +import type { StreamChunk } from '@tanstack/ai/client' +import type { ChainStepEventValue } from '@tanstack/ai' +import type { + ConnectConnectionAdapter, + RunAgentInputContext, +} from './connection-adapters' +import type { + GenerationClientState, + GenerationFetcher, +} from './generation-types' +import type { + ChainClientOptions, + ChainStepState, + ChainSteps, +} from './chain-types' + +interface ChainCallbacks { + onResult?: ((result: TResult) => TOutput | null | void) | undefined + onError?: ((error: Error) => void) | undefined + onChunk?: ((chunk: StreamChunk) => void) | undefined + onStep?: ((step: ChainStepState, steps: ChainSteps) => void) | undefined + onResultChange?: ((result: TOutput | null) => void) | undefined + onLoadingChange?: ((isLoading: boolean) => void) | undefined + onErrorChange?: ((error: Error | undefined) => void) | undefined + onStatusChange?: ((status: GenerationClientState) => void) | undefined + onStepsChange?: ((steps: ChainSteps) => void) | undefined +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null +} + +function isChainStepEvent(value: unknown): value is ChainStepEventValue { + return ( + isRecord(value) && + typeof value.step === 'string' && + (value.status === 'started' || + value.status === 'done' || + value.status === 'error') + ) +} + +/** + * Client for a server-side `chain()` activity streamed as one AG-UI run. + * + * Understands the same outer lifecycle as {@link GenerationClient} + * (`generation:result`, RUN_*), plus demuxes `chain:step` CUSTOM events into + * a reactive {@link ChainSteps} map for progressive UIs. + * + * @template TInput - Chain input (e.g. `{ topic: string }`) + * @template TResult - Final `generation:result` payload + * @template TOutput - Stored result after optional transform + * + * @example + * ```ts + * const client = new ChainClient<{ topic: string }, BlogResult>({ + * fetcher: (input, { signal }) => + * createBlogPostChainFn({ data: input, signal }), + * onStepsChange: setSteps, + * onResultChange: setResult, + * }) + * await client.run({ topic: 'urban foxes' }) + * client.getSteps()['draft']?.status // 'done' + * ``` + */ +export class ChainClient< + TInput extends Record, + TResult, + TOutput = TResult, +> { + private readonly connection: ConnectConnectionAdapter | undefined + private readonly fetcher: GenerationFetcher | undefined + private readonly uniqueId: string + private readonly threadId: string + private body: Record + private result: TOutput | null = null + private steps: ChainSteps = {} + private isLoading = false + private error: Error | undefined = undefined + private status: GenerationClientState = 'idle' + private abortController: AbortController | null = null + private readonly callbacksRef: ChainCallbacks + + constructor( + options: ChainClientOptions & + ( + | { connection: ConnectConnectionAdapter; fetcher?: never } + | { + fetcher: GenerationFetcher + connection?: never + } + ), + ) { + this.uniqueId = options.id ?? this.generateUniqueId('chain') + this.threadId = this.uniqueId + this.connection = options.connection + this.fetcher = options.fetcher + this.body = options.body ?? {} + + this.callbacksRef = { + onResult: options.onResult, + onError: options.onError, + onChunk: options.onChunk, + onStep: options.onStep, + onResultChange: options.onResultChange, + onLoadingChange: options.onLoadingChange, + onErrorChange: options.onErrorChange, + onStatusChange: options.onStatusChange, + onStepsChange: options.onStepsChange, + } + } + + /** + * Run the chain. Only one run at a time; concurrent calls are no-ops. + * Resolves with the stored result (after optional `onResult` transform), + * or `null` if aborted / no `generation:result` arrived. + */ + async run(input: TInput): Promise { + if (this.isLoading) return this.result + + this.setSteps({}) + this.setIsLoading(true) + this.setStatus('generating') + this.setError(undefined) + + const abortController = new AbortController() + this.abortController = abortController + const { signal } = abortController + + try { + if (this.fetcher) { + const result = await this.fetcher(input, { signal }) + if (signal.aborted) return null + if (result instanceof Response) { + await this.processStream(parseSSEResponse(result, signal)) + } else { + this.setResult(result) + this.setStatus('success') + } + } else if (this.connection) { + const mergedData = { ...this.body, ...input } + const stream = this.connection.connect( + [], + mergedData, + signal, + this.createRunContext(this.generateUniqueId('run')), + ) + await this.processStream(stream) + } else { + throw new Error( + 'ChainClient requires either a connection or fetcher option', + ) + } + return signal.aborted ? null : this.result + } catch (err: unknown) { + if (signal.aborted) return null + const error = err instanceof Error ? err : new Error(String(err)) + this.setError(error) + this.setStatus('error') + this.callbacksRef.onError?.(error) + return null + } finally { + this.abortController = null + this.setIsLoading(false) + } + } + + private async processStream(source: AsyncIterable): Promise { + for await (const chunk of source) { + if (this.abortController?.signal.aborted) break + + this.callbacksRef.onChunk?.(chunk) + + // eslint-disable-next-line @typescript-eslint/switch-exhaustiveness-check -- only handle lifecycle + chain events + switch (chunk.type) { + case 'CUSTOM': { + if (chunk.name === CHAIN_EVENTS.STEP) { + this.applyStepEvent(chunk.value) + } else if (chunk.name === GENERATION_EVENTS.RESULT) { + this.setResult(chunk.value as TResult) + } + break + } + case 'RUN_FINISHED': { + this.setStatus('success') + break + } + case 'RUN_ERROR': { + const msg = + (chunk.message as string | undefined) || + chunk.error?.message || + 'An error occurred' + // Treat client abort as cancellation, not an error surface. + if (msg === 'Aborted') { + return + } + throw new Error(msg) + } + default: + break + } + } + + // Stream ended without RUN_FINISHED but with a result — still success. + if (this.status === 'generating' && this.result !== null) { + this.setStatus('success') + } + } + + private applyStepEvent(value: unknown): void { + if (!isChainStepEvent(value)) return + + const key = chainStepKey(value.step, value.branch) + const prev = this.steps[key] + let next: ChainStepState + + if (value.status === 'started') { + next = { + step: value.step, + ...(value.branch !== undefined && { branch: value.branch }), + index: value.index, + status: 'active', + } + } else if (value.status === 'error') { + next = { + step: value.step, + ...(value.branch !== undefined && { branch: value.branch }), + index: value.index, + status: 'error', + error: value.error, + ...(prev?.result !== undefined && { result: prev.result }), + } + } else { + next = { + step: value.step, + ...(value.branch !== undefined && { branch: value.branch }), + index: value.index, + status: 'done', + result: value.result, + } + } + + this.steps = { ...this.steps, [key]: next } + this.callbacksRef.onStepsChange?.(this.steps) + this.callbacksRef.onStep?.(next, this.steps) + } + + stop(): void { + if (this.abortController) { + this.abortController.abort() + this.abortController = null + } + this.setIsLoading(false) + if (this.status === 'generating') { + this.setStatus('idle') + } + } + + reset(): void { + this.stop() + this.setResult(null) + this.setSteps({}) + this.setError(undefined) + this.setStatus('idle') + } + + updateOptions( + options: Partial< + Pick< + ChainClientOptions, + 'body' | 'onResult' | 'onError' | 'onChunk' | 'onStep' + > + >, + ): void { + if (options.body !== undefined) { + this.body = options.body ?? {} + } + if (options.onResult !== undefined) { + this.callbacksRef.onResult = options.onResult + } + if (options.onError !== undefined) { + this.callbacksRef.onError = options.onError + } + if (options.onChunk !== undefined) { + this.callbacksRef.onChunk = options.onChunk + } + if (options.onStep !== undefined) { + this.callbacksRef.onStep = options.onStep + } + } + + dispose(): void { + this.stop() + } + + getResult(): TOutput | null { + return this.result + } + + getSteps(): ChainSteps { + return this.steps + } + + getStep(step: string, branch?: string): ChainStepState | undefined { + return this.steps[chainStepKey(step, branch)] + } + + getIsLoading(): boolean { + return this.isLoading + } + + getError(): Error | undefined { + return this.error + } + + getStatus(): GenerationClientState { + return this.status + } + + private setResult(rawResult: TResult | null): void { + if (rawResult === null) { + this.result = null + this.callbacksRef.onResultChange?.(null) + return + } + + if (this.callbacksRef.onResult) { + const transformed = this.callbacksRef.onResult(rawResult) + if (transformed === null) { + return + } + if (transformed !== undefined) { + this.result = transformed + this.callbacksRef.onResultChange?.(this.result) + return + } + } + + // eslint-disable-next-line no-restricted-syntax -- TOutput defaults to TResult when no onResult is supplied + this.result = rawResult as unknown as TOutput + this.callbacksRef.onResultChange?.(this.result) + } + + private setSteps(steps: ChainSteps): void { + this.steps = steps + this.callbacksRef.onStepsChange?.(steps) + } + + private setIsLoading(isLoading: boolean): void { + this.isLoading = isLoading + this.callbacksRef.onLoadingChange?.(isLoading) + } + + private setError(error: Error | undefined): void { + this.error = error + this.callbacksRef.onErrorChange?.(error) + } + + private setStatus(status: GenerationClientState): void { + this.status = status + this.callbacksRef.onStatusChange?.(status) + } + + private generateUniqueId(prefix: string): string { + return `${prefix}-${Date.now()}-${Math.random().toString(36).substring(7)}` + } + + private createRunContext(runId: string): RunAgentInputContext { + return { + threadId: this.threadId, + runId, + } + } +} diff --git a/packages/ai-client/src/chain-types.ts b/packages/ai-client/src/chain-types.ts new file mode 100644 index 000000000..5367eb66f --- /dev/null +++ b/packages/ai-client/src/chain-types.ts @@ -0,0 +1,87 @@ +import type { StreamChunk } from '@tanstack/ai/client' +import type { ConnectConnectionAdapter } from './connection-adapters' +import type { + GenerationClientState, + GenerationFetcher, + InferGenerationOutputFromReturn, +} from './generation-types' + +export type { InferGenerationOutputFromReturn } + +/** Status of a single chain step (or parallel branch). */ +export type ChainStepStatus = 'pending' | 'active' | 'done' | 'error' + +/** + * Reactive state for one step (or one branch of a parallel fan-out). + * Keys in the steps map use {@link chainStepKey}. + */ +export interface ChainStepState { + step: string + branch?: string + index?: number + status: ChainStepStatus + /** Present when `status === 'done'`. */ + result?: unknown + /** Present when `status === 'error'`. */ + error?: string +} + +/** + * Map of step/branch key → state. + * Values are optional so `steps.draft` / `steps['media/hero']` type-check as + * possibly missing before that step has started. + */ +export type ChainSteps = { + readonly [key: string]: ChainStepState | undefined +} + +/** + * Build the stable key used in {@link ChainSteps}. + * Sequential steps: `"draft"`. Parallel branches: `"media/hero"`. + */ +export function chainStepKey(step: string, branch?: string): string { + return branch ? `${step}/${branch}` : step +} + +/** + * Options for {@link ChainClient}. + * + * @template TInput - Chain input + * @template TResult - Final `generation:result` payload type + * @template TOutput - Stored result after optional `onResult` transform + */ +// eslint-disable-next-line @typescript-eslint/naming-convention -- _TInput is part of the public positional generic API +export interface ChainClientOptions<_TInput, TResult, TOutput = TResult> { + /** Unique id for this client instance */ + id?: string + /** Extra body fields merged into connect-adapter requests */ + body?: Record + + /** + * Transform the final `generation:result` payload before storage. + * - Return a value → store it + * - Return `null` → keep previous result + * - Return `void` → store the raw result + */ + onResult?: (result: TResult) => TOutput | null | void + onError?: (error: Error) => void + onChunk?: (chunk: StreamChunk) => void + /** Fired for every `chain:step` event (after the steps map is updated). */ + onStep?: (step: ChainStepState, steps: ChainSteps) => void + + // Framework state callbacks (set by hooks) + /** @internal */ + onResultChange?: (result: TOutput | null) => void + /** @internal */ + onLoadingChange?: (isLoading: boolean) => void + /** @internal */ + onErrorChange?: (error: Error | undefined) => void + /** @internal */ + onStatusChange?: (status: GenerationClientState) => void + /** @internal */ + onStepsChange?: (steps: ChainSteps) => void +} + +export type ChainTransport = + | { connection: ConnectConnectionAdapter; fetcher?: never } + | { fetcher: GenerationFetcher; connection?: never } diff --git a/packages/ai-client/src/index.ts b/packages/ai-client/src/index.ts index bd9270ff6..77d22ce4b 100644 --- a/packages/ai-client/src/index.ts +++ b/packages/ai-client/src/index.ts @@ -10,7 +10,16 @@ export { createMcpAppBridge } from './mcp-app-bridge' export type { McpAppBridge, CreateMcpAppBridgeOptions } from './mcp-app-bridge' export { RealtimeClient } from './realtime-client' export { GenerationClient } from './generation-client' +export { ChainClient } from './chain-client' export { VideoGenerationClient } from './video-generation-client' +export { chainStepKey } from './chain-types' +export type { + ChainClientOptions, + ChainStepState, + ChainStepStatus, + ChainSteps, + ChainTransport, +} from './chain-types' export type { // Core message types (re-exported from @tanstack/ai via types.ts) UIMessage, diff --git a/packages/ai-client/tests/chain-client.test.ts b/packages/ai-client/tests/chain-client.test.ts new file mode 100644 index 000000000..f8bbd38d2 --- /dev/null +++ b/packages/ai-client/tests/chain-client.test.ts @@ -0,0 +1,238 @@ +import { describe, expect, it, vi } from 'vitest' +import { CHAIN_EVENTS, EventType } from '@tanstack/ai' +import { ChainClient, chainStepKey } from '../src' +import type { StreamChunk } from '@tanstack/ai/client' +import type { ConnectConnectionAdapter } from '../src/connection-adapters' + +function createMockConnection( + chunks: Array, +): ConnectConnectionAdapter { + return { + async *connect() { + for (const chunk of chunks) { + yield chunk + } + }, + } +} + +function chainChunks(result: { + post: { title: string } + hero: { id: string } +}): Array { + return [ + { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: CHAIN_EVENTS.STEP, + value: { step: 'draft', index: 0, status: 'started' }, + timestamp: Date.now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm1', + delta: '{"title":', + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: CHAIN_EVENTS.STEP, + value: { + step: 'draft', + index: 0, + status: 'done', + result: result.post, + }, + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: CHAIN_EVENTS.STEP, + value: { + step: 'media', + index: 1, + branch: 'hero', + status: 'started', + }, + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: CHAIN_EVENTS.STEP, + value: { + step: 'media', + index: 1, + branch: 'hero', + status: 'done', + result: result.hero, + }, + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: 'generation:result', + value: result, + timestamp: Date.now(), + }, + { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + }, + ] +} + +describe('chainStepKey', () => { + it('joins step and branch', () => { + expect(chainStepKey('draft')).toBe('draft') + expect(chainStepKey('media', 'hero')).toBe('media/hero') + }) +}) + +describe('ChainClient', () => { + it('demuxes chain:step events and stores generation:result', async () => { + const final = { + post: { title: 'Foxes' }, + hero: { id: 'img-1' }, + } + const onStepsChange = vi.fn() + const onStep = vi.fn() + const onChunk = vi.fn() + const onResultChange = vi.fn() + + const client = new ChainClient({ + connection: createMockConnection(chainChunks(final)), + onStepsChange, + onStep, + onChunk, + onResultChange, + }) + + const out = await client.run({ topic: 'foxes' }) + + expect(out).toEqual(final) + expect(client.getResult()).toEqual(final) + expect(client.getStatus()).toBe('success') + expect(client.getIsLoading()).toBe(false) + + const draft = client.getStep('draft') + expect(draft?.status).toBe('done') + expect(draft?.result).toEqual(final.post) + + const hero = client.getStep('media', 'hero') + expect(hero?.status).toBe('done') + expect(hero?.result).toEqual(final.hero) + expect(client.getSteps()['media/hero']?.status).toBe('done') + + expect(onStep).toHaveBeenCalled() + expect(onStepsChange).toHaveBeenCalled() + expect(onResultChange).toHaveBeenCalledWith(final) + expect(onChunk).toHaveBeenCalledWith( + expect.objectContaining({ type: EventType.TEXT_MESSAGE_CONTENT }), + ) + }) + + it('records step errors without losing earlier done steps', async () => { + const client = new ChainClient({ + connection: createMockConnection([ + { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: CHAIN_EVENTS.STEP, + value: { + step: 'draft', + index: 0, + status: 'done', + result: { title: 'ok' }, + }, + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: CHAIN_EVENTS.STEP, + value: { + step: 'media', + index: 1, + branch: 'hero', + status: 'error', + error: 'image failed', + }, + timestamp: Date.now(), + }, + { + type: EventType.RUN_ERROR, + message: 'image failed', + timestamp: Date.now(), + } as StreamChunk, + ]), + }) + + await client.run({ topic: 'x' }) + + expect(client.getStatus()).toBe('error') + expect(client.getError()?.message).toBe('image failed') + expect(client.getStep('draft')?.status).toBe('done') + expect(client.getStep('media', 'hero')?.status).toBe('error') + expect(client.getStep('media', 'hero')?.error).toBe('image failed') + }) + + it('supports fetcher that returns a plain result', async () => { + const final = { post: { title: 'direct' } } + const client = new ChainClient({ + fetcher: async () => final, + }) + + const out = await client.run({ topic: 't' }) + expect(out).toEqual(final) + expect(client.getStatus()).toBe('success') + }) + + it('resets steps and result', async () => { + const final = { post: { title: 'a' }, hero: { id: '1' } } + const client = new ChainClient({ + connection: createMockConnection(chainChunks(final)), + }) + + await client.run({ topic: 't' }) + expect(client.getStep('draft')).toBeDefined() + + client.reset() + expect(client.getResult()).toBeNull() + expect(client.getSteps()).toEqual({}) + expect(client.getStatus()).toBe('idle') + }) + + it('does not allow concurrent runs', async () => { + let release!: () => void + const gate = new Promise((resolve) => { + release = resolve + }) + let calls = 0 + + const client = new ChainClient({ + fetcher: async () => { + calls++ + await gate + return { ok: true } + }, + }) + + const p1 = client.run({ topic: 'a' }) + const p2 = client.run({ topic: 'b' }) + release() + await p1 + await p2 + expect(calls).toBe(1) + }) +}) diff --git a/packages/ai-react/src/index.ts b/packages/ai-react/src/index.ts index f8842333a..41bffefbc 100644 --- a/packages/ai-react/src/index.ts +++ b/packages/ai-react/src/index.ts @@ -21,6 +21,9 @@ export type { UseGenerationReturn, } from './use-generation' +export { useChain } from './use-chain' +export type { UseChainOptions, UseChainReturn } from './use-chain' + export { useGenerateImage } from './use-generate-image' export type { UseGenerateImageOptions, @@ -83,6 +86,10 @@ export { type XhrConnectionOptions, type InferChatMessages, type GenerationClientState, + type ChainStepState, + type ChainSteps, + type ChainStepStatus, + chainStepKey, type ImageGenerateInput, type AudioGenerateInput, type SpeechGenerateInput, diff --git a/packages/ai-react/src/use-chain.ts b/packages/ai-react/src/use-chain.ts new file mode 100644 index 000000000..1a1e946a4 --- /dev/null +++ b/packages/ai-react/src/use-chain.ts @@ -0,0 +1,179 @@ +import { ChainClient } from '@tanstack/ai-client' +import { useCallback, useEffect, useId, useMemo, useRef, useState } from 'react' +import type { StreamChunk } from '@tanstack/ai' +import type { + ChainStepState, + ChainSteps, + ConnectConnectionAdapter, + GenerationClientState, + GenerationFetcher, + InferGenerationOutputFromReturn, +} from '@tanstack/ai-client' + +/** + * Options for {@link useChain}. + * + * Accepts either a `connection` (streaming transport) or a `fetcher` that + * returns a `Response` with an SSE body (typical for `createServerFn` + + * `toServerSentEventsResponse(chain.stream(...))`). + */ +export interface UseChainOptions { + connection?: ConnectConnectionAdapter + fetcher?: GenerationFetcher + id?: string + body?: Record + onResult?: (result: TResult) => TOutput | null | void + onError?: (error: Error) => void + onChunk?: (chunk: StreamChunk) => void + /** Fired after each `chain:step` updates the steps map. */ + onStep?: (step: ChainStepState, steps: ChainSteps) => void +} + +export interface UseChainReturn { + /** Start a chain run. Resolves to the final result (or `null` if aborted/failed). */ + run: (input: TInput) => Promise + /** Final `generation:result` payload (after optional transform). */ + result: TOutput | null + /** + * Live step map. Keys are step names (`"draft"`) or parallel branch paths + * (`"media/hero"`). See `chainStepKey` from `@tanstack/ai-client`. + */ + steps: ChainSteps + isLoading: boolean + error: Error | undefined + status: GenerationClientState + stop: () => void + reset: () => void + /** Lookup a step (or parallel branch) by name. */ + getStep: (step: string, branch?: string) => ChainStepState | undefined +} + +/** + * React hook for a server-side `chain()` activity. + * + * Demuxes one SSE run into: + * - `result` — terminal `generation:result` + * - `steps` — progressive `chain:step` state for UI (drafting / hero / …) + * - live chunks via `onChunk` (e.g. structured-output text deltas) + * + * @example + * ```tsx + * const chain = useChain<{ topic: string }, BlogStudioChainResult>({ + * fetcher: (input, { signal }) => + * createBlogPostChainFn({ data: input, signal }), + * }) + * + * await chain.run({ topic: 'urban foxes' }) + * chain.steps['draft']?.status // 'done' + * chain.steps['media/hero']?.result + * chain.result // { post, hero, narration } + * ``` + */ +export function useChain< + TInput extends Record, + TResult, + TTransformed = void, +>( + options: Omit, 'onResult'> & { + onResult?: (result: TResult) => TTransformed + }, +): UseChainReturn> { + type TOutput = InferGenerationOutputFromReturn + const hookId = useId() + const clientId = options.id || hookId + + const [result, setResult] = useState(null) + const [steps, setSteps] = useState({}) + const [isLoading, setIsLoading] = useState(false) + const [error, setError] = useState(undefined) + const [status, setStatus] = useState('idle') + + const optionsRef = useRef(options) + optionsRef.current = options + + const client = useMemo(() => { + const opts = optionsRef.current + + const clientOptions = { + id: clientId, + body: opts.body, + onResult: ((r: TResult) => optionsRef.current.onResult?.(r)) as ( + result: TResult, + ) => TOutput | null | void, + onError: (e: Error) => { + optionsRef.current.onError?.(e) + }, + onChunk: (c: StreamChunk) => { + optionsRef.current.onChunk?.(c) + }, + onStep: (step: ChainStepState, all: ChainSteps) => { + optionsRef.current.onStep?.(step, all) + }, + onResultChange: setResult, + onStepsChange: setSteps, + onLoadingChange: setIsLoading, + onErrorChange: setError, + onStatusChange: setStatus, + } + + if (opts.connection) { + return new ChainClient({ + ...clientOptions, + connection: opts.connection, + }) + } + + if (opts.fetcher) { + return new ChainClient({ + ...clientOptions, + fetcher: opts.fetcher, + }) + } + + throw new Error('useChain requires either a connection or fetcher option') + }, [clientId]) + + useEffect(() => { + client.updateOptions({ + ...(options.body !== undefined && { body: options.body }), + }) + }, [client, options.body]) + + useEffect(() => { + return () => { + client.dispose() + } + }, [client]) + + const run = useCallback( + async (input: TInput) => { + return client.run(input) + }, + [client], + ) + + const stop = useCallback(() => { + client.stop() + }, [client]) + + const reset = useCallback(() => { + client.reset() + }, [client]) + + const getStep = useCallback( + (step: string, branch?: string) => client.getStep(step, branch), + [client], + ) + + return { + run, + result, + steps, + isLoading, + error, + status, + stop, + reset, + getStep, + } +} diff --git a/packages/ai-react/tests/use-chain.test.tsx b/packages/ai-react/tests/use-chain.test.tsx new file mode 100644 index 000000000..dbcb09309 --- /dev/null +++ b/packages/ai-react/tests/use-chain.test.tsx @@ -0,0 +1,144 @@ +import { act, renderHook, waitFor } from '@testing-library/react' +import { describe, expect, it, vi } from 'vitest' +import { CHAIN_EVENTS, EventType } from '@tanstack/ai' +import { useChain } from '../src/use-chain' +import { createMockConnectionAdapter } from './test-utils' +import type { StreamChunk } from '@tanstack/ai' + +function createChainChunks(result: { + post: { title: string } + hero: { id: string } +}): Array { + return [ + { + type: EventType.RUN_STARTED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: CHAIN_EVENTS.STEP, + value: { step: 'draft', index: 0, status: 'started' }, + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: CHAIN_EVENTS.STEP, + value: { + step: 'draft', + index: 0, + status: 'done', + result: result.post, + }, + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: CHAIN_EVENTS.STEP, + value: { + step: 'media', + index: 1, + branch: 'hero', + status: 'done', + result: result.hero, + }, + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: 'generation:result', + value: result, + timestamp: Date.now(), + }, + { + type: EventType.RUN_FINISHED, + runId: 'run-1', + threadId: 'thread-1', + timestamp: Date.now(), + }, + ] +} + +describe('useChain', () => { + it('initializes idle with empty steps', () => { + const adapter = createMockConnectionAdapter({ chunks: [] }) + const { result } = renderHook(() => useChain({ connection: adapter })) + + expect(result.current.result).toBeNull() + expect(result.current.steps).toEqual({}) + expect(result.current.isLoading).toBe(false) + expect(result.current.status).toBe('idle') + }) + + it('tracks steps and final result from a stream', async () => { + const final = { + post: { title: 'Foxes' }, + hero: { id: 'img-1' }, + } + const adapter = createMockConnectionAdapter({ + chunks: createChainChunks(final), + }) + const onChunk = vi.fn() + + const { result } = renderHook(() => + useChain<{ topic: string }, typeof final>({ + connection: adapter, + onChunk, + }), + ) + + await act(async () => { + const out = await result.current.run({ topic: 'foxes' }) + expect(out).toEqual(final) + }) + + await waitFor(() => { + expect(result.current.status).toBe('success') + }) + + expect(result.current.result).toEqual(final) + expect(result.current.steps.draft?.status).toBe('done') + expect(result.current.steps.draft?.result).toEqual(final.post) + expect(result.current.steps['media/hero']?.result).toEqual(final.hero) + expect(result.current.getStep('media', 'hero')?.status).toBe('done') + }) + + it('supports fetcher mode', async () => { + const final = { post: { title: 'direct' } } + const { result } = renderHook(() => + useChain({ + fetcher: async () => final, + }), + ) + + await act(async () => { + await result.current.run({ topic: 't' }) + }) + + await waitFor(() => { + expect(result.current.result).toEqual(final) + }) + }) + + it('reset clears result and steps', async () => { + const final = { post: { title: 'a' }, hero: { id: '1' } } + const adapter = createMockConnectionAdapter({ + chunks: createChainChunks(final), + }) + const { result } = renderHook(() => useChain({ connection: adapter })) + + await act(async () => { + await result.current.run({ topic: 't' }) + }) + await waitFor(() => expect(result.current.result).not.toBeNull()) + + act(() => { + result.current.reset() + }) + + expect(result.current.result).toBeNull() + expect(result.current.steps).toEqual({}) + expect(result.current.status).toBe('idle') + }) +}) From ff75e74ae36d7b63d7b97060d6271a1693ebd4de Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 04:50:59 +0000 Subject: [PATCH 11/13] ci: apply automated fixes --- packages/ai-client/src/chain-client.ts | 4 +++- packages/ai-react/src/use-chain.ts | 5 ++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/packages/ai-client/src/chain-client.ts b/packages/ai-client/src/chain-client.ts index 39de8a70f..da623d55c 100644 --- a/packages/ai-client/src/chain-client.ts +++ b/packages/ai-client/src/chain-client.ts @@ -169,7 +169,9 @@ export class ChainClient< } } - private async processStream(source: AsyncIterable): Promise { + private async processStream( + source: AsyncIterable, + ): Promise { for await (const chunk of source) { if (this.abortController?.signal.aborted) break diff --git a/packages/ai-react/src/use-chain.ts b/packages/ai-react/src/use-chain.ts index 1a1e946a4..9ca4bb05a 100644 --- a/packages/ai-react/src/use-chain.ts +++ b/packages/ai-react/src/use-chain.ts @@ -77,7 +77,10 @@ export function useChain< options: Omit, 'onResult'> & { onResult?: (result: TResult) => TTransformed }, -): UseChainReturn> { +): UseChainReturn< + TInput, + InferGenerationOutputFromReturn +> { type TOutput = InferGenerationOutputFromReturn const hookId = useId() const clientId = options.id || hookId From 6406b97464952fe8aac713de4508f866a98f825f Mon Sep 17 00:00:00 2001 From: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Date: Thu, 16 Jul 2026 14:56:45 +1000 Subject: [PATCH 12/13] feat(ai-client): stream structured-output partials on useChain steps Demux native structured-output.start + JSON text deltas into steps[name].partial via parsePartialJSON, and render live draft fields in the Blog Studio chain example. --- .changeset/use-chain-hook.md | 2 +- .../src/routes/blog-studio-chain.tsx | 65 +++++---- packages/ai-client/src/chain-client.ts | 124 +++++++++++++++--- packages/ai-client/src/chain-types.ts | 7 + packages/ai-client/tests/chain-client.test.ts | 107 ++++++++++++++- packages/ai-react/src/use-chain.ts | 7 +- 6 files changed, 265 insertions(+), 47 deletions(-) diff --git a/.changeset/use-chain-hook.md b/.changeset/use-chain-hook.md index b7ab7a698..69c7c868d 100644 --- a/.changeset/use-chain-hook.md +++ b/.changeset/use-chain-hook.md @@ -3,4 +3,4 @@ '@tanstack/ai-react': minor --- -Add `ChainClient` and React `useChain` for server-side `chain()` activities: demux `chain:step` progress into a reactive steps map, surface `generation:result`, and support stop/reset with connection or fetcher transports. +Add `ChainClient` and React `useChain` for server-side `chain()` activities: demux `chain:step` progress into a reactive steps map, stream native structured-output partials onto the active step (`steps[name].partial`), surface `generation:result`, and support stop/reset with connection or fetcher transports. diff --git a/examples/ts-react-chat/src/routes/blog-studio-chain.tsx b/examples/ts-react-chat/src/routes/blog-studio-chain.tsx index f45951f6f..8c284e9e2 100644 --- a/examples/ts-react-chat/src/routes/blog-studio-chain.tsx +++ b/examples/ts-react-chat/src/routes/blog-studio-chain.tsx @@ -14,7 +14,6 @@ import { Volume2, Wand2, } from 'lucide-react' -import { EventType } from '@tanstack/ai' import { useChain } from '@tanstack/ai-react' import { forNarration, heroPromptFor } from '../lib/blog-studio' import { createBlogPostChainFn } from '../lib/blog-studio-chain-server-fns' @@ -76,10 +75,15 @@ function toUiStep(status: ChainStepStatus | undefined): StepState { return 'pending' } +function stringField(value: unknown, key: string): string | undefined { + if (!isRecord(value)) return undefined + const field = value[key] + return typeof field === 'string' ? field : undefined +} + function BlogStudioChain() { const [topic, setTopic] = useState('') const [hasRun, setHasRun] = useState(false) - const [draftChars, setDraftChars] = useState(0) // Local overrides after "Regenerate hero" / "Re-narrate" (one-shot server fns). const [heroOverride, setHeroOverride] = @@ -89,16 +93,12 @@ function BlogStudioChain() { const [narrationBusy, setNarrationBusy] = useState(false) const [touchUpError, setTouchUpError] = useState(null) + // Structured draft partials stream natively: server emits + // structured-output.start + JSON TEXT_MESSAGE_CONTENT + .complete; useChain + // puts progressive parsePartialJSON onto steps.draft.partial. const chain = useChain<{ topic: string }, BlogStudioChainResult>({ fetcher: (input, options) => createBlogPostChainFn({ data: input, signal: options?.signal }), - onChunk: (chunk) => { - if (chunk.type !== EventType.TEXT_MESSAGE_CONTENT) return - const delta = chunk.delta - if (typeof delta === 'string') { - setDraftChars((n) => n + delta.length) - } - }, }) const draftStep = chain.getStep('draft') @@ -106,12 +106,18 @@ function BlogStudioChain() { const narrationStepMeta = chain.getStep('media', 'narration') const draftResult = draftStep ? draftStep.result : undefined + const draftPartial = draftStep ? draftStep.partial : undefined const post: BlogPost | null = isBlogPost(draftResult) ? draftResult : chain.result && isBlogPost(chain.result.post) ? chain.result.post : null + // Live article fields: validated post wins; otherwise progressive partial. + const title = post?.title ?? stringField(draftPartial, 'title') + const subtitle = post?.subtitle ?? stringField(draftPartial, 'subtitle') + const body = post?.body ?? stringField(draftPartial, 'body') ?? '' + const heroFromStep = heroStepMeta ? heroStepMeta.result : undefined const heroFromResult = chain.result ? chain.result.hero : undefined const hero: ImageGenerationResult | null = @@ -148,7 +154,7 @@ function BlogStudioChain() { const error = touchUpError ?? (chain.error ? chain.error.message : null) const imageUrl = imageUrlOf(hero) const audioSrc = audioSrcOf(audio) - const showArticle = Boolean(post) + const showArticle = Boolean(post) || Boolean(title) || Boolean(body) async function onSubmit(e: FormEvent) { e.preventDefault() @@ -156,7 +162,6 @@ function BlogStudioChain() { if (!nextTopic || isRunning) return setHasRun(true) - setDraftChars(0) setHeroOverride(null) setAudioOverride(null) setTouchUpError(null) @@ -219,9 +224,11 @@ function BlogStudioChain() { useChain:{' '} live{' '} chain:step{' '} - progress, draft deltas via{' '} - onChunk, - then{' '} + progress, native structured-output partials on{' '} + + steps.draft.partial + + , then{' '} generation:result @@ -301,8 +308,8 @@ function BlogStudioChain() { icon={} state={writingStep} detail={ - writingStep === 'active' && draftChars > 0 - ? `${draftChars} chars drafted` + writingStep === 'active' && body.length > 0 + ? `${body.length} chars drafted` : undefined } /> @@ -361,13 +368,13 @@ function BlogStudioChain() {
- {showArticle && post ? ( + {showArticle ? (
{imageUrl && !heroBusy && heroStep !== 'active' ? ( {post.title} ) : ( @@ -393,10 +400,10 @@ function BlogStudioChain() {

- {post.title} + {title ?? (writingStep === 'active' ? 'Drafting…' : 'Untitled')}

- {post.subtitle && ( -

{post.subtitle}

+ {subtitle && ( +

{subtitle}

)}
@@ -423,9 +430,13 @@ function BlogStudioChain() {
- - {post.body} - + {body ? ( + + {body} + + ) : writingStep === 'active' ? ( +

Streaming structured draft…

+ ) : null}
@@ -435,9 +446,7 @@ function BlogStudioChain() {

{writingStep === 'active' - ? draftChars > 0 - ? `Drafting… ${draftChars} characters so far.` - : 'Drafting the article…' + ? 'Drafting the article…' : 'Illustrating and recording voice-over…'}

@@ -447,7 +456,7 @@ function BlogStudioChain() {

- Your finished post streams in step-by-step via{' '} + Your post streams live via structured-output partials on{' '} useChain diff --git a/packages/ai-client/src/chain-client.ts b/packages/ai-client/src/chain-client.ts index da623d55c..f7c9b26cf 100644 --- a/packages/ai-client/src/chain-client.ts +++ b/packages/ai-client/src/chain-client.ts @@ -1,4 +1,5 @@ import { CHAIN_EVENTS } from '@tanstack/ai' +import { parsePartialJSON } from '@tanstack/ai/client' import { GENERATION_EVENTS } from './generation-types' import { parseSSEResponse } from './sse-parser' import { chainStepKey } from './chain-types' @@ -48,24 +49,17 @@ function isChainStepEvent(value: unknown): value is ChainStepEventValue { * Client for a server-side `chain()` activity streamed as one AG-UI run. * * Understands the same outer lifecycle as {@link GenerationClient} - * (`generation:result`, RUN_*), plus demuxes `chain:step` CUSTOM events into - * a reactive {@link ChainSteps} map for progressive UIs. + * (`generation:result`, RUN_*), demuxes `chain:step` CUSTOM events into a + * reactive {@link ChainSteps} map, and mirrors chat's native structured-output + * protocol onto the active step: + * + * 1. `structured-output.start` — begin accumulating JSON for the current step + * 2. `TEXT_MESSAGE_CONTENT` deltas — progressive `parsePartialJSON` → `step.partial` + * 3. `structured-output.complete` / `chain:step` done — terminal object on `step.result` * * @template TInput - Chain input (e.g. `{ topic: string }`) * @template TResult - Final `generation:result` payload * @template TOutput - Stored result after optional transform - * - * @example - * ```ts - * const client = new ChainClient<{ topic: string }, BlogResult>({ - * fetcher: (input, { signal }) => - * createBlogPostChainFn({ data: input, signal }), - * onStepsChange: setSteps, - * onResultChange: setResult, - * }) - * await client.run({ topic: 'urban foxes' }) - * client.getSteps()['draft']?.status // 'done' - * ``` */ export class ChainClient< TInput extends Record, @@ -85,6 +79,13 @@ export class ChainClient< private abortController: AbortController | null = null private readonly callbacksRef: ChainCallbacks + /** Most recently `started` step key — structured deltas attach here. */ + private lastActiveStepKey: string | null = null + /** Step currently receiving structured-output JSON (after `.start`). */ + private structuredTargetKey: string | null = null + /** Accumulated raw JSON for the in-flight structured stream. */ + private structuredRaw = '' + constructor( options: ChainClientOptions & ( @@ -122,6 +123,7 @@ export class ChainClient< async run(input: TInput): Promise { if (this.isLoading) return this.result + this.clearStructuredStream() this.setSteps({}) this.setIsLoading(true) this.setStatus('generating') @@ -165,6 +167,7 @@ export class ChainClient< return null } finally { this.abortController = null + this.clearStructuredStream() this.setIsLoading(false) } } @@ -177,16 +180,30 @@ export class ChainClient< this.callbacksRef.onChunk?.(chunk) - // eslint-disable-next-line @typescript-eslint/switch-exhaustiveness-check -- only handle lifecycle + chain events + // eslint-disable-next-line @typescript-eslint/switch-exhaustiveness-check -- only handle lifecycle + chain + structured-output events switch (chunk.type) { case 'CUSTOM': { if (chunk.name === CHAIN_EVENTS.STEP) { this.applyStepEvent(chunk.value) + } else if (chunk.name === 'structured-output.start') { + this.beginStructuredStream() + } else if (chunk.name === 'structured-output.complete') { + this.finishStructuredStream(chunk.value) } else if (chunk.name === GENERATION_EVENTS.RESULT) { this.setResult(chunk.value as TResult) } break } + case 'TEXT_MESSAGE_CONTENT': { + if ( + this.structuredTargetKey && + typeof chunk.delta === 'string' && + chunk.delta.length > 0 + ) { + this.appendStructuredDelta(chunk.delta) + } + break + } case 'RUN_FINISHED': { this.setStatus('success') break @@ -213,6 +230,63 @@ export class ChainClient< } } + private beginStructuredStream(): void { + // Attach to the step that most recently started (typically the chat step + // that declared outputSchema). Without an active step, ignore. + this.structuredTargetKey = this.lastActiveStepKey + this.structuredRaw = '' + if (!this.structuredTargetKey) return + + const prev = this.steps[this.structuredTargetKey] + if (!prev || prev.status !== 'active') return + + // Clear any stale partial from a previous structured attempt on this step. + this.patchStep(this.structuredTargetKey, { + ...prev, + partial: undefined, + }) + } + + private appendStructuredDelta(delta: string): void { + const key = this.structuredTargetKey + if (!key) return + + const prev = this.steps[key] + if (!prev || prev.status !== 'active') return + + this.structuredRaw += delta + const progressive = parsePartialJSON(this.structuredRaw) + // Keep the last good partial when the buffer isn't a parseable prefix yet + // (same behavior as chat's appendStructuredOutputDelta). + const nextPartial = + progressive !== undefined && progressive !== null + ? progressive + : prev.partial + + this.patchStep(key, { + ...prev, + ...(nextPartial !== undefined ? { partial: nextPartial } : {}), + }) + } + + private finishStructuredStream(value: unknown): void { + const key = this.structuredTargetKey + this.clearStructuredStream() + if (!key) return + + const prev = this.steps[key] + if (!prev || prev.status !== 'active') return + + // Prefer the validated object on complete so partial jumps to final shape + // before the outer chain:step done event arrives. + if (isRecord(value) && 'object' in value) { + this.patchStep(key, { + ...prev, + partial: value.object, + }) + } + } + private applyStepEvent(value: unknown): void { if (!isChainStepEvent(value)) return @@ -221,6 +295,7 @@ export class ChainClient< let next: ChainStepState if (value.status === 'started') { + this.lastActiveStepKey = key next = { step: value.step, ...(value.branch !== undefined && { branch: value.branch }), @@ -228,6 +303,9 @@ export class ChainClient< status: 'active', } } else if (value.status === 'error') { + if (this.structuredTargetKey === key) { + this.clearStructuredStream() + } next = { step: value.step, ...(value.branch !== undefined && { branch: value.branch }), @@ -235,27 +313,42 @@ export class ChainClient< status: 'error', error: value.error, ...(prev?.result !== undefined && { result: prev.result }), + ...(prev?.partial !== undefined && { partial: prev.partial }), } } else { + if (this.structuredTargetKey === key) { + this.clearStructuredStream() + } next = { step: value.step, ...(value.branch !== undefined && { branch: value.branch }), index: value.index, status: 'done', result: value.result, + // Drop partial once the validated result is on the step. } } + this.patchStep(key, next) + } + + private patchStep(key: string, next: ChainStepState): void { this.steps = { ...this.steps, [key]: next } this.callbacksRef.onStepsChange?.(this.steps) this.callbacksRef.onStep?.(next, this.steps) } + private clearStructuredStream(): void { + this.structuredTargetKey = null + this.structuredRaw = '' + } + stop(): void { if (this.abortController) { this.abortController.abort() this.abortController = null } + this.clearStructuredStream() this.setIsLoading(false) if (this.status === 'generating') { this.setStatus('idle') @@ -264,6 +357,7 @@ export class ChainClient< reset(): void { this.stop() + this.lastActiveStepKey = null this.setResult(null) this.setSteps({}) this.setError(undefined) diff --git a/packages/ai-client/src/chain-types.ts b/packages/ai-client/src/chain-types.ts index 5367eb66f..db792ddbc 100644 --- a/packages/ai-client/src/chain-types.ts +++ b/packages/ai-client/src/chain-types.ts @@ -24,6 +24,13 @@ export interface ChainStepState { result?: unknown /** Present when `status === 'error'`. */ error?: string + /** + * Progressive structured-output object while this step streams + * (`structured-output.start` + `TEXT_MESSAGE_CONTENT` JSON deltas, parsed + * with `parsePartialJSON`). Cleared when the step finishes or a new run + * starts. Same protocol as `useChat` / `useAssistant` chat surfaces. + */ + partial?: unknown } /** diff --git a/packages/ai-client/tests/chain-client.test.ts b/packages/ai-client/tests/chain-client.test.ts index f8bbd38d2..20e9832d4 100644 --- a/packages/ai-client/tests/chain-client.test.ts +++ b/packages/ai-client/tests/chain-client.test.ts @@ -174,7 +174,7 @@ describe('ChainClient', () => { type: EventType.RUN_ERROR, message: 'image failed', timestamp: Date.now(), - } as StreamChunk, + }, ]), }) @@ -235,4 +235,109 @@ describe('ChainClient', () => { await p2 expect(calls).toBe(1) }) + + it('streams structured-output partials onto the active step', async () => { + const finalPost = { + title: 'Urban Foxes', + subtitle: 'Quiet comeback', + body: 'Full article.', + } + const partials: Array = [] + + const client = new ChainClient({ + connection: createMockConnection([ + { + type: EventType.RUN_STARTED, + runId: 'r1', + threadId: 't1', + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: CHAIN_EVENTS.STEP, + value: { step: 'draft', index: 0, status: 'started' }, + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: 'structured-output.start', + value: { messageId: 'msg-1' }, + timestamp: Date.now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'msg-1', + delta: '{"title":"Urban', + timestamp: Date.now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'msg-1', + delta: ' Foxes","subtitle":"Quiet', + timestamp: Date.now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'msg-1', + delta: ' comeback","body":"Full article."}', + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: 'structured-output.complete', + value: { object: finalPost, raw: JSON.stringify(finalPost) }, + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: CHAIN_EVENTS.STEP, + value: { + step: 'draft', + index: 0, + status: 'done', + result: finalPost, + }, + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: 'generation:result', + value: { post: finalPost }, + timestamp: Date.now(), + }, + { + type: EventType.RUN_FINISHED, + runId: 'r1', + threadId: 't1', + timestamp: Date.now(), + }, + ]), + onStepsChange: (steps) => { + const p = steps.draft?.partial + if (p !== undefined) partials.push(p) + }, + }) + + await client.run({ topic: 'foxes' }) + + // Progressive partials should have included the title before the step finished. + expect( + partials.some( + (p) => + isRecord(p) && + typeof p.title === 'string' && + p.title.includes('Urban'), + ), + ).toBe(true) + + const done = client.getStep('draft') + expect(done?.status).toBe('done') + expect(done?.result).toEqual(finalPost) + // partial is cleared once the validated result is on the step + expect(done?.partial).toBeUndefined() + }) }) + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null +} diff --git a/packages/ai-react/src/use-chain.ts b/packages/ai-react/src/use-chain.ts index 9ca4bb05a..c5dba20e5 100644 --- a/packages/ai-react/src/use-chain.ts +++ b/packages/ai-react/src/use-chain.ts @@ -54,7 +54,9 @@ export interface UseChainReturn { * Demuxes one SSE run into: * - `result` — terminal `generation:result` * - `steps` — progressive `chain:step` state for UI (drafting / hero / …) - * - live chunks via `onChunk` (e.g. structured-output text deltas) + * - `steps[name].partial` — live structured-output object while a step streams + * (native `structured-output.start` + JSON text deltas + `.complete`) + * - live chunks via `onChunk` for anything else you want to handle yourself * * @example * ```tsx @@ -64,7 +66,8 @@ export interface UseChainReturn { * }) * * await chain.run({ topic: 'urban foxes' }) - * chain.steps['draft']?.status // 'done' + * chain.steps['draft']?.partial // { title?: string, body?: string, ... } + * chain.steps['draft']?.result // full validated object when the step finishes * chain.steps['media/hero']?.result * chain.result // { post, hero, narration } * ``` From b830a684bd7613ba7458bf38620b9d5da80cad5d Mon Sep 17 00:00:00 2001 From: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Date: Thu, 16 Jul 2026 14:59:26 +1000 Subject: [PATCH 13/13] fix(ai-client): keep ChainClient off the server @tanstack/ai barrel Import CHAIN_EVENTS from the client-safe export so React Native Metro does not resolve BaseTextAdapter declare fields through useChain. --- packages/ai-client/src/chain-client.ts | 6 ++---- packages/ai-client/src/index.ts | 4 ++++ packages/ai-client/tests/chain-client.test.ts | 2 +- packages/ai-react/src/index.ts | 2 ++ packages/ai-react/tests/use-chain.test.tsx | 2 +- packages/ai/src/client.ts | 10 ++++++++++ 6 files changed, 20 insertions(+), 6 deletions(-) diff --git a/packages/ai-client/src/chain-client.ts b/packages/ai-client/src/chain-client.ts index f7c9b26cf..fa35cfc08 100644 --- a/packages/ai-client/src/chain-client.ts +++ b/packages/ai-client/src/chain-client.ts @@ -1,10 +1,8 @@ -import { CHAIN_EVENTS } from '@tanstack/ai' -import { parsePartialJSON } from '@tanstack/ai/client' +import { CHAIN_EVENTS, parsePartialJSON } from '@tanstack/ai/client' import { GENERATION_EVENTS } from './generation-types' import { parseSSEResponse } from './sse-parser' import { chainStepKey } from './chain-types' -import type { StreamChunk } from '@tanstack/ai/client' -import type { ChainStepEventValue } from '@tanstack/ai' +import type { ChainStepEventValue, StreamChunk } from '@tanstack/ai/client' import type { ConnectConnectionAdapter, RunAgentInputContext, diff --git a/packages/ai-client/src/index.ts b/packages/ai-client/src/index.ts index 77d22ce4b..9a6faa9db 100644 --- a/packages/ai-client/src/index.ts +++ b/packages/ai-client/src/index.ts @@ -20,6 +20,10 @@ export type { ChainSteps, ChainTransport, } from './chain-types' +// Re-export the wire protocol constant so consumers don't need the server +// package just to match `chain:step` events. +export { CHAIN_EVENTS } from '@tanstack/ai/client' +export type { ChainStepEventValue } from '@tanstack/ai/client' export type { // Core message types (re-exported from @tanstack/ai via types.ts) UIMessage, diff --git a/packages/ai-client/tests/chain-client.test.ts b/packages/ai-client/tests/chain-client.test.ts index 20e9832d4..69392e301 100644 --- a/packages/ai-client/tests/chain-client.test.ts +++ b/packages/ai-client/tests/chain-client.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it, vi } from 'vitest' -import { CHAIN_EVENTS, EventType } from '@tanstack/ai' +import { CHAIN_EVENTS, EventType } from '@tanstack/ai/client' import { ChainClient, chainStepKey } from '../src' import type { StreamChunk } from '@tanstack/ai/client' import type { ConnectConnectionAdapter } from '../src/connection-adapters' diff --git a/packages/ai-react/src/index.ts b/packages/ai-react/src/index.ts index 41bffefbc..d1819ceaf 100644 --- a/packages/ai-react/src/index.ts +++ b/packages/ai-react/src/index.ts @@ -89,7 +89,9 @@ export { type ChainStepState, type ChainSteps, type ChainStepStatus, + type ChainStepEventValue, chainStepKey, + CHAIN_EVENTS, type ImageGenerateInput, type AudioGenerateInput, type SpeechGenerateInput, diff --git a/packages/ai-react/tests/use-chain.test.tsx b/packages/ai-react/tests/use-chain.test.tsx index dbcb09309..562921489 100644 --- a/packages/ai-react/tests/use-chain.test.tsx +++ b/packages/ai-react/tests/use-chain.test.tsx @@ -1,6 +1,6 @@ import { act, renderHook, waitFor } from '@testing-library/react' import { describe, expect, it, vi } from 'vitest' -import { CHAIN_EVENTS, EventType } from '@tanstack/ai' +import { CHAIN_EVENTS, EventType } from '@tanstack/ai/client' import { useChain } from '../src/use-chain' import { createMockConnectionAdapter } from './test-utils' import type { StreamChunk } from '@tanstack/ai' diff --git a/packages/ai/src/client.ts b/packages/ai/src/client.ts index 254daf736..56a7c07b0 100644 --- a/packages/ai/src/client.ts +++ b/packages/ai/src/client.ts @@ -141,3 +141,13 @@ export type { RealtimeToolResultPart, VADConfig, } from './realtime/types' + +// Client-safe chain protocol constants/types (no server-side chain() runtime). +// Importing these must not pull BaseTextAdapter or other declare-field sources +// into React Native Metro source resolution. +export { CHAIN_EVENTS } from './activities/chain/types' +export type { + ChainStepEventValue, + InferChainInput, + InferChainOutput, +} from './activities/chain/types'