Codegen improvements for using multiple entities through persistence …

src/Http/Wolverine.Http.Tests/Marten/working_against_multiple_streams.cs

@@ -100,6 +100,64 @@ await Scenario(x =>
         (await fetchAmount(from)).ShouldBe(850);
         (await fetchAmount(to)).ShouldBe(250);
     }
+
+    [Fact]
+    public async Task happy_path_found_both_accounts_and_able_to_use_each_in_Before_method()
+    {
+        var from = await createAccount(1000);
+        var to = await createAccount(100);
+
+        await Scenario(x =>
+        {
+            x.Post.Json(new TransferMoney(from, to, 150)).ToUrl("/accounts/transfer4");
+            x.StatusCodeShouldBe(204);
+        });
+
+        (await fetchAmount(from)).ShouldBe(850);
+        (await fetchAmount(to)).ShouldBe(250);
+
+        TransferMoneyEndpointWithBefore.From.Id.ShouldBe(from);
+        TransferMoneyEndpointWithBefore.To.Id.ShouldBe(to);
+    }
+
+    [Fact]
+    public async Task happy_path_with_aggregate_attribute_and_before_method()
+    {
+        var from = await createAccount(1000);
+        var to = await createAccount(100);
+
+        await Scenario(x =>
+        {
+            x.Post.Json(new TransferMoney(from, to, 150)).ToUrl("/accounts/transfer5");
+            x.StatusCodeShouldBe(204);
+        });
+
+        (await fetchAmount(from)).ShouldBe(850);
+        (await fetchAmount(to)).ShouldBe(250);
+
+        TransferMoneyEndpointWithBeforeAggregate.From.Id.ShouldBe(from);
+        TransferMoneyEndpointWithBeforeAggregate.To.Id.ShouldBe(to);
+    }
+
+    [Fact]
+    public async Task happy_path_with_mixed_write_and_read_aggregate_and_before_method()
+    {
+        var from = await createAccount(1000);
+        var to = await createAccount(100);
+
+        await Scenario(x =>
+        {
+            x.Post.Json(new TransferMoney(from, to, 150)).ToUrl("/accounts/transfer6");
+            x.StatusCodeShouldBe(204);
+        });
+
+        (await fetchAmount(from)).ShouldBe(850);
+        // toAccount is read-only, so amount should not change
+        (await fetchAmount(to)).ShouldBe(100);
+
+        TransferMoneyEndpointWithBeforeMixed.From.Id.ShouldBe(from);
+        TransferMoneyEndpointWithBeforeMixed.To.Id.ShouldBe(to);
+    }
 }
 
 #region sample_when_transfering_money

src/Http/WolverineWebApi/Accounts/AccountCode.cs

@@ -95,3 +95,88 @@ public static void Handle(
     }
 }
 
+
+public static class TransferMoneyEndpointWithBefore
+{
+    public static void Before(Account fromAccount, Account toAccount)
+    {
+        From = fromAccount;
+        To = toAccount;
+    }
+
+    public static Account To { get; set; }
+
+    public static Account From { get; set; }
+
+    [WolverinePost("/accounts/transfer4")]
+    public static void Handle(
+        TransferMoney command,
+
+        [WriteAggregate(nameof(TransferMoney.FromId), LoadStyle = ConcurrencyStyle.Exclusive)] IEventStream<Account> fromAccount,
+
+        [WriteAggregate(nameof(TransferMoney.ToId))] IEventStream<Account> toAccount)
+    {
+        // Would already 404 if either referenced account does not exist
+        if (fromAccount.Aggregate.Amount >= command.Amount)
+        {
+            fromAccount.AppendOne(new Withdrawn(command.Amount));
+            toAccount.AppendOne(new Debited(command.Amount));
+        }
+    }
+}
+
+public static class TransferMoneyEndpointWithBeforeAggregate
+{
+    public static void Before(Account fromAccount, Account toAccount)
+    {
+        From = fromAccount;
+        To = toAccount;
+    }
+
+    public static Account To { get; set; }
+
+    public static Account From { get; set; }
+
+    [WolverinePost("/accounts/transfer5")]
+    public static void Handle(
+        TransferMoney command,
+
+        [Aggregate(nameof(TransferMoney.FromId))] IEventStream<Account> fromAccount,
+
+        [Aggregate(nameof(TransferMoney.ToId))] IEventStream<Account> toAccount)
+    {
+        if (fromAccount.Aggregate.Amount >= command.Amount)
+        {
+            fromAccount.AppendOne(new Withdrawn(command.Amount));
+            toAccount.AppendOne(new Debited(command.Amount));
+        }
+    }
+}
+
+public static class TransferMoneyEndpointWithBeforeMixed
+{
+    public static void Before(Account fromAccount, Account toAccount)
+    {
+        From = fromAccount;
+        To = toAccount;
+    }
+
+    public static Account To { get; set; }
+
+    public static Account From { get; set; }
+
+    [WolverinePost("/accounts/transfer6")]
+    public static void Handle(
+        TransferMoney command,
+
+        [WriteAggregate(nameof(TransferMoney.FromId))] IEventStream<Account> fromAccount,
+
+        [ReadAggregate(nameof(TransferMoney.ToId))] Account toAccount)
+    {
+        if (fromAccount.Aggregate.Amount >= command.Amount)
+        {
+            fromAccount.AppendOne(new Withdrawn(command.Amount));
+        }
+    }
+}
+

