Handle mixed schema definitions in codegen

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: stephentoub

scripts/codegen/csharp.ts

@@ -326,7 +326,7 @@ let sessionDefinitions: Record<string, JSONSchema7Definition> = {};
 }
 
 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];

scripts/codegen/go.ts

@@ -23,6 +23,7 @@ import {
     postProcessSchema,
     writeGeneratedFile,
     collectDefinitions,
+    withSharedDefinitions,
     refTypeName,
     resolveRef,
     type ApiSchema,
@@ -881,10 +882,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) {
         if (isVoidSchema(method.result)) {
@@ -918,13 +921,11 @@ async function generateRpc(schemaPath?: string): Promise<void> {
     // Generate types via quicktype
     const schemaInput = new JSONSchemaInput(new FetchingJSONSchemaStore());
     for (const [name, def] of Object.entries(rootDefinitions)) {
-        await schemaInput.addSource({
-            name,
-            schema: JSON.stringify({
-                ...def,
-                definitions: allDefinitions,
-            }),
-        });
+        const schemaWithDefs = withSharedDefinitions(
+            typeof def === "object" ? (def as JSONSchema7) : {},
+            allDefinitions
+        );
+        await schemaInput.addSource({ name, schema: JSON.stringify(schemaWithDefs) });
     }
 
     const inputData = new InputData();

scripts/codegen/python.ts

@@ -22,6 +22,8 @@ import {
     isNodeFullyExperimental,
     postProcessSchema,
     writeGeneratedFile,
+    collectDefinitions,
+    withSharedDefinitions,
     type ApiSchema,
     type RpcMethod,
 } from "./utils.js";
@@ -1275,10 +1277,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) {
         if (!isVoidSchema(method.result)) {
@@ -1308,13 +1312,11 @@ async function generateRpc(schemaPath?: string): Promise<void> {
     // Generate types via quicktype
     const schemaInput = new JSONSchemaInput(new FetchingJSONSchemaStore());
     for (const [name, def] of Object.entries(rootDefinitions)) {
-        await schemaInput.addSource({
-            name,
-            schema: JSON.stringify({
-                ...def,
-                definitions: allDefinitions,
-            }),
-        });
+        const schemaWithDefs = withSharedDefinitions(
+            typeof def === "object" ? (def as JSONSchema7) : {},
+            allDefinitions
+        );
+        await schemaInput.addSource({ name, schema: JSON.stringify(schemaWithDefs) });
     }
 
     const inputData = new InputData();

scripts/codegen/typescript.ts

@@ -16,6 +16,7 @@ import {
     postProcessSchema,
     writeGeneratedFile,
     collectDefinitions,
+    withSharedDefinitions,
     isRpcMethod,
     isNodeFullyExperimental,
     isVoidSchema,
@@ -180,11 +181,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>();

scripts/codegen/utils.ts

@@ -21,6 +21,17 @@ const __dirname = path.dirname(__filename);
 /** Root of the copilot-sdk repo */
 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> {
@@ -51,16 +62,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];
@@ -93,12 +95,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") {
@@ -339,6 +343,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,
+    };
 }