Fix C# codegen for runtime schema unions

nodejs/test/csharp-codegen.test.ts

@@ -0,0 +1,255 @@
+import type { JSONSchema7 } from "json-schema";
+import { describe, expect, it } from "vitest";
+
+import { generateRpcCode } from "../../scripts/codegen/csharp.ts";
+import type { ApiSchema } from "../../scripts/codegen/utils.ts";
+
+describe("C# RPC codegen", () => {
+    it.each([
+        ["anyOf", false],
+        ["anyOf", true],
+        ["oneOf", false],
+        ["oneOf", true],
+    ] as const)(
+        "preserves the %s hierarchy when adding a variant (nullable: %s)",
+        (keyword, nullable) => {
+            const jsonSchemaVariant: JSONSchema7 = {
+                type: "object",
+                properties: {
+                    type: { type: "string", const: "json_schema" },
+                    jsonSchema: { $ref: "#/definitions/JsonSchemaResponseFormat" },
+                },
+                required: ["type", "jsonSchema"],
+            };
+            const nullVariants: JSONSchema7[] = nullable ? [{ type: "null" }] : [];
+            const responseFormat: JSONSchema7 = {
+                title: "ResponseFormat",
+                description: "A provider-native output format.",
+                [keyword]: [jsonSchemaVariant, ...nullVariants],
+            };
+            const schema: ApiSchema = {
+                session: {
+                    send: {
+                        rpcMethod: "session.send",
+                        params: {
+                            type: "object",
+                            title: "SendRequest",
+                            properties: {
+                                responseFormat: { $ref: "#/definitions/ResponseFormat" },
+                                requiredFormat: { $ref: "#/definitions/ResponseFormat" },
+                            },
+                            required: ["requiredFormat"],
+                        },
+                    },
+                },
+                definitions: {
+                    ResponseFormat: responseFormat,
+                    JsonSchemaResponseFormat: {
+                        type: "object",
+                        properties: {
+                            name: { type: "string" },
+                            schema: { "x-opaque-json": true } as JSONSchema7,
+                            strict: { type: "boolean" },
+                        },
+                        required: ["name", "schema"],
+                    },
+                },
+            };
+            const code = generateRpcCode(schema);
+            const futureCode = generateRpcCode({
+                ...schema,
+                definitions: {
+                    ...schema.definitions,
+                    ResponseFormat: {
+                        ...responseFormat,
+                        [keyword]: [
+                            jsonSchemaVariant,
+                            {
+                                type: "object",
+                                properties: { type: { type: "string", const: "text" } },
+                                required: ["type"],
+                            },
+                            ...nullVariants,
+                        ],
+                    },
+                },
+            });
+
+            for (const generated of [code, futureCode]) {
+                expect(generated).toContain("public partial class ResponseFormat\n");
+                expect(generated).toContain("A provider-native output format.");
+                expect(generated).toContain(
+                    '[JsonPolymorphic(\n    TypeDiscriminatorPropertyName = "type",'
+                );
+                expect(generated).toContain(
+                    '[JsonDerivedType(typeof(ResponseFormatJsonSchema), "json_schema")]'
+                );
+                expect(generated).toContain(
+                    "public partial class ResponseFormatJsonSchema : ResponseFormat"
+                );
+                expect(generated).toContain('public override string Type => "json_schema";');
+                expect(generated).toContain("public ResponseFormat? ResponseFormat");
+                expect(generated).toContain(
+                    `public ResponseFormat${nullable ? "?" : ""} RequiredFormat`
+                );
+                expect(generated).toContain("public required JsonSchemaResponseFormat JsonSchema");
+                expect(generated).toContain("public JsonElement Schema");
+                expect(generated).toContain("public bool? Strict");
+                expect(generated).toContain("[JsonSerializable(typeof(ResponseFormat))]");
+                expect(generated.match(/public partial class ResponseFormat\b/g)).toHaveLength(1);
+            }
+            for (const name of ["ResponseFormat", "ResponseFormatJsonSchema"]) {
+                const declaration = new RegExp(
+                    `public partial class ${name}\\b[^\\n]*\\n\\{[\\s\\S]*?\\n\\}`
+                );
+                expect(code.match(declaration)?.[0]).toBe(futureCode.match(declaration)?.[0]);
+            }
+            expect(futureCode).toContain('[JsonDerivedType(typeof(ResponseFormatText), "text")]');
+            expect(futureCode).toContain(
+                "public partial class ResponseFormatText : ResponseFormat"
+            );
+        }
+    );
+
+    it.each(["anyOf", "oneOf"] as const)(
+        "preserves referenced enum discriminators and nested %s unions",
+        (keyword) => {
+            const code = generateRpcCode({
+                session: {
+                    configure: {
+                        rpcMethod: "session.configure",
+                        stability: "experimental",
+                        params: {
+                            type: "object",
+                            properties: {
+                                systemMessage: { $ref: "#/definitions/SystemMessage" },
+                                requiredMessage: { $ref: "#/definitions/SystemMessage" },
+                            },
+                            required: ["requiredMessage"],
+                        },
+                    },
+                },
+                definitions: {
+                    SystemMessage: {
+                        description: "System message configuration.",
+                        [keyword]: [
+                            { $ref: "#/definitions/AppendConfig" },
+                            { $ref: "#/definitions/ReplaceConfig" },
+                            { $ref: "#/definitions/CustomizeConfig" },
+                        ],
+                    },
+                    AppendMode: { type: "string", enum: ["append"] },
+                    ReplaceMode: { type: "string", enum: ["replace"] },
+                    CustomizeMode: { type: "string", enum: ["customize"] },
+                    AppendConfig: {
+                        type: "object",
+                        properties: {
+                            mode: { $ref: "#/definitions/AppendMode" },
+                            content: { type: "string" },
+                        },
+                    },
+                    ReplaceConfig: {
+                        type: "object",
+                        properties: {
+                            mode: { $ref: "#/definitions/ReplaceMode" },
+                            content: { type: "string" },
+                        },
+                        required: ["mode", "content"],
+                    },
+                    CustomizeConfig: {
+                        type: "object",
+                        properties: {
+                            mode: { $ref: "#/definitions/CustomizeMode" },
+                            content: { type: "string" },
+                            sections: {
+                                type: "object",
+                                additionalProperties: {
+                                    $ref: "#/definitions/SectionOverride",
+                                },
+                            },
+                        },
+                        required: ["mode"],
+                    },
+                    SectionOverride: {
+                        [keyword]: [
+                            { $ref: "#/definitions/StaticSectionOverride" },
+                            { $ref: "#/definitions/MarkerSectionOverride" },
+                        ],
+                    },
+                    StaticSectionAction: {
+                        type: "string",
+                        enum: ["replace", "remove", "append", "prepend"],
+                    },
+                    StaticSectionOverride: {
+                        type: "object",
+                        properties: {
+                            action: { $ref: "#/definitions/StaticSectionAction" },
+                            content: { type: "string" },
+                        },
+                        required: ["action"],
+                    },
+                    MarkerSectionOverride: {
+                        [keyword]: ["transform", "preserve"].map((action) => ({
+                            type: "object",
+                            properties: { action: { type: "string", const: action } },
+                            required: ["action"],
+                        })),
+                    },
+                },
+            });
+
+            expect(code).toContain("public sealed partial class SystemMessage");
+            expect(code).toContain(
+                "[Experimental(Diagnostics.Experimental)]\n[JsonConverter(typeof(Converter))]\npublic sealed partial class SystemMessage"
+            );
+            expect(code).toContain("System message configuration.");
+            expect(code).toContain("public SystemMessage? SystemMessage");
+            expect(code).toContain("public SystemMessage RequiredMessage");
+            expect(code.match(/public sealed partial class SystemMessage\b/g)).toHaveLength(1);
+            expect(code).toContain("public AppendConfig? AppendConfig { get; }");
+            expect(code).toContain("public ReplaceConfig? ReplaceConfig { get; }");
+            expect(code).toContain("public CustomizeConfig? CustomizeConfig { get; }");
+            expect(code).toContain("public AppendMode? Mode");
+            expect(code).toContain("public ReplaceMode Mode");
+            expect(code).toContain("public IDictionary<string, SectionOverride>? Sections");
+            expect(code).toContain("public StaticSectionOverride? StaticSectionOverride { get; }");
+            expect(code).toContain("public MarkerSectionOverride? MarkerSectionOverride { get; }");
+            expect(code).toContain("public StaticSectionAction Action");
+            expect(code).toContain(
+                '!element.TryGetProperty("mode", out _) || (element.TryGetProperty("mode", out _)'
+            );
+            for (const mode of ["append", "replace", "customize"]) {
+                expect(code).toContain(`element.GetProperty("mode").GetString() == "${mode}"`);
+            }
+            for (const action of [
+                "replace",
+                "remove",
+                "append",
+                "prepend",
+                "transform",
+                "preserve",
+            ]) {
+                expect(code).toContain(`element.GetProperty("action").GetString() == "${action}"`);
+            }
+            expect(code).toContain('element.GetProperty("mode").ValueKind == JsonValueKind.String');
+            expect(code).not.toContain("catch (JsonException)");
+            for (const type of [
+                "SystemMessage",
+                "AppendConfig",
+                "ReplaceConfig",
+                "CustomizeConfig",
+                "SectionOverride",
+                "StaticSectionOverride",
+                "MarkerSectionOverride",
+            ]) {
+                expect(code).toContain(`[JsonSerializable(typeof(${type}))]`);
+            }
+            expect(code).toContain(
+                "JsonSerializer.Deserialize(element, RpcJsonContext.Default.AppendConfig)"
+            );
+            expect(code).toContain(
+                "JsonSerializer.Serialize(writer, appendConfig, RpcJsonContext.Default.AppendConfig)"
+            );
+        }
+    );
+});

scripts/codegen/csharp.ts

@@ -1057,6 +1057,7 @@ interface JsonUnionVariant {
     typeName: string;
     propertyName: string;
     schema?: JSONSchema7;
+    matchExpression?: string;
 }
 
 function getUnionMembers(schema: JSONSchema7): JSONSchema7[] | undefined {
@@ -1107,9 +1108,10 @@ function getJsonUnionMatchExpression(variant: JsonUnionVariant, variants: JsonUn
     ].join(" && ");
 }
 
-function generateJsonUnionClass(className: string, variants: JsonUnionVariant[], description: string | undefined, jsonContextType: string, isInternal: boolean): string {
+function generateJsonUnionClass(className: string, variants: JsonUnionVariant[], description: string | undefined, jsonContextType: string, isInternal: boolean, experimental = false): string {
     const lines: string[] = [];
     lines.push(...xmlDocCommentWithFallback(description, `JSON union data type for <c>${escapeXml(className)}</c>.`, ""));
+    if (experimental) pushExperimentalAttribute(lines);
     lines.push(`[JsonConverter(typeof(Converter))]`);
     lines.push(`${isInternal ? "internal" : "public"} sealed partial class ${className}`);
     lines.push(`{`);
@@ -1147,7 +1149,7 @@ function generateJsonUnionClass(className: string, variants: JsonUnionVariant[],
 
     const fallbackVariants: JsonUnionVariant[] = [];
     for (const variant of variants) {
-        const matchExpression = getJsonUnionMatchExpression(variant, variants);
+        const matchExpression = variant.matchExpression ?? getJsonUnionMatchExpression(variant, variants);
         if (!matchExpression) {
             fallbackVariants.push(variant);
             continue;
@@ -1663,6 +1665,39 @@ function stableStringify(value: unknown): string {
     return JSON.stringify(value);
 }
 
+// Match literal values before deserialization: an optional discriminator must not swallow another variant.
+function getRpcUnionMatchExpression(schema: JSONSchema7, seenRefs: ReadonlySet<string> = new Set()): string | undefined {
+    if (schema.$ref) {
+        if (seenRefs.has(schema.$ref)) return undefined;
+        seenRefs = new Set([...seenRefs, schema.$ref]);
+    }
+    const resolved = resolveSchema(schema, rpcDefinitions) ?? schema;
+    const members = getUnionMembers(resolved);
+    if (members) {
+        const expressions = members.map((member) => getRpcUnionMatchExpression(member, seenRefs));
+        return expressions.every((expression) => expression !== undefined)
+            ? `(${expressions.join(" || ")})`
+            : undefined;
+    }
+
+    const expressions: string[] = [];
+    for (const [name, property] of Object.entries(resolved.properties ?? {})) {
+        if (typeof property !== "object") continue;
+        const propSchema = resolveSchema(property, rpcDefinitions) ?? property;
+        const values = propSchema.const !== undefined ? [propSchema.const] : propSchema.enum;
+        if (!values?.length || !values.every((value) => typeof value === "string")) continue;
+        const propertyName = escapeCSharpStringLiteral(name);
+        const present = `element.TryGetProperty("${propertyName}", out _)`;
+        const value = `element.GetProperty("${propertyName}")`;
+        const matches = values.map((entry) => `${value}.GetString() == "${escapeCSharpStringLiteral(entry as string)}"`);
+        const match = `${present} && ${value}.ValueKind == JsonValueKind.String && (${matches.join(" || ")})`;
+        expressions.push(resolved.required?.includes(name) ? `(${match})` : `(!${present} || (${match}))`);
+    }
+    return expressions.length > 0
+        ? `element.ValueKind == JsonValueKind.Object && ${expressions.join(" && ")}`
+        : undefined;
+}
+
 function resolveRpcType(schema: JSONSchema7, isRequired: boolean, parentClassName: string, propName: string, classes: string[]): string {
     if (isOpaqueJson(schema)) {
         return isRequired ? "JsonElement" : "JsonElement?";
@@ -1686,17 +1721,13 @@ function resolveRpcType(schema: JSONSchema7, isRequired: boolean, parentClassNam
             return isRequired ? typeName : `${typeName}?`;
         }
 
-        return resolveRpcType(refSchema, isRequired, parentClassName, propName, classes);
+        return resolveRpcType({ ...refSchema, title: refSchema.title ?? typeName }, isRequired, parentClassName, propName, classes);
     }
-    // Handle anyOf: [T, null/{not:{}}] → T? (nullable typed property)
-    const nullableInner = getNullableInner(schema);
-    if (nullableInner) {
-        return resolveRpcType(nullableInner, false, parentClassName, propName, classes);
-    }
-    // Discriminated union: anyOf with multiple variants sharing a const discriminator
-    if (schema.anyOf && Array.isArray(schema.anyOf)) {
-        const nonNull = schema.anyOf.filter((s) => typeof s === "object" && s !== null && (s as JSONSchema7).type !== "null");
-        if (nonNull.length > 1) {
+    const unionVariants = schema.anyOf ?? schema.oneOf;
+    // Keep the same polymorphic API even when a discriminated union has only one variant.
+    if (unionVariants) {
+        const nonNull = getNonNullUnionMembers(schema);
+        if (nonNull.length > 0) {
             const variants = (nonNull as JSONSchema7[]).map((v) => {
                 if (v.$ref) {
                     const resolved = resolveRef(v.$ref, rpcDefinitions);
@@ -1706,7 +1737,7 @@ function resolveRpcType(schema: JSONSchema7, isRequired: boolean, parentClassNam
             });
             const discriminatorInfo = findDiscriminator(variants);
             if (discriminatorInfo) {
-                const hasNull = schema.anyOf.length > nonNull.length;
+                const hasNull = unionVariants.length > nonNull.length;
                 const baseClassName = (schema.title as string) ?? `${parentClassName}${propName}`;
                 if (!emittedRpcClassSchemas.has(baseClassName)) {
                     emittedRpcClassSchemas.set(baseClassName, "polymorphic");
@@ -1727,6 +1758,32 @@ function resolveRpcType(schema: JSONSchema7, isRequired: boolean, parentClassNam
             }
         }
     }
+    // Preserve nullable references without introducing another wrapper around their declared type.
+    const nullableInner = getNullableInner(schema);
+    if (nullableInner) {
+        return resolveRpcType(nullableInner, false, parentClassName, propName, classes);
+    }
+    if (unionVariants && getNonNullUnionMembers(schema).length > 0) {
+        const members = getNonNullUnionMembers(schema);
+        const matchExpressions = members.map((member) => getRpcUnionMatchExpression(member));
+        if (matchExpressions.every((expression) => expression !== undefined)) {
+            const className = schema.title ?? `${parentClassName}${propName}`;
+            if (!emittedRpcClassSchemas.has(className)) {
+                emittedRpcClassSchemas.set(className, "union");
+                const usedNames = new Set<string>();
+                const variants = members.map((member, index) => {
+                    const typeName = resolveRpcType(member, true, className, `Variant${index + 1}`, classes);
+                    return {
+                        typeName,
+                        propertyName: toUnionVariantPropertyName(typeName, usedNames),
+                        matchExpression: matchExpressions[index],
+                    };
+                });
+                classes.push(generateJsonUnionClass(className, variants, schema.description, "RpcJsonContext", isSchemaInternal(schema), isSchemaExperimental(schema) || experimentalRpcTypes.has(className)));
+            }
+            return isRequired && members.length === unionVariants.length ? className : `${className}?`;
+        }
+    }
     // Handle enums (string unions like "interactive" | "plan" | "autopilot")
     if (schema.enum && Array.isArray(schema.enum)) {
         const explicitName = schema.title as string | undefined;
@@ -2607,7 +2664,7 @@ function emitClientGlobalApiRegistration(clientSchema: Record<string, unknown>,
     return lines;
 }
 
-function generateRpcCode(
+export function generateRpcCode(
     schema: ApiSchema,
     externalJsonSerializableRefs: Map<string, Set<string>> = new Map(),
     externalValueTypes: Set<string> = new Set()