GH-4825: construct the internal agent messages' routers directly for Native AOT

src/Testing/CoreTests/Runtime/framework_router_factories_cover_every_agent_message.cs

@@ -0,0 +1,83 @@
+using System.Reflection;
+using JasperFx.Core.Reflection;
+using Shouldly;
+using Wolverine.Runtime;
+using Wolverine.Runtime.Agents;
+using Xunit;
+
+namespace CoreTests.Runtime;
+
+/// <summary>
+///     GH-4825. <c>PrepopulateRoutingCache</c> walks every message type the handler graph discovered, and
+///     <c>RoutingFor</c>'s cache-miss path closes <c>MessageRouter&lt;T&gt;</c> / <c>EmptyMessageRouter&lt;T&gt;</c>
+///     over each one reflectively. An application with a durable message store has the framework's own agent
+///     messages in that set — and they are <c>internal</c>, so
+///     <c>HandlerRegistryCodeFile.onlyPublic()</c> drops them and the emitted <c>[DynamicDependency]</c>
+///     rooting block <b>cannot</b> name them: generated code cannot put a non-public type inside a
+///     <c>typeof()</c>. In a native image ILC has therefore trimmed exactly those instantiations.
+///
+///     <para>Direct construction in <c>_frameworkRouterFactories</c> is the only fix available, which is
+///     what GH-4287 already did for <see cref="IAgentCommand" /> itself.</para>
+/// </summary>
+/// <remarks>
+///     Asserted as COVERAGE of the real set rather than as a list of the eight types that were missing.
+///     A list would be satisfied the day it was written and say nothing when the next agent message is
+///     added — and the cost of missing one is a native image that fails at startup, which no analyzer
+///     reports and only a Balanced-mode publish reproduces.
+/// </remarks>
+public class framework_router_factories_cover_every_agent_message
+{
+    [Fact]
+    public void every_internal_agent_message_has_a_directly_constructed_router_factory()
+    {
+        var factories = frameworkRouterFactoryKeys();
+
+        // Not vacuous: if the reflection above ever stops finding the field, this fails rather than
+        // passing over an empty set.
+        factories.ShouldNotBeEmpty();
+
+        var uncovered = agentMessageTypes()
+            .Where(x => !factories.Contains(x))
+            .OrderBy(x => x.FullName, StringComparer.Ordinal)
+            .ToArray();
+
+        uncovered.ShouldBeEmpty(
+            $"These agent message types are internal, so no emitted rooting block can name them, and nothing constructs their routers directly -- a native image with a durable message store will fail in PrepopulateRoutingCache: {string.Join(", ", uncovered.Select(x => x.Name))}");
+    }
+
+    [Fact]
+    public void the_agent_message_scan_actually_finds_something()
+    {
+        // The guard for the test above. agentMessageTypes() returning nothing would make it green while
+        // covering no types at all, which is precisely the failure mode this whole class of AOT bug has.
+        agentMessageTypes().ShouldNotBeEmpty();
+
+        // And they really are non-public, which is the reason they cannot be rooted from generated code.
+        foreach (var type in agentMessageTypes())
+        {
+            type.IsPublic.ShouldBeFalse($"{type.FullName} is public, so the rooting block could name it");
+        }
+    }
+
+    private static Type[] agentMessageTypes()
+    {
+        return typeof(IAgentCommand).Assembly
+            .GetTypes()
+            .Where(x => x is { IsAbstract: false, IsInterface: false, IsGenericTypeDefinition: false })
+            .Where(x => x.CanBeCastTo<IAgentCommand>() || x.CanBeCastTo<IDeferredAgentWork>())
+            .Where(x => !x.IsPublic)
+            .ToArray();
+    }
+
+    private static HashSet<Type> frameworkRouterFactoryKeys()
+    {
+        var field = typeof(WolverineRuntime).GetField("_frameworkRouterFactories",
+            BindingFlags.NonPublic | BindingFlags.Static);
+
+        field.ShouldNotBeNull("WolverineRuntime._frameworkRouterFactories was renamed; this test needs updating");
+
+        var dictionary = (System.Collections.IDictionary)field!.GetValue(null)!;
+
+        return dictionary.Keys.Cast<Type>().ToHashSet();
+    }
+}

