fix(rpc): stream v2 chunk frames with backpressure and recover stale page cursors
Two fixes on top of the paged transport: - Near-limit v2 framing no longer materializes the full base64 transport: chunk lines are generated lazily from a single serialization, the 64 MiB reassembly ceiling is enforced via Buffer.byteLength before any full-payload allocation, and RPC stdout writes drain with backpressure one physical line at a time. Peak RSS for a 63 MiB response drops ~686 MB -> ~521 MB; a rejected 80 MiB response drops ~507 MB -> ~259 MB (parity with the v1 path). - get_messages_page errors now carry a machine-readable code (session_busy | stale_cursor). Both bundled clients' high-level getMessages() drains discard partial pages and fall back to the legacy snapshot on either code — previously a cursor invalidated by a background mutation (e.g. an appended bash message) threw instead of falling back. Direct page calls remain strict.
This commit is contained in:
@@ -118,6 +118,8 @@ _MAX_RPC_FRAME_BYTES = 1024 * 1024
|
||||
_MAX_RPC_REASSEMBLED_BYTES = 64 * 1024 * 1024
|
||||
_RPC_CHUNK_PAYLOAD_BYTES = 256 * 1024
|
||||
_RPC_MESSAGES_PAGE_BUSY_ERROR = "Cannot page messages while the session is changing"
|
||||
_RPC_MESSAGES_PAGE_STALE_ERROR = "RPC message cursor is stale"
|
||||
_RPC_MESSAGES_PAGE_FALLBACK_CODES = frozenset({"session_busy", "stale_cursor"})
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -313,12 +315,16 @@ class RpcConcurrencyError(RpcError):
|
||||
|
||||
|
||||
class RpcCommandError(RpcError):
|
||||
"""Raised when the RPC server returns `success: false`."""
|
||||
"""Raised when the RPC server returns `success: false`.
|
||||
|
||||
def __init__(self, command: str, error: str):
|
||||
`code` carries the server's machine-readable error code when present.
|
||||
"""
|
||||
|
||||
def __init__(self, command: str, error: str, code: str | None = None):
|
||||
super().__init__(f"{command}: {error}")
|
||||
self.command = command
|
||||
self.error = error
|
||||
self.code = code
|
||||
|
||||
|
||||
class RpcProtocolError(RpcError):
|
||||
@@ -1036,9 +1042,13 @@ class RpcClient:
|
||||
)
|
||||
return tuple(messages)
|
||||
except RpcCommandError as error:
|
||||
if (
|
||||
error.command != "get_messages_page"
|
||||
or error.error != _RPC_MESSAGES_PAGE_BUSY_ERROR
|
||||
if error.command != "get_messages_page" or not (
|
||||
error.code in _RPC_MESSAGES_PAGE_FALLBACK_CODES
|
||||
or error.error
|
||||
in (
|
||||
_RPC_MESSAGES_PAGE_BUSY_ERROR,
|
||||
_RPC_MESSAGES_PAGE_STALE_ERROR,
|
||||
)
|
||||
):
|
||||
raise
|
||||
payload = self._request("get_messages")
|
||||
@@ -1385,9 +1395,11 @@ class RpcClient:
|
||||
raise response
|
||||
|
||||
if not bool(response.get("success", False)):
|
||||
raw_code = response.get("code")
|
||||
raise RpcCommandError(
|
||||
command=str(response.get("command", command_type)),
|
||||
error=str(response.get("error", "")),
|
||||
code=raw_code if isinstance(raw_code, str) else None,
|
||||
)
|
||||
|
||||
data = response.get("data")
|
||||
|
||||
@@ -519,6 +519,34 @@ V2_MESSAGES_SERVER = textwrap.dedent(
|
||||
"command": command_type,
|
||||
"success": False,
|
||||
"error": "Cannot page messages while the session is changing",
|
||||
"code": "session_busy",
|
||||
}
|
||||
)
|
||||
continue
|
||||
if os.environ.get("V2_MESSAGES_STALE") == "1":
|
||||
if command.get("cursor") is not None:
|
||||
emit(
|
||||
{
|
||||
"id": request_id,
|
||||
"type": "response",
|
||||
"command": command_type,
|
||||
"success": False,
|
||||
"error": "RPC message cursor is stale",
|
||||
"code": "stale_cursor",
|
||||
}
|
||||
)
|
||||
continue
|
||||
emit(
|
||||
{
|
||||
"id": request_id,
|
||||
"type": "response",
|
||||
"command": command_type,
|
||||
"success": True,
|
||||
"data": {
|
||||
"messages": [message],
|
||||
"totalMessages": 2,
|
||||
"nextCursor": "page-two",
|
||||
},
|
||||
}
|
||||
)
|
||||
continue
|
||||
@@ -1036,6 +1064,17 @@ class RpcClientTests(unittest.TestCase):
|
||||
self.assertEqual(len(messages), 1)
|
||||
self.assertEqual(messages[0]["content"][0]["text"], "streaming snapshot")
|
||||
|
||||
def test_protocol_v2_get_messages_discards_stale_page_walk(self) -> None:
|
||||
with self.make_client(
|
||||
server=V2_MESSAGES_SERVER, env={"V2_MESSAGES_STALE": "1"}
|
||||
) as client:
|
||||
with self.assertRaisesRegex(RpcCommandError, "RPC message cursor is stale"):
|
||||
client.get_messages_page(cursor="page-two")
|
||||
messages = client.get_messages()
|
||||
|
||||
self.assertEqual(len(messages), 1)
|
||||
self.assertEqual(messages[0]["content"][0]["text"], "streaming snapshot")
|
||||
|
||||
def test_collect_events_returns_turn_events(self) -> None:
|
||||
with self.make_client() as client:
|
||||
client.prompt("slow")
|
||||
|
||||
Reference in New Issue
Block a user