phoenix-channels-python-client 0.2.4__py3-none-any.whl → 0.3.0__py3-none-any.whl

phoenix_channels_python_client/__init__.py

@@ -1,5 +1,4 @@
-"""
-Phoenix Channels Python Client
+"""Phoenix Channels Python Client.
 
 A Python client library for connecting to Phoenix Channels.
 """
@@ -12,14 +11,15 @@ from phoenix_channels_python_client.client import (
     PHXChannelsClient,
     ReconnectCallback,
     ReconnectPolicy,
+    TopicLostCallback,
 )
 from phoenix_channels_python_client.protocol_handler import (
-    PHXProtocolHandler,
     PhoenixChannelsProtocolVersion,
+    PHXProtocolHandler,
 )
 from phoenix_channels_python_client.utils import setup_logging
 
-__version__ = "0.2.4"
+__version__ = "0.3.0"
 __author__ = "Phoenix Channels Python Client"
 
 __all__ = [
@@ -30,5 +30,6 @@ __all__ = [
     "PhoenixChannelsProtocolVersion",
     "ReconnectCallback",
     "ReconnectPolicy",
+    "TopicLostCallback",
     "setup_logging",
 ]

phoenix_channels_python_client/client.py

@@ -5,26 +5,24 @@ import logging
 from collections import deque
 from collections.abc import Awaitable, Callable
 from types import TracebackType
+from typing import Self
 from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
 
 from websockets import ClientConnection
 
 from phoenix_channels_python_client.client_state_machine import transition_client_state
-from phoenix_channels_python_client.client_types import (
-    ClientState,
-    ReconnectPolicy,
-    reconnect_policy_is_invalid,
-    validate_reconnect_policy,
-)
+from phoenix_channels_python_client.client_types import ClientState, ReconnectPolicy
 from phoenix_channels_python_client.exceptions import PHXConnectionError
 from phoenix_channels_python_client.protocol_handler import (
-    PHXProtocolHandler,
+    DEFAULT_PROTOCOL_VERSION,
     PhoenixChannelsProtocolVersion,
+    PHXProtocolHandler,
 )
 from phoenix_channels_python_client.reconnect_controller import ReconnectControllerMixin
 from phoenix_channels_python_client.supervisor import SupervisorMixin
 from phoenix_channels_python_client.topic_runtime import TopicRuntimeMixin
 from phoenix_channels_python_client.topic_subscription import TopicSubscription
+from phoenix_channels_python_client.utils import cancel_and_wait
 
 logger = logging.getLogger(__name__)
 
@@ -34,6 +32,7 @@ DisconnectCallback = Callable[[Exception | None], Awaitable[None]]
 # on the message-routing hot path, so blocking work here stalls heartbeat
 # sends and message routing for every connection.
 HeartbeatAckCallback = Callable[[], None]
+TopicLostCallback = Callable[[str, Exception], Awaitable[None]]
 
 
 def _build_channel_socket_urls(
@@ -73,12 +72,13 @@ def _build_channel_socket_urls(
 
 
 class PHXChannelsClient(SupervisorMixin, TopicRuntimeMixin, ReconnectControllerMixin):
-    def __init__(
+    def __init__(  # noqa: PLR0913  # the public options
         self,
         websocket_url: str,
         api_key: str,
         *,
-        protocol_version: PhoenixChannelsProtocolVersion = PhoenixChannelsProtocolVersion.V2,
+        protocol_version: PhoenixChannelsProtocolVersion
+        | str = DEFAULT_PROTOCOL_VERSION,
         auto_reconnect: bool = True,
         reconnect_policy: ReconnectPolicy | None = None,
         join_timeout_s: float = 10.0,
@@ -89,8 +89,9 @@ class PHXChannelsClient(SupervisorMixin, TopicRuntimeMixin, ReconnectControllerM
         on_reconnect: ReconnectCallback | None = None,
         on_disconnect: DisconnectCallback | None = None,
         on_heartbeat_ack: HeartbeatAckCallback | None = None,
+        on_topic_lost: TopicLostCallback | None = None,
         additional_headers: dict[str, str] | None = None,
-    ):
+    ) -> None:
         self.logger = logger
 
         if heartbeat_interval_s is not None and heartbeat_interval_s <= 0:
@@ -104,20 +105,12 @@ class PHXChannelsClient(SupervisorMixin, TopicRuntimeMixin, ReconnectControllerM
             raise ValueError("max_topic_queue_size must be > 0")
         if callback_drain_timeout_s <= 0:
             raise ValueError("callback_drain_timeout_s must be > 0")
-        try:
-            validate_reconnect_policy(reconnect_policy or ReconnectPolicy())
-        except ValueError as exc:
-            raise ValueError("Invalid reconnect policy configuration") from exc
-
-        vsn = (
-            "2.0.0"
-            if protocol_version == PhoenixChannelsProtocolVersion.V2
-            else "1.0.0"
-        )
+
+        self._protocol_handler = PHXProtocolHandler(protocol_version)
         connect_url, redacted_url = _build_channel_socket_urls(
             websocket_url=websocket_url,
             api_key=api_key,
-            vsn=vsn,
+            vsn=self._protocol_handler.protocol_version.vsn,
         )
         self.channel_socket_url = connect_url
         self.channel_socket_url_redacted = redacted_url
@@ -137,7 +130,6 @@ class PHXChannelsClient(SupervisorMixin, TopicRuntimeMixin, ReconnectControllerM
 
         self.connection: ClientConnection | None = None
         self._topic_subscriptions: dict[str, TopicSubscription] = {}
-        self._protocol_handler = PHXProtocolHandler(protocol_version)
         self._ref_counter = 0
         self._state = ClientState.CLOSED
         self._topics_lock = asyncio.Lock()
@@ -145,6 +137,7 @@ class PHXChannelsClient(SupervisorMixin, TopicRuntimeMixin, ReconnectControllerM
         self._connected_event = asyncio.Event()
         self._conn_generation = 0
         self._supervisor_task: asyncio.Task[None] | None = None
+        self._shutdown_task: asyncio.Task[None] | None = None
         self._message_routing_task: asyncio.Task[None] | None = None
         self._initial_connection_future: asyncio.Future[None] | None = None
         self._rapid_disconnects: deque[float] = deque()
@@ -155,21 +148,33 @@ class PHXChannelsClient(SupervisorMixin, TopicRuntimeMixin, ReconnectControllerM
         self._on_reconnect = on_reconnect
         self._on_disconnect = on_disconnect
         self._on_heartbeat_ack = on_heartbeat_ack
+        self._on_topic_lost = on_topic_lost
+        self._lifecycle_tasks: set[asyncio.Task[None]] = set()
         self._forced_close_pending = False
 
-    @staticmethod
-    def reconnect_policy_is_invalid(policy: ReconnectPolicy) -> bool:
-        return reconnect_policy_is_invalid(policy)
+    @property
+    def _active_shutdown(self) -> asyncio.Task[None] | None:
+        task = self._shutdown_task
+        return task if task is not None and not task.done() else None
 
-    async def __aenter__(self) -> PHXChannelsClient:
+    @property
+    def _is_fully_closed(self) -> bool:
+        return (
+            self._state == ClientState.CLOSED
+            and not self._topic_subscriptions
+            and self.connection is None
+        )
+
+    async def __aenter__(self) -> Self:
         self.logger.debug("Entering PHXChannelsClient context")
-        if self._state != ClientState.CLOSED:
+        if self._state != ClientState.CLOSED or self._active_shutdown is not None:
             raise PHXConnectionError("Client is already running")
 
         self._shutdown_event.clear()
         self._connected_event.clear()
         self._rapid_disconnects.clear()
         self._terminal_error = None
+        self._conn_generation = 0
         self._initial_connection_future = asyncio.get_running_loop().create_future()
         self._transition_state(ClientState.CONNECTING)
 
@@ -192,57 +197,58 @@ class PHXChannelsClient(SupervisorMixin, TopicRuntimeMixin, ReconnectControllerM
         self.logger.debug("Leaving PHXChannelsClient context")
         await self.shutdown("Leaving PHXChannelsClient context")
 
-    async def shutdown(
-        self,
-        reason: str,
-    ) -> None:
-        if (
-            self._state == ClientState.CLOSED
-            and not self._topic_subscriptions
-            and self.connection is None
-        ):
-            return
-
+    async def shutdown(self, reason: str) -> None:
+        task = self._active_shutdown
+        if task is None:
+            if self._is_fully_closed:
+                return
+            task = self._shutdown_task = asyncio.create_task(self._shutdown(reason))
+        # Every caller waits for the same shutdown, and one caller's
+        # cancellation never aborts it.
+        await asyncio.shield(task)
+
+    async def _shutdown(self, reason: str) -> None:
         self.logger.info("Event loop shutting down! reason=%s", reason)
 
         if self._state not in (ClientState.SHUTTING_DOWN, ClientState.CLOSED):
             self._transition_state(ClientState.SHUTTING_DOWN)
-
         self._shutdown_event.set()
 
-        topics_to_unsubscribe = list(self._topic_subscriptions.keys())
-        if topics_to_unsubscribe:
-            unsubscribe_tasks = [
-                self.unsubscribe_from_topic(topic, _allow_disconnected=True)
-                for topic in topics_to_unsubscribe
-            ]
-            results = await asyncio.gather(*unsubscribe_tasks, return_exceptions=True)
-
-            for topic, result in zip(topics_to_unsubscribe, results, strict=False):
-                if isinstance(result, Exception):
-                    self.logger.warning(
-                        "Failed to unsubscribe from topic %s during shutdown: %s",
-                        topic,
-                        result,
-                    )
-
-        if self._supervisor_task and not self._supervisor_task.done():
-            self._supervisor_task.cancel()
-            try:
-                await self._supervisor_task
-            except asyncio.CancelledError:
-                self.logger.debug("Supervisor task cancelled during shutdown")
-
+        await self._unsubscribe_all()
+        # Includes on_topic_lost callbacks still running for topics already gone.
+        await cancel_and_wait(*self._lifecycle_tasks)
+        await self._stop_supervisor()
         await self._cleanup_connection()
         self._connected_event.clear()
         self._transition_state(ClientState.CLOSED)
 
+    async def _unsubscribe_all(self) -> None:
+        topics = list(self._topic_subscriptions)
+        results = await asyncio.gather(
+            *(
+                self.unsubscribe_from_topic(topic, _allow_disconnected=True)
+                for topic in topics
+            ),
+            return_exceptions=True,
+        )
+        for topic, result in zip(topics, results, strict=True):
+            if isinstance(result, Exception):
+                self.logger.warning(
+                    "Failed to unsubscribe from topic %s during shutdown: %s",
+                    topic,
+                    result,
+                )
+
+    async def _stop_supervisor(self) -> None:
+        # Cleared, so run_forever() after the stop reports the client isn't running.
+        supervisor, self._supervisor_task = self._supervisor_task, None
+        if supervisor is not None and not supervisor.done():
+            await cancel_and_wait(supervisor)
+
     def _transition_state(self, new_state: ClientState) -> None:
         if self._state == new_state:
             return
 
         transitioned_state = transition_client_state(self._state, new_state)
-        self.logger.debug(
-            "Client state transition: %s -> %s", self._state.value, new_state.value
-        )
+        self.logger.debug("Client state transition: %s -> %s", self._state, new_state)
         self._state = transitioned_state

phoenix_channels_python_client/client_state_machine.py

@@ -31,8 +31,6 @@ def transition_client_state(
     }
 
     if new_state not in allowed_transitions[current]:
-        raise RuntimeError(
-            f"Invalid state transition {current.value} -> {new_state.value}"
-        )
+        raise RuntimeError(f"Invalid state transition {current} -> {new_state}")
 
     return new_state

phoenix_channels_python_client/client_types.py

@@ -1,12 +1,22 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from enum import Enum
+from enum import StrEnum
+from typing import Self
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    NonNegativeFloat,
+    NonNegativeInt,
+    PositiveFloat,
+    model_validator,
+)
 
 from phoenix_channels_python_client.exceptions import PHXConnectionError
 
 
-class ClientState(Enum):
+class ClientState(StrEnum):
     CONNECTING = "connecting"
     CONNECTED = "connected"
     RECONNECTING = "reconnecting"
@@ -14,28 +24,45 @@ class ClientState(Enum):
     CLOSED = "closed"
 
 
-@dataclass(frozen=True)
-class ReconnectPolicy:
-    base_delay_s: float = 0.5
-    factor: float = 2.0
-    max_delay_s: float = 30.0
-    stable_reset_s: float = 60.0
+# Each (low, high) pair of fields where high must not be below low.
+_ORDERED_POLICY_FIELDS = (
+    ("service_restart_min_delay_s", "service_restart_max_delay_s"),
+    ("try_again_later_min_delay_s", "try_again_later_max_delay_s"),
+    ("rapid_cooldown_base_s", "rapid_cooldown_max_s"),
+    ("rapid_hold_down_jitter_low_ratio", "rapid_hold_down_jitter_high_ratio"),
+)
+
+
+class ReconnectPolicy(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    base_delay_s: NonNegativeFloat = 0.5
+    factor: PositiveFloat = 2.0
+    max_delay_s: NonNegativeFloat = 30.0
+    stable_reset_s: PositiveFloat = 60.0
     reconnect_on_normal_close: bool = False
     policy_violation_is_terminal: bool = True
-    service_restart_min_delay_s: float = 1.0
-    service_restart_max_delay_s: float = 5.0
-    try_again_later_min_delay_s: float = 30.0
-    try_again_later_max_delay_s: float = 60.0
-    rapid_disconnect_uptime_s: float = 5.0
-    rapid_window_s: float = 60.0
-    rapid_first_min_delay_s: float = 2.0
-    rapid_second_min_delay_s: float = 10.0
-    rapid_cooldown_base_s: float = 60.0
-    rapid_cooldown_step_s: float = 30.0
-    rapid_cooldown_max_s: float = 300.0
-    rapid_suppress_disconnect_count: int = 10
-    rapid_hold_down_jitter_low_ratio: float = 0.25
-    rapid_hold_down_jitter_high_ratio: float = 1.0
+    service_restart_min_delay_s: NonNegativeFloat = 1.0
+    service_restart_max_delay_s: NonNegativeFloat = 5.0
+    try_again_later_min_delay_s: NonNegativeFloat = 30.0
+    try_again_later_max_delay_s: NonNegativeFloat = 60.0
+    rapid_disconnect_uptime_s: NonNegativeFloat = 5.0
+    rapid_window_s: PositiveFloat = 60.0
+    rapid_first_min_delay_s: NonNegativeFloat = 2.0
+    rapid_second_min_delay_s: NonNegativeFloat = 10.0
+    rapid_cooldown_base_s: NonNegativeFloat = 60.0
+    rapid_cooldown_step_s: NonNegativeFloat = 30.0
+    rapid_cooldown_max_s: NonNegativeFloat = 300.0
+    rapid_suppress_disconnect_count: NonNegativeInt = 10
+    rapid_hold_down_jitter_low_ratio: NonNegativeFloat = 0.25
+    rapid_hold_down_jitter_high_ratio: NonNegativeFloat = 1.0
+
+    @model_validator(mode="after")
+    def _check_ordering(self) -> Self:
+        for low, high in _ORDERED_POLICY_FIELDS:
+            if getattr(self, high) < getattr(self, low):
+                raise ValueError(f"{high} must be >= {low}")
+        return self
 
 
 @dataclass(frozen=True)
@@ -44,59 +71,3 @@ class ReconnectDecision:
     min_delay_s: float | None = None
     max_delay_s: float | None = None
     terminal_error: PHXConnectionError | None = None
-
-
-def validate_reconnect_policy(policy: ReconnectPolicy) -> None:
-    if policy.base_delay_s < 0:
-        raise ValueError("base_delay_s must be >= 0")
-    if policy.factor <= 0:
-        raise ValueError("factor must be > 0")
-    if policy.max_delay_s < 0:
-        raise ValueError("max_delay_s must be >= 0")
-    if policy.stable_reset_s <= 0:
-        raise ValueError("stable_reset_s must be > 0")
-    if policy.service_restart_min_delay_s < 0:
-        raise ValueError("service_restart_min_delay_s must be >= 0")
-    if policy.service_restart_max_delay_s < policy.service_restart_min_delay_s:
-        raise ValueError(
-            "service_restart_max_delay_s must be >= service_restart_min_delay_s"
-        )
-    if policy.try_again_later_min_delay_s < 0:
-        raise ValueError("try_again_later_min_delay_s must be >= 0")
-    if policy.try_again_later_max_delay_s < policy.try_again_later_min_delay_s:
-        raise ValueError(
-            "try_again_later_max_delay_s must be >= try_again_later_min_delay_s"
-        )
-    if policy.rapid_disconnect_uptime_s < 0:
-        raise ValueError("rapid_disconnect_uptime_s must be >= 0")
-    if policy.rapid_window_s <= 0:
-        raise ValueError("rapid_window_s must be > 0")
-    if policy.rapid_first_min_delay_s < 0:
-        raise ValueError("rapid_first_min_delay_s must be >= 0")
-    if policy.rapid_second_min_delay_s < 0:
-        raise ValueError("rapid_second_min_delay_s must be >= 0")
-    if policy.rapid_cooldown_base_s < 0:
-        raise ValueError("rapid_cooldown_base_s must be >= 0")
-    if policy.rapid_cooldown_step_s < 0:
-        raise ValueError("rapid_cooldown_step_s must be >= 0")
-    if policy.rapid_cooldown_max_s < policy.rapid_cooldown_base_s:
-        raise ValueError("rapid_cooldown_max_s must be >= rapid_cooldown_base_s")
-    if policy.rapid_suppress_disconnect_count < 0:
-        raise ValueError("rapid_suppress_disconnect_count must be >= 0")
-    if policy.rapid_hold_down_jitter_low_ratio < 0:
-        raise ValueError("rapid_hold_down_jitter_low_ratio must be >= 0")
-    if (
-        policy.rapid_hold_down_jitter_high_ratio
-        < policy.rapid_hold_down_jitter_low_ratio
-    ):
-        raise ValueError(
-            "rapid_hold_down_jitter_high_ratio must be >= rapid_hold_down_jitter_low_ratio"
-        )
-
-
-def reconnect_policy_is_invalid(policy: ReconnectPolicy) -> bool:
-    try:
-        validate_reconnect_policy(policy)
-    except ValueError:
-        return True
-    return False

phoenix_channels_python_client/exceptions.py

@@ -8,5 +8,3 @@ class PHXTopicError(PHXClientError):
 
 class PHXConnectionError(PHXClientError):
     """Raised when there's an error connecting to the Phoenix WebSocket server."""
