From a9f31e8c624715187a04b919a4987ea6ced63966 Mon Sep 17 00:00:00 2001 From: Tianning Li Date: Tue, 3 Nov 2026 20:39:25 +0800 Subject: [PATCH 1/2] feat(coding-agent): add tool_search for deferred tool discovery A session that registers many tools pays for all of them in every prompt, even when a turn needs one. Add `tool_search`: a BM25 ranker over tool metadata that activates the matches so the next model call declares them. Activation goes through the active tool set, so the change is recorded on the tool result like any other activation and survives resume and fork. Tools opt out of the default selection with `defaultActive: false`; they stay registered and reachable through the tool settings or another tool's activation. Explicit selection still wins: a name in the `--tools` allowlist is always activated. `tool_search` ships registered but undeclared, enabled per session through `defaultTools`, `--tools`, or the tool settings. Step's permission policy treats it as read-only: it loads tool definitions and touches no files. --- apps/cli/src/bootstrap/extensions.ts | 2 + packages/coding-agent/docs/extensions.md | 2 + .../coding-agent/src/core/agent-session.ts | 10 +- .../coding-agent/src/core/extensions/types.ts | 7 + .../src/core/tools/tool-search.ts | 293 ++++++++++++++++++ .../src/features/step-tool-search.ts | 14 + packages/coding-agent/src/index.ts | 1 + packages/coding-agent/src/step/permissions.ts | 2 + .../test/tool-search-integration.test.ts | 170 ++++++++++ .../coding-agent/test/tool-search.test.ts | 209 +++++++++++++ 10 files changed, 708 insertions(+), 2 deletions(-) create mode 100644 packages/coding-agent/src/core/tools/tool-search.ts create mode 100644 packages/coding-agent/src/features/step-tool-search.ts create mode 100644 packages/coding-agent/test/tool-search-integration.test.ts create mode 100644 packages/coding-agent/test/tool-search.test.ts diff --git a/apps/cli/src/bootstrap/extensions.ts b/apps/cli/src/bootstrap/extensions.ts index 95d8fe77c..537e595c2 100644 --- a/apps/cli/src/bootstrap/extensions.ts +++ b/apps/cli/src/bootstrap/extensions.ts @@ -12,6 +12,7 @@ import { createStepCronExtension, createStepExtensionInline, createStepGoalExtension, + createStepToolSearchExtension, type InlineExtension, type StepExtensionOptions, type StepTelemetryReporter, @@ -66,6 +67,7 @@ export function createStepExtensionFactories(deps: StepExtensionFactoryDeps): In }), createStepCronExtension({ telemetry: deps.telemetry }), createStepGoalExtension({ telemetry: deps.telemetry }), + createStepToolSearchExtension(), { name: "connect-lark", factory: createLarkConnectionExtension(), hidden: true }, ...(deps.stepCodeProviderExtension ? [deps.stepCodeProviderExtension] : []), ]; diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index f0242d217..94bc94969 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -1405,6 +1405,8 @@ Register a custom tool callable by the LLM. See [Custom Tools](#custom-tools) fo Use `pi.setActiveTools()` to enable or disable tools (including dynamically added tools) at runtime. +Set `defaultActive: false` on a definition to register a tool without declaring it: it stays in `pi.getAllTools()` and can be activated by the tool settings, `--tools`, or another tool's activation. Step's `tool_search` is registered that way. + Use `promptSnippet` to opt a custom tool into a one-line entry in `Available tools`, and `promptGuidelines` to append tool-specific bullets to the default `Guidelines` section when the tool is active. **Important:** `promptGuidelines` bullets are appended flat to the `Guidelines` section with no tool name prefix. Each guideline must name the tool it refers to — avoid "Use this tool when..." because the LLM cannot tell which tool "this" means. Write "Use my_tool when..." instead. diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 0d9a6d490..59c2a8b9d 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -4076,6 +4076,10 @@ export class AgentSession { const excludedToolNames = this._excludedToolNames; const isAllowedTool = (name: string): boolean => (!allowedToolNames || allowedToolNames.has(name)) && !excludedToolNames?.has(name); + // A tool that opts out of the default selection stays registered but is not declared until + // the session activates it, through the tool settings or another tool's activation. + const isDefaultActive = (name: string): boolean => + this._toolDefinitions.get(name)?.definition.defaultActive !== false; const registeredTools = this._extensionRunner.getAllRegisteredTools(); const allCustomTools = [ @@ -4142,6 +4146,8 @@ export class AgentSession { ).filter((name) => isAllowedTool(name)); if (allowedToolNames && options?.includeAllExtensionTools) { + // The allowlist is an explicit selection, so it activates a tool that opted out of the + // default set. for (const toolName of this._toolRegistry.keys()) { if (allowedToolNames.has(toolName)) { nextActiveToolNames.push(toolName); @@ -4149,13 +4155,13 @@ export class AgentSession { } } else if (options?.includeAllExtensionTools) { for (const tool of wrappedExtensionTools) { - nextActiveToolNames.push(tool.name); + if (isDefaultActive(tool.name)) nextActiveToolNames.push(tool.name); } } else { // Preserve selections for known tools during refresh/reload; only genuinely // new registrations opt in automatically after initial runtime creation. for (const toolName of this._toolRegistry.keys()) { - if (!previousRegistryNames.has(toolName)) { + if (!previousRegistryNames.has(toolName) && isDefaultActive(toolName)) { nextActiveToolNames.push(toolName); } } diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 07d8d9490..2aa27c71f 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -579,6 +579,13 @@ export interface ToolDefinition Static; diff --git a/packages/coding-agent/src/core/tools/tool-search.ts b/packages/coding-agent/src/core/tools/tool-search.ts new file mode 100644 index 000000000..9b54ab7fe --- /dev/null +++ b/packages/coding-agent/src/core/tools/tool-search.ts @@ -0,0 +1,293 @@ +/** + * Tool discovery: a BM25 ranker over tool metadata, and the `tool_search` tool built on it. + * + * `tool_search` searches the tools that are registered but not declared to the model and activates + * the matches, so the next model call declares them. Activation goes through the active tool set, + * so it is recorded in the transcript like any other tool change and survives resume and fork on + * that branch. + */ +import { type Static, Type } from "typebox"; +import type { AgentToolResult, ToolDefinition, ToolInfo } from "../extensions/types.ts"; + +export const TOOL_SEARCH_TOOL_NAME = "tool_search"; + +export const DEFAULT_TOOL_SEARCH_LIMIT = 8; + +/** The part of the extension runtime `tool_search` needs, so it stays testable without a session. */ +export interface ToolSearchCatalog { + getActiveTools(): string[]; + getAllTools(): ToolInfo[]; + setActiveTools(toolNames: string[]): void; +} + +export interface ToolSearchResult { + name: string; + description: string; +} + +const wordSegmenter = new Intl.Segmenter(undefined, { granularity: "word" }); + +const STOP_WORDS = new Set([ + "a", + "an", + "and", + "are", + "as", + "at", + "be", + "by", + "for", + "from", + "in", + "is", + "it", + "of", + "on", + "or", + "that", + "the", + "this", + "to", + "with", +]); + +/** Naive singular form, so `issues` matches `issue` and `searches` matches `search`. */ +function stem(term: string): string { + if (term.length > 4 && term.endsWith("ies")) return `${term.slice(0, -3)}y`; + if (term.length > 4 && /(ches|shes|sses|xes|zes)$/.test(term)) return term.slice(0, -2); + if (term.length > 3 && term.endsWith("s") && !term.endsWith("ss")) return term.slice(0, -1); + return term; +} + +/** Unicode word segmentation with identifier boundaries, English stop words, and singular forms. */ +export function tokenize(text: string): string[] { + const normalized = text + .normalize("NFC") + .replace(/([a-z0-9])([A-Z])/g, "$1 $2") + .replace(/([A-Z]+)([A-Z][a-z])/g, "$1 $2") + .toLowerCase() + .replace(/[^\p{L}\p{M}\p{N}]+/gu, " "); + // The segmenter costs about 15 times a split. Plain ASCII splits into the same words at spaces, so only other + // text, such as Chinese written without spaces, goes through it. + const words = /^[a-z0-9 ]*$/.test(normalized) + ? normalized.split(" ") + : Array.from(wordSegmenter.segment(normalized), (segment) => (segment.isWordLike ? segment.segment : "")); + return words + .filter((word) => word.length > 0 && !STOP_WORDS.has(word)) + .map((word) => (/^[a-z0-9]+$/.test(word) ? stem(word) : word)); +} + +function isObject(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +/** Schema descriptions and property names, recursively. */ +function schemaText(schema: unknown, parts: string[]): void { + if (!isObject(schema)) return; + if (typeof schema.description === "string") parts.push(schema.description); + if (isObject(schema.properties)) { + for (const [name, property] of Object.entries(schema.properties)) { + parts.push(name); + schemaText(property, parts); + } + } + schemaText(schema.items, parts); + for (const key of ["anyOf", "oneOf", "allOf"]) { + const variants = schema[key]; + if (Array.isArray(variants)) for (const variant of variants) schemaText(variant, parts); + } +} + +/** Search text of a tool: its description and its schema descriptions and property names. */ +export function createToolSearchDocument(tool: Pick): { + name: string; + text: string; +} { + const parts = [tool.description]; + schemaText(tool.parameters, parts); + return { name: tool.name, text: parts.filter((part) => part.trim()).join(" ") }; +} + +/** Lowercase with characters invalid in script identifiers as `_`, so `dev-radius` matches `dev_radius`. */ +function normalizeName(name: string): string { + return name + .trim() + .toLowerCase() + .replace(/[^a-z0-9_$]/g, "_"); +} + +/** + * Whether the whole query is the tool's name: its full name or a suffix after a `__` separator, so + * `thread_read` and `app__thread_read` both name `mcp__app__thread_read`. + */ +function isExactNameQuery(query: string, name: string): boolean { + const normalized = normalizeName(name); + return normalized === query || normalized.endsWith(`__${query}`); +} + +interface IndexedDocument { + name: string; + nameTerms: Set; + localNameTerms: string[]; + textCounts: Map; + textLength: number; +} + +function indexDocument(document: { name: string; text: string }): IndexedDocument { + const textCounts = new Map(); + for (const term of tokenize(document.text)) textCounts.set(term, (textCounts.get(term) ?? 0) + 1); + const segments = document.name.split("__"); + return { + name: document.name, + nameTerms: new Set(tokenize(document.name)), + localNameTerms: [...new Set(tokenize(segments[segments.length - 1]))], + textCounts, + textLength: [...textCounts.values()].reduce((sum, count) => sum + count, 0), + }; +} + +/** + * A query term in the name counts this many times its IDF. The text part of a term is at most + * `k1 + 1` (2.2) times its IDF and about 1 for a single mention, so a name match beats any number of + * mentions in the description. + */ +const NAME_WEIGHT = 3; + +export interface ScoredDocument { + name: string; + score: number; +} + +/** + * Okapi BM25 over the description, with the name as a separate field: + * - a tool that matches more of the query's terms, in any field, ranks above one that matches + * fewer, whatever their scores: for `screenshot window`, a tool whose description has both words + * beats `window`, whose name has one; + * - a query term in the name scores more than any number of mentions in the description; + * - a tool whose whole local name is in the query (`app` in `peekaboo app` for + * `mcp__peekaboo__app`) gets a bonus worth one more match of every query term, so among tools + * matching as many terms it mostly wins over tools that only mention the words; + * - a query that is exactly a tool's name ranks that tool first, before coverage. + * + * The description alone is length-normalized, so a long name does not lower a tool's score. Ties + * keep document order. + */ +export class Bm25Ranker { + private readonly k1: number; + private readonly b: number; + + constructor(options: { k1?: number; b?: number } = {}) { + this.k1 = options.k1 ?? 1.2; + this.b = options.b ?? 0.75; + } + + rank(query: string, documents: readonly { name: string; text: string }[], limit: number): ScoredDocument[] { + const queryTerms = new Set(tokenize(query)); + if (queryTerms.size === 0 || documents.length === 0 || limit <= 0) return []; + const indexed = documents.map(indexDocument); + const averageLength = indexed.reduce((sum, document) => sum + document.textLength, 0) / indexed.length || 1; + const idf = new Map( + [...queryTerms].map((term) => { + const frequency = indexed.filter( + (document) => document.nameTerms.has(term) || document.textCounts.has(term), + ).length; + return [term, Math.log(1 + (indexed.length - frequency + 0.5) / (frequency + 0.5))]; + }), + ); + const fullNameBonus = [...idf.values()].reduce((sum, value) => sum + value, 0); + const exactQuery = normalizeName(query); + const matches: { name: string; score: number; covered: number; exact: boolean }[] = []; + for (const document of indexed) { + const norm = this.k1 * (1 - this.b + (this.b * document.textLength) / averageLength); + let score = 0; + let covered = 0; + for (const term of queryTerms) { + const count = document.textCounts.get(term) ?? 0; + const weight = (document.nameTerms.has(term) ? NAME_WEIGHT : 0) + (count * (this.k1 + 1)) / (count + norm); + if (weight > 0) covered++; + score += (idf.get(term) ?? 0) * weight; + } + if (score === 0) continue; + const local = document.localNameTerms; + if (local.length > 0 && local.every((term) => queryTerms.has(term))) score += fullNameBonus; + matches.push({ name: document.name, score, covered, exact: isExactNameQuery(exactQuery, document.name) }); + } + return matches + .sort((a, b) => Number(b.exact) - Number(a.exact) || b.covered - a.covered || b.score - a.score) + .slice(0, limit) + .map(({ name, score }) => ({ name, score })); + } +} + +export const toolSearchSchema = Type.Object({ + query: Type.String({ description: "Search query for deferred tools." }), + limit: Type.Optional( + Type.Number({ description: `Maximum number of tools to return. Defaults to ${DEFAULT_TOOL_SEARCH_LIMIT}.` }), + ), +}); + +/** Rank the tools that are registered but not active yet, and activate the matches. */ +export function searchAndLoad(catalog: ToolSearchCatalog, query: string, limit: number): ToolSearchResult[] { + const active = catalog.getActiveTools(); + const inactive = new Set(active); + const candidates = catalog.getAllTools().filter((tool) => !inactive.has(tool.name)); + const documents = candidates.map((tool) => createToolSearchDocument(tool)); + const matches = new Bm25Ranker().rank(query, documents, limit); + if (matches.length === 0) return []; + catalog.setActiveTools([...active, ...matches.map((match) => match.name)]); + const descriptions = new Map(candidates.map((tool) => [tool.name, tool.description])); + return matches.map((match) => ({ name: match.name, description: descriptions.get(match.name) ?? "" })); +} + +/** + * The `tool_search` description. It does not list the searchable tools, so it stays the same while + * tools are registered, for example when MCP servers connect. + */ +export const TOOL_SEARCH_DESCRIPTION = `# Tool discovery + +Searches the tools that are not loaded yet and declares the matches for the next model call. + +A tool that is not in your tool list may still be registered, for example one you turned off or +one that a plugin publishes. Use this tool (\`${TOOL_SEARCH_TOOL_NAME}\`) to look for it instead of +telling the user the tool does not exist.`; + +export interface ToolSearchDetails { + loaded: string[]; +} + +/** + * The tool definition with a two-argument `execute`: `tool_search` needs no abort signal, progress + * callback, or extension context, and a narrower signature still registers as a tool. + */ +export type ToolSearchToolDefinition = Omit< + ToolDefinition, + "parameters" | "execute" +> & { + parameters: typeof toolSearchSchema; + execute(toolCallId: string, params: Static): Promise>; +}; + +export function createToolSearchToolDefinition(catalog: ToolSearchCatalog): ToolSearchToolDefinition { + return { + name: TOOL_SEARCH_TOOL_NAME, + label: "Tool Search", + description: TOOL_SEARCH_DESCRIPTION, + promptSnippet: "Search for tools that are not loaded yet and load the matches", + parameters: toolSearchSchema, + // Inactive until the session asks for it; see `_refreshToolRegistry`. + defaultActive: false, + async execute(_toolCallId, { query, limit }) { + if (query.trim() === "") throw new Error("query must not be empty"); + const max = limit ?? DEFAULT_TOOL_SEARCH_LIMIT; + if (!Number.isInteger(max) || max <= 0) throw new Error("limit must be a positive integer"); + const found = searchAndLoad(catalog, query, max); + const text = + found.length === 0 + ? "No matching tools found." + : `Loaded ${found.length} tool${found.length === 1 ? "" : "s"}. They are available from your next call:\n${found + .map((tool) => `- ${tool.name}: ${tool.description.trim().split(/\r?\n/)[0]}`) + .join("\n")}`; + return { content: [{ type: "text", text }], details: { loaded: found.map((tool) => tool.name) } }; + }, + }; +} diff --git a/packages/coding-agent/src/features/step-tool-search.ts b/packages/coding-agent/src/features/step-tool-search.ts new file mode 100644 index 000000000..91504a8fd --- /dev/null +++ b/packages/coding-agent/src/features/step-tool-search.ts @@ -0,0 +1,14 @@ +import type { ExtensionAPI, ExtensionFactory } from "../core/extensions/types.ts"; +import { createToolSearchToolDefinition } from "../core/tools/tool-search.ts"; + +/** + * Mount `tool_search`: it is registered inactive and finds the tools that are registered but not + * declared to the model, so a session can keep a large tool catalogue out of the prompt. + * + * Enable it per session with `defaultTools`, `--tools`, or the tool settings. + */ +export function createStepToolSearchExtension(): ExtensionFactory { + return (pi: ExtensionAPI): void => { + pi.registerTool(createToolSearchToolDefinition(pi)); + }; +} diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 51134e1f3..3fd652af4 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -526,6 +526,7 @@ export { stepScheduleExtensionInline, UpdateGoalParams, } from "./features/step-schedule.ts"; +export { createStepToolSearchExtension } from "./features/step-tool-search.ts"; export { type CompletedSubagentResponse, readSubagentLaneSnapshot, diff --git a/packages/coding-agent/src/step/permissions.ts b/packages/coding-agent/src/step/permissions.ts index edfaf6f31..ff7a29fe5 100644 --- a/packages/coding-agent/src/step/permissions.ts +++ b/packages/coding-agent/src/step/permissions.ts @@ -82,6 +82,8 @@ const READ_ONLY_TOOLS = new Set([ "search_web", "read_file", "find_tools", + // Loads tool definitions for the next model call; it touches no files and runs no command. + "tool_search", // Keep native names safe when a caller explicitly enables a native tool. "ls", "find", diff --git a/packages/coding-agent/test/tool-search-integration.test.ts b/packages/coding-agent/test/tool-search-integration.test.ts new file mode 100644 index 000000000..a8962afca --- /dev/null +++ b/packages/coding-agent/test/tool-search-integration.test.ts @@ -0,0 +1,170 @@ +import { existsSync, mkdirSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { fauxAssistantMessage, fauxToolCall, registerFauxProvider } from "@step-harness/providers/compat"; +import { Type } from "typebox"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { AuthStorage } from "../src/core/auth-storage.ts"; +import type { ExtensionAPI } from "../src/core/extensions/types.ts"; +import { ModelRuntime } from "../src/core/model-runtime.ts"; +import { DefaultResourceLoader } from "../src/core/resource-loader.ts"; +import { createAgentSession } from "../src/core/sdk.ts"; +import { SessionManager } from "../src/core/session-manager.ts"; +import { SettingsManager } from "../src/core/settings-manager.ts"; +import { createStepToolSearchExtension } from "../src/features/step-tool-search.ts"; + +/** + * A plugin tool the model never sees until `tool_search` loads it: registered, but opted out of + * the default active set. + */ +function deferredTool(name: string, description: string, query: string) { + return { + name, + label: name, + description, + defaultActive: false, + parameters: Type.Object({ query: Type.String() }), + async execute(_toolCallId: string) { + return { content: [{ type: "text" as const, text: `${name}:${query}` }], details: {} }; + }, + }; +} + +describe("tool_search over a real model turn", () => { + let tempDir: string; + let faux: ReturnType; + const cleanups: Array<() => void | Promise> = []; + + beforeEach(() => { + tempDir = join(tmpdir(), `pi-tool-search-e2e-${Date.now()}-${Math.random().toString(36).slice(2)}`); + mkdirSync(tempDir, { recursive: true }); + }); + + afterEach(async () => { + while (cleanups.length > 0) await cleanups.pop()?.(); + faux?.unregister(); + if (tempDir && existsSync(tempDir)) rmSync(tempDir, { recursive: true, force: true }); + }); + + async function createSession(extensionFactories: ((pi: ExtensionAPI) => void)[]) { + faux = registerFauxProvider(); + const authStorage = AuthStorage.inMemory(); + await authStorage.modify(faux.getModel().provider, async () => ({ type: "api_key", key: "faux-key" })); + const modelRuntime = await ModelRuntime.create({ + credentials: authStorage, + modelsPath: join(tempDir, "models.json"), + }); + const model = faux.getModel(); + modelRuntime.registerProvider(model.provider, { + baseUrl: model.baseUrl, + api: model.api, + models: [ + { + id: model.id, + name: model.name, + api: model.api, + reasoning: model.reasoning, + input: model.input, + cost: model.cost, + contextWindow: model.contextWindow, + maxTokens: model.maxTokens, + baseUrl: model.baseUrl, + }, + ], + }); + const settingsManager = SettingsManager.inMemory({ + defaultTools: ["read", "bash", "edit", "write", "tool_search"], + }); + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir: tempDir, + settingsManager, + extensionFactories, + noSkills: true, + noPromptTemplates: true, + noThemes: true, + }); + await resourceLoader.reload(); + const { session } = await createAgentSession({ + cwd: tempDir, + agentDir: tempDir, + model, + modelRuntime, + settingsManager, + sessionManager: SessionManager.inMemory(tempDir), + resourceLoader, + }); + await session.bindExtensions({}); + cleanups.push(() => session.dispose()); + return session; + } + + it("declares the loaded tool on the next request", async () => { + const declaredToolNames: string[][] = []; + const toolResultTexts: string[] = []; + const session = await createSession([ + (pi) => { + pi.registerTool(deferredTool("issue_search", "Lists open issues in a repository", "query")); + pi.registerTool(deferredTool("calendar_read", "Reads calendar entries", "query")); + }, + createStepToolSearchExtension(), + (pi) => { + pi.on("tool_result", (event) => { + toolResultTexts.push( + event.content + .map((block: { type: string; text?: string }) => (block.type === "text" ? block.text : "")) + .join(""), + ); + }); + }, + ]); + faux.setResponses([ + (context) => { + declaredToolNames.push((context.tools ?? []).map((tool: { name: string }) => tool.name).sort()); + return fauxAssistantMessage([fauxToolCall("tool_search", { query: "issues" })]); + }, + (context) => { + declaredToolNames.push((context.tools ?? []).map((tool: { name: string }) => tool.name).sort()); + return fauxAssistantMessage("done"); + }, + ]); + + await session.prompt("list the open issues"); + + expect(declaredToolNames).toHaveLength(2); + // First request: the deferred tools are registered but not declared. + expect(declaredToolNames[0]).toContain("tool_search"); + expect(declaredToolNames[0]).not.toContain("issue_search"); + expect(declaredToolNames[0]).not.toContain("calendar_read"); + // Second request: the loaded tool is declared, the unrelated one still is not. + expect(declaredToolNames[1]).toContain("issue_search"); + expect(declaredToolNames[1]).not.toContain("calendar_read"); + expect(toolResultTexts.join("\n")).toContain("Loaded 1 tool"); + expect(toolResultTexts.join("\n")).toContain("- issue_search: Lists open issues in a repository"); + }); + + it("records the loaded tools as an active-tool change on the tool result", async () => { + const session = await createSession([ + (pi) => { + pi.registerTool(deferredTool("issue_search", "Lists open issues in a repository", "query")); + }, + createStepToolSearchExtension(), + ]); + const addedToolNames: string[][] = []; + faux.setResponses([ + () => fauxAssistantMessage(fauxToolCall("tool_search", { query: "issues" }), { stopReason: "toolUse" }), + (context) => { + addedToolNames.push( + context.messages + .filter((message) => message.role === "toolResult") + .flatMap((message) => message.addedToolNames ?? []), + ); + return fauxAssistantMessage("done"); + }, + ]); + + await session.prompt("list the open issues"); + + expect(addedToolNames).toEqual([["issue_search"]]); + }); +}); diff --git a/packages/coding-agent/test/tool-search.test.ts b/packages/coding-agent/test/tool-search.test.ts new file mode 100644 index 000000000..e04aa2082 --- /dev/null +++ b/packages/coding-agent/test/tool-search.test.ts @@ -0,0 +1,209 @@ +import { existsSync, mkdirSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { Type } from "typebox"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import type { ToolInfo } from "../src/core/extensions/types.ts"; +import { DefaultResourceLoader } from "../src/core/resource-loader.ts"; +import { createAgentSession } from "../src/core/sdk.ts"; +import { SessionManager } from "../src/core/session-manager.ts"; +import { SettingsManager } from "../src/core/settings-manager.ts"; +import { createSyntheticSourceInfo } from "../src/core/source-info.ts"; +import { + Bm25Ranker, + createToolSearchDocument, + createToolSearchToolDefinition, + DEFAULT_TOOL_SEARCH_LIMIT, + searchAndLoad, + type ToolSearchCatalog, + tokenize, +} from "../src/core/tools/tool-search.ts"; +import { createStepToolSearchExtension } from "../src/features/step-tool-search.ts"; +import { decideStepToolCall, stepPermissionStateForPreset } from "../src/step/permissions.ts"; +import { stepModel } from "./utilities.ts"; + +function toolInfo(name: string, description: string, parameters = Type.Object({})): ToolInfo { + return { + name, + description, + parameters, + sourceInfo: createSyntheticSourceInfo(``, { source: "extension" }), + }; +} + +interface FakeCatalog extends ToolSearchCatalog { + active: string[]; + activations: string[][]; +} + +function fakeCatalog(tools: ToolInfo[], active: string[]): FakeCatalog { + const state: FakeCatalog = { + active, + activations: [], + getActiveTools: () => [...state.active], + getAllTools: () => tools, + setActiveTools: (names) => { + state.active = [...names]; + state.activations.push([...names]); + }, + }; + return state; +} + +describe("tool search ranker", () => { + it("splits identifiers, drops stop words, and folds plurals", () => { + expect(tokenize("Search the issue_tracker for Issues")).toEqual(["search", "issue", "tracker", "issue"]); + }); + + it("ranks a name match above a description mention", () => { + const documents = [ + createToolSearchDocument(toolInfo("report", "Reads calendar entries and issues")), + createToolSearchDocument(toolInfo("issues", "Lists open items")), + ]; + expect(new Bm25Ranker().rank("issue", documents, 2).map((match) => match.name)).toEqual(["issues", "report"]); + }); + + it("ranks an exact name query first", () => { + const documents = [ + createToolSearchDocument(toolInfo("issue_search", "Searches issues")), + createToolSearchDocument(toolInfo("search", "Search anything, including issues")), + ]; + expect(new Bm25Ranker().rank("search", documents, 3)[0]?.name).toBe("search"); + }); + + it("matches a tool by its suffix after the namespace separator", () => { + const documents = [createToolSearchDocument(toolInfo("mcp__app__thread_read", "Reads a thread"))]; + expect(new Bm25Ranker().rank("thread_read", documents, 1).map((match) => match.name)).toEqual([ + "mcp__app__thread_read", + ]); + }); + + it("ignores a query made only of stop words", () => { + const documents = [createToolSearchDocument(toolInfo("issues", "Lists open items"))]; + expect(new Bm25Ranker().rank("the of and", documents, 3)).toEqual([]); + }); +}); + +describe("searchAndLoad", () => { + const inactive = toolInfo("issues", "Lists open issues in a repository"); + const active = toolInfo("read", "Reads a file"); + + it("activates the matches and keeps the active tools", () => { + const tools = [active, inactive]; + const catalog = fakeCatalog(tools, ["read"]); + const found = searchAndLoad(catalog, "open issues", DEFAULT_TOOL_SEARCH_LIMIT); + expect(found.map((tool) => tool.name)).toEqual(["issues"]); + expect(found[0]?.description).toBe("Lists open issues in a repository"); + expect(catalog.active).toEqual(["read", "issues"]); + }); + + it("leaves the active set alone when nothing matches", () => { + const catalog = fakeCatalog([active, inactive], ["read"]); + expect(searchAndLoad(catalog, "kubernetes", DEFAULT_TOOL_SEARCH_LIMIT)).toEqual([]); + expect(catalog.activations).toEqual([]); + }); + + it("searches only the tools that are not active yet", () => { + const catalog = fakeCatalog([active, inactive], ["read", "issues"]); + expect(searchAndLoad(catalog, "issues", DEFAULT_TOOL_SEARCH_LIMIT)).toEqual([]); + }); +}); + +describe("tool_search tool", () => { + it("reports the loaded tools", async () => { + const catalog = fakeCatalog([toolInfo("issues", "Lists open issues")], []); + const definition = createToolSearchToolDefinition(catalog); + const result = await definition.execute("call-1", { query: "issues" }); + expect(result.content).toEqual([ + { type: "text", text: "Loaded 1 tool. They are available from your next call:\n- issues: Lists open issues" }, + ]); + expect(result.details).toEqual({ loaded: ["issues"] }); + }); + + it("reports an empty result without activating anything", async () => { + const catalog = fakeCatalog([toolInfo("issues", "Lists open issues")], []); + const result = await createToolSearchToolDefinition(catalog).execute("call-1", { query: "kubernetes" }); + expect(result.content).toEqual([{ type: "text", text: "No matching tools found." }]); + expect(catalog.activations).toEqual([]); + }); + + it("rejects an empty query and a bad limit", async () => { + const definition = createToolSearchToolDefinition(fakeCatalog([], [])); + await expect(definition.execute("call-1", { query: " " })).rejects.toThrow("query must not be empty"); + await expect(definition.execute("call-1", { query: "x", limit: 0 })).rejects.toThrow( + "limit must be a positive integer", + ); + }); +}); + +describe("tool_search permissions", () => { + it("runs without approval because it touches no files", () => { + const decision = decideStepToolCall("tool_search", { query: "weather" }, stepPermissionStateForPreset("approve-for-me")); + expect(decision).toMatchObject({ action: "allow", hazardous: false }); + }); +}); + +describe("tool_search session wiring", () => { + let tempDir: string; + let agentDir: string; + + beforeEach(() => { + tempDir = join(tmpdir(), `pi-tool-search-${Date.now()}-${Math.random().toString(36).slice(2)}`); + agentDir = join(tempDir, "agent"); + mkdirSync(agentDir, { recursive: true }); + }); + + afterEach(() => { + if (tempDir && existsSync(tempDir)) rmSync(tempDir, { recursive: true, force: true }); + }); + + async function createSession(defaultTools?: string[], options: { tools?: string[] } = {}) { + const settingsManager = SettingsManager.inMemory(defaultTools ? { defaultTools } : {}); + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir, + settingsManager, + extensionFactories: [createStepToolSearchExtension()], + }); + await resourceLoader.reload(); + return ( + await createAgentSession({ + cwd: tempDir, + agentDir, + model: stepModel(), + settingsManager, + sessionManager: SessionManager.inMemory(tempDir), + resourceLoader, + tools: options.tools, + }) + ).session; + } + + it("registers tool_search without declaring it", async () => { + const session = await createSession(); + expect(session.getAllTools().map((tool) => tool.name)).toContain("tool_search"); + expect(session.getActiveToolNames()).not.toContain("tool_search"); + expect(session.systemPrompt).not.toContain("- tool_search:"); + session.dispose(); + }); + + it("declares tool_search when the tool settings ask for it", async () => { + const session = await createSession(["read", "bash", "edit", "write", "tool_search"]); + expect(session.getActiveToolNames()).toContain("tool_search"); + expect(session.systemPrompt).toContain("- tool_search:"); + session.dispose(); + }); + + it("declares tool_search when the tool selector names it", async () => { + const session = await createSession(undefined, { tools: ["read", "tool_search"] }); + expect(session.getActiveToolNames()).toContain("tool_search"); + session.dispose(); + }); + + it("keeps tool_search active once a session enables it", async () => { + const session = await createSession(); + session.setActiveToolsByName(["read", "bash", "edit", "write", "tool_search"]); + expect(session.getActiveToolNames()).toContain("tool_search"); + session.dispose(); + }); +}); From 672da5ae2af09d222231878cc9df6db54785eaaf Mon Sep 17 00:00:00 2001 From: Tianning Li Date: Tue, 3 Nov 2026 21:11:37 +0800 Subject: [PATCH 2/2] feat(step): load deferred MCP tools from find_tools A session that connects MCP servers pays for every server tool in every prompt. Add tool exposure so a tool can stay registered but undeclared, and make the Step tool profile's `find_tools` the loader for those tools. - `ToolDefinition.exposure`: `direct` (default), `deferred` (registered, undeclared, loaded on demand), `hidden` (registered, unreachable and left out of `getAllTools()`). Explicit selection still wins: a name in the tool settings or the `--tools` allowlist activates a deferred tool. - `find_tools` ranks deferred and active tools with a BM25 ranker and activates the deferred matches, which the next request declares. The activation is recorded on the tool result, so it survives resume and fork. - MCP servers publish their tools as `deferred` by default; `exposure` and `description` in `config.toml` set that per server. - The MCP extension adds an `mcp_servers` system-prompt section naming the servers whose tools are deferred, because the tool list and `find_tools` both stay silent about them. - `ExtensionContext` exposes the live tool set (`getActiveTools`, `getAllTools`, `setActiveTools`) so a profile tool can activate others. `find_tools` is already classified read-only by the permission policy. --- apps/cli/src/bootstrap/extensions.ts | 2 - packages/coding-agent/docs/extensions.md | 10 +- packages/coding-agent/docs/mcp-catalog.md | 25 +- .../coding-agent/src/core/agent-session.ts | 13 +- .../src/core/extensions/runner.ts | 12 + .../coding-agent/src/core/extensions/types.ts | 23 +- .../src/core/tools/tool-search.ts | 94 ++------ .../src/features/step-tool-search.ts | 14 -- packages/coding-agent/src/index.ts | 1 - packages/coding-agent/src/step/config-toml.ts | 4 + .../coding-agent/src/step/mcp-catalog.test.ts | 7 +- .../coding-agent/src/step/mcp-startup.test.ts | 67 ++++++ packages/coding-agent/src/step/mcp.ts | 58 ++++- packages/coding-agent/src/step/permissions.ts | 3 +- .../coding-agent/src/step/tool-profile.ts | 90 ++++--- .../test/step-tool-profile.test.ts | 2 +- .../test/tool-search-integration.test.ts | 77 +++--- .../coding-agent/test/tool-search-mcp.test.ts | 223 ++++++++++++++++++ .../coding-agent/test/tool-search.test.ts | 129 +++++----- 19 files changed, 613 insertions(+), 241 deletions(-) delete mode 100644 packages/coding-agent/src/features/step-tool-search.ts create mode 100644 packages/coding-agent/test/tool-search-mcp.test.ts diff --git a/apps/cli/src/bootstrap/extensions.ts b/apps/cli/src/bootstrap/extensions.ts index 537e595c2..95d8fe77c 100644 --- a/apps/cli/src/bootstrap/extensions.ts +++ b/apps/cli/src/bootstrap/extensions.ts @@ -12,7 +12,6 @@ import { createStepCronExtension, createStepExtensionInline, createStepGoalExtension, - createStepToolSearchExtension, type InlineExtension, type StepExtensionOptions, type StepTelemetryReporter, @@ -67,7 +66,6 @@ export function createStepExtensionFactories(deps: StepExtensionFactoryDeps): In }), createStepCronExtension({ telemetry: deps.telemetry }), createStepGoalExtension({ telemetry: deps.telemetry }), - createStepToolSearchExtension(), { name: "connect-lark", factory: createLarkConnectionExtension(), hidden: true }, ...(deps.stepCodeProviderExtension ? [deps.stepCodeProviderExtension] : []), ]; diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 94bc94969..a9a83e434 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -1405,7 +1405,15 @@ Register a custom tool callable by the LLM. See [Custom Tools](#custom-tools) fo Use `pi.setActiveTools()` to enable or disable tools (including dynamically added tools) at runtime. -Set `defaultActive: false` on a definition to register a tool without declaring it: it stays in `pi.getAllTools()` and can be activated by the tool settings, `--tools`, or another tool's activation. Step's `tool_search` is registered that way. +Set `exposure` on a definition to control how it reaches the model: + +| `exposure` | Behavior | +| --- | --- | +| `direct` (default) | Declared like any other tool, and active at session start. | +| `deferred` | Registered but not declared. The Step tool profile's `find_tools` loads it, and naming it in the tool settings, `--tools`, or `pi.setActiveTools()` activates it. | +| `hidden` | Registered but unreachable: left out of `pi.getAllTools()` and never activated automatically. | + +The tools of MCP servers are `deferred` by default, so a session can keep a large catalogue out of the prompt; `find_tools` loads the ones a turn needs. An extension that publishes deferred tools should also make the deferred list discoverable, as the MCP extension does with its system-prompt section. Use `promptSnippet` to opt a custom tool into a one-line entry in `Available tools`, and `promptGuidelines` to append tool-specific bullets to the default `Guidelines` section when the tool is active. diff --git a/packages/coding-agent/docs/mcp-catalog.md b/packages/coding-agent/docs/mcp-catalog.md index fd45bcee4..687b25574 100644 --- a/packages/coding-agent/docs/mcp-catalog.md +++ b/packages/coding-agent/docs/mcp-catalog.md @@ -10,7 +10,30 @@ Step listens for `notifications/tools/list_changed` from the start of the connec A closed connection removes its tools and changes its status to failed. Session shutdown removes its MCP registrations and closes the clients. Results that settle after cancellation or connection closure cannot publish a catalog. Calls already in progress retain their captured tool definition when a live catalog is updated or a tool is removed. This behavior does not add an automatic reconnect policy. -`find_tools` reads the currently active session tools on every invocation, so late MCP and extension registrations become discoverable. Inactive and removed tools disappear from search. Each match includes its callable name, description, and JSON parameter schema. The optional `ExtensionContext.getToolCatalog()` accessor supplies this metadata through the runner's existing session actions. Direct SDK callers whose context does not provide that accessor retain the builtin Step catalog fallback. +## Tool exposure + +Every server tool is registered as `__`. A server's `exposure` in `config.toml` decides how the model reaches it: + +| `exposure` | Behavior | +| --- | --- | +| `deferred` (default) | Registered but not declared. `find_tools` loads the matches, which become callable on the next turn. | +| `direct` | Declared like any other active tool, for small and frequently called servers. | +| `hidden` | Registered but unreachable; left out of the tool catalog. | + +The optional `description` becomes the one-line summary in the `mcp_servers` system-prompt section, which lists the servers whose tools are deferred and names `find_tools` as the way to reach them. Without that section the model cannot know the servers exist, because neither `find_tools` nor the declared tool list names them. + +```toml +[mcp_servers.github] +url = "https://api.githubcopilot.com/mcp/" +description = "GitHub issues and pull requests" +# exposure = "direct" +``` + +## find_tools + +`find_tools` ranks the deferred tools plus the active ones by intent, tool name, description, and parameter names, and activates the deferred matches so the next request declares them. Activation is recorded on the tool result, so it survives resume and fork on that branch. Removed tools disappear from search. Each match includes its callable name, description, and JSON parameter schema. + +The optional `ExtensionContext.getActiveTools()`, `getAllTools()`, and `setActiveTools()` accessors supply the catalog and activation through the runner's existing session actions. Direct SDK callers whose context provides only `getToolCatalog()` get the ranked answer without activation, and hosts with neither retain the builtin Step catalog fallback. Extensions can replace part of their own catalog with the existing batch API: diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 59c2a8b9d..d5313cdab 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -1909,12 +1909,13 @@ export class AgentSession { getAllTools(): ToolInfo[] { const unavailable = this._extensionRunner.getUnavailableToolNames(); return Array.from(this._toolDefinitions.values()) - .filter(({ definition }) => !unavailable.has(definition.name)) + .filter(({ definition }) => !unavailable.has(definition.name) && definition.exposure !== "hidden") .map(({ definition, sourceInfo }) => ({ name: definition.name, description: definition.description, parameters: definition.parameters, promptGuidelines: definition.promptGuidelines, + exposure: definition.exposure, sourceInfo, })); } @@ -4076,10 +4077,12 @@ export class AgentSession { const excludedToolNames = this._excludedToolNames; const isAllowedTool = (name: string): boolean => (!allowedToolNames || allowedToolNames.has(name)) && !excludedToolNames?.has(name); - // A tool that opts out of the default selection stays registered but is not declared until - // the session activates it, through the tool settings or another tool's activation. - const isDefaultActive = (name: string): boolean => - this._toolDefinitions.get(name)?.definition.defaultActive !== false; + // A tool that is deferred or hidden stays registered but is not declared until the session + // activates it, through the tool settings, an allowlist, or another tool's activation. + const isDefaultActive = (name: string): boolean => { + const exposure = this._toolDefinitions.get(name)?.definition.exposure; + return exposure !== "deferred" && exposure !== "hidden"; + }; const registeredTools = this._extensionRunner.getAllRegisteredTools(); const allCustomTools = [ diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 8febc6080..3360e79ab 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -856,6 +856,18 @@ export class ExtensionRunner { const active = new Set(runner.runtime.getActiveTools()); return runner.runtime.getAllTools().filter((tool) => active.has(tool.name)); }, + getActiveTools: () => { + runner.assertActive(); + return runner.runtime.getActiveTools(); + }, + getAllTools: () => { + runner.assertActive(); + return runner.runtime.getAllTools(); + }, + setActiveTools: (toolNames) => { + runner.assertActive(); + runner.runtime.setActiveTools(toolNames); + }, get autoRetryEnabled() { runner.assertActive(); return runner.getAutoRetryEnabledFn?.(); diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 2aa27c71f..28db0b5e9 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -447,6 +447,12 @@ export interface ExtensionContext { getSystemPrompt(): string; /** Read the currently active tool definitions at call time, when the host provides a live catalog. */ getToolCatalog?(): readonly ToolInfo[]; + /** Currently active tool names, when the host exposes the live tool set. */ + getActiveTools?(): string[]; + /** Every registered tool, including deferred and hidden ones, when the host exposes it. */ + getAllTools?(): ToolInfo[]; + /** Activate tools by name, when the host allows a tool to load others. */ + setActiveTools?(toolNames: string[]): void; /** Whether Pi's native provider retry loop is enabled, when exposed by the host. */ readonly autoRetryEnabled?: boolean; /** Toggle Pi's native provider retry loop, when exposed by the host. */ @@ -558,6 +564,11 @@ export interface ToolRenderContext { isError: boolean; } +/** + * How a registered tool reaches the model. See `ToolDefinition.exposure`. + */ +export type ToolExposure = "direct" | "deferred" | "hidden"; + /** * Tool definition for registerTool(). */ @@ -580,11 +591,13 @@ export interface ToolDefinition Static; @@ -1826,7 +1839,7 @@ export type GetSessionNameHandler = () => string | undefined; export type GetActiveToolsHandler = () => string[]; /** Tool info with name, description, parameter schema, prompt guidelines, and source metadata. */ -export type ToolInfo = Pick & { +export type ToolInfo = Pick & { sourceInfo: SourceInfo; }; diff --git a/packages/coding-agent/src/core/tools/tool-search.ts b/packages/coding-agent/src/core/tools/tool-search.ts index 9b54ab7fe..fc84d4adf 100644 --- a/packages/coding-agent/src/core/tools/tool-search.ts +++ b/packages/coding-agent/src/core/tools/tool-search.ts @@ -1,19 +1,15 @@ /** - * Tool discovery: a BM25 ranker over tool metadata, and the `tool_search` tool built on it. + * Tool discovery: a BM25 ranker over tool metadata for the model-facing tool finder. * - * `tool_search` searches the tools that are registered but not declared to the model and activates - * the matches, so the next model call declares them. Activation goes through the active tool set, - * so it is recorded in the transcript like any other tool change and survives resume and fork on - * that branch. + * The finder ranks the tools that are registered but not declared to the model and activates the + * matches, so the next model call declares them. Activation goes through the active tool set, so + * it is recorded in the transcript like any other tool change and survives resume and fork. */ -import { type Static, Type } from "typebox"; -import type { AgentToolResult, ToolDefinition, ToolInfo } from "../extensions/types.ts"; - -export const TOOL_SEARCH_TOOL_NAME = "tool_search"; +import type { ToolInfo } from "../extensions/types.ts"; export const DEFAULT_TOOL_SEARCH_LIMIT = 8; -/** The part of the extension runtime `tool_search` needs, so it stays testable without a session. */ +/** The part of the host tool set the finder needs, kept narrow so it is testable without a session. */ export interface ToolSearchCatalog { getActiveTools(): string[]; getAllTools(): ToolInfo[]; @@ -219,75 +215,17 @@ export class Bm25Ranker { } } -export const toolSearchSchema = Type.Object({ - query: Type.String({ description: "Search query for deferred tools." }), - limit: Type.Optional( - Type.Number({ description: `Maximum number of tools to return. Defaults to ${DEFAULT_TOOL_SEARCH_LIMIT}.` }), - ), -}); - -/** Rank the tools that are registered but not active yet, and activate the matches. */ -export function searchAndLoad(catalog: ToolSearchCatalog, query: string, limit: number): ToolSearchResult[] { +/** + * Rank the deferred tools that are not declared yet, and activate the matches. Hidden tools and + * tools that are already active are never candidates. + */ +export function searchAndLoad(catalog: ToolSearchCatalog, query: string, limit: number): ToolInfo[] { const active = catalog.getActiveTools(); - const inactive = new Set(active); - const candidates = catalog.getAllTools().filter((tool) => !inactive.has(tool.name)); - const documents = candidates.map((tool) => createToolSearchDocument(tool)); - const matches = new Bm25Ranker().rank(query, documents, limit); + const declared = new Set(active); + const candidates = catalog.getAllTools().filter((tool) => tool.exposure === "deferred" && !declared.has(tool.name)); + const matches = new Bm25Ranker().rank(query, candidates.map(createToolSearchDocument), limit); if (matches.length === 0) return []; catalog.setActiveTools([...active, ...matches.map((match) => match.name)]); - const descriptions = new Map(candidates.map((tool) => [tool.name, tool.description])); - return matches.map((match) => ({ name: match.name, description: descriptions.get(match.name) ?? "" })); -} - -/** - * The `tool_search` description. It does not list the searchable tools, so it stays the same while - * tools are registered, for example when MCP servers connect. - */ -export const TOOL_SEARCH_DESCRIPTION = `# Tool discovery - -Searches the tools that are not loaded yet and declares the matches for the next model call. - -A tool that is not in your tool list may still be registered, for example one you turned off or -one that a plugin publishes. Use this tool (\`${TOOL_SEARCH_TOOL_NAME}\`) to look for it instead of -telling the user the tool does not exist.`; - -export interface ToolSearchDetails { - loaded: string[]; -} - -/** - * The tool definition with a two-argument `execute`: `tool_search` needs no abort signal, progress - * callback, or extension context, and a narrower signature still registers as a tool. - */ -export type ToolSearchToolDefinition = Omit< - ToolDefinition, - "parameters" | "execute" -> & { - parameters: typeof toolSearchSchema; - execute(toolCallId: string, params: Static): Promise>; -}; - -export function createToolSearchToolDefinition(catalog: ToolSearchCatalog): ToolSearchToolDefinition { - return { - name: TOOL_SEARCH_TOOL_NAME, - label: "Tool Search", - description: TOOL_SEARCH_DESCRIPTION, - promptSnippet: "Search for tools that are not loaded yet and load the matches", - parameters: toolSearchSchema, - // Inactive until the session asks for it; see `_refreshToolRegistry`. - defaultActive: false, - async execute(_toolCallId, { query, limit }) { - if (query.trim() === "") throw new Error("query must not be empty"); - const max = limit ?? DEFAULT_TOOL_SEARCH_LIMIT; - if (!Number.isInteger(max) || max <= 0) throw new Error("limit must be a positive integer"); - const found = searchAndLoad(catalog, query, max); - const text = - found.length === 0 - ? "No matching tools found." - : `Loaded ${found.length} tool${found.length === 1 ? "" : "s"}. They are available from your next call:\n${found - .map((tool) => `- ${tool.name}: ${tool.description.trim().split(/\r?\n/)[0]}`) - .join("\n")}`; - return { content: [{ type: "text", text }], details: { loaded: found.map((tool) => tool.name) } }; - }, - }; + const byName = new Map(candidates.map((tool) => [tool.name, tool])); + return matches.map((match) => byName.get(match.name)!).filter(Boolean); } diff --git a/packages/coding-agent/src/features/step-tool-search.ts b/packages/coding-agent/src/features/step-tool-search.ts deleted file mode 100644 index 91504a8fd..000000000 --- a/packages/coding-agent/src/features/step-tool-search.ts +++ /dev/null @@ -1,14 +0,0 @@ -import type { ExtensionAPI, ExtensionFactory } from "../core/extensions/types.ts"; -import { createToolSearchToolDefinition } from "../core/tools/tool-search.ts"; - -/** - * Mount `tool_search`: it is registered inactive and finds the tools that are registered but not - * declared to the model, so a session can keep a large tool catalogue out of the prompt. - * - * Enable it per session with `defaultTools`, `--tools`, or the tool settings. - */ -export function createStepToolSearchExtension(): ExtensionFactory { - return (pi: ExtensionAPI): void => { - pi.registerTool(createToolSearchToolDefinition(pi)); - }; -} diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 3fd652af4..51134e1f3 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -526,7 +526,6 @@ export { stepScheduleExtensionInline, UpdateGoalParams, } from "./features/step-schedule.ts"; -export { createStepToolSearchExtension } from "./features/step-tool-search.ts"; export { type CompletedSubagentResponse, readSubagentLaneSnapshot, diff --git a/packages/coding-agent/src/step/config-toml.ts b/packages/coding-agent/src/step/config-toml.ts index e4737368b..6a675b7a2 100644 --- a/packages/coding-agent/src/step/config-toml.ts +++ b/packages/coding-agent/src/step/config-toml.ts @@ -20,6 +20,10 @@ export interface StepMcpServerConfig { tool_timeout_sec?: number; enabled_tools?: string[]; disabled_tools?: string[]; + /** How this server's tools reach the model. Defaults to `deferred`: loaded by `find_tools`. */ + exposure?: "direct" | "deferred" | "hidden"; + /** One-line summary shown in the system prompt when the tools are deferred. */ + description?: string; oauth?: { client_id?: string; client_secret?: string; scopes?: string[]; callback_port?: number }; } diff --git a/packages/coding-agent/src/step/mcp-catalog.test.ts b/packages/coding-agent/src/step/mcp-catalog.test.ts index 7d1919d68..8c82de7d6 100644 --- a/packages/coding-agent/src/step/mcp-catalog.test.ts +++ b/packages/coding-agent/src/step/mcp-catalog.test.ts @@ -823,11 +823,14 @@ test("find_tools sees a late MCP catalog and stops returning tools removed by th }; expect(await search()).toBe("(no matching tools)"); initial.release(); - await vi.waitFor(() => expect(session.getActiveToolNames()).toContain("live__calendar")); + // The catalog is published deferred, so the search both finds and loads it. + await vi.waitFor(() => expect(session.getAllTools().map((tool) => tool.name)).toContain("live__calendar")); + expect(session.getActiveToolNames()).not.toContain("live__calendar"); expect(await search()).toContain("live__calendar"); + expect(session.getActiveToolNames()).toContain("live__calendar"); expect(await search()).toContain('"calendar_id"'); await server.notify(); - await vi.waitFor(() => expect(session.getActiveToolNames()).not.toContain("live__calendar")); + await vi.waitFor(() => expect(session.getAllTools().map((tool) => tool.name)).not.toContain("live__calendar")); expect(await search()).toBe("(no matching tools)"); }); diff --git a/packages/coding-agent/src/step/mcp-startup.test.ts b/packages/coding-agent/src/step/mcp-startup.test.ts index fb321ed55..59784e331 100644 --- a/packages/coding-agent/src/step/mcp-startup.test.ts +++ b/packages/coding-agent/src/step/mcp-startup.test.ts @@ -21,6 +21,7 @@ import { discoverStepMcpServers, expandHeaderTemplate, getStepMcpStatuses, + renderMcpServersSection, resolveStepMcpServer, } from "./mcp.ts"; @@ -695,3 +696,69 @@ test("shutdown suppresses late warnings from manifest preinstallation", async () await shutdown; expect(harness.notify).not.toHaveBeenCalled(); }); + +test("publishes server tools as deferred and names the server in the prompt", async () => { + const server = await slowServer(2); + config.value = { mcp_servers: { slow: { url: server.url, description: "Searches tickets" } } }; + const harness = await setup("tui"); + await harness.start(); + server.release(); + await vi.waitFor(() => expect(harness.extension.tools.size).toBe(2)); + const published = [...harness.extension.tools.values()].map((tool) => tool.definition.exposure); + expect(published).toEqual(["deferred", "deferred"]); + + const prompt = harness.extension.handlers.get("before_agent_start")![0]!( + { type: "before_agent_start", prompt: "hi", systemPrompt: "BASE", systemPromptOptions: {} }, + { cwd: process.cwd(), mode: "tui", isProjectTrusted: () => true, ui: { notify: harness.notify } }, + ) as { systemPrompt?: string } | undefined; + + expect(prompt?.systemPrompt).toContain("BASE\n\n# MCP servers"); + expect(prompt?.systemPrompt).toContain("- slow: Searches tickets (2 tools)"); +}); + +test("honors a direct or hidden server exposure", async () => { + const direct = await slowServer(1); + const hidden = await slowServer(1); + config.value = { + mcp_servers: { + direct: { url: direct.url, exposure: "direct" }, + hidden: { url: hidden.url, exposure: "hidden" }, + }, + }; + const harness = await setup("tui"); + await harness.start(); + direct.release(); + hidden.release(); + await vi.waitFor(() => expect(harness.extension.tools.size).toBe(2)); + expect(harness.extension.tools.get("direct__tool_0")?.definition.exposure).toBe("direct"); + expect(harness.extension.tools.get("hidden__tool_0")?.definition.exposure).toBe("hidden"); + + // Neither server belongs in the deferred section. + const prompt = harness.extension.handlers.get("before_agent_start")![0]!( + { type: "before_agent_start", prompt: "hi", systemPrompt: "BASE", systemPromptOptions: {} }, + { cwd: process.cwd(), mode: "tui", isProjectTrusted: () => true, ui: { notify: harness.notify } }, + ) as { systemPrompt?: string } | undefined; + expect(prompt?.systemPrompt).toBeUndefined(); +}); + +test("renderMcpServersSection lists only deferred servers with tools", () => { + expect(renderMcpServersSection([])).toBeUndefined(); + expect( + renderMcpServersSection([ + { name: "empty", exposure: "deferred", tools: [] }, + { name: "direct", exposure: "direct", tools: [{ name: "a" } as never] }, + ]), + ).toBeUndefined(); + const section = renderMcpServersSection([ + { name: "b", exposure: "deferred", description: "Second\nignored line", tools: [{ name: "a" } as never] }, + { name: "a", exposure: "deferred", tools: [{ name: "x" } as never, { name: "y" } as never] }, + ]); + expect(section).toBe( + [ + "# MCP servers", + "These servers are connected, but their tools are not declared. Search for one with `find_tools` before telling the user it is unavailable.", + "- a: (2 tools)", + "- b: Second (1 tool)", + ].join("\n"), + ); +}); diff --git a/packages/coding-agent/src/step/mcp.ts b/packages/coding-agent/src/step/mcp.ts index e6680b2d6..c4ca3d633 100644 --- a/packages/coding-agent/src/step/mcp.ts +++ b/packages/coding-agent/src/step/mcp.ts @@ -7,7 +7,7 @@ import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js" import { StreamableHTTPClientTransport, StreamableHTTPError } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; import { type Tool as McpTool, ToolListChangedNotificationSchema } from "@modelcontextprotocol/sdk/types.js"; import type { AgentToolResult } from "@step-harness/agent-core"; -import type { ExtensionAPI, ExtensionFactory } from "../core/extensions/types.ts"; +import type { ExtensionAPI, ExtensionFactory, ToolExposure } from "../core/extensions/types.ts"; import { theme } from "../theme/theme.ts"; import { readGlobalStepConfig } from "./config-toml.ts"; import { createMcpToolCaller, listAllMcpTools } from "./mcp-client.ts"; @@ -42,6 +42,10 @@ interface ConnectedServer { readonly client: Client; readonly transport: StdioClientTransport | StreamableHTTPClientTransport; tools: McpTool[]; + /** How this server's tools reach the model. */ + exposure: ToolExposure; + /** One-line summary for the system prompt, from the declaration. */ + description?: string; /** Ends when this connection closes or its session shuts down. */ readonly signal: AbortSignal; readonly catalogTimeoutMs: number; @@ -63,6 +67,14 @@ interface ServerDeclaration { tool_timeout_sec?: number; enabled_tools?: string[]; disabled_tools?: string[]; + /** + * How this server's tools reach the model. Defaults to `deferred`: they stay undeclared until + * `find_tools` loads them. Use `direct` for a small, frequently called server and `hidden` for + * one that must stay unreachable. + */ + exposure?: ToolExposure; + /** One-line summary shown in the system prompt when the tools are deferred. */ + description?: string; oauth?: { client_id?: string; client_secret?: string; @@ -270,6 +282,14 @@ export function createStepMcpExtension(): ExtensionFactory { if (ctx.mode !== "tui") await startup; }); + // The deferred tool list is invisible to the model until the prompt names the servers that + // hold it. Recomputed per turn so a server that connects later becomes searchable. + pi.on("before_agent_start", (event) => { + const section = renderMcpServersSection(servers); + if (!section) return; + return { systemPrompt: `${event.systemPrompt}\n\n${section}` }; + }); + pi.on("session_shutdown", async () => { cancellation?.abort(); await startup; @@ -563,6 +583,10 @@ function normalizeDeclaration(value: Record): ServerDeclaration for (const key of ["enabled_tools", "disabled_tools"] as const) { if (Array.isArray(value[key])) declaration[key] = value[key].filter(isString); } + if (value.exposure === "direct" || value.exposure === "deferred" || value.exposure === "hidden") + declaration.exposure = value.exposure; + if (typeof value.description === "string" && value.description.trim()) + declaration.description = value.description.trim(); if (isRecord(value.oauth)) { declaration.oauth = { ...(typeof value.oauth.client_id === "string" ? { client_id: value.oauth.client_id } : {}), @@ -634,6 +658,8 @@ export async function connectStepMcpServer( client, transport, tools: selectDeclaredTools(tools, input.declaration), + exposure: serverExposure(input.declaration), + description: input.declaration.description, signal: lifetime, catalogTimeoutMs: timeout, callTimeoutMs, @@ -768,6 +794,35 @@ function remoteToolName(server: Pick, remote: McpTool): return `${server.name}${PLUGIN_SERVER_SEPARATOR}${sanitizeName(remote.name)}`; } +/** A server's tool exposure: deferred unless its declaration says otherwise. */ +function serverExposure(declaration: ServerDeclaration): ToolExposure { + return declaration.exposure ?? "deferred"; +} + +/** + * The `mcp_servers` system-prompt section: servers whose tools are deferred, with the one-line + * summary their declaration carries. Without it the model cannot know they exist, because neither + * `find_tools` nor the tool list names them. + */ +export function renderMcpServersSection( + servers: readonly Pick[], +): string | undefined { + const listed = servers + .filter((server) => server.exposure === "deferred" && server.tools.length > 0) + .sort((a, b) => a.name.localeCompare(b.name)); + if (listed.length === 0) return undefined; + const lines = listed.map((server) => { + const summary = server.description?.split("\n", 1)[0]?.trim(); + const count = `${server.tools.length} tool${server.tools.length === 1 ? "" : "s"}`; + return `- ${server.name}: ${summary ? `${summary} ` : ""}(${count})`; + }); + return [ + "# MCP servers", + `These servers are connected, but their tools are not declared. Search for one with \`find_tools\` before telling the user it is unavailable.`, + ...lines, + ].join("\n"); +} + function createRemoteTool(server: ConnectedServer, remote: McpTool) { const name = remoteToolName(server, remote); const call = createMcpToolCaller(server.client, remote); @@ -778,6 +833,7 @@ function createRemoteTool(server: ConnectedServer, remote: McpTool) { // The provider boundary accepts JSON Schema directly. Reconstructing it // as TypeBox types loses constraints, local references and combinators. parameters: remote.inputSchema, + exposure: server.exposure, execute: async ( _toolCallId: string, params: unknown, diff --git a/packages/coding-agent/src/step/permissions.ts b/packages/coding-agent/src/step/permissions.ts index ff7a29fe5..86056659b 100644 --- a/packages/coding-agent/src/step/permissions.ts +++ b/packages/coding-agent/src/step/permissions.ts @@ -81,9 +81,8 @@ const READ_ONLY_TOOLS = new Set([ "search_files", "search_web", "read_file", + // Loads deferred tool definitions for the next model call; it touches no files and runs no command. "find_tools", - // Loads tool definitions for the next model call; it touches no files and runs no command. - "tool_search", // Keep native names safe when a caller explicitly enables a native tool. "ls", "find", diff --git a/packages/coding-agent/src/step/tool-profile.ts b/packages/coding-agent/src/step/tool-profile.ts index 22d8f5eca..c0da4a83b 100644 --- a/packages/coding-agent/src/step/tool-profile.ts +++ b/packages/coding-agent/src/step/tool-profile.ts @@ -52,6 +52,7 @@ import { import type { LsToolOptions } from "../core/tools/ls.ts"; import { pathExists, resolveReadPathAsync, resolveToCwd } from "../core/tools/path-utils.ts"; import type { ReadToolOptions } from "../core/tools/read.ts"; +import { Bm25Ranker, searchAndLoad } from "../core/tools/tool-search.ts"; import type { WriteToolOptions } from "../core/tools/write.ts"; import { detectSupportedImageMimeTypeFromFile } from "../utils/mime.ts"; import { resolvePath } from "../utils/paths.ts"; @@ -1474,41 +1475,76 @@ function createFindToolsDefinition( return { name: "find_tools", label: "find_tools", - description: "Search registered tools by natural-language intent, tool name, description, and parameter names.", + description: + "Search registered tools by natural-language intent, tool name, description, and parameter names. Tools that are not loaded yet, such as the tools of MCP servers, are loaded by this call and become callable on your next turn.", promptSnippet: "Find a tool by describing the operation you need", parameters: findToolsSchema, execute: async (_toolCallId, args: FindToolsInput, _signal, _onUpdate, ctx) => { - const definitions = ctx?.getToolCatalog?.() ?? fallbackDefinitions; - const queryTokens = args.query.toLowerCase().match(/[\p{L}\p{N}_]+/gu) ?? []; const limit = Math.max(1, Math.min(20, args.limit ?? 8)); - const matches = definitions - .map((definition) => { - const haystack = - `${definition.name} ${definition.description} ${JSON.stringify(definition.parameters)}`.toLowerCase(); - const score = queryTokens.reduce((total, token) => total + (haystack.includes(token) ? 1 : 0), 0); - return { definition, score }; - }) - .filter((entry) => entry.score > 0) - .sort( - (left, right) => right.score - left.score || left.definition.name.localeCompare(right.definition.name), - ) - .slice(0, limit); - const content = matches.length - ? matches - .map( - ({ definition, score }, index) => - `${index + 1}. ${definition.name} [score=${score}]\ndescription: ${definition.description}\nparameters: ${JSON.stringify(definition.parameters)}`, - ) - .join("\n\n") - : "(no matching tools)"; - return { - content: [{ type: "text", text: content }], - details: undefined, - }; + const catalog = + ctx?.getActiveTools && ctx.getAllTools && ctx.setActiveTools + ? { + getActiveTools: () => ctx.getActiveTools!(), + getAllTools: () => ctx.getAllTools!(), + setActiveTools: (names: string[]) => ctx.setActiveTools!(names), + } + : undefined; + if (!catalog) + return describeFallbackMatches(args.query, limit, ctx?.getToolCatalog?.() ?? fallbackDefinitions); + + // Load the deferred matches first, then describe everything the query matches, so the + // answer covers both the tools this call made available and the ones already loaded. + const loaded = searchAndLoad(catalog, args.query, limit); + const active = new Set(catalog.getActiveTools()); + const visible = catalog.getAllTools().filter((tool) => active.has(tool.name)); + return describeFoundTools( + visible.map((tool) => ({ + name: tool.name, + description: tool.description, + parameters: tool.parameters, + })), + args.query, + limit, + loaded.map((tool) => tool.name), + ); }, }; } +/** The same ranking and answer text, without a host tool set to activate against. */ +function describeFallbackMatches( + query: string, + limit: number, + definitions: readonly { name: string; description: string; parameters: unknown }[], +): { content: [{ type: "text"; text: string }]; details: undefined } { + return describeFoundTools(definitions, query, limit, []); +} + +/** Answer text: ranked matches with their schema, plus the names this call activated. */ +function describeFoundTools( + definitions: readonly { name: string; description: string; parameters: unknown }[], + query: string, + limit: number, + loaded: readonly string[], +): { content: [{ type: "text"; text: string }]; details: undefined } { + const documents = definitions.map((definition) => ({ + name: definition.name, + text: `${definition.description} ${JSON.stringify(definition.parameters)}`, + })); + const matches = new Bm25Ranker().rank(query, documents, limit); + const byName = new Map(definitions.map((definition) => [definition.name, definition])); + const lines = matches.length + ? matches + .map((match, index) => { + const definition = byName.get(match.name)!; + return `${index + 1}. ${definition.name}\ndescription: ${definition.description}\nparameters: ${JSON.stringify(definition.parameters)}`; + }) + .join("\n\n") + : "(no matching tools)"; + const loadedLines = loaded.length ? `\n\nNewly available from your next call: ${loaded.join(", ")}` : ""; + return { content: [{ type: "text", text: `${lines}${loadedLines}` }], details: undefined }; +} + /** * Detached run for long-lived processes (dev servers, watchers). The child * outlives the tool call but not the session: its pid stays in the detached diff --git a/packages/coding-agent/test/step-tool-profile.test.ts b/packages/coding-agent/test/step-tool-profile.test.ts index 298b7ae01..b61b2fb09 100644 --- a/packages/coding-agent/test/step-tool-profile.test.ts +++ b/packages/coding-agent/test/step-tool-profile.test.ts @@ -78,7 +78,7 @@ describe("Step tool profile", () => { }, find_tools: { description: - "Search registered tools by natural-language intent, tool name, description, and parameter names.", + "Search registered tools by natural-language intent, tool name, description, and parameter names. Tools that are not loaded yet, such as the tools of MCP servers, are loaded by this call and become callable on your next turn.", promptSnippet: "Find a tool by describing the operation you need", }, }; diff --git a/packages/coding-agent/test/tool-search-integration.test.ts b/packages/coding-agent/test/tool-search-integration.test.ts index a8962afca..73e076e87 100644 --- a/packages/coding-agent/test/tool-search-integration.test.ts +++ b/packages/coding-agent/test/tool-search-integration.test.ts @@ -5,38 +5,35 @@ import { fauxAssistantMessage, fauxToolCall, registerFauxProvider } from "@step- import { Type } from "typebox"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; -import type { ExtensionAPI } from "../src/core/extensions/types.ts"; +import type { InlineExtension } from "../src/core/extensions/types.ts"; import { ModelRuntime } from "../src/core/model-runtime.ts"; import { DefaultResourceLoader } from "../src/core/resource-loader.ts"; import { createAgentSession } from "../src/core/sdk.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; -import { createStepToolSearchExtension } from "../src/features/step-tool-search.ts"; +import { createStepToolProfile } from "../src/step/tool-profile.ts"; -/** - * A plugin tool the model never sees until `tool_search` loads it: registered, but opted out of - * the default active set. - */ -function deferredTool(name: string, description: string, query: string) { +/** A plugin tool the model never sees until `find_tools` loads it. */ +function deferredTool(name: string, description: string) { return { name, label: name, description, - defaultActive: false, + exposure: "deferred" as const, parameters: Type.Object({ query: Type.String() }), async execute(_toolCallId: string) { - return { content: [{ type: "text" as const, text: `${name}:${query}` }], details: {} }; + return { content: [{ type: "text" as const, text: `${name} ran` }], details: {} }; }, }; } -describe("tool_search over a real model turn", () => { +describe("find_tools over a real model turn", () => { let tempDir: string; let faux: ReturnType; - const cleanups: Array<() => void | Promise> = []; + const cleanups: Array<() => unknown> = []; beforeEach(() => { - tempDir = join(tmpdir(), `pi-tool-search-e2e-${Date.now()}-${Math.random().toString(36).slice(2)}`); + tempDir = join(tmpdir(), `pi-find-tools-e2e-${Date.now()}-${Math.random().toString(36).slice(2)}`); mkdirSync(tempDir, { recursive: true }); }); @@ -46,7 +43,7 @@ describe("tool_search over a real model turn", () => { if (tempDir && existsSync(tempDir)) rmSync(tempDir, { recursive: true, force: true }); }); - async function createSession(extensionFactories: ((pi: ExtensionAPI) => void)[]) { + async function createSession(extensionFactories: InlineExtension[]) { faux = registerFauxProvider(); const authStorage = AuthStorage.inMemory(); await authStorage.modify(faux.getModel().provider, async () => ({ type: "api_key", key: "faux-key" })); @@ -72,9 +69,7 @@ describe("tool_search over a real model turn", () => { }, ], }); - const settingsManager = SettingsManager.inMemory({ - defaultTools: ["read", "bash", "edit", "write", "tool_search"], - }); + const settingsManager = SettingsManager.inMemory({}); const resourceLoader = new DefaultResourceLoader({ cwd: tempDir, agentDir: tempDir, @@ -93,6 +88,9 @@ describe("tool_search over a real model turn", () => { settingsManager, sessionManager: SessionManager.inMemory(tempDir), resourceLoader, + // The Step profile supplies `find_tools` and turns the native built-ins off. + customTools: createStepToolProfile(tempDir), + noTools: "builtin", }); await session.bindExtensions({}); cleanups.push(() => session.dispose()); @@ -100,59 +98,50 @@ describe("tool_search over a real model turn", () => { } it("declares the loaded tool on the next request", async () => { - const declaredToolNames: string[][] = []; + const declared: string[][] = []; const toolResultTexts: string[] = []; const session = await createSession([ (pi) => { - pi.registerTool(deferredTool("issue_search", "Lists open issues in a repository", "query")); - pi.registerTool(deferredTool("calendar_read", "Reads calendar entries", "query")); - }, - createStepToolSearchExtension(), - (pi) => { + pi.registerTool(deferredTool("issue_search", "Lists open issues in a repository")); + pi.registerTool(deferredTool("calendar_read", "Reads calendar entries")); pi.on("tool_result", (event) => { - toolResultTexts.push( - event.content - .map((block: { type: string; text?: string }) => (block.type === "text" ? block.text : "")) - .join(""), - ); + toolResultTexts.push(event.content.map((block) => (block.type === "text" ? block.text : "")).join("")); }); }, ]); faux.setResponses([ (context) => { - declaredToolNames.push((context.tools ?? []).map((tool: { name: string }) => tool.name).sort()); - return fauxAssistantMessage([fauxToolCall("tool_search", { query: "issues" })]); + declared.push((context.tools ?? []).map((tool: { name: string }) => tool.name).sort()); + return fauxAssistantMessage(fauxToolCall("find_tools", { query: "open issues" }), { + stopReason: "toolUse", + }); }, (context) => { - declaredToolNames.push((context.tools ?? []).map((tool: { name: string }) => tool.name).sort()); + declared.push((context.tools ?? []).map((tool: { name: string }) => tool.name).sort()); return fauxAssistantMessage("done"); }, ]); await session.prompt("list the open issues"); - expect(declaredToolNames).toHaveLength(2); // First request: the deferred tools are registered but not declared. - expect(declaredToolNames[0]).toContain("tool_search"); - expect(declaredToolNames[0]).not.toContain("issue_search"); - expect(declaredToolNames[0]).not.toContain("calendar_read"); + expect(declared[0]).toContain("find_tools"); + expect(declared[0]).not.toContain("issue_search"); + expect(declared[0]).not.toContain("calendar_read"); // Second request: the loaded tool is declared, the unrelated one still is not. - expect(declaredToolNames[1]).toContain("issue_search"); - expect(declaredToolNames[1]).not.toContain("calendar_read"); - expect(toolResultTexts.join("\n")).toContain("Loaded 1 tool"); - expect(toolResultTexts.join("\n")).toContain("- issue_search: Lists open issues in a repository"); + expect(declared[1]).toContain("issue_search"); + expect(declared[1]).not.toContain("calendar_read"); + expect(toolResultTexts.join("\n")).toContain("Newly available from your next call: issue_search"); + expect(toolResultTexts.join("\n")).toContain("1. issue_search"); }); it("records the loaded tools as an active-tool change on the tool result", async () => { + const addedToolNames: string[][] = []; const session = await createSession([ - (pi) => { - pi.registerTool(deferredTool("issue_search", "Lists open issues in a repository", "query")); - }, - createStepToolSearchExtension(), + (pi) => pi.registerTool(deferredTool("issue_search", "Lists open issues in a repository")), ]); - const addedToolNames: string[][] = []; faux.setResponses([ - () => fauxAssistantMessage(fauxToolCall("tool_search", { query: "issues" }), { stopReason: "toolUse" }), + () => fauxAssistantMessage(fauxToolCall("find_tools", { query: "open issues" }), { stopReason: "toolUse" }), (context) => { addedToolNames.push( context.messages diff --git a/packages/coding-agent/test/tool-search-mcp.test.ts b/packages/coding-agent/test/tool-search-mcp.test.ts new file mode 100644 index 000000000..3b97937e8 --- /dev/null +++ b/packages/coding-agent/test/tool-search-mcp.test.ts @@ -0,0 +1,223 @@ +import { mkdirSync, rmSync } from "node:fs"; +import { createServer } from "node:http"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { fauxAssistantMessage, fauxToolCall, registerFauxProvider } from "@step-harness/providers/compat"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { AuthStorage } from "../src/core/auth-storage.ts"; +import type { InlineExtension } from "../src/core/extensions/types.ts"; +import { ModelRuntime } from "../src/core/model-runtime.ts"; +import { DefaultResourceLoader } from "../src/core/resource-loader.ts"; +import { createAgentSession } from "../src/core/sdk.ts"; +import { SessionManager } from "../src/core/session-manager.ts"; +import { SettingsManager } from "../src/core/settings-manager.ts"; +import type { StepConfigDocument } from "../src/step/config-toml.ts"; +import { createStepMcpExtension } from "../src/step/mcp.ts"; +import { createStepToolProfile } from "../src/step/tool-profile.ts"; + +const config = vi.hoisted(() => ({ value: {} as StepConfigDocument })); +vi.mock("../src/step/config-toml.ts", () => ({ readGlobalStepConfig: () => config.value })); +vi.mock("../src/step/plugins.ts", () => ({ + defaultStepPluginsDir: () => "/unused-test-plugins", + ensureBuiltinPluginsInstalled: async () => ({ installed: [], warnings: [] }), + listStepPluginDirectories: async () => [], + provisionInstallCommand: () => undefined, + readStepPluginManifest: async () => ({ manifest: undefined, errors: [] }), +})); +vi.mock("../src/step/steppage-provision.ts", () => ({ + ensureStepPageReady: async () => ({ command: undefined }), + isManagedStepPageCommand: () => false, +})); +vi.mock("../src/step/mcp-oauth.ts", () => ({ hasStoredMcpOAuthCredential: () => false })); + +/** A one-tool MCP server over streamable HTTP: handshake, list, and one call. */ +async function mcpServer() { + const calls: Array> = []; + const server = createServer(async (req, res) => { + if (req.method !== "POST") { + res.writeHead(405).end(); + return; + } + const chunks: Buffer[] = []; + for await (const chunk of req) chunks.push(Buffer.from(chunk)); + const message = JSON.parse(Buffer.concat(chunks).toString()) as { + id?: number; + method: string; + params?: { name?: string; arguments?: Record }; + }; + if (message.id === undefined) { + res.writeHead(202).end(); + return; + } + let result: unknown; + if (message.method === "initialize") { + result = { + protocolVersion: "2025-03-26", + capabilities: { tools: {} }, + serverInfo: { name: "tickets", version: "1" }, + }; + } else if (message.method === "tools/list") { + result = { + tools: [ + { + name: "search_tickets", + description: "Searches support tickets by text.", + inputSchema: { + type: "object", + properties: { text: { type: "string", description: "Text to search for" } }, + required: ["text"], + }, + }, + ], + }; + } else { + calls.push(message.params?.arguments ?? {}); + result = { content: [{ type: "text", text: `ticket about ${message.params?.arguments?.text ?? ""}` }] }; + } + res.writeHead(200, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ jsonrpc: "2.0", id: message.id, result })); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const address = server.address(); + if (!address || typeof address === "string") throw new Error("Missing test port"); + return { + url: `http://127.0.0.1:${address.port}/mcp`, + calls, + close: async () => { + server.closeAllConnections(); + await new Promise((resolve) => server.close(() => resolve())); + }, + }; +} + +describe("MCP tools behind tool_search", () => { + let tempDir: string; + let faux: ReturnType; + const cleanups: Array<() => unknown> = []; + + beforeEach(() => { + config.value = {}; + tempDir = join(tmpdir(), `pi-tool-search-mcp-${Date.now()}-${Math.random().toString(36).slice(2)}`); + mkdirSync(tempDir, { recursive: true }); + }); + + afterEach(async () => { + while (cleanups.length > 0) await cleanups.pop()?.(); + faux?.unregister(); + rmSync(tempDir, { recursive: true, force: true }); + }); + + async function createSession(extensionFactories: InlineExtension[], defaultTools?: string[]) { + faux = registerFauxProvider(); + const authStorage = AuthStorage.inMemory(); + await authStorage.modify(faux.getModel().provider, async () => ({ type: "api_key", key: "faux-key" })); + const modelRuntime = await ModelRuntime.create({ + credentials: authStorage, + modelsPath: join(tempDir, "models.json"), + }); + const model = faux.getModel(); + modelRuntime.registerProvider(model.provider, { + baseUrl: model.baseUrl, + api: model.api, + models: [ + { + id: model.id, + name: model.name, + api: model.api, + reasoning: model.reasoning, + input: model.input, + cost: model.cost, + contextWindow: model.contextWindow, + maxTokens: model.maxTokens, + baseUrl: model.baseUrl, + }, + ], + }); + const settingsManager = SettingsManager.inMemory(defaultTools ? { defaultTools } : {}); + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir: tempDir, + settingsManager, + extensionFactories, + noSkills: true, + noPromptTemplates: true, + noThemes: true, + }); + await resourceLoader.reload(); + const { session } = await createAgentSession({ + cwd: tempDir, + agentDir: tempDir, + model, + modelRuntime, + settingsManager, + sessionManager: SessionManager.inMemory(tempDir), + resourceLoader, + // The Step profile supplies `find_tools` and turns the native built-ins off. + customTools: createStepToolProfile(tempDir), + noTools: "builtin", + }); + await session.bindExtensions({}); + cleanups.push(() => session.dispose()); + return session; + } + + it("finds an MCP tool the model cannot see, then calls it", async () => { + const server = await mcpServer(); + cleanups.push(server.close); + config.value = { mcp_servers: { tickets: { url: server.url, description: "Support tickets" } } }; + + const session = await createSession([createStepMcpExtension()]); + const declared: string[][] = []; + const prompts: string[] = []; + faux.setResponses([ + (context) => { + declared.push((context.tools ?? []).map((tool: { name: string }) => tool.name).sort()); + prompts.push(context.systemPrompt ?? ""); + return fauxAssistantMessage(fauxToolCall("find_tools", { query: "search support tickets" }), { + stopReason: "toolUse", + }); + }, + (context) => { + declared.push((context.tools ?? []).map((tool: { name: string }) => tool.name).sort()); + return fauxAssistantMessage(fauxToolCall("tickets__search_tickets", { text: "login fails" }), { + stopReason: "toolUse", + }); + }, + () => fauxAssistantMessage("done"), + ]); + + await session.prompt("find the ticket about the login failure"); + + // The server's tools stay undeclared until `find_tools` loads them. + expect(declared[0]).toContain("find_tools"); + expect(declared[0]?.some((name) => name.startsWith("tickets__"))).toBe(false); + // The prompt names the server so the model knows to search. + expect(prompts[0]).toContain("# MCP servers"); + expect(prompts[0]).toContain("- tickets: Support tickets (1 tool)"); + // The next request declares the loaded tool, which the model then calls. + expect(declared[1]).toContain("tickets__search_tickets"); + expect(server.calls).toEqual([{ text: "login fails" }]); + expect(session.messages.at(-1)).toMatchObject({ role: "assistant" }); + }); + + it("declares a server's tools directly when its declaration says so", async () => { + const server = await mcpServer(); + cleanups.push(server.close); + config.value = { mcp_servers: { tickets: { url: server.url, exposure: "direct" } } }; + + const session = await createSession([createStepMcpExtension()]); + const declared: string[][] = []; + faux.setResponses([ + (context) => { + declared.push((context.tools ?? []).map((tool: { name: string }) => tool.name).sort()); + return fauxAssistantMessage("done"); + }, + ]); + + await session.prompt("hello"); + + expect(declared[0]).toContain("tickets__search_tickets"); + // No deferred server, so the prompt carries no server section. + expect(session.systemPrompt).not.toContain("# MCP servers"); + }); +}); diff --git a/packages/coding-agent/test/tool-search.test.ts b/packages/coding-agent/test/tool-search.test.ts index e04aa2082..9097580e8 100644 --- a/packages/coding-agent/test/tool-search.test.ts +++ b/packages/coding-agent/test/tool-search.test.ts @@ -3,7 +3,7 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { Type } from "typebox"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; -import type { ToolInfo } from "../src/core/extensions/types.ts"; +import type { InlineExtension, ToolInfo } from "../src/core/extensions/types.ts"; import { DefaultResourceLoader } from "../src/core/resource-loader.ts"; import { createAgentSession } from "../src/core/sdk.ts"; import { SessionManager } from "../src/core/session-manager.ts"; @@ -12,21 +12,25 @@ import { createSyntheticSourceInfo } from "../src/core/source-info.ts"; import { Bm25Ranker, createToolSearchDocument, - createToolSearchToolDefinition, DEFAULT_TOOL_SEARCH_LIMIT, searchAndLoad, type ToolSearchCatalog, tokenize, } from "../src/core/tools/tool-search.ts"; -import { createStepToolSearchExtension } from "../src/features/step-tool-search.ts"; import { decideStepToolCall, stepPermissionStateForPreset } from "../src/step/permissions.ts"; import { stepModel } from "./utilities.ts"; -function toolInfo(name: string, description: string, parameters = Type.Object({})): ToolInfo { +function toolInfo( + name: string, + description: string, + exposure: ToolInfo["exposure"] = "deferred", + parameters = Type.Object({}), +): ToolInfo { return { name, description, parameters, + exposure, sourceInfo: createSyntheticSourceInfo(``, { source: "extension" }), }; } @@ -86,11 +90,10 @@ describe("tool search ranker", () => { describe("searchAndLoad", () => { const inactive = toolInfo("issues", "Lists open issues in a repository"); - const active = toolInfo("read", "Reads a file"); + const active = toolInfo("read", "Reads a file", "direct"); it("activates the matches and keeps the active tools", () => { - const tools = [active, inactive]; - const catalog = fakeCatalog(tools, ["read"]); + const catalog = fakeCatalog([active, inactive], ["read"]); const found = searchAndLoad(catalog, "open issues", DEFAULT_TOOL_SEARCH_LIMIT); expect(found.map((tool) => tool.name)).toEqual(["issues"]); expect(found[0]?.description).toBe("Lists open issues in a repository"); @@ -107,48 +110,32 @@ describe("searchAndLoad", () => { const catalog = fakeCatalog([active, inactive], ["read", "issues"]); expect(searchAndLoad(catalog, "issues", DEFAULT_TOOL_SEARCH_LIMIT)).toEqual([]); }); -}); - -describe("tool_search tool", () => { - it("reports the loaded tools", async () => { - const catalog = fakeCatalog([toolInfo("issues", "Lists open issues")], []); - const definition = createToolSearchToolDefinition(catalog); - const result = await definition.execute("call-1", { query: "issues" }); - expect(result.content).toEqual([ - { type: "text", text: "Loaded 1 tool. They are available from your next call:\n- issues: Lists open issues" }, - ]); - expect(result.details).toEqual({ loaded: ["issues"] }); - }); - - it("reports an empty result without activating anything", async () => { - const catalog = fakeCatalog([toolInfo("issues", "Lists open issues")], []); - const result = await createToolSearchToolDefinition(catalog).execute("call-1", { query: "kubernetes" }); - expect(result.content).toEqual([{ type: "text", text: "No matching tools found." }]); - expect(catalog.activations).toEqual([]); - }); - it("rejects an empty query and a bad limit", async () => { - const definition = createToolSearchToolDefinition(fakeCatalog([], [])); - await expect(definition.execute("call-1", { query: " " })).rejects.toThrow("query must not be empty"); - await expect(definition.execute("call-1", { query: "x", limit: 0 })).rejects.toThrow( - "limit must be a positive integer", - ); + it("skips a direct or hidden tool", () => { + const direct = toolInfo("issues_direct", "Lists open issues directly", "direct"); + const hidden = toolInfo("issues_hidden", "Lists open issues secretly", "hidden"); + const catalog = fakeCatalog([direct, hidden], []); + expect(searchAndLoad(catalog, "issues", DEFAULT_TOOL_SEARCH_LIMIT)).toEqual([]); }); }); -describe("tool_search permissions", () => { +describe("find_tools permissions", () => { it("runs without approval because it touches no files", () => { - const decision = decideStepToolCall("tool_search", { query: "weather" }, stepPermissionStateForPreset("approve-for-me")); + const decision = decideStepToolCall( + "find_tools", + { query: "weather" }, + stepPermissionStateForPreset("approve-for-me"), + ); expect(decision).toMatchObject({ action: "allow", hazardous: false }); }); }); -describe("tool_search session wiring", () => { +describe("exposure and the session tool set", () => { let tempDir: string; let agentDir: string; beforeEach(() => { - tempDir = join(tmpdir(), `pi-tool-search-${Date.now()}-${Math.random().toString(36).slice(2)}`); + tempDir = join(tmpdir(), `pi-tool-exposure-${Date.now()}-${Math.random().toString(36).slice(2)}`); agentDir = join(tempDir, "agent"); mkdirSync(agentDir, { recursive: true }); }); @@ -157,16 +144,16 @@ describe("tool_search session wiring", () => { if (tempDir && existsSync(tempDir)) rmSync(tempDir, { recursive: true, force: true }); }); - async function createSession(defaultTools?: string[], options: { tools?: string[] } = {}) { - const settingsManager = SettingsManager.inMemory(defaultTools ? { defaultTools } : {}); + async function createSession(factories: InlineExtension[] = []) { + const settingsManager = SettingsManager.inMemory({}); const resourceLoader = new DefaultResourceLoader({ cwd: tempDir, agentDir, settingsManager, - extensionFactories: [createStepToolSearchExtension()], + extensionFactories: factories, }); await resourceLoader.reload(); - return ( + const session = ( await createAgentSession({ cwd: tempDir, agentDir, @@ -174,36 +161,64 @@ describe("tool_search session wiring", () => { settingsManager, sessionManager: SessionManager.inMemory(tempDir), resourceLoader, - tools: options.tools, }) ).session; + await session.bindExtensions({}); + return session; } - it("registers tool_search without declaring it", async () => { - const session = await createSession(); - expect(session.getAllTools().map((tool) => tool.name)).toContain("tool_search"); - expect(session.getActiveToolNames()).not.toContain("tool_search"); - expect(session.systemPrompt).not.toContain("- tool_search:"); + function pluginTool(name: string, exposure: "direct" | "deferred" | "hidden") { + return { + name, + label: name, + description: `${name} plugin tool`, + exposure, + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text" as const, text: "ok" }], details: {} }), + }; + } + + it("leaves a hidden tool out of the catalog", async () => { + const session = await createSession([(pi) => void pi.registerTool(pluginTool("secret_tool", "hidden"))]); + expect(session.getAllTools().map((tool) => tool.name)).not.toContain("secret_tool"); + expect(session.getActiveToolNames()).not.toContain("secret_tool"); session.dispose(); }); - it("declares tool_search when the tool settings ask for it", async () => { - const session = await createSession(["read", "bash", "edit", "write", "tool_search"]); - expect(session.getActiveToolNames()).toContain("tool_search"); - expect(session.systemPrompt).toContain("- tool_search:"); + it("keeps a deferred tool registered but undeclared", async () => { + const session = await createSession([(pi) => void pi.registerTool(pluginTool("plugin_tool", "deferred"))]); + expect(session.getAllTools().map((tool) => tool.name)).toContain("plugin_tool"); + expect(session.getActiveToolNames()).not.toContain("plugin_tool"); session.dispose(); }); - it("declares tool_search when the tool selector names it", async () => { - const session = await createSession(undefined, { tools: ["read", "tool_search"] }); - expect(session.getActiveToolNames()).toContain("tool_search"); + it("declares a direct tool like any other", async () => { + const session = await createSession([(pi) => void pi.registerTool(pluginTool("direct_tool", "direct"))]); + expect(session.getActiveToolNames()).toContain("direct_tool"); session.dispose(); }); - it("keeps tool_search active once a session enables it", async () => { - const session = await createSession(); - session.setActiveToolsByName(["read", "bash", "edit", "write", "tool_search"]); - expect(session.getActiveToolNames()).toContain("tool_search"); + it("activates a deferred tool the settings name", async () => { + const settingsManager = SettingsManager.inMemory({ defaultTools: ["read", "plugin_tool"] }); + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir, + settingsManager, + extensionFactories: [(pi) => void pi.registerTool(pluginTool("plugin_tool", "deferred"))], + }); + await resourceLoader.reload(); + const session = ( + await createAgentSession({ + cwd: tempDir, + agentDir, + model: stepModel(), + settingsManager, + sessionManager: SessionManager.inMemory(tempDir), + resourceLoader, + }) + ).session; + await session.bindExtensions({}); + expect(session.getActiveToolNames()).toContain("plugin_tool"); session.dispose(); }); });