diff --git a/packages/coding-agent/src/extensions/tool-search/tool.ts b/packages/coding-agent/src/extensions/tool-search/tool.ts index 1841ebd97..b9fa4ade0 100644 --- a/packages/coding-agent/src/extensions/tool-search/tool.ts +++ b/packages/coding-agent/src/extensions/tool-search/tool.ts @@ -115,7 +115,28 @@ export function createToolSearchDocument( return { name: tool.name, text: parts.filter((part) => part.trim()).join(" ") }; } -/** Okapi BM25 with the usual parameters. Ties keep document order. */ +/** 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 `query` names the tool: its full name or a suffix after a `__` separator, so `thread_read` + * and `app__thread_read` both name `mcp__app__thread_read`. + */ +function isNameQuery(query: string, name: string): boolean { + const normalized = normalizeName(name); + return normalized === query || normalized.endsWith(`__${query}`); +} + +/** + * Okapi BM25 with the usual parameters. A query that names a tool ranks that tool first, because + * length normalization can otherwise rank a short description sharing the name's words above it. + * Ties keep document order. + */ export class Bm25Ranker implements ToolRanker { private readonly k1: number; private readonly b: number; @@ -141,7 +162,8 @@ export class Bm25Ranker implements ToolRanker { return [term, Math.log(1 + (documents.length - frequency + 0.5) / (frequency + 0.5))] as const; }), ); - const matches: ToolSearchMatch[] = []; + const nameQuery = normalizeName(query); + const matches: (ToolSearchMatch & { named: boolean })[] = []; documents.forEach((document, index) => { let score = 0; for (const term of queryTerms) { @@ -150,9 +172,12 @@ export class Bm25Ranker implements ToolRanker { const norm = this.k1 * (1 - this.b + (this.b * lengths[index]) / averageLength); score += (idf.get(term) ?? 0) * ((count * (this.k1 + 1)) / (count + norm)); } - if (score > 0) matches.push({ name: document.name, score }); + if (score > 0) matches.push({ name: document.name, score, named: isNameQuery(nameQuery, document.name) }); }); - return matches.sort((a, b) => b.score - a.score).slice(0, limit); + return matches + .sort((a, b) => Number(b.named) - Number(a.named) || b.score - a.score) + .slice(0, limit) + .map(({ name, score }) => ({ name, score })); } } diff --git a/packages/coding-agent/test/tool-search.test.ts b/packages/coding-agent/test/tool-search.test.ts index ba07448a0..fd55b5c56 100644 --- a/packages/coding-agent/test/tool-search.test.ts +++ b/packages/coding-agent/test/tool-search.test.ts @@ -53,6 +53,29 @@ describe("Bm25Ranker", () => { expect(ranker.rank("tickets", documents, 8)).toEqual([]); }); + // https://github.com/earendil-works/pi/issues/11331 + it("ranks a tool first when the query is its name", () => { + const named = [ + tool("mcp__app__queue_read", "Read a queued message of this thread."), + tool( + "mcp__dev-app__thread_read", + "Read durable state and a paginated timeline from any thread. The default view returns user messages, assistant messages and plans. Continue with afterPosition until hasMore is false.", + ), + ].map((entry) => createToolSearchDocument(entry)); + const ranker = new Bm25Ranker(); + for (const query of [ + "thread_read", + "dev_app__thread_read", + "mcp__dev-app__thread_read", + "Thread_Read", + "thread read", + ]) { + expect(ranker.rank(query, named, 1).map((match) => match.name)).toEqual(["mcp__dev-app__thread_read"]); + } + // Partial names still rank by BM25. + expect(ranker.rank("read", named, 1).map((match) => match.name)).toEqual(["mcp__app__queue_read"]); + }); + it("includes the namespace in the search text", () => { const document = createToolSearchDocument(tool("mcp__x__run", "Run it."), { name: "mcp__x",