Marten IDocumentSession scope priming (GH-3001, Marten slice)

src/Persistence/MartenTests/service_location_document_session.cs

@@ -0,0 +1,81 @@
+using IntegrationTests;
+using JasperFx.CodeGeneration.Model;
+using Marten;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+using Shouldly;
+using Wolverine;
+using Wolverine.Marten;
+using Wolverine.Tracking;
+using Xunit;
+
+namespace MartenTests;
+
+/// <summary>
+/// GH-3001: when a handler chain falls back to service location, a dependency that takes
+/// <see cref="IDocumentSession"/> must receive the SAME outbox-enrolled session the handler is
+/// using — not a separate, un-enrolled one (which would defeat the transaction boundary). Proven via
+/// reference equality against the handler's own session.
+/// </summary>
+public class service_location_document_session : PostgresqlContext
+{
+    [Fact]
+    public async Task service_located_session_is_same_instance_as_the_handler_session()
+    {
+        using var host = await Host.CreateDefaultBuilder()
+            .UseWolverine(opts =>
+            {
+                opts.Services.AddMarten(Servers.PostgresConnectionString).IntegrateWithWolverine();
+
+                opts.Discovery.DisableConventionalDiscovery().IncludeType(typeof(SessionProbeCommandHandler));
+
+                opts.ServiceLocationPolicy = ServiceLocationPolicy.AllowedButWarn;
+                opts.Services.AddScoped<SessionCapturingService>();
+
+                // Force the capturing service to be resolved via service location so the chain creates
+                // a child scope — the path GH-3001 primes.
+                opts.CodeGeneration.AlwaysUseServiceLocationFor<SessionCapturingService>();
+            }).StartAsync();
+
+        SessionIdentityProbe.Reset();
+
+        await host.InvokeMessageAndWaitAsync(new SessionProbeCommand());
+
+        SessionIdentityProbe.HandlerSession.ShouldNotBeNull();
+        SessionIdentityProbe.ServiceLocatedSession.ShouldNotBeNull();
+
+        // Reference equality — the service-located session IS the handler's outbox-enrolled session.
+        ReferenceEquals(SessionIdentityProbe.HandlerSession, SessionIdentityProbe.ServiceLocatedSession)
+            .ShouldBeTrue();
+    }
+}
+
+public record SessionProbeCommand;
+
+public static class SessionIdentityProbe
+{
+    public static IDocumentSession? HandlerSession;
+    public static IDocumentSession? ServiceLocatedSession;
+
+    public static void Reset()
+    {
+        HandlerSession = null;
+        ServiceLocatedSession = null;
+    }
+}
+
+public class SessionCapturingService(IDocumentSession session)
+{
+    public IDocumentSession Capture() => session;
+}
+
+public static class SessionProbeCommandHandler
+{
+    public static void Handle(SessionProbeCommand command, IDocumentSession handlerSession, IServiceProvider services)
+    {
+        SessionIdentityProbe.HandlerSession = handlerSession;
+        SessionIdentityProbe.ServiceLocatedSession = services
+            .GetRequiredService<SessionCapturingService>()
+            .Capture();
+    }
+}

src/Persistence/Wolverine.Marten/Codegen/PrimeScopedDocumentSessionFrame.cs

@@ -0,0 +1,59 @@
+using JasperFx.CodeGeneration;
+using JasperFx.CodeGeneration.Frames;
+using JasperFx.CodeGeneration.Model;
+using JasperFx.CodeGeneration.Services;
+using JasperFx.Core.Reflection;
+using Marten;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Wolverine.Marten.Codegen;
+
+/// <summary>
+/// Emitted (via <c>IScopedContainerCreation.AddPostProcessor</c>) immediately after a handler's
+/// service-location child scope is created. Primes that scope's <see cref="ScopedDocumentSessionHolder"/>
+/// with the handler's outbox-enrolled <see cref="IDocumentSession"/>, so any service-located
+/// <see cref="IDocumentSession"/> / <see cref="IQuerySession"/> resolves to that single enrolled
+/// session rather than a separate one. See GH-3001.
+/// </summary>
+internal sealed class PrimeScopedDocumentSessionFrame : SyncFrame, IUsesServiceProviderFrame
+{
+    private Variable? _session;
+    private Variable? _scopedProvider;
+
+    // The parent ScopedContainerCreation hands us the scoped IServiceProvider variable before we
+    // resolve our other variables (avoiding a bi-directional dependency with the scope line).
+    public void UseServiceProvider(Variable serviceProvider) => _scopedProvider = serviceProvider;
+
+    public override IEnumerable<Variable> FindVariables(IMethodVariables chain)
+    {
+        // The enrolled session created by CreateDocumentSessionFrame (NotServices: never the
+        // container's own scoped IDocumentSession registration).
+        _session = chain.TryFindVariable(typeof(IDocumentSession), VariableSource.NotServices);
+        if (_session != null)
+        {
+            yield return _session;
+        }
+    }
+
+    public override void GenerateCode(GeneratedMethod method, ISourceWriter writer)
+    {
+        if (_session != null)
+        {
+            writer.Write(
+                $"{typeof(ServiceProviderServiceExtensions).FullNameInCode()}.{nameof(ServiceProviderServiceExtensions.GetRequiredService)}<{typeof(ScopedDocumentSessionHolder).FullNameInCode()}>({_scopedProvider!.Usage}).{nameof(ScopedDocumentSessionHolder.Session)} = {_session.Usage};");
+        }
+
+        Next?.GenerateCode(method, writer);
+    }
+
+    public override void GenerateFSharpCode(GeneratedMethod method, ISourceWriter writer)
+    {
+        if (_session != null)
+        {
+            writer.Write(
+                $"{typeof(ServiceProviderServiceExtensions).FSharpName()}.{nameof(ServiceProviderServiceExtensions.GetRequiredService)}<{typeof(ScopedDocumentSessionHolder).FSharpName()}>({_scopedProvider!.Usage}).{nameof(ScopedDocumentSessionHolder.Session)} <- {_session.Usage}");
+        }
+
+        Next?.GenerateFSharpCode(method, writer);
+    }
+}

src/Persistence/Wolverine.Marten/MartenIntegration.cs

@@ -63,6 +63,11 @@ public void Configure(WolverineOptions options)
 
         options.CodeGeneration.Sources.Add(new MartenBackedPersistenceMarker());
 
+        // GH-3001: prime the service-location child scope with the handler's outbox-enrolled
+        // IDocumentSession so a service-located IDocumentSession / IQuerySession resolves to that same
+        // session. The frame self-guards (no-op when the chain has no Marten session).
+        options.ScopingFrameSources.Add(() => new PrimeScopedDocumentSessionFrame());
+
         options.CodeGeneration.InsertFirstPersistenceStrategy<MartenPersistenceFrameProvider>();
         options.CodeGeneration.Sources.Add(new SessionVariableSource());
         options.CodeGeneration.Sources.Add(new DocumentOperationsSource());

src/Persistence/Wolverine.Marten/ScopedDocumentSessionHolder.cs

@@ -0,0 +1,19 @@
+using Marten;
+
+namespace Wolverine.Marten;
+
+/// <summary>
+/// Scope-local carrier for the outbox-enrolled <see cref="IDocumentSession"/> a handler is using.
+/// When a handler falls back to service location, Wolverine's generated code creates a child
+/// <see cref="IServiceScope"/> off the root provider; the Marten scoping frame primes this holder in
+/// that scope so a service-located <see cref="IDocumentSession"/> / <see cref="IQuerySession"/>
+/// resolves to the SAME session enrolled with the active outbox instead of a separate, un-enrolled
+/// one (which would defeat the transaction boundary). See GH-3001.
+///
+/// The holder is empty in non-handler scopes (hosted services, admin tools, raw resolution), where
+/// the decorated registration falls back to Marten's own session factory.
+/// </summary>
+public sealed class ScopedDocumentSessionHolder
+{
+    public IDocumentSession? Session { get; set; }
+}

src/Persistence/Wolverine.Marten/WolverineOptionsMartenExtensions.cs

@@ -88,6 +88,17 @@ public static MartenServiceCollectionExtensions.MartenConfigurationExpression In
 
         expression.Services.AddScoped<IMartenOutbox, MartenOutbox>();
 
+        // GH-3001: structural scope priming for Marten sessions. When a handler falls back to service
+        // location, the generated code primes the child scope's ScopedDocumentSessionHolder with the
+        // handler's outbox-enrolled IDocumentSession (PrimeScopedDocumentSessionFrame). Decorate
+        // Marten's own IDocumentSession / IQuerySession scoped registrations so service-located
+        // resolution prefers that primed session — enrolled with the active outbox — instead of a
+        // separate, un-enrolled session. Non-handler scopes (the holder is empty) fall back to
+        // Marten's original session factory.
+        expression.Services.AddScoped<ScopedDocumentSessionHolder>();
+        preferScopedSession<IDocumentSession>(expression.Services);
+        preferScopedSession<IQuerySession>(expression.Services);
+
         // Gotta have at least a placeholder just in case a user also has
         // EF Core
         expression.Services.AddSingleton<DatabaseSettings>(s =>
@@ -154,6 +165,48 @@ public static MartenServiceCollectionExtensions.MartenConfigurationExpression In
         return expression;
     }
 
+    // GH-3001: replace Marten's scoped session registration with one that prefers a scope-primed
+    // session (the outbox-enrolled session the handler is using), falling back to Marten's original
+    // factory when the holder is empty (non-handler scopes). Preserving the original factory keeps
+    // Marten's exact session-building (options, tenancy) for the fall-back path.
+    private static void preferScopedSession<T>(IServiceCollection services) where T : class
+    {
+        var descriptor = services.LastOrDefault(x => x.ServiceType == typeof(T));
+        if (descriptor == null)
+        {
+            return;
+        }
+
+        Func<IServiceProvider, object> original;
+        if (descriptor.ImplementationFactory != null)
+        {
+            original = descriptor.ImplementationFactory;
+        }
+        else if (descriptor.ImplementationInstance != null)
+        {
+            original = _ => descriptor.ImplementationInstance;
+        }
+        else if (descriptor.ImplementationType != null)
+        {
+            original = sp => ActivatorUtilities.CreateInstance(sp, descriptor.ImplementationType);
+        }
+        else
+        {
+            return;
+        }
+
+        services.Remove(descriptor);
+        services.AddScoped<T>(sp =>
+        {
+            if (sp.GetRequiredService<ScopedDocumentSessionHolder>().Session is T primed)
+            {
+                return primed;
+            }
+
+            return (T)original(sp);
+        });
+    }
+
     internal static NpgsqlDataSource findMasterDataSource(
         DocumentStore store,
         IWolverineRuntime runtime,

src/Wolverine/Configuration/Chain.cs

@@ -497,9 +497,10 @@ protected internal void tryApplyResponseAware()
     ///
     /// When a chain service-locates, the generated code creates a child scope, and Wolverine primes
     /// that scope so a service-located <see cref="IMessageContext"/> / <see cref="IMessageBus"/>
-    /// (and <see cref="IRequireScopingFrame"/> contributors such as Marten's
-    /// <c>IDocumentSession</c>) resolves to the same instance the handler received rather than a
-    /// duplicate. See GH-3001 (which replaced the earlier AsyncLocal handoff from GH-2583).
+    /// (and integration-contributed instances such as Marten's <c>IDocumentSession</c>, via
+    /// <c>WolverineOptions.ScopingFrameSources</c>) resolves to the same instance the handler
+    /// received rather than a duplicate. See GH-3001 (which replaced the earlier AsyncLocal handoff
+    /// from GH-2583).
     /// </summary>
     public bool UsesServiceLocation { get; private set; }
 

src/Wolverine/Runtime/Handlers/HandlerChain.cs

@@ -674,11 +674,13 @@ internal virtual List<Frame> DetermineFrames(GenerationRules rules, IServiceCont
             : Array.Empty<Frame>();
 
         // GH-3001: when this chain falls back to service location, prime the child scope so
-        // service-located IMessageContext / IMessageBus (and any IRequireScopingFrame contributor)
-        // resolve to the same instances the handler uses instead of duplicates. The activator emits
-        // nothing unless a service-location scope is actually created for the chain.
+        // service-located IMessageContext / IMessageBus (and integration-registered instances like
+        // Marten's outbox-enrolled IDocumentSession) resolve to the same instances the handler uses
+        // instead of duplicates. The MessageContext frame is always added; integrations contribute
+        // more via WolverineOptions.ScopingFrameSources. Every scoping frame self-guards, and the
+        // activator emits nothing unless a service-location scope is actually created for the chain.
         var scopingFrames = new List<SyncFrame> { new PrimeScopedMessageContextFrame() };
-        scopingFrames.AddRange(collectScopingFrames(handlerReturnValueFrames));
+        scopingFrames.AddRange(options?.ScopingFrameSources.Select(x => x()) ?? []);
         var scopeActivator = new ScopePrimingActivatorFrame(scopingFrames);
 
         // The Enqueue cascading needs to happen before the post processors because of the
@@ -696,28 +698,6 @@ internal virtual List<Frame> DetermineFrames(GenerationRules rules, IServiceCont
             .ToList();
     }
 
-    // GH-3001: collect a scoping frame from every IRequireScopingFrame in the chain (persistence
-    // providers contribute IDocumentSession / IQuerySession priming), de-duplicated so two frames
-    // never prime the same type. The MessageContext priming frame is added unconditionally by the
-    // caller, so it is excluded here.
-    private IEnumerable<SyncFrame> collectScopingFrames(IEnumerable<Frame> handlerReturnValueFrames)
-    {
-        var candidates = Middleware
-            .Concat(Handlers)
-            .Concat(handlerReturnValueFrames)
-            .Concat(Postprocessors)
-            .OfType<IRequireScopingFrame>();
-
-        var seen = new HashSet<Type>();
-        foreach (var candidate in candidates)
-        {
-            var frame = candidate.BuildScopingFrame();
-            if (frame != null && seen.Add(candidate.GetType()))
-            {
-                yield return frame;
-            }
-        }
-    }
 
     protected void applyCustomizations(GenerationRules rules, IServiceContainer container)
     {

src/Wolverine/Runtime/Handlers/ScopePrimingActivatorFrame.cs

@@ -9,9 +9,10 @@ namespace Wolverine.Runtime.Handlers;
 /// Codegen-time activator (GH-3001). When a chain falls back to service location, the generated code
 /// creates a child <c>IServiceScope</c> off the root provider. This frame detects that scope at
 /// arrangement time and registers scoping postprocessors on it — always the
-/// <see cref="PrimeScopedMessageContextFrame"/>, plus one from every collected
-/// <see cref="Wolverine.Configuration.IRequireScopingFrame"/> — so they run immediately after the
-/// scope is created and prime it with the correct already-resolved instances. Emits no code itself.
+/// <see cref="PrimeScopedMessageContextFrame"/>, plus the frames produced by
+/// <c>WolverineOptions.ScopingFrameSources</c> (integration-contributed, e.g. Marten's session
+/// priming) — so they run immediately after the scope is created and prime it with the correct
+/// already-resolved instances. Emits no code itself.
 ///
 /// If no service-location scope is created for the chain (the IServiceProvider variable's Creator is
 /// not an <see cref="IScopedContainerCreation"/>), nothing is attached.

src/Wolverine/WolverineOptions.cs

@@ -332,6 +332,17 @@ public MultipleHandlerBehavior MultipleHandlerBehavior
     /// </summary>
     public List<IEnvelopeRule> MetadataRules { get; } = new();
 
+    /// <summary>
+    /// GH-3001 extension point. Factories for codegen frames that prime a handler's service-location
+    /// child scope with an already-resolved "singleton-per-message" instance (e.g. Marten's
+    /// outbox-enrolled IDocumentSession), so service-located dependencies resolve to that instance
+    /// rather than a duplicate. Each produced frame is emitted right after the scope is created and
+    /// must self-guard (no-op when its target variable is absent from the chain). Integrations
+    /// (Wolverine.Marten / Wolverine.Polecat) register a factory here; the MessageContext priming
+    /// frame is always added by the runtime in addition to these.
+    /// </summary>
+    public List<Func<JasperFx.CodeGeneration.Frames.SyncFrame>> ScopingFrameSources { get; } = new();
+
 
     /// For advanced usages, this gives you the ability to register pre-canned message handling
     /// that does not require any code generation.