fix: type registerTool structured output

Author: Grmiade

.changeset/register-tool-output-schema.md

@@ -0,0 +1,5 @@
+---
+'@modelcontextprotocol/server': patch
+---
+
+Type `McpServer.registerTool` callbacks against `outputSchema` so returned `structuredContent` is checked when present.

packages/server/src/server/mcp.ts

@@ -942,7 +942,10 @@ export class McpServer {
      * );
      * ```
      */
-    registerTool<OutputArgs extends StandardSchemaWithJSON, InputArgs extends StandardSchemaWithJSON | undefined = undefined>(
+    registerTool<
+        InputArgs extends StandardSchemaWithJSON | ZodRawShape | undefined = undefined,
+        OutputArgs extends StandardSchemaWithJSON | ZodRawShape | undefined = undefined
+    >(
         name: string,
         config: {
             title?: string;
@@ -953,21 +956,7 @@ export class McpServer {
             icons?: Icon[];
             _meta?: Record<string, unknown>;
         },
-        cb: ToolCallback<InputArgs>
-    ): RegisteredTool;
-    /** @deprecated Wrap with `z.object({...})` instead. Raw-shape form: `inputSchema`/`outputSchema` may be a plain `{ field: z.string() }` record; it is auto-wrapped with `z.object()`. */
-    registerTool<InputArgs extends ZodRawShape, OutputArgs extends ZodRawShape | StandardSchemaWithJSON | undefined = undefined>(
-        name: string,
-        config: {
-            title?: string;
-            description?: string;
-            inputSchema?: InputArgs;
-            outputSchema?: OutputArgs;
-            annotations?: ToolAnnotations;
-            icons?: Icon[];
-            _meta?: Record<string, unknown>;
-        },
-        cb: LegacyToolCallback<InputArgs>
+        cb: ToolCallback<InputArgs, OutputArgs>
     ): RegisteredTool;
     registerTool(
         name: string,
@@ -980,7 +969,9 @@ export class McpServer {
             icons?: Icon[];
             _meta?: Record<string, unknown>;
         },
-        cb: ToolCallback<StandardSchemaWithJSON | undefined> | LegacyToolCallback<ZodRawShape>
+        cb:
+            | ToolCallback<StandardSchemaWithJSON | undefined, StandardSchemaWithJSON | undefined>
+            | LegacyToolCallback<ZodRawShape, StandardSchemaWithJSON | ZodRawShape | undefined>
     ): RegisteredTool {
         if (this._registeredTools[name]) {
             throw new Error(`Tool ${name} is already registered`);
@@ -1211,12 +1202,10 @@ export type ZodRawShape = Record<string, z.ZodType>;
 export type InferRawShape<S extends ZodRawShape> = z.infer<z.ZodObject<S>>;
 
 /** {@linkcode ToolCallback} variant used when `inputSchema` is a {@linkcode ZodRawShape}. */
-export type LegacyToolCallback<Args extends ZodRawShape | undefined> = Args extends ZodRawShape
-    ? (
-          args: InferRawShape<Args>,
-          ctx: ServerContext
-      ) => CallToolResult | InputRequiredResult | Promise<CallToolResult | InputRequiredResult>
-    : (ctx: ServerContext) => CallToolResult | InputRequiredResult | Promise<CallToolResult | InputRequiredResult>;
+export type LegacyToolCallback<
+    Args extends ZodRawShape | undefined,
+    OutputArgs extends StandardSchemaWithJSON | ZodRawShape | undefined = undefined
+> = ToolCallback<Args, OutputArgs>;
 
 /** {@linkcode PromptCallback} variant used when `argsSchema` is a {@linkcode ZodRawShape}. */
 export type LegacyPromptCallback<Args extends ZodRawShape | undefined> = Args extends ZodRawShape
@@ -1229,16 +1218,30 @@ export type LegacyPromptCallback<Args extends ZodRawShape | undefined> = Args ex
 export type BaseToolCallback<
     SendResultT extends Result,
     Ctx extends ServerContext,
-    Args extends StandardSchemaWithJSON | undefined
+    Args extends StandardSchemaWithJSON | ZodRawShape | undefined
 > = Args extends StandardSchemaWithJSON
     ? (args: StandardSchemaWithJSON.InferOutput<Args>, ctx: Ctx) => SendResultT | Promise<SendResultT>
-    : (ctx: Ctx) => SendResultT | Promise<SendResultT>;
+    : Args extends ZodRawShape
+      ? (args: InferRawShape<Args>, ctx: Ctx) => SendResultT | Promise<SendResultT>
+      : (ctx: Ctx) => SendResultT | Promise<SendResultT>;
 
 /**
  * Callback for a tool handler registered with {@linkcode McpServer.registerTool}.
  */
-export type ToolCallback<Args extends StandardSchemaWithJSON | undefined = undefined> = BaseToolCallback<
-    CallToolResult | InputRequiredResult,
+export type ToolCallback<
+    Args extends StandardSchemaWithJSON | ZodRawShape | undefined = undefined,
+    OutputArgs extends StandardSchemaWithJSON | ZodRawShape | undefined = undefined
+> = BaseToolCallback<
+    | (OutputArgs extends StandardSchemaWithJSON
+          ? Omit<CallToolResult, 'structuredContent'> & {
+                structuredContent?: StandardSchemaWithJSON.InferOutput<OutputArgs>;
+            }
+          : OutputArgs extends ZodRawShape
+            ? Omit<CallToolResult, 'structuredContent'> & {
+                  structuredContent?: InferRawShape<OutputArgs>;
+              }
+            : CallToolResult)
+    | InputRequiredResult,
     ServerContext,
     Args
 >;

packages/server/test/server/mcp.compat.test.ts

@@ -135,15 +135,57 @@ describe('SEP-2106: registerTool with non-object outputSchema (type-level)', ()
             content: [],
             structuredContent: [n, n + 1] satisfies number[]
         }));
-        // NOTE (SEP-2106 PR-B verification item): the OutputArgs generic on registerTool is
-        // captured but does NOT currently flow into the callback's return type — ToolCallback's
-        // SendResultT is `CallToolResult | InputRequiredResult` (structuredContent: unknown), so
-        // a wrong-typed structuredContent ALSO compiles. Runtime validation (validateToolOutput)
-        // is the guard. Tightening the generic is out of this commit's scope.
-        server.registerTool('arr-loose', { outputSchema: z.array(z.number()) }, async () => ({
+        server.registerTool('arr-loose', { outputSchema: z.array(z.number()) }, () => ({
             content: [],
-            structuredContent: 'not-an-array' // compiles: structuredContent is `unknown`
+            // @ts-expect-error structuredContent must match outputSchema when present
+            structuredContent: 'not-an-array'
         }));
         expectTypeOf<number[]>().toMatchTypeOf<z.infer<ReturnType<typeof z.array<z.ZodNumber>>>>();
     });
 });
+
+describe('registerTool outputSchema typing', () => {
+    it('checks object structuredContent', () => {
+        const server = new McpServer({ name: 's', version: '1' });
+        const outputSchema = z.object({ id: z.string() });
+
+        server.registerTool('object-ok', { inputSchema: z.object({}), outputSchema }, async () => ({
+            content: [{ type: 'text' as const, text: 'ok' }],
+            structuredContent: { id: 'abc' }
+        }));
+
+        server.registerTool('object-bad', { inputSchema: z.object({}), outputSchema }, () => ({
+            content: [{ type: 'text' as const, text: 'bad' }],
+            structuredContent: {
+                // @ts-expect-error structuredContent must match outputSchema when present
+                id: 123
+            }
+        }));
+    });
+
+    it('checks raw-shape structuredContent', () => {
+        const server = new McpServer({ name: 's', version: '1' });
+
+        server.registerTool(
+            'raw-ok',
+            { inputSchema: { count: z.number() }, outputSchema: { id: z.string(), total: z.number().optional() } },
+            async ({ count }) => ({
+                content: [],
+                structuredContent: { id: String(count), total: count }
+            })
+        );
+
+        server.registerTool(
+            'raw-bad',
+            { inputSchema: { count: z.number() }, outputSchema: { id: z.string(), total: z.number().optional() } },
+            () => ({
+                content: [],
+                structuredContent: {
+                    id: 'abc',
+                    // @ts-expect-error structuredContent must match raw-shape outputSchema when present
+                    total: 'nope'
+                }
+            })
+        );
+    });
+});

test/e2e/scenarios/standard-schema.test.ts

@@ -139,7 +139,15 @@ verifies('standardschema:tool:output-schema-validation', async ({ transport }: T
             'get-server-status-corrupt',
             { inputSchema: type({}), outputSchema },
             // intentionally nonconforming structuredContent (server-side output validation must reject it)
-            () => ({ structuredContent: { healthy: 'definitely', uptimeSeconds: 'a while' }, content: [] })
+            () => ({
+                structuredContent: {
+                    // @ts-expect-error intentionally nonconforming structuredContent
+                    healthy: 'definitely',
+                    // @ts-expect-error intentionally nonconforming structuredContent
+                    uptimeSeconds: 'a while'
+                },
+                content: []
+            })
         );
         return s;
     };

test/e2e/scenarios/tools.test.ts

@@ -92,7 +92,13 @@ function schemaServer(): McpServer {
         'structured-mismatch',
         { inputSchema: z.object({}), outputSchema: z.object({ value: z.number() }) },
         // intentionally invalid structuredContent (tests server-side validation rejects it)
-        () => ({ structuredContent: { value: 'not-a-number' }, content: [] })
+        () => ({
+            structuredContent: {
+                // @ts-expect-error intentionally nonconforming structuredContent
+                value: 'not-a-number'
+            },
+            content: []
+        })
     );
     s.registerTool('structured-missing', { inputSchema: z.object({}), outputSchema: z.object({ value: z.number() }) }, () => ({
         content: [{ type: 'text', text: 'handler-body-no-structured' }]