Centralize union member names in mutation engine

Author: FionaBronwen

packages/graphql/src/components/types/enum-type.tsx

@@ -1,4 +1,4 @@
-import { type Enum, getDoc, getDeprecationDetails } from "@typespec/compiler";
+import { type Enum, getDeprecationDetails } from "@typespec/compiler";
 import * as gql from "@alloy-js/graphql";
 import { useTsp } from "@typespec/emitter-framework";
 
@@ -11,14 +11,14 @@ export interface EnumTypeProps {
  * Renders a GraphQL enum type declaration with members
  */
 export function EnumType(props: EnumTypeProps) {
-  const { program } = useTsp();
-  const doc = getDoc(program, props.type);
+  const { $, program } = useTsp();
+  const doc = $.type.getDoc(props.type);
   const members = Array.from(props.type.members.values());
 
   return (
     <gql.EnumType name={props.type.name} description={doc}>
       {members.map((member) => {
-        const memberDoc = getDoc(program, member);
+        const memberDoc = $.type.getDoc(member);
         const deprecation = getDeprecationDetails(program, member);
 
         return (

packages/graphql/src/components/types/scalar-type.tsx

@@ -1,4 +1,4 @@
-import { type Scalar, getDoc } from "@typespec/compiler";
+import type { Scalar } from "@typespec/compiler";
 import * as gql from "@alloy-js/graphql";
 import { useTsp } from "@typespec/emitter-framework";
 import { useGraphQLSchema } from "../../context/index.js";
@@ -12,9 +12,9 @@ export interface ScalarTypeProps {
  * Renders a GraphQL scalar type declaration with optional @specifiedBy directive
  */
 export function ScalarType(props: ScalarTypeProps) {
-  const { program } = useTsp();
+  const { $ } = useTsp();
   const { scalarSpecifications } = useGraphQLSchema();
-  const doc = getDoc(program, props.type);
+  const doc = $.type.getDoc(props.type);
   const specificationUrl = scalarSpecifications.get(props.type.name);
 
   return (

packages/graphql/src/components/types/union-type.tsx

@@ -1,50 +1,34 @@
-import { type Union, getDoc } from "@typespec/compiler";
+import type { Union } from "@typespec/compiler";
 import * as gql from "@alloy-js/graphql";
 import { useTsp } from "@typespec/emitter-framework";
-import { getUnionName, isScalarLikeType, toTypeName } from "../../lib/type-utils.js";
+import { useGraphQLSchema } from "../../context/index.js";
+import { getUnionName } from "../../lib/type-utils.js";
 
 export interface UnionTypeProps {
   /** The union type to render */
   type: Union;
 }
 
 /**
- * Renders a GraphQL union type declaration
- * Scalars are wrapped in object types since GraphQL unions can only contain object types
- * This wrapping is done by the mutation engine
+ * Renders a GraphQL union type declaration.
+ *
+ * Member names (wrapper names for scalar variants, GraphQL type names for
+ * object variants) are produced by the mutation engine and read from
+ * `GraphQLSchemaContextValue.unionMembers`. The component only decides
+ * how to render, not what the members are.
  */
 export function UnionType(props: UnionTypeProps) {
-  const { program } = useTsp();
+  const { $, program } = useTsp();
+  const { unionMembers } = useGraphQLSchema();
   const name = getUnionName(props.type, program);
-  const doc = getDoc(program, props.type);
-  const variants = Array.from(props.type.variants.values());
+  const doc = $.type.getDoc(props.type);
+  const members = unionMembers.get(props.type);
 
-  // Build the union member list, using wrapper names for scalars
-  // The wrapper models are created by the mutation engine
-  const unionMembers = variants.map((variant) => {
-    const variantName =
-      typeof variant.name === "string" ? variant.name : String(variant.name);
-
-    if (isScalarLikeType(variant.type)) {
-      // Reference the wrapper type for scalars (created by mutation engine)
-      // Include union name to match wrapper model naming convention
-      return toTypeName(name) + toTypeName(variantName) + "UnionVariant";
-    } else {
-      // For non-scalars, use the type name directly
-      if (variant.type.kind === "Model") {
-        return variant.type.name;
-      } else if (
-        "name" in variant.type &&
-        typeof variant.type.name === "string"
-      ) {
-        return variant.type.name;
-      }
-      throw new Error(
-        `Unexpected union variant type kind "${variant.type.kind}" in union "${name}". ` +
-          `This is a bug in the GraphQL emitter.`,
-      );
-    }
-  });
-
-  return <gql.UnionType name={name} description={doc} members={unionMembers} />;
+  return (
+    <gql.UnionType
+      name={name}
+      description={doc}
+      members={members ? [...members] : []}
+    />
+  );
 }

packages/graphql/src/context/graphql-schema-context.tsx

@@ -58,7 +58,11 @@ export interface ScalarVariant {
 export interface GraphQLSchemaContextValue {
   /** Classified types for schema generation */
   classifiedTypes: ClassifiedTypes;
-  /** Ordered member names for each union, keyed by the mutated Union */
+  /**
+   * Ordered member names for each output union, keyed by the mutated Union.
+   * Consumed by UnionType to avoid reconstructing the naming convention
+   * the mutation engine already applied.
+   */
   unionMembers: Map<Union, readonly string[]>;
   /** Scalar specification URLs for @specifiedBy directives */
   scalarSpecifications: Map<string, string>;

packages/graphql/src/lib/type-utils.ts

@@ -318,11 +318,6 @@ function getTemplateStringInternal(
   return args.length > 0 ? args.map(toTypeName).join(options.conjunction) : "";
 }
 
-/** Check if a type is a scalar (built-in or custom) or an intrinsic type like `unknown`. */
-export function isScalarLikeType(type: Type): boolean {
-  return type.kind === "Scalar" || type.kind === "Intrinsic";
-}
-
 /** Check if a model should be emitted as a GraphQL object type (not an array, record, or never). */
 export function isTrueModel(model: Model): boolean {
   return !(

packages/graphql/src/mutation-engine/mutations/union.ts

@@ -35,6 +35,7 @@ function variantNameToString(name: string | symbol): string {
 export class GraphQLUnionMutation extends UnionMutation<MutationOptions, any, MutationEngine<any>> {
   #mutationNode: UnionMutationNode;
   #wrapperModels: Model[] = [];
+  #memberNames: string[] = [];
   #flattenedUnion: Union | null = null;
 
   constructor(
@@ -76,6 +77,15 @@ export class GraphQLUnionMutation extends UnionMutation<MutationOptions, any, Mu
     return this.#wrapperModels;
   }
 
+  /**
+   * Ordered member names for the output union (wrapper name for scalars,
+   * GraphQL type name for object variants). Empty for input/@oneOf unions
+   * and nullable-union replacements.
+   */
+  get memberNames(): readonly string[] {
+    return this.#memberNames;
+  }
+
   /** Creates a half-edge for bidirectional variant mutation updates. */
   protected startVariantEdge(): MutationHalfEdge<
     GraphQLUnionMutation,
@@ -145,7 +155,9 @@ export class GraphQLUnionMutation extends UnionMutation<MutationOptions, any, Mu
       setNullable(this.mutatedType);
     }
 
-    // GraphQL unions can only contain object types — wrap scalars in synthetic models
+    // GraphQL unions can only contain object types — wrap scalars in synthetic models.
+    // Record member names in variant order so the renderer doesn't have to
+    // reconstruct the naming convention.
     for (const variant of flattenedVariants) {
       const isScalar = variant.type.kind === "Scalar" || variant.type.kind === "Intrinsic";
 
@@ -166,6 +178,9 @@ export class GraphQLUnionMutation extends UnionMutation<MutationOptions, any, Mu
         });
 
         this.#wrapperModels.push(wrapperModel);
+        this.#memberNames.push(wrapperName);
+      } else if ("name" in variant.type && typeof variant.type.name === "string") {
+        this.#memberNames.push(sanitizeNameForGraphQL(variant.type.name));
       }
     }
   }

packages/graphql/src/mutation-engine/schema-mutator.ts

@@ -63,6 +63,12 @@ export interface MutatedSchema {
   // Derived metadata
   /** Synthetic wrapper models created by union mutations for scalar variants */
   wrapperModels: Model[];
+  /**
+   * Ordered member names for each output union, keyed by the mutated Union.
+   * Entries are the GraphQL names that appear on the right-hand side of the
+   * `union X = A | B | ...` declaration.
+   */
+  unionMembers: Map<Union, readonly string[]>;
   /** Encoded stdlib scalars mapped to GraphQL custom scalars (e.g. bytes + base64 → Bytes) */
   scalarVariants: ScalarVariant[];
   /** `@specifiedBy` URLs indexed by GraphQL scalar name */
@@ -105,6 +111,9 @@ export function mutateSchema(
   // Synthetic wrapper models from union mutation (always output).
   const wrapperModels: Model[] = [];
 
+  // Ordered member names for each output union, keyed by the mutated Union.
+  const unionMembers = new Map<Union, readonly string[]>();
+
   // Void-returning operations — collected here so the emitter can report
   // the diagnostic. Mutation shapes the graph; it doesn't warn.
   const voidOperations: Operation[] = [];
@@ -236,8 +245,10 @@ export function mutateSchema(
       if (unwrapNullableUnion(node) !== undefined) return;
       if (typeUsage.isUnreachable(node)) return;
       const mutation = engine.mutateUnion(node, GraphQLTypeContext.Output);
-      unions.push(mutation.mutatedType as Union);
+      const mutatedUnion = mutation.mutatedType as Union;
+      unions.push(mutatedUnion);
       wrapperModels.push(...mutation.wrapperModels);
+      unionMembers.set(mutatedUnion, mutation.memberNames);
     },
     operation: (node: Operation) => {
       // Operations pass through unmutated. classifyOperation collects void
@@ -309,6 +320,7 @@ export function mutateSchema(
     subscriptions,
     voidOperations,
     wrapperModels,
+    unionMembers,
     scalarVariants: Array.from(scalarVariantsMap.values()),
     scalarSpecifications,
   };

packages/graphql/test/components/component-test-utils.tsx

@@ -31,7 +31,7 @@ export function renderComponentToSDL(
       mutations: [],
       subscriptions: [],
     },
-    modelVariants: { outputModels: new Map(), inputModels: new Map() },
+    unionMembers: new Map(),
     scalarSpecifications: new Map(),
     ...contextOverrides,
   };

packages/graphql/test/components/union-type.test.tsx

@@ -52,11 +52,13 @@ describe("UnionType component", () => {
         </gql.ObjectType>
         <UnionType type={mutatedUnion} />
       </>,
+      { unionMembers: new Map([[mutatedUnion, mutation.memberNames]]) },
     );
 
     expect(sdl).toContain("union Pet =");
-    expect(sdl).toContain("Cat");
-    expect(sdl).toContain("Dog");
+    for (const member of mutation.memberNames) {
+      expect(sdl).toContain(member);
+    }
   });
 
   it("renders a union with doc comment description", async () => {
@@ -84,6 +86,7 @@ describe("UnionType component", () => {
         </gql.ObjectType>
         <UnionType type={mutatedUnion} />
       </>,
+      { unionMembers: new Map([[mutatedUnion, mutation.memberNames]]) },
     );
 
     expect(sdl).toContain("The result of an operation");
@@ -118,12 +121,13 @@ describe("UnionType component", () => {
         </gql.ObjectType>
         <UnionType type={mutatedUnion} />
       </>,
+      { unionMembers: new Map([[mutatedUnion, mutation.memberNames]]) },
     );
 
     expect(sdl).toContain("union Shape =");
-    expect(sdl).toContain("Circle");
-    expect(sdl).toContain("Square");
-    expect(sdl).toContain("Triangle");
+    for (const member of mutation.memberNames) {
+      expect(sdl).toContain(member);
+    }
   });
 
   it("references wrapper type names for scalar variants", async () => {
@@ -138,23 +142,27 @@ describe("UnionType component", () => {
     const mutation = engine.mutateUnion(Mixed, GraphQLTypeContext.Output);
     const mutatedUnion = assertUnionResult(mutation);
 
-    // Register Cat and the wrapper type that the union will reference
+    // Register the mutation's wrapper models so buildSchema resolves the members.
     const sdl = renderComponentToSDL(
       tester.program,
       <>
         <gql.ObjectType name="Cat">
           <gql.Field name="name" type={gql.String} nonNull />
         </gql.ObjectType>
-        <gql.ObjectType name="MixedTextUnionVariant">
-          <gql.Field name="value" type={gql.String} nonNull />
-        </gql.ObjectType>
+        {mutation.wrapperModels.map((wrapper) => (
+          <gql.ObjectType name={wrapper.name}>
+            <gql.Field name="value" type={gql.String} nonNull />
+          </gql.ObjectType>
+        ))}
         <UnionType type={mutatedUnion} />
       </>,
+      { unionMembers: new Map([[mutatedUnion, mutation.memberNames]]) },
     );
 
-    // Scalar variant should reference wrapper type name
+    // Every member the engine reported should land in the SDL.
     expect(sdl).toContain("union Mixed =");
-    expect(sdl).toContain("MixedTextUnionVariant");
-    expect(sdl).toContain("Cat");
+    for (const member of mutation.memberNames) {
+      expect(sdl).toContain(member);
+    }
   });
 });