fix: update ToolCallback type to include output arguments for better type safety

Author: sushichan044

src/server/mcp.test.ts

@@ -1248,8 +1248,7 @@ describe("tool()", () => {
           processedInput: input,
           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) {
@@ -760,7 +760,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,
@@ -917,7 +917,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`);
@@ -932,7 +932,7 @@ export class McpServer {
       inputSchema,
       outputSchema,
       annotations,
-      cb as ToolCallback<ZodRawShape | undefined>
+      cb as ToolCallback<ZodRawShape | undefined, ZodRawShape | undefined>
     );
   }
 
@@ -1126,6 +1126,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().
  *
@@ -1136,36 +1146,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 = {