Keep the service key when resolving keyed services via service location

src/CodegenTests/Services/keyed_service_location.cs

@@ -0,0 +1,105 @@
+using JasperFx;
+using JasperFx.CodeGeneration;
+using JasperFx.CodeGeneration.Frames;
+using JasperFx.CodeGeneration.Services;
+using JasperFx.RuntimeCompiler;
+using Microsoft.Extensions.DependencyInjection;
+using Shouldly;
+using Xunit.Abstractions;
+
+namespace CodegenTests.Services;
+
+// GH-2878 (JasperFx/wolverine): keyed services resolved through the scoped-IServiceProvider
+// "service location" path lost their key — the frame always emitted GetRequiredService<T>(provider)
+// instead of GetRequiredKeyedService<T>(provider, key). Service location kicks in whenever something
+// in the same generated method forces it (a directly-injected IServiceProvider, or an opaque lambda
+// registration like the ones MS Graph adds), which then drags every sibling dependency — including
+// keyed ones — onto the service-location path where the key was being dropped.
+public class keyed_service_location
+{
+    private readonly ITestOutputHelper _output;
+    private readonly ServiceCollection theServices = new();
+
+    public keyed_service_location(ITestOutputHelper output)
+    {
+        _output = output;
+    }
+
+    // Generates a harness whose Build() calls KeyedHandler.Handle(...), exactly how Wolverine
+    // generates a handler: the keyed IWidget parameter is resolved via [FromKeyedServices] during
+    // code generation, and the opaque IScopedLambda sibling forces the method onto service location.
+    private (string Code, ServiceContainer Graph, GeneratedType Type) generateHandlerCall()
+    {
+        var graph = new ServiceContainer(theServices, theServices.BuildServiceProvider());
+
+        var assembly = new GeneratedAssembly(new GenerationRules());
+        var type = assembly.AddType("KeyedHandlerHarness", typeof(ServiceHarness<WidgetResult>));
+        var buildMethod = type.MethodFor("Build");
+
+        var call = new MethodCall(typeof(KeyedHandler), nameof(KeyedHandler.Handle));
+        buildMethod.Frames.Add(call);
+        buildMethod.Frames.Code("return {0};", call.ReturnVariable!);
+
+        var source = new ServiceCollectionServerVariableSource(graph);
+        source.StartNewType();
+        source.StartNewMethod();
+
+        var code = assembly.GenerateCode(source);
+        _output.WriteLine(code);
+        return (code, graph, type);
+    }
+
+    [Fact]
+    public void keyed_concrete_service_dragged_onto_location_keeps_its_key()
+    {
+        theServices.AddKeyedScoped<IWidget, BWidget>("blue");
+        theServices.AddScoped<IScopedLambda>(_ => new ScopedLambda());
+
+        var (code, _, _) = generateHandlerCall();
+
+        code.ShouldContain("GetRequiredKeyedService");
+        code.ShouldContain("\"blue\"");
+        code.ShouldNotContain("GetRequiredService<CodegenTests.Services.IWidget>");
+    }
+
+    [Fact]
+    public void keyed_opaque_registration_keeps_its_key()
+    {
+        theServices.AddKeyedScoped<IWidget>("blue", (_, _) => new BWidget());
+        theServices.AddScoped<IScopedLambda>(_ => new ScopedLambda());
+
+        var (code, _, _) = generateHandlerCall();
+
+        code.ShouldContain("GetRequiredKeyedService");
+        code.ShouldContain("\"blue\"");
+    }
+
+    [Fact]
+    public void keyed_service_resolves_correctly_at_runtime()
+    {
+        // End-to-end: compile the generated harness and prove the keyed service actually resolves
+        // (before the fix this threw "No service for type IWidget has been registered").
+        theServices.AddKeyedScoped<IWidget, BWidget>("blue");
+        theServices.AddScoped<IScopedLambda>(_ => new ScopedLambda());
+
+        var (code, graph, _) = generateHandlerCall();
+
+        var compiler = new AssemblyGenerator();
+        compiler.ReferenceAssembly(GetType().Assembly);
+        var builtAssembly = compiler.Generate(code);
+        var builtType = builtAssembly.ExportedTypes.Single();
+
+        var result = ((ServiceHarness<WidgetResult>)graph.BuildFromType(builtType)).Build();
+        result.Widget.ShouldBeOfType<BWidget>();
+    }
+}
+
+public class KeyedHandler
+{
+    public static WidgetResult Handle([FromKeyedServices("blue")] IWidget widget, IScopedLambda opaque)
+    {
+        return new WidgetResult(widget);
+    }
+}
+
+public record WidgetResult(IWidget Widget);