src/Persistence/Wolverine.Marten/AggregateHandling.cs

@@ -241,21 +241,42 @@ internal Variable RelayAggregateToHandlerMethod(Variable eventStream, IChain cha
             // If the handle method is on the aggregate itself
             firstCall.Target = aggregateVariable;
         }
+        else if (Parameter != null && Parameter.ParameterType.Closes(typeof(IEventStream<>)))
+        {
+            // When the handler parameter is IEventStream<T>, set the stream directly by name
+            var index = firstCall.Method.GetParameters().IndexOf(x => x.Name == Parameter.Name);
+            firstCall.Arguments[index] = eventStream;
+        }
+        else if (Parameter != null)
+        {
+            // Use name-based matching to avoid accidentally setting the wrong same-type parameter
+            firstCall.TrySetArgument(Parameter.Name, aggregateVariable);
+        }
         else
         {
-            if (!firstCall.TrySetArgument(aggregateVariable))
-            {
-                if (Parameter != null && Parameter.ParameterType.Closes(typeof(IEventStream<>)))
-                {
-                    var index = firstCall.Method.GetParameters().IndexOf(x => x.Name == Parameter.Name);
-                    firstCall.Arguments[index] = eventStream;
-                }
-            };
+            firstCall.TrySetArgument(aggregateVariable);
+        }
+
+        // Store deferred assignment for middleware methods added later (Before/After)
+        if (Parameter != null)
+        {
+            StoreDeferredMiddlewareVariable(chain, Parameter.Name, aggregateVariable);
         }
 
+        // Also do immediate relay for any middleware already present
         foreach (var methodCall in chain.Middleware.OfType<MethodCall>())
         {
-            methodCall.TrySetArgument(aggregateVariable);
+            if (Parameter != null)
+            {
+                if (!methodCall.TrySetArgument(Parameter.Name, aggregateVariable))
+                {
+                    methodCall.TrySetArgument(aggregateVariable);
+                }
+            }
+            else
+            {
+                methodCall.TrySetArgument(aggregateVariable);
+            }
         }
 
         return aggregateVariable;
@@ -295,4 +316,15 @@ internal static MemberInfo DetermineVersionMember(Type aggregateType)
             _versioningBaseType.CloseAndBuildAs<IAggregateVersioning>(AggregationScope.SingleStream, aggregateType);
         return versioning.VersionMember;
     }
+
+    internal static void StoreDeferredMiddlewareVariable(IChain chain, string parameterName, Variable variable)
+    {
+        const string key = "DeferredMiddlewareVariables";
+        if (!chain.Tags.TryGetValue(key, out var raw))
+        {
+            raw = new List<(string Name, Variable Variable)>();
+            chain.Tags[key] = raw;
+        }
+        ((List<(string Name, Variable Variable)>)raw).Add((parameterName, variable));
+    }
 }

src/Persistence/Wolverine.Marten/ReadAggregateAttribute.cs

@@ -65,20 +65,28 @@ public override Variable Modify(IChain chain, ParameterInfo parameter, IServiceC
         }
 
         var frame = new FetchLatestAggregateFrame(parameter.ParameterType, identity);
-
+        frame.Aggregate.OverrideName(parameter.Name);
+
+        Variable returnVariable;
         if (Required)
         {
             var otherFrames = chain.AddStopConditionIfNull(frame.Aggregate, identity, this);
-            
+
             var block = new LoadEntityFrameBlock(frame.Aggregate, otherFrames);
             chain.Middleware.Add(block);
 
-            return block.Mirror;
+            returnVariable = block.Mirror;
         }
-
-        chain.Middleware.Add(frame);
+        else
+        {
+            chain.Middleware.Add(frame);
+            returnVariable = frame.Aggregate;
+        }
+
+        // Store deferred assignment for middleware methods added later (Before/After)
+        AggregateHandling.StoreDeferredMiddlewareVariable(chain, parameter.Name, returnVariable);
 
-        return frame.Aggregate;
+        return returnVariable;
     }
 }
 

src/Wolverine/Configuration/Chain.cs

@@ -211,7 +211,7 @@
 
         if (type.IsArray)
         {
             return isMaybeServiceDependency(type.GetElementType());
         }
 
         if (ServiceContainer.IsEnumerable(type))
@@ -279,7 +279,7 @@
     {
         if (_appliedImpliedMiddleware) return;
         _appliedImpliedMiddleware = true;
-        
+
         var handlerTypes = HandlerCalls().Select(x => x.HandlerType).Distinct();
         foreach (var handlerType in handlerTypes)
         {
@@ -304,6 +304,26 @@
                 }
             }
         }
+
+        applyDeferredMiddlewareVariables();
+    }
+
+    private void applyDeferredMiddlewareVariables()
+    {
+        const string key = "DeferredMiddlewareVariables";
+        if (!Tags.TryGetValue(key, out var raw)) return;
+        if (raw is not List<(string Name, Variable Variable)> assignments) return;
+
+        foreach (var (name, variable) in assignments)
+        {
+            foreach (var methodCall in Middleware.OfType<MethodCall>())
+            {
+                if (!methodCall.TrySetArgument(name, variable))
+                {
+                    methodCall.TrySetArgument(variable);
+                }
+            }
+        }
     }
 
     public void AddMiddleware(GenerationRules generationRules, MethodCall frame)

src/Wolverine/Persistence/EntityAttribute.cs

@@ -161,24 +161,44 @@ public override Variable Modify(IChain chain, ParameterInfo parameter, IServiceC
         var frame = provider.DetermineLoadFrame(container, parameter.ParameterType, identity);
 
         var entity = frame.Creates.First(x => x.VariableType == parameter.ParameterType);
+        entity.OverrideName(parameter.Name);
 
         if (MaybeSoftDeleted is false)
         {
             var softDeleteFrames = provider.DetermineFrameToNullOutMaybeSoftDeleted(entity);
             chain.Middleware.AddRange(softDeleteFrames);
         }
 
+        Variable returnVariable;
         if (Required)
         {
             var otherFrames = chain.AddStopConditionIfNull(entity, identity, this);
 
             var block = new LoadEntityFrameBlock(entity, otherFrames);
             chain.Middleware.Add(block);
 
-            return block.Mirror;
+            returnVariable = block.Mirror;
+        }
+        else
+        {
+            chain.Middleware.Add(frame);
+            returnVariable = entity;
         }
 
-        chain.Middleware.Add(frame);
-        return entity;
+        // Store deferred assignment for middleware methods added later (Before/After)
+        StoreDeferredMiddlewareVariable(chain, parameter.Name, returnVariable);
+
+        return returnVariable;
+    }
+
+    internal static void StoreDeferredMiddlewareVariable(IChain chain, string parameterName, Variable variable)
+    {
+        const string key = "DeferredMiddlewareVariables";
+        if (!chain.Tags.TryGetValue(key, out var raw))
+        {
+            raw = new List<(string Name, Variable Variable)>();
+            chain.Tags[key] = raw;
+        }
+        ((List<(string Name, Variable Variable)>)raw).Add((parameterName, variable));
     }
 }