Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 98 additions & 0 deletions packages/effect-codex-app-server/src/protocol.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<CodexError.CodexAppServerError>();
const rawLines: Array<unknown> = [];
const decoded: Array<unknown> = [];
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<CodexError.CodexAppServerError>();
const rawLines: Array<unknown> = [];
const decoded: Array<unknown> = [];
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) =>
Expand Down
67 changes: 58 additions & 9 deletions packages/effect-codex-app-server/src/protocol.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand All @@ -39,6 +44,7 @@ export interface CodexAppServerPatchedProtocolOptions {
readonly terminationError?: Effect.Effect<CodexError.CodexAppServerError>;
readonly logIncoming?: boolean;
readonly logOutgoing?: boolean;
readonly maxIncomingMessageBytes?: number;
readonly logger?: (event: CodexAppServerProtocolLogEvent) => Effect.Effect<void, never>;
readonly onNotification?: (
notification: CodexAppServerIncomingNotification,
Expand Down Expand Up @@ -164,7 +170,9 @@ export const makeCodexAppServerPatchedProtocol = Effect.fn("makeCodexAppServerPa
yield* Queue.sliding<CodexAppServerIncomingRequest>(MAX_BUFFERED_RAW_MESSAGES);
const pending = yield* Ref.make(new Map<string, CodexAppServerPendingRequest>());
const nextRequestId = yield* Ref.make(1);
const maxIncomingMessageBytes = options.maxIncomingMessageBytes ?? MAX_INCOMING_MESSAGE_BYTES;
const remainder: Array<string> = [];
let remainderBytes = 0;
const terminationHandled = yield* Ref.make(false);
const terminationFailure = yield* Ref.make(Option.none<CodexError.CodexAppServerError>());
const terminationSignal = yield* Deferred.make<void>();
Expand Down Expand Up @@ -399,38 +407,79 @@ export const makeCodexAppServerPatchedProtocol = Effect.fn("makeCodexAppServerPa
Stream.interruptWhen(Deferred.await(terminationSignal)),
Stream.decodeText(),
Stream.runForEach((chunk) =>
Effect.sync(() => {
Effect.suspend(() => {
const lines: Array<string> = [];
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) =>
handleTermination(() =>
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: () =>
Expand Down
Loading