diff --git a/agent/turn_context.py b/agent/turn_context.py index bff4a8f094..691b1de9c0 100644 --- a/agent/turn_context.py +++ b/agent/turn_context.py @@ -635,13 +635,18 @@ def build_turn_context( # Between-turns MCP refresh: an MCP server that finished connecting since # the previous turn (slow HTTP/OAuth servers routinely take 2-6s on a cold # connect, missing the bounded startup wait) lands in THIS turn's tool - # snapshot. This is cache-safe by construction: it runs in the per-turn + # snapshot. Timing is cache-safe by construction: it runs in the per-turn # prologue, before this turn's first API call assembles ``tools=``, so it - # only ever extends a fresh request prefix — it never mutates the cached - # prefix of an in-flight turn. No-op when no MCP servers are registered - # (the common case, gated by the cheap ``has_registered_mcp_tools`` check) - # or when the tool set is unchanged (``refresh_agent_mcp_tools`` diffs by - # name and leaves the snapshot untouched on no-change). + # never mutates the prefix of an in-flight turn. ``preserve_prefix`` makes + # the *content* cache-safe too (#100336): a plain rebuild re-derives the + # array from live availability, so a flapping ``check_fn`` silently drops a + # tool and a late arrival splices into sorted position — either one forks + # the tool block and re-prefills the whole history behind it, every turn it + # happens. With the flag the live order is authoritative and the array + # only ever grows. No-op when no MCP servers are registered (the common + # case, gated by the cheap ``has_registered_mcp_tools`` check) or when the + # tool set is unchanged (``refresh_agent_mcp_tools`` diffs by name and + # leaves the snapshot untouched on no-change). try: if not getattr(agent, "_skip_mcp_refresh", False): # Import-cost gate: ``tools.mcp_tool`` pulls in the whole ``mcp`` @@ -656,7 +661,9 @@ def build_turn_context( if "tools.mcp_tool" in _sys.modules: from tools.mcp_tool import has_registered_mcp_tools, refresh_agent_mcp_tools if has_registered_mcp_tools(): - refresh_agent_mcp_tools(agent, quiet_mode=True) + refresh_agent_mcp_tools( + agent, quiet_mode=True, preserve_prefix=True, + ) except Exception: logger.debug("between-turns MCP tool refresh skipped", exc_info=True) diff --git a/tests/tools/test_refresh_agent_mcp_tools.py b/tests/tools/test_refresh_agent_mcp_tools.py index b1aa95e12f..22203d93af 100644 --- a/tests/tools/test_refresh_agent_mcp_tools.py +++ b/tests/tools/test_refresh_agent_mcp_tools.py @@ -224,3 +224,65 @@ def test_wait_returns_instantly_when_no_discovery_thread(monkeypatch): t0 = time.time() mcp_startup.wait_for_mcp_discovery() assert time.time() - t0 < 0.2 # never blocks on the bound when nothing's pending + + +# --------------------------------------------------------------------------- +# preserve_prefix: the tool array is a cached request prefix (#100336) +# --------------------------------------------------------------------------- + + +def _registered(monkeypatch, names): + """Make the registry report exactly *names* as still registered.""" + from tools import registry as registry_mod + + entries = [types.SimpleNamespace(name=n) for n in names] + monkeypatch.setattr( + registry_mod.registry, "get_all_entries", lambda: entries, raising=False + ) + + +def _serve(monkeypatch, defs): + import model_tools + + monkeypatch.setattr(model_tools, "get_tool_definitions", lambda **kw: list(defs)) + + +def test_preserve_prefix_carries_a_flapping_tool_forward(monkeypatch): + """A check_fn flip must not shrink a live session's tool prefix. + + ``browser_navigate``'s availability probe fails this turn (headless box, + expired credential, docker blip) so ``get_tool_definitions`` omits it. The + tool is still *registered* — only its probe flapped — so the snapshot must + keep it, byte-for-byte, instead of forking the cached prefix. + """ + agent = _agent(["read_file", "browser_navigate", "terminal"]) + before = list(agent.tools) + + _serve(monkeypatch, [_tool("read_file"), _tool("terminal")]) + _registered(monkeypatch, ["read_file", "browser_navigate", "terminal"]) + + added = mcp_tool.refresh_agent_mcp_tools(agent, preserve_prefix=True) + + assert added == set() + assert agent.tools == before + assert "browser_navigate" in agent.valid_tool_names + + +def test_preserve_prefix_appends_late_arrivals_at_the_tail(monkeypatch): + """``get_definitions`` sorts by name, so a late tool can splice in at 0. + + Under ``preserve_prefix`` the live order is authoritative and the new tool + extends the array, leaving every earlier byte where the provider cached it. + """ + agent = _agent(["read_file", "terminal"]) + + # Sorted order would put the new tool first. + _serve(monkeypatch, [_tool("aaa_mcp_late"), _tool("read_file"), _tool("terminal")]) + _registered(monkeypatch, ["aaa_mcp_late", "read_file", "terminal"]) + + added = mcp_tool.refresh_agent_mcp_tools(agent, preserve_prefix=True) + + assert added == {"aaa_mcp_late"} + assert [t["function"]["name"] for t in agent.tools] == [ + "read_file", "terminal", "aaa_mcp_late", + ] diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index dab35f4777..cc69786304 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -8344,6 +8344,7 @@ def refresh_agent_mcp_tools( disabled_override=None, quiet_mode: bool = True, content_aware: bool = False, + preserve_prefix: bool = False, ) -> set: """Re-derive an already-built agent's tool snapshot from the live registry. @@ -8372,6 +8373,22 @@ def refresh_agent_mcp_tools( under ``_agent_tools_lock`` so a concurrent reader never sees a cross-attribute half-swap. + ``preserve_prefix`` is for the callers that rebuild inside a live + conversation (the between-turns prologue). There the tool array is a + cached request prefix: every provider that renders ``tools`` ahead of the + messages re-prefills the entire history behind any byte that moves. A + plain rebuild moves two kinds of bytes — it drops a tool whose ``check_fn`` + merely flapped (a headless browser probe, an expired credential, a docker + blip), and it splices a late-landing tool into sorted position, which can + be index 0. With ``preserve_prefix`` the live order is authoritative: + existing tools keep their slot (schemas still refresh), a tool that is + still *registered* but momentarily unavailable is carried forward, a tool + that genuinely left the registry is still dropped, and new tools are + appended at the tail so the prefix only ever grows. Carrying an + unavailable tool forward changes nothing about dispatch — ``check_fn`` + gates exposure at snapshot time, never invocation, and every handler + already owns its own unavailability error. + Returns the set of newly-added tool names (empty when nothing changed), so callers can decide whether to notify the user / re-emit session info. The caller owns the prompt-cache contract: this helper does NOT check turn state, @@ -8427,6 +8444,18 @@ def refresh_agent_mcp_tools( # this rebuild actually appended (matching agent_init's dedup-aware add). staged_engine_names = _reinject_post_build_tools(agent, new_defs, new_names) + # Snapshot registry membership OUTSIDE ``_agent_tools_lock`` — it is the + # only input ``preserve_prefix`` needs beyond the two tool lists, and + # taking ``registry._lock`` under the tools lock would be the first place + # in the process to nest those two. + registered_names: set = set() + if preserve_prefix: + try: + registered_names = {entry.name for entry in registry.get_all_entries()} + except Exception: # noqa: BLE001 + # Fail open to the plain rebuild rather than pinning a stale list. + preserve_prefix = False + # Single atomic read-diff-publish so the returned ``added`` is consistent # with what was actually published, even under concurrent callers, and a # stale (older-generation) rebuild can't overwrite a newer published one. @@ -8440,10 +8469,12 @@ def refresh_agent_mcp_tools( if snapshot_generation < published_gen: # A newer snapshot already won; our set is stale — drop it. return set() - current = { - t["function"]["name"] - for t in (getattr(agent, "tools", None) or []) - } + current_defs = list(getattr(agent, "tools", None) or []) + current = {t["function"]["name"] for t in current_defs} + if preserve_prefix: + new_defs, new_names = _merge_preserving_prefix( + current_defs, new_defs, registered_names, + ) if new_names == current: # Same NAME set. For MCP-reload callers that is "no change" — # leave the live snapshot untouched (no churn). Content-aware @@ -8481,6 +8512,40 @@ def refresh_agent_mcp_tools( return new_names - current +def _merge_preserving_prefix( + current_defs: list, new_defs: list, registered_names: set, +) -> tuple[list, set]: + """Fold a fresh tool snapshot into a live one without moving existing bytes. + + The live tool array is a cached request prefix, so the merge is ordered by + ``current_defs``, not by the fresh list: + + * a name in both keeps its slot and takes the fresh schema (dynamic + overrides — delegate_task limits, execute_code stubs — still land); + * a name only in the live list is carried forward when it is still + registered (its ``check_fn`` flapped) and dropped when it is not (the + MCP server or plugin genuinely went away); + * a name only in the fresh list is appended at the tail, so a late-landing + MCP tool extends the prefix instead of splicing into sorted position. + """ + fresh = {} + for entry in new_defs: + name = (entry.get("function") or {}).get("name", "") + if name: + fresh[name] = entry + + merged = [] + for entry in current_defs: + name = (entry.get("function") or {}).get("name", "") + replacement = fresh.pop(name, None) + if replacement is not None: + merged.append(replacement) + elif name and name in registered_names: + merged.append(entry) + merged.extend(fresh.values()) + return merged, {(t.get("function") or {}).get("name", "") for t in merged} + + def _reinject_post_build_tools(agent, tools_list: list, name_set: set) -> set: """Append memory-provider and context-engine tools onto staged locals.