diff --git a/packages/effect-codex-app-server/src/protocol.test.ts b/packages/effect-codex-app-server/src/protocol.test.ts index 7249afff1071..9d916b9133e8 100644 --- a/packages/effect-codex-app-server/src/protocol.test.ts +++ b/packages/effect-codex-app-server/src/protocol.test.ts @@ -325,6 +325,53 @@ it.layer(NodeServices.layer)("effect-codex-app-server protocol", (it) => { }), ); + it.effect("rejects an oversized fragmented message before join or parse", () => + Effect.gen(function* () { + const { stdio, input, output } = yield* makeInMemoryStdio(); + const termination = yield* Deferred.make(); + const rawLines: Array = []; + const decoded: Array = []; + let notificationCount = 0; + const transport = yield* CodexProtocol.makeCodexAppServerPatchedProtocol({ + stdio, + maxIncomingMessageBytes: 32, + logIncoming: true, + logger: (event) => + Effect.sync(() => { + if (event.stage === "raw") { + rawLines.push(event.payload); + } + if (event.stage === "decoded") { + decoded.push(event.payload); + } + }), + onNotification: () => Effect.sync(() => notificationCount++).pipe(Effect.asVoid), + onTermination: (error) => Deferred.succeed(termination, error).pipe(Effect.asVoid), + }); + const pending = yield* transport.request("thread/read", {}).pipe(Effect.forkScoped); + yield* Queue.take(output); + + yield* Queue.offer(input, encoder.encode('{"method":"x/huge"')); + yield* Queue.offer(input, encoder.encode(',"params":"xxxxxxxx')); + yield* Queue.offer(input, encoder.encode('xxxxxxxx"}\n')); + + const error = yield* Deferred.await(termination); + assert.instanceOf(error, CodexError.CodexAppServerTransportError); + assert.equal(error.operation, "read-input-stream"); + assert.equal(rawLines.length, 0); + assert.equal(decoded.length, 0); + assert.equal(notificationCount, 0); + + 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); + }), + ); + 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 4a32973a988b..b7111d8f4d6e 100644 --- a/packages/effect-codex-app-server/src/protocol.ts +++ b/packages/effect-codex-app-server/src/protocol.ts @@ -16,6 +16,10 @@ const isJsonRpcId = Schema.is(JsonRpcId); const isJsonRpcResponseEnvelope = Schema.is(JsonRpcResponseEnvelope); const isCodexAppServerError = Schema.is(CodexError.CodexAppServerError); const MAX_BUFFERED_RAW_MESSAGES = 32; +// Decoded remainder size before join/parse. 128 MiB sits above observed Codex +// diffs (~49M characters) and Effect ndjson's 16 MiB default, and well below a +// V8 heap-threatening line. Tests inject a smaller ceiling. +const MAX_INCOMING_MESSAGE_BYTES = 128 * 1024 * 1024; export interface CodexAppServerProtocolLogEvent { readonly direction: "incoming" | "outgoing"; @@ -39,6 +43,7 @@ export interface CodexAppServerPatchedProtocolOptions { readonly terminationError?: Effect.Effect; readonly logIncoming?: boolean; readonly logOutgoing?: boolean; + readonly maxIncomingMessageBytes?: number; readonly logger?: (event: CodexAppServerProtocolLogEvent) => Effect.Effect; readonly onNotification?: ( notification: CodexAppServerIncomingNotification, @@ -139,6 +144,12 @@ const normalizeIncomingError = ( cause: error, }); +const incomingMessageTooLarge = (maxIncomingMessageBytes: number) => + new CodexError.CodexAppServerTransportError({ + operation: "read-input-stream", + cause: new Error(`Incoming message exceeded ${String(maxIncomingMessageBytes)} bytes.`), + }); + const toProtocolMessage = ( requestId: string | number, fields: { @@ -164,7 +175,9 @@ 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 = []; + let remainderBytes = 0; const terminationHandled = yield* Ref.make(false); const terminationFailure = yield* Ref.make(Option.none()); const terminationSignal = yield* Deferred.make(); @@ -399,25 +412,39 @@ export const makeCodexAppServerPatchedProtocol = Effect.fn("makeCodexAppServerPa Stream.interruptWhen(Deferred.await(terminationSignal)), Stream.decodeText(), Stream.runForEach((chunk) => - Effect.sync(() => { + Effect.suspend(() => { const lines: Array = []; let start = 0; + const retainRange = (from: number, to: number) => { + const fragmentLength = to - from; + if (remainderBytes + fragmentLength > maxIncomingMessageBytes) { + remainder.length = 0; + remainderBytes = 0; + return false; + } + remainder.push(chunk.slice(from, to)); + remainderBytes += fragmentLength; + return true; + }; for ( let newline = chunk.indexOf("\n"); newline !== -1; newline = chunk.indexOf("\n", start) ) { - remainder.push(chunk.slice(start, newline)); + if (!retainRange(start, newline)) { + return Effect.fail(incomingMessageTooLarge(maxIncomingMessageBytes)); + } lines.push(remainder.join("").replace(/\r$/, "")); remainder.length = 0; + remainderBytes = 0; start = newline + 1; } // Keep unfinished lines in fragments so each chunk is scanned only once. - if (start < chunk.length) { - remainder.push(chunk.slice(start)); + if (start < chunk.length && !retainRange(start, chunk.length)) { + return Effect.fail(incomingMessageTooLarge(maxIncomingMessageBytes)); } - 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 +452,17 @@ export const makeCodexAppServerPatchedProtocol = Effect.fn("makeCodexAppServerPa Effect.succeed(normalizeIncomingError(error, "read-input-stream")), ), onSuccess: () => - Effect.sync(() => { + Effect.suspend(() => { + if (remainderBytes > maxIncomingMessageBytes) { + remainder.length = 0; + remainderBytes = 0; + return Effect.fail(incomingMessageTooLarge(maxIncomingMessageBytes)); + } const line = remainder.join(""); remainder.length = 0; - return line; + remainderBytes = 0; + return handleLine(line); }).pipe( - Effect.flatMap(handleLine), Effect.matchEffect({ onFailure: (error) => handleTermination(() => Effect.succeed(error)), onSuccess: () =>