Add type safe returns

Author: louislepper

modules/openapi-generator/src/main/java/org/openapitools/codegen/languages/KotlinSpringServerCodegen.java

@@ -97,6 +97,7 @@ public class KotlinSpringServerCodegen extends AbstractKotlinCodegen
     public static final String USE_REQUEST_MAPPING_ON_CONTROLLER = "useRequestMappingOnController";
     public static final String USE_REQUEST_MAPPING_ON_INTERFACE = "useRequestMappingOnInterface";
     public static final String AUTO_X_SPRING_PAGINATED = "autoXSpringPaginated";
+    public static final String USE_SEALED_RESPONSE_INTERFACES = "useSealedResponseInterfaces";
 
     @Getter
     public enum DeclarativeInterfaceReactiveMode {
@@ -161,13 +162,19 @@ public String getDescription() {
     @Setter private DeclarativeInterfaceReactiveMode declarativeInterfaceReactiveMode = DeclarativeInterfaceReactiveMode.coroutines;
     @Setter private boolean useResponseEntity = true;
     @Setter private boolean autoXSpringPaginated = false;
+    @Setter private boolean useSealedResponseInterfaces = false;
 
     @Getter @Setter
     protected boolean useSpringBoot3 = false;
     protected RequestMappingMode requestMappingMode = RequestMappingMode.controller;
     private DocumentationProvider documentationProvider;
     private AnnotationLibrary annotationLibrary;
 
+    // Map to track which models implement which sealed response interfaces
+    private Map<String, List<String>> modelToSealedInterfaces = new HashMap<>();
+    private Map<String, String> sealedInterfaceToOperationId = new HashMap<>();
+    private boolean sealedInterfacesFileWritten = false;
+
     public KotlinSpringServerCodegen() {
         super();
 
@@ -250,6 +257,9 @@ public KotlinSpringServerCodegen() {
         addSwitch(USE_RESPONSE_ENTITY,
                 "Whether (when false) to return actual type (e.g. List<Fruit>) and handle non-happy path responses via exceptions flow or (when true) return entire ResponseEntity (e.g. ResponseEntity<List<Fruit>>). If disabled, method are annotated using a @ResponseStatus annotation, which has the status of the first response declared in the Api definition",
                 useResponseEntity);
+        addSwitch(USE_SEALED_RESPONSE_INTERFACES,
+                "Generate sealed interfaces for endpoint responses that all possible response types implement. Allows controllers to return any valid response type in a type-safe manner (e.g., sealed interface CreateUserResponse implemented by User, ConflictResponse, ErrorResponse)",
+                useSealedResponseInterfaces);
         addOption(X_KOTLIN_IMPLEMENTS_SKIP, "A list of fully qualified interfaces that should NOT be implemented despite their presence in vendor extension `x-kotlin-implements`. Example: yaml `xKotlinImplementsSkip: [com.some.pack.WithPhotoUrls]` skips implementing the interface in any schema", "empty list");
         addOption(X_KOTLIN_IMPLEMENTS_FIELDS_SKIP, "A list of fields per schema name that should NOT be created with `override` keyword despite their presence in vendor extension `x-kotlin-implements-fields` for the schema. Example: yaml `xKotlinImplementsFieldsSkip: Pet: [photoUrls]` skips `override` for `photoUrls` in schema `Pet`", "empty map");
         addOption(SCHEMA_IMPLEMENTS, "A map of single interface or a list of interfaces per schema name that should be implemented (serves similar purpose as `x-kotlin-implements`, but is fully decoupled from the api spec). Example: yaml `schemaImplements: {Pet: com.some.pack.WithId, Category: [com.some.pack.CategoryInterface], Dog: [com.some.pack.Canine, com.some.pack.OtherInterface]}` implements interfaces in schemas `Pet` (interface `com.some.pack.WithId`), `Category` (interface `com.some.pack.CategoryInterface`), `Dog`(interfaces `com.some.pack.Canine`, `com.some.pack.OtherInterface`)", "empty map");
@@ -496,6 +506,12 @@ public void processOpts() {
             this.setUseResponseEntity(Boolean.parseBoolean(additionalProperties.get(USE_RESPONSE_ENTITY).toString()));
         }
         writePropertyBack(USE_RESPONSE_ENTITY, useResponseEntity);
+
+        if(additionalProperties.containsKey(USE_SEALED_RESPONSE_INTERFACES)) {
+            this.setUseSealedResponseInterfaces(Boolean.parseBoolean(additionalProperties.get(USE_SEALED_RESPONSE_INTERFACES).toString()));
+        }
+        writePropertyBack(USE_SEALED_RESPONSE_INTERFACES, useSealedResponseInterfaces);
+
         additionalProperties.put("springHttpStatus", new SpringHttpStatusLambda());
 
         // Set basePackage from invokerPackage
@@ -1024,6 +1040,34 @@ public void preprocessOpenAPI(OpenAPI openAPI) {
             this.additionalProperties.put(SERVER_PORT, URLPathUtils.getPort(url, 8080));
         }
 
+        // Build modelToSealedInterfaces map early, before models are processed
+        // Note: We check additionalProperties here because processOpts() hasn't been called yet
+        boolean shouldUseSealedInterfaces = additionalProperties.containsKey(USE_SEALED_RESPONSE_INTERFACES)
+            && Boolean.parseBoolean(additionalProperties.get(USE_SEALED_RESPONSE_INTERFACES).toString());
+        if (shouldUseSealedInterfaces && openAPI.getPaths() != null) {
+            openAPI.getPaths().forEach((pathName, pathItem) -> {
+                pathItem.readOperations().forEach(operation -> {
+                    if (operation.getOperationId() != null && operation.getResponses() != null) {
+                        String sealedInterfaceName = camelize(operation.getOperationId()) + "Response";
+                        sealedInterfaceToOperationId.put(sealedInterfaceName, operation.getOperationId());
+
+                        operation.getResponses().forEach((statusCode, response) -> {
+                            if (response.getContent() != null) {
+                                response.getContent().forEach((mediaType, content) -> {
+                                    if (content.getSchema() != null && content.getSchema().get$ref() != null) {
+                                        String ref = content.getSchema().get$ref();
+                                        String modelName = ModelUtils.getSimpleRef(ref);
+                                        modelToSealedInterfaces.computeIfAbsent(modelName, k -> new ArrayList<>())
+                                            .add(sealedInterfaceName);
+                                    }
+                                });
+                            }
+                        });
+                    }
+                });
+            });
+        }
+
         // TODO: Handle tags
     }
 
@@ -1079,6 +1123,53 @@ public ModelsMap postProcessModelsEnum(ModelsMap objs) {
                     imports.add(itemJsonProperty);
                 });
 
+        // Add sealed interface implementations to models if enabled
+        // Note: We check additionalProperties here because processOpts() may not have been called yet
+        boolean shouldUseSealedInterfaces = additionalProperties.containsKey(USE_SEALED_RESPONSE_INTERFACES)
+            && Boolean.parseBoolean(additionalProperties.get(USE_SEALED_RESPONSE_INTERFACES).toString());
+        if (shouldUseSealedInterfaces) {
+            objs.getModels().stream()
+                    .map(ModelMap::getModel)
+                    .forEach(cm -> {
+                        String modelName = cm.classname;
+                        if (modelToSealedInterfaces.containsKey(modelName)) {
+                            List<String> sealedInterfaces = modelToSealedInterfaces.get(modelName);
+                            cm.vendorExtensions.put("x-implements-sealed-interfaces", sealedInterfaces);
+
+                            // Add imports for each sealed interface
+                            for (String sealedInterface : sealedInterfaces) {
+                                String importStatement = modelPackage + "." + sealedInterface;
+                                cm.imports.add(sealedInterface);
+                                Map<String, String> item = new HashMap<>();
+                                item.put("import", importStatement);
+                                imports.add(item);
+                            }
+                        }
+                    });
+
+            // Write sealed interfaces file once
+            if (!sealedInterfacesFileWritten && !sealedInterfaceToOperationId.isEmpty()) {
+                List<Map<String, String>> sealedInterfacesList = new ArrayList<>();
+                sealedInterfaceToOperationId.forEach((sealedInterfaceName, operationId) -> {
+                    Map<String, String> sealedInterface = new HashMap<>();
+                    sealedInterface.put("name", sealedInterfaceName);
+                    sealedInterface.put("operationId", operationId);
+                    sealedInterfacesList.add(sealedInterface);
+                });
+
+                Map<String, Object> sealedInterfacesData = new HashMap<>();
+                sealedInterfacesData.put("package", modelPackage);
+                sealedInterfacesData.put("sealedInterfaces", sealedInterfacesList);
+
+                additionalProperties.put("sealedInterfacesData", sealedInterfacesData);
+                supportingFiles.add(new SupportingFile("sealedResponseInterfaces.mustache",
+                        (sourceFolder + File.separator + modelPackage).replace(".", File.separator),
+                        "SealedResponseInterfaces.kt"));
+
+                sealedInterfacesFileWritten = true;
+            }
+        }
+
         return objs;
     }
 
@@ -1147,10 +1238,54 @@ public void setReturnContainer(final String returnContainer) {
                         operation.returnContainer = returnContainer;
                     }
                 });
