fix(pkce): keep the redirect port bound until the callback server owns it

CHANGELOG.md

@@ -24,6 +24,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `_gemini` now forwards the tool definitions and converts `FunctionCall` parts back into `ToolCallRequest`; the reply is read from `candidates[0].content.parts` instead of `resp.text`, which raises when the reply has no text part (a pure function call has none)
 - gemini's `Schema` proto rejects unknown keywords rather than ignoring them (`ValueError: Unknown field for Schema: <keyword>`), so tool schemas are reduced to the keywords the proto carries — ten of the twenty-three shipped tool schemas include a JSON-Schema `default`, which would have raised on the first tool-using call
 - 13 tests in `tests/test_gemini_tool_calling.py`, offline: the SDK's own request builder runs for real and only the transport client is faked
+**The OAuth redirect port is no longer released between picking it and listening on it**
+- `_find_free_port()` bound a socket, read the port and closed it again, so any process on the machine could claim that port before the callback server bound it — the browser's redirect then landed on a port nothing was listening on and the flow died with a connection refused
+- the flow now binds the socket once (`_bind_callback_socket`) and hands it to the server, which adopts it instead of binding a port of its own; the socket is closed when the flow ends rather than staying bound for the rest of the process
+- 6 tests in `tests/test_pkce_callback_port.py`, including the end-to-end arm that follows the redirect the flow advertises
 
 ### Added
 

devagent/mcp/auth/pkce.py

@@ -26,6 +26,7 @@
 import http.server
 import json
 import secrets
+import socket
 import threading
 import urllib.parse
 import urllib.request
@@ -87,19 +88,39 @@ def log_message(self, fmt: str, *args: object) -> None:
     return _Handler
 
 
-def _start_callback_server(port: int) -> tuple[http.server.HTTPServer, _CallbackResult]:
+def _start_callback_server(sock: socket.socket) -> tuple[http.server.HTTPServer, _CallbackResult]:
+    """Serve on *sock*, which must already be bound to the redirect address.
+
+    The socket is handed over rather than looked up again: the port is never free
+    between picking it and listening on it, so another process cannot claim it in
+    that window and leave the browser's redirect with nothing to connect to.
+    """
     result = _CallbackResult()
-    server = http.server.HTTPServer(("127.0.0.1", port), _make_handler(result))
+    server = http.server.HTTPServer(
+        sock.getsockname(), _make_handler(result), bind_and_activate=False
+    )
+    server.socket = sock
+    server.server_activate()  # listen() on the socket we already hold
+    server.server_name = "127.0.0.1"
+    server.server_port = sock.getsockname()[1]
     thread = threading.Thread(target=server.serve_forever, daemon=True)
     thread.start()
     return server, result
 
 
-def _find_free_port() -> int:
-    import socket
-    with socket.socket() as s:
-        s.bind(("127.0.0.1", 0))
-        return s.getsockname()[1]
+def _bind_callback_socket() -> socket.socket:
+    """Bind a loopback socket on an ephemeral port and keep it open.
+
+    The caller reads the port off this socket for the redirect URI and passes the
+    same socket to :func:`_start_callback_server` once the browser is on its way.
+    """
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    try:
+        sock.bind(("127.0.0.1", 0))
+    except OSError:
+        sock.close()
+        raise
+    return sock
 
 
 # ---------------------------------------------------------------------------
