fix(http): make SSE close and response delivery deterministic

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<void>();
+    const allowDelete = createDeferred<void>();
+    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<void> | 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<void>();
     const requests: RecordedRequest[] = [];

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();
     }

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<AnyMessage, AnyMessage>();
+    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<Uint8Array>;
+
+    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<void> {
+  await Promise.resolve();
+  await Promise.resolve();
+}
+
+function decodeText(chunk: Uint8Array): string {
+  return new TextDecoder().decode(chunk);
+}

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<Uint8Array> {
+  return new ReadableStream<Uint8Array>(createSseBodySource(subscription));
+}
+
+/** @internal */
+export function createSseBodySource(
+  subscription: OutboundSubscription,
+): UnderlyingDefaultSource<Uint8Array> {
+  const encoder = new TextEncoder();
+  const replay = [...subscription.replay];
+  let keepAliveTimer: ReturnType<typeof setInterval> | undefined;
+  let reader: ReadableStreamDefaultReader<AnyMessage> | undefined;
+  let isClosed = false;
+
+  const clearKeepAlive = (): void => {
+    if (keepAliveTimer) {
+      clearInterval(keepAliveTimer);
+      keepAliveTimer = undefined;
+    }
+  };
+
+  const enqueueText = (
+    controller: ReadableStreamDefaultController<Uint8Array>,
+    text: string,
+  ): boolean => {
+    try {
+      controller.enqueue(encoder.encode(text));
+      return true;
+    } catch {
+      return false;
+    }
+  };
+
+  const hasDemand = (
+    controller: ReadableStreamDefaultController<Uint8Array>,
+  ): boolean => controller.desiredSize !== null && controller.desiredSize > 0;
+
+  const closeBody = (
+    controller: ReadableStreamDefaultController<Uint8Array>,
+  ): 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();
+      }
+    },
+  };
+}