fix(runtime)!: configure grain context before construction

src/Orleans.Runtime/Activation/ActivationDataActivatorProvider.cs

@@ -76,14 +76,19 @@ public ActivationDataActivator(
             _startActivation = state => ((ActivationData)state!).Start(_grainActivator);
         }
 
-        public IGrainContext CreateContext(GrainAddress activationAddress)
+        public IGrainContext CreateContext(GrainAddress activationAddress, IConfigureGrainContext[] configureActions)
         {
             var context = new ActivationData(
                 activationAddress,
                 _createWorkItemGroup,
                 _serviceProvider,
                 _sharedComponents);
 
+            foreach (var configure in configureActions)
+            {
+                configure.Configure(context);
+            }
+
             using var ecSuppressor = ExecutionContext.SuppressFlow();
             _ = Task.Factory.StartNew(
                 _startActivation,
@@ -98,5 +103,6 @@ public IGrainContext CreateContext(GrainAddress activationAddress)
 
 internal class StatelessWorkerActivator(StatelessWorkerGrainTypeSharedContext sharedContext, IGrainContextActivator innerActivator) : IGrainContextActivator
 {
-    public IGrainContext CreateContext(GrainAddress address) => new StatelessWorkerGrainContext(address, sharedContext, innerActivator);
+    public IGrainContext CreateContext(GrainAddress address, IConfigureGrainContext[] configureActions)
+        => new StatelessWorkerGrainContext(address, sharedContext, innerActivator, configureActions);
 }

src/Orleans.Runtime/Activation/IGrainContextActivator.cs

@@ -64,13 +64,7 @@ public IGrainContext CreateInstance(GrainAddress address)
                 activator = this.CreateActivator(grainId.Type);
             }
 
-            var result = activator.Activator.CreateContext(address);
-            foreach (var configure in activator.ConfigureActions)
-            {
-                configure.Configure(result);
-            }
-
-            return result;
+            return activator.Activator.CreateContext(address, activator.ConfigureActions);
         }
 
         private (IGrainContextActivator, IConfigureGrainContext[]) CreateActivator(GrainType grainType)
@@ -135,8 +129,9 @@ public interface IGrainContextActivator
         /// Creates a grain context for the given grain address.
         /// </summary>
         /// <param name="address">The grain address.</param>
+        /// <param name="configureActions">The actions which must be used to configure the context before grain construction begins.</param>
         /// <returns>The newly created grain context.</returns>
-        public IGrainContext CreateContext(GrainAddress address);
+        public IGrainContext CreateContext(GrainAddress address, IConfigureGrainContext[] configureActions);
     }
 
     /// <summary>

src/Orleans.Runtime/Catalog/StatelessWorkerGrainContext.cs

