Skip to content
Closed
Changes from 1 commit
Commits
File filter

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Use runtime guard for Zod types that works across bundles
  • Loading branch information
bohdan-romanchenko committed May 9, 2025
commit e0c35b608fc0d59a9ad3c997b32aae69727040c3
61 changes: 44 additions & 17 deletions src/server/mcp.ts
Original file line number Diff line number Diff line change
Expand Up @@ -659,15 +659,31 @@ 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;

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;
Expand All @@ -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;
Expand All @@ -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;
}

/**
Expand Down