[OpenAPI] Add @responseBody directive

src/HotChocolate/Adapters/src/Adapters.OpenApi.AspNetCore/DynamicOpenApiDocumentTransformer.cs

@@ -21,7 +21,9 @@
 
 namespace HotChocolate.Adapters.OpenApi;
 
-internal sealed class DynamicOpenApiDocumentTransformer : IOpenApiDocumentTransformer, IDynamicOpenApiDocumentTransformer
+internal sealed class DynamicOpenApiDocumentTransformer
+    : IOpenApiDocumentTransformer,
+      IDynamicOpenApiDocumentTransformer
 {
     private const string JsonContentType = "application/json";
 
@@ -301,18 +303,13 @@ private EndpointDescriptor CreateEndpointDescriptor(OpenApiEndpointDefinition en
                 operation.RequestBody = requestBody;
             }
 
-            var operationType = _schema.GetOperationType(endpoint.OperationDefinition.Operation);
+            var responseBodySelection = endpoint.GetResponseBodySelection(_schema);
+            var fieldType = responseBodySelection.FieldType
+                ?? throw new InvalidOperationException("Expected to resolve the response field type.");
 
-            if (endpoint.OperationDefinition.SelectionSet.Selections is not [FieldNode rootField])
-            {
-                throw new InvalidOperationException("Expected to have a single field selection on the root");
-            }
-
-            var fieldType = operationType.Fields[rootField.Name.Value].Type;
-
-            var responseSchema = rootField.SelectionSet is not null
+            var responseSchema = responseBodySelection.SelectionSet is not null
                 ? CreateOpenApiSchemaForSelectionSet(
-                    rootField.SelectionSet,
+                    responseBodySelection.SelectionSet,
                     fieldType,
                     endpoint.LocalFragmentsByName)
                 : CreateOpenApiSchemaForType(fieldType);

src/HotChocolate/Adapters/src/Adapters.OpenApi.Core/Definitions/OpenApiEndpointDefinition.cs

@@ -3,49 +3,14 @@
 
 namespace HotChocolate.Adapters.OpenApi;
 
-public sealed class OpenApiEndpointDefinition : IOpenApiDefinition
-{
-    public OpenApiEndpointDefinition(
-        string httpMethod,
-        string route,
-        string? description,
-        ImmutableArray<OpenApiEndpointDefinitionParameter> routeParameters,
-        ImmutableArray<OpenApiEndpointDefinitionParameter> queryParameters,
-        string? bodyVariableName,
-        DocumentNode document,
-        OperationDefinitionNode operationDefinition,
-        Dictionary<string, FragmentDefinitionNode> localFragmentsByName,
-        HashSet<string> externalFragmentReferences)
-    {
-        HttpMethod = httpMethod;
-        Route = route;
-        Description = description;
-        RouteParameters = routeParameters;
-        QueryParameters = queryParameters;
-        BodyVariableName = bodyVariableName;
-        Document = document;
-        OperationDefinition = operationDefinition;
-        LocalFragmentsByName = localFragmentsByName;
-        ExternalFragmentReferences = externalFragmentReferences;
-    }
-
-    public string HttpMethod { get; }
-
-    public string Route { get; }
-
-    public string? Description { get; }
-
-    public ImmutableArray<OpenApiEndpointDefinitionParameter> RouteParameters { get; }
-
-    public ImmutableArray<OpenApiEndpointDefinitionParameter> QueryParameters { get; }
-
-    public string? BodyVariableName { get; }
-
-    public DocumentNode Document { get; }
-
-    public OperationDefinitionNode OperationDefinition { get; }
-
-    public Dictionary<string, FragmentDefinitionNode> LocalFragmentsByName { get; }
-
-    public HashSet<string> ExternalFragmentReferences { get; }
-}
+public sealed record OpenApiEndpointDefinition(
+    string HttpMethod,
+    string Route,
+    string? Description,
+    ImmutableArray<OpenApiEndpointDefinitionParameter> RouteParameters,
+    ImmutableArray<OpenApiEndpointDefinitionParameter> QueryParameters,
+    string? BodyVariableName,
+    DocumentNode Document,
+    OperationDefinitionNode OperationDefinition,
+    Dictionary<string, FragmentDefinitionNode> LocalFragmentsByName,
+    HashSet<string> ExternalFragmentReferences) : IOpenApiDefinition;

src/HotChocolate/Adapters/src/Adapters.OpenApi.Core/Definitions/OpenApiModelDefinition.cs

@@ -2,33 +2,10 @@
 
 namespace HotChocolate.Adapters.OpenApi;
 
-public sealed record OpenApiModelDefinition : IOpenApiDefinition
-{
-    public OpenApiModelDefinition(
-        string name,
-        string? description,
-        DocumentNode document,
-        FragmentDefinitionNode fragmentDefinition,
-        Dictionary<string, FragmentDefinitionNode> localFragmentsByName,
-        HashSet<string> externalFragmentReferences)
-    {
-        Name = name;
-        Description = description;
-        Document = document;
-        FragmentDefinition = fragmentDefinition;
-        LocalFragmentsByName = localFragmentsByName;
-        ExternalFragmentReferences = externalFragmentReferences;
-    }
-
-    public string Name { get; }
-
-    public string? Description { get; }
-
-    public DocumentNode Document { get; }
-
-    public FragmentDefinitionNode FragmentDefinition { get; }
-
-    public Dictionary<string, FragmentDefinitionNode> LocalFragmentsByName { get; }
-
-    public HashSet<string> ExternalFragmentReferences { get; }
-}
+public sealed record OpenApiModelDefinition(
+    string Name,
+    string? Description,
+    DocumentNode Document,
+    FragmentDefinitionNode FragmentDefinition,
+    Dictionary<string, FragmentDefinitionNode> LocalFragmentsByName,
+    HashSet<string> ExternalFragmentReferences) : IOpenApiDefinition;

src/HotChocolate/Adapters/src/Adapters.OpenApi.Core/Execution/OpenApiEndpointDefinitionExtensions.cs

@@ -0,0 +1,129 @@
+using System.Collections.Immutable;
+using System.Runtime.InteropServices;
+using HotChocolate.Language;
+using HotChocolate.Language.Visitors;
+using HotChocolate.Types;
+
+namespace HotChocolate.Adapters.OpenApi;
+
+internal static class OpenApiEndpointDefinitionExtensions
+{
+    private static readonly ResponseBodySelectionFinder s_responseBodySelectionFinder = new();
+
+    public static OpenApiResponseBodySelection GetResponseBodySelection(
+        this OpenApiEndpointDefinition endpoint,
+        ISchemaDefinition schema)
+    {
+        var operation = endpoint.OperationDefinition;
+        var rootType = schema.GetOperationType(operation.Operation);
+
+        return FindResponseBody(operation.SelectionSet, rootType)
+            ?? CreateDefaultResponseBody(operation.SelectionSet, rootType);
+    }
+
+    private static OpenApiResponseBodySelection CreateDefaultResponseBody(
+        SelectionSetNode selectionSet,
+        IOutputType? rootType)
+    {
+        var rootField = selectionSet.Selections.FirstOrDefault() as FieldNode
+            ?? throw new InvalidOperationException("Expected to have a response field.");
+
+        return new OpenApiResponseBodySelection(
+            [rootField.Alias?.Value ?? rootField.Name.Value],
+            rootField.SelectionSet,
+            ResolveFieldType(rootField, rootType));
+    }
+
+    private static OpenApiResponseBodySelection? FindResponseBody(
+        SelectionSetNode selectionSet,
+        IOutputType? rootType)
+    {
+        var context = new ResponseBodySelectionFinderContext(rootType);
+        s_responseBodySelectionFinder.Visit(selectionSet, context);
+        return context.ResponseBodySelection;
+    }
+
+    private static IOutputType? ResolveFieldType(FieldNode field, IOutputType? declaringType)
+    {
+        if (declaringType?.NamedType() is IComplexTypeDefinition complexType
+            && complexType.Fields.TryGetField(field.Name.Value, out var fieldDefinition))
+        {
+            return fieldDefinition.Type;
+        }
+
+        return null;
+    }
+
+    private sealed class ResponseBodySelectionFinder
+        : SyntaxWalker<ResponseBodySelectionFinderContext>
+    {
+        protected override ISyntaxVisitorAction Enter(
+            FieldNode node,
+            ResponseBodySelectionFinderContext context)
+        {
+            var declaringType = context.Path.Count == 0
+                ? context.RootType
+                : context.Path[^1].FieldType;
+            var fieldType = ResolveFieldType(node, declaringType);
+            context.Path.Add(
+                new ResponseBodyPathSegment(
+                    node.Alias?.Value ?? node.Name.Value,
+                    fieldType));
+
+            if (node.Directives.Any(
+                    d => d.Name.Value == WellKnownDirectiveNames.ResponseBody))
+            {
+                context.ResponseBodySelection = new OpenApiResponseBodySelection(
+                    CreateResponseNamePath(context.Path),
+                    node.SelectionSet,
+                    fieldType);
+                return Break;
+            }
+
+            return Continue;
+        }
+
+        protected override ISyntaxVisitorAction Leave(
+            FieldNode node,
+            ResponseBodySelectionFinderContext context)
+        {
+            context.Path.RemoveAt(context.Path.Count - 1);
+            return Continue;
+        }
+
+        protected override ISyntaxVisitorAction Enter(
+            InlineFragmentNode node,
+            ResponseBodySelectionFinderContext context)
+            => node.TypeCondition is null ? Continue : Skip;
+
+        private static ImmutableArray<string> CreateResponseNamePath(
+            List<ResponseBodyPathSegment> path)
+        {
+            var responseNamePath = new string[path.Count];
+            for (var i = 0; i < path.Count; i++)
+            {
+                responseNamePath[i] = path[i].ResponseName;
+            }
+
+            return ImmutableCollectionsMarshal.AsImmutableArray(responseNamePath);
+        }
+    }
+
+    private sealed class ResponseBodySelectionFinderContext(IOutputType? rootType)
+    {
+        public IOutputType? RootType { get; } = rootType;
+
+        public List<ResponseBodyPathSegment> Path { get; } = [];
+
+        public OpenApiResponseBodySelection? ResponseBodySelection { get; set; }
+    }
+
+    private readonly record struct ResponseBodyPathSegment(
+        string ResponseName,
+        IOutputType? FieldType);
+}
+
+internal sealed record OpenApiResponseBodySelection(
+    ImmutableArray<string> ResponseNamePath,
+    SelectionSetNode? SelectionSet,
+    IOutputType? FieldType);

src/HotChocolate/Adapters/src/Adapters.OpenApi.Core/Execution/OpenApiEndpointDescriptor.cs

@@ -11,7 +11,7 @@ internal sealed record OpenApiEndpointDescriptor(
     RoutePattern Route,
     VariableValueInsertionTrie ParameterTrie,
     string? VariableFilledThroughBody,
-    string ResponseNameToExtract);
+    OpenApiResponseBodySelection ResponseBodySelection);
 
 internal interface IVariableValueInsertionTrieSegment;
 