-
-    pass

phoenix_channels_python_client/phx_messages.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
 from enum import Enum, unique
 from functools import cached_property
-from typing import Any, NewType
+from typing import Annotated, Any, NewType
 
+from pydantic import BaseModel, ConfigDict, Field
 
 PHOENIX_TOPIC = "phoenix"
 
@@ -21,16 +21,25 @@ class PHXEvent(Enum):
         return self.value
 
 
+# The client handles these itself: a crash rejoins, a close loses the topic.
+CLIENT_LIFECYCLE_EVENTS = frozenset({PHXEvent.error, PHXEvent.close})
+
 UserEvent = NewType("UserEvent", str)
 # Compatibility alias for existing imports and call sites.
 Event = UserEvent
+# Not a PHXEvent member: that would reroute make_message and change a public enum.
+HEARTBEAT_EVENT = Event("heartbeat")
 ChannelEvent = PHXEvent | UserEvent
 
+# Servers may send a ref as a JSON number; the client matches refs as strings.
+Ref = Annotated[str | None, Field(coerce_numbers_to_str=True)]
+
+
+class BasePHXMessage(BaseModel):
+    model_config = ConfigDict(frozen=True)
 
-@dataclass(frozen=True)
-class BasePHXMessage:
     topic: str
-    ref: str | None
+    ref: Ref
     payload: dict[str, Any]
 
     @cached_property
@@ -41,16 +50,14 @@ class BasePHXMessage:
         return subtopic
 
 
-@dataclass(frozen=True)
 class PHXMessage(BasePHXMessage):
     event: Event
-    join_ref: str | None = None
+    join_ref: Ref = None
 
 
-@dataclass(frozen=True)
 class PHXEventMessage(BasePHXMessage):
     event: PHXEvent
-    join_ref: str | None = None
+    join_ref: Ref = None
 
 
 ChannelMessage = PHXMessage | PHXEventMessage