feat: port to SDK v2 via concrete Protocol (typescript-sdk#1891)

package.json

@@ -76,7 +76,7 @@
   "author": "Olivier Chafik",
   "devDependencies": {
     "@boneskull/typedoc-plugin-mermaid": "^0.2.0",
-    "@modelcontextprotocol/sdk": "^1.29.0",
+    "@modelcontextprotocol/sdk": "^2.0.0-alpha.2",
     "@playwright/test": "1.57.0",
     "@types/bun": "^1.3.2",
     "@types/node": "20.19.27",
@@ -107,7 +107,7 @@
     "zod": "^4.1.13"
   },
   "peerDependencies": {
-    "@modelcontextprotocol/sdk": "^1.29.0",
+    "@modelcontextprotocol/sdk": "^2.0.0-alpha.2",
     "react": "^17.0.0 || ^18.0.0 || ^19.0.0",
     "react-dom": "^17.0.0 || ^18.0.0 || ^19.0.0",
     "zod": "^3.25.0 || ^4.0.0"

src/app-bridge.ts

@@ -1325,7 +1325,7 @@ export class AppBridge extends ProtocolWithEvents<
    * Verify that the guest supports the capability required for the given request method.
    * @internal
    */
-  assertCapabilityForMethod(method: AppRequest["method"]): void {
+  assertCapabilityForMethod(method: string): void {
     // TODO
   }
 
@@ -1341,7 +1341,7 @@ export class AppBridge extends ProtocolWithEvents<
    * Verify that the host supports the capability required for the given notification method.
    * @internal
    */
-  assertNotificationCapability(method: AppNotification["method"]): void {
+  assertNotificationCapability(method: string): void {
     // TODO
   }
 

src/app.ts

@@ -762,7 +762,7 @@ export class App extends ProtocolWithEvents<
    * Verify that the host supports the capability required for the given request method.
    * @internal
    */
-  assertCapabilityForMethod(method: AppRequest["method"]): void {
+  assertCapabilityForMethod(method: string): void {
     // TODO
   }
 
@@ -792,7 +792,7 @@ export class App extends ProtocolWithEvents<
    * Verify that the app supports the capability required for the given notification method.
    * @internal
    */
-  assertNotificationCapability(method: AppNotification["method"]): void {
+  assertNotificationCapability(method: string): void {
     // TODO
   }
 

src/events.ts

@@ -3,10 +3,12 @@ import {
   Request,
   Notification,
   Result,
+  type BaseContext,
+  type LegacyContextFields,
+  type ZodLikeRequestSchema,
 } from "@modelcontextprotocol/sdk/types.js";
-import { ZodLiteral, ZodObject } from "zod/v4";
 
-type MethodSchema = ZodObject<{ method: ZodLiteral<string> }>;
+type AppContext = BaseContext & LegacyContextFields;
 
 /**
  * Per-event state: a singular `on*` handler (replace semantics) plus a
@@ -61,7 +63,10 @@ export abstract class ProtocolWithEvents<
   SendNotificationT extends Notification,
   SendResultT extends Result,
   EventMap extends Record<string, unknown>,
-> extends Protocol<SendRequestT, SendNotificationT, SendResultT> {
+> extends Protocol<AppContext> {
+  protected buildContext(ctx: AppContext): AppContext {
+    return ctx;
+  }
   private _registeredMethods = new Set<string>();
   private _eventSlots = new Map<keyof EventMap, EventSlot>();
 
@@ -71,7 +76,7 @@ export abstract class ProtocolWithEvents<
    * schema on first use.
    */
   protected abstract readonly eventSchemas: {
-    [K in keyof EventMap]: MethodSchema;
+    [K in keyof EventMap]: ZodLikeRequestSchema;
   };
 
   /**
@@ -195,12 +200,15 @@ export abstract class ProtocolWithEvents<
 
   // ── Handler registration with double-set protection ─────────────────
 
-  // The two overrides below are arrow-function class fields rather than
-  // prototype methods so that Protocol's constructor — which registers its
-  // own ping/cancelled/progress handlers via `this.setRequestHandler`
-  // before our fields initialize — hits the base implementation and skips
-  // tracking. Converting these to proper methods would crash with
-  // `_registeredMethods` undefined during super().
+  // These overrides are prototype methods, so Protocol's constructor (which
+  // registers built-in ping/cancelled/progress handlers via
+  // `this.setRequestHandler` during super()) dispatches here before our own
+  // fields have initialized. The `_registeredMethods === undefined` guard
+  // skips tracking during that window.
+  //
+  // The base method has four overloads in v2; we expose only the Zod-schema
+  // form here since that is all this package uses. Subclasses retain the
+  // typed (request, ctx) handler signature.
 
   /**
    * Registers a request handler. Throws if a handler for the same method
@@ -209,14 +217,24 @@ export abstract class ProtocolWithEvents<
    *
    * @throws {Error} if a handler for this method is already registered.
    */
-  override setRequestHandler: Protocol<
-    SendRequestT,
-    SendNotificationT,
-    SendResultT
-  >["setRequestHandler"] = (schema, handler) => {
+  override setRequestHandler<T extends ZodLikeRequestSchema>(
+    requestSchema: T,
+    handler: (
+      request: ReturnType<T["parse"]>,
+      ctx: AppContext,
+    ) => Result | Promise<Result>,
+  ): void;
+  override setRequestHandler(
+    schema: ZodLikeRequestSchema,
+    handler: (request: unknown, ctx: AppContext) => Result | Promise<Result>,
+  ): void {
+    if (this._registeredMethods === undefined) {
+      super.setRequestHandler(schema, handler);
+      return;
+    }
     this._assertMethodNotRegistered(schema, "setRequestHandler");
     super.setRequestHandler(schema, handler);
-  };
+  }
 
   /**
    * Registers a notification handler. Throws if a handler for the same
@@ -225,14 +243,21 @@ export abstract class ProtocolWithEvents<
    *
    * @throws {Error} if a handler for this method is already registered.
    */
-  override setNotificationHandler: Protocol<
-    SendRequestT,
-    SendNotificationT,
-    SendResultT
-  >["setNotificationHandler"] = (schema, handler) => {
+  override setNotificationHandler<T extends ZodLikeRequestSchema>(
+    notificationSchema: T,
+    handler: (notification: ReturnType<T["parse"]>) => void | Promise<void>,
+  ): void;
+  override setNotificationHandler(
+    schema: ZodLikeRequestSchema,
+    handler: (notification: unknown) => void | Promise<void>,
+  ): void {
+    if (this._registeredMethods === undefined) {
+      super.setNotificationHandler(schema, handler);
+      return;
+    }
     this._assertMethodNotRegistered(schema, "setNotificationHandler");
     super.setNotificationHandler(schema, handler);
-  };
+  }
 
   /**
    * Warn if a request handler `on*` setter is replacing a previously-set
@@ -255,18 +280,30 @@ export abstract class ProtocolWithEvents<
    * Replace a request handler, bypassing double-set protection. Used by
    * `on*` request-handler setters that need replace semantics.
    */
-  protected replaceRequestHandler: Protocol<
-    SendRequestT,
-    SendNotificationT,
-    SendResultT
-  >["setRequestHandler"] = (schema, handler) => {
-    const method = (schema as MethodSchema).shape.method.value;
-    this._registeredMethods.add(method);
+  protected replaceRequestHandler<T extends ZodLikeRequestSchema>(
+    requestSchema: T,
+    handler: (
+      request: ReturnType<T["parse"]>,
+      ctx: AppContext,
+    ) => Result | Promise<Result>,
+  ): void;
+  protected replaceRequestHandler(
+    schema: ZodLikeRequestSchema,
+    handler: (request: unknown, ctx: AppContext) => Result | Promise<Result>,
+  ): void {
+    this._registeredMethods.add(this._methodOf(schema));
     super.setRequestHandler(schema, handler);
-  };
+  }
+
+  private _methodOf(arg: ZodLikeRequestSchema | string): string {
+    return typeof arg === "string" ? arg : arg.shape.method.value;
+  }
 
-  private _assertMethodNotRegistered(schema: unknown, via: string): void {
-    const method = (schema as MethodSchema).shape.method.value;
+  private _assertMethodNotRegistered(
+    schema: ZodLikeRequestSchema | string,
+    via: string,
+  ): void {
+    const method = this._methodOf(schema);
     if (this._registeredMethods.has(method)) {
       throw new Error(
         `Handler for "${method}" already registered (via ${via}). ` +

src/server/index.ts

@@ -41,6 +41,7 @@ import {
 } from "../app.js";
 import type {
   BaseToolCallback,
+  LegacyToolCallback,
   McpServer,
   RegisteredTool,
   ResourceMetadata,
@@ -51,7 +52,9 @@ import type {
 import type {
   AnySchema,
   ZodRawShapeCompat,
+  StandardSchemaWithJSON,
 } from "@modelcontextprotocol/sdk/server/zod-compat.js";
+import type { ZodRawShape } from "@modelcontextprotocol/sdk";
 import type {
   ClientCapabilities,
   ReadResourceResult,
@@ -60,7 +63,7 @@ import type {
 
 // Re-exports for convenience
 export { RESOURCE_URI_META_KEY, RESOURCE_MIME_TYPE };
-export type { ResourceMetadata, ToolCallback };
+export type { ResourceMetadata, ToolCallback, LegacyToolCallback };
 
 /**
  * Base tool configuration matching the standard MCP server tool options.
@@ -214,8 +217,8 @@ export interface McpUiAppResourceConfig extends ResourceMetadata {
  * @see {@link registerAppResource `registerAppResource`} to register the HTML resource referenced by the tool
  */
 export function registerAppTool<
-  OutputArgs extends ZodRawShapeCompat | AnySchema,
-  InputArgs extends undefined | ZodRawShapeCompat | AnySchema = undefined,
+  OutputArgs extends StandardSchemaWithJSON,
+  InputArgs extends StandardSchemaWithJSON | undefined = undefined,
 >(
   server: Pick<McpServer, "registerTool">,
   name: string,
@@ -224,6 +227,25 @@ export function registerAppTool<
     outputSchema?: OutputArgs;
   },
   cb: ToolCallback<InputArgs>,
+): RegisteredTool;
+/** Raw-shape form: `inputSchema` may be a plain `{ field: z.string() }` record. */
+export function registerAppTool<
+  InputArgs extends ZodRawShape,
+  OutputArgs extends ZodRawShape | StandardSchemaWithJSON | undefined = undefined,
+>(
+  server: Pick<McpServer, "registerTool">,
+  name: string,
+  config: McpUiAppToolConfig & {
+    inputSchema: InputArgs;
+    outputSchema?: OutputArgs;
+  },
+  cb: LegacyToolCallback<InputArgs>,
+): RegisteredTool;
+export function registerAppTool(
+  server: Pick<McpServer, "registerTool">,
+  name: string,
+  config: McpUiAppToolConfig,
+  cb: (...args: never) => ReturnType<ToolCallback<undefined>>,
 ): RegisteredTool {
   // Normalize metadata for backward compatibility:
   // - If _meta.ui.resourceUri is set, also set the legacy flat key
@@ -241,6 +263,10 @@ export function registerAppTool<
     normalizedMeta = { ...meta, ui: { ...uiMeta, resourceUri: legacyUri } };
   }
 
+  // The two public overloads above guarantee (config.inputSchema, cb) match one
+  // of registerTool's overloads. The impl signature loses that pairing, so this
+  // forward needs a single suppression rather than three casts.
+  // @ts-expect-error -- forwarding overload-paired args through one impl signature
   return server.registerTool(name, { ...config, _meta: normalizedMeta }, cb);
 }