sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

fix(durable): rank tool search like the coding agent (#11631)

Browse source 
Ports the coding agent's tool ranking fixes for #11331 (#11440, #11625)
to durable, which keeps its own copy of the ranker in `src/search.ts`.

Problem: durable's `rankDocuments()` was still the coding agent's ranker
from before both fixes. `thread_read` with `limit: 1` returned
`queue_read`, and a tool name among other words (`peekaboo app`) lost to
tools that only mention the words in their descriptions.

Change:
- `src/search.ts`: same scoring as the coding agent's `Bm25Ranker`, plus
coverage first: a tool matching more of the query's terms (in any field)
ranks above one matching fewer, whatever the scores. Without it, name
weighting regressed `screenshot window` (1 to 2, `window` beat `image`)
and `describe app` (1 to 2, `app` beat `see`). The name is a separate
field (3x IDF per term, above any number of description mentions), a
tool whose whole local name is in the query gets a bonus worth one more
match of every query term, a query that is exactly a tool's name ranks
it first, and namespace text is a separate 0.5x field that is not
length-normalized.
- `tool_search`: the search text no longer repeats the name; the ranker
reads it from `name`.
- `searchTools()`: same, and the namespace text now includes the
namespace instructions, as in the coding agent.
- `test/search.test.ts`: the coding agent's ranker cases.

Remaining difference from the coding agent: durable's `tool_search` has
no namespaces (they are a code mode option), and `searchTools()` ranks
on the tool sample from the start data (description plus TypeScript
declaration) rather than description plus schema text.

`searchTools()` results are journaled, so replaying a script started
before this change returns the old results; no `CODEMODE_MODULE_VERSION`
bump.


## Before/after

`searchTools()` ranking on a made-up set of 16 MCP tools from 5 servers.
Each tool is searched with its description; each server is a code mode
namespace (prefix `mcp__<server>__`) with the description and
instructions below. Rank of the wanted tool among all matches, with the
top result when it is another tool.

| Server | Description | Instructions | Tools: description |
|---|---|---|---|
| `github` | GitHub repositories, issues, and pull requests | Use
owner/repo for repositories. Issues and pull requests share numbers. |
`list_issues`: List issues in a repository, filtered by state, labels,
or assignee.<br>`get_issue`: Get an issue with its comments by
number.<br>`create_issue`: Open a new issue in a
repository.<br>`search_code`: Search code across
repositories.<br>`create_pull_request`: Open a pull request from a
branch. |
| `linear` | Linear issue tracker | Search issues before creating new
ones. Issues belong to teams. | `search_issues`: Search Linear issues by
text.<br>`create_issue`: Create a Linear issue in a team. |
| `peekaboo` | macOS automation | Peekaboo controls macOS apps. Use app
names or bundle IDs. Peekaboo can capture any app window. | `app`:
Launch, quit, relaunch, focus, hide, or switch applications. Targets an
application by name, bundle identifier, or process ID, and can wait
until it is ready.<br>`see`: Capture the app window and describe the UI
elements of the app.<br>`window`: Move, resize, or close a window of an
app.<br>`menu`: Click a menu item in an app.<br>`image`: Take a
screenshot of the screen, an app, or a window. |
| `app` | Thread and queue state of the app | Threads hold messages; the
queue holds pending messages of a thread. | `queue_read`: Read a queued
message of this thread.<br>`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. |
| `docs` | Project documentation | Search the docs before answering
questions about the API. | `search`: Search the
documentation.<br>`read_page`: Read a documentation page by path. |

Before, a tool's search text was its name, the name with `_` as spaces,
its description, and the namespace name and description (no
instructions). After, the name and the namespace (with instructions) are
separate fields.

| Query | Wanted | Before: rank | After: rank | Change |
|---|---|---|---|---|
| `thread_read` | `app__thread_read` | 2 (top: `app__queue_read`) | 1 |
$\color{green}{\textsf{better}}$ |
| `mcp__app__thread_read` | `app__thread_read` | 2 (top:
`app__queue_read`) | 1 | $\color{green}{\textsf{better}}$ |
| `peekaboo app` | `peekaboo__app` | 5 (top: `peekaboo__see`) | 1 |
$\color{green}{\textsf{better}}$ |
| `app` | `peekaboo__app` | 4 (top: `app__queue_read`) | 1 |
$\color{green}{\textsf{better}}$ |
| `read thread messages` | `app__thread_read` | 2 (top:
`app__queue_read`) | 1 | $\color{green}{\textsf{better}}$ |
| `open the peekaboo app tool` | `peekaboo__app` | 5 (top:
`peekaboo__see`) | 1 | $\color{green}{\textsf{better}}$ |
| `launch app` | `peekaboo__app` | 1 | 1 | $\color{gray}{\textsf{same}}$
|
| `create linear issue` | `linear__create_issue` | 1 | 1 |
$\color{gray}{\textsf{same}}$ |
| `search documentation` | `docs__search` | 1 | 1 |
$\color{gray}{\textsf{same}}$ |
| `read docs` | `docs__read_page` | 1 | 1 |
$\color{gray}{\textsf{same}}$ |
| `pull request` | `github__create_pull_request` | 1 | 1 |
$\color{gray}{\textsf{same}}$ |
| `get issue comments` | `github__get_issue` | 1 | 1 |
$\color{gray}{\textsf{same}}$ |
| `screenshot window` | `peekaboo__image` | 1 | 1 |
$\color{gray}{\textsf{same}}$ |
| `describe app` | `peekaboo__see` | 1 | 1 |
$\color{gray}{\textsf{same}}$ |
| `create issue` | `github__create_issue` | 2 (top:
`linear__create_issue`) | 2 (top: `linear__create_issue`) |
$\color{gray}{\textsf{same}}$ |
| `search github issues` | `github__list_issues` | 5 (top:
`github__search_code`) | 5 (top: `github__search_code`) |
$\color{gray}{\textsf{same}}$ |

No regressions on this set. `search github issues` stays at 5:
`list_issues` does not contain "search", so it ties on coverage (2 of 3
terms) with `search_code` and `linear__search_issues`, and their scores
win.
This commit is contained in:
Petr Houška 2026-10-30 23:08:57 +01:00 • committed by GitHub
parent e033f255f0
commit c7372af420
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 251 additions and 55 deletions
Download patch file Download diff file Expand all files Collapse all files

1
packages/durable/CHANGELOG.md
View file

@ -181,6 +181,7 @@
- Fixed a host that crashes during a model request resending it forever, adding an aborted entry each time: the interrupted request is a failed attempt under the retry policy, and the run ends with `runFailure` reason `interrupted` once no retry is left; other work cut off `restartLimit` times ends instead of keeping its conversation busy ([#11139](https://github.com/earendil-works/pi/issues/11139))
- Fixed the built-in `read` resolving to `""` for an image in Code Mode scripts, so `image(await tools.read(...))` threw: `read` declares a `structuredOutputSchema`, a text file's text or an image's `ImageContent`, as the coding agent's `read` does ([#11015](https://github.com/earendil-works/pi/issues/11015))
- Fixed the Code Mode extension loading pi-codemode's WASM QuickJS runtime (`quickjs-wasi`) although it only runs scripts in the sandbox the host passes: it imports `toCodemodeIdentifier` from `@earendil-works/pi-codemode/declarations` instead of `/portable` ([#11254](https://github.com/earendil-works/pi/issues/11254))
- Fixed `tool_search` and Code Mode's `searchTools()` ranking a tool below tools that only mention its name, as the coding agent did before [#11440](https://github.com/earendil-works/pi/pull/11440) and [#11625](https://github.com/earendil-works/pi/pull/11625): `thread_read` with `limit: 1` returned `queue_read`, and `peekaboo app` missed `mcp__peekaboo__app`. Both now rank like the coding agent: a query that is exactly a tool's name ranks it first, a tool matching more of the query's words ranks above one matching fewer, a query word in the name outweighs any number of mentions in the description, and a tool whose whole name is in the query gets a large bonus. `searchTools()` also searches namespace instructions, as a low-weight field that does not lower a tool's score ([#11331](https://github.com/earendil-works/pi/issues/11331))
## [1.4.0] - 2026-10-21

10
packages/durable/docs/spec.md
View file

@ -3808,8 +3808,14 @@ declares them, built from the resolved tools. Loaded tools follow the other
tools in load order, so a load appends declarations to the transcript and the
cached prefix holds. `@earendil-works/pi-durable/tool-search` provides the
coding agent's `tool_search` (`createToolSearchExtension()`): it ranks
`Agent.deferred` with BM25 over the tools' names, descriptions, and parameter
names and descriptions, and loads the matches. The `codemode` section does not
`Agent.deferred` with the coding agent's tool ranking and loads the matches.
The ranking is BM25 over the tools' descriptions and parameter names and
descriptions, with the name as a separate field: a tool matching more of the
query's terms ranks above one matching fewer, a query term in the name
outweighs any number of mentions in the description, a tool whose whole name
after its last `__` is in the query gets a large bonus, and a query that is
exactly a tool's name ranks it first. Code mode's `searchTools()` ranks the
same way, with each tool's namespace text as a further, low-weight field. The `codemode` section does not
list deferred tools; scripts find them with `searchTools()`.
Tools reach files and processes only through `api.env`, never through an

12
packages/durable/src/codemode/codemode.ts
View file

@ -946,14 +946,10 @@ class CodemodeCall {
const documents = tools.flatMap((tool) => {
const own = namespaceOfTool(tool.name);
if (typeof namespace === "string" && (own === undefined || !matchesNamespace(own, namespace))) return [];
const text = [
tool.name,
tool.name.replaceAll("_", " "),
tool.description,
own?.name ?? "",
own?.description ?? "",
];
return [{ name: tool.name, text: text.join(" ") }];
const namespaceText =
own === undefined ? "" : [own.name, own.description ?? "", own.instructions ?? ""].join(" ");
// `description` is the tool's sample: its description and declaration, with the parameters' names and descriptions.
return [{ name: tool.name, text: tool.description, namespaceText }];
});
return rankDocuments(query, documents, limit).map((name) => {
const tool = tools.find((candidate) => candidate.name === name)!;

173
packages/durable/src/search.ts
View file

@ -3,7 +3,113 @@
* ranking: durable keeps its own copy of what it needs from the coding agent.
*/
export type SearchDocument = { readonly name: string; readonly text: string };
/** A tool as the ranker sees it: its name, its text, and the text of its namespace. */
export type SearchDocument = {
readonly name: string;
/** Description and parameter names and descriptions; not the name, which the ranker reads from `name`. */
readonly text: string;
/** Name, description, and instructions of the tool's namespace. */
readonly namespaceText?: string;
};
/**
* Okapi BM25 (k1 1.2, b 0.75) over `text`, with the name and the namespace as separate fields: the names of the best
* matches, best first; ties keep document order.
* - 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 text.
* - A tool whose whole local name (after the last `__`) is in the query gets a bonus worth one more match of every
* query term, so among tools matching as many terms, `peekaboo app` finds `mcp__peekaboo__app` over tools that
* only mention the words.
* - A query that is exactly a tool's name, or its name after a `__`, ranks that tool first, before coverage.
*
* Only `text` is length-normalized, so a long name or long namespace instructions do not lower a tool's score.
*/
export function rankDocuments(query: string, documents: readonly SearchDocument[], limit: number): string[] {
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) || document.namespaceTerms.has(term),
).length;
return [term, Math.log(1 + (indexed.length - frequency + 0.5) / (frequency + 0.5))] as const;
}),
);
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 = K1 * (1 - B + (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) +
(document.namespaceTerms.has(term) ? NAMESPACE_WEIGHT : 0) +
(count * (K1 + 1)) / (count + norm);
if (weight > 0) covered++;
score += idf.get(term)! * 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(
(left, right) =>
Number(right.exact) - Number(left.exact) || right.covered - left.covered || right.score - left.score,
)
.slice(0, limit)
.map((match) => match.name);
}
const K1 = 1.2;
const B = 0.75;
/**
* 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 text.
*/
const NAME_WEIGHT = 3;
/**
* A query term in the namespace text counts this many times its IDF. Every tool of the namespace shares it, so it helps
* find the namespace but not the tool in it.
*/
const NAMESPACE_WEIGHT = 0.5;
/** A document split into the fields the ranker weighs differently. */
type IndexedDocument = {
name: string;
/** Terms of the name: `mcp`, `app`, `thread`, `read` for `mcp__app__thread_read`. */
nameTerms: ReadonlySet<string>;
/** Terms of the name after its last `__`, the part models say: `thread`, `read`. */
localNameTerms: readonly string[];
textCounts: ReadonlyMap<string, number>;
textLength: number;
namespaceTerms: ReadonlySet<string>;
};
function indexDocument(document: SearchDocument): IndexedDocument {
const textCounts = new Map<string, number>();
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),
namespaceTerms: new Set(tokenize(document.namespaceText ?? "")),
};
}
const STOP_WORDS: ReadonlySet<string> = new Set([
"a",
@ -29,14 +135,6 @@ const STOP_WORDS: ReadonlySet<string> = new Set([
"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;
}
/** Lowercase terms, split at camelCase boundaries and non-alphanumerics, without stop words. */
export function tokenize(text: string): string[] {
return text
@ -48,38 +146,27 @@ export function tokenize(text: string): string[] {
.map(stem);
}
/** Okapi BM25 (k1 1.2, b 0.75): the names of the best matches, best first; ties keep document order. */
export function rankDocuments(query: string, documents: readonly SearchDocument[], limit: number): string[] {
const k1 = 1.2;
const b = 0.75;
const queryTerms = [...new Set(tokenize(query))];
if (queryTerms.length === 0 || documents.length === 0 || limit <= 0) return [];
const termCounts = documents.map((document) => {
const counts = new Map<string, number>();
for (const term of tokenize(document.text)) counts.set(term, (counts.get(term) ?? 0) + 1);
return counts;
});
const lengths = termCounts.map((counts) => [...counts.values()].reduce((sum, count) => sum + count, 0));
const averageLength = lengths.reduce((sum, length) => sum + length, 0) / documents.length || 1;
const idf = new Map(
queryTerms.map((term) => {
const frequency = termCounts.filter((counts) => counts.has(term)).length;
return [term, Math.log(1 + (documents.length - frequency + 0.5) / (frequency + 0.5))] as const;
}),
);
const matches: { name: string; score: number }[] = [];
documents.forEach((document, index) => {
let score = 0;
for (const term of queryTerms) {
const count = termCounts[index]!.get(term);
if (!count) continue;
const norm = k1 * (1 - b + (b * lengths[index]!) / averageLength);
score += (idf.get(term) ?? 0) * ((count * (k1 + 1)) / (count + norm));
}
if (score > 0) matches.push({ name: document.name, score });
});
return matches
.sort((left, right) => right.score - left.score)
.slice(0, limit)
.map((match) => match.name);
/** 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;
}
/** 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}`);
}

