fixes and regresison

Author: jackyzha0

python-client/river/client.py

@@ -13,6 +13,7 @@
 
 from typing_extensions import TypedDict
 
+from river.session import SessionState
 from river.streams import Readable, Writable
 from river.transport import WebSocketClientTransport
 from river.types import (
@@ -259,8 +260,30 @@ def _handle_proc(
         if self._connect_on_invoke:
             transport.connect(to)
 
-        # Get the session and a send function
+        # Get the session and a send function.
+        # If connect() couldn't start (retry budget exhausted, transport
+        # closing, etc.) the session will be in NO_CONNECTION with no
+        # connect task in flight — fail immediately instead of hanging.
         session = transport._get_or_create_session(to)
+        connect_task = transport._connect_tasks.get(to)
+        has_active_connect = connect_task is not None and not connect_task.done()
+        if session.state == SessionState.NO_CONNECTION and not has_active_connect:
+            transport._delete_session(to, emit_closing=False)
+            res_readable = Readable()
+            res_readable._push_value(
+                err_result(
+                    UNEXPECTED_DISCONNECT_CODE,
+                    f"{to} connection failed",
+                )
+            )
+            res_readable._trigger_close()
+            req_writable = Writable(write_cb=lambda _: None, close_cb=None)
+            req_writable._closed = True
+            return {
+                "res_readable": res_readable,
+                "req_writable": req_writable,
+            }
+
         session_id = session.id
         try:
             send_fn = transport.get_session_bound_send_fn(to, session_id)

python-client/river/codegen/emitter.py

@@ -132,6 +132,17 @@ def render_service_client(svc: ServiceDef, ir: SchemaIR, import_prefix: str) ->
     has_upload = "upload" in proc_types
     has_subscription = "subscription" in proc_types
 
+    # Check if any annotation references Literal (e.g. const schemas)
+    all_annotations = []
+    for p in svc.procedures:
+        all_annotations.append(p.init_type.annotation)
+        all_annotations.append(p.output_type.annotation)
+        if p.input_type:
+            all_annotations.append(p.input_type.annotation)
+        if p.error_type:
+            all_annotations.append(p.error_type.annotation)
+    needs_literal = any("Literal[" in a for a in all_annotations)
+
     return _env.get_template("service_client.py.j2").render(
         service=svc,
         type_names=type_names,
@@ -140,6 +151,7 @@ def render_service_client(svc: ServiceDef, ir: SchemaIR, import_prefix: str) ->
         has_stream=has_stream,
         has_upload=has_upload,
         has_subscription=has_subscription,
+        needs_literal=needs_literal,
     )
 
 

python-client/river/codegen/schema.py

@@ -151,6 +151,7 @@ def __init__(self) -> None:
 
     def convert(self, raw: dict) -> SchemaIR:
         """Convert the top-level serialized schema dict to IR."""
+        self._typedicts = []
         services: list[ServiceDef] = []
         for svc_name, svc_data in raw.get("services", {}).items():
             svc_def = self._convert_service(svc_name, svc_data)
@@ -247,7 +248,11 @@ def _schema_to_typeref(self, schema: dict, name_hint: str) -> TypeRef:
         if "const" in schema:
             val = schema["const"]
             if isinstance(val, str):
-                return TypeRef(annotation=f'Literal["{val}"]')
+                # Use repr to handle all escaping (quotes, backslashes,
+                # control chars) then unwrap the outer quotes and re-wrap
+                # with double quotes for Literal["..."] syntax.
+                escaped = repr(val)[1:-1].replace('"', '\\"')
+                return TypeRef(annotation=f'Literal["{escaped}"]')
             return TypeRef(annotation=f"Literal[{val!r}]")
 
         # anyOf (union)

python-client/river/codegen/templates/service_client.py.j2

@@ -3,7 +3,11 @@
 from __future__ import annotations
 
 import asyncio
+{% if needs_literal %}
+from typing import Any, Literal
+{% else %}
 from typing import Any
+{% endif %}
 
 from river.client import (
     ErrResult,

python-client/river/streams.py

@@ -55,7 +55,7 @@ def is_readable(self) -> bool:
         return not self._locked and not self._broken
 
     def is_closed(self) -> bool:
-        """Whether the stream has been closed."""
+        """Whether the stream is fully consumed (closed and queue drained)."""
         return self._closed and len(self._queue) == 0
 
     def _has_values_in_queue(self) -> bool:

python-client/river/transport.py

@@ -235,6 +235,7 @@ def connect(self, to: str | None = None) -> None:
         session.state = SessionState.BACKING_OFF
 
         async def _do_connect():
+            ws = None
             try:
                 if backoff_ms > 0:
                     await asyncio.sleep(backoff_ms / 1000.0)
@@ -252,7 +253,9 @@ async def _do_connect():
                 session.state = SessionState.HANDSHAKING
                 await self._do_handshake(session, ws, to)
             except asyncio.CancelledError:
-                pass
+                # Clean up socket if we got cancelled mid-handshake
+                if ws is not None and session._ws is not ws:
+                    await ws.close()
             except Exception as e:
                 logger.debug("Connection attempt failed for %s: %s", to, e)
                 if not session._destroyed:
@@ -484,9 +487,12 @@ def _on_connection_failed(self, to: str) -> None:
 
         # Transition to NoConnection with grace period so the session
         # is eventually destroyed if reconnect doesn't succeed.
+        # Only start the grace period if one isn't already running,
+        # so repeated failures don't keep extending the deadline.
         loop = self._get_loop()
         session.state = SessionState.NO_CONNECTION
-        session.start_grace_period(loop)
+        if session._grace_period_task is None or session._grace_period_task.done():
+            session.start_grace_period(loop)
 
         if self._reconnect_on_connection_drop:
             self._try_reconnecting(to)

python-client/tests/conftest.py

@@ -8,6 +8,7 @@
 
 import os
 import re
+import selectors
 import signal
 import subprocess
 import sys
@@ -119,20 +120,44 @@ def _start_server(
     port = None
     deadline = time.monotonic() + 30
     assert proc.stdout is not None
-    while time.monotonic() < deadline:
-        line = proc.stdout.readline().decode("utf-8").strip()
-        if not line:
-            if proc.poll() is not None:
-                stderr = proc.stderr.read().decode("utf-8") if proc.stderr else ""
-                raise RuntimeError(
-                    f"{label} exited with code {proc.returncode}.\nstderr: {stderr}"
-                )
-            time.sleep(0.1)
-            continue
-        m = re.match(r"RIVER_PORT=(\d+)", line)
-        if m:
-            port = int(m.group(1))
-            break
+    sel = selectors.DefaultSelector()
+    sel.register(proc.stdout, selectors.EVENT_READ)
+    buf = b""
+    try:
+        while time.monotonic() < deadline:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                break
+            ready = sel.select(timeout=min(remaining, 1.0))
+            if not ready:
+                if proc.poll() is not None:
+                    stderr = proc.stderr.read().decode("utf-8") if proc.stderr else ""
+                    raise RuntimeError(
+                        f"{label} exited with code {proc.returncode}.\nstderr: {stderr}"
+                    )
+                continue
+            chunk = proc.stdout.read1(4096)  # type: ignore[union-attr]
+            if not chunk:
+                # EOF — child closed stdout (likely exited)
+                if proc.poll() is not None:
+                    stderr = proc.stderr.read().decode("utf-8") if proc.stderr else ""
+                    raise RuntimeError(
+                        f"{label} exited with code {proc.returncode}.\nstderr: {stderr}"
+                    )
+                continue
+            buf += chunk
+            while b"\n" in buf:
+                line_bytes, buf = buf.split(b"\n", 1)
+                line = line_bytes.decode("utf-8").strip()
+                m = re.match(r"RIVER_PORT=(\d+)", line)
+                if m:
+                    port = int(m.group(1))
+                    break
+            if port is not None:
+                break
+    finally:
+        sel.unregister(proc.stdout)
+        sel.close()
 
     if port is None:
         proc.kill()

python-client/tests/test_e2e.py

@@ -1751,3 +1751,52 @@ def test_enable_transparent_reconnects_option(self):
             options=opts,
         )
         assert transport.reconnect_on_connection_drop is False
+
+    def test_literal_const_escaping(self):
+        """String consts with quotes/backslashes/control chars are escaped."""
+        from river.codegen.schema import SchemaConverter
+
+        converter = SchemaConverter()
+        schema = {"const": 'a"b'}
+        ref = converter._schema_to_typeref(schema, "Test")
+        assert ref.annotation == 'Literal["a\\"b"]'
+
+        schema2 = {"const": "a\\b"}
+        ref2 = converter._schema_to_typeref(schema2, "Test")
+        assert ref2.annotation == 'Literal["a\\\\b"]'
+
+        # Control characters must be escaped
+        schema3 = {"const": "line1\nline2"}
+        ref3 = converter._schema_to_typeref(schema3, "Test")
+        assert ref3.annotation == 'Literal["line1\\nline2"]'
+
+        schema4 = {"const": "a\tb"}
+        ref4 = converter._schema_to_typeref(schema4, "Test")
+        assert ref4.annotation == 'Literal["a\\tb"]'
+
+    def test_is_closed_with_buffered_data(self):
+        """is_closed() is False when closed but queue has data."""
+        from river.streams import Readable
+
+        r: Readable = Readable()
+        r._push_value({"val": 1})
+        r._trigger_close()
+        # Closed but not fully consumed
+        assert r.is_closed() is False
+        assert r._closed is True
+
+    @pytest.mark.asyncio
+    async def test_close_cancels_inflight_connect(self, server_url: str):
+        """close() during handshake doesn't leak the websocket."""
+        transport = WebSocketClientTransport(
+            ws_url=server_url,
+            client_id=None,
+            server_id="SERVER",
+            codec=NaiveJsonCodec(),
+        )
+        transport.connect("SERVER")
+        # Let connection start but don't wait for completion
+        await asyncio.sleep(0)
+        await transport.close()
+        # No leaked sessions
+        assert len(transport.sessions) == 0

python-client/tests/test_session.py

@@ -390,3 +390,112 @@ async def test_connect_on_invoke_false_no_reconnect(self, server_url: str):
         finally:
             # transport already closed above
             pass
+
+
+# =====================================================================
+# Regression: stale connect-task must not block fail-fast
+# =====================================================================
+
+
+class TestStaleConnectTask:
+    @pytest.mark.asyncio
+    async def test_done_connect_task_does_not_block_failfast(self):
+        """A completed (done) connect task in _connect_tasks must not
+        prevent the fail-fast path from firing.
+
+        Regression: previously the check was `to not in _connect_tasks`,
+        so a done task kept the entry alive and calls would hang instead
+        of failing immediately when retries were exhausted.
+        """
+        transport = WebSocketClientTransport(
+            ws_url="ws://127.0.0.1:1",  # unreachable
+            client_id=None,
+            server_id="STALE",
+            codec=NaiveJsonCodec(),
+            options=SessionOptions(
+                connection_timeout_ms=100,
+                handshake_timeout_ms=100,
+                session_disconnect_grace_ms=200,
+            ),
+        )
+        transport.reconnect_on_connection_drop = False
+        try:
+            # Trigger a connect that will fail
+            transport.connect("STALE")
+            await wait_for_session_gone(transport, "STALE")
+
+            # The done task is still in _connect_tasks
+            assert "STALE" in transport._connect_tasks
+            assert transport._connect_tasks["STALE"].done()
+
+            # Exhaust the retry budget so connect() is a no-op
+            transport._retry_budget.budget_consumed = (
+                transport._retry_budget.attempt_budget_capacity
+            )
+
+            # RPC must fail immediately, not hang
+            client = RiverClient(
+                transport, server_id="STALE", connect_on_invoke=True
+            )
+            result = await asyncio.wait_for(
+                client.rpc("test", "add", {"n": 1}), timeout=1.0
+            )
+            assert result["ok"] is False
+            assert result["payload"]["code"] == "UNEXPECTED_DISCONNECT"
+        finally:
+            await transport.close()
+
+
+# =====================================================================
+# Regression: grace period must not reset on each failed reconnect
+# =====================================================================
+
+
+class TestGracePeriodNotResetOnRetry:
+    @pytest.mark.asyncio
+    async def test_grace_period_not_extended_by_retries(self, server_url: str):
+        """Repeated connection failures must not restart the grace timer.
+
+        Regression: _on_connection_failed() unconditionally called
+        start_grace_period(), which cancelled and restarted the timer
+        on every retry, extending session lifetime far beyond
+        session_disconnect_grace_ms.
+        """
+        grace_ms = 400
+        transport = WebSocketClientTransport(
+            ws_url="ws://127.0.0.1:1",  # unreachable
+            client_id=None,
+            server_id="GRACE",
+            codec=NaiveJsonCodec(),
+            options=SessionOptions(
+                connection_timeout_ms=100,
+                handshake_timeout_ms=100,
+                session_disconnect_grace_ms=grace_ms,
+            ),
+        )
+        try:
+            transport.connect("GRACE")
+
+            # Wait for at least one connection failure to set the grace period
+            await wait_for(
+                lambda: (
+                    (s := transport.sessions.get("GRACE")) is not None
+                    and s._grace_period_task is not None
+                ),
+                timeout=2.0,
+            )
+
+            session = transport.sessions["GRACE"]
+            original_expiry = session._grace_expiry_time
+            assert original_expiry is not None
+
+            # After further retries, the expiry time must not have moved forward
+            await asyncio.sleep(0.2)
+            session2 = transport.sessions.get("GRACE")
+            if session2 is not None and session2._grace_expiry_time is not None:
+                assert session2._grace_expiry_time <= original_expiry
+
+            # Session should be gone within grace_ms + generous margin
+            await wait_for_session_gone(transport, "GRACE", timeout=3.0)
+        finally:
+            await transport.close()