fix(openapi): restore named discriminated union branches (#2785)

packages/openapi/src/OpenAPIGenerator.ts

@@ -67,7 +67,7 @@ export class OpenAPIGenerator {
 		return result.document;
 	}
 	public async generateWithStats(): Promise<OpenAPIGenerationResult> {
-		const schemaRegistry = new SchemaRegistry(this.settings.schemaTarget);
+		const schemaRegistry = new SchemaRegistry(this.settings.schemaTarget, this.settings.routeScope !== 'admin');
 		const controllerFiles = discoverControllerFiles(`${this.settings.basePath}/fluxer_api`);
 		const routes = this.filterRoutesForScope(extractRoutesFromControllers(controllerFiles));
 		await loadSchemasIntoRegistry(this.settings.basePath, schemaRegistry);

packages/openapi/src/converters/ZodToOpenAPI.ts

@@ -5,7 +5,7 @@ import type {OpenAPISchemaTarget} from '@fluxer/openapi/src/OpenAPIGenerationTyp
 import {visitOpenAPISchemaObjects} from '@fluxer/openapi/src/OpenAPISchemaVisitor';
 import type {OpenAPIDocument, OpenAPIRef, OpenAPISchema} from '@fluxer/openapi/src/OpenAPITypes';
 import {schemaMetadata} from '@fluxer/schema/src/SchemaMetadata';
-import {type core, z} from 'zod';
+import {core, z} from 'zod';
 
 export type SchemaIO = 'input' | 'output';
 
@@ -76,6 +76,25 @@ function renameComponentRefs(value: unknown, renames: Map<string, string>): void
 	});
 }
 