@@ -152,7 +173,8 @@ def pkce_authorize(
     """
     verifier = generate_code_verifier()
     challenge = compute_code_challenge(verifier)
-    port = _find_free_port()
+    callback_sock = _bind_callback_socket()
+    port = callback_sock.getsockname()[1]
     redirect_uri = f"http://127.0.0.1:{port}/callback"
 
     params = {
@@ -166,7 +188,7 @@ def pkce_authorize(
     }
     auth_url = f"{authorization_url}?{urllib.parse.urlencode(params)}"
 
-    server, result = _start_callback_server(port)
+    server, result = _start_callback_server(callback_sock)
     try:
         webbrowser.open(auth_url)
         if not result.done.wait(timeout=timeout):
@@ -176,6 +198,7 @@ def pkce_authorize(
         return _exchange_code(token_url, client_id, result.code, verifier, redirect_uri)
     finally:
         server.shutdown()
+        server.server_close()
 
 
 def refresh_token_flow(

tests/test_pkce_callback_port.py

@@ -0,0 +1,152 @@
+"""The OAuth redirect port must not be free between picking it and listening on it.
+
+``_find_free_port`` bound a socket, read the port and closed it again, so anything on the
+machine could claim that port before the callback server bound it -- the redirect from the
+browser then landed on a port nothing was listening on. The flow now binds the socket once
+and hands it to the server, so the port is never released in between.
+
+Offline: loopback sockets only, no browser and no network beyond 127.0.0.1.
+"""
+
+from __future__ import annotations
+
+import socket
+import urllib.parse
+import urllib.request
+
+import pytest
+
+from devagent.mcp.auth import pkce
+
+
+def _probe_bind(port: int) -> None:
+    """Try to take *port* the way an unrelated process would, without round trips."""
+    probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    try:
+        probe.bind(("127.0.0.1", port))
+    finally:
+        probe.close()
+
+
+def _nothing_is_listening(port: int) -> bool:
+    """Whether a fresh connection to *port* is refused, i.e. no server holds it.
+
+    A refused connection is the fact under test -- once the flow returns (or once
+    ``server_close()`` has run) nothing may be listening on the callback port.
+    Probing with a ``bind`` measures something else, and on Linux it measures it
+    wrongly: the port is also unbindable -- even with ``SO_REUSEADDR`` -- while the
+    connection the last run accepted lingers (the client never closed its end), and
+    that is not the server still holding the port. A connect to a port with no
+    listener is refused whatever state the old connection is in, on either OS.
+    """
+    try:
+        with socket.create_connection(("127.0.0.1", port), timeout=2):
+            return False
+    except OSError:
+        return True
+
+
+class TestThePortIsHeld:
+    def test_another_socket_cannot_take_the_port(self) -> None:
+        sock = pkce._bind_callback_socket()
+        try:
+            port = sock.getsockname()[1]
+            with pytest.raises(OSError):
+                _probe_bind(port)
+        finally:
+            sock.close()
+
+    def test_the_callback_server_serves_on_the_same_socket(self) -> None:
+        sock = pkce._bind_callback_socket()
+        port = sock.getsockname()[1]
+        server, result = pkce._start_callback_server(sock)
+        try:
+            with urllib.request.urlopen(
+                f"http://127.0.0.1:{port}/callback?code=abc", timeout=5
+            ) as resp:
+                body = resp.read()
+                status = resp.status
+            assert status == 200
+            assert b"Authorization complete" in body
+            assert result.code == "abc"
+            assert result.error == ""
+            assert result.done.is_set()
+        finally:
+            server.shutdown()
+            server.server_close()
+
+    def test_an_oauth_error_is_captured_too(self) -> None:
+        sock = pkce._bind_callback_socket()
+        port = sock.getsockname()[1]
+        server, result = pkce._start_callback_server(sock)
+        try:
+            urllib.request.urlopen(
+                f"http://127.0.0.1:{port}/callback?error=access_denied", timeout=5
+            ).read()
+            assert result.error == "access_denied"
+            assert result.code == ""
+        finally:
+            server.shutdown()
+            server.server_close()
+
+    def test_closing_the_server_releases_the_port(self) -> None:
+        sock = pkce._bind_callback_socket()
+        port = sock.getsockname()[1]
+        server, _ = pkce._start_callback_server(sock)
+        server.shutdown()
+        server.server_close()
+
+        assert _nothing_is_listening(port)
+
+
+class TestTheFlowUsesThatSocket:
+    def test_the_redirect_uri_points_at_the_listening_port(self, monkeypatch) -> None:
+        """The browser's redirect has to reach the server the flow just started."""
+        seen: dict[str, str] = {}
+
+        def fake_open(url: str) -> None:
+            query = dict(urllib.parse.parse_qsl(urllib.parse.urlparse(url).query))
+            seen["redirect_uri"] = query["redirect_uri"]
+            # stand in for the browser following the provider's redirect
+            with urllib.request.urlopen(f"{query['redirect_uri']}?code=CODE123", timeout=5) as resp:
+                resp.read()
+
+        def fake_exchange(token_url, client_id, code, verifier, redirect_uri):
+            return {"access_token": "t", "code": code, "redirect_uri": redirect_uri}
+
+        monkeypatch.setattr(pkce.webbrowser, "open", fake_open)
+        monkeypatch.setattr(pkce, "_exchange_code", fake_exchange)
+
+        tokens = pkce.pkce_authorize(
+            authorization_url="https://auth.example.com/authorize",
+            token_url="https://auth.example.com/token",
+            client_id="devagent",
+            scopes=["mcp.read"],
+            timeout=10,
+        )
+
+        assert tokens["code"] == "CODE123"
+        assert tokens["redirect_uri"] == seen["redirect_uri"]
+        assert seen["redirect_uri"].startswith("http://127.0.0.1:")
+
+    def test_the_port_is_released_when_the_flow_ends(self, monkeypatch) -> None:
+        held: dict[str, int] = {}
+
+        def fake_open(url: str) -> None:
+            query = dict(urllib.parse.parse_qsl(urllib.parse.urlparse(url).query))
+            held["port"] = int(urllib.parse.urlparse(query["redirect_uri"]).port)
+            with urllib.request.urlopen(f"{query['redirect_uri']}?code=C", timeout=5) as resp:
+                resp.read()
+
+        monkeypatch.setattr(pkce.webbrowser, "open", fake_open)
+        monkeypatch.setattr(pkce, "_exchange_code", lambda *a, **k: {"access_token": "t"})
+
+        pkce.pkce_authorize(
+            authorization_url="https://auth.example.com/authorize",
+            token_url="https://auth.example.com/token",
+            client_id="devagent",
+            scopes=["mcp.read"],
+            timeout=10,
+        )
+
+        assert _nothing_is_listening(held["port"])