From 418752a6aacc95cadd0dc4b1958175e90a609413 Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Thu, 24 Sep 2026 20:27:30 +0200 Subject: [PATCH] fix: bound per-connection memory in AcpServer Session IDs a client sends no longer create lasting per-connection state, and agent output a client has not read is held to a fixed budget by pausing the agent rather than queuing without limit. - Create session streams on demand, drop them once empty and unreceived, and close a connection that buffers output for more than maxBufferedSessionStreams sessions. - Reject session IDs and request IDs longer than maxIdLength. - Flow-control agent output to maxBufferedBytes per connection: the router pauses, held-back client requests wait, and SSE bodies pull on demand. - Close connections whose client reads nothing for maxOutputStallMs, and bound WebSocket frames waiting to be handled. - End streams, held requests and sockets when a connection shuts down. --- src/connection.test.ts | 219 +++++++- src/connection.ts | 707 ++++++++++++++++++++++--- src/examples/http-server.ts | 7 +- src/node-adapter.test.ts | 118 +++++ src/node-adapter.ts | 42 +- src/server-sse.test.ts | 76 ++- src/server-sse.ts | 130 +++-- src/server-websocket-upgrade.test.ts | 367 ++++++++++++- src/server.test.ts | 752 +++++++++++++++++++++++++++ src/server.ts | 100 +++- src/sse.ts | 7 +- src/test-support/until.ts | 18 + src/ws-server.ts | 195 ++++++- src/ws-utils.ts | 6 + 14 files changed, 2576 insertions(+), 168 deletions(-) create mode 100644 src/test-support/until.ts diff --git a/src/connection.test.ts b/src/connection.test.ts index 0cea6225..95b3d379 100644 --- a/src/connection.test.ts +++ b/src/connection.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it, vi } from "vitest"; import { + ConnectionClosedError, ConnectionRegistry, OutboundMailbox, type OutboundLease, @@ -9,7 +10,7 @@ import { messageIdKey } from "./protocol.js"; import { createTestAgentApp } from "./test-support/test-agent.js"; import type { InitializeResponse } from "./acp.js"; -import type { AnyMessage, AnyWireMessage } from "./jsonrpc.js"; +import type { AnyWireMessage } from "./jsonrpc.js"; import type { WireStream } from "./stream.js"; const initializeRequest = { @@ -131,7 +132,9 @@ describe("ConnectionRegistry", () => { ); try { const lease = connection.allOutbound.tryAcquire(); - expect(lease).toBeDefined(); + if (!lease) { + throw new Error("Expected outbound mailbox lease"); + } connection.startRouter(); @@ -157,10 +160,147 @@ describe("ConnectionRegistry", () => { writer.releaseLock(); } - expect(await lease?.receive()).toEqual({ + // The router can no longer deliver output, so the connection closes. + await expect(withTimeout(lease.receive())).rejects.toThrow( + "AcpServer transports do not support outbound JSON-RPC batch messages", + ); + await withTimeout(connection.closed); + expect(connection.isClosed).toBe(true); + } finally { + error.mockRestore(); + await registry.closeAll(); + } + }); + + it("makes agent sends wait while a receiver is behind", async () => { + let agentStream: WireStream | undefined; + const registry = new ConnectionRegistry({ maxBufferedBytes: 50 }); + const connection = registry.createConnection({ + connect(stream) { + agentStream = stream; + }, + }); + try { + const lease = connection.connectionStream.tryAcquire(); + if (!lease) { + throw new Error("Expected outbound mailbox lease"); + } + + connection.startRouter(); + + if (!agentStream) { + throw new Error("Expected agent stream"); + } + + // The first two messages are 39 characters of JSON each, which + // together fill the limit. + const writer = agentStream.writable.getWriter(); + await withTimeout(writer.write(messageOne)); + await withTimeout(writer.write(messageTwo)); + let isThirdWritten = false; + const third = writer.write(messageThree).then(() => { + isThirdWritten = true; + }); + await delay(20); + expect(isThirdWritten).toBe(false); + + expect(await readJsonLease(lease)).toEqual(messageOne); + await delay(20); + expect(isThirdWritten).toBe(false); + + expect(await readJsonLease(lease)).toEqual(messageTwo); + await withTimeout(third); + expect(await readJsonLease(lease)).toEqual(messageThree); + writer.releaseLock(); + } finally { + await registry.closeAll(); + } + }); + + it("delivers what the agent wrote before closing while a receiver is behind", async () => { + let agentStream: WireStream | undefined; + const agentClosed = createDeferred(); + const registry = new ConnectionRegistry({ maxBufferedBytes: 50 }); + const connection = registry.createConnection({ + connect(stream) { + agentStream = stream; + return { closed: agentClosed.promise }; + }, + }); + try { + const lease = connection.connectionStream.tryAcquire(); + if (!lease) { + throw new Error("Expected outbound mailbox lease"); + } + + connection.startRouter(); + + if (!agentStream) { + throw new Error("Expected agent stream"); + } + + // The receiver is behind, so the later writes are still waiting when + // the agent closes without closing its stream. + const writer = agentStream.writable.getWriter(); + const messages = [messageOne, messageTwo, messageThree, messageFour]; + const writes = messages.map((message) => writer.write(message)); + await delay(20); + agentClosed.resolve(); + await Promise.all(writes.map((write) => withTimeout(write))); + + for (const message of messages) { + expect(await readJsonLease(lease)).toEqual(message); + } + await expect(withTimeout(lease.receive())).resolves.toEqual({ done: true, value: undefined, }); + } finally { + await registry.closeAll(); + } + }); + + it("closes the connection and its agent when agent output cannot be serialized", async () => { + let agentStream: WireStream | undefined; + const error = vi + .spyOn(console, "error") + .mockImplementation(() => undefined); + const registry = new ConnectionRegistry(); + const connection = registry.createConnection({ + connect(stream) { + agentStream = stream; + }, + }); + try { + const lease = connection.connectionStream.tryAcquire(); + if (!lease) { + throw new Error("Expected outbound mailbox lease"); + } + + connection.startRouter(); + + if (!agentStream) { + throw new Error("Expected agent stream"); + } + + const writer = agentStream.writable.getWriter(); + try { + await writer.write({ + jsonrpc: "2.0", + method: "_vendor/acme/notification", + params: { value: 1n }, + } as unknown as AnyWireMessage); + } finally { + writer.releaseLock(); + } + + await expect(withTimeout(lease.receive())).rejects.toThrow(TypeError); + await withTimeout(connection.closed); + expect(registry.get(connection.connectionId)).toBeUndefined(); + + // The agent sees its input fail instead of waiting on a dead connection. + const reader = agentStream.readable.getReader(); + await expect(withTimeout(reader.read())).rejects.toThrow(TypeError); } finally { error.mockRestore(); await registry.closeAll(); @@ -215,9 +355,9 @@ describe("ConnectionRegistry", () => { writer.releaseLock(); } - expect(await readLease(sessionOutbound)).toEqual(batch[0]); - expect(await readLease(sessionOutbound)).toEqual(batch[2]); - expect(await readLease(connectionOutbound)).toEqual(batch[1]); + expect(await readJsonLease(sessionOutbound)).toEqual(batch[0]); + expect(await readJsonLease(sessionOutbound)).toEqual(batch[2]); + expect(await readJsonLease(connectionOutbound)).toEqual(batch[1]); expect(connection.clientResponseRoutes).toEqual( new Map([ ["number:10", { session: sessionId }], @@ -303,6 +443,26 @@ describe("ConnectionRegistry", () => { await registry.closeAll(); }); + it("opens no new session stream once its agent has closed", async () => { + const agentClosed = createDeferred(); + const registry = new ConnectionRegistry(); + const connection = registry.createConnection({ + connect() { + return { closed: agentClosed.promise }; + }, + }); + connection.startRouter(); + + agentClosed.resolve(); + await withTimeout(connection.closed); + + // Nothing could ever reach it, so it would stay open forever. + expect(() => connection.acquireSessionStream("session-1")).toThrow( + ConnectionClosedError, + ); + await registry.closeAll(); + }); + it("waits for active and pending connection shutdowns before closeAll resolves", async () => { const registry = new ConnectionRegistry(); const active = registry.createConnection(createTestAgentApp()); @@ -360,7 +520,7 @@ describe("ConnectionRegistry", () => { await writeInbound(connection.inboundTx, sessionNewRequest); - const connectionMessage = await readLease(connectionLease); + const connectionMessage = await readJsonLease(connectionLease); expect(connectionMessage).toMatchObject({ jsonrpc: "2.0", @@ -385,7 +545,7 @@ describe("ConnectionRegistry", () => { await writeInbound(connection.inboundTx, sessionNewRequest); - expect(await readLease(lease)).toMatchObject({ + expect(await readJsonLease(lease)).toMatchObject({ jsonrpc: "2.0", id: sessionNewRequest.id, result: { @@ -432,7 +592,7 @@ describe("ConnectionRegistry", () => { await writeInbound(connection.inboundTx, promptRequest); - expect(await readLease(sessionLease)).toMatchObject({ + expect(await readJsonLease(sessionLease)).toMatchObject({ jsonrpc: "2.0", method: "session/update", params: { @@ -445,7 +605,7 @@ describe("ConnectionRegistry", () => { }, }, }); - expect(await readLease(sessionLease)).toMatchObject({ + expect(await readJsonLease(sessionLease)).toMatchObject({ jsonrpc: "2.0", id: promptRequest.id, result: { @@ -457,6 +617,28 @@ describe("ConnectionRegistry", () => { await registry.closeAll(); }); + + it("rejects inbound writes still waiting on the agent when it shuts down", async () => { + const registry = new ConnectionRegistry(); + // The agent never reads, so the first write waits and the second queues. + const connection = registry.createConnection({ connect() {} }); + const writes = [ + connection.writeInbound(initializeRequest), + connection.writeInbound({ ...initializeRequest, id: 2 }), + ]; + await flushMicrotasks(); + expect(connection.inboundTx.locked).toBe(true); + + await connection.shutdown(); + + for (const write of writes) { + await expect(withTimeout(write)).rejects.toThrow( + "ACP connection is closed", + ); + } + + await registry.closeAll(); + }); }); describe("OutboundMailbox", () => { @@ -558,7 +740,7 @@ async function writeInbound( } } -async function readLease( +async function readLease( lease: OutboundLease | undefined, ): Promise { if (!lease) { @@ -573,9 +755,16 @@ async function readLease( return result.value; } -async function readLeaseOrUndefined( - lease: OutboundLease | undefined, -): Promise { +/** Reads a message from an HTTP stream, which queues messages as JSON text. */ +async function readJsonLease( + lease: OutboundLease | undefined, +): Promise { + return JSON.parse(await readLease(lease)); +} + +async function readLeaseOrUndefined( + lease: OutboundLease | undefined, +): Promise { if (!lease) { throw new Error("Expected outbound mailbox lease"); } @@ -603,7 +792,7 @@ function createDeferred(): { async function withTimeout( promise: Promise, - timeoutMs = 100, + timeoutMs = 1_000, ): Promise { let timer: ReturnType | undefined; diff --git a/src/connection.ts b/src/connection.ts index 708dc4a7..b5c7af09 100644 --- a/src/connection.ts +++ b/src/connection.ts @@ -8,6 +8,111 @@ import { import type { AnyMessage, AnyResponse, AnyWireMessage } from "./jsonrpc.js"; import type { WireStream } from "./stream.js"; +/** Default for {@link ConnectionLimits.maxBufferedBytes}: 64 MiB. */ +export const DEFAULT_MAX_BUFFERED_BYTES = 64 * 1024 * 1024; +/** Default for {@link ConnectionLimits.maxOutputStallMs}: 60 seconds. */ +export const DEFAULT_MAX_OUTPUT_STALL_MS = 60_000; +/** Default for {@link ConnectionLimits.maxBufferedSessionStreams}: 1024. */ +export const DEFAULT_MAX_BUFFERED_SESSION_STREAMS = 1024; +/** Default for {@link ConnectionLimits.maxIdLength}: 1024. */ +export const DEFAULT_MAX_ID_LENGTH = 1024; + +/** The longest a timer can wait; timers fire at once past it. */ +const MAX_TIMER_MS = 2 ** 31 - 1; + +/** Per-connection limits on the state an ACP server keeps for a client. */ +export interface ConnectionLimits { + /** + * Maximum agent output a connection holds for its client before the client + * reads it. When a connection reaches the limit, it stops taking output + * from the agent and holds back client requests until the client reads; + * see `maxOutputStallMs` for a client that does not. Only an agent that + * awaits its sends is slowed down, and output can pass the limit by a + * message or two. + * + * Streamable HTTP connections count the output queued on all their + * streams, including streams the client has not opened, by the length of + * its JSON text in UTF-16 code units. Text outside Latin-1 takes twice that + * in memory. Held-back requests stay open with their bodies, so limit how + * many requests a client may have open in the HTTP server or proxy. + * Requests the agent has accepted but not answered are not counted; bound + * those in the agent. + * + * WebSocket connections count the bytes the socket reports as not yet sent + * (`bufferedAmount`), and do not limit output if the socket does not report + * it. They also close when client messages waiting to be handled exceed the + * limit, each counted as its length plus 1 KiB, so keep the WebSocket + * server's maximum message size below it. + * Defaults to {@link DEFAULT_MAX_BUFFERED_BYTES}. + */ + readonly maxBufferedBytes?: number; + /** + * Maximum time, in milliseconds, a connection holding `maxBufferedBytes` + * waits for its client to read any of it. The connection then closes, + * which frees the agent, and client requests held back, from a client that + * stopped reading or never opens the stream its output is for; memory + * stays bounded meanwhile. On WebSocket connections, any bytes the socket + * sends count as reading. At most 2^31 - 1 (about 24.8 days), the longest + * a timer can wait. Defaults to {@link DEFAULT_MAX_OUTPUT_STALL_MS}. + */ + readonly maxOutputStallMs?: number; + /** + * Maximum sessions per Streamable HTTP connection holding messages that no + * client is receiving, such as responses for a session whose stream was + * never opened. The connection closes when agent output would buffer + * another session, or when a client leaves a session stream with messages + * still queued past the limit. Open session streams are not counted: each + * is backed by an open request, so the embedding HTTP server or proxy must + * limit how many a client may hold. WebSocket connections deliver + * everything over the socket and never buffer per session. Defaults to + * {@link DEFAULT_MAX_BUFFERED_SESSION_STREAMS}. + */ + readonly maxBufferedSessionStreams?: number; + /** + * Maximum length, in UTF-16 code units, of the session IDs and request IDs + * a connection keeps track of. Client messages carrying a longer session ID + * or request ID are rejected (HTTP 400, or WebSocket close code 1008), and + * a connection whose agent issues a longer session ID is closed. Defaults + * to {@link DEFAULT_MAX_ID_LENGTH}. + */ + readonly maxIdLength?: number; +} + +/** Raised when a client or agent exceeds a {@link ConnectionLimits} limit. */ +export class ConnectionLimitError extends Error { + constructor(message: string) { + super(message); + this.name = "ConnectionLimitError"; + } +} + +/** Error for a connection that holds more than a limit allows. */ +export function connectionLimitExceeded( + limit: "maxBufferedBytes" | "maxBufferedSessionStreams", + limits: Required, +): ConnectionLimitError { + return new ConnectionLimitError( + `Connection exceeds ${limit} (${limits[limit]})`, + ); +} + +/** Error for a connection whose client read none of its output in time. */ +export function connectionOutputStalled( + limits: Required, +): ConnectionLimitError { + return new ConnectionLimitError( + `Connection output stalled for maxOutputStallMs (${limits.maxOutputStallMs})`, + ); +} + +/** Raised when a message is sent to a connection that has shut down. */ +export class ConnectionClosedError extends Error { + constructor() { + super("ACP connection is closed"); + this.name = "ConnectionClosedError"; + } +} + export interface AgentConnectOptions { readonly deferConnectHandlers?: boolean; } @@ -26,18 +131,41 @@ export interface AgentConnector { export type ResponseRoute = "connection" | { readonly session: string }; -export interface OutboundLease { +export interface OutboundLease { + /** + * Aborted when the stream is shut down while this lease holds it, with + * the error as its reason or a `ConnectionClosedError` if there was none. + * A receiver waiting on `receive()` learns this from its result; one that + * is not, such as a response body whose client stopped reading, must stop + * on this signal. + */ + readonly stopped: AbortSignal; receive(): Promise>; release(): void; } -export class OutboundMailbox { - private readonly queue: Message[] = []; +/** Follows what a stream queues, for its connection's accounting. */ +export interface OutboundQueueObserver { + /** Called as a message is queued. */ + queued?(message: Message): void; + /** Called as a message leaves the queue, taken by a receiver or dropped. */ + dequeued(message: Message): void; +} + +export class OutboundMailbox { + /** Queued messages start at `head`; see `dequeue`. */ + private queue: (Message | undefined)[] = []; + private head = 0; private activeLease: MailboxLease | undefined; private isFinished = false; private isAborted = false; + private abortError: unknown; - constructor(private readonly enabled = true) {} + constructor( + private readonly enabled = true, + private onReceiverChange?: () => void, + private observer?: OutboundQueueObserver, + ) {} push(message: Message): void { if (!this.enabled || this.isFinished) { @@ -45,9 +173,18 @@ export class OutboundMailbox { } this.queue.push(message); + this.observer?.queued?.(message); this.activeLease?.wake(); } + get hasReceiver(): boolean { + return this.activeLease !== undefined; + } + + get hasQueuedMessages(): boolean { + return this.head < this.queue.length; + } + tryAcquire(): OutboundLease | undefined { if (!this.enabled || this.isAborted || this.activeLease) { return undefined; @@ -55,6 +192,7 @@ export class OutboundMailbox { const lease = new MailboxLease(this); this.activeLease = lease; + this.onReceiverChange?.(); return lease; } @@ -67,15 +205,23 @@ export class OutboundMailbox { this.activeLease?.wake(); } - abort(): void { + abort(error?: unknown): void { if (this.isAborted) { return; } this.isAborted = true; + this.abortError = error; this.isFinished = true; - this.queue.length = 0; - this.activeLease?.wake(); + while (this.hasQueuedMessages) { + this.dequeue(); + } + + // Whatever still holds this stream, such as a stalled response body, + // must not keep its connection reachable. + this.observer = undefined; + this.onReceiverChange = undefined; + this.activeLease?.stop(error ?? new ConnectionClosedError()); } /** @internal */ @@ -83,15 +229,19 @@ export class OutboundMailbox { lease: MailboxLease, ): Promise> { for (;;) { - if (this.isAborted || lease.released || this.activeLease !== lease) { + if (lease.released || this.activeLease !== lease) { + return { done: true, value: undefined }; + } + + if (this.isAborted) { + if (this.abortError !== undefined) { + throw this.abortError; + } return { done: true, value: undefined }; } - if (this.queue.length > 0) { - return { - done: false, - value: this.queue.shift() as Message, - }; + if (this.hasQueuedMessages) { + return { done: false, value: this.dequeue() }; } if (this.isFinished) { @@ -110,6 +260,30 @@ export class OutboundMailbox { this.activeLease = undefined; lease.markReleased(); + this.onReceiverChange?.(); + } + + /** + * Takes the oldest queued message. `Array.prototype.shift` copies the rest + * of a large array, which would make draining a full queue quadratic, so + * this advances `head` instead and drops the taken prefix once the queue + * empties or is mostly taken. + */ + private dequeue(): Message { + const message = this.queue[this.head] as Message; + this.queue[this.head] = undefined; + this.head += 1; + + if (this.head === this.queue.length) { + this.queue = []; + this.head = 0; + } else if (this.head >= 1024 && this.head * 2 >= this.queue.length) { + this.queue = this.queue.slice(this.head); + this.head = 0; + } + + this.observer?.dequeued(message); + return message; } } @@ -119,16 +293,37 @@ export class ConnectionState { readonly connectionId: string; readonly inboundTx: WritableStream; readonly outboundRx: ReadableStream; - readonly connectionStream: OutboundMailbox; + /** Connection-level HTTP stream, queuing each message as JSON text. */ + readonly connectionStream: OutboundMailbox; + /** Every outbound message, for WebSocket connections. */ readonly allOutbound: OutboundMailbox; - readonly sessionStreams = new Map(); + /** Session-level HTTP streams, queuing each message as JSON text. */ + readonly sessionStreams = new Map>(); readonly pendingRoutes = new Map(); readonly clientResponseRoutes = new Map(); readonly closed: Promise; private readonly agentConnection: AgentConnectionLifecycle | unknown; + /** Length of the JSON text queued on HTTP streams; see `maxBufferedBytes`. */ + private bufferedBytes = 0; + private readonly streamObserver: OutboundQueueObserver = { + queued: (json) => { + this.bufferedBytes += json.length; + }, + dequeued: (json) => { + this.bufferedBytes -= json.length; + this.notifyOutboundProgress(); + }, + }; + /** Session streams without a receiver; see `maxBufferedSessionStreams`. */ + private bufferedSessionStreams = 0; + /** Called when a client takes output, or the connection stops routing. */ + private readonly progressWaiters = new Set<() => void>(); private supportsBatches = false; private hasStartedRouter = false; + private hasAgentOutputEnded = false; + private hasFinishedRouting = false; + private hasStartedShutdown = false; private inboundWriteChain: Promise = Promise.resolve(); private initialReader: ReadableStreamDefaultReader | undefined; @@ -137,33 +332,59 @@ export class ConnectionState { private routerPromise: Promise | undefined; private shutdownPromise: Promise | undefined; private hasResolvedClosed = false; - private finishAgentOutbound: () => void = () => {}; + private abortAgentInbound: (error: unknown) => void = () => {}; + private endAgentOutbound: () => void = () => {}; private resolveClosed: () => void = () => {}; constructor( agent: AgentConnector, - private readonly transport: ConnectionTransport = "http", + private readonly transport: ConnectionTransport, + private readonly limits: Required, ) { this.connectionId = globalThis.crypto.randomUUID(); - this.connectionStream = new OutboundMailbox(transport === "http"); + this.connectionStream = new OutboundMailbox( + transport === "http", + undefined, + this.streamObserver, + ); + // WebSocket output waits here only until the socket pump takes it; the + // socket itself holds what the client has not read. this.allOutbound = new OutboundMailbox( transport === "websocket", + undefined, + { + dequeued: () => { + this.notifyOutboundProgress(); + }, + }, ); this.closed = new Promise((resolve) => { this.resolveClosed = resolve; }); - const inbound = new TransformStream(); - const outbound = createBufferedOutboundChannel((message) => { - if (!this.supportsBatches && Array.isArray(message)) { - throw new TypeError( - "AcpServer transports do not support outbound JSON-RPC batch messages", - ); - } + const inbound = new TransformStream({ + start: (controller) => { + this.abortAgentInbound = (error) => controller.error(error); + }, }); + const outbound = createOutboundChannel( + (message) => { + if (!this.supportsBatches && Array.isArray(message)) { + throw new TypeError( + "AcpServer transports do not support outbound JSON-RPC batch messages", + ); + } + }, + () => { + // Nothing more will come from the agent, so let the router finish + // instead of waiting for clients to make room. + this.hasAgentOutputEnded = true; + this.notifyOutboundProgress(); + }, + ); this.inboundTx = inbound.writable; this.outboundRx = outbound.readable; - this.finishAgentOutbound = outbound.finish; + this.endAgentOutbound = outbound.end; const stream: WireStream = { readable: inbound.readable, @@ -248,41 +469,123 @@ export class ConnectionState { return this.supportsBatches; } - ensureSession(sessionId: string): OutboundMailbox { + /** Whether shutdown has begun; a closed connection accepts no more work. */ + get isClosed(): boolean { + return this.hasStartedShutdown; + } + + /** + * Whether the connection takes no more agent output for now. HTTP + * connections hold up to `maxBufferedBytes` in their streams. WebSocket + * connections hand each message to the socket pump, which waits while the + * socket holds `maxBufferedBytes`, so they wait until the pump takes it. + */ + get hasOutboundBacklog(): boolean { + return this.transport === "http" + ? this.bufferedBytes >= this.limits.maxBufferedBytes + : this.allOutbound.hasQueuedMessages; + } + + /** Throws `ConnectionLimitError` for a session ID over `maxIdLength`. */ + validateSessionId(sessionId: string): void { + this.validateIdLength("Session ID", sessionId); + } + + /** Throws `ConnectionLimitError` for a request ID over `maxIdLength`. */ + validateRequestId(id: unknown): void { + if (typeof id === "string") { + this.validateIdLength("Request ID", id); + } + } + + /** + * Resolves once the connection can take more agent output, or has stopped, + * and rejects if `signal` aborts first. Client requests wait on this, so a + * client that is not reading cannot make the agent produce more. + */ + async waitForOutboundCapacity(signal?: AbortSignal): Promise { + await this.waitForRoom(() => this.hasFinishedRouting, signal); + } + + /** + * Returns the stream that delivers messages for a session, creating one + * that buffers them until a receiver attaches. Throws + * `ConnectionLimitError` for an over-long ID or when no further session + * may buffer. + */ + ensureSession(sessionId: string): OutboundMailbox { + this.validateSessionId(sessionId); + + // WebSocket connections deliver every message over the socket. + if (this.transport === "websocket") { + return this.connectionStream; + } + const existing = this.sessionStreams.get(sessionId); if (existing) { return existing; } - const stream = new OutboundMailbox(this.transport === "http"); - this.sessionStreams.set(sessionId, stream); + // A new stream has no receiver yet, so its messages are buffered. + if (this.bufferedSessionStreams >= this.limits.maxBufferedSessionStreams) { + throw connectionLimitExceeded("maxBufferedSessionStreams", this.limits); + } - return stream; + return this.createSessionStream(sessionId); } - async shutdown(): Promise { + /** + * Attaches a receiver to a session's stream, or returns `undefined` if the + * stream cannot take one. The open request backing the receiver bounds it, + * so this is not limited by `maxBufferedSessionStreams`. Throws + * `ConnectionLimitError` for an over-long ID, and `ConnectionClosedError` + * for a new stream once the connection can no longer route output to it. + */ + acquireSessionStream(sessionId: string): OutboundLease | undefined { + this.validateSessionId(sessionId); + + if (this.transport === "websocket") { + return undefined; + } + + const existing = this.sessionStreams.get(sessionId); + if (existing) { + return existing.tryAcquire(); + } + + if (this.isClosed || this.hasFinishedRouting) { + throw new ConnectionClosedError(); + } + + return this.createSessionStream(sessionId).tryAcquire(); + } + + async shutdown(error?: unknown): Promise { if (!this.shutdownPromise) { - this.shutdownPromise = this.runShutdown(); + this.shutdownPromise = this.runShutdown(error); } return this.shutdownPromise; } - private async runShutdown(): Promise { + private async runShutdown(error?: unknown): Promise { try { - this.connectionStream.abort(); - this.allOutbound.abort(); + this.hasStartedShutdown = true; + this.notifyOutboundProgress(); + this.connectionStream.abort(error); + this.allOutbound.abort(error); for (const stream of this.sessionStreams.values()) { - stream.abort(); + stream.abort(error); } this.sessionStreams.clear(); + this.bufferedSessionStreams = 0; this.pendingRoutes.clear(); this.clientResponseRoutes.clear(); await Promise.allSettled([ - this.inboundTx.close(), + this.closeInbound(error), this.cancelOutboundReader(), ]); } finally { @@ -306,7 +609,7 @@ export class ConnectionState { return; } - this.finishAgentOutbound(); + this.endAgentOutbound(); await this.routerPromise; }); } @@ -320,11 +623,33 @@ export class ConnectionState { return this.outboundRx.cancel(); } + private closeInbound(error?: unknown): Promise { + if (error !== undefined || this.inboundTx.locked) { + this.abortAgentInbound(error ?? new ConnectionClosedError()); + return Promise.resolve(); + } + + return this.inboundTx.close(); + } + private async writeInboundMessage(message: AnyWireMessage): Promise { + // Once routing has finished, the agent can no longer answer, such as a + // request that waited for room while the agent exited. + if (this.isClosed || this.hasFinishedRouting) { + throw new ConnectionClosedError(); + } + const writer = this.inboundTx.getWriter(); try { await writer.write(message); + } catch (error) { + // Shutdown aborts a write that was still waiting on the agent. + if (this.isClosed) { + throw new ConnectionClosedError(); + } + + throw error; } finally { writer.releaseLock(); } @@ -338,15 +663,32 @@ export class ConnectionState { while (true) { const result = await reader.read(); - if (result.done) { + if (result.done || this.isClosed) { return; } this.routeOutbound(result.value); + // Take no more output, which makes the agent's sends wait, until + // clients read what is queued. Once the agent's output has ended, + // nothing waits on the router, so it routes what is left. + await this.waitForRoom(() => this.hasAgentOutputEnded); } } catch (error) { - console.error("ACP connection router stopped unexpectedly:", error); + // Output can no longer be routed, so close the connection rather than + // leave the agent and its client waiting on it. + if (error instanceof ConnectionLimitError) { + this.closeForLimit(error); + } else { + console.error( + `ACP connection ${this.connectionId} router stopped unexpectedly:`, + error, + ); + void this.shutdown(error); + } } finally { + this.hasFinishedRouting = true; + this.notifyOutboundProgress(); + if (this.outboundReader === reader) { this.outboundReader = undefined; } @@ -372,17 +714,84 @@ export class ConnectionState { this.resolveClosed(); } - private routeOutbound(message: AnyWireMessage): void { - this.allOutbound.push(message); + /** + * Waits while the connection holds `maxBufferedBytes` of output, until + * `stopWaiting()` holds or the connection closes, and rejects if `signal` + * aborts first. An HTTP client that takes none of the output for + * `maxOutputStallMs`, having stopped reading or never opened the stream it + * is for, is not coming back for it, so the connection closes rather than + * keep the agent or the client's own requests waiting. A WebSocket session + * times out its socket itself, where it can see each byte the client reads. + */ + private async waitForRoom( + stopWaiting: () => boolean, + signal?: AbortSignal, + ): Promise { + while (!this.isClosed && !stopWaiting() && this.hasOutboundBacklog) { + if (!(await this.waitForProgress(signal))) { + this.closeForLimit(connectionOutputStalled(this.limits)); + return; + } + } + } + + /** + * Resolves `true` the next time a client takes output or the connection + * stops routing, or on HTTP, `false` if `maxOutputStallMs` passes first. + * Rejects if `signal` aborts first, leaving nothing behind. + */ + private waitForProgress(signal?: AbortSignal): Promise { + const timeoutMs = + this.transport === "http" ? this.limits.maxOutputStallMs : undefined; + + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(signal.reason); + return; + } + + let timer: ReturnType | undefined; + const settle = (): void => { + this.progressWaiters.delete(onProgress); + signal?.removeEventListener("abort", onAbort); + clearTimeout(timer); + }; + const onProgress = (): void => { + settle(); + resolve(true); + }; + const onAbort = (): void => { + settle(); + reject(signal?.reason); + }; + + this.progressWaiters.add(onProgress); + signal?.addEventListener("abort", onAbort, { once: true }); + if (timeoutMs !== undefined) { + timer = setTimeout(() => { + settle(); + resolve(false); + }, timeoutMs); + } + }); + } + private notifyOutboundProgress(): void { + for (const onProgress of [...this.progressWaiters]) { + onProgress(); + } + } + + private routeOutbound(message: AnyWireMessage): void { if (Array.isArray(message)) { for (const item of message) { this.routeOutboundMessage(item); } - return; + } else { + this.routeOutboundMessage(message as AnyMessage); } - this.routeOutboundMessage(message as AnyMessage); + this.allOutbound.push(message); } private routeOutboundMessage(message: AnyMessage): void { @@ -400,7 +809,8 @@ export class ConnectionState { const sessionId = sessionIdFromResponseResult(message); if (sessionId) { - this.ensureSession(sessionId); + // Never hand out a session ID this connection would refuse to route. + this.validateSessionId(sessionId); } if (key) { @@ -410,16 +820,80 @@ export class ConnectionState { this.pushToRoute(route ?? "connection", message); } + private createSessionStream(sessionId: string): OutboundMailbox { + const stream: OutboundMailbox = new OutboundMailbox( + true, + () => { + this.onSessionReceiverChange(sessionId, stream); + }, + this.streamObserver, + ); + this.sessionStreams.set(sessionId, stream); + this.bufferedSessionStreams += 1; + + return stream; + } + + /** + * Keeps `bufferedSessionStreams` in step as receivers attach and leave. A + * stream a receiver leaves is kept only while messages still wait on it. + */ + private onSessionReceiverChange( + sessionId: string, + stream: OutboundMailbox, + ): void { + // Shutdown discards every stream and resets the count itself. + if (this.isClosed) { + return; + } + + if (stream.hasReceiver) { + this.bufferedSessionStreams -= 1; + return; + } + + if (!stream.hasQueuedMessages) { + this.sessionStreams.delete(sessionId); + stream.abort(); + return; + } + + this.bufferedSessionStreams += 1; + // Once routing has finished, nothing more can be buffered. + if ( + !this.hasFinishedRouting && + this.bufferedSessionStreams > this.limits.maxBufferedSessionStreams + ) { + this.closeForLimit( + connectionLimitExceeded("maxBufferedSessionStreams", this.limits), + ); + } + } + + private closeForLimit(error: ConnectionLimitError): void { + console.warn(`Closing ACP connection ${this.connectionId}:`, error.message); + void this.shutdown(error); + } + + private validateIdLength(label: string, id: string): void { + if (id.length > this.limits.maxIdLength) { + throw new ConnectionLimitError( + `${label} exceeds maxIdLength (${this.limits.maxIdLength})`, + ); + } + } + private routeOutboundRequestOrNotification(message: AnyMessage): void { const sessionId = sessionIdFromMessageParams(message); if (sessionId) { + const stream = this.ensureSession(sessionId); this.trackClientResponseRoute(message, { session: sessionId }); - this.ensureSession(sessionId).push(message); + this.deliver(stream, message); return; } this.trackClientResponseRoute(message, "connection"); - this.connectionStream.push(message); + this.deliver(this.connectionStream, message); } private trackClientResponseRoute( @@ -438,23 +912,39 @@ export class ConnectionState { private pushToRoute(route: ResponseRoute, message: AnyMessage): void { if (route === "connection") { - this.connectionStream.push(message); + this.deliver(this.connectionStream, message); return; } - this.ensureSession(route.session).push(message); + this.deliver(this.ensureSession(route.session), message); + } + + /** + * Queues a message on an HTTP stream as JSON text, which is all a stream + * delivers and all `maxBufferedBytes` needs to count. WebSocket + * connections deliver through `allOutbound` instead. + */ + private deliver(stream: OutboundMailbox, message: AnyMessage): void { + if (this.transport === "http") { + stream.push(JSON.stringify(message)); + } } } export class ConnectionRegistry { + readonly limits: Required; private readonly connections = new Map(); private readonly pendingConnections = new Map(); + constructor(options: ConnectionLimits = {}) { + this.limits = resolveConnectionLimits(options); + } + createConnection( agent: AgentConnector, transport: ConnectionTransport = "http", ): ConnectionState { - const connection = new ConnectionState(agent, transport); + const connection = new ConnectionState(agent, transport, this.limits); this.connections.set(connection.connectionId, connection); this.trackConnectionClose(connection); return connection; @@ -464,7 +954,7 @@ export class ConnectionRegistry { agent: AgentConnector, transport: ConnectionTransport = "websocket", ): ConnectionState { - const connection = new ConnectionState(agent, transport); + const connection = new ConnectionState(agent, transport, this.limits); this.pendingConnections.set(connection.connectionId, connection); this.trackConnectionClose(connection); return connection; @@ -476,11 +966,14 @@ export class ConnectionRegistry { } get(connectionId: string): ConnectionState | undefined { - return this.connections.get(connectionId); + const connection = this.connections.get(connectionId); + return connection?.isClosed ? undefined : connection; } remove(connectionId: string): ConnectionState | undefined { - const connection = this.get(connectionId); + // Unlike `get`, include a connection that is already shutting down, so + // DELETE still succeeds until its shutdown completes. + const connection = this.connections.get(connectionId); if (!connection) { return undefined; @@ -531,17 +1024,44 @@ export class ConnectionRegistry { } } -class MailboxLease< - Message extends AnyWireMessage, -> implements OutboundLease { +function resolveConnectionLimits( + options: ConnectionLimits, +): Required { + const limits = { + maxBufferedBytes: options.maxBufferedBytes ?? DEFAULT_MAX_BUFFERED_BYTES, + maxOutputStallMs: options.maxOutputStallMs ?? DEFAULT_MAX_OUTPUT_STALL_MS, + maxBufferedSessionStreams: + options.maxBufferedSessionStreams ?? DEFAULT_MAX_BUFFERED_SESSION_STREAMS, + maxIdLength: options.maxIdLength ?? DEFAULT_MAX_ID_LENGTH, + }; + + for (const [name, value] of Object.entries(limits)) { + if (!Number.isSafeInteger(value) || value <= 0) { + throw new RangeError(`${name} must be a positive safe integer`); + } + } + + if (limits.maxOutputStallMs > MAX_TIMER_MS) { + throw new RangeError(`maxOutputStallMs must be at most ${MAX_TIMER_MS}`); + } + + return limits; +} + +class MailboxLease implements OutboundLease { released = false; + private readonly stoppedController = new AbortController(); private receiving = false; private wakePromise: Promise | undefined; private resolveWake: (() => void) | undefined; constructor(private readonly mailbox: OutboundMailbox) {} + get stopped(): AbortSignal { + return this.stoppedController.signal; + } + async receive(): Promise> { if (this.receiving) { throw new Error( @@ -581,18 +1101,37 @@ class MailboxLease< this.released = true; this.wake(); } + + stop(reason: unknown): void { + this.stoppedController.abort(reason); + this.wake(); + } } -function createBufferedOutboundChannel( +/** + * Carries the agent's output to the router. A write resolves at once if the + * router is waiting for a message, and otherwise once the router takes it + * and asks for the next one, so an agent that awaits its sends goes no + * faster than the router, which pauses while clients are behind. + */ +function createOutboundChannel( validate: (message: AnyWireMessage) => void, + onEnd: () => void, ): { readonly readable: ReadableStream; readonly writable: WritableStream; - readonly finish: () => void; + readonly end: () => void; } { let controller: ReadableStreamDefaultController | undefined; let isFinished = false; + let isEnding = false; + let resumeWriter: (() => void) | undefined; + const resume = (): void => { + const resolve = resumeWriter; + resumeWriter = undefined; + resolve?.(); + }; const finish = (): void => { if (isFinished) { return; @@ -604,6 +1143,8 @@ function createBufferedOutboundChannel( } catch { // The router may already have cancelled the readable side. } + resume(); + onEnd(); }; const fail = (error: unknown): void => { if (isFinished) { @@ -616,19 +1157,44 @@ function createBufferedOutboundChannel( } catch { // The router may already have cancelled the readable side. } + resume(); + onEnd(); + }; + /** + * Ends the channel for an agent that closed without closing its stream, as + * a connector whose `closed` resolves on its own may. Writes it made before + * closing can still be queued in the writable stream, which hands them + * over one per microtask, so they stop waiting for the router and the + * channel finishes after them. + */ + const end = (): void => { + if (isFinished || isEnding) { + return; + } + + isEnding = true; + resume(); + setTimeout(finish, 0); }; return { - readable: new ReadableStream({ - start(readableController) { - controller = readableController; + readable: new ReadableStream( + { + start(readableController) { + controller = readableController; + }, + // With no high-water mark, this runs only once the router is waiting + // on an empty queue, which means it took the last message. + pull: resume, + cancel() { + isFinished = true; + resume(); + }, }, - cancel() { - isFinished = true; - }, - }), + { highWaterMark: 0 }, + ), writable: new WritableStream({ - write(message) { + async write(message) { if (isFinished) { throw new Error("ACP outbound channel is closed"); } @@ -640,11 +1206,18 @@ function createBufferedOutboundChannel( fail(error); throw error; } + + // The message went straight to a waiting router unless it is queued. + if (!isFinished && !isEnding && (controller?.desiredSize ?? 0) < 0) { + await new Promise((resolve) => { + resumeWriter = resolve; + }); + } }, close: finish, abort: fail, }), - finish, + end, }; } diff --git a/src/examples/http-server.ts b/src/examples/http-server.ts index 201a0183..6ad56ab4 100644 --- a/src/examples/http-server.ts +++ b/src/examples/http-server.ts @@ -121,7 +121,12 @@ const agent = acp const acpServer = new AcpServer({ agent }); const acpHttpHandler = createNodeHttpHandler(acpServer); -const webSocketServer = new WebSocketServer({ noServer: true }); +// Match the HTTP handler's 16 MiB request limit, which keeps messages well +// below AcpServer's maxBufferedBytes. +const webSocketServer = new WebSocketServer({ + noServer: true, + maxPayload: 16 * 1024 * 1024, +}); // Use the ACP upgrade helper so the 101 response includes Acp-Connection-Id. const acpWebSocketUpgradeHandler = createNodeWebSocketUpgradeHandler( acpServer, diff --git a/src/node-adapter.test.ts b/src/node-adapter.test.ts index b4429232..f89ad546 100644 --- a/src/node-adapter.test.ts +++ b/src/node-adapter.test.ts @@ -450,8 +450,106 @@ describe("createNodeHttpHandler", () => { expect(response.ended).toBe(false); }); + + it("destroys a started response when its body fails", async () => { + const acpServer = new AcpServer({ + createAgent: () => createTestAgentApp(), + }); + const response = new CapturingServerResponse(); + const body = failingBody(); + acpServer.handleRequest = () => + Promise.resolve(eventStreamResponse(body.stream)); + + createNodeHttpHandler(acpServer)( + fakeRequest(), + response as unknown as ServerResponse, + ); + + await response.wroteChunk; + body.fail(new Error("Connection exceeds maxBufferedBytes (1024)")); + await response.destroyCalled; + + expect(response.writableEnded).toBe(false); + expect(response.chunks.join("")).toBe("data: one\n\n"); + }); + + it("destroys a backpressured response when its body fails", async () => { + const acpServer = new AcpServer({ + createAgent: () => createTestAgentApp(), + }); + const response = new BackpressuredServerResponse(); + const body = failingBody(); + acpServer.handleRequest = () => + Promise.resolve(eventStreamResponse(body.stream)); + + createNodeHttpHandler(acpServer)( + fakeRequest(), + response as unknown as ServerResponse, + ); + + // The write is now waiting for the socket to drain, with no read pending. + await response.wroteChunk; + body.fail(new Error("Connection exceeds maxBufferedBytes (1024)")); + await response.destroyCalled; + + expect(response.ended).toBe(false); + expect(response.listenerCount("drain")).toBe(0); + }); + + it("cancels the body when its response fails", async () => { + const acpServer = new AcpServer({ + createAgent: () => createTestAgentApp(), + }); + const response = new BackpressuredServerResponse(); + const cancelled = createDeferred(); + const body = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("data: one\n\n")); + }, + cancel(reason) { + cancelled.resolve(reason); + }, + }); + acpServer.handleRequest = () => Promise.resolve(eventStreamResponse(body)); + + createNodeHttpHandler(acpServer)( + fakeRequest(), + response as unknown as ServerResponse, + ); + + // The socket fails while the write waits for it to drain. + await response.wroteChunk; + response.emit("error", new Error("socket hang up")); + + await expect(cancelled.promise).resolves.toBeInstanceOf(Error); + await response.destroyCalled; + }); }); +function failingBody(): { + readonly stream: ReadableStream; + readonly fail: (error: Error) => void; +} { + let fail: (error: Error) => void = () => {}; + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("data: one\n\n")); + fail = (error) => controller.error(error); + }, + }); + + return { stream, fail: (error) => fail(error) }; +} + +function eventStreamResponse(body: ReadableStream): Response { + return new Response(body, { + status: 200, + headers: { + "Content-Type": "text/event-stream", + }, + }); +} + describe("createNodeWebSocketUpgradeHandler", () => { it("destroys the upgrade socket when WebSocket preparation throws", async () => { const error = new Error("factory failed"); @@ -525,6 +623,10 @@ class CapturingServerResponse extends EventEmitter { private readonly finishDeferred = createDeferred(); readonly finished = this.finishDeferred.promise; + private readonly writeDeferred = createDeferred(); + readonly wroteChunk = this.writeDeferred.promise; + private readonly destroyDeferred = createDeferred(); + readonly destroyCalled = this.destroyDeferred.promise; setHeader(): void {} @@ -536,9 +638,17 @@ class CapturingServerResponse extends EventEmitter { this.chunks.push( typeof chunk === "string" ? chunk : Buffer.from(chunk).toString("utf8"), ); + this.writeDeferred.resolve(); return true; } + destroy(): this { + this.destroyed = true; + this.emit("close"); + this.destroyDeferred.resolve(); + return this; + } + end(chunk?: Uint8Array | string): void { if (chunk !== undefined) { this.write(chunk); @@ -558,6 +668,8 @@ class BackpressuredServerResponse extends EventEmitter { private readonly writeDeferred = createDeferred(); readonly wroteChunk = this.writeDeferred.promise; + private readonly destroyDeferred = createDeferred(); + readonly destroyCalled = this.destroyDeferred.promise; setHeader(): void {} @@ -579,6 +691,12 @@ class BackpressuredServerResponse extends EventEmitter { this.destroyed = true; this.emit("close"); } + + destroy(): this { + this.close(); + this.destroyDeferred.resolve(); + return this; + } } function fakeRequest( diff --git a/src/node-adapter.ts b/src/node-adapter.ts index f5a9c81c..8daed493 100644 --- a/src/node-adapter.ts +++ b/src/node-adapter.ts @@ -111,6 +111,14 @@ async function handleNodeRequest( return; } + if (res.headersSent) { + // The status is already committed, so abort the response and its + // socket; appending an error message would make a failed stream look + // like a complete one. + res.destroy(); + return; + } + writePlainTextErrorResponse( res, 500, @@ -476,7 +484,14 @@ async function writeNodeResponse( } const reader = responseBody.getReader(); + // The body can fail while a write waits for the socket to drain, when no + // read is pending to observe it, so surface the failure to writes as well. + const bodyFailure = new AbortController(); + reader.closed.catch((error: unknown) => { + bodyFailure.abort(error); + }); let cancelReader: Promise | undefined; + let isBodyDone = false; const onClose = (): void => { cancelReader = reader @@ -491,6 +506,7 @@ async function writeNodeResponse( const result = await reader.read(); if (result.done) { + isBodyDone = true; res.off("close", onClose); if (!isNodeResponseClosed(res)) { @@ -500,7 +516,7 @@ async function writeNodeResponse( return; } - await writeChunk(res, result.value); + await writeChunk(res, result.value, bodyFailure.signal); } } catch (error) { if (error instanceof NodeResponseClosedError) { @@ -510,6 +526,11 @@ async function writeNodeResponse( throw error; } finally { res.off("close", onClose); + // The response can also fail without closing first, such as on a socket + // error, so stop the body however the response ended. + if (!isBodyDone && !cancelReader) { + onClose(); + } await cancelReader; reader.releaseLock(); } @@ -545,7 +566,11 @@ function getSetCookieHeaders(headers: Headers): string[] | undefined { : undefined; } -function writeChunk(res: ServerResponse, chunk: Uint8Array): Promise { +function writeChunk( + res: ServerResponse, + chunk: Uint8Array, + bodyFailure: AbortSignal, +): Promise { return new Promise((resolve, reject) => { let isSettled = false; @@ -558,6 +583,7 @@ function writeChunk(res: ServerResponse, chunk: Uint8Array): Promise { res.off("close", onClose); res.off("drain", onDrain); res.off("error", onError); + bodyFailure.removeEventListener("abort", onBodyFailure); callback(); }; @@ -577,13 +603,25 @@ function writeChunk(res: ServerResponse, chunk: Uint8Array): Promise { }); }; + const onBodyFailure = (): void => { + settle(() => { + reject(bodyFailure.reason); + }); + }; + if (isNodeResponseClosed(res)) { reject(new NodeResponseClosedError()); return; } + if (bodyFailure.aborted) { + reject(bodyFailure.reason); + return; + } + res.once("close", onClose); res.once("error", onError); + bodyFailure.addEventListener("abort", onBodyFailure, { once: true }); if (res.write(chunk)) { settle(resolve); diff --git a/src/server-sse.test.ts b/src/server-sse.test.ts index 4bf5d280..917984bf 100644 --- a/src/server-sse.test.ts +++ b/src/server-sse.test.ts @@ -1,8 +1,8 @@ -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { OutboundMailbox } from "./connection.js"; -import { createSseBodySource } from "./server-sse.js"; -import { serializeSseEvent } from "./sse.js"; +import { createSseBody, createSseBodySource } from "./server-sse.js"; +import { serializeSseEvent, serializeSseKeepAlive } from "./sse.js"; import type { AnyMessage } from "./jsonrpc.js"; @@ -16,7 +16,7 @@ const message = { describe("createSseBodySource", () => { it("enqueues a subscription message after it has been read even if demand changed", async () => { - const mailbox = new OutboundMailbox(); + const mailbox = new OutboundMailbox(); const lease = mailbox.tryAcquire(); if (!lease) { throw new Error("Expected outbound mailbox lease"); @@ -41,17 +41,79 @@ describe("createSseBodySource", () => { }, } as ReadableStreamDefaultController; - const pull = Promise.resolve(source.pull?.(controller)); + await source.pull?.(controller); await flushMicrotasks(); desiredSize = 0; - mailbox.push(message); - await pull; + mailbox.push(JSON.stringify(message)); + await new Promise((resolve) => setTimeout(resolve, 0)); expect(enqueued.map(decodeText)).toEqual([serializeSseEvent(message)]); }); }); +describe("createSseBody", () => { + it("leaves messages queued until its reader asks for one", async () => { + const mailbox = new OutboundMailbox(); + const lease = mailbox.tryAcquire(); + if (!lease) { + throw new Error("Expected outbound mailbox lease"); + } + const reader = createSseBody(lease).getReader(); + + mailbox.push(JSON.stringify(message)); + await flushMicrotasks(); + expect(mailbox.hasQueuedMessages).toBe(true); + + const { value } = await reader.read(); + expect(value && decodeText(value)).toBe(serializeSseEvent(message)); + expect(mailbox.hasQueuedMessages).toBe(false); + await reader.cancel(); + }); + + it("ends cleanly when its stream shuts down while the reader waits", async () => { + const mailbox = new OutboundMailbox(); + const lease = mailbox.tryAcquire(); + if (!lease) { + throw new Error("Expected outbound mailbox lease"); + } + const reader = createSseBody(lease).getReader(); + const read = reader.read(); + await flushMicrotasks(); + + mailbox.abort(); + + await expect(read).resolves.toEqual({ done: true, value: undefined }); + }); + + it("sends keep-alives only while its reader waits, and aborts once stopped otherwise", async () => { + vi.useFakeTimers(); + + try { + const mailbox = new OutboundMailbox(); + const lease = mailbox.tryAcquire(); + if (!lease) { + throw new Error("Expected outbound mailbox lease"); + } + const reader = createSseBody(lease).getReader(); + + // The reader takes a keep-alive, then stops asking, as if stuck writing + // it to a client that stopped reading. + const read = reader.read(); + await vi.advanceTimersByTimeAsync(15_000); + const { value } = await read; + expect(value && decodeText(value)).toBe(serializeSseKeepAlive()); + await vi.advanceTimersByTimeAsync(60_000); + + // Even a clean shutdown aborts the body, and no keep-alives piled up. + mailbox.abort(); + await expect(reader.read()).rejects.toThrow("ACP connection is closed"); + } finally { + vi.useRealTimers(); + } + }); +}); + async function flushMicrotasks(): Promise { await Promise.resolve(); await Promise.resolve(); diff --git a/src/server-sse.ts b/src/server-sse.ts index 256eec09..6f6af30e 100644 --- a/src/server-sse.ts +++ b/src/server-sse.ts @@ -1,19 +1,32 @@ -import { serializeSseEvent, serializeSseKeepAlive } from "./sse.js"; +import { serializeSseJson, serializeSseKeepAlive } from "./sse.js"; import type { OutboundLease } from "./connection.js"; +/** + * Streams the JSON text of each message a lease yields as an SSE event. The + * body takes a message only when its reader asks for one, so messages its + * client has not read stay queued, and counted, in their connection. + */ export function createSseBody( - lease: OutboundLease, + lease: OutboundLease, ): ReadableStream { - return new ReadableStream(createSseBodySource(lease)); + return new ReadableStream(createSseBodySource(lease), { + highWaterMark: 0, + }); } /** @internal */ export function createSseBodySource( - lease: OutboundLease, + lease: OutboundLease, ): UnderlyingDefaultSource { const encoder = new TextEncoder(); let keepAliveTimer: ReturnType | undefined; + /** + * Whether the reader has asked for a chunk it has not been given. With no + * high-water mark, a reader that is not waiting may still be writing the + * last chunk out to a client that stopped reading. + */ + let isReaderWaiting = false; let isReceiving = false; let isClosed = false; @@ -28,6 +41,9 @@ export function createSseBodySource( controller: ReadableStreamDefaultController, text: string, ): boolean => { + // The reader takes this chunk at once; pull() runs again when it asks for + // another. + isReaderWaiting = false; try { controller.enqueue(encoder.encode(text)); return true; @@ -36,10 +52,6 @@ export function createSseBodySource( } }; - const hasDemand = ( - controller: ReadableStreamDefaultController, - ): boolean => controller.desiredSize !== null && controller.desiredSize > 0; - const closeBody = ( controller: ReadableStreamDefaultController, ): void => { @@ -58,26 +70,46 @@ export function createSseBodySource( } }; - return { - start(controller) { - keepAliveTimer = setInterval(() => { - if (isClosed || !hasDemand(controller)) { - return; - } + const errorBody = ( + controller: ReadableStreamDefaultController, + error: unknown, + ): void => { + if (isClosed) { + return; + } - if (!enqueueText(controller, serializeSseKeepAlive())) { - closeBody(controller); - } - }, 15_000); - }, - async pull(controller) { - if (isClosed || isReceiving || !hasDemand(controller)) { - return; - } + isClosed = true; + clearKeepAlive(); + lease.release(); + controller.error(error); + }; + + /** + * Once the stream stops, the body can only end. A waiting reader has taken + * every chunk, so the pending receive() ends the body, cleanly or with the + * error. Any other reader may be stuck on a chunk its client will never + * take, so abort the body rather than leave its response open. + */ + const endIfStopped = ( + controller: ReadableStreamDefaultController, + ): void => { + if (lease.stopped.aborted && !isReaderWaiting) { + errorBody(controller, lease.stopped.reason); + } + }; - isReceiving = true; + /** Hands the reader messages for as long as it waits for them. */ + const deliver = async ( + controller: ReadableStreamDefaultController, + ): Promise => { + if (isReceiving) { + return; + } + + isReceiving = true; - try { + try { + while (isReaderWaiting && !isClosed) { const result = await lease.receive(); if (isClosed) { @@ -89,18 +121,48 @@ export function createSseBodySource( return; } - if (!enqueueText(controller, serializeSseEvent(result.value))) { + if (!enqueueText(controller, serializeSseJson(result.value))) { closeBody(controller); + return; } - } catch (error) { - if (!isClosed) { - isClosed = true; - clearKeepAlive(); - controller.error(error); - } - } finally { - isReceiving = false; + + // The stream may have stopped while this message was on its way. + endIfStopped(controller); } + } catch (error) { + errorBody(controller, error); + } finally { + isReceiving = false; + } + }; + + return { + start(controller) { + lease.stopped.addEventListener( + "abort", + () => { + endIfStopped(controller); + }, + { once: true }, + ); + + keepAliveTimer = setInterval(() => { + // Only a waiting reader needs a keep-alive. + if (isClosed || !isReaderWaiting) { + return; + } + + if (!enqueueText(controller, serializeSseKeepAlive())) { + closeBody(controller); + } + }, 15_000); + }, + // With no high-water mark, this runs only when the reader asks for a + // chunk. It returns at once, so the stream calls it again whenever the + // reader asks while a receive() is still pending. + pull(controller) { + isReaderWaiting = true; + void deliver(controller); }, cancel() { isClosed = true; diff --git a/src/server-websocket-upgrade.test.ts b/src/server-websocket-upgrade.test.ts index 40f5d073..61a939a9 100644 --- a/src/server-websocket-upgrade.test.ts +++ b/src/server-websocket-upgrade.test.ts @@ -15,6 +15,7 @@ import { import { HEADER_CONNECTION_ID, JSON_MIME_TYPE } from "./protocol.js"; import { AcpServer } from "./server.js"; import { createTestAgentApp, TestAgent } from "./test-support/test-agent.js"; +import { until } from "./test-support/until.js"; import { handleWebSocketConnection } from "./ws-server.js"; import type { InitializeResponse } from "./acp.js"; @@ -687,6 +688,347 @@ describe("AcpServer prepared WebSocket upgrades", () => { } }); + it("closes the socket for session IDs and request IDs longer than maxIdLength", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const registry = new ConnectionRegistry({ maxIdLength: 8 }); + const agent = createTestAgentApp(); + const tooLong = "s".repeat(9); + const frames = [ + { + jsonrpc: "2.0", + method: "session/cancel", + params: { sessionId: tooLong }, + }, + { ...sessionNewRequest, id: tooLong }, + ]; + + try { + for (const [index, frame] of frames.entries()) { + const socket = new FakeServerSocket(); + const session = handleWebSocketConnection(socket, { registry, agent }); + socket.receive(JSON.stringify(initializeRequest)); + await readSentMessage(socket); + + socket.receive(JSON.stringify(frame)); + await session.closed; + + expect(socket.closeCode).toBe(1008); + expect(socket.closeReason).toBe( + `${index === 0 ? "Session" : "Request"} ID exceeds maxIdLength (8)`, + ); + } + } finally { + warn.mockRestore(); + await registry.closeAll(); + } + }); + + it("holds back client requests while the socket holds maxBufferedBytes", async () => { + const registry = new ConnectionRegistry({ maxBufferedBytes: 1024 }); + const agent = createTestAgentApp(); + const socket = new FakeServerSocket(); + handleWebSocketConnection(socket, { registry, agent }); + + try { + socket.receive(JSON.stringify(initializeRequest)); + await readSentMessage(socket); + + // The client has not read what the socket holds, so the request waits. + socket.bufferedAmount = 1024; + const reads = socket.bufferedAmountReads; + socket.receive(JSON.stringify(sessionNewRequest)); + await until(() => socket.bufferedAmountReads > reads); + await delay(20); + expect(socket.sent).toEqual([]); + + socket.bufferedAmount = 0; + await expect(readSentMessage(socket)).resolves.toMatchObject({ + id: sessionNewRequest.id, + result: { sessionId: expect.any(String) }, + }); + expect(socket.closeCount).toBe(0); + } finally { + socket.close(); + await registry.closeAll(); + } + }); + + it("holds back batches the agent answers while the socket holds maxBufferedBytes", async () => { + const registry = new ConnectionRegistry({ maxBufferedBytes: 4096 }); + const agent = createProtocolAgent(2); + const socket = new FakeServerSocket(); + handleWebSocketConnection(socket, { registry, agent }); + + try { + socket.receive(JSON.stringify(v2InitializeRequest)); + await readSentMessage(socket); + + // Each of these gets an error reply, and none is a request. + socket.bufferedAmount = 4096; + const reads = socket.bufferedAmountReads; + socket.receive(JSON.stringify([])); + socket.receive( + JSON.stringify([ + { id: 1, padding: "p" }, + { jsonrpc: "2.0", method: "_vendor/acme/notification" }, + ]), + ); + await until(() => socket.bufferedAmountReads > reads); + await delay(20); + expect(socket.sent).toEqual([]); + + socket.bufferedAmount = 0; + await expect(readSentWireMessage(socket)).resolves.toMatchObject({ + error: { code: -32600 }, + }); + await expect(readSentWireMessage(socket)).resolves.toMatchObject([ + { error: { code: -32600 } }, + ]); + } finally { + socket.close(); + await registry.closeAll(); + } + }); + + it("paces agent output to what the socket has sent", async () => { + const registry = new ConnectionRegistry({ maxBufferedBytes: 4096 }); + const sent: number[] = []; + const agent = createAgentApp({ name: "burst-agent" }) + .onRequest(methods.agent.initialize, () => ({ + protocolVersion: PROTOCOL_VERSION, + agentCapabilities: {}, + })) + .onRequest(methods.agent.session.prompt, async (c) => { + for (let index = 0; index < 50; index++) { + await c.client.notify(methods.client.session.update, { + sessionId: c.params.sessionId, + update: { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text: "x".repeat(1000) }, + }, + }); + sent.push(index); + } + + return { stopReason: "end_turn" }; + }); + const socket = new FakeServerSocket(); + handleWebSocketConnection(socket, { registry, agent }); + + try { + socket.receive(JSON.stringify(initializeRequest)); + await readSentMessage(socket); + + // The client stops reading, so everything sent stays in the socket. + socket.holdsSentData = true; + socket.receive( + JSON.stringify({ + jsonrpc: "2.0", + id: 5, + method: "session/prompt", + params: { + sessionId: "session-1", + prompt: [{ type: "text", text: "go" }], + }, + }), + ); + await until(() => socket.bufferedAmount >= 4096); + await delay(20); + expect(sent.length).toBeGreaterThan(0); + expect(sent.length).toBeLessThan(10); + expect(socket.sent.length).toBeLessThan(10); + + // Once the client reads, the agent finishes. + socket.holdsSentData = false; + socket.bufferedAmount = 0; + await until( + () => socket.sent.some((data) => JSON.parse(data).id === 5), + 2_000, + ); + expect(sent).toHaveLength(50); + expect(socket.closeCount).toBe(0); + } finally { + socket.close(); + await registry.closeAll(); + } + }); + + it("closes the socket when its client reads none of its output for maxOutputStallMs", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const registry = new ConnectionRegistry({ + maxBufferedBytes: 1024, + maxOutputStallMs: 50, + }); + // A reply larger than the socket may hold, after which the agent has + // nothing more to send. + const agent = createTestAgentApp({ + newSession: () => ({ + sessionId: "session-1", + _meta: { padding: "p".repeat(2000) }, + }), + }); + const socket = new FakeServerSocket(); + const session = handleWebSocketConnection(socket, { registry, agent }); + + try { + socket.receive(JSON.stringify(initializeRequest)); + await readSentMessage(socket); + + socket.holdsSentData = true; + socket.receive(JSON.stringify(sessionNewRequest)); + await session.closed; + + expect(socket.sent).toHaveLength(1); + expect(socket.closeCode).toBe(1008); + expect(socket.closeReason).toBe( + "Connection output stalled for maxOutputStallMs (50)", + ); + } finally { + warn.mockRestore(); + await registry.closeAll(); + } + }); + + it("closes a full socket once its connection shuts down", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const registry = new ConnectionRegistry({ + maxBufferedBytes: 1024, + maxIdLength: 8, + }); + let agentWriter: WritableStreamDefaultWriter | undefined; + const agent: AgentConnector = { + connect(stream) { + void (async () => { + const reader = stream.readable.getReader(); + const { value } = await reader.read(); + agentWriter = stream.writable.getWriter(); + await agentWriter.write({ + jsonrpc: "2.0", + id: (value as { id: number }).id, + result: { + protocolVersion: PROTOCOL_VERSION, + agentCapabilities: {}, + }, + }); + })(); + }, + }; + const socket = new FakeServerSocket(); + const session = handleWebSocketConnection(socket, { registry, agent }); + + try { + socket.receive(JSON.stringify(initializeRequest)); + await readSentMessage(socket); + if (!agentWriter) { + throw new Error("Expected the agent to have answered initialize"); + } + + // The client stops reading, so this fills the socket. + socket.holdsSentData = true; + void agentWriter.write({ + jsonrpc: "2.0", + method: "_vendor/acme/notification", + params: { padding: "p".repeat(2000) }, + }); + await readSentMessage(socket); + + // A session ID the connection refuses shuts it down, which closes the + // socket without waiting for the client to read. + void agentWriter.write({ + jsonrpc: "2.0", + method: "_vendor/acme/notification", + params: { sessionId: "s".repeat(9) }, + }); + await session.closed; + + expect(socket.closeCode).toBe(1008); + expect(socket.closeReason).toBe("Session ID exceeds maxIdLength (8)"); + } finally { + warn.mockRestore(); + await registry.closeAll(); + } + }); + + it("closes the socket when the client sends more than maxBufferedBytes while earlier messages wait", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const initialize = createDeferred(); + const registry = new ConnectionRegistry({ maxBufferedBytes: 4096 }); + const agent = createTestAgentApp({ initialize: () => initialize.promise }); + const connection = registry.createPendingConnection(agent); + const socket = new FakeServerSocket(); + const session = handleWebSocketConnection(socket, { + registry, + agent, + connection, + }); + const padding = JSON.stringify({ + jsonrpc: "2.0", + method: "_vendor/acme/padding", + params: { padding: "p".repeat(1500) }, + }); + + try { + // Frames wait behind the initialize request until the agent answers. + socket.receive(JSON.stringify(initializeRequest)); + await delay(0); + socket.receive(padding); + expect(socket.closeCount).toBe(0); + socket.receive(padding); + await session.closed; + + expect(socket.closeCode).toBe(1008); + expect(socket.closeReason).toBe( + "Connection exceeds maxBufferedBytes (4096)", + ); + expect(warn).toHaveBeenCalledWith( + `Closing ACP connection ${connection.connectionId}:`, + "Connection exceeds maxBufferedBytes (4096)", + ); + } finally { + initialize.resolve({ + protocolVersion: PROTOCOL_VERSION, + agentCapabilities: {}, + }); + warn.mockRestore(); + await registry.closeAll(); + } + }); + + it("counts empty frames waiting to be handled toward maxBufferedBytes", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const initialize = createDeferred(); + const registry = new ConnectionRegistry({ maxBufferedBytes: 4096 }); + const agent = createTestAgentApp({ initialize: () => initialize.promise }); + const socket = new FakeServerSocket(); + const session = handleWebSocketConnection(socket, { registry, agent }); + + try { + // Frames wait behind the initialize request until the agent answers. + // Each also holds more than its text, which counts 1 KiB, so four fit. + socket.receive(JSON.stringify(initializeRequest)); + await delay(0); + let frames = 0; + while (socket.closeCount === 0 && frames < 100) { + socket.receive(""); + frames += 1; + } + + expect(frames).toBe(5); + await session.closed; + expect(socket.closeCode).toBe(1008); + expect(socket.closeReason).toBe( + "Connection exceeds maxBufferedBytes (4096)", + ); + } finally { + initialize.resolve({ + protocolVersion: PROTOCOL_VERSION, + agentCapabilities: {}, + }); + warn.mockRestore(); + await registry.closeAll(); + } + }); + it("forwards post-initialize batches and sends one response-array frame", async () => { const registry = new ConnectionRegistry(); const agent = createProtocolAgent(2); @@ -785,7 +1127,7 @@ describe("AcpServer prepared WebSocket upgrades", () => { await loadStarted.promise; expect(connection.pendingRoutes.get("number:30")).toBe("connection"); - expect(connection.sessionStreams.has(sessionId)).toBe(true); + expect(connection.sessionStreams.has(sessionId)).toBe(false); finishLoad.resolve(); const response = await readSentWireMessage(socket); @@ -903,7 +1245,7 @@ describe("AcpServer prepared WebSocket upgrades", () => { [`number:${secondRequest.id}`, "connection"], ]), ); - expect(connection.sessionStreams.has(sessionId)).toBe(true); + expect(connection.sessionStreams.has(sessionId)).toBe(false); socket.receive( JSON.stringify([ @@ -1100,6 +1442,10 @@ function createDeferred(): { return { promise, resolve, reject }; } +function delay(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + function readSentMessage(socket: FakeServerSocket): Promise { return readSentWireMessage(socket).then((message) => { if (Array.isArray(message)) { @@ -1131,12 +1477,29 @@ class FakeServerSocket implements WebSocketServerSocket { readonly sent: string[] = []; readonly listeners = new Map void>>(); onSend: ((data: string) => void) | undefined; + /** Whether sent data stays in `bufferedAmount`, as if the peer stopped reading. */ + holdsSentData = false; + /** How many times `bufferedAmount` was read, such as by a session waiting on it. */ + bufferedAmountReads = 0; closeCount = 0; closeCode: number | undefined; closeReason: string | undefined; + private unsentBytes = 0; + + get bufferedAmount(): number { + this.bufferedAmountReads += 1; + return this.unsentBytes; + } + + set bufferedAmount(value: number) { + this.unsentBytes = value; + } send(data: string): void { this.sent.push(data); + if (this.holdsSentData) { + this.unsentBytes += data.length; + } this.onSend?.(data); this.onSend = undefined; } diff --git a/src/server.test.ts b/src/server.test.ts index 8c471fdb..e1ed18ec 100644 --- a/src/server.test.ts +++ b/src/server.test.ts @@ -15,7 +15,9 @@ import { AcpServer } from "./server.js"; import { parseSseStream } from "./sse.js"; import { createTestAgentApp, TestAgent } from "./test-support/test-agent.js"; import { startTestServer } from "./test-support/test-http-server.js"; +import { until } from "./test-support/until.js"; +import type { ConnectionRegistry, ConnectionState } from "./connection.js"; import type { AnyMessage } from "./jsonrpc.js"; const initializeRequest = { @@ -1162,6 +1164,717 @@ describe("AcpServer", () => { }); }); +describe("AcpServer connection limits", () => { + it("does not keep session streams after their receivers leave", async () => { + const server = new AcpServer({ + createAgent: () => createTestAgentApp(), + }); + + try { + const connectionId = await initializeDirect(server); + + for (let index = 0; index < 100; index++) { + const response = await server.handleRequest( + sseRequest(connectionId, `unknown-session-${index}`), + ); + expect(response.status).toBe(200); + await response.body?.cancel(); + } + + expect(connectionState(server, connectionId).sessionStreams.size).toBe(0); + } finally { + await server.close(); + } + }); + + it("does not track routes for messages the agent never answers", async () => { + const error = vi + .spyOn(console, "error") + .mockImplementation(() => undefined); + const server = new AcpServer({ + createAgent: () => createTestAgentApp(), + }); + + try { + const connectionId = await initializeDirect(server); + const unanswered = [ + { id: "bare-id" }, + { jsonrpc: "2.0", id: "no-result" }, + { + jsonrpc: "2.0", + id: "result-and-error", + result: {}, + error: { code: -32603, message: "Internal error" }, + }, + { + id: "no-version", + method: "session/cancel", + params: { sessionId: "session-1" }, + }, + ]; + + for (const message of unanswered) { + const response = await server.handleRequest( + jsonRequest(message, sessionHeaders(connectionId, "session-1")), + ); + expect(response.status).toBe(202); + } + + expect(connectionState(server, connectionId).pendingRoutes.size).toBe(0); + } finally { + error.mockRestore(); + await server.close(); + } + }); + + it("rejects session IDs and request IDs longer than maxIdLength", async () => { + const server = new AcpServer({ + createAgent: () => createTestAgentApp(), + maxIdLength: 8, + }); + + try { + const connectionId = await initializeDirect(server); + const tooLong = "s".repeat(9); + + const stream = await server.handleRequest( + sseRequest(connectionId, tooLong), + ); + const prompt = await server.handleRequest( + jsonRequest( + promptRequestFor(4, tooLong), + sessionHeaders(connectionId, tooLong), + ), + ); + const longRequestId = await server.handleRequest( + jsonRequest( + { ...sessionNewRequest, id: tooLong }, + { [HEADER_CONNECTION_ID]: connectionId }, + ), + ); + expect(stream.status).toBe(400); + expect(prompt.status).toBe(400); + expect(longRequestId.status).toBe(400); + expect(await longRequestId.text()).toBe( + "Request ID exceeds maxIdLength (8)", + ); + + const allowed = await server.handleRequest( + sseRequest(connectionId, "s".repeat(8)), + ); + expect(allowed.status).toBe(200); + await allowed.body?.cancel(); + } finally { + await server.close(); + } + }); + + it("closes the connection when the agent issues an over-long session ID", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const server = new AcpServer({ + createAgent: () => + createTestAgentApp({ + newSession: () => ({ sessionId: "s".repeat(9) }), + }), + maxIdLength: 8, + }); + + try { + const connectionId = await initializeDirect(server); + const connection = connectionState(server, connectionId); + const stream = await server.handleRequest(sseRequest(connectionId)); + if (!stream.body) { + throw new Error("Expected SSE response body"); + } + const reader = stream.body.getReader(); + + await server.handleRequest( + jsonRequest(sessionNewRequest, { + [HEADER_CONNECTION_ID]: connectionId, + }), + ); + + await expect(withTimeout(reader.read())).rejects.toThrow( + "Session ID exceeds maxIdLength (8)", + ); + await withTimeout(connection.closed); + expect(warn).toHaveBeenCalledWith( + `Closing ACP connection ${connectionId}:`, + "Session ID exceeds maxIdLength (8)", + ); + } finally { + warn.mockRestore(); + await server.close(); + } + }); + + it("closes only the connection that buffers too many session streams", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const server = new AcpServer({ + createAgent: () => createTestAgentApp(), + maxBufferedSessionStreams: 2, + }); + + try { + const connectionId = await initializeDirect(server); + const otherConnectionId = await initializeDirect(server); + const connection = connectionState(server, connectionId); + + // Nothing receives these sessions' streams, so their output buffers. + for (const [index, sessionId] of ["a", "b", "c"].entries()) { + await server.handleRequest( + jsonRequest( + promptRequestFor(10 + index, sessionId), + sessionHeaders(connectionId, sessionId), + ), + ); + } + + await withTimeout(connection.closed); + expect(warn).toHaveBeenCalledWith( + `Closing ACP connection ${connectionId}:`, + "Connection exceeds maxBufferedSessionStreams (2)", + ); + + const other = await server.handleRequest(sseRequest(otherConnectionId)); + expect(other.status).toBe(200); + await other.body?.cancel(); + } finally { + warn.mockRestore(); + await server.close(); + } + }); + + it("closes a connection whose clients leave too many session streams with output queued", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const server = new AcpServer({ + createAgent: () => createTestAgentApp(), + maxBufferedSessionStreams: 2, + }); + + try { + const connectionId = await initializeDirect(server); + const connection = connectionState(server, connectionId); + const sessionIds = ["a", "b", "c"]; + + // Open streams are not limited, and these read none of their output. + const streams: Response[] = []; + for (const [index, sessionId] of sessionIds.entries()) { + streams.push( + await server.handleRequest(sseRequest(connectionId, sessionId)), + ); + await server.handleRequest( + jsonRequest( + promptRequestFor(10 + index, sessionId), + sessionHeaders(connectionId, sessionId), + ), + ); + } + await until(() => + sessionIds.every( + (sessionId) => + connection.sessionStreams.get(sessionId)?.hasQueuedMessages, + ), + ); + expect(streams.map((stream) => stream.status)).toEqual([200, 200, 200]); + expect(connection.isClosed).toBe(false); + + // Leaving them buffers their output, one session too many. + for (const stream of streams) { + await stream.body?.cancel(); + } + + await withTimeout(connection.closed); + expect(warn).toHaveBeenCalledWith( + `Closing ACP connection ${connectionId}:`, + "Connection exceeds maxBufferedSessionStreams (2)", + ); + } finally { + warn.mockRestore(); + await server.close(); + } + }); + + it("closes a full connection whose output no client takes for maxOutputStallMs", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const server = new AcpServer({ + createAgent: () => createTestAgentApp(), + maxBufferedBytes: 1024, + maxOutputStallMs: 50, + }); + + try { + const connectionId = await initializeDirect(server); + const connection = connectionState(server, connectionId); + + // Nothing receives the connection stream, so its output never drains. + for (let index = 0; index < 50 && !connection.isClosed; index++) { + await server.handleRequest( + jsonRequest( + { ...sessionNewRequest, id: 100 + index }, + { [HEADER_CONNECTION_ID]: connectionId }, + ), + ); + } + + await withTimeout(connection.closed); + expect(warn).toHaveBeenCalledWith( + `Closing ACP connection ${connectionId}:`, + "Connection output stalled for maxOutputStallMs (50)", + ); + } finally { + warn.mockRestore(); + await server.close(); + } + }); + + it("delivers a burst queued before the client opens its session stream", async () => { + const server = new AcpServer({ + createAgent: () => createBurstAgent(50, 1000), + maxBufferedBytes: 4096, + }); + + try { + const connectionId = await initializeDirect(server); + const connection = connectionState(server, connectionId); + await server.handleRequest( + jsonRequest( + promptRequestFor(10, "session-1"), + sessionHeaders(connectionId, "session-1"), + ), + ); + await until(() => connection.hasOutboundBacklog); + + const stream = await server.handleRequest( + sseRequest(connectionId, "session-1"), + ); + expect(await readSseMessages(stream, 51)).toHaveLength(51); + expect(connection.isClosed).toBe(false); + } finally { + await server.close(); + } + }); + + it("keeps a full connection open while its receiver reconnects", async () => { + const server = new AcpServer({ + createAgent: () => createBurstAgent(50, 1000), + maxBufferedBytes: 4096, + }); + + try { + const connectionId = await initializeDirect(server); + const first = await server.handleRequest( + sseRequest(connectionId, "session-1"), + ); + await server.handleRequest( + jsonRequest( + promptRequestFor(10, "session-1"), + sessionHeaders(connectionId, "session-1"), + ), + ); + await until( + () => connectionState(server, connectionId).hasOutboundBacklog, + ); + await first.body?.cancel(); + + const second = await server.handleRequest( + sseRequest(connectionId, "session-1"), + ); + if (!second.body) { + throw new Error("Expected SSE response body"); + } + const messages = parseSseStream(second.body)[Symbol.asyncIterator](); + for (;;) { + const next = await withTimeout(messages.next()); + if (next.done) { + throw new Error("Expected the prompt response"); + } + if ((next.value as { id?: unknown }).id === 10) { + break; + } + } + await messages.return?.(); + expect(connectionState(server, connectionId).isClosed).toBe(false); + } finally { + await server.close(); + } + }); + + it("leaves nothing waiting for requests aborted while held back", async () => { + const server = new AcpServer({ + createAgent: () => createBurstAgent(50, 1000), + maxBufferedBytes: 4096, + }); + + try { + const connectionId = await initializeDirect(server); + const stream = await server.handleRequest( + sseRequest(connectionId, "session-1"), + ); + await server.handleRequest( + jsonRequest( + promptRequestFor(10, "session-1"), + sessionHeaders(connectionId, "session-1"), + ), + ); + const connection = connectionState(server, connectionId); + await until(() => connection.hasOutboundBacklog); + + // Held-back requests wait alongside the router, which is paused. + const waiters = (): number => + (connection as unknown as { progressWaiters: Set }) + .progressWaiters.size; + const waitersBefore = waiters(); + for (let index = 0; index < 20; index++) { + const abort = new AbortController(); + const response = server.handleRequest( + jsonRequest( + promptRequestFor(100 + index, "session-1"), + sessionHeaders(connectionId, "session-1"), + abort.signal, + ), + ); + await delay(1); + abort.abort(); + expect((await withTimeout(response)).status).toBe(499); + } + + expect(waiters()).toBe(waitersBefore); + await stream.body?.cancel(); + } finally { + await server.close(); + } + }); + + it("delivers an agent burst larger than maxBufferedBytes to a client that reads it", async () => { + const server = new AcpServer({ + createAgent: () => createBurstAgent(50, 1000), + maxBufferedBytes: 4096, + }); + + try { + const connectionId = await initializeDirect(server); + const stream = await server.handleRequest( + sseRequest(connectionId, "session-1"), + ); + if (!stream.body) { + throw new Error("Expected SSE response body"); + } + const messages = parseSseStream(stream.body)[Symbol.asyncIterator](); + + await server.handleRequest( + jsonRequest( + promptRequestFor(10, "session-1"), + sessionHeaders(connectionId, "session-1"), + ), + ); + + for (let index = 0; index < 50; index++) { + const next = await withTimeout(messages.next()); + expect(next.value).toMatchObject({ method: "session/update" }); + } + expect((await withTimeout(messages.next())).value).toMatchObject({ + id: 10, + result: { stopReason: "end_turn" }, + }); + + await messages.return?.(); + expect(connectionState(server, connectionId).isClosed).toBe(false); + } finally { + await server.close(); + } + }); + + it("holds back agent output and client requests while a receiver is behind", async () => { + const sent: number[] = []; + const server = new AcpServer({ + createAgent: () => + createBurstAgent(50, 1000, (index) => { + sent.push(index); + }), + maxBufferedBytes: 4096, + }); + + try { + const connectionId = await initializeDirect(server); + const stream = await server.handleRequest( + sseRequest(connectionId, "session-1"), + ); + if (!stream.body) { + throw new Error("Expected SSE response body"); + } + + await server.handleRequest( + jsonRequest( + promptRequestFor(10, "session-1"), + sessionHeaders(connectionId, "session-1"), + ), + ); + await until( + () => connectionState(server, connectionId).hasOutboundBacklog, + ); + await delay(20); + + // Nothing has read the stream, so the agent is waiting on its sends + // and another request waits too. + expect(sent.length).toBeGreaterThan(0); + expect(sent.length).toBeLessThan(10); + let isNextPromptAccepted = false; + const nextPrompt = server + .handleRequest( + jsonRequest( + promptRequestFor(11, "session-1"), + sessionHeaders(connectionId, "session-1"), + ), + ) + .then((response) => { + isNextPromptAccepted = true; + return response; + }); + await delay(50); + expect(isNextPromptAccepted).toBe(false); + + // Reading lets both carry on. + const messages = parseSseStream(stream.body)[Symbol.asyncIterator](); + for (;;) { + const next = await withTimeout(messages.next()); + if (next.done) { + throw new Error("Expected the prompt response"); + } + if ((next.value as { id?: unknown }).id === 10) { + break; + } + } + expect(sent.length).toBeGreaterThanOrEqual(50); + expect((await withTimeout(nextPrompt)).status).toBe(202); + + await messages.return?.(); + } finally { + await server.close(); + } + }); + + it("answers a held-back request with 404 once its agent exits", async () => { + const agentClosed = createDeferred(); + let agentWriter: WritableStreamDefaultWriter | undefined; + const server = new AcpServer({ + agent: { + connect(stream) { + void (async () => { + const reader = stream.readable.getReader(); + const { value } = await reader.read(); + agentWriter = stream.writable.getWriter(); + await agentWriter.write({ + jsonrpc: "2.0", + id: (value as { id: number }).id, + result: { protocolVersion: 1, agentCapabilities: {} }, + }); + })(); + return { closed: agentClosed.promise }; + }, + }, + maxBufferedBytes: 64, + }); + + try { + const connectionId = await initializeDirect(server); + const connection = connectionState(server, connectionId); + const stream = await server.handleRequest(sseRequest(connectionId)); + if (!agentWriter) { + throw new Error("Expected the agent to have answered initialize"); + } + + // The connection stream is open but unread, so this fills the limit. + void agentWriter.write({ + jsonrpc: "2.0", + method: "_vendor/acme/notification", + params: { padding: "p".repeat(100) }, + }); + await until(() => connection.hasOutboundBacklog); + const request = server.handleRequest( + jsonRequest( + { ...sessionNewRequest, id: 20 }, + { [HEADER_CONNECTION_ID]: connectionId }, + ), + ); + await delay(20); + + agentClosed.resolve(); + expect((await withTimeout(request)).status).toBe(404); + await stream.body?.cancel(); + } finally { + await server.close(); + } + }); + + it("closes a connection whose invented session stalls its output", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); + const server = new AcpServer({ + createAgent: () => createTestAgentApp(), + maxBufferedBytes: 1024, + maxOutputStallMs: 50, + }); + + try { + const connectionId = await initializeDirect(server); + const connection = connectionState(server, connectionId); + + // One session, never received, so its output fills the limit and stays. + for (let index = 0; index < 50 && !connection.isClosed; index++) { + await server.handleRequest( + jsonRequest( + promptRequestFor(100 + index, "unknown-session"), + sessionHeaders(connectionId, "unknown-session"), + ), + ); + } + + await withTimeout(connection.closed); + expect(warn).toHaveBeenCalledWith( + `Closing ACP connection ${connectionId}:`, + "Connection output stalled for maxOutputStallMs (50)", + ); + } finally { + warn.mockRestore(); + await server.close(); + } + }); + + it("ends the stream of a client that stopped reading when its connection closes", async () => { + const server = new AcpServer({ + createAgent: () => createTestAgentApp(), + }); + + try { + const connectionId = await initializeDirect(server); + const stream = await server.handleRequest(sseRequest(connectionId)); + if (!stream.body) { + throw new Error("Expected SSE response body"); + } + + await server.handleRequest( + jsonRequest(sessionNewRequest, { + [HEADER_CONNECTION_ID]: connectionId, + }), + ); + const deleted = await server.handleRequest( + new Request("http://127.0.0.1/acp", { + method: "DELETE", + headers: { [HEADER_CONNECTION_ID]: connectionId }, + }), + ); + expect(deleted.status).toBe(202); + + // A body left waiting for its reader would keep the response open. + await expect(withTimeout(stream.body.getReader().read())).rejects.toThrow( + "ACP connection is closed", + ); + } finally { + await server.close(); + } + }); + + it("rejects connection limits that are not positive safe integers", () => { + const createAgent = () => createTestAgentApp(); + + for (const value of [0, -1, 1.5, Number.NaN, Number.POSITIVE_INFINITY]) { + expect( + () => new AcpServer({ createAgent, maxBufferedBytes: value }), + ).toThrow(RangeError); + expect( + () => new AcpServer({ createAgent, maxOutputStallMs: value }), + ).toThrow(RangeError); + expect( + () => new AcpServer({ createAgent, maxBufferedSessionStreams: value }), + ).toThrow(RangeError); + expect(() => new AcpServer({ createAgent, maxIdLength: value })).toThrow( + RangeError, + ); + } + + // Timers cannot wait longer than this. + expect( + () => new AcpServer({ createAgent, maxOutputStallMs: 2 ** 31 }), + ).toThrow("maxOutputStallMs must be at most 2147483647"); + }); + + it("does not count delivered output against maxBufferedBytes", async () => { + const server = new AcpServer({ + createAgent: () => createTestAgentApp(), + maxBufferedBytes: 1024, + }); + + try { + const connectionId = await initializeDirect(server); + const stream = await server.handleRequest(sseRequest(connectionId)); + if (!stream.body) { + throw new Error("Expected SSE response body"); + } + const messages = parseSseStream(stream.body)[Symbol.asyncIterator](); + + // Together these responses exceed the limit, but each is delivered first. + for (let index = 0; index < 50; index++) { + await server.handleRequest( + jsonRequest( + { ...sessionNewRequest, id: 100 + index }, + { [HEADER_CONNECTION_ID]: connectionId }, + ), + ); + const next = await withTimeout(messages.next()); + expect(next.value).toMatchObject({ id: 100 + index }); + } + + await messages.return?.(); + expect(connectionState(server, connectionId).isClosed).toBe(false); + } finally { + await server.close(); + } + }); +}); + +function connectionState( + server: AcpServer, + connectionId: string, +): ConnectionState { + const connection = ( + server as unknown as { registry: ConnectionRegistry } + ).registry.get(connectionId); + if (!connection) { + throw new Error("Expected an open connection"); + } + + return connection; +} + +function sseRequest(connectionId: string, sessionId?: string): Request { + return new Request("http://127.0.0.1/acp", { + method: "GET", + headers: { + Accept: EVENT_STREAM_MIME_TYPE, + [HEADER_CONNECTION_ID]: connectionId, + ...(sessionId === undefined ? {} : { [HEADER_SESSION_ID]: sessionId }), + }, + }); +} + +function sessionHeaders( + connectionId: string, + sessionId: string, +): Record { + return { + [HEADER_CONNECTION_ID]: connectionId, + [HEADER_SESSION_ID]: sessionId, + }; +} + +function promptRequestFor(id: number, sessionId: string) { + return { + ...promptRequest, + id, + params: { ...promptRequest.params, sessionId }, + }; +} + async function initializeDirect(server: AcpServer): Promise { const response = await server.handleRequest(jsonRequest(initializeRequest)); const connectionId = response.headers.get(HEADER_CONNECTION_ID); @@ -1322,6 +2035,45 @@ function createBackpressureAgent(onPromptDone: () => void) { .onNotification(methods.agent.session.cancel, () => {}); } +/** + * An agent that answers each prompt with `count` updates of about `size` + * characters, calling `onSent` as each send resolves. + */ +function createBurstAgent( + count: number, + size: number, + onSent: (index: number) => void = () => {}, +) { + return createAgentApp({ name: "burst-agent" }) + .onRequest(methods.agent.initialize, () => ({ + protocolVersion: 1, + agentCapabilities: { + loadSession: false, + }, + })) + .onRequest(methods.agent.session.new, () => ({ sessionId: "session-1" })) + .onRequest(methods.agent.authenticate, () => ({})) + .onRequest(methods.agent.session.prompt, async (c) => { + for (let index = 0; index < count; index++) { + await c.client.notify(methods.client.session.update, { + sessionId: c.params.sessionId, + update: { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text: "x".repeat(size) }, + }, + }); + onSent(index); + } + + return { stopReason: "end_turn" }; + }) + .onNotification(methods.agent.session.cancel, () => {}); +} + +function delay(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + async function waitForConnectionNotFound( server: AcpServer, connectionId: string, diff --git a/src/server.ts b/src/server.ts index 05eb9075..f33fb252 100644 --- a/src/server.ts +++ b/src/server.ts @@ -1,4 +1,8 @@ -import { ConnectionRegistry } from "./connection.js"; +import { + ConnectionClosedError, + ConnectionRegistry, + ConnectionLimitError, +} from "./connection.js"; import { EVENT_STREAM_MIME_TYPE, HEADER_CONNECTION_ID, @@ -9,7 +13,12 @@ import { methodRequiresSessionHeader, sessionIdFromParams, } from "./protocol.js"; -import { isRecord, isResponseMessage } from "./jsonrpc.js"; +import { + isNotificationMessage, + isRecord, + isRequestMessage, + isResponseMessage, +} from "./jsonrpc.js"; import { AGENT_METHODS } from "./schema/index.js"; import { createSseBody } from "./server-sse.js"; import { handleWebSocketConnection } from "./ws-server.js"; @@ -24,6 +33,7 @@ import type { ConnectionState, OutboundLease, ResponseRoute, + ConnectionLimits, } from "./connection.js"; import type { AnyMessage, @@ -34,6 +44,14 @@ import type { import type { Agent } from "./acp.js"; import type { Stream } from "./stream.js"; +export { + DEFAULT_MAX_BUFFERED_BYTES, + DEFAULT_MAX_BUFFERED_SESSION_STREAMS, + DEFAULT_MAX_ID_LENGTH, + DEFAULT_MAX_OUTPUT_STALL_MS, +} from "./connection.js"; +export type { ConnectionLimits } from "./connection.js"; + export type AgentFactory = () => AgentConnector; /** @deprecated Prefer {@link AgentFactory}. */ export type LegacyAgentFactory = (conn: AgentSideConnection) => Agent; @@ -87,7 +105,7 @@ type OptionalAgentOption = }; /** Options for creating an ACP server transport. */ -export type AcpServerOptions = AgentOption; +export type AcpServerOptions = AgentOption & ConnectionLimits; export type HandleRequestOptions = OptionalAgentOption; @@ -108,11 +126,12 @@ export interface PreparedWebSocketUpgrade { */ export class AcpServer { private readonly agent: AgentConnector; - private readonly registry = new ConnectionRegistry(); + private readonly registry: ConnectionRegistry; private readonly webSocketSessions = new Set(); constructor(options: AcpServerOptions) { this.agent = resolveAgent(options); + this.registry = new ConnectionRegistry(options); } /** Handles one Streamable HTTP ACP request. */ @@ -120,19 +139,38 @@ export class AcpServer { req: Request, options: HandleRequestOptions = {}, ): Promise { - if (req.method === "POST") { - return await this.handlePost(req, options); - } + try { + if (req.method === "POST") { + return await this.handlePost(req, options); + } - if (req.method === "GET") { - return this.handleGet(req); - } + if (req.method === "GET") { + return this.handleGet(req); + } - if (req.method === "DELETE") { - return this.handleDelete(req); - } + if (req.method === "DELETE") { + return this.handleDelete(req); + } + + return textResponse("Method Not Allowed", 405); + } catch (error) { + // Requests only hit connection limits through over-long IDs. + if (error instanceof ConnectionLimitError) { + return textResponse(error.message, 400); + } + + // A connection that shut down mid-request is answered as if it were gone. + if (error instanceof ConnectionClosedError) { + return textResponse("Unknown Acp-Connection-Id", 404); + } + + // The client gave up on a request still waiting for its connection. + if (error instanceof RequestAbortedError) { + return textResponse("Request aborted", 499); + } - return textResponse("Method Not Allowed", 405); + throw error; + } } /** Creates a WebSocket connection before accepting the HTTP upgrade. */ @@ -233,6 +271,7 @@ export class AcpServer { connection, message, req.headers, + req.signal, ); if (!forwarded.ok) { return textResponse(forwarded.message, forwarded.status); @@ -265,10 +304,9 @@ export class AcpServer { } const sessionId = req.headers.get(HEADER_SESSION_ID); - const mailbox = sessionId - ? connection.ensureSession(sessionId) - : connection.connectionStream; - const lease = mailbox.tryAcquire(); + const lease = sessionId + ? connection.acquireSessionStream(sessionId) + : connection.connectionStream.tryAcquire(); if (!lease) { return textResponse( @@ -360,12 +398,18 @@ export class AcpServer { connection: ConnectionState, message: AnyMessage, headers: Headers, + signal: AbortSignal, ): Promise { if (isResponseMessage(message)) { return await forwardClientResponse(connection, message, headers); } - return await forwardClientMethodMessage(connection, message, headers); + return await forwardClientMethodMessage( + connection, + message, + headers, + signal, + ); } } @@ -525,6 +569,7 @@ async function forwardClientMethodMessage( connection: ConnectionState, message: ClientMethodMessage, headers: Headers, + signal: AbortSignal, ): Promise { const route = determineRoute(message, headers); @@ -533,10 +578,21 @@ async function forwardClientMethodMessage( } if (route.value !== "connection") { - connection.ensureSession(route.value.session); + connection.validateSessionId(route.value.session); + } + + // Anything but a notification makes the agent reply, so wait until the + // client has read enough of its output to make room. + if (!isNotificationMessage(message)) { + connection.validateRequestId((message as { id?: unknown }).id); + await connection.waitForOutboundCapacity(signal).catch((error: unknown) => { + throw signal.aborted ? new RequestAbortedError() : error; + }); } - const key = "id" in message ? messageIdKey(message.id) : undefined; + // Only a valid request is answered with its own ID, which is what removes + // its route; the agent answers anything else without one, or not at all. + const key = isRequestMessage(message) ? messageIdKey(message.id) : undefined; if (key) { connection.pendingRoutes.set( @@ -644,7 +700,7 @@ function isJsonContentType(contentType: string | null): boolean { return contentType?.split(";", 1)[0]?.trim().toLowerCase() === JSON_MIME_TYPE; } -function sseResponse(lease: OutboundLease): Response { +function sseResponse(lease: OutboundLease): Response { return new Response(createSseBody(lease), { status: 200, headers: { diff --git a/src/sse.ts b/src/sse.ts index c8313f3f..e30b99b2 100644 --- a/src/sse.ts +++ b/src/sse.ts @@ -12,7 +12,12 @@ const dataSeparator = new Uint8Array([0x0a]); const maxLineOverhead = 9; // UTF-8 BOM + "data: "; LineBuffer strips LF/CRLF. export function serializeSseEvent(msg: AnyMessage): string { - return `data: ${JSON.stringify(msg)}\n\n`; + return serializeSseJson(JSON.stringify(msg)); +} + +/** Frames a message's JSON text, which has no line breaks, as an SSE event. */ +export function serializeSseJson(json: string): string { + return `data: ${json}\n\n`; } export function serializeSseKeepAlive(): string { diff --git a/src/test-support/until.ts b/src/test-support/until.ts new file mode 100644 index 00000000..76a2aebe --- /dev/null +++ b/src/test-support/until.ts @@ -0,0 +1,18 @@ +/** + * Resolves once `condition` holds, checking every few milliseconds, and + * rejects if it still does not after `timeoutMs`. + */ +export async function until( + condition: () => boolean, + timeoutMs = 1_000, +): Promise { + const deadline = Date.now() + timeoutMs; + + while (!condition()) { + if (Date.now() > deadline) { + throw new Error("Timed out waiting for condition"); + } + + await new Promise((resolve) => setTimeout(resolve, 5)); + } +} diff --git a/src/ws-server.ts b/src/ws-server.ts index bd500ea6..72b476bb 100644 --- a/src/ws-server.ts +++ b/src/ws-server.ts @@ -3,6 +3,7 @@ import { isNotificationMessage, isRecord, isRequestMessage, + isResponseBatch, isResponseShapedMessage, protocolErrorResponse, } from "./jsonrpc.js"; @@ -13,6 +14,12 @@ import { } from "./protocol.js"; import { AGENT_METHODS } from "./schema/index.js"; import { onWebSocket, webSocketMessageToString } from "./ws-utils.js"; +import { + ConnectionClosedError, + ConnectionLimitError, + connectionLimitExceeded, + connectionOutputStalled, +} from "./connection.js"; import type { AgentConnector, ConnectionRegistry, @@ -26,6 +33,16 @@ import type { WebSocketLike } from "./ws-utils.js"; /** WebSocket shape accepted by prepared ACP WebSocket upgrades. */ export type WebSocketServerSocket = WebSocketLike; +// Sockets report no drain event, so a session polls `bufferedAmount` while +// its client is behind, backing off while the client stays behind. +const SOCKET_POLL_INITIAL_MS = 10; +const SOCKET_POLL_MAX_MS = 1000; + +// Besides its text, a frame waiting to be handled holds the promises that +// queue it, about 350 bytes in V8, so each one counts this much more toward +// `maxBufferedBytes`. Otherwise small or empty frames would barely count. +const WAITING_FRAME_OVERHEAD = 1024; + type ForwardResult = | { ok: true; @@ -61,6 +78,12 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { private outboundLease: OutboundLease | undefined; private inboundWriteChain: Promise = Promise.resolve(); private messageChain: Promise = Promise.resolve(); + /** + * Size of the frames waiting in `messageChain` behind the one being + * handled, counting `WAITING_FRAME_OVERHEAD` for each; see + * `maxBufferedBytes`. + */ + private waitingInboundBytes = 0; private isClosed = false; private readonly closedPromise: Promise; private resolveClosed: () => void = () => {}; @@ -105,18 +128,6 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { } private enqueueSocketMessage(args: unknown[]): void { - const handled = this.messageChain.then(() => - this.handleSocketMessage(args), - ); - this.messageChain = handled.catch((error) => { - if (!this.isClosed) { - console.error("ACP WebSocket message handling failed:", error); - void this.shutdown(1011, "Message handling failed"); - } - }); - } - - private async handleSocketMessage(args: unknown[]): Promise { if (this.isClosed) { return; } @@ -127,15 +138,56 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { return; } + // Frames wait here while an earlier one is handled, which can take a + // while, such as during a slow initialize or while the client is behind + // on its output, and the socket keeps delivering them meanwhile. + const size = text.length + WAITING_FRAME_OVERHEAD; + const { limits } = this.options.registry; + if ( + this.waitingInboundBytes > 0 && + size > limits.maxBufferedBytes - this.waitingInboundBytes + ) { + this.closeForLimit(connectionLimitExceeded("maxBufferedBytes", limits)); + return; + } + + this.waitingInboundBytes += size; + const handled = this.messageChain.then(() => { + this.waitingInboundBytes -= size; + return this.handleSocketMessage(text); + }); + this.messageChain = handled.catch((error) => { + // When the connection shuts down, its outbound pump closes the socket. + if (this.isClosed || error instanceof ConnectionClosedError) { + return; + } + + if (error instanceof ConnectionLimitError) { + this.closeForLimit(error); + return; + } + + console.error("ACP WebSocket message handling failed:", error); + void this.shutdown(1011, "Message handling failed"); + }); + } + + private async handleSocketMessage(text: string): Promise { + if (this.isClosed) { + return; + } + let value: unknown; try { value = JSON.parse(text); } catch { + await this.waitForSocketCapacity(); this.send(protocolErrorResponse(RequestError.parseError())); return; } if (!Array.isArray(value) && !isRecord(value)) { + await this.waitForSocketCapacity(); this.send(protocolErrorResponse(RequestError.invalidRequest(value))); return; } @@ -161,6 +213,7 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { } if (!isRequestMessage(message) && !isNotificationMessage(message)) { + await this.waitForSocketCapacity(); this.send(protocolErrorResponse(RequestError.invalidRequest(message))); return; } @@ -274,11 +327,11 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { this.trackInboundRoutes(item); } - await this.writeInbound(message); - return { ok: true }; + return await this.forwardToAgent(connection, message); } if (isRequestMessage(message) && isInitializeRequest(message)) { + await this.waitForSocketCapacity(); this.send({ jsonrpc: "2.0", id: message.id, @@ -294,10 +347,55 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { } this.trackInboundRoutes(message as AnyMessage); + return await this.forwardToAgent(connection, message); + } + + private async forwardToAgent( + connection: ConnectionState, + message: AnyWireMessage, + ): Promise { + // A reply adds output, so first wait until the client has read enough to + // make room: a client that is not reading cannot make the agent produce + // more. + if (needsReply(message)) { + await connection.waitForOutboundCapacity(); + await this.waitForSocketCapacity(); + } + await this.writeInbound(message); return { ok: true }; } + /** + * Resolves once the socket holds less than `maxBufferedBytes` that it has + * not sent, or the session closes. A client that reads none of it for + * `maxOutputStallMs` is not coming back for it, so the session closes. + */ + private async waitForSocketCapacity(): Promise { + const { limits } = this.options.registry; + let pollMs = SOCKET_POLL_INITIAL_MS; + let unsentBytes = this.socket.bufferedAmount ?? 0; + let stalledMs = 0; + + while (!this.isClosed && unsentBytes >= limits.maxBufferedBytes) { + if (stalledMs >= limits.maxOutputStallMs) { + this.closeForLimit(connectionOutputStalled(limits)); + return; + } + + await new Promise((resolve) => setTimeout(resolve, pollMs)); + stalledMs += pollMs; + pollMs = Math.min(pollMs * 2, SOCKET_POLL_MAX_MS); + + const previousUnsentBytes = unsentBytes; + unsentBytes = this.socket.bufferedAmount ?? 0; + if (unsentBytes < previousUnsentBytes) { + // The client read some of it. + stalledMs = 0; + } + } + } + private trackInboundRoutes(message: unknown): void { const connection = this.connection; @@ -309,8 +407,9 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { const route = determineWebSocketRoute(message); if (route !== "connection") { - connection.ensureSession(route.session); + connection.validateSessionId(route.session); } + connection.validateRequestId(message.id); const key = messageIdKey(message.id); if (key) { @@ -325,7 +424,7 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { if (isNotificationMessage(message)) { const route = determineWebSocketRoute(message); if (route !== "connection") { - connection.ensureSession(route.session); + connection.validateSessionId(route.session); } return; } @@ -369,6 +468,16 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { } this.outboundLease = lease; + // The connection stops the lease when it shuts down, even while the pump + // waits for the client to read, so close the socket then. + lease.stopped.addEventListener( + "abort", + () => { + void this.closeForStoppedConnection(lease.stopped.reason); + }, + { once: true }, + ); + void (async () => { try { while (!this.isClosed) { @@ -381,11 +490,18 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { if (!this.send(result.value)) { return; } + + // Take nothing more from the router while the client is behind, + // which in turn makes the agent's sends wait. + await this.waitForSocketCapacity(); } } catch (error) { - if (!this.isClosed) { + // receive() fails once the connection stops the lease, and the + // connection logs why; anything else is unexpected. + if (!lease.stopped.aborted) { console.error("ACP WebSocket outbound pump failed:", error); } + await this.closeForStoppedConnection(error); } finally { if (this.outboundLease === lease) { this.outboundLease = undefined; @@ -415,6 +531,34 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { } } + private closeForLimit(error: ConnectionLimitError): void { + const connectionId = + this.connection?.connectionId ?? this.preparedConnection?.connectionId; + const label = connectionId + ? `ACP connection ${connectionId}` + : "ACP connection"; + console.warn(`Closing ${label}:`, error.message); + void this.shutdown(1008, error.message); + } + + /** + * Closes the socket for the reason its connection shut down, which the + * connection has already logged if it broke a limit. + */ + private async closeForStoppedConnection(reason: unknown): Promise { + if (this.isClosed) { + return; + } + + if (reason instanceof ConnectionLimitError) { + await this.shutdown(1008, reason.message); + } else if (reason instanceof ConnectionClosedError) { + await this.shutdown(); + } else { + await this.shutdown(1011, "Internal error"); + } + } + private async shutdownIfUninitialized( code?: number, reason?: string, @@ -490,6 +634,23 @@ function determineWebSocketRoute(message: AnyCall): ResponseRoute { return "connection"; } +/** + * Whether the agent replies to a message, so forwarding it adds output. This + * follows how the agent answers: an empty batch is an error, and a batch of + * calls gets a reply for each entry that is not a notification. + */ +function needsReply(message: unknown): boolean { + if (Array.isArray(message)) { + return ( + message.length === 0 || + (!isResponseBatch(message) && + !message.every((entry) => isNotificationMessage(entry))) + ); + } + + return !isNotificationMessage(message) && !isResponseShapedMessage(message); +} + function isDuplicateInitializeRequest(message: unknown): boolean { return isRequestMessage(message) && isInitializeRequest(message); } diff --git a/src/ws-utils.ts b/src/ws-utils.ts index a77082b2..98a12438 100644 --- a/src/ws-utils.ts +++ b/src/ws-utils.ts @@ -1,6 +1,12 @@ /** Minimal browser/Node-compatible WebSocket shape used by ACP transports. */ export interface WebSocketLike { readonly readyState?: number; + /** + * Bytes passed to `send()` that the socket has not sent yet. When a server + * socket reports it, `AcpServer` stops sending while it holds + * `maxBufferedBytes`, which pauses the agent until the client reads. + */ + readonly bufferedAmount?: number; send(data: string): void; close(code?: number, reason?: string): void; addEventListener?(type: string, listener: (event: unknown) => void): void;