+function toPascalCase(value: string): string {
+	return value
+		.split(/[^A-Za-z0-9]+/)
+		.filter((word) => word.length > 0)
+		.map((word) => word.charAt(0).toUpperCase() + word.slice(1).toLowerCase())
+		.join('');
+}
+
+function branchNamePrefix(discriminator: core.JSONSchema._JSONSchema | undefined): string | undefined {
+	if (typeof discriminator !== 'object') return undefined;
+	const enumNames = discriminator['x-enumNames'];
+	const value = Array.isArray(enumNames)
+		? enumNames[0]
+		: discriminator.enum?.length === 1
+			? discriminator.enum[0]
+			: discriminator.const;
+	return typeof value === 'string' ? toPascalCase(value) : undefined;
+}
+
 function removeRedundantReferenceProperties(value: unknown, schemas: Record<string, OpenAPISchema>): void {
 	visitOpenAPISchemaObjects(value, (schema) => {
 		const refs = [schema.$ref, ...(schema.allOf ?? []).map((branch) => branch.$ref)];
@@ -99,7 +118,10 @@ export class ZodOpenAPIConverter {
 	private readonly worklist: Array<PendingSchema> = [];
 	private nextPendingIndex = 0;
 
-	constructor(private readonly target: OpenAPISchemaTarget) {}
+	constructor(
+		private readonly target: OpenAPISchemaTarget,
+		private readonly nameUnionBranches = false,
+	) {}
 
 	register(name: string, schema: core.$ZodType): void {
 		validateOpenAPIComponentName(name);
@@ -164,6 +186,24 @@ export class ZodOpenAPIConverter {
 		removeRedundantReferenceProperties(document, document.components.schemas);
 	}
 
+	private nameDiscriminatedUnionBranches(
+		unionName: string,
+		schema: core.$ZodType,
+		json: OpenAPISchema,
+		io: SchemaIO,
+	): void {
+		if (!(schema instanceof core.$ZodDiscriminatedUnion) || !json.oneOf) return;
+		const {discriminator, options} = schema._zod.def;
+		const usedPrefixes = new Set<string>();
+		json.oneOf = json.oneOf.map((branch, index) => {
+			if (branch.$ref) return branch;
+			const base = branchNamePrefix(branch.properties?.[discriminator]) ?? `Variant${index}`;
+			const prefix = usedPrefixes.has(base) ? `${base}${index}` : base;
+			usedPrefixes.add(prefix);
+			return this.getRef(`${prefix}${unionName}`, options[index], io);
+		});
+	}
+
 	private convertPending(): void {
 		while (this.nextPendingIndex < this.worklist.length) {
 			const pending = this.worklist[this.nextPendingIndex];
@@ -193,6 +233,8 @@ export class ZodOpenAPIConverter {
 								const ref = this.getRef(name, zodSchema, io);
 								replaceSchema(jsonSchema, description ? {...ref, description} : ref);
 							}
+						} else if (this.nameUnionBranches) {
+							this.nameDiscriminatedUnionBranches(pending.name, zodSchema, jsonSchema, io);
 						}
 						if (this.target === 'openapi-3.0' && jsonSchema.$ref && Object.keys(jsonSchema).length > 1) {
 							const ref = {$ref: jsonSchema.$ref};

packages/openapi/src/converters/__tests__/ZodToOpenAPI.test.ts

@@ -2,7 +2,7 @@ import {ZodOpenAPIConverter} from '@fluxer/openapi/src/converters/ZodToOpenAPI';
 import type {OpenAPIDocument} from '@fluxer/openapi/src/OpenAPITypes';
 import {SudoVerificationSchema} from '@fluxer/schema/src/domains/auth/AuthSchemas';
 import {HarvestArchiveResponse} from '@fluxer/schema/src/domains/user/UserHarvestSchemas';
-import {SnowflakeType, withOpenApiType} from '@fluxer/schema/src/primitives/SchemaPrimitives';
+import {createNamedLiteral, SnowflakeType, withOpenApiType} from '@fluxer/schema/src/primitives/SchemaPrimitives';
 import {describe, expect, it} from 'vitest';
 import {z} from 'zod';
 
@@ -122,6 +122,34 @@ describe('native Zod conversion', () => {
 		);
 	});
 
+	it('names discriminated union branches after their discriminator value', () => {
+		const union = z.discriminatedUnion('kind', [
+			z.object({kind: createNamedLiteral(0, 'GUILD_TEXT'), name: z.string()}),
+			z.object({kind: z.literal('refresh_token'), token: z.string()}),
+			z.object({kind: z.literal(7), size: z.number()}),
+		]);
+		const converter = new ZodOpenAPIConverter('draft-2020-12', true);
+		converter.getRef('Request', union, 'input');
+		const document = documentFor(converter);
+		converter.normalizeComponents(document);
+		expect(document.components.schemas.Request.oneOf).toEqual([
+			{$ref: '#/components/schemas/GuildTextRequest'},
+			{$ref: '#/components/schemas/RefreshTokenRequest'},
+			{$ref: '#/components/schemas/Variant2Request'},
+		]);
+		expect(document.components.schemas.GuildTextRequest.properties).toHaveProperty('name');
+		expect(document.components.schemas.RefreshTokenRequest.properties).toHaveProperty('token');
+		expect(document.components.schemas.Variant2Request.properties).toHaveProperty('size');
+	});
+
+	it('keeps discriminated union branches inline unless asked to name them', () => {
+		const union = z.discriminatedUnion('kind', [z.object({kind: z.literal('a')}), z.object({kind: z.literal('b')})]);
+		const converter = new ZodOpenAPIConverter('openapi-3.0');
+		const schema = converter.getSchema('Request', union, 'input');
+		expect(schema.oneOf?.map((branch) => branch.$ref)).toEqual([undefined, undefined]);
+		expect(Object.keys(converter.getAllSchemas())).toEqual(['RequestInput']);
+	});
+
 	it('converts newly discovered components after an earlier conversion pass', () => {
 		const converter = new ZodOpenAPIConverter('draft-2020-12');
 		converter.getSchema('First', z.string(), 'input');

packages/openapi/src/registry/SchemaRegistry.ts

@@ -12,8 +12,8 @@ export class SchemaRegistry {
 	private readonly emptyObjectAcceptance = new WeakMap<core.$ZodType, boolean>();
 	private readonly converter: ZodOpenAPIConverter;
 
-	constructor(target: OpenAPISchemaTarget = 'draft-2020-12') {
-		this.converter = new ZodOpenAPIConverter(target);
+	constructor(target: OpenAPISchemaTarget = 'draft-2020-12', nameUnionBranches = false) {
+		this.converter = new ZodOpenAPIConverter(target, nameUnionBranches);
 	}
 
 	register(name: string, schema: OpenAPISchema): void {