+
+                // Generate sealed response interface metadata if enabled
+                if (useSealedResponseInterfaces && responses != null && !responses.isEmpty()) {
+                    // Generate sealed interface name from operation ID
+                    String sealedInterfaceName = camelize(operation.operationId) + "Response";
+                    operation.vendorExtensions.put("x-sealed-response-interface", sealedInterfaceName);
+
+                    // Collect all unique response base types (models)
+                    List<String> responseTypes = responses.stream()
+                            .map(r -> r.baseType)
+                            .filter(baseType -> baseType != null && !baseType.isEmpty())
+                            .distinct()
+                            .collect(Collectors.toList());
+
+                    operation.vendorExtensions.put("x-sealed-response-types", responseTypes);
+
+                    // Track which models should implement this sealed interface
+                    for (String responseType : responseTypes) {
+                        modelToSealedInterfaces.computeIfAbsent(responseType, k -> new ArrayList<>())
+                                .add(sealedInterfaceName);
+                    }
+                }
+
 //                if(implicitHeaders){
 //                    removeHeadersFromAllParams(operation.allParams);
 //                }
             });
+
+            // Add imports for sealed interfaces if feature is enabled
+            if (useSealedResponseInterfaces) {
+                Set<String> sealedInterfacesToImport = new HashSet<>();
+                ops.forEach(operation -> {
+                    if (operation.vendorExtensions.containsKey("x-sealed-response-interface")) {
+                        String sealedInterfaceName = (String) operation.vendorExtensions.get("x-sealed-response-interface");
+                        operation.imports.add(sealedInterfaceName);
+                        sealedInterfacesToImport.add(sealedInterfaceName);
+                    }
+                });
+
+                // Add import statements to the operations imports map
+                List<Map<String, String>> imports = objs.getImports();
+                for (String sealedInterfaceName : sealedInterfacesToImport) {
+                    String importStatement = modelPackage + "." + sealedInterfaceName;
+                    Map<String, String> item = new HashMap<>();
+                    item.put("import", importStatement);
+                    imports.add(item);
+                }
+            }
         }
 
         return objs;
