Fixed Batch Consumer Middleware Pipeline

src/Mocha/src/Mocha.Inbox/ConsumeInboxMiddleware.cs

@@ -61,6 +61,12 @@ public async ValueTask InvokeAsync(IConsumeContext context, ConsumerDelegate nex
             return;
         }
 
+        if (context is IBatchConsumeContext batchContext)
+        {
+            await InvokeBatchAsync(batchContext, context, next, activity);
+            return;
+        }
+
         var messageId = context.MessageId;
 
         if (messageId is null)
@@ -133,6 +139,95 @@ public async ValueTask InvokeAsync(IConsumeContext context, ConsumerDelegate nex
         await next(context);
     }
 
+    private async ValueTask InvokeBatchAsync(
+        IBatchConsumeContext batchContext,
+        IConsumeContext context,
+        ConsumerDelegate next,
+        Activity? activity)
+    {
+        var consumerType = GetConsumerType(context);
+        var inbox = context.Services.GetRequiredService<IMessageInbox>();
+        var claimed = 0;
+        var skipped = 0;
+        var untracked = 0;
+
+        var batch = batchContext.Message;
+
+        for (var i = 0; i < batch.Count; i++)
+        {
+            var itemContext = batch.GetContext(i);
+            var messageId = itemContext.MessageId;
+
+            if (messageId is null)
+            {
+                untracked++;
+                continue;
+            }
+
+            if (itemContext.Envelope is not { } envelope)
+            {
+                try
+                {
+                    if (await inbox.ExistsAsync(messageId, consumerType, context.CancellationToken))
+                    {
+                        logger.MessageSkippedDueToInboxExists(messageId, consumerType);
+                        skipped++;
+                    }
+                    else
+                    {
+                        untracked++;
+                    }
+                }
+                catch (OperationCanceledException)
+                {
+                    throw;
+                }
+                catch (Exception ex)
+                {
+                    logger.InboxExistsCheckFailed(messageId, ex);
+                    untracked++;
+                }
+
+                continue;
+            }
+
+            try
+            {
+                if (await inbox.TryClaimAsync(envelope, consumerType, context.CancellationToken))
+                {
+                    claimed++;
+                }
+                else
+                {
+                    logger.MessageSkippedDueToInboxExists(messageId, consumerType);
+                    skipped++;
+                }
+            }
+            catch (OperationCanceledException)
+            {
+                throw;
+            }
+            catch (Exception ex)
+            {
+                logger.InboxClaimFailed(messageId, ex);
+                untracked++;
+            }
+        }
+
+        activity?.SetTag("inbox.batch.claimed", claimed);
+        activity?.SetTag("inbox.batch.skipped", skipped);
+        activity?.SetTag("inbox.batch.untracked", untracked);
+
+        if (claimed == 0 && untracked == 0)
+        {
+            activity?.SetTag("inbox.claimed", false);
+            return;
+        }
+
+        activity?.SetTag("inbox.claimed", true);
+        await next(context);
+    }
+
     /// <summary>
     /// Gets the consumer type name from the current context by reading the
     /// <see cref="ReceiveConsumerFeature.CurrentConsumer"/> identity.

src/Mocha/src/Mocha/Abstractions/Context/IBatchConsumeContext.cs

@@ -0,0 +1,17 @@
+namespace Mocha;
+
+/// <summary>
+/// Represents the consume context for a batch of messages.
+/// </summary>
+public interface IBatchConsumeContext : IConsumeContext<IMessageBatch>
+{
+    /// <summary>
+    /// Gets the logical identifier of this batch operation.
+    /// </summary>
+    string BatchId { get; }
+
+    /// <summary>
+    /// Gets the message type metadata for each item in the batch.
+    /// </summary>
+    MessageType? ItemMessageType { get; }
+}

src/Mocha/src/Mocha/Abstractions/Context/IBatchConsumeContext~1.cs

@@ -0,0 +1,18 @@
+namespace Mocha;
+
+/// <summary>
+/// Represents the consume context for a typed batch of messages.
+/// </summary>
+/// <typeparam name="TMessage">The message type contained in the batch.</typeparam>
+public interface IBatchConsumeContext<out TMessage> :
+    IBatchConsumeContext,
+    IConsumeContext<IMessageBatch<TMessage>>
+{
+    /// <summary>
+    /// Gets the typed message batch.
+    /// </summary>
+    new IMessageBatch<TMessage> Message { get; }
+
+    /// <inheritdoc />
+    IMessageBatch IConsumeContext<IMessageBatch>.Message => Message;
+}

src/Mocha/src/Mocha/Abstractions/IMessageBatch.cs

@@ -1,21 +1,24 @@
 namespace Mocha;
 
 /// <summary>
-/// An immutable batch of messages delivered to an <see cref="IBatchEventHandler{TEvent}"/>.
-/// Provides indexed access, count, enumeration, and batch metadata.
+/// An immutable batch of messages delivered to a batch handler.
 /// </summary>
-/// <typeparam name="TEvent">The type of event in the batch.</typeparam>
-public interface IMessageBatch<out TEvent> : IReadOnlyList<TEvent>
+public interface IMessageBatch
 {
     /// <summary>
     /// Gets the reason this batch was dispatched.
     /// </summary>
     BatchCompletionMode CompletionMode { get; }
 
+    /// <summary>
+    /// Gets the number of messages in the batch.
+    /// </summary>
+    int Count { get; }
+
     /// <summary>
     /// Gets the consume context for a specific message in the batch.
     /// </summary>
     /// <param name="index">The zero-based index of the message.</param>
     /// <returns>The consume context for the message.</returns>
-    IConsumeContext<TEvent> GetContext(int index);
+    IConsumeContext GetContext(int index);
 }