src/HotChocolate/Adapters/src/Adapters.OpenApi.Core/Execution/OpenApiEndpointFactory.cs

@@ -58,43 +58,10 @@ public static OpenApiEndpointDescriptor CreateEndpointDescriptor(
         IDictionary<string, OpenApiModelDefinition> modelsByName,
         ISchemaDefinition schema)
     {
-        List<IExecutableDefinitionNode> definitions = [
-            ..endpointDefinition.Document.Definitions.OfType<IExecutableDefinitionNode>()
-        ];
-
-        var externalFragmentReferencesQueue = new Queue<string>(endpointDefinition.ExternalFragmentReferences);
-        var processedFragmentReferences = new HashSet<string>();
-
-        while (externalFragmentReferencesQueue.TryDequeue(out var referencedFragmentName))
-        {
-            if (!processedFragmentReferences.Add(referencedFragmentName))
-            {
-                continue;
-            }
-
-            if (!modelsByName.TryGetValue(referencedFragmentName, out var model))
-            {
-                continue;
-            }
-
-            foreach (var definition in model.Document.Definitions.OfType<IExecutableDefinitionNode>())
-            {
-                definitions.Add(definition);
-            }
-
-            foreach (var externalFragmentReference in model.ExternalFragmentReferences)
-            {
-                externalFragmentReferencesQueue.Enqueue(externalFragmentReference);
-            }
-        }
+        var responseBodySelection = endpointDefinition.GetResponseBodySelection(schema);
+        var document = ComposeExecutionDocument(endpointDefinition, modelsByName);
 
-        var document = new DocumentNode(definitions);
-
-        var rootField = endpointDefinition.OperationDefinition.SelectionSet.Selections
-            .OfType<FieldNode>()
-            .First();
-
-        var responseNameToExtract = rootField.Alias?.Value ?? rootField.Name.Value;
+        document = ResponseBodyDirectiveRewriter.Instance.Rewrite(document);
 
         var route = RoutePatternFactory.Parse(endpointDefinition.Route);
 
@@ -115,7 +82,7 @@ public static OpenApiEndpointDescriptor CreateEndpointDescriptor(
             route,
             parameterTrie,
             endpointDefinition.BodyVariableName,
-            responseNameToExtract);
+            responseBodySelection);
 
         void InsertParametersIntoTrie(
             IEnumerable<OpenApiEndpointDefinitionParameter> parameters,
@@ -186,6 +153,36 @@ void InsertParametersIntoTrie(
         }
     }
 
+    private static DocumentNode ComposeExecutionDocument(
+        OpenApiEndpointDefinition endpoint,
+        IDictionary<string, OpenApiModelDefinition> modelsByName)
+    {
+        List<IExecutableDefinitionNode> definitions =
+        [
+            .. endpoint.Document.Definitions.OfType<IExecutableDefinitionNode>()
+        ];
+        var pendingFragments = new Queue<string>(endpoint.ExternalFragmentReferences);
+        var processedFragments = new HashSet<string>();
+
+        while (pendingFragments.TryDequeue(out var fragmentName))
+        {
+            if (!processedFragments.Add(fragmentName)
+                || !modelsByName.TryGetValue(fragmentName, out var model))
+            {
+                continue;
+            }
+
+            definitions.AddRange(model.Document.Definitions.OfType<IExecutableDefinitionNode>());
+
+            foreach (var externalFragmentReference in model.ExternalFragmentReferences)
+            {
+                pendingFragments.Enqueue(externalFragmentReference);
+            }
+        }
+
+        return new DocumentNode(definitions);
+    }
+
     private static (ITypeDefinition Type, bool HasDefaultValue, bool IsNonNullType) GetParameterDetails(
         OpenApiEndpointDefinitionParameter parameter,
         OperationDefinitionNode operation,

src/HotChocolate/Adapters/src/Adapters.OpenApi.Core/Execution/ResponseBodyDirectiveRewriter.cs

@@ -0,0 +1,27 @@
+using HotChocolate.Language;
+using HotChocolate.Language.Visitors;
+
+namespace HotChocolate.Adapters.OpenApi;
+
+internal sealed class ResponseBodyDirectiveRewriter : SyntaxRewriter<object?>
+{
+    public static ResponseBodyDirectiveRewriter Instance { get; } = new();
+
+    public DocumentNode Rewrite(DocumentNode document)
+        => (DocumentNode)Rewrite(document, null)!;
+
+    protected override FieldNode? RewriteField(FieldNode node, object? context)
+    {
+        var rewritten = base.RewriteField(node, context)!;
+        if (!rewritten.Directives.Any(
+                d => d.Name.Value == WellKnownDirectiveNames.ResponseBody))
+        {
+            return rewritten;
+        }
+
+        return rewritten.WithDirectives(
+            rewritten.Directives
+                .Where(d => d.Name.Value != WellKnownDirectiveNames.ResponseBody)
+                .ToArray());
+    }
+}