Fix Mocha saga fault handling so OnReplyFault fires

src/Mocha/src/Mocha/Consumers/Extensions/ConsumeContextExtensions.cs

@@ -7,6 +7,14 @@ namespace Mocha;
 /// </summary>
 internal static class ConsumeContextExtensions
 {
+    /// <summary>
+    /// Creates reply options from the incoming message metadata when a fault channel is available.
+    /// </summary>
+    public static bool TryCreateFaultOptions(this IMessageContext context, out ReplyOptions options)
+    {
+        return TryCreateReplyOptions(context, context.FaultAddress, out options);
+    }
+
     /// <summary>
     /// Creates reply options from the incoming message metadata when a response channel is available.
     /// </summary>
@@ -15,10 +23,17 @@ internal static class ConsumeContextExtensions
     /// workflows (for example saga headers) keep working. The correlation id is echoed when present,
     /// so callers that correlate by a different mechanism (such as a saga header) are still supported.
     /// </remarks>
-    public static bool TryCreateResponseOptions(this IConsumeContext context, out ReplyOptions options)
+    public static bool TryCreateResponseOptions(this IMessageContext context, out ReplyOptions options)
+    {
+        return TryCreateReplyOptions(context, context.ResponseAddress, out options);
+    }
+
+    private static bool TryCreateReplyOptions(
+        IMessageContext context,
+        Uri? replyTo,
+        out ReplyOptions options)
     {
         options = ReplyOptions.Default;
-        var replyTo = context.ResponseAddress;
         if (replyTo is null)
         {
             return false;

src/Mocha/src/Mocha/Consumers/Implementations/ReplyConsumer.cs

@@ -1,3 +1,5 @@
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
 using Mocha.Events;
 
 namespace Mocha;
@@ -14,6 +16,8 @@ namespace Mocha;
 // TODO Not sure if this really has to be consumer. could also just be a middleware
 public sealed class ReplyConsumer(DeferredResponseManager responseManager) : Consumer
 {
+    private ILogger<ReplyConsumer> _logger = default!;
+
     protected override void Configure(IConsumerDescriptor descriptor)
     {
         descriptor.Name("Reply");
@@ -22,29 +26,29 @@ protected override void Configure(IConsumerDescriptor descriptor)
     protected override void OnAfterInitialize(IMessagingSetupContext context)
     {
         base.OnAfterInitialize(context);
+
+        _logger = context.Services.GetRequiredService<ILogger<ReplyConsumer>>();
     }
 
     protected override ValueTask ConsumeAsync(IConsumeContext context)
     {
         if (context.CorrelationId is not { } correlationId)
         {
-            // TODO logs!
-            // Replies without correlation cannot be matched to a pending request promise.
+            // Dispatch always stamps a correlation id, so a reply without one came from elsewhere.
+            // It cannot match a promise, and only matters when no other consumer claimed it.
+            ReportUnmatchedReply(context, null);
             return default;
         }
 
         try
         {
-            var message = context.GetMessage();
-
-            if (message is null)
+            if (context.GetMessage() is not { } message)
             {
                 throw ThrowHelper.ResponseBodyNotSet();
             }
 
             if (message is NotAcknowledgedEvent failure)
             {
-                // Fault replies complete the pending request with a remote exception.
                 responseManager.SetException(
                     correlationId,
                     new RemoteErrorException(
@@ -53,18 +57,53 @@ protected override ValueTask ConsumeAsync(IConsumeContext context)
                         failure.MessageId,
                         failure.CorrelationId));
             }
-            else if (!responseManager.CompletePromise(context.CorrelationId, message))
+            else if (!responseManager.CompletePromise(correlationId, message))
             {
-                // A late/unknown reply indicates there is no active waiter for this correlation id.
-                throw ThrowHelper.PromiseNotFound();
+                ReportUnmatchedReply(context, correlationId);
             }
         }
         catch (Exception ex)
         {
-            // TODO logs!
+            // Fault the waiting requester rather than leave it to time out.
             responseManager.SetException(correlationId, ex);
+            _logger.ReplyProcessingFailed(ex, correlationId, context.MessageId);
         }
 
         return default;
     }
+
+    /// <summary>
+    /// Reports a reply that completed no promise, distinguishing one that another consumer on the
+    /// same message owns from one that nothing handled.
+    /// </summary>
+    private void ReportUnmatchedReply(IConsumeContext context, string? correlationId)
+    {
+        // A saga reply route selects the saga consumer alongside this one, so a second consumer on
+        // the message means the reply is owned there and no promise was expected.
+        if (context.Features.Get<ReceiveConsumerFeature>()?.Consumers is { Count: > 1 })
+        {
+            return;
+        }
+
+        _logger.ReplyDiscarded(correlationId, context.MessageId);
+    }
+}
+
+internal static partial class ReplyConsumerLogs
+{
+    [LoggerMessage(
+        LogLevel.Warning,
+        "Discarded a reply that no pending request and no consumer claimed "
+            + "(correlation id {CorrelationId}, message id {MessageId})")]
+    public static partial void ReplyDiscarded(this ILogger logger, string? correlationId, string? messageId);
+
+    [LoggerMessage(
+        LogLevel.Error,
+        "Failed to process a reply "
+            + "(correlation id {CorrelationId}, message id {MessageId})")]
+    public static partial void ReplyProcessingFailed(
+        this ILogger logger,
+        Exception exception,
+        string correlationId,
+        string? messageId);
 }

src/Mocha/src/Mocha/Execution/PublishOptions.cs

@@ -1,7 +1,8 @@
 namespace Mocha;
 
 /// <summary>
-/// Options controlling the behavior of a publish operation, such as scheduling, expiration, and custom headers.
+/// Options controlling the behavior of a publish operation, including scheduling, expiration, fault routing,
+/// and custom headers.
 /// </summary>
 public readonly struct PublishOptions
 {
@@ -20,6 +21,11 @@ public readonly struct PublishOptions
     /// </summary>
     public DateTimeOffset? ExpirationTime { get; init; }
 
+    /// <summary>
+    /// Gets the endpoint address where fault notifications should be sent, or <c>null</c> to use the default.
+    /// </summary>
+    public Uri? FaultEndpoint { get; init; }
+
     /// <summary>
     /// Gets custom headers to include with the published message, or <c>null</c> if none.
     /// </summary>

src/Mocha/src/Mocha/Execution/ReplyOptions.cs

@@ -25,6 +25,11 @@ public readonly struct ReplyOptions
     /// </summary>
     public Dictionary<string, object?>? Headers { get; init; }
 
+    /// <summary>
+    /// Gets the message kind stamped on the reply, or <c>null</c> to use <see cref="Mocha.MessageKind.Reply"/>.
+    /// </summary>
+    public string? MessageKind { get; init; }
+
     /// <summary>
     /// Gets the default reply options with no overrides.
     /// </summary>

src/Mocha/src/Mocha/MessageTypes/Conditions/NotFaultCondition.cs

@@ -0,0 +1,21 @@
+using Mocha.Middlewares;
+
+namespace Mocha;
+
+/// <summary>
+/// Matches messages that are not fault notifications.
+/// </summary>
+internal sealed class NotFaultCondition : RouteCondition
+{
+    public static NotFaultCondition Instance { get; } = new();
+
+    private NotFaultCondition() { }
+
+    /// <inheritdoc />
+    public override bool Matches(IReceiveContext context)
+        => context.Headers.GetMessageKind() != MessageKind.Fault;
+
+    /// <inheritdoc />
+    public override RouteConditionDescription Describe()
+        => new("NotFault", null, []);
+}

src/Mocha/src/Mocha/Middlewares/DefaultMessageBus.cs

@@ -46,12 +46,12 @@ public async ValueTask PublishAsync<T>(T message, CancellationToken cancellation
     /// </summary>
     /// <remarks>
     /// The message is routed through the publish endpoint resolved by the runtime's router for the
-    /// given message type. Custom headers and expiration time from <paramref name="options"/> are
-    /// applied to the dispatch context before pipeline execution.
+    /// given message type. Fault routing, custom headers, and expiration from <paramref name="options"/>
+    /// are applied to the dispatch context before pipeline execution.
     /// </remarks>
     /// <typeparam name="T">The type of the message to publish.</typeparam>
     /// <param name="message">The message instance to publish. Must not be <see langword="null"/>.</param>
-    /// <param name="options">Options controlling headers and expiration for this publish operation.</param>
+    /// <param name="options">Options controlling fault routing, headers, and expiration for this publish operation.</param>
     /// <param name="cancellationToken">A token to cancel the publish operation.</param>
     public async ValueTask PublishAsync<T>(T message, PublishOptions options, CancellationToken cancellationToken)
         where T : notnull
@@ -67,6 +67,7 @@ public async ValueTask PublishAsync<T>(T message, PublishOptions options, Cancel
             context.Message = message;
             context.AddHeaders(options.Headers);
             context.Headers.SetMessageKind(MessageKind.Publish);
+            context.FaultAddress = options.FaultEndpoint;
             context.ScheduledTime = options.ScheduledTime;
             context.DeliverBy = options.ExpirationTime;
 
@@ -239,14 +240,13 @@ public async ValueTask ReplyAsync<TResponse>(
             context.Message = response;
 
             context.AddHeaders(headers);
-            context.Headers.SetMessageKind(MessageKind.Reply);
+            context.Headers.SetMessageKind(options.MessageKind ?? MessageKind.Reply);
 
             await replyEndpoint.ExecuteAsync(context);
         }
-        catch
+        finally
         {
             _contextPool.Return(context);
-            throw;
         }
     }
 
@@ -333,6 +333,7 @@ public async ValueTask<SchedulingResult> SchedulePublishAsync<T>(
             context.Message = message;
             context.AddHeaders(options.Headers);
             context.Headers.SetMessageKind(MessageKind.Publish);
+            context.FaultAddress = options.FaultEndpoint;
             context.ScheduledTime = scheduledTime;
             context.DeliverBy = options.ExpirationTime;
 

src/Mocha/src/Mocha/Middlewares/Receive/ReceiveFaultMiddleware.cs

@@ -10,13 +10,8 @@ namespace Mocha.Middlewares;
 /// the original message.
 /// </summary>
 /// <remarks>
-/// The middleware follows a two-path failure contract:
-/// request/response flows receive a direct negative acknowledgement on the response address, while
-/// non-request flows are forwarded to the error endpoint with fault metadata in headers.
-/// This keeps failure observable for both callers and operations, similar to fault-event + error
-/// queue patterns used in broker-centric systems.
-/// Without this middleware, callers often only see timeouts, and operators lose structured error
-/// context tied to the original envelope.
+/// Faults are sent to the fault address when present, otherwise to the response address. Messages
+/// with neither address are forwarded to the error endpoint with fault metadata in headers.
 /// </remarks>
 internal sealed class ReceiveFaultMiddleware(
     TimeProvider provider,
@@ -33,72 +28,41 @@ public async ValueTask InvokeAsync(IReceiveContext context, ReceiveDelegate next
         }
         catch (Exception ex)
         {
-            var envelope = context.Envelope;
-
             var fault = FaultInfo.From(Guid.NewGuid(), provider.GetUtcNow(), ex);
 
-            // A requester expecting a reply should get an explicit negative acknowledgement first.
-            if (envelope?.ResponseAddress is { } responseAddress
-                && Uri.TryCreate(responseAddress, UriKind.Absolute, out var responseAddressUri))
+            if (context.TryCreateFaultOptions(out var options)
+                || context.TryCreateResponseOptions(out options))
             {
-                await ReplyToSenderAsync(context, responseAddressUri, envelope, fault);
+                await ReplyToSenderAsync(context, options, fault);
             }
             else
             {
-                await SendToErrorEndpointAsync(context, envelope, fault);
+                await SendToErrorEndpointAsync(context, context.Envelope, fault);
             }
 
             feature.MessageConsumed = true;
         }
     }
 
-    private async ValueTask ReplyToSenderAsync(
+    private static async ValueTask ReplyToSenderAsync(
         IReceiveContext context,
-        Uri responseAddress,
-        MessageEnvelope envelope,
+        ReplyOptions options,
         FaultInfo fault)
     {
-        var replyEndpoint = context.Runtime.GetTransport(responseAddress)?.ReplyDispatchEndpoint;
-        if (replyEndpoint is null)
-        {
-            // TODO critical error! (Poision Pill)
-            throw ThrowHelper.NoReplyEndpointFound(responseAddress.ToString());
-        }
+        var exceptionType = fault.Exceptions.FirstOrDefault()?.ExceptionType;
 
-        var messageType = context.Runtime.GetMessageType(typeof(NotAcknowledgedEvent));
+        var notAcknowledged = new NotAcknowledgedEvent(
+            context.CorrelationId,
+            context.MessageId,
+            fault.ErrorCode,
+            $"The message faulted with an exception of type {exceptionType}");
 
-        var dispatchContext = pools.DispatchContext.Get();
-        try
-        {
-            dispatchContext.CorrelationId = envelope?.CorrelationId;
-            dispatchContext.ConversationId = envelope?.ConversationId;
-            dispatchContext.DestinationAddress = responseAddress;
-            dispatchContext.SourceAddress = replyEndpoint.Address;
+        var bus = context.Services.GetRequiredService<IMessageBus>();
 
-            dispatchContext.Initialize(
-                context.Services,
-                replyEndpoint,
-                context.Runtime,
-                messageType,
-                context.CancellationToken);
-
-            var exceptionType = fault.Exceptions.FirstOrDefault()?.ExceptionType;
-            var message = $"The message faulted with an exception of type {exceptionType}";
-
-            dispatchContext.Headers.SetMessageKind(MessageKind.Fault);
-
-            dispatchContext.Message = new NotAcknowledgedEvent(
-                envelope!.CorrelationId,
-                envelope.MessageId,
-                fault.ErrorCode,
-                message);
-
-            await replyEndpoint.ExecuteAsync(dispatchContext);
-        }
-        finally
-        {
-            pools.DispatchContext.Return(dispatchContext);
-        }
+        await bus.ReplyAsync(
+            notAcknowledged,
+            options with { MessageKind = MessageKind.Fault },
+            context.CancellationToken);
     }
 
     private async ValueTask SendToErrorEndpointAsync(

src/Mocha/src/Mocha/Sagas/Descriptors/SagaStateDescriptorExtensions.cs

@@ -9,7 +9,8 @@ namespace Mocha.Sagas;
 public static class SagaStateDescriptorExtensions
 {
     /// <summary>
-    /// Registers a transition triggered by a fault (not-acknowledged) event.
+    /// Registers a transition triggered by a fault notification for a message the saga sent or published.
+    /// It does not cover failures of the saga's own transitions.
     /// </summary>
     /// <typeparam name="TState">The saga state type.</typeparam>
     /// <param name="descriptor">The state descriptor to configure.</param>
@@ -18,11 +19,12 @@ public static ISagaTransitionDescriptor<TState, NotAcknowledgedEvent> OnFault<TS
         this ISagaStateDescriptor<TState> descriptor)
         where TState : SagaStateBase
     {
-        return descriptor.OnEvent<NotAcknowledgedEvent>();
+        return descriptor.OnReply<NotAcknowledgedEvent>();
     }
 
     /// <summary>
-    /// Registers a transition triggered by any reply message.
+    /// Registers a transition triggered by any successful reply. Fault notifications are handled
+    /// separately through <c>OnFault</c>.
     /// </summary>
     /// <typeparam name="TState">The saga state type.</typeparam>
     /// <param name="descriptor">The state descriptor to configure.</param>