src/Testing/Wolverine.AotSmoke.Marten/Wolverine.AotSmoke.Marten.csproj

@@ -21,6 +21,9 @@
              not the Configuration.Applier<T> that the GH-4778 rooting block covers.
           3. IntegrateWithWolverine registers Envelope with Schema.For<Envelope>(), and in a native image
              Marten's DocumentMapping cannot find Envelope.Id.
+          4. With a durable message store, PrepopulateRoutingCache closes MessageRouter<T> over the
+             framework's internal agent messages (StopAgents and its siblings), which nothing rooted.
+             NOT reproduced by this lane, which runs Solo, where agent messaging is never used (GH-4843).
 
         The endpoint shapes are the ones the reporting application uses. Boots, serves both endpoints and
         reads the started stream back. Needs the postgresql service from docker-compose.yml.
@@ -29,7 +32,9 @@
 
           dotnet run -f net9.0 codegen write
 
-        This project carries NO hand-written roots, so the lane fails until the library itself is fixed.
+        This project carries NO hand-written roots, so a regression in the library's own rooting fails the
+        lane outright. (It said "the lane fails until the library itself is fixed" while that was true; the
+        library was fixed in #4828 and the lane passes.)
         -->
         <PublishAot>true</PublishAot>
         <StripSymbols>false</StripSymbols>

src/Wolverine/Runtime/WolverineRuntime.Routing.cs

@@ -227,7 +227,36 @@ private static readonly Dictionary<Type, Func<WolverineRuntime, List<IMessageRou
             [typeof(Envelope)] = routerFactory<Envelope>(),
             [typeof(Acknowledgement)] = routerFactory<Acknowledgement>(),
             [typeof(FailureAcknowledgement)] = routerFactory<FailureAcknowledgement>(),
-            [typeof(IAgentCommand)] = routerFactory<IAgentCommand>()
+            [typeof(IAgentCommand)] = routerFactory<IAgentCommand>(),
+
+            // GH-4825. The concrete agent messages the framework itself handles. They are internal, so the
+            // `codegen write` rooting block cannot name them -- HandlerRegistryCodeFile.onlyPublic() drops
+            // every non-public type, because generated code cannot put one inside a typeof(). An application
+            // with a durable message store walks all of them through PrepopulateRoutingCache at startup, so
+            // in a native image the reflective close in RoutingFor fires against instantiations nothing
+            // could root. Direct construction is the only fix available, exactly as GH-4287 did for
+            // IAgentCommand above.
+            //
+            // Held to COVERAGE of the real set by framework_router_factories_cover_every_agent_message,
+            // not to this list: the eight types the reported failure named were the ones that application's
+            // lane happened to reach, and that test immediately found six more with the identical defect --
+            // the Balanced-mode leader's own assignment messages, built in
+            // NodeAgentController.EvaluateAssignments and AssignmentGrid. Add a new agent message and that
+            // test fails until it is covered here.
+            [typeof(AgentPresenceReport)] = routerFactory<AgentPresenceReport>(),
+            [typeof(AgentsStarted)] = routerFactory<AgentsStarted>(),
+            [typeof(AgentsStopped)] = routerFactory<AgentsStopped>(),
+            [typeof(AssignAgent)] = routerFactory<AssignAgent>(),
+            [typeof(AssignAgents)] = routerFactory<AssignAgents>(),
+            [typeof(QueryAgentPresence)] = routerFactory<QueryAgentPresence>(),
+            [typeof(ReassignAgent)] = routerFactory<ReassignAgent>(),
+            [typeof(ReassignAgents)] = routerFactory<ReassignAgents>(),
+            [typeof(StartAgent)] = routerFactory<StartAgent>(),
+            [typeof(StartAgents)] = routerFactory<StartAgents>(),
+            [typeof(StopAgent)] = routerFactory<StopAgent>(),
+            [typeof(StopAgents)] = routerFactory<StopAgents>(),
+            [typeof(StopRemoteAgent)] = routerFactory<StopRemoteAgent>(),
+            [typeof(StopRemoteAgents)] = routerFactory<StopRemoteAgents>()
         };
 
     private static Func<WolverineRuntime, List<IMessageRoute>, IMessageRouter> routerFactory<T>()