diff --git a/packages/effect-codex-app-server/src/protocol.test.ts b/packages/effect-codex-app-server/src/protocol.test.ts index 7249afff10713..be538199457ad 100644 --- a/packages/effect-codex-app-server/src/protocol.test.ts +++ b/packages/effect-codex-app-server/src/protocol.test.ts @@ -34,6 +34,59 @@ const decodeConsumeRateLimitResetCreditResponse = Schema.decodeUnknownEffect( CodexRpc.CLIENT_REQUEST_RESPONSES["account/rateLimitResetCredit/consume"], ); +const lineFramings = [ + { label: "LF", frame: (line: string) => [`${line}\n`] }, + { label: "CRLF", frame: (line: string) => [`${line}\r\n`] }, + { label: "CRLF split across chunks", frame: (line: string) => [`${line}\r`, "\n"] }, + { label: "CRLF split by an empty chunk", frame: (line: string) => [`${line}\r`, "", "\n"] }, +]; + +// REMAINDER_BLOCK_LENGTH in protocol.ts: unfinished fragments shorter than this +// many UTF-16 code units are joined into one entry once they add up to it. +const remainderBlockLength = 1024; + +// Text that differs by position, so a fragment kept out of order changes it. +const positionalText = (length: number) => + Array.from({ length }, (_, index) => index.toString(36)) + .join("") + .slice(0, length); + +const makeLimitedProtocol = (maxIncomingMessageBytes: number) => + Effect.gen(function* () { + const { stdio, input, output } = yield* makeInMemoryStdio(); + const termination = yield* Deferred.make(); + const terminations: Array = []; + const notifications: Array = []; + const events: Array = []; + const transport = yield* CodexProtocol.makeCodexAppServerPatchedProtocol({ + stdio, + maxIncomingMessageBytes, + logIncoming: true, + logger: (event) => + Effect.sync(() => { + events.push(event); + }), + onNotification: (notification) => + Effect.sync(() => { + notifications.push(notification); + }), + onTermination: (error) => + Effect.sync(() => terminations.push(error)).pipe( + Effect.andThen(Deferred.succeed(termination, error)), + Effect.asVoid, + ), + }); + return { input, output, transport, termination, terminations, notifications, events }; + }); + +const awaitMessageLimitFailure = (termination: Deferred.Deferred) => + Effect.gen(function* () { + const error = yield* Deferred.await(termination); + assert.instanceOf(error, CodexError.CodexAppServerTransportError); + assert.equal(error.operation, "read-input-stream"); + return error; + }); + it.layer(NodeServices.layer)("effect-codex-app-server protocol", (it) => { it.effect("maps account usage responses to the upstream token usage schema", () => Effect.gen(function* () { @@ -325,6 +378,393 @@ it.layer(NodeServices.layer)("effect-codex-app-server protocol", (it) => { }), ); + it.effect.each(lineFramings)( + "accepts a message exactly at the byte limit with $label framing", + ({ frame }) => + Effect.gen(function* () { + const message = { method: "x/limit", params: { text: "at the limit" } }; + const line = encodeUnknownJsonString(message); + const { input, termination, notifications } = yield* makeLimitedProtocol( + encoder.encode(line).length, + ); + + for (const chunk of frame(line)) { + yield* Queue.offer(input, encoder.encode(chunk)); + } + yield* Queue.end(input); + + assert.instanceOf( + yield* Deferred.await(termination), + CodexError.CodexAppServerInputStreamEndedError, + ); + assert.deepEqual(notifications, [message]); + }), + ); + + it.effect.each(lineFramings)( + "rejects a message one byte over the limit with $label framing", + ({ frame }) => + Effect.gen(function* () { + const line = encodeUnknownJsonString({ method: "x/over", params: { text: "one over" } }); + const { input, output, transport, termination, terminations, notifications, events } = + yield* makeLimitedProtocol(encoder.encode(line).length - 1); + const pending = yield* transport.request("thread/read", {}).pipe(Effect.forkScoped); + yield* Queue.take(output); + + for (const chunk of frame(line)) { + yield* Queue.offer(input, encoder.encode(chunk)); + } + yield* Queue.end(input); + + const error = yield* awaitMessageLimitFailure(termination); + const pendingError = yield* Fiber.join(pending).pipe( + Effect.match({ + onFailure: (failure) => failure, + onSuccess: () => assert.fail("Expected the oversized message to fail the request"), + }), + ); + assert.strictEqual(pendingError, error); + assert.lengthOf(terminations, 1); + assert.strictEqual(terminations[0], error); + assert.deepEqual(notifications, []); + assert.deepEqual(events, []); + }), + ); + + it.effect( + "rejects an over-limit message from small fragments before its newline arrives", + () => + Effect.gen(function* () { + const { input, output, transport, termination, notifications, events } = + yield* makeLimitedProtocol(1024); + const pending = yield* transport.request("thread/read", {}).pipe(Effect.forkScoped); + yield* Queue.take(output); + + yield* Queue.offer(input, encoder.encode('{"method":"x/huge","params":{"text":"')); + for (let index = 0; index < 512; index++) { + yield* Queue.offer(input, encoder.encode("xxxxxxxx")); + } + + const error = yield* awaitMessageLimitFailure(termination); + const pendingError = yield* Fiber.join(pending).pipe( + Effect.match({ + onFailure: (failure) => failure, + onSuccess: () => assert.fail("Expected the oversized message to fail the request"), + }), + ); + assert.strictEqual(pendingError, error); + assert.deepEqual(notifications, []); + assert.deepEqual(events, []); + }), + 5_000, + ); + + it.effect.each([ + { label: "LF after several blocks", length: 5 * remainderBlockLength + 300, terminator: "\n" }, + { + label: "CRLF whose \\r completes a block", + length: 4 * remainderBlockLength - 1, + terminator: "\r\n", + }, + ])( + "routes a message sent one byte per chunk at the byte limit with $label", + ({ length, terminator }) => + Effect.gen(function* () { + const empty = encodeUnknownJsonString({ method: "x/bytes", params: { text: "" } }); + const message = { + method: "x/bytes", + params: { text: positionalText(length - empty.length) }, + }; + const line = encodeUnknownJsonString(message); + assert.lengthOf(line, length); + const { input, termination, notifications } = yield* makeLimitedProtocol(length); + + const bytes = encoder.encode(`${line}${terminator}`); + for (let offset = 0; offset < bytes.length; offset++) { + yield* Queue.offer(input, bytes.subarray(offset, offset + 1)); + } + yield* Queue.end(input); + + assert.instanceOf( + yield* Deferred.await(termination), + CodexError.CodexAppServerInputStreamEndedError, + ); + assert.deepEqual(notifications, [message]); + }), + ); + + it.effect("keeps one-byte chunks in order around a chunk longer than a block", () => + Effect.gen(function* () { + const message = { + method: "x/mixed", + params: { text: positionalText(4 * remainderBlockLength) }, + }; + const bytes = encodeJsonl(message); + const { input, termination, notifications } = yield* makeLimitedProtocol(bytes.length - 1); + + const longEnd = 100 + 2 * remainderBlockLength; + const chunks = [ + ...Array.from({ length: 100 }, (_, offset) => bytes.subarray(offset, offset + 1)), + bytes.subarray(100, longEnd), + ...Array.from({ length: 100 }, (_, index) => + bytes.subarray(longEnd + index, longEnd + index + 1), + ), + bytes.subarray(longEnd + 100), + ]; + for (const chunk of chunks) { + yield* Queue.offer(input, chunk); + } + yield* Queue.end(input); + + assert.instanceOf( + yield* Deferred.await(termination), + CodexError.CodexAppServerInputStreamEndedError, + ); + assert.deepEqual(notifications, [message]); + }), + ); + + it.effect( + "rejects an over-limit message sent one byte per chunk before its newline arrives", + () => + Effect.gen(function* () { + const limit = 3 * remainderBlockLength + 5; + const { input, output, transport, termination, notifications, events } = + yield* makeLimitedProtocol(limit); + const pending = yield* transport.request("thread/read", {}).pipe(Effect.forkScoped); + yield* Queue.take(output); + + const bytes = encoder.encode(`{"method":"x/huge","params":{"text":"${"x".repeat(limit)}`); + for (let offset = 0; offset < bytes.length; offset++) { + yield* Queue.offer(input, bytes.subarray(offset, offset + 1)); + } + + const error = yield* awaitMessageLimitFailure(termination); + const pendingError = yield* Fiber.join(pending).pipe( + Effect.match({ + onFailure: (failure) => failure, + onSuccess: () => assert.fail("Expected the oversized message to fail the request"), + }), + ); + assert.strictEqual(pendingError, error); + assert.deepEqual(notifications, []); + assert.deepEqual(events, []); + }), + 5_000, + ); + + it.effect.each([1, 7])( + "keeps multi-byte characters intact across remainder blocks with %i-byte input chunks", + (chunkSize) => + Effect.gen(function* () { + const head = '{"method":"x/unicode","params":{"text":"'; + // The emoji's surrogate pair straddles the first block edge; 3-byte + // characters fill the blocks after it. + const text = `${"Γ©".repeat(remainderBlockLength - 1 - head.length)}πŸ™‚${"ζœ€".repeat(3 * remainderBlockLength)}`; + const message = { method: "x/unicode", params: { text } }; + const line = encodeUnknownJsonString(message); + assert.strictEqual(line.indexOf("πŸ™‚"), remainderBlockLength - 1); + const bytes = encoder.encode(`${line}\n`); + const { input, termination, notifications } = yield* makeLimitedProtocol(bytes.length - 1); + + for (let offset = 0; offset < bytes.length; offset += chunkSize) { + yield* Queue.offer(input, bytes.subarray(offset, offset + chunkSize)); + } + yield* Queue.end(input); + + assert.instanceOf( + yield* Deferred.await(termination), + CodexError.CodexAppServerInputStreamEndedError, + ); + assert.deepEqual(notifications, [message]); + }), + ); + + it.effect.each([ + { label: "short", prefixLength: 64 }, + { label: "long", prefixLength: remainderBlockLength + 76 }, + ])( + "routes a message with a $label first fragment after complete lines in a large chunk", + ({ prefixLength }) => + Effect.gen(function* () { + const before = { method: "x/before", params: { index: 1 } }; + const message = { + method: "x/after", + params: { text: `ζœ€πŸ™‚${positionalText(4 * remainderBlockLength)}` }, + }; + const line = encodeUnknownJsonString(message); + const lineLimit = 64 * 1024; + const { input, termination, notifications } = yield* makeLimitedProtocol(lineLimit); + + // Whitespace lines make the chunk far longer than the fragment kept from its end. + const blankLines = `${" ".repeat(lineLimit)}\n`.repeat(16); + yield* Queue.offer( + input, + encoder.encode( + `${blankLines}${encodeUnknownJsonString(before)}\n${line.slice(0, prefixLength)}`, + ), + ); + yield* Queue.offer(input, encoder.encode(`${line.slice(prefixLength)}\n`)); + yield* Queue.end(input); + + assert.instanceOf( + yield* Deferred.await(termination), + CodexError.CodexAppServerInputStreamEndedError, + ); + assert.deepEqual(notifications, [before, message]); + }), + ); + + it.effect("routes an unterminated final message at the byte limit when input ends", () => + Effect.gen(function* () { + const message = { method: "x/final", params: { text: "no newline" } }; + const line = encodeUnknownJsonString(message); + const { input, termination, notifications } = yield* makeLimitedProtocol( + encoder.encode(line).length, + ); + + yield* Queue.offer(input, encoder.encode(line)); + yield* Queue.end(input); + + assert.instanceOf( + yield* Deferred.await(termination), + CodexError.CodexAppServerInputStreamEndedError, + ); + assert.deepEqual(notifications, [message]); + }), + ); + + it.effect.each([ + { label: "an unterminated final message", suffix: "" }, + { label: "a final message ending in a bare \\r", suffix: "\r" }, + ])("rejects $label one byte over the limit when input ends", ({ suffix }) => + Effect.gen(function* () { + const line = encodeUnknownJsonString({ method: "x/final", params: { text: "no newline" } }); + const { input, termination, notifications, events } = yield* makeLimitedProtocol( + encoder.encode(line).length + suffix.length - 1, + ); + + yield* Queue.offer(input, encoder.encode(`${line}${suffix}`)); + yield* Queue.end(input); + + yield* awaitMessageLimitFailure(termination); + assert.deepEqual(notifications, []); + assert.deepEqual(events, []); + }), + ); + + it.effect.each([1, 7])( + "counts UTF-8 bytes of characters split across %i-byte input chunks", + (chunkSize) => + Effect.gen(function* () { + const message = { method: "x/unicode", params: { text: "ζœ€εΎŒπŸ™‚πŸ™‚ζœ€εΎŒ" } }; + const line = encodeUnknownJsonString(message); + const bytes = encoder.encode(`${line}\n`); + const lineBytes = bytes.length - 1; + // Counting UTF-16 code units instead of bytes would accept the over-limit message. + assert.isBelow(line.length, lineBytes - 1); + + const atLimit = yield* makeLimitedProtocol(lineBytes); + const overLimit = yield* makeLimitedProtocol(lineBytes - 1); + for (let offset = 0; offset < bytes.length; offset += chunkSize) { + yield* Queue.offer(atLimit.input, bytes.subarray(offset, offset + chunkSize)); + yield* Queue.offer(overLimit.input, bytes.subarray(offset, offset + chunkSize)); + } + yield* Queue.end(atLimit.input); + yield* Queue.end(overLimit.input); + + assert.instanceOf( + yield* Deferred.await(atLimit.termination), + CodexError.CodexAppServerInputStreamEndedError, + ); + assert.deepEqual(atLimit.notifications, [message]); + yield* awaitMessageLimitFailure(overLimit.termination); + assert.deepEqual(overLimit.notifications, []); + assert.deepEqual(overLimit.events, []); + }), + ); + + it.effect("counts an invalid UTF-8 byte as the three bytes of its U+FFFD replacement", () => + Effect.gen(function* () { + const message = { method: "x/invalid", params: { text: "a\uFFFDb" } }; + const line = encodeUnknownJsonString(message); + const lineBytes = encoder.encode(line).length; + // 0xff is never valid UTF-8, so the decoder replaces it with U+FFFD. + const [head, tail] = line.split("\uFFFD"); + const bytes = Uint8Array.from([ + ...encoder.encode(head), + 0xff, + ...encoder.encode(`${tail}\n`), + ]); + // Counting raw input bytes instead would accept the over-limit message. + assert.strictEqual(bytes.length - 1, lineBytes - 2); + + const atLimit = yield* makeLimitedProtocol(lineBytes); + const overLimit = yield* makeLimitedProtocol(lineBytes - 1); + yield* Queue.offer(atLimit.input, bytes); + yield* Queue.offer(overLimit.input, bytes); + yield* Queue.end(atLimit.input); + yield* Queue.end(overLimit.input); + + assert.instanceOf( + yield* Deferred.await(atLimit.termination), + CodexError.CodexAppServerInputStreamEndedError, + ); + assert.deepEqual(atLimit.notifications, [message]); + yield* awaitMessageLimitFailure(overLimit.termination); + assert.deepEqual(overLimit.notifications, []); + assert.deepEqual(overLimit.events, []); + }), + ); + + it.effect("keeps an over-limit message out of protocol logs and the termination error", () => + Effect.gen(function* () { + const secret = "codex-oversized-secret-sentinel"; + const small = { method: "x/small", params: {} }; + const huge = encodeJsonl({ method: "x/huge", params: { text: secret } }); + const { input, termination, notifications, events } = yield* makeLimitedProtocol(64); + + yield* Queue.offer(input, encodeJsonl(small)); + yield* Queue.offer(input, huge); + yield* Queue.end(input); + + const error = yield* awaitMessageLimitFailure(termination); + assert.deepEqual(notifications, [small]); + assert.deepEqual( + events.map(({ stage }) => stage), + ["raw", "decoded"], + ); + assert.notInclude(encodeUnknownJsonString(events), secret); + assert.notInclude(error.message, secret); + assert.notInclude(String(error.cause), secret); + assert.include(String(error.cause), `reached ${huge.length - 1} bytes`); + assert.include(String(error.cause), "over the 64-byte limit"); + }), + ); + + it.effect.each([ + { label: "a terminated", newline: "\n" }, + { label: "an unterminated", newline: "" }, + ])("dispatches complete lines that precede $label over-limit line in one chunk", ({ newline }) => + Effect.gen(function* () { + const first = { method: "x/first", params: { index: 1 } }; + const second = { method: "x/second", params: { index: 2 } }; + const oversized = { method: "x/huge", params: { text: "y".repeat(64) } }; + const { input, termination, notifications } = yield* makeLimitedProtocol(64); + + yield* Queue.offer( + input, + encoder.encode( + `${[first, second, oversized].map((message) => encodeUnknownJsonString(message)).join("\n")}${newline}`, + ), + ); + yield* Queue.end(input); + + yield* awaitMessageLimitFailure(termination); + assert.deepEqual(notifications, [first, second]); + }), + ); + it.effect.each([1, 7, 1024])( "preserves JSONL framing and UTF-8 across %i-byte input chunks", (chunkSize) => diff --git a/packages/effect-codex-app-server/src/protocol.ts b/packages/effect-codex-app-server/src/protocol.ts index 4a32973a988bd..7d3b05e0f5dcd 100644 --- a/packages/effect-codex-app-server/src/protocol.ts +++ b/packages/effect-codex-app-server/src/protocol.ts @@ -16,6 +16,36 @@ const isJsonRpcId = Schema.is(JsonRpcId); const isJsonRpcResponseEnvelope = Schema.is(JsonRpcResponseEnvelope); const isCodexAppServerError = Schema.is(CodexError.CodexAppServerError); const MAX_BUFFERED_RAW_MESSAGES = 32; +// Largest incoming JSONL message, in UTF-8 bytes without its line terminator. +// Codex diffs reach ~49M characters and grow when JSON-escaped; a line far +// past this ceiling would exhaust the heap during join and parse. +const MAX_INCOMING_MESSAGE_BYTES = 128 * 1024 * 1024; +// Unfinished-message fragments shorter than this many UTF-16 code units are +// joined into one entry once they add up to it, so a peer that sends one byte +// per read holds one entry per block rather than one per byte. A whole longer +// read, such as a 64 KiB pipe read, is kept as it is without a copy. +const REMAINDER_BLOCK_LENGTH = 1024; + +// `text.slice(start)` as a new string that does not keep `text` alive. V8 +// returns a slice of 13 or more code units, and a one-item join of that slice, +// as a view that holds the whole of `text`; a join of two items builds a new +// flat string. +const sliceDetached = (text: string, start: number) => + [text.slice(start, start + 1), text.slice(start + 1)].join(""); + +// UTF-8 size of `text.slice(start, end)` without encoding a copy. Text from a +// streaming TextDecoder holds no lone surrogates, so each surrogate is one half +// of a four-byte sequence. +const utf8ByteLength = (text: string, start: number, end: number) => { + let bytes = 0; + for (let index = start; index < end; index++) { + const code = text.charCodeAt(index); + if (code < 0x80) bytes += 1; + else if (code < 0x800 || (code >= 0xd800 && code <= 0xdfff)) bytes += 2; + else bytes += 3; + } + return bytes; +}; export interface CodexAppServerProtocolLogEvent { readonly direction: "incoming" | "outgoing"; @@ -39,6 +69,8 @@ export interface CodexAppServerPatchedProtocolOptions { readonly terminationError?: Effect.Effect; readonly logIncoming?: boolean; readonly logOutgoing?: boolean; + // Test override; production uses the 128 MiB default. + readonly maxIncomingMessageBytes?: number; readonly logger?: (event: CodexAppServerProtocolLogEvent) => Effect.Effect; readonly onNotification?: ( notification: CodexAppServerIncomingNotification, @@ -164,7 +196,11 @@ export const makeCodexAppServerPatchedProtocol = Effect.fn("makeCodexAppServerPa yield* Queue.sliding(MAX_BUFFERED_RAW_MESSAGES); const pending = yield* Ref.make(new Map()); const nextRequestId = yield* Ref.make(1); + const maxIncomingMessageBytes = options.maxIncomingMessageBytes ?? MAX_INCOMING_MESSAGE_BYTES; const remainder: Array = []; + const remainderTail: Array = []; + let remainderTailLength = 0; + let remainderBytes = 0; const terminationHandled = yield* Ref.make(false); const terminationFailure = yield* Ref.make(Option.none()); const terminationSignal = yield* Deferred.make(); @@ -395,29 +431,85 @@ export const makeCodexAppServerPatchedProtocol = Effect.fn("makeCodexAppServerPa ); }; + // Fragments shorter than a block wait in the tail, which stays under a block + // of code units and so under that many entries. Packing joins the tail into + // one remainder entry, so a short fragment is copied here and once more when + // its line is joined, plus once on arrival if it follows a newline. + const packRemainderTail = () => { + if (remainderTail.length > 0) { + remainder.push(remainderTail.join("")); + remainderTail.length = 0; + remainderTailLength = 0; + } + }; + const retainFragment = (fragment: string) => { + if (fragment.length >= REMAINDER_BLOCK_LENGTH) { + packRemainderTail(); + remainder.push(fragment); + return; + } + remainderTail.push(fragment); + remainderTailLength += fragment.length; + if (remainderTailLength >= REMAINDER_BLOCK_LENGTH) { + packRemainderTail(); + } + }; + const clearRemainder = () => { + remainder.length = 0; + remainderTail.length = 0; + remainderTailLength = 0; + remainderBytes = 0; + }; + yield* options.stdio.stdin.pipe( Stream.interruptWhen(Deferred.await(terminationSignal)), Stream.decodeText(), Stream.runForEach((chunk) => - Effect.sync(() => { + Effect.suspend(() => { const lines: Array = []; let start = 0; - for ( - let newline = chunk.indexOf("\n"); - newline !== -1; - newline = chunk.indexOf("\n", start) - ) { + while (start < chunk.length) { + const newline = chunk.indexOf("\n", start); + const end = newline === -1 ? chunk.length : newline; + const bytes = remainderBytes + utf8ByteLength(chunk, start, end); + // A \r ending the fragment may be half of a CRLF terminator, so it does + // not count yet. An empty fragment adds nothing to a remainder that fits. + const overLimit = + end > start && + bytes - (chunk.charCodeAt(end - 1) === 0x0d ? 1 : 0) > maxIncomingMessageBytes; + if (overLimit) { + clearRemainder(); + // Complete lines from earlier in this chunk are still handled, in order. + return Effect.forEach(lines, handleLine, { discard: true }).pipe( + Effect.andThen( + Effect.fail( + new CodexError.CodexAppServerTransportError({ + operation: "read-input-stream", + cause: new Error( + `Incoming message reached ${bytes} bytes, over the ${maxIncomingMessageBytes}-byte limit.`, + ), + }), + ), + ), + ); + } + if (newline === -1) { + // Keep unfinished lines in fragments so each chunk is scanned only once. + // A fragment after a newline is copied so it does not keep the lines + // before it alive. Only a message's first fragment can follow a + // newline, so a message pays for at most one such copy of one read. + retainFragment(start === 0 ? chunk : sliceDetached(chunk, start)); + remainderBytes = bytes; + break; + } + packRemainderTail(); remainder.push(chunk.slice(start, newline)); lines.push(remainder.join("").replace(/\r$/, "")); - remainder.length = 0; + clearRemainder(); start = newline + 1; } - // Keep unfinished lines in fragments so each chunk is scanned only once. - if (start < chunk.length) { - remainder.push(chunk.slice(start)); - } - return lines; - }).pipe(Effect.flatMap((lines) => Effect.forEach(lines, handleLine, { discard: true }))), + return Effect.forEach(lines, handleLine, { discard: true }); + }), ), Effect.matchEffect({ onFailure: (error) => @@ -425,12 +517,23 @@ export const makeCodexAppServerPatchedProtocol = Effect.fn("makeCodexAppServerPa Effect.succeed(normalizeIncomingError(error, "read-input-stream")), ), onSuccess: () => - Effect.sync(() => { + Effect.suspend(() => { + // With no terminator left to strip, a final \r counts toward the limit. + if (remainderBytes > maxIncomingMessageBytes) { + const error = new CodexError.CodexAppServerTransportError({ + operation: "read-input-stream", + cause: new Error( + `Incoming message reached ${remainderBytes} bytes, over the ${maxIncomingMessageBytes}-byte limit.`, + ), + }); + clearRemainder(); + return Effect.fail(error); + } + packRemainderTail(); const line = remainder.join(""); - remainder.length = 0; - return line; + clearRemainder(); + return handleLine(line); }).pipe( - Effect.flatMap(handleLine), Effect.matchEffect({ onFailure: (error) => handleTermination(() => Effect.succeed(error)), onSuccess: () =>