fix(server): preserve inputSchema for z.discriminatedUnion / z.union (#1643)

normalizeObjectSchema returned undefined for any schema whose root was not
z.object(...), so registerTool silently advertised an empty schema in
tools/list when given z.discriminatedUnion(...). Tool calls still
validated correctly via the fallback in validateToolInput, which masked
the bug from end-to-end tests.

Pass discriminated unions and unions through unchanged in
normalizeObjectSchema; let toJsonSchemaCompat handle the actual
conversion (zod-to-json-schema for v3, Mini's z.toJSONSchema for v4).
Both emit oneOf/anyOf branches.

Default a top-level type: "object" on the emitted JSON Schema in mcp.ts
so the wire payload is spec compliant when the converter does not include
one (the discriminated-union and union case).

Adds three regression tests under "Tools with union and intersection
schemas" exercising tools/list, callTool, and outputSchema for v3 and v4.

Author: ankitvirdi4

.changeset/fix-discriminated-union-input-schema.md

@@ -0,0 +1,13 @@
+---
+'@modelcontextprotocol/sdk': patch
+---
+
+Fix `registerTool` / `registerPrompt` silently dropping `inputSchema` and `outputSchema` when given a `z.discriminatedUnion(...)` or `z.union(...)` of objects.
+
+`normalizeObjectSchema` previously returned `undefined` for any schema whose root was not `z.object(...)`, so the schema never reached `toJsonSchemaCompat` and `tools/list` advertised an empty schema. Tool calls still validated correctly via the fallback in `validateToolInput`,
+which masked the bug.
+
+`normalizeObjectSchema` now passes discriminated unions and unions through unchanged. The `tools/list` payload is also given a top-level `type: "object"` when missing so the emitted JSON Schema satisfies the MCP spec for tool input/output schemas (Zod emits `oneOf` / `anyOf`
+without a root type for these cases).
+
+Closes #1643.

src/server/mcp.ts

@@ -147,12 +147,17 @@ export class McpServer {
                             description: tool.description,
                             inputSchema: (() => {
                                 const obj = normalizeObjectSchema(tool.inputSchema);
-                                return obj
-                                    ? (toJsonSchemaCompat(obj, {
-                                          strictUnions: true,
-                                          pipeStrategy: 'input'
-                                      }) as Tool['inputSchema'])
-                                    : EMPTY_OBJECT_JSON_SCHEMA;
+                                if (!obj) return EMPTY_OBJECT_JSON_SCHEMA;
+                                const json = toJsonSchemaCompat(obj, {
+                                    strictUnions: true,
+                                    pipeStrategy: 'input'
+                                });
+                                // MCP requires `type: "object"` at the root of
+                                // tool inputSchema. Discriminated unions and
+                                // unions of objects produce `oneOf` / `anyOf`
+                                // without a top-level `type`; default it so the
+                                // emitted schema is spec compliant.
+                                return ensureObjectRoot(json) as Tool['inputSchema'];
                             })(),
                             annotations: tool.annotations,
                             execution: tool.execution,
@@ -162,10 +167,11 @@ export class McpServer {
                         if (tool.outputSchema) {
                             const obj = normalizeObjectSchema(tool.outputSchema);
                             if (obj) {
-                                toolDefinition.outputSchema = toJsonSchemaCompat(obj, {
+                                const json = toJsonSchemaCompat(obj, {
                                     strictUnions: true,
                                     pipeStrategy: 'output'
-                                }) as Tool['outputSchema'];
+                                });
+                                toolDefinition.outputSchema = ensureObjectRoot(json) as Tool['outputSchema'];
                             }
                         }
 
@@ -1331,6 +1337,22 @@ const EMPTY_OBJECT_JSON_SCHEMA = {
     properties: {}
 };
 
+/**
+ * Ensures a JSON Schema produced from a Zod schema has a top-level
+ * `type: "object"` per the MCP spec for tool input/output schemas.
+ *
+ * Plain `z.object(...)` already emits `type: "object"`, so this is a
+ * no-op for the common case. Discriminated unions and unions of objects
+ * emit `oneOf` / `anyOf` without a root `type`; we default it here so
+ * the wire payload is spec compliant. Schemas with an explicit non-object
+ * root `type` are left alone (they will fail downstream validation, which
+ * is the right signal for the user).
+ */
+function ensureObjectRoot(json: Record<string, unknown>): Record<string, unknown> {
+    if (json.type !== undefined) return json;
+    return { type: 'object', ...json };
+}
+
 /**
  * Checks if a value looks like a Zod schema by checking for parse/safeParse methods.
  */

src/server/zod-compat.ts

@@ -133,9 +133,18 @@ export function getObjectShape(schema: AnyObjectSchema | undefined): Record<stri
 
 // --- Schema normalization ---
 /**
- * Normalizes a schema to an object schema. Handles both:
+ * Normalizes a schema for use as an object-shaped tool/prompt input or output.
+ * Handles:
  * - Already-constructed object schemas (v3 or v4)
+ * - Discriminated unions and unions whose branches are object schemas
+ *   (e.g. `z.discriminatedUnion('action', [...])`); these convert to a
+ *   valid JSON Schema (`oneOf` / `anyOf`) via `toJsonSchemaCompat` and
+ *   parse correctly via `safeParse`, so they pass through unchanged.
  * - Raw shapes that need to be wrapped into object schemas
+ *
+ * Returns `undefined` for schemas whose root is not object-shaped (e.g.
+ * `z.string()`, `z.array(...)`), since those cannot satisfy the MCP spec's
+ * requirement that tool input/output schemas describe objects.
  */
 export function normalizeObjectSchema(schema: AnySchema | ZodRawShapeCompat | undefined): AnyObjectSchema | undefined {
     if (!schema) return undefined;
@@ -169,20 +178,35 @@ export function normalizeObjectSchema(schema: AnySchema | ZodRawShapeCompat | un
     }
 
     // If we get here, it should be an AnySchema (not a raw shape)
-    // Check if it's already an object schema
+    // Check if it's already an object schema, a discriminated union, or a
+    // union — all three convert cleanly to a JSON Schema with type:object
+    // (via `toJsonSchemaCompat`) and validate via `safeParse`.
     if (isZ4Schema(schema as AnySchema)) {
-        // Check if it's a v4 object
         const v4Schema = schema as unknown as ZodV4Internal;
         const def = v4Schema._zod?.def;
-        if (def && (def.type === 'object' || def.shape !== undefined)) {
-            return schema as AnyObjectSchema;
+        if (def) {
+            if (def.type === 'object' || def.shape !== undefined) {
+                return schema as AnyObjectSchema;
+            }
+            // v4 reports both `z.union(...)` and `z.discriminatedUnion(...)`
+            // as `def.type === 'union'`. Pass them through; downstream
+            // `toJsonSchemaCompat` (Mini's `z.toJSONSchema`) emits a valid
+            // `oneOf` / `anyOf` JSON Schema and `safeParse` handles them.
+            if (def.type === 'union') {
+                return schema as AnyObjectSchema;
+            }
         }
     } else {
-        // Check if it's a v3 object
         const v3Schema = schema as unknown as ZodV3Internal;
         if (v3Schema.shape !== undefined) {
             return schema as AnyObjectSchema;
         }
+        // v3 distinguishes the two; both serialise to a JSON Schema with
+        // `oneOf` / `anyOf` via the vendored `zodToJsonSchema` converter.
+        const typeName = v3Schema._def?.typeName;
+        if (typeName === 'ZodDiscriminatedUnion' || typeName === 'ZodUnion') {
+            return schema as AnyObjectSchema;
+        }
     }
 
     return undefined;

test/server/mcp.test.ts

@@ -5023,6 +5023,112 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                 ])
             );
         });
+
+        // Regression for https://github.com/modelcontextprotocol/typescript-sdk/issues/1643
+        // Before the fix, normalizeObjectSchema returned undefined for
+        // discriminated unions and unions, so registerTool silently dropped
+        // the schema in tools/list and emitted EMPTY_OBJECT_JSON_SCHEMA. Tool
+        // calls still validated correctly via the fallback in validateToolInput.
+        test('should expose a discriminated union inputSchema in tools/list', async () => {
+            const server = new McpServer({ name: 'test', version: '1.0.0' });
+            const client = new Client({ name: 'test-client', version: '1.0.0' });
+
+            const inputSchema = z.discriminatedUnion('action', [
+                z.object({ action: z.literal('create'), name: z.string() }),
+                z.object({ action: z.literal('delete'), id: z.string() })
+            ]);
+
+            server.registerTool('mutate', { inputSchema }, async args => ({
+                content: [{ type: 'text' as const, text: JSON.stringify(args) }]
+            }));
+
+            const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+            await server.connect(serverTransport);
+            await client.connect(clientTransport);
+
+            const list = await client.listTools();
+            expect(list.tools).toHaveLength(1);
+            const advertised = list.tools[0].inputSchema as Record<string, unknown>;
+            expect(advertised.type).toBe('object');
+            // Both v3 (zod-to-json-schema) and v4 (Mini) emit oneOf or anyOf.
+            const branches = (advertised.oneOf ?? advertised.anyOf) as Array<Record<string, unknown>> | undefined;
+            expect(branches).toBeDefined();
+            expect(branches).toHaveLength(2);
+
+            // Tool calls keep working.
+            const ok = await client.callTool({
+                name: 'mutate',
+                arguments: { action: 'create', name: 'foo' }
+            });
+            expect(ok.isError).toBeFalsy();
+
+            const bad = await client.callTool({
+                name: 'mutate',
+                arguments: { action: 'create' }
+            });
+            expect(bad.isError).toBe(true);
+        });
+
+        test('should expose a union of objects inputSchema in tools/list', async () => {
+            const server = new McpServer({ name: 'test', version: '1.0.0' });
+            const client = new Client({ name: 'test-client', version: '1.0.0' });
+
+            const inputSchema = z.union([
+                z.object({ kind: z.literal('a'), x: z.string() }),
+                z.object({ kind: z.literal('b'), y: z.number() })
+            ]);
+
+            server.registerTool('pick', { inputSchema }, async () => ({
+                content: [{ type: 'text' as const, text: 'ok' }]
+            }));
+
+            const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+            await server.connect(serverTransport);
+            await client.connect(clientTransport);
+
+            const list = await client.listTools();
+            const advertised = list.tools[0].inputSchema as Record<string, unknown>;
+            expect(advertised.type).toBe('object');
+            expect(advertised.oneOf ?? advertised.anyOf).toBeDefined();
+        });
+
+        test('should expose a discriminated union outputSchema in tools/list', async () => {
+            const server = new McpServer({ name: 'test', version: '1.0.0' });
+            const client = new Client({ name: 'test-client', version: '1.0.0' });
+
+            const outputSchema = z.discriminatedUnion('kind', [
+                z.object({ kind: z.literal('ok'), data: z.string() }),
+                z.object({ kind: z.literal('err'), message: z.string() })
+            ]);
+
+            server.registerTool(
+                'maybe',
+                {
+                    inputSchema: z.object({ should_fail: z.boolean() }),
+                    outputSchema
+                },
+                async ({ should_fail }) => {
+                    const structured = should_fail ? { kind: 'err' as const, message: 'oops' } : { kind: 'ok' as const, data: 'fine' };
+                    return {
+                        content: [{ type: 'text' as const, text: JSON.stringify(structured) }],
+                        structuredContent: structured
+                    };
+                }
+            );
+
+            const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+            await server.connect(serverTransport);
+            await client.connect(clientTransport);
+
+            const list = await client.listTools();
+            const advertised = list.tools[0].outputSchema as Record<string, unknown> | undefined;
+            expect(advertised).toBeDefined();
+            expect(advertised!.type).toBe('object');
+            expect(advertised!.oneOf ?? advertised!.anyOf).toBeDefined();
+
+            const ok = await client.callTool({ name: 'maybe', arguments: { should_fail: false } });
+            expect(ok.isError).toBeFalsy();
+        });
     });
 
     describe('Tools with transformation schemas', () => {