fix(websocket): move Phoenix auth to x-api-key header

.pre-commit-config.yaml

@@ -28,7 +28,7 @@ repos:
 
       - id: pyrefly
         name: pyrefly
-        entry: uv run --extra examples pyrefly check
+        entry: uv run --extra dev pyrefly check phoenix_channels_python_client tests --project-excludes **/__pycache__ --disable-project-excludes-heuristics=true
         language: system
         types: [python]
         pass_filenames: false

phoenix_channels_python_client/client.py

@@ -32,40 +32,23 @@
 DisconnectCallback = Callable[[Exception | None], Awaitable[None]]
 
 
-def _build_channel_socket_urls(
-    websocket_url: str, api_key: str, vsn: str
-) -> tuple[str, str]:
+def _build_channel_socket_url(websocket_url: str, vsn: str) -> str:
     split_url = urlsplit(websocket_url)
     query_params = parse_qsl(split_url.query, keep_blank_values=True)
     filtered = [
         (key, value) for key, value in query_params if key not in {"api_key", "vsn"}
     ]
-    with_auth = [*filtered, ("api_key", api_key), ("vsn", vsn)]
+    query = urlencode([*filtered, ("vsn", vsn)])
 
-    connect_url = urlunsplit(
+    return urlunsplit(
         (
             split_url.scheme,
             split_url.netloc,
             split_url.path,
-            urlencode(with_auth),
+            query,
             split_url.fragment,
         )
     )
-    redacted_url = urlunsplit(
-        (
-            split_url.scheme,
-            split_url.netloc,
-            split_url.path,
-            urlencode(
-                [
-                    (key, "***" if key == "api_key" else value)
-                    for key, value in with_auth
-                ]
-            ),
-            split_url.fragment,
-        )
-    )
-    return connect_url, redacted_url
 
 
 class PHXChannelsClient(SupervisorMixin, TopicRuntimeMixin, ReconnectControllerMixin):
@@ -108,13 +91,12 @@ def __init__(
             if protocol_version == PhoenixChannelsProtocolVersion.V2
             else "1.0.0"
         )
-        connect_url, redacted_url = _build_channel_socket_urls(
+        self.channel_socket_url = _build_channel_socket_url(
             websocket_url=websocket_url,
-            api_key=api_key,
             vsn=vsn,
         )
-        self.channel_socket_url = connect_url
-        self.channel_socket_url_redacted = redacted_url
+        self.channel_socket_url_redacted = self.channel_socket_url
+        self.channel_socket_headers = {"x-api-key": api_key}
 
         self.auto_reconnect = auto_reconnect
         self.reconnect_policy = reconnect_policy or ReconnectPolicy()

phoenix_channels_python_client/supervisor.py

@@ -45,6 +45,7 @@ class SupervisorMixin:
     logger: logging.Logger
     channel_socket_url: str
     channel_socket_url_redacted: str
+    channel_socket_headers: dict[str, str]
     auto_reconnect: bool
     reconnect_policy: ReconnectPolicy
     connection: ClientConnection | None
@@ -137,7 +138,10 @@ async def _supervisor_loop(self) -> None:
                     self.channel_socket_url,
                 )
                 try:
-                    connection = await connect(self.channel_socket_url)
+                    connection = await connect(
+                        self.channel_socket_url,
+                        additional_headers=self.channel_socket_headers,
+                    )
                 except asyncio.CancelledError:
                     raise
                 except Exception as exc:

pyproject.toml

@@ -11,7 +11,7 @@ license = "MIT"
 authors = [{ name = "Phoenix Channels Python Client" }]
 requires-python = ">=3.11"
 dependencies = [
-    "websockets>=10.0",
+    "websockets>=16.0",
 ]
 classifiers = [
     "Development Status :: 3 - Alpha",
@@ -36,6 +36,7 @@ test = [
     "pytest",
     "pytest-asyncio",
 ]
+examples = []
 
 # Package discovery for flat layout (phoenix_channels_python_client/ in root)
 [tool.setuptools.packages.find]

tests/test_internal_components.py

@@ -126,8 +126,9 @@ def __init__(self) -> None:
 class _SupervisorHarness(SupervisorMixin):
     def __init__(self) -> None:
         self.logger = logging.getLogger(__name__)
-        self.channel_socket_url = "ws://unit-test/socket"
-        self.channel_socket_url_redacted = "ws://unit-test/socket?api_key=***"
+        self.channel_socket_url = "ws://unit-test/socket?vsn=2.0.0"
+        self.channel_socket_url_redacted = "ws://unit-test/socket?vsn=2.0.0"
+        self.channel_socket_headers = {"x-api-key": "test-key"}
         self.auto_reconnect = True
         self.reconnect_policy = ReconnectPolicy(stable_reset_s=0.0)
         self.connection: ClientConnection | None = None
@@ -719,7 +720,7 @@ async def test_supervisor_connect_failures_and_terminal_suppression(
     harness = _SupervisorHarness()
     harness.auto_reconnect = False
 
-    async def fail_connect(_: str) -> ClientConnection:
+    async def fail_connect(_url: str, **_kwargs: object) -> ClientConnection:
         raise RuntimeError("connect fail")
 
     monkeypatch.setattr(
@@ -747,7 +748,7 @@ async def test_supervisor_initial_connect_retries_before_failing_enter(
     harness = _SupervisorHarness()
     attempts = 0
 
-    async def fail_then_connect(_: str) -> ClientConnection:
+    async def fail_then_connect(_url: str, **_kwargs: object) -> ClientConnection:
         nonlocal attempts
         attempts += 1
         if attempts == 1:
@@ -789,7 +790,7 @@ async def test_supervisor_routing_failure_disconnect_decisions_and_cleanup(
 
     socket = _FakeSocket(close_code=1012, close_reason="restart")
 
-    async def connect_once(_: str) -> ClientConnection:
+    async def connect_once(_url: str, **_kwargs: object) -> ClientConnection:
         return cast(ClientConnection, socket)
 
     monkeypatch.setattr(
@@ -807,7 +808,7 @@ async def connect_once(_: str) -> ClientConnection:
     no_reconnect.disconnect_decision = ReconnectDecision(should_reconnect=False)
     monkeypatch.setattr(
         "phoenix_channels_python_client.supervisor.connect",
-        lambda _: asyncio.sleep(0, result=cast(ClientConnection, _FakeSocket())),
+        lambda _, **__: asyncio.sleep(0, result=cast(ClientConnection, _FakeSocket())),
     )
     await no_reconnect._supervisor_loop()
     assert ClientState.CLOSED in no_reconnect.transition_history
@@ -910,7 +911,7 @@ async def on_disconnect(error: Exception | None) -> None:
 
     monkeypatch.setattr(
         "phoenix_channels_python_client.supervisor.connect",
-        lambda _: asyncio.sleep(0, result=cast(ClientConnection, socket)),
+        lambda _, **__: asyncio.sleep(0, result=cast(ClientConnection, socket)),
     )
     await harness._supervisor_loop()
 
@@ -933,7 +934,7 @@ async def on_reconnect() -> None:
 
     harness._on_reconnect = on_reconnect
 
-    async def connect_and_shutdown(_: str) -> ClientConnection:
+    async def connect_and_shutdown(_url: str, **_kwargs: object) -> ClientConnection:
         nonlocal connect_count
         connect_count += 1
         if connect_count >= 2:
@@ -972,7 +973,7 @@ async def bad_reconnect() -> None:
 
     monkeypatch.setattr(
         "phoenix_channels_python_client.supervisor.connect",
-        lambda _: asyncio.sleep(0, result=cast(ClientConnection, _FakeSocket())),
+        lambda _, **__: asyncio.sleep(0, result=cast(ClientConnection, _FakeSocket())),
     )
 
     await harness._supervisor_loop()

tests/test_reconnect_policy_invariants.py

@@ -25,11 +25,26 @@ def test_invalid_reconnect_policy_is_rejected() -> None:
         _ = _make_client(bad_policy)
 
 
-def test_client_maintains_redacted_socket_url_for_logging() -> None:
+def test_client_keeps_api_key_out_of_socket_urls() -> None:
     client = _make_client()
-    assert "api_key=test-key" in client.channel_socket_url
-    assert "api_key=%2A%2A%2A" in client.channel_socket_url_redacted
+    assert "vsn=2.0.0" in client.channel_socket_url
+    assert "api_key" not in client.channel_socket_url
+    assert "test-key" not in client.channel_socket_url
+    assert "api_key" not in client.channel_socket_url_redacted
     assert "test-key" not in client.channel_socket_url_redacted
+    assert client.channel_socket_headers == {"x-api-key": "test-key"}
+
+
+def test_client_strips_stale_api_key_from_configured_socket_url() -> None:
+    client = PHXChannelsClient(
+        "ws://example.invalid/socket/websocket?api_key=stale-key&debug=true",
+        api_key="test-key",
+    )
+    assert "debug=true" in client.channel_socket_url
+    assert "vsn=2.0.0" in client.channel_socket_url
+    assert "api_key" not in client.channel_socket_url
+    assert "stale-key" not in client.channel_socket_url
+    assert "test-key" not in client.channel_socket_url
 
 
 def test_close_code_classification_uses_expected_semantics() -> None:

tests/test_v1_protocol/conftest.py

@@ -2,7 +2,7 @@
 
 import json
 from collections.abc import AsyncGenerator, Mapping
-from urllib.parse import parse_qs, urlparse
+from urllib.parse import urlparse
 
 import pytest_asyncio
 from websockets.asyncio.server import Server, ServerConnection, serve
@@ -23,6 +23,8 @@ def __init__(self, host: str = "localhost", port: int = 8765):
         self._client_ids: dict[ServerConnection, int] = {}
         self._client_path: dict[ServerConnection, str] = {}
         self._client_api_key: dict[ServerConnection, str] = {}
+        self._request_paths: list[str] = []
+        self._request_api_keys: list[str] = []
         self._next_client_id = 1
         self.close_on_join_ids: set[int] = set()
         self.close_on_join_code = 1012
@@ -45,21 +47,23 @@ def _extract_request_path(websocket: ServerConnection) -> str:
         return str(getattr(request, "path", ""))
 
     @staticmethod
-    def _extract_api_key(request_path: str) -> str:
-        parsed = urlparse(request_path)
-        query = parse_qs(parsed.query)
-        values = query.get("api_key")
-        if not values:
+    def _extract_api_key(websocket: ServerConnection) -> str:
+        request = getattr(websocket, "request", None)
+        headers = getattr(request, "headers", None)
+        if headers is None:
             return ""
-        return values[0]
+        value = headers.get("x-api-key")
+        return value if isinstance(value, str) else ""
 
     async def handler(self, websocket: ServerConnection) -> None:
         client_id = self._next_client_id
         self._next_client_id += 1
         request_path = self._extract_request_path(websocket)
         path_only = urlparse(request_path).path
-        api_key = self._extract_api_key(request_path)
+        api_key = self._extract_api_key(websocket)
 
+        self._request_paths.append(request_path)
+        self._request_api_keys.append(api_key)
         self.client_websocket = websocket
         self._clients.add(websocket)
         self._client_ids[websocket] = client_id
@@ -209,6 +213,12 @@ def get_connection_attempts(self, path: str) -> int:
     def list_client_connections(self) -> list[ServerConnection]:
         return list(self._clients)
 
+    def list_request_paths(self) -> list[str]:
+        return list(self._request_paths)
+
+    def list_request_api_keys(self) -> list[str]:
+        return list(self._request_api_keys)
+
     def get_client_id_for_path(self, path: str) -> int | None:
         for websocket, websocket_path in self._client_path.items():
             if websocket_path == path:

tests/test_v1_protocol/test_client_v1.py

@@ -62,6 +62,24 @@ def _test_reconnect_policy() -> ReconnectPolicy:
     )
 
 
+async def test_websocket_auth_uses_header_not_query_param(
+    phoenix_server: FakePhoenixServer,
+):
+    async with V1Client(
+        f"{phoenix_server.url}?api_key=stale&debug=true",
+        api_key="test_key",
+    ):
+        pass
+
+    assert phoenix_server.list_request_api_keys() == ["test_key"]
+    paths = phoenix_server.list_request_paths()
+    assert len(paths) == 1
+    assert "debug=true" in paths[0]
+    assert "vsn=1.0.0" in paths[0]
+    assert "api_key" not in paths[0]
+    assert "test_key" not in paths[0]
+
+
 async def test_subscribe_to_topic_succeeds_when_subscribing_to_valid_topic(
     phoenix_server: FakePhoenixServer,
 ):
@@ -448,6 +466,14 @@ async def callback(message: Message) -> None:
             ].current_join_ready.done(),
             timeout_s=2.0,
         )
+        assert len(phoenix_server.list_request_paths()) >= 2
+        assert all(
+            "api_key" not in path for path in phoenix_server.list_request_paths()
+        )
+        assert all(
+            "test_key" not in path for path in phoenix_server.list_request_paths()
+        )
+        assert phoenix_server.list_request_api_keys()[-2:] == ["test_key", "test_key"]
 
         current_join_ref = client.get_current_subscriptions()["test-topic"].join_ref
 

tests/test_v2_protocol/conftest.py

@@ -2,7 +2,7 @@
 
 import json
 from collections.abc import AsyncGenerator, Mapping
-from urllib.parse import parse_qs, urlparse
+from urllib.parse import urlparse
 
 import pytest_asyncio
 from websockets.asyncio.server import Server, ServerConnection, serve
@@ -23,6 +23,8 @@ def __init__(self, host: str = "localhost", port: int = 8765):
         self._client_ids: dict[ServerConnection, int] = {}
         self._client_path: dict[ServerConnection, str] = {}
         self._client_api_key: dict[ServerConnection, str] = {}
+        self._request_paths: list[str] = []
+        self._request_api_keys: list[str] = []
         self._next_client_id = 1
         self.close_on_join_ids: set[int] = set()
         self.close_on_join_code = 1012
@@ -46,22 +48,24 @@ def _extract_request_path(websocket: ServerConnection) -> str:
         return str(getattr(request, "path", ""))
 
     @staticmethod
-    def _extract_api_key(request_path: str) -> str:
-        parsed = urlparse(request_path)
-        query = parse_qs(parsed.query)
-        values = query.get("api_key")
-        if not values:
+    def _extract_api_key(websocket: ServerConnection) -> str:
+        request = getattr(websocket, "request", None)
+        headers = getattr(request, "headers", None)
+        if headers is None:
             return ""
-        return values[0]
+        value = headers.get("x-api-key")
+        return value if isinstance(value, str) else ""
 
     async def handler(self, websocket: ServerConnection) -> None:
         """Handle WebSocket connections and messages."""
         client_id = self._next_client_id
         self._next_client_id += 1
         request_path = self._extract_request_path(websocket)
         path_only = urlparse(request_path).path
-        api_key = self._extract_api_key(request_path)
+        api_key = self._extract_api_key(websocket)
 
+        self._request_paths.append(request_path)
+        self._request_api_keys.append(api_key)
         self.client_websocket = websocket
         self._clients.add(websocket)
         self._client_ids[websocket] = client_id
@@ -212,6 +216,12 @@ def get_connection_attempts(self, path: str) -> int:
     def list_client_connections(self) -> list[ServerConnection]:
         return list(self._clients)
 
+    def list_request_paths(self) -> list[str]:
+        return list(self._request_paths)
+
+    def list_request_api_keys(self) -> list[str]:
+        return list(self._request_api_keys)
+
     def get_client_id_for_path(self, path: str) -> int | None:
         for websocket, websocket_path in self._client_path.items():
             if websocket_path == path: