From 31087ed6d4be5a0f0e0e90dd8945a512787be5a9 Mon Sep 17 00:00:00 2001
From: Stephen Toub <stoub@microsoft.com>
Date: Fri, 10 Apr 2026 18:33:04 -0400
Subject: [PATCH 1/3] Add $ref support to all four language code generators

Enable JSON Schema $ref for type deduplication across all SDK code
generators (TypeScript, Python, Go, C#).

Changes:
- utils.ts: Add resolveRef(), refTypeName(), collectDefinitions() helpers;
  normalize $defs to definitions in postProcessSchema
- typescript.ts: Build combined schema with shared definitions and compile
  once via unreachableDefinitions, instead of per-method compilation
- python.ts/go.ts: Include all definitions alongside SessionEvent for
  quicktype  resolution; include shared API defs in RPC combined schema
- csharp.ts: Add  handling to resolveSessionPropertyType and
  resolveRpcType; generate classes for referenced types on demand

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 scripts/codegen/csharp.ts     | 60 ++++++++++++++++++++++++++++++++++-
 scripts/codegen/go.ts         | 36 ++++++++++++++++++---
 scripts/codegen/python.ts     | 29 +++++++++++++----
 scripts/codegen/typescript.ts | 55 ++++++++++++++++++++++++--------
 scripts/codegen/utils.ts      | 37 +++++++++++++++++++++
 5 files changed, 190 insertions(+), 27 deletions(-)

diff --git a/scripts/codegen/csharp.ts b/scripts/codegen/csharp.ts
index e6042eae5..eb70e0b4a 100644
--- a/scripts/codegen/csharp.ts
+++ b/scripts/codegen/csharp.ts
@@ -10,11 +10,14 @@ import { execFile } from "child_process";
 import fs from "fs/promises";
 import path from "path";
 import { promisify } from "util";
-import type { JSONSchema7 } from "json-schema";
+import type { JSONSchema7, JSONSchema7Definition } from "json-schema";
 import {
     getSessionEventsSchemaPath,
     getApiSchemaPath,
     writeGeneratedFile,
+    collectDefinitions,
+    resolveRef,
+    refTypeName,
     isRpcMethod,
     isNodeFullyExperimental,
     EXCLUDED_EVENT_TYPES,
@@ -297,6 +300,9 @@ interface EventVariant {
 
 let generatedEnums = new Map<string, { enumName: string; values: string[] }>();
 
+/** Schema definitions available during session event generation (for $ref resolution). */
+let sessionDefinitions: Record<string, JSONSchema7Definition> = {};
+
 function getOrCreateEnum(parentClassName: string, propName: string, values: string[], enumOutput: string[], description?: string): string {
     const valuesKey = [...values].sort().join("|");
     for (const [, existing] of generatedEnums) {
@@ -504,6 +510,21 @@ function resolveSessionPropertyType(
     nestedClasses: Map<string, string>,
     enumOutput: string[]
 ): string {
+    // Handle $ref by resolving against schema definitions
+    if (propSchema.$ref) {
+        const typeName = refTypeName(propSchema.$ref);
+        const className = typeToClassName(typeName);
+        if (!nestedClasses.has(className)) {
+            const refSchema = resolveRef(propSchema.$ref, sessionDefinitions);
+            if (refSchema) {
+                if (refSchema.enum && Array.isArray(refSchema.enum)) {
+                    return getOrCreateEnum(className, "", refSchema.enum as string[], enumOutput);
+                }
+                nestedClasses.set(className, generateNestedClass(className, refSchema, knownTypes, nestedClasses, enumOutput));
+            }
+        }
+        return isRequired ? className : `${className}?`;
+    }
     if (propSchema.anyOf) {
         const hasNull = propSchema.anyOf.some((s) => typeof s === "object" && (s as JSONSchema7).type === "null");
         const nonNull = propSchema.anyOf.filter((s) => typeof s === "object" && (s as JSONSchema7).type !== "null");
@@ -535,6 +556,18 @@ function resolveSessionPropertyType(
     }
     if (propSchema.type === "array" && propSchema.items) {
         const items = propSchema.items as JSONSchema7;
+        // Handle $ref in array items
+        if (items.$ref) {
+            const typeName = refTypeName(items.$ref);
+            const className = typeToClassName(typeName);
+            if (!nestedClasses.has(className)) {
+                const refSchema = resolveRef(items.$ref, sessionDefinitions);
+                if (refSchema) {
+                    nestedClasses.set(className, generateNestedClass(className, refSchema, knownTypes, nestedClasses, enumOutput));
+                }
+            }
+            return isRequired ? `${className}[]` : `${className}[]?`;
+        }
         // Array of discriminated union (anyOf with shared discriminator)
         if (items.anyOf && Array.isArray(items.anyOf)) {
             const variants = items.anyOf.filter((v): v is JSONSchema7 => typeof v === "object");
@@ -595,6 +628,7 @@ function generateDataClass(variant: EventVariant, knownTypes: Map<string, string
 
 function generateSessionEventsCode(schema: JSONSchema7): string {
     generatedEnums.clear();
+    sessionDefinitions = schema.definitions as Record<string, JSONSchema7Definition> || {};
     const variants = extractEventVariants(schema);
     const knownTypes = new Map<string, string>();
     const nestedClasses = new Map<string, string>();
@@ -706,6 +740,9 @@ let experimentalRpcTypes = new Set<string>();
 let rpcKnownTypes = new Map<string, string>();
 let rpcEnumOutput: string[] = [];
 
+/** Schema definitions available during RPC generation (for $ref resolution). */
+let rpcDefinitions: Record<string, JSONSchema7Definition> = {};
+
 function singularPascal(s: string): string {
     const p = toPascalCase(s);
     if (p.endsWith("ies")) return `${p.slice(0, -3)}y`;
@@ -734,6 +771,16 @@ function stableStringify(value: unknown): string {
 }
 
 function resolveRpcType(schema: JSONSchema7, isRequired: boolean, parentClassName: string, propName: string, classes: string[]): string {
+    // Handle $ref by resolving against schema definitions and generating the referenced class
+    if (schema.$ref) {
+        const typeName = refTypeName(schema.$ref);
+        const refSchema = resolveRef(schema.$ref, rpcDefinitions);
+        if (refSchema && !emittedRpcClasses.has(typeName)) {
+            const cls = emitRpcClass(typeName, refSchema, "public", classes);
+            if (cls) classes.push(cls);
+        }
+        return isRequired ? typeName : `${typeName}?`;
+    }
     // Handle anyOf: [T, null] → T? (nullable typed property)
     if (schema.anyOf) {
         const hasNull = schema.anyOf.some((s) => typeof s === "object" && (s as JSONSchema7).type === "null");
@@ -754,6 +801,16 @@ function resolveRpcType(schema: JSONSchema7, isRequired: boolean, parentClassNam
     }
     if (schema.type === "array" && schema.items) {
         const items = schema.items as JSONSchema7;
+        // Handle $ref in array items
+        if (items.$ref) {
+            const typeName = refTypeName(items.$ref);
+            const refSchema = resolveRef(items.$ref, rpcDefinitions);
+            if (refSchema && !emittedRpcClasses.has(typeName)) {
+                const cls = emitRpcClass(typeName, refSchema, "public", classes);
+                if (cls) classes.push(cls);
+            }
+            return isRequired ? `List<${typeName}>` : `List<${typeName}>?`;
+        }
         if (items.type === "object" && items.properties) {
             const itemClass = (items.title as string) ?? singularPascal(propName);
             classes.push(emitRpcClass(itemClass, items, "public", classes));
@@ -1197,6 +1254,7 @@ function generateRpcCode(schema: ApiSchema): string {
     rpcKnownTypes.clear();
     rpcEnumOutput = [];
     generatedEnums.clear(); // Clear shared enum deduplication map
+    rpcDefinitions = collectDefinitions(schema as Record<string, unknown>);
     const classes: string[] = [];
 
     let serverRpcParts: string[] = [];
diff --git a/scripts/codegen/go.ts b/scripts/codegen/go.ts
index 101702f18..90e49d302 100644
--- a/scripts/codegen/go.ts
+++ b/scripts/codegen/go.ts
@@ -8,7 +8,7 @@
 
 import { execFile } from "child_process";
 import fs from "fs/promises";
-import type { JSONSchema7 } from "json-schema";
+import type { JSONSchema7, JSONSchema7Definition } from "json-schema";
 import { FetchingJSONSchemaStore, InputData, JSONSchemaInput, quicktype } from "quicktype-core";
 import { promisify } from "util";
 import {
@@ -19,6 +19,9 @@ import {
     isRpcMethod,
     postProcessSchema,
     writeGeneratedFile,
+    collectDefinitions,
+    refTypeName,
+    resolveRef,
     type ApiSchema,
     type RpcMethod,
 } from "./utils.js";
@@ -152,6 +155,7 @@ interface GoCodegenCtx {
     enums: string[];
     enumsByValues: Map<string, string>; // sorted-values-key → enumName
     generatedNames: Set<string>;
+    definitions?: Record<string, JSONSchema7Definition>;
 }
 
 function extractGoEventVariants(schema: JSONSchema7): GoEventVariant[] {
@@ -257,6 +261,21 @@ function resolveGoPropertyType(
 ): string {
     const nestedName = parentTypeName + toGoFieldName(jsonPropName);
 
+    // Handle $ref — resolve the reference and generate the referenced type
+    if (propSchema.$ref && typeof propSchema.$ref === "string") {
+        const typeName = toGoFieldName(refTypeName(propSchema.$ref));
+        const resolved = resolveRef(propSchema.$ref, ctx.definitions);
+        if (resolved) {
+            if (resolved.enum) {
+                return getOrCreateGoEnum(typeName, resolved.enum as string[], ctx, resolved.description);
+            }
+            emitGoStruct(typeName, resolved, ctx);
+            return isRequired ? typeName : `*${typeName}`;
+        }
+        // Fallback: use the type name directly
+        return isRequired ? typeName : `*${typeName}`;
+    }
+
     // Handle anyOf
     if (propSchema.anyOf) {
         const nonNull = (propSchema.anyOf as JSONSchema7[]).filter((s) => s.type !== "null");
@@ -514,6 +533,7 @@ function generateGoSessionEventsCode(schema: JSONSchema7): string {
         enums: [],
         enumsByValues: new Map(),
         generatedNames: new Set(),
+        definitions: schema.definitions as Record<string, JSONSchema7Definition> | undefined,
     };
 
     // Generate per-event data structs
@@ -802,10 +822,12 @@ async function generateRpc(schemaPath?: string): Promise<void> {
         ...collectRpcMethods(schema.clientSession || {}),
     ];
 
-    // Build a combined schema for quicktype - prefix types to avoid conflicts
+    // Build a combined schema for quicktype — prefix types to avoid conflicts.
+    // Include shared definitions from the API schema for $ref resolution.
+    const sharedDefs = collectDefinitions(schema as Record<string, unknown>);
     const combinedSchema: JSONSchema7 = {
         $schema: "http://json-schema.org/draft-07/schema#",
-        definitions: {},
+        definitions: { ...sharedDefs },
     };
 
     for (const method of allMethods) {
@@ -832,10 +854,14 @@ async function generateRpc(schemaPath?: string): Promise<void> {
         }
     }
 
-    // Generate types via quicktype
+    // Generate types via quicktype — include all definitions in each source for $ref resolution
     const schemaInput = new JSONSchemaInput(new FetchingJSONSchemaStore());
     for (const [name, def] of Object.entries(combinedSchema.definitions!)) {
-        await schemaInput.addSource({ name, schema: JSON.stringify(def) });
+        const schemaWithDefs: JSONSchema7 = {
+            ...(typeof def === "object" ? (def as JSONSchema7) : {}),
+            definitions: combinedSchema.definitions,
+        };
+        await schemaInput.addSource({ name, schema: JSON.stringify(schemaWithDefs) });
     }
 
     const inputData = new InputData();
diff --git a/scripts/codegen/python.ts b/scripts/codegen/python.ts
index 2aa593c5d..41ce7c2b1 100644
--- a/scripts/codegen/python.ts
+++ b/scripts/codegen/python.ts
@@ -15,6 +15,7 @@ import {
     isRpcMethod,
     postProcessSchema,
     writeGeneratedFile,
+    collectDefinitions,
     isRpcMethod,
     isNodeFullyExperimental,
     type ApiSchema,
@@ -151,11 +152,20 @@ async function generateSessionEvents(schemaPath?: string): Promise<void> {
 
     const resolvedPath = schemaPath ?? (await getSessionEventsSchemaPath());
     const schema = JSON.parse(await fs.readFile(resolvedPath, "utf-8")) as JSONSchema7;
-    const resolvedSchema = (schema.definitions?.SessionEvent as JSONSchema7) || schema;
-    const processed = postProcessSchema(resolvedSchema);
+    const processed = postProcessSchema(schema);
+
+    // Extract SessionEvent as root but keep all other definitions for $ref resolution
+    const sessionEventDef = (processed.definitions?.SessionEvent as JSONSchema7) || processed;
+    const otherDefs = Object.fromEntries(
+        Object.entries(processed.definitions || {}).filter(([key]) => key !== "SessionEvent")
+    );
+    const schemaForQuicktype: JSONSchema7 = {
+        ...sessionEventDef,
+        ...(Object.keys(otherDefs).length > 0 ? { definitions: otherDefs } : {}),
+    };
 
     const schemaInput = new JSONSchemaInput(new FetchingJSONSchemaStore());
-    await schemaInput.addSource({ name: "SessionEvent", schema: JSON.stringify(processed) });
+    await schemaInput.addSource({ name: "SessionEvent", schema: JSON.stringify(schemaForQuicktype) });
 
     const inputData = new InputData();
     inputData.addInput(schemaInput);
@@ -214,10 +224,11 @@ async function generateRpc(schemaPath?: string): Promise<void> {
         ...collectRpcMethods(schema.clientSession || {}),
     ];
 
-    // Build a combined schema for quicktype
+    // Build a combined schema for quicktype, including shared definitions from the API schema
+    const sharedDefs = collectDefinitions(schema as Record<string, unknown>);
     const combinedSchema: JSONSchema7 = {
         $schema: "http://json-schema.org/draft-07/schema#",
-        definitions: {},
+        definitions: { ...sharedDefs },
     };
 
     for (const method of allMethods) {
@@ -243,10 +254,14 @@ async function generateRpc(schemaPath?: string): Promise<void> {
         }
     }
 
-    // Generate types via quicktype
+    // Generate types via quicktype — include all definitions in each source for $ref resolution
     const schemaInput = new JSONSchemaInput(new FetchingJSONSchemaStore());
     for (const [name, def] of Object.entries(combinedSchema.definitions!)) {
-        await schemaInput.addSource({ name, schema: JSON.stringify(def) });
+        const schemaWithDefs: JSONSchema7 = {
+            ...(typeof def === "object" ? (def as JSONSchema7) : {}),
+            definitions: combinedSchema.definitions,
+        };
+        await schemaInput.addSource({ name, schema: JSON.stringify(schemaWithDefs) });
     }
 
     const inputData = new InputData();
diff --git a/scripts/codegen/typescript.ts b/scripts/codegen/typescript.ts
index e5e82bdc6..5d6c49d26 100644
--- a/scripts/codegen/typescript.ts
+++ b/scripts/codegen/typescript.ts
@@ -14,6 +14,7 @@ import {
     getApiSchemaPath,
     postProcessSchema,
     writeGeneratedFile,
+    collectDefinitions,
     isRpcMethod,
     isNodeFullyExperimental,
     type ApiSchema,
@@ -88,32 +89,58 @@ import type { MessageConnection } from "vscode-jsonrpc/node.js";
     const allMethods = [...collectRpcMethods(schema.server || {}), ...collectRpcMethods(schema.session || {})];
     const clientSessionMethods = collectRpcMethods(schema.clientSession || {});
 
+    // Build a single combined schema with shared definitions and all method types.
+    // This ensures $ref-referenced types are generated exactly once.
+    const sharedDefs = collectDefinitions(schema as Record<string, unknown>);
+    const combinedSchema: JSONSchema7 = {
+        $schema: "http://json-schema.org/draft-07/schema#",
+        type: "object",
+        definitions: { ...sharedDefs },
+    };
+
+    // Track which type names come from experimental methods for JSDoc annotations.
+    const experimentalTypes = new Set<string>();
+
     for (const method of [...allMethods, ...clientSessionMethods]) {
         if (method.result) {
-            const compiled = await compile(method.result, resultTypeName(method.rpcMethod), {
-                bannerComment: "",
-                additionalProperties: false,
-            });
+            combinedSchema.definitions![resultTypeName(method.rpcMethod)] = method.result;
             if (method.stability === "experimental") {
-                lines.push("/** @experimental */");
+                experimentalTypes.add(resultTypeName(method.rpcMethod));
             }
-            lines.push(compiled.trim());
-            lines.push("");
         }
 
         if (method.params?.properties && Object.keys(method.params.properties).length > 0) {
-            const paramsCompiled = await compile(method.params, paramsTypeName(method.rpcMethod), {
-                bannerComment: "",
-                additionalProperties: false,
-            });
+            combinedSchema.definitions![paramsTypeName(method.rpcMethod)] = method.params;
             if (method.stability === "experimental") {
-                lines.push("/** @experimental */");
+                experimentalTypes.add(paramsTypeName(method.rpcMethod));
             }
-            lines.push(paramsCompiled.trim());
-            lines.push("");
         }
     }
 
+    const compiled = await compile(combinedSchema, "_RpcSchemaRoot", {
+        bannerComment: "",
+        additionalProperties: false,
+        unreachableDefinitions: true,
+    });
+
+    // Strip the placeholder root type and keep only the definition-generated types
+    const strippedTs = compiled
+        .replace(/export interface _RpcSchemaRoot\s*\{[^}]*\}\s*/g, "")
+        .trim();
+
+    if (strippedTs) {
+        // Add @experimental JSDoc annotations for types from experimental methods
+        let annotatedTs = strippedTs;
+        for (const expType of experimentalTypes) {
+            annotatedTs = annotatedTs.replace(
+                new RegExp(`(^|\\n)(export (?:interface|type) ${expType}\\b)`, "m"),
+                `$1/** @experimental */\n$2`
+            );
+        }
+        lines.push(annotatedTs);
+        lines.push("");
+    }
+
     // Generate factory functions
     if (schema.server) {
         lines.push(`/** Create typed server-scoped RPC methods (no session required). */`);
diff --git a/scripts/codegen/utils.ts b/scripts/codegen/utils.ts
index 1e95b4dd4..a568bcf9c 100644
--- a/scripts/codegen/utils.ts
+++ b/scripts/codegen/utils.ts
@@ -56,6 +56,15 @@ export function postProcessSchema(schema: JSONSchema7): JSONSchema7 {
 
     const processed: JSONSchema7 = { ...schema };
 
+    // Normalize $defs → definitions for draft 2019+ compatibility
+    if ("$defs" in processed && !processed.definitions) {
+        processed.definitions = (processed as Record<string, unknown>).$defs as Record<
+            string,
+            JSONSchema7Definition
+        >;
+        delete (processed as Record<string, unknown>).$defs;
+    }
+
     if ("const" in processed && typeof processed.const === "boolean") {
         processed.enum = [processed.const];
         delete processed.const;
@@ -130,6 +139,8 @@ export interface RpcMethod {
 }
 
 export interface ApiSchema {
+    definitions?: Record<string, JSONSchema7Definition>;
+    $defs?: Record<string, JSONSchema7Definition>;
     server?: Record<string, unknown>;
     session?: Record<string, unknown>;
     clientSession?: Record<string, unknown>;
@@ -153,3 +164,29 @@ export function isNodeFullyExperimental(node: Record<string, unknown>): boolean
     })(node);
     return methods.length > 0 && methods.every(m => m.stability === "experimental");
 }
+
+// ── $ref resolution ─────────────────────────────────────────────────────────
+
+/** Extract the type name from a `$ref` path (e.g. "#/definitions/Model" → "Model"). */
+export function refTypeName(ref: string): string {
+    return ref.split("/").pop()!;
+}
+
+/** Resolve a `$ref` path against a definitions map, returning the referenced schema. */
+export function resolveRef(
+    ref: string,
+    definitions: Record<string, JSONSchema7Definition> | undefined
+): JSONSchema7 | undefined {
+    const match = ref.match(/^#\/(definitions|\$defs)\/(.+)$/);
+    if (!match || !definitions) return undefined;
+    const def = definitions[match[2]];
+    return typeof def === "object" ? (def as JSONSchema7) : undefined;
+}
+
+/** Collect the shared definitions from a schema (handles both `definitions` and `$defs`). */
+export function collectDefinitions(
+    schema: Record<string, unknown>
+): Record<string, JSONSchema7Definition> {
+    const defs = (schema.definitions ?? schema.$defs ?? {}) as Record<string, JSONSchema7Definition>;
+    return { ...defs }
+}

From 17ba6bcea00da66a0070a2bb1a8f4a3e3d6ce1bf Mon Sep 17 00:00:00 2001
From: Stephen Toub <stoub@microsoft.com>
Date: Sun, 12 Apr 2026 23:06:02 -0400
Subject: [PATCH 2/3] Align rebased ref generator changes

Keep the rebased $ref generator follow-up aligned with the latest C# typing changes and clean up the Python/TypeScript generator adjustments.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 scripts/codegen/csharp.ts     | 120 +++++++++++++---------------------
 scripts/codegen/python.ts     |   1 -
 scripts/codegen/typescript.ts |   5 ++
 3 files changed, 52 insertions(+), 74 deletions(-)

diff --git a/scripts/codegen/csharp.ts b/scripts/codegen/csharp.ts
index eb70e0b4a..9c4b26f55 100644
--- a/scripts/codegen/csharp.ts
+++ b/scripts/codegen/csharp.ts
@@ -16,6 +16,7 @@ import {
     getApiSchemaPath,
     writeGeneratedFile,
     collectDefinitions,
+    postProcessSchema,
     resolveRef,
     refTypeName,
     isRpcMethod,
@@ -512,18 +513,25 @@ function resolveSessionPropertyType(
 ): string {
     // Handle $ref by resolving against schema definitions
     if (propSchema.$ref) {
-        const typeName = refTypeName(propSchema.$ref);
-        const className = typeToClassName(typeName);
-        if (!nestedClasses.has(className)) {
-            const refSchema = resolveRef(propSchema.$ref, sessionDefinitions);
-            if (refSchema) {
-                if (refSchema.enum && Array.isArray(refSchema.enum)) {
-                    return getOrCreateEnum(className, "", refSchema.enum as string[], enumOutput);
-                }
+        const className = typeToClassName(refTypeName(propSchema.$ref));
+        const refSchema = resolveRef(propSchema.$ref, sessionDefinitions);
+        if (!refSchema) {
+            return isRequired ? className : `${className}?`;
+        }
+
+        if (refSchema.enum && Array.isArray(refSchema.enum)) {
+            const enumName = getOrCreateEnum(className, "", refSchema.enum as string[], enumOutput, refSchema.description);
+            return isRequired ? enumName : `${enumName}?`;
+        }
+
+        if (refSchema.type === "object" && refSchema.properties) {
+            if (!nestedClasses.has(className)) {
                 nestedClasses.set(className, generateNestedClass(className, refSchema, knownTypes, nestedClasses, enumOutput));
             }
+            return isRequired ? className : `${className}?`;
         }
-        return isRequired ? className : `${className}?`;
+
+        return resolveSessionPropertyType(refSchema, parentClassName, propName, isRequired, knownTypes, nestedClasses, enumOutput);
     }
     if (propSchema.anyOf) {
         const hasNull = propSchema.anyOf.some((s) => typeof s === "object" && (s as JSONSchema7).type === "null");
@@ -556,40 +564,15 @@ function resolveSessionPropertyType(
     }
     if (propSchema.type === "array" && propSchema.items) {
         const items = propSchema.items as JSONSchema7;
-        // Handle $ref in array items
-        if (items.$ref) {
-            const typeName = refTypeName(items.$ref);
-            const className = typeToClassName(typeName);
-            if (!nestedClasses.has(className)) {
-                const refSchema = resolveRef(items.$ref, sessionDefinitions);
-                if (refSchema) {
-                    nestedClasses.set(className, generateNestedClass(className, refSchema, knownTypes, nestedClasses, enumOutput));
-                }
-            }
-            return isRequired ? `${className}[]` : `${className}[]?`;
-        }
-        // Array of discriminated union (anyOf with shared discriminator)
-        if (items.anyOf && Array.isArray(items.anyOf)) {
-            const variants = items.anyOf.filter((v): v is JSONSchema7 => typeof v === "object");
-            const discriminatorInfo = findDiscriminator(variants);
-            if (discriminatorInfo) {
-                const baseClassName = `${parentClassName}${propName}Item`;
-                const renamedBase = applyTypeRename(baseClassName);
-                const polymorphicCode = generatePolymorphicClasses(baseClassName, discriminatorInfo.property, variants, knownTypes, nestedClasses, enumOutput, items.description);
-                nestedClasses.set(renamedBase, polymorphicCode);
-                return isRequired ? `${renamedBase}[]` : `${renamedBase}[]?`;
-            }
-        }
-        if (items.type === "object" && items.properties) {
-            const itemClassName = `${parentClassName}${propName}Item`;
-            nestedClasses.set(itemClassName, generateNestedClass(itemClassName, items, knownTypes, nestedClasses, enumOutput));
-            return isRequired ? `${itemClassName}[]` : `${itemClassName}[]?`;
-        }
-        if (items.enum && Array.isArray(items.enum)) {
-            const enumName = getOrCreateEnum(parentClassName, `${propName}Item`, items.enum as string[], enumOutput, items.description);
-            return isRequired ? `${enumName}[]` : `${enumName}[]?`;
-        }
-        const itemType = schemaTypeToCSharp(items, true, knownTypes);
+        const itemType = resolveSessionPropertyType(
+            items,
+            parentClassName,
+            `${propName}Item`,
+            true,
+            knownTypes,
+            nestedClasses,
+            enumOutput
+        );
         return isRequired ? `${itemType}[]` : `${itemType}[]?`;
     }
     return schemaTypeToCSharp(propSchema, isRequired, knownTypes);
@@ -725,7 +708,8 @@ export async function generateSessionEvents(schemaPath?: string): Promise<void>
     console.log("C#: generating session-events...");
     const resolvedPath = schemaPath ?? (await getSessionEventsSchemaPath());
     const schema = JSON.parse(await fs.readFile(resolvedPath, "utf-8")) as JSONSchema7;
-    const code = generateSessionEventsCode(schema);
+    const processed = postProcessSchema(schema);
+    const code = generateSessionEventsCode(processed);
     const outPath = await writeGeneratedFile("dotnet/src/Generated/SessionEvents.cs", code);
     console.log(`  ✓ ${outPath}`);
     await formatCSharpFile(outPath);
@@ -773,13 +757,24 @@ function stableStringify(value: unknown): string {
 function resolveRpcType(schema: JSONSchema7, isRequired: boolean, parentClassName: string, propName: string, classes: string[]): string {
     // Handle $ref by resolving against schema definitions and generating the referenced class
     if (schema.$ref) {
-        const typeName = refTypeName(schema.$ref);
+        const typeName = typeToClassName(refTypeName(schema.$ref));
         const refSchema = resolveRef(schema.$ref, rpcDefinitions);
-        if (refSchema && !emittedRpcClasses.has(typeName)) {
+        if (!refSchema) {
+            return isRequired ? typeName : `${typeName}?`;
+        }
+
+        if (refSchema.enum && Array.isArray(refSchema.enum)) {
+            const enumName = getOrCreateEnum(typeName, "", refSchema.enum as string[], rpcEnumOutput, refSchema.description);
+            return isRequired ? enumName : `${enumName}?`;
+        }
+
+        if (refSchema.type === "object" && refSchema.properties) {
             const cls = emitRpcClass(typeName, refSchema, "public", classes);
             if (cls) classes.push(cls);
+            return isRequired ? typeName : `${typeName}?`;
         }
-        return isRequired ? typeName : `${typeName}?`;
+
+        return resolveRpcType(refSchema, isRequired, parentClassName, propName, classes);
     }
     // Handle anyOf: [T, null] → T? (nullable typed property)
     if (schema.anyOf) {
@@ -801,32 +796,17 @@ function resolveRpcType(schema: JSONSchema7, isRequired: boolean, parentClassNam
     }
     if (schema.type === "array" && schema.items) {
         const items = schema.items as JSONSchema7;
-        // Handle $ref in array items
-        if (items.$ref) {
-            const typeName = refTypeName(items.$ref);
-            const refSchema = resolveRef(items.$ref, rpcDefinitions);
-            if (refSchema && !emittedRpcClasses.has(typeName)) {
-                const cls = emitRpcClass(typeName, refSchema, "public", classes);
-                if (cls) classes.push(cls);
-            }
-            return isRequired ? `List<${typeName}>` : `List<${typeName}>?`;
-        }
         if (items.type === "object" && items.properties) {
             const itemClass = (items.title as string) ?? singularPascal(propName);
             classes.push(emitRpcClass(itemClass, items, "public", classes));
             return isRequired ? `IList<${itemClass}>` : `IList<${itemClass}>?`;
         }
-        const itemType = schemaTypeToCSharp(items, true, rpcKnownTypes);
+        const itemType = resolveRpcType(items, true, parentClassName, `${propName}Item`, classes);
         return isRequired ? `IList<${itemType}>` : `IList<${itemType}>?`;
     }
     if (schema.type === "object" && schema.additionalProperties && typeof schema.additionalProperties === "object") {
         const vs = schema.additionalProperties as JSONSchema7;
-        if (vs.type === "object" && vs.properties) {
-            const valClass = `${parentClassName}${propName}Value`;
-            classes.push(emitRpcClass(valClass, vs, "public", classes));
-            return isRequired ? `IDictionary<string, ${valClass}>` : `IDictionary<string, ${valClass}>?`;
-        }
-        const valueType = schemaTypeToCSharp(vs, true, rpcKnownTypes);
+        const valueType = resolveRpcType(vs, true, parentClassName, `${propName}Value`, classes);
         return isRequired ? `IDictionary<string, ${valueType}>` : `IDictionary<string, ${valueType}>?`;
     }
     return schemaTypeToCSharp(schema, isRequired, rpcKnownTypes);
@@ -1007,15 +987,9 @@ function emitServerInstanceMethod(
         if (typeof pSchema !== "object") continue;
         const isReq = requiredSet.has(pName);
         const jsonSchema = pSchema as JSONSchema7;
-        let csType: string;
-        // If the property has an enum, resolve to the generated enum type
-        if (jsonSchema.enum && Array.isArray(jsonSchema.enum) && requestClassName) {
-            const valuesKey = [...jsonSchema.enum].sort().join("|");
-            const match = [...generatedEnums.values()].find((e) => [...e.values].sort().join("|") === valuesKey);
-            csType = match ? (isReq ? match.enumName : `${match.enumName}?`) : schemaTypeToCSharp(jsonSchema, isReq, rpcKnownTypes);
-        } else {
-            csType = schemaTypeToCSharp(jsonSchema, isReq, rpcKnownTypes);
-        }
+        const csType = requestClassName
+            ? resolveRpcType(jsonSchema, isReq, requestClassName, toPascalCase(pName), classes)
+            : schemaTypeToCSharp(jsonSchema, isReq, rpcKnownTypes);
         sigParams.push(`${csType} ${pName}${isReq ? "" : " = null"}`);
         bodyAssignments.push(`${toPascalCase(pName)} = ${pName}`);
     }
diff --git a/scripts/codegen/python.ts b/scripts/codegen/python.ts
index 41ce7c2b1..d9c64a390 100644
--- a/scripts/codegen/python.ts
+++ b/scripts/codegen/python.ts
@@ -16,7 +16,6 @@ import {
     postProcessSchema,
     writeGeneratedFile,
     collectDefinitions,
-    isRpcMethod,
     isNodeFullyExperimental,
     type ApiSchema,
     type RpcMethod,
diff --git a/scripts/codegen/typescript.ts b/scripts/codegen/typescript.ts
index 5d6c49d26..840cce640 100644
--- a/scripts/codegen/typescript.ts
+++ b/scripts/codegen/typescript.ts
@@ -125,7 +125,12 @@ import type { MessageConnection } from "vscode-jsonrpc/node.js";
 
     // Strip the placeholder root type and keep only the definition-generated types
     const strippedTs = compiled
+        .replace(
+            /\/\*\*\n \* This (?:interface|type) was referenced by `_RpcSchemaRoot`'s JSON-Schema\n \* via the `definition` "[^"]+"\.\n \*\/\n/g,
+            "\n"
+        )
         .replace(/export interface _RpcSchemaRoot\s*\{[^}]*\}\s*/g, "")
+        .replace(/export type _RpcSchemaRoot = [^;]+;\s*/g, "")
         .trim();
 
     if (strippedTs) {

From 7ab5c81e30256480e38505e06df90a69f94e3aef Mon Sep 17 00:00:00 2001
From: Stephen Toub <stoub@microsoft.com>
Date: Sun, 12 Apr 2026 23:14:06 -0400
Subject: [PATCH 3/3] Handle mixed schema definitions in codegen

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 scripts/codegen/csharp.ts     |  7 +++---
 scripts/codegen/go.ts         | 19 +++++++++-------
 scripts/codegen/python.ts     | 19 +++++++++-------
 scripts/codegen/typescript.ts | 13 +++++++----
 scripts/codegen/utils.ts      | 43 +++++++++++++++++++++++------------
 5 files changed, 63 insertions(+), 38 deletions(-)

diff --git a/scripts/codegen/csharp.ts b/scripts/codegen/csharp.ts
index 9c4b26f55..21701703e 100644
--- a/scripts/codegen/csharp.ts
+++ b/scripts/codegen/csharp.ts
@@ -325,7 +325,7 @@ function getOrCreateEnum(parentClassName: string, propName: string, values: stri
 }
 
 function extractEventVariants(schema: JSONSchema7): EventVariant[] {
-    const sessionEvent = schema.definitions?.SessionEvent as JSONSchema7;
+    const sessionEvent = collectDefinitions(schema as Record<string, unknown>).SessionEvent as JSONSchema7;
     if (!sessionEvent?.anyOf) throw new Error("Schema must have SessionEvent definition with anyOf");
 
     return sessionEvent.anyOf
@@ -611,14 +611,15 @@ function generateDataClass(variant: EventVariant, knownTypes: Map<string, string
 
 function generateSessionEventsCode(schema: JSONSchema7): string {
     generatedEnums.clear();
-    sessionDefinitions = schema.definitions as Record<string, JSONSchema7Definition> || {};
+    sessionDefinitions = collectDefinitions(schema as Record<string, unknown>);
     const variants = extractEventVariants(schema);
     const knownTypes = new Map<string, string>();
     const nestedClasses = new Map<string, string>();
     const enumOutput: string[] = [];
 
     // Extract descriptions for base class properties from the first variant
-    const firstVariant = (schema.definitions?.SessionEvent as JSONSchema7)?.anyOf?.[0];
+    const sessionEventDefinition = sessionDefinitions.SessionEvent;
+    const firstVariant = typeof sessionEventDefinition === "object" ? (sessionEventDefinition as JSONSchema7).anyOf?.[0] : undefined;
     const baseProps = typeof firstVariant === "object" && firstVariant?.properties ? firstVariant.properties : {};
     const baseDesc = (name: string) => {
         const prop = baseProps[name];
diff --git a/scripts/codegen/go.ts b/scripts/codegen/go.ts
index 90e49d302..378f0a79a 100644
--- a/scripts/codegen/go.ts
+++ b/scripts/codegen/go.ts
@@ -20,6 +20,7 @@ import {
     postProcessSchema,
     writeGeneratedFile,
     collectDefinitions,
+    withSharedDefinitions,
     refTypeName,
     resolveRef,
     type ApiSchema,
@@ -825,10 +826,12 @@ async function generateRpc(schemaPath?: string): Promise<void> {
     // Build a combined schema for quicktype — prefix types to avoid conflicts.
     // Include shared definitions from the API schema for $ref resolution.
     const sharedDefs = collectDefinitions(schema as Record<string, unknown>);
-    const combinedSchema: JSONSchema7 = {
-        $schema: "http://json-schema.org/draft-07/schema#",
-        definitions: { ...sharedDefs },
-    };
+    const combinedSchema = withSharedDefinitions(
+        {
+            $schema: "http://json-schema.org/draft-07/schema#",
+        },
+        sharedDefs
+    );
 
     for (const method of allMethods) {
         const baseName = toPascalCase(method.rpcMethod);
@@ -857,10 +860,10 @@ async function generateRpc(schemaPath?: string): Promise<void> {
     // Generate types via quicktype — include all definitions in each source for $ref resolution
     const schemaInput = new JSONSchemaInput(new FetchingJSONSchemaStore());
     for (const [name, def] of Object.entries(combinedSchema.definitions!)) {
-        const schemaWithDefs: JSONSchema7 = {
-            ...(typeof def === "object" ? (def as JSONSchema7) : {}),
-            definitions: combinedSchema.definitions,
-        };
+        const schemaWithDefs = withSharedDefinitions(
+            typeof def === "object" ? (def as JSONSchema7) : {},
+            combinedSchema.definitions
+        );
         await schemaInput.addSource({ name, schema: JSON.stringify(schemaWithDefs) });
     }
 
diff --git a/scripts/codegen/python.ts b/scripts/codegen/python.ts
index d9c64a390..5d84b0538 100644
--- a/scripts/codegen/python.ts
+++ b/scripts/codegen/python.ts
@@ -16,6 +16,7 @@ import {
     postProcessSchema,
     writeGeneratedFile,
     collectDefinitions,
+    withSharedDefinitions,
     isNodeFullyExperimental,
     type ApiSchema,
     type RpcMethod,
@@ -225,10 +226,12 @@ async function generateRpc(schemaPath?: string): Promise<void> {
 
     // Build a combined schema for quicktype, including shared definitions from the API schema
     const sharedDefs = collectDefinitions(schema as Record<string, unknown>);
-    const combinedSchema: JSONSchema7 = {
-        $schema: "http://json-schema.org/draft-07/schema#",
-        definitions: { ...sharedDefs },
-    };
+    const combinedSchema = withSharedDefinitions(
+        {
+            $schema: "http://json-schema.org/draft-07/schema#",
+        },
+        sharedDefs
+    );
 
     for (const method of allMethods) {
         const baseName = toPascalCase(method.rpcMethod);
@@ -256,10 +259,10 @@ async function generateRpc(schemaPath?: string): Promise<void> {
     // Generate types via quicktype — include all definitions in each source for $ref resolution
     const schemaInput = new JSONSchemaInput(new FetchingJSONSchemaStore());
     for (const [name, def] of Object.entries(combinedSchema.definitions!)) {
-        const schemaWithDefs: JSONSchema7 = {
-            ...(typeof def === "object" ? (def as JSONSchema7) : {}),
-            definitions: combinedSchema.definitions,
-        };
+        const schemaWithDefs = withSharedDefinitions(
+            typeof def === "object" ? (def as JSONSchema7) : {},
+            combinedSchema.definitions
+        );
         await schemaInput.addSource({ name, schema: JSON.stringify(schemaWithDefs) });
     }
 
diff --git a/scripts/codegen/typescript.ts b/scripts/codegen/typescript.ts
index 840cce640..3032d83d7 100644
--- a/scripts/codegen/typescript.ts
+++ b/scripts/codegen/typescript.ts
@@ -15,6 +15,7 @@ import {
     postProcessSchema,
     writeGeneratedFile,
     collectDefinitions,
+    withSharedDefinitions,
     isRpcMethod,
     isNodeFullyExperimental,
     type ApiSchema,
@@ -92,11 +93,13 @@ import type { MessageConnection } from "vscode-jsonrpc/node.js";
     // Build a single combined schema with shared definitions and all method types.
     // This ensures $ref-referenced types are generated exactly once.
     const sharedDefs = collectDefinitions(schema as Record<string, unknown>);
-    const combinedSchema: JSONSchema7 = {
-        $schema: "http://json-schema.org/draft-07/schema#",
-        type: "object",
-        definitions: { ...sharedDefs },
-    };
+    const combinedSchema = withSharedDefinitions(
+        {
+            $schema: "http://json-schema.org/draft-07/schema#",
+            type: "object",
+        },
+        sharedDefs
+    );
 
     // Track which type names come from experimental methods for JSDoc annotations.
     const experimentalTypes = new Set<string>();
diff --git a/scripts/codegen/utils.ts b/scripts/codegen/utils.ts
index a568bcf9c..c64590fe8 100644
--- a/scripts/codegen/utils.ts
+++ b/scripts/codegen/utils.ts
@@ -24,6 +24,15 @@ export const REPO_ROOT = path.resolve(__dirname, "../..");
 /** Event types to exclude from generation (internal/legacy types) */
 export const EXCLUDED_EVENT_TYPES = new Set(["session.import_legacy"]);
 
+export interface JSONSchema7WithDefs extends JSONSchema7 {
+    $defs?: Record<string, JSONSchema7Definition>;
+}
+
+export type SchemaWithSharedDefinitions<T extends JSONSchema7 = JSONSchema7> = T & {
+    definitions: Record<string, JSONSchema7Definition>;
+    $defs: Record<string, JSONSchema7Definition>;
+};
+
 // ── Schema paths ────────────────────────────────────────────────────────────
 
 export async function getSessionEventsSchemaPath(): Promise<string> {
@@ -54,16 +63,7 @@ export async function getApiSchemaPath(cliArg?: string): Promise<string> {
 export function postProcessSchema(schema: JSONSchema7): JSONSchema7 {
     if (typeof schema !== "object" || schema === null) return schema;
 
-    const processed: JSONSchema7 = { ...schema };
-
-    // Normalize $defs → definitions for draft 2019+ compatibility
-    if ("$defs" in processed && !processed.definitions) {
-        processed.definitions = (processed as Record<string, unknown>).$defs as Record<
-            string,
-            JSONSchema7Definition
-        >;
-        delete (processed as Record<string, unknown>).$defs;
-    }
+    const processed = { ...schema } as JSONSchema7WithDefs;
 
     if ("const" in processed && typeof processed.const === "boolean") {
         processed.enum = [processed.const];
@@ -105,12 +105,14 @@ export function postProcessSchema(schema: JSONSchema7): JSONSchema7 {
         }
     }
 
-    if (processed.definitions) {
+    const definitions = collectDefinitions(processed as Record<string, unknown>);
+    if (Object.keys(definitions).length > 0) {
         const newDefs: Record<string, JSONSchema7Definition> = {};
-        for (const [key, value] of Object.entries(processed.definitions)) {
+        for (const [key, value] of Object.entries(definitions)) {
             newDefs[key] = typeof value === "object" ? postProcessSchema(value as JSONSchema7) : value;
         }
         processed.definitions = newDefs;
+        processed.$defs = newDefs;
     }
 
     if (typeof processed.additionalProperties === "object") {
@@ -187,6 +189,19 @@ export function resolveRef(
 export function collectDefinitions(
     schema: Record<string, unknown>
 ): Record<string, JSONSchema7Definition> {
-    const defs = (schema.definitions ?? schema.$defs ?? {}) as Record<string, JSONSchema7Definition>;
-    return { ...defs }
+    const legacyDefinitions = (schema.definitions ?? {}) as Record<string, JSONSchema7Definition>;
+    const draft2019Definitions = (schema.$defs ?? {}) as Record<string, JSONSchema7Definition>;
+    return { ...draft2019Definitions, ...legacyDefinitions };
+}
+
+export function withSharedDefinitions<T extends JSONSchema7>(
+    schema: T,
+    definitions: Record<string, JSONSchema7Definition>
+): SchemaWithSharedDefinitions<T> {
+    const sharedDefinitions = { ...definitions };
+    return {
+        ...schema,
+        definitions: sharedDefinitions,
+        $defs: sharedDefinitions,
+    };
 }