feat(core)!: Improve DI robustness by removing unsafe casts

src/DispatchR/Configuration/ServiceRegistrator.cs

@@ -1,7 +1,6 @@
-using Microsoft.Extensions.DependencyInjection;
-using System.Runtime.CompilerServices;
-using DispatchR.Abstractions.Notification;
+using DispatchR.Abstractions.Notification;
 using DispatchR.Abstractions.Send;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace DispatchR.Configuration
 {
@@ -143,10 +142,17 @@ public static void RegisterHandlers(IServiceCollection services, List<Type> allT
                         }
                     }
 
-                    services.AddScoped(handlerInterface, sp =>
+                services.AddScoped(handlerInterface, sp =>
+                {
+                    var keyedServices = sp.GetKeyedServices<IRequestHandler>(key);
+
+                    var pipelinesWithHandler = keyedServices as IRequestHandler[] ?? keyedServices.ToArray();
+
+                    // Single handler - no pipeline chaining needed
+                    if (pipelinesWithHandler.Length == 1)
                     {
-                        var pipelinesWithHandler = Unsafe
-                            .As<IRequestHandler[]>(sp.GetKeyedServices<IRequestHandler>(key));
+                        return pipelinesWithHandler[0];
+                    }
 
                         IRequestHandler lastPipeline = pipelinesWithHandler[0];
                         for (int i = 1; i < pipelinesWithHandler.Length; i++)

src/DispatchR/IMediator.cs

@@ -1,9 +1,9 @@
-using System.Runtime.CompilerServices;
-using DispatchR.Abstractions.Notification;
+using DispatchR.Abstractions.Notification;
 using DispatchR.Abstractions.Send;
 using DispatchR.Abstractions.Stream;
 using DispatchR.Exceptions;
 using Microsoft.Extensions.DependencyInjection;
+using System.Runtime.CompilerServices;
 
 namespace DispatchR;
 
@@ -17,7 +17,7 @@ IAsyncEnumerable<TResponse> CreateStream<TRequest, TResponse>(IStreamRequest<TRe
 
     ValueTask Publish<TNotification>(TNotification request, CancellationToken cancellationToken)
         where TNotification : INotification;
-    
+
     /// <summary>
     /// This method is not recommended for performance-critical scenarios.  
     /// Use it only if it is strictly necessary, as its performance is lower compared  
@@ -28,8 +28,8 @@ ValueTask Publish<TNotification>(TNotification request, CancellationToken cancel
     /// </param>
     /// <param name="cancellationToken"></param>
     /// <returns></returns>
-    [Obsolete(message: "This method has performance issues. Use only if strictly necessary", 
-        error: false, 
+    [Obsolete(message: "This method has performance issues. Use only if strictly necessary",
+        error: false,
         DiagnosticId = Constants.DiagnosticPerformanceIssue)]
     ValueTask Publish(object request, CancellationToken cancellationToken);
 }
@@ -60,13 +60,29 @@ public IAsyncEnumerable<TResponse> CreateStream<TRequest, TResponse>(IStreamRequ
     public async ValueTask Publish<TNotification>(TNotification request, CancellationToken cancellationToken)
         where TNotification : INotification
     {
-        var notificationsInDi = serviceProvider.GetRequiredService<IEnumerable<INotificationHandler<TNotification>>>();
+        var handlers = serviceProvider.GetRequiredService<IEnumerable<INotificationHandler<TNotification>>>();
+
+        if (handlers is INotificationHandler<TNotification>[] handlerArray)
+        {
+            foreach (var handler in handlerArray)
+            {
+                await ProcessHandlerAsync(handler);
+            }
+        }
+        else
+        {
+            foreach (var handler in handlers)
+            {
+                await ProcessHandlerAsync(handler);
+            }
+        }
+
+        return;
 
-        var notifications = Unsafe.As<INotificationHandler<TNotification>[]>(notificationsInDi);
-        foreach (var notification in notifications)
+        async ValueTask ProcessHandlerAsync(INotificationHandler<TNotification> handler)
         {
-            var valueTask = notification.Handle(request, cancellationToken);
-            if (valueTask.IsCompletedSuccessfully is false)
+            var valueTask = handler.Handle(request, cancellationToken);
+            if (!valueTask.IsCompletedSuccessfully)
             {
                 await valueTask;
             }

tests/DispatchR.IntegrationTest/NotificationTests.cs

@@ -74,7 +74,7 @@
 
         // Act
         object notificationObject = new MultiHandlersNotification(Guid.Empty);
         await mediator.Publish(notificationObject, CancellationToken.None);
 
         // Assert
         spyPipelineOneMock.Verify(p => p.Handle(It.IsAny<MultiHandlersNotification>(), It.IsAny<CancellationToken>()), Times.Exactly(1));
@@ -103,4 +103,144 @@
         Assert.Contains(handlers1, h => h is MultiNotificationHandler);
         Assert.Contains(handlers2, h => h is MultiNotificationHandler);
     }
+
+    [Fact]
+    public async Task Publish_CallsSingleHandler_WhenOnlyOneHandlerIsRegistered()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+        services.AddDispatchR(cfg =>
+        {
+            cfg.Assemblies.Add(typeof(Fixture).Assembly);
+            cfg.RegisterPipelines = false;
+            cfg.RegisterNotifications = false;
+        });
+
+        var spyHandlerMock = new Mock<INotificationHandler<MultiHandlersNotification>>();
+        spyHandlerMock.Setup(p => p.Handle(It.IsAny<MultiHandlersNotification>(), It.IsAny<CancellationToken>()))
+            .Returns(ValueTask.CompletedTask);
+
+        services.AddScoped<INotificationHandler<MultiHandlersNotification>>(sp => spyHandlerMock.Object);
+
+        var serviceProvider = services.BuildServiceProvider();
+        var mediator = serviceProvider.GetRequiredService<IMediator>();
+
+        // Act
+        await mediator.Publish(new MultiHandlersNotification(Guid.Empty), CancellationToken.None);
+
+        // Assert
+        spyHandlerMock.Verify(p => p.Handle(It.IsAny<MultiHandlersNotification>(), It.IsAny<CancellationToken>()), Times.Exactly(1));
+    }
+
+    [Fact]
+    public async Task Publish_CallsAsyncHandlers_WhenHandlersRequireAwaiting()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+        services.AddDispatchR(cfg =>
+        {
+            cfg.Assemblies.Add(typeof(Fixture).Assembly);
+            cfg.RegisterPipelines = false;
+            cfg.RegisterNotifications = true;
+            cfg.IncludeHandlers = [typeof(NotificationOneHandler)];
+        });
+
+        var serviceProvider = services.BuildServiceProvider();
+        var mediator = serviceProvider.GetRequiredService<IMediator>();
+
+        // Act
+        await mediator.Publish(new MultiHandlersNotification(Guid.Empty), CancellationToken.None);
+
+        // Assert - if this completes without exception, the async handler was properly awaited
+        Assert.True(true);
+    }
+
+    [Fact]
+    public async Task Publish_CallsSyncHandlers_WhenHandlersAreAlreadyCompleted()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+        services.AddDispatchR(cfg =>
+        {
+            cfg.Assemblies.Add(typeof(Fixture).Assembly);
+            cfg.RegisterPipelines = false;
+            cfg.RegisterNotifications = true;
+            cfg.IncludeHandlers = [typeof(NotificationTwoHandler), typeof(NotificationThreeHandler)];
+        });
+
+        var serviceProvider = services.BuildServiceProvider();
+        var mediator = serviceProvider.GetRequiredService<IMediator>();
+
+        // Act
+        await mediator.Publish(new MultiHandlersNotification(Guid.Empty), CancellationToken.None);
+
+        // Assert - if this completes without exception, the sync handlers were properly handled
+        Assert.True(true);
+    }
+
+    [Fact]
+    public async Task Publish_HandlesNonArrayEnumerable_WhenHandlersAreNotArray()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+        services.AddDispatchR(cfg =>
+        {
+            cfg.Assemblies.Add(typeof(Fixture).Assembly);
+            cfg.RegisterPipelines = false;
+            cfg.RegisterNotifications = false;
+        });
+
+        var handler1Mock = new Mock<INotificationHandler<MultiHandlersNotification>>();
+        var handler2Mock = new Mock<INotificationHandler<MultiHandlersNotification>>();
+
+        handler1Mock.Setup(p => p.Handle(It.IsAny<MultiHandlersNotification>(), It.IsAny<CancellationToken>()))
+            .Returns(ValueTask.CompletedTask);
+        handler2Mock.Setup(p => p.Handle(It.IsAny<MultiHandlersNotification>(), It.IsAny<CancellationToken>()))
+            .Returns(ValueTask.CompletedTask);
+
+        // Register a custom service that returns a non-array IEnumerable
+        services.AddScoped<IEnumerable<INotificationHandler<MultiHandlersNotification>>>(sp =>
+        {
+            var list = new List<INotificationHandler<MultiHandlersNotification>>
+            {
+                handler1Mock.Object,
+                handler2Mock.Object
+            };
+            // Return as IEnumerable (not array) by using LINQ
+            return list.Where(h => h != null);
+        });
+
+        var serviceProvider = services.BuildServiceProvider();
+        var mediator = serviceProvider.GetRequiredService<IMediator>();
+
+        // Act
+        await mediator.Publish(new MultiHandlersNotification(Guid.Empty), CancellationToken.None);
+
+        // Assert
+        handler1Mock.Verify(p => p.Handle(It.IsAny<MultiHandlersNotification>(), It.IsAny<CancellationToken>()), Times.Exactly(1));
+        handler2Mock.Verify(p => p.Handle(It.IsAny<MultiHandlersNotification>(), It.IsAny<CancellationToken>()), Times.Exactly(1));
+    }
+
+    [Fact]
+    public async Task Publish_HandlesMixedAsyncAndSyncHandlers_WhenMultipleHandlersAreRegistered()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+        services.AddDispatchR(cfg =>
+        {
+            cfg.Assemblies.Add(typeof(Fixture).Assembly);
+            cfg.RegisterPipelines = false;
+            cfg.RegisterNotifications = true;
+            cfg.IncludeHandlers = [typeof(NotificationOneHandler), typeof(NotificationTwoHandler), typeof(NotificationThreeHandler)];
+        });
+
+        var serviceProvider = services.BuildServiceProvider();
+        var mediator = serviceProvider.GetRequiredService<IMediator>();
+
+        // Act
+        await mediator.Publish(new MultiHandlersNotification(Guid.Empty), CancellationToken.None);
+
+        // Assert - if this completes without exception, all handlers (async and sync) were properly handled
+        Assert.True(true);
+    }
 }