diff --git a/.changeset/lazy-tool-discovery.md b/.changeset/lazy-tool-discovery.md new file mode 100644 index 000000000..6c7a2c802 --- /dev/null +++ b/.changeset/lazy-tool-discovery.md @@ -0,0 +1,7 @@ +--- +'@tanstack/ai': minor +--- + +feat: Add lazy tool discovery for `chat()` + +Tools marked with `lazy: true` are not sent to the LLM upfront. Instead, a synthetic discovery tool lets the LLM selectively discover lazy tools by name, receiving their descriptions and schemas on demand. Discovered tools are dynamically injected as normal tools. This reduces token usage and improves response quality when applications have many tools. diff --git a/.changeset/sixty-lions-accept.md b/.changeset/sixty-lions-accept.md new file mode 100644 index 000000000..e92d8175f --- /dev/null +++ b/.changeset/sixty-lions-accept.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-openrouter': patch +--- + +Fix issue with tool calling diff --git a/docs/config.json b/docs/config.json index 8860b6667..1677e856f 100644 --- a/docs/config.json +++ b/docs/config.json @@ -46,6 +46,10 @@ "label": "Tool Approval Flow", "to": "guides/tool-approval" }, + { + "label": "Lazy Tool Discovery", + "to": "guides/lazy-tool-discovery" + }, { "label": "Agentic Cycle", "to": "guides/agentic-cycle" diff --git a/docs/guides/lazy-tool-discovery.md b/docs/guides/lazy-tool-discovery.md new file mode 100644 index 000000000..0d8e671d1 --- /dev/null +++ b/docs/guides/lazy-tool-discovery.md @@ -0,0 +1,222 @@ +--- +title: Lazy Tool Discovery +id: lazy-tool-discovery +order: 6 +--- + +When an application has many tools, sending all tool definitions to the LLM on every request wastes tokens and can degrade response quality. Lazy tool discovery lets the LLM selectively discover only the tools it needs for the current task. + +## How It Works + +Tools marked with `lazy: true` are **not** sent to the LLM upfront. Instead, a synthetic `__lazy__tool__discovery__` tool is created whose description lists all available lazy tool names. The LLM can call this discovery tool with the names of tools it wants to learn about, and receives their full descriptions and argument schemas in return. + +Once discovered, lazy tools are dynamically injected as normal tools — the LLM calls them directly like any other tool. + +```mermaid +sequenceDiagram + participant LLM + participant Server + participant Discovery Tool + participant Lazy Tool + + Note over LLM: Sees __lazy__tool__discovery__
with available tool names + + LLM->>Server: Call __lazy__tool__discovery__
{toolNames: ["searchProducts"]} + Server->>Discovery Tool: Execute discovery + Discovery Tool-->>Server: Return description + schema + Server-->>LLM: Tool result with schema + + Note over LLM: searchProducts now available
as a normal tool + + LLM->>Server: Call searchProducts
{query: "red shoes"} + Server->>Lazy Tool: Execute searchProducts + Lazy Tool-->>Server: Return results + Server-->>LLM: Tool result +``` + +## Marking Tools as Lazy + +Add `lazy: true` to any tool definition: + +```typescript +import { toolDefinition } from "@tanstack/ai"; +import { z } from "zod"; + +const searchProductsDef = toolDefinition({ + name: "searchProducts", + description: "Search products by keyword in name or description", + inputSchema: z.object({ + query: z.string().describe("Search keyword or phrase"), + }), + outputSchema: z.object({ + results: z.array( + z.object({ + id: z.number(), + name: z.string(), + price: z.number(), + }) + ), + }), + lazy: true, // This tool won't be sent to the LLM upfront +}); + +const searchProducts = searchProductsDef.server(async ({ query }) => { + const results = await db.products.search(query); + return { results }; +}); +``` + +Then pass it to `chat()` alongside your other tools: + +```typescript +import { chat, toServerSentEventsResponse } from "@tanstack/ai"; +import { openaiText } from "@tanstack/ai-openai"; + +const stream = chat({ + adapter: openaiText("gpt-4o"), + messages, + tools: [ + getProducts, // Normal tool — sent to LLM immediately + searchProducts, // Lazy tool — discovered on demand + compareProducts, // Lazy tool — discovered on demand + ], +}); + +return toServerSentEventsResponse(stream); +``` + +## When to Use Lazy Tools + +Lazy tools are useful when: + +- **You have many tools** and want to reduce token usage per request +- **Some tools are rarely needed** — secondary features like comparison, financing, or advanced search +- **Tool descriptions are large** — lazy tools keep the initial prompt lean + +Tools that are called in most conversations should remain eager (the default). + +## Discovery Flow + +1. The LLM sees `__lazy__tool__discovery__` with a list of available tool names in its description +2. Based on the user's request, the LLM decides which tools it needs and calls the discovery tool +3. The discovery tool returns the full description and JSON Schema for each requested tool +4. The discovered tools are injected as normal tools for the next iteration +5. The LLM calls the discovered tools directly + +The LLM can discover one or many tools in a single call: + +``` +// LLM calls: +__lazy__tool__discovery__({ toolNames: ["searchProducts", "compareProducts"] }) +``` + +## Multi-Turn Conversations + +Lazy tool discovery works across multiple turns. When a tool is discovered in one turn, it remains available in subsequent turns within the same conversation — the LLM does not need to re-discover it. + +This is handled automatically by scanning the message history for previous discovery tool results on each `chat()` call. + +## Self-Correction + +If the LLM tries to call a lazy tool that hasn't been discovered yet, it receives an error message: + +``` +Error: Tool 'searchProducts' must be discovered first. +Call __lazy__tool__discovery__ with toolNames: ['searchProducts'] to discover it. +``` + +The LLM then self-corrects by calling the discovery tool first, then retrying the original tool call. + +## Zero Overhead + +If none of your tools have `lazy: true`, no discovery tool is created and the behavior is identical to the default. There is no performance or token cost when lazy discovery is not in use. + +When all lazy tools have been discovered, the discovery tool is automatically removed from the active tool set. + +## Example + +Here's a complete example with a mix of eager and lazy tools: + +```typescript +import { toolDefinition, chat, toServerSentEventsResponse, maxIterations } from "@tanstack/ai"; +import { openaiText } from "@tanstack/ai-openai"; +import { z } from "zod"; + +// Eager tool — always available +const getProductsDef = toolDefinition({ + name: "getProducts", + description: "Get all products from the catalog", + inputSchema: z.object({}), + outputSchema: z.array( + z.object({ + id: z.number(), + name: z.string(), + price: z.number(), + }) + ), +}); + +const getProducts = getProductsDef.server(async () => { + return await db.products.findMany(); +}); + +// Lazy tool — discovered on demand +const compareProductsDef = toolDefinition({ + name: "compareProducts", + description: "Compare two or more products side by side", + inputSchema: z.object({ + productIds: z.array(z.number()).min(2), + }), + lazy: true, +}); + +const compareProducts = compareProductsDef.server(async ({ productIds }) => { + const products = await db.products.findMany({ + where: { id: { in: productIds } }, + }); + return { products }; +}); + +// Lazy tool — discovered on demand +const calculateFinancingDef = toolDefinition({ + name: "calculateFinancing", + description: "Calculate monthly payment plans for a product", + inputSchema: z.object({ + productId: z.number(), + months: z.number(), + }), + lazy: true, +}); + +const calculateFinancing = calculateFinancingDef.server(async ({ productId, months }) => { + const product = await db.products.findUnique({ where: { id: productId } }); + const monthlyPayment = product.price / months; + return { monthlyPayment, totalPrice: product.price, months }; +}); + +// Use in chat +export async function POST(request: Request) { + const { messages } = await request.json(); + + const stream = chat({ + adapter: openaiText("gpt-4o"), + messages, + tools: [getProducts, compareProducts, calculateFinancing], + agentLoopStrategy: maxIterations(20), + }); + + return toServerSentEventsResponse(stream); +} +``` + +With this setup: +- The LLM always sees `getProducts` and `__lazy__tool__discovery__` +- When a user asks to compare products, the LLM discovers `compareProducts` first, then calls it +- When a user asks about financing, the LLM discovers `calculateFinancing` first, then calls it + +## Next Steps + +- [Tools Overview](./tools) - Basic tool concepts +- [Server Tools](./server-tools) - Server-side tool execution +- [Tool Architecture](./tool-architecture) - Deep dive into the tool system +- [Agentic Cycle](./agentic-cycle) - How the agent loop works diff --git a/examples/ts-react-chat/README.md b/examples/ts-react-chat/README.md index 7d49f43fc..572680cc2 100644 --- a/examples/ts-react-chat/README.md +++ b/examples/ts-react-chat/README.md @@ -39,6 +39,41 @@ An example chat application built with TanStack Start, TanStack Store, and **Tan OPENAI_API_KEY=your_openai_api_key ``` +## Trying Out Lazy Tool Discovery + +This example includes three **lazy tools** — tools that are not sent to the LLM upfront. Instead, the LLM sees a `__lazy__tool__discovery__` tool that lists their names. When the LLM needs one, it discovers it first (getting the full description and schema), then calls it normally. + +The lazy tools are: `compareGuitars`, `calculateFinancing`, and `searchGuitars`. + +### Test Prompts + +**Compare guitars** — triggers discovery of `compareGuitars`: + +> "Can you compare the Motherboard Guitar and the Racing Guitar for me?" + +**Financing** — triggers discovery of `calculateFinancing`: + +> "How much would it cost per month if I financed the Superhero Guitar over 12 months?" + +**Search** — triggers discovery of `searchGuitars`: + +> "Do you have any guitars with LED lights or tech features?" + +**Multi-discovery** — triggers discovery of multiple lazy tools at once: + +> "I'm looking for acoustic guitars. Can you search for them, then compare the ones you find, and show me financing options for the cheapest one?" + +**Self-correction** — the LLM may try calling a lazy tool directly without discovering it first. It will get an error telling it to discover first, then self-correct: + +> "Compare guitar 1 and guitar 4 right now" + +### What to watch for + +- The LLM calls `__lazy__tool__discovery__` before using a lazy tool for the first time +- After discovery, the tool is called directly like any normal tool +- In multi-turn conversations, previously discovered tools are usable immediately without re-discovery +- If the LLM skips discovery, it gets an error and self-corrects + ## ✨ Features ### AI Capabilities diff --git a/examples/ts-react-chat/src/lib/guitar-tools.ts b/examples/ts-react-chat/src/lib/guitar-tools.ts index 32c9a4ee3..b984f5ee0 100644 --- a/examples/ts-react-chat/src/lib/guitar-tools.ts +++ b/examples/ts-react-chat/src/lib/guitar-tools.ts @@ -88,3 +88,154 @@ export const addToCartToolDef = toolDefinition({ }), needsApproval: true, }) + +// --- Lazy tools (discovered on demand) --- + +// Compare two or more guitars side by side +export const compareGuitarsToolDef = toolDefinition({ + name: 'compareGuitars', + description: + 'Compare two or more guitars side by side, showing their differences in price, type, and features.', + inputSchema: z.object({ + guitarIds: z + .array(z.number()) + .min(2) + .describe('Array of guitar IDs to compare (minimum 2)'), + }), + outputSchema: z.object({ + comparison: z.array( + z.object({ + id: z.number(), + name: z.string(), + price: z.number(), + description: z.string(), + priceDifference: z.string(), + }), + ), + cheapest: z.string(), + mostExpensive: z.string(), + }), + lazy: true, +}) + +export const compareGuitars = compareGuitarsToolDef.server((args) => { + const selected = args.guitarIds + .map((id) => guitars.find((g) => g.id === id)) + .filter(Boolean) as (typeof guitars)[number][] + + const prices = selected.map((g) => g.price) + const minPrice = Math.min(...prices) + const maxPrice = Math.max(...prices) + + return { + comparison: selected.map((g) => ({ + id: g.id, + name: g.name, + price: g.price, + description: g.shortDescription, + priceDifference: + g.price === minPrice + ? 'Cheapest' + : `+$${g.price - minPrice} more than cheapest`, + })), + cheapest: selected.find((g) => g.price === minPrice)!.name, + mostExpensive: selected.find((g) => g.price === maxPrice)!.name, + } +}) + +// Calculate monthly financing for a guitar +export const calculateFinancingToolDef = toolDefinition({ + name: 'calculateFinancing', + description: + 'Calculate monthly payment plans for a guitar purchase. Supports 6, 12, and 24 month terms.', + inputSchema: z.object({ + guitarId: z + .number() + .describe('The ID of the guitar to calculate financing for'), + months: z + .number() + .describe('Number of months for the payment plan (6, 12, or 24)'), + }), + outputSchema: z.object({ + guitarName: z.string(), + totalPrice: z.number(), + months: z.number(), + monthlyPayment: z.number(), + apr: z.number(), + totalWithInterest: z.number(), + }), + lazy: true, +}) + +export const calculateFinancing = calculateFinancingToolDef.server((args) => { + const guitar = guitars.find((g) => g.id === args.guitarId) + if (!guitar) { + throw new Error(`Guitar with ID ${args.guitarId} not found`) + } + + const apr = args.months <= 6 ? 0 : args.months <= 12 ? 5.9 : 9.9 + const monthlyRate = apr / 100 / 12 + const monthlyPayment = + apr === 0 + ? guitar.price / args.months + : (guitar.price * monthlyRate) / + (1 - Math.pow(1 + monthlyRate, -args.months)) + + return { + guitarName: guitar.name, + totalPrice: guitar.price, + months: args.months, + monthlyPayment: Math.round(monthlyPayment * 100) / 100, + apr, + totalWithInterest: Math.round(monthlyPayment * args.months * 100) / 100, + } +}) + +// Search guitars by features/keywords +export const searchGuitarsToolDef = toolDefinition({ + name: 'searchGuitars', + description: + 'Search guitars by keyword in their name or description. Useful for finding guitars matching specific features like "acoustic", "electric", "LED", "vintage", etc.', + inputSchema: z.object({ + query: z + .string() + .describe( + 'Search keyword or phrase to match against guitar names and descriptions', + ), + }), + outputSchema: z.object({ + results: z.array( + z.object({ + id: z.number(), + name: z.string(), + price: z.number(), + shortDescription: z.string(), + matchedIn: z.string(), + }), + ), + totalFound: z.number(), + }), + lazy: true, +}) + +export const searchGuitars = searchGuitarsToolDef.server((args) => { + const query = args.query.toLowerCase() + const results = guitars + .filter( + (g) => + g.name.toLowerCase().includes(query) || + g.description.toLowerCase().includes(query), + ) + .map((g) => ({ + id: g.id, + name: g.name, + price: g.price, + shortDescription: g.shortDescription, + matchedIn: g.name.toLowerCase().includes(query) ? 'name' : 'description', + })) + + return { + results, + totalFound: results.length, + } +}) diff --git a/examples/ts-react-chat/src/routes/api.tanchat.ts b/examples/ts-react-chat/src/routes/api.tanchat.ts index 086bd00ee..a1eb8ee02 100644 --- a/examples/ts-react-chat/src/routes/api.tanchat.ts +++ b/examples/ts-react-chat/src/routes/api.tanchat.ts @@ -16,9 +16,12 @@ import { groqText } from '@tanstack/ai-groq' import { addToCartToolDef, addToWishListToolDef, + calculateFinancing, + compareGuitars, getGuitars, getPersonalGuitarPreferenceToolDef, recommendGuitarToolDef, + searchGuitars, } from '@/lib/guitar-tools' type Provider = @@ -49,8 +52,9 @@ IMPORTANT: Example workflow: User: "I want an acoustic guitar" Step 1: Call getGuitars() -Step 2: Call recommendGuitar(id: "6") +Step 2: Call recommendGuitar(id: "6") Step 3: Done - do NOT add any text after calling recommendGuitar + ` const addToCartToolServer = addToCartToolDef.server((args, context) => { context?.emitCustomEvent('tool:progress', { @@ -200,6 +204,10 @@ export const Route = createFileRoute('/api/tanchat')({ addToCartToolServer, addToWishListToolDef, getPersonalGuitarPreferenceToolDef, + // Lazy tools - discovered on demand + compareGuitars, + calculateFinancing, + searchGuitars, ], middleware: [loggingMiddleware], systemPrompts: [SYSTEM_PROMPT], diff --git a/packages/typescript/ai-openrouter/src/adapters/text.ts b/packages/typescript/ai-openrouter/src/adapters/text.ts index 387110b8b..76733e913 100644 --- a/packages/typescript/ai-openrouter/src/adapters/text.ts +++ b/packages/typescript/ai-openrouter/src/adapters/text.ts @@ -31,7 +31,6 @@ import type { OpenRouterMessageMetadataByModality, } from '../message-types' import type { - ChatCompletionFinishReason, ChatGenerationParams, ChatGenerationTokenUsage, ChatMessageContentItem, @@ -99,8 +98,6 @@ export class OpenRouterTextAdapter< let accumulatedContent = '' let responseId: string | null = null let currentModel = options.model - let lastFinishReason: ChatCompletionFinishReason | undefined - // AG-UI lifecycle tracking const aguiState: AGUIState = { runId: this.generateId(), @@ -149,9 +146,6 @@ export class OpenRouterTextAdapter< } for (const choice of chunk.choices) { - if (chunk.choices[0]?.finishReason) { - lastFinishReason = chunk.choices[0].finishReason - } yield* this.processChoice( choice, toolCallBuffers, @@ -165,7 +159,6 @@ export class OpenRouterTextAdapter< accumulatedReasoning = r accumulatedContent = c }, - lastFinishReason, chunk.usage, aguiState, ) @@ -288,7 +281,6 @@ export class OpenRouterTextAdapter< meta: { id: string; model: string; timestamp: number }, accumulated: { reasoning: string; content: string }, updateAccumulated: (reasoning: string, content: string) => void, - lastFinishReason: ChatCompletionFinishReason | undefined, usage: ChatGenerationTokenUsage | undefined, aguiState: AGUIState, ): Iterable { @@ -447,7 +439,8 @@ export class OpenRouterTextAdapter< } if (finishReason) { - if (finishReason === 'tool_calls') { + // Emit all completed tool calls when finish reason indicates tool usage + if (finishReason === 'tool_calls' || toolCallBuffers.size > 0) { for (const [, tc] of toolCallBuffers.entries()) { // Parse arguments for TOOL_CALL_END let parsedInput: unknown = {} @@ -470,12 +463,11 @@ export class OpenRouterTextAdapter< toolCallBuffers.clear() } - } - if (usage) { + const computedFinishReason = - lastFinishReason === 'tool_calls' + finishReason === 'tool_calls' ? 'tool_calls' - : lastFinishReason === 'length' + : finishReason === 'length' ? 'length' : 'stop' @@ -495,11 +487,13 @@ export class OpenRouterTextAdapter< runId: aguiState.runId, model: meta.model, timestamp: meta.timestamp, - usage: { - promptTokens: usage.promptTokens || 0, - completionTokens: usage.completionTokens || 0, - totalTokens: usage.totalTokens || 0, - }, + usage: usage + ? { + promptTokens: usage.promptTokens || 0, + completionTokens: usage.completionTokens || 0, + totalTokens: usage.totalTokens || 0, + } + : undefined, finishReason: computedFinishReason, } } diff --git a/packages/typescript/ai/src/activities/chat/index.ts b/packages/typescript/ai/src/activities/chat/index.ts index 74c5d5ce5..c7ec866d6 100644 --- a/packages/typescript/ai/src/activities/chat/index.ts +++ b/packages/typescript/ai/src/activities/chat/index.ts @@ -7,6 +7,7 @@ import { devtoolsMiddleware } from '@tanstack/ai-event-client' import { streamToText } from '../../stream-to-response.js' +import { LazyToolManager } from './tools/lazy-tool-manager' import { MiddlewareAbortError, ToolCallManager, @@ -237,7 +238,8 @@ class TextEngine< private systemPrompts: Array private tools: Array private readonly loopStrategy: AgentLoopStrategy - private readonly toolCallManager: ToolCallManager + private toolCallManager: ToolCallManager + private readonly lazyToolManager: LazyToolManager private readonly initialMessageCount: number private readonly requestId: string private readonly streamId: string @@ -273,10 +275,8 @@ class TextEngine< this.adapter = config.adapter this.params = config.params this.systemPrompts = config.params.systemPrompts || [] - this.tools = config.params.tools || [] this.loopStrategy = config.params.agentLoopStrategy || maxIterationsStrategy(5) - this.toolCallManager = new ToolCallManager(this.tools) this.initialMessageCount = config.params.messages.length // Extract client state (approvals, client tool results) from original messages BEFORE conversion @@ -293,6 +293,14 @@ class TextEngine< this.messages = convertMessagesToModelMessages( config.params.messages as Array, ) + + // Initialize lazy tool manager after messages are converted (needs message history for scanning) + this.lazyToolManager = new LazyToolManager( + config.params.tools || [], + this.messages, + ) + this.tools = this.lazyToolManager.getActiveTools() + this.toolCallManager = new ToolCallManager(this.tools) this.requestId = this.createId('chat') this.streamId = this.createId('stream') this.effectiveRequest = config.params.abortController @@ -643,10 +651,42 @@ class TextEngine< const finishEvent = this.createSyntheticFinishedEvent() + // Handle undiscovered lazy tool calls with self-correcting error messages + const undiscoveredLazyResults: Array = [] + const executablePendingCalls = pendingToolCalls.filter((tc) => { + if (this.lazyToolManager.isUndiscoveredLazyTool(tc.function.name)) { + undiscoveredLazyResults.push({ + toolCallId: tc.id, + toolName: tc.function.name, + result: { + error: this.lazyToolManager.getUndiscoveredToolError( + tc.function.name, + ), + }, + state: 'output-error', + }) + return false + } + return true + }) + + if (undiscoveredLazyResults.length > 0) { + for (const chunk of this.buildToolResultChunks( + undiscoveredLazyResults, + finishEvent, + )) { + yield chunk + } + } + + if (executablePendingCalls.length === 0) { + return 'continue' + } + const { approvals, clientToolResults } = this.collectClientState() const generator = executeToolCalls( - pendingToolCalls, + executablePendingCalls, this.tools, approvals, clientToolResults, @@ -750,12 +790,47 @@ class TextEngine< this.addAssistantToolCallMessage(toolCalls) + // Handle undiscovered lazy tool calls with self-correcting error messages + const undiscoveredLazyResults: Array = [] + const executableToolCalls = toolCalls.filter((tc) => { + if (this.lazyToolManager.isUndiscoveredLazyTool(tc.function.name)) { + undiscoveredLazyResults.push({ + toolCallId: tc.id, + toolName: tc.function.name, + result: { + error: this.lazyToolManager.getUndiscoveredToolError( + tc.function.name, + ), + }, + state: 'output-error', + }) + return false + } + return true + }) + + if (undiscoveredLazyResults.length > 0) { + const finishEvt = this.finishedEvent! + for (const chunk of this.buildToolResultChunks( + undiscoveredLazyResults, + finishEvt, + )) { + yield chunk + } + } + + if (executableToolCalls.length === 0) { + // All tool calls were undiscovered lazy tools — errors emitted, continue loop + this.toolCallManager.clear() + this.setToolPhase('continue') + return + } this.middlewareCtx.phase = 'beforeTools' const { approvals, clientToolResults } = this.collectClientState() const generator = executeToolCalls( - toolCalls, + executableToolCalls, this.tools, approvals, clientToolResults, @@ -842,6 +917,14 @@ class TextEngine< yield chunk } + // Refresh tools if lazy tools were discovered in this batch + if (this.lazyToolManager.hasNewlyDiscoveredTools()) { + this.tools = this.lazyToolManager.getActiveTools() + this.toolCallManager = new ToolCallManager(this.tools) + this.setToolPhase('continue') + return + } + this.toolCallManager.clear() this.setToolPhase('continue') diff --git a/packages/typescript/ai/src/activities/chat/tools/lazy-tool-manager.ts b/packages/typescript/ai/src/activities/chat/tools/lazy-tool-manager.ts new file mode 100644 index 000000000..d4aa50ca0 --- /dev/null +++ b/packages/typescript/ai/src/activities/chat/tools/lazy-tool-manager.ts @@ -0,0 +1,254 @@ +import { convertSchemaToJsonSchema } from './schema-converter' +import type { Tool } from '../../../types' + +const DISCOVERY_TOOL_NAME = '__lazy__tool__discovery__' + +/** + * Manages lazy tool discovery for the chat agent loop. + * + * Lazy tools are not sent to the LLM initially. Instead, a synthetic + * "discovery tool" is provided that lets the LLM discover lazy tools + * by name, receiving their full descriptions and schemas on demand. + */ +export class LazyToolManager { + private readonly eagerTools: ReadonlyArray + private readonly lazyToolMap: Map + private readonly discoveredTools: Set + private hasNewDiscoveries: boolean + private readonly discoveryTool: Tool | null + + constructor( + tools: ReadonlyArray, + messages: ReadonlyArray<{ + role: string + content?: any + toolCalls?: Array<{ + id: string + type: string + function: { name: string; arguments: string } + }> + toolCallId?: string + }>, + ) { + const eager: Array = [] + this.lazyToolMap = new Map() + this.discoveredTools = new Set() + this.hasNewDiscoveries = false + + // Separate tools into eager and lazy + for (const tool of tools) { + if (tool.lazy) { + this.lazyToolMap.set(tool.name, tool) + } else { + eager.push(tool) + } + } + this.eagerTools = eager + + // If no lazy tools, no discovery tool needed + if (this.lazyToolMap.size === 0) { + this.discoveryTool = null + return + } + + // Scan message history to pre-populate discoveredTools + this.scanMessageHistory(messages) + + // Create the synthetic discovery tool + this.discoveryTool = this.createDiscoveryTool() + } + + /** + * Returns the set of tools that should be sent to the LLM: + * eager tools + discovered lazy tools + discovery tool (if undiscovered tools remain). + * Resets the hasNewDiscoveries flag. + */ + getActiveTools(): Array { + this.hasNewDiscoveries = false + + const active: Array = [...this.eagerTools] + + // Add discovered lazy tools + for (const name of this.discoveredTools) { + const tool = this.lazyToolMap.get(name) + if (tool) { + active.push(tool) + } + } + + // Add discovery tool if there are still undiscovered lazy tools + if ( + this.discoveryTool && + this.discoveredTools.size < this.lazyToolMap.size + ) { + active.push(this.discoveryTool) + } + + return active + } + + /** + * Returns whether new tools have been discovered since the last getActiveTools() call. + */ + hasNewlyDiscoveredTools(): boolean { + return this.hasNewDiscoveries + } + + /** + * Returns true if the given name is a lazy tool that has not yet been discovered. + */ + isUndiscoveredLazyTool(name: string): boolean { + return this.lazyToolMap.has(name) && !this.discoveredTools.has(name) + } + + /** + * Returns a helpful error message for when an undiscovered lazy tool is called. + */ + getUndiscoveredToolError(name: string): string { + return `Error: Tool '${name}' must be discovered first. Call ${DISCOVERY_TOOL_NAME} with toolNames: ['${name}'] to discover it.` + } + + /** + * Scans message history to find previously discovered lazy tools. + * Looks for assistant messages with discovery tool calls and their + * corresponding tool result messages. + */ + private scanMessageHistory( + messages: ReadonlyArray<{ + role: string + content?: any + toolCalls?: Array<{ + id: string + type: string + function: { name: string; arguments: string } + }> + toolCallId?: string + }>, + ): void { + // Collect tool call IDs for discovery tool invocations + const discoveryCallIds = new Set() + + for (const msg of messages) { + if (msg.role === 'assistant' && msg.toolCalls) { + for (const tc of msg.toolCalls) { + if (tc.function.name === DISCOVERY_TOOL_NAME) { + discoveryCallIds.add(tc.id) + } + } + } + } + + if (discoveryCallIds.size === 0) return + + // Find corresponding tool result messages + for (const msg of messages) { + if ( + msg.role === 'tool' && + msg.toolCallId && + discoveryCallIds.has(msg.toolCallId) + ) { + try { + const content = + typeof msg.content === 'string' + ? msg.content + : JSON.stringify(msg.content) + const parsed = JSON.parse(content) + if (parsed && Array.isArray(parsed.tools)) { + for (const tool of parsed.tools) { + if ( + tool && + typeof tool.name === 'string' && + this.lazyToolMap.has(tool.name) + ) { + this.discoveredTools.add(tool.name) + } + } + } + } catch { + // Malformed JSON — skip gracefully + } + } + } + } + + /** + * Creates the synthetic discovery tool that the LLM can call + * to discover lazy tools' descriptions and schemas. + */ + private createDiscoveryTool(): Tool { + const undiscoveredNames = (): Array => { + const names: Array = [] + for (const [name] of this.lazyToolMap) { + if (!this.discoveredTools.has(name)) { + names.push(name) + } + } + return names + } + + const lazyToolMap = this.lazyToolMap + + // Build the static description with all lazy tool names + const allLazyNames = Array.from(this.lazyToolMap.keys()) + const description = `You have access to additional tools that can be discovered. Available tools: [${allLazyNames.join(', ')}]. Call this tool with a list of tool names to discover their full descriptions and argument schemas before using them.` + + // Use the arrow function to capture `this` context + const manager = this + + return { + name: DISCOVERY_TOOL_NAME, + description, + inputSchema: { + type: 'object', + properties: { + toolNames: { + type: 'array', + items: { type: 'string' }, + description: + 'List of tool names to discover. Each name must match one of the available tools.', + }, + }, + required: ['toolNames'], + }, + execute: (args: { toolNames: Array }) => { + const tools: Array<{ + name: string + description: string + inputSchema?: any + }> = [] + const errors: Array = [] + + for (const name of args.toolNames) { + const tool = lazyToolMap.get(name) + if (tool) { + manager.discoveredTools.add(name) + manager.hasNewDiscoveries = true + const jsonSchema = tool.inputSchema + ? convertSchemaToJsonSchema(tool.inputSchema) + : undefined + tools.push({ + name: tool.name, + description: tool.description, + ...(jsonSchema ? { inputSchema: jsonSchema } : {}), + }) + } else { + errors.push( + `Unknown tool: '${name}'. Available tools: [${undiscoveredNames().join(', ')}]`, + ) + } + } + + const result: { + tools: typeof tools + errors?: Array + } = { tools } + + if (errors.length > 0) { + result.errors = errors + } + + return result + }, + } + } +} diff --git a/packages/typescript/ai/src/activities/chat/tools/tool-definition.ts b/packages/typescript/ai/src/activities/chat/tools/tool-definition.ts index ca82b8cd3..c12c61898 100644 --- a/packages/typescript/ai/src/activities/chat/tools/tool-definition.ts +++ b/packages/typescript/ai/src/activities/chat/tools/tool-definition.ts @@ -32,6 +32,7 @@ export interface ClientTool< inputSchema?: TInput outputSchema?: TOutput needsApproval?: boolean + lazy?: boolean metadata?: Record execute?: ( args: InferSchemaType, @@ -96,6 +97,7 @@ export interface ToolDefinitionConfig< inputSchema?: TInput outputSchema?: TOutput needsApproval?: boolean + lazy?: boolean metadata?: Record } diff --git a/packages/typescript/ai/src/types.ts b/packages/typescript/ai/src/types.ts index 49e800fc5..c197ed21a 100644 --- a/packages/typescript/ai/src/types.ts +++ b/packages/typescript/ai/src/types.ts @@ -493,6 +493,9 @@ export interface Tool< /** If true, tool execution requires user approval before running. Works with both server and client tools. */ needsApproval?: boolean + /** If true, this tool is lazy and will only be sent to the LLM after being discovered via the lazy tool discovery mechanism. Only meaningful when used with chat(). */ + lazy?: boolean + /** Additional metadata for adapters or custom extensions */ metadata?: Record } diff --git a/packages/typescript/ai/tests/chat.test.ts b/packages/typescript/ai/tests/chat.test.ts index 1b9b8eec4..f213c5b36 100644 --- a/packages/typescript/ai/tests/chat.test.ts +++ b/packages/typescript/ai/tests/chat.test.ts @@ -9,6 +9,16 @@ import { clientTool, } from './test-utils' +/** Lazy server tool (has execute, lazy: true). */ +function lazyServerTool(name: string, executeFn: (args: any) => any): Tool { + return { + name, + description: `Lazy tool: ${name}`, + execute: executeFn, + lazy: true, + } +} + // ============================================================================ // Tests // ============================================================================ @@ -1078,4 +1088,199 @@ describe('chat()', () => { expect(result).toBe('Hello') }) }) + + // ========================================================================== + // Lazy tool discovery + // ========================================================================== + describe('lazy tool discovery', () => { + it('should create discovery tool when lazy tools are provided', async () => { + const weatherExecute = vi.fn().mockReturnValue({ temp: 72 }) + + let callCount = 0 + const { adapter, calls } = createMockAdapter({ + chatStreamFn: (opts: any) => { + callCount++ + const toolNames = opts.tools?.map((t: any) => t.name) || [] + + if (callCount === 1) { + // First call: only discovery tool available, LLM discovers getWeather + return (async function* () { + yield ev.runStarted() + yield ev.toolStart('call_disc', '__lazy__tool__discovery__') + yield ev.toolArgs( + 'call_disc', + JSON.stringify({ toolNames: ['getWeather'] }), + ) + yield ev.runFinished('tool_calls') + })() + } else if (callCount === 2 && toolNames.includes('getWeather')) { + // Second call: getWeather is now available, LLM calls it + return (async function* () { + yield ev.runStarted() + yield ev.toolStart('call_weather', 'getWeather') + yield ev.toolArgs('call_weather', '{"city":"NYC"}') + yield ev.runFinished('tool_calls') + })() + } else { + // Third call: final text after tool execution + return (async function* () { + yield ev.runStarted() + yield ev.textStart() + yield ev.textContent('It is 72F in NYC.') + yield ev.textEnd() + yield ev.runFinished('stop') + })() + } + }, + }) + + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'Weather in NYC?' }], + tools: [lazyServerTool('getWeather', weatherExecute)], + }) + + const chunks = await collectChunks(stream as AsyncIterable) + + // First adapter call should have __lazy__tool__discovery__ but NOT getWeather + const firstCallToolNames = calls[0].tools.map((t: any) => t.name) + expect(firstCallToolNames).toContain('__lazy__tool__discovery__') + expect(firstCallToolNames).not.toContain('getWeather') + + // Second adapter call should have getWeather (after discovery) + const secondCallToolNames = calls[1].tools.map((t: any) => t.name) + expect(secondCallToolNames).toContain('getWeather') + + // TOOL_CALL_END chunks should exist for both discovery and getWeather + const toolEndChunks = chunks.filter((c) => c.type === 'TOOL_CALL_END') + expect(toolEndChunks.length).toBeGreaterThanOrEqual(2) + + // getWeather should have been executed + expect(weatherExecute).toHaveBeenCalledTimes(1) + }) + + it('should work with mix of eager and lazy tools', async () => { + const eagerExecute = vi.fn().mockReturnValue({ result: 'eager' }) + const lazyExecute = vi.fn().mockReturnValue({ result: 'lazy' }) + + const { adapter, calls } = createMockAdapter({ + iterations: [ + [ev.runStarted(), ev.textContent('Hello'), ev.runFinished('stop')], + ], + }) + + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'Hi' }], + tools: [ + serverTool('eagerTool', eagerExecute), + lazyServerTool('lazyTool', lazyExecute), + ], + }) + + await collectChunks(stream as AsyncIterable) + + // First adapter call should have eager tool + discovery tool, but NOT lazyTool + const firstCallToolNames = calls[0].tools.map((t: any) => t.name) + expect(firstCallToolNames).toContain('eagerTool') + expect(firstCallToolNames).toContain('__lazy__tool__discovery__') + expect(firstCallToolNames).not.toContain('lazyTool') + }) + + it('should handle undiscovered lazy tool call with self-correcting error', async () => { + const weatherExecute = vi.fn().mockReturnValue({ temp: 72 }) + + let callCount = 0 + const { adapter } = createMockAdapter({ + chatStreamFn: (opts: any) => { + callCount++ + const toolNames = opts.tools?.map((t: any) => t.name) || [] + + if (callCount === 1) { + // First call: LLM tries to call getWeather without discovering it + return (async function* () { + yield ev.runStarted() + yield ev.toolStart('call_weather_bad', 'getWeather') + yield ev.toolArgs('call_weather_bad', '{"city":"NYC"}') + yield ev.runFinished('tool_calls') + })() + } else if (callCount === 2) { + // Second call: LLM discovers getWeather + return (async function* () { + yield ev.runStarted() + yield ev.toolStart('call_disc', '__lazy__tool__discovery__') + yield ev.toolArgs( + 'call_disc', + JSON.stringify({ toolNames: ['getWeather'] }), + ) + yield ev.runFinished('tool_calls') + })() + } else if (callCount === 3 && toolNames.includes('getWeather')) { + // Third call: LLM now calls getWeather successfully + return (async function* () { + yield ev.runStarted() + yield ev.toolStart('call_weather_ok', 'getWeather') + yield ev.toolArgs('call_weather_ok', '{"city":"NYC"}') + yield ev.runFinished('tool_calls') + })() + } else { + // Fourth call: final text + return (async function* () { + yield ev.runStarted() + yield ev.textStart() + yield ev.textContent('72F in NYC') + yield ev.textEnd() + yield ev.runFinished('stop') + })() + } + }, + }) + + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'Weather in NYC?' }], + tools: [lazyServerTool('getWeather', weatherExecute)], + }) + + const chunks = await collectChunks(stream as AsyncIterable) + + // The first tool call result should contain a "must be discovered first" error + const toolEndChunks = chunks.filter( + (c) => c.type === 'TOOL_CALL_END', + ) as Array + const errorResult = toolEndChunks.find( + (c: any) => + c.toolName === 'getWeather' && + c.result && + c.result.includes('must be discovered first'), + ) + expect(errorResult).toBeDefined() + + // Eventually getWeather should be executed successfully + expect(weatherExecute).toHaveBeenCalledTimes(1) + }) + + it('should not create discovery tool when no lazy tools exist', async () => { + const executeSpy = vi.fn().mockReturnValue({ result: 'ok' }) + + const { adapter, calls } = createMockAdapter({ + iterations: [ + [ev.runStarted(), ev.textContent('Hi'), ev.runFinished('stop')], + ], + }) + + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'Hi' }], + tools: [serverTool('normalTool', executeSpy)], + }) + + await collectChunks(stream as AsyncIterable) + + // No __lazy__tool__discovery__ should appear in the tools sent to the adapter + const toolNames = calls[0].tools.map((t: any) => t.name) + expect(toolNames).not.toContain('__lazy__tool__discovery__') + expect(toolNames).toContain('normalTool') + }) + }) }) diff --git a/packages/typescript/ai/tests/lazy-tool-manager.test.ts b/packages/typescript/ai/tests/lazy-tool-manager.test.ts new file mode 100644 index 000000000..c62830da6 --- /dev/null +++ b/packages/typescript/ai/tests/lazy-tool-manager.test.ts @@ -0,0 +1,364 @@ +import { describe, it, expect } from 'vitest' +import type { Tool } from '../src/types' +import { LazyToolManager } from '../src/activities/chat/tools/lazy-tool-manager' + +const DISCOVERY_TOOL_NAME = '__lazy__tool__discovery__' + +function makeTool( + name: string, + opts?: { lazy?: boolean; execute?: (args: any) => any }, +): Tool { + return { + name, + description: `Tool: ${name}`, + execute: opts?.execute ?? (async (args: any) => args), + lazy: opts?.lazy, + } +} + +/** Find the discovery tool from the active tools list */ +function findDiscoveryTool(manager: LazyToolManager): Tool | undefined { + return manager.getActiveTools().find((t) => t.name === DISCOVERY_TOOL_NAME) +} + +describe('LazyToolManager', () => { + describe('construction and separation', () => { + it('returns all tools unchanged when none are lazy', () => { + const tools = [makeTool('a'), makeTool('b')] + const manager = new LazyToolManager(tools, []) + + const active = manager.getActiveTools() + expect(active).toHaveLength(2) + expect(active.map((t) => t.name)).toEqual(['a', 'b']) + }) + + it('does not create discovery tool when no tools are lazy', () => { + const tools = [makeTool('a'), makeTool('b')] + const manager = new LazyToolManager(tools, []) + + const discovery = findDiscoveryTool(manager) + expect(discovery).toBeUndefined() + }) + + it('separates lazy tools from eager tools', () => { + const tools = [ + makeTool('eager1'), + makeTool('lazy1', { lazy: true }), + makeTool('eager2'), + makeTool('lazy2', { lazy: true }), + ] + const manager = new LazyToolManager(tools, []) + + const active = manager.getActiveTools() + const names = active.map((t) => t.name) + + // Eager tools should be present + expect(names).toContain('eager1') + expect(names).toContain('eager2') + + // Lazy tools should NOT be present (not yet discovered) + expect(names).not.toContain('lazy1') + expect(names).not.toContain('lazy2') + + // Discovery tool should be present + expect(names).toContain(DISCOVERY_TOOL_NAME) + + // Total: 2 eager + 1 discovery = 3 + expect(active).toHaveLength(3) + }) + }) + + describe('discovery tool', () => { + it('has correct name and description listing tool names', () => { + const tools = [ + makeTool('lazyA', { lazy: true }), + makeTool('lazyB', { lazy: true }), + ] + const manager = new LazyToolManager(tools, []) + + const discovery = findDiscoveryTool(manager) + expect(discovery).toBeDefined() + expect(discovery!.name).toBe(DISCOVERY_TOOL_NAME) + expect(discovery!.description).toContain('lazyA') + expect(discovery!.description).toContain('lazyB') + }) + + it('discovers valid tools and returns descriptions + schemas', async () => { + const tools = [ + { + name: 'weather', + description: 'Get weather info', + lazy: true, + execute: async () => ({}), + inputSchema: { + type: 'object', + properties: { + location: { type: 'string', description: 'City name' }, + }, + required: ['location'], + }, + } satisfies Tool, + ] + const manager = new LazyToolManager(tools, []) + + const discovery = findDiscoveryTool(manager) + const result = await discovery!.execute!({ + toolNames: ['weather'], + }) + + expect(result.tools).toHaveLength(1) + expect(result.tools[0].name).toBe('weather') + expect(result.tools[0].description).toBe('Get weather info') + expect(result.tools[0].inputSchema).toBeDefined() + expect(result.errors).toBeUndefined() + }) + + it('returns errors for unknown tool names', async () => { + const tools = [makeTool('lazyA', { lazy: true })] + const manager = new LazyToolManager(tools, []) + + const discovery = findDiscoveryTool(manager) + const result = await discovery!.execute!({ + toolNames: ['nonexistent'], + }) + + expect(result.tools).toHaveLength(0) + expect(result.errors).toBeDefined() + expect(result.errors).toHaveLength(1) + expect(result.errors![0]).toContain('nonexistent') + }) + + it('handles mix of valid and invalid tool names', async () => { + const tools = [ + makeTool('lazyA', { lazy: true }), + makeTool('lazyB', { lazy: true }), + ] + const manager = new LazyToolManager(tools, []) + + const discovery = findDiscoveryTool(manager) + const result = await discovery!.execute!({ + toolNames: ['lazyA', 'bogus', 'lazyB'], + }) + + expect(result.tools).toHaveLength(2) + expect(result.tools.map((t: any) => t.name)).toEqual(['lazyA', 'lazyB']) + expect(result.errors).toHaveLength(1) + expect(result.errors![0]).toContain('bogus') + }) + }) + + describe('state tracking', () => { + it('includes discovered tools in getActiveTools after discovery', async () => { + const tools = [ + makeTool('eager1'), + makeTool('lazyA', { lazy: true }), + makeTool('lazyB', { lazy: true }), + ] + const manager = new LazyToolManager(tools, []) + + // Discover lazyA via the discovery tool + const discovery = findDiscoveryTool(manager) + await discovery!.execute!({ toolNames: ['lazyA'] }) + + const active = manager.getActiveTools() + const names = active.map((t) => t.name) + + expect(names).toContain('eager1') + expect(names).toContain('lazyA') + expect(names).not.toContain('lazyB') + // Discovery tool should still be present since lazyB remains undiscovered + expect(names).toContain(DISCOVERY_TOOL_NAME) + }) + + it('resets hasNewlyDiscoveredTools after getActiveTools', async () => { + const tools = [makeTool('lazyA', { lazy: true })] + const manager = new LazyToolManager(tools, []) + + expect(manager.hasNewlyDiscoveredTools()).toBe(false) + + const discovery = findDiscoveryTool(manager) + await discovery!.execute!({ toolNames: ['lazyA'] }) + expect(manager.hasNewlyDiscoveredTools()).toBe(true) + + manager.getActiveTools() + expect(manager.hasNewlyDiscoveredTools()).toBe(false) + }) + + it('removes discovery tool when all lazy tools discovered', async () => { + const tools = [ + makeTool('lazyA', { lazy: true }), + makeTool('lazyB', { lazy: true }), + ] + const manager = new LazyToolManager(tools, []) + + const discovery = findDiscoveryTool(manager) + await discovery!.execute!({ + toolNames: ['lazyA', 'lazyB'], + }) + + const active = manager.getActiveTools() + const names = active.map((t) => t.name) + + expect(names).toContain('lazyA') + expect(names).toContain('lazyB') + expect(names).not.toContain(DISCOVERY_TOOL_NAME) + }) + + it('correctly identifies undiscovered lazy tools via isUndiscoveredLazyTool', async () => { + const tools = [ + makeTool('eager1'), + makeTool('lazyA', { lazy: true }), + makeTool('lazyB', { lazy: true }), + ] + const manager = new LazyToolManager(tools, []) + + expect(manager.isUndiscoveredLazyTool('lazyA')).toBe(true) + expect(manager.isUndiscoveredLazyTool('lazyB')).toBe(true) + expect(manager.isUndiscoveredLazyTool('eager1')).toBe(false) + expect(manager.isUndiscoveredLazyTool('nonexistent')).toBe(false) + + const discovery = findDiscoveryTool(manager) + await discovery!.execute!({ toolNames: ['lazyA'] }) + + expect(manager.isUndiscoveredLazyTool('lazyA')).toBe(false) + expect(manager.isUndiscoveredLazyTool('lazyB')).toBe(true) + }) + + it('provides helpful error message for undiscovered tools', () => { + const tools = [makeTool('lazyA', { lazy: true })] + const manager = new LazyToolManager(tools, []) + + const errorMsg = manager.getUndiscoveredToolError('lazyA') + expect(errorMsg).toContain('lazyA') + expect(errorMsg).toContain(DISCOVERY_TOOL_NAME) + expect(errorMsg).toContain('must be discovered first') + }) + }) + + describe('message history scanning', () => { + it('pre-populates discovered tools from message history', () => { + const tools = [ + makeTool('lazyA', { lazy: true }), + makeTool('lazyB', { lazy: true }), + makeTool('lazyC', { lazy: true }), + ] + const messages = [ + { role: 'user', content: 'hello' }, + { + role: 'assistant', + content: null, + toolCalls: [ + { + id: 'tc_1', + type: 'function', + function: { + name: DISCOVERY_TOOL_NAME, + arguments: JSON.stringify({ toolNames: ['lazyA'] }), + }, + }, + ], + }, + { + role: 'tool', + content: JSON.stringify({ + tools: [ + { + name: 'lazyA', + description: 'Tool: lazyA', + inputSchema: {}, + }, + ], + }), + toolCallId: 'tc_1', + }, + ] + + const manager = new LazyToolManager(tools, messages as any) + + const active = manager.getActiveTools() + const names = active.map((t) => t.name) + + // lazyA should already be discovered from history + expect(names).toContain('lazyA') + // lazyB and lazyC still undiscovered + expect(names).not.toContain('lazyB') + expect(names).not.toContain('lazyC') + // Discovery tool should still exist + expect(names).toContain(DISCOVERY_TOOL_NAME) + }) + + it('handles malformed discovery results in history gracefully', () => { + const tools = [makeTool('lazyA', { lazy: true })] + const messages = [ + { + role: 'assistant', + content: null, + toolCalls: [ + { + id: 'tc_bad', + type: 'function', + function: { + name: DISCOVERY_TOOL_NAME, + arguments: '{}', + }, + }, + ], + }, + { + role: 'tool', + content: 'THIS IS NOT VALID JSON{{{', + toolCallId: 'tc_bad', + }, + ] + + // Should not throw + expect(() => new LazyToolManager(tools, messages as any)).not.toThrow() + + const manager = new LazyToolManager(tools, messages as any) + // lazyA should still be undiscovered + expect(manager.isUndiscoveredLazyTool('lazyA')).toBe(true) + }) + + it('ignores discovery results for tools not in current tool list', () => { + const tools = [makeTool('lazyA', { lazy: true })] + const messages = [ + { + role: 'assistant', + content: null, + toolCalls: [ + { + id: 'tc_1', + type: 'function', + function: { + name: DISCOVERY_TOOL_NAME, + arguments: JSON.stringify({ + toolNames: ['removedTool'], + }), + }, + }, + ], + }, + { + role: 'tool', + content: JSON.stringify({ + tools: [ + { + name: 'removedTool', + description: 'This tool no longer exists', + inputSchema: {}, + }, + ], + }), + toolCallId: 'tc_1', + }, + ] + + const manager = new LazyToolManager(tools, messages as any) + + // removedTool should not appear in discovered set (it's not in lazyToolMap) + expect(manager.isUndiscoveredLazyTool('removedTool')).toBe(false) + // lazyA is still undiscovered + expect(manager.isUndiscoveredLazyTool('lazyA')).toBe(true) + }) + }) +}) diff --git a/packages/typescript/ai/tests/tool-definition.test.ts b/packages/typescript/ai/tests/tool-definition.test.ts index 2d37178c9..642992c68 100644 --- a/packages/typescript/ai/tests/tool-definition.test.ts +++ b/packages/typescript/ai/tests/tool-definition.test.ts @@ -131,6 +131,37 @@ describe('toolDefinition', () => { expect(clientTool.needsApproval).toBe(true) }) + it('should preserve lazy flag', () => { + const tool = toolDefinition({ + name: 'discoverableWeather', + description: 'Get weather (lazy)', + lazy: true, + inputSchema: z.object({ + location: z.string(), + }), + }) + + expect(tool.lazy).toBe(true) + + const serverTool = tool.server(async (_args) => ({ temp: 72 })) + expect(serverTool.lazy).toBe(true) + + const clientTool = tool.client() + expect(clientTool.lazy).toBe(true) + }) + + it('should default lazy to undefined when not specified', () => { + const tool = toolDefinition({ + name: 'eagerTool', + description: 'A normal tool', + }) + + expect(tool.lazy).toBeUndefined() + + const serverTool = tool.server(async () => ({})) + expect(serverTool.lazy).toBeUndefined() + }) + it('should preserve metadata', () => { const tool = toolDefinition({ name: 'customTool',