fix(server): preserve routed outbound messages

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", () => {

src/protocol.ts

@@ -10,11 +10,23 @@ export const JSON_MIME_TYPE = "application/json";
 const SESSION_SCOPED_METHODS = new Set<string>([
   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 {

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

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<AnyMessage, AnyMessage>();
-    const writer = backend.writable.getWriter();
-    const source = createSseBodySource({
-      replay: [],
-      stream: backend.readable,
-    });
+    const mailbox = new OutboundMailbox<AnyMessage>();
+    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();
   });
 });
 

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<Uint8Array> {
-  return new ReadableStream<Uint8Array>(createSseBodySource(subscription));
+  return new ReadableStream<Uint8Array>(createSseBodySource(lease));
 }
 
 /** @internal */
 export function createSseBodySource(
-  subscription: OutboundSubscription,
+  lease: OutboundLease,
 ): UnderlyingDefaultSource<Uint8Array> {
   const encoder = new TextEncoder();
-  const replay = [...subscription.replay];
   let keepAliveTimer: ReturnType<typeof setInterval> | undefined;
-  let reader: ReadableStreamDefaultReader<AnyMessage> | 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();
     },
   };
 }