diff --git a/packages/effect-acp/src/agent.ts b/packages/effect-acp/src/agent.ts index 578509f9c85d9..ec6e2379342e6 100644 --- a/packages/effect-acp/src/agent.ts +++ b/packages/effect-acp/src/agent.ts @@ -35,6 +35,11 @@ export class AcpAgent extends Context.Service< AcpAgent, { readonly raw: { + /** + * Stream of inbound ACP notifications observed on the connection. Unread + * ones are kept best-effort: the newest 32 whose JSON lines fit in 16Mi + * characters, always including the latest. Ends when the connection does. + */ readonly notifications: Stream.Stream; readonly request: ( method: string, diff --git a/packages/effect-acp/src/client.ts b/packages/effect-acp/src/client.ts index ca1dba902ee8d..1d8f06c862def 100644 --- a/packages/effect-acp/src/client.ts +++ b/packages/effect-acp/src/client.ts @@ -60,6 +60,11 @@ export interface AcpClientOptions { } type AcpClientRaw = { + /** + * Inbound notifications. Unread ones are kept best-effort: the newest 32 whose + * JSON lines fit in 16Mi characters, always including the latest. Ends with the + * connection. + */ readonly notifications: Stream.Stream; readonly request: (method: string, payload: unknown) => Effect.Effect; readonly notify: (method: string, payload: unknown) => Effect.Effect; diff --git a/packages/effect-acp/src/protocol.test.ts b/packages/effect-acp/src/protocol.test.ts index eae3394a1653f..cbb27b9508c23 100644 --- a/packages/effect-acp/src/protocol.test.ts +++ b/packages/effect-acp/src/protocol.test.ts @@ -62,6 +62,16 @@ const JsonRpcErrorResponse = Schema.Struct({ }); const encodeUnknownJsonString = Schema.encodeUnknownSync(Schema.fromJsonString(Schema.Unknown)); const encoder = new TextEncoder(); +const MIB = 1024 * 1024; +const decodeIndex = Schema.decodeUnknownSync(Schema.Struct({ index: Schema.Number })); +const largeNotification = (index: number, textLength: number) => + encoder.encode( + `${encodeUnknownJsonString({ + jsonrpc: "2.0", + method: "x/large", + params: { index, text: "x".repeat(textLength) }, + })}\n`, + ); const mockPeerPath = Effect.map(Effect.service(Path.Path), (path) => path.join(import.meta.dirname, "../test/fixtures/acp-mock-peer.ts"), ); @@ -148,6 +158,173 @@ it.layer(NodeServices.layer)("effect-acp protocol", (it) => { }), ); + it.effect("keeps only recent raw notifications after their callbacks run", () => + Effect.gen(function* () { + const { stdio, input } = yield* makeInMemoryStdio(); + const handled = yield* Deferred.make(); + let handledCount = 0; + const transport = yield* AcpProtocol.makeAcpPatchedProtocol({ + stdio, + serverRequestMethods: new Set(), + onNotification: () => + Effect.sync(() => ++handledCount).pipe( + Effect.flatMap((count) => + count === 64 ? Deferred.succeed(handled, undefined).pipe(Effect.asVoid) : Effect.void, + ), + ), + }); + + const messages = Array.from({ length: 64 }, (_, index) => + encodeUnknownJsonString({ + jsonrpc: "2.0", + method: "x/performance", + params: { index }, + }), + ); + yield* Queue.offer(input, encoder.encode(`${messages.join("\n")}\n`)); + yield* Deferred.await(handled); + + const retained = yield* transport.incoming.pipe(Stream.take(32), Stream.runCollect); + + assert.equal(handledCount, 64); + assert.equal(retained.length, 32); + assert.deepEqual(retained[0]?.params, { index: 32 }); + assert.deepEqual(retained[31]?.params, { index: 63 }); + }), + ); + + it.effect("keeps the newest unread raw notifications whose JSON fits in 16Mi characters", () => + Effect.gen(function* () { + const { stdio, input } = yield* makeInMemoryStdio(); + const handled = yield* Queue.unbounded(); + const transport = yield* AcpProtocol.makeAcpPatchedProtocol({ + stdio, + serverRequestMethods: new Set(), + onNotification: (notification) => + Queue.offer(handled, decodeIndex(notification.params).index).pipe(Effect.asVoid), + }); + + // Ten 3 MiB notifications: only the newest five fit in the budget. + for (let index = 0; index < 10; index++) { + yield* Queue.offer(input, largeNotification(index, 3 * MIB)); + assert.equal(yield* Queue.take(handled), index); + } + + const retained = yield* transport.incoming.pipe(Stream.take(5), Stream.runCollect); + assert.deepEqual( + retained.map((notification) => decodeIndex(notification.params).index), + [5, 6, 7, 8, 9], + ); + }), + ); + + it.effect("keeps a single unread raw notification that alone exceeds the budget", () => + Effect.gen(function* () { + const { stdio, input } = yield* makeInMemoryStdio(); + const handled = yield* Queue.unbounded(); + const transport = yield* AcpProtocol.makeAcpPatchedProtocol({ + stdio, + serverRequestMethods: new Set(), + onNotification: (notification) => + Queue.offer(handled, decodeIndex(notification.params).index).pipe(Effect.asVoid), + }); + + yield* Queue.offer(input, largeNotification(0, 1)); + yield* Queue.take(handled); + yield* Queue.offer(input, largeNotification(1, 17 * MIB)); + yield* Queue.take(handled); + + const [retained] = yield* transport.incoming.pipe(Stream.take(1), Stream.runCollect); + assert.equal(decodeIndex(retained?.params).index, 1); + }), + ); + + it.effect("delivers every raw notification to a reader that keeps up, then ends", () => + Effect.gen(function* () { + const { stdio, input } = yield* makeInMemoryStdio(); + const transport = yield* AcpProtocol.makeAcpPatchedProtocol({ + stdio, + serverRequestMethods: new Set(), + }); + const received = yield* Queue.unbounded(); + const reader = yield* transport.incoming.pipe( + Stream.map((notification) => decodeIndex(notification.params).index), + Stream.tap((index) => Queue.offer(received, index)), + Stream.runCollect, + Effect.forkScoped, + ); + + // 30 MiB in total, well past the unread budget. + for (let index = 0; index < 10; index++) { + yield* Queue.offer(input, largeNotification(index, 3 * MIB)); + assert.equal(yield* Queue.take(received), index); + } + yield* Queue.end(input); + + assert.deepEqual(yield* Fiber.join(reader), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + }), + ); + + it.effect("ends the raw notification stream when the connection scope closes", () => + Effect.gen(function* () { + const { stdio } = yield* makeInMemoryStdio(); + const connection = yield* Scope.make(); + const transport = yield* AcpProtocol.makeAcpPatchedProtocol({ + stdio, + serverRequestMethods: new Set(), + }).pipe(Scope.provide(connection)); + const reader = yield* transport.incoming.pipe(Stream.runCollect, Effect.forkScoped); + + yield* Scope.close(connection, Exit.void); + + assert.deepEqual(yield* Fiber.join(reader), []); + }), + ); + + it.effect("keeps the ended raw backlog intact when notifications arrive after termination", () => + Effect.gen(function* () { + const { stdio, input } = yield* makeInMemoryStdio(); + const handled = yield* Queue.unbounded(); + const terminated = yield* Deferred.make(); + const transport = yield* AcpProtocol.makeAcpPatchedProtocol({ + stdio: Stdio.make({ + args: stdio.args, + stdin: stdio.stdin, + stdout: () => Sink.forEach(() => Effect.die("stdout failed")), + stderr: () => Sink.drain, + }), + serverRequestMethods: new Set(), + onNotification: (notification) => + Queue.offer(handled, decodeIndex(notification.params).index).pipe(Effect.asVoid), + onTermination: () => Deferred.succeed(terminated, undefined).pipe(Effect.asVoid), + }); + const deliver = (index: number) => + Queue.offer( + input, + encoder.encode( + `${encodeUnknownJsonString({ jsonrpc: "2.0", method: "x/late", params: { index } })}\n`, + ), + ).pipe(Effect.andThen(Queue.take(handled))); + + for (let index = 0; index < 32; index++) { + assert.equal(yield* deliver(index), index); + } + // A failed write terminates the connection while the reader keeps running. + yield* transport.notify("x/fail", {}).pipe(Effect.exit); + yield* Deferred.await(terminated); + assert.equal(yield* deliver(32), 32); + + const retained = yield* transport.incoming.pipe( + Stream.map((notification) => decodeIndex(notification.params).index), + Stream.runCollect, + ); + assert.deepEqual( + retained, + Array.from({ length: 32 }, (_, index) => index), + ); + }), + ); + // clanka once sent a usage_update without `used` and `size`. One bad update // must not end a session whose later updates are fine. it.effect("drops a session/update that fails to decode and keeps reading", () => diff --git a/packages/effect-acp/src/protocol.ts b/packages/effect-acp/src/protocol.ts index 8763e66361e12..ec3bb48606874 100644 --- a/packages/effect-acp/src/protocol.ts +++ b/packages/effect-acp/src/protocol.ts @@ -128,6 +128,17 @@ interface AcpPendingRequest { readonly method: string; } +/** A decoded wire message and the length of the JSON line it arrived on. */ +interface AcpDecodedWireMessage { + readonly message: RpcMessage.FromClientEncoded | RpcMessage.FromServerEncoded; + readonly lineLength: number; +} + +interface AcpBufferedRawNotification { + readonly notification: AcpIncomingNotification; + readonly lineLength: number; +} + interface AcpOutgoingWrite { readonly acknowledgement?: Deferred.Deferred; readonly payload: string | Uint8Array; @@ -168,7 +179,8 @@ const makeStrictNdJsonRpcParser = () => { const decoder = new TextDecoder(); let buffer = ""; return { - decode: (bytes: Uint8Array | string): ReadonlyArray => { + // Every message in a batch line is charged the whole line. + decode: (bytes: Uint8Array | string): ReadonlyArray => { const chunk = typeof bytes === "string" ? bytes : decoder.decode(bytes, { stream: true }); const lastNewline = chunk.lastIndexOf("\n"); // Splitting the whole buffer per chunk is quadratic for multi-megabyte lines. @@ -178,7 +190,14 @@ const makeStrictNdJsonRpcParser = () => { } const lines = (buffer + chunk.slice(0, lastNewline)).split("\n"); buffer = chunk.slice(lastNewline + 1); - return lines.flatMap((line) => codec.decode(line)); + return lines.flatMap((line) => + codec.decode(line).map((message) => ({ + message: normalizeAcpJsonRpcError( + message as RpcMessage.FromClientEncoded | RpcMessage.FromServerEncoded, + ), + lineLength: line.length, + })), + ); }, encode: (response: Parameters[0]) => { const encoded = codec.encode(response); @@ -186,7 +205,11 @@ const makeStrictNdJsonRpcParser = () => { }, }; }; +// Unread raw notifications are best-effort. The newest is always kept; older +// ones are dropped once the backlog passes either limit. The length limit counts +// JSON line characters (UTF-16 code units), not bytes or heap size. const MAX_BUFFERED_RAW_NOTIFICATIONS = 32; +const MAX_BUFFERED_RAW_NOTIFICATION_JSON_LENGTH = 16 * 1024 * 1024; const decodeFailureDiagnostics = (error: AcpError.AcpProtocolParseError) => ({ operation: error.operation, @@ -248,9 +271,14 @@ export const makeAcpPatchedProtocol = Effect.fn("makeAcpPatchedProtocol")(functi const parser = makeStrictNdJsonRpcParser(); const serverQueue = yield* Queue.unbounded(); const clientQueue = yield* Queue.unbounded(); - const notificationQueue = yield* Queue.sliding( - MAX_BUFFERED_RAW_NOTIFICATIONS, - ); + const rawNotifications = yield* Queue.unbounded>(); + // JSON length of the notifications offered since the queue was last empty. + // `incoming` takes the whole backlog on every pull, so an empty queue holds + // nothing, and a read can only make this an overestimate. + let rawNotificationsJsonLength = 0; + // Raw readers finish with the connection: termination ends the queue, and so + // does closing the scope, which interrupts the reader before it terminates. + yield* Effect.addFinalizer(() => Queue.end(rawNotifications)); const disconnects = yield* Queue.unbounded(); const outgoing = yield* Queue.unbounded>(); const outgoingWriterState = yield* Ref.make({ @@ -420,8 +448,27 @@ export const makeAcpPatchedProtocol = Effect.fn("makeAcpPatchedProtocol")(functi ), ); - const dispatchNotification = (notification: AcpIncomingNotification) => - Queue.offer(notificationQueue, notification).pipe( + // Offer first so an ended queue rejects the notification without evicting the + // backlog readers are still draining. Offer and eviction share one sync block, + // so readers never see more than the limits allow. + const bufferRawNotification = (notification: AcpIncomingNotification, lineLength: number) => + Effect.sync(() => { + if (Queue.sizeUnsafe(rawNotifications) === 0) rawNotificationsJsonLength = 0; + if (!Queue.offerUnsafe(rawNotifications, { notification, lineLength })) return; + rawNotificationsJsonLength += lineLength; + while ( + Queue.sizeUnsafe(rawNotifications) > MAX_BUFFERED_RAW_NOTIFICATIONS || + (Queue.sizeUnsafe(rawNotifications) > 1 && + rawNotificationsJsonLength > MAX_BUFFERED_RAW_NOTIFICATION_JSON_LENGTH) + ) { + const dropped = Queue.takeUnsafe(rawNotifications); + if (dropped === undefined || Exit.isFailure(dropped)) break; + rawNotificationsJsonLength -= dropped.value.lineLength; + } + }); + + const dispatchNotification = (notification: AcpIncomingNotification, lineLength: number) => + bufferRawNotification(notification, lineLength).pipe( Effect.andThen( // A failing or dying handler must not stop the reader, or every later // message on the connection goes unanswered. @@ -451,6 +498,7 @@ export const makeAcpPatchedProtocol = Effect.fn("makeAcpPatchedProtocol")(functi return [ Effect.gen(function* () { yield* Queue.offer(disconnects, 0); + yield* Queue.end(rawNotifications); const error = yield* classify(); yield* Deferred.fail(terminationFailure, error); yield* failAllExtPending(error); @@ -559,18 +607,21 @@ export const makeAcpPatchedProtocol = Effect.fn("makeAcpPatchedProtocol")(functi ); }; - const handleRequestEncoded = (message: RpcMessage.RequestEncoded) => { + const handleRequestEncoded = (message: RpcMessage.RequestEncoded, lineLength: number) => { if (message.id === "") { if (message.tag === CLIENT_METHODS.session_update) { return decodeSessionUpdate(message.payload).pipe( Effect.matchEffect({ onFailure: dropUndecodableNotification(CLIENT_METHODS.session_update), onSuccess: (params) => - dispatchNotification({ - _tag: "SessionUpdate", - method: CLIENT_METHODS.session_update, - params: options.transformSessionUpdate?.(params) ?? params, - }), + dispatchNotification( + { + _tag: "SessionUpdate", + method: CLIENT_METHODS.session_update, + params: options.transformSessionUpdate?.(params) ?? params, + }, + lineLength, + ), }), ); } @@ -583,15 +634,14 @@ export const makeAcpPatchedProtocol = Effect.fn("makeAcpPatchedProtocol")(functi Effect.matchEffect({ onFailure: dropUndecodableNotification(method), onSuccess: (params) => - dispatchNotification({ _tag: "ElicitationComplete", method, params }), + dispatchNotification({ _tag: "ElicitationComplete", method, params }, lineLength), }), ); } - return dispatchNotification({ - _tag: "ExtNotification", - method: message.tag, - params: message.payload, - }); + return dispatchNotification( + { _tag: "ExtNotification", method: message.tag, params: message.payload }, + lineLength, + ); } const observeIncoming = @@ -667,12 +717,13 @@ export const makeAcpPatchedProtocol = Effect.fn("makeAcpPatchedProtocol")(functi }), ); - const routeDecodedMessage = ( - message: RpcMessage.FromClientEncoded | RpcMessage.FromServerEncoded, - ): Effect.Effect => { + const routeDecodedMessage = ({ + message, + lineLength, + }: AcpDecodedWireMessage): Effect.Effect => { switch (message._tag) { case "Request": - return handleRequestEncoded(message); + return handleRequestEncoded(message, lineLength); case "Exit": return handleExitEncoded(message); case "Chunk": @@ -714,12 +765,7 @@ export const makeAcpPatchedProtocol = Effect.fn("makeAcpPatchedProtocol")(functi ).pipe( Effect.flatMap(() => Effect.try({ - try: () => - ( - parser.decode(data) as ReadonlyArray< - RpcMessage.FromClientEncoded | RpcMessage.FromServerEncoded - > - ).map(normalizeAcpJsonRpcError), + try: () => parser.decode(data), catch: (cause) => new AcpError.AcpProtocolParseError({ operation: "decode-wire-message", @@ -727,12 +773,14 @@ export const makeAcpPatchedProtocol = Effect.fn("makeAcpPatchedProtocol")(functi }), }), ), - Effect.tap((messages) => - logProtocol({ - direction: "incoming", - stage: "decoded", - payload: messages, - }), + Effect.tap((decoded) => + options.logIncoming + ? logProtocol({ + direction: "incoming", + stage: "decoded", + payload: decoded.map(({ message }) => message), + }) + : Effect.void, ), Effect.tapErrorTag("AcpProtocolParseError", (error) => logProtocol({ @@ -741,8 +789,8 @@ export const makeAcpPatchedProtocol = Effect.fn("makeAcpPatchedProtocol")(functi payload: decodeFailureDiagnostics(error), }), ), - Effect.flatMap((messages) => - Effect.forEach(messages, routeDecodedMessage, { + Effect.flatMap((decoded) => + Effect.forEach(decoded, routeDecodedMessage, { discard: true, }), ), @@ -1009,7 +1057,9 @@ export const makeAcpPatchedProtocol = Effect.fn("makeAcpPatchedProtocol")(functi clientProtocol, serverProtocol, get incoming() { - return Stream.fromQueue(notificationQueue); + return Stream.fromQueue(rawNotifications).pipe( + Stream.map(({ notification }) => notification), + ); }, request: sendRequest, notify: sendNotification,