fix(omp): negotiate RPC protocol v2 so large model catalogs don't overflow (#3184)

* fix(omp): negotiate RPC protocol v2 so large responses don't overflow

OMP caps protocol-v1 single-line RPC frames at 1 MiB. get_available_models
returns the full effective model catalog (hundreds of models), which exceeds
that cap, so OMP replies with an overflow error ("RPC response exceeded the
transport limit") and the OMP provider fails to start sessions.

Negotiate RPC protocol v2 after OMP's startup `ready` frame when it is
supported: v2 lifts the single-frame ceiling to 64 MiB by chunking oversized
frames. Add rpc_chunk reassembly to the shared JSONL transport so chunked
responses are reconstructed before routing, and gate the negotiation on the
ready frame so it is in place before the first request.

Protocol v2 is negotiated only when the peer advertises support in its ready
frame; peers that stay on v1 (e.g. Pi's --mode rpc) are unaffected.

* fix(omp): make protocol startup atomic

---------

Co-authored-by: pi3123 <pi3123@users.noreply.github.com>
Co-authored-by: Mohamed Boudra <boudra.moha@gmail.com>
This commit is contained in:
pi3123
2026-08-11 19:19:24 +02:00
committed by GitHub
co-authored by pi3123 Mohamed Boudra
parent cdb1163145
commit 163e7d1cc4
4 changed files with 561 additions and 3 deletions
@@ -219,4 +219,123 @@ describe("JsonlRpcProcess", () => {
await rejection;
});
test("reassembles protocol v2 chunked responses into the logical frame", async () => {
const child = createInMemoryChildProcess();
const transport = startProcess({ child });
try {
let buffer = "";
const sentRequest = new Promise<Record<string, unknown>>((resolve) => {
child.stdin.on("data", (chunk) => {
buffer += chunk.toString();
const newlineIndex = buffer.indexOf("\n");
if (newlineIndex === -1) return;
resolve(JSON.parse(buffer.slice(0, newlineIndex)) as Record<string, unknown>);
});
});
const request = transport.request({ type: "chunked" });
const command = await sentRequest;
// Emit a logical response split into protocol v2 chunk frames.
const value = "x".repeat(1024 * 1024);
const logical = JSON.stringify({
id: command.id,
type: "response",
command: command.type,
success: true,
data: { value },
});
const bytes = Buffer.from(logical, "utf8");
expect(bytes.byteLength).toBeGreaterThan(1024 * 1024);
const chunkPayload = 256 * 1024;
const count = Math.ceil(bytes.length / chunkPayload);
expect(count).toBeGreaterThan(1);
for (let index = 0; index < count; index++) {
child.stdout.write(
`${JSON.stringify({
type: "rpc_chunk",
chunkId: "seq1",
index,
count,
byteLength: bytes.length,
data: bytes
.subarray(index * chunkPayload, (index + 1) * chunkPayload)
.toString("base64"),
})}\n`,
);
}
await expect(request).resolves.toEqual({ value });
} finally {
await transport.close();
}
});
test("rejects protocol v2 chunks larger than the transport payload limit", async () => {
const child = createInMemoryChildProcess();
const transport = startProcess({ child });
try {
let buffer = "";
const sentRequest = new Promise<Record<string, unknown>>((resolve) => {
child.stdin.on("data", (chunk) => {
buffer += chunk.toString();
const newlineIndex = buffer.indexOf("\n");
if (newlineIndex === -1) return;
resolve(JSON.parse(buffer.slice(0, newlineIndex)) as Record<string, unknown>);
});
});
const request = transport.request({ type: "oversized-chunks" });
const command = await sentRequest;
const logicalResponse = (value: string): Buffer =>
Buffer.from(
JSON.stringify({
id: command.id,
type: "response",
command: command.type,
success: true,
data: { value },
}),
"utf8",
);
const oversized = logicalResponse(`invalid-${"x".repeat(1024 * 1024)}`);
const split = Math.ceil(oversized.byteLength / 2);
expect(split).toBeGreaterThan(256 * 1024);
for (let index = 0; index < 2; index += 1) {
child.stdout.write(
`${JSON.stringify({
type: "rpc_chunk",
chunkId: "oversized",
index,
count: 2,
byteLength: oversized.byteLength,
data: oversized.subarray(index * split, (index + 1) * split).toString("base64"),
})}\n`,
);
}
const valid = logicalResponse(`valid-${"x".repeat(1024 * 1024)}`);
const count = Math.ceil(valid.byteLength / (256 * 1024));
for (let index = 0; index < count; index += 1) {
child.stdout.write(
`${JSON.stringify({
type: "rpc_chunk",
chunkId: "valid",
index,
count,
byteLength: valid.byteLength,
data: valid.subarray(index * 256 * 1024, (index + 1) * 256 * 1024).toString("base64"),
})}\n`,
);
}
const response = (await request) as { value: string };
expect(response.value.startsWith("valid-")).toBe(true);
} finally {
await transport.close();
}
});
});
@@ -15,6 +15,11 @@ export const JSONL_RPC_NO_TIMEOUT = null;
const STDERR_BUFFER_LIMIT = 8192;
const GRACEFUL_SHUTDOWN_TIMEOUT_MS = 2_000;
const FORCE_SHUTDOWN_TIMEOUT_MS = 1_000;
const RPC_V2_MAX_FRAME_BYTES = 1024 * 1024;
const RPC_V2_MAX_REASSEMBLED_BYTES = 64 * 1024 * 1024;
const RPC_V2_CHUNK_PAYLOAD_BYTES = 256 * 1024;
const RPC_V2_MAX_ENCODED_CHUNK_LENGTH = Math.ceil(RPC_V2_CHUNK_PAYLOAD_BYTES / 3) * 4;
const RPC_V2_MAX_CHUNK_COUNT = Math.ceil(RPC_V2_MAX_REASSEMBLED_BYTES / RPC_V2_CHUNK_PAYLOAD_BYTES);
export interface JsonlRpcLaunch {
command: string;
@@ -38,6 +43,30 @@ interface PendingRequest {
timer: NodeJS.Timeout | null;
}
/**
* Reassembly state for a protocol v2 chunked frame sequence. Some agents that
* speak this JSONL RPC protocol (OMP) cap single-line frames at 1 MiB and emit
* larger frames as ordered `rpc_chunk` frames; the transport reassembles them
* into the original logical frame before routing it.
*/
interface PendingChunkSequence {
chunkId: string;
count: number;
byteLength: number;
nextIndex: number;
chunks: Buffer[];
receivedBytes: number;
}
/** Validated metadata fields of a protocol v2 `rpc_chunk` frame. */
interface RpcChunkMetadata {
chunkId: string;
index: number;
count: number;
byteLength: number;
data: string;
}
export interface JsonlRpcExit {
code: number | null;
signal: NodeJS.Signals | null;
@@ -51,6 +80,16 @@ export interface JsonlRpcProcessOptions {
spawn?: (launch: JsonlRpcLaunch) => ChildProcessWithoutNullStreams;
}
export function supportsJsonlRpcProtocolV2(message: Record<string, unknown>): boolean {
return (
message.type === "ready" &&
Array.isArray(message.supportedProtocolVersions) &&
message.supportedProtocolVersions.includes(2) &&
message.maxFrameBytes === RPC_V2_MAX_FRAME_BYTES &&
message.maxReassembledFrameBytes === RPC_V2_MAX_REASSEMBLED_BYTES
);
}
function assertChildWithPipes(
child: ChildProcess,
): asserts child is ChildProcessWithoutNullStreams {
@@ -79,6 +118,7 @@ export class JsonlRpcProcess {
private nextRequestId = 1;
private disposed = false;
private stdoutBuffer = "";
private pendingChunks: PendingChunkSequence | null = null;
constructor(private readonly options: JsonlRpcProcessOptions) {
this.diagnosticName = options.diagnosticName ?? "JSONL RPC";
@@ -218,6 +258,20 @@ export class JsonlRpcProcess {
return;
}
const message = parsed as Record<string, unknown>;
if (message.type === "rpc_chunk") {
const reassembled = this.consumeRpcChunk(message);
if (reassembled !== undefined) {
// The reassembled frame is a complete logical frame; route it through
// the same path as any other parsed line.
this.handleLine(JSON.stringify(reassembled));
}
return;
}
this.dispatchFrame(message);
}
/** Route a complete logical frame to its consumer (response or subscriber). */
private dispatchFrame(message: Record<string, unknown>): void {
if (message.type === "response") {
this.handleResponse(message as unknown as JsonlRpcResponse);
return;
@@ -227,6 +281,173 @@ export class JsonlRpcProcess {
}
}
/**
* Reassemble a protocol v2 `rpc_chunk` frame. Returns the reconstructed
* logical frame once the final chunk of a sequence arrives, `undefined` while
* a sequence is still in progress, and drops (returning `undefined`) any
* malformed or out-of-order chunk after logging it.
*/
private consumeRpcChunk(chunk: Record<string, unknown>): Record<string, unknown> | undefined {
const metadata = this.parseChunkMetadata(chunk);
if (metadata === undefined) {
return undefined;
}
const bytes = this.decodeChunkPayload(metadata.data);
if (bytes === undefined) {
return undefined;
}
return this.appendChunk(metadata, bytes);
}
/** Validate a `rpc_chunk` frame's metadata, or reset the sequence and return undefined. */
private parseChunkMetadata(chunk: Record<string, unknown>): RpcChunkMetadata | undefined {
const chunkId = chunk.chunkId;
const data = chunk.data;
const index = chunk.index;
const count = chunk.count;
const byteLength = chunk.byteLength;
if (
typeof chunkId !== "string" ||
chunkId.length === 0 ||
chunkId.length > 128 ||
typeof data !== "string" ||
data.length === 0 ||
typeof index !== "number" ||
!Number.isSafeInteger(index) ||
index < 0 ||
typeof count !== "number" ||
!Number.isSafeInteger(count) ||
count < 2 ||
count > RPC_V2_MAX_CHUNK_COUNT ||
index >= count ||
typeof byteLength !== "number" ||
!Number.isSafeInteger(byteLength) ||
byteLength < RPC_V2_MAX_FRAME_BYTES ||
byteLength > RPC_V2_MAX_REASSEMBLED_BYTES
) {
this.options.logger.warn(
{ chunk },
`Ignoring malformed ${this.diagnosticName} rpc_chunk frame`,
);
this.pendingChunks = null;
return undefined;
}
return { chunkId, index, count, byteLength, data };
}
/** Decode a chunk payload, or reset the sequence and return undefined on bad base64. */
private decodeChunkPayload(data: string): Buffer | undefined {
try {
if (data.length > RPC_V2_MAX_ENCODED_CHUNK_LENGTH) {
throw new Error("encoded payload exceeds the transport limit");
}
const bytes = Buffer.from(data, "base64");
if (bytes.byteLength > RPC_V2_CHUNK_PAYLOAD_BYTES || bytes.toString("base64") !== data) {
throw new Error("invalid base64 payload");
}
return bytes;
} catch {
this.options.logger.warn({}, `Ignoring ${this.diagnosticName} rpc_chunk with invalid base64`);
this.pendingChunks = null;
return undefined;
}
}
/** Append a chunk to the in-progress sequence, finalizing it once complete. */
private appendChunk(
metadata: RpcChunkMetadata,
bytes: Buffer,
): Record<string, unknown> | undefined {
let pending = this.pendingChunks;
if (pending === null) {
if (metadata.index !== 0) {
this.options.logger.warn(
{},
`Ignoring ${this.diagnosticName} rpc_chunk sequence that did not start at index 0`,
);
return undefined;
}
pending = {
chunkId: metadata.chunkId,
count: metadata.count,
byteLength: metadata.byteLength,
nextIndex: 0,
chunks: [],
receivedBytes: 0,
};
this.pendingChunks = pending;
}
if (
pending.chunkId !== metadata.chunkId ||
pending.count !== metadata.count ||
pending.byteLength !== metadata.byteLength ||
pending.nextIndex !== metadata.index
) {
this.options.logger.warn(
{},
`Ignoring out-of-order ${this.diagnosticName} rpc_chunk; dropping sequence`,
);
this.pendingChunks = null;
return undefined;
}
pending.chunks.push(bytes);
pending.receivedBytes += bytes.byteLength;
pending.nextIndex += 1;
if (pending.receivedBytes > pending.byteLength) {
this.options.logger.warn(
{},
`Ignoring ${this.diagnosticName} rpc_chunk sequence exceeding its declared length`,
);
this.pendingChunks = null;
return undefined;
}
if (pending.nextIndex < pending.count) {
return undefined;
}
if (pending.receivedBytes !== pending.byteLength) {
this.options.logger.warn(
{},
`Ignoring ${this.diagnosticName} rpc_chunk sequence with a length mismatch`,
);
this.pendingChunks = null;
return undefined;
}
this.pendingChunks = null;
return this.finalizeChunks(pending.chunks);
}
/** Reassemble and parse a complete chunk sequence into its logical frame. */
private finalizeChunks(chunks: Buffer[]): Record<string, unknown> | undefined {
let json: string;
try {
json = new TextDecoder("utf-8", { fatal: true }).decode(Buffer.concat(chunks));
} catch {
this.options.logger.warn(
{},
`Ignoring ${this.diagnosticName} rpc_chunk sequence with invalid UTF-8`,
);
return undefined;
}
let frame: unknown;
try {
frame = JSON.parse(json);
} catch {
this.options.logger.warn(
{},
`Ignoring ${this.diagnosticName} rpc_chunk sequence with invalid JSON`,
);
return undefined;
}
if (!frame || typeof frame !== "object" || Array.isArray(frame)) {
return undefined;
}
return frame as Record<string, unknown>;
}
private handleResponse(response: JsonlRpcResponse): void {
if (!response.id) {
return;
@@ -14,7 +14,11 @@ type OmpChild = ChildProcessWithoutNullStreams & {
killedSignals: Array<NodeJS.Signals | number | undefined>;
};
function createOmpChild(): OmpChild {
function createOmpChild(options?: {
supportedProtocolVersions?: number[];
emitReady?: boolean;
maxFrameBytes?: number;
}): OmpChild {
const child = Object.assign(new EventEmitter(), {
stdin: new PassThrough(),
stdout: new PassThrough(),
@@ -28,6 +32,20 @@ function createOmpChild(): OmpChild {
queueMicrotask(() => child.emit("exit", null, signal ?? null));
return true;
}) as ChildProcessWithoutNullStreams["kill"];
// Real OMP writes a `ready` frame immediately after launch; the runtime waits
// for it before negotiating the RPC protocol. Advertise v1-only by default so
// the session stays on protocol v1 and command streams are unaffected.
if (options?.emitReady !== false) {
child.stdout.write(
`${JSON.stringify({
type: "ready",
protocolVersion: 1,
supportedProtocolVersions: options?.supportedProtocolVersions ?? [1],
maxFrameBytes: options?.maxFrameBytes ?? 1024 * 1024,
maxReassembledFrameBytes: 64 * 1024 * 1024,
})}\n`,
);
}
return child;
}
@@ -243,4 +261,136 @@ describe("OMP CLI runtime", () => {
await expect(session.prompt("hello")).resolves.toEqual({ requestId: "req_1" });
});
test("negotiates RPC protocol v2 when OMP advertises it", async () => {
const child = createOmpChild({ supportedProtocolVersions: [1, 2] });
const commands: Record<string, unknown>[] = [];
replyToCommands(child, (command) => {
commands.push(command);
if (command.type === "negotiate_protocol") {
return { protocolVersion: command.protocolVersion };
}
if (command.type === "get_available_models") {
return {
models: [{ provider: "opencode-go", id: "deepseek-v4-flash", name: "DeepSeek V4 Flash" }],
};
}
return undefined;
});
const session = await createRuntime(child).startSession({ cwd: "/workspace/project" });
await expect(session.getAvailableModels()).resolves.toEqual([
expect.objectContaining({ provider: "opencode-go", id: "deepseek-v4-flash" }),
]);
expect(commands.map(withoutRequestId)).toContainEqual({
type: "negotiate_protocol",
protocolVersion: 2,
});
});
test("stays on protocol v1 when OMP advertises incompatible framing limits", async () => {
const child = createOmpChild({
supportedProtocolVersions: [1, 2],
maxFrameBytes: 512 * 1024,
});
const commands: Record<string, unknown>[] = [];
replyToCommands(child, (command) => {
commands.push(command);
return { models: [] };
});
const session = await createRuntime(child).startSession({ cwd: "/workspace/project" });
await session.getAvailableModels();
expect(commands.map(withoutRequestId)).toEqual([{ type: "get_available_models" }]);
await session.close();
});
test("rejects startup when OMP exits before advertising readiness", async () => {
const child = createOmpChild({ emitReady: false });
const startup = createRuntime(child).startSession({ cwd: "/workspace/project" });
child.stderr.write("startup exploded");
child.emit("exit", 7, null);
await expect(startup).rejects.toThrow("startup exploded");
});
test("rejects startup when OMP exits during protocol negotiation", async () => {
const child = createOmpChild({ supportedProtocolVersions: [1, 2] });
child.stdin.on("data", () => {
child.stderr.write("negotiation exploded");
child.emit("exit", 8, null);
});
await expect(createRuntime(child).startSession({ cwd: "/workspace/project" })).rejects.toThrow(
"negotiation exploded",
);
});
test("reassembles chunked protocol v2 responses for large payloads", async () => {
const child = createOmpChild({ supportedProtocolVersions: [1, 2] });
const models = Array.from({ length: 2_000 }, (_, index) => ({
provider: "p",
id: `model-${index}`,
name: `model-${index}-${"x".repeat(512)}`,
}));
let buffer = "";
child.stdin.on("data", (chunk) => {
buffer += chunk.toString();
for (;;) {
const newlineIndex = buffer.indexOf("\n");
if (newlineIndex === -1) break;
const line = buffer.slice(0, newlineIndex);
buffer = buffer.slice(newlineIndex + 1);
const command = JSON.parse(line) as Record<string, unknown>;
if (command.type === "negotiate_protocol") {
child.stdout.write(
`${JSON.stringify({
id: command.id,
type: "response",
command: command.type,
success: true,
data: { protocolVersion: 2 },
})}\n`,
);
continue;
}
if (command.type === "get_available_models") {
// Emit a logical response larger than the 1 MiB protocol-v1 cap as a
// chunked (protocol v2) frame sequence, like real OMP does.
const logical = JSON.stringify({
id: command.id,
type: "response",
command: command.type,
success: true,
data: { models },
});
const bytes = Buffer.from(logical, "utf8");
expect(bytes.byteLength).toBeGreaterThan(1024 * 1024);
const chunkPayload = 256 * 1024;
const count = Math.ceil(bytes.length / chunkPayload);
for (let index = 0; index < count; index++) {
child.stdout.write(
`${JSON.stringify({
type: "rpc_chunk",
chunkId: "c1",
index,
count,
byteLength: bytes.length,
data: bytes
.subarray(index * chunkPayload, (index + 1) * chunkPayload)
.toString("base64"),
})}\n`,
);
}
}
}
});
const session = await createRuntime(child).startSession({ cwd: "/workspace/project" });
const result = await session.getAvailableModels();
expect(result).toHaveLength(2_000);
expect(result[0]).toEqual(expect.objectContaining({ id: "model-0" }));
});
});
@@ -2,7 +2,12 @@ import type { ChildProcessWithoutNullStreams } from "node:child_process";
import type { Logger } from "pino";
import type { ProviderRuntimeSettings } from "../../provider-launch-config.js";
import { JsonlRpcProcess, type JsonlRpcLaunch } from "../jsonl-rpc-process.js";
import {
JSONL_RPC_DEFAULT_TIMEOUT_MS,
JsonlRpcProcess,
supportsJsonlRpcProtocolV2,
type JsonlRpcLaunch,
} from "../jsonl-rpc-process.js";
import {
buildOmpLaunch,
type OmpRuntime,
@@ -40,6 +45,8 @@ import {
const DEFAULT_OMP_COMMAND: [string, ...string[]] = [process.env.OMP_COMMAND ?? "omp"];
const DEFAULT_COMMANDS_RPC_NAME = "get_available_commands";
/** How long to wait for OMP's startup `ready` frame before failing startup. */
const OMP_READY_TIMEOUT_MS = 10_000;
export interface OmpCliRuntimeOptions {
logger: Logger;
@@ -81,10 +88,71 @@ export class OmpCliRuntime implements OmpRuntime {
...(spawn ? { spawn: () => spawn(launch) } : {}),
};
const process = new JsonlRpcProcess(processOptions);
return new OmpCliRuntimeSession(process, this.commandsRpcName);
try {
await negotiateOmpProtocolV2(process, this.options.logger);
return new OmpCliRuntimeSession(process, this.commandsRpcName);
} catch (error) {
const startupError = error instanceof Error ? error : new Error(String(error));
await process.close(startupError);
throw startupError;
}
}
}
/**
* Wait for OMP to advertise readiness and negotiate RPC protocol v2 when it is
* supported. OMP caps protocol-v1 single-line frames at 1 MiB; `get_available_models`
* (and other large payloads) can exceed that and are returned as an overflow error.
* Protocol v2 lifts the ceiling to 64 MiB by chunking oversized frames, which the
* JSONL transport reassembles. Supported OMP versions send a `ready` frame immediately
* after launch; startup fails if the process exits or never becomes ready.
*/
async function negotiateOmpProtocolV2(process: JsonlRpcProcess, logger: Logger): Promise<void> {
const ready = await waitForOmpReadyFrame(process);
if (!supportsJsonlRpcProtocolV2(ready)) {
return;
}
const response = (await process.request(
{ type: "negotiate_protocol", protocolVersion: 2 },
JSONL_RPC_DEFAULT_TIMEOUT_MS,
)) as { protocolVersion?: unknown } | undefined;
if (response?.protocolVersion !== 2) {
throw new Error("OMP did not accept RPC protocol v2");
}
logger.debug({}, "Negotiated OMP RPC protocol v2 (chunked frame transport)");
}
function waitForOmpReadyFrame(process: JsonlRpcProcess): Promise<Record<string, unknown>> {
return new Promise((resolve, reject) => {
let settled = false;
let unsubscribeMessage = (): void => {};
let unsubscribeExit = (): void => {};
let timer: NodeJS.Timeout | undefined;
const finish = (result: Record<string, unknown> | Error): void => {
if (settled) return;
settled = true;
if (timer) clearTimeout(timer);
unsubscribeMessage();
unsubscribeExit();
if (result instanceof Error) {
reject(result);
} else {
resolve(result);
}
};
unsubscribeMessage = process.onMessage((message) => {
if (message.type === "ready") {
finish(message);
}
});
unsubscribeExit = process.onExit(({ error }) => finish(error));
timer = setTimeout(
() => finish(new Error("Timed out waiting for OMP to become ready")),
OMP_READY_TIMEOUT_MS,
);
});
}
class OmpCliRuntimeSession implements OmpRuntimeSession {
private readonly subscribers = new Set<(event: OmpRuntimeEvent) => void>();
activeBranchEntryId?: string;