Fix @semanticNonNull on mutation types (#8239)

Author: tobias-tengler

src/HotChocolate/Core/src/Types/SemanticNonNullTypeInterceptor.cs

@@ -15,6 +15,7 @@ namespace HotChocolate;
 internal sealed class SemanticNonNullTypeInterceptor : TypeInterceptor
 {
     private ITypeInspector _typeInspector = null!;
+    private ObjectTypeDefinition? _mutationDef;
 
     public override bool IsEnabled(IDescriptorContext context)
         => context.Options.EnableSemanticNonNull;
@@ -29,52 +30,16 @@ internal override void InitializeContext(
         _typeInspector = context.TypeInspector;
     }
 
-    /// <summary>
-    /// After the root types have been resolved, we go through all the fields of the mutation type
-    /// and undo semantic non-nullability. This is because mutations can be chained and we want to retain
-    /// the null-bubbling so execution is aborted once one non-null mutation field produces an error.
-    /// We have to do this in a different hook because the mutation type is not yet fully resolved in the
-    /// <see cref="OnAfterCompleteName"/> hook.
-    /// </summary>
-    public override void OnAfterResolveRootType(
-        ITypeCompletionContext completionContext,
-        ObjectTypeDefinition definition,
+    public override void OnAfterResolveRootType(ITypeCompletionContext completionContext, ObjectTypeDefinition definition,
         OperationType operationType)
     {
-        if (operationType == OperationType.Mutation)
+        if (operationType is OperationType.Mutation)
         {
-            foreach (var field in definition.Fields)
-            {
-                if (field.IsIntrospectionField)
-                {
-                    continue;
-                }
-
-                if (!field.HasDirectives)
-                {
-                    continue;
-                }
-
-                var semanticNonNullDirective =
-                    field.Directives.FirstOrDefault(d => d.Value is SemanticNonNullDirective);
-
-                if (semanticNonNullDirective is not null)
-                {
-                    field.Directives.Remove(semanticNonNullDirective);
-                }
-
-                var semanticNonNullFormatterDefinition =
-                    field.FormatterDefinitions.FirstOrDefault(fd => fd.Key == WellKnownMiddleware.SemanticNonNull);
-
-                if (semanticNonNullFormatterDefinition is not null)
-                {
-                    field.FormatterDefinitions.Remove(semanticNonNullFormatterDefinition);
-                }
-            }
+            _mutationDef = definition;
         }
     }
 
-    public override void OnAfterCompleteName(ITypeCompletionContext completionContext, DefinitionBase definition)
+    public override void OnBeforeCompleteType(ITypeCompletionContext completionContext, DefinitionBase definition)
     {
         if (completionContext.IsIntrospectionType)
         {
@@ -88,6 +53,14 @@ public override void OnAfterCompleteName(ITypeCompletionContext completionContex
                 return;
             }
 
+            // We undo semantic non-nullability on each mutation field, since mutations can be chained
+            // and we want to retain the null-bubbling so execution is aborted if a non-null mutation field
+            // produces an error.
+            if (objectDef == _mutationDef)
+            {
+                return;
+            }
+
             foreach (var field in objectDef.Fields)
             {
                 if (field.IsIntrospectionField)

src/HotChocolate/Core/test/Types.Tests/SemanticNonNullTests.cs

@@ -35,8 +35,9 @@ public async Task MutationConventions()
                 o.StrictValidation = false;
                 o.EnableSemanticNonNull = true;
             })
-            .AddMutationConventions()
+            .AddMutationConventions(applyToAllMutations: false)
             .AddMutationType<Mutation>()
+            .AddTypeExtension<MutationExtensions>()
             .BuildSchemaAsync()
             .MatchSnapshotAsync();
     }
@@ -319,6 +320,12 @@ public class Mutation
         public bool DoSomething() => true;
     }
 
+    [ExtendObjectType<Mutation>]
+    public class MutationExtensions
+    {
+        public bool DoSomethingElse() => true;
+    }
+
     public class MyException : Exception;
 
     public class QueryWithPagination

src/HotChocolate/Core/test/Types.Tests/__snapshots__/SemanticNonNullTests.MutationConventions.snap

@@ -13,6 +13,7 @@ type DoSomethingPayload {
 
 type Mutation {
   doSomething: DoSomethingPayload!
+  doSomethingElse: Boolean!
 }
 
 type MyError implements Error {