diff --git a/.changeset/ai-gateway-tools.md b/.changeset/ai-gateway-tools.md new file mode 100644 index 000000000..1bdc347a8 --- /dev/null +++ b/.changeset/ai-gateway-tools.md @@ -0,0 +1,79 @@ +--- +"@tailor-platform/app-shell": minor +--- + +Add local tool support to `useAIChat()` and the AI Gateway transport layer. Streamed function names and arguments are assembled across chunks independently for each tool call. + +You can now register **local tools** via the `tools` option on `useAIChat()`. Tool-call and tool-result protocol messages stay internal to the hook — the public `messages` array remains user/assistant text only. Provider-native tools are not supported. + +Stopped or failed requests retain user messages, partial assistant text, and completed tool exchanges in subsequent request history. Unanswered tool calls are excluded. Stopping does not undo tool side effects or guarantee safe retries of tools that were already running. + +**Low-level API compatibility:** `AIGatewayChatMessage` now includes `role: "tool"`, and `AIChatCompletionEvent` now includes `type: "tool-call"`. Exhaustive handlers for these exported unions must handle the new cases, even when tools are not enabled. The high-level `useAIChat().messages` role union remains `"user" | "assistant"`. + +## Local tools + +Define tools that run inside AppShell using `defineAIChatTool`. The model decides when to call them, AppShell validates arguments with the schema, executes the handler, and feeds results back into the next model turn automatically. + +```ts +import { defineAIChatTool } from "@tailor-platform/app-shell"; +import { z } from "zod/v4"; + +const lookupCustomer = defineAIChatTool({ + description: "Look up a customer by ID and return their profile", + schema: z.object({ + customerId: z.string().describe("Customer ID"), + includeInactive: z.boolean().optional(), + }), + async execute({ customerId, includeInactive }, { signal }) { + const res = await fetch(`/api/customers/${customerId}?inactive=${includeInactive ?? false}`, { + signal, + }); + return res.json(); + }, +}); +``` + +Tool schemas must implement both [Standard Schema](https://standardschema.dev/) validation and Standard JSON Schema generation; Zod 4 does this directly. + +### Tool context + +The `execute` function receives a second argument with: + +- `signal` — the `AbortSignal` for the in-flight chat request +- `messages` — the public user/assistant message history at the time of execution + +Stopping a request while tool arguments are being validated prevents the tool from starting. Cancellation is propagated rather than returned to the model as a tool error. + +## Registering tools with `useAIChat` + +Pass tools as a record to the `tools` option. The key becomes the tool name sent to the model. + +```tsx +import { useAIChat, defineAIChatTool } from "@tailor-platform/app-shell"; +import { z } from "zod/v4"; + +const lookupCustomer = defineAIChatTool({ + description: "Look up a customer by ID", + schema: z.object({ + customerId: z.string(), + }), + async execute({ customerId }) { + return { customerId, name: "Acme Corp", plan: "enterprise" }; + }, +}); + +function ChatPanel({ client }) { + const { messages, status, sendMessage } = useAIChat({ + client, + model: "gpt-5-mini", + tools: { + lookupCustomer, + }, + }); + + // messages only contains user/assistant text — tool calls are handled internally + // status transitions: "submitted" → "streaming" → ("submitted" during tool rounds) → "ready" +} +``` + +The hook runs up to 8 tool rounds per user message, then allows a final model response using the eighth round's results. If the model requests a ninth tool batch, the request fails before those tools execute. diff --git a/.changeset/chat-progress-tool-names.md b/.changeset/chat-progress-tool-names.md new file mode 100644 index 000000000..b0f134e88 --- /dev/null +++ b/.changeset/chat-progress-tool-names.md @@ -0,0 +1,7 @@ +--- +"@tailor-platform/app-shell": minor +--- + +Show a localized waiting indicator in `AIChat.Conversation` while a response or tool round is pending. + +Expose requested tool names in a transient `toolCalls` list from `useAIChat()`. Render them with `toolCalls.map(call => )`; the localized labels stack vertically and the list clears on completion, stop, or error without changing the message history. Call metadata does not imply successful execution, and tool arguments/results remain internal. diff --git a/docs-manifest.json b/docs-manifest.json index ac5a6abef..94dba88bf 100644 --- a/docs-manifest.json +++ b/docs-manifest.json @@ -74,16 +74,17 @@ "AIChatConversationProps", "AIChatHeaderProps", "AIChatProps", + "AIChatToolCallProps", "ChainOfThoughtStepStatus", "ChatHistoryGroupData", "ChatHistoryItemData", "ToolState" ], "hashes": { - "typeSurface": "7cb8a34c1c36eb02", - "outline": "797e012e34847b80", - "snapshot": "a21f0e323ea4774e", - "outputMd": "824d1aa36bb7f194", + "typeSurface": "932ff7674cc6df99", + "outline": "0d84350ee1e55312", + "snapshot": "39edd53bc6562561", + "outputMd": "82313d5b6fe140d2", "examples": null } }, @@ -326,9 +327,9 @@ "symbols": [], "hashes": { "typeSurface": null, - "outline": "171ba8a5b6eda15b", + "outline": "43a13487e8f3e6bd", "snapshot": null, - "outputMd": "7eaff97521a14ab0", + "outputMd": "ed0db6372d25e2f9", "examples": null } }, @@ -1534,9 +1535,9 @@ "symbols": [], "hashes": { "typeSurface": null, - "outline": "4cf294ec6ae1492f", + "outline": "807ee6ee2b33d5de", "snapshot": null, - "outputMd": "f4a88f636000fd60", + "outputMd": "3bcaceab3290ca92", "examples": null } }, @@ -1773,7 +1774,7 @@ "packages/core/skills/app-shell-patterns/references/concepts/styling-theming.md": "1ed3e6d92e582e26", "packages/core/skills/app-shell-patterns/references/components/action-panel.md": "dc861082aa9610af", "packages/core/skills/app-shell-patterns/references/components/activity-card.md": "4f386e49380fb0b0", - "packages/core/skills/app-shell-patterns/references/components/ai-chat.md": "824d1aa36bb7f194", + "packages/core/skills/app-shell-patterns/references/components/ai-chat.md": "82313d5b6fe140d2", "packages/core/skills/app-shell-patterns/references/components/alert.md": "db2c4388ca783a75", "packages/core/skills/app-shell-patterns/references/components/app-shell.md": "081e4c7bb61db191", "packages/core/skills/app-shell-patterns/references/components/appearance-switcher.md": "b55f1d538651d467", @@ -1819,7 +1820,7 @@ "packages/core/skills/app-shell-patterns/references/components/tooltip.md": "113c95b2995917ed", "packages/core/skills/app-shell-patterns/references/components/with-guard.md": "d79aaf40cb97073f", "packages/core/skills/app-shell-patterns/references/api/aggregations.md": "daf6c93e7285ee88", - "packages/core/skills/app-shell-patterns/references/api/create-ai-gateway-client.md": "7eaff97521a14ab0", + "packages/core/skills/app-shell-patterns/references/api/create-ai-gateway-client.md": "ed0db6372d25e2f9", "packages/core/skills/app-shell-patterns/references/api/create-typed-paths.md": "4c3a858d3eb02a2a", "packages/core/skills/app-shell-patterns/references/api/define-i18n-labels.md": "3035dea9ac31b7f5", "packages/core/skills/app-shell-patterns/references/api/define-module.md": "199295d8e7a9f4fa", @@ -1831,7 +1832,7 @@ "packages/core/skills/app-shell-patterns/references/api/router/use-location.md": "ddb44cc51af636cc", "packages/core/skills/app-shell-patterns/references/api/router/use-navigate.md": "d018cb13a7ddc38b", "packages/core/skills/app-shell-patterns/references/api/router/use-params.md": "535f701e77dfc4e6", - "packages/core/skills/app-shell-patterns/references/api/use-ai-chat.md": "f4a88f636000fd60", + "packages/core/skills/app-shell-patterns/references/api/use-ai-chat.md": "3bcaceab3290ca92", "packages/core/skills/app-shell-patterns/references/api/use-app-shell.md": "ce989e06790aade8", "packages/core/skills/app-shell-patterns/references/api/use-app-shell-config.md": "1a0a937f3a412c13", "packages/core/skills/app-shell-patterns/references/api/use-app-shell-data.md": "2d2eaf5aa04d95d3", @@ -1856,6 +1857,6 @@ "packages/core/skills/app-shell-patterns/references/patterns/list-dense-scan.md": "59c93617dc135033", "packages/core/skills/app-shell-patterns/references/pages/document-detail.md": "5a6612349cb50ef1", "packages/core/skills/app-shell-patterns/references/migrations.md": "e0286979718e2769", - "packages/core/skills/app-shell-patterns/SKILL.md": "de250e21921fc3e8" + "packages/core/skills/app-shell-patterns/SKILL.md": "2c900e3b4d203b31" } } diff --git a/docs-src/components/ai-chat.docs.outline.md b/docs-src/components/ai-chat.docs.outline.md index f457a9954..a80ab866c 100644 --- a/docs-src/components/ai-chat.docs.outline.md +++ b/docs-src/components/ai-chat.docs.outline.md @@ -86,6 +86,17 @@ Other props spread onto the root `
`. Other props spread onto the region's outer `
`. +While the root's `status` is `"submitted"`, the conversation automatically appends a [`Spinner`](./spinner.md) and localized “Thinking…” status at the end. This covers the initial reply wait, tool execution, and the next reply wait after a tool result. The indicator disappears when text starts streaming, or when the request finishes, fails, or is stopped. It participates in the conversation's normal auto-scroll behavior; no separate loader state is needed. + +### `AIChat.ToolCall` + +A compact label for one model-requested tool call. Each label occupies its own row, so multiple calls stack vertically. It does not imply execution success, render tool arguments/results, or manage history. + +| Prop | Type | Default | Description | +| ----------- | ------------------------------ | -------- | ------------------------------------------------ | +| `call` | `{ id: string; name: string }` | required | A `toolCalls` element returned by `useAIChat()`. | +| `className` | `string` | - | Merged onto the outer row, not the badge. | + ### `AIChat.Composer` | Prop | Type | Default | Description | @@ -114,17 +125,18 @@ Most props stay inside their region. These `AIChat.Composer` props land on **ano `onValueChange` is wrapped rather than forwarded — it is called from the `Textarea`'s `onChange`, and again with `""` after a submit. `defaultValue` seeds the composer's own state and is never forwarded. Nothing on `AIChat`, `AIChat.Header`, or `AIChat.Conversation` reaches another AppShell component. -Three attached parts wrap an AppShell component. They expose only the props the part's job needs rather than inheriting the wrapped component's surface, so its visual treatment is fixed and not yours to change: +Four attached parts wrap an AppShell component. They expose only the props the part's job needs rather than inheriting the wrapped component's surface, so its visual treatment is fixed and not yours to change: | Part | Wraps | Accepts | Fixed | | ----------------------------------- | ---------------------------------------------------------- | ------------------------------------------------------------- | ---------------------------------- | | `AIChat.Suggestion` | [`Button`](./button.md) | `suggestion`, `onSelect`, `className`, `disabled`, `children` | `variant="secondary"`, `size="sm"` | | `AIChat.Action` | [`Button`](./button.md) inside a [`Tooltip`](./tooltip.md) | `label`, `onClick`, `className`, `disabled`, `children` | `variant="ghost"`, `size="icon"` | | `AIChat.ChainOfThoughtSearchResult` | [`Badge`](./badge.md) | `className`, `children`, `variant` | — | +| `AIChat.ToolCall` | [`Badge`](./badge.md) | `call`, `className` | `variant="neutral"` | `variant` stays open on `ChainOfThoughtSearchResult` because conveying status is what a `Badge` is for. `AIChat.Action`'s `label` becomes both the accessible name and the tooltip text. -One thing to know about `className` on these: it is merged with `cn()`, which is not configured for AppShell's internal prefix, so it cannot reliably override a base utility from the wrapped component. Add properties the component doesn't set, and reach for the wrapped component's own props or variants to change ones it does — see [Styling AppShell components](../concepts/styling-theming.md#styling-appshell-components). +`ToolCall` applies `className` to its outer row, not the inner badge. For `Suggestion`, `Action`, and `ChainOfThoughtSearchResult`, `className` is merged with `cn()`, which is not configured for AppShell's internal prefix, so it cannot reliably override a base utility from the wrapped component. Add properties the component doesn't set, and reach for the wrapped component's own props or variants to change ones it does — see [Styling AppShell components](../concepts/styling-theming.md#styling-appshell-components). ## Filling the page @@ -263,6 +275,28 @@ Structured step timeline for agentic turns — the discrete plan being executed, `status` is `"complete" | "active" | "pending"`. +### Tool calls during the active flow + +`useAIChat()` returns a transient `toolCalls` list separately from `messages`. Render it with `AIChat.ToolCall` after the message history inside `AIChat.Conversation`: + +```tsx +{ + toolCalls.length > 0 ? ( + + {toolCalls.map((call) => ( + + ))} + + ) : null; +} +``` + +Each call renders on its own row, in the order supplied. Labels are localized automatically. + +The list clears when the response finishes, is stopped, or fails. It does not create tool messages or persist in the conversation log. + +These names record model requests, not execution success. Tool arguments and results stay internal. Do not infer a `ToolHeader` success or running state from name-only metadata. + ### Tool Collapsible tool-call card. Maps 1:1 to the AI SDK's tool-part lifecycle, so a streaming loop can update a call in place as parts arrive. @@ -281,7 +315,7 @@ Collapsible tool-call card. Maps 1:1 to the AI SDK's tool-part lifecycle, so a s ### Sources -Compact "Used N sources" citation trigger under a grounded answer. +Compact "Used N sources" citation trigger under a grounded answer. Supply source data from your application; `useAIChat()` does not return citation metadata. ```tsx diff --git a/docs-src/hooks/create-ai-gateway-client.docs.outline.md b/docs-src/hooks/create-ai-gateway-client.docs.outline.md index 421885841..64463135d 100644 --- a/docs-src/hooks/create-ai-gateway-client.docs.outline.md +++ b/docs-src/hooks/create-ai-gateway-client.docs.outline.md @@ -7,7 +7,7 @@ description: Create a low-level AI Gateway client that reuses AppShell authentic # createAIGatewayClient -Creates a small AI Gateway transport client for text-only chat completions. +Creates a small AI Gateway transport client for chat completions and local function tool calls. ## Signature @@ -43,6 +43,7 @@ interface AIGatewayClient { The iterable yields completion events: - `text-delta` — append `event.text` to build the assistant response +- `tool-call` — a local function tool call requested by the model - `done` — terminal event with an optional `finishReason` ## Related Types @@ -56,11 +57,33 @@ type AIGatewayChatMessage = | { role: "assistant"; content?: string; + toolCalls?: AIGatewayToolCall[]; + } + | { + role: "tool"; + toolCallId: string; + content: string; }; +interface AIGatewayToolCall { + id: string; + name: string; + argumentsText: string; +} + +interface AIGatewayFunctionTool { + type: "function"; + function: { + name: string; + description?: string; + parameters: Record; + }; +} + interface AIGatewayChatRequest { model: string; messages: AIGatewayChatMessage[]; + tools?: AIGatewayFunctionTool[]; signal?: AbortSignal; } @@ -69,6 +92,12 @@ type AIChatCompletionEvent = type: "text-delta"; text: string; } + | { + type: "tool-call"; + toolCallId: string; + toolName: string; + argumentsText: string; + } | { type: "done"; finishReason?: string; @@ -106,7 +135,8 @@ console.log(text); ## Notes - AppShell chooses the appropriate AI Gateway transport automatically -- The low-level API is intentionally narrow: text deltas plus completion metadata +- Local tools are sent as normalized `type: "function"` definitions; provider-native tools are not supported +- The low-level API stays event-based so hooks can layer streaming UI and tool loops on top - `request.signal` is passed through so callers can abort in-flight work ## Related diff --git a/docs-src/hooks/use-ai-chat.docs.outline.md b/docs-src/hooks/use-ai-chat.docs.outline.md index a21918ca9..25f4d8913 100644 --- a/docs-src/hooks/use-ai-chat.docs.outline.md +++ b/docs-src/hooks/use-ai-chat.docs.outline.md @@ -2,18 +2,23 @@ kind: prose group: use-ai-chat title: useAIChat -description: Simple text-only chat hook for AI Gateway +description: AI Gateway chat hook with optional local tools --- # useAIChat -React hook for simple text-only chat on top of `createAIGatewayClient`. +React hook for AI Gateway chat on top of `createAIGatewayClient`, with optional local tool support. ## Signature ```typescript -const useAIChat: (config: { client: AIGatewayClient; model: string }) => { +const useAIChat: (config: { + client: AIGatewayClient; + model: string; + tools?: Record; +}) => { messages: AIChatMessage[]; + toolCalls: { id: string; name: string }[]; status: "ready" | "submitted" | "streaming" | "error"; error?: Error; sendMessage: (message: string) => Promise; @@ -33,10 +38,17 @@ interface AIChatMessage { } ``` +### `toolCalls` + +- **Type:** `{ id: string; name: string }[]` +- **Description:** Tool calls requested by the model in the active flow, including earlier tool rounds of the current request. This transient list is cleared when the response finishes, is stopped, or fails; it is never added to `messages` or conversation history. + +Calls become available before local tool execution. They record requested tool names, not successful execution. Arguments and results remain internal. + ### `status` - **Type:** `"ready" | "submitted" | "streaming" | "error"` -- **Description:** Current request state +- **Description:** Current request state. `submitted` covers waiting for the first text, local tool execution, and waiting for the model's next reply after a tool result. `streaming` begins when text arrives. `AIChat.Conversation` automatically shows a waiting indicator while `submitted`. ### `error` @@ -54,10 +66,20 @@ interface AIChatMessage { - **Type:** `() => void` - **Description:** Aborts the current request if one is in progress +### `tools` + +Register local tools under a single object using `defineAIChatTool(...)`. Their schemas must implement both [Standard Schema](https://standardschema.dev/) validation and Standard JSON Schema generation (for example, Zod 4). Provider-native tools are not supported. + ## Usage ```tsx -import { createAuthClient, createAIGatewayClient, useAIChat } from "@tailor-platform/app-shell"; +import { + createAuthClient, + createAIGatewayClient, + defineAIChatTool, + useAIChat, +} from "@tailor-platform/app-shell"; +import { z } from "zod/v4"; const authClient = createAuthClient({ clientId: "your-client-id", @@ -69,10 +91,23 @@ const aiClient = createAIGatewayClient({ authClient, }); +const lookupCustomer = defineAIChatTool({ + description: "Look up a customer in the current workspace", + schema: z.object({ + customerId: z.string(), + }), + async execute({ customerId }) { + return { customerId, name: "Acme Corp" }; + }, +}); + export function ChatScreen() { - const { messages, sendMessage, status, stop, error } = useAIChat({ + const { messages, toolCalls, sendMessage, status, stop, error } = useAIChat({ client: aiClient, model: "gpt-5-mini", + tools: { + lookupCustomer, + }, }); return ( @@ -82,6 +117,9 @@ export function ChatScreen() { {message.role}: {message.content}
))} + {toolCalls.map((call) => ( +
Tool call: {call.name}
+ ))}
))} + {toolCalls.map((call) => ( +
Tool call: {call.name}
+ ))} "`; + exports[`AIChat > snapshots > default 1`] = `"

