Reset OpenApiSchemaTransformerContext before each invocation (#56905)

* Reset OpenApiSchemaTransformerContext before each invocation

* Only call UpdateJsonTypeInfo once

Author: captainsafia

src/OpenApi/src/Services/Schemas/OpenApiSchemaService.cs

@@ -155,6 +155,7 @@ internal async Task ApplySchemaTransformersAsync(OpenApiSchema schema, Type type
         };
         for (var i = 0; i < _openApiOptions.SchemaTransformers.Count; i++)
         {
+            // Reset context object to base state before running each transformer.
             var transformer = _openApiOptions.SchemaTransformers[i];
             // If the transformer is a type-based transformer, we need to initialize and finalize it
             // once in the context of the top-level assembly and not the child properties we are invoking
@@ -164,7 +165,7 @@ internal async Task ApplySchemaTransformersAsync(OpenApiSchema schema, Type type
                 var initializedTransformer = typeBasedTransformer.InitializeTransformer(serviceProvider);
                 try
                 {
-                    await InnerApplySchemaTransformersAsync(schema, jsonTypeInfo, context, initializedTransformer, cancellationToken);
+                    await InnerApplySchemaTransformersAsync(schema, jsonTypeInfo, null, context, initializedTransformer, cancellationToken);
                 }
                 finally
                 {
@@ -173,17 +174,19 @@ internal async Task ApplySchemaTransformersAsync(OpenApiSchema schema, Type type
             }
             else
             {
-                await InnerApplySchemaTransformersAsync(schema, jsonTypeInfo, context, transformer, cancellationToken);
+                await InnerApplySchemaTransformersAsync(schema, jsonTypeInfo, null, context, transformer, cancellationToken);
             }
         }
     }
 
     private async Task InnerApplySchemaTransformersAsync(OpenApiSchema schema,
         JsonTypeInfo jsonTypeInfo,
+        JsonPropertyInfo? jsonPropertyInfo,
         OpenApiSchemaTransformerContext context,
         IOpenApiSchemaTransformer transformer,
         CancellationToken cancellationToken = default)
     {
+        context.UpdateJsonTypeInfo(jsonTypeInfo, jsonPropertyInfo);
         await transformer.TransformAsync(schema, context, cancellationToken);
 
         // Only apply transformers on polymorphic schemas where we can resolve the derived
@@ -194,31 +197,28 @@ private async Task InnerApplySchemaTransformersAsync(OpenApiSchema schema,
             foreach (var derivedType in jsonTypeInfo.PolymorphismOptions.DerivedTypes)
             {
                 var derivedJsonTypeInfo = _jsonSerializerOptions.GetTypeInfo(derivedType.DerivedType);
-                context.UpdateJsonTypeInfo(derivedJsonTypeInfo, null);
                 if (schema.AnyOf.Count <= anyOfIndex)
                 {
                     break;
                 }
-                await InnerApplySchemaTransformersAsync(schema.AnyOf[anyOfIndex], derivedJsonTypeInfo, context, transformer, cancellationToken);
+                await InnerApplySchemaTransformersAsync(schema.AnyOf[anyOfIndex], derivedJsonTypeInfo, null, context, transformer, cancellationToken);
                 anyOfIndex++;
             }
         }
 
         if (schema.Items is not null)
         {
             var elementTypeInfo = _jsonSerializerOptions.GetTypeInfo(jsonTypeInfo.ElementType!);
-            context.UpdateJsonTypeInfo(elementTypeInfo, null);
-            await InnerApplySchemaTransformersAsync(schema.Items, elementTypeInfo, context, transformer, cancellationToken);
+            await InnerApplySchemaTransformersAsync(schema.Items, elementTypeInfo, null, context, transformer, cancellationToken);
         }
 
         if (schema.Properties is { Count: > 0 })
         {
             foreach (var propertyInfo in jsonTypeInfo.Properties)
             {
-                context.UpdateJsonTypeInfo(_jsonSerializerOptions.GetTypeInfo(propertyInfo.PropertyType), propertyInfo);
                 if (schema.Properties.TryGetValue(propertyInfo.Name, out var propertySchema))
                 {
-                    await InnerApplySchemaTransformersAsync(propertySchema, _jsonSerializerOptions.GetTypeInfo(propertyInfo.PropertyType), context, transformer, cancellationToken);
+                    await InnerApplySchemaTransformersAsync(propertySchema, _jsonSerializerOptions.GetTypeInfo(propertyInfo.PropertyType), propertyInfo, context, transformer, cancellationToken);
                 }
             }
         }

src/OpenApi/test/Transformers/SchemaTransformerTests.cs

@@ -19,7 +19,28 @@ public async Task SchemaTransformer_CanAccessTypeAndParameterDescriptionForParam
         builder.MapPost("/todo", (Todo todo) => { });
 
         var options = new OpenApiOptions();
+        var firstInvocationOnSecondTransformer = true;
         options.AddSchemaTransformer((schema, context, cancellationToken) =>
+        {
+            ValidateContext(context);
+            return Task.CompletedTask;
+        })
+        .AddSchemaTransformer((schema, context, cancellationToken) =>
+        {
+            // Coverage for https://github.com/dotnet/aspnetcore/issues/56899
+            if (firstInvocationOnSecondTransformer)
+            {
+                Assert.Equal(typeof(Todo), context.JsonTypeInfo.Type);
+                firstInvocationOnSecondTransformer = false;
+            }
+            // Rest of the state is still consistent
+            ValidateContext(context);
+            return Task.CompletedTask;
+        });
+
+        await VerifyOpenApiDocument(builder, options, document => { });
+
+        static void ValidateContext(OpenApiSchemaTransformerContext context)
         {
             if (context.JsonPropertyInfo == null)
             {
@@ -42,10 +63,7 @@ public async Task SchemaTransformer_CanAccessTypeAndParameterDescriptionForParam
             {
                 Assert.Equal(typeof(DateTime), context.JsonTypeInfo.Type);
             }
-            return Task.CompletedTask;
-        });
-
-        await VerifyOpenApiDocument(builder, options, document => { });
+        }
     }
 
     [Fact]