src/Mocha/src/Mocha/Abstractions/IMessageBatch~1.cs

@@ -0,0 +1,24 @@
+namespace Mocha;
+
+/// <summary>
+/// An immutable batch of messages delivered to an <see cref="IBatchEventHandler{TEvent}"/>.
+/// Provides indexed access, count, enumeration, and batch metadata.
+/// </summary>
+/// <typeparam name="TEvent">The type of event in the batch.</typeparam>
+public interface IMessageBatch<out TEvent> : IMessageBatch, IReadOnlyList<TEvent>
+{
+    /// <summary>
+    /// Gets the number of messages in the batch.
+    /// </summary>
+    new int Count { get; }
+
+    /// <summary>
+    /// Gets the typed consume context for a specific message in the batch.
+    /// </summary>
+    /// <param name="index">The zero-based index of the message.</param>
+    /// <returns>The consume context for the message.</returns>
+    new IConsumeContext<TEvent> GetContext(int index);
+
+    /// <inheritdoc />
+    IConsumeContext IMessageBatch.GetContext(int index) => GetContext(index);
+}

src/Mocha/src/Mocha/Consumers/Batching/BatchConsumeContext.cs

@@ -0,0 +1,81 @@
+using Mocha.Features;
+using Mocha.Middlewares;
+
+namespace Mocha;
+
+internal sealed class BatchConsumeContext<TMessage> : IBatchConsumeContext<TMessage>
+{
+    private readonly Headers _headers = new(0);
+    private readonly FeatureCollection _features = new();
+
+    public BatchConsumeContext(
+        MessageBatch<TMessage> message,
+        IServiceProvider services,
+        IConsumeContext firstContext,
+        string batchId,
+        MessageType? itemMessageType,
+        CancellationToken cancellationToken)
+    {
+        Message = message;
+        Services = services;
+        Runtime = firstContext.Runtime;
+        Transport = firstContext.Transport;
+        Endpoint = firstContext.Endpoint;
+        Host = firstContext.Host;
+        BatchId = batchId;
+        ItemMessageType = itemMessageType;
+        CancellationToken = cancellationToken;
+    }
+
+    public IMessageBatch<TMessage> Message { get; }
+
+    public string BatchId { get; }
+
+    public MessageType? ItemMessageType { get; }
+
+    public IFeatureCollection Features => _features;
+
+    public MessagingTransport Transport { get; set; }
+
+    public ReceiveEndpoint Endpoint { get; set; }
+
+    public string? MessageId { get; set; }
+
+    public string? CorrelationId { get; set; }
+
+    public string? ConversationId { get; set; }
+
+    public string? CausationId { get; set; }
+
+    public Uri? SourceAddress { get; set; }
+
+    public Uri? DestinationAddress { get; set; }
+
+    public Uri? ResponseAddress { get; set; }
+
+    public Uri? FaultAddress { get; set; }
+
+    public MessageContentType? ContentType { get; set; }
+
+    public MessageType? MessageType { get; set; }
+
+    public IReadOnlyHeaders Headers => _headers;
+
+    public DateTimeOffset? SentAt { get; set; }
+
+    public DateTimeOffset? DeliverBy { get; set; }
+
+    public int? DeliveryCount { get; set; }
+
+    public ReadOnlyMemory<byte> Body => ReadOnlyMemory<byte>.Empty;
+
+    public MessageEnvelope? Envelope { get; set; }
+
+    public IRemoteHostInfo Host { get; set; }
+
+    public IMessagingRuntime Runtime { get; set; }
+
+    public CancellationToken CancellationToken { get; set; }
+
+    public IServiceProvider Services { get; set; }
+}

src/Mocha/src/Mocha/Consumers/Consumer.cs

@@ -56,7 +56,7 @@ protected virtual void Configure(IConsumerDescriptor descriptor) { }
     /// </summary>
     public Type Identity { get; private set; } = null!;
 
-    private ConsumerDelegate _pipeline = null!;
+    private protected ConsumerDelegate Pipeline { get; private set; } = null!;
 
     /// <summary>
     /// Handles an incoming message after the consume middleware pipeline has completed.
@@ -84,14 +84,14 @@ protected virtual void Configure(IConsumerDescriptor descriptor) { }
     /// <exception cref="InvalidOperationException">
     /// Thrown when the context does not implement <see cref="IConsumeContext"/>.
     /// </exception>
-    public async ValueTask ProcessAsync(IReceiveContext context)
+    public virtual async ValueTask ProcessAsync(IReceiveContext context)
     {
         if (context is not IConsumeContext handlerContext)
         {
             throw ThrowHelper.InvalidHandlerContext();
         }
 
-        await _pipeline(handlerContext);
+        await Pipeline(handlerContext);
     }
 
     /// <summary>
@@ -176,7 +176,7 @@ internal void Complete(IMessagingSetupContext context)
             Consumer = this
         };
 
-        _pipeline = MiddlewareCompiler.CompileHandler(
+        Pipeline = MiddlewareCompiler.CompileHandler(
             middlewareFactoryContext,
             ConsumeAsync,
             [context.GetConsumerMiddlewares(), Configuration!.ConsumerMiddlewares],