From cf3ae7c59c27e229622490566f79ee7c95eff9a4 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Fri, 17 Jul 2026 01:33:59 -0700 Subject: [PATCH] fix(mcp): preserve live OAuth state during reauth --- hermes_cli/web_server.py | 34 +++++++++++-------- tests/tools/test_mcp_dashboard_oauth.py | 43 +++++++++++++++++++++++++ tests/tools/test_mcp_oauth_manager.py | 14 ++++++++ tools/mcp_dashboard_oauth.py | 1 + tools/mcp_oauth.py | 13 ++++++-- tools/mcp_oauth_manager.py | 18 +++++++++-- 6 files changed, 106 insertions(+), 17 deletions(-) diff --git a/hermes_cli/web_server.py b/hermes_cli/web_server.py index e84456ed2c26..e2dce9451df9 100644 --- a/hermes_cli/web_server.py +++ b/hermes_cli/web_server.py @@ -11398,10 +11398,15 @@ def _run_dashboard_mcp_oauth(flow, cfg: dict) -> None: try: transaction = _mcp_oauth_transaction(flow) with transaction, force_interactive_oauth(), dashboard_oauth_flow(flow): + manager = get_manager() storage = HermesTokenStorage(flow.server_name) backup = storage.snapshot() + previous_entry = None try: - get_manager().remove(flow.server_name, hermes_home=flow.hermes_home) + previous_entry = manager.remove( + flow.server_name, + hermes_home=flow.hermes_home, + ) tools = _probe_single_server( flow.server_name, cfg, @@ -11415,8 +11420,18 @@ def _run_dashboard_mcp_oauth(flow, cfg: dict) -> None: _save_mcp_server(flow.server_name, cfg) flow.tools = [{"name": t, "description": d} for t, d in tools] flow.mark_approved() + if flow.reconnect_live: + from tools.mcp_tool import reconnect_mcp_server + + reconnect_mcp_server(flow.server_name) except Exception: - storage.restore(backup) + manager.evict(flow.server_name, hermes_home=flow.hermes_home) + storage.restore(backup, only_if_absent=True) + manager.restore_entry( + flow.server_name, + previous_entry, + hermes_home=flow.hermes_home, + ) raise finally: reset_secret_scope(secret_token) @@ -11438,15 +11453,6 @@ def _run_dashboard_mcp_oauth(flow, cfg: dict) -> None: ) flow.mark_error(msg) finally: - # Dashboard auth builds a provider with a public callback URI and bridge - # handlers. Evict that one-shot provider after completion; persisted - # tokens/client registration remain for the normal runtime rebuild. - try: - from tools.mcp_oauth_manager import get_manager - - get_manager().evict(flow.server_name, hermes_home=flow.hermes_home) - except Exception: - pass flow.mark_worker_done() @@ -11458,10 +11464,11 @@ async def auth_mcp_server(name: str, request: Request, profile: Optional[str] = _require_token(request) _gc_mcp_oauth_flows() + from hermes_constants import get_hermes_home + + process_home = str(get_hermes_home().expanduser().resolve(strict=False)) with _profile_scope(profile): servers = _get_mcp_servers() - from hermes_constants import get_hermes_home - flow_home = str(get_hermes_home().expanduser().resolve(strict=False)) if name not in servers: raise HTTPException(status_code=404, detail=f"Server '{name}' not found") @@ -11480,6 +11487,7 @@ async def auth_mcp_server(name: str, request: Request, profile: Optional[str] = hermes_home=flow_home, redirect_uri=(cfg.get("oauth") or {}).get("redirect_uri") or _mcp_oauth_callback_url(request, name), + reconnect_live=flow_home == process_home, ) with _mcp_oauth_flows_lock: pending = sum( diff --git a/tests/tools/test_mcp_dashboard_oauth.py b/tests/tools/test_mcp_dashboard_oauth.py index 586eb13ebe9b..a4c5ea6522b9 100644 --- a/tests/tools/test_mcp_dashboard_oauth.py +++ b/tests/tools/test_mcp_dashboard_oauth.py @@ -246,3 +246,46 @@ def test_reconnect_mcp_server_signals_live_task(monkeypatch): assert mcp_tool.reconnect_mcp_server("reports") is True assert server._reconnect_event.called is True + + +def test_reconnect_mcp_server_keeps_manager_entry_until_live_task_rebuilds( + tmp_path, monkeypatch +): + from tools import mcp_tool + from tools.mcp_oauth_manager import MCPOAuthManager, _ProviderEntry + + class Event: + called = False + + def set(self): + self.called = True + + class Server: + _reconnect_event = Event() + + server = Server() + manager = MCPOAuthManager() + manager._entries[manager._key("reports", tmp_path)] = _ProviderEntry( + server_url="https://mcp.example/mcp", oauth_config={} + ) + monkeypatch.setitem(mcp_tool._servers, "reports", server) + monkeypatch.setattr(mcp_tool, "_mcp_loop", None) + + assert mcp_tool.reconnect_mcp_server("reports") is True + assert manager._key("reports", tmp_path) in manager._entries + + +def test_failed_reauth_rollback_preserves_newer_oauth_state(tmp_path, monkeypatch): + from tools.mcp_oauth import HermesTokenStorage + + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + storage = HermesTokenStorage("reports") + storage._tokens_path().parent.mkdir(parents=True) + storage._tokens_path().write_text("OLD") + backup = storage.snapshot() + storage.remove() + + storage._tokens_path().write_text("FRESH") + storage.restore(backup, only_if_absent=True) + + assert storage._tokens_path().read_text() == "FRESH" diff --git a/tests/tools/test_mcp_oauth_manager.py b/tests/tools/test_mcp_oauth_manager.py index 6343ef8ff423..3c47239f6f2d 100644 --- a/tests/tools/test_mcp_oauth_manager.py +++ b/tests/tools/test_mcp_oauth_manager.py @@ -74,6 +74,20 @@ def test_manager_explicit_home_removes_only_that_profiles_tokens(tmp_path): assert paths[0].exists() assert not paths[1].exists() + +def test_manager_can_restore_removed_entry_after_failed_reauth(tmp_path, monkeypatch): + from tools.mcp_oauth_manager import MCPOAuthManager + + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + _set_interactive_stdin(monkeypatch) + manager = MCPOAuthManager() + provider = manager.get_or_build_provider("shared", "https://mcp.example", {}) + + entry = manager.remove("shared") + manager.restore_entry("shared", entry) + + assert manager.get_or_build_provider("shared", "https://mcp.example", {}) is provider + pytest.importorskip( "mcp.client.auth.oauth2", reason="MCP SDK 1.26.0+ required for OAuth support", diff --git a/tools/mcp_dashboard_oauth.py b/tools/mcp_dashboard_oauth.py index df8dd73f1b53..31663e0f37b0 100644 --- a/tools/mcp_dashboard_oauth.py +++ b/tools/mcp_dashboard_oauth.py @@ -25,6 +25,7 @@ class DashboardOAuthFlow: profile: str | None hermes_home: str redirect_uri: str + reconnect_live: bool = False created_at: float = field(default_factory=time.time) status: str = "starting" authorization_url: str | None = None diff --git a/tools/mcp_oauth.py b/tools/mcp_oauth.py index ef85ed265318..a42a11a7e8d5 100644 --- a/tools/mcp_oauth.py +++ b/tools/mcp_oauth.py @@ -522,8 +522,17 @@ class HermesTokenStorage: pass return snap - def restore(self, snapshot: dict[str, bytes]) -> None: - """Revert to a ``snapshot()`` capture (dropping any newer partial state).""" + def restore(self, snapshot: dict[str, bytes], *, only_if_absent: bool = False) -> None: + """Revert to a snapshot without overwriting a concurrent successful write.""" + if only_if_absent and any( + path.exists() + for path in (self._tokens_path(), self._client_info_path(), self._meta_path()) + ): + logger.info( + "Skipping OAuth rollback for %s because newer state exists", + self._server_name, + ) + return self.remove() if not snapshot: return diff --git a/tools/mcp_oauth_manager.py b/tools/mcp_oauth_manager.py index d5f86a203cdc..9d2a625c2632 100644 --- a/tools/mcp_oauth_manager.py +++ b/tools/mcp_oauth_manager.py @@ -587,14 +587,14 @@ class MCPOAuthManager: server_name: str, *, hermes_home: str | Path | None = None, - ) -> None: + ) -> _ProviderEntry | None: """Evict the provider from cache AND delete tokens from disk. Called by ``hermes mcp remove `` and (indirectly) by ``hermes mcp login `` during forced re-auth. """ with self._entries_lock: - self._entries.pop(self._key(server_name, hermes_home), None) + entry = self._entries.pop(self._key(server_name, hermes_home), None) from tools.mcp_oauth import remove_oauth_tokens remove_oauth_tokens(server_name, hermes_home=hermes_home) @@ -602,6 +602,20 @@ class MCPOAuthManager: "MCP OAuth '%s': evicted from cache and removed from disk", server_name, ) + return entry + + def restore_entry( + self, + server_name: str, + entry: _ProviderEntry | None, + *, + hermes_home: str | Path | None = None, + ) -> None: + """Restore a provider entry removed for a failed reauthorization.""" + if entry is None: + return + with self._entries_lock: + self._entries.setdefault(self._key(server_name, hermes_home), entry) def evict( self,