feat(ai/core): lazy tool discovery for token-efficient large tool sets

packages/ai/src/generate-text/generate-text.ts

@@ -86,6 +86,7 @@ import { TypedToolError } from './tool-error';
 import { ToolOutput } from './tool-output';
 import { TypedToolResult } from './tool-result';
 import type { ToolSet } from '@ai-sdk/provider-utils';
+import { buildEffectiveTools } from './load-tool-schema';
 
 const originalGenerateId = createIdGenerator({
   prefix: 'aitxt',
@@ -466,6 +467,9 @@ export async function generateText<
     };
   }): Promise<GenerateTextResult<TOOLS, CONTEXT, OUTPUT>> {
   const model = resolveLanguageModel(modelArg);
+
+  const effectiveTools = buildEffectiveTools(tools);
+
   const createGlobalTelemetry = getGlobalTelemetryIntegration<any, OUTPUT>();
   const stopConditions = asArray(stopWhen);
 
@@ -565,7 +569,7 @@ export async function generateText<
         toolCalls: localApprovedToolApprovals.map(
           toolApproval => toolApproval.toolCall,
         ),
-        tools: tools as TOOLS,
+        tools: effectiveTools as TOOLS,
         telemetry,
         callId,
         messages: initialMessages,
@@ -605,7 +609,7 @@ export async function generateText<
         const modelOutput = await createToolModelOutput({
           toolCallId: output.toolCallId,
           input: output.input,
-          tool: tools?.[output.toolName],
+          tool: effectiveTools?.[output.toolName],
           output: output.type === 'tool-result' ? output.output : output.error,
           errorMode: output.type === 'tool-error' ? 'text' : 'none',
         });
@@ -717,7 +721,7 @@ export async function generateText<
           prepareStepResult?.experimental_context ?? experimental_context;
 
         const stepActiveTools = filterActiveTools({
-          tools,
+          tools: effectiveTools,
           activeTools: prepareStepResult?.activeTools ?? activeTools,
         });
 
@@ -807,7 +811,7 @@ export async function generateText<
             .map(toolCall =>
               parseToolCall({
                 toolCall,
-                tools,
+                tools: effectiveTools,
                 repairToolCall,
                 system,
                 messages: stepInputMessages,
@@ -825,7 +829,7 @@ export async function generateText<
             continue; // ignore invalid tool calls
           }
 
-          const tool = tools?.[toolCall.toolName];
+          const tool = effectiveTools?.[toolCall.toolName];
 
           if (tool == null) {
             // ignore tool calls for tools that are not available,
@@ -883,15 +887,15 @@ export async function generateText<
           toolCall => !toolCall.providerExecuted,
         );
 
-        if (tools != null) {
+        if (effectiveTools != null) {
           clientToolOutputs.push(
             ...(await executeTools({
               toolCalls: clientToolCalls.filter(
                 toolCall =>
                   !toolCall.invalid &&
                   toolApprovalRequests[toolCall.toolCallId] == null,
               ),
-              tools,
+              tools: effectiveTools,
               telemetry,
               callId,
               messages: stepInputMessages,
@@ -932,7 +936,7 @@ export async function generateText<
         // the client tool's result is sent back.
         for (const toolCall of stepToolCalls) {
           if (!toolCall.providerExecuted) continue;
-          const tool = tools?.[toolCall.toolName];
+          const tool = effectiveTools?.[toolCall.toolName];
           if (tool?.type === 'provider' && tool.supportsDeferredResults) {
             // Check if this tool call already has a result in the current response
             const hasResultInResponse = currentModelResponse.content.some(
@@ -961,14 +965,14 @@ export async function generateText<
           toolCalls: stepToolCalls,
           toolOutputs: clientToolOutputs,
           toolApprovalRequests: Object.values(toolApprovalRequests),
-          tools,
+          tools: effectiveTools,
         });
 
         // append to messages for potential next step:
         responseMessages.push(
           ...(await toResponseMessages({
             content: stepContent,
-            tools,
+            tools: effectiveTools,
           })),
         );
 

packages/ai/src/generate-text/index.ts

@@ -19,6 +19,7 @@ export {
   type GenerateTextOnToolCallStartCallback,
 } from './generate-text';
 export type { GenerateTextResult } from './generate-text-result';
+export { LOAD_TOOL_SCHEMA_NAME } from './load-tool-schema';
 export {
   DefaultGeneratedFile,
   type GeneratedFile as Experimental_GeneratedImage, // Image for backwards compatibility, TODO remove in v7

packages/ai/src/generate-text/load-tool-schema.test.ts

@@ -0,0 +1,228 @@
+import { describe, expect, it } from 'vitest';
+import { tool } from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+import {
+  buildEffectiveTools,
+  createLoadToolSchemaTool,
+  LOAD_TOOL_SCHEMA_NAME,
+} from './load-tool-schema';
+
+describe('LOAD_TOOL_SCHEMA_NAME', () => {
+  it("equals '__load_tool_schema__'", () => {
+    expect(LOAD_TOOL_SCHEMA_NAME).toBe('__load_tool_schema__');
+  });
+});
+
+describe('createLoadToolSchemaTool', () => {
+  it('returns tool with description listing lazy tool names', () => {
+    const lazyTools = {
+      getWeather: tool({
+        description: 'Get weather',
+        inputSchema: z.object({ city: z.string() }),
+        lazy: true,
+      }),
+      searchDocs: tool({
+        description: 'Search docs',
+        inputSchema: z.object({ query: z.string() }),
+        lazy: true,
+      }),
+    };
+
+    const loadTool = createLoadToolSchemaTool(lazyTools);
+
+    expect(loadTool.description).toContain('getWeather');
+    expect(loadTool.description).toContain('searchDocs');
+  });
+
+  it('execute returns full inputSchema for requested tools', async () => {
+    const lazyTools = {
+      getWeather: tool({
+        description: 'Get weather',
+        inputSchema: z.object({ city: z.string() }),
+        lazy: true,
+      }),
+    };
+
+    const loadTool = createLoadToolSchemaTool(lazyTools);
+
+    const result = await loadTool.execute!({ toolNames: ['getWeather'] }, {
+      toolCallId: 'test-id',
+      messages: [],
+      abortSignal: undefined,
+      experimental_context: undefined,
+    } as any);
+
+    expect(result).toBeDefined();
+    const parsed = typeof result === 'string' ? JSON.parse(result) : result;
+
+    // Should contain the full schema for getWeather with city property
+    const weatherSchema = parsed.getWeather ?? parsed['getWeather'];
+    expect(weatherSchema).toBeDefined();
+    expect(weatherSchema.inputSchema).toBeDefined();
+    expect(weatherSchema.inputSchema.properties).toHaveProperty('city');
+  });
+
+  it('execute returns skill text when present', async () => {
+    const lazyTools = {
+      myTool: tool({
+        description: 'My tool',
+        inputSchema: z.object({ x: z.number() }),
+        lazy: true,
+        skill: 'detailed usage info',
+      }),
+    };
+
+    const loadTool = createLoadToolSchemaTool(lazyTools);
+
+    const result = await loadTool.execute!({ toolNames: ['myTool'] }, {
+      toolCallId: 'test-id',
+      messages: [],
+      abortSignal: undefined,
+      experimental_context: undefined,
+    } as any);
+
+    const parsed = typeof result === 'string' ? JSON.parse(result) : result;
+    const toolResult = parsed.myTool ?? parsed['myTool'];
+    expect(toolResult).toBeDefined();
+    expect(toolResult.skill).toBe('detailed usage info');
+  });
+
+  it('execute returns inputExamples when present', async () => {
+    const lazyTools = {
+      myTool: tool({
+        description: 'My tool',
+        inputSchema: z.object({ city: z.string() }),
+        lazy: true,
+        inputExamples: [{ input: { city: 'NYC' } }],
+      }),
+    };
+
+    const loadTool = createLoadToolSchemaTool(lazyTools);
+
+    const result = await loadTool.execute!({ toolNames: ['myTool'] }, {
+      toolCallId: 'test-id',
+      messages: [],
+      abortSignal: undefined,
+      experimental_context: undefined,
+    } as any);
+
+    const parsed = typeof result === 'string' ? JSON.parse(result) : result;
+    const toolResult = parsed.myTool ?? parsed['myTool'];
+    expect(toolResult).toBeDefined();
+    expect(toolResult.inputExamples).toEqual([{ input: { city: 'NYC' } }]);
+  });
+
+  it('execute returns error entry for unknown tool names', async () => {
+    const lazyTools = {
+      myTool: tool({
+        description: 'My tool',
+        inputSchema: z.object({ x: z.number() }),
+        lazy: true,
+      }),
+    };
+
+    const loadTool = createLoadToolSchemaTool(lazyTools);
+
+    const result = await loadTool.execute!({ toolNames: ['nonexistent'] }, {
+      toolCallId: 'test-id',
+      messages: [],
+      abortSignal: undefined,
+      experimental_context: undefined,
+    } as any);
+
+    const parsed = typeof result === 'string' ? JSON.parse(result) : result;
+
+    expect(parsed.nonexistent).toEqual({
+      error: "Tool 'nonexistent' is not a lazy tool or does not exist",
+    });
+  });
+
+  it('execute handles mixed known and unknown names', async () => {
+    const lazyTools = {
+      validTool: tool({
+        description: 'Valid tool',
+        inputSchema: z.object({ name: z.string() }),
+        lazy: true,
+      }),
+    };
+
+    const loadTool = createLoadToolSchemaTool(lazyTools);
+
+    const result = await loadTool.execute!(
+      { toolNames: ['validTool', 'invalidTool'] },
+      {
+        toolCallId: 'test-id',
+        messages: [],
+        abortSignal: undefined,
+        experimental_context: undefined,
+      } as any,
+    );
+
+    const parsed = typeof result === 'string' ? JSON.parse(result) : result;
+
+    const validResult = parsed.validTool ?? parsed['validTool'];
+    expect(validResult).toBeDefined();
+    expect(validResult.inputSchema).toBeDefined();
+
+    expect(parsed.invalidTool).toEqual({
+      error: "Tool 'invalidTool' is not a lazy tool or does not exist",
+    });
+  });
+});
+
+describe('buildEffectiveTools', () => {
+  it('returns undefined when tools is undefined', () => {
+    expect(buildEffectiveTools(undefined)).toBeUndefined();
+  });
+
+  it('returns tools unchanged when no lazy tools exist', () => {
+    const tools = {
+      eagerTool: tool({
+        description: 'Eager',
+        inputSchema: z.object({}),
+      }),
+    };
+
+    const result = buildEffectiveTools(tools);
+    expect(result).toBe(tools);
+  });
+
+  it('injects __load_tool_schema__ when lazy tools exist', () => {
+    const tools = {
+      eagerTool: tool({
+        description: 'Eager',
+        inputSchema: z.object({}),
+      }),
+      lazyTool: tool({
+        description: 'Lazy',
+        inputSchema: z.object({ q: z.string() }),
+        lazy: true,
+      }),
+    };
+
+    const result = buildEffectiveTools(tools) as Record<string, unknown>;
+    expect(result).toBeDefined();
+    expect(result[LOAD_TOOL_SCHEMA_NAME]).toBeDefined();
+    expect(result.eagerTool).toBe(tools.eagerTool);
+    expect(result.lazyTool).toBe(tools.lazyTool);
+  });
+
+  it('does not override user-provided __load_tool_schema__', () => {
+    const customLoader = tool({
+      description: 'Custom loader',
+      inputSchema: z.object({}),
+    });
+
+    const tools = {
+      [LOAD_TOOL_SCHEMA_NAME]: customLoader,
+      lazyTool: tool({
+        description: 'Lazy',
+        inputSchema: z.object({}),
+        lazy: true,
+      }),
+    };
+
+    const result = buildEffectiveTools(tools);
+    expect(result![LOAD_TOOL_SCHEMA_NAME]).toBe(customLoader);
+  });
+});