From a263d38429575020d9f3e28c96d743ebd913077f Mon Sep 17 00:00:00 2001 From: Ben Wachman Date: Sat, 24 Oct 2026 14:55:14 -0700 Subject: [PATCH 1/4] fix(agent): count each query word once in lexical search --- src/agent/lexical-rank.test.ts | 22 ++++++++++++++++++++++ src/agent/lexical-rank.ts | 5 ++++- src/agent/tool-search.test.ts | 34 ++++++++++++++++++++++++++++++++++ 3 files changed, 60 insertions(+), 1 deletion(-) create mode 100644 src/agent/lexical-rank.test.ts diff --git a/src/agent/lexical-rank.test.ts b/src/agent/lexical-rank.test.ts new file mode 100644 index 000000000..073e74181 --- /dev/null +++ b/src/agent/lexical-rank.test.ts @@ -0,0 +1,22 @@ +import { describe, expect, test } from "bun:test"; + +import { + lexicalFields, + scoreLexical, + tokenizeLexical, +} from "./lexical-rank.js"; + +describe("scoreLexical", () => { + test("counts a word repeated in the query once", () => { + const fields = lexicalFields( + "mcp__linear__list_release_notes", + "List release notes in the workspace", + ); + // Same raw query on both sides so only the token list differs; the raw + // query feeds a separate whole-query substring bonus. + const rawQuery = "release notes release"; + expect( + scoreLexical(fields, tokenizeLexical("release notes release"), rawQuery), + ).toBe(scoreLexical(fields, tokenizeLexical("release notes"), rawQuery)); + }); +}); diff --git a/src/agent/lexical-rank.ts b/src/agent/lexical-rank.ts index 1b8da588a..1a878219a 100644 --- a/src/agent/lexical-rank.ts +++ b/src/agent/lexical-rank.ts @@ -30,7 +30,10 @@ export function scoreLexical( const name = fields.name.toLowerCase(); const blob = fields.blob.toLowerCase(); let total = 0; - for (const token of queryTokens) { + // A model repeats a word because its sentence needs it ("release ID or + // release slug"), not to weight it; counting each copy multiplies the + // bonus for every tool that matches that one word. + for (const token of new Set(queryTokens)) { if (fields.nameTokens.includes(token)) total += 3; else if (fields.blobTokens.has(token)) total += 1; else if (name.includes(token)) total += 0.75; diff --git a/src/agent/tool-search.test.ts b/src/agent/tool-search.test.ts index bf380feb4..6848950da 100644 --- a/src/agent/tool-search.test.ts +++ b/src/agent/tool-search.test.ts @@ -112,6 +112,33 @@ const linearCatalog: ToolDefinition[] = [ const linearIndex = createToolIndex(() => linearCatalog); +// Trimmed from the live Linear MCP catalog: release-named tools competing with +// list_issues, the only one that lists a release's issues. +const linearReleaseCatalog: ToolDefinition[] = [ + mcpDef( + "mcp__linear__list_issues", + '[linear] List issues in the user\'s Linear workspace, including active Triage Intelligence suggestions for issues in triage. For issues assigned to me, use "me" as the assignee. For issues created by me, use "me" as the creator. Use "null" for no assignee.', + ), + mcpDef( + "mcp__linear__get_release", + "[linear] Retrieve details of a release by ID or slug.", + ), + mcpDef( + "mcp__linear__get_release_note", + "[linear] Retrieve release notes by ID or slug, including markdown content.", + ), + mcpDef( + "mcp__linear__list_release_notes", + "[linear] List release notes in the workspace, optionally filtered by pipeline or covered release.", + ), + mcpDef( + "mcp__linear__list_release_pipelines", + "[linear] List release pipelines in the workspace.", + ), +]; + +const linearReleaseIndex = createToolIndex(() => linearReleaseCatalog); + const wideCatalog: ToolDefinition[] = Array.from({ length: 12 }, (_, i) => mcpDef( `mcp__linear__op_${String(i).padStart(2, "0")}`, @@ -311,6 +338,13 @@ describe("createToolIndex", () => { expect(ranked.length).toBeGreaterThanOrEqual(1); expect(ranked.length).toBeLessThanOrEqual(2); }); + + test("a word repeated in the query does not bury the tool that answers it", () => { + const ranked = linearReleaseIndex.search( + "Linear list issues filtered by release ID or release slug", + ); + expect(ranked).toContain("mcp__linear__list_issues"); + }); }); describe("createToolSearchTool skills", () => { From 1ff0696127ff1aa47d19df9262bf68c9c815b1a5 Mon Sep 17 00:00:00 2001 From: Ben Wachman Date: Sat, 24 Oct 2026 15:06:38 -0700 Subject: [PATCH 2/4] refactor(agent): build reconnect stub definitions in one place --- src/agent/tools-mcp-disconnect.test.ts | 67 ++++++++++++++++++++++++++ src/agent/tools.ts | 21 ++++---- 2 files changed, 79 insertions(+), 9 deletions(-) diff --git a/src/agent/tools-mcp-disconnect.test.ts b/src/agent/tools-mcp-disconnect.test.ts index a5529b238..de46ab6da 100644 --- a/src/agent/tools-mcp-disconnect.test.ts +++ b/src/agent/tools-mcp-disconnect.test.ts @@ -380,6 +380,8 @@ const { MCP_RECONNECTING_TOOL_ERROR, isDegradedMcpState } = await import("../mcp/plugin.js"); const { mcpNotice } = await import("../tui/runtime-notices.js"); const { mcpReconnectDelayMs } = await import("./tools.js"); +const { TOOL_SEARCH_PENDING_WAIT_MS, TOOL_SEARCH_RECONNECT_WAIT_MS } = + await import("./tool-search.js"); // connectMCPServer is mocked; tests simulate a dead transport by invoking the // onDisconnect hook the real client wires to transport.onclose. @@ -773,6 +775,71 @@ describe("unintentional disconnect and automatic reconnect", () => { await toolset.dispose(); } }); + + const widgetTool = { + name: "list", + description: "List widgets", + inputSchema: { type: "object", required: ["q"] }, + }; + + test("reconnect stubs keep the server's name, description, and schema", async () => { + mock.connectedTools = [widgetTool]; + const toolset = await makeToolset(); + try { + await toolset.connectMCPServer(acme, callbacks([])); + killTransport(); + const stub = toolset.dynamicRunner + .currentDefinitions() + .find((d) => d.name === "mcp__acme__list"); + expect(stub).toEqual({ + name: "mcp__acme__list", + description: "[acme] List widgets", + inputSchema: widgetTool.inputSchema, + }); + } finally { + await toolset.dispose(); + } + }); + + // Resolves to whether tool_search was still waiting once the tier-1 miss + // wait elapsed, i.e. whether a retained tool matched the query. + async function searchWaitsForRedial(query: string): Promise { + mock.connectedTools = [widgetTool]; + const toolset = await makeToolset(); + jest.useFakeTimers(); + try { + await toolset.connectMCPServer(acme, callbacks([])); + mock.mode = "deferred"; + killTransport(); + await advanceUntil(() => mock.connectOptions.length === 2); + let settled = false; + const search = toolset.dynamicRunner + .run( + { id: "s1", name: "tool_search", arguments: { query } }, + new AbortController().signal, + ) + .then(() => { + settled = true; + }); + await advanceAndFlush(TOOL_SEARCH_PENDING_WAIT_MS); + const waiting = !settled; + await advanceAndFlush(TOOL_SEARCH_RECONNECT_WAIT_MS); + await search; + return waiting; + } finally { + mock.releaseDeferredConnect?.(); + jest.useRealTimers(); + await toolset.dispose(); + } + } + + test("tool_search waits for a redial whose retained tool matches by description", async () => { + expect(await searchWaitsForRedial("widgets")).toBe(true); + }); + + test("tool_search does not wait for a redial whose retained tools do not match", async () => { + expect(await searchWaitsForRedial("zebra")).toBe(false); + }); }); describe("MCP handshake bounds", () => { diff --git a/src/agent/tools.ts b/src/agent/tools.ts index a4a048ddb..dffcecc68 100644 --- a/src/agent/tools.ts +++ b/src/agent/tools.ts @@ -833,11 +833,9 @@ export async function createAgentToolset( if (queryTokens.length === 0) return false; for (const [serverName, state] of reconnectingServers) { for (const tool of state.tools) { + const definition = reconnectStubDefinition(serverName, tool); const score = scoreLexical( - lexicalFields( - mcpToolName(serverName, tool.name), - `[${serverName}] ${tool.description}`, - ), + lexicalFields(definition.name, definition.description), queryTokens, rawQuery, ); @@ -1031,6 +1029,15 @@ export async function createAgentToolset( } const reconnectingServers = new Map(); + const reconnectStubDefinition = ( + serverName: string, + tool: McpReconnectTool, + ): ToolDefinition => ({ + name: mcpToolName(serverName, tool.name), + description: `[${serverName}] ${tool.description}`, + inputSchema: tool.inputSchema, + }); + const mountReconnectingStubs = ( name: string, tools: readonly McpReconnectTool[], @@ -1039,11 +1046,7 @@ export async function createAgentToolset( dynamicRunner.addTools( tools.map((tool): AgentTool => ({ kind: "full", - definition: { - name: mcpToolName(name, tool.name), - description: `[${name}] ${tool.description}`, - inputSchema: tool.inputSchema, - }, + definition: reconnectStubDefinition(name, tool), handler: async (call: ToolCall) => ({ callId: call.id, content: MCP_RECONNECTING_TOOL_ERROR, From 6f2ee7d4379cc3dbefa0ffed52474978cb611db5 Mon Sep 17 00:00:00 2001 From: Ben Wachman Date: Sat, 24 Oct 2026 15:27:49 -0700 Subject: [PATCH 3/4] refactor(agent): validate input schema properties with arktype --- src/agent/tool-search.test.ts | 41 +++++++++++++++++++++++++++++++++++ src/agent/tool-search.ts | 15 +++++++------ 2 files changed, 49 insertions(+), 7 deletions(-) diff --git a/src/agent/tool-search.test.ts b/src/agent/tool-search.test.ts index 6848950da..243b8a8db 100644 --- a/src/agent/tool-search.test.ts +++ b/src/agent/tool-search.test.ts @@ -628,6 +628,47 @@ describe("createToolSearchTool", () => { ); }); + async function schemaCardLine( + inputSchema: Record, + ): Promise { + const def: ToolDefinition = { + name: "mcp__linear__big", + description: "Linear issue with a large schema", + inputSchema, + }; + const tool = createToolSearchTool({ + search: () => [def.name], + lookup: () => def, + }); + const out = await call(tool, { query: "linear" }); + return out + .split("\n") + .find((line) => line.startsWith(" {")) + ?.trim(); + } + + test("an oversized input schema's card lists its property names", async () => { + const line = await schemaCardLine({ + type: "object", + properties: { + alpha: { type: "string", description: "x".repeat(400) }, + beta: { type: "string", description: "y".repeat(400) }, + }, + required: [], + }); + expect(line).toBe('{"properties":["alpha","beta"],"required":[]}'); + }); + + test("an oversized card lists no property names when properties is not an object", async () => { + const line = await schemaCardLine({ + type: "object", + properties: null, + description: "z".repeat(TOOL_SEARCH_SCHEMA_MAX), + required: ["alpha"], + }); + expect(line).toBe('{"properties":[],"required":["alpha"]}'); + }); + test("search does not promote matches and keeps description lines short", async () => { const tool = createToolSearchTool({ search: (q) => linearIndex.search(q), diff --git a/src/agent/tool-search.ts b/src/agent/tool-search.ts index 49c8dc3df..cb74be76e 100644 --- a/src/agent/tool-search.ts +++ b/src/agent/tool-search.ts @@ -438,18 +438,19 @@ function isEmptyInputSchema(schema: unknown): boolean { ); } +const InputSchemaProperties = type({ properties: "Record" }); + +function schemaProperties(schema: unknown): Record { + const parsed = InputSchemaProperties(schema); + return parsed instanceof type.errors ? {} : parsed.properties; +} + function compactInputSchema(schema: unknown): string | undefined { if (isEmptyInputSchema(schema)) return undefined; const json = JSON.stringify(schema); if (json.length <= TOOL_SEARCH_SCHEMA_MAX) return json; const obj = schema as Record; - const properties = obj.properties; - const keys = - properties !== undefined && - typeof properties === "object" && - properties !== null - ? Object.keys(properties) - : []; + const keys = Object.keys(schemaProperties(schema)); const required = Array.isArray(obj.required) ? obj.required : []; const stub = JSON.stringify({ properties: keys, required }); if (stub.length <= TOOL_SEARCH_SCHEMA_MAX) return stub; From eef2bd2798ece24a655f3442116e2ec8fe8cf7e0 Mon Sep 17 00:00:00 2001 From: Ben Wachman Date: Sat, 24 Oct 2026 15:38:13 -0700 Subject: [PATCH 4/4] fix(agent): match tool_search queries against tool parameters --- src/agent/agent-search.ts | 3 +- src/agent/lexical-rank.test.ts | 79 +++++++++++++ src/agent/lexical-rank.ts | 68 ++++++++++-- src/agent/skill-search.ts | 3 +- src/agent/tool-search.test.ts | 148 ++++++++++++++++++++++++- src/agent/tool-search.ts | 36 ++++-- src/agent/tools-mcp-disconnect.test.ts | 12 +- src/agent/tools.ts | 12 +- 8 files changed, 327 insertions(+), 34 deletions(-) diff --git a/src/agent/agent-search.ts b/src/agent/agent-search.ts index 95eceaafe..a4cfe7fff 100644 --- a/src/agent/agent-search.ts +++ b/src/agent/agent-search.ts @@ -5,6 +5,7 @@ import { type } from "arktype"; import { scrubSecretShapedContent } from "../plugins/tool-result-secret-scrub.js"; import { lexicalFields, + NO_LEXICAL_PARAMS, rankAndCut, scoreLexical, tokenizeLexical, @@ -35,7 +36,7 @@ export function createAgentIndex( rawQuery: string, ): number => scoreLexical( - lexicalFields(profile.id, profileSearchText(profile)), + lexicalFields(profile.id, profileSearchText(profile), NO_LEXICAL_PARAMS), queryTokens, rawQuery, ) + diff --git a/src/agent/lexical-rank.test.ts b/src/agent/lexical-rank.test.ts index 073e74181..c6e032893 100644 --- a/src/agent/lexical-rank.test.ts +++ b/src/agent/lexical-rank.test.ts @@ -2,6 +2,7 @@ import { describe, expect, test } from "bun:test"; import { lexicalFields, + NO_LEXICAL_PARAMS, scoreLexical, tokenizeLexical, } from "./lexical-rank.js"; @@ -11,6 +12,7 @@ describe("scoreLexical", () => { const fields = lexicalFields( "mcp__linear__list_release_notes", "List release notes in the workspace", + NO_LEXICAL_PARAMS, ); // Same raw query on both sides so only the token list differs; the raw // query feeds a separate whole-query substring bonus. @@ -19,4 +21,81 @@ describe("scoreLexical", () => { scoreLexical(fields, tokenizeLexical("release notes release"), rawQuery), ).toBe(scoreLexical(fields, tokenizeLexical("release notes"), rawQuery)); }); + + // Two-word raw query that no fixture contains whole, so the whole-query + // substring bonus never applies and only the per-token tiers are compared. + const rawQuery = "release filter"; + const queryTokens = tokenizeLexical(rawQuery); + + test("ranks a name substring above a parameter name above a parameter description", () => { + // Every fixture matches "filter" in its blob, so the parameter tiers have + // a direct match to boost. + const nameSubstring = lexicalFields( + "mcp__acme__releaselog", + "Filter logs", + NO_LEXICAL_PARAMS, + ); + const paramName = lexicalFields("mcp__acme__records", "Filter records", { + names: ["release"], + blob: "", + }); + const paramDescription = lexicalFields( + "mcp__acme__records", + "Filter records", + { + names: ["id"], + blob: "Release ID or slug", + }, + ); + const scores = [nameSubstring, paramName, paramDescription].map((fields) => + scoreLexical(fields, queryTokens, rawQuery), + ); + expect(scores[0]).toBeGreaterThan(scores[1] ?? 0); + expect(scores[1]).toBeGreaterThan(scores[2] ?? 0); + expect(scores[2]).toBeGreaterThan(0); + }); + + test("ranks a parameter name above a description substring of the same word", () => { + // "release" sits only inside "prereleases", so the description offers + // just a raw-substring hit. + const without = lexicalFields( + "mcp__acme__records", + "Read prereleases", + NO_LEXICAL_PARAMS, + ); + const withParam = lexicalFields("mcp__acme__records", "Read prereleases", { + names: ["release"], + blob: "", + }); + expect(scoreLexical(withParam, queryTokens, rawQuery)).toBeGreaterThan( + scoreLexical(without, queryTokens, rawQuery), + ); + }); + + test("a parameter match alone scores nothing", () => { + const fields = lexicalFields("mcp__acme__records", "Read records", { + names: ["release"], + blob: "Release ID or slug", + }); + expect(scoreLexical(fields, queryTokens, rawQuery)).toBe(0); + }); + + test("a parameter match adds nothing to a word already matched by name or description", () => { + const params = { names: ["release"], blob: "Release ID or slug" }; + for (const [name, blob] of [ + ["mcp__acme__release", "Read logs"], + ["mcp__acme__records", "Read a release"], + ["mcp__acme__releaselog", "Read logs"], + ] as const) { + expect( + scoreLexical(lexicalFields(name, blob, params), queryTokens, rawQuery), + ).toBe( + scoreLexical( + lexicalFields(name, blob, NO_LEXICAL_PARAMS), + queryTokens, + rawQuery, + ), + ); + } + }); }); diff --git a/src/agent/lexical-rank.ts b/src/agent/lexical-rank.ts index 1a878219a..4cf45dbd8 100644 --- a/src/agent/lexical-rank.ts +++ b/src/agent/lexical-rank.ts @@ -1,7 +1,12 @@ // One shared dependency-free lexical ranker behind tool_search, // search_agents, and skill_search. Exact name-token hits weigh most, then -// blob-token hits, then raw-substring matches (so "linear" finds -// mcp__linear__* even though it is not a whole token there). +// blob-token hits, then a raw name substring (so "linear" finds +// mcp__linear__* even though it is not a whole token there), then a raw blob +// substring. Parameter name and description tokens (tool_search only) weigh +// less than a name substring and only boost an item that already matches by +// name or blob: a filter makes the tool that carries it rank higher, but a +// query that matches nothing except parameter text stays a miss, so +// tool_search still waits for servers that are connecting. export function tokenizeLexical(text: string): string[] { return text.toLowerCase().match(/[a-z0-9]+/g) ?? []; } @@ -11,17 +16,61 @@ export interface LexicalFields { readonly blob: string; readonly nameTokens: readonly string[]; readonly blobTokens: ReadonlySet; + readonly paramNameTokens: ReadonlySet; + readonly paramBlobTokens: ReadonlySet; } -export function lexicalFields(name: string, blob: string): LexicalFields { +interface LexicalParams { + readonly names: readonly string[]; + readonly blob: string; +} + +export const NO_LEXICAL_PARAMS: LexicalParams = { names: [], blob: "" }; + +// JSON Schema keys are usually camelCase, and an undescribed flag such as +// `includeArchived` carries its word only inside the key. +function splitCamelCase(key: string): string { + return key + .replace(/([a-z0-9])([A-Z])/g, "$1 $2") + .replace(/([A-Z])([A-Z][a-z])/g, "$1 $2"); +} + +export function lexicalFields( + name: string, + blob: string, + params: LexicalParams, +): LexicalFields { return { name, blob, nameTokens: tokenizeLexical(name), blobTokens: new Set(tokenizeLexical(blob)), + paramNameTokens: new Set( + tokenizeLexical(params.names.map(splitCamelCase).join(" ")), + ), + paramBlobTokens: new Set(tokenizeLexical(params.blob)), }; } +function directTokenScore( + fields: LexicalFields, + name: string, + blob: string, + token: string, +): number { + if (fields.nameTokens.includes(token)) return 3; + if (fields.blobTokens.has(token)) return 1; + if (name.includes(token)) return 0.75; + if (blob.includes(token)) return 0.25; + return 0; +} + +function paramTokenScore(fields: LexicalFields, token: string): number { + if (fields.paramNameTokens.has(token)) return 0.5; + if (fields.paramBlobTokens.has(token)) return 0.25; + return 0; +} + export function scoreLexical( fields: LexicalFields, queryTokens: readonly string[], @@ -29,18 +78,17 @@ export function scoreLexical( ): number { const name = fields.name.toLowerCase(); const blob = fields.blob.toLowerCase(); - let total = 0; + let matched = name.includes(rawQuery) ? 1 : 0; + let total = matched; // A model repeats a word because its sentence needs it ("release ID or // release slug"), not to weight it; counting each copy multiplies the // bonus for every tool that matches that one word. for (const token of new Set(queryTokens)) { - if (fields.nameTokens.includes(token)) total += 3; - else if (fields.blobTokens.has(token)) total += 1; - else if (name.includes(token)) total += 0.75; - else if (blob.includes(token)) total += 0.25; + const direct = directTokenScore(fields, name, blob, token); + matched += direct; + total += Math.max(direct, paramTokenScore(fields, token)); } - if (name.includes(rawQuery)) total += 1; - return total; + return matched > 0 ? total : 0; } // One shared rank→filter→sort→slice→map cut behind tool_search, diff --git a/src/agent/skill-search.ts b/src/agent/skill-search.ts index 22a50d069..76e59f820 100644 --- a/src/agent/skill-search.ts +++ b/src/agent/skill-search.ts @@ -6,6 +6,7 @@ import { type } from "arktype"; import type { SkillSummary } from "../extensions/skills.js"; import { lexicalFields, + NO_LEXICAL_PARAMS, rankAndCut, scoreLexical, tokenizeLexical, @@ -67,7 +68,7 @@ function scoreSkill( rawQuery: string, ): number { return scoreLexical( - lexicalFields(skill.name, skill.description), + lexicalFields(skill.name, skill.description, NO_LEXICAL_PARAMS), queryTokens, rawQuery, ); diff --git a/src/agent/tool-search.test.ts b/src/agent/tool-search.test.ts index 243b8a8db..b4bf9a0bb 100644 --- a/src/agent/tool-search.test.ts +++ b/src/agent/tool-search.test.ts @@ -118,22 +118,63 @@ const linearReleaseCatalog: ToolDefinition[] = [ mcpDef( "mcp__linear__list_issues", '[linear] List issues in the user\'s Linear workspace, including active Triage Intelligence suggestions for issues in triage. For issues assigned to me, use "me" as the assignee. For issues created by me, use "me" as the creator. Use "null" for no assignee.', + { + assignee: { + type: ["string", "null"], + description: 'User ID, name, email, or "me"', + }, + cycle: { type: "string", description: "Cycle name, number, or ID" }, + label: { type: "string", description: "Label name or ID" }, + project: { + type: "string", + description: "Project name, ID, identifier (e.g., P-ENG-123), or slug", + }, + release: { type: "string", description: "Release ID or slug" }, + state: { type: "string", description: "State type, name, or ID" }, + team: { type: "string", description: "Team name or ID" }, + }, ), mcpDef( "mcp__linear__get_release", "[linear] Retrieve details of a release by ID or slug.", + { + id: { type: "string", description: "Release ID or slug" }, + includeReleaseNotes: { + type: "boolean", + description: "Include associated release notes", + }, + }, ), mcpDef( "mcp__linear__get_release_note", "[linear] Retrieve release notes by ID or slug, including markdown content.", + { + id: { type: "string", description: "Release notes ID or slug" }, + includeReleases: { + type: "boolean", + description: "Include associated releases", + }, + }, ), mcpDef( "mcp__linear__list_release_notes", "[linear] List release notes in the workspace, optionally filtered by pipeline or covered release.", + { + pipeline: { + type: "string", + description: "Release pipeline ID, slug, or exact name", + }, + query: { type: "string", description: "Search release notes title" }, + release: { type: "string", description: "Release ID or slug" }, + }, ), mcpDef( "mcp__linear__list_release_pipelines", "[linear] List release pipelines in the workspace.", + { + query: { type: "string", description: "Search pipeline name" }, + team: { type: "string", description: "Team name or ID" }, + }, ), ]; @@ -340,11 +381,116 @@ describe("createToolIndex", () => { }); test("a word repeated in the query does not bury the tool that answers it", () => { - const ranked = linearReleaseIndex.search( + // Parameters stripped so only counting the repeated word once keeps + // list_issues above the cut. + const namesOnlyIndex = createToolIndex(() => + linearReleaseCatalog.map((def) => ({ ...def, inputSchema: {} })), + ); + const ranked = namesOnlyIndex.search( "Linear list issues filtered by release ID or release slug", ); expect(ranked).toContain("mcp__linear__list_issues"); }); + + test("a release filter lifts list_issues above the tools named for releases", () => { + expect(linearReleaseIndex.search("linear issues release filter")[0]).toBe( + "mcp__linear__list_issues", + ); + }); + + test("a release filter narrows a release query to the tool that takes it", () => { + expect(linearReleaseIndex.search("issues for release")).toEqual([ + "mcp__linear__list_issues", + ]); + }); + + // Identical names and descriptions apart from the filter they take, with + // the cycle tool first so catalog order alone would rank it ahead. + const recordsCatalog: ToolDefinition[] = [ + mcpDef("mcp__acme__list_records_v1", "List records", { + cycle: { type: "string", description: "Cycle name, number, or ID" }, + }), + mcpDef("mcp__acme__list_records_v2", "List records", { + release: { type: "string", description: "Release ID or slug" }, + }), + ]; + const recordsIndex = createToolIndex(() => recordsCatalog); + + test("a parameter name lifts the tool that takes it", () => { + expect(recordsIndex.search("list records by release")[0]).toBe( + "mcp__acme__list_records_v2", + ); + }); + + test("a word inside a camelCase parameter name lifts the tool that takes it", () => { + // Undescribed boolean flags, a common JSON Schema shape, so only the key + // carries the word; the closed flag comes first so catalog order alone + // would rank it ahead. + const flagsIndex = createToolIndex(() => [ + mcpDef("mcp__acme__list_records_v1", "List records", { + includeClosed: { type: "boolean" }, + }), + mcpDef("mcp__acme__list_records_v2", "List records", { + includeArchived: { type: "boolean" }, + }), + ]); + expect(flagsIndex.search("list records archived")[0]).toBe( + "mcp__acme__list_records_v2", + ); + }); + + test("a word in a parameter description lifts the tool that takes it", () => { + expect(recordsIndex.search("list records by slug")[0]).toBe( + "mcp__acme__list_records_v2", + ); + }); + + test("a query matching only parameter text is a miss", () => { + expect(recordsIndex.search("release")).toEqual([]); + }); + + test("a tool named for the query keeps a parameter-only match off the shortlist", () => { + const ranked = linearReleaseIndex.search("linear release"); + expect(ranked).toContain("mcp__linear__get_release"); + expect(ranked).not.toContain("mcp__linear__list_issues"); + }); + + test("a parameter shared across tools does not displace the tool named for it", () => { + const teamIndex = createToolIndex(() => [ + ...linearReleaseCatalog, + mcpDef( + "mcp__linear__list_teams", + "[linear] List teams in the user's Linear workspace, including their parent IDs and hierarchy names", + { query: { type: "string", description: "Search query" } }, + ), + ]); + expect(teamIndex.search("linear team")[0]).toBe("mcp__linear__list_teams"); + }); + + test("a malformed input schema still scores on name and description", () => { + const malformedIndex = createToolIndex(() => [ + { + name: "mcp__acme__null_props", + description: "Archive widgets", + inputSchema: { type: "object", properties: null }, + }, + { + name: "mcp__acme__numeric_description", + description: "Archive gadgets", + inputSchema: { + type: "object", + properties: { release: { type: "string", description: 42 } }, + }, + }, + ]); + expect(malformedIndex.search("archive")).toEqual([ + "mcp__acme__null_props", + "mcp__acme__numeric_description", + ]); + expect(malformedIndex.search("archive release")[0]).toBe( + "mcp__acme__numeric_description", + ); + }); }); describe("createToolSearchTool skills", () => { diff --git a/src/agent/tool-search.ts b/src/agent/tool-search.ts index cb74be76e..1285c0195 100644 --- a/src/agent/tool-search.ts +++ b/src/agent/tool-search.ts @@ -305,6 +305,22 @@ export const TOOL_SEARCH_DESC_MAX = 160; export const TOOL_SEARCH_SCHEMA_MAX = 600; export const TOOL_SEARCH_SCHEMA_CARDS = 5; +export function scoreToolDefinition( + def: ToolDefinition, + queryTokens: readonly string[], + rawQuery: string, +): number { + const properties = schemaProperties(def.inputSchema); + return scoreLexical( + lexicalFields(def.name, def.description, { + names: Object.keys(properties), + blob: Object.values(properties).map(propertyDescription).join(" "), + }), + queryTokens, + rawQuery, + ); +} + export function createToolIndex( getDefs: () => readonly ToolDefinition[], advertisedNames: readonly string[] = ADVERTISED_TOOL_NAMES, @@ -312,17 +328,6 @@ export function createToolIndex( // tools. allow?: readonly string[] | undefined, ): ToolIndex { - const score = ( - def: ToolDefinition, - queryTokens: string[], - rawQuery: string, - ): number => - scoreLexical( - lexicalFields(def.name, def.description ?? ""), - queryTokens, - rawQuery, - ); - return { search(query: string, limit = TOOL_SEARCH_MAX_RESULTS): string[] { const rawQuery = query.toLowerCase().trim(); @@ -342,7 +347,7 @@ export function createToolIndex( ); return rankAndCut( candidates, - (def) => score(def, queryTokens, rawQuery), + (def) => scoreToolDefinition(def, queryTokens, rawQuery), limit, TOOL_SEARCH_SCORE_RATIO, ).map((def) => def.name); @@ -445,6 +450,13 @@ function schemaProperties(schema: unknown): Record { return parsed instanceof type.errors ? {} : parsed.properties; } +const PropertyDescription = type({ description: "string" }); + +function propertyDescription(property: unknown): string { + const parsed = PropertyDescription(property); + return parsed instanceof type.errors ? "" : parsed.description; +} + function compactInputSchema(schema: unknown): string | undefined { if (isEmptyInputSchema(schema)) return undefined; const json = JSON.stringify(schema); diff --git a/src/agent/tools-mcp-disconnect.test.ts b/src/agent/tools-mcp-disconnect.test.ts index de46ab6da..d5c327fa3 100644 --- a/src/agent/tools-mcp-disconnect.test.ts +++ b/src/agent/tools-mcp-disconnect.test.ts @@ -779,7 +779,13 @@ describe("unintentional disconnect and automatic reconnect", () => { const widgetTool = { name: "list", description: "List widgets", - inputSchema: { type: "object", required: ["q"] }, + inputSchema: { + type: "object", + properties: { + region: { type: "string", description: "Warehouse region" }, + }, + required: ["region"], + }, }; test("reconnect stubs keep the server's name, description, and schema", async () => { @@ -837,6 +843,10 @@ describe("unintentional disconnect and automatic reconnect", () => { expect(await searchWaitsForRedial("widgets")).toBe(true); }); + test("tool_search does not wait for a redial whose retained tool matches only by parameter", async () => { + expect(await searchWaitsForRedial("region")).toBe(false); + }); + test("tool_search does not wait for a redial whose retained tools do not match", async () => { expect(await searchWaitsForRedial("zebra")).toBe(false); }); diff --git a/src/agent/tools.ts b/src/agent/tools.ts index dffcecc68..7c8c27b2b 100644 --- a/src/agent/tools.ts +++ b/src/agent/tools.ts @@ -112,14 +112,11 @@ import { searchSkillCatalog } from "./skill-search.js"; import { createToolIndex, createToolSearchTool, + scoreToolDefinition, TOOL_SEARCH_PENDING_WAIT_MS, toolSearchDefinition, } from "./tool-search.js"; -import { - lexicalFields, - scoreLexical, - tokenizeLexical, -} from "./lexical-rank.js"; +import { tokenizeLexical } from "./lexical-rank.js"; import { createSearchAgentsTool } from "./agent-search.js"; import { createReadAgentTraceTool } from "../subagent/trace-tool.js"; import { errorMessage } from "./error-message.js"; @@ -833,9 +830,8 @@ export async function createAgentToolset( if (queryTokens.length === 0) return false; for (const [serverName, state] of reconnectingServers) { for (const tool of state.tools) { - const definition = reconnectStubDefinition(serverName, tool); - const score = scoreLexical( - lexicalFields(definition.name, definition.description), + const score = scoreToolDefinition( + reconnectStubDefinition(serverName, tool), queryTokens, rawQuery, );