diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index f0242d21..a9a83e43 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -1405,6 +1405,16 @@ 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 `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. **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/docs/mcp-catalog.md b/packages/coding-agent/docs/mcp-catalog.md index fd45bcee..687b2557 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 0d9a6d49..d5313cda 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,6 +4077,12 @@ export class AgentSession { const excludedToolNames = this._excludedToolNames; const isAllowedTool = (name: string): boolean => (!allowedToolNames || allowedToolNames.has(name)) && !excludedToolNames?.has(name); + // 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 = [ @@ -4142,6 +4149,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 +4158,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/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 8febc608..3360e79a 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 07d8d949..28db0b5e 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(). */ @@ -579,6 +590,15 @@ export interface ToolDefinition Static; @@ -1819,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 new file mode 100644 index 00000000..fc84d4ad --- /dev/null +++ b/packages/coding-agent/src/core/tools/tool-search.ts @@ -0,0 +1,231 @@ +/** + * Tool discovery: a BM25 ranker over tool metadata for the model-facing tool finder. + * + * 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 { ToolInfo } from "../extensions/types.ts"; + +export const DEFAULT_TOOL_SEARCH_LIMIT = 8; + +/** 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[]; + 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 })); + } +} + +/** + * 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 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 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/step/config-toml.ts b/packages/coding-agent/src/step/config-toml.ts index e4737368..6a675b7a 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 7d1919d6..8c82de7d 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 fb321ed5..59784e33 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 e6680b2d..c4ca3d63 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 edfaf6f3..86056659 100644 --- a/packages/coding-agent/src/step/permissions.ts +++ b/packages/coding-agent/src/step/permissions.ts @@ -81,6 +81,7 @@ 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", // Keep native names safe when a caller explicitly enables a native tool. "ls", diff --git a/packages/coding-agent/src/step/tool-profile.ts b/packages/coding-agent/src/step/tool-profile.ts index 22d8f5ec..c0da4a83 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 298b7ae0..b61b2fb0 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 new file mode 100644 index 00000000..73e076e8 --- /dev/null +++ b/packages/coding-agent/test/tool-search-integration.test.ts @@ -0,0 +1,159 @@ +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 { 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 { createStepToolProfile } from "../src/step/tool-profile.ts"; + +/** A plugin tool the model never sees until `find_tools` loads it. */ +function deferredTool(name: string, description: string) { + return { + name, + label: name, + description, + exposure: "deferred" as const, + parameters: Type.Object({ query: Type.String() }), + async execute(_toolCallId: string) { + return { content: [{ type: "text" as const, text: `${name} ran` }], details: {} }; + }, + }; +} + +describe("find_tools over a real model turn", () => { + let tempDir: string; + let faux: ReturnType; + const cleanups: Array<() => unknown> = []; + + beforeEach(() => { + tempDir = join(tmpdir(), `pi-find-tools-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: InlineExtension[]) { + 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({}); + 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("declares the loaded tool on the next request", async () => { + const declared: string[][] = []; + const toolResultTexts: string[] = []; + const session = await createSession([ + (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) => (block.type === "text" ? block.text : "")).join("")); + }); + }, + ]); + faux.setResponses([ + (context) => { + declared.push((context.tools ?? []).map((tool: { name: string }) => tool.name).sort()); + return fauxAssistantMessage(fauxToolCall("find_tools", { query: "open issues" }), { + stopReason: "toolUse", + }); + }, + (context) => { + declared.push((context.tools ?? []).map((tool: { name: string }) => tool.name).sort()); + return fauxAssistantMessage("done"); + }, + ]); + + await session.prompt("list the open issues"); + + // First request: the deferred tools are registered but not declared. + 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(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")), + ]); + faux.setResponses([ + () => fauxAssistantMessage(fauxToolCall("find_tools", { query: "open 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-mcp.test.ts b/packages/coding-agent/test/tool-search-mcp.test.ts new file mode 100644 index 00000000..3b97937e --- /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 new file mode 100644 index 00000000..9097580e --- /dev/null +++ b/packages/coding-agent/test/tool-search.test.ts @@ -0,0 +1,224 @@ +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 { 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"; +import { SettingsManager } from "../src/core/settings-manager.ts"; +import { createSyntheticSourceInfo } from "../src/core/source-info.ts"; +import { + Bm25Ranker, + createToolSearchDocument, + DEFAULT_TOOL_SEARCH_LIMIT, + searchAndLoad, + type ToolSearchCatalog, + tokenize, +} from "../src/core/tools/tool-search.ts"; +import { decideStepToolCall, stepPermissionStateForPreset } from "../src/step/permissions.ts"; +import { stepModel } from "./utilities.ts"; + +function toolInfo( + name: string, + description: string, + exposure: ToolInfo["exposure"] = "deferred", + parameters = Type.Object({}), +): ToolInfo { + return { + name, + description, + parameters, + exposure, + 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", "direct"); + + it("activates the matches and keeps the active tools", () => { + 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"); + 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([]); + }); + + 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("find_tools permissions", () => { + it("runs without approval because it touches no files", () => { + const decision = decideStepToolCall( + "find_tools", + { query: "weather" }, + stepPermissionStateForPreset("approve-for-me"), + ); + expect(decision).toMatchObject({ action: "allow", hazardous: false }); + }); +}); + +describe("exposure and the session tool set", () => { + let tempDir: string; + let agentDir: string; + + beforeEach(() => { + tempDir = join(tmpdir(), `pi-tool-exposure-${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(factories: InlineExtension[] = []) { + const settingsManager = SettingsManager.inMemory({}); + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir, + settingsManager, + extensionFactories: factories, + }); + await resourceLoader.reload(); + const session = ( + await createAgentSession({ + cwd: tempDir, + agentDir, + model: stepModel(), + settingsManager, + sessionManager: SessionManager.inMemory(tempDir), + resourceLoader, + }) + ).session; + await session.bindExtensions({}); + return session; + } + + 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("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 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("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(); + }); +});