fix(ai): keep lazy-tool discovery callable when re-requested after all tools are discovered (#795)

* fix(ai): keep lazy-tool discovery callable when re-requested after all tools are discovered

Once every lazy tool has been discovered, the synthetic
`__lazy__tool__discovery__` tool is dropped from the set advertised to the
model (a deliberate token optimization). But if the model re-requested
discovery anyway — common in long-context turns or when it overlooks that a
tool is already available — the call fell through to tool execution and came
back as `Unknown tool: __lazy__tool__discovery__`.

Decouple the advertised tool set from the executable one: the discovery tool
is now kept executable for the turn whenever a pending call references it, even
after it leaves the advertised set, so re-discovery returns the schemas again
instead of erroring. The advertised set is unchanged. Discovery is also now
idempotent — re-requesting an already-discovered tool returns its schema
without triggering a redundant tool-list refresh.

Also export `DISCOVERY_TOOL_NAME` so custom message-compaction logic can
reference the discovery tool by constant instead of hard-coding the string.

Closes #788

* test(ai): use DISCOVERY_TOOL_NAME constant and fix import order in lazy-tool tests

Address review feedback on #795: reorder the value import ahead of the
type-only import in lazy-tool-manager.test.ts, and replace the hard-coded
'__lazy__tool__discovery__' literals in the re-discovery regression tests with
the exported DISCOVERY_TOOL_NAME constant.

* ci: apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

Author: season179

.changeset/lazy-tool-rediscovery.md

@@ -0,0 +1,23 @@
+---
+'@tanstack/ai': patch
+---
+
+fix: don't error when an already-discovered lazy tool's discovery is re-requested
+
+Lazy tools use progressive disclosure: a synthetic `__lazy__tool__discovery__`
+tool is advertised so the model can reveal a lazy tool's schema on demand. Once
+**every** lazy tool has been discovered, that discovery tool is (intentionally)
+dropped from the set advertised to the model to save tokens. But if the model
+then re-requested discovery anyway — common in long-context or when it overlooks
+that a tool is already available — the call fell through to tool execution and
+came back as `Unknown tool: __lazy__tool__discovery__`.
+
+The discovery tool is now kept _executable_ for the turn whenever a pending call
+references it, even after it leaves the advertised set, so re-discovery returns
+the schemas again instead of erroring. The advertised set is unchanged. The
+discovery tool is also now idempotent — re-requesting an already-discovered tool
+returns its schema without triggering a redundant tool-list refresh.
+
+Additionally, `DISCOVERY_TOOL_NAME` is now exported from `@tanstack/ai` so custom
+message-compaction / history-trimming logic can reference the discovery tool by
+constant instead of hard-coding the string.

packages/ai/src/activities/chat/index.ts

@@ -1202,6 +1202,23 @@ class TextEngine<
     }
   }
 
+  /**
+   * Tools available for execution this turn. The discovery tool is dropped
+   * from the advertised set (`this.tools`) once every lazy tool is discovered,
+   * but a model may still re-request discovery; this widens execution lookup
+   * to include it so such calls don't fail with "Unknown tool". Centralised so
+   * both execution sites (`processToolCalls` and `checkForPendingToolCalls`)
+   * stay in sync.
+   */
+  private resolveExecutableTools(
+    toolCalls: ReadonlyArray<ToolCall>,
+  ): ReadonlyArray<AnyTool> {
+    return this.lazyToolManager.getExecutableTools(
+      this.tools,
+      toolCalls.map((tc) => tc.function.name),
+    )
+  }
+
   private async *checkForPendingToolCalls(): AsyncGenerator<
     StreamChunk,
     ToolPhaseResult,
