Skip to content

Commit 1e1ff30

Browse files
committed
fix(agent): match tool_search queries against tool parameters
1 parent 6f2ee7d commit 1e1ff30

8 files changed

Lines changed: 300 additions & 34 deletions

‎src/agent/agent-search.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { type } from "arktype";
55
import { scrubSecretShapedContent } from "../plugins/tool-result-secret-scrub.js";
66
import {
77
lexicalFields,
8+
NO_LEXICAL_PARAMS,
89
rankAndCut,
910
scoreLexical,
1011
tokenizeLexical,
@@ -35,7 +36,7 @@ export function createAgentIndex(
3536
rawQuery: string,
3637
): number =>
3738
scoreLexical(
38-
lexicalFields(profile.id, profileSearchText(profile)),
39+
lexicalFields(profile.id, profileSearchText(profile), NO_LEXICAL_PARAMS),
3940
queryTokens,
4041
rawQuery,
4142
) +

‎src/agent/lexical-rank.test.ts‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { describe, expect, test } from "bun:test";
22

33
import {
44
lexicalFields,
5+
NO_LEXICAL_PARAMS,
56
scoreLexical,
67
tokenizeLexical,
78
} from "./lexical-rank.js";
@@ -11,6 +12,7 @@ describe("scoreLexical", () => {
1112
const fields = lexicalFields(
1213
"mcp__linear__list_release_notes",
1314
"List release notes in the workspace",
15+
NO_LEXICAL_PARAMS,
1416
);
1517
// Same raw query on both sides so only the token list differs; the raw
1618
// query feeds a separate whole-query substring bonus.
@@ -19,4 +21,81 @@ describe("scoreLexical", () => {
1921
scoreLexical(fields, tokenizeLexical("release notes release"), rawQuery),
2022
).toBe(scoreLexical(fields, tokenizeLexical("release notes"), rawQuery));
2123
});
24+
25+
// Two-word raw query that no fixture contains whole, so the whole-query
26+
// substring bonus never applies and only the per-token tiers are compared.
27+
const rawQuery = "release filter";
28+
const queryTokens = tokenizeLexical(rawQuery);
29+
30+
test("ranks a name substring above a parameter name above a parameter description", () => {
31+
// Every fixture matches "filter" in its blob, so the parameter tiers have
32+
// a direct match to boost.
33+
const nameSubstring = lexicalFields(
34+
"mcp__acme__releaselog",
35+
"Filter logs",
36+
NO_LEXICAL_PARAMS,
37+
);
38+
const paramName = lexicalFields("mcp__acme__records", "Filter records", {
39+
names: ["release"],
40+
blob: "",
41+
});
42+
const paramDescription = lexicalFields(
43+
"mcp__acme__records",
44+
"Filter records",
45+
{
46+
names: ["id"],
47+
blob: "Release ID or slug",
48+
},
49+
);
50+
const scores = [nameSubstring, paramName, paramDescription].map((fields) =>
51+
scoreLexical(fields, queryTokens, rawQuery),
52+
);
53+
expect(scores[0]).toBeGreaterThan(scores[1] ?? 0);
54+
expect(scores[1]).toBeGreaterThan(scores[2] ?? 0);
55+
expect(scores[2]).toBeGreaterThan(0);
56+
});
57+
58+
test("ranks a parameter name above a description substring of the same word", () => {
59+
// "release" sits only inside "prereleases", so the description offers
60+
// just a raw-substring hit.
61+
const without = lexicalFields(
62+
"mcp__acme__records",
63+
"Read prereleases",
64+
NO_LEXICAL_PARAMS,
65+
);
66+
const withParam = lexicalFields("mcp__acme__records", "Read prereleases", {
67+
names: ["release"],
68+
blob: "",
69+
});
70+
expect(scoreLexical(withParam, queryTokens, rawQuery)).toBeGreaterThan(
71+
scoreLexical(without, queryTokens, rawQuery),
72+
);
73+
});
74+
75+
test("a parameter match alone scores nothing", () => {
76+
const fields = lexicalFields("mcp__acme__records", "Read records", {
77+
names: ["release"],
78+
blob: "Release ID or slug",
79+
});
80+
expect(scoreLexical(fields, queryTokens, rawQuery)).toBe(0);
81+
});
82+
83+
test("a parameter match adds nothing to a word already matched by name or description", () => {
84+
const params = { names: ["release"], blob: "Release ID or slug" };
85+
for (const [name, blob] of [
86+
["mcp__acme__release", "Read logs"],
87+
["mcp__acme__records", "Read a release"],
88+
["mcp__acme__releaselog", "Read logs"],
89+
] as const) {
90+
expect(
91+
scoreLexical(lexicalFields(name, blob, params), queryTokens, rawQuery),
92+
).toBe(
93+
scoreLexical(
94+
lexicalFields(name, blob, NO_LEXICAL_PARAMS),
95+
queryTokens,
96+
rawQuery,
97+
),
98+
);
99+
}
100+
});
22101
});

