Bind MCP OAuth handlers per provider

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."""

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

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:<port>/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} <user>@<this-host>\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} <user>@<this-host>\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)),
     )