Hello there.

"`; exports[`AIChat > snapshots > with header 1`] = `"
Assistant
"`; diff --git a/packages/core/__snapshots__/src__components__ai-chat__tool-call.test.tsx.snap b/packages/core/__snapshots__/src__components__ai-chat__tool-call.test.tsx.snap new file mode 100644 index 000000000..7f64063a8 --- /dev/null +++ b/packages/core/__snapshots__/src__components__ai-chat__tool-call.test.tsx.snap @@ -0,0 +1,3 @@ +// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html + +exports[`AIChat.ToolCall > renders a neutral tool-call label without implying success 1`] = `"
Tool call: add
"`; diff --git a/packages/core/package.json b/packages/core/package.json index c9d7f8a96..6b08791dc 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -98,7 +98,8 @@ "typescript": "5.9.3", "vite": "catalog:", "vite-plugin-dts": "5.1.2", - "vitest": "catalog:" + "vitest": "catalog:", + "zod": "4.4.3" }, "peerDependencies": { "react": ">=19.2.7", diff --git a/packages/core/src/ai/assistant-loop.ts b/packages/core/src/ai/assistant-loop.ts new file mode 100644 index 000000000..cab8aa351 --- /dev/null +++ b/packages/core/src/ai/assistant-loop.ts @@ -0,0 +1,161 @@ +import type { + AIGatewayChatMessage, + AIGatewayClient, + AIGatewayFunctionTool, + AIGatewayToolCall, +} from "./client"; +import type { AILocalTool } from "./tools"; +import { deriveVisibleMessages, resolveToolCalls } from "./tool-execution"; + +const MAX_TOOL_ROUNDS = 8; + +/** + * Splits the public tools object into: + * - normalized AI Gateway tool definitions sent over the wire + * - local tool executors kept for in-process validation and execution + */ +function normalizeConfiguredTools(tools: Record | undefined): { + gatewayTools: AIGatewayFunctionTool[]; + localTools: Map; +} { + if (!tools) { + return { + gatewayTools: [], + localTools: new Map(), + }; + } + + const localTools = new Map(); + const gatewayTools = Object.entries(tools).map(([name, tool]) => { + localTools.set(name, tool); + return { + type: "function", + function: { + name, + ...(tool.description ? { description: tool.description } : {}), + parameters: tool.schema["~standard"].jsonSchema.input({ target: "draft-07" }), + }, + } satisfies AIGatewayFunctionTool; + }); + + return { + gatewayTools, + localTools, + }; +} + +export interface AIChatAssistantTurnResult { + gatewayAssistantMessage: Extract; + toolCalls: AIGatewayToolCall[]; +} + +export type AssistantLoopEvent = + | { type: "text-delta"; delta: string } + | { type: "turn-end"; turn: AIChatAssistantTurnResult } + | { type: "tool-resolution-start" } + | { + type: "tool-results"; + messages: Extract[]; + } + | { type: "complete" }; + +/** + * Pure async generator that runs the multi-turn assistant loop. + * Yields events for the consumer to translate into state changes. + * Manages transcript internally — the caller only provides the initial transcript. + */ +export async function* runAssistantLoop(input: { + client: AIGatewayClient; + model: string; + tools: Record | undefined; + transcript: AIGatewayChatMessage[]; + signal: AbortSignal; +}): AsyncGenerator { + const { gatewayTools, localTools } = normalizeConfiguredTools(input.tools); + const transcript = [...input.transcript]; + + for (let round = 0; round <= MAX_TOOL_ROUNDS; round += 1) { + const turn: AIChatAssistantTurnResult = yield* streamAssistantTurn({ + client: input.client, + model: input.model, + tools: gatewayTools, + transcript, + signal: input.signal, + }); + + if (turn.gatewayAssistantMessage.content || turn.toolCalls.length > 0) { + transcript.push(turn.gatewayAssistantMessage); + } + + yield { type: "turn-end", turn }; + + if (turn.toolCalls.length === 0) { + yield { type: "complete" }; + return; + } + + if (round === MAX_TOOL_ROUNDS) { + break; + } + + yield { type: "tool-resolution-start" }; + + const toolMessages = await resolveToolCalls({ + toolCalls: turn.toolCalls, + localTools, + signal: input.signal, + visibleMessages: deriveVisibleMessages(transcript), + }); + + transcript.push(...toolMessages); + yield { type: "tool-results", messages: toolMessages }; + } + + throw new Error(`AI chat exceeded the maximum number of tool rounds (${MAX_TOOL_ROUNDS}).`); +} + +/** + * Streams one assistant turn from the AI Gateway. + * Yields text-delta events as they arrive. + * Returns the collected turn result (content + tool calls) via generator return. + */ +async function* streamAssistantTurn(input: { + client: AIGatewayClient; + model: string; + tools: AIGatewayFunctionTool[]; + transcript: AIGatewayChatMessage[]; + signal: AbortSignal; +}): AsyncGenerator { + let content = ""; + const toolCalls: AIGatewayToolCall[] = []; + + for await (const event of input.client.streamChatCompletion({ + model: input.model, + messages: input.transcript, + ...(input.tools.length > 0 ? { tools: input.tools } : {}), + signal: input.signal, + })) { + if (event.type === "text-delta") { + content = `${content}${event.text}`; + yield { type: "text-delta", delta: event.text }; + continue; + } + + if (event.type === "tool-call") { + toolCalls.push({ + id: event.toolCallId, + name: event.toolName, + argumentsText: event.argumentsText, + }); + } + } + + return { + gatewayAssistantMessage: { + role: "assistant", + ...(content ? { content } : {}), + ...(toolCalls.length > 0 ? { toolCalls } : {}), + }, + toolCalls, + }; +} diff --git a/packages/core/src/ai/client.test.ts b/packages/core/src/ai/client.test.ts index f747297cb..b0c713de8 100644 --- a/packages/core/src/ai/client.test.ts +++ b/packages/core/src/ai/client.test.ts @@ -178,6 +178,188 @@ describe("createAIGatewayClient", () => { }); }); + it("emits tool calls from streaming responses", async () => { + const authClient = createMockAuthClient( + new Response( + createSSEStream([ + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"lookupCustomer","arguments":"{\\"customerId\\":\\""}}]}}]}\n\n', + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"cust-1\\"}"}}]}}]}\n\n', + 'data: {"choices":[{"finish_reason":"tool_calls"}]}\n\n', + "data: [DONE]\n\n", + ]), + { + status: 200, + headers: { "Content-Type": "text/event-stream" }, + }, + ), + ); + const client = createAIGatewayClient({ + gatewayUri: "https://gateway.example.com", + authClient, + }); + + await expect( + collectEvents( + client, + createRequest({ + tools: [ + { + type: "function", + function: { + name: "lookupCustomer", + parameters: { type: "object", properties: { customerId: { type: "string" } } }, + }, + }, + ], + }), + ), + ).resolves.toEqual([ + { + type: "tool-call", + toolCallId: "call_1", + toolName: "lookupCustomer", + argumentsText: '{"customerId":"cust-1"}', + }, + { type: "done", finishReason: "tool_calls" }, + ]); + + const [input, init] = getRequestCall(authClient); + expect(JSON.parse(String(await readRequestBody(input, init)))).toEqual({ + model: "gpt-5-mini", + messages: [{ role: "user", content: "Hello" }], + tools: [ + { + type: "function", + function: { + name: "lookupCustomer", + parameters: { type: "object", properties: { customerId: { type: "string" } } }, + }, + }, + ], + stream: true, + }); + }); + + it("accumulates split function names independently by tool-call index", async () => { + const authClient = createMockAuthClient( + new Response( + createSSEStream([ + 'data: {"choices":[{"delta":{"tool_calls":[{"index":1,"id":"call_2","type":"function","function":{"name":"lookup","arguments":"{\\"orderId\\":\\""}}]}}]}\n\n', + 'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"lookup","arguments":"{\\"customerId\\":\\""}}]}}]}\n\n', + 'data: {"choices":[{"delta":{"tool_calls":[{"index":1,"function":{"name":"Order","arguments":"order-1\\"}"}},{"index":0,"function":{"name":"Customer","arguments":"cust-1\\"}"}}]}}]}\n\n', + 'data: {"choices":[{"finish_reason":"tool_calls"}]}\n\n', + "data: [DONE]\n\n", + ]), + { + status: 200, + headers: { "Content-Type": "text/event-stream" }, + }, + ), + ); + const client = createAIGatewayClient({ + gatewayUri: "https://gateway.example.com", + authClient, + }); + + await expect(collectEvents(client, createRequest())).resolves.toEqual([ + { + type: "tool-call", + toolCallId: "call_1", + toolName: "lookupCustomer", + argumentsText: '{"customerId":"cust-1"}', + }, + { + type: "tool-call", + toolCallId: "call_2", + toolName: "lookupOrder", + argumentsText: '{"orderId":"order-1"}', + }, + { type: "done", finishReason: "tool_calls" }, + ]); + }); + + it("emits local function tool calls from json responses", async () => { + const authClient = createMockAuthClient( + new Response( + JSON.stringify({ + choices: [ + { + message: { + content: null, + tool_calls: [ + { + id: "call_1", + type: "function", + function: { name: "lookupCustomer", arguments: '{"customerId":"cust-1"}' }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + }), + { + status: 200, + headers: { "Content-Type": "application/json" }, + }, + ), + ); + const client = createAIGatewayClient({ + gatewayUri: "https://gateway.example.com", + authClient, + }); + const tools = [ + { + type: "function" as const, + function: { + name: "lookupCustomer", + parameters: { type: "object", properties: { customerId: { type: "string" } } }, + }, + }, + ]; + + await expect( + collectEvents(client, createRequest({ model: "gemini-2.5-flash", tools })), + ).resolves.toEqual([ + { + type: "tool-call", + toolCallId: "call_1", + toolName: "lookupCustomer", + argumentsText: '{"customerId":"cust-1"}', + }, + { type: "done", finishReason: "tool_calls" }, + ]); + + const [input, init] = getRequestCall(authClient); + expect(JSON.parse(String(await readRequestBody(input, init)))).toMatchObject({ + tools, + stream: false, + }); + }); + + it.each(["gpt-5-mini", "gemini-2.5-flash"])( + "rejects provider tools before sending a request for %s", + async (model) => { + const authClient = createMockAuthClient(new Response()); + const client = createAIGatewayClient({ + gatewayUri: "https://gateway.example.com", + authClient, + }); + + await expect( + collectEvents( + client, + createRequest({ + model, + // @ts-expect-error Provider tools are not supported. + tools: [{ type: "provider", provider: "openai", name: "web_search" }], + }), + ), + ).rejects.toThrow("AI Gateway only supports local function tools."); + expect(authClient.fetch).not.toHaveBeenCalled(); + }, + ); + it("emits done even when json responses do not include assistant text", async () => { const authClient = createMockAuthClient( new Response( diff --git a/packages/core/src/ai/client.ts b/packages/core/src/ai/client.ts index 5f9fca4e9..931d2857a 100644 --- a/packages/core/src/ai/client.ts +++ b/packages/core/src/ai/client.ts @@ -6,6 +6,14 @@ import type { } from "openai/resources/chat/completions"; import type { AuthClient } from "@tailor-platform/auth-public-client"; +/** Normalized local function tool call emitted by the AI Gateway transport. */ +export interface AIGatewayToolCall { + id: string; + name: string; + argumentsText: string; +} + +/** Internal transcript message shape sent to the AI Gateway transport layer. */ export type AIGatewayChatMessage = | { role: "system" | "user"; @@ -14,19 +22,43 @@ export type AIGatewayChatMessage = | { role: "assistant"; content?: string; + toolCalls?: AIGatewayToolCall[]; + } + | { + role: "tool"; + toolCallId: string; + content: string; }; +/** Normalized local function tool definition understood by the AI Gateway. */ +export interface AIGatewayFunctionTool { + type: "function"; + function: { + name: string; + description?: string; + parameters: Record; + }; +} + export interface AIGatewayChatRequest { model: string; messages: AIGatewayChatMessage[]; + tools?: AIGatewayFunctionTool[]; signal?: AbortSignal; } +/** Low-level event stream consumed by `useAIChat()`. */ export type AIChatCompletionEvent = | { type: "text-delta"; text: string; } + | { + type: "tool-call"; + toolCallId: string; + toolName: string; + argumentsText: string; + } | { type: "done"; finishReason?: string; @@ -103,29 +135,36 @@ async function* streamOpenAICompatibleResponse(input: { request: AIGatewayChatRequest; }): AsyncGenerator { try { - const stream = await input.client.chat.completions.create( - { - model: input.request.model, - messages: input.request.messages.map(toOpenAIMessage), - stream: true, - }, + const stream = (await input.client.chat.completions.create( + buildChatRequestBody(input.request, true), { signal: input.request.signal, }, - ); + )) as AsyncIterable; let finishReason: string | undefined; + const toolCalls = new Map>(); for await (const chunk of stream) { const choice = chunk.choices[0]; const delta = extractStreamingText(choice); finishReason = extractFinishReason(choice?.finish_reason) ?? finishReason; + collectStreamingToolCalls(toolCalls, choice); if (delta) { yield { type: "text-delta", text: delta }; } } + for (const toolCall of finalizeStreamingToolCalls(toolCalls)) { + yield { + type: "tool-call", + toolCallId: toolCall.id, + toolName: toolCall.name, + argumentsText: toolCall.argumentsText, + }; + } + yield createDoneEvent(finishReason); } catch (error) { throw normalizeGatewayError(error, input.request.signal); @@ -137,16 +176,12 @@ async function* streamJSONResponse(input: { request: AIGatewayChatRequest; }): AsyncGenerator { try { - const completion = await input.client.chat.completions.create( - { - model: input.request.model, - messages: input.request.messages.map(toOpenAIMessage), - stream: false, - }, + const completion = (await input.client.chat.completions.create( + buildChatRequestBody(input.request, false), { signal: input.request.signal, }, - ); + )) as ChatCompletion; const choice = completion.choices[0]; const text = extractFinalText(choice); @@ -155,21 +190,71 @@ async function* streamJSONResponse(input: { yield { type: "text-delta", text }; } + for (const toolCall of extractToolCalls(choice)) { + yield { + type: "tool-call", + toolCallId: toolCall.id, + toolName: toolCall.name, + argumentsText: toolCall.argumentsText, + }; + } + yield createDoneEvent(extractFinishReason(choice?.finish_reason)); } catch (error) { throw normalizeGatewayError(error, input.request.signal); } } +/** + * Builds the narrow OpenAI-compatible request body that Envoy AI Gateway + * expects, while preserving AppShell's normalized tool contract internally. + */ +function buildChatRequestBody( + request: AIGatewayChatRequest, + stream: boolean, +): Parameters[0] { + if (request.tools?.some((tool) => tool.type !== "function")) { + throw new Error("AI Gateway only supports local function tools."); + } + + return { + model: request.model, + messages: request.messages.map(toOpenAIMessage), + ...(request.tools?.length ? { tools: request.tools } : {}), + stream, + } as Parameters[0]; +} + function withTrailingSlash(value: string): string { return value.endsWith("/") ? value : `${value}/`; } function toOpenAIMessage(message: AIGatewayChatMessage): ChatCompletionMessageParam { if (message.role === "assistant") { - return message.content === undefined - ? { role: "assistant" } - : { role: "assistant", content: message.content }; + return { + role: "assistant", + ...(message.content !== undefined ? { content: message.content } : { content: null }), + ...(message.toolCalls?.length + ? { + tool_calls: message.toolCalls.map((toolCall) => ({ + id: toolCall.id, + type: "function", + function: { + name: toolCall.name, + arguments: toolCall.argumentsText, + }, + })), + } + : {}), + } as ChatCompletionMessageParam; + } + + if (message.role === "tool") { + return { + role: "tool", + tool_call_id: message.toolCallId, + content: message.content, + } as ChatCompletionMessageParam; } return { @@ -178,6 +263,63 @@ function toOpenAIMessage(message: AIGatewayChatMessage): ChatCompletionMessagePa }; } +/** + * Accumulates streamed tool-call deltas into complete normalized tool calls. + * + * OpenAI-compatible streams may split ids, names, and JSON arguments across + * many chunks, so this keeps a per-index assembly buffer until the stream ends. + */ +function collectStreamingToolCalls( + toolCalls: Map>, + choice: ChatCompletionChunk["choices"][number] | undefined, +): void { + for (const toolCallDelta of choice?.delta?.tool_calls ?? []) { + const existing = toolCalls.get(toolCallDelta.index) ?? { argumentsText: "" }; + + if (toolCallDelta.id) { + existing.id = toolCallDelta.id; + } + + if (toolCallDelta.function?.name) { + existing.name = `${existing.name ?? ""}${toolCallDelta.function.name}`; + } + + if (toolCallDelta.function?.arguments) { + existing.argumentsText = `${existing.argumentsText ?? ""}${toolCallDelta.function.arguments}`; + } + + toolCalls.set(toolCallDelta.index, existing); + } +} + +/** Finalizes buffered streamed tool calls into stable ordered results. */ +function finalizeStreamingToolCalls( + toolCalls: Map>, +): AIGatewayToolCall[] { + return [...toolCalls.entries()] + .toSorted(([leftIndex], [rightIndex]) => leftIndex - rightIndex) + .map(([, toolCall]) => ({ + id: toolCall.id ?? crypto.randomUUID(), + name: toolCall.name ?? "", + argumentsText: toolCall.argumentsText ?? "", + })) + .filter((toolCall) => toolCall.name.length > 0); +} + +/** Extracts complete function tool calls from a non-streaming completion. */ +function extractToolCalls( + choice: ChatCompletion["choices"][number] | undefined, +): AIGatewayToolCall[] { + return (choice?.message?.tool_calls ?? []) + .filter((toolCall) => toolCall.type === "function") + .map((toolCall) => ({ + id: toolCall.id ?? crypto.randomUUID(), + name: toolCall.function.name, + argumentsText: toolCall.function.arguments, + })) + .filter((toolCall) => toolCall.name.length > 0); +} + function extractStreamingText(choice: ChatCompletionChunk["choices"][number] | undefined): string { return extractText(choice?.delta?.content); } diff --git a/packages/core/src/ai/tool-execution.test.ts b/packages/core/src/ai/tool-execution.test.ts new file mode 100644 index 000000000..2967ed154 --- /dev/null +++ b/packages/core/src/ai/tool-execution.test.ts @@ -0,0 +1,72 @@ +import { describe, expect, it, vi } from "vitest"; +import { z } from "zod/v4"; +import type { AILocalTool } from "./tools"; +import { resolveToolCalls } from "./tool-execution"; + +const toolCall = { id: "call-1", name: "lookup", argumentsText: "{}" }; + +describe("resolveToolCalls", () => { + it.each([false, true])( + "propagates cancellation while validation is pending (validation rejects: %s)", + async (rejectValidation) => { + let releaseValidation!: () => void; + const validationGate = new Promise((resolve) => { + releaseValidation = resolve; + }); + const schema = z.object({}); + const validate = vi.spyOn(schema["~standard"], "validate").mockImplementation(async () => { + await validationGate; + if (rejectValidation) throw new Error("validation failed"); + return { value: {} }; + }); + const execute = vi.fn(() => "found"); + const controller = new AbortController(); + const pending = resolveToolCalls({ + toolCalls: [toolCall], + localTools: new Map([["lookup", { kind: "local", schema, execute }]]), + signal: controller.signal, + visibleMessages: [], + }); + + expect(validate).toHaveBeenCalledOnce(); + controller.abort(); + releaseValidation(); + + await expect(pending).rejects.toMatchObject({ name: "AbortError" }); + expect(execute).not.toHaveBeenCalled(); + }, + ); + + it("propagates cancellation errors from the executor", async () => { + const error = new DOMException("Stopped", "AbortError"); + const execute = vi.fn().mockRejectedValue(error); + + await expect( + resolveToolCalls({ + toolCalls: [toolCall], + localTools: new Map([ + ["lookup", { kind: "local", schema: z.object({}), execute }], + ]), + signal: new AbortController().signal, + visibleMessages: [], + }), + ).rejects.toBe(error); + }); + + it("still returns ordinary executor failures as tool results", async () => { + const execute = vi.fn().mockRejectedValue(new Error("lookup failed")); + + await expect( + resolveToolCalls({ + toolCalls: [toolCall], + localTools: new Map([ + ["lookup", { kind: "local", schema: z.object({}), execute }], + ]), + signal: new AbortController().signal, + visibleMessages: [], + }), + ).resolves.toEqual([ + { role: "tool", toolCallId: "call-1", content: '{"error":"lookup failed"}' }, + ]); + }); +}); diff --git a/packages/core/src/ai/tool-execution.ts b/packages/core/src/ai/tool-execution.ts new file mode 100644 index 000000000..533482e3f --- /dev/null +++ b/packages/core/src/ai/tool-execution.ts @@ -0,0 +1,140 @@ +import { isAbortError, type AIGatewayChatMessage, type AIGatewayToolCall } from "./client"; +import type { AIChatMessage } from "./use-ai-chat"; +import type { AIChatToolContext, AILocalTool } from "./tools"; + +/** + * Executes requested local tools in call order and converts their outputs into + * internal `role: "tool"` transcript messages for the next model round. + */ +export async function resolveToolCalls(input: { + toolCalls: AIGatewayToolCall[]; + localTools: Map; + signal: AbortSignal; + visibleMessages: AIChatToolContext["messages"]; +}): Promise[]> { + const messages: Extract[] = []; + + for (const toolCall of input.toolCalls) { + if (input.signal.aborted) { + throw createAbortError(); + } + + messages.push( + await resolveToolCall({ + toolCall, + localTools: input.localTools, + context: { + signal: input.signal, + messages: input.visibleMessages, + }, + }), + ); + } + + return messages; +} + +/** + * Validates, executes, and stringifies one local tool call. + * + * Tool failures are converted into tool result payloads so the model can + * recover in-band instead of failing the entire chat request. + */ +async function resolveToolCall(input: { + toolCall: AIGatewayToolCall; + localTools: Map; + context: AIChatToolContext; +}): Promise> { + const tool = input.localTools.get(input.toolCall.name); + + if (!tool) { + return { + role: "tool", + toolCallId: input.toolCall.id, + content: JSON.stringify({ error: `Unknown tool: ${input.toolCall.name}` }), + }; + } + + try { + const rawArguments = parseToolArguments(input.toolCall.argumentsText); + const result = await tool.schema["~standard"].validate(rawArguments); + + if (input.context.signal.aborted) { + throw createAbortError(); + } + + if (result.issues) { + return { + role: "tool", + toolCallId: input.toolCall.id, + content: JSON.stringify({ error: result.issues.map((issue) => issue.message).join("; ") }), + }; + } + + const output = await tool.execute(result.value, input.context); + return { + role: "tool", + toolCallId: input.toolCall.id, + content: stringifyToolResult(output), + }; + } catch (error) { + if (input.context.signal.aborted) { + throw createAbortError(); + } + if (isAbortError(error)) { + throw error; + } + + return { + role: "tool", + toolCallId: input.toolCall.id, + content: JSON.stringify({ error: error instanceof Error ? error.message : String(error) }), + }; + } +} + +/** + * Derives the user/assistant visible messages from the full transcript. + * Used to provide conversation context to tool executors. + */ +export function deriveVisibleMessages(transcript: AIGatewayChatMessage[]): AIChatMessage[] { + const result: AIChatMessage[] = []; + + for (const message of transcript) { + if (message.role === "user") { + // ponytail: ids are placeholders — tool executors only read content, not ids + result.push({ id: "", role: "user", content: message.content }); + } else if (message.role === "assistant" && message.content) { + result.push({ id: "", role: "assistant", content: message.content }); + } + } + + return result; +} + +function parseToolArguments(argumentsText: string): unknown { + const text = argumentsText.trim(); + return text ? JSON.parse(text) : {}; +} + +function stringifyToolResult(value: unknown): string { + if (typeof value === "string") { + return value; + } + + if (value === undefined) { + return "null"; + } + + return JSON.stringify(value); +} + +function createAbortError(): Error { + if (typeof DOMException !== "undefined") { + return new DOMException("The operation was aborted.", "AbortError"); + } + + const error = new Error("The operation was aborted."); + error.name = "AbortError"; + return error; +} diff --git a/packages/core/src/ai/tools.test-d.ts b/packages/core/src/ai/tools.test-d.ts new file mode 100644 index 000000000..738d4b790 --- /dev/null +++ b/packages/core/src/ai/tools.test-d.ts @@ -0,0 +1,47 @@ +import { describe, expectTypeOf, it } from "vitest"; +import { z } from "zod/v4"; +import { + defineAIChatTool, + type AIGatewayChatRequest, + type AILocalTool, + type useAIChat, +} from "../index"; + +describe("local-only tool contract", () => { + it("accepts only local tools in useAIChat", () => { + type ConfiguredTool = NonNullable[0]["tools"]>[string]; + + expectTypeOf().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf<"local">(); + }); + + it("accepts heterogeneous schemas without weakening schema or executor types", () => { + const text = defineAIChatTool({ + schema: z.object({ value: z.string() }), + execute: ({ value }) => { + expectTypeOf(value).toEqualTypeOf(); + return value; + }, + }); + const number = defineAIChatTool({ + schema: z.object({ value: z.number() }), + execute: ({ value }) => value, + }); + const tools = { text, number }; + const widerSchema = z.object({ value: z.union([z.string(), z.number()]) }); + + expectTypeOf(tools).toMatchTypeOf[0]["tools"]>>(); + expectTypeOf<{}>().not.toMatchTypeOf(); + expectTypeOf(text).not.toMatchTypeOf>(); + }); + + it("accepts only function tools in AI Gateway requests", () => { + type GatewayTool = NonNullable[number]; + + expectTypeOf().toEqualTypeOf<"function">(); + }); + + it("does not export provider tool helpers", () => { + expectTypeOf().not.toHaveProperty("aiProviderTool"); + }); +}); diff --git a/packages/core/src/ai/tools.test.ts b/packages/core/src/ai/tools.test.ts new file mode 100644 index 000000000..2070b78a8 --- /dev/null +++ b/packages/core/src/ai/tools.test.ts @@ -0,0 +1,38 @@ +import { describe, expect, it } from "vitest"; +import { z } from "zod/v4"; +import { defineAIChatTool } from "./tools"; + +describe("AI chat tool helpers", () => { + it("accepts Zod schemas for local tool definitions", async () => { + const lookupCustomer = defineAIChatTool({ + description: "Look up a customer", + schema: z.object({ + customerId: z.string(), + includeInactive: z.boolean().optional(), + }), + async execute({ customerId, includeInactive }) { + return { customerId, includeInactive, name: "Acme" }; + }, + }); + + expect(await lookupCustomer.schema["~standard"].validate({ customerId: "cust-1" })).toEqual({ + value: { customerId: "cust-1" }, + }); + expect(lookupCustomer.schema["~standard"].jsonSchema.input({ target: "draft-07" })).toEqual({ + $schema: "http://json-schema.org/draft-07/schema#", + type: "object", + properties: { + customerId: { type: "string" }, + includeInactive: { type: "boolean" }, + }, + required: ["customerId"], + }); + + await expect( + lookupCustomer.execute( + { customerId: "cust-1" }, + { signal: new AbortController().signal, messages: [] }, + ), + ).resolves.toEqual({ customerId: "cust-1", includeInactive: undefined, name: "Acme" }); + }); +}); diff --git a/packages/core/src/ai/tools.ts b/packages/core/src/ai/tools.ts new file mode 100644 index 000000000..4388b4e74 --- /dev/null +++ b/packages/core/src/ai/tools.ts @@ -0,0 +1,89 @@ +import type { StandardJSONSchemaV1, StandardSchemaV1 } from "@standard-schema/spec"; +import type { AIChatMessage } from "./use-ai-chat"; + +/** + * Schema contract for local AI chat tools. + * + * A schema must validate runtime inputs and generate JSON Schema for the AI + * Gateway request. Zod 4 implements both Standard Schema contracts directly. + * + * @example + * ```ts + * import { z } from "zod/v4"; + * + * const schema = z.object({ + * customerId: z.string(), + * includeInactive: z.boolean().optional(), + * }); + * ``` + */ +export type AIChatToolSchema = StandardSchemaV1 & + StandardJSONSchemaV1; + +/** Context passed to local tool executors. */ +export interface AIChatToolContext { + /** Abort signal for the in-flight chat request. */ + signal: AbortSignal; + /** Public user/assistant transcript visible to the current chat turn. */ + messages: AIChatMessage[]; +} + +/** + * A tool executed locally inside AppShell. + * + * Local tools are validated with Standard Schema, executed in-process, and then + * written back into the internal AI Gateway transcript as `role: "tool"` + * messages for the next model round. + * Schema and executor are split so heterogeneous schemas compare structurally. + */ +export type AILocalTool = AIChatToolSchema> = { + kind: "local"; + description?: string; + schema: TSchema; +} & { + execute: ( + args: StandardSchemaV1.InferOutput, + context: AIChatToolContext, + ) => unknown | Promise; +}; + +/** + * Small typed helper for defining local AI chat tools. + * + * This is mostly an identity function, but it keeps `execute(args)` inferred + * from the provided schema when tools are declared inline. + * + * @example + * ```ts + * import { defineAIChatTool } from "@tailor-platform/app-shell"; + * import { z } from "zod/v4"; + * + * const lookupCustomer = defineAIChatTool({ + * description: "Look up a customer in the current workspace", + * schema: z.object({ + * customerId: z.string(), + * }), + * async execute({ customerId }, { signal }) { + * const customer = await fetchCustomer(customerId, { signal }); + * return { + * id: customer.id, + * name: customer.name, + * status: customer.status, + * }; + * }, + * }); + * ``` + */ +export function defineAIChatTool>(tool: { + description?: string; + schema: TSchema; + execute: ( + args: StandardSchemaV1.InferOutput, + context: AIChatToolContext, + ) => unknown | Promise; +}): AILocalTool { + return { + kind: "local", + ...tool, + }; +} diff --git a/packages/core/src/ai/use-ai-chat.test.tsx b/packages/core/src/ai/use-ai-chat.test.tsx index 54c96f22f..79ffff395 100644 --- a/packages/core/src/ai/use-ai-chat.test.tsx +++ b/packages/core/src/ai/use-ai-chat.test.tsx @@ -1,7 +1,9 @@ import { act, renderHook, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { z } from "zod/v4"; import { useAIChat } from "./use-ai-chat"; -import type { AIGatewayClient } from "./client"; +import type { AIGatewayChatRequest, AIGatewayClient } from "./client"; +import { defineAIChatTool } from "./tools"; function createAbortError(): Error { if (typeof DOMException !== "undefined") { @@ -31,11 +33,16 @@ describe("useAIChat", () => { const { result } = renderHook(() => useAIChat({ client, model: "gpt-5-mini" })); expect(result.current.messages).toEqual([]); + expect(result.current.toolCalls).toEqual([]); expect(result.current.status).toBe("ready"); expect(result.current.error).toBeUndefined(); }); it("appends user and assistant messages while streaming", async () => { + let releaseFirstChunk!: () => void; + const firstChunkGate = new Promise((resolve) => { + releaseFirstChunk = resolve; + }); let releaseSecondChunk!: () => void; const secondChunkGate = new Promise((resolve) => { releaseSecondChunk = resolve; @@ -44,6 +51,7 @@ describe("useAIChat", () => { const client = { streamChatCompletion: vi.fn(async function* ({ signal }) { expect(signal).toBeInstanceOf(AbortSignal); + await firstChunkGate; yield { type: "text-delta", text: "Hello" } as const; await secondChunkGate; yield { type: "text-delta", text: " world" } as const; @@ -58,6 +66,13 @@ describe("useAIChat", () => { sendPromise = result.current.sendMessage("Hi"); }); + expect(result.current.status).toBe("submitted"); + expect(result.current.messages).toEqual([{ id: "id-1", role: "user", content: "Hi" }]); + + await act(async () => { + releaseFirstChunk(); + }); + await waitFor(() => { expect(result.current.status).toBe("streaming"); expect(result.current.messages).toEqual([ @@ -250,6 +265,593 @@ describe("useAIChat", () => { }); }); + it.each(["stop", "error"] as const)( + "includes retained text and earlier tool history in the next request after %s", + async (interruption) => { + let releaseReply!: () => void; + const replyGate = new Promise((resolve) => { + releaseReply = resolve; + }); + const earlierToolCall = { id: "earlier-call", name: "lookup", argumentsText: "{}" }; + const earlierHistory = [ + { role: "user", content: "Earlier" }, + { role: "assistant", toolCalls: [earlierToolCall] }, + { role: "tool", toolCallId: "earlier-call", content: "found" }, + { role: "assistant", content: "Earlier answer" }, + ]; + const client = { + streamChatCompletion: vi.fn(async function* ({ messages }) { + const lastMessage = messages.at(-1); + if (lastMessage?.content === "Earlier") { + yield { + type: "tool-call", + toolCallId: earlierToolCall.id, + toolName: earlierToolCall.name, + argumentsText: "{}", + } as const; + } else if (lastMessage?.role === "tool") { + yield { type: "text-delta", text: "Earlier answer" } as const; + } else if (lastMessage?.content === "First") { + yield { type: "text-delta", text: "Partial" } as const; + await replyGate; + if (interruption === "error") throw new Error("reply failed"); + yield { type: "text-delta", text: " stale" } as const; + } else { + expect(messages).toEqual([ + ...earlierHistory, + { role: "user", content: "First" }, + { role: "assistant", content: "Partial" }, + { role: "user", content: "Continue" }, + ]); + yield { type: "text-delta", text: "Fresh" } as const; + } + yield { type: "done" } as const; + }), + } satisfies AIGatewayClient; + const lookup = defineAIChatTool({ schema: z.object({}), execute: () => "found" }); + const { result } = renderHook(() => + useAIChat({ client, model: "gpt-5-mini", tools: { lookup } }), + ); + await act(async () => { + await expect(result.current.sendMessage("Earlier")).resolves.toBe(true); + }); + + let interruptedSend!: Promise; + await act(async () => { + interruptedSend = result.current.sendMessage("First"); + }); + expect(result.current.messages.at(-1)?.content).toBe("Partial"); + + let nextSend!: Promise; + if (interruption === "stop") { + act(() => { + result.current.stop(); + nextSend = result.current.sendMessage("Continue"); + }); + } else { + await act(async () => { + releaseReply(); + await expect(interruptedSend).resolves.toBe(false); + }); + act(() => { + nextSend = result.current.sendMessage("Continue"); + }); + } + await act(async () => { + await expect(nextSend).resolves.toBe(true); + }); + if (interruption === "stop") { + await act(async () => { + releaseReply(); + await expect(interruptedSend).resolves.toBe(false); + }); + } + expect(result.current.status).toBe("ready"); + expect(result.current.messages.at(-1)?.content).toBe("Fresh"); + expect(result.current.error).toBeUndefined(); + }, + ); + + it("executes local tools and continues the conversation", async () => { + const lookupCustomer = defineAIChatTool({ + description: "Look up a customer", + schema: z.object({ + customerId: z.string(), + }), + async execute({ customerId }) { + return { customerId, name: "Acme Corp" }; + }, + }); + + const client = { + streamChatCompletion: vi.fn(async function* ({ messages, tools }) { + const lastMessage = messages.at(-1); + + if (lastMessage?.role === "user") { + expect(tools).toEqual([ + { + type: "function", + function: { + name: "lookupCustomer", + description: "Look up a customer", + parameters: { + $schema: "http://json-schema.org/draft-07/schema#", + type: "object", + properties: { + customerId: { type: "string" }, + }, + required: ["customerId"], + }, + }, + }, + ]); + + yield { + type: "tool-call", + toolCallId: "call-1", + toolName: "lookupCustomer", + argumentsText: '{"customerId":"cust-1"}', + } as const; + yield { type: "done", finishReason: "tool_calls" } as const; + return; + } + + expect(messages).toEqual([ + { role: "user", content: "Find customer" }, + { + role: "assistant", + toolCalls: [ + { + id: "call-1", + name: "lookupCustomer", + argumentsText: '{"customerId":"cust-1"}', + }, + ], + }, + { + role: "tool", + toolCallId: "call-1", + content: '{"customerId":"cust-1","name":"Acme Corp"}', + }, + ]); + + yield { type: "text-delta", text: "Customer: Acme Corp" } as const; + yield { type: "done", finishReason: "stop" } as const; + }), + } satisfies AIGatewayClient; + + const { result } = renderHook(() => + useAIChat({ + client, + model: "gpt-5-mini", + tools: { + lookupCustomer, + }, + }), + ); + + await act(async () => { + await expect(result.current.sendMessage("Find customer")).resolves.toBe(true); + }); + + await waitFor(() => { + expect(result.current.status).toBe("ready"); + expect(result.current.messages).toEqual([ + { id: "id-1", role: "user", content: "Find customer" }, + { id: "id-2", role: "assistant", content: "Customer: Acme Corp" }, + ]); + expect(result.current.toolCalls).toEqual([]); + }); + + expect(client.streamChatCompletion).toHaveBeenCalledTimes(2); + }); + + it("exposes calls only during the active flow across text and tool-only rounds", async () => { + let releaseTool!: () => void; + const toolGate = new Promise((resolve) => { + releaseTool = resolve; + }); + let releaseReply!: () => void; + const replyGate = new Promise((resolve) => { + releaseReply = resolve; + }); + let releaseFinalReply!: () => void; + const finalReplyGate = new Promise((resolve) => { + releaseFinalReply = resolve; + }); + const echo = defineAIChatTool({ + schema: z.object({ value: z.string() }), + async execute({ value }) { + if (value === "first") await toolGate; + return value; + }, + }); + let round = 0; + const client = { + streamChatCompletion: vi.fn(async function* () { + round += 1; + if (round === 1) { + yield { type: "text-delta", text: "Checking" } as const; + for (const value of ["first", "second"]) { + yield { + type: "tool-call", + toolCallId: value, + toolName: "echo", + argumentsText: JSON.stringify({ value }), + } as const; + } + } else if (round === 2) { + await replyGate; + yield { + type: "tool-call", + toolCallId: "third", + toolName: "echo", + argumentsText: '{"value":"third"}', + } as const; + } else { + await finalReplyGate; + yield { type: "text-delta", text: "Done" } as const; + } + yield { type: "done" } as const; + }), + } satisfies AIGatewayClient; + const { result } = renderHook(() => + useAIChat({ client, model: "gpt-5-mini", tools: { echo } }), + ); + + let sendPromise: Promise | undefined; + await act(async () => { + sendPromise = result.current.sendMessage("Check three values"); + }); + + const firstCalls = [ + { id: "first", name: "echo" }, + { id: "second", name: "echo" }, + ]; + expect(result.current.status).toBe("submitted"); + expect(client.streamChatCompletion).toHaveBeenCalledTimes(1); + expect(result.current.toolCalls).toEqual(firstCalls); + expect(result.current.messages).toEqual([ + { id: "id-1", role: "user", content: "Check three values" }, + { id: "id-2", role: "assistant", content: "Checking" }, + ]); + + await act(async () => { + await expect(result.current.sendMessage("Ignored")).resolves.toBe(false); + }); + expect(result.current.toolCalls).toEqual(firstCalls); + + await act(async () => { + releaseTool(); + }); + await waitFor(() => expect(client.streamChatCompletion).toHaveBeenCalledTimes(2)); + expect(result.current.status).toBe("submitted"); + expect(result.current.toolCalls).toEqual(firstCalls); + + await act(async () => { + releaseReply(); + }); + await waitFor(() => expect(client.streamChatCompletion).toHaveBeenCalledTimes(3)); + expect(result.current.toolCalls).toEqual([...firstCalls, { id: "third", name: "echo" }]); + expect(result.current.messages).toHaveLength(2); + expect(result.current.status).toBe("submitted"); + + await act(async () => { + releaseFinalReply(); + await expect(sendPromise).resolves.toBe(true); + }); + + expect(result.current.status).toBe("ready"); + expect(result.current.toolCalls).toEqual([]); + expect(result.current.messages).toEqual([ + { id: "id-1", role: "user", content: "Check three values" }, + { id: "id-2", role: "assistant", content: "Checking" }, + { id: "id-3", role: "assistant", content: "Done" }, + ]); + }); + + it.each([false, true])( + "drops the unfinished tool batch on stop (has text: %s)", + async (hasText) => { + let releaseTool!: () => void; + const toolGate = new Promise((resolve) => { + releaseTool = resolve; + }); + const executePending = vi.fn(async () => { + await toolGate; + return "late result"; + }); + const completedCall = { id: "completed", name: "lookup", argumentsText: "{}" }; + const lookup = defineAIChatTool({ schema: z.object({}), execute: () => "found" }); + const pending = defineAIChatTool({ schema: z.object({}), execute: executePending }); + const client = { + streamChatCompletion: vi.fn(async function* ({ messages }) { + const lastMessage = messages.at(-1); + if (lastMessage?.content === "First") { + yield { type: "text-delta", text: "Checking" } as const; + yield { + type: "tool-call", + toolCallId: completedCall.id, + toolName: completedCall.name, + argumentsText: "{}", + } as const; + } else if (lastMessage?.role === "tool") { + if (hasText) yield { type: "text-delta", text: "Checking again" } as const; + yield { + type: "tool-call", + toolCallId: "pending", + toolName: "pending", + argumentsText: "{}", + } as const; + } else { + expect(messages).toEqual([ + { role: "user", content: "First" }, + { role: "assistant", content: "Checking", toolCalls: [completedCall] }, + { role: "tool", toolCallId: "completed", content: "found" }, + ...(hasText ? [{ role: "assistant", content: "Checking again" }] : []), + { role: "user", content: "Continue" }, + ]); + yield { type: "text-delta", text: "Fresh" } as const; + } + yield { type: "done" } as const; + }), + } satisfies AIGatewayClient; + const { result } = renderHook(() => + useAIChat({ client, model: "gpt-5-mini", tools: { lookup, pending } }), + ); + let interruptedSend!: Promise; + await act(async () => { + interruptedSend = result.current.sendMessage("First"); + }); + expect(executePending).toHaveBeenCalledOnce(); + + let nextSend!: Promise; + act(() => { + result.current.stop(); + nextSend = result.current.sendMessage("Continue"); + }); + await act(async () => { + await expect(nextSend).resolves.toBe(true); + releaseTool(); + await expect(interruptedSend).resolves.toBe(false); + }); + expect(result.current.messages.map(({ content }) => content)).toEqual([ + "First", + "Checking", + ...(hasText ? ["Checking again"] : []), + "Continue", + "Fresh", + ]); + expect(result.current.toolCalls).toEqual([]); + expect(result.current.status).toBe("ready"); + expect(client.streamChatCompletion).toHaveBeenCalledTimes(3); + }, + ); + + it("clears calls on stop and ignores late events while the next flow is active", async () => { + let releaseStoppedRequest!: () => void; + const stoppedRequestGate = new Promise((resolve) => { + releaseStoppedRequest = resolve; + }); + let releaseSecondReply!: () => void; + const secondReplyGate = new Promise((resolve) => { + releaseSecondReply = resolve; + }); + const lookup = defineAIChatTool({ schema: z.object({}), execute: () => "found" }); + const client = { + streamChatCompletion: vi.fn(async function* ({ messages }: AIGatewayChatRequest) { + const lastMessage = messages.at(-1); + if (lastMessage?.role === "user") { + expect(messages).toEqual( + lastMessage.content === "First" + ? [{ role: "user", content: "First" }] + : [ + { role: "user", content: "First" }, + { + role: "assistant", + toolCalls: [{ id: "observed", name: "lookup", argumentsText: "{}" }], + }, + { role: "tool", toolCallId: "observed", content: "found" }, + { role: "user", content: "Second" }, + ], + ); + yield { + type: "tool-call", + toolCallId: lastMessage.content === "First" ? "observed" : "fresh", + toolName: "lookup", + argumentsText: "{}", + } as const; + } else if (messages.findLast((message) => message.role === "user")?.content === "First") { + await stoppedRequestGate; + yield { + type: "tool-call", + toolCallId: "late", + toolName: "lookup", + argumentsText: "{}", + } as const; + } else { + await secondReplyGate; + yield { type: "text-delta", text: "Fresh" } as const; + } + yield { type: "done" } as const; + }), + } satisfies AIGatewayClient; + const { result } = renderHook(() => + useAIChat({ client, model: "gpt-5-mini", tools: { lookup } }), + ); + + let firstSendPromise: Promise | undefined; + await act(async () => { + firstSendPromise = result.current.sendMessage("First"); + }); + expect(result.current.status).toBe("submitted"); + expect(result.current.toolCalls).toEqual([{ id: "observed", name: "lookup" }]); + expect(result.current.messages).toEqual([{ id: "id-1", role: "user", content: "First" }]); + + act(() => result.current.stop()); + expect(result.current.status).toBe("ready"); + expect(result.current.toolCalls).toEqual([]); + + let secondSendPromise: Promise | undefined; + await act(async () => { + secondSendPromise = result.current.sendMessage("Second"); + }); + expect(result.current.toolCalls).toEqual([{ id: "fresh", name: "lookup" }]); + expect(result.current.messages).toEqual([ + { id: "id-1", role: "user", content: "First" }, + { id: "id-2", role: "user", content: "Second" }, + ]); + + await act(async () => { + releaseStoppedRequest(); + await expect(firstSendPromise).resolves.toBe(false); + }); + expect(result.current.toolCalls).toEqual([{ id: "fresh", name: "lookup" }]); + expect(result.current.status).toBe("submitted"); + + await act(async () => { + releaseSecondReply(); + await expect(secondSendPromise).resolves.toBe(true); + }); + expect(result.current.toolCalls).toEqual([]); + expect(result.current.messages).toEqual([ + { id: "id-1", role: "user", content: "First" }, + { id: "id-2", role: "user", content: "Second" }, + { id: "id-3", role: "assistant", content: "Fresh" }, + ]); + expect(result.current.status).toBe("ready"); + expect(result.current.error).toBeUndefined(); + expect(client.streamChatCompletion).toHaveBeenCalledTimes(4); + }); + + it("clears active calls but retains completed tool history after a reply error", async () => { + let releaseFailure!: () => void; + const failureGate = new Promise((resolve) => { + releaseFailure = resolve; + }); + const failingTool = defineAIChatTool({ + schema: z.object({}), + async execute() { + await failureGate; + throw new Error("tool failed"); + }, + }); + const client = { + streamChatCompletion: vi.fn(async function* ({ messages }) { + if (messages.at(-1)?.content === "Retry") { + expect(messages).toEqual([ + { role: "user", content: "Try the tool" }, + { + role: "assistant", + toolCalls: [{ id: "failed-call", name: "failingTool", argumentsText: "{}" }], + }, + { role: "tool", toolCallId: "failed-call", content: '{"error":"tool failed"}' }, + { role: "user", content: "Retry" }, + ]); + yield { type: "text-delta", text: "Recovered" } as const; + } else if (messages.at(-1)?.role === "tool") { + expect(messages.at(-1)?.content).toBe('{"error":"tool failed"}'); + throw new Error("reply failed"); + } else { + yield { + type: "tool-call", + toolCallId: "failed-call", + toolName: "failingTool", + argumentsText: "{}", + } as const; + } + yield { type: "done" } as const; + }), + } satisfies AIGatewayClient; + const { result } = renderHook(() => + useAIChat({ client, model: "gpt-5-mini", tools: { failingTool } }), + ); + + let sendPromise: Promise | undefined; + await act(async () => { + sendPromise = result.current.sendMessage("Try the tool"); + }); + expect(result.current.toolCalls).toEqual([{ id: "failed-call", name: "failingTool" }]); + expect(result.current.messages).toEqual([ + { id: "id-1", role: "user", content: "Try the tool" }, + ]); + + await act(async () => { + releaseFailure(); + await expect(sendPromise).resolves.toBe(false); + }); + expect(result.current.status).toBe("error"); + expect(result.current.error?.message).toBe("reply failed"); + expect(result.current.toolCalls).toEqual([]); + expect(result.current.messages).toEqual([ + { id: "id-1", role: "user", content: "Try the tool" }, + ]); + + await act(async () => { + await expect(result.current.sendMessage("Retry")).resolves.toBe(true); + }); + expect(result.current.toolCalls).toEqual([]); + expect(result.current.messages).toEqual([ + { id: "id-1", role: "user", content: "Try the tool" }, + { id: "id-2", role: "user", content: "Retry" }, + { id: "id-3", role: "assistant", content: "Recovered" }, + ]); + }); + + it.each([false, true])( + "allows a final answer after eight tool batches but rejects a ninth batch (requests ninth: %s)", + async (requestNinthBatch) => { + const execute = vi.fn(({ value }: { value: string }) => ({ value })); + const loop = defineAIChatTool({ + schema: z.object({ value: z.string() }), + execute, + }); + let callCount = 0; + const client = { + streamChatCompletion: vi.fn(async function* ({ messages }) { + callCount += 1; + if (callCount === 9) { + expect(messages.at(-1)).toEqual({ + role: "tool", + toolCallId: "call-8", + content: '{"value":"again"}', + }); + if (!requestNinthBatch) { + yield { type: "text-delta", text: "Finished" } as const; + yield { type: "done" } as const; + return; + } + } + yield { + type: "tool-call", + toolCallId: `call-${callCount}`, + toolName: "loop", + argumentsText: '{"value":"again"}', + } as const; + yield { type: "done", finishReason: "tool_calls" } as const; + }), + } satisfies AIGatewayClient; + const { result } = renderHook(() => + useAIChat({ client, model: "gpt-5-mini", tools: { loop } }), + ); + + await act(async () => { + await expect(result.current.sendMessage("Start looping")).resolves.toBe(!requestNinthBatch); + }); + + expect(client.streamChatCompletion).toHaveBeenCalledTimes(9); + expect(execute).toHaveBeenCalledTimes(8); + expect(result.current.toolCalls).toEqual([]); + if (requestNinthBatch) { + expect(result.current.status).toBe("error"); + expect(result.current.error?.message).toContain("maximum number of tool rounds"); + } else { + expect(result.current.status).toBe("ready"); + expect(result.current.messages.at(-1)?.content).toBe("Finished"); + } + }, + ); + it("ignores blank and concurrent sends", async () => { let releaseRequest!: () => void; const requestGate = new Promise((resolve) => { diff --git a/packages/core/src/ai/use-ai-chat.ts b/packages/core/src/ai/use-ai-chat.ts index c551b89a0..4c7d475cb 100644 --- a/packages/core/src/ai/use-ai-chat.ts +++ b/packages/core/src/ai/use-ai-chat.ts @@ -1,7 +1,14 @@ import { useCallback, useEffect, useRef, useState } from "react"; import { isAbortError } from "./client"; import type { AIGatewayChatMessage, AIGatewayClient } from "./client"; - +import { + runAssistantLoop, + type AIChatAssistantTurnResult, + type AssistantLoopEvent, +} from "./assistant-loop"; +import type { AILocalTool } from "./tools"; + +/** Public chat message shape exposed by `useAIChat()`. */ export interface AIChatMessage { id: string; role: "user" | "assistant"; @@ -10,8 +17,84 @@ export interface AIChatMessage { export type AIChatStatus = "ready" | "submitted" | "streaming" | "error"; -export function useAIChat(config: { client: AIGatewayClient; model: string }): { +type ActiveToolCall = { id: string; name: string }; + +/** Active request handle stored in the ref for staleness checks + cancellation. */ +interface ActiveHandle { + readonly controller: AbortController; + readonly ctx: ChatRequestContext; +} + +/** + * Encapsulates per-request mutable state: transcript staging + streaming turn buffer. + * + * Success commits the full transcript; interruption preserves completed tool + * exchanges and visible text without retaining unanswered tool calls. + */ +class ChatRequestContext { + transcript: AIGatewayChatMessage[]; + private turnMessageId: string | null = null; + private turnContent = ""; + + constructor(initialTranscript: AIGatewayChatMessage[]) { + this.transcript = [...initialTranscript]; + } + + /** Appends a text delta to the current turn buffer, lazily creating a message id. */ + appendDelta(delta: string): { messageId: string; content: string; isNew: boolean } { + this.turnContent = `${this.turnContent}${delta}`; + const isNew = !this.turnMessageId; + if (isNew) { + this.turnMessageId = crypto.randomUUID(); + } + return { messageId: this.turnMessageId!, content: this.turnContent, isNew }; + } + + commitTurn(turn: AIChatAssistantTurnResult): void { + if (turn.gatewayAssistantMessage.content || turn.toolCalls.length > 0) { + this.transcript.push(turn.gatewayAssistantMessage); + } + this.turnMessageId = null; + this.turnContent = ""; + } + + commitToolResults(messages: Extract[]): void { + this.transcript.push(...messages); + } + + snapshotInterruptedTranscript(): AIGatewayChatMessage[] { + const transcript = [...this.transcript]; + const last = transcript.at(-1); + + // Tool results arrive as a complete batch, so only the tail can be unanswered. + if (last?.role === "assistant" && last.toolCalls?.length) { + transcript.pop(); + if (last.content) { + transcript.push({ role: "assistant", content: last.content }); + } + } + if (this.turnContent) { + transcript.push({ role: "assistant", content: this.turnContent }); + } + + return transcript; + } +} + +/** + * React hook for AI Gateway chat with optional local tools. + * + * Messages stay text-first. Requested tool names are exposed only for the + * active request; tool-call and tool-result protocol messages remain internal. + */ +export function useAIChat(config: { + client: AIGatewayClient; + model: string; + tools?: Record; +}): { messages: AIChatMessage[]; + /** Requested tools in the active flow; cleared on completion, stop, or error. */ + toolCalls: ActiveToolCall[]; status: AIChatStatus; error?: Error; /** @@ -22,51 +105,40 @@ export function useAIChat(config: { client: AIGatewayClient; model: string }): { stop: () => void; } { const [messages, setMessages] = useState([]); + const [toolCalls, setToolCalls] = useState([]); const [status, setStatus] = useState("ready"); const [error, setError] = useState(undefined); - const messagesRef = useRef([]); - const activeRequestRef = useRef(null); - const abortControllerRef = useRef(null); + const transcriptRef = useRef([]); + const activeRef = useRef(null); const updateMessages = useCallback((updater: (previous: AIChatMessage[]) => AIChatMessage[]) => { - setMessages((previous) => { - const next = updater(previous); - messagesRef.current = next; - return next; - }); + setMessages((previous) => updater(previous)); }, []); const stop = useCallback(() => { - const controller = abortControllerRef.current; - - if (!controller) { - return; - } + const active = activeRef.current; + if (!active) return; - activeRequestRef.current = null; - abortControllerRef.current = null; - controller.abort(); + transcriptRef.current = active.ctx.snapshotInterruptedTranscript(); + activeRef.current = null; + active.controller.abort(); + setToolCalls([]); setStatus("ready"); }, []); useEffect(() => { return () => { - activeRequestRef.current = null; - abortControllerRef.current?.abort(); - abortControllerRef.current = null; + activeRef.current?.controller.abort(); + activeRef.current = null; }; }, []); const sendMessage = useCallback( - async (message: string) => { - if (activeRequestRef.current) { - return false; - } + async (message: string): Promise => { + if (activeRef.current) return false; const text = message.trim(); - if (!text) { - return false; - } + if (!text) return false; const userMessage: AIChatMessage = { id: crypto.randomUUID(), @@ -74,67 +146,44 @@ export function useAIChat(config: { client: AIGatewayClient; model: string }): { content: text, }; - const nextMessages = [...messagesRef.current, userMessage]; - messagesRef.current = nextMessages; - setMessages(nextMessages); + updateMessages((previous) => [...previous, userMessage]); + setToolCalls([]); setError(undefined); setStatus("submitted"); + const ctx = new ChatRequestContext([ + ...transcriptRef.current, + { role: "user", content: text }, + ]); const controller = new AbortController(); - const requestId = Symbol(); - activeRequestRef.current = requestId; - abortControllerRef.current = controller; - let assistantMessageId: string | null = null; - const isActive = () => activeRequestRef.current === requestId; + const handle: ActiveHandle = { controller, ctx }; + activeRef.current = handle; try { - for await (const event of config.client.streamChatCompletion({ + for await (const event of runAssistantLoop({ + client: config.client, model: config.model, - messages: nextMessages.map(toGatewayMessage), + tools: config.tools, + transcript: ctx.transcript, signal: controller.signal, })) { - if (!isActive()) { - return false; - } - - if (event.type !== "text-delta" || !event.text) { - continue; - } - - setStatus("streaming"); - - if (!assistantMessageId) { - assistantMessageId = crypto.randomUUID(); - updateMessages((previous) => [ - ...previous, - { - id: assistantMessageId!, - role: "assistant", - content: event.text, - }, - ]); - continue; - } - - updateMessages((previous) => - previous.map((entry) => - entry.id === assistantMessageId - ? { ...entry, content: `${entry.content}${event.text}` } - : entry, - ), - ); + if (activeRef.current !== handle) break; + applyEvent(event, ctx, updateMessages, setStatus, setToolCalls); } - if (!isActive()) { - return false; - } + if (activeRef.current !== handle) return false; + // Commit staged transcript to persistent history + transcriptRef.current = ctx.transcript; + activeRef.current = null; + setToolCalls([]); setStatus("ready"); return true; } catch (caughtError) { - if (!isActive()) { - return false; - } + if (activeRef.current !== handle) return false; + transcriptRef.current = ctx.snapshotInterruptedTranscript(); + activeRef.current = null; + setToolCalls([]); if (isAbortError(caughtError)) { setStatus("ready"); @@ -144,33 +193,60 @@ export function useAIChat(config: { client: AIGatewayClient; model: string }): { setError(toError(caughtError)); setStatus("error"); return false; - } finally { - if (activeRequestRef.current === requestId) { - activeRequestRef.current = null; - } - - if (abortControllerRef.current === controller) { - abortControllerRef.current = null; - } } }, - [config.client, config.model, updateMessages], + [config.client, config.model, config.tools, updateMessages], ); - return { - messages, - status, - error, - sendMessage, - stop, - }; + return { messages, toolCalls, status, error, sendMessage, stop }; } -function toGatewayMessage(message: AIChatMessage): AIGatewayChatMessage { - return { - role: message.role, - content: message.content, - }; +/** Single dispatch point: loop events → React state updates + context mutations. */ +function applyEvent( + event: AssistantLoopEvent, + ctx: ChatRequestContext, + updateMessages: (updater: (previous: AIChatMessage[]) => AIChatMessage[]) => void, + setStatus: (status: AIChatStatus) => void, + setToolCalls: (updater: (previous: ActiveToolCall[]) => ActiveToolCall[]) => void, +): void { + switch (event.type) { + case "text-delta": { + const { messageId, content, isNew } = ctx.appendDelta(event.delta); + setStatus("streaming"); + + if (isNew) { + updateMessages((previous) => [...previous, { id: messageId, role: "assistant", content }]); + } else { + updateMessages((previous) => + previous.map((entry) => (entry.id === messageId ? { ...entry, content } : entry)), + ); + } + break; + } + + case "turn-end": { + if (event.turn.toolCalls.length > 0) { + const calls = event.turn.toolCalls.map(({ id, name }) => ({ id, name })); + setToolCalls((previous) => [...previous, ...calls]); + } + ctx.commitTurn(event.turn); + break; + } + + case "tool-resolution-start": { + setStatus("submitted"); + break; + } + + case "tool-results": { + ctx.commitToolResults(event.messages); + break; + } + + case "complete": { + break; + } + } } function toError(error: unknown): Error { diff --git a/packages/core/src/components/ai-chat/ai-chat.test-d.ts b/packages/core/src/components/ai-chat/ai-chat.test-d.ts index 2bb76bbc2..4725e8438 100644 --- a/packages/core/src/components/ai-chat/ai-chat.test-d.ts +++ b/packages/core/src/components/ai-chat/ai-chat.test-d.ts @@ -19,6 +19,7 @@ import { type AIChatHeaderProps, type AIChatProps, } from "./ai-chat"; +import type { AIChatToolCallProps, useAIChat } from "../../index"; describe("region ownership", () => { it("keeps the root to status (plus div props) — the regions own the rest", () => { @@ -58,10 +59,18 @@ describe("callback names that collide with DOM handlers", () => { }); }); -// These three wrap a Button or a Badge. They Pick what the part's job needs +// These parts wrap a Button or a Badge. They expose what the part's job needs // rather than inheriting, so the wrapped component's surface cannot leak out // and become contract we did not intend to offer. describe("wrapper parts expose only picked props", () => { + it("accepts a call from useAIChat directly with only call and className", () => { + expectTypeOf>().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf<"call" | "className">(); + expectTypeOf().toEqualTypeOf< + ReturnType["toolCalls"][number] + >(); + }); + it("keeps Suggestion to its own props plus className/disabled/children", () => { expectTypeOf>().toEqualTypeOf< "className" | "disabled" | "children" | "suggestion" | "onSelect" diff --git a/packages/core/src/components/ai-chat/ai-chat.test.tsx b/packages/core/src/components/ai-chat/ai-chat.test.tsx index f7a854e25..808c9cd25 100644 --- a/packages/core/src/components/ai-chat/ai-chat.test.tsx +++ b/packages/core/src/components/ai-chat/ai-chat.test.tsx @@ -77,6 +77,18 @@ describe("AIChat", () => { }); }); + it("shows the submitted state", () => { + const { container } = render( + + + Hello + + + , + ); + expect(container.innerHTML).toMatchSnapshot(); + }); + describe("region placement", () => { it("renders Header, Conversation, Composer in that order regardless of source order", () => { const { container } = render( diff --git a/packages/core/src/components/ai-chat/ai-chat.tsx b/packages/core/src/components/ai-chat/ai-chat.tsx index c280def42..c91398702 100644 --- a/packages/core/src/components/ai-chat/ai-chat.tsx +++ b/packages/core/src/components/ai-chat/ai-chat.tsx @@ -30,6 +30,7 @@ import { Response } from "./response"; import { Source, Sources, SourcesContent, SourcesTrigger } from "./sources"; import { Suggestion, Suggestions } from "./suggestion"; import { Tool, ToolContent, ToolHeader, ToolInput, ToolOutput } from "./tool"; +import { ToolCall } from "./tool-call"; type AIChatProps = Omit, "children"> & { /** @@ -177,6 +178,7 @@ type AIChatComponent = typeof AIChatRoot & { ChainOfThoughtSearchResults: typeof ChainOfThoughtSearchResults; ChainOfThoughtSearchResult: typeof ChainOfThoughtSearchResult; Tool: typeof Tool; + ToolCall: typeof ToolCall; ToolHeader: typeof ToolHeader; ToolContent: typeof ToolContent; ToolInput: typeof ToolInput; @@ -209,6 +211,7 @@ const AIChat: AIChatComponent = Object.assign(AIChatRoot, { ChainOfThoughtSearchResults, ChainOfThoughtSearchResult, Tool, + ToolCall, ToolHeader, ToolContent, ToolInput, @@ -225,5 +228,6 @@ export type { AIChatHeaderProps } from "./header"; export type { AIChatConversationProps } from "./conversation"; export type { AIChatComposerProps } from "./composer"; export type { ToolState } from "./tool"; +export type { AIChatToolCallProps } from "./tool-call"; export type { ChainOfThoughtStepStatus } from "./chain-of-thought"; export type { ChatHistoryGroupData, ChatHistoryItemData } from "./chat-history"; diff --git a/packages/core/src/components/ai-chat/conversation.test.tsx b/packages/core/src/components/ai-chat/conversation.test.tsx index 1c33b389e..489c5d30c 100644 --- a/packages/core/src/components/ai-chat/conversation.test.tsx +++ b/packages/core/src/components/ai-chat/conversation.test.tsx @@ -3,6 +3,8 @@ import { cleanup, fireEvent, render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { Conversation, ConversationScrollButton } from "./conversation"; +import { AIChatContext } from "./ai-chat-context"; +import type { AIChatStatus } from "@/ai/use-ai-chat"; afterEach(() => { cleanup(); @@ -31,7 +33,31 @@ function renderConversation() { const scrollButton = () => screen.queryByRole("button", { name: "Scroll to latest message" }); +const conversationWithStatus = (status: AIChatStatus) => ( + + content + +); + describe("Conversation", () => { + it("shows the waiting indicator only while submitted, including between tool rounds", () => { + const { rerender } = render(conversationWithStatus("submitted")); + + expect(screen.getByRole("status").textContent).toBe("Thinking…"); + expect(screen.getByRole("status").querySelector('[data-slot="spinner"]')).not.toBeNull(); + + for (const status of ["streaming", "submitted", "ready", "submitted", "error"] as const) { + rerender(conversationWithStatus(status)); + expect(screen.queryByRole("status") !== null).toBe(status === "submitted"); + } + }); + + it("remains usable outside AIChat without a waiting indicator", () => { + renderConversation(); + expect(screen.getByRole("log").textContent).toBe("content"); + expect(screen.queryByRole("status")).toBeNull(); + }); + it("hides the scroll button while the reader is at the bottom", () => { renderConversation(); expect(scrollButton()).toBeNull(); diff --git a/packages/core/src/components/ai-chat/conversation.tsx b/packages/core/src/components/ai-chat/conversation.tsx index a01c30037..7239af6e5 100644 --- a/packages/core/src/components/ai-chat/conversation.tsx +++ b/packages/core/src/components/ai-chat/conversation.tsx @@ -11,8 +11,11 @@ import { } from "react"; import { Button } from "@/components/button"; +import { Spinner } from "@/components/spinner"; import { useT } from "@/i18n-labels"; import { cn } from "@/lib/utils"; +import { AIChatContext } from "./ai-chat-context"; +import { Message } from "./message"; import { Text } from "../text"; // How close (px) to the bottom still counts as "pinned". @@ -58,6 +61,8 @@ function Conversation({ children, ...props }: AIChatConversationProps) { + const status = useContext(AIChatContext)?.status; + const t = useT(); const viewportRef = useRef(null); const contentRef = useRef(null); const pinnedRef = useRef(true); @@ -124,7 +129,20 @@ function Conversation({ against the relative root above, so it stays pinned while content scrolls underneath it. */}
- {children} + + {children} + {status === "submitted" ? ( + + + + {t("aiChatThinking")} + + + ) : null} +
diff --git a/packages/core/src/components/ai-chat/index.ts b/packages/core/src/components/ai-chat/index.ts index 21a273493..8da82b9d8 100644 --- a/packages/core/src/components/ai-chat/index.ts +++ b/packages/core/src/components/ai-chat/index.ts @@ -4,6 +4,7 @@ export { type AIChatHeaderProps, type AIChatConversationProps, type AIChatComposerProps, + type AIChatToolCallProps, type ToolState, type ChainOfThoughtStepStatus, type ChatHistoryGroupData, diff --git a/packages/core/src/components/ai-chat/tool-call.test.tsx b/packages/core/src/components/ai-chat/tool-call.test.tsx new file mode 100644 index 000000000..282dc84b8 --- /dev/null +++ b/packages/core/src/components/ai-chat/tool-call.test.tsx @@ -0,0 +1,34 @@ +import { afterEach, describe, expect, it } from "vitest"; +import { cleanup, render, screen } from "@testing-library/react"; + +import { AIChat } from "./index"; + +afterEach(() => { + cleanup(); +}); + +describe("AIChat.ToolCall", () => { + it("renders a neutral tool-call label without implying success", () => { + const { container } = render(); + + expect(screen.getByText("Tool call: add")).toBeDefined(); + expect(screen.queryByText("Done")).toBeNull(); + expect(screen.queryByRole("button")).toBeNull(); + expect(container.querySelector("svg")?.getAttribute("aria-hidden")).toBe("true"); + expect(container.innerHTML).toMatchSnapshot(); + }); + + it("renders multiple calls as separate rows in order and forwards className to the row", () => { + const { container } = render( + <> + + + , + ); + const rows = Array.from(container.querySelectorAll('[data-slot="ai-chat-tool-call"]')); + + expect(rows.map((row) => row.textContent)).toEqual(["Tool call: first", "Tool call: second"]); + expect(rows.every((row) => row.classList.contains("astw:flex"))).toBe(true); + expect(rows[0]?.classList.contains("consumer-class")).toBe(true); + }); +}); diff --git a/packages/core/src/components/ai-chat/tool-call.tsx b/packages/core/src/components/ai-chat/tool-call.tsx new file mode 100644 index 000000000..1b93f7ff5 --- /dev/null +++ b/packages/core/src/components/ai-chat/tool-call.tsx @@ -0,0 +1,28 @@ +import { Wrench } from "lucide-react"; + +import type { useAIChat } from "@/ai/use-ai-chat"; +import { Badge } from "@/components/badge"; +import { useT } from "@/i18n-labels"; +import { cn } from "@/lib/utils"; + +type AIChatToolCallProps = { + /** A requested tool from useAIChat().toolCalls, not an execution success receipt. */ + call: ReturnType["toolCalls"][number]; + className?: string; +}; + +/** Compact tool-call label for the active flow, rendered as its own row. */ +function ToolCall({ call, className }: AIChatToolCallProps) { + const t = useT(); + + return ( +
+ + + {t("aiChatToolCall", { name: call.name })} + +
+ ); +} + +export { ToolCall, type AIChatToolCallProps }; diff --git a/packages/core/src/i18n-labels.ts b/packages/core/src/i18n-labels.ts index 902fe2128..40519f369 100644 --- a/packages/core/src/i18n-labels.ts +++ b/packages/core/src/i18n-labels.ts @@ -43,6 +43,7 @@ export const labels = defineI18nLabels({ aiChatThoughtFor: (props: { seconds: number }) => `Thought for ${props.seconds}s`, aiChatReasoning: "Reasoning", aiChatChainOfThought: "Chain of thought", + aiChatToolCall: (props: { name: string }) => `Tool call: ${props.name}`, aiChatToolPreparing: "Preparing", aiChatToolRunning: "Running", aiChatToolDone: "Done", @@ -94,6 +95,7 @@ export const labels = defineI18nLabels({ aiChatThoughtFor: (props: { seconds: number }) => `${props.seconds}秒 思考しました`, aiChatReasoning: "思考", aiChatChainOfThought: "思考の流れ", + aiChatToolCall: (props: { name: string }) => `ツール呼び出し: ${props.name}`, aiChatToolPreparing: "準備中", aiChatToolRunning: "実行中", aiChatToolDone: "完了", diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 7abbb0f7e..ed65dfd63 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -72,15 +72,24 @@ export { type AIGatewayClient, type AIGatewayChatMessage, type AIGatewayChatRequest, + type AIGatewayToolCall, + type AIGatewayFunctionTool, type AIChatCompletionEvent, } from "./ai/client"; export { useAIChat, type AIChatMessage, type AIChatStatus } from "./ai/use-ai-chat"; +export { + defineAIChatTool, + type AIChatToolContext, + type AIChatToolSchema, + type AILocalTool, +} from "./ai/tools"; export { AIChat, type AIChatProps, type AIChatHeaderProps, type AIChatConversationProps, type AIChatComposerProps, + type AIChatToolCallProps, type ToolState, type ChainOfThoughtStepStatus, type ChatHistoryGroupData, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 821219830..04fc3fbe6 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -511,6 +511,9 @@ importers: vitest: specifier: 'catalog:' version: 4.1.11(@opentelemetry/api@1.9.1)(@types/node@25.9.8)(@vitest/browser-playwright@4.1.11)(@vitest/coverage-v8@4.1.11)(happy-dom@21.0.4)(vite@8.2.2(@types/node@25.9.8)(esbuild@0.28.1)(jiti@2.7.0)(tsx@4.23.1)(yaml@2.9.0)) + zod: + specifier: 4.4.3 + version: 4.4.3 packages/docs-kit: dependencies: