Fix Mocha saga fault handling so OnReplyFault fires

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

@@ -15,7 +15,7 @@ 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)
     {
         options = ReplyOptions.Default;
         var replyTo = context.ResponseAddress;

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;
+
     protected override void Configure(IConsumerDescriptor descriptor)
     {
         descriptor.Name("Reply");
@@ -22,49 +26,83 @@ 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, correlationId: null);
             return default;
         }
 
         try
         {
-            var message = context.GetMessage();
-
-            if (message is null)
-            {
-                throw ThrowHelper.ResponseBodyNotSet();
-            }
+            var message = context.GetMessage() ?? throw ThrowHelper.ResponseBodyNotSet();
 
-            if (message is NotAcknowledgedEvent failure)
-            {
-                // Fault replies complete the pending request with a remote exception.
-                responseManager.SetException(
+            var matched = message is NotAcknowledgedEvent failure
+                ? responseManager.SetException(
                     correlationId,
                     new RemoteErrorException(
                         failure.ErrorCode,
                         failure.ErrorMessage,
                         failure.MessageId,
-                        failure.CorrelationId));
-            }
-            else if (!responseManager.CompletePromise(context.CorrelationId, message))
+                        failure.CorrelationId))
+                : responseManager.CompletePromise(correlationId, message);
+
+            if (!matched)
             {
-                // 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!
-            responseManager.SetException(correlationId, ex);
+            // Fault the waiting requester rather than leave it to time out.
+            if (!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 with no pending request "
+        + "(correlation id {CorrelationId}, message id {MessageId})")]
+    public static partial void ReplyProcessingFailed(
+        this ILogger logger,
+        Exception exception,
+        string correlationId,
+        string? messageId);
 }

src/Mocha/src/Mocha/DeferredResponseManager.cs

@@ -58,14 +58,18 @@ public sealed class DeferredResponseManager(TimeProvider timeProvider)
     /// </summary>
     /// <param name="correlationId">The correlation identifier of the promise to fault.</param>
     /// <param name="exception">The exception to propagate to the waiting caller.</param>
-    public void SetException(string correlationId, Exception exception)
+    /// <returns><c>true</c> if a matching promise was found and faulted; <c>false</c> if no promise was registered for the correlation identifier.</returns>
+    public bool SetException(string correlationId, Exception exception)
     {
         if (_matches.TryRemove(correlationId, out var promise))
         {
             promise.Cts.Cancel();
             promise.Cts.Dispose();
             promise.TaskCompletionSource.SetException(exception);
+            return true;
         }
+
+        return false;
     }
 
     /// <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/Middlewares/DefaultMessageBus.cs

@@ -239,14 +239,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;
         }
     }
 

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

@@ -33,72 +33,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.TryCreateResponseOptions(out var 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,9 @@ 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 reply to a message the saga sent. It covers neither
+    /// failures of events the saga published nor failures of the saga's own transitions, since neither
+    /// produces a reply.
     /// </summary>
     /// <typeparam name="TState">The saga state type.</typeparam>
     /// <param name="descriptor">The state descriptor to configure.</param>
@@ -18,11 +20,13 @@ 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. A state that declares this must also
+    /// handle fault replies through <c>OnFault</c>, either on the state itself or through
+    /// <c>DuringAny</c>.
     /// </summary>
     /// <typeparam name="TState">The saga state type.</typeparam>
     /// <param name="descriptor">The state descriptor to configure.</param>

src/Mocha/src/Mocha/Sagas/Initialization/SagaValidator.cs

@@ -1,3 +1,5 @@
+using Mocha.Events;
+
 namespace Mocha.Sagas;
 
 internal static class SagaValidator
@@ -34,6 +36,8 @@ public static void ValidateStateMachine(Saga saga)
                 finalStates.Add(stateName);
             }
 
+            ValidateFaultHandling(saga, stateName, sagaState);
+
             foreach (var transition in sagaState.Transitions.Values)
             {
                 if (!allStateNames.Contains(transition.TransitionTo))
@@ -87,4 +91,23 @@ public static void ValidateStateMachine(Saga saga)
                 "The following states cannot reach a final state: " + unreachableList);
         }
     }
+
+    /// <summary>
+    /// Rejects a state that handles any reply without handling faults, because a fault reply would
+    /// otherwise select the catch-all transition through its base type.
+    /// </summary>
+    private static void ValidateFaultHandling(Saga saga, string stateName, SagaState state)
+    {
+        if (!state.Transitions.TryGetValue(typeof(object), out var anyReply)
+            || anyReply.TransitionKind is not SagaTransitionKind.Reply
+            || state.Transitions.ContainsKey(typeof(NotAcknowledgedEvent)))
+        {
+            return;
+        }
+
+        throw new SagaInitializationException(
+            saga,
+            $"State '{stateName}' handles any reply but does not handle faults. "
+            + "Add '.OnFault()' to this state, or '.DuringAny().OnFault()' to the saga.");
+    }
 }