Merge pull request #670 from sushichan044/feat/structured-content-with-type-safety

fix: add type safety for tool output schemas in ToolCallback

Author: bhosmer-ant

src/examples/server/mcpServerOutputSchema.ts

@@ -43,7 +43,14 @@ server.registerTool(
     void country;
     // Simulate weather API call
     const temp_c = Math.round((Math.random() * 35 - 5) * 10) / 10;
-    const conditions = ["sunny", "cloudy", "rainy", "stormy", "snowy"][Math.floor(Math.random() * 5)];
+    const conditionCandidates = [
+      "sunny",
+      "cloudy",
+      "rainy",
+      "stormy",
+      "snowy",
+    ] as const;
+    const conditions = conditionCandidates[Math.floor(Math.random() * conditionCandidates.length)];
 
     const structuredContent = {
       temperature: {
@@ -77,4 +84,4 @@ async function main() {
 main().catch((error) => {
   console.error("Server error:", error);
   process.exit(1);
-});
+});

src/server/mcp.test.ts

@@ -1312,7 +1312,7 @@ describe("tool()", () => {
           resultType: "structured",
           // Missing required 'timestamp' field
           someExtraField: "unexpected" // Extra field not in schema
-        },
+        } as unknown as { processedInput: string; resultType: string; timestamp: string }, // Type assertion to bypass TypeScript validation for testing purposes
       })
     );
 

src/server/mcp.ts

@@ -169,7 +169,7 @@ export class McpServer {
           }
 
           const args = parseResult.data;
-          const cb = tool.callback as ToolCallback<ZodRawShape>;
+          const cb = tool.callback as ToolCallback<ZodRawShape, ZodRawShape>;
           try {
             result = await Promise.resolve(cb(args, extra));
           } catch (error) {
@@ -184,7 +184,7 @@ export class McpServer {
             };
           }
         } else {
-          const cb = tool.callback as ToolCallback<undefined>;
+          const cb = tool.callback as ToolCallback<undefined, ZodRawShape>;
           try {
             result = await Promise.resolve(cb(extra));
           } catch (error) {
@@ -772,7 +772,7 @@ export class McpServer {
     inputSchema: ZodRawShape | undefined,
     outputSchema: ZodRawShape | undefined,
     annotations: ToolAnnotations | undefined,
-    callback: ToolCallback<ZodRawShape | undefined>
+    callback: ToolCallback<ZodRawShape | undefined, ZodRawShape | undefined>
   ): RegisteredTool {
     const registeredTool: RegisteredTool = {
       title,
@@ -929,7 +929,7 @@ export class McpServer {
       outputSchema?: OutputArgs;
       annotations?: ToolAnnotations;
     },
-    cb: ToolCallback<InputArgs>
+    cb: ToolCallback<InputArgs, OutputArgs>
   ): RegisteredTool {
     if (this._registeredTools[name]) {
       throw new Error(`Tool ${name} is already registered`);
@@ -944,7 +944,7 @@ export class McpServer {
       inputSchema,
       outputSchema,
       annotations,
-      cb as ToolCallback<ZodRawShape | undefined>
+      cb as ToolCallback<ZodRawShape | undefined, ZodRawShape | undefined>
     );
   }
 
@@ -1138,6 +1138,16 @@ export class ResourceTemplate {
   }
 }
 
+/**
+ * Type helper to create a strongly-typed CallToolResult with structuredContent
+ */
+type TypedCallToolResult<OutputArgs extends undefined | ZodRawShape> =
+  OutputArgs extends ZodRawShape
+    ? CallToolResult & {
+        structuredContent?: z.objectOutputType<OutputArgs, ZodTypeAny>;
+      }
+    : CallToolResult;
+
 /**
  * Callback for a tool handler registered with Server.tool().
  *
@@ -1148,36 +1158,46 @@ export class ResourceTemplate {
  * - `content` if the tool does not have an outputSchema
  * - Both fields are optional but typically one should be provided
  */
-export type ToolCallback<Args extends undefined | ZodRawShape = undefined> =
-  Args extends ZodRawShape
+export type ToolCallback<
+  InputArgs extends undefined | ZodRawShape = undefined,
+  OutputArgs extends undefined | ZodRawShape = undefined
+> = InputArgs extends ZodRawShape
   ? (
-    args: z.objectOutputType<Args, ZodTypeAny>,
-    extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
-  ) => CallToolResult | Promise<CallToolResult>
-  : (extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => CallToolResult | Promise<CallToolResult>;
+      args: z.objectOutputType<InputArgs, ZodTypeAny>,
+      extra: RequestHandlerExtra<ServerRequest, ServerNotification>
+    ) =>
+      | TypedCallToolResult<OutputArgs>
+      | Promise<TypedCallToolResult<OutputArgs>>
+  : (
+      extra: RequestHandlerExtra<ServerRequest, ServerNotification>
+    ) =>
+      | TypedCallToolResult<OutputArgs>
+      | Promise<TypedCallToolResult<OutputArgs>>;
 
 export type RegisteredTool = {
   title?: string;
   description?: string;
   inputSchema?: AnyZodObject;
   outputSchema?: AnyZodObject;
   annotations?: ToolAnnotations;
-  callback: ToolCallback<undefined | ZodRawShape>;
+  callback: ToolCallback<ZodRawShape | undefined, ZodRawShape | undefined>;
   enabled: boolean;
   enable(): void;
   disable(): void;
-  update<InputArgs extends ZodRawShape, OutputArgs extends ZodRawShape>(
-    updates: {
-      name?: string | null,
-      title?: string,
-      description?: string,
-      paramsSchema?: InputArgs,
-      outputSchema?: OutputArgs,
-      annotations?: ToolAnnotations,
-      callback?: ToolCallback<InputArgs>,
-      enabled?: boolean
-    }): void
-  remove(): void
+  update<
+    InputArgs extends ZodRawShape,
+    OutputArgs extends ZodRawShape
+  >(updates: {
+    name?: string | null;
+    title?: string;
+    description?: string;
+    paramsSchema?: InputArgs;
+    outputSchema?: OutputArgs;
+    annotations?: ToolAnnotations;
+    callback?: ToolCallback<InputArgs, OutputArgs>
+    enabled?: boolean
+  }): void;
+  remove(): void;
 };
 
 const EMPTY_OBJECT_JSON_SCHEMA = {