fix(tool-search): bridge batch barrier, listing truncation, source indexing (salvage #92693, part 1)

Four fixes to the tool-search deferral layer, split from PR #92693 (the
availability-cache staleness fix ships separately):

1. The parallel batch planner now peels the tool_call bridge wrapper and
   decides admission on the underlying tool — supports_parallel_tool_calls
   works again when deferral is active. Unparseable wrappers stay
   sequential barriers; bridged calls get exactly the admission the same
   call gets direct. tool_search/tool_describe lookups batch concurrently.
2. _short_desc no longer truncates listing lines at 'e.g.', hostnames, or
   version strings — a sentence terminator must be followed by whitespace.
3. BM25 indexes the source label (e.g. 'linear' for mcp-linear), so
   service-name queries reach tools whose own name omits the service; the
   dead 'mcp' prefix token is stripped.
4. Substring-fallback docstring corrected (token misses, not zero-IDF).

Salvaged from #92693 by @alt-glitch with authorship preserved.
This commit is contained in:
alt-glitch
2026-08-25 14:16:35 -07:00
committed by Teknium
parent 95668f5eab
commit 62b2d78025
5 changed files with 435 additions and 27 deletions
+49 -4
View File
@@ -114,6 +114,38 @@ def _is_mcp_tool_parallel_safe(tool_name: str) -> bool:
return False
# Read-only bridge lookups: dispatch_tool_search / dispatch_tool_describe are
# stateless catalog reads (the catalog is rebuilt from the current tool-defs
# list on every call), so a batch of them can run concurrently.
_PARALLEL_SAFE_BRIDGE_LOOKUPS = frozenset({"tool_search", "tool_describe"})
def _peel_bridge_call(tool_name: str, function_args: dict) -> tuple[str, dict]:
"""Resolve a ``tool_call`` bridge invocation to its underlying tool.
The batch planner admits calls to a parallel run by tool NAME, but when
tool search is active the model emits the literal name ``tool_call`` for
every deferred tool — so a server opted in via
``supports_parallel_tool_calls: true`` silently lost concurrency the
moment the bridge activated. Peel the wrapper here so admission is
decided on the underlying tool, exactly like the executors' unwrap.
Returns ``(underlying_name, underlying_args)`` when the wrapper parses
cleanly, else ``(tool_name, function_args)`` unchanged — an unparseable
bridge call stays a sequential barrier and fails at dispatch as before.
"""
try:
from tools.tool_search import TOOL_CALL_NAME, resolve_underlying_call
if tool_name != TOOL_CALL_NAME:
return tool_name, function_args
underlying, underlying_args, err = resolve_underlying_call(function_args)
if err is not None or not underlying:
return tool_name, function_args
return underlying, underlying_args
except Exception:
return tool_name, function_args
def _plan_tool_batch_segments(tool_calls, *, execution_cwd: Optional[Path] = None) -> List[tuple]:
"""Split a tool-call batch into ordered ``(kind, calls)`` segments.
@@ -193,14 +225,23 @@ def _plan_tool_batch_segments(tool_calls, *, execution_cwd: Optional[Path] = Non
_add_sequential(tool_call)
continue
if tool_name in _PATH_SCOPED_TOOLS:
# Bridge unwrap: admission is decided on the UNDERLYING tool, not on
# the literal wrapper name the model emitted. Read-only bridge
# lookups (tool_search / tool_describe) are parallel-safe as-is.
effective_name, effective_args = _peel_bridge_call(tool_name, function_args)
if effective_name in _NEVER_PARALLEL_TOOLS:
_add_sequential(tool_call)
continue
if effective_name in _PATH_SCOPED_TOOLS:
scoped_paths = _extract_parallel_scope_paths(
tool_name, function_args, execution_cwd=execution_cwd
effective_name, effective_args, execution_cwd=execution_cwd
)
if not scoped_paths:
_add_sequential(tool_call)
continue
is_writer = tool_name in _PATH_SCOPED_WRITERS
is_writer = effective_name in _PATH_SCOPED_WRITERS
if any(
(is_writer or existing_is_writer)
and _paths_overlap(scoped_path, existing)
@@ -216,7 +257,11 @@ def _plan_tool_batch_segments(tool_calls, *, execution_cwd: Optional[Path] = Non
current.append(tool_call)
continue
if tool_name in _PARALLEL_SAFE_TOOLS or _is_mcp_tool_parallel_safe(tool_name):
if (
effective_name in _PARALLEL_SAFE_TOOLS
or effective_name in _PARALLEL_SAFE_BRIDGE_LOOKUPS
or _is_mcp_tool_parallel_safe(effective_name)
):
current.append(tool_call)
continue
+4
View File
@@ -2832,6 +2832,10 @@ DEFAULT_CONFIG = {
# model knows which domains are reachable; individual tools
# discoverable through tool_search only.
# "auto"/"on" — activate when at least one deferrable tool exists.
# Today "auto" is an alias of "on"; it stays the default so a
# future budget-gated mode (inline schemas when they fit, defer
# only when they don't) can land on "auto" without changing
# behavior for anyone who pinned "on" or "off" explicitly.
# "off" — disable entirely. Tools-array assembly is a pass-through.
"enabled": "auto",
# Listing budget as a percentage of the active model's context
+325
View File
@@ -0,0 +1,325 @@
"""Deferral-layer fixes: behavior regression suite.
Each test class pins one user-visible behavior that was broken while the
tool_search bridge was active. Tests assert at public seams (planner
segment shapes, search results, listing lines, get_tool_definitions
output) — not private implementation details — so refactors that keep
the behavior keep the tests.
The bugs, as reproduced before the fix:
1. ``_plan_tool_batch_segments`` classified the literal name ``tool_call``
as a sequential barrier, so a server opted in via
``supports_parallel_tool_calls: true`` silently lost all concurrency
the moment the bridge activated (every deferred call arrives wrapped).
2. ``_short_desc`` cut at the first ``.`` anywhere, so "e.g.", "v1.2",
and "api.github.com" truncated catalog listing lines to garbage.
3. The BM25 document didn't include the tool's source, so a query naming
the service ("linear") missed tools whose own name omits it.
4. (docstring-only) the substring fallback documented a zero-IDF case
that cannot occur with the Lucene IDF variant.
"""
import json
import time
import uuid
from types import SimpleNamespace
import pytest
from agent.tool_dispatch_helpers import _plan_tool_batch_segments
from tools.tool_search import _short_desc, build_catalog, search_catalog
def _tc(name, arguments="{}", call_id=None):
return SimpleNamespace(
id=call_id or f"call_{uuid.uuid4().hex[:8]}",
type="function",
function=SimpleNamespace(name=name, arguments=arguments),
)
def _bridge_tc(underlying, arguments=None, call_id=None):
"""A tool_call bridge invocation as the model emits it."""
return _tc(
"tool_call",
json.dumps({"name": underlying, "arguments": arguments or {}}),
call_id=call_id,
)
def _td(name, desc="", params=None, required=None):
parameters = {"type": "object", "properties": params or {}}
if required:
parameters["required"] = required
return {
"type": "function",
"function": {"name": name, "description": desc, "parameters": parameters},
}
def _kinds(segments):
return [kind for kind, _ in segments]
def _flatten_ids(segments):
return [tc.id for _, calls in segments for tc in calls]
@pytest.fixture
def mcp_pair(monkeypatch):
"""Two tools on a parallel-opted-in MCP server, registered for real.
Registers via the actual registry (so ``resolve_underlying_call``'s
deferability check passes) and marks the server parallel-safe through
the real provenance maps in ``tools.mcp_tool``.
"""
from tools import mcp_tool
from tools.registry import registry
names = ["mcp__pytestsrv__alpha_read", "mcp__pytestsrv__beta_read"]
for n in names:
registry.register(
name=n,
toolset="mcp-pytestsrv",
schema=_td(n, "Read-only test tool.")["function"],
handler=lambda args, **kw: json.dumps({"ok": True}),
)
with mcp_tool._lock:
for n in names:
mcp_tool._mcp_tool_server_names[n] = "pytestsrv"
mcp_tool._parallel_safe_servers.add("pytestsrv")
yield names
with mcp_tool._lock:
mcp_tool._parallel_safe_servers.discard("pytestsrv")
for n in names:
mcp_tool._mcp_tool_server_names.pop(n, None)
for n in names:
registry.deregister(n)
class TestBridgePeelInPlanner:
"""Fix 1: batch admission is decided on the underlying tool."""
def test_two_bridged_parallel_safe_mcp_calls_run_parallel(self, mcp_pair):
alpha, beta = mcp_pair
calls = [_bridge_tc(alpha, call_id="a"), _bridge_tc(beta, call_id="b")]
segments = _plan_tool_batch_segments(calls)
assert _kinds(segments) == ["parallel"]
assert _flatten_ids(segments) == ["a", "b"]
def test_bridged_call_to_non_opted_in_tool_stays_sequential(self, mcp_pair):
from tools import mcp_tool
with mcp_tool._lock:
mcp_tool._parallel_safe_servers.discard("pytestsrv")
try:
alpha, beta = mcp_pair
calls = [_bridge_tc(alpha, call_id="a"), _bridge_tc(beta, call_id="b")]
segments = _plan_tool_batch_segments(calls)
assert _kinds(segments) == ["sequential"]
finally:
with mcp_tool._lock:
mcp_tool._parallel_safe_servers.add("pytestsrv")
def test_bridge_lookups_are_parallel_safe(self):
calls = [
_tc("tool_search", '{"query": "issues"}', call_id="s1"),
_tc("tool_search", '{"query": "pages"}', call_id="s2"),
_tc("tool_describe", '{"name": "mcp__x__y"}', call_id="d1"),
]
segments = _plan_tool_batch_segments(calls)
assert _kinds(segments) == ["parallel"]
assert _flatten_ids(segments) == ["s1", "s2", "d1"]
def test_malformed_bridge_call_stays_a_barrier(self):
calls = [
_tc("tool_call", '{"arguments": {}}', call_id="bad"), # no name
_tc("web_search", '{"query": "x"}', call_id="r1"),
_tc("web_search", '{"query": "y"}', call_id="r2"),
]
segments = _plan_tool_batch_segments(calls)
assert _kinds(segments) == ["sequential", "parallel"]
assert [tc.id for tc in segments[0][1]] == ["bad"]
def test_emission_order_survives_the_peel(self, mcp_pair):
alpha, beta = mcp_pair
calls = [
_bridge_tc(alpha, call_id="a"),
_tc("terminal", '{"command": "make"}', call_id="t"),
_bridge_tc(beta, call_id="b"),
]
segments = _plan_tool_batch_segments(calls)
assert _flatten_ids(segments) == ["a", "t", "b"]
def test_bridged_mcp_admission_matches_direct_admission(self, mcp_pair, tmp_path, monkeypatch):
"""The peel restores PARITY, not extra permissiveness: a bridged call
to an opted-in MCP tool gets exactly the admission the same tool gets
when called directly. Opted-in MCP tools have always shared parallel
runs with core path-scoped tools (the server opt-in is the owner's
declared contract; the planner has never had per-MCP-tool resource
scopes) — the bridge must not silently upgrade OR downgrade that."""
monkeypatch.chdir(tmp_path)
alpha, _ = mcp_pair
direct = _plan_tool_batch_segments([
_tc(alpha, "{}", call_id="m1"),
_tc("write_file", '{"path":"x.py","content":"a"}', call_id="w1"),
])
bridged = _plan_tool_batch_segments([
_bridge_tc(alpha, {}, call_id="m1"),
_tc("write_file", '{"path":"x.py","content":"a"}', call_id="w1"),
])
assert [(k, [c.id for c in cs]) for k, cs in direct] == \
[(k, [c.id for c in cs]) for k, cs in bridged]
def test_core_file_tools_cannot_be_smuggled_through_the_bridge(self):
"""Wrapped core file tools remain sequential because they are not deferrable."""
calls = [
_bridge_tc("write_file", {"path": "a.py", "content": "x"}, call_id="w"),
_bridge_tc("read_file", {"path": "a.py"}, call_id="r"),
]
segments = _plan_tool_batch_segments(calls)
assert _kinds(segments) == ["sequential"]
assert _flatten_ids(segments) == ["w", "r"]
class TestShortDescSentenceBoundary:
"""Fix 2: listing lines survive abbreviations, versions, hostnames."""
def test_clean_two_sentence_case_still_clips_at_first(self):
assert _short_desc("Open an issue. Second sentence dropped.") == "Open an issue."
def test_abbreviation_does_not_truncate(self):
s = _short_desc("Create an issue (e.g. a bug report) in a repository.")
assert s.startswith("Create an issue (e.g. a bug report)")
def test_hostname_does_not_truncate(self):
s = _short_desc("Fetch a page from api.github.com and return the JSON body.")
assert "api.github.com" in s
def test_version_string_does_not_truncate(self):
s = _short_desc("Upgrade to v1.2 of the schema and migrate all rows.")
assert "v1.2" in s
def test_exclamation_terminator_is_kept(self):
assert _short_desc("List repos! Supports pagination.") == "List repos!"
def test_question_terminator_is_kept(self):
s = _short_desc("What does this do? It lists channels.")
assert s == "What does this do?"
def test_long_text_still_clips_with_ellipsis(self):
s = _short_desc("word " * 40)
assert len(s) <= 61
assert s.endswith("…")
def test_empty_is_empty(self):
assert _short_desc("") == ""
class TestSourceNameIndexing:
"""Fix 3: a query naming the service finds that source's tools."""
@staticmethod
def _register(name, toolset, desc):
from tools.registry import registry
registry.register(
name=name,
toolset=toolset,
schema=_td(name, desc)["function"],
handler=lambda args, **kw: json.dumps({"ok": True}),
)
return name
def test_service_query_reaches_tool_without_service_in_name(self):
"""A plugin tool named ``create_issue`` in toolset ``mcp-linear``
must be reachable by the query "linear"."""
from tools.registry import registry
names = [
self._register("create_issue", "mcp-linear", "Create a new issue in a team."),
self._register("post_message", "mcp-slack", "Post a message to a channel."),
]
try:
defs = [_td(n, d) for n, d in
[("create_issue", "Create a new issue in a team."),
("post_message", "Post a message to a channel.")]]
catalog = build_catalog(defs)
hits = search_catalog(catalog, "linear")
assert [h.name for h in hits] == ["create_issue"]
finally:
for n in names:
registry.deregister(n)
def test_mcp_prefix_is_not_a_matchable_token(self):
"""The shared ``mcp`` prefix used to sit in every native MCP document
as a near-zero-IDF token: a query containing "mcp" matched EVERY
tool, drowning the discriminating terms. Now "mcp" contributes
nothing to ranking, so the discriminating term decides alone."""
from tools.registry import registry
names = [
self._register("mcp__linear__create_issue", "mcp-linear", "Create an issue."),
self._register("mcp__slack__post_message", "mcp-slack", "Post a message."),
]
try:
defs = [_td("mcp__linear__create_issue", "Create an issue."),
_td("mcp__slack__post_message", "Post a message.")]
catalog = build_catalog(defs)
hits = search_catalog(catalog, "mcp message")
# Before the fix "mcp" BM25-matched both docs, so both came
# back and the order was decided by document length, not by
# the term the model actually meant.
assert [h.name for h in hits] == ["mcp__slack__post_message"]
finally:
for n in names:
registry.deregister(n)
def test_source_label_is_indexed_once_for_native_and_plugin_names(self):
from tools.registry import registry
source_label = "catalogsource"
names = [
self._register(
"mcp__catalogsource__native_action",
"mcp-catalogsource",
"Perform a native action.",
),
self._register(
"plugin_action",
"mcp-catalogsource",
"Perform a plugin action.",
),
]
try:
catalog = build_catalog([
_td("mcp__catalogsource__native_action", "Perform a native action."),
_td("plugin_action", "Perform a plugin action."),
])
tokens_by_name = {entry.name: entry._tokens for entry in catalog}
assert tokens_by_name[names[0]].count(source_label) == 1
assert tokens_by_name[names[1]].count(source_label) == 1
finally:
for name in names:
registry.deregister(name)
def test_substring_fallback_covers_token_misses(self):
""""hub" is a substring of github but never a token — the fallback
(not BM25) must return the github tools."""
from tools.registry import registry
names = [
self._register("github_create_issue", "mcp-github", "Create an issue."),
self._register("github_merge_pr", "mcp-github", "Merge a pull request."),
]
try:
defs = [_td("github_create_issue", "Create an issue."),
_td("github_merge_pr", "Merge a pull request.")]
catalog = build_catalog(defs)
hits = search_catalog(catalog, "hub")
assert {h.name for h in hits} == {"github_create_issue", "github_merge_pr"}
assert search_catalog(catalog, "zzzz") == []
finally:
for n in names:
registry.deregister(n)
+44 -17
View File
@@ -81,7 +81,7 @@ CHARS_PER_TOKEN = 4.0
class ToolSearchConfig:
"""Resolved, validated tool-search configuration for a single assembly."""
enabled: str # "auto" | "on" | "off"
enabled: str # "auto" | "on" | "off" — "auto" is an alias of "on" today
# Listing budget as a percentage of the model's context window. Under
# tiered disclosure this no longer gates *activation* (any deferrable
# tool activates the bridge) — it bounds how much context the embedded
@@ -295,6 +295,13 @@ def should_activate(
``"off"`` skips unconditionally. ``"on"`` and ``"auto"`` activate whenever
at least one deferrable tool exists (there's no point swapping a no-op).
``"auto"`` is an ALIAS of ``"on"`` under tiered disclosure — it is kept
as the shipped default so that a future budget-gated mode ("inline the
schemas when they fit, defer only when they don't") can change ``auto``'s
behavior without breaking users who explicitly pinned ``on`` or ``off``.
Do not add behavior that distinguishes them without that design; see the
config reference for the user-facing statement of this contract.
Tiered-disclosure semantics (July 2026): the presence of ANY MCP/plugin
tool activates the bridge — schemas always defer. What the threshold now
controls is the *listing budget* (see :func:`listing_token_budget`), not
@@ -353,22 +360,34 @@ def _tokenize(text: str) -> List[str]:
return [t.lower() for t in _TOKEN_RE.findall(text)]
def _entry_search_text(td: Dict[str, Any]) -> str:
def _entry_search_text(td: Dict[str, Any], source_label: str = "") -> str:
"""Build the search-text blob for a deferrable tool.
Includes the tool name (with underscores broken into words so BM25 can
match against query terms), the description, and the names of the
top-level parameters. Schema bodies are deliberately excluded —
indexing them adds noise without improving recall in our measurement.
match against query terms), the source label (the MCP server / plugin
toolset the tool belongs to, e.g. ``linear`` for toolset ``mcp-linear``),
the description, and the names of the top-level parameters. Schema
bodies are deliberately excluded — indexing them adds noise without
improving recall in our measurement.
The ``mcp__`` name prefix is stripped before splitting: ``mcp`` appears
in every native MCP tool document, so its IDF collapses to near zero —
it is dead weight in every document and useless as a query term.
Indexing the source label is what makes a service-name query ("linear")
reach a tool whose NAME does not carry the service (a plugin tool named
``create_issue``, or any catalog whose naming omits the vendor).
"""
fn = td.get("function") or {}
name = fn.get("name", "")
if name.startswith("mcp__"):
name = name[len("mcp__"):]
desc = fn.get("description", "") or ""
params = ((fn.get("parameters") or {}).get("properties") or {})
param_names = " ".join(params.keys())
# Break snake_case and dotted names into words for BM25.
name_words = name.replace("_", " ").replace(".", " ").replace("-", " ").replace(":", " ")
return f"{name_words} {desc} {param_names}"
extra = source_label if source_label and source_label not in name_words.split() else ""
return f"{name_words} {extra} {desc} {param_names}"
def _classify_source(name: str) -> Tuple[str, str]:
@@ -399,13 +418,17 @@ def build_catalog(tool_defs: List[Dict[str, Any]]) -> List[CatalogEntry]:
continue
desc = fn.get("description", "") or ""
source, source_name = _classify_source(name)
# Index the human-facing group label ("linear", not "mcp-linear") so
# a service-name query matches tools from that source even when the
# tool's own name omits the service.
source_label = _listing_group_label(source_name) if source_name else ""
entry = CatalogEntry(
name=name,
description=desc,
schema=td,
source=source,
source_name=source_name,
_tokens=_tokenize(_entry_search_text(td)),
_tokens=_tokenize(_entry_search_text(td, source_label)),
)
catalog.append(entry)
return catalog
@@ -445,11 +468,13 @@ def _bm25_score(query_tokens: List[str], doc_tokens: List[str],
def search_catalog(catalog: List[CatalogEntry], query: str, limit: int = 5) -> List[CatalogEntry]:
"""Return the top-``limit`` catalog entries for ``query`` by BM25.
Falls back to a stable name-substring match when BM25 yields no hits
above zero. That ensures a query like ``"github"`` against a catalog
where every tool is named ``github_*`` still returns results — BM25
can underperform when query and document share only one token that
appears in every document (zero IDF).
Falls back to a stable name-substring match when every query token
misses every document — e.g. the query ``"hub"`` against ``github_*``
tools ("hub" is a substring of the name but never a token, so BM25
scores nothing). The IDF variant used here,
``log(1 + (N - df + 0.5) / (df + 0.5))``, is strictly positive even
when a term appears in every document, so the fallback only runs when
no query token appears in any document.
"""
if not catalog or limit <= 0:
return []
@@ -490,22 +515,24 @@ def search_catalog(catalog: List[CatalogEntry], query: str, limit: int = 5) -> L
# ---------------------------------------------------------------------------
_SENTENCE_END_RE = re.compile(r"[.!?\n]")
# A sentence ends at ., !, or ? followed by whitespace or end-of-string, but
# not at the end of a common dotted abbreviation.
_SENTENCE_END_RE = re.compile(r"(?<!\be\.g)(?<!\bi\.e)(?<!\betc)[.!?](?=\s|$)")
def _short_desc(description: str, max_chars: int = 60) -> str:
"""First sentence of a tool description, clipped to ``max_chars``.
Mirrors the skills-listing convention: one terse line per capability.
Whitespace is collapsed; a hard clip never cuts mid-word unless the
first word itself exceeds the budget.
A terminator must be followed by whitespace or end-of-string; ``e.g.``,
``i.e.``, and ``etc.`` do not end a sentence. Whitespace normalization and
the unbounded regex search both remain linear-time on hostile input.
"""
text = " ".join((description or "").split())
if not text:
return ""
m = _SENTENCE_END_RE.search(text)
if m:
text = text[:m.start() + (1 if text[m.start()] == "." else 0)]
text = text[:m.end()]
if len(text) <= max_chars:
return text
clipped = text[:max_chars]
@@ -85,7 +85,7 @@ tools:
| Key | Default | Meaning |
| --- | --- | --- |
| `enabled` | `auto` | `auto`/`on` activate whenever at least one deferrable tool exists; `off` disables entirely (everything stays eager). |
| `enabled` | `auto` | `auto`/`on` activate whenever at least one deferrable tool exists; `off` disables entirely (everything stays eager). `auto` is currently an alias of `on` — it is reserved for a future mode that inlines schemas when they fit the context and defers only when they don't. Pin `on` or `off` if you want today's behavior guaranteed across upgrades. |
| `threshold_pct` | `5` | Listing budget as a percentage of the active model's context length. Range 0–100. |
| `search_default_limit` | `5` | Hits returned when the model calls `tool_search` without a `limit`. |
| `max_search_limit` | `20` | Hard upper bound the model can request via `limit`. Range 1–50. |
@@ -148,11 +148,18 @@ to any progressive-disclosure design, not specific to this implementation:
## Implementation details
- **Retrieval:** BM25 over tokenized tool name + description + parameter
names. Falls back to a literal substring match on the tool name when
BM25 returns no positive-score hits, which protects against
zero-IDF degenerate cases (e.g. searching `"github"` against a
catalog where every tool name contains "github").
- **Retrieval:** BM25 over tokenized tool name, source name (the MCP
server or plugin toolset the tool belongs to, so searching `"linear"`
finds that server's tools even when a tool's own name doesn't carry
the service), description, and parameter names. Falls back to a
literal substring match on the tool name when no query token matches
any document (e.g. searching `"hub"` where the token is `github`).
- **Parallel execution unwraps the bridge.** The batch planner decides
concurrency on the *underlying* tool of a `tool_call`, not on the
literal bridge name — so an MCP server opted in via
`supports_parallel_tool_calls: true` keeps its concurrency when its
tools are called through the bridge, and `tool_search` /
`tool_describe` lookups batch concurrently like any read-only tool.
- **Catalog is stateless across turns.** It rebuilds from the current
tool-defs list every assembly — no session-keyed `Map`. This avoids
the class of bug where a stored catalog drifts out of sync with the