fix(v2): enforce connection lifecycle

src/examples/dual-version-agent.ts

@@ -105,7 +105,7 @@ const v2Agent = v2
     await cancelV2Turn(session);
     session.active = false;
   })
-  .onRequest(v2.methods.agent.session.prompt, ({ params, client }) => {
+  .onRequest(v2.methods.agent.session.prompt, ({ params, client, accept }) => {
     const session = requireActiveV2Session(params.sessionId);
     if (session.turn) {
       throw v2.RequestError.invalidRequest(
@@ -118,13 +118,7 @@ const v2Agent = v2
       controller,
       done: Promise.resolve(),
     };
-    // The framework queues the prompt response after this handler returns.
-    // Start work in the next event-loop task so that response is queued before
-    // any session updates from the turn.
-    const responseQueued = new Promise<void>((resolve) => {
-      setTimeout(resolve, 0);
-    });
-    turn.done = responseQueued
+    turn.done = accept()
       .then(() => runV2Turn(params, client, session, controller.signal))
       .catch((error) => {
         console.error("v2 example turn failed", error);

src/jsonrpc.ts

@@ -684,6 +684,13 @@ export class RequestResponder<Resp = unknown> {
      */
     public readonly signal: AbortSignal = new AbortController().signal,
     private finishRequest?: () => void,
+    /**
+     * Number of entries in the containing wire batch, when this request was
+     * received in a batch.
+     *
+     * @internal
+     */
+    public readonly batchSize?: number,
   ) {}
 
   /**
@@ -1325,6 +1332,7 @@ export class Connection {
       const processing = this.receiveMessage(
         message,
         isRequestMessage(message) ? collectResponse : undefined,
+        batch.length,
       );
       if (isNotificationMessage(message)) {
         void processing.finally(() => {
@@ -1338,6 +1346,7 @@ export class Connection {
   private receiveMessage(
     message: AnyMessage,
     sendResponse?: (response: AnyResponse) => Promise<void>,
+    batchSize?: number,
   ): Promise<void> {
     if (this.abortController.signal.aborted) {
       return Promise.resolve();
@@ -1355,7 +1364,7 @@ export class Connection {
         this.handleProtocolNotification(message);
       }
       return this.processIncomingMessage(
-        this.toIncomingMessage(message, sendResponse),
+        this.toIncomingMessage(message, sendResponse, batchSize),
       ).catch((error) => this.close(error));
     } else if ("id" in message) {
       this.handleResponse(message);
@@ -1428,6 +1437,7 @@ export class Connection {
   private toIncomingMessage(
     message: AnyRequest | AnyNotification,
     sendResponse?: (response: AnyResponse) => Promise<void>,
+    batchSize?: number,
   ): IncomingMessage {
     if ("id" in message) {
       const abortController = new AbortController();
@@ -1458,6 +1468,7 @@ export class Connection {
           },
           abortController.signal,
           finishRequest,
+          batchSize,
         ),
       };
     }

src/protocol-router.test.ts

@@ -62,6 +62,62 @@ describe("AgentProtocolRouter", () => {
     await second.close();
   });
 
+  it.each([
+    [
+      1,
+      {
+        clientCapabilities: {
+          _future: { nullable: null, entries: [{ value: 1, extra: true }] },
+        },
+        clientInfo: implementation("v1-client"),
+        futureTopLevel: { empty: [], explicitNull: null },
+      },
+    ],
+    [
+      2,
+      {
+        info: implementation("v2-client"),
+        capabilities: {
+          _future: { nullable: null, entries: [{ value: 2, extra: true }] },
+        },
+        futureTopLevel: { empty: [], explicitNull: null },
+      },
+    ],
+  ])(
+    "validates and preserves same-version v%s initialize params",
+    async (version, params) => {
+      const selected = new MockAgentConnector();
+      const router =
+        version === 1
+          ? new AgentProtocolRouter().withV1(selected)
+          : new AgentProtocolRouter().withV2(selected);
+      const initialize = initializeRequest(version, params);
+      const connection = await openRoutedConnection(router, initialize);
+
+      expect(await selected.nextMessage()).toEqual(initialize);
+      await connection.close();
+    },
+  );
+
+  it("still validates same-version initialize params before routing", async () => {
+    const v2 = new MockAgentConnector();
+    const router = new AgentProtocolRouter().withV2(v2);
+    const { response, closed } = await rejectedConnection(
+      router,
+      initializeRequest(2, { capabilities: {} }),
+    );
+
+    expect(response).toMatchObject({
+      id: 1,
+      error: {
+        code: -32602,
+        data: expect.stringContaining("invalid initialize params"),
+      },
+    });
+    expect(v2.connectionCount).toBe(0);
+    await closed;
+  });
+
   it("routes a future protocol version to v2 and normalizes initialize", async () => {
     const v1 = new MockAgentConnector();
     const v2 = new MockAgentConnector();

src/protocol-router.ts

@@ -43,8 +43,9 @@ const MAX_PROTOCOL_VERSION = 0xffff;
  * The router consumes the first wire item, which must be an `initialize`
  * request or a single-entry batch containing one. It selects the highest
  * configured protocol version that does not exceed the client's requested
- * version. Only the initialize params and optional batch framing are
- * normalized; every later wire item is forwarded unchanged.
+ * version. Same-version initialize params are validated and forwarded
+ * unchanged. Version selection and downgrade may normalize initialize params;
+ * every later wire item is forwarded unchanged.
  *
  * @experimental
  */
@@ -560,6 +561,15 @@ function rewriteInitializeParams(
   requested: number,
   selected: AgentProtocol,
 ): JsonObject {
+  if (requested === selected) {
+    if (selected === PROTOCOL_V1) {
+      zV1InitializeRequest.parse(params);
+    } else {
+      zV2InitializeRequest.parse(params);
+    }
+    return params;
+  }
+
   if (selected === PROTOCOL_V1) {
     return requested >= PROTOCOL_V2
       ? v2InitializeToV1(zV2InitializeRequest.parse(params))