@@ -43,12 +43,18 @@ internal partial class StatelessWorkerGrainContext : IGrainContext, IAsyncDispos
     public StatelessWorkerGrainContext(
         GrainAddress address,
         StatelessWorkerGrainTypeSharedContext sharedContext,
-        IGrainContextActivator innerActivator)
+        IGrainContextActivator innerActivator,
+        IConfigureGrainContext[] configureActions)
     {
         Address = address;
         _shared = sharedContext;
         _innerActivator = innerActivator;
 
+        foreach (var configure in configureActions)
+        {
+            configure.Configure(this);
+        }
+
         if (_shared.RemoveIdleWorkers)
         {
             _inspectionTimer = new Timer(
@@ -320,7 +326,7 @@ private ActivationData CreateWorker(object? message)
     {
         Debug.Assert(!_terminated, "CreateWorker must not be called on a terminated stateless worker context.");
         var address = GrainAddress.GetAddress(Address.SiloAddress, Address.GrainId, ActivationId.NewId());
-        var newWorker = (ActivationData)_innerActivator.CreateContext(address);
+        var newWorker = (ActivationData)_innerActivator.CreateContext(address, []);
 
         // Observe the create/destroy lifecycle of the activation
         newWorker.SetComponent<IActivationLifecycleObserver>(this);

src/api/Orleans.Runtime/Orleans.Runtime.cs

@@ -749,7 +749,7 @@ public partial interface IGrainActivator
 
     public partial interface IGrainContextActivator
     {
-        IGrainContext CreateContext(GrainAddress address);
+        IGrainContext CreateContext(GrainAddress address, IConfigureGrainContext[] configureActions);
     }
 
     public partial interface IGrainContextActivatorProvider

test/Orleans.Core.Tests/Runtime/GrainContextActivatorTests.cs

@@ -0,0 +1,68 @@
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using NSubstitute;
+using Orleans.Metadata;
+using Orleans.Runtime;
+using TestExtensions;
+using Xunit;
+
+namespace UnitTests.Runtime;
+
+public class GrainContextActivatorTests
+{
+    [Fact, TestCategory("BVT")]
+    public void CreateInstance_ConfiguresContextBeforeStartingActivation()
+    {
+        var events = new List<string>();
+        var context = Substitute.For<IGrainContext>();
+        var contextActivator = new TestGrainContextActivator(context, events);
+        var activator = new GrainContextActivator(
+            [new TestGrainContextActivatorProvider(contextActivator)],
+            [new TestConfigureGrainContextProvider(events)],
+            new GrainPropertiesResolver(Substitute.For<IClusterManifestProvider>()));
+        var address = new GrainAddress { GrainId = GrainId.Create("test", "grain") };
+
+        Assert.Same(context, activator.CreateInstance(address));
+        Assert.Equal(["configure", "activate"], events);
+    }
+
+    private sealed class TestGrainContextActivatorProvider(IGrainContextActivator activator) : IGrainContextActivatorProvider
+    {
+        public bool TryGet(GrainType grainType, [NotNullWhen(true)] out IGrainContextActivator? result)
+        {
+            result = activator;
+            return true;
+        }
+    }
+
+    private sealed class TestConfigureGrainContextProvider(List<string> events) : IConfigureGrainContextProvider
+    {
+        public bool TryGetConfigurator(
+            GrainType grainType,
+            GrainProperties properties,
+            [NotNullWhen(true)] out IConfigureGrainContext? configurator)
+        {
+            configurator = new TestConfigureGrainContext(events);
+            return true;
+        }
+    }
+
+    private sealed class TestConfigureGrainContext(List<string> events) : IConfigureGrainContext
+    {
+        public void Configure(IGrainContext context) => events.Add("configure");
+    }
+
+    private sealed class TestGrainContextActivator(IGrainContext context, List<string> events) : IGrainContextActivator
+    {
+        public IGrainContext CreateContext(GrainAddress address, IConfigureGrainContext[] configureActions)
+        {
+            foreach (var configure in configureActions)
+            {
+                configure.Configure(context);
+            }
+
+            events.Add("activate");
+            return context;
+        }
+    }
+}

test/Orleans.Runtime.Tests/GrainActivatorTests.cs

@@ -1,12 +1,13 @@
+using System.Diagnostics.CodeAnalysis;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging.Abstractions;
+using Orleans.Metadata;
 using Orleans.Runtime;
 using Orleans.TestingHost;
 using TestExtensions;
 using UnitTests.GrainInterfaces;
 using UnitTests.Grains;
 using Xunit;
-using Orleans.Metadata;
 
 namespace UnitTests.General
 {
@@ -48,6 +49,7 @@ public void Configure(ISiloBuilder hostBuilder)
                         // Register our custom grain activator as a grain type component configurator
                         // This allows it to selectively apply to specific grain types
                         services.AddSingleton<IConfigureGrainTypeComponents, HardcodedGrainActivator>();
+                        services.AddSingleton<IConfigureGrainContextProvider, ActivationOrderingConfiguratorProvider>();
                     });
                 }
             }
@@ -96,6 +98,20 @@ public async Task CanUseCustomGrainActivatorToReleaseGrains()
             Assert.Equal(initialReleasedInstances + 1, finalReleasedInstances);
         }
 
