fixes and stablization fror test

Author: KSemenenko

AGENTS.md

@@ -109,6 +109,7 @@ If no new rule is detected -> do not update the file.
   - run broader suites only when you have something real to verify (avoid re-running the same command without changes)
   - run coverage once per change (it is heavier than tests)
 - Run tests in layers: new/changed -> related suite -> broader regressions
+- Final verification for code changes requires the full `dotnet test -c Debug` suite to pass; targeted test runs are only for iteration, not for closing the task
 - After tests pass: run format
 - After format: run build (final check)
 - Summarize changes and test results before marking complete
@@ -185,6 +186,7 @@ If no new rule is detected -> do not update the file.
 - Avoid `this.` qualification; prefer predefined types (`int` over `Int32`)
 - Never use `ConfigureAwait(false)`
 - No magic literals - extract to constants, enums, config
+- In concurrency-sensitive paths like `OrleansHubLifetimeManager`, prefer the smallest behavior-preserving fix; avoid widening the async/concurrency shape unless tests prove it is necessary, because this code is easy to make unsafe
 
 ### Diagnostics
 

ManagedCode.Orleans.SignalR.Client/Extensions/OrleansHubGroupExtensions.cs

@@ -34,6 +34,13 @@ private static IOrleansGroupManager<THub> ResolveGroupManager<THub>(THub hub) wh
     {
         var serviceProvider = (hub.Context.GetHttpContext()?.RequestServices) ?? throw new InvalidOperationException("Unable to resolve SignalR services for the current hub connection.");
 
-        return serviceProvider.GetRequiredService<IOrleansGroupManager<THub>>();
+        var groupManager = serviceProvider.GetService<IOrleansGroupManager<THub>>();
+        if (groupManager is not null)
+        {
+            return groupManager;
+        }
+
+        var lifetimeManager = serviceProvider.GetRequiredService<HubLifetimeManager<THub>>();
+        return new OrleansGroupManager<THub>(lifetimeManager);
     }
 }

ManagedCode.Orleans.SignalR.Core/SignalR/OrleansHubLifetimeManager.cs

@@ -371,15 +371,16 @@ public async Task AddToGroupsAsync(string connectionId, IReadOnlyList<string> gr
         if (_orleansSignalOptions.Value.GroupPartitionCount > 1)
         {
             var coordinatorGrain = NameHelperGenerator.GetGroupCoordinatorGrain<THub>(_clusterClient);
-            var partitionIds = await Task.Run(
-                () => coordinatorGrain.AddConnectionToGroups(uniqueGroupNames, connectionId, subscription.Reference),
-                cancellationToken);
-
-            foreach (var partitionId in partitionIds)
+            foreach (var groupName in uniqueGroupNames)
             {
+                var partitionId = await Task.Run(() => coordinatorGrain.GetPartitionForGroup(groupName), cancellationToken);
                 var partitionGrain = NameHelperGenerator.GetGroupPartitionGrain<THub>(_clusterClient, partitionId);
                 subscription.AddGrain(partitionGrain);
             }
+
+            await Task.Run(
+                () => coordinatorGrain.AddConnectionToGroups(uniqueGroupNames, connectionId, subscription.Reference),
+                cancellationToken);
         }
         else
         {

ManagedCode.Orleans.SignalR.Server/Extensions/OrleansDependencyInjectionExtensions.cs

@@ -24,6 +24,7 @@ public static ISignalRServerBuilder AddOrleans(this ISignalRServerBuilder signal
         signalrBuilder.Services.AddOptions<OrleansSignalROptions>().Configure(options);
         signalrBuilder.Services.AddSingleton(typeof(HubLifetimeManager<>), typeof(OrleansHubLifetimeManager<>));
         signalrBuilder.Services.AddSingleton(typeof(IOrleansHubContext<,>), typeof(OrleansHubContext<,>));
+        signalrBuilder.Services.AddSingleton(typeof(IOrleansGroupManager<>), typeof(OrleansGroupManager<>));
 
         return signalrBuilder;
     }

ManagedCode.Orleans.SignalR.Server/SignalRGroupCoordinatorGrain.cs

@@ -408,7 +408,7 @@ private Dictionary<int, List<string>> GetPartitionsForGroups(string[] groupNames
             }
             else
             {
-                continue;
+                partition = PartitionHelper.GetPartitionId(groupName, (uint)_currentPartitionCount);
             }
 
             ref var list = ref CollectionsMarshal.GetValueRefOrAddDefault(groupsByPartition, partition, out var exists);

ManagedCode.Orleans.SignalR.Tests/GrainPersistenceTests.cs

@@ -224,6 +224,37 @@ public async Task GroupCoordinatorBatchMethodsShouldTrackMembershipAcrossTouched
         }
     }
 
