diff --git a/agent/tool_executor.py b/agent/tool_executor.py index 8500baf48e..962ce4df09 100644 --- a/agent/tool_executor.py +++ b/agent/tool_executor.py @@ -360,34 +360,23 @@ def _append_skipped_tool_results( def _tool_search_scoped_names(agent) -> frozenset: """Deferrable tool names the session may invoke via ``tool_call``; the unwrap bypasses the bridge's scope check in ``model_tools.handle_function_call``, so restricted sessions - validate against this set. Cached on the agent, keyed by registry scope/generation.""" + validate against this set. Read through the ``get_tool_definitions`` memo on every call + and never copied onto the agent: the memo expires with the check_fn verdicts it was + filtered on, a per-agent copy keyed on registry generation would not, and a recovered + MCP server's tools would stay uncallable for the agent's lifetime.""" try: import model_tools from tools import tool_search as _ts - from tools.registry import registry as _registry except Exception: return frozenset() - - enabled = getattr(agent, "enabled_toolsets", None) - disabled = getattr(agent, "disabled_toolsets", None) - cache_key = ( - _registry.current_scope_key(), - getattr(_registry, "_generation", 0), - frozenset(enabled) if enabled is not None else None, - frozenset(disabled) if disabled is not None else None, - ) - cached = getattr(agent, "_tool_search_scope_cache", None) - if cached is not None and cached[0] == cache_key: - return cached[1] try: - names = _ts.scoped_deferrable_names(model_tools.get_tool_definitions( - enabled_toolsets=enabled, disabled_toolsets=disabled, quiet_mode=True, skip_tool_search_assembly=True, + return _ts.scoped_deferrable_names(model_tools.get_tool_definitions( + enabled_toolsets=getattr(agent, "enabled_toolsets", None), + disabled_toolsets=getattr(agent, "disabled_toolsets", None), + quiet_mode=True, skip_tool_search_assembly=True, ) or []) except Exception: - names = frozenset() - with contextlib.suppress(Exception): - agent._tool_search_scope_cache = (cache_key, names) - return names + return frozenset() def _canonical_tool_name(function_name: str) -> str: diff --git a/model_tools.py b/model_tools.py index 5010146333..b9864a454b 100644 --- a/model_tools.py +++ b/model_tools.py @@ -18,7 +18,9 @@ import threading import time from typing import Dict, Any, List, Optional, Tuple -from tools.registry import CHECK_FN_CACHE_BYPASS, check_fn_cache_scope, discover_builtin_tools, registry, tool_error +from tools.registry import ( + CHECK_FN_CACHE_BYPASS, check_fn_cache_epoch, check_fn_cache_scope, check_fn_verdict_window, discover_builtin_tools, + registry, tool_error) from tools.registry import _MAX_TOOL_ERROR_CHARS as _TOOL_ERROR_MAX_LEN from toolsets import resolve_toolset, validate_toolset from tools.arg_coercion import coerce_tool_args @@ -190,9 +192,13 @@ _LEGACY_TOOLSET_MAP = { # Memo for get_tool_definitions(), active only with quiet_mode=True (the # non-quiet path prints). Hot callers (gateway runner, AIAgent.__init__) hit it # every turn; a miss costs ~7 ms of registry walk + check_fn probing. The key -# includes registry._generation (bumped on register/deregister/alias) so -# invalidation is transparent; check_fn drift is handled by registry.py's 30 s TTL. +# includes registry._generation (bumped on register/deregister/alias) and the +# check_fn cache epoch (bumped by invalidate_check_fn_cache). An entry expires with +# the earliest check_fn verdict it was filtered on (registry.check_fn_verdict_window): +# a hit never re-runs the checks, so without that expiry an MCP server that dropped +# and reconnected without re-registering would stay hidden for the life of the process. _tool_defs_cache: dict[tuple, list[dict[str, Any]]] = {} +_tool_defs_expiry: dict[tuple, float] = {} # cache key -> time.monotonic() deadline _tool_defs_cache_lock = threading.Lock() # FIFO cap: 8 covers a long-lived gateway's warm set of platform/toolset combos. # Hard cap on memoized get_tool_definitions() results. A long-lived Gateway process sees many distinct @@ -207,6 +213,15 @@ def _clear_tool_defs_cache() -> None: """Drop memoized results when a dynamic-schema dependency changes (discord caps, sandbox mode).""" with _tool_defs_cache_lock: _tool_defs_cache.clear() + _tool_defs_expiry.clear() + + +def _live_tool_defs_entry(cache_key: tuple) -> Optional[list[dict[str, Any]]]: + """The memoized list for *cache_key* unless missing or past its verdict expiry. Caller holds the lock.""" + cached = _tool_defs_cache.get(cache_key) + if cached is None or time.monotonic() >= _tool_defs_expiry.get(cache_key, float("inf")): + return None + return cached def get_tool_definitions(enabled_toolsets: Optional[list[str]] = None, disabled_toolsets: Optional[list[str]] = None, @@ -232,17 +247,21 @@ def get_tool_definitions(enabled_toolsets: Optional[list[str]] = None, disabled_ # Gateway process doesn't accumulate entries unboundedly across the many distinct toolset/config # fingerprints it sees over its lifetime (#19251). with _tool_defs_cache_lock: - cached = _tool_defs_cache.get(cache_key) if cache_key is not None else None + cached = _live_tool_defs_entry(cache_key) if cache_key is not None else None if cached is None: - result = compute() + with check_fn_verdict_window() as verdict_expiry: + result = compute() if cache_key is None: return list(result) with _tool_defs_cache_lock: - cached = _tool_defs_cache.get(cache_key) # another thread may have filled it meanwhile + cached = _live_tool_defs_entry(cache_key) # another thread may have filled it meanwhile if cached is None: - if len(_tool_defs_cache) >= _TOOL_DEFS_CACHE_MAX: - _tool_defs_cache.pop(next(iter(_tool_defs_cache))) + if cache_key not in _tool_defs_cache and len(_tool_defs_cache) >= _TOOL_DEFS_CACHE_MAX: + evicted = next(iter(_tool_defs_cache)) + _tool_defs_cache.pop(evicted) + _tool_defs_expiry.pop(evicted, None) _tool_defs_cache[cache_key] = cached = result + _tool_defs_expiry[cache_key] = verdict_expiry[0] else: global _last_resolved_tool_names _last_resolved_tool_names = [t["function"]["name"] for t in cached] @@ -257,8 +276,9 @@ def _tool_defs_cache_key( """Memo key for get_tool_definitions, or None when caching must be bypassed. Covers every argument plus everything that changes the result without one: - registry generation, config.yaml stat signature (dynamic schemas), kanban - context, profile scope. check_fn results are TTL-cached in the registry. + registry generation, check_fn cache epoch, config.yaml stat signature (dynamic + schemas), kanban context, profile scope. check_fn verdicts expire the entry + itself (see get_tool_definitions). """ profile_scope = check_fn_cache_scope() if profile_scope == CHECK_FN_CACHE_BYPASS: @@ -273,7 +293,7 @@ def _tool_defs_cache_key( registry.current_scope_key(), frozenset(enabled_toolsets) if enabled_toolsets is not None else None, frozenset(disabled_toolsets) if disabled_toolsets else None, registry._generation, cfg_fp, bool(os.environ.get("HERMES_KANBAN_TASK")), bool(skip_tool_search_assembly), - _is_delegated_child_context(), _is_dispatcher_owned_worker(), profile_scope, + _is_delegated_child_context(), _is_dispatcher_owned_worker(), profile_scope, check_fn_cache_epoch(), ) diff --git a/tests/tools/test_tool_defs_memo_availability.py b/tests/tools/test_tool_defs_memo_availability.py new file mode 100644 index 0000000000..7ba861a3c7 --- /dev/null +++ b/tests/tools/test_tool_defs_memo_availability.py @@ -0,0 +1,208 @@ +"""The quiet-mode ``get_tool_definitions`` memo must not outlive the check_fn verdicts it filtered on. + +An MCP server whose connection dropped (check_fn False) and came back without re-registering its +tools left the memo holding the omission: ``tool_search`` / ``tool_describe`` read the memo, so the +recovered server stayed invisible for the life of the process, ``/new`` included. +""" +import json + +import pytest + +import model_tools +from model_tools import get_tool_definitions +from tests.module_clock import install_module_clock +from tools import registry as registry_mod +from tools.mcp_tool_agent import reprobe_tool_availability +from tools.registry import invalidate_check_fn_cache, registry + +MAIL = "mcp__probe_mail__list_messages" +ANCHOR = "mcp__probe_anchor__ping" +BRIDGE_KW = {"quiet_mode": True, "skip_tool_search_assembly": True} +BOTH = ["mcp-probe-mail", "mcp-probe-anchor"] + + +def _schema(name, description): + return {"name": name, "description": description, + "parameters": {"type": "object", "properties": {}}} + + +def _names(defs): + return {d["function"]["name"] for d in defs} + + +@pytest.fixture +def mail_state(): + """A recovering MCP-style tool (mail) next to an always-healthy one (anchor).""" + state = {"up": False, "probes": 0} + + def mail_connected(): + state["probes"] += 1 + return state["up"] + + registry.register(name=MAIL, toolset="mcp-probe-mail", + schema=_schema(MAIL, "List messages in the probe mailbox"), + handler=lambda args, **kw: "MAIL_OK", check_fn=mail_connected) + registry.register(name=ANCHOR, toolset="mcp-probe-anchor", + schema=_schema(ANCHOR, "Ping the anchor server"), + handler=lambda args, **kw: "ANCHOR_OK", check_fn=lambda: True) + reprobe_tool_availability() + yield state + registry.deregister(MAIL) + registry.deregister(ANCHOR) + reprobe_tool_availability() + + +@pytest.fixture +def clock(monkeypatch): + """One fake monotonic clock for the verdict cache and the definitions memo.""" + module_clock = install_module_clock(monkeypatch, registry_mod) + monkeypatch.setattr(model_tools, "time", module_clock) + return module_clock + + +def test_reporter_reproducer_explicit_invalidation_reaches_quiet_memo(mail_state): + kwargs = {"enabled_toolsets": ["mcp-probe-mail"], **BRIDGE_KW} + assert get_tool_definitions(**kwargs) == [] + + mail_state["up"] = True + invalidate_check_fn_cache() # e.g. hermes_platform.declaration.on_change + + assert _names(get_tool_definitions(**kwargs)) == {MAIL} + + +def test_false_to_true_recovers_on_natural_ttl_expiry(mail_state, clock): + kwargs = {"enabled_toolsets": BOTH, **BRIDGE_KW} + assert _names(get_tool_definitions(**kwargs)) == {ANCHOR} + + mail_state["up"] = True # the server reconnected; nothing re-registered, nothing invalidated + clock.now += registry_mod._CHECK_FN_TTL_SECONDS - 1 + assert _names(get_tool_definitions(**kwargs)) == {ANCHOR} # the False verdict is still within its TTL + + clock.now += 2 + assert _names(get_tool_definitions(**kwargs)) == {ANCHOR, MAIL} + + +def test_memo_hits_within_verdict_ttl(mail_state, clock, monkeypatch): + computes = [] + real_compute = model_tools._compute_tool_definitions + + def counting_compute(*args, **kwargs): + computes.append(1) + return real_compute(*args, **kwargs) + + monkeypatch.setattr(model_tools, "_compute_tool_definitions", counting_compute) + kwargs = {"enabled_toolsets": BOTH, **BRIDGE_KW} + for _ in range(5): + get_tool_definitions(**kwargs) + clock.now += 5 + assert len(computes) == 1 + assert mail_state["probes"] == 1 + + clock.now += registry_mod._CHECK_FN_TTL_SECONDS + get_tool_definitions(**kwargs) + get_tool_definitions(**kwargs) + assert len(computes) == 2 + assert mail_state["probes"] == 2 + + +def test_true_to_false_respects_ttl_then_failure_grace(mail_state, clock): + ttl, grace = registry_mod._CHECK_FN_TTL_SECONDS, registry_mod._CHECK_FN_FAILURE_GRACE_SECONDS + kwargs = {"enabled_toolsets": BOTH, **BRIDGE_KW} + mail_state["up"] = True + start = clock.now + assert MAIL in _names(get_tool_definitions(**kwargs)) + + mail_state["up"] = False # connection dropped + clock.now = start + ttl - 1 + assert MAIL in _names(get_tool_definitions(**kwargs)) # cached True verdict + assert mail_state["probes"] == 1 + + clock.now = start + ttl + 1 + assert MAIL in _names(get_tool_definitions(**kwargs)) # failure inside the grace window: a flake + clock.now += 1 + assert MAIL in _names(get_tool_definitions(**kwargs)) + assert mail_state["probes"] == 3 # the grace verdict is not memoized: every call re-probes + + clock.now = start + grace + 1 + assert _names(get_tool_definitions(**kwargs)) == {ANCHOR} # grace over: the failure is honored + + +def test_restricted_toolset_never_gains_the_recovered_tool(mail_state, clock): + restricted = {"enabled_toolsets": ["mcp-probe-anchor"], **BRIDGE_KW} + assert _names(get_tool_definitions(**restricted)) == {ANCHOR} + mail_state["up"] = True + clock.now += registry_mod._CHECK_FN_TTL_SECONDS + 1 + invalidate_check_fn_cache() + assert _names(get_tool_definitions(**restricted)) == {ANCHOR} + result, redispatch = model_tools._dispatch_bridge_tool( + "tool_call", {"calls": [{"name": MAIL, "arguments": {}}]}, ["mcp-probe-anchor"], None) + assert redispatch is None and "error" in json.loads(result) + + +def test_profile_scoped_verdicts_stay_separate(mail_state, monkeypatch): + profile = {"current": "/profiles/a"} + verdicts = {"/profiles/a": True, "/profiles/b": False} + mail_state["up"] = True + for module in (registry_mod, model_tools): + monkeypatch.setattr(module, "check_fn_cache_scope", lambda: profile["current"]) + registry.register(name=MAIL, toolset="mcp-probe-mail", schema=_schema(MAIL, "List messages"), + handler=lambda args, **kw: "MAIL_OK", check_fn=lambda: verdicts[profile["current"]]) + kwargs = {"enabled_toolsets": BOTH, **BRIDGE_KW} + assert MAIL in _names(get_tool_definitions(**kwargs)) + profile["current"] = "/profiles/b" + assert _names(get_tool_definitions(**kwargs)) == {ANCHOR} + profile["current"] = "/profiles/a" + assert MAIL in _names(get_tool_definitions(**kwargs)) + + monkeypatch.setattr(model_tools, "check_fn_cache_scope", lambda: registry_mod.CHECK_FN_CACHE_BYPASS) + model_tools._clear_tool_defs_cache() + get_tool_definitions(**kwargs) + assert not model_tools._tool_defs_cache # unresolved profile identity still bypasses the memo + + +def _tool_results(server): + """``{tool_call_id: content}`` for every tool message the agent sent back.""" + out = {} + for body in server.main_requests(): + for msg in body.get("messages", []): + if msg.get("role") == "tool": + out[msg.get("tool_call_id")] = msg.get("content") + return out + + +def test_long_running_agent_bridge_recovers_after_mcp_reconnect(mail_state, clock): + """Composed: one AIAgent across turns, real bridge dispatch, the mail server recovering + between turns without re-registering (the reconnect path keeps its registered names).""" + from run_agent import AIAgent + from tests.fakes.fake_llm_provider import MODEL_ID, FakeLLMServer, Text, ToolCall + + script = [ + ToolCall("tool_call", {"calls": [{"name": ANCHOR, "arguments": {}}]}, + parallel=[("tool_describe", {"names": [MAIL]})]), + Text("turn one done"), + ToolCall("tool_search", {"queries": ["probe mailbox messages"]}, + parallel=[("tool_describe", {"names": [MAIL]})]), + ToolCall("tool_call", {"calls": [{"name": MAIL, "arguments": {}}]}), + Text("turn two done"), + ] + with FakeLLMServer(script) as server: + agent = AIAgent(model=MODEL_ID, provider="openai", api_key="unused-test-key", + base_url=server.base_url, enabled_toolsets=BOTH, quiet_mode=True, + skip_context_files=True, skip_memory=True, save_trajectories=False) + bridge = {(t.get("function") or {}).get("name") for t in agent.tools} + assert {"tool_search", "tool_describe", "tool_call"} <= bridge + + first = agent.run_conversation("turn one") + turn_one = list(_tool_results(server).values()) + assert any("ANCHOR_OK" in c for c in turn_one), turn_one + assert any(MAIL in c and "not_found" in c for c in turn_one), turn_one + + mail_state["up"] = True + clock.now += registry_mod._CHECK_FN_TTL_SECONDS + 1 # the False verdict has expired + agent.run_conversation("turn two", conversation_history=first.get("messages")) + turn_two = list(_tool_results(server).values())[len(turn_one):] + + search, describe, call = turn_two + assert MAIL in search, search + assert "not_found" not in describe and "List messages in the probe mailbox" in describe, describe + assert "MAIL_OK" in call, call diff --git a/tools/registry.py b/tools/registry.py index f272b26bc0..0a8be8d0be 100644 --- a/tools/registry.py +++ b/tools/registry.py @@ -267,6 +267,10 @@ _check_fn_core_drop_warned: set[tuple[Callable, Optional[str]]] = set() # once- _check_fn_cache_lock = threading.Lock() CHECK_FN_CACHE_BYPASS = "" _NO_CACHE_CHECK_FNS: set[Callable] = set() +# Bumped by invalidate_check_fn_cache(): anything derived from the dropped verdicts is void. +_check_fn_epoch = 0 +# Earliest expiry of the verdicts consulted inside check_fn_verdict_window() (None outside one). +_verdict_expiry: ContextVar[Optional[list]] = ContextVar("check_fn_verdict_expiry", default=None) def no_cache_check_fn(fn: Callable) -> Callable: @@ -275,6 +279,32 @@ def no_cache_check_fn(fn: Callable) -> Callable: return fn +def check_fn_cache_epoch() -> int: + """Changes whenever ``invalidate_check_fn_cache()`` drops the cached verdicts.""" + return _check_fn_epoch + + +@contextmanager +def check_fn_verdict_window(): + """Yield a one-item list holding the earliest expiry (``time.monotonic``) of the cached + ``check_fn`` verdicts consulted inside the block, ``inf`` when none was. A result filtered on + those verdicts must not be reused past it: reusing it would replay a verdict this cache has + already expired. A grace-served True expires at once, because the registry re-probes it on the + next call. ``no_cache_check_fn`` probes are config-backed and set no expiry.""" + sink = [float("inf")] + token = _verdict_expiry.set(sink) + try: + yield sink + finally: + _verdict_expiry.reset(token) + + +def _note_verdict_expiry(deadline: float) -> None: + sink = _verdict_expiry.get() + if sink is not None and deadline < sink[0]: + sink[0] = deadline + + def _fn_label(fn: Callable) -> object: return getattr(fn, "__qualname__", fn) @@ -352,6 +382,7 @@ def _check_fn_cached(fn: Callable) -> bool: _prune_check_fn_caches(now) # leaves only entries within TTL cached = _check_fn_cache.get(cache_key) if cached is not None: + _note_verdict_expiry(cached[0] + _CHECK_FN_TTL_SECONDS) return cached[1] exc_info = None try: @@ -369,10 +400,12 @@ def _check_fn_cached(fn: Callable) -> bool: _check_fn_last_good[cache_key] = now _check_fn_ever_good.add(cache_key) _check_fn_cache[cache_key] = (now, True) + _note_verdict_expiry(now + _CHECK_FN_TTL_SECONDS) return True last_good = _check_fn_last_good.get(cache_key) if last_good is not None and now - last_good < _CHECK_FN_FAILURE_GRACE_SECONDS: # Recent success → flake: serve last-good True, do NOT cache (next call re-probes). + _note_verdict_expiry(now) logger.warning( "check_fn %s failed (%s) within %.0fs of last success; " "treating as transient and keeping tool(s) available", @@ -399,6 +432,7 @@ def _check_fn_cached(fn: Callable) -> bool: "check_fn %s %s; dependent tools will be unavailable this turn", _fn_label(fn), outcome, exc_info=exc_info) _check_fn_cache[cache_key] = (now, False) + _note_verdict_expiry(now + _CHECK_FN_TTL_SECONDS) return False @@ -421,7 +455,9 @@ def _memo_check(fn: Callable, memo: dict[Callable, bool]) -> bool: def invalidate_check_fn_cache() -> None: """Drop all cached ``check_fn`` results (after config changes like ``hermes tools enable``).""" + global _check_fn_epoch with _check_fn_cache_lock: + _check_fn_epoch += 1 _check_fn_cache.clear() _check_fn_last_good.clear() _check_fn_ever_good.clear()