diff --git a/agent/turn_context.py b/agent/turn_context.py index 8041eabdb7f..3e107e0900c 100644 --- a/agent/turn_context.py +++ b/agent/turn_context.py @@ -112,6 +112,23 @@ def build_turn_context( # Restore the primary runtime if the previous turn activated fallback. agent._restore_primary_runtime() + # 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 + # 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). + try: + 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) + except Exception: + logger.debug("between-turns MCP tool refresh skipped", exc_info=True) + # Sanitize surrogate characters from user input. if isinstance(user_message, str): user_message = sanitize_surrogates(user_message) diff --git a/tests/agent/test_turn_context.py b/tests/agent/test_turn_context.py index 52aef95ed96..c475c4fb145 100644 --- a/tests/agent/test_turn_context.py +++ b/tests/agent/test_turn_context.py @@ -47,6 +47,8 @@ class _FakeAgent: self.max_iterations = 90 self.tools = [] self.valid_tool_names = set() + self.enabled_toolsets = None + self.disabled_toolsets = None self.compression_enabled = False self.context_compressor = types.SimpleNamespace( protect_first_n=2, protect_last_n=2 @@ -185,3 +187,59 @@ def test_no_review_when_memory_disabled(): agent = _FakeAgent() ctx = _build(agent) assert ctx.should_review_memory is False + + +# ── Between-turns MCP refresh (cache-safe late-binding) ────────────────────── +# +# A slow MCP server that connects after the agent's build-time tool snapshot +# must become callable by the user's NEXT turn — without mutating an in-flight +# turn's cached request prefix. The prologue is exactly that boundary, so the +# refresh hook lives here. These assert the contract (R1/R2/R6 in the spec), +# not timing permutations. + + +def test_between_turns_refresh_adds_late_tool_when_servers_registered(): + """R1: a tool that registered since build lands in this turn's snapshot.""" + agent = _FakeAgent() + + new_def = {"type": "function", "function": {"name": "mcp_x_tool", "description": "", "parameters": {}}} + + import model_tools + with patch("tools.mcp_tool.has_registered_mcp_tools", return_value=True), \ + patch.object(model_tools, "get_tool_definitions", return_value=[new_def]): + _build(agent) + + assert "mcp_x_tool" in agent.valid_tool_names + assert any(t["function"]["name"] == "mcp_x_tool" for t in agent.tools) + + +def test_between_turns_refresh_skipped_when_no_servers(): + """R6: the common case (no MCP servers) never walks the registry.""" + agent = _FakeAgent() + import model_tools + + with patch("tools.mcp_tool.has_registered_mcp_tools", return_value=False), \ + patch.object(model_tools, "get_tool_definitions") as gtd: + _build(agent) + + gtd.assert_not_called() + + +def test_between_turns_refresh_no_churn_when_unchanged(): + """R2: an unchanged tool set leaves the snapshot object identity intact + (no needless swap → nothing for the next request prefix to diff against).""" + agent = _FakeAgent() + same = [{"type": "function", "function": {"name": "a", "description": "", "parameters": {}}}] + agent.tools = same + agent.valid_tool_names = {"a"} + + import model_tools + with patch("tools.mcp_tool.has_registered_mcp_tools", return_value=True), \ + patch.object( + model_tools, "get_tool_definitions", + return_value=[{"type": "function", "function": {"name": "a", "description": "", "parameters": {}}}], + ): + _build(agent) + + assert agent.tools is same # not replaced → no churn + diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index 5858d34cfb3..1611d8394f6 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -4236,6 +4236,17 @@ def probe_mcp_server_tools() -> Dict[str, List[tuple]]: _agent_tools_lock = threading.Lock() +def has_registered_mcp_tools() -> bool: + """True if any MCP server currently has tools registered in the registry. + + Cheap — checks the live server map under ``_lock``, no registry walk. Used + by the per-turn refresh hook so a session with no MCP servers configured + (the common case) skips the ``get_tool_definitions`` rebuild entirely. + """ + with _lock: + return bool(_servers) + + def refresh_agent_mcp_tools(agent, *, quiet_mode: bool = True) -> set: """Re-derive an already-built agent's tool snapshot from the live registry.