+    [Fact]
+    public async Task GroupCoordinatorBatchRemoveShouldCleanupPartitionWhenAssignmentMetadataIsMissingAsync()
+    {
+        var client = _cluster.Cluster.Client;
+        var coordinator = NameHelperGenerator.GetGroupCoordinatorGrain<SimpleTestHub>(client);
+        var connectionId = $"group-cleanup-{Guid.NewGuid():N}";
+        var groupName = $"group-drift-{Guid.NewGuid():N}";
+        var observer = client.CreateObjectReference<ISignalRObserver>(new SignalRObserver(_ => Task.CompletedTask));
+
+        var partitionId = await coordinator.GetPartitionForGroup(groupName);
+        var partition = NameHelperGenerator.GetGroupPartitionGrain<SimpleTestHub>(client, partitionId);
+
+        try
+        {
+            var addedPartitions = await coordinator.AddConnectionToGroups([groupName], connectionId, observer);
+            addedPartitions.ShouldContain(partitionId);
+            (await partition.HasConnection(connectionId)).ShouldBeTrue();
+
+            await coordinator.NotifyGroupRemoved(groupName);
+            (await partition.HasConnection(connectionId)).ShouldBeTrue();
+
+            var removedPartitions = await coordinator.RemoveConnectionFromGroups([groupName], connectionId, observer);
+            removedPartitions.ShouldContain(partitionId);
+            (await partition.HasConnection(connectionId)).ShouldBeFalse();
+        }
+        finally
+        {
+            await coordinator.RemoveConnectionFromGroups([groupName], connectionId, observer);
+        }
+    }
+
     private static async Task AssertRoutedAsync(Func<Task<bool>> sendAction, string reason)
     {
         var deadline = DateTime.UtcNow + TimeSpan.FromSeconds(5);

ManagedCode.Orleans.SignalR.Tests/HubSmokeTests.cs

@@ -89,6 +89,52 @@ await WaitUntilAsync(() =>
         await DisposeAsync(connection1, connection2);
     }
 
+    [Fact]
+    public async Task BatchGroupHelpersShouldWorkWithoutOrleansRegistrationAsync()
+    {
+        using var app = new TestWebApplication(_cluster, 8089, useOrleans: false, loggerAccessor: _loggerAccessor);
+        var groupNames = new[]
+        {
+            $"plain-group-alpha-{Guid.NewGuid():N}",
+            $"plain-group-beta-{Guid.NewGuid():N}"
+        };
+
+        var received = new ConcurrentQueue<string>();
+        var member = await StartConnectionAsync(app, received.Enqueue, null, _output);
+        var sender = await StartConnectionAsync(app, _ => { }, null, _output);
+
+        try
+        {
+            await member.InvokeAsync("AddToGroups", groupNames);
+
+            foreach (var groupName in groupNames)
+            {
+                await sender.InvokeAsync("GroupSendAsync", groupName, $"before:{groupName}");
+            }
+
+            await WaitUntilAsync(() => received.Count >= groupNames.Length, TimeSpan.FromSeconds(10));
+
+            received.Count.ShouldBe(groupNames.Length);
+
+            await member.InvokeAsync("RemoveFromGroups", groupNames);
+            while (received.TryDequeue(out _))
+            {
+            }
+
+            foreach (var groupName in groupNames)
+            {
+                await sender.InvokeAsync("GroupSendAsync", groupName, $"after:{groupName}");
+            }
+
+            await Task.Delay(500);
+            received.ShouldBeEmpty();
+        }
+        finally
+        {
+            await DisposeAsync(member, sender);
+        }
+    }
+
     [Fact]
     public async Task UserMessageIsDeliveredToSpecificUserAsync()
     {

ManagedCode.Orleans.SignalR.Tests/PartitioningTests.cs

@@ -425,6 +425,105 @@ public async Task BatchGroupMembershipShouldAddAndRemoveAcrossMultipleGroupsAsyn
         }
     }
 
