diff --git a/docs/chat/streaming.md b/docs/chat/streaming.md index 2c799a772..0e1401d9e 100644 --- a/docs/chat/streaming.md +++ b/docs/chat/streaming.md @@ -78,6 +78,71 @@ TanStack AI implements the [AG-UI Protocol](https://docs.ag-ui.com/introduction) > **Tip:** Some models expose their internal reasoning as thinking content that streams before the response. See [Thinking & Reasoning](./thinking-content). +### Type-Safe Tool Call Events + +When you pass typed tools (defined with `toolDefinition()` and Zod schemas) to `chat()`, the stream chunks automatically carry type information for tool call events. The `toolName` field narrows to the union of your tool name literals, and the `input` field on `TOOL_CALL_END` events is typed as the union of your tool input schemas: + +```typescript +import { chat, toolDefinition } from "@tanstack/ai"; +import { openaiText } from "@tanstack/ai-openai"; +import { z } from "zod"; + +const weatherTool = toolDefinition({ + name: "get_weather", + description: "Get weather for a location", + inputSchema: z.object({ + location: z.string(), + unit: z.enum(["celsius", "fahrenheit"]).optional(), + }), +}); + +const stream = chat({ + adapter: openaiText("gpt-5.2"), + messages, + tools: [weatherTool], +}); + +for await (const chunk of stream) { + if (chunk.type === "TOOL_CALL_END") { + chunk.toolName; // ✅ typed as "get_weather" (not string) + chunk.input; // ✅ typed as { location: string; unit?: "celsius" | "fahrenheit" } | undefined + } +} +``` + +Without typed tools, `toolName` defaults to `string` and `input` defaults to `unknown` — the same behavior as before. The type narrowing is automatic when you use `toolDefinition()` with Zod schemas. + +When multiple tools are provided, tool call events form a **discriminated union** — checking `toolName` narrows `input` to that specific tool's type: + +```typescript +const searchTool = toolDefinition({ + name: "search", + description: "Search the web", + inputSchema: z.object({ query: z.string() }), +}); + +const stream = chat({ + adapter: openaiText("gpt-5.2"), + messages, + tools: [weatherTool, searchTool], +}); + +for await (const chunk of stream) { + if (chunk.type === "TOOL_CALL_END") { + if (chunk.toolName === "get_weather") { + // ✅ input is narrowed to { location: string; unit?: "celsius" | "fahrenheit" } + console.log(`Weather in ${chunk.input?.location}`); + } + if (chunk.toolName === "search") { + // ✅ input is narrowed to { query: string } + console.log(`Searched for: ${chunk.input?.query}`); + } + } +} +``` + +> **Tip:** The typed stream chunk type is exported as `TypedStreamChunk` if you need to annotate variables or function parameters. When used without type arguments, `TypedStreamChunk` is equivalent to `StreamChunk`. + ### Thinking Chunks Thinking/reasoning is represented by AG-UI events `STEP_STARTED` and `STEP_FINISHED`. They stream separately from the final response text: diff --git a/docs/reference/type-aliases/StreamChunk.md b/docs/reference/type-aliases/StreamChunk.md index 4c0fb5cdb..f3de1b3c2 100644 --- a/docs/reference/type-aliases/StreamChunk.md +++ b/docs/reference/type-aliases/StreamChunk.md @@ -9,7 +9,42 @@ title: StreamChunk type StreamChunk = AGUIEvent; ``` -Defined in: [types.ts:976](https://github.com/TanStack/ai/blob/main/packages/typescript/ai/src/types.ts#L976) +Defined in: [types.ts:990](https://github.com/TanStack/ai/blob/main/packages/typescript/ai/src/types.ts#L990) Chunk returned by the SDK during streaming chat completions. Uses the AG-UI protocol event format. + +# Type Alias: TypedStreamChunk + +```ts +type TypedStreamChunk> = ReadonlyArray>> +``` + +Defined in: [types.ts:1065](https://github.com/TanStack/ai/blob/main/packages/typescript/ai/src/types.ts#L1065) + +A variant of `StreamChunk` parameterized by the tools array. When specific tool types are provided (e.g. from `chat({ tools: [myTool] })`): + +- `TOOL_CALL_START` and `TOOL_CALL_END` events form a **discriminated union** over tool names. +- Checking `toolName === 'x'` narrows `input` to that specific tool's input type. +- `TOOL_CALL_END` events have `input` typed per-tool via Standard Schema inference. + +When tools are untyped or absent, `TypedStreamChunk` degrades to the same type as `StreamChunk`. + +This is the type returned by `chat()` when streaming is enabled (the default). You don't typically need to reference it directly unless annotating function parameters or return types. + +```ts +import type { TypedStreamChunk } from "@tanstack/ai"; +import { toolDefinition } from "@tanstack/ai"; + +// Given tools created with toolDefinition(): +const weatherTool = toolDefinition({ name: "get_weather", description: "...", inputSchema: /* Zod schema */ }); +const searchTool = toolDefinition({ name: "search", description: "...", inputSchema: /* Zod schema */ }); + +// Without type args — equivalent to StreamChunk +type Chunk = TypedStreamChunk; + +// With specific tools — tool call events are typed +type TypedChunk = TypedStreamChunk<[typeof weatherTool, typeof searchTool]>; +``` + +See [Streaming - Type-Safe Tool Call Events](../../chat/streaming) for a practical walkthrough. diff --git a/docs/tools/server-tools.md b/docs/tools/server-tools.md index bcae69ecf..5f07ac2b0 100644 --- a/docs/tools/server-tools.md +++ b/docs/tools/server-tools.md @@ -299,6 +299,8 @@ const getUserData = getUserDataDef.server(async (args) => { > **Note:** JSON Schema tools skip runtime validation. Zod schemas are recommended for full type safety and validation. +> **Tip:** When you pass typed tools (server, client, or definition) to `chat()`, the returned stream is fully typed — `toolName` narrows to your tool name literals and `input` narrows per-tool when you check the name. See [Type-Safe Tool Call Events](../chat/streaming#type-safe-tool-call-events). + ## Best Practices 1. **Keep tools focused** - Each tool should do one thing well diff --git a/docs/tools/tools.md b/docs/tools/tools.md index c0a651a95..ac18636c2 100644 --- a/docs/tools/tools.md +++ b/docs/tools/tools.md @@ -78,6 +78,8 @@ const inputSchema: JSONSchema = { > **Note:** When using JSON Schema, TypeScript will infer `any` for input/output types since JSON Schema cannot provide compile-time type information. Zod schemas are recommended for full type safety. +> **Tip:** Type safety from Zod schemas extends beyond tool execution — when you iterate over the stream returned by `chat()`, tool call events have typed `toolName` and `input` fields too. See [Type-Safe Tool Call Events](../chat/streaming#type-safe-tool-call-events). + ## Tool Definition Tools are defined using `toolDefinition()` from `@tanstack/ai`: diff --git a/examples/ts-react-chat/src/routes/api.tanchat.ts b/examples/ts-react-chat/src/routes/api.tanchat.ts index a1eb8ee02..2649cde38 100644 --- a/examples/ts-react-chat/src/routes/api.tanchat.ts +++ b/examples/ts-react-chat/src/routes/api.tanchat.ts @@ -108,6 +108,72 @@ const loggingMiddleware: ChatMiddleware = { }, } +// =========================== +// TypedStreamChunk showcase — type-safe tool call events +// =========================== +// +// When `chat()` receives tools with typed schemas, the returned stream +// carries type information on TOOL_CALL_START and TOOL_CALL_END events. +// No casts, no `as any` — just narrow by `chunk.type` and everything is typed. + +const tools = [ + getGuitars, + recommendGuitarToolDef, + addToCartToolServer, + addToWishListToolDef, + getPersonalGuitarPreferenceToolDef, + compareGuitars, + calculateFinancing, + searchGuitars, +] as const + +async function typedStreamShowcase() { + const stream = chat({ + adapter: openaiText('gpt-4o'), + messages: [ + { role: 'user' as const, content: 'Recommend an acoustic guitar' }, + ], + tools, + }) + + for await (const chunk of stream) { + switch (chunk.type) { + case 'TOOL_CALL_START': + // ✅ chunk.toolName is typed as the union of all tool name literals: + // 'getGuitars' | 'recommendGuitar' | 'addToCart' | 'addToWishList' + // | 'getPersonalGuitarPreference' | 'compareGuitars' + // | 'calculateFinancing' | 'searchGuitars' + // + // ❌ Without TypedStreamChunk, this would just be `string` + console.log(`Tool call started: ${chunk.toolName}`) + break + + case 'TOOL_CALL_END': + // ✅ Discriminated union — checking toolName narrows input to that tool's type + if (chunk.toolName === 'searchGuitars') { + // ✅ chunk.input is { query: string } (not the full union) + console.log(`Searching for: ${chunk.input?.query}`) + } else if (chunk.toolName === 'calculateFinancing') { + // ✅ chunk.input is { guitarId: number; months: number } + console.log( + `Financing guitar ${chunk.input?.guitarId} for ${chunk.input?.months} months`, + ) + } else { + console.log(`Tool call ended: ${chunk.toolName}`, chunk.input) + } + break + + case 'TEXT_MESSAGE_CONTENT': + // Non-tool events are unaffected — still fully typed + console.log(chunk.delta) + break + } + } +} + +// Suppress unused warning — this is a showcase, not called at runtime +void typedStreamShowcase + export const Route = createFileRoute('/api/tanchat')({ server: { handlers: { diff --git a/packages/typescript/ai/src/activities/chat/index.ts b/packages/typescript/ai/src/activities/chat/index.ts index 24cc41529..18932fea2 100644 --- a/packages/typescript/ai/src/activities/chat/index.ts +++ b/packages/typescript/ai/src/activities/chat/index.ts @@ -43,6 +43,7 @@ import type { ToolCallArgsEvent, ToolCallEndEvent, ToolCallStartEvent, + TypedStreamChunk, } from '../../types' import type { ChatMiddleware, @@ -69,11 +70,15 @@ export const kind = 'text' as const * @template TAdapter - The text adapter type (created by a provider function) * @template TSchema - Optional Standard Schema for structured output * @template TStream - Whether to stream the output (default: true) + * @template TTools - The tools array type for type-safe tool call events in the stream */ export interface TextActivityOptions< TAdapter extends AnyTextAdapter, TSchema extends SchemaInput | undefined, TStream extends boolean, + TTools extends ReadonlyArray> = ReadonlyArray< + Tool + >, > { /** The text adapter to use (created by a provider function like openaiText('gpt-4o')) */ adapter: TAdapter @@ -87,7 +92,7 @@ export interface TextActivityOptions< /** System prompts to prepend to the conversation */ systemPrompts?: TextOptions['systemPrompts'] /** Tools for function calling (auto-executed when called) */ - tools?: TextOptions['tools'] + tools?: TTools /** Controls the randomness of the output. Higher values make output more random. Range: [0.0, 2.0] */ temperature?: TextOptions['temperature'] /** Nucleus sampling parameter. The model considers tokens with topP probability mass. */ @@ -125,7 +130,7 @@ export interface TextActivityOptions< outputSchema?: TSchema /** * Whether to stream the text result. - * When true (default), returns an AsyncIterable for streaming output. + * When true (default), returns an AsyncIterable> for streaming output. * When false, returns a Promise with the collected text content. * * Note: If outputSchema is provided, this option is ignored and the result @@ -186,9 +191,12 @@ export function createChatOptions< TAdapter extends AnyTextAdapter, TSchema extends SchemaInput | undefined = undefined, TStream extends boolean = true, + TTools extends ReadonlyArray> = ReadonlyArray< + Tool + >, >( - options: TextActivityOptions, -): TextActivityOptions { + options: TextActivityOptions, +): TextActivityOptions { return options } @@ -200,16 +208,22 @@ export function createChatOptions< * Result type for the text activity. * - If outputSchema is provided: Promise> * - If stream is false: Promise - * - Otherwise (stream is true, default): AsyncIterable + * - Otherwise (stream is true, default): AsyncIterable> + * + * When tools with typed schemas are provided, the stream chunks include + * type-safe `toolName` and `input` fields on tool call events. */ export type TextActivityResult< TSchema extends SchemaInput | undefined, TStream extends boolean = true, + TTools extends ReadonlyArray> = ReadonlyArray< + Tool + >, > = TSchema extends SchemaInput ? Promise> : TStream extends false ? Promise - : AsyncIterable + : AsyncIterable> // =========================== // ChatEngine Implementation @@ -1374,9 +1388,12 @@ export function chat< TAdapter extends AnyTextAdapter, TSchema extends SchemaInput | undefined = undefined, TStream extends boolean = true, + TTools extends ReadonlyArray> = ReadonlyArray< + Tool + >, >( - options: TextActivityOptions, -): TextActivityResult { + options: TextActivityOptions, +): TextActivityResult { const { outputSchema, stream } = options // If outputSchema is provided, run agentic structured output @@ -1387,7 +1404,7 @@ export function chat< SchemaInput, boolean >, - ) as TextActivityResult + ) as TextActivityResult } // If stream is explicitly false, run non-streaming text @@ -1398,13 +1415,13 @@ export function chat< undefined, false >, - ) as TextActivityResult + ) as TextActivityResult } // Otherwise, run streaming text (default) return runStreamingText( options as unknown as TextActivityOptions, - ) as TextActivityResult + ) as TextActivityResult } /** diff --git a/packages/typescript/ai/src/types.ts b/packages/typescript/ai/src/types.ts index 984e15125..ed69050fb 100644 --- a/packages/typescript/ai/src/types.ts +++ b/packages/typescript/ai/src/types.ts @@ -840,13 +840,19 @@ export interface TextMessageEndEvent extends BaseAGUIEvent { /** * Emitted when a tool call starts. + * + * @typeParam TToolName - Constrained tool name type. Defaults to `string` (untyped). + * When the stream is returned from `chat()` with typed tools, this narrows to + * the union of tool name literals. */ -export interface ToolCallStartEvent extends BaseAGUIEvent { +export interface ToolCallStartEvent< + TToolName extends string = string, +> extends BaseAGUIEvent { type: 'TOOL_CALL_START' /** Unique identifier for this tool call */ toolCallId: string /** Name of the tool being called */ - toolName: string + toolName: TToolName /** ID of the parent message that initiated this tool call */ parentMessageId?: string /** Index for parallel tool calls */ @@ -870,15 +876,23 @@ export interface ToolCallArgsEvent extends BaseAGUIEvent { /** * Emitted when a tool call completes. + * + * @typeParam TToolName - Constrained tool name type. Defaults to `string` (untyped). + * @typeParam TInput - Constrained input arguments type. Defaults to `unknown` (untyped). + * When the stream is returned from `chat()` with typed tools, these narrow to + * the union of tool name literals and the union of tool input types respectively. */ -export interface ToolCallEndEvent extends BaseAGUIEvent { +export interface ToolCallEndEvent< + TToolName extends string = string, + TInput = unknown, +> extends BaseAGUIEvent { type: 'TOOL_CALL_END' /** Tool call identifier */ toolCallId: string /** Name of the tool */ - toolName: string + toolName: TToolName /** Final parsed input arguments */ - input?: unknown + input?: TInput /** Tool execution result (if executed) */ result?: string } @@ -975,6 +989,94 @@ export type AGUIEvent = */ export type StreamChunk = AGUIEvent +// ============================================================================ +// Typed Stream Chunks (tool-aware) +// ============================================================================ + +/** + * Detect the `any` type. Returns `true` for `any`, `false` for everything else. + * @internal + */ +type IsAny = 0 extends 1 & T ? true : false + +/** + * Check whether the tools array carries typed tool definitions. + * Returns `false` for empty arrays or arrays with generic `string` names. + * @internal + */ +type HasTypedTools>> = [ + TTools[number], +] extends [never] + ? false + : string extends TTools[number]['name'] + ? false + : true + +/** + * Safely infer input type for a single tool, guarding against `any` leaks. + * Returns `unknown` when the tool has no inputSchema or when InferSchemaType + * produces `any` (e.g. for plain JSON Schema tools). + * @internal + */ +type SafeToolInput = T extends { + inputSchema?: infer TInput +} + ? IsAny>> extends true + ? unknown + : InferSchemaType> + : unknown + +/** + * Distribute over each tool to create a per-tool `ToolCallStartEvent`. + * This produces a discriminated union — one variant per tool name literal. + * @internal + */ +type DistributedToolCallStart< + TTools extends ReadonlyArray>, +> = TTools[number] extends infer T + ? T extends Tool + ? ToolCallStartEvent + : never + : never + +/** + * Distribute over each tool to create a per-tool `ToolCallEndEvent`. + * Each variant pairs the tool's name literal with its specific input type, + * enabling discriminated narrowing: checking `toolName === 'x'` narrows `input`. + * @internal + */ +type DistributedToolCallEnd>> = + TTools[number] extends infer T + ? T extends Tool + ? ToolCallEndEvent> + : never + : never + +/** + * Stream chunk type parameterized by the tools array for type-safe tool call events. + * + * When specific tool types are provided (e.g. from `chat({ tools: [myTool] })`): + * - `TOOL_CALL_START` and `TOOL_CALL_END` events form a **discriminated union** + * over tool names — checking `toolName === 'x'` narrows `input` to that tool's type. + * - `TOOL_CALL_END` events have `input` typed per-tool via Standard Schema inference. + * + * When tools are untyped or absent, degrades to the same type as `StreamChunk`. + */ +export type TypedStreamChunk< + TTools extends ReadonlyArray> = ReadonlyArray< + Tool + >, +> = + HasTypedTools extends true + ? + | Exclude< + StreamChunk, + { type: 'TOOL_CALL_START' } | { type: 'TOOL_CALL_END' } + > + | DistributedToolCallStart + | DistributedToolCallEnd + : StreamChunk + // Simple streaming format for basic text completions // Converted to StreamChunk format by convertTextCompletionStream() export interface TextCompletionChunk { diff --git a/packages/typescript/ai/tests/type-check.test.ts b/packages/typescript/ai/tests/type-check.test.ts index 124710beb..ed6cee51d 100644 --- a/packages/typescript/ai/tests/type-check.test.ts +++ b/packages/typescript/ai/tests/type-check.test.ts @@ -1,13 +1,26 @@ /** - * Type-level tests for TextActivityOptions + * Type-level tests for TextActivityOptions and TypedStreamChunk * These should fail to compile if the types are incorrect */ import { describe, it, expectTypeOf } from 'vitest' -import { createChatOptions } from '../src' +import { z } from 'zod' +import { chat, createChatOptions, toolDefinition } from '../src' +import type { + JSONSchema, + StreamChunk, + Tool, + ToolCallArgsEvent, + ToolCallStartEvent, + ToolCallEndEvent, + TypedStreamChunk, +} from '../src' import type { TextAdapter } from '../src/activities/chat/adapter' -// Mock adapter for testing - simulates OpenAI adapter +// =========================== +// Mock adapter (inline — needed for typeof in generic args) +// =========================== + type MockAdapter = TextAdapter< 'test-model', { validOption: string; anotherOption?: number }, @@ -29,6 +42,8 @@ const mockAdapter = { providerOptions: {} as { validOption: string; anotherOption?: number }, inputModalities: ['text', 'image'] as const, messageMetadataByModality: { + // These `as unknown` casts are necessary — TextAdapter requires all 5 + // modality keys but the mock doesn't have real metadata types for them. text: undefined as unknown, image: undefined as unknown, audio: undefined as unknown, @@ -40,9 +55,77 @@ const mockAdapter = { structuredOutput: async () => ({ data: {}, rawText: '{}' }), } satisfies MockAdapter +// =========================== +// Tool definitions for type tests +// =========================== + +const weatherTool = toolDefinition({ + name: 'get_weather', + description: 'Get weather', + inputSchema: z.object({ + location: z.string(), + unit: z.enum(['celsius', 'fahrenheit']).optional(), + }), + outputSchema: z.object({ + temperature: z.number(), + conditions: z.string(), + }), +}) + +const searchTool = toolDefinition({ + name: 'search', + description: 'Search the web', + inputSchema: z.object({ + query: z.string(), + }), +}) + +const weatherServerTool = weatherTool.server(async () => ({ + temperature: 72, + conditions: 'sunny', +})) + +const searchClientTool = searchTool.client(async () => 'results') + +const noInputTool = toolDefinition({ + name: 'get_time', + description: 'Get the current time', +}) + +const jsonSchemaTool: Tool = { + name: 'json_tool', + description: 'A tool with plain JSON Schema', + inputSchema: { + type: 'object', + properties: { key: { type: 'string' } }, + }, +} + +// =========================== +// Type-level helpers to reduce Extract repetition +// =========================== + +/** Extract the TOOL_CALL_START event from a chunk union */ +type StartEventOf = Extract + +/** Extract the TOOL_CALL_END event from a chunk union */ +type EndEventOf = Extract + +/** Extract the chunk type from an AsyncIterable (e.g. chat() return) */ +type ChunkOf = T extends AsyncIterable ? C : never + +/** Build the full TypedStreamChunk and extract both event types at once */ +type ToolEventsOf>> = { + start: StartEventOf> + end: EndEventOf> +} + +// =========================== +// TextActivityOptions type checking (pre-existing) +// =========================== + describe('TextActivityOptions type checking', () => { it('should allow valid options', () => { - // This should type-check successfully const options = createChatOptions({ adapter: mockAdapter, messages: [{ role: 'user', content: 'Hello' }], @@ -75,3 +158,427 @@ describe('TextActivityOptions type checking', () => { }) }) }) + +// =========================== +// TypedStreamChunk: tool name and input typing +// =========================== + +describe('TypedStreamChunk tool call type safety', () => { + describe('tool name typing', () => { + it('should narrow toolName to literal union on both START and END events', () => { + type E = ToolEventsOf<[typeof weatherTool, typeof searchTool]> + + expectTypeOf().toEqualTypeOf< + 'get_weather' | 'search' + >() + expectTypeOf().toEqualTypeOf< + 'get_weather' | 'search' + >() + }) + + it('should narrow toolName to a single literal with one tool', () => { + type E = ToolEventsOf<[typeof weatherTool]> + + expectTypeOf().toEqualTypeOf<'get_weather'>() + expectTypeOf().toEqualTypeOf<'get_weather'>() + }) + }) + + describe('tool input typing', () => { + it('should type input as the union of tool input types', () => { + type E = ToolEventsOf<[typeof weatherTool, typeof searchTool]> + + type ExpectedInput = + | { location: string; unit?: 'celsius' | 'fahrenheit' } + | { query: string } + expectTypeOf< + Exclude + >().toEqualTypeOf() + }) + + it('should type input correctly with a single tool', () => { + type E = ToolEventsOf<[typeof searchTool]> + + expectTypeOf>().toEqualTypeOf<{ + query: string + }>() + }) + + it('should produce unknown input for tools without inputSchema', () => { + type E = ToolEventsOf<[typeof noInputTool]> + + // Use toBeUnknown() instead of toEqualTypeOf() — + // the latter can't distinguish `any` from `unknown` in vitest. + expectTypeOf>().toBeUnknown() + }) + + it('should produce unknown input for plain JSON Schema tools', () => { + type E = ToolEventsOf<[typeof jsonSchemaTool]> + + expectTypeOf>().toBeUnknown() + }) + + it('should preserve tool names when mixing Zod and no-schema tools', () => { + type E = ToolEventsOf<[typeof searchTool, typeof noInputTool]> + + expectTypeOf().toEqualTypeOf< + 'search' | 'get_time' + >() + }) + }) + + describe('discriminated union narrowing', () => { + it('should narrow input to specific tool type when checking toolName', () => { + type Chunk = TypedStreamChunk<[typeof weatherTool, typeof searchTool]> + type End = Extract + + // Narrowing by toolName should give the specific tool's input type + type WeatherEnd = Extract + expectTypeOf>().toEqualTypeOf<{ + location: string + unit?: 'celsius' | 'fahrenheit' + }>() + + type SearchEnd = Extract + expectTypeOf>().toEqualTypeOf<{ + query: string + }>() + }) + + it('should narrow START events by toolName', () => { + type Chunk = TypedStreamChunk<[typeof weatherTool, typeof searchTool]> + type Start = Extract + + type WeatherStart = Extract + expectTypeOf().toEqualTypeOf<'get_weather'>() + + type SearchStart = Extract + expectTypeOf().toEqualTypeOf<'search'>() + }) + + it('should narrow input with three or more tools', () => { + type Chunk = TypedStreamChunk< + [typeof weatherTool, typeof searchTool, typeof noInputTool] + > + type End = Extract + + type WeatherEnd = Extract + expectTypeOf>().toEqualTypeOf<{ + location: string + unit?: 'celsius' | 'fahrenheit' + }>() + + type SearchEnd = Extract + expectTypeOf>().toEqualTypeOf<{ + query: string + }>() + + type TimeEnd = Extract + expectTypeOf>().toBeUnknown() + }) + + it('should narrow input through chat() return type', () => { + const stream = chat({ + adapter: mockAdapter, + messages: [], + tools: [weatherTool, searchTool], + }) + type Chunk = ChunkOf + type End = Extract + + type WeatherEnd = Extract + expectTypeOf>().toEqualTypeOf<{ + location: string + unit?: 'celsius' | 'fahrenheit' + }>() + + type SearchEnd = Extract + expectTypeOf>().toEqualTypeOf<{ + query: string + }>() + }) + + it('should narrow input with server tool variants', () => { + type Chunk = TypedStreamChunk< + [typeof weatherServerTool, typeof searchClientTool] + > + type End = Extract + + type WeatherEnd = Extract + expectTypeOf>().toEqualTypeOf<{ + location: string + unit?: 'celsius' | 'fahrenheit' + }>() + + type SearchEnd = Extract + // .client() preserves the original inputSchema type from the base definition + expectTypeOf>().toEqualTypeOf<{ + query: string + }>() + }) + }) + + describe('server and client tool variants', () => { + it('should type ServerTool name and input from .server()', () => { + type E = ToolEventsOf<[typeof weatherServerTool]> + + expectTypeOf().toEqualTypeOf<'get_weather'>() + expectTypeOf>().toEqualTypeOf<{ + location: string + unit?: 'celsius' | 'fahrenheit' + }>() + }) + + it('should type ClientTool name from .client()', () => { + type E = ToolEventsOf<[typeof searchClientTool]> + + expectTypeOf().toEqualTypeOf<'search'>() + }) + + it('should deduplicate names across definition, server, and client variants', () => { + type E = ToolEventsOf< + [typeof weatherTool, typeof weatherServerTool, typeof searchClientTool] + > + + expectTypeOf().toEqualTypeOf< + 'get_weather' | 'search' + >() + }) + + it('should narrow input through chat() with server/client tools', () => { + const stream = chat({ + adapter: mockAdapter, + messages: [], + tools: [weatherServerTool, searchClientTool], + }) + type Chunk = ChunkOf + type End = Extract + + type WeatherEnd = Extract + expectTypeOf>().toEqualTypeOf<{ + location: string + unit?: 'celsius' | 'fahrenheit' + }>() + + type SearchEnd = Extract + expectTypeOf>().toEqualTypeOf<{ + query: string + }>() + }) + }) + + describe('mixed schema types', () => { + it('should narrow per-tool when mixing Zod and JSON Schema tools', () => { + type Chunk = TypedStreamChunk<[typeof searchTool, typeof jsonSchemaTool]> + type End = Extract + + type SearchEnd = Extract + expectTypeOf>().toEqualTypeOf<{ + query: string + }>() + + type JsonEnd = Extract + expectTypeOf>().toBeUnknown() + }) + }) + + describe('non-tool events are preserved', () => { + it('should include all non-tool-call AG-UI events in the union', () => { + type Chunk = TypedStreamChunk<[typeof weatherTool]> + + // Every AG-UI event type should still be extractable + expectTypeOf>().not.toBeNever() + expectTypeOf>().not.toBeNever() + expectTypeOf>().not.toBeNever() + expectTypeOf< + Extract + >().not.toBeNever() + expectTypeOf< + Extract + >().not.toBeNever() + expectTypeOf< + Extract + >().not.toBeNever() + expectTypeOf>().not.toBeNever() + expectTypeOf>().not.toBeNever() + expectTypeOf< + Extract + >().not.toBeNever() + expectTypeOf>().not.toBeNever() + expectTypeOf>().not.toBeNever() + expectTypeOf>().not.toBeNever() + }) + + it('should keep ToolCallArgsEvent unparameterized (string delta, no toolName)', () => { + type Chunk = TypedStreamChunk<[typeof weatherTool]> + type ArgsEvent = Extract + + expectTypeOf().not.toBeNever() + expectTypeOf().toEqualTypeOf() + expectTypeOf().toMatchTypeOf() + }) + }) +}) + +// =========================== +// chat() return type integration +// =========================== + +describe('chat() tool type inference', () => { + it('should infer typed tool names through chat() return type', () => { + type Chunk = ChunkOf< + ReturnType< + typeof chat< + typeof mockAdapter, + undefined, + true, + [typeof weatherTool, typeof searchTool] + > + > + > + + expectTypeOf['toolName']>().toEqualTypeOf< + 'get_weather' | 'search' + >() + expectTypeOf['toolName']>().toEqualTypeOf< + 'get_weather' | 'search' + >() + }) + + it('should infer TTools from options.tools without explicit type args', () => { + // This is the actual user-facing API — if inference breaks, users silently + // get `string` for toolName even when passing typed tools. + const stream = chat({ + adapter: mockAdapter, + messages: [], + tools: [weatherTool, searchTool], + }) + type Chunk = ChunkOf + + expectTypeOf['toolName']>().toEqualTypeOf< + 'get_weather' | 'search' + >() + expectTypeOf['toolName']>().toEqualTypeOf< + 'get_weather' | 'search' + >() + }) + + it('should return Promise when stream: false, regardless of tools', () => { + type Result = ReturnType< + typeof chat + > + + expectTypeOf().toEqualTypeOf>() + }) + + it('should return Promise when outputSchema is provided', () => { + const schema = z.object({ summary: z.string() }) + type Result = ReturnType< + typeof chat + > + + expectTypeOf().toEqualTypeOf>() + }) +}) + +// =========================== +// createChatOptions() preserves TTools +// =========================== + +describe('createChatOptions() tool type preservation', () => { + it('should preserve specific tool types through options helper', () => { + const opts = createChatOptions({ + adapter: mockAdapter, + tools: [weatherTool, searchTool], + }) + + type ToolsType = Exclude + + // Use union check — tuple ordering is not guaranteed across TS versions + expectTypeOf().toEqualTypeOf< + 'get_weather' | 'search' + >() + }) +}) + +// =========================== +// Fallback / default behavior +// =========================== + +describe('TypedStreamChunk fallback behavior', () => { + it('should fallback to string/unknown with no tools (default generic)', () => { + type Chunk = ChunkOf>> + + expectTypeOf['toolName']>().toEqualTypeOf() + expectTypeOf['toolName']>().toEqualTypeOf() + expectTypeOf['input'], undefined>>().toBeUnknown() + }) + + it('should fallback to string/unknown with empty tools array', () => { + type E = ToolEventsOf<[]> + + expectTypeOf().toEqualTypeOf() + expectTypeOf().toEqualTypeOf() + expectTypeOf>().toBeUnknown() + }) + + it('should fallback to string/unknown when used without type args', () => { + type E = { + start: StartEventOf + end: EndEventOf + } + + expectTypeOf().toEqualTypeOf() + expectTypeOf().toEqualTypeOf() + expectTypeOf>().toBeUnknown() + }) + + it('should handle readonly tools array (as const)', () => { + const tools = [weatherTool, searchTool] as const + type E = ToolEventsOf + + expectTypeOf().toEqualTypeOf< + 'get_weather' | 'search' + >() + }) +}) + +// =========================== +// Backward compatibility +// =========================== + +describe('backward compatibility', () => { + it('should preserve unparameterized ToolCallStartEvent/ToolCallEndEvent defaults', () => { + expectTypeOf().toEqualTypeOf() + expectTypeOf().toEqualTypeOf() + expectTypeOf>().toBeUnknown() + }) + + it('should treat explicit defaults as identical to unparameterized', () => { + expectTypeOf< + ToolCallStartEvent + >().toEqualTypeOf() + expectTypeOf< + ToolCallEndEvent + >().toEqualTypeOf() + }) + + it('should make typed events assignable to untyped events', () => { + expectTypeOf< + ToolCallStartEvent<'get_weather'> + >().toMatchTypeOf() + expectTypeOf< + ToolCallEndEvent<'get_weather', { location: string }> + >().toMatchTypeOf() + }) + + it('should make TypedStreamChunk assignable to StreamChunk', () => { + type Typed = TypedStreamChunk<[typeof weatherTool]> + expectTypeOf().toMatchTypeOf() + }) + + it('should keep StreamChunk itself unchanged', () => { + type Start = Extract + expectTypeOf().toEqualTypeOf() + }) +})