@@ -1159,6 +1294,12 @@ public void setReturnContainer(final String returnContainer) {
     @Override
     public Map<String, Object> postProcessSupportingFileData(Map<String, Object> objs) {
         generateYAMLSpecFile(objs);
+
+        // Add sealed interfaces data if available
+        if (additionalProperties.containsKey("sealedInterfacesData")) {
+            objs.putAll((Map<String, Object>) additionalProperties.get("sealedInterfacesData"));
+        }
+
         return objs;
     }
 

modules/openapi-generator/src/main/resources/kotlin-spring/apiInterface.mustache

@@ -116,7 +116,7 @@ interface {{classname}} {
         {{/hasParams}}{{#swagger1AnnotationLibrary}}@ApiParam(hidden = true) {{/swagger1AnnotationLibrary}}{{#swagger2AnnotationLibrary}}@Parameter(hidden = true) {{/swagger2AnnotationLibrary}}{{#reactive}}exchange: org.springframework.web.server.ServerWebExchange{{/reactive}}{{^reactive}}request: {{javaxPackage}}.servlet.http.HttpServletRequest{{/reactive}}{{/includeHttpRequestContext}}{{#vendorExtensions.x-spring-paginated}}{{#hasParams}},
         {{/hasParams}}{{^hasParams}}{{#includeHttpRequestContext}},
         {{/includeHttpRequestContext}}{{/hasParams}}{{#swagger1AnnotationLibrary}}@ApiParam(hidden = true) {{/swagger1AnnotationLibrary}}{{#swagger2AnnotationLibrary}}@Parameter(hidden = true) {{/swagger2AnnotationLibrary}}pageable: Pageable{{/vendorExtensions.x-spring-paginated}}{{#hasParams}}
-    {{/hasParams}}): {{#useResponseEntity}}ResponseEntity<{{/useResponseEntity}}{{>returnTypes}}{{#useResponseEntity}}>{{/useResponseEntity}}{{^skipDefaultApiInterface}} {
+    {{/hasParams}}): {{#useResponseEntity}}ResponseEntity<{{/useResponseEntity}}{{#useSealedResponseInterfaces}}{{#vendorExtensions.x-sealed-response-interface}}{{vendorExtensions.x-sealed-response-interface}}{{/vendorExtensions.x-sealed-response-interface}}{{/useSealedResponseInterfaces}}{{^useSealedResponseInterfaces}}{{>returnTypes}}{{/useSealedResponseInterfaces}}{{#useResponseEntity}}>{{/useResponseEntity}}{{^skipDefaultApiInterface}} {
         {{^isDelegate}}
         return {{>returnValue}}
         {{/isDelegate}}

modules/openapi-generator/src/main/resources/kotlin-spring/dataClass.mustache

@@ -20,12 +20,17 @@
 ){{/discriminator}}{{! no newline
 }}{{#parent}} : {{{.}}}{{#isMap}}(){{/isMap}}{{! no newline
       }}{{#vendorExtensions.x-kotlin-implements}}, {{{.}}}{{/vendorExtensions.x-kotlin-implements}}{{! <- serializableModel is also handled via x-kotlin-implements
+      }}{{#vendorExtensions.x-implements-sealed-interfaces}}{{#.}}, {{{.}}}{{/.}}{{/vendorExtensions.x-implements-sealed-interfaces}}{{! <- add sealed interface implementations
 }}{{/parent}}{{! no newline
 }}{{^parent}}{{! no newline
     }}{{#vendorExtensions.x-kotlin-implements}}{{! no newline
     }}{{#-first}} : {{{.}}}{{/-first}}{{! no newline
       }}{{^-first}}, {{{.}}}{{/-first}}{{! no newline
     }}{{/vendorExtensions.x-kotlin-implements}}{{! no newline
+    }}{{#vendorExtensions.x-implements-sealed-interfaces}}{{! no newline
+    }}{{#-first}}{{^vendorExtensions.x-kotlin-implements}} : {{{.}}}{{/vendorExtensions.x-kotlin-implements}}{{#vendorExtensions.x-kotlin-implements}}, {{{.}}}{{/vendorExtensions.x-kotlin-implements}}{{/-first}}{{! no newline
+      }}{{^-first}}, {{{.}}}{{/-first}}{{! no newline
+    }}{{/vendorExtensions.x-implements-sealed-interfaces}}{{! no newline
 }}{{/parent}} {
 {{#discriminator}}
     {{#requiredVars}}

modules/openapi-generator/src/main/resources/kotlin-spring/sealedResponseInterfaces.mustache

@@ -0,0 +1,9 @@
+package {{package}}
+
+{{#sealedInterfaces}}
+/**
+ * Sealed interface for all possible responses from {{operationId}}
+ */
+sealed interface {{name}}
+
+{{/sealedInterfaces}}