‎src/agent/lexical-rank.ts‎

Lines changed: 48 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,12 @@
11
// One shared dependency-free lexical ranker behind tool_search,
22
// search_agents, and skill_search. Exact name-token hits weigh most, then
3-
// blob-token hits, then raw-substring matches (so "linear" finds
4-
// mcp__linear__* even though it is not a whole token there).
3+
// blob-token hits, then a raw name substring (so "linear" finds
4+
// mcp__linear__* even though it is not a whole token there), then a raw blob
5+
// substring. Parameter name and description tokens (tool_search only) weigh
6+
// less than a name substring and only boost an item that already matches by
7+
// name or blob: a filter makes the tool that carries it rank higher, but a
8+
// query that matches nothing except parameter text stays a miss, so
9+
// tool_search still waits for servers that are connecting.
510
export function tokenizeLexical(text: string): string[] {
611
return text.toLowerCase().match(/[a-z0-9]+/g) ?? [];
712
}
@@ -11,36 +16,69 @@ export interface LexicalFields {
1116
readonly blob: string;
1217
readonly nameTokens: readonly string[];
1318
readonly blobTokens: ReadonlySet<string>;
19+
readonly paramNameTokens: ReadonlySet<string>;
20+
readonly paramBlobTokens: ReadonlySet<string>;
1421
}
1522

