diff --git a/packages/effect-codex-app-server/src/protocol.test.ts b/packages/effect-codex-app-server/src/protocol.test.ts index 7249afff1071..a27a898e5e58 100644 --- a/packages/effect-codex-app-server/src/protocol.test.ts +++ b/packages/effect-codex-app-server/src/protocol.test.ts @@ -325,6 +325,104 @@ 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("rejects an oversized multibyte message counted in UTF-8 bytes", () => + Effect.gen(function* () { + const { stdio, input, output } = yield* makeInMemoryStdio(); + const termination = yield* Deferred.make(); + const rawLines: Array = []; + const decoded: Array = []; + let notificationCount = 0; + const line = '{"method":"x","params":"你好你好"}'; + assert.ok(line.length <= 32); + assert.ok(encoder.encode(line).byteLength > 32); + 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"')); + yield* Queue.offer(input, encoder.encode(',"params":"你好')); + yield* Queue.offer(input, encoder.encode('你好"}\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 multibyte 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..52386f2d1c2f 100644 --- a/packages/effect-codex-app-server/src/protocol.ts +++ b/packages/effect-codex-app-server/src/protocol.ts @@ -16,6 +16,11 @@ const isJsonRpcId = Schema.is(JsonRpcId); const isJsonRpcResponseEnvelope = Schema.is(JsonRpcResponseEnvelope); const isCodexAppServerError = Schema.is(CodexError.CodexAppServerError); const MAX_BUFFERED_RAW_MESSAGES = 32; +// UTF-8 byte size of decoded remainder 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; +const utf8 = new TextEncoder(); export interface CodexAppServerProtocolLogEvent { readonly direction: "incoming" | "outgoing"; @@ -39,6 +44,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, @@ -164,7 +170,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 +407,54 @@ 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 fragment = chunk.slice(from, to); + const fragmentLength = utf8.encode(fragment).byteLength; + if (remainderBytes + fragmentLength > maxIncomingMessageBytes) { + remainder.length = 0; + remainderBytes = 0; + return false; + } + remainder.push(fragment); + 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( + new CodexError.CodexAppServerTransportError({ + operation: "read-input-stream", + cause: new Error( + `Incoming message exceeded ${String(maxIncomingMessageBytes)} bytes.`, + ), + }), + ); + } 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( + new CodexError.CodexAppServerTransportError({ + operation: "read-input-stream", + cause: new Error( + `Incoming message exceeded ${String(maxIncomingMessageBytes)} bytes.`, + ), + }), + ); } - 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 +462,24 @@ 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( + new CodexError.CodexAppServerTransportError({ + operation: "read-input-stream", + cause: new Error( + `Incoming message exceeded ${String(maxIncomingMessageBytes)} bytes.`, + ), + }), + ); + } 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: () =>