diff --git a/CHANGELOG.md b/CHANGELOG.md index 1f086d79c..2626cf0f3 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/opensquilla/cli/gateway_client.py b/src/opensquilla/cli/gateway_client.py index 7392d473e..8b9e9005a 100644 --- a/src/opensquilla/cli/gateway_client.py +++ b/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() diff --git a/tests/test_cli/test_gateway_client_keepalive.py b/tests/test_cli/test_gateway_client_keepalive.py index 617f9c693..07d73ddf8 100644 --- a/tests/test_cli/test_gateway_client_keepalive.py +++ b/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