diff --git a/src/http-stream.test.ts b/src/http-stream.test.ts index 93c7a5e..c96de71 100644 --- a/src/http-stream.test.ts +++ b/src/http-stream.test.ts @@ -830,6 +830,57 @@ describe("createHttpStream", () => { } }); + it("aborts active SSE requests before awaiting DELETE on close", async () => { + const deleteStarted = createDeferred(); + const allowDelete = createDeferred(); + let sseSignal: AbortSignal | null | undefined; + const fetch: typeof globalThis.fetch = async (_input, init) => { + const method = init?.method ?? "GET"; + const headers = new Headers(init?.headers); + + if (method === "POST" && !headers.has(HEADER_CONNECTION_ID)) { + return jsonResponse(initializeResponse, 200, { + [HEADER_CONNECTION_ID]: "connection-1", + }); + } + + if (method === "GET") { + sseSignal = init?.signal; + return sseResponse(sseSignal); + } + + if (method === "DELETE") { + deleteStarted.resolve(); + await allowDelete.promise; + return new Response(null, { status: 202 }); + } + + throw new Error(`Unexpected ${method} request`); + }; + const stream = createHttpStream("https://agent.example/acp", { fetch }); + const writer = stream.writable.getWriter(); + const reader = stream.readable.getReader(); + let closePromise: Promise | undefined; + + try { + await writer.write(initializeRequest); + await readMessage(reader); + + closePromise = writer.close(); + await deleteStarted.promise; + + expect(sseSignal?.aborted).toBe(true); + allowDelete.resolve(); + await closePromise; + } finally { + allowDelete.resolve(); + reader.releaseLock(); + writer.releaseLock(); + await closePromise?.catch(() => undefined); + await closeStream(stream); + } + }); + it("aborts in-flight connected POSTs when closed", async () => { const connectedPostAborted = createDeferred(); const requests: RecordedRequest[] = []; diff --git a/src/http-stream.ts b/src/http-stream.ts index 3a89da7..5da8265 100644 --- a/src/http-stream.ts +++ b/src/http-stream.ts @@ -472,11 +472,11 @@ class HttpStreamTransport { } this.isClosed = true; + this.abortController.abort(); try { await this.deleteConnection(); } finally { - this.abortController.abort(); this.clearOwnedCookieStore(); this.closeReadable(); } diff --git a/src/server-sse.test.ts b/src/server-sse.test.ts new file mode 100644 index 0000000..55098ef --- /dev/null +++ b/src/server-sse.test.ts @@ -0,0 +1,61 @@ +import { describe, expect, it } from "vitest"; + +import { createSseBodySource } from "./server-sse.js"; +import { serializeSseEvent } from "./sse.js"; + +import type { AnyMessage } from "./jsonrpc.js"; + +const message = { + jsonrpc: "2.0", + id: 1, + result: { + ok: true, + }, +} satisfies AnyMessage; + +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 enqueued: Uint8Array[] = []; + let desiredSize: number | null = 1; + const controller = { + get desiredSize() { + return desiredSize; + }, + enqueue(chunk: Uint8Array) { + enqueued.push(chunk); + }, + close() { + return undefined; + }, + error(error?: unknown) { + if (error) { + throw error; + } + }, + } as ReadableStreamDefaultController; + + const pull = Promise.resolve(source.pull?.(controller)); + await flushMicrotasks(); + desiredSize = 0; + + await Promise.all([pull, writer.write(message)]); + + expect(enqueued.map(decodeText)).toEqual([serializeSseEvent(message)]); + writer.releaseLock(); + }); +}); + +async function flushMicrotasks(): Promise { + await Promise.resolve(); + await Promise.resolve(); +} + +function decodeText(chunk: Uint8Array): string { + return new TextDecoder().decode(chunk); +} diff --git a/src/server-sse.ts b/src/server-sse.ts new file mode 100644 index 0000000..09e8a3a --- /dev/null +++ b/src/server-sse.ts @@ -0,0 +1,131 @@ +import { serializeSseEvent, serializeSseKeepAlive } from "./sse.js"; + +import type { OutboundSubscription } from "./connection.js"; +import type { AnyMessage } from "./jsonrpc.js"; + +export function createSseBody( + subscription: OutboundSubscription, +): ReadableStream { + return new ReadableStream(createSseBodySource(subscription)); +} + +/** @internal */ +export function createSseBodySource( + subscription: OutboundSubscription, +): UnderlyingDefaultSource { + const encoder = new TextEncoder(); + const replay = [...subscription.replay]; + let keepAliveTimer: ReturnType | undefined; + let reader: ReadableStreamDefaultReader | undefined; + let isClosed = false; + + const clearKeepAlive = (): void => { + if (keepAliveTimer) { + clearInterval(keepAliveTimer); + keepAliveTimer = undefined; + } + }; + + const enqueueText = ( + controller: ReadableStreamDefaultController, + text: string, + ): boolean => { + try { + controller.enqueue(encoder.encode(text)); + return true; + } catch { + return false; + } + }; + + const hasDemand = ( + controller: ReadableStreamDefaultController, + ): boolean => controller.desiredSize !== null && controller.desiredSize > 0; + + const closeBody = ( + controller: ReadableStreamDefaultController, + ): void => { + if (isClosed) { + return; + } + + isClosed = true; + clearKeepAlive(); + + try { + controller.close(); + } catch { + // Stream may already be cancelled by the consumer. + } + }; + + return { + start(controller) { + keepAliveTimer = setInterval(() => { + if (isClosed || !hasDemand(controller)) { + return; + } + + if (!enqueueText(controller, serializeSseKeepAlive())) { + closeBody(controller); + } + }, 15_000); + }, + async pull(controller) { + if (isClosed || reader || !hasDemand(controller)) { + return; + } + + const replayMessage = replay.shift(); + if (replayMessage) { + if (!enqueueText(controller, serializeSseEvent(replayMessage))) { + closeBody(controller); + } + + return; + } + + const currentReader = subscription.stream.getReader(); + reader = currentReader; + + try { + const result = await currentReader.read(); + + if (isClosed) { + return; + } + + if (result.done) { + closeBody(controller); + return; + } + + if (!enqueueText(controller, serializeSseEvent(result.value))) { + closeBody(controller); + } + } catch (error) { + if (!isClosed) { + isClosed = true; + clearKeepAlive(); + controller.error(error); + } + } finally { + if (reader === currentReader) { + reader = undefined; + } + + currentReader.releaseLock(); + } + }, + cancel() { + isClosed = true; + clearKeepAlive(); + + if (reader) { + void reader.cancel(); + } else { + void subscription.stream.cancel(); + } + }, + }; +} diff --git a/src/server.ts b/src/server.ts index 95443aa..3d7a175 100644 --- a/src/server.ts +++ b/src/server.ts @@ -11,7 +11,7 @@ import { } from "./protocol.js"; import { isJsonRpcMessage, isResponseMessage } from "./jsonrpc.js"; import { AGENT_METHODS } from "./schema/index.js"; -import { serializeSseEvent, serializeSseKeepAlive } from "./sse.js"; +import { createSseBody } from "./server-sse.js"; import { handleWebSocketConnection } from "./ws-server.js"; import { AgentSideConnection } from "./acp.js"; import type { @@ -639,143 +639,6 @@ function sseResponse(subscription: OutboundSubscription): Response { }); } -function createSseBody( - subscription: OutboundSubscription, -): ReadableStream { - const encoder = new TextEncoder(); - const replay = [...subscription.replay]; - let keepAliveTimer: ReturnType | undefined; - let reader: ReadableStreamDefaultReader | undefined; - let pendingMessage: AnyMessage | undefined; - let isClosed = false; - - const clearKeepAlive = (): void => { - if (keepAliveTimer) { - clearInterval(keepAliveTimer); - keepAliveTimer = undefined; - } - }; - - const enqueueText = ( - controller: ReadableStreamDefaultController, - text: string, - ): boolean => { - try { - controller.enqueue(encoder.encode(text)); - return true; - } catch { - return false; - } - }; - - const hasDemand = ( - controller: ReadableStreamDefaultController, - ): boolean => controller.desiredSize !== null && controller.desiredSize > 0; - - const closeBody = ( - controller: ReadableStreamDefaultController, - ): void => { - if (isClosed) { - return; - } - - isClosed = true; - clearKeepAlive(); - - try { - controller.close(); - } catch { - // Stream may already be cancelled by the consumer. - } - }; - - return new ReadableStream({ - start(controller) { - keepAliveTimer = setInterval(() => { - if (isClosed || pendingMessage || !hasDemand(controller)) { - return; - } - - if (!enqueueText(controller, serializeSseKeepAlive())) { - closeBody(controller); - } - }, 15_000); - }, - async pull(controller) { - if (isClosed || reader || !hasDemand(controller)) { - return; - } - - const replayMessage = replay.shift(); - if (replayMessage) { - if (!enqueueText(controller, serializeSseEvent(replayMessage))) { - closeBody(controller); - } - - return; - } - - if (pendingMessage) { - const message = pendingMessage; - pendingMessage = undefined; - - if (!enqueueText(controller, serializeSseEvent(message))) { - closeBody(controller); - } - - return; - } - - const currentReader = subscription.stream.getReader(); - reader = currentReader; - - try { - const result = await currentReader.read(); - - if (isClosed) { - return; - } - - if (result.done) { - closeBody(controller); - return; - } - - if (!hasDemand(controller)) { - pendingMessage = result.value; - return; - } - - if (!enqueueText(controller, serializeSseEvent(result.value))) { - closeBody(controller); - } - } catch (error) { - if (!isClosed) { - isClosed = true; - clearKeepAlive(); - controller.error(error); - } - } finally { - if (reader === currentReader) { - reader = undefined; - } - - currentReader.releaseLock(); - } - }, - cancel() { - isClosed = true; - clearKeepAlive(); - - if (reader) { - void reader.cancel(); - } else { - void subscription.stream.cancel(); - } - }, - }); -} - function jsonResponse( value: unknown, status: number,