From 399238f2c2a7c3d29c8a4b6970cb9a0ccca698f1 Mon Sep 17 00:00:00 2001 From: Amandeep Khurana Date: Sat, 12 Sep 2026 14:32:28 -0700 Subject: [PATCH] fix(mcp): isolate OAuth connections by profile --- .../test_mcp_multiplex_connection_keys.py | 20 ++++++++++++++++++ tools/mcp_tool_registration.py | 21 +++++++++++++++---- 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/tests/tools/test_mcp_multiplex_connection_keys.py b/tests/tools/test_mcp_multiplex_connection_keys.py index 9acff00127..8a2b36bf54 100644 --- a/tests/tools/test_mcp_multiplex_connection_keys.py +++ b/tests/tools/test_mcp_multiplex_connection_keys.py @@ -89,6 +89,26 @@ def test_same_named_server_with_other_credentials_is_a_separate_connection(two_p assert handlers._check_circuit_breaker("x") is None +def test_oauth_server_is_not_adopted_across_profiles(two_profiles): + from tools import mcp_tool_discovery as disc + from tools import mcp_tool_registration as reg + from tools.registry import registry + + cfg = {"url": "https://mcp.example/x", "auth": "oauth"} + + two_profiles("a") + srv_a = _server("x", cfg) + disc._adopt_server("x", srv_a) + srv_a._registered_tool_names = reg._register_server_tools("x", srv_a, cfg) + assert reg.register_connected_into_current_scope({"x": dict(cfg)}) == 0 + assert registry.get_tool_names_for_toolset("mcp-x") == ["mcp__x__t"] + + two_profiles("b") + assert reg.register_connected_into_current_scope({"x": dict(cfg)}) == 0 + assert registry.get_tool_names_for_toolset("mcp-x") == [] + assert "x" in disc._select_new_servers({"x": dict(cfg)}) + + def test_owner_reload_reregisters_profiles_that_adopted_its_connection(two_profiles): import tools.mcp_tool as core from tools import mcp_tool_discovery as disc, mcp_tool_lifecycle as lifecycle diff --git a/tools/mcp_tool_registration.py b/tools/mcp_tool_registration.py index df3a4aad33..aff3544dfa 100644 --- a/tools/mcp_tool_registration.py +++ b/tools/mcp_tool_registration.py @@ -413,8 +413,19 @@ def _connection_identity(config: dict) -> tuple: (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 _same_server_route(server: Any, config: dict, *, cross_profile: bool = False) -> bool: + """Whether *server* matches *config*, with OAuth connections never reusable across profiles. + + OAuth credentials live in the owning profile's token storage rather than the static config, + so identical OAuth configs cannot prove that two profiles authenticate as the same account. + """ + server_config = getattr(server, "_config", {}) or {} + if cross_profile and any( + (candidate.get("auth") or "").lower().strip() == "oauth" + for candidate in (server_config, config) + ): + return False + return _connection_identity(server_config) == _connection_identity(config) def register_connected_into_current_scope(servers: dict) -> int: @@ -447,8 +458,10 @@ def _register_connected_into_current_scope(servers: dict) -> int: continue server = _core._servers.get(key) config = servers.get(_key_name(key)) + cross_profile = key != _server_key(_key_name(key), scope, current=False) if (config is None or not _server_enabled(config) or server is None - or getattr(server, "session", None) is None or not _same_server_route(server, config)): + or getattr(server, "session", None) is None + or not _same_server_route(server, config, cross_profile=cross_profile)): stale.append(key) for key in stale: _remove_server_scope(key, scope) @@ -463,7 +476,7 @@ def _register_connected_into_current_scope(servers: dict) -> int: # Any other profile's live connection with the same route AND credentials is shareable. shared = [(key, live) for key, live in _core._servers.items() if _key_name(key) == name and getattr(live, "session", None) is not None - and _same_server_route(live, config)] + and _same_server_route(live, config, cross_profile=True)] if not shared: continue key, server = shared[0]