Add $ref support to all four language code generators

scripts/codegen/csharp.ts

@@ -10,11 +10,15 @@ 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,
+    postProcessSchema,
+    resolveRef,
+    refTypeName,
     isRpcMethod,
     isNodeFullyExperimental,
     EXCLUDED_EVENT_TYPES,
@@ -297,6 +301,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) {
@@ -318,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
@@ -504,6 +511,28 @@ function resolveSessionPropertyType(
     nestedClasses: Map<string, string>,
     enumOutput: string[]
 ): string {
+    // Handle $ref by resolving against schema definitions
+    if (propSchema.$ref) {
+        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 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");
         const nonNull = propSchema.anyOf.filter((s) => typeof s === "object" && (s as JSONSchema7).type !== "null");
@@ -535,28 +564,15 @@ function resolveSessionPropertyType(
     }
     if (propSchema.type === "array" && propSchema.items) {
         const items = propSchema.items as JSONSchema7;
-        // 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);
@@ -595,13 +611,15 @@ function generateDataClass(variant: EventVariant, knownTypes: Map<string, string
 
 function generateSessionEventsCode(schema: JSONSchema7): string {
     generatedEnums.clear();
+    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];
@@ -691,7 +709,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);
@@ -706,6 +725,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 +756,27 @@ 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 = typeToClassName(refTypeName(schema.$ref));
+        const refSchema = resolveRef(schema.$ref, rpcDefinitions);
+        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 resolveRpcType(refSchema, isRequired, parentClassName, propName, classes);
+    }
     // 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");
@@ -759,17 +802,12 @@ function resolveRpcType(schema: JSONSchema7, isRequired: boolean, parentClassNam
             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);
@@ -950,15 +988,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}`);
     }
@@ -1197,6 +1229,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[] = [];

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,10 @@ import {
     isRpcMethod,
     postProcessSchema,
     writeGeneratedFile,
+    collectDefinitions,
+    withSharedDefinitions,
+    refTypeName,
+    resolveRef,
     type ApiSchema,
     type RpcMethod,
 } from "./utils.js";
@@ -152,6 +156,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 +262,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 +534,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,11 +823,15 @@ async function generateRpc(schemaPath?: string): Promise<void> {
         ...collectRpcMethods(schema.clientSession || {}),
     ];
 
-    // Build a combined schema for quicktype - prefix types to avoid conflicts
-    const combinedSchema: JSONSchema7 = {
-        $schema: "http://json-schema.org/draft-07/schema#",
-        definitions: {},
-    };
+    // 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 = withSharedDefinitions(
+        {
+            $schema: "http://json-schema.org/draft-07/schema#",
+        },
+        sharedDefs
+    );
 
     for (const method of allMethods) {
         const baseName = toPascalCase(method.rpcMethod);
@@ -832,10 +857,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 = withSharedDefinitions(
+            typeof def === "object" ? (def as JSONSchema7) : {},
+            combinedSchema.definitions
+        );
+        await schemaInput.addSource({ name, schema: JSON.stringify(schemaWithDefs) });
     }
 
     const inputData = new InputData();

scripts/codegen/python.ts

@@ -15,7 +15,8 @@ import {
     isRpcMethod,
     postProcessSchema,
     writeGeneratedFile,
-    isRpcMethod,
+    collectDefinitions,
+    withSharedDefinitions,
     isNodeFullyExperimental,
     type ApiSchema,
     type RpcMethod,
@@ -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,11 +224,14 @@ async function generateRpc(schemaPath?: string): Promise<void> {
         ...collectRpcMethods(schema.clientSession || {}),
     ];
 
-    // Build a combined schema for quicktype
-    const combinedSchema: JSONSchema7 = {
-        $schema: "http://json-schema.org/draft-07/schema#",
-        definitions: {},
-    };
+    // Build a combined schema for quicktype, including shared definitions from the API schema
+    const sharedDefs = collectDefinitions(schema as Record<string, unknown>);
+    const combinedSchema = withSharedDefinitions(
+        {
+            $schema: "http://json-schema.org/draft-07/schema#",
+        },
+        sharedDefs
+    );
 
     for (const method of allMethods) {
         const baseName = toPascalCase(method.rpcMethod);
@@ -243,10 +256,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 = withSharedDefinitions(
+            typeof def === "object" ? (def as JSONSchema7) : {},
+            combinedSchema.definitions
+        );
+        await schemaInput.addSource({ name, schema: JSON.stringify(schemaWithDefs) });
     }
 
     const inputData = new InputData();