[typespec-ts] fix extensible enum model emission

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

Author: maorleger

packages/typespec-ts/src/codegen/models.ts

@@ -1,9 +1,12 @@
 import { Project, SourceFile } from "ts-morph";
-import type { TSCodeModel } from "../codemodel/index.js";
+import type { TSCodeModel, TSEnum } from "../codemodel/index.js";
 import type { SdkContext } from "../utils/interfaces.js";
+import { addDeclaration } from "../framework/declaration.js";
+import { refkey } from "../framework/refkey.js";
 import { buildHelperTypeLookup } from "../tcgcadapter/helperTypes.js";
 import {
   addSerializationFunctions,
+  buildEnumTypes,
   emitType,
   getModelNamespaces,
   getModelsPath
@@ -65,7 +68,7 @@ export function emitModelFiles(
       codeModel.settings.sourceRoot,
       enumType.namespace
     );
-    emitType(sdkContext, rawEnum, sourceFile);
+    emitEnumFromCodeModel(sdkContext, enumType, rawEnum as any, sourceFile);
   }
 
   for (const unionType of codeModel.unions) {
@@ -133,6 +136,30 @@ function getTypeKey(name: string, namespace: string[]): string {
   return `${namespace.join("/")}:${name}`;
 }
 
+function emitEnumFromCodeModel(
+  sdkContext: SdkContext,
+  enumType: TSEnum,
+  rawEnum: any,
+  sourceFile: SourceFile
+): void {
+  const [enumDeclaration, knownValuesEnum] = buildEnumTypes(
+    sdkContext,
+    rawEnum,
+    false,
+    enumType.isExtensible
+  );
+
+  if (enumDeclaration.name.startsWith("_")) {
+    return;
+  }
+
+  if (enumType.isExtensible) {
+    addDeclaration(sourceFile, knownValuesEnum, refkey(rawEnum, "knownValues"));
+  }
+
+  addDeclaration(sourceFile, enumDeclaration, rawEnum);
+}
+
 function getOrCreateModelsFile(
   project: Project,
   sourceRoot: string,

packages/typespec-ts/src/modular/emitModels.ts

@@ -601,7 +601,8 @@ function buildNullableType(context: SdkContext, type: SdkNullableType) {
 export function buildEnumTypes(
   context: SdkContext,
   type: SdkEnumType,
-  reportMemberNameDiagnostic = false // if reportMemberNameDiagnostic is true, it will report diagnostic for enum member name
+  reportMemberNameDiagnostic = false, // if reportMemberNameDiagnostic is true, it will report diagnostic for enum member name
+  treatAsExtensible = isExtensibleEnum(context, type)
 ): [TypeAliasDeclarationStructure, EnumDeclarationStructure] {
   const rawMembers = type.values.map((value) =>
     emitEnumMember(context, value, reportMemberNameDiagnostic)
@@ -617,14 +618,14 @@ export function buildEnumTypes(
     kind: StructureKind.TypeAlias,
     name: normalizeModelName(context, type),
     isExported: true,
-    type: !isExtensibleEnum(context, type)
+    type: !treatAsExtensible
       ? type.values.map((v) => getTypeExpression(context, v)).join(" | ")
       : getTypeExpression(context, type.valueType)
   };
 
   const docs = type.doc ? type.doc : "Type of " + enumAsUnion.name;
   enumAsUnion.docs =
-    isExtensibleEnum(context, type) && type.doc
+    treatAsExtensible && type.doc
       ? [getExtensibleEnumDescription(context, type) ?? docs]
       : [docs];
   enumDeclaration.docs = type.doc

packages/typespec-ts/src/tcgcadapter/adapter.ts

@@ -1203,9 +1203,6 @@ function hasModelUsage(usage: UsageFlags | undefined): boolean {
   );
 }
 
-// TODO(strategy-b/U1): extensible enum filtering is a separate adaptEnums fix.
-// Keep this migration scoped to helper-type registration so U1 can land in a
-// focused follow-up.
 function shouldAdaptEnum(
   sdkContext: SdkContext,
   enumType: SdkEnumType

packages/typespec-ts/test/modularUnit/models-extensible-enums.spec.ts

@@ -0,0 +1,77 @@
+import { describe, expect, it } from "vitest";
+import { Project } from "ts-morph";
+import { useContext } from "../../src/contextManager.js";
+import { useBinder } from "../../src/framework/hooks/binder.js";
+import { adaptToCodeModel } from "../../src/tcgcadapter/adapter.js";
+import { transformModularEmitterOptions } from "../../src/modular/buildModularOptions.js";
+import { emitModelFiles } from "../../src/codegen/models.js";
+import { renameClientName } from "../../src/index.js";
+import { createDpgContextTestHelper, rlcEmitterFor } from "../util/testUtil.js";
+
+async function emitModels(body: string): Promise<string> {
+  const typeSpec = `
+    import "@typespec/http";
+    import "@typespec/rest";
+    import "@typespec/versioning";
+    import "@azure-tools/typespec-client-generator-core";
+    import "@azure-tools/typespec-azure-core";
+
+    using Http;
+    using Rest;
+    using Versioning;
+    using Azure.ClientGenerator.Core;
+    using Azure.Core;
+    using Azure.Core.Traits;
+
+    @service(#{ title: "Azure TypeScript Testing" })
+    namespace Azure.TypeScript.Testing {
+      ${body}
+    }
+  `;
+
+  const host = await rlcEmitterFor(typeSpec, { withRawContent: true });
+  const sdkContext = await createDpgContextTestHelper(host.program, false, {
+    isModularLibrary: true
+  });
+  sdkContext.rlcOptions!.isModularLibrary = true;
+
+  const emitterOptions = transformModularEmitterOptions(sdkContext, "", {
+    casing: "camel"
+  });
+  for (const client of sdkContext.sdkPackage.clients) {
+    await renameClientName(client, emitterOptions);
+  }
+
+  const codeModel = adaptToCodeModel({ sdkContext, emitterOptions });
+  const project = useContext("outputProject") as Project;
+  emitModelFiles(project, codeModel, sdkContext);
+  useBinder().resolveAllReferences(codeModel.settings.sourceRoot);
+
+  return project
+    .getSourceFiles(`${codeModel.settings.sourceRoot}/models/**/*.ts`)
+    .map((file) => file.getFullText())
+    .join("\n");
+}
+
+describe("models extensible enums", () => {
+  it("emits KnownXxx declarations from enum IR", async () => {
+    const modelsText = await emitModels(`
+      union PetKind {
+        dog: "dog",
+        cat: "cat",
+        string,
+      }
+
+      model PetEnvelope {
+        kind: PetKind;
+      }
+
+      @route("/pets")
+      @get
+      op getPet(): PetEnvelope;
+    `);
+
+    expect(modelsText).toContain("export enum KnownPetKind");
+    expect(modelsText).toContain("export type PetKind = string;");
+  });
+});