diff --git a/tests/tools/test_mcp_oauth.py b/tests/tools/test_mcp_oauth.py index e43bf0a1851df..3089a35425edb 100644 --- a/tests/tools/test_mcp_oauth.py +++ b/tests/tools/test_mcp_oauth.py @@ -21,6 +21,8 @@ _is_interactive, _wait_for_callback, _make_callback_handler, + _make_redirect_handler, + _make_wait_for_callback, _redirect_handler, _paste_callback_reader, ) @@ -301,6 +303,19 @@ def test_no_ssh_hint_when_port_not_set(self, monkeypatch, capsys): err = capsys.readouterr().err assert "ssh -N -L" not in err + def test_factory_uses_closed_over_port(self, monkeypatch, capsys): + import tools.mcp_oauth as mco + monkeypatch.setattr(mco, "_oauth_port", 49999) + monkeypatch.setenv("SSH_CLIENT", "1.2.3.4 1234 22") + monkeypatch.delenv("SSH_TTY", raising=False) + monkeypatch.setattr(mco, "_can_open_browser", lambda: False) + + self._run(_make_redirect_handler(49203)("https://example.com/auth")) + + err = capsys.readouterr().err + assert "49203" in err + assert "49999" not in err + # --------------------------------------------------------------------------- # Path traversal protection @@ -475,6 +490,33 @@ async def instant_sleep(_seconds): with pytest.raises(OAuthNonInteractiveError, match="callback timed out"): asyncio.run(_wait_for_callback()) + def test_factory_binds_closed_over_port(self, monkeypatch): + import tools.mcp_oauth as mod + bound: dict[str, int] = {} + + class _FakeServer: + def __init__(self, addr, _handler_cls): + bound["port"] = addr[1] + + def handle_request(self): + return None + + def server_close(self): + return None + + mod._oauth_port = 49998 + monkeypatch.setattr(mod, "HTTPServer", _FakeServer) + monkeypatch.setattr(mod, "_is_interactive", lambda: False) + + async def instant_sleep(_seconds): + pass + + with patch.object(mod.asyncio, "sleep", instant_sleep): + with pytest.raises(OAuthNonInteractiveError, match="callback timed out"): + asyncio.run(_make_wait_for_callback(49204)()) + + assert bound["port"] == 49204 + class TestBuildOAuthAuthNonInteractive: """build_oauth_auth() in non-interactive mode.""" diff --git a/tests/tools/test_mcp_oauth_manager.py b/tests/tools/test_mcp_oauth_manager.py index 2a66449cbd2ea..42c87f608edee 100644 --- a/tests/tools/test_mcp_oauth_manager.py +++ b/tests/tools/test_mcp_oauth_manager.py @@ -139,3 +139,45 @@ def test_manager_builds_hermes_provider_subclass(tmp_path, monkeypatch): assert isinstance(provider, _HERMES_PROVIDER_CLS) assert provider._hermes_server_name == "srv" + +def test_manager_builds_provider_with_per_server_oauth_handlers(tmp_path, monkeypatch): + """Manager must pass per-provider redirect/callback closures, not globals.""" + import tools.mcp_oauth as oauth_mod + import tools.mcp_oauth_manager as manager_mod + from tools.mcp_oauth_manager import reset_manager_for_tests + + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + reset_manager_for_tests() + + captured_kwargs: dict[str, object] = {} + seen_ports: dict[str, int] = {} + redirect_handler = object() + callback_handler = object() + + class _FakeProvider: + def __init__(self, **kwargs): + captured_kwargs.update(kwargs) + + def fake_make_redirect_handler(port): + seen_ports["redirect"] = port + return redirect_handler + + def fake_make_wait_for_callback(port): + seen_ports["callback"] = port + return callback_handler + + monkeypatch.setattr(manager_mod, "_HERMES_PROVIDER_CLS", _FakeProvider) + monkeypatch.setattr(oauth_mod, "_make_redirect_handler", fake_make_redirect_handler) + monkeypatch.setattr(oauth_mod, "_make_wait_for_callback", fake_make_wait_for_callback) + + mgr = manager_mod.MCPOAuthManager() + provider = mgr.get_or_build_provider( + "srv", + "https://example.com/mcp", + {"redirect_port": 43123}, + ) + + assert isinstance(provider, _FakeProvider) + assert seen_ports == {"redirect": 43123, "callback": 43123} + assert captured_kwargs["redirect_handler"] is redirect_handler + assert captured_kwargs["callback_handler"] is callback_handler diff --git a/tools/mcp_oauth.py b/tools/mcp_oauth.py index 832a6f5945f4a..19dc0d9ffb3ec 100644 --- a/tools/mcp_oauth.py +++ b/tools/mcp_oauth.py @@ -397,142 +397,123 @@ def log_message(self, fmt: str, *args: Any) -> None: # --------------------------------------------------------------------------- -async def _redirect_handler(authorization_url: str) -> None: - """Show the authorization URL to the user. - - Opens the browser automatically when possible; always prints the URL - as a fallback for headless/SSH/gateway environments. - """ - msg = ( - f"\n MCP OAuth: authorization required.\n" - f" Open this URL in your browser:\n\n" - f" {authorization_url}\n" - ) - print(msg, file=sys.stderr) - - # On a remote SSH session the OAuth provider redirects to - # http://127.0.0.1:/callback, which reaches the callback server on - # the *remote* machine — not the user's local machine where the browser - # opened. Two ways out: paste the redirect URL back (default fallback, - # offered by _wait_for_callback on interactive TTYs), or set up an SSH - # port forward so the redirect tunnels through. - if _oauth_port and (os.getenv("SSH_CLIENT") or os.getenv("SSH_TTY")): - print( - f" Remote session detected. After you authorize, the provider redirects to\n" - f" http://127.0.0.1:{_oauth_port}/callback\n" - f" which only the listener on THIS machine can receive. Two options:\n" - f"\n" - f" 1. Easiest — when your browser shows a connection error after\n" - f" authorizing, copy the full URL from the address bar and paste\n" - f" it at the prompt below. The pasted ``code=...&state=...`` is\n" - f" enough to complete the flow.\n" - f"\n" - f" 2. Or forward the port first in a separate terminal:\n" - f" ssh -N -L {_oauth_port}:127.0.0.1:{_oauth_port} @\n" - f" then open the URL above and let it redirect normally.\n" - f"\n" - f" See: https://hermes-agent.nousresearch.com/docs/guides/oauth-over-ssh\n", - file=sys.stderr, +def _make_redirect_handler(callback_port: int | None): + """Return a redirect handler bound to one OAuth flow's callback port.""" + async def _handler(authorization_url: str) -> None: + msg = ( + f"\n MCP OAuth: authorization required.\n" + f" Open this URL in your browser:\n\n" + f" {authorization_url}\n" ) + print(msg, file=sys.stderr) + + # Bind the SSH hint to the flow's own callback port instead of the + # mutable module-level fallback. Concurrent OAuth providers can then + # render the correct port even if another flow starts in between. + if callback_port and (os.getenv("SSH_CLIENT") or os.getenv("SSH_TTY")): + print( + f" Remote session detected. After you authorize, the provider redirects to\n" + f" http://127.0.0.1:{callback_port}/callback\n" + f" which only the listener on THIS machine can receive. Two options:\n" + f"\n" + f" 1. Easiest — when your browser shows a connection error after\n" + f" authorizing, copy the full URL from the address bar and paste\n" + f" it at the prompt below. The pasted ``code=...&state=...`` is\n" + f" enough to complete the flow.\n" + f"\n" + f" 2. Or forward the port first in a separate terminal:\n" + f" ssh -N -L {callback_port}:127.0.0.1:{callback_port} @\n" + f" then open the URL above and let it redirect normally.\n" + f"\n" + f" See: https://hermes-agent.nousresearch.com/docs/guides/oauth-over-ssh\n", + file=sys.stderr, + ) - if _can_open_browser(): - try: - opened = webbrowser.open(authorization_url) - if opened: - print(" (Browser opened automatically.)\n", file=sys.stderr) - else: + if _can_open_browser(): + try: + opened = webbrowser.open(authorization_url) + if opened: + print(" (Browser opened automatically.)\n", file=sys.stderr) + else: + print(" (Could not open browser — please open the URL manually.)\n", file=sys.stderr) + except Exception: print(" (Could not open browser — please open the URL manually.)\n", file=sys.stderr) - except Exception: - print(" (Could not open browser — please open the URL manually.)\n", file=sys.stderr) - else: - print(" (Headless environment detected — open the URL manually.)\n", file=sys.stderr) + else: + print(" (Headless environment detected — open the URL manually.)\n", file=sys.stderr) + return _handler -async def _wait_for_callback() -> tuple[str, str | None]: - """Wait for the OAuth callback to arrive on the local callback server. - - Uses the module-level ``_oauth_port`` which is set by ``build_oauth_auth`` - before this is ever called. Polls for the result without blocking the - event loop. - - On an interactive TTY, races the HTTP listener against a stdin paste - fallback so users without an SSH tunnel can copy the redirect URL (or - just the ``code=...&state=...`` query string) from a browser on another - machine and paste it back. The HTTP listener wins when the redirect - reaches it first; the paste fallback wins when it doesn't. - - Raises: - OAuthNonInteractiveError: If the callback times out (no user present - to complete the browser auth). - RuntimeError: If ``_oauth_port`` has not been set, which would indicate - that ``build_oauth_auth`` was skipped — the asserting form below - was a silent bug when running Python with ``-O``/``-OO``. - """ - if _oauth_port is None: - raise RuntimeError( - "OAuth callback port not set — build_oauth_auth must be called " - "before _wait_for_oauth_callback" - ) - # The callback server is already running (started in build_oauth_auth). - # We just need to poll for the result. - handler_cls, result = _make_callback_handler() +async def _redirect_handler(authorization_url: str) -> None: + """Back-compat wrapper for call sites that still rely on ``_oauth_port``.""" + await _make_redirect_handler(_oauth_port)(authorization_url) - # Start a temporary server on the known port - try: - server = HTTPServer(("127.0.0.1", _oauth_port), handler_cls) - except OSError: - # Port already in use — the server from build_oauth_auth is running. - # Fall back to polling the server started by build_oauth_auth. - raise OAuthNonInteractiveError( - "OAuth callback timed out — could not bind callback port. " - "Complete the authorization in a browser first, then retry." - ) - server_thread = threading.Thread(target=server.handle_request, daemon=True) - server_thread.start() +def _make_wait_for_callback(callback_port: int | None): + """Return a callback waiter bound to one OAuth flow's callback port.""" + async def _handler() -> tuple[str, str | None]: + if callback_port is None: + raise RuntimeError( + "OAuth callback port not set — build_oauth_auth must be called " + "before _wait_for_oauth_callback" + ) - # Optional paste-fallback thread: only on interactive TTYs. Reads one - # line from stdin and writes the parsed code/state into the shared - # result dict. The HTTP listener and this thread race for the result; - # whichever fills it first wins. - paste_thread: threading.Thread | None = None - if _is_interactive(): - print( - "\n Or paste the redirect URL here (or the ``?code=...&state=...`` " - "portion) and press Enter. Type ``skip`` + Enter to continue " - "without this server:", - file=sys.stderr, - flush=True, - ) - paste_thread = threading.Thread( - target=_paste_callback_reader, args=(result,), daemon=True - ) - paste_thread.start() + handler_cls, result = _make_callback_handler() - timeout = 300.0 - poll_interval = 0.5 - elapsed = 0.0 - try: - while elapsed < timeout: - if result["auth_code"] is not None or result["error"] is not None: - break - await asyncio.sleep(poll_interval) - elapsed += poll_interval - finally: - server.server_close() - - if result["error"] == _USER_SKIPPED_SENTINEL: - raise OAuthNonInteractiveError("user_skipped") - if result["error"]: - raise RuntimeError(f"OAuth authorization failed: {result['error']}") - if result["auth_code"] is None: - raise OAuthNonInteractiveError( - "OAuth callback timed out — no authorization code received. " - "Ensure you completed the browser authorization flow." - ) + try: + server = HTTPServer(("127.0.0.1", callback_port), handler_cls) + except OSError: + raise OAuthNonInteractiveError( + "OAuth callback timed out — could not bind callback port. " + "Complete the authorization in a browser first, then retry." + ) - return result["auth_code"], result["state"] + server_thread = threading.Thread(target=server.handle_request, daemon=True) + server_thread.start() + + if _is_interactive(): + print( + "\n Or paste the redirect URL here (or the ``?code=...&state=...`` " + "portion) and press Enter. Type ``skip`` + Enter to continue " + "without this server:", + file=sys.stderr, + flush=True, + ) + paste_thread = threading.Thread( + target=_paste_callback_reader, args=(result,), daemon=True + ) + paste_thread.start() + + timeout = 300.0 + poll_interval = 0.5 + elapsed = 0.0 + try: + while elapsed < timeout: + if result["auth_code"] is not None or result["error"] is not None: + break + await asyncio.sleep(poll_interval) + elapsed += poll_interval + finally: + server.server_close() + + if result["error"] == _USER_SKIPPED_SENTINEL: + raise OAuthNonInteractiveError("user_skipped") + if result["error"]: + raise RuntimeError(f"OAuth authorization failed: {result['error']}") + if result["auth_code"] is None: + raise OAuthNonInteractiveError( + "OAuth callback timed out — no authorization code received. " + "Ensure you completed the browser authorization flow." + ) + + return result["auth_code"], result["state"] + + return _handler + + +async def _wait_for_callback() -> tuple[str, str | None]: + """Back-compat wrapper for call sites that still rely on ``_oauth_port``.""" + return await _make_wait_for_callback(_oauth_port)() def _paste_callback_reader(result: dict) -> None: @@ -766,11 +747,13 @@ def build_oauth_auth( client_metadata = _build_client_metadata(cfg) _maybe_preregister_client(storage, cfg, client_metadata) + callback_port = cfg["_resolved_port"] + return OAuthClientProvider( server_url=server_url, client_metadata=client_metadata, storage=storage, - redirect_handler=_redirect_handler, - callback_handler=_wait_for_callback, + redirect_handler=_make_redirect_handler(callback_port), + callback_handler=_make_wait_for_callback(callback_port), timeout=float(cfg.get("timeout", 300)), ) diff --git a/tools/mcp_oauth_manager.py b/tools/mcp_oauth_manager.py index 6a4573a8677db..de1e027cd3a21 100644 --- a/tools/mcp_oauth_manager.py +++ b/tools/mcp_oauth_manager.py @@ -412,9 +412,9 @@ def _build_provider( _build_client_metadata, _configure_callback_port, _is_interactive, + _make_redirect_handler, + _make_wait_for_callback, _maybe_preregister_client, - _redirect_handler, - _wait_for_callback, ) if not _OAUTH_AVAILABLE: @@ -432,6 +432,7 @@ def _build_provider( ) _configure_callback_port(cfg) + callback_port = cfg["_resolved_port"] client_metadata = _build_client_metadata(cfg) _maybe_preregister_client(storage, cfg, client_metadata) @@ -440,8 +441,8 @@ def _build_provider( server_url=entry.server_url, client_metadata=client_metadata, storage=storage, - redirect_handler=_redirect_handler, - callback_handler=_wait_for_callback, + redirect_handler=_make_redirect_handler(callback_port), + callback_handler=_make_wait_for_callback(callback_port), timeout=float(cfg.get("timeout", 300)), )