diff --git a/packages/server/src/server/agent/providers/jsonl-rpc-process.test.ts b/packages/server/src/server/agent/providers/jsonl-rpc-process.test.ts index 6d5928ac6..025f1e7b8 100644 --- a/packages/server/src/server/agent/providers/jsonl-rpc-process.test.ts +++ b/packages/server/src/server/agent/providers/jsonl-rpc-process.test.ts @@ -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>((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); + }); + }); + + 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>((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); + }); + }); + + 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(); + } + }); }); diff --git a/packages/server/src/server/agent/providers/jsonl-rpc-process.ts b/packages/server/src/server/agent/providers/jsonl-rpc-process.ts index 2e56a0659..ddcbaae8a 100644 --- a/packages/server/src/server/agent/providers/jsonl-rpc-process.ts +++ b/packages/server/src/server/agent/providers/jsonl-rpc-process.ts @@ -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): 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; + 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): 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): Record | 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): 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 | 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 | 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; + } + private handleResponse(response: JsonlRpcResponse): void { if (!response.id) { return; diff --git a/packages/server/src/server/agent/providers/omp/cli-runtime.test.ts b/packages/server/src/server/agent/providers/omp/cli-runtime.test.ts index da195f057..66cd3cb0c 100644 --- a/packages/server/src/server/agent/providers/omp/cli-runtime.test.ts +++ b/packages/server/src/server/agent/providers/omp/cli-runtime.test.ts @@ -14,7 +14,11 @@ type OmpChild = ChildProcessWithoutNullStreams & { killedSignals: Array; }; -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[] = []; + 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[] = []; + 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; + 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" })); + }); }); diff --git a/packages/server/src/server/agent/providers/omp/cli-runtime.ts b/packages/server/src/server/agent/providers/omp/cli-runtime.ts index b080a541b..22acf58ce 100644 --- a/packages/server/src/server/agent/providers/omp/cli-runtime.ts +++ b/packages/server/src/server/agent/providers/omp/cli-runtime.ts @@ -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 { + 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> { + return new Promise((resolve, reject) => { + let settled = false; + let unsubscribeMessage = (): void => {}; + let unsubscribeExit = (): void => {}; + let timer: NodeJS.Timeout | undefined; + const finish = (result: Record | 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;