+    [Fact]
+    public async Task BatchGroupMembershipShouldCleanupTouchedPartitionsWhenDisconnectHappensMidJoinAsync()
+    {
+        var groupNames = Enumerable.Range(0, 512)
+            .Select(index => $"batch-disconnect-group-{index}-{Guid.NewGuid():N}")
+            .ToArray();
+
+        var connection = _apps[0].CreateSignalRClient(nameof(SimpleTestHub));
+
+        try
+        {
+            await connection.StartAsync();
+            var connectionId = connection.ConnectionId ?? throw new InvalidOperationException("ConnectionId was not initialized.");
+
+            var coordinator = NameHelperGenerator.GetGroupCoordinatorGrain<SimpleTestHub>(_siloCluster.Cluster.Client);
+            var partitionIds = (await Task.WhenAll(groupNames.Select(coordinator.GetPartitionForGroup)))
+                .Distinct()
+                .ToArray();
+            partitionIds.Length.ShouldBeGreaterThan(0);
+
+            var partitions = partitionIds
+                .Select(partitionId => NameHelperGenerator.GetGroupPartitionGrain<SimpleTestHub>(_siloCluster.Cluster.Client, partitionId))
+                .ToArray();
+
+            var joinTask = connection.InvokeAsync("AddToGroups", groupNames);
+
+            var joinStarted = await WaitForBatchJoinToStartAsync(partitions, connectionId, joinTask, TimeSpan.FromSeconds(3));
+
+            joinStarted.ShouldBeTrue();
+
+            await connection.StopAsync();
+
+            try
+            {
+                await joinTask;
+            }
+            catch (Exception ex) when (ex is HubException or InvalidOperationException or TaskCanceledException)
+            {
+                _testOutputHelper.WriteLine($"Join invocation ended after disconnect: {ex.GetType().Name}: {ex.Message}");
+            }
+
+            var released = await WaitUntilAsync(
+                "all touched partitions to release disconnected connection after batch join",
+                async () => await GetTrackedPartitionCountAsync(partitions, connectionId) == 0,
+                progress: async () =>
+                    $"tracked={await GetTrackedPartitionCountAsync(partitions, connectionId)}/{partitions.Length}",
+                timeout: TimeSpan.FromSeconds(15));
+
+            released.ShouldBeTrue();
+        }
+        finally
+        {
+            await connection.DisposeAsync();
+        }
+    }
+
+    private static async Task<int> GetTrackedPartitionCountAsync(
+        IReadOnlyCollection<Core.Interfaces.ISignalRGroupPartitionGrain> partitions,
+        string connectionId)
+    {
+        var tracked = 0;
+
+        foreach (var partition in partitions)
+        {
+            if (await partition.HasConnection(connectionId))
+            {
+                tracked++;
+            }
+        }
+
+        return tracked;
+    }
+
+    private static async Task<bool> WaitForBatchJoinToStartAsync(
+        IReadOnlyCollection<Core.Interfaces.ISignalRGroupPartitionGrain> partitions,
+        string connectionId,
+        Task joinTask,
+        TimeSpan timeout)
+    {
+        var deadline = DateTime.UtcNow + timeout;
+
+        while (DateTime.UtcNow < deadline)
+        {
+            if (await GetTrackedPartitionCountAsync(partitions, connectionId) > 0 && !joinTask.IsCompleted)
+            {
+                return true;
+            }
+
+            if (joinTask.IsCompleted)
+            {
+                return false;
+            }
+
+            await Task.Delay(TimeSpan.FromMilliseconds(5));
+        }
+
+        return await GetTrackedPartitionCountAsync(partitions, connectionId) > 0 && !joinTask.IsCompleted;
+    }
+
     private async Task<bool> WaitUntilAsync(
         string description,
         Func<Task<bool>> condition,

docs/Features/Group-Partitioning.md

@@ -23,6 +23,9 @@ Group operations are routed through a coordinator grain that assigns group names
 - [x] Add package-level batch group APIs for hub and host-service callers.
 - [x] Batch coordinator and partition updates so one request does not force sequential writes per group.
 - [x] Cover batch add/remove with integration tests and direct coordinator verification.
+- [x] Restore disconnect-safe subscription tracking for partitioned batch adds.
+- [x] Restore the hashed cleanup fallback when batch removals see missing coordinator assignments.
+- [x] Add regression tests for disconnect cleanup and degraded coordinator state.
 
 ## Main flow
 
@@ -44,6 +47,8 @@ flowchart TD
 - Partition grains hold `group -> connection -> observer` mappings and emit fan-out to observers.
 - Empty groups trigger cleanup so partitions can shed state.
 - Batch membership operations collapse repeated coordinator writes into one persistence step per request and one partition write per touched partition.
+- Partitioned batch adds pre-register touched partitions in the connection subscription before the coordinator write finishes so disconnect cleanup still reaches every touched partition.
+- Batch removals fall back to the hashed partition when coordinator assignment metadata is missing, which preserves cleanup in degraded-state recovery scenarios.
 
 ## Configuration knobs
 

docs/Features/Hub-Lifetime-Manager-Integration.md

@@ -19,6 +19,8 @@ The ASP.NET Core host swaps the default SignalR hub lifetime manager with `Orlea
 - [x] Restore fire-and-forget fan-out for multi-group and multi-user sends.
 - [x] Ensure per-target send failures are logged without blocking hub execution.
 - [x] Route package-specific batch group membership calls through the Orleans lifetime manager.
+- [x] Keep batch group helper calls usable when the host is running plain `AddSignalR()`.
+- [x] Add regression coverage for the batch helper path without Orleans registration.
 
 ## Main flow
 
@@ -38,6 +40,7 @@ flowchart TD
 - `AddOrleans()` registers `OrleansHubLifetimeManager<THub>` as the `HubLifetimeManager` implementation.
 - The lifetime manager creates a per-connection `Subscription` and registers observers with connection/group/user grains.
 - Package-specific batch group operations (`AddToGroupsAsync` / `RemoveFromGroupsAsync`) also route through the lifetime manager instead of looping over sequential single-group writes.
+- Hub batch helpers fall back to the registered `HubLifetimeManager<THub>` when `IOrleansGroupManager<THub>` is not explicitly registered, so the API still works on plain `AddSignalR()` hosts.
 - Detailed batching behavior and partition persistence rules live in `docs/Features/Group-Partitioning.md`.
 
 ## Configuration knobs