From 4ddbcbd35ef03325fbabb10f6477a188b43339cd Mon Sep 17 00:00:00 2001 From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com> Date: Wed, 9 Sep 2026 11:42:23 +0530 Subject: [PATCH] fix(mcp): a profile only adopts a shared connection whose credentials match its own config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _same_server_route compared config_fingerprint alone, which by design excludes env/headers/auth (so the schema cache survives a token rotation). Profile B with the same URL but different headers/env therefore adopted profile A's live connection and called tools as A. _connection_identity = route fingerprint + env + headers + auth mode, used by both the adopt and the stale-removal checks. Also collapses the three writers of _server_tool_scopes to two: the adoption loop re-implemented in mcp_tool_discovery._select_new_servers is dropped — register_connected_into_current_scope (which runs first in register_mcp_servers) is the single adopter, and _register_candidates records scope for freshly registered tools. --- tests/gateway/test_multiplex_mcp_discovery.py | 25 +++++++++++-------- tools/mcp_tool_discovery.py | 13 ---------- tools/mcp_tool_registration.py | 25 ++++++++++++++----- 3 files changed, 34 insertions(+), 29 deletions(-) diff --git a/tests/gateway/test_multiplex_mcp_discovery.py b/tests/gateway/test_multiplex_mcp_discovery.py index 28905a5144..95be0c1c9d 100644 --- a/tests/gateway/test_multiplex_mcp_discovery.py +++ b/tests/gateway/test_multiplex_mcp_discovery.py @@ -126,7 +126,8 @@ async def test_reload_mcp_reports_a_shared_server_to_a_non_owner_profile( get_or_create_session=MagicMock(side_effect=RuntimeError("skip transcript")), ) - live_server = SimpleNamespace(session=object(), _config={}) + live_server = SimpleNamespace(session=object(), _config={}, _tools=[], tool_timeout=30, + initialize_result=None, _registered_tool_names=[]) monkeypatch.setattr(mcp_tool, "_servers", {"shared": live_server}) monkeypatch.setattr(mcp_tool, "_server_scope_keys", {"shared": launch_scope}) monkeypatch.setattr(mcp_tool, "_server_tool_scopes", {"shared": {launch_scope}}, raising=False) @@ -136,7 +137,8 @@ async def test_reload_mcp_reports_a_shared_server_to_a_non_owner_profile( monkeypatch.setattr(mcp_tool, "_mcp_registry_scope", lambda: worker_scope) def fake_discover() -> list[str]: - _mcp_discovery._select_new_servers({"shared": {}}) + from tools import mcp_tool_registration as _mcp_registration + _mcp_registration.register_connected_into_current_scope({"shared": {}}) return ["mcp__shared__tool"] monkeypatch.setattr(_mcp_lifecycle, "shutdown_mcp_servers", lambda **_kwargs: None) @@ -157,11 +159,18 @@ async def test_reload_mcp_reports_a_shared_server_to_a_non_owner_profile( assert mcp_tool._server_tool_scopes["shared"] == {launch_scope, worker_scope} -def test_scope_visibility_rejects_a_foreign_server_route( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch +@pytest.mark.parametrize("worker_cfg", [ + {"url": "https://worker.example/mcp"}, # different route + {"url": "https://default.example/mcp", "headers": {"Authorization": "Bearer worker"}}, # same route, other credentials + {"url": "https://default.example/mcp", "env": {"API_TOKEN": "worker"}}, +]) +def test_scope_visibility_rejects_a_foreign_or_differently_authenticated_route( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, worker_cfg: dict ) -> None: + """A profile may only see a live connection whose route AND credentials match its own config; + otherwise it would call tools as the owning profile's identity.""" from tools import mcp_tool - from tools import mcp_tool_discovery as _mcp_discovery + from tools import mcp_tool_registration as _mcp_registration worker_scope = hermes_home_key(tmp_path / "worker") launch_scope = hermes_home_key(tmp_path / "default") @@ -169,13 +178,9 @@ def test_scope_visibility_rejects_a_foreign_server_route( monkeypatch.setattr(mcp_tool, "_servers", {"shared": live_server}) monkeypatch.setattr(mcp_tool, "_server_scope_keys", {"shared": launch_scope}) monkeypatch.setattr(mcp_tool, "_server_tool_scopes", {"shared": {launch_scope}}, raising=False) - monkeypatch.setattr(mcp_tool, "_server_connecting", set()) - monkeypatch.setattr(mcp_tool, "_server_connect_errors", {}) - monkeypatch.setattr(mcp_tool, "_lazy_server_configs", {}) monkeypatch.setattr(mcp_tool, "_mcp_registry_scope", lambda: worker_scope) - _mcp_discovery._select_new_servers({"shared": {"url": "https://worker.example/mcp"}}) - + assert _mcp_registration.register_connected_into_current_scope({"shared": worker_cfg}) == 0 assert mcp_tool._server_tool_scopes["shared"] == {launch_scope} diff --git a/tools/mcp_tool_discovery.py b/tools/mcp_tool_discovery.py index 686ead1c1d..7f25687caa 100644 --- a/tools/mcp_tool_discovery.py +++ b/tools/mcp_tool_discovery.py @@ -247,19 +247,6 @@ def _select_new_servers(servers: Dict[str, dict]) -> Dict[str, dict]: for srv_name in new_servers: _core._server_scope_keys[srv_name] = current_scope _core._server_connect_errors.pop(srv_name, None) - # A shared connection can already be owned by another profile. Record the current - # profile's visibility separately so its scoped reload sees the connection without - # moving teardown ownership away from the profile that opened it. - if current_scope is not None: - from tools.mcp_schema_cache import config_fingerprint - for srv_name, srv_cfg in servers.items(): - if not _enabled(srv_cfg): - continue - server = _core._servers.get(srv_name) - if (server is not None and getattr(server, "session", None) is not None - and config_fingerprint(getattr(server, "_config", {}) or {}) - == config_fingerprint(srv_cfg)): - _core._server_tool_scopes.setdefault(srv_name, set()).add(current_scope) # Track which servers opt-in to parallel tool calls (idempotent). for srv_name, srv_cfg in servers.items(): if _parse_boolish(srv_cfg.get("supports_parallel_tool_calls", False), default=False): diff --git a/tools/mcp_tool_registration.py b/tools/mcp_tool_registration.py index 98cef3064a..abe0d4b025 100644 --- a/tools/mcp_tool_registration.py +++ b/tools/mcp_tool_registration.py @@ -3,6 +3,7 @@ include/exclude filtering, trust-tier metadata capture, utility-tool selection, resolution and the schema-cache write-through. Both entry points (``_register_server_tools`` live, ``_register_from_cache_sync`` lazy) build ``_Candidate`` records for ``_register_candidates``.""" +import json import logging import threading from dataclasses import dataclass @@ -376,10 +377,22 @@ def _server_enabled(config: dict) -> bool: return _parse_boolish(config.get("enabled", True), default=True) -def _same_server_route(server: Any, config: dict) -> bool: +def _connection_identity(config: dict) -> tuple: + """What makes one live connection reusable for another profile: the route fingerprint PLUS + everything that authenticates it (``config_fingerprint`` deliberately excludes credentials so + the schema cache survives a token rotation). Two profiles pointing at the same URL with different + headers/env/auth are two identities; borrowing across them would call tools as the other user.""" from tools.mcp_schema_cache import config_fingerprint - return config_fingerprint(getattr(server, "_config", {}) or {}) == config_fingerprint(config) + def _frozen(value): + return json.dumps(value or {}, sort_keys=True, default=str) + + return (config_fingerprint(config), _frozen(config.get("env")), _frozen(config.get("headers")), + (config.get("auth") or "").lower().strip()) + + +def _same_server_route(server: Any, config: dict) -> bool: + return _connection_identity(getattr(server, "_config", {}) or {}) == _connection_identity(config) def register_connected_into_current_scope(servers: dict) -> int: @@ -405,8 +418,6 @@ def _register_connected_into_current_scope(servers: dict) -> int: if scope is None: return 0 - from tools.mcp_schema_cache import config_fingerprint - with _core._lock: stale = [] for name, scopes in _core._server_tool_scopes.items(): @@ -415,8 +426,7 @@ def _register_connected_into_current_scope(servers: dict) -> int: server = _core._servers.get(name) config = servers.get(name) if (config is None or not _server_enabled(config) or server is None - or getattr(server, "session", None) is None - or config_fingerprint(getattr(server, "_config", {}) or {}) != config_fingerprint(config)): + or getattr(server, "session", None) is None or not _same_server_route(server, config)): stale.append(name) for name in stale: _remove_server_scope(name, scope) @@ -429,6 +439,9 @@ def _register_connected_into_current_scope(servers: dict) -> int: server = _core._servers.get(name) if server is None or getattr(server, "session", None) is None or not _same_server_route(server, config): continue + # Visibility for this profile: the owner keeps teardown, this scope sees the connection. + with _core._lock: + _core._server_tool_scopes.setdefault(name, set()).add(scope) if registry.get_tool_names_for_toolset(f"mcp-{name}"): continue candidates = _tool_candidates(name, server._tools, _make_tool_filter(name, config), server.tool_timeout)