fix(v2): enforce method contracts

src/v2/acp.test.ts

@@ -1,4 +1,4 @@
-import { describe, expect, it } from "vitest";
+import { describe, expect, expectTypeOf, it } from "vitest";
 
 import {
   PROTOCOL_VERSION,
@@ -19,10 +19,12 @@ import {
 import type {
   AgentContext,
   Annotations,
+  ClientContext,
   DiffPatch,
   InitializeResponse,
   McpServer,
   NewSessionRequest,
+  NewSessionResponse,
   SessionInfo,
   SessionInfoUpdate,
   SessionUpdate,
@@ -31,6 +33,148 @@ import type {
 const clientInfo = { name: "test-client", version: "1.0.0" };
 const agentInfo = { name: "test-agent", version: "1.0.0" };
 
+function assertV2MethodTypes(
+  agentContext: ClientContext,
+  clientContext: AgentContext,
+): void {
+  // @ts-expect-error Built-in methods must not fall through the extension overload.
+  agentContext.request(methods.agent.session.new, { sessionId: "wrong" });
+  // @ts-expect-error Built-in notifications must not fall through the extension overload.
+  agentContext.notify(methods.agent.session.cancel, {});
+  agent().onRequest(
+    // @ts-expect-error Built-in handlers cannot replace their generated params parser.
+    methods.agent.session.new,
+    (params: unknown) => params,
+    () => ({ sessionId: "wrong-parser" }),
+  );
+  // @ts-expect-error Request methods cannot be sent as notifications.
+  agentContext.notify(methods.agent.session.new, {});
+  // @ts-expect-error Notification methods cannot be sent as requests.
+  agentContext.request(methods.agent.session.cancel, {});
+  // @ts-expect-error Client-directed methods cannot be sent to an agent.
+  agentContext.request(methods.client.mcp.disconnect, {
+    connectionId: "connection-1",
+  });
+
+  const parseValue = (params: unknown): { value: string } =>
+    params as { value: string };
+  void agentContext.request("session/load", { sessionId: "session-1" });
+  void agentContext.request<
+    { value: string },
+    { sessionId: string },
+    "session/load"
+  >("session/load", { sessionId: "session-1" });
+  void agentContext.request<{ value: string }, { value: string }>(
+    "_vendor/acme/echo",
+    { value: "request" },
+  );
+  void agentContext.notify("session/set_model", { modelId: "model-1" });
+  agent().onRequest("session/load", parseValue, ({ params }) => params);
+  agent().onNotification("session/set_model", parseValue, () => {});
+
+  const dynamicMethod: string = "method/from-another-draft";
+  void agentContext.request(dynamicMethod, {});
+  void agentContext.notify(dynamicMethod, {});
+  agent().onRequest(dynamicMethod, parseValue, ({ params }) => params);
+  agent().onNotification(dynamicMethod, parseValue, () => {});
+
+  const outputs = agentContext.batch([
+    batchRequest(methods.agent.session.new, {
+      cwd: "/workspace",
+      mcpServers: [],
+    }),
+    batchNotification(methods.agent.session.cancel, {
+      sessionId: "session-1",
+    }),
+  ] as const);
+  expectTypeOf(outputs).toEqualTypeOf<Promise<[NewSessionResponse, void]>>();
+  void agentContext.notify(methods.protocol.cancelRequest, { requestId: 1 });
+
+  const clientRequest = batchRequest(methods.client.mcp.disconnect, {
+    connectionId: "connection-1",
+  });
+  // @ts-expect-error Notification methods cannot be used as batch requests.
+  batchRequest(methods.agent.session.cancel, { sessionId: "session-1" });
+  // @ts-expect-error Request methods cannot be used as batch notifications.
+  batchNotification(methods.agent.session.new, {
+    cwd: "/workspace",
+    mcpServers: [],
+  });
+  // @ts-expect-error Client-directed methods cannot be sent in an agent-directed batch.
+  agentContext.batch([clientRequest] as const);
+  agentContext.batch([
+    // @ts-expect-error Raw built-in batch entries must use method-specific params.
+    {
+      kind: "request",
+      method: methods.agent.session.new,
+      params: { sessionId: "wrong" },
+    },
+  ] as const);
+  agentContext.batch([
+    batchRequest("session/load", { sessionId: "session-1" }),
+    batchNotification("session/set_model", { modelId: "model-1" }),
+  ] as const);
+  const legacyAgentEntry: sdk.AgentBatchEntry = batchRequest("session/load", {
+    sessionId: "session-1",
+  });
+  const legacyClientEntry: sdk.ClientBatchEntry = batchNotification(
+    "authentication/status",
+    {},
+  );
+  void legacyAgentEntry;
+  void legacyClientEntry;
+  agentContext.batch([
+    { kind: "request", method: dynamicMethod, params: {} },
+    { kind: "notification", method: dynamicMethod, params: {} },
+  ] as const);
+
+  clientContext.batch([clientRequest] as const);
+}
+
+void assertV2MethodTypes;
+
+function memoryWireStreamPair(): [sdk.Stream, sdk.Stream] {
+  const leftToRight = new TransformStream<sdk.AnyWireMessage>();
+  const rightToLeft = new TransformStream<sdk.AnyWireMessage>();
+  return [
+    {
+      readable: rightToLeft.readable,
+      writable: leftToRight.writable,
+    },
+    {
+      readable: leftToRight.readable,
+      writable: rightToLeft.writable,
+    },
+  ];
+}
+
+async function respondToNextRequest(
+  stream: sdk.Stream,
+  result: unknown,
+): Promise<void> {
+  const reader = stream.readable.getReader();
+  const request = await reader.read();
+  reader.releaseLock();
+  if (
+    request.done ||
+    Array.isArray(request.value) ||
+    !("id" in request.value)
+  ) {
+    throw new Error("Expected one JSON-RPC request");
+  }
+
+  const writer = stream.writable.getWriter();
+  try {
+    await writer.write({
+      jsonrpc: "2.0",
+      id: request.value.id,
+      result,
+    });
+  } finally {
+    writer.releaseLock();
+  }
+}
+
 describe("experimental v2 date-time schemas", () => {
   it("preserves RFC 3339 timestamps with timezone offsets as strings", () => {
     const timestamp = "2026-07-20T01:00:00+01:00";
@@ -583,74 +727,172 @@ describe("experimental v2 app API", () => {
     ).rejects.toMatchObject({ code: -32600 });
   });
 
-  it("requires and preserves underscore-prefixed extension methods", async () => {
+  it("validates every built-in direct response before returning it", async () => {
+    const [clientStream, peerStream] = memoryWireStreamPair();
+    const response = client().connectWith(clientStream, (agentContext) =>
+      agentContext.request(methods.agent.session.new, {
+        cwd: "/workspace",
+        mcpServers: [],
+      }),
+    );
+
+    await respondToNextRequest(peerStream, { sessionId: 42 });
+    await expect(response).rejects.toThrow();
+  });
+
+  it("validates built-in batch responses before applying caller mappings", async () => {
+    const [clientStream, peerStream] = memoryWireStreamPair();
+    let mapped = false;
+    const response = client().connectWith(clientStream, (agentContext) =>
+      agentContext.batch([
+        batchRequest(
+          methods.agent.session.new,
+          { cwd: "/workspace", mcpServers: [] },
+          (session) => {
+            mapped = true;
+            return session.sessionId;
+          },
+        ),
+      ] as const),
+    );
+
+    const reader = peerStream.readable.getReader();
+    const request = await reader.read();
+    reader.releaseLock();
+    if (
+      request.done ||
+      !Array.isArray(request.value) ||
+      request.value.length !== 1 ||
+      !("id" in request.value[0])
+    ) {
+      throw new Error("Expected one JSON-RPC batch request");
+    }
+    const writer = peerStream.writable.getWriter();
+    try {
+      await writer.write([
+        {
+          jsonrpc: "2.0",
+          id: request.value[0].id,
+          result: { sessionId: 42 },
+        },
+      ]);
+    } finally {
+      writer.releaseLock();
+    }
+
+    await expect(response).rejects.toThrow();
+    expect(mapped).toBe(false);
+  });
+
+  it("rejects peer null for empty responses but preserves local void handlers", async () => {
+    const [clientStream, peerStream] = memoryWireStreamPair();
+    const invalidResponse = client().connectWith(clientStream, (agentContext) =>
+      agentContext.request(methods.agent.session.delete, {
+        sessionId: "session-1",
+      }),
+    );
+
+    await respondToNextRequest(peerStream, null);
+    await expect(invalidResponse).rejects.toThrow();
+
+    await expect(
+      client().connectWith(
+        agent().onRequest(methods.agent.session.delete, () => {}),
+        (agentContext) =>
+          agentContext.request(methods.agent.session.delete, {
+            sessionId: "session-1",
+          }),
+      ),
+    ).resolves.toEqual({});
+  });
+
+  it("supports unrecognized protocol methods and underscore extensions", async () => {
     const parseValue = (params: unknown): { value: string } =>
       params as { value: string };
     const returnValue = ({ params }: { params: { value: string } }) => params;
 
-    expect(() =>
-      agent().onRequest("vendor/echo", parseValue, returnValue),
-    ).toThrow("must start with '_'");
-    expect(() =>
-      agent().onNotification("vendor/event", parseValue, () => {}),
-    ).toThrow("must start with '_'");
-    expect(() =>
-      client().onRequest("vendor/echo", parseValue, returnValue),
-    ).toThrow("must start with '_'");
-    expect(() =>
-      client().onNotification("vendor/event", parseValue, () => {}),
-    ).toThrow("must start with '_'");
-
-    let notificationValue: string | undefined;
-    const clientApp = client().onNotification(
-      "_vendor/acme/event",
-      parseValue,
-      ({ params }) => {
-        notificationValue = params.value;
-      },
-    );
+    let agentNotificationValue: string | undefined;
+    let clientNotificationValue: string | undefined;
+    const clientApp = client()
+      .onRequest("authentication/logout", parseValue, returnValue)
+      .onNotification("authentication/status", parseValue, ({ params }) => {
+        clientNotificationValue = params.value;
+      });
     const agentApp = agent()
+      .onRequest("session/load", parseValue, async ({ params, client }) => {
+        const response = await client.request<
+          { value: string },
+          { value: string },
+          "authentication/logout"
+        >("authentication/logout", params);
+        await client.notify("authentication/status", params);
+        return response;
+      })
       .onRequest("_vendor/acme/echo", parseValue, returnValue)
-      .onRequest(
-        methods.agent.initialize,
-        async ({ client: clientContext }) => {
-          expect(() =>
-            clientContext.request("vendor/client-request", {}),
-          ).toThrow("must start with '_'");
-          expect(() =>
-            clientContext.notify("vendor/client-notification", {}),
-          ).toThrow("must start with '_'");
-          await clientContext.notify("_vendor/acme/event", {
-            value: "notification",
-          });
-          return { protocolVersion: PROTOCOL_VERSION, info: agentInfo };
-        },
-      );
+      .onNotification("session/set_model", parseValue, ({ params }) => {
+        agentNotificationValue = params.value;
+      })
+      .onRequest(methods.agent.initialize, () => ({
+        protocolVersion: PROTOCOL_VERSION,
+        info: agentInfo,
+      }));
 
     await clientApp.connectWith(agentApp, async (agentContext) => {
-      expect(() => agentContext.request("vendor/request", {})).toThrow(
-        "must start with '_'",
-      );
-      expect(() => agentContext.notify("vendor/notification", {})).toThrow(
-        "must start with '_'",
-      );
+      const wrongDirection: string = methods.client.mcp.disconnect;
+      expect(() =>
+        agentContext.request(wrongDirection, { connectionId: "connection-1" }),
+      ).toThrow("not valid in this direction");
       expect(() =>
         agentContext.batch([
-          batchNotification("vendor/batch-notification", {}),
+          {
+            kind: "request",
+            method: wrongDirection,
+            params: { connectionId: "connection-1" },
+          },
         ] as const),
-      ).toThrow("must start with '_'");
+      ).toThrow("not valid in this direction");
 
       await agentContext.request(methods.agent.initialize, {
         protocolVersion: PROTOCOL_VERSION,
         info: clientInfo,
       });
+      await expect(
+        agentContext.request<
+          { value: string },
+          { value: string },
+          "session/load"
+        >("session/load", { value: "legacy response" }),
+      ).resolves.toEqual({ value: "legacy response" });
+      expect(clientNotificationValue).toBe("legacy response");
+
+      await agentContext.notify("session/set_model", {
+        value: "direct notification",
+      });
+      expect(agentNotificationValue).toBe("direct notification");
+
+      const [batchResponse] = await agentContext.batch([
+        batchRequest<{ value: string }, { value: string }, "session/load">(
+          "session/load",
+          { value: "batch response" },
+        ),
+        batchNotification("session/set_model", {
+          value: "batch notification",
+        }),
+      ] as const);
+      expect(batchResponse).toEqual({ value: "batch response" });
+      expect(agentNotificationValue).toBe("batch notification");
+
       await expect(
         agentContext.request<{ value: string }, { value: string }>(
           "_vendor/acme/echo",
           { value: "response" },
         ),
       ).resolves.toEqual({ value: "response" });
     });
-    expect(notificationValue).toBe("notification");
+
+    const dynamicBuiltIn: string = methods.agent.session.new;
+    expect(() =>
+      agent().onRequest(dynamicBuiltIn, parseValue, returnValue),
+    ).toThrow("Cannot replace the built-in");
   });
 });