fix(cli): bound Gateway RPC response waits

CHANGELOG.md

@@ -6,6 +6,11 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [Unreleased]
 
+### Fixed
+
+- CLI chat requests now time out cleanly when a connected Gateway stops
+  answering, and a lost unsubscribe response no longer blocks later turns.
+
 ## [0.5.4] - 2026-08-25
 
 ### Added

src/opensquilla/cli/gateway_client.py

@@ -355,7 +355,8 @@ def _close_from_client(self) -> None:
 class GatewayClient:
     """WebSocket client for connecting to OpenSquilla gateway daemon."""
 
-    def __init__(self) -> None:
+    def __init__(self, *, request_timeout_s: float | None = 30.0) -> None:
+        self.request_timeout_s = request_timeout_s
         self._ws: Any = None
         self._recv_queue: asyncio.Queue[dict] = asyncio.Queue()
         self._pending: dict[str, asyncio.Future[dict]] = {}
@@ -629,10 +630,31 @@ async def _send_ping(self, ws: Any | None = None) -> None:
             raise ConnectionError("WebSocket is not connected")
         await target.send('{"type":"ping"}')
 
-    async def _call(self, method: str, params: dict | None = None) -> Any:
-        """Send a JSON-RPC request and await its response."""
+    def _discard_pending_response(
+        self,
+        req_id: str,
+        fut: asyncio.Future[dict],
+    ) -> None:
+        """Remove one request without disturbing a newer entry for the same id."""
+
+        if self._pending.get(req_id) is fut:
+            self._pending.pop(req_id, None)
+        if not fut.done():
+            fut.cancel()
+
+    async def _start_call(
+        self,
+        method: str,
+        params: dict | None = None,
+    ) -> tuple[str, asyncio.Future[dict]]:
+        """Send one RPC request and return the future for its response."""
+
         if self._connection_error is not None:
             raise self._connection_error
+        if self._closing:
+            raise ConnectionError(
+                "Gateway connection is closing; reconnect before sending another command."
+            )
         if self._ws is None:
             raise ConnectionError(
                 "Gateway connection lost; restart chat or reconnect before sending another command."
@@ -646,13 +668,36 @@ async def _call(self, method: str, params: dict | None = None) -> Any:
                 json.dumps({"type": "req", "id": req_id, "method": method, "params": params})
             )
         except asyncio.CancelledError:
-            self._pending.pop(req_id, None)
+            self._discard_pending_response(req_id, fut)
             raise
         except Exception as exc:
-            self._pending.pop(req_id, None)
+            self._discard_pending_response(req_id, fut)
             err = self._mark_connection_failed(exc)
             raise err from exc
-        res = await fut
+        return req_id, fut
+
+    async def _await_call_response(
+        self,
+        method: str,
+        req_id: str,
+        fut: asyncio.Future[dict],
+    ) -> Any:
+        """Await and decode one response with the configured RPC deadline."""
+
+        try:
+            if self.request_timeout_s is None:
+                res = await fut
+            else:
+                res = await asyncio.wait_for(fut, timeout=self.request_timeout_s)
+        except TimeoutError as exc:
+            raise TimeoutError(
+                f"{method} timed out after {self.request_timeout_s:g}s"
+            ) from exc
+        finally:
+            # The listener normally pops the entry first. Timeout and caller
+            # cancellation arrive without a response, so clean those paths here.
+            self._discard_pending_response(req_id, fut)
+
         if not res.get("ok"):
             err = res.get("error", {})
             raw_details = err.get("data")
@@ -670,6 +715,12 @@ async def _call(self, method: str, params: dict | None = None) -> Any:
         payload = res.get("payload")
         return {} if payload is None else payload
 
+    async def _call(self, method: str, params: dict | None = None) -> Any:
+        """Send a JSON-RPC request and await its bounded response."""
+
+        req_id, fut = await self._start_call(method, params)
+        return await self._await_call_response(method, req_id, fut)
+
     async def call(self, method: str, params: dict | None = None) -> Any:
         """Public thin wrapper for CLI commands that intentionally use RPC names."""
 
