Add support for AddOpenApiOperationTransformer API (#60566)

* Add support for AddOpenApiOperationTransformer API

* Apply all transformers in order instead of last-one-wins

* Add cancellation and async tests to transformers

* Fix Timespan declaration

* Avoid delays in async and cancellation tests

* Avoid delays in async and cancellation tests

Author: captainsafia

src/OpenApi/src/Extensions/OpenApiEndpointConventionBuilderExtensions.cs

@@ -121,4 +121,17 @@ private static void AddAndConfigureOperationForEndpoint(EndpointBuilder endpoint
             }
         }
     }
+
+    /// <summary>
+    /// Adds an OpenAPI operation transformer to the <see cref="EndpointBuilder.Metadata" /> associated
+    /// with the current endpoint.
+    /// </summary>
+    /// <param name="builder">The <see cref="IEndpointConventionBuilder"/>.</param>
+    /// <param name="transformer">The <see cref="Func{OpenApiOperation, OpenApiOperationTransformerContext, CancellationToken, Task}"/> that modifies the operation in the <see cref="OpenApiDocument"/>.</param>
+    /// <returns>A <see cref="IEndpointConventionBuilder"/> that can be used to further customize the endpoint.</returns>
+    public static TBuilder AddOpenApiOperationTransformer<TBuilder>(this TBuilder builder, Func<OpenApiOperation, OpenApiOperationTransformerContext, CancellationToken, Task> transformer) where TBuilder : IEndpointConventionBuilder
+    {
+        builder.WithMetadata(new DelegateOpenApiOperationTransformer(transformer));
+        return builder;
+    }
 }

src/OpenApi/src/PublicAPI.Unshipped.txt

@@ -1 +1,2 @@
 #nullable enable
+static Microsoft.AspNetCore.Builder.OpenApiEndpointConventionBuilderExtensions.AddOpenApiOperationTransformer<TBuilder>(this TBuilder builder, System.Func<Microsoft.OpenApi.Models.OpenApiOperation!, Microsoft.AspNetCore.OpenApi.OpenApiOperationTransformerContext!, System.Threading.CancellationToken, System.Threading.Tasks.Task!>! transformer) -> TBuilder

src/OpenApi/src/Services/OpenApiDocumentService.cs

@@ -268,6 +268,15 @@ private async Task<Dictionary<OperationType, OpenApiOperation>> GetOperationsAsy
                 var transformer = operationTransformers[i];
                 await transformer.TransformAsync(operation, operationContext, cancellationToken);
             }
+
+            // Apply any endpoint-specific operation transformers registered via
+            // the AddOpenApiOperationTransformer extension method.
+            var endpointOperationTransformers = description.ActionDescriptor.EndpointMetadata
+                .OfType<DelegateOpenApiOperationTransformer>();
+            foreach (var endpointOperationTransformer in endpointOperationTransformers)
+            {
+                await endpointOperationTransformer.TransformAsync(operation, operationContext, cancellationToken);
+            }
         }
         return operations;
     }

src/OpenApi/test/Microsoft.AspNetCore.OpenApi.Build.Tests/GenerateAdditionalXmlFilesForOpenApiTests.cs

@@ -9,7 +9,7 @@ namespace Microsoft.AspNetCore.OpenApi.Build.Tests;
 
 public class GenerateAdditionalXmlFilesForOpenApiTests
 {
-    private static readonly TimeSpan _defaultProcessTimeout = TimeSpan.FromSeconds(120);
+    private static readonly TimeSpan _defaultProcessTimeout = TimeSpan.FromMinutes(2);
 
     [Fact]
     public void VerifiesTargetGeneratesXmlFiles()

src/OpenApi/test/Microsoft.AspNetCore.OpenApi.Tests/Services/OpenApiDocumentServiceTestsBase.cs

@@ -3,12 +3,10 @@
 
 using System.Reflection;
 using System.Text;
-using System.Text.Json;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting.Server;
 using Microsoft.AspNetCore.Hosting.Server.Features;
 using Microsoft.AspNetCore.Http.Features;
-using Microsoft.AspNetCore.Http.Json;
 using Microsoft.AspNetCore.Mvc;
 using Microsoft.AspNetCore.Mvc.Abstractions;
 using Microsoft.AspNetCore.Mvc.ActionConstraints;
@@ -29,8 +27,8 @@
 
 public abstract class OpenApiDocumentServiceTestBase
 {
-    public static async Task VerifyOpenApiDocument(IEndpointRouteBuilder builder, Action<OpenApiDocument> verifyOpenApiDocument, CancellationToken cancellationToken = default)
-        => await VerifyOpenApiDocument(builder, new OpenApiOptions(), verifyOpenApiDocument, cancellationToken);
+    public static async Task VerifyOpenApiDocument(IEndpointRouteBuilder builder, Action<OpenApiDocument> verifyOpenApiDocument)
+        => await VerifyOpenApiDocument(builder, new OpenApiOptions(), verifyOpenApiDocument);
 
     public static async Task VerifyOpenApiDocument(IEndpointRouteBuilder builder, OpenApiOptions openApiOptions, Action<OpenApiDocument> verifyOpenApiDocument, CancellationToken cancellationToken = default)
     {
@@ -40,12 +38,12 @@ public static async Task VerifyOpenApiDocument(IEndpointRouteBuilder builder, Op
         verifyOpenApiDocument(document);
     }
 
-    public static async Task VerifyOpenApiDocument(ActionDescriptor action, Action<OpenApiDocument> verifyOpenApiDocument)
+    public static async Task VerifyOpenApiDocument(ActionDescriptor action, Action<OpenApiDocument> verifyOpenApiDocument, CancellationToken cancellationToken = default)
     {
         var builder = CreateBuilder();
         var documentService = CreateDocumentService(builder, action);
         var scopedService = ((TestServiceProvider)builder.ServiceProvider).CreateScope();
-        var document = await documentService.GetOpenApiDocumentAsync(scopedService.ServiceProvider);
+        var document = await documentService.GetOpenApiDocumentAsync(scopedService.ServiceProvider, cancellationToken);
         verifyOpenApiDocument(document);
     }
 

src/OpenApi/test/Microsoft.AspNetCore.OpenApi.Tests/Transformers/DocumentTransformerTests.cs

@@ -240,6 +240,89 @@ public async Task DocumentTransformer_CanAccessTransientServiceFromContextApplic
         Assert.Equal(2, Dependency.InstantiationCount);
     }
 
+    [Fact]
+    public async Task DocumentTransformer_RespectsOperationCancellation()
+    {
+        var builder = CreateBuilder();
+        builder.MapGet("/todo", () => { });
+
+        var options = new OpenApiOptions();
+        var transformerCalled = false;
+        var exceptionThrown = false;
+        var tcs = new TaskCompletionSource();
+
+        options.AddDocumentTransformer(async (document, context, cancellationToken) =>
+        {
+            transformerCalled = true;
+            try
+            {
+                await tcs.Task.WaitAsync(cancellationToken);
+                document.Info.Description = "Should not be set";
+            }
+            catch (OperationCanceledException)
+            {
+                exceptionThrown = true;
+                throw;
+            }
+        });
+
+        using var cts = new CancellationTokenSource();
+        cts.CancelAfter(1);
+
+        await Assert.ThrowsAsync<TaskCanceledException>(async () =>
+        {
+            await VerifyOpenApiDocument(builder, options, _ => { }, cts.Token);
+        });
+
+        Assert.True(transformerCalled);
+        Assert.True(exceptionThrown);
+    }
+
+    [Fact]
+    public async Task DocumentTransformer_ExecutesAsynchronously()
+    {
+        var builder = CreateBuilder();
+        builder.MapGet("/todo", () => { });
+
+        var options = new OpenApiOptions();
+        var transformerOrder = new List<int>();
+        var tcs1 = new TaskCompletionSource();
+        var tcs2 = new TaskCompletionSource();
+
+        options.AddDocumentTransformer(async (document, context, cancellationToken) =>
+        {
+            await tcs1.Task;
+            transformerOrder.Add(1);
+            document.Info.Title = "First";
+        });
+
+        options.AddDocumentTransformer((document, context, cancellationToken) =>
+        {
+            transformerOrder.Add(2);
+            document.Info.Title += " Second";
+            tcs2.TrySetResult();
+            return Task.CompletedTask;
+        });
+
+        options.AddDocumentTransformer(async (document, context, cancellationToken) =>
+        {
+            await tcs2.Task;
+            transformerOrder.Add(3);
+            document.Info.Title += " Third";
+        });
+
+        var documentTask = VerifyOpenApiDocument(builder, options, document =>
+        {
+            Assert.Equal("First Second Third", document.Info.Title);
+        });
+
+        tcs1.TrySetResult();
+
+        await documentTask;
+
+        Assert.Equal([1, 2, 3], transformerOrder);
+    }
+
     private class ActivatedTransformer : IOpenApiDocumentTransformer
     {
         public Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransformerContext context, CancellationToken cancellationToken)

src/OpenApi/test/Microsoft.AspNetCore.OpenApi.Tests/Transformers/OperationTransformerTests.cs

@@ -478,6 +478,185 @@ public async Task OperationTransformer_CanAccessTransientServiceFromContextAppli
         Assert.Equal(4, Dependency.InstantiationCount);
     }
 
+    [Fact]
+    public async Task AddOpenApiOperationTransformer_CanApplyTransformer()
+    {
+        var builder = CreateBuilder();
+
+        builder.MapGet("/", () => { })
+            .AddOpenApiOperationTransformer((operation, context, cancellationToken) =>
+            {
+                operation.Description = "Operation Description";
+                return Task.CompletedTask;
+            });
+
+        await VerifyOpenApiDocument(builder, document =>
+        {
+            Assert.Collection(document.Paths.OrderBy(p => p.Key),
+                path =>
+                {
+                    Assert.Equal("/", path.Key);
+                    var operation = Assert.Single(path.Value.Operations.Values);
+                    Assert.Equal("Operation Description", operation.Description);
+                });
+        });
+    }
+
+    [Fact]
+    public async Task AddOpenApiOperationTransformer_TransformerRunsAfterOtherTransformers()
+    {
+        var builder = CreateBuilder();
+
+        builder.MapGet("/", () => { })
+            .AddOpenApiOperationTransformer((operation, context, cancellationToken) =>
+            {
+                operation.Description = "Operation Description";
+                return Task.CompletedTask;
+            });
+
+        var options = new OpenApiOptions();
+        options.AddOperationTransformer((operation, context, cancellationToken) =>
+        {
+            operation.Description = "Operation Description 2";
+            return Task.CompletedTask;
+        });
+
+        await VerifyOpenApiDocument(builder, document =>
+        {
+            Assert.Collection(document.Paths.OrderBy(p => p.Key),
+                path =>
+                {
+                    Assert.Equal("/", path.Key);
+                    var operation = Assert.Single(path.Value.Operations.Values);
+                    Assert.Equal("Operation Description", operation.Description);
+                });
+        });
+    }
+
+    [Fact]
+    public async Task AddOpenApiOperationTransformer_SupportsMultipleTransformers()
+    {
+        var builder = CreateBuilder();
+
+        builder.MapGet("/", () => { })
+            .AddOpenApiOperationTransformer((operation, context, cancellationToken) =>
+            {
+                operation.Description = "Operation Description";
+                return Task.CompletedTask;
+            })
+            .AddOpenApiOperationTransformer((operation, context, cancellationToken) =>
+            {
+                operation.Description += " 2";
+                operation.Deprecated = true;
+                return Task.CompletedTask;
+            })
+            .AddOpenApiOperationTransformer((operation, context, cancellationToken) =>
+            {
+                operation.Description += " 3";
+                operation.OperationId = "OperationId";
+                return Task.CompletedTask;
+            });
+
+        await VerifyOpenApiDocument(builder, document =>
+        {
+            Assert.Collection(document.Paths.OrderBy(p => p.Key),
+                path =>
+                {
+                    Assert.Equal("/", path.Key);
+                    var operation = Assert.Single(path.Value.Operations.Values);
+                    Assert.Equal("Operation Description 2 3", operation.Description);
+                    Assert.True(operation.Deprecated);
+                    Assert.Equal("OperationId", operation.OperationId);
+                });
+        });
+    }
+
+    [Fact]
+    public async Task OperationTransformer_RespectsOperationCancellation()
+    {
+        var builder = CreateBuilder();
+        builder.MapGet("/todo", () => { });
+
+        var options = new OpenApiOptions();
+        var transformerCalled = false;
+        var exceptionThrown = false;
+        var tcs = new TaskCompletionSource();
+
+        options.AddOperationTransformer(async (operation, context, cancellationToken) =>
+        {
+            transformerCalled = true;
+            try
+            {
+                await tcs.Task.WaitAsync(cancellationToken);
+                operation.Description = "Should not be set";
+            }
+            catch (OperationCanceledException)
+            {
+                exceptionThrown = true;
+                throw;
+            }
+        });
+
+        using var cts = new CancellationTokenSource();
+        cts.CancelAfter(1);
+
+        await Assert.ThrowsAsync<TaskCanceledException>(async () =>
+        {
+            await VerifyOpenApiDocument(builder, options, _ => { }, cts.Token);
+        });
+
+        Assert.True(transformerCalled);
+        Assert.True(exceptionThrown);
+    }
+
+    [Fact]
+    public async Task OperationTransformer_ExecutesAsynchronously()
+    {
+        var builder = CreateBuilder();
+        builder.MapGet("/todo", () => { });
+
+        var options = new OpenApiOptions();
+        var transformerOrder = new List<int>();
+        var tcs1 = new TaskCompletionSource();
+        var tcs2 = new TaskCompletionSource();
+
+        options.AddOperationTransformer(async (operation, context, cancellationToken) =>
+        {
+            await tcs1.Task;
+            transformerOrder.Add(1);
+            operation.Description = "First";
+        });
+
+        options.AddOperationTransformer((operation, context, cancellationToken) =>
+        {
+            transformerOrder.Add(2);
+            operation.Description += " Second";
+            tcs2.TrySetResult();
+            return Task.CompletedTask;
+        });
+
+        options.AddOperationTransformer(async (operation, context, cancellationToken) =>
+        {
+            await tcs2.Task;
+            transformerOrder.Add(3);
+            operation.Description += " Third";
+        });
+
+        var documentTask = VerifyOpenApiDocument(builder, options, document =>
+        {
+            var operation = Assert.Single(document.Paths["/todo"].Operations.Values);
+            Assert.Equal("First Second Third", operation.Description);
+        });
+
+        tcs1.TrySetResult();
+
+        await documentTask;
+
+        // Verify transformers executed in the correct order, once for each transformer
+        // since there is a single operation in the document.
+        Assert.Equal([1, 2, 3], transformerOrder);
+    }
+
     private class ActivatedTransformer : IOpenApiOperationTransformer
     {
         public Task TransformAsync(OpenApiOperation operation, OpenApiOperationTransformerContext context, CancellationToken cancellationToken)