Patch circular import

Author: blast-hardcheese

src/replit_river/rate_limiter.py

@@ -4,12 +4,18 @@
 from contextvars import Context
 from typing import Literal
 
+from replit_river.error_schema import RiverException
 from replit_river.transport_options import ConnectionRetryOptions
-from replit_river.v2.client_transport import BudgetExhaustedException
 
 logger = logging.getLogger(__name__)
 
 
+class BudgetExhaustedException(RiverException):
+    def __init__(self, code: str, message: str, client_id: str) -> None:
+        super().__init__(code, message)
+        self.client_id = client_id
+
+
 class LeakyBucketRateLimit:
     """Asynchronous leaky bucket rate limiter.
 

src/replit_river/v2/client_transport.py

@@ -26,12 +26,6 @@ def __init__(self, code: str, message: str, client_id: str) -> None:
         self.client_id = client_id
 
 
-class BudgetExhaustedException(RiverException):
-    def __init__(self, code: str, message: str, client_id: str) -> None:
-        super().__init__(code, message)
-        self.client_id = client_id
-
-
 class ClientTransport(Generic[HandshakeMetadataType]):
     _session: Session | None
 

src/replit_river/v2/session.py

@@ -70,9 +70,6 @@
     TransportOptions,
     UriAndMetadata,
 )
-from replit_river.v2.client_transport import (
-    PROTOCOL_VERSION,
-)
 
 STREAM_CANCEL_BIT_TYPE = Literal[0b00100]
 STREAM_CANCEL_BIT: STREAM_CANCEL_BIT_TYPE = 0b00100
@@ -233,6 +230,7 @@ async def ensure_connected[HandshakeMetadata](
         uri_and_metadata_factory: Callable[
             [], Awaitable[UriAndMetadata[HandshakeMetadata]]
         ],  # noqa: E501
+        protocol_version: str,
     ) -> None:
         """
         Either return immediately or establish a websocket connection and return
@@ -266,7 +264,7 @@ async def ensure_connected[HandshakeMetadata](
                             HandshakeMetadata
                         ](  # noqa: E501
                             type="HANDSHAKE_REQ",
-                            protocolVersion=PROTOCOL_VERSION,
+                            protocolVersion=protocol_version,
                             sessionId=self.session_id,
                             metadata=uri_and_metadata["metadata"],
                             expectedSessionState=expectedSessionState,