Preserve session routing after failed resume

Capture the wrapper displaced before session.resume and restore it only if the failed replacement is still registered.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: c71188a1-1445-46aa-9faf-3b73cf6a6dd9

Author: roji

dotnet/src/Client.cs

@@ -774,8 +774,10 @@ private CopilotSession InitializeSession(
         Dictionary<string, Func<string, Task<string>>>? transformCallbacks,
         bool hasHooks,
         string callerName,
-        bool replaceExisting = false)
+        bool replaceExisting,
+        out CopilotSession? replacedSession)
     {
+        replacedSession = null;
         var setupTimestamp = Stopwatch.GetTimestamp();
         var session = new CopilotSession(
             sessionId,
@@ -810,7 +812,16 @@ private CopilotSession InitializeSession(
         session.RegisterBearerTokenProviders(BuildBearerTokenCallbacks(config));
         if (replaceExisting)
         {
-            _sessions[session.SessionId] = session;
+            CopilotSession? displacedSession = null;
+            _sessions.AddOrUpdate(
+                session.SessionId,
+                session,
+                (_, current) =>
+                {
+                    displacedSession = current;
+                    return session;
+                });
+            replacedSession = displacedSession;
         }
         else if (!_sessions.TryAdd(session.SessionId, session))
         {
@@ -1128,7 +1139,9 @@ public async Task<CopilotSession> CreateSessionAsync(SessionConfig config, Cance
                 config,
                 transformCallbacks,
                 hasHooks,
-                "CopilotClient.CreateSessionAsync");
+                "CopilotClient.CreateSessionAsync",
+                replaceExisting: false,
+                out _);
         }
         try
         {
@@ -1228,7 +1241,9 @@ public async Task<CopilotSession> CreateSessionAsync(SessionConfig config, Cance
                         config,
                         transformCallbacks,
                         hasHooks,
-                        "CopilotClient.CreateSessionAsync");
+                        "CopilotClient.CreateSessionAsync",
+                        replaceExisting: false,
+                        out _);
                 }
             };
 
@@ -1344,7 +1359,8 @@ public async Task<CopilotSession> ResumeSessionAsync(string sessionId, ResumeSes
             transformCallbacks,
             hasHooks,
             "CopilotClient.ResumeSessionAsync",
-            replaceExisting: true);
+            replaceExisting: true,
+            out var previousSession);
         try
         {
             var (traceparent, tracestate) = TelemetryHelpers.GetTraceContext();
@@ -1443,7 +1459,15 @@ public async Task<CopilotSession> ResumeSessionAsync(string sessionId, ResumeSes
         }
         catch (Exception ex)
         {
-            session.RemoveFromClient();
+            if (previousSession is null)
+            {
+                session.RemoveFromClient();
+            }
+            else
+            {
+                _sessions.TryUpdate(sessionId, previousSession, session);
+            }
+
             if (ex is not OperationCanceledException)
             {
                 LoggingHelpers.LogTiming(_logger, LogLevel.Warning, ex,

dotnet/test/Unit/ClientSessionLifetimeTests.cs

@@ -216,7 +216,7 @@ public async Task ResumeSessionAsync_Replaces_Session_Tracked_By_Same_Client()
     }
 
     [Fact]
-    public async Task Failed_ResumeSessionAsync_Removes_Replacement_Registration()
+    public async Task Failed_ResumeSessionAsync_Restores_Previous_Registration()
     {
         await using var server = await FakeCopilotServer.StartAsync();
         await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) });
@@ -234,8 +234,8 @@ public async Task Failed_ResumeSessionAsync_Removes_Replacement_Registration()
             OnPermissionRequest = PermissionHandler.ApproveAll
         }));
 
-        AssertSessionCount(client, sessions: 0);
-        Assert.Null(GetTrackedSession(client, sessionId));
+        AssertSessionCount(client, sessions: 1);
+        Assert.Same(session, GetTrackedSession(client, sessionId));
         Assert.Equal("message-1", await session.SendAsync("The original session remains active."));
 
         await session.DisposeAsync();