@@ -1095,16 +1146,41 @@ async def _remove_event_subscription(
             return
         if any(item.session_key == session_key for item in self._event_subscriptions.values()):
             return
+        pending_unsubscribe: tuple[str, asyncio.Future[dict]] | None = None
         async with self._subscription_lock:
+            # A replacement subscription may have been registered while this
+            # coroutine waited for the lock. In that case the server subscription
+            # must remain active.
+            if any(
+                item.session_key == session_key
+                for item in self._event_subscriptions.values()
+            ):
+                return
             if session_key not in self._server_session_subscriptions:
                 return
             self._server_session_subscriptions.discard(session_key)
             if self._closing or self._connection_error is not None or self._ws is None:
                 return
             try:
-                await self._call("sessions.messages.unsubscribe", {"key": session_key})
-            except (ConnectionError, GatewayRPCError):
+                # Send under the lock so a replacement subscribe is ordered after
+                # this frame, but release the lock before waiting for its response.
+                pending_unsubscribe = await self._start_call(
+                    "sessions.messages.unsubscribe",
+                    {"key": session_key},
+                )
+            except ConnectionError:
                 return
+        if pending_unsubscribe is None:
+            return
+        req_id, fut = pending_unsubscribe
+        try:
+            await self._await_call_response(
+                "sessions.messages.unsubscribe",
+                req_id,
+                fut,
+            )
+        except (ConnectionError, GatewayRPCError, TimeoutError):
+            return
 
     def _preserve_foreign_event(
         self,
@@ -1233,6 +1309,9 @@ async def send_message(
     async def close(self) -> None:
         """Close the WebSocket connection."""
         self._closing = True
+        self._fail_pending_requests(
+            ConnectionError("Gateway connection closed before the RPC response was received")
+        )
         for task in (self._heartbeat_task, self._listener_task):
             if task is None:
                 continue
@@ -1254,6 +1333,13 @@ async def close(self) -> None:
             subscription._close_from_client()
         self._event_subscriptions.clear()
 
+    def _fail_pending_requests(self, error: BaseException) -> None:
+        pending = tuple(self._pending.values())
+        self._pending.clear()
+        for fut in pending:
+            if not fut.done():
+                fut.set_exception(error)
+
     def _mark_connection_failed(self, exc: BaseException) -> ConnectionError:
         if isinstance(exc, ConnectionError) and str(exc).startswith("Gateway connection lost"):
             err = exc
@@ -1267,10 +1353,7 @@ def _mark_connection_failed(self, exc: BaseException) -> ConnectionError:
             self._connection_error = err
         else:
             err = self._connection_error
-        for fut in self._pending.values():
-            if not fut.done():
-                fut.set_exception(err)
-        self._pending.clear()
+        self._fail_pending_requests(err)
         for subscription in tuple(self._event_subscriptions.values()):
             subscription._fail(err)
         current_task = asyncio.current_task()

tests/test_cli/test_gateway_client_keepalive.py

@@ -60,6 +60,11 @@ async def send(self, payload: str) -> None:
         raise RuntimeError("socket already closed")
 
 
+def test_gateway_client_uses_a_bounded_default_rpc_timeout() -> None:
+    assert GatewayClient().request_timeout_s == 30.0
+    assert GatewayClient(request_timeout_s=None).request_timeout_s is None
+
+
 async def _wait_for(predicate, *, timeout: float = 1.0) -> None:
     deadline = asyncio.get_running_loop().time() + timeout
     while not predicate():
@@ -289,6 +294,155 @@ async def test_call_after_send_failure_raises_clear_connection_error() -> None:
     assert client._pending == {}  # noqa: SLF001
 
 
+@pytest.mark.asyncio
+async def test_call_times_out_without_marking_the_connection_failed() -> None:
+    ws = _FakeWebSocket()
+    client = GatewayClient(request_timeout_s=0.01)
+    client._ws = ws  # noqa: SLF001
+
+    with pytest.raises(
+        TimeoutError,
+        match=r"sessions\.list timed out after 0\.01s",
+    ):
+        await client._call("sessions.list", {"limit": 1})  # noqa: SLF001
+
+    assert client._pending == {}  # noqa: SLF001
+    assert client._connection_error is None  # noqa: SLF001
+    assert client._ws is ws  # noqa: SLF001
+
+
+@pytest.mark.asyncio
+async def test_call_cancellation_cleans_up_pending_response() -> None:
+    ws = _FakeWebSocket()
+    client = GatewayClient()
+    client._ws = ws  # noqa: SLF001
+
+    call_task = asyncio.create_task(
+        client._call("sessions.list", {"limit": 1})  # noqa: SLF001
+    )
+    await _wait_for(lambda: bool(ws.sent))
+    request_id = json.loads(ws.sent[0])["id"]
+    response_future = client._pending[request_id]  # noqa: SLF001
+
+    call_task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await call_task
+
+    assert request_id not in client._pending  # noqa: SLF001
+    assert response_future.cancelled()
+
+
+@pytest.mark.asyncio
+async def test_close_wakes_pending_rpc() -> None:
+    ws = _FakeWebSocket()
+    client = GatewayClient()
+    client._ws = ws  # noqa: SLF001
+
+    call_task = asyncio.create_task(
+        client._call("sessions.list", {"limit": 1})  # noqa: SLF001
+    )
+    await _wait_for(lambda: bool(ws.sent))
+
+    await client.close()
+
+    with pytest.raises(ConnectionError, match="closed before the RPC response"):
+        await call_task
+    assert client._pending == {}  # noqa: SLF001
+    assert ws.closed is True
+
+
+@pytest.mark.asyncio
+async def test_late_response_after_timeout_is_ignored_and_connection_remains_usable() -> None:
+    ws = _FakeWebSocket()
+    client = GatewayClient(request_timeout_s=0.01)
+    client._ws = ws  # noqa: SLF001
+    listener_task = asyncio.create_task(client._listen())  # noqa: SLF001
+    client._listener_task = listener_task  # noqa: SLF001
+
+    try:
+        with pytest.raises(TimeoutError, match=r"sessions\.list timed out"):
+            await client._call("sessions.list", {"limit": 1})  # noqa: SLF001
+        timed_out_request = json.loads(ws.sent[0])
+
+        await ws.iter_queue.put(
+            json.dumps(
+                {
+                    "type": "res",
+                    "id": timed_out_request["id"],
+                    "ok": True,
+                    "payload": {"late": True},
+                }
+            )
+        )
+
+        client.request_timeout_s = None
+        next_call = asyncio.create_task(
+            client._call("sessions.list", {"limit": 2})  # noqa: SLF001
+        )
+        await _wait_for(lambda: len(ws.sent) == 2)
+        next_request = json.loads(ws.sent[1])
+        await ws.iter_queue.put(
+            json.dumps(
+                {
+                    "type": "res",
+                    "id": next_request["id"],
+                    "ok": True,
+                    "payload": {"sessions": []},
+                }
+            )
+        )
+
+        assert await asyncio.wait_for(next_call, timeout=0.1) == {"sessions": []}
+        assert client._pending == {}  # noqa: SLF001
+        assert client._connection_error is None  # noqa: SLF001
+        assert listener_task.done() is False
+    finally:
+        client._closing = True  # noqa: SLF001
+        await ws.iter_queue.put(_STOP)
+        await listener_task
+
+
+@pytest.mark.asyncio
+async def test_unsubscribe_response_wait_does_not_block_replacement_subscribe() -> None:
+    ws = _FakeWebSocket()
+    client = GatewayClient(request_timeout_s=None)
+    client._ws = ws  # noqa: SLF001
+    session_key = "agent:main:replacement"
+    client._server_session_subscriptions.add(session_key)  # noqa: SLF001
+    original = client._new_event_subscription(session_key=session_key)  # noqa: SLF001
+
+    close_task = asyncio.create_task(original.close())
+    await _wait_for(lambda: len(ws.sent) == 1)
+    replacement_task = asyncio.create_task(client.subscribe_session_events(session_key))
+    await _wait_for(lambda: len(ws.sent) == 2)
+
+    requests = [json.loads(payload) for payload in ws.sent]
+    assert [request["method"] for request in requests] == [
+        "sessions.messages.unsubscribe",
+        "sessions.messages.subscribe",
+    ]
+    subscribe_id = requests[1]["id"]
+    client._pending[subscribe_id].set_result(  # noqa: SLF001
+        {
+            "type": "res",
+            "id": subscribe_id,
+            "ok": True,
+            "payload": {"replay_complete": True, "current_stream_seq": 0},
+        }
+    )
+    replacement = await asyncio.wait_for(replacement_task, timeout=0.1)
+
+    unsubscribe_id = requests[0]["id"]
+    client._pending[unsubscribe_id].set_result(  # noqa: SLF001
+        {"type": "res", "id": unsubscribe_id, "ok": True, "payload": {}}
+    )
+    await asyncio.wait_for(close_task, timeout=0.1)
+    assert client._pending == {}  # noqa: SLF001
+
+    client._closing = True  # noqa: SLF001
+    await replacement.close()
+
+
 @pytest.mark.asyncio
 async def test_call_preserves_gateway_error_details_for_safe_fallback_decisions() -> None:
     ws = _FakeWebSocket()
@@ -321,6 +475,7 @@ async def test_call_preserves_gateway_error_details_for_safe_fallback_decisions(
         "fallback_safe": False,
         "orphan_message_id": "message-orphan",
     }
+    assert client._pending == {}  # noqa: SLF001
 
 
 @pytest.mark.asyncio