16-
export function lexicalFields(name: string, blob: string): LexicalFields {
23+
interface LexicalParams {
24+
readonly names: readonly string[];
25+
readonly blob: string;
26+
}
27+
28+
export const NO_LEXICAL_PARAMS: LexicalParams = { names: [], blob: "" };
29+
30+
export function lexicalFields(
31+
name: string,
32+
blob: string,
33+
params: LexicalParams,
34+
): LexicalFields {
1735
return {
1836
name,
1937
blob,
2038
nameTokens: tokenizeLexical(name),
2139
blobTokens: new Set(tokenizeLexical(blob)),
40+
paramNameTokens: new Set(tokenizeLexical(params.names.join(" "))),
41+
paramBlobTokens: new Set(tokenizeLexical(params.blob)),
2242
};
2343
}
2444

45+
function directTokenScore(
46+
fields: LexicalFields,
47+
name: string,
48+
blob: string,
49+
token: string,
50+
): number {
51+
if (fields.nameTokens.includes(token)) return 3;
52+
if (fields.blobTokens.has(token)) return 1;
53+
if (name.includes(token)) return 0.75;
54+
if (blob.includes(token)) return 0.25;
55+
return 0;
56+
}
57+
58+
function paramTokenScore(fields: LexicalFields, token: string): number {
59+
if (fields.paramNameTokens.has(token)) return 0.5;
60+
if (fields.paramBlobTokens.has(token)) return 0.25;
61+
return 0;
62+
}
63+
2564
export function scoreLexical(
2665
fields: LexicalFields,
2766
queryTokens: readonly string[],
2867
rawQuery: string,
2968
): number {
3069
const name = fields.name.toLowerCase();
3170
const blob = fields.blob.toLowerCase();
32-
let total = 0;
71+
let matched = name.includes(rawQuery) ? 1 : 0;
72+
let total = matched;
3373
// A model repeats a word because its sentence needs it ("release ID or
3474
// release slug"), not to weight it; counting each copy multiplies the
3575
// bonus for every tool that matches that one word.
3676
for (const token of new Set(queryTokens)) {
37-
if (fields.nameTokens.includes(token)) total += 3;
38-
else if (fields.blobTokens.has(token)) total += 1;
39-
else if (name.includes(token)) total += 0.75;
40-
else if (blob.includes(token)) total += 0.25;
77+
const direct = directTokenScore(fields, name, blob, token);
78+
matched += direct;
79+
total += Math.max(direct, paramTokenScore(fields, token));
4180
}
42-
if (name.includes(rawQuery)) total += 1;
43-
return total;
81+
return matched > 0 ? total : 0;
4482
}
4583

4684
// One shared rank→filter→sort→slice→map cut behind tool_search,

‎src/agent/skill-search.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { type } from "arktype";
66
import type { SkillSummary } from "../extensions/skills.js";
77
import {
88
lexicalFields,
9+
NO_LEXICAL_PARAMS,
910
rankAndCut,
1011
scoreLexical,
1112
tokenizeLexical,
@@ -67,7 +68,7 @@ function scoreSkill(
6768
rawQuery: string,
6869
): number {
6970
return scoreLexical(
70-
lexicalFields(skill.name, skill.description),
71+
lexicalFields(skill.name, skill.description, NO_LEXICAL_PARAMS),
7172
queryTokens,
7273
rawQuery,
7374
);

‎src/agent/tool-search.test.ts‎

Lines changed: 130 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,22 +118,63 @@ const linearReleaseCatalog: ToolDefinition[] = [
118118
mcpDef(
119119
"mcp__linear__list_issues",
120120
'[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.',
121+
{
122+
assignee: {
123+
type: ["string", "null"],
124+
description: 'User ID, name, email, or "me"',
125+
},
126+
cycle: { type: "string", description: "Cycle name, number, or ID" },
127+
label: { type: "string", description: "Label name or ID" },
128+
project: {
129+
type: "string",
130+
description: "Project name, ID, identifier (e.g., P-ENG-123), or slug",
131+
},
132+
release: { type: "string", description: "Release ID or slug" },
133+
state: { type: "string", description: "State type, name, or ID" },
134+
team: { type: "string", description: "Team name or ID" },
135+
},
121136
),
122137
mcpDef(
123138
"mcp__linear__get_release",
124139
"[linear] Retrieve details of a release by ID or slug.",
140+
{
141+
id: { type: "string", description: "Release ID or slug" },
142+
includeReleaseNotes: {
143+
type: "boolean",
144+
description: "Include associated release notes",
145+
},
146+
},
125147
),
126148
mcpDef(
127149
"mcp__linear__get_release_note",
128150
"[linear] Retrieve release notes by ID or slug, including markdown content.",
151+
{
152+
id: { type: "string", description: "Release notes ID or slug" },
153+
includeReleases: {
154+
type: "boolean",
155+
description: "Include associated releases",
156+
},
157+
},
129158
),
130159
mcpDef(
131160
"mcp__linear__list_release_notes",
132161
"[linear] List release notes in the workspace, optionally filtered by pipeline or covered release.",
162+
{
163+
pipeline: {
164+
type: "string",
165+
description: "Release pipeline ID, slug, or exact name",
166+
},
167+
query: { type: "string", description: "Search release notes title" },
168+
release: { type: "string", description: "Release ID or slug" },
169+
},
133170
),
134171
mcpDef(
135172
"mcp__linear__list_release_pipelines",
136173
"[linear] List release pipelines in the workspace.",
174+
{
175+
query: { type: "string", description: "Search pipeline name" },
176+
team: { type: "string", description: "Team name or ID" },
177+
},
137178
),
138179
];
139180

@@ -340,11 +381,99 @@ describe("createToolIndex", () => {
340381
});
341382

342383
test("a word repeated in the query does not bury the tool that answers it", () => {
343-
const ranked = linearReleaseIndex.search(
384+
// Parameters stripped so only counting the repeated word once keeps
385+
// list_issues above the cut.
386+
const namesOnlyIndex = createToolIndex(() =>
387+
linearReleaseCatalog.map((def) => ({ ...def, inputSchema: {} })),
388+
);
389+
const ranked = namesOnlyIndex.search(
344390
"Linear list issues filtered by release ID or release slug",
345391
);
346392
expect(ranked).toContain("mcp__linear__list_issues");
347393
});
394+
395+
test("a release filter lifts list_issues above the tools named for releases", () => {
396+
expect(linearReleaseIndex.search("linear issues release filter")[0]).toBe(
397+
"mcp__linear__list_issues",
398+
);
399+
});
400+
401+
test("a release filter narrows a release query to the tool that takes it", () => {
402+
expect(linearReleaseIndex.search("issues for release")).toEqual([
403+
"mcp__linear__list_issues",
404+
]);
405+
});
406+
407+
// Identical names and descriptions apart from the filter they take, with
408+
// the cycle tool first so catalog order alone would rank it ahead.
409+
const recordsCatalog: ToolDefinition[] = [
410+
mcpDef("mcp__acme__list_records_v1", "List records", {
411+
cycle: { type: "string", description: "Cycle name, number, or ID" },
412+
}),
413+
mcpDef("mcp__acme__list_records_v2", "List records", {
414+
release: { type: "string", description: "Release ID or slug" },
415+
}),
416+
];
417+
const recordsIndex = createToolIndex(() => recordsCatalog);
418+
419+
test("a parameter name lifts the tool that takes it", () => {
420+
expect(recordsIndex.search("list records by release")[0]).toBe(
421+
"mcp__acme__list_records_v2",
422+
);
423+
});
424+
425+
test("a word in a parameter description lifts the tool that takes it", () => {
426+
expect(recordsIndex.search("list records by slug")[0]).toBe(
427+
"mcp__acme__list_records_v2",
428+
);
429+
});
430+
431+
test("a query matching only parameter text is a miss", () => {
432+
expect(recordsIndex.search("release")).toEqual([]);
433+
});
434+
435+
test("a tool named for the query keeps a parameter-only match off the shortlist", () => {
436+
const ranked = linearReleaseIndex.search("linear release");
437+
expect(ranked).toContain("mcp__linear__get_release");
438+
expect(ranked).not.toContain("mcp__linear__list_issues");
439+
});
440+
441+
test("a parameter shared across tools does not displace the tool named for it", () => {
442+
const teamIndex = createToolIndex(() => [
443+
...linearReleaseCatalog,
444+
mcpDef(
445+
"mcp__linear__list_teams",
446+
"[linear] List teams in the user's Linear workspace, including their parent IDs and hierarchy names",
447+
{ query: { type: "string", description: "Search query" } },
448+
),
449+
]);
450+
expect(teamIndex.search("linear team")[0]).toBe("mcp__linear__list_teams");
451+
});
452+
453+
test("a malformed input schema still scores on name and description", () => {
454+
const malformedIndex = createToolIndex(() => [
455+
{
456+
name: "mcp__acme__null_props",
457+
description: "Archive widgets",
458+
inputSchema: { type: "object", properties: null },
459+
},
460+
{
461+
name: "mcp__acme__numeric_description",
462+
description: "Archive gadgets",
463+
inputSchema: {
464+
type: "object",
465+
properties: { release: { type: "string", description: 42 } },
466+
},
467+
},
468+
]);
469+
expect(malformedIndex.search("archive")).toEqual([
470+
"mcp__acme__null_props",
471+
"mcp__acme__numeric_description",
472+
]);
473+
expect(malformedIndex.search("archive release")[0]).toBe(
474+
"mcp__acme__numeric_description",
475+
);
476+
});
348477
});
349478

350479
describe("createToolSearchTool skills", () => {

0 commit comments

Comments
 (0)