+        [Fact, TestCategory("BVT")]
+        public async Task GrainContextIsConfiguredBeforeGrainConstruction()
+        {
+            var state = ActivationOrderingState.Instance;
+            state.Arm();
+            var grain = this.fixture.GrainFactory.GetGrain<ISimpleDIGrain>(
+                GetRandomGrainId(),
+                grainClassNamePrefix: "UnitTests.Grains.ExplicitlyRegistered");
+
+            await grain.GetStringValue();
+
+            Assert.True(state.WasConfiguredAtConstruction);
+        }
+
         /// <summary>
         /// Custom grain activator that bypasses dependency injection entirely.
         /// Implements both IGrainActivator (for creation/disposal) and IConfigureGrainTypeComponents
@@ -119,14 +135,16 @@ public void Configure(GrainType grainType, GrainProperties properties, GrainType
             {
                 // Selectively register this activator only for ExplicitlyRegisteredSimpleDIGrain types
                 // Other grain types will continue using the default DI-based activator
-                if (_grainClassMap.TryGetGrainClass(grainType, out var grainClass) && grainClass.IsAssignableFrom(typeof(ExplicitlyRegisteredSimpleDIGrain)))
+                if (_grainClassMap.TryGetGrainClass(grainType, out var grainClass) && grainClass == typeof(ExplicitlyRegisteredSimpleDIGrain))
                 {
                     shared.SetComponent<IGrainActivator>(this);
                 }
             }
 
             public object CreateInstance(IGrainContext context)
             {
+                ActivationOrderingState.Instance.ObserveConstruction(context);
+
                 // Custom instantiation logic - creates grain with hardcoded dependencies
                 // In real scenarios, this could get objects from a pool, perform complex
                 // initialization, or integrate with external systems
@@ -141,5 +159,67 @@ public ValueTask DisposeInstance(IGrainContext context, object instance)
                 return default;
             }
         }
+
+        private sealed class ActivationOrderingConfiguratorProvider(GrainClassMap grainClassMap) : IConfigureGrainContextProvider
+        {
+            public bool TryGetConfigurator(
+                GrainType grainType,
+                GrainProperties properties,
+                [NotNullWhen(true)] out IConfigureGrainContext? configurator)
+            {
+                if (grainClassMap.TryGetGrainClass(grainType, out var grainClass)
+                    && grainClass == typeof(ExplicitlyRegisteredSimpleDIGrain))
+                {
+                    configurator = ActivationOrderingState.Instance;
+                    return true;
+                }
+
+                configurator = null;
+                return false;
+            }
+        }
+
+        private sealed class ActivationOrderingState : IConfigureGrainContext
+        {
+            private int _armed;
+            private int _wasConfiguredAtConstruction;
+
+            public static ActivationOrderingState Instance { get; } = new();
+
+            public bool WasConfiguredAtConstruction => Volatile.Read(ref _wasConfiguredAtConstruction) != 0;
+
+            public void Arm()
+            {
+                Volatile.Write(ref _wasConfiguredAtConstruction, 0);
+                Volatile.Write(ref _armed, 1);
+            }
+
+            public void Configure(IGrainContext context)
+            {
+                if (Volatile.Read(ref _armed) == 0)
+                {
+                    return;
+                }
+
+                context.SetComponent(ConfiguredContextMarker.Instance);
+            }
+
+            public void ObserveConstruction(IGrainContext context)
+            {
+                if (Interlocked.Exchange(ref _armed, 0) == 0)
+                {
+                    return;
+                }
+
+                Volatile.Write(
+                    ref _wasConfiguredAtConstruction,
+                    context.GetComponent<ConfiguredContextMarker>() is not null ? 1 : 0);
+            }
+        }
+
+        private sealed class ConfiguredContextMarker
+        {
+            public static ConfiguredContextMarker Instance { get; } = new();
+        }
     }
 }