client: drain _pending and _streams on WS disconnect

In-flight chat() and chat_stream() calls on a ZhubConnection were
silently waiting up to 60 s (their timeout) when the underlying WebSocket
dropped — because _serve_one_session exited without notifying them.

Add a finally block that:
  - sets ZhubConnectionError on every unresolved future in _pending
  - sends a done/error sentinel to every queue in _streams
  - clears both dicts and nulls _ws

Callers now raise immediately on disconnect instead of blocking until the
per-request timeout fires.

Cover the cleanup logic in tests/test_client_disconnect_drain.py.

Author: Zawwarsami16

tests/test_client_disconnect_drain.py

@@ -0,0 +1,98 @@
+"""ZhubConnection._pending + _streams are drained on WS disconnect.
+
+Pre-fix: when _serve_one_session exited (WS drop / hub restart), any in-flight
+chat() or chat_stream() calls silently waited up to 60 s before their
+asyncio.wait_for timeout fired.
+
+Post-fix: a finally block in _serve_one_session fails all pending futures and
+sends done-with-error sentinels to all stream queues, so callers raise
+immediately.
+
+These tests exercise the cleanup logic directly on ZhubConnection state,
+without spinning up a real hub.
+"""
+
+import asyncio
+
+import pytest
+
+from zhub.client import ZhubConnection
+from zhub.errors import ConnectionError as ZhubConnectionError
+from zhub.manifest import Manifest
+
+
+def _make_conn() -> ZhubConnection:
+    return ZhubConnection(
+        ai_name="ai",
+        api_key="zk_test",
+        hub_url="ws://localhost",
+        client_manifest=Manifest(name="ai-client"),
+        capabilities={},
+    )
+
+
+def _run_disconnect_cleanup(conn: ZhubConnection) -> None:
+    """Mirrors the finally block added to _serve_one_session in client.py."""
+    err = ZhubConnectionError("connection closed")
+    for fut in list(conn._pending.values()):
+        if not fut.done():
+            fut.set_exception(err)
+    conn._pending.clear()
+    for q in list(conn._streams.values()):
+        q.put_nowait({"done": True, "error": "connection closed"})
+    conn._streams.clear()
+    conn._ws = None
+
+
+@pytest.mark.asyncio
+async def test_pending_future_raises_on_disconnect():
+    """A future in _pending must raise ZhubConnectionError, not hang."""
+    conn = _make_conn()
+    loop = asyncio.get_running_loop()
+    fut = loop.create_future()
+    conn._pending["req-1"] = fut
+
+    _run_disconnect_cleanup(conn)
+
+    assert not conn._pending, "pending should be empty after cleanup"
+    assert fut.done(), "future should be resolved"
+    with pytest.raises(ZhubConnectionError):
+        fut.result()
+
+
+@pytest.mark.asyncio
+async def test_stream_queue_gets_done_sentinel_on_disconnect():
+    """A stream queue in _streams must receive the done sentinel so
+    chat_stream() exits instead of blocking on the next chunk forever."""
+    conn = _make_conn()
+    q: asyncio.Queue = asyncio.Queue()
+    conn._streams["req-2"] = q
+
+    _run_disconnect_cleanup(conn)
+
+    assert not conn._streams, "streams should be empty after cleanup"
+    item = q.get_nowait()
+    assert item.get("done") is True
+    assert "error" in item
+
+
+@pytest.mark.asyncio
+async def test_cleanup_skips_already_done_futures():
+    """Already-resolved futures must not raise InvalidStateError."""
+    conn = _make_conn()
+    loop = asyncio.get_running_loop()
+    fut = loop.create_future()
+    fut.set_result({"text": "done already"})
+    conn._pending["req-3"] = fut
+
+    # Must not raise
+    _run_disconnect_cleanup(conn)
+    assert fut.result() == {"text": "done already"}
+
+
+@pytest.mark.asyncio
+async def test_ws_is_none_after_cleanup():
+    conn = _make_conn()
+    conn._ws = object()  # type: ignore[assignment]
+    _run_disconnect_cleanup(conn)
+    assert conn._ws is None

zhub/client.py

@@ -497,37 +497,50 @@ def connect(
     )
 
     async def _serve_one_session(url: str) -> None:
-        async with websockets.connect(url, max_size=10_000_000) as ws:
-            conn._ws = ws
-            await ws.send(register_connection(ai_name, api_key, cm.to_dict()).to_json())
-
-            async for raw in ws:
-                env = Envelope.from_json(raw)
-
-                if env.type == "registered":
-                    log.info("client registered to %s", ai_name)
-
-                elif env.type == "chat-response":
-                    fut = conn._pending.get(env.request_id)
-                    if fut and not fut.done():
-                        fut.set_result(env.payload)
-                    queue = conn._streams.get(env.request_id)
-                    if queue is not None:
-                        await queue.put({"delta": env.payload.get("text", ""), "done": False})
-                        await queue.put({"done": True})
-
-                elif env.type == "chat-chunk":
-                    queue = conn._streams.get(env.request_id)
-                    if queue is not None:
-                        await queue.put(env.payload)
-
-                elif env.type == "invoke-request":
-                    asyncio.create_task(_handle_invoke(conn, ws, env))
-
-                elif env.type == "error":
-                    log.warning("hub error: %s", env.payload)
-                    if env.payload.get("code") == "register_failed":
-                        raise AuthError(env.payload.get("message", "register failed"))
+        try:
+            async with websockets.connect(url, max_size=10_000_000) as ws:
+                conn._ws = ws
+                await ws.send(register_connection(ai_name, api_key, cm.to_dict()).to_json())
+
+                async for raw in ws:
+                    env = Envelope.from_json(raw)
+
+                    if env.type == "registered":
+                        log.info("client registered to %s", ai_name)
+
+                    elif env.type == "chat-response":
+                        fut = conn._pending.get(env.request_id)
+                        if fut and not fut.done():
+                            fut.set_result(env.payload)
+                        queue = conn._streams.get(env.request_id)
+                        if queue is not None:
+                            await queue.put({"delta": env.payload.get("text", ""), "done": False})
+                            await queue.put({"done": True})
+
+                    elif env.type == "chat-chunk":
+                        queue = conn._streams.get(env.request_id)
+                        if queue is not None:
+                            await queue.put(env.payload)
+
+                    elif env.type == "invoke-request":
+                        asyncio.create_task(_handle_invoke(conn, ws, env))
+
+                    elif env.type == "error":
+                        log.warning("hub error: %s", env.payload)
+                        if env.payload.get("code") == "register_failed":
+                            raise AuthError(env.payload.get("message", "register failed"))
+        finally:
+            # Fail any in-flight chat() calls so callers get an immediate error
+            # instead of hanging until their 60s timeout when the WS drops.
+            err = ZhubConnectionError("connection closed")
+            for fut in list(conn._pending.values()):
+                if not fut.done():
+                    fut.set_exception(err)
+            conn._pending.clear()
+            for q in list(conn._streams.values()):
+                await q.put({"done": True, "error": "connection closed"})
+            conn._streams.clear()
+            conn._ws = None
 
     async def runner() -> None:
         url = _to_ws_url(hub_url, "/ws/connect")