From 62fb646ff233b4154c330ff5e9802a832faae1cd Mon Sep 17 00:00:00 2001 From: Darvell Date: Tue, 19 May 2026 14:25:47 -0700 Subject: [PATCH 1/2] fix: move Phoenix auth to x-api-key header --- .pre-commit-config.yaml | 2 +- phoenix_channels_python_client/client.py | 32 +++++--------------- phoenix_channels_python_client/supervisor.py | 6 +++- pyproject.toml | 3 +- tests/test_internal_components.py | 19 ++++++------ tests/test_reconnect_policy_invariants.py | 21 +++++++++++-- tests/test_v1_protocol/conftest.py | 26 +++++++++++----- tests/test_v1_protocol/test_client_v1.py | 26 ++++++++++++++++ tests/test_v2_protocol/conftest.py | 26 +++++++++++----- tests/test_v2_protocol/test_client_v2.py | 28 +++++++++++++++++ uv.lock | 4 +-- 11 files changed, 135 insertions(+), 58 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9d6fa3e..8095ab2 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -28,7 +28,7 @@ repos: - id: pyrefly name: pyrefly - entry: uv run --extra examples pyrefly check + entry: uv run --extra examples pyrefly check phoenix_channels_python_client tests --project-excludes **/__pycache__ --disable-project-excludes-heuristics=true language: system types: [python] pass_filenames: false \ No newline at end of file diff --git a/phoenix_channels_python_client/client.py b/phoenix_channels_python_client/client.py index 6f70699..1939eb8 100644 --- a/phoenix_channels_python_client/client.py +++ b/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() diff --git a/phoenix_channels_python_client/supervisor.py b/phoenix_channels_python_client/supervisor.py index b4729f3..0e7d67d 100644 --- a/phoenix_channels_python_client/supervisor.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index 8e97990..8c72cab 100644 --- a/pyproject.toml +++ b/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] diff --git a/tests/test_internal_components.py b/tests/test_internal_components.py index ac245c8..6348b0d 100644 --- a/tests/test_internal_components.py +++ b/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() diff --git a/tests/test_reconnect_policy_invariants.py b/tests/test_reconnect_policy_invariants.py index 79dcb45..ba4ffaa 100644 --- a/tests/test_reconnect_policy_invariants.py +++ b/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: diff --git a/tests/test_v1_protocol/conftest.py b/tests/test_v1_protocol/conftest.py index 7414fc8..78c3cd1 100644 --- a/tests/test_v1_protocol/conftest.py +++ b/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: diff --git a/tests/test_v1_protocol/test_client_v1.py b/tests/test_v1_protocol/test_client_v1.py index 93a3626..58f6ff5 100644 --- a/tests/test_v1_protocol/test_client_v1.py +++ b/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 diff --git a/tests/test_v2_protocol/conftest.py b/tests/test_v2_protocol/conftest.py index 70873a5..4e36b59 100644 --- a/tests/test_v2_protocol/conftest.py +++ b/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,13 +48,13 @@ 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.""" @@ -60,8 +62,10 @@ async def handler(self, websocket: ServerConnection) -> None: 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: diff --git a/tests/test_v2_protocol/test_client_v2.py b/tests/test_v2_protocol/test_client_v2.py index 207c2a2..e432c0f 100644 --- a/tests/test_v2_protocol/test_client_v2.py +++ b/tests/test_v2_protocol/test_client_v2.py @@ -40,6 +40,26 @@ async def wait_for_condition( return False +@pytest.mark.asyncio +async def test_websocket_auth_uses_header_not_query_param( + phoenix_server: FakePhoenixServerV2, +): + async with PHXChannelsClient( + f"{phoenix_server.url}?api_key=stale&debug=true", + api_key="test_key", + protocol_version=PhoenixChannelsProtocolVersion.V2, + ): + 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=2.0.0" in paths[0] + assert "api_key" not in paths[0] + assert "test_key" not in paths[0] + + @pytest.mark.asyncio async def test_subscribe_to_topic_succeeds_when_subscribing_to_valid_topic( phoenix_server: FakePhoenixServerV2, @@ -806,6 +826,14 @@ async def message_callback(message: ChannelMessage): interval=0.05, ) assert result, "Client did not reconnect and rejoin topic in time" + 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"] # Verify topic was re-subscribed assert "test-topic" in client.get_current_subscriptions() diff --git a/uv.lock b/uv.lock index 96ff824..862e787 100644 --- a/uv.lock +++ b/uv.lock @@ -318,9 +318,9 @@ requires-dist = [ { name = "pytest-asyncio", marker = "extra == 'dev'" }, { name = "pytest-asyncio", marker = "extra == 'test'" }, { name = "ruff", marker = "extra == 'dev'" }, - { name = "websockets", specifier = ">=10.0" }, + { name = "websockets", specifier = ">=16.0" }, ] -provides-extras = ["dev", "test"] +provides-extras = ["dev", "test", "examples"] [[package]] name = "platformdirs" From 2b3c4e2d270e5dc76a91e7522569c1b59cfc6267 Mon Sep 17 00:00:00 2001 From: Darvell <1046915+darvell@users.noreply.github.com> Date: Tue, 26 May 2026 18:24:17 -0700 Subject: [PATCH 2/2] fix: run pyrefly pre-commit with dev extra --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8095ab2..6afbbe9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -28,7 +28,7 @@ repos: - id: pyrefly name: pyrefly - entry: uv run --extra examples pyrefly check phoenix_channels_python_client tests --project-excludes **/__pycache__ --disable-project-excludes-heuristics=true + 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 \ No newline at end of file