src/JasperFx/CodeGeneration/Services/GetServiceFromScopedContainerFrame.cs

@@ -8,17 +8,19 @@ namespace JasperFx.CodeGeneration.Services;
 public class GetServiceFromScopedContainerFrame : SyncFrame
 {
     private readonly Variable _scoped;
+    private readonly object? _serviceKey;
 
 
-    public GetServiceFromScopedContainerFrame(Variable scoped, Type serviceType)
+    public GetServiceFromScopedContainerFrame(Variable scoped, Type serviceType, object? serviceKey = null)
     {
         if (scoped.VariableType != typeof(IServiceProvider))
         {
             throw new ArgumentOutOfRangeException(nameof(scoped),
                 $"Wrong type for the variable. Expected {typeof(IServiceProvider).FullNameInCode()} but got {scoped.VariableType.FullNameInCode()}");
         }
-        
+
         _scoped = scoped;
+        _serviceKey = serviceKey;
         uses.Add(_scoped);
 
         Variable = new Variable(serviceType, this);
@@ -60,8 +62,19 @@ public override void GenerateCode(GeneratedMethod method, ISourceWriter writer)
             Header.Write(writer);
         }
 
-        writer.Write(
-            $"var {Variable.Usage} = {typeof(ServiceProviderServiceExtensions).FullNameInCode()}.{nameof(ServiceProviderServiceExtensions.GetRequiredService)}<{Variable.VariableType.FullNameInCode()}>({_scoped.Usage});");
+        if (_serviceKey != null)
+        {
+            // Keyed service resolved through service location. Without the key this would emit
+            // GetRequiredService<T> and either throw or resolve the wrong registration. See GH-2878.
+            writer.Write(
+                $"var {Variable.Usage} = {typeof(ServiceProviderKeyedServiceExtensions).FullNameInCode()}.{nameof(ServiceProviderKeyedServiceExtensions.GetRequiredKeyedService)}<{Variable.VariableType.FullNameInCode()}>({_scoped.Usage}, {CodeFormatter.Write(_serviceKey)});");
+        }
+        else
+        {
+            writer.Write(
+                $"var {Variable.Usage} = {typeof(ServiceProviderServiceExtensions).FullNameInCode()}.{nameof(ServiceProviderServiceExtensions.GetRequiredService)}<{Variable.VariableType.FullNameInCode()}>({_scoped.Usage});");
+        }
+
         Next?.GenerateCode(method, writer);
     }
 }

src/JasperFx/CodeGeneration/Services/ServiceCollectionServerVariableSource.cs

@@ -153,7 +153,11 @@ private void useServiceProvider(IMethodVariables method)
         var written = false;
         foreach (var standin in _standins)
         {
-            var frame = new GetServiceFromScopedContainerFrame(_scoped, standin.VariableType);
+            // Keyed services must keep their key when dragged onto the service-location path,
+            // otherwise the generated code emits GetRequiredService<T> and loses the key. See GH-2878.
+            var descriptor = standin.Plan.Descriptor;
+            var serviceKey = descriptor is { IsKeyedService: true } ? descriptor.ServiceKey : null;
+            var frame = new GetServiceFromScopedContainerFrame(_scoped, standin.VariableType, serviceKey);
             var variable = frame.Variable;
 
             // Write description of why this had to use the nested container

src/JasperFx/CodeGeneration/Services/ServiceLocationPlan.cs

@@ -40,7 +40,8 @@ public override string WhyRequireServiceProvider(IMethodVariables method)
 
     public override Variable CreateVariable(ServiceVariables resolverVariables)
     {
-        return new GetServiceFromScopedContainerFrame(resolverVariables.ServiceProvider, Descriptor.ServiceType)
+        var serviceKey = Descriptor.IsKeyedService ? Descriptor.ServiceKey : null;
+        return new GetServiceFromScopedContainerFrame(resolverVariables.ServiceProvider, Descriptor.ServiceType, serviceKey)
             .Variable;
     }
 }