4
packages/durable/src/tool-search/tool-search.ts
View file

@ -45,9 +45,9 @@ function schemaText(schema: unknown, parts: string[]): void {
}
}
/** Search text of a tool: its name, the name with `_` as spaces, its description, and its parameters' text. */
/** Search text of a tool: its description and its parameters' names and descriptions; the ranker reads the name. */
function searchText(tool: ToolRegistration): string {
const parts = [tool.name, tool.name.replaceAll("_", " "), tool.description];
const parts = [tool.description];
schemaText(tool.parameters, parts);
return parts.filter((part) => part.trim() !== "").join(" ");
}

106
packages/durable/test/search.test.ts Normal file
View file

@ -0,0 +1,106 @@
import { describe, expect, it } from "vitest";
import { rankDocuments, type SearchDocument, tokenize } from "../src/search.ts";
function doc(name: string, text: string, namespaceText?: string): SearchDocument {
return namespaceText === undefined ? { name, text } : { name, text, namespaceText };
}
describe("tokenize", () => {
it("splits camelCase and snake_case, drops stop words, and folds plurals", () => {
expect(tokenize("listIssues for the GitHub_repo")).toEqual(["list", "issue", "git", "hub", "repo"]);
expect(tokenize("searches queries HTTPServer")).toEqual(["search", "query", "http", "server"]);
});
});
describe("rankDocuments", () => {
const documents = [
doc("mcp__github__list_issues", "List issues in a repository. state open or closed"),
doc("mcp__github__create_pull_request", "Open a pull request."),
doc("mcp__linear__search_issues", "Search Linear issues by text."),
doc("mcp__docs__search", "Search the documentation."),
];
it("ranks by term relevance and respects the limit", () => {
expect(rankDocuments("issue", documents, 8)).toEqual(["mcp__linear__search_issues", "mcp__github__list_issues"]);
expect(rankDocuments("pull requests", documents, 8)[0]).toBe("mcp__github__create_pull_request");
expect(rankDocuments("search", documents, 1)).toHaveLength(1);
expect(rankDocuments("closed", documents, 8)).toEqual(["mcp__github__list_issues"]);
});
it("returns nothing for unknown or empty queries", () => {
expect(rankDocuments("kubernetes", documents, 8)).toEqual([]);
expect(rankDocuments("the", documents, 8)).toEqual([]);
});
it("searches the namespace text", () => {
const namespaced = [doc("mcp__x__run", "Run it.", "mcp__x Kubernetes cluster tools")];
expect(rankDocuments("kubernetes", namespaced, 8)).toEqual(["mcp__x__run"]);
});
// https://github.com/earendil-works/pi/issues/11331
it("ranks a tool first when the query is its name", () => {
const named = [
doc("mcp__app__queue_read", "Read a queued message of this thread."),
doc(
"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.",
),
];
for (const query of ["thread_read", "dev_app__thread_read", "mcp__dev-app__thread_read", "Thread_Read"]) {
expect(rankDocuments(query, named, 1)).toEqual(["mcp__dev-app__thread_read"]);
}
// Partial names still rank by BM25.
expect(rankDocuments("read", named, 1)).toEqual(["mcp__app__queue_read"]);
});
// https://github.com/earendil-works/pi/issues/11331
it("ranks a name match above mentions in descriptions", () => {
const ci = [
doc("mcp__ci__status", "Deploy status: last deploy, deploy log, deploy errors."),
doc(
"mcp__ci__deploy_release",
"Start a release of the current branch to an environment, run the migrations, switch traffic over, and wait until the new version is live and healthy.",
),
];
expect(rankDocuments("deploy", ci, 1)).toEqual(["mcp__ci__deploy_release"]);
});
// https://github.com/earendil-works/pi/issues/11331
it("ranks a tool first when the query contains its whole name", () => {
const peekaboo = "mcp__peekaboo Peekaboo controls macOS apps by app name.";
const named = [
doc("mcp__peekaboo__see", "Capture the app window of the peekaboo app and describe the app.", peekaboo),
doc("mcp__peekaboo__menu", "Click a menu item of an app.", peekaboo),
doc(
"mcp__peekaboo__app",
"Launch, quit, relaunch, focus, hide, or switch applications. Targets an application by name, bundle identifier, or process ID, and can wait until it is ready.",
peekaboo,
),
doc("mcp__app__queue_read", "Read a queued message of this thread."),
doc(
"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.",
),
];
expect(rankDocuments("peekaboo app", named, 1)).toEqual(["mcp__peekaboo__app"]);
expect(rankDocuments("open the peekaboo app tool", named, 1)).toEqual(["mcp__peekaboo__app"]);
expect(rankDocuments("read thread_read state", named, 1)).toEqual(["mcp__dev-app__thread_read"]);
// Without the whole name, mentions in descriptions decide.
expect(rankDocuments("describe window", named, 1)).toEqual(["mcp__peekaboo__see"]);
});
// https://github.com/earendil-works/pi/issues/11331
it("ranks a tool matching more query terms first, whatever the scores", () => {
const peekaboo = [
doc("mcp__peekaboo__app", "Launch, quit, focus, or switch applications by name or bundle identifier."),
doc("mcp__peekaboo__see", "Capture the app window and describe the UI elements of the app."),
doc("mcp__peekaboo__window", "Move, resize, or close a window of an app."),
doc("mcp__peekaboo__image", "Take a screenshot of the screen, an app, or a window."),
];
// `window` and `app` are whole names, but each covers one of the two terms.
expect(rankDocuments("screenshot window", peekaboo, 1)).toEqual(["mcp__peekaboo__image"]);
expect(rankDocuments("describe app", peekaboo, 1)).toEqual(["mcp__peekaboo__see"]);
// An exact name still ranks first.
expect(rankDocuments("window", peekaboo, 1)).toEqual(["mcp__peekaboo__window"]);
});
});