Skip to content

Commit fbd3762

Browse files
authored
feat(code-mode): lazy tool support (progressive disclosure) (#726)
* feat(ai): add shared LazyToolsConfig and thread it through chat lazy discovery Introduces a single optional LazyToolsConfig ({ includeDescription: 'none' | 'first-sentence' | 'full' }, default 'none') plus shared renderLazyCatalogEntry/firstSentence helpers. chat() now accepts lazyToolsConfig and threads it into LazyToolManager, which renders the discovery-tool catalog accordingly. The 'none' default is byte-identical to the previous names-only output. * feat(code-mode): add lazy tool support with discover_tools catalog Tools marked `lazy: true` are kept out of the execute_typescript system prompt (their type stubs omitted) and listed in a "Discoverable APIs" catalog instead. A new discover_tools companion tool returns each lazy tool signature on demand; lookups tolerate the optional external_ prefix so the model can pass the catalog name verbatim. All tool bindings are still injected into the sandbox (documentation-only lazy). createCodeMode now returns { tool, discoveryTool, tools, systemPrompt } (additive; tool/systemPrompt unchanged) and honors the shared lazyToolsConfig. * test(e2e): cover chat lazyToolsConfig discovery catalog Adds a wire-journal E2E (api.lazy-tools-wire route + fixture + spec) that asserts the lazy discovery tool's wire description renders names-only for 'none', name plus first sentence for 'first-sentence', and the full description for 'full'. routeTree.gen.ts is the router plugin's auto-regen for the new route. * docs: add code-mode lazy tools page, update skills, add changeset New docs/code-mode/lazy-tools.md (with config.json nav entry) covering lazy tools, the discover_tools flow, and lazyToolsConfig.includeDescription for both Code Mode and plain chat(). Updates the ai-code-mode and ai-core/tool-calling agent skills to document the new API, and adds the changeset (minor for @tanstack/ai and @tanstack/ai-code-mode). * docs: document lazyToolsConfig on the lazy tool discovery page The chat lazy-tool discovery page now documents the new optional lazyToolsConfig.includeDescription ('none' default / 'first-sentence' / 'full') that tunes the pre-discovery catalog, with a cross-link to the Code Mode lazy tools page. Bumps the example model ids to gpt-5.2 and sets updatedAt on the docs config entry. * chore: address PR review feedback - Use gpt-5.5 (newest OpenAI chat model in model-meta) in docs, skill, and the e2e route instead of gpt-5.2. - e2e: drop the beforeEach DELETE /v1/_requests journal reset; the spec already isolates per-test via the X-Test-Id header, so the global reset only risked racing adjacent parallel specs. - Alphabetize vitest named imports in the two new test files (sort-imports). - Add `text` language identifiers to the Discoverable APIs example fences in the lazy-tools doc and code-mode skill (markdownlint). * fix(docs): make lazy-tools snippets type-check under kiira after main merge Main introduced the kiira docs type-checker. Make the lazy-tools doc snippets self-contained (real imports, endpoint-wrapped messages, relative tool imports) and group the sequential server-setup example. Exclude docs/superpowers/** planning artifacts (illustrative pseudo-code, not published examples).
1 parent e3ee4ae commit fbd3762

27 files changed

Lines changed: 1177 additions & 32 deletions

.changeset/code-mode-lazy-tools.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
---
2+
'@tanstack/ai': minor
3+
'@tanstack/ai-code-mode': minor
4+
---
5+
6+
Add lazy tool support (progressive disclosure) to Code Mode. Tools marked `lazy: true` are kept out of the `execute_typescript` system prompt and listed in a discoverable catalog; the model fetches their TypeScript signatures on demand via a new `discover_tools` tool. A shared optional `lazyToolsConfig` (`includeDescription: 'none' | 'first-sentence' | 'full'`) tunes the catalog detail for both `chat()` and `createCodeMode()`. `createCodeMode` now also returns `discoveryTool` and a `tools` array (backward compatible — `tool` and `systemPrompt` are unchanged).

docs/code-mode/lazy-tools.md

Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
---
2+
title: Lazy Tools
3+
id: lazy-tools
4+
order: 5
5+
description: "Keep large tool catalogs out of the Code Mode system prompt with lazy tools — the model fetches TypeScript signatures on demand via a discover_tools call."
6+
keywords:
7+
- tanstack ai
8+
- code mode
9+
- lazy tools
10+
- discover_tools
11+
- progressive disclosure
12+
- prompt size
13+
- tool catalog
14+
---
15+
16+
Large tool catalogs bloat the `execute_typescript` system prompt. Every tool you pass to `createCodeMode` becomes a full TypeScript type stub in that prompt — and at 50+ tools, those stubs can push the effective prompt into the tens of thousands of tokens before the model has even seen your user message.
17+
18+
Lazy tools fix this with **progressive disclosure**: mark rarely-used tools `lazy: true` and they are withheld from the initial system prompt. The model sees only their names in a short "Discoverable APIs" catalog. When it needs one, it calls the `discover_tools` sibling tool to fetch the TypeScript signature on demand, then uses it inside `execute_typescript`. All sandbox bindings are always injected — lazy only defers _documentation_, not callability.
19+
20+
## Marking a Tool Lazy
21+
22+
Add `lazy: true` to the `toolDefinition` config for any tool you want to defer:
23+
24+
```typescript group=lazy-tools
25+
import { toolDefinition } from "@tanstack/ai";
26+
import { z } from "zod";
27+
28+
// Always eager — documented upfront
29+
const fetchWeather = toolDefinition({
30+
name: "fetchWeather",
31+
description: "Get current weather for a city",
32+
inputSchema: z.object({ location: z.string() }),
33+
outputSchema: z.object({ temperature: z.number(), condition: z.string() }),
34+
}).server(async ({ location }) => {
35+
const res = await fetch(`https://api.weather.example/v1?city=${location}`);
36+
return res.json();
37+
});
38+
39+
// Lazy — kept out of the system prompt until discovered
40+
const fetchArchive = toolDefinition({
41+
name: "fetchArchive",
42+
description: "Retrieve historical weather archive data for a date range",
43+
inputSchema: z.object({
44+
location: z.string(),
45+
from: z.string(),
46+
to: z.string(),
47+
}),
48+
outputSchema: z.array(z.object({ date: z.string(), temperature: z.number() })),
49+
lazy: true,
50+
}).server(async ({ location, from, to }) => {
51+
const res = await fetch(
52+
`https://api.weather.example/v1/archive?city=${location}&from=${from}&to=${to}`
53+
);
54+
return res.json();
55+
});
56+
```
57+
58+
Eager tools continue to receive full type stubs in the system prompt. Lazy tools appear only by name.
59+
60+
## Server Setup
61+
62+
Pass both eager and lazy tools to `createCodeMode`. When at least one tool is lazy, `createCodeMode` also returns a `discover_tools` sibling tool — include it in the `tools` array you pass to `chat()`:
63+
64+
```typescript group=lazy-tools
65+
// server/route.ts
66+
import { chat, maxIterations, toServerSentEventsStream } from "@tanstack/ai";
67+
import { createCodeMode } from "@tanstack/ai-code-mode";
68+
import { createNodeIsolateDriver } from "@tanstack/ai-isolate-node";
69+
import { openaiText } from "@tanstack/ai-openai";
70+
71+
const { tools, systemPrompt } = createCodeMode({
72+
driver: createNodeIsolateDriver(),
73+
tools: [fetchWeather, fetchArchive], // fetchArchive is lazy
74+
});
75+
76+
// tools is [execute_typescript, discover_tools]
77+
// — discover_tools is included automatically because fetchArchive is lazy
78+
79+
export async function POST(req: Request) {
80+
const { messages } = await req.json();
81+
82+
const stream = chat({
83+
adapter: openaiText("gpt-5.5"),
84+
systemPrompts: ["You are a helpful weather assistant.", systemPrompt],
85+
tools: [...tools],
86+
messages,
87+
agentLoopStrategy: maxIterations(10),
88+
});
89+
90+
return toServerSentEventsStream(stream);
91+
}
92+
```
93+
94+
`createCodeMode` returns `{ tool, discoveryTool, tools, systemPrompt }`:
95+
96+
| Field | Type | Description |
97+
|-------|------|-------------|
98+
| `tool` | `ServerTool` | The `execute_typescript` tool (backward compatible) |
99+
| `discoveryTool` | `ServerTool \| null` | The `discover_tools` tool, or `null` when there are no lazy tools |
100+
| `tools` | `Array<ServerTool>` | `[tool]` or `[tool, discoveryTool]` — spread into `chat({ tools })` |
101+
| `systemPrompt` | `string` | The matching system prompt |
102+
103+
If no tools are lazy, `discoveryTool` is `null` and `tools` contains only `execute_typescript`.
104+
105+
## The `discover_tools` Flow
106+
107+
When the model encounters a task that requires a lazy tool, it:
108+
109+
1. Calls `discover_tools` with the tool name (bare name, no `external_` prefix).
110+
2. Receives the TypeScript type stub and description for that tool.
111+
3. Writes `execute_typescript` code using the now-documented `external_fetchArchive(...)` call.
112+
113+
The bindings are always injected into the sandbox — discovering a tool only retrieves documentation, it does not enable the binding. The model could call `external_fetchArchive` without discovering it first, but it would be writing blind without the type signature.
114+
115+
## Tuning the Discoverable APIs Catalog
116+
117+
By default, lazy tools appear in the system prompt as bare names with no description:
118+
119+
```text
120+
### Discoverable APIs
121+
122+
- external_fetchArchive
123+
- external_runReport
124+
- external_exportData
125+
```
126+
127+
If you want the model to have a hint about what each tool does before deciding whether to discover it, use `lazyToolsConfig.includeDescription`:
128+
129+
```typescript
130+
import { createCodeMode } from "@tanstack/ai-code-mode";
131+
import { createNodeIsolateDriver } from "@tanstack/ai-isolate-node";
132+
import {
133+
fetchWeather,
134+
fetchArchive,
135+
runReport,
136+
exportData,
137+
} from "./tools";
138+
139+
const { tools, systemPrompt } = createCodeMode({
140+
driver: createNodeIsolateDriver(),
141+
tools: [fetchWeather, fetchArchive, runReport, exportData],
142+
lazyToolsConfig: {
143+
includeDescription: "first-sentence", // 'none' | 'first-sentence' | 'full'
144+
},
145+
});
146+
```
147+
148+
With `'first-sentence'` the catalog becomes:
149+
150+
```text
151+
### Discoverable APIs
152+
153+
- external_fetchArchive — Retrieve historical weather archive data for a date range.
154+
- external_runReport — Generate a summary report for a given time period.
155+
- external_exportData — Export query results to CSV or JSON format.
156+
```
157+
158+
| Value | Effect |
159+
|-------|--------|
160+
| `'none'` (default) | Bare names only — smallest possible prompt addition |
161+
| `'first-sentence'` | Name plus the first sentence of the tool's description |
162+
| `'full'` | Name plus the complete description |
163+
164+
The full type stub and input/output schema are always returned on discovery — `includeDescription` only affects the pre-discovery catalog.
165+
166+
## Lazy Tools with Plain `chat()`
167+
168+
The same `lazyToolsConfig` option works for lazy tools used directly with `chat()`, outside of Code Mode. Tools marked `lazy: true` are withheld from the `__lazy__tool__discovery__` catalog description until the model calls for them. Pass `lazyToolsConfig` directly to `chat()`:
169+
170+
```typescript
171+
import { chat, maxIterations, toServerSentEventsStream } from "@tanstack/ai";
172+
import { openaiText } from "@tanstack/ai-openai";
173+
import { fetchWeather, fetchArchive, runReport } from "./tools";
174+
175+
// Non-code-mode: lazy tools in a regular chat agent
176+
export async function POST(req: Request) {
177+
const { messages } = await req.json();
178+
179+
const stream = chat({
180+
adapter: openaiText("gpt-5.5"),
181+
messages,
182+
tools: [fetchWeather, fetchArchive, runReport],
183+
lazyToolsConfig: {
184+
includeDescription: "first-sentence",
185+
},
186+
agentLoopStrategy: maxIterations(10),
187+
});
188+
189+
return toServerSentEventsStream(stream);
190+
}
191+
```
192+
193+
The `includeDescription` behavior is identical — `'none'` lists bare tool names, `'first-sentence'` appends the first sentence, `'full'` appends the complete description.
194+
195+
## Tips
196+
197+
- **Start with `'none'`.** The bare-names catalog is enough for models that reason well about tool names. Add `'first-sentence'` only if the model frequently discovers irrelevant tools.
198+
- **Lazy tools are always callable.** Their `external_*` bindings are injected into the sandbox regardless of whether the model has called `discover_tools`. Discovery only reveals documentation.
199+
- **Use `discoveryTool` for observability.** You can inspect `discoveryTool.name` (`"discover_tools"`) to confirm the tool is wired up, or log its calls for analytics.
200+
- **Partition by frequency, not capability.** Mark tools lazy when they are rarely needed for a typical request. Core tools that most requests use should stay eager.
201+
202+
## Next Steps
203+
204+
- [Code Mode](./code-mode) — Core Code Mode setup and API reference
205+
- [Code Mode with Skills](./code-mode-with-skills) — Persistent reusable skill libraries
206+
- [Isolate Drivers](./code-mode-isolates) — Compare Node, QuickJS, and Cloudflare sandbox runtimes

docs/config.json

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,8 @@
107107
{
108108
"label": "Lazy Tool Discovery",
109109
"to": "tools/lazy-tool-discovery",
110-
"addedAt": "2026-04-15"
110+
"addedAt": "2026-04-15",
111+
"updatedAt": "2026-06-08"
111112
}
112113
]
113114
},
@@ -220,6 +221,11 @@
220221
"label": "Code Mode Isolate Drivers",
221222
"to": "code-mode/code-mode-isolates",
222223
"addedAt": "2026-04-15"
224+
},
225+
{
226+
"label": "Lazy Tools",
227+
"to": "code-mode/lazy-tools",
228+
"addedAt": "2026-06-08"
223229
}
224230
]
225231
},

docs/tools/lazy-tool-discovery.md

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,49 @@ async function handleRequest(request: Request) {
9797
}
9898
```
9999

100+
## Controlling the Discovery Catalog
101+
102+
By default, the `__lazy__tool__discovery__` tool's description lists only the
103+
**names** of available lazy tools. The optional `lazyToolsConfig` on `chat()`
104+
controls how much of each lazy tool's description appears in that pre-discovery
105+
catalog:
106+
107+
```typescript
108+
import { chat, toServerSentEventsResponse } from "@tanstack/ai";
109+
import { openaiText } from "@tanstack/ai-openai";
110+
import { getProducts, searchProducts, compareProducts } from "./tools";
111+
112+
async function handleRequest(request: Request) {
113+
const { messages } = await request.json();
114+
const stream = chat({
115+
adapter: openaiText("gpt-5.5"),
116+
messages,
117+
tools: [getProducts, searchProducts, compareProducts],
118+
lazyToolsConfig: {
119+
// 'none' (default) | 'first-sentence' | 'full'
120+
includeDescription: "first-sentence",
121+
},
122+
});
123+
124+
return toServerSentEventsResponse(stream);
125+
}
126+
```
127+
128+
| `includeDescription` | Catalog entry for `searchProducts` |
129+
| -------------------- | ----------------------------------------------- |
130+
| `'none'` (default) | `searchProducts` |
131+
| `'first-sentence'` | `searchProducts — Search products by keyword.` |
132+
| `'full'` | `searchProducts — <full description>` |
133+
134+
This only affects the **pre-discovery** catalog. Regardless of the setting, the
135+
discovery tool's result always returns each tool's full description and argument
136+
schema — `includeDescription` just tunes how much the LLM sees before it
137+
decides what to discover. The default `'none'` keeps the catalog as lean as
138+
possible.
139+
140+
`lazyToolsConfig` is optional and the same option is accepted by Code Mode's
141+
`createCodeMode()` — see [Code Mode Lazy Tools](../code-mode/lazy-tools).
142+
100143
## When to Use Lazy Tools
101144

102145
Lazy tools are useful when:

kiira.config.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,10 @@ export default defineConfig({
1414
// ai-angular exposes a source-resolvable entry.
1515
'docs/api/ai-angular.md',
1616
'docs/getting-started/quick-start-angular.md',
17+
// docs/superpowers/** are internal planning/spec artifacts (design docs and
18+
// implementation plans), not published, curated examples. Their snippets are
19+
// illustrative pseudo-code, not meant to compile against package source.
20+
'docs/superpowers/**',
1721
],
1822
defaultValidate: 'type',
1923
languages: ['ts', 'tsx', 'js', 'jsx'],

packages/ai-code-mode/skills/ai-code-mode/SKILL.md

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ sources:
1515
- 'TanStack/ai:docs/code-mode/code-mode-isolates.md'
1616
- 'TanStack/ai:docs/code-mode/code-mode-with-skills.md'
1717
- 'TanStack/ai:docs/code-mode/client-integration.md'
18+
- 'TanStack/ai:docs/code-mode/lazy-tools.md'
1819
---
1920

2021
> **Note**: This skill requires familiarity with ai-core and ai-core/chat-experience. Code Mode is always used on top of a chat experience.
@@ -328,6 +329,84 @@ Skill-specific events (when using `codeModeWithSkills`):
328329
| `code_mode:skill_error` | Skill failed | `skill`, `error`, `duration` |
329330
| `skill:registered` | New skill saved | `id`, `name`, `description` |
330331

332+
### 4. Lazy Tools
333+
334+
When a large tool catalog would bloat the `execute_typescript` system prompt, mark low-priority tools `lazy: true`. Lazy tools are kept out of the full type-stub documentation and listed in a compact "Discoverable APIs" catalog instead. All sandbox bindings are always injected — `lazy` defers documentation, not callability.
335+
336+
**Marking a tool lazy:**
337+
338+
```typescript
339+
import { toolDefinition } from '@tanstack/ai'
340+
import { z } from 'zod'
341+
342+
const rarelyUsedTool = toolDefinition({
343+
name: 'fetchStocks',
344+
description: 'Get stock prices for a ticker. Returns a price quote.',
345+
inputSchema: z.object({ ticker: z.string() }),
346+
outputSchema: z.object({ price: z.number() }),
347+
lazy: true, // <-- opt out of full system-prompt documentation
348+
}).server(async ({ ticker }) => {
349+
// ...
350+
return { price: 0 }
351+
})
352+
```
353+
354+
**`createCodeMode` return shape:**
355+
356+
`createCodeMode()` returns `{ tool, discoveryTool, tools, systemPrompt }`. When lazy tools are present `discoveryTool` is a `discover_tools` server tool; otherwise it is `null`. Always spread `tools` (not just `tool`) into `chat()` so the discovery tool is registered:
357+
358+
```typescript
359+
import { chat } from '@tanstack/ai'
360+
import { createCodeMode } from '@tanstack/ai-code-mode'
361+
import { createNodeIsolateDriver } from '@tanstack/ai-isolate-node'
362+
import { openaiText } from '@tanstack/ai-openai'
363+
364+
const { tools, systemPrompt } = createCodeMode({
365+
driver: createNodeIsolateDriver(),
366+
tools: [eagerTool, rarelyUsedTool], // rarelyUsedTool has lazy: true
367+
})
368+
369+
const stream = chat({
370+
adapter: openaiText('gpt-5.5'),
371+
systemPrompts: ['You are a helpful assistant.', systemPrompt],
372+
tools: [...tools, ...otherTools], // spread tools, not just tool
373+
messages,
374+
})
375+
```
376+
377+
`tools` equals `[tool]` when there are no lazy tools (backward compatible) and `[tool, discoveryTool]` when lazy tools exist.
378+
379+
**`discover_tools` flow:**
380+
381+
When the model encounters a lazy tool it has not seen before, it calls `discover_tools` with the bare name (no `external_` prefix). The tool returns each requested tool's TypeScript type stub and description. The model then writes correctly-typed `external_<name>` calls inside `execute_typescript`.
382+
383+
```text
384+
Model sees: "Discoverable APIs: external_fetchStocks"
385+
Model calls: discover_tools({ toolNames: ["fetchStocks"] })
386+
Response: { tools: [{ name: "external_fetchStocks", description: "...", typeStub: "declare function external_fetchStocks(...)" }] }
387+
Model writes inside execute_typescript: const result = await external_fetchStocks({ ticker: "AAPL" })
388+
```
389+
390+
**`lazyToolsConfig.includeDescription`:**
391+
392+
Control how much of each lazy tool's description appears in the Discoverable APIs catalog (the pre-discovery list):
393+
394+
| Value | Catalog entry |
395+
| ------------------ | ----------------------------------------------------------------- |
396+
| `'none'` | `external_fetchStocks` (name only — default) |
397+
| `'first-sentence'` | `external_fetchStocks — Get stock prices.` |
398+
| `'full'` | `external_fetchStocks — Get stock prices. Returns a price quote.` |
399+
400+
```typescript
401+
const { tools, systemPrompt } = createCodeMode({
402+
driver: createNodeIsolateDriver(),
403+
tools: [eagerTool, rarelyUsedTool],
404+
lazyToolsConfig: { includeDescription: 'first-sentence' },
405+
})
406+
```
407+
408+
The same `lazyToolsConfig` option is accepted by plain `chat()` for its own lazy-tool discovery catalog (see `ai-core/tool-calling/SKILL.md`).
409+
331410
## Common Mistakes
332411

333412
### CRITICAL: Passing API keys or secrets to the sandbox environment

0 commit comments

Comments
 (0)