feat: add discriminatedUnion transform for tagged unions

Add DISCRIMINATED_UNIONS config to convert z.union() to z.discriminatedUnion()
for better performance and error messages in tagged unions:
- ContentBlockSchema -> discriminatedUnion('type')
- SamplingMessageContentBlockSchema -> discriminatedUnion('type')

Re-export SamplingMessageContentBlockSchema from generated (now compatible).
Types.ts reduced from ~1275 to ~1267 lines.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: ochafik

scripts/generate-schemas.ts

@@ -100,6 +100,16 @@ const STRICT_SCHEMAS = [
     'EmptyResultSchema',
 ];
 
+/**
+ * Schemas that should use z.discriminatedUnion instead of z.union for better performance.
+ * Maps schema name to the discriminator field name.
+ */
+const DISCRIMINATED_UNIONS: Record<string, string> = {
+    'SamplingContentSchema': 'type',
+    'SamplingMessageContentBlockSchema': 'type',
+    'ContentBlockSchema': 'type',
+};
+
 // =============================================================================
 // Pre-processing: Transform spec types to SDK-compatible hierarchy
 // =============================================================================
@@ -273,6 +283,7 @@ const AST_TRANSFORMS: Transform[] = [
     transformUnionToEnum,
     applyFieldOverrides,
     addStrictToSchemas,
+    convertToDiscriminatedUnion,
 ];
 
 /**
@@ -503,6 +514,30 @@ function addStrictToSchemas(sourceFile: SourceFile): void {
     }
 }
 
+/**
+ * Convert z.union() to z.discriminatedUnion() for specified schemas.
+ * This provides better performance and error messages for tagged unions.
+ */
+function convertToDiscriminatedUnion(sourceFile: SourceFile): void {
+    for (const [schemaName, discriminator] of Object.entries(DISCRIMINATED_UNIONS)) {
+        const varDecl = sourceFile.getVariableDeclaration(schemaName);
+        if (!varDecl) continue;
+
+        const initializer = varDecl.getInitializer();
+        if (!initializer) continue;
+
+        const text = initializer.getText();
+
+        // Match z.union([...]) pattern and convert to z.discriminatedUnion('discriminator', [...])
+        const unionMatch = text.match(/^z\.union\(\s*\[([\s\S]*)\]\s*\)$/);
+        if (unionMatch) {
+            const members = unionMatch[1];
+            varDecl.setInitializer(`z.discriminatedUnion('${discriminator}', [${members}])`);
+            console.log(`    ✓ Converted ${schemaName} to discriminatedUnion('${discriminator}')`);
+        }
+    }
+}
+
 // =============================================================================
 // Main
 // =============================================================================

src/generated/sdk.schemas.ts

@@ -2115,7 +2115,7 @@ export const ResourceLinkSchema = ResourceSchema.extend({
 /**
  * @category Content
  */
-export const ContentBlockSchema = z.union([
+export const ContentBlockSchema = z.discriminatedUnion('type', [
     TextContentSchema,
     ImageContentSchema,
     AudioContentSchema,
@@ -2384,7 +2384,7 @@ export const PromptMessageSchema = z.object({
     content: ContentBlockSchema
 });
 
-export const SamplingMessageContentBlockSchema = z.union([
+export const SamplingMessageContentBlockSchema = z.discriminatedUnion('type', [
     TextContentSchema,
     ImageContentSchema,
     AudioContentSchema,

src/types.ts

@@ -192,6 +192,8 @@ import {
     UnsubscribeRequestParamsSchema,
     ResourceUpdatedNotificationParamsSchema,
     SetLevelRequestParamsSchema,
+    // Sampling schemas (now using discriminatedUnion)
+    SamplingMessageContentBlockSchema,
 } from './generated/sdk.schemas.js';
 
 // Alias RequestParamsSchema to BaseRequestParamsSchema for internal use
@@ -336,6 +338,7 @@ export {
     UnsubscribeRequestParamsSchema,
     ResourceUpdatedNotificationParamsSchema,
     SetLevelRequestParamsSchema,
+    SamplingMessageContentBlockSchema,
 };
 
 export const LATEST_PROTOCOL_VERSION = '2025-11-25';
@@ -820,19 +823,8 @@ export type ListChangedHandlers = {
  */
 export const SamplingContentSchema = z.discriminatedUnion('type', [TextContentSchema, ImageContentSchema, AudioContentSchema]);
 
-/**
- * Content block types allowed in sampling messages.
- * This includes text, image, audio, tool use requests, and tool results.
- */
-export const SamplingMessageContentBlockSchema = z.discriminatedUnion('type', [
-    TextContentSchema,
-    ImageContentSchema,
-    AudioContentSchema,
-    ToolUseContentSchema,
-    ToolResultContentSchema
-]);
-
-// Note: SamplingMessageSchema, CreateMessageRequestParamsSchema, CreateMessageRequestSchema,
+// Note: SamplingMessageContentBlockSchema (using discriminatedUnion),
+// SamplingMessageSchema, CreateMessageRequestParamsSchema, CreateMessageRequestSchema,
 // CreateMessageResultSchema are re-exported from generated.
 
 /**