diff --git a/openhands-agent-server/openhands/agent_server/mcp_oauth_store.py b/openhands-agent-server/openhands/agent_server/mcp_oauth_store.py index 56f6708722..609e80201a 100644 --- a/openhands-agent-server/openhands/agent_server/mcp_oauth_store.py +++ b/openhands-agent-server/openhands/agent_server/mcp_oauth_store.py @@ -24,7 +24,10 @@ MCPOAuthTokenStorageField, MCPServer, ) -from openhands.sdk.mcp.utils import ToolsChangedCallback, create_mcp_tools +from openhands.sdk.mcp.utils import ( + ToolsChangedCallback, + create_mcp_tools, +) logger = get_logger(__name__) diff --git a/openhands-sdk/openhands/sdk/agent/base.py b/openhands-sdk/openhands/sdk/agent/base.py index 623dee7b32..c2d210235a 100644 --- a/openhands-sdk/openhands/sdk/agent/base.py +++ b/openhands-sdk/openhands/sdk/agent/base.py @@ -3,6 +3,7 @@ import os import re import sys +import threading from abc import ABC, abstractmethod from collections import Counter from collections.abc import Generator, Iterable, Sequence @@ -27,8 +28,9 @@ from openhands.sdk.llm import LLM from openhands.sdk.llm.utils.model_prompt_spec import get_model_prompt_spec from openhands.sdk.logger import get_logger +from openhands.sdk.mcp.client import MCPClient from openhands.sdk.mcp.config import MCPServer -from openhands.sdk.mcp.tool import MCPToolExecutor +from openhands.sdk.mcp.tool import MCPToolDefinition, MCPToolExecutor from openhands.sdk.tool import ( BUILT_IN_TOOL_CLASSES, BUILT_IN_TOOLS, @@ -300,6 +302,7 @@ def _validate_system_prompt_fields(cls, data: Any) -> Any: # Runtime materialized tools; private and non-serializable _tools: dict[str, ToolDefinition] = PrivateAttr(default_factory=dict) + _tools_lock: threading.RLock = PrivateAttr(default_factory=threading.RLock) _initialized: bool = PrivateAttr(default=False) @property @@ -561,6 +564,7 @@ def _initialize( if self.filter_tools_regex: pattern = re.compile(self.filter_tools_regex) tools = [tool for tool in tools if pattern.match(tool.name)] + tool_names = [tool.name for tool in tools] logger.info("Filtered to %d tools after applying regex filter", len(tools)) # Include default tools from include_default_tools; not subject to regex @@ -867,13 +871,14 @@ def add_runtime_tools(self, tools: Sequence[ToolDefinition]) -> None: name for name, count in Counter(tool_names).items() if count > 1 } raise ValueError(f"Duplicate runtime tool names found: {duplicates}") - existing = set(self._tools) & set(tool_names) - if existing: - raise ValueError(f"Duplicate tool names found: {existing}") + with self._tools_lock: + existing = set(self._tools) & set(tool_names) + if existing: + raise ValueError(f"Duplicate tool names found: {existing}") - # AgentBase is frozen, so update its mutable tool map in place. - for tool in tools: - self._tools[tool.name] = tool + # AgentBase is frozen, so update its mutable tool map in place. + for tool in tools: + self._tools[tool.name] = tool def _on_mcp_tools_changed(self, tools: Sequence[ToolDefinition]) -> None: """Handle dynamically advertised MCP tools. @@ -897,35 +902,36 @@ def _on_mcp_tools_changed(self, tools: Sequence[ToolDefinition]) -> None: } raise ValueError(f"Duplicate MCP tool names found: {duplicates}") - additions: list[ToolDefinition] = [] - replacements: list[ToolDefinition] = [] - conflicts: set[str] = set() - for tool in tools: - existing = self._tools.get(tool.name) - if existing is None: - additions.append(tool) - continue - - existing_executor = existing.executor - replacement_executor = tool.executor - if ( - isinstance(existing_executor, MCPToolExecutor) - and isinstance(replacement_executor, MCPToolExecutor) - and existing_executor.client is replacement_executor.client - ): - replacements.append(tool) - else: - conflicts.add(tool.name) - - if conflicts: - raise ValueError( - "Dynamically advertised MCP tools conflict with existing runtime " - f"tools: {sorted(conflicts)}" - ) + with self._tools_lock: + additions: list[ToolDefinition] = [] + replacements: list[ToolDefinition] = [] + conflicts: set[str] = set() + for tool in tools: + existing = self._tools.get(tool.name) + if existing is None: + additions.append(tool) + continue + + existing_executor = existing.executor + replacement_executor = tool.executor + if ( + isinstance(existing_executor, MCPToolExecutor) + and isinstance(replacement_executor, MCPToolExecutor) + and existing_executor.client is replacement_executor.client + ): + replacements.append(tool) + else: + conflicts.add(tool.name) + + if conflicts: + raise ValueError( + "Dynamically advertised MCP tools conflict with existing runtime " + f"tools: {sorted(conflicts)}" + ) - self.add_runtime_tools(additions) - for tool in replacements: - self._tools[tool.name] = tool + self.add_runtime_tools(additions) + for tool in replacements: + self._tools[tool.name] = tool if additions: logger.info( @@ -940,6 +946,46 @@ def _on_mcp_tools_changed(self, tools: Sequence[ToolDefinition]) -> None: ", ".join(tool.name for tool in replacements), ) + def _on_mcp_tools_reconciled( + self, + client: MCPClient, + tools: Sequence[MCPToolDefinition], + ) -> None: + """Replace this MCP client's tools with its current server snapshot.""" + tool_names = [tool.name for tool in tools] + if len(tool_names) != len(set(tool_names)): + duplicates = { + name for name, count in Counter(tool_names).items() if count > 1 + } + raise ValueError(f"Duplicate MCP tool names found: {duplicates}") + + if self.filter_tools_regex: + pattern = re.compile(self.filter_tools_regex) + tools = [tool for tool in tools if pattern.match(tool.name)] + tool_names = [tool.name for tool in tools] + + with self._tools_lock: + owned_names = { + name + for name, tool in self._tools.items() + if isinstance(tool.executor, MCPToolExecutor) + and tool.executor.client is client + } + conflicts = (set(tool_names) & set(self._tools)) - owned_names + if conflicts: + raise ValueError( + "Dynamically advertised MCP tools conflict with existing runtime " + f"tools: {sorted(conflicts)}" + ) + + reconciled = { + name: tool + for name, tool in self._tools.items() + if name not in owned_names + } + reconciled.update((tool.name, tool) for tool in tools) + object.__setattr__(self, "_tools", reconciled) + @property def tools_map(self) -> dict[str, ToolDefinition]: """Get the initialized tools map. @@ -949,7 +995,8 @@ def tools_map(self) -> dict[str, ToolDefinition]: if not self._initialized: raise RuntimeError("Agent not initialized; call _initialize() before use") # Isolate readers from background MCP tool updates. - return dict(self._tools) + with self._tools_lock: + return dict(self._tools) # -- Capability helpers ----------------------------------------------- # Downstream code should branch on these properties rather than doing diff --git a/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py b/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py index e2c9ff7cc1..02c55f0523 100644 --- a/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py +++ b/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py @@ -59,16 +59,19 @@ from openhands.sdk.llm.llm_registry import LLMRegistry from openhands.sdk.logger import get_logger from openhands.sdk.marketplace.registry import MarketplaceRegistry +from openhands.sdk.mcp.client import MCPClient from openhands.sdk.mcp.config import ( MCPServer, coerce_mcp_config, dump_mcp_config, enabled_mcp_servers, ) +from openhands.sdk.mcp.tool import MCPToolDefinition from openhands.sdk.mcp.utils import ( DefaultMCPToolProvider, MCPToolProvider, ToolsChangedCallback, + ToolsReconciledCallback, ) from openhands.sdk.observability.laminar import observe from openhands.sdk.plugin import ( @@ -1277,6 +1280,7 @@ def _runtime_mcp_tools( mcp_config: dict[str, MCPServer], *, on_tools_changed: ToolsChangedCallback | None = None, + on_tools_reconciled: ToolsReconciledCallback | None = None, ) -> list[ToolDefinition]: # Servers the user switched off stay in the settings map but must not # be connected to. Filter before the emptiness check so an all-disabled @@ -1289,14 +1293,23 @@ def _runtime_mcp_tools( _RUNTIME_MCP_TIMEOUT_SECS, on_tools_changed=on_tools_changed, ) + client._tools_reconciled_callback = on_tools_reconciled return list(client.tools) + def _on_mcp_tools_reconciled( + self, + client: MCPClient, + tools: Sequence[MCPToolDefinition], + ) -> None: + self.agent._on_mcp_tools_reconciled(client, tools) + def _runtime_mcp_tools_for_agent(self) -> list[ToolDefinition]: if not self.agent.supports_openhands_tools or not self.agent.mcp_config: return [] return self._runtime_mcp_tools( self.agent.mcp_config, - on_tools_changed=self.agent._on_mcp_tools_changed, + on_tools_changed=lambda tools: self.agent._on_mcp_tools_changed(tools), + on_tools_reconciled=self._on_mcp_tools_reconciled, ) def _runtime_skill_tools_for_agent(self) -> list[ToolDefinition]: @@ -1366,7 +1379,13 @@ def load_plugin(self, plugin_ref: str) -> None: ) merged_mcp = coerce_mcp_config(expanded_mcp["mcpServers"]) runtime_mcp_tools = ( - self._runtime_mcp_tools(runtime_plugin_mcp) if self._agent_ready else [] + self._runtime_mcp_tools( + runtime_plugin_mcp, + on_tools_changed=lambda tools: self.agent._on_mcp_tools_changed(tools), + on_tools_reconciled=self._on_mcp_tools_reconciled, + ) + if self._agent_ready + else [] ) with self._state: diff --git a/openhands-sdk/openhands/sdk/mcp/client.py b/openhands-sdk/openhands/sdk/mcp/client.py index aa9bd53b79..898e85e16f 100644 --- a/openhands-sdk/openhands/sdk/mcp/client.py +++ b/openhands-sdk/openhands/sdk/mcp/client.py @@ -2,7 +2,7 @@ import asyncio import inspect -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Sequence from typing import TYPE_CHECKING, Any from fastmcp import Client as AsyncMCPClient @@ -15,6 +15,12 @@ from openhands.sdk.mcp.tool import MCPToolDefinition +ToolsReconciledCallback = Callable[ + ["MCPClient", Sequence["MCPToolDefinition"]], + None, +] + + class MCPClient(AsyncMCPClient): """MCP client with sync helpers and lifecycle management. @@ -35,12 +41,14 @@ class MCPClient(AsyncMCPClient): _executor: AsyncExecutor _closed: bool _tools: "list[MCPToolDefinition]" + _tools_reconciled_callback: ToolsReconciledCallback | None def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._executor = AsyncExecutor() self._closed = False self._tools = [] + self._tools_reconciled_callback = None @property def tools(self) -> "list[MCPToolDefinition]": diff --git a/openhands-sdk/openhands/sdk/mcp/tool.py b/openhands-sdk/openhands/sdk/mcp/tool.py index 78fae7f115..f81b509d3d 100644 --- a/openhands-sdk/openhands/sdk/mcp/tool.py +++ b/openhands-sdk/openhands/sdk/mcp/tool.py @@ -1,6 +1,7 @@ """Utility functions for MCP integration.""" import copy +import json import re from collections.abc import Sequence from typing import TYPE_CHECKING, Any @@ -195,7 +196,7 @@ def close(self) -> None: self.client.sync_close() -_mcp_dynamic_action_type: dict[str, type[Schema]] = {} +_mcp_dynamic_action_type: dict[tuple[str, str], type[Schema]] = {} def _create_mcp_action_type(action_type: mcp.types.Tool) -> type[Schema]: @@ -213,14 +214,17 @@ def _create_mcp_action_type(action_type: mcp.types.Tool) -> type[Schema]: to openai tool schema. """ - # Tool.name should be unique, so we can cache the created types. - mcp_action_type = _mcp_dynamic_action_type.get(action_type.name) + cache_key = ( + action_type.name, + json.dumps(action_type.inputSchema, sort_keys=True, separators=(",", ":")), + ) + mcp_action_type = _mcp_dynamic_action_type.get(cache_key) if mcp_action_type: return mcp_action_type model_name = f"MCP{to_camel_case(action_type.name)}Action" mcp_action_type = Schema.from_mcp_schema(model_name, action_type.inputSchema) - _mcp_dynamic_action_type[action_type.name] = mcp_action_type + _mcp_dynamic_action_type[cache_key] = mcp_action_type return mcp_action_type diff --git a/openhands-sdk/openhands/sdk/mcp/utils.py b/openhands-sdk/openhands/sdk/mcp/utils.py index a7ab70ecfd..46a740935a 100644 --- a/openhands-sdk/openhands/sdk/mcp/utils.py +++ b/openhands-sdk/openhands/sdk/mcp/utils.py @@ -13,7 +13,7 @@ from key_value.aio.protocols import AsyncKeyValue from openhands.sdk.logger import get_logger -from openhands.sdk.mcp.client import MCPClient +from openhands.sdk.mcp.client import MCPClient, ToolsReconciledCallback from openhands.sdk.mcp.config import ( MCPOAuthAuthCredential, MCPOAuthAuthentication, @@ -33,9 +33,7 @@ OAuth | None, ] -# Callback invoked when an MCP server signals that its tool list changed. -# Receives the *newly added* tool definitions; removed tools are dropped from -# the owning client's tool list but are not reported here. +# Backward-compatible callback that reports only newly added tools. ToolsChangedCallback = Callable[[Sequence[MCPToolDefinition]], None] @@ -174,15 +172,15 @@ async def _connect_and_list_tools(client: MCPClient) -> None: async def _refresh_tools( client: MCPClient, on_tools_changed: ToolsChangedCallback | None = None, + on_tools_reconciled: ToolsReconciledCallback | None = None, ) -> None: """Re-list tools from the server and reconcile ``client._tools``. Called after the initial connection and whenever the server sends a ``notifications/tools/list_changed`` notification. When an - ``on_tools_changed`` callback is supplied, newly discovered tools are - reported so a running agent can register them via ``add_runtime_tools``. - Tools that are no longer advertised are dropped from ``client._tools`` but - are not proactively removed from an agent's tool map. + ``on_tools_changed`` preserves the original additions-only callback contract. + ``on_tools_reconciled`` receives the complete current snapshot so a running + agent can add, replace, and remove tools owned by this client. """ mcp_type_tools: list[mcp.types.Tool] = await client.list_tools() existing_by_name = {tool.name: tool for tool in client._tools} @@ -190,16 +188,18 @@ async def _refresh_tools( reconciled: list[MCPToolDefinition] = [] added: list[MCPToolDefinition] = [] + updated: list[MCPToolDefinition] = [] for mcp_tool in mcp_type_tools: prior = existing_by_name.get(mcp_tool.name) - if prior is not None: - # Preserve the existing definition so its executor (and the - # shared MCPClient it closes on shutdown) stays wired up. + if prior is not None and prior.mcp_tool == mcp_tool: reconciled.append(prior) continue tool_sequence = MCPToolDefinition.create(mcp_tool=mcp_tool, mcp_client=client) reconciled.extend(tool_sequence) - added.extend(tool_sequence) + if prior is None: + added.extend(tool_sequence) + else: + updated.extend(tool_sequence) # Drop tools the server no longer advertises. Reassign atomically so # concurrent readers iterating client.tools never observe mid-update state. @@ -208,6 +208,11 @@ async def _refresh_tools( ] if removed: logger.info("MCP server removed tools: %s", ", ".join(sorted(removed))) + if updated: + logger.info( + "MCP server updated tools: %s", + ", ".join(sorted(tool.name for tool in updated)), + ) client._tools = reconciled if added and on_tools_changed is not None: @@ -220,6 +225,15 @@ async def _refresh_tools( exc_info=True, ) + if (added or updated or removed) and on_tools_reconciled is not None: + try: + on_tools_reconciled(client, reconciled) + except Exception: + logger.warning( + "on_tools_reconciled callback failed for MCP tool refresh", + exc_info=True, + ) + class _ToolListChangedHandler(MessageHandler): """Message handler that refreshes tools on ``tools/list_changed``. @@ -261,7 +275,11 @@ async def _refresh_tools(self) -> None: async with self._refresh_lock: if client._closed: return - await _refresh_tools(client, self._on_tools_changed) + await _refresh_tools( + client, + self._on_tools_changed, + client._tools_reconciled_callback, + ) except Exception: logger.warning( "Failed to refresh MCP tools after list_changed notification", @@ -274,6 +292,7 @@ def create_mcp_tools( timeout: float = 30.0, *, on_tools_changed: ToolsChangedCallback | None = None, + on_tools_reconciled: ToolsReconciledCallback | None = None, mcp_oauth_token_storage: AsyncKeyValue | None = None, mcp_oauth_factory: MCPOAuthFactory | None = None, ) -> MCPClient: @@ -289,9 +308,10 @@ def create_mcp_tools( The client subscribes to ``notifications/tools/list_changed`` and reconciles its tool list whenever the server signals a change. When ``on_tools_changed`` is provided, the client invokes it with newly added - tool definitions so progressive-disclosure servers can surface them to an - agent. The callback runs on the client's background event-loop thread, so - callers must ensure it is thread-safe (e.g. ``Agent.add_runtime_tools``). + tool definitions, preserving the original callback contract. When + ``on_tools_reconciled`` is provided, it receives the client and complete + current tool snapshot after additions, updates, or removals. Callbacks run + on the client's background event-loop thread and must be thread-safe. """ mcp_config = _require_native_mcp_config(mcp_config) requested = mcp_config @@ -313,6 +333,7 @@ def create_mcp_tools( ) client = MCPClient(config, log_handler=log_handler, message_handler=handler) handler._client = client + client._tools_reconciled_callback = on_tools_reconciled try: client.call_async_from_sync( diff --git a/tests/sdk/conversation/test_local_conversation_mcp.py b/tests/sdk/conversation/test_local_conversation_mcp.py index b45bb079d9..4a78ceccd2 100644 --- a/tests/sdk/conversation/test_local_conversation_mcp.py +++ b/tests/sdk/conversation/test_local_conversation_mcp.py @@ -3,19 +3,31 @@ from pathlib import Path from typing import Any, cast +import mcp.types as mcp_types from pydantic import SecretStr from openhands.sdk import LLM, Agent from openhands.sdk.conversation.impl.local_conversation import LocalConversation from openhands.sdk.mcp.client import MCPClient from openhands.sdk.mcp.config import MCPServer, coerce_mcp_config +from openhands.sdk.mcp.tool import MCPToolDefinition + + +class EmptyMCPClient: + def __init__(self) -> None: + self.tools: list[MCPToolDefinition] = [] + self._tools_reconciled_callback: Any = None + + def sync_close(self) -> None: + pass class RecordingMCPToolProvider: """Records every attempt to open an MCP connection.""" - def __init__(self) -> None: + def __init__(self, client: EmptyMCPClient | None = None) -> None: self.calls: list[dict[str, MCPServer]] = [] + self.client = client if client is not None else EmptyMCPClient() def create_tools( self, @@ -25,7 +37,7 @@ def create_tools( on_tools_changed: Any = None, ) -> MCPClient: self.calls.append(mcp_config) - return cast(MCPClient, type("EmptyMCPClient", (), {"tools": []})()) + return cast(MCPClient, self.client) def test_disabling_every_server_skips_the_mcp_connection(tmp_path: Path) -> None: @@ -47,3 +59,44 @@ def test_disabling_every_server_skips_the_mcp_connection(tmp_path: Path) -> None assert provider.calls == [] conversation.close() + + +def test_reconciliation_targets_replaced_agent(tmp_path: Path) -> None: + client = EmptyMCPClient() + initial = MCPToolDefinition.create( + mcp_tool=mcp_types.Tool( + name="initial", + description="initial", + inputSchema={"type": "object", "properties": {}}, + ), + mcp_client=cast(MCPClient, client), + )[0] + client.tools = [initial] + conversation = LocalConversation( + agent=Agent( + llm=LLM(model="test-model", api_key=SecretStr("test-key")), + tools=[], + include_default_tools=[], + mcp_config=coerce_mcp_config({"fake": {"command": "true"}}), + ), + workspace=str(tmp_path), + visualizer=None, + mcp_tool_provider=RecordingMCPToolProvider(client), + ) + conversation._ensure_agent_ready() + old_agent = conversation.agent + conversation.agent = old_agent.model_copy() + replacement = MCPToolDefinition.create( + mcp_tool=mcp_types.Tool( + name="replacement", + description="replacement", + inputSchema={"type": "object", "properties": {}}, + ), + mcp_client=cast(MCPClient, client), + )[0] + + client._tools_reconciled_callback(cast(MCPClient, client), [replacement]) + + assert set(conversation.agent.tools_map) == {"replacement"} + assert set(old_agent.tools_map) == {"initial"} + conversation.close() diff --git a/tests/sdk/conversation/test_local_conversation_plugins.py b/tests/sdk/conversation/test_local_conversation_plugins.py index 88c382e63c..4413eea924 100644 --- a/tests/sdk/conversation/test_local_conversation_plugins.py +++ b/tests/sdk/conversation/test_local_conversation_plugins.py @@ -794,6 +794,7 @@ class RuntimeOnlyTool(ThinkTool): class RuntimeMCPClient: def __init__(self): self.tools = [runtime_tool] + self._tools_reconciled_callback: Any = None marketplace_dir = create_test_marketplace( tmp_path / "marketplace", @@ -817,13 +818,14 @@ def __init__(self): ] ), ) + runtime_client = RuntimeMCPClient() conversation = LocalConversation( agent=agent, workspace=workspace, visualizer=None, mcp_tool_provider=RecordingMCPToolProvider( mcp_tools_created, - RuntimeMCPClient(), + runtime_client, state_locked=lambda: conversation.state.locked(), ), ) @@ -835,6 +837,7 @@ def __init__(self): for name, tool in existing_tools.items(): assert conversation.agent.tools_map[name] is tool assert conversation.agent.tools_map[runtime_tool.name] is runtime_tool + assert callable(runtime_client._tools_reconciled_callback) assert "runtime-server" in conversation.agent.mcp_config assert len(mcp_tools_created) == 1 created_config, state_locked = mcp_tools_created[0] diff --git a/tests/sdk/mcp/test_mcp_tool_list_changed.py b/tests/sdk/mcp/test_mcp_tool_list_changed.py index 5dc6e50d8b..fa4e57d1f9 100644 --- a/tests/sdk/mcp/test_mcp_tool_list_changed.py +++ b/tests/sdk/mcp/test_mcp_tool_list_changed.py @@ -25,6 +25,7 @@ import pytest from fastmcp import FastMCP from fastmcp.server.dependencies import get_context +from pydantic import ValidationError from openhands.sdk.agent.base import AgentBase from openhands.sdk.llm import TextContent @@ -85,6 +86,7 @@ def __init__(self, tools: list[mcp_types.Tool]): self._server_tools = list(tools) self._tools: list[MCPToolDefinition] = [] self._closed = False + self._tools_reconciled_callback = None async def list_tools(self) -> list[mcp_types.Tool]: return list(self._server_tools) @@ -102,6 +104,7 @@ def __init__(self, _initialized: bool, _tools): # noqa: ANN001 # Skip pydantic validation; set the attributes the helpers read. object.__setattr__(self, "_initialized", _initialized) object.__setattr__(self, "_tools", _tools) + object.__setattr__(self, "_tools_lock", threading.RLock()) object.__setattr__(self, "filter_tools_regex", None) def step(self, conversation, on_event, on_token=None): # noqa: ARG002, ANN001 @@ -147,6 +150,61 @@ async def run(): assert {t.name for t in client._tools} == {"a"} +def test_refresh_tools_reconciles_updates_and_removals(): + """The full snapshot callback receives updated definitions and removals.""" + old_tool = mcp_types.Tool( + name="changing", + description="old schema", + inputSchema={ + "type": "object", + "properties": {"old": {"type": "string"}}, + "required": ["old"], + }, + ) + new_tool = mcp_types.Tool( + name="changing", + description="new schema", + inputSchema={ + "type": "object", + "properties": {"new": {"type": "integer"}}, + "required": ["new"], + }, + ) + client = _FakeClient([new_tool]) + old_definition = MCPToolDefinition.create( + mcp_tool=old_tool, + mcp_client=cast(MCPClient, client), + )[0] + old_definition.action_from_arguments({"old": "value"}) + client._tools = [ + old_definition, + MCPToolDefinition.create( + mcp_tool=_make_mcp_tool("gone"), + mcp_client=cast(MCPClient, client), + )[0], + ] + received: list[tuple[object, list[MCPToolDefinition]]] = [] + + async def run(): + await _refresh_tools( + cast(MCPClient, client), + on_tools_reconciled=lambda owner, tools: received.append( + (owner, list(tools)) + ), + ) + + asyncio.new_event_loop().run_until_complete(run()) + + assert len(received) == 1 + owner, tools = received[0] + assert owner is client + assert [tool.name for tool in tools] == ["changing"] + assert tools[0].description == "new schema" + tools[0].action_from_arguments({"new": 42}) + with pytest.raises(ValidationError): + tools[0].action_from_arguments({"old": "value"}) + + def test_refresh_tools_no_callback_still_reconciles(): """Without a callback the client tool list is still kept in sync.""" client = _FakeClient([_make_mcp_tool("a"), _make_mcp_tool("b")]) @@ -301,12 +359,15 @@ def test_list_changed_notification_reconciles_readded_agent_tool( def on_tools_changed(tools): # noqa: ANN001 received.extend(tool.name for tool in tools) - agent._on_mcp_tools_changed(tools) + + def on_tools_reconciled(client, tools): # noqa: ANN001 + agent._on_mcp_tools_reconciled(client, tools) with create_mcp_tools( config, timeout=10.0, on_tools_changed=on_tools_changed, + on_tools_reconciled=on_tools_reconciled, ) as client: agent.add_runtime_tools(client.tools) initial_names = {t.name for t in client.tools} @@ -350,13 +411,13 @@ def on_tools_changed(tools): # noqa: ANN001 time.sleep(0.1) assert all(tool.name != "extra" for tool in client.tools) - assert agent.tools_map["extra"] is first_agent_extra + assert "extra" not in agent.tools_map register_observation = register_tool(register_tool.action_from_arguments({})) assert not register_observation.is_error deadline = time.time() + 10.0 - while time.time() < deadline and agent.tools_map["extra"] is first_agent_extra: + while time.time() < deadline and "extra" not in agent.tools_map: time.sleep(0.1) readded_agent_extra = agent.tools_map["extra"] @@ -391,3 +452,121 @@ def test_on_mcp_tools_changed_skips_when_not_initialized(): # Must not raise even though add_runtime_tools would warn. agent._on_mcp_tools_changed([]) # type: ignore[arg-type] + + +def test_on_mcp_tools_reconciled_does_not_remove_other_client_tools(): + """A client snapshot only replaces tools owned by that client.""" + first_client = _FakeClient([]) + second_client = _FakeClient([]) + first_tool = MCPToolDefinition.create( + mcp_tool=_make_mcp_tool("first"), + mcp_client=cast(MCPClient, first_client), + )[0] + second_tool = MCPToolDefinition.create( + mcp_tool=_make_mcp_tool("second"), + mcp_client=cast(MCPClient, second_client), + )[0] + replacement = MCPToolDefinition.create( + mcp_tool=_make_mcp_tool("replacement"), + mcp_client=cast(MCPClient, first_client), + )[0] + agent = _ConcreteAgent( + _initialized=True, + _tools={"first": first_tool, "second": second_tool}, + ) + + agent._on_mcp_tools_reconciled( + cast(MCPClient, first_client), + [replacement], + ) + + assert set(agent.tools_map) == {"replacement", "second"} + assert agent.tools_map["second"] is second_tool + + +def test_on_mcp_tools_reconciled_filters_before_conflict_check(): + first_client = _FakeClient([]) + second_client = _FakeClient([]) + blocked = MCPToolDefinition.create( + mcp_tool=_make_mcp_tool("blocked"), + mcp_client=cast(MCPClient, first_client), + )[0] + filtered_conflict = MCPToolDefinition.create( + mcp_tool=_make_mcp_tool("blocked"), + mcp_client=cast(MCPClient, second_client), + )[0] + allowed = MCPToolDefinition.create( + mcp_tool=_make_mcp_tool("allowed"), + mcp_client=cast(MCPClient, second_client), + )[0] + agent = _ConcreteAgent(_initialized=True, _tools={"blocked": blocked}) + object.__setattr__(agent, "filter_tools_regex", r"^allowed$") + + agent._on_mcp_tools_reconciled( + cast(MCPClient, second_client), + [filtered_conflict, allowed], + ) + + assert set(agent.tools_map) == {"blocked", "allowed"} + + +def test_on_mcp_tools_reconciled_serializes_client_updates(): + first_snapshot = threading.Event() + second_snapshot = threading.Event() + + class CoordinatedDict(dict[str, MCPToolDefinition]): + def __init__(self, values: dict[str, MCPToolDefinition]): + super().__init__(values) + self.local = threading.local() + + def items(self): # type: ignore[override] # noqa: ANN201 + snapshot = list(super().items()) + count = getattr(self.local, "count", 0) + 1 + self.local.count = count + if count == 2: + if first_snapshot.is_set(): + second_snapshot.set() + else: + first_snapshot.set() + second_snapshot.wait(0.2) + return snapshot + + first_client = _FakeClient([]) + second_client = _FakeClient([]) + first_old = MCPToolDefinition.create( + mcp_tool=_make_mcp_tool("first_old"), + mcp_client=cast(MCPClient, first_client), + )[0] + second_old = MCPToolDefinition.create( + mcp_tool=_make_mcp_tool("second_old"), + mcp_client=cast(MCPClient, second_client), + )[0] + first_new = MCPToolDefinition.create( + mcp_tool=_make_mcp_tool("first_new"), + mcp_client=cast(MCPClient, first_client), + )[0] + second_new = MCPToolDefinition.create( + mcp_tool=_make_mcp_tool("second_new"), + mcp_client=cast(MCPClient, second_client), + )[0] + agent = _ConcreteAgent( + _initialized=True, + _tools=CoordinatedDict({"first_old": first_old, "second_old": second_old}), + ) + threads = [ + threading.Thread( + target=agent._on_mcp_tools_reconciled, + args=(cast(MCPClient, first_client), [first_new]), + ), + threading.Thread( + target=agent._on_mcp_tools_reconciled, + args=(cast(MCPClient, second_client), [second_new]), + ), + ] + + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert set(agent.tools_map) == {"first_new", "second_new"}