From e0c35b608fc0d59a9ad3c997b32aae69727040c3 Mon Sep 17 00:00:00 2001
From: Bohdan Romanchenko <bohdan@bettorfaster.com>
Date: Fri, 9 May 2025 15:53:41 +0200
Subject: [PATCH] Use runtime guard for Zod types that works across bundles

---
 src/server/mcp.ts | 61 ++++++++++++++++++++++++++++++++++-------------
 1 file changed, 44 insertions(+), 17 deletions(-)

diff --git a/src/server/mcp.ts b/src/server/mcp.ts
index dccfa04e..921c67f3 100644
--- a/src/server/mcp.ts
+++ b/src/server/mcp.ts
@@ -659,7 +659,21 @@ export class McpServer {
     if (this._registeredTools[name]) {
       throw new Error(`Tool ${name} is already registered`);
     }
-    
+
+    // Runtime type guard for Zod types that works across bundles/realms
+    const isZodType = (value: unknown): value is ZodTypeAny => {
+      if (typeof value !== "object" || value === null) return false;
+
+      if (!("_def" in value)) return false;
+
+      const def = (value as { _def: unknown })._def;
+      if (typeof def !== "object" || def === null) return false;
+
+      if (!("typeName" in def)) return false;
+
+      return typeof (def as { typeName: unknown }).typeName === "string";
+    };
+
     // Helper to check if an object is a Zod schema (ZodRawShape)
     const isZodRawShape = (obj: unknown): obj is ZodRawShape => {
       if (typeof obj !== "object" || obj === null) return false;
@@ -667,7 +681,9 @@ export class McpServer {
       const isEmptyObject = z.object({}).strict().safeParse(obj).success;
 
       // Check if object is empty or at least one property is a ZodType instance
-      return isEmptyObject || Object.values(obj as object).some(v => v instanceof ZodType);
+      return (
+        isEmptyObject || Object.values(obj as object).some((v) => isZodType(v))
+      );
     };
 
     let description: string | undefined;
@@ -677,18 +693,23 @@ export class McpServer {
 
     let paramsSchema: ZodRawShape | undefined;
     let annotations: ToolAnnotations | undefined;
-    
+
     // Handle the different overload combinations
     if (rest.length > 1) {
       // We have at least two more args before the callback
       const firstArg = rest[0];
-      
+
       if (isZodRawShape(firstArg)) {
         // We have a params schema as the first arg
         paramsSchema = rest.shift() as ZodRawShape;
-        
-        // Check if the next arg is potentially annotations  
-        if (rest.length > 1 && typeof rest[0] === "object" && rest[0] !== null && !(isZodRawShape(rest[0]))) {
+
+        // Check if the next arg is potentially annotations
+        if (
+          rest.length > 1 &&
+          typeof rest[0] === "object" &&
+          rest[0] !== null &&
+          !isZodRawShape(rest[0])
+        ) {
           // Case: tool(name, paramsSchema, annotations, cb)
           // Or: tool(name, description, paramsSchema, annotations, cb)
           annotations = rest.shift() as ToolAnnotations;
@@ -714,23 +735,29 @@ export class McpServer {
       remove: () => registeredTool.update({ name: null }),
       update: (updates) => {
         if (typeof updates.name !== "undefined" && updates.name !== name) {
-          delete this._registeredTools[name]
-          if (updates.name) this._registeredTools[updates.name] = registeredTool
+          delete this._registeredTools[name];
+          if (updates.name)
+            this._registeredTools[updates.name] = registeredTool;
         }
-        if (typeof updates.description !== "undefined") registeredTool.description = updates.description
-        if (typeof updates.paramsSchema !== "undefined") registeredTool.inputSchema = z.object(updates.paramsSchema)
-        if (typeof updates.callback !== "undefined") registeredTool.callback = updates.callback
-        if (typeof updates.annotations !== "undefined") registeredTool.annotations = updates.annotations
-        if (typeof updates.enabled !== "undefined") registeredTool.enabled = updates.enabled
-        this.sendToolListChanged()
+        if (typeof updates.description !== "undefined")
+          registeredTool.description = updates.description;
+        if (typeof updates.paramsSchema !== "undefined")
+          registeredTool.inputSchema = z.object(updates.paramsSchema);
+        if (typeof updates.callback !== "undefined")
+          registeredTool.callback = updates.callback;
+        if (typeof updates.annotations !== "undefined")
+          registeredTool.annotations = updates.annotations;
+        if (typeof updates.enabled !== "undefined")
+          registeredTool.enabled = updates.enabled;
+        this.sendToolListChanged();
       },
     };
     this._registeredTools[name] = registeredTool;
 
     this.setToolRequestHandlers();
-    this.sendToolListChanged()
+    this.sendToolListChanged();
 
-    return registeredTool
+    return registeredTool;
   }
 
   /**