@@ -1250,7 +1267,7 @@ class TextEngine<
 
     const generator = executeToolCalls(
       executablePendingCalls,
-      this.tools,
+      this.resolveExecutableTools(executablePendingCalls),
       approvals,
       clientToolResults,
       (eventName, data) => this.createCustomEventChunk(eventName, data),
@@ -1412,7 +1429,7 @@ class TextEngine<
 
     const generator = executeToolCalls(
       executableToolCalls,
-      this.tools,
+      this.resolveExecutableTools(executableToolCalls),
       approvals,
       clientToolResults,
       (eventName, data) => this.createCustomEventChunk(eventName, data),

packages/ai/src/activities/chat/tools/lazy-tool-manager.ts

@@ -1,7 +1,14 @@
 import { convertSchemaToJsonSchema } from './schema-converter'
-import type { Tool } from '../../../types'
+import type { AnyTool, Tool } from '../../../types'
 
-const DISCOVERY_TOOL_NAME = '__lazy__tool__discovery__'
+/**
+ * Name of the synthetic tool the LLM calls to discover lazy tools.
+ *
+ * Exported so callers building custom message-compaction / history-trimming
+ * logic can reference the discovery tool by constant instead of hard-coding
+ * the string (which is an internal contract that could change).
+ */
+export const DISCOVERY_TOOL_NAME = '__lazy__tool__discovery__'
 
 /**
  * Manages lazy tool discovery for the chat agent loop.
@@ -87,6 +94,35 @@ export class LazyToolManager {
     return active
   }
 
+  /**
+   * Returns the tools that should be available for *execution* this turn.
+   *
+   * This is the advertised set (`getActiveTools()`, passed in as `activeTools`)
+   * plus the discovery tool when a pending call references it but it is no
+   * longer advertised. Once every lazy tool has been discovered the discovery
+   * tool is dropped from the advertised set, but a model may still re-request
+   * discovery (long context / hallucination); keeping it executable lets that
+   * call return the schemas again instead of failing with "Unknown tool".
+   *
+   * The advertised set is intentionally left unchanged — only execution lookup
+   * is widened. Operates on the already-built `activeTools`: it must NOT call
+   * `getActiveTools()`, which would reset `hasNewDiscoveries` before the
+   * post-execution refresh check in the agent loop.
+   */
+  getExecutableTools(
+    activeTools: ReadonlyArray<AnyTool>,
+    pendingToolCallNames: ReadonlyArray<string>,
+  ): ReadonlyArray<AnyTool> {
+    if (
+      this.discoveryTool &&
+      pendingToolCallNames.includes(DISCOVERY_TOOL_NAME) &&
+      !activeTools.some((t) => t.name === DISCOVERY_TOOL_NAME)
+    ) {
+      return [...activeTools, this.discoveryTool]
+    }
+    return activeTools
+  }
+
   /**
    * Returns whether new tools have been discovered since the last getActiveTools() call.
    */
@@ -221,8 +257,14 @@ export class LazyToolManager {
         for (const name of args.toolNames) {
           const tool = lazyToolMap.get(name)
           if (tool) {
-            manager.discoveredTools.add(name)
-            manager.hasNewDiscoveries = true
+            // Only flag a refresh for genuinely new discoveries. Re-requesting
+            // an already-discovered tool still returns its schema below (the
+            // model asked for it), but must not trigger a redundant tool-list
+            // refresh + continue in the agent loop.
+            if (!manager.discoveredTools.has(name)) {
+              manager.discoveredTools.add(name)
+              manager.hasNewDiscoveries = true
+            }
             const jsonSchema = tool.inputSchema
               ? convertSchemaToJsonSchema(tool.inputSchema)
               : undefined

packages/ai/src/index.ts

@@ -82,6 +82,10 @@ export {
 // Tool call management
 export { ToolCallManager } from './activities/chat/tools/tool-calls'
 
+// Lazy tool discovery (name of the synthetic discovery tool, for custom
+// message-compaction logic that needs to reference it)
+export { DISCOVERY_TOOL_NAME } from './activities/chat/tools/lazy-tool-manager'
+
 // Provider tool type
 export type { ProviderTool } from './tools/provider-tool'
 export { brandProviderTool } from './tools/provider-tool'

packages/ai/tests/chat.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, it, vi } from 'vitest'
 import { chat, createChatOptions } from '../src/activities/chat/index'
+import { DISCOVERY_TOOL_NAME } from '../src/activities/chat/tools/lazy-tool-manager'
 import { EventType } from '../src/types'
 import type { StreamChunk, Tool } from '../src/types'
 import {
@@ -1767,6 +1768,158 @@ describe('chat()', () => {
       expect(toolNames).not.toContain('__lazy__tool__discovery__')
       expect(toolNames).toContain('normalTool')
     })
+
+    it('should not error when the model re-requests discovery after all lazy tools are discovered (#788)', async () => {
+      const weatherExecute = vi.fn().mockReturnValue({ temp: 72 })
+      const toolNamesPerCall: Array<Array<string>> = []
+
+      let callCount = 0
+      const { adapter } = createMockAdapter({
+        chatStreamFn: (opts: any) => {
+          callCount++
+          toolNamesPerCall.push((opts.tools ?? []).map((t: any) => t.name))
+
+          if (callCount === 1) {
+            // Discover the only lazy tool -> all discovered, so the discovery
+            // tool is dropped from the advertised set.
+            return (async function* () {
+              yield ev.runStarted()
+              yield ev.toolStart('c1', DISCOVERY_TOOL_NAME)
+              yield ev.toolArgs(
+                'c1',
+                JSON.stringify({ toolNames: ['getWeather'] }),
+              )
+              yield ev.runFinished('tool_calls')
+            })()
+          } else if (callCount === 2) {
+            // Model overlooks that getWeather is already available and asks to
+            // discover it again, even though the discovery tool is no longer
+            // advertised.
+            return (async function* () {
+              yield ev.runStarted()
+              yield ev.toolStart('c2', DISCOVERY_TOOL_NAME)
+              yield ev.toolArgs(
+                'c2',
+                JSON.stringify({ toolNames: ['getWeather'] }),
+              )
+              yield ev.runFinished('tool_calls')
+            })()
+          }
+          return (async function* () {
+            yield ev.runStarted()
+            yield ev.textStart()
+            yield ev.textContent('done')
+            yield ev.textEnd()
+            yield ev.runFinished('stop')
+          })()
+        },
+      })
+
+      const stream = chat({
+        adapter,
+        messages: [{ role: 'user', content: 'Weather?' }],
+        tools: [lazyServerTool('getWeather', weatherExecute)],
+      })
+
+      const chunks = await collectChunks(stream as AsyncIterable<StreamChunk>)
+
+      // The discovery tool was removed from the advertised set on the 2nd call,
+      // proving we fixed execution without re-advertising it.
+      expect(toolNamesPerCall[1]).not.toContain(DISCOVERY_TOOL_NAME)
+      expect(toolNamesPerCall[1]).toContain('getWeather')
+
+      // Re-requesting discovery must NOT produce an "Unknown tool" error.
+      const toolResults = chunks.filter(
+        (c) => c.type === 'TOOL_CALL_RESULT',
+      ) as Array<any>
+      const unknownToolError = toolResults.find(
+        (c: any) =>
+          typeof c.content === 'string' && c.content.includes('Unknown tool'),
+      )
+      expect(unknownToolError).toBeUndefined()
+
+      // The run progressed past the re-discovery turn to the final answer.
+      const text = chunks
+        .filter((c) => c.type === 'TEXT_MESSAGE_CONTENT')
+        .map((c: any) => c.delta)
+        .join('')
+      expect(text).toContain('done')
+    })
+
+    it('should handle a discovery call batched with an already-available tool in one turn', async () => {
+      const lazyAExecute = vi.fn().mockReturnValue({ a: 1 })
+      const lazyBExecute = vi.fn().mockReturnValue({ b: 2 })
+      const toolNamesPerCall: Array<Array<string>> = []
+
+      let callCount = 0
+      const { adapter } = createMockAdapter({
+        chatStreamFn: (opts: any) => {
+          callCount++
+          toolNamesPerCall.push((opts.tools ?? []).map((t: any) => t.name))
+
+          if (callCount === 1) {
+            // Discover lazyA only (lazyB stays undiscovered).
+            return (async function* () {
+              yield ev.runStarted()
+              yield ev.toolStart('d1', DISCOVERY_TOOL_NAME)
+              yield ev.toolArgs('d1', JSON.stringify({ toolNames: ['lazyA'] }))
+              yield ev.runFinished('tool_calls')
+            })()
+          } else if (callCount === 2) {
+            // One batch: call the already-discovered lazyA AND discover lazyB.
+            return (async function* () {
+              yield ev.runStarted()
+              yield ev.toolStart('a1', 'lazyA')
+              yield ev.toolArgs('a1', '{}')
+              yield ev.toolStart('d2', DISCOVERY_TOOL_NAME)
+              yield ev.toolArgs('d2', JSON.stringify({ toolNames: ['lazyB'] }))
+              yield ev.runFinished('tool_calls')
+            })()
+          } else if (callCount === 3) {
+            // lazyB is now available -> call it.
+            return (async function* () {
+              yield ev.runStarted()
+              yield ev.toolStart('b1', 'lazyB')
+              yield ev.toolArgs('b1', '{}')
+              yield ev.runFinished('tool_calls')
+            })()
+          }
+          return (async function* () {
+            yield ev.runStarted()
+            yield ev.textStart()
+            yield ev.textContent('ok')
+            yield ev.textEnd()
+            yield ev.runFinished('stop')
+          })()
+        },
+      })
+
+      const stream = chat({
+        adapter,
+        messages: [{ role: 'user', content: 'go' }],
+        tools: [
+          lazyServerTool('lazyA', lazyAExecute),
+          lazyServerTool('lazyB', lazyBExecute),
+        ],
+      })
+
+      const chunks = await collectChunks(stream as AsyncIterable<StreamChunk>)
+
+      // lazyA executed despite sharing a batch with a discovery call.
+      expect(lazyAExecute).toHaveBeenCalledTimes(1)
+      // lazyB became available after the batched discovery and executed.
+      expect(toolNamesPerCall[2]).toContain('lazyB')
+      expect(lazyBExecute).toHaveBeenCalledTimes(1)
+
+      const toolResults = chunks.filter(
+        (c) => c.type === 'TOOL_CALL_RESULT',
+      ) as Array<any>
+      const unknownToolError = toolResults.find(
+        (c: any) =>
+          typeof c.content === 'string' && c.content.includes('Unknown tool'),
+      )
+      expect(unknownToolError).toBeUndefined()
+    })
   })
 
   // ==========================================================================