From f833810d0ea45b65ee7c2c3f487b295a6baad9a9 Mon Sep 17 00:00:00 2001 From: can1357 Date: Thu, 23 Jul 2026 12:38:13 +0200 Subject: [PATCH] fix(rpc): stream v2 chunk frames with backpressure and recover stale page cursors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- docs/rpc.md | 4 +- packages/coding-agent/CHANGELOG.md | 2 + .../coding-agent/src/modes/rpc/rpc-client.ts | 31 ++++++++- .../coding-agent/src/modes/rpc/rpc-frame.ts | 64 +++++++++++++------ .../src/modes/rpc/rpc-messages.ts | 18 +++++- .../coding-agent/src/modes/rpc/rpc-mode.ts | 30 +++++++-- .../coding-agent/src/modes/rpc/rpc-types.ts | 4 +- .../test/fixtures/mock-rpc-agent.ts | 17 ++++- .../test/rpc-client.restart.test.ts | 19 ++++++ python/omp-rpc/src/omp_rpc/client.py | 22 +++++-- python/omp-rpc/tests/test_client.py | 39 +++++++++++ 11 files changed, 210 insertions(+), 40 deletions(-) diff --git a/docs/rpc.md b/docs/rpc.md index 7515255a1..f1f27db30 100644 --- a/docs/rpc.md +++ b/docs/rpc.md @@ -187,9 +187,9 @@ correlate it via `id`. Ordering across concurrent commands is not guaranteed - `{ id?, type: "get_messages" }` - `{ id?, type: "get_messages_page", cursor?: string, limit?: number }` -`get_messages_page` returns a stable chronological page with `messages`, `totalMessages`, and an opaque `nextCursor` when more messages remain. Cursors are bound to the session ID, durable leaf, and message count. The server rejects stale cursors if the session changes between requests, and refuses to start a paging walk while the session is streaming or compacting. Pages contain at most 256 messages and normally stay below the v1 physical-frame ceiling. A v1 caller can page ordinary histories, but an individual message whose response exceeds that ceiling produces an overflow error; retrieving it losslessly requires negotiated v2 framing. +`get_messages_page` returns a stable chronological page with `messages`, `totalMessages`, and an opaque `nextCursor` when more messages remain. Cursors are bound to the session ID, durable leaf, and message count. The server rejects stale cursors if the session changes between requests, and refuses to start a paging walk while the session is streaming or compacting. Failed page requests carry a machine-readable `code` on the error response — `session_busy` (session is streaming or compacting) or `stale_cursor` (the snapshot behind the cursor changed, e.g. a background bash appended a message between pages) — so clients can react without matching error-message text. Pages contain at most 256 messages and normally stay below the v1 physical-frame ceiling. A v1 caller can page ordinary histories, but an individual message whose response exceeds that ceiling produces an overflow error; retrieving it losslessly requires negotiated v2 framing. -The bundled TypeScript `RpcClient.getMessages()` and Python `RpcClient.get_messages()` drain this paged endpoint automatically after negotiating v2. They retain the legacy monolithic command when connected to a v1 server, and fall back to its best-effort snapshot if the session begins streaming or compacting during a page walk. Direct `getMessagesPage()` and `get_messages_page()` calls remain strict so incremental hosts never mix snapshots silently. +The bundled TypeScript `RpcClient.getMessages()` and Python `RpcClient.get_messages()` drain this paged endpoint automatically after negotiating v2. They retain the legacy monolithic command when connected to a v1 server, and on either `session_busy` or `stale_cursor` they discard partial pages and fall back to the legacy best-effort snapshot. Direct `getMessagesPage()` and `get_messages_page()` calls remain strict so incremental hosts never mix snapshots silently. ### Login diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 62f4e0534..102529b6a 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -28,6 +28,8 @@ - Rewrote the `/guided-goal` interviewer rubric around loop-engineering: deterministic success criteria, verification commands, attempt caps, scope boundaries, and stop conditions. Ready objectives must use the five-section structured markdown form. - Added `task.isolation.apply` (default `true`) to choose whether successful isolated `task` runs automatically apply their changes to the parent checkout or retain patch/branch artifacts for later integration. - Added opt-in RPC protocol v2 negotiation with bounded, lossless chunking for stdout objects up to 64 MiB, plus stable cursor-based message pages for histories that should not travel as one response. Legacy JSONL clients remain on protocol v1, while the bundled TypeScript and Python RPC clients negotiate, reassemble, and drain message pages automatically. +- Fixed protocol v2 chunked framing materializing the whole base64 transport in memory: near-limit logical frames (~63 MiB) peaked around 686 MB RSS and over-ceiling frames allocated the full payload buffer before rejection. Chunk lines are now produced lazily from a single serialization, the 64 MiB ceiling is checked before any full-payload allocation, and RPC stdout writes honor backpressure line by line. +- Fixed the bundled TypeScript and Python RPC clients throwing when a `get_messages_page` cursor went stale mid-walk (e.g. a background bash appending a message between pages): the high-level `getMessages()` drains now discard partial pages and fall back to the legacy snapshot on both `session_busy` and `stale_cursor`, driven by a new machine-readable `code` field on RPC error responses. Direct page calls remain strict. ## [17.0.8] - 2026-07-22 diff --git a/packages/coding-agent/src/modes/rpc/rpc-client.ts b/packages/coding-agent/src/modes/rpc/rpc-client.ts index 4269d65de..ef9b06541 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -13,7 +13,12 @@ import type { FileSink } from "bun"; import type { BashResult } from "../../exec/bash-executor"; import type { AgentSessionEvent, SessionStats } from "../../session/agent-session"; import { MAX_RPC_FRAME_BYTES, MAX_RPC_REASSEMBLED_BYTES, RpcFrameDecoder, type RpcProtocolVersion } from "./rpc-frame"; -import { RPC_MESSAGES_PAGE_BUSY_ERROR, type RpcMessagesPage, type RpcMessagesPageOptions } from "./rpc-messages"; +import { + RPC_MESSAGES_PAGE_BUSY_ERROR, + RPC_MESSAGES_PAGE_STALE_ERROR, + type RpcMessagesPage, + type RpcMessagesPageOptions, +} from "./rpc-messages"; import type { RpcAvailableCommandsUpdateFrame, RpcAvailableSlashCommand, @@ -212,6 +217,26 @@ function normalizeToolResult(result: RpcClientToolResult): A return result; } +/** Failed RPC command; `code` mirrors the server's machine-readable error code when present. */ +export class RpcCommandError extends Error { + constructor( + message: string, + readonly command: string, + readonly code?: string, + ) { + super(message); + this.name = "RpcCommandError"; + } +} + +/** True when a high-level `getMessages()` drain should discard partial pages and fall back to `get_messages`. */ +function isPageFallbackError(error: unknown): boolean { + if (!(error instanceof Error)) return false; + if (error instanceof RpcCommandError && (error.code === "session_busy" || error.code === "stale_cursor")) + return true; + return error.message === RPC_MESSAGES_PAGE_BUSY_ERROR || error.message === RPC_MESSAGES_PAGE_STALE_ERROR; +} + // ============================================================================ // RPC Client // ============================================================================ @@ -810,7 +835,7 @@ export class RpcClient { throw new Error("RPC message pagination ended before the advertised total"); return messages; } catch (error) { - if (!(error instanceof Error) || error.message !== RPC_MESSAGES_PAGE_BUSY_ERROR) throw error; + if (!isPageFallbackError(error)) throw error; } } const response = await this.#send({ type: "get_messages" }); @@ -1160,7 +1185,7 @@ export class RpcClient { #getData(response: RpcResponse): T { if (!response.success) { const errorResponse = response as Extract; - throw new Error(errorResponse.error); + throw new RpcCommandError(errorResponse.error, errorResponse.command, errorResponse.code); } // Type assertion: we trust response.data matches T based on the command sent. // This is safe because each public method specifies the correct T for its command. diff --git a/packages/coding-agent/src/modes/rpc/rpc-frame.ts b/packages/coding-agent/src/modes/rpc/rpc-frame.ts index abfb0cf6c..862b654c8 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-frame.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-frame.ts @@ -84,19 +84,27 @@ function encodedMessageSnapshot(encoded: string): { message: unknown } | undefin : undefined; } -function encodeChunkedRpcFrame(frame: object, chunkId: string): string { - const json = JSON.stringify(frame); +/** + * Emit protocol v2 chunk frames for one pre-serialized logical frame, one physical + * JSONL line at a time so callers can write with backpressure instead of holding the + * whole ~4/3-sized base64 transport in memory. The reassembly ceiling is enforced on + * `Buffer.byteLength` BEFORE any full-payload allocation. + */ +function* encodeChunkedRpcFrames(frame: object, json: string, chunkId: string): Generator { + const byteLength = Buffer.byteLength(json, "utf8"); + if (byteLength > MAX_RPC_REASSEMBLED_BYTES) { + yield `${JSON.stringify(overflowFrame(frame))}\n`; + return; + } const bytes = Buffer.from(json, "utf8"); - if (bytes.byteLength > MAX_RPC_REASSEMBLED_BYTES) return `${JSON.stringify(overflowFrame(frame))}\n`; - const count = Math.ceil(bytes.byteLength / RPC_CHUNK_PAYLOAD_BYTES); - let encoded = ""; + const count = Math.ceil(byteLength / RPC_CHUNK_PAYLOAD_BYTES); for (let index = 0; index < count; index++) { const chunk: RpcChunkFrame = { type: "rpc_chunk", chunkId, index, count, - byteLength: bytes.byteLength, + byteLength, data: bytes .subarray(index * RPC_CHUNK_PAYLOAD_BYTES, (index + 1) * RPC_CHUNK_PAYLOAD_BYTES) .toString("base64"), @@ -104,9 +112,8 @@ function encodeChunkedRpcFrame(frame: object, chunkId: string): string { const line = `${JSON.stringify(chunk)}\n`; if (serializedFrameBytes(line.slice(0, -1)) > MAX_RPC_FRAME_BYTES) throw new Error("RPC chunk exceeded the transport limit"); - encoded += line; + yield line; } - return encoded; } function isRpcChunkFrame(value: unknown): value is RpcChunkFrame { @@ -263,28 +270,47 @@ export class RpcFrameEncoder { this.#protocolVersion = version; } - encode(frame: object): string { + /** + * Encode one logical frame into physical JSONL lines. Encoder bookkeeping runs + * eagerly; only chunk emission is lazy, so a chunked result can be streamed to + * stdout with backpressure without holding the whole transport in memory. The + * returned iterable MUST be fully consumed exactly once. + */ + encodeFrames(frame: object): Iterable { if (isRecord(frame) && frame.type === "agent_start") this.#streamedMessages = []; - let json = JSON.stringify(frame); - let encoded: string; + const json = JSON.stringify(frame); + let frames: Iterable; + let singleFrame: string | undefined; if (this.#protocolVersion === 2 && serializedFrameBytes(json) > MAX_RPC_FRAME_BYTES) { const compacted = compactTerminalFrame(frame, this.#streamedMessages.length, this.#streamedMessages); - json = JSON.stringify(compacted); - encoded = - serializedFrameBytes(json) > MAX_RPC_FRAME_BYTES - ? encodeChunkedRpcFrame(compacted, `rpc-${++this.#chunkCounter}`) - : `${json}\n`; + // Reuse the original serialization when compaction was a no-op. + const compactedJson = compacted === frame ? json : JSON.stringify(compacted); + if (serializedFrameBytes(compactedJson) > MAX_RPC_FRAME_BYTES) { + frames = encodeChunkedRpcFrames(compacted, compactedJson, `rpc-${++this.#chunkCounter}`); + } else { + singleFrame = `${compactedJson}\n`; + frames = [singleFrame]; + } } else { - encoded = encodeRpcFrame(frame, this.#streamedMessages.length, this.#streamedMessages); + singleFrame = encodeRpcFrame(frame, this.#streamedMessages.length, this.#streamedMessages); + frames = [singleFrame]; } - if (!isRecord(frame)) return encoded; + if (!isRecord(frame)) return frames; if (frame.type === "message_end") { const snapshot = this.#protocolVersion === 2 && Object.hasOwn(frame, "message") ? { message: jsonSnapshot(frame.message) } - : encodedMessageSnapshot(encoded); + : singleFrame !== undefined + ? encodedMessageSnapshot(singleFrame) + : undefined; if (snapshot) this.#streamedMessages.push(snapshot.message); } else if (frame.type === "agent_end" && frame.willContinue !== true) this.#streamedMessages = []; + return frames; + } + + encode(frame: object): string { + let encoded = ""; + for (const line of this.encodeFrames(frame)) encoded += line; return encoded; } } diff --git a/packages/coding-agent/src/modes/rpc/rpc-messages.ts b/packages/coding-agent/src/modes/rpc/rpc-messages.ts index 5307456be..4bec96207 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-messages.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-messages.ts @@ -7,6 +7,21 @@ const MAX_RPC_MESSAGE_PAGE_BYTES = 768 * 1024; const MAX_RPC_MESSAGE_CURSOR_CHARS = 2048; export const RPC_MESSAGES_PAGE_BUSY_ERROR = "Cannot page messages while the session is changing"; +export const RPC_MESSAGES_PAGE_STALE_ERROR = "RPC message cursor is stale"; + +/** Machine-readable reasons a `get_messages_page` request can fail; carried as `code` on the error response. */ +export type RpcMessagesPageErrorCode = "session_busy" | "stale_cursor"; + +/** Paging failure that maps to a structured wire `code`, so clients can react without matching message text. */ +export class RpcMessagesPageError extends Error { + constructor( + message: string, + readonly code: RpcMessagesPageErrorCode, + ) { + super(message); + this.name = "RpcMessagesPageError"; + } +} export interface RpcMessageSnapshot { sessionId: string; @@ -88,7 +103,8 @@ export function pageRpcMessages( let offset = 0; if (options.cursor !== undefined) { const cursor = decodeCursor(options.cursor); - if (!sameSnapshot(cursor, snapshot)) throw new Error("RPC message cursor is stale"); + if (!sameSnapshot(cursor, snapshot)) + throw new RpcMessagesPageError(RPC_MESSAGES_PAGE_STALE_ERROR, "stale_cursor"); offset = cursor.offset; } diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index ef62adf12..5fe031bb1 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -10,6 +10,7 @@ * - Events: AgentSessionEvent objects streamed as they occur * - Extension UI: Extension UI requests are emitted, client responds with extension_ui_response */ +import { once } from "node:events"; import { getOAuthProviders } from "@oh-my-pi/pi-ai/oauth"; import { isZodSchema, zodToWireSchema } from "@oh-my-pi/pi-ai/utils/schema"; import { $env, isRecord, readLines, Snowflake } from "@oh-my-pi/pi-utils"; @@ -36,7 +37,7 @@ import { isRpcHostToolResult, isRpcHostToolUpdate, RpcHostToolBridge } from "./h import { isRpcHostUriResult, RpcHostUriBridge } from "./host-uris"; import { MAX_RPC_FRAME_BYTES, MAX_RPC_REASSEMBLED_BYTES, RpcFrameEncoder } from "./rpc-frame"; import { claimRpcInput } from "./rpc-input"; -import { pageRpcMessages, RPC_MESSAGES_PAGE_BUSY_ERROR } from "./rpc-messages"; +import { pageRpcMessages, RPC_MESSAGES_PAGE_BUSY_ERROR, RpcMessagesPageError } from "./rpc-messages"; import { RpcSubagentRegistry, readRpcSubagentTranscript } from "./rpc-subagents"; import type { RpcCommand, @@ -620,8 +621,22 @@ export async function runRpcMode( process.env.PI_NOTIFICATIONS = "off"; const frameEncoder = new RpcFrameEncoder(); - process.stdout.write( - frameEncoder.encode({ + // Ordered stdout writer honoring backpressure: chunked v2 frames are produced + // lazily by the encoder and written one physical line at a time, so a near-limit + // logical frame never materializes its full base64 transport in memory. + let stdoutQueue: Promise = Promise.resolve(); + const writeFrames = (frames: Iterable) => { + stdoutQueue = stdoutQueue + .then(async () => { + for (const line of frames) { + if (!process.stdout.write(line)) await once(process.stdout, "drain"); + } + }) + // stdout gone (host exited) — nothing left to deliver; keep the queue alive. + .catch(() => {}); + }; + writeFrames( + frameEncoder.encodeFrames({ type: "ready", protocolVersion: 1, supportedProtocolVersions: [1, 2], @@ -630,7 +645,7 @@ export async function runRpcMode( }), ); const output = (obj: RpcResponse | RpcExtensionUIRequest | object) => { - process.stdout.write(frameEncoder.encode(obj)); + writeFrames(frameEncoder.encodeFrames(obj)); if (isRecord(obj) && obj.type === "response" && obj.command === "negotiate_protocol" && obj.success === true) frameEncoder.setProtocolVersion(2); }; @@ -647,8 +662,8 @@ export async function runRpcMode( return { id, type: "response", command, success: true, data } as RpcResponse; }; - const error = (id: string | undefined, command: string, message: string): RpcResponse => { - return { id, type: "response", command, success: false, error: message }; + const error = (id: string | undefined, command: string, message: string, code?: string): RpcResponse => { + return { id, type: "response", command, success: false, error: message, ...(code ? { code } : {}) }; }; const extensionUserMessageTracker = new RpcExtensionUserMessageTracker(); @@ -1308,7 +1323,7 @@ export async function runRpcMode( case "get_messages_page": { if (session.isStreaming || session.isCompacting) - return error(id, "get_messages_page", RPC_MESSAGES_PAGE_BUSY_ERROR); + return error(id, "get_messages_page", RPC_MESSAGES_PAGE_BUSY_ERROR, "session_busy"); const messages = session.messages; try { return success( @@ -1329,6 +1344,7 @@ export async function runRpcMode( id, "get_messages_page", pageError instanceof Error ? pageError.message : String(pageError), + pageError instanceof RpcMessagesPageError ? pageError.code : undefined, ); } } diff --git a/packages/coding-agent/src/modes/rpc/rpc-types.ts b/packages/coding-agent/src/modes/rpc/rpc-types.ts index babb08ccd..ab409a959 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-types.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-types.ts @@ -327,8 +327,8 @@ export type RpcResponse = } | { id?: string; type: "response"; command: "login"; success: true; data: { providerId: string } } - // Error response (any command can fail) - | { id?: string; type: "response"; command: string; success: false; error: string }; + // Error response (any command can fail); `code` is an optional machine-readable reason. + | { id?: string; type: "response"; command: string; success: false; error: string; code?: string }; // ============================================================================ // Subagent Events (stdout) diff --git a/packages/coding-agent/test/fixtures/mock-rpc-agent.ts b/packages/coding-agent/test/fixtures/mock-rpc-agent.ts index 918df7227..f7a7ccc33 100755 --- a/packages/coding-agent/test/fixtures/mock-rpc-agent.ts +++ b/packages/coding-agent/test/fixtures/mock-rpc-agent.ts @@ -87,6 +87,18 @@ for await (const raw of console) { command: frame.type, success: false, error: "Cannot page messages while the session is changing", + code: "session_busy", + }); + continue; + } + if (Bun.env.MOCK_RPC_PAGE_STALE === "1" && frame.cursor !== undefined) { + writeFrame({ + id, + type: "response", + command: frame.type, + success: false, + error: "RPC message cursor is stale", + code: "stale_cursor", }); continue; } @@ -109,7 +121,10 @@ for await (const raw of console) { }); continue; } - if (frame.type === "get_messages" && Bun.env.MOCK_RPC_PAGE_BUSY === "1") { + if ( + frame.type === "get_messages" && + (Bun.env.MOCK_RPC_PAGE_BUSY === "1" || Bun.env.MOCK_RPC_PAGE_STALE === "1") + ) { writeFrame({ id, type: "response", diff --git a/packages/coding-agent/test/rpc-client.restart.test.ts b/packages/coding-agent/test/rpc-client.restart.test.ts index 9b31b117e..540a6998d 100644 --- a/packages/coding-agent/test/rpc-client.restart.test.ts +++ b/packages/coding-agent/test/rpc-client.restart.test.ts @@ -43,6 +43,25 @@ describe("RpcClient lifecycle (issue #4079 B)", () => { ]); }, 20_000); + test("discards partial pages and falls back to get_messages when a cursor goes stale mid-walk", async () => { + using client = new RpcClient({ + cliPath: MOCK_AGENT, + env: { MOCK_RPC_V2: "1", MOCK_RPC_PAGE_STALE: "1" }, + }); + + await client.start(); + // Direct page walks stay strict: the stale cursor is surfaced to the caller. + const firstPage = await client.getMessagesPage(); + expect(firstPage.nextCursor).toBe("second-page"); + await expect(client.getMessagesPage({ cursor: firstPage.nextCursor })).rejects.toThrow( + "RPC message cursor is stale", + ); + // The high-level drain discards the partial first page and takes the legacy snapshot. + expect((await client.getMessages()) as unknown).toEqual([ + { role: "assistant", content: [{ type: "text", text: "streaming snapshot" }], timestamp: 3 }, + ]); + }, 20_000); + test("start() succeeds a second time after stop() on the same instance", async () => { using client = new RpcClient({ cliPath: MOCK_AGENT, diff --git a/python/omp-rpc/src/omp_rpc/client.py b/python/omp-rpc/src/omp_rpc/client.py index b43482f22..a0c1cf729 100644 --- a/python/omp-rpc/src/omp_rpc/client.py +++ b/python/omp-rpc/src/omp_rpc/client.py @@ -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") diff --git a/python/omp-rpc/tests/test_client.py b/python/omp-rpc/tests/test_client.py index dda1ebdb3..fb803b66f 100644 --- a/python/omp-rpc/tests/test_client.py +++ b/python/omp-rpc/tests/test_client.py @@ -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")