diff --git a/src/connection.test.ts b/src/connection.test.ts index f163ab1..0cea622 100644 --- a/src/connection.test.ts +++ b/src/connection.test.ts @@ -1,7 +1,8 @@ import { describe, expect, it, vi } from "vitest"; import { ConnectionRegistry, - OutboundStream, + OutboundMailbox, + type OutboundLease, type ResponseRoute, } from "./connection.js"; import { messageIdKey } from "./protocol.js"; @@ -120,13 +121,17 @@ describe("ConnectionRegistry", () => { .spyOn(console, "error") .mockImplementation(() => undefined); const registry = new ConnectionRegistry(); - const connection = registry.createConnection({ - connect(stream) { - agentStream = stream; + const connection = registry.createConnection( + { + connect(stream) { + agentStream = stream; + }, }, - }); + "websocket", + ); try { - const subscription = connection.allOutbound.subscribe(); + const lease = connection.allOutbound.tryAcquire(); + expect(lease).toBeDefined(); connection.startRouter(); @@ -152,7 +157,7 @@ describe("ConnectionRegistry", () => { writer.releaseLock(); } - expect(await readNextResult(subscription.stream)).toEqual({ + expect(await lease?.receive()).toEqual({ done: true, value: undefined, }); @@ -193,9 +198,10 @@ describe("ConnectionRegistry", () => { ] as const; try { - const allOutbound = connection.allOutbound.subscribe(); - const sessionOutbound = connection.ensureSession(sessionId).subscribe(); - const connectionOutbound = connection.connectionStream.subscribe(); + const sessionOutbound = connection.ensureSession(sessionId).tryAcquire(); + const connectionOutbound = connection.connectionStream.tryAcquire(); + expect(sessionOutbound).toBeDefined(); + expect(connectionOutbound).toBeDefined(); connection.startRouter(); if (!agentStream) { @@ -209,10 +215,9 @@ describe("ConnectionRegistry", () => { writer.releaseLock(); } - expect(await readNext(allOutbound.stream)).toEqual(batch); - expect(await readNext(sessionOutbound.stream)).toEqual(batch[0]); - expect(await readNext(sessionOutbound.stream)).toEqual(batch[2]); - expect(await readNext(connectionOutbound.stream)).toEqual(batch[1]); + expect(await readLease(sessionOutbound)).toEqual(batch[0]); + expect(await readLease(sessionOutbound)).toEqual(batch[2]); + expect(await readLease(connectionOutbound)).toEqual(batch[1]); expect(connection.clientResponseRoutes).toEqual( new Map([ ["number:10", { session: sessionId }], @@ -260,6 +265,44 @@ describe("ConnectionRegistry", () => { } }); + it("drains the final agent frame before natural connection close", async () => { + const agentClosed = createDeferred(); + let agentStream: WireStream | undefined; + const registry = new ConnectionRegistry(); + const connection = registry.createConnection( + { + connect(stream) { + agentStream = stream; + return { closed: agentClosed.promise }; + }, + }, + "websocket", + ); + const lease = connection.allOutbound.tryAcquire(); + expect(lease).toBeDefined(); + connection.startRouter(); + + if (!agentStream) { + throw new Error("Expected agent stream"); + } + + const writer = agentStream.writable.getWriter(); + try { + await writer.write(messageOne); + } finally { + writer.releaseLock(); + } + agentClosed.resolve(); + await connection.closed; + + expect(await readLease(lease)).toEqual(messageOne); + expect(await lease?.receive()).toEqual({ + done: true, + value: undefined, + }); + await registry.closeAll(); + }); + it("waits for active and pending connection shutdowns before closeAll resolves", async () => { const registry = new ConnectionRegistry(); const active = registry.createConnection(createTestAgentApp()); @@ -302,14 +345,14 @@ describe("ConnectionRegistry", () => { expect(closeResolved).toBe(true); }); - it("routes pending responses to the connection stream and all outbound stream", async () => { + it("routes pending responses to the connection mailbox", async () => { const registry = new ConnectionRegistry(); const connection = registry.createConnection(createTestAgentApp()); await initializeConnection(connection); - const connectionSubscription = connection.connectionStream.subscribe(); - const allOutboundSubscription = connection.allOutbound.subscribe(); + const connectionLease = connection.connectionStream.tryAcquire(); + expect(connectionLease).toBeDefined(); const key = messageIdKey(sessionNewRequest.id); expect(key).toBe("number:2"); @@ -317,8 +360,7 @@ describe("ConnectionRegistry", () => { await writeInbound(connection.inboundTx, sessionNewRequest); - const connectionMessage = await readNext(connectionSubscription.stream); - const allOutboundMessage = await readNext(allOutboundSubscription.stream); + const connectionMessage = await readLease(connectionLease); expect(connectionMessage).toMatchObject({ jsonrpc: "2.0", @@ -327,7 +369,6 @@ describe("ConnectionRegistry", () => { sessionId: expect.stringMatching(/^[0-9a-f-]{36}$/), }, }); - expect(allOutboundMessage).toMatchObject(connectionMessage); expect(connection.pendingRoutes.has(key ?? "")).toBe(false); await registry.closeAll(); @@ -339,11 +380,12 @@ describe("ConnectionRegistry", () => { await initializeConnection(connection); - const subscription = connection.connectionStream.subscribe(); + const lease = connection.connectionStream.tryAcquire(); + expect(lease).toBeDefined(); await writeInbound(connection.inboundTx, sessionNewRequest); - expect(await readNext(subscription.stream)).toMatchObject({ + expect(await readLease(lease)).toMatchObject({ jsonrpc: "2.0", id: sessionNewRequest.id, result: { @@ -379,8 +421,10 @@ describe("ConnectionRegistry", () => { await initializeConnection(connection); - const sessionSubscription = connection.ensureSession(sessionId).subscribe(); - const connectionSubscription = connection.connectionStream.subscribe(); + const sessionLease = connection.ensureSession(sessionId).tryAcquire(); + const connectionLease = connection.connectionStream.tryAcquire(); + expect(sessionLease).toBeDefined(); + expect(connectionLease).toBeDefined(); const key = messageIdKey(promptRequest.id); expect(key).toBe("number:3"); @@ -388,7 +432,7 @@ describe("ConnectionRegistry", () => { await writeInbound(connection.inboundTx, promptRequest); - expect(await readNext(sessionSubscription.stream)).toMatchObject({ + expect(await readLease(sessionLease)).toMatchObject({ jsonrpc: "2.0", method: "session/update", params: { @@ -401,7 +445,7 @@ describe("ConnectionRegistry", () => { }, }, }); - expect(await readNext(sessionSubscription.stream)).toMatchObject({ + expect(await readLease(sessionLease)).toMatchObject({ jsonrpc: "2.0", id: promptRequest.id, result: { @@ -409,77 +453,87 @@ describe("ConnectionRegistry", () => { }, }); expect(connection.pendingRoutes.has(key ?? "")).toBe(false); - expect( - await readNextOrUndefined(connectionSubscription.stream), - ).toBeUndefined(); + expect(await readLeaseOrUndefined(connectionLease)).toBeUndefined(); await registry.closeAll(); }); }); -describe("OutboundStream", () => { - it("replays buffered messages to the first subscriber", () => { - const stream = new OutboundStream(); - - stream.push(messageOne); - stream.push(messageTwo); - - expect(stream.subscribe().replay).toEqual([messageOne, messageTwo]); - }); - - it("does not replay buffered messages to later subscribers", async () => { - const stream = new OutboundStream(); - - stream.push(messageOne); +describe("OutboundMailbox", () => { + it("buffers an ordered burst without a receiver", async () => { + const mailbox = new OutboundMailbox(); + const count = 1_025; - const first = stream.subscribe(); - const second = stream.subscribe(); - - expect(first.replay).toEqual([messageOne]); - expect(second.replay).toEqual([]); - - stream.push(messageTwo); + for (let index = 0; index < count; index++) { + mailbox.push({ + jsonrpc: "2.0", + id: index, + result: `message-${index}`, + }); + } - expect(await readNext(first.stream)).toEqual(messageTwo); - expect(await readNext(second.stream)).toEqual(messageTwo); + const lease = mailbox.tryAcquire(); + expect(lease).toBeDefined(); + for (let index = 0; index < count; index++) { + expect(await readLease(lease)).toEqual({ + jsonrpc: "2.0", + id: index, + result: `message-${index}`, + }); + } }); - it("evicts oldest replay messages when capacity is exceeded", () => { - const stream = new OutboundStream(2); - - stream.push(messageOne); - stream.push(messageTwo); - stream.push(messageThree); - - expect(stream.subscribe().replay).toEqual([messageTwo, messageThree]); + it("allows one active receiver and preserves FIFO across handoff", async () => { + const mailbox = new OutboundMailbox(); + mailbox.push(messageOne); + mailbox.push(messageTwo); + mailbox.push(messageThree); + + const first = mailbox.tryAcquire(); + expect(first).toBeDefined(); + expect(mailbox.tryAcquire()).toBeUndefined(); + expect(await readLease(first)).toEqual(messageOne); + + first?.release(); + mailbox.push(messageFour); + + const resumed = mailbox.tryAcquire(); + expect(resumed).toBeDefined(); + expect(await readLease(resumed)).toEqual(messageTwo); + expect(await readLease(resumed)).toEqual(messageThree); + expect(await readLease(resumed)).toEqual(messageFour); }); - it("drops oldest queued live messages for lagging subscribers", async () => { - const warn = vi.spyOn(console, "warn").mockImplementation(() => undefined); - const stream = new OutboundStream(2); - const subscription = stream.subscribe(); + it("drains queued messages before a finished mailbox ends", async () => { + const mailbox = new OutboundMailbox(); + const lease = mailbox.tryAcquire(); + expect(lease).toBeDefined(); - stream.push(messageOne); - stream.push(messageTwo); - stream.push(messageThree); - stream.push(messageFour); + mailbox.push(messageOne); + mailbox.push(messageTwo); + mailbox.finish(); - expect(await readNext(subscription.stream)).toEqual(messageOne); - expect(await readNext(subscription.stream)).toEqual(messageThree); - expect(await readNext(subscription.stream)).toEqual(messageFour); - expect(warn).toHaveBeenCalledOnce(); - - warn.mockRestore(); + expect(await readLease(lease)).toEqual(messageOne); + expect(await readLease(lease)).toEqual(messageTwo); + expect(await lease?.receive()).toEqual({ + done: true, + value: undefined, + }); }); - it("closes subscriber streams", async () => { - const stream = new OutboundStream(); - const reader = stream.subscribe().stream.getReader(); + it("discards queued messages on abortive shutdown", async () => { + const mailbox = new OutboundMailbox(); + const lease = mailbox.tryAcquire(); + expect(lease).toBeDefined(); - stream.close(); + mailbox.push(messageOne); + mailbox.abort(); - expect(await reader.read()).toEqual({ done: true, value: undefined }); - reader.releaseLock(); + expect(await lease?.receive()).toEqual({ + done: true, + value: undefined, + }); + expect(mailbox.tryAcquire()).toBeUndefined(); }); }); @@ -504,49 +558,32 @@ async function writeInbound( } } -async function readNext( - stream: ReadableStream, +async function readLease( + lease: OutboundLease | undefined, ): Promise { - const reader = stream.getReader(); - - try { - const result = await reader.read(); - - if (result.done) { - throw new Error("Expected stream message"); - } - - return result.value; - } finally { - reader.releaseLock(); + if (!lease) { + throw new Error("Expected outbound mailbox lease"); } -} - -async function readNextResult( - stream: ReadableStream, -): Promise> { - const reader = stream.getReader(); - try { - return await reader.read(); - } finally { - reader.releaseLock(); + const result = await lease.receive(); + if (result.done) { + throw new Error("Expected mailbox message"); } + + return result.value; } -async function readNextOrUndefined( - stream: ReadableStream, +async function readLeaseOrUndefined( + lease: OutboundLease | undefined, ): Promise { - const reader = stream.getReader(); - - try { - return await Promise.race([ - reader.read().then((result) => (result.done ? undefined : result.value)), - delay(50).then(() => undefined), - ]); - } finally { - reader.releaseLock(); + if (!lease) { + throw new Error("Expected outbound mailbox lease"); } + + return await Promise.race([ + lease.receive().then((result) => (result.done ? undefined : result.value)), + delay(50).then(() => undefined), + ]); } function createDeferred(): { diff --git a/src/connection.ts b/src/connection.ts index 621c051..708dc4a 100644 --- a/src/connection.ts +++ b/src/connection.ts @@ -26,88 +26,102 @@ export interface AgentConnector { export type ResponseRoute = "connection" | { readonly session: string }; -export interface OutboundSubscription< - Message extends AnyWireMessage = AnyMessage, -> { - readonly replay: readonly Message[]; - readonly stream: ReadableStream; +export interface OutboundLease { + receive(): Promise>; + release(): void; } -export class OutboundStream { - private readonly subscribers = new Set>(); - private replayBuffer: Message[] = []; - private hasSubscriber = false; - private isClosed = false; +export class OutboundMailbox { + private readonly queue: Message[] = []; + private activeLease: MailboxLease | undefined; + private isFinished = false; + private isAborted = false; - constructor(private readonly capacity = 1024) {} + constructor(private readonly enabled = true) {} push(message: Message): void { - if (this.isClosed) { + if (!this.enabled || this.isFinished) { return; } - if (!this.hasSubscriber) { - this.replayBuffer.push(message); + this.queue.push(message); + this.activeLease?.wake(); + } - if (this.replayBuffer.length > this.capacity) { - this.replayBuffer.shift(); - } + tryAcquire(): OutboundLease | undefined { + if (!this.enabled || this.isAborted || this.activeLease) { + return undefined; + } + const lease = new MailboxLease(this); + this.activeLease = lease; + return lease; + } + + finish(): void { + if (this.isFinished) { return; } - for (const subscriber of this.subscribers) { - subscriber.push(message); + this.isFinished = true; + this.activeLease?.wake(); + } + + abort(): void { + if (this.isAborted) { + return; } + + this.isAborted = true; + this.isFinished = true; + this.queue.length = 0; + this.activeLease?.wake(); } - subscribe(): OutboundSubscription { - const replay = this.hasSubscriber ? [] : [...this.replayBuffer]; - this.replayBuffer = []; - this.hasSubscriber = true; + /** @internal */ + async receive( + lease: MailboxLease, + ): Promise> { + for (;;) { + if (this.isAborted || lease.released || this.activeLease !== lease) { + return { done: true, value: undefined }; + } - const subscriber = new OutboundSubscriber( - this.capacity, - (item) => { - this.subscribers.delete(item); - }, - ); + if (this.queue.length > 0) { + return { + done: false, + value: this.queue.shift() as Message, + }; + } - this.subscribers.add(subscriber); + if (this.isFinished) { + return { done: true, value: undefined }; + } - if (this.isClosed) { - subscriber.close(); + await lease.wait(); } - - return { - replay, - stream: subscriber.stream, - }; } - close(): void { - if (this.isClosed) { + /** @internal */ + release(lease: MailboxLease): void { + if (this.activeLease !== lease) { return; } - this.isClosed = true; - this.replayBuffer = []; - - for (const subscriber of this.subscribers) { - subscriber.close(); - } - - this.subscribers.clear(); + this.activeLease = undefined; + lease.markReleased(); } } +export type ConnectionTransport = "http" | "websocket"; + export class ConnectionState { readonly connectionId: string; readonly inboundTx: WritableStream; readonly outboundRx: ReadableStream; - readonly connectionStream = new OutboundStream(); - readonly allOutbound = new OutboundStream(); - readonly sessionStreams = new Map(); + readonly connectionStream: OutboundMailbox; + readonly allOutbound: OutboundMailbox; + readonly sessionStreams = new Map(); readonly pendingRoutes = new Map(); readonly clientResponseRoutes = new Map(); readonly closed: Promise; @@ -120,29 +134,36 @@ export class ConnectionState { ReadableStreamDefaultReader | undefined; private outboundReader: ReadableStreamDefaultReader | undefined; + private routerPromise: Promise | undefined; private shutdownPromise: Promise | undefined; + private hasResolvedClosed = false; + private finishAgentOutbound: () => void = () => {}; private resolveClosed: () => void = () => {}; - constructor(agent: AgentConnector) { + constructor( + agent: AgentConnector, + private readonly transport: ConnectionTransport = "http", + ) { this.connectionId = globalThis.crypto.randomUUID(); + this.connectionStream = new OutboundMailbox(transport === "http"); + this.allOutbound = new OutboundMailbox( + transport === "websocket", + ); this.closed = new Promise((resolve) => { this.resolveClosed = resolve; }); const inbound = new TransformStream(); - const outbound = new TransformStream({ - transform: (message, controller) => { - if (!this.supportsBatches && Array.isArray(message)) { - throw new TypeError( - "AcpServer transports do not support outbound JSON-RPC batch messages", - ); - } - - controller.enqueue(message); - }, + const outbound = createBufferedOutboundChannel((message) => { + if (!this.supportsBatches && Array.isArray(message)) { + throw new TypeError( + "AcpServer transports do not support outbound JSON-RPC batch messages", + ); + } }); this.inboundTx = inbound.writable; this.outboundRx = outbound.readable; + this.finishAgentOutbound = outbound.finish; const stream: WireStream = { readable: inbound.readable, @@ -204,7 +225,7 @@ export class ConnectionState { } this.hasStartedRouter = true; - void this.runRouter(); + this.routerPromise = this.runRouter(); } startConnectHandlers(): void { @@ -227,13 +248,13 @@ export class ConnectionState { return this.supportsBatches; } - ensureSession(sessionId: string): OutboundStream { + ensureSession(sessionId: string): OutboundMailbox { const existing = this.sessionStreams.get(sessionId); if (existing) { return existing; } - const stream = new OutboundStream(); + const stream = new OutboundMailbox(this.transport === "http"); this.sessionStreams.set(sessionId, stream); return stream; @@ -249,11 +270,11 @@ export class ConnectionState { private async runShutdown(): Promise { try { - this.connectionStream.close(); - this.allOutbound.close(); + this.connectionStream.abort(); + this.allOutbound.abort(); for (const stream of this.sessionStreams.values()) { - stream.close(); + stream.abort(); } this.sessionStreams.clear(); @@ -265,7 +286,7 @@ export class ConnectionState { this.cancelOutboundReader(), ]); } finally { - this.resolveClosed(); + this.resolveClosedOnce(); } } @@ -279,8 +300,14 @@ export class ConnectionState { return; } - void Promise.resolve(this.agentConnection.closed).finally(() => { - void this.shutdown(); + void Promise.resolve(this.agentConnection.closed).finally(async () => { + if (!this.hasStartedRouter) { + await this.shutdown(); + return; + } + + this.finishAgentOutbound(); + await this.routerPromise; }); } @@ -325,15 +352,26 @@ export class ConnectionState { } reader.releaseLock(); - this.connectionStream.close(); - this.allOutbound.close(); + this.connectionStream.finish(); + this.allOutbound.finish(); for (const stream of this.sessionStreams.values()) { - stream.close(); + stream.finish(); } + + this.resolveClosedOnce(); } } + private resolveClosedOnce(): void { + if (this.hasResolvedClosed) { + return; + } + + this.hasResolvedClosed = true; + this.resolveClosed(); + } + private routeOutbound(message: AnyWireMessage): void { this.allOutbound.push(message); @@ -412,15 +450,21 @@ export class ConnectionRegistry { private readonly connections = new Map(); private readonly pendingConnections = new Map(); - createConnection(agent: AgentConnector): ConnectionState { - const connection = new ConnectionState(agent); + createConnection( + agent: AgentConnector, + transport: ConnectionTransport = "http", + ): ConnectionState { + const connection = new ConnectionState(agent, transport); this.connections.set(connection.connectionId, connection); this.trackConnectionClose(connection); return connection; } - createPendingConnection(agent: AgentConnector): ConnectionState { - const connection = new ConnectionState(agent); + createPendingConnection( + agent: AgentConnector, + transport: ConnectionTransport = "websocket", + ): ConnectionState { + const connection = new ConnectionState(agent, transport); this.pendingConnections.set(connection.connectionId, connection); this.trackConnectionClose(connection); return connection; @@ -487,92 +531,121 @@ export class ConnectionRegistry { } } -class OutboundSubscriber { - readonly stream: ReadableStream; +class MailboxLease< + Message extends AnyWireMessage, +> implements OutboundLease { + released = false; - private controller: ReadableStreamDefaultController | undefined; - private queue: Message[] = []; - private isClosed = false; - private hasWarnedAboutOverflow = false; + private receiving = false; + private wakePromise: Promise | undefined; + private resolveWake: (() => void) | undefined; - constructor( - private readonly capacity: number, - private readonly onCancel: ( - subscriber: OutboundSubscriber, - ) => void, - ) { - this.stream = new ReadableStream({ - start: (controller) => { - this.controller = controller; - this.flush(); - }, - pull: () => { - this.flush(); - }, - cancel: () => { - this.cancel(); - }, - }); - } + constructor(private readonly mailbox: OutboundMailbox) {} - push(message: Message): void { - if (this.isClosed) { - return; + async receive(): Promise> { + if (this.receiving) { + throw new Error( + "ACP outbound mailbox lease already has a pending receive", + ); } - this.queue.push(message); - - if (this.queue.length > this.capacity) { - this.queue.shift(); - - if (!this.hasWarnedAboutOverflow) { - console.warn("ACP outbound subscriber lagged; dropping oldest message"); - this.hasWarnedAboutOverflow = true; - } + this.receiving = true; + try { + return await this.mailbox.receive(this); + } finally { + this.receiving = false; } + } - this.flush(); + release(): void { + this.mailbox.release(this); } - close(): void { - if (this.isClosed) { - return; + wake(): void { + this.resolveWake?.(); + this.wakePromise = undefined; + this.resolveWake = undefined; + } + + wait(): Promise { + if (!this.wakePromise) { + this.wakePromise = new Promise((resolve) => { + this.resolveWake = resolve; + }); } - this.isClosed = true; - this.queue = []; - this.controller?.close(); + return this.wakePromise; } - private cancel(): void { - this.isClosed = true; - this.queue = []; - this.onCancel(this); + markReleased(): void { + this.released = true; + this.wake(); } +} - private flush(): void { - if (!this.controller) { +function createBufferedOutboundChannel( + validate: (message: AnyWireMessage) => void, +): { + readonly readable: ReadableStream; + readonly writable: WritableStream; + readonly finish: () => void; +} { + let controller: ReadableStreamDefaultController | undefined; + let isFinished = false; + + const finish = (): void => { + if (isFinished) { return; } - while ( - this.queue.length > 0 && - this.controller.desiredSize !== null && - this.controller.desiredSize > 0 - ) { - const message = this.queue.shift(); - - if (!message) { - return; - } - - this.controller.enqueue(message); + isFinished = true; + try { + controller?.close(); + } catch { + // The router may already have cancelled the readable side. + } + }; + const fail = (error: unknown): void => { + if (isFinished) { + return; } - if (this.queue.length === 0) { - this.hasWarnedAboutOverflow = false; + isFinished = true; + try { + controller?.error(error); + } catch { + // The router may already have cancelled the readable side. } - } + }; + + return { + readable: new ReadableStream({ + start(readableController) { + controller = readableController; + }, + cancel() { + isFinished = true; + }, + }), + writable: new WritableStream({ + write(message) { + if (isFinished) { + throw new Error("ACP outbound channel is closed"); + } + + try { + validate(message); + controller?.enqueue(message); + } catch (error) { + fail(error); + throw error; + } + }, + close: finish, + abort: fail, + }), + finish, + }; } function isMatchingResponse( diff --git a/src/protocol.test.ts b/src/protocol.test.ts index 1192a43..18ccec7 100644 --- a/src/protocol.test.ts +++ b/src/protocol.test.ts @@ -27,6 +27,10 @@ describe("protocol transport helpers", () => { true, ); expect(methodRequiresSessionHeader(AGENT_METHODS.session_close)).toBe(true); + expect(methodRequiresSessionHeader(AGENT_METHODS.session_delete)).toBe( + true, + ); + expect(methodRequiresSessionHeader(AGENT_METHODS.session_fork)).toBe(true); expect(methodRequiresSessionHeader(AGENT_METHODS.session_load)).toBe(true); expect(methodRequiresSessionHeader(AGENT_METHODS.session_prompt)).toBe( true, @@ -40,16 +44,33 @@ describe("protocol transport helpers", () => { expect(methodRequiresSessionHeader(AGENT_METHODS.session_set_mode)).toBe( true, ); + expect(methodRequiresSessionHeader("session/set_model")).toBe(true); + expect(methodRequiresSessionHeader(AGENT_METHODS.nes_suggest)).toBe(true); + expect(methodRequiresSessionHeader(AGENT_METHODS.nes_accept)).toBe(true); + expect(methodRequiresSessionHeader(AGENT_METHODS.nes_reject)).toBe(true); + expect(methodRequiresSessionHeader(AGENT_METHODS.nes_close)).toBe(true); + expect(methodRequiresSessionHeader(AGENT_METHODS.document_did_open)).toBe( + true, + ); + expect(methodRequiresSessionHeader(AGENT_METHODS.document_did_change)).toBe( + true, + ); + expect(methodRequiresSessionHeader(AGENT_METHODS.document_did_close)).toBe( + true, + ); + expect(methodRequiresSessionHeader(AGENT_METHODS.document_did_save)).toBe( + true, + ); + expect(methodRequiresSessionHeader(AGENT_METHODS.document_did_focus)).toBe( + true, + ); }); - it("does not require a session header for connection-level or unsupported methods", () => { + it("does not require a session header for connection-level methods", () => { expect(methodRequiresSessionHeader(AGENT_METHODS.initialize)).toBe(false); expect(methodRequiresSessionHeader(AGENT_METHODS.session_new)).toBe(false); expect(methodRequiresSessionHeader(AGENT_METHODS.session_list)).toBe(false); - expect(methodRequiresSessionHeader(AGENT_METHODS.session_fork)).toBe(false); expect(methodRequiresSessionHeader(AGENT_METHODS.nes_start)).toBe(false); - expect(methodRequiresSessionHeader(AGENT_METHODS.nes_suggest)).toBe(false); - expect(methodRequiresSessionHeader(AGENT_METHODS.nes_close)).toBe(false); }); it("extracts a top-level string session ID from params", () => { diff --git a/src/protocol.ts b/src/protocol.ts index 18116b7..54ee051 100644 --- a/src/protocol.ts +++ b/src/protocol.ts @@ -10,11 +10,23 @@ export const JSON_MIME_TYPE = "application/json"; const SESSION_SCOPED_METHODS = new Set([ AGENT_METHODS.session_cancel, AGENT_METHODS.session_close, + AGENT_METHODS.session_delete, + AGENT_METHODS.session_fork, AGENT_METHODS.session_load, AGENT_METHODS.session_prompt, AGENT_METHODS.session_resume, AGENT_METHODS.session_set_config_option, AGENT_METHODS.session_set_mode, + "session/set_model", + AGENT_METHODS.nes_suggest, + AGENT_METHODS.nes_accept, + AGENT_METHODS.nes_reject, + AGENT_METHODS.nes_close, + AGENT_METHODS.document_did_open, + AGENT_METHODS.document_did_change, + AGENT_METHODS.document_did_close, + AGENT_METHODS.document_did_save, + AGENT_METHODS.document_did_focus, ]); export function methodRequiresSessionHeader(method: string): boolean { diff --git a/src/server-session-sse.test.ts b/src/server-session-sse.test.ts index e1af341..c469d70 100644 --- a/src/server-session-sse.test.ts +++ b/src/server-session-sse.test.ts @@ -322,7 +322,47 @@ describe("AcpServer session SSE", () => { } }); - it("routes non-required session methods using params.sessionId when the session header is absent", async () => { + it.each([ + "session/delete", + "session/fork", + "nes/suggest", + "nes/accept", + "nes/reject", + "nes/close", + "document/didOpen", + "document/didChange", + "document/didClose", + "document/didSave", + "document/didFocus", + "_vendor/session-method", + ])( + "rejects %s with params.sessionId but no session header", + async (method) => { + const server = await startTestServer(); + + try { + const connectionId = await initialize(server.url); + const response = await postJson( + server.url, + { + jsonrpc: "2.0", + id: 30, + method, + params: { sessionId: "session-1" }, + }, + { + [HEADER_CONNECTION_ID]: connectionId, + }, + ); + + expect(response.status).toBe(400); + } finally { + await server.close(); + } + }, + ); + + it("rejects session/fork without a session header", async () => { const server = await startTestServer(); try { @@ -336,6 +376,27 @@ describe("AcpServer session SSE", () => { }, ); + expect(response.status).toBe(400); + } finally { + await server.close(); + } + }); + + it("accepts session/fork with a matching session header", async () => { + const server = await startTestServer(); + + try { + const connectionId = await initialize(server.url); + const sessionId = await createSession(server.url, connectionId); + const response = await postJson( + server.url, + createForkRequest(3, sessionId), + { + [HEADER_CONNECTION_ID]: connectionId, + [HEADER_SESSION_ID]: sessionId, + }, + ); + expect(response.status).toBe(202); } finally { await server.close(); diff --git a/src/server-sse.test.ts b/src/server-sse.test.ts index 55098ef..4bf5d28 100644 --- a/src/server-sse.test.ts +++ b/src/server-sse.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from "vitest"; +import { OutboundMailbox } from "./connection.js"; import { createSseBodySource } from "./server-sse.js"; import { serializeSseEvent } from "./sse.js"; @@ -15,12 +16,12 @@ const message = { describe("createSseBodySource", () => { it("enqueues a subscription message after it has been read even if demand changed", async () => { - const backend = new TransformStream(); - const writer = backend.writable.getWriter(); - const source = createSseBodySource({ - replay: [], - stream: backend.readable, - }); + const mailbox = new OutboundMailbox(); + const lease = mailbox.tryAcquire(); + if (!lease) { + throw new Error("Expected outbound mailbox lease"); + } + const source = createSseBodySource(lease); const enqueued: Uint8Array[] = []; let desiredSize: number | null = 1; const controller = { @@ -44,10 +45,10 @@ describe("createSseBodySource", () => { await flushMicrotasks(); desiredSize = 0; - await Promise.all([pull, writer.write(message)]); + mailbox.push(message); + await pull; expect(enqueued.map(decodeText)).toEqual([serializeSseEvent(message)]); - writer.releaseLock(); }); }); diff --git a/src/server-sse.ts b/src/server-sse.ts index 09e8a3a..256eec0 100644 --- a/src/server-sse.ts +++ b/src/server-sse.ts @@ -1,22 +1,20 @@ import { serializeSseEvent, serializeSseKeepAlive } from "./sse.js"; -import type { OutboundSubscription } from "./connection.js"; -import type { AnyMessage } from "./jsonrpc.js"; +import type { OutboundLease } from "./connection.js"; export function createSseBody( - subscription: OutboundSubscription, + lease: OutboundLease, ): ReadableStream { - return new ReadableStream(createSseBodySource(subscription)); + return new ReadableStream(createSseBodySource(lease)); } /** @internal */ export function createSseBodySource( - subscription: OutboundSubscription, + lease: OutboundLease, ): UnderlyingDefaultSource { const encoder = new TextEncoder(); - const replay = [...subscription.replay]; let keepAliveTimer: ReturnType | undefined; - let reader: ReadableStreamDefaultReader | undefined; + let isReceiving = false; let isClosed = false; const clearKeepAlive = (): void => { @@ -51,6 +49,7 @@ export function createSseBodySource( isClosed = true; clearKeepAlive(); + lease.release(); try { controller.close(); @@ -72,24 +71,14 @@ export function createSseBodySource( }, 15_000); }, async pull(controller) { - if (isClosed || reader || !hasDemand(controller)) { + if (isClosed || isReceiving || !hasDemand(controller)) { return; } - const replayMessage = replay.shift(); - if (replayMessage) { - if (!enqueueText(controller, serializeSseEvent(replayMessage))) { - closeBody(controller); - } - - return; - } - - const currentReader = subscription.stream.getReader(); - reader = currentReader; + isReceiving = true; try { - const result = await currentReader.read(); + const result = await lease.receive(); if (isClosed) { return; @@ -110,22 +99,13 @@ export function createSseBodySource( controller.error(error); } } finally { - if (reader === currentReader) { - reader = undefined; - } - - currentReader.releaseLock(); + isReceiving = false; } }, cancel() { isClosed = true; clearKeepAlive(); - - if (reader) { - void reader.cancel(); - } else { - void subscription.stream.cancel(); - } + lease.release(); }, }; } diff --git a/src/server.test.ts b/src/server.test.ts index cb41ec8..8c471fd 100644 --- a/src/server.test.ts +++ b/src/server.test.ts @@ -683,8 +683,66 @@ describe("AcpServer", () => { } }); - it("does not drain outbound subscriptions faster than SSE body demand", async () => { - const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + it("allows one active receiver per route and preserves queued messages across handoff", async () => { + const server = new AcpServer({ + createAgent: () => createTestAgentApp(), + }); + const open = (connectionId: string, sessionId?: string) => + server.handleRequest( + new Request("http://127.0.0.1/acp", { + method: "GET", + headers: { + Accept: EVENT_STREAM_MIME_TYPE, + [HEADER_CONNECTION_ID]: connectionId, + ...(sessionId ? { [HEADER_SESSION_ID]: sessionId } : {}), + }, + }), + ); + const bodies: Array | null> = []; + + try { + const connectionId = await initializeDirect(server); + const connectionStream = await open(connectionId); + bodies.push(connectionStream.body); + expect(connectionStream.status).toBe(200); + expect((await open(connectionId)).status).toBe(409); + + const firstSession = await open(connectionId, "session-1"); + bodies.push(firstSession.body); + expect(firstSession.status).toBe(200); + expect((await open(connectionId, "session-1")).status).toBe(409); + + const secondSession = await open(connectionId, "session-2"); + bodies.push(secondSession.body); + expect(secondSession.status).toBe(200); + + await connectionStream.body?.cancel(); + const accepted = await server.handleRequest( + jsonRequest(sessionNewRequest, { + [HEADER_CONNECTION_ID]: connectionId, + }), + ); + expect(accepted.status).toBe(202); + + const resumed = await open(connectionId); + bodies.push(resumed.body); + expect(resumed.status).toBe(200); + expect(await readFirstSseMessage(resumed)).toMatchObject({ + jsonrpc: "2.0", + id: sessionNewRequest.id, + result: { + sessionId: expect.stringMatching(/^[0-9a-f-]{36}$/), + }, + }); + } finally { + await Promise.all( + bodies.map((body) => body?.cancel().catch(() => undefined)), + ); + await server.close(); + } + }); + + it("preserves an outbound burst for a slow SSE receiver", async () => { let resolvePromptDone: () => void = () => {}; const promptDone = new Promise((resolve) => { resolvePromptDone = resolve; @@ -725,22 +783,27 @@ describe("AcpServer", () => { expect(accepted.status).toBe(202); expect(chunkText(await firstMessage)).toBe("chunk-1"); await promptDone; - await waitFor(() => warnSpy.mock.calls.length > 0); const chunkNumbers = [1]; - for (let index = 0; index < 128; index++) { + for (let index = 1; index < 1_100; index++) { const chunk = chunkText(await iterator.next()); chunkNumbers.push(Number(chunk.slice("chunk-".length))); } - expect(hasGap(chunkNumbers)).toBe(true); - expect(warnSpy).toHaveBeenCalledWith( - "ACP outbound subscriber lagged; dropping oldest message", + expect(chunkNumbers).toEqual( + Array.from({ length: 1_100 }, (_unused, index) => index + 1), ); + expect(await iterator.next()).toMatchObject({ + done: false, + value: { + jsonrpc: "2.0", + id: promptRequest.id, + result: { stopReason: "end_turn" }, + }, + }); } finally { await iterator?.return?.(); await body?.cancel().catch(() => undefined); - warnSpy.mockRestore(); await server.close(); } }); @@ -1229,16 +1292,6 @@ function chunkText(result: IteratorResult): string { return text; } -function hasGap(values: readonly number[]): boolean { - return values.some((value, index) => { - if (index === 0) { - return false; - } - - return value !== values[index - 1] + 1; - }); -} - function createBackpressureAgent(onPromptDone: () => void) { return createAgentApp({ name: "backpressure-agent" }) .onRequest(methods.agent.initialize, () => ({ @@ -1269,18 +1322,6 @@ function createBackpressureAgent(onPromptDone: () => void) { .onNotification(methods.agent.session.cancel, () => {}); } -async function waitFor(callback: () => boolean): Promise { - const deadline = Date.now() + 1_000; - - while (!callback()) { - if (Date.now() > deadline) { - throw new Error("Timed out waiting for condition"); - } - - await new Promise((resolve) => setTimeout(resolve, 1)); - } -} - async function waitForConnectionNotFound( server: AcpServer, connectionId: string, diff --git a/src/server.ts b/src/server.ts index b54f814..05eb907 100644 --- a/src/server.ts +++ b/src/server.ts @@ -22,7 +22,7 @@ import type { import type { AgentConnector, ConnectionState, - OutboundSubscription, + OutboundLease, ResponseRoute, } from "./connection.js"; import type { @@ -265,11 +265,19 @@ export class AcpServer { } const sessionId = req.headers.get(HEADER_SESSION_ID); - if (sessionId) { - return sseResponse(connection.ensureSession(sessionId).subscribe()); + const mailbox = sessionId + ? connection.ensureSession(sessionId) + : connection.connectionStream; + const lease = mailbox.tryAcquire(); + + if (!lease) { + return textResponse( + "Outbound stream already has an active receiver", + 409, + ); } - return sseResponse(connection.connectionStream.subscribe()); + return sseResponse(lease); } private handleDelete(req: Request): Response { @@ -588,7 +596,11 @@ function determineRoute( const headerSessionId = headers.get(HEADER_SESSION_ID); const paramsSessionId = sessionIdFromParams(message.params); - if (methodRequiresSessionHeader(message.method) && !headerSessionId) { + if ( + (methodRequiresSessionHeader(message.method) || + paramsSessionId !== undefined) && + !headerSessionId + ) { return { ok: false, status: 400, @@ -632,8 +644,8 @@ function isJsonContentType(contentType: string | null): boolean { return contentType?.split(";", 1)[0]?.trim().toLowerCase() === JSON_MIME_TYPE; } -function sseResponse(subscription: OutboundSubscription): Response { - return new Response(createSseBody(subscription), { +function sseResponse(lease: OutboundLease): Response { + return new Response(createSseBody(lease), { status: 200, headers: { "Content-Type": EVENT_STREAM_MIME_TYPE, diff --git a/src/ws-server.ts b/src/ws-server.ts index 4f3397a..0d4efa5 100644 --- a/src/ws-server.ts +++ b/src/ws-server.ts @@ -15,6 +15,7 @@ import type { AgentConnector, ConnectionRegistry, ConnectionState, + OutboundLease, ResponseRoute, } from "./connection.js"; import type { AnyCall, AnyMessage, AnyWireMessage } from "./jsonrpc.js"; @@ -55,8 +56,7 @@ export function handleWebSocketConnection( class WebSocketServerSession implements WebSocketServerSessionHandle { private connection: ConnectionState | undefined; private preparedConnection: ConnectionState | undefined; - private outboundReader: - ReadableStreamDefaultReader | undefined; + private outboundLease: OutboundLease | undefined; private inboundWriteChain: Promise = Promise.resolve(); private messageChain: Promise = Promise.resolve(); private isClosed = false; @@ -353,20 +353,20 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { } private startOutboundPump(connection: ConnectionState): void { - const subscription = connection.allOutbound.subscribe(); - const reader = subscription.stream.getReader(); - this.outboundReader = reader; + const lease = connection.allOutbound.tryAcquire(); + if (!lease) { + void this.shutdown( + 1011, + "Outbound stream already has an active receiver", + ); + return; + } + this.outboundLease = lease; void (async () => { try { - for (const message of subscription.replay) { - if (!this.send(message)) { - return; - } - } - while (!this.isClosed) { - const result = await reader.read(); + const result = await lease.receive(); if (result.done) { return; @@ -381,11 +381,11 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { console.error("ACP WebSocket outbound pump failed:", error); } } finally { - if (this.outboundReader === reader) { - this.outboundReader = undefined; + if (this.outboundLease === lease) { + this.outboundLease = undefined; } - reader.releaseLock(); + lease.release(); if (!this.isClosed) { void this.shutdown(); @@ -445,12 +445,10 @@ class WebSocketServerSession implements WebSocketServerSessionHandle { detach(); } - const outboundReader = this.outboundReader; - this.outboundReader = undefined; + const outboundLease = this.outboundLease; + this.outboundLease = undefined; - if (outboundReader) { - await outboundReader.cancel(); - } + outboundLease?.release(); if (this.connection) { this.options.registry.discard(this.connection.connectionId);