Adding method for registering stream behaviors

Author: jbogard

src/MediatR/MicrosoftExtensionsDI/MediatrServiceConfiguration.cs

@@ -47,6 +47,11 @@ public class MediatRServiceConfiguration
     /// </summary>
     public List<ServiceDescriptor> BehaviorsToRegister { get; } = new();
 
+    /// <summary>
+    /// List of stream behaviors to register in specific order
+    /// </summary>
+    public List<ServiceDescriptor> StreamBehaviorsToRegister { get; } = new();
+
     /// <summary>
     /// Register various handlers from assembly containing given type
     /// </summary>
@@ -140,4 +145,28 @@ public MediatRServiceConfiguration AddOpenBehavior(Type openBehaviorType, Servic
 
         return this;
     }
+
+    /// <summary>
+    /// Register a closed stream behavior type
+    /// </summary>
+    /// <typeparam name="TServiceType">Closed stream behavior interface type</typeparam>
+    /// <typeparam name="TImplementationType">Closed stream behavior implementation type</typeparam>
+    /// <param name="serviceLifetime">Optional service lifetime, defaults to <see cref="ServiceLifetime.Transient"/>.</param>
+    /// <returns>This</returns>
+    public MediatRServiceConfiguration AddStreamBehavior<TServiceType, TImplementationType>(ServiceLifetime serviceLifetime = ServiceLifetime.Transient)
+        => AddStreamBehavior(typeof(TServiceType), typeof(TImplementationType), serviceLifetime);
+    
+    /// <summary>
+    /// Register a closed stream behavior type
+    /// </summary>
+    /// <param name="serviceType">Closed stream behavior interface type</param>
+    /// <param name="implementationType">Closed stream behavior implementation type</param>
+    /// <param name="serviceLifetime">Optional service lifetime, defaults to <see cref="ServiceLifetime.Transient"/>.</param>
+    /// <returns>This</returns>
+    public MediatRServiceConfiguration AddStreamBehavior(Type serviceType, Type implementationType, ServiceLifetime serviceLifetime = ServiceLifetime.Transient)
+    {
+        StreamBehaviorsToRegister.Add(new ServiceDescriptor(serviceType, implementationType, serviceLifetime));
+
+        return this;
+    }
 }

src/MediatR/Registration/ServiceRegistrar.cs

@@ -225,6 +225,11 @@ public static void AddRequiredServices(IServiceCollection services, MediatRServi
         services.TryAdd(notificationPublisherServiceDescriptor);
 
         foreach (var serviceDescriptor in serviceConfiguration.BehaviorsToRegister)
+        {
+            services.TryAddEnumerable(serviceDescriptor);
+        } 
+        
+        foreach (var serviceDescriptor in serviceConfiguration.StreamBehaviorsToRegister)
         {
             services.TryAddEnumerable(serviceDescriptor);
         }

test/MediatR.Tests/MicrosoftExtensionsDI/StreamPipelineTests.cs

@@ -84,4 +84,38 @@ public async Task Should_wrap_with_behavior()
             "Outer after"
         });
     }
+   
+    [Fact]
+    public async Task Should_register_and_wrap_with_behavior()
+    {
+        var output = new Logger();
+        IServiceCollection services = new ServiceCollection();
+        services.AddSingleton(output);
+        services.AddMediatR(cfg =>
+        {
+            cfg.RegisterServicesFromAssembly(typeof(Ping).Assembly);
+            cfg.AddStreamBehavior<IStreamPipelineBehavior<StreamPing, Pong>, OuterBehavior>();
+            cfg.AddStreamBehavior<IStreamPipelineBehavior<StreamPing, Pong>, InnerBehavior>();
+        });
+        var provider = services.BuildServiceProvider();
+
+        var mediator = provider.GetRequiredService<IMediator>();
+
+        var stream = mediator.CreateStream(new StreamPing { Message = "Ping" });
+
+        await foreach (var response in stream)
+        {
+            response.Message.ShouldBe("Ping Pang");
+        }
+
+        output.Messages.ShouldBe(new[]
+        {
+            "Outer before",
+            "Inner before",
+            "Handler",
+            "Inner after",
+            "Outer after"
+        });
+    }
+
 }