Fix abort recovery test event subscription race

Subscribe before sending the recovery prompt so an ephemeral session.idle cannot be lost. Exercise the shared abort scenario through the fake RPC server with events delivered before replies, retaining the original ordering and timeout budget.

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

Author: stephentoub

dotnet/test/E2E/SessionE2ETests.cs

@@ -349,11 +349,15 @@ await Assert.ThrowsAsync<IOException>(() =>
     [Fact]
     public async Task Should_Abort_A_Session()
     {
-        var session = await CreateSessionAsync();
+        await using var session = await CreateSessionAsync();
+        await AssertAbortAndRecoveryAsync(session, TimeSpan.FromSeconds(120));
+    }
 
+    internal static async Task AssertAbortAndRecoveryAsync(CopilotSession session, TimeSpan timeout)
+    {
         // Set up wait for tool execution to start BEFORE sending
-        var toolStartTask = TestHelper.GetNextEventOfTypeAsync<ToolExecutionStartEvent>(session);
-        var sessionIdleTask = TestHelper.GetNextEventOfTypeAsync<SessionIdleEvent>(session);
+        var toolStartTask = TestHelper.GetNextEventOfTypeAsync<ToolExecutionStartEvent>(session, timeout);
+        var sessionIdleTask = TestHelper.GetNextEventOfTypeAsync<SessionIdleEvent>(session, timeout);
 
         // Send a message that will take some time to process
         await session.SendAsync(new MessageOptions
@@ -375,8 +379,8 @@ await session.SendAsync(new MessageOptions
         // Verify an abort event exists in messages
         Assert.Contains(messages, m => m is AbortEvent);
 
-        await session.SendAsync(new MessageOptions { Prompt = "What is 2+2?" });
-        var recoveryMessage = await TestHelper.GetFinalAssistantMessageAsync(session);
+        // Subscribe before sending: session.idle is ephemeral and cannot be backfilled.
+        var recoveryMessage = await session.SendAndWaitAsync(new MessageOptions { Prompt = "What is 2+2?" }, timeout);
         Assert.NotNull(recoveryMessage);
         Assert.Contains("4", recoveryMessage.Data.Content ?? string.Empty);
     }

dotnet/test/Unit/ClientSessionLifetimeTests.cs

@@ -1717,6 +1717,87 @@ private static void AssertMessageSource(JsonElement request, string? source)
         Assert.False(request.TryGetProperty("wait", out _));
     }
 
+    [Theory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public async Task Abort_Recovery_Observes_Early_Events(bool recoveryCompletesBeforeReply)
+    {
+        await using var server = await FakeCopilotServer.StartAsync();
+        await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) });
+        await using var session = await client.CreateSessionAsync(new SessionConfig());
+        var timeout = TimeSpan.FromSeconds(5);
+        var sendCount = 0;
+        server.BeforeResponseAsync = async (request, cancellationToken) =>
+        {
+            if (request.Method == "session.send")
+            {
+                sendCount++;
+                await server.SendSessionEventAsync(session.SessionId, "user.message", new()
+                {
+                    ["content"] = request.Params.GetProperty("prompt").GetString()
+                });
+                if (sendCount == 1)
+                {
+                    await SendAndDrainAsync("tool.execution_start", new()
+                    {
+                        ["toolCallId"] = "slow-tool",
+                        ["toolName"] = "shell"
+                    }, cancellationToken);
+                }
+                else
+                {
+                    Assert.Equal(2, sendCount);
+                    await SendAndDrainAsync("assistant.message", new()
+                    {
+                        ["messageId"] = "recovery-message",
+                        ["content"] = "4"
+                    }, cancellationToken);
+                    if (recoveryCompletesBeforeReply)
+                    {
+                        await SendAndDrainAsync("session.idle", new(), cancellationToken);
+                    }
+                }
+            }
+            else if (request.Method == "session.abort")
+            {
+                Assert.Equal(1, sendCount);
+                await server.SendSessionEventAsync(session.SessionId, "abort", new()
+                {
+                    ["reason"] = "user"
+                });
+                await SendAndDrainAsync("session.idle", new() { ["aborted"] = true }, cancellationToken);
+            }
+        };
+        server.AfterResponseAsync = async (request, cancellationToken) =>
+        {
+            if (request.Method == "session.send" && sendCount == 2 && !recoveryCompletesBeforeReply)
+            {
+                await SendAndDrainAsync("session.idle", new(), cancellationToken);
+            }
+        };
+
+        // Exercise the E2E test's actual ordering and assertions, without launching a CLI.
+        await E2E.SessionE2ETests.AssertAbortAndRecoveryAsync(session, timeout);
+
+        Assert.Equal(
+            ["session.send", "session.abort", "session.send"],
+            server.Requests.Select(request => request.Method)
+                .Where(method => method is "session.send" or "session.abort"));
+        var history = await session.GetEventsAsync();
+        Assert.DoesNotContain(history, evt => evt is SessionIdleEvent);
+        Assert.Equal("4", Assert.Single(history.OfType<AssistantMessageEvent>()).Data.Content);
+
+        async Task SendAndDrainAsync(string type, Dictionary<string, object?> data, CancellationToken cancellationToken)
+        {
+            var drained = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+            using var subscription = session.On<SessionTitleChangedEvent>(_ => drained.TrySetResult());
+            await server.SendSessionEventAsync(session.SessionId, type, data);
+            // A later event is a fence: every subscriber has finished handling the target event.
+            await server.SendSessionEventAsync(session.SessionId, "session.title_changed", new() { ["title"] = "fence" });
+            await drained.Task.WaitAsync(timeout, cancellationToken);
+        }
+    }
+
     [Fact]
     public async Task SendAndWaitAsync_Skips_Autopilot_Continuation_Idle()
     {
@@ -2192,6 +2273,7 @@ private sealed class FakeCopilotServer : IAsyncDisposable
         private readonly TaskCompletionSource _allowDestroy = new(TaskCreationOptions.RunContinuationsAsynchronously);
         private readonly Task _serverTask;
         private readonly List<RpcRequestRecord> _requests = [];
+        private readonly ConcurrentQueue<object?> _sessionEvents = new();
         private readonly object _requestsLock = new();
         private readonly ConcurrentDictionary<int, TaskCompletionSource<JsonElement>> _pendingRequests = new();
         private NetworkStream? _stream;
@@ -2228,6 +2310,10 @@ public static Task<FakeCopilotServer> StartAsync()
 
         public int RuntimeShutdownCount { get; private set; }
 
+        public Func<RpcRequestRecord, CancellationToken, Task>? BeforeResponseAsync { get; set; }
+
+        public Func<RpcRequestRecord, CancellationToken, Task>? AfterResponseAsync { get; set; }
+
         public IReadOnlyList<RpcRequestRecord> Requests
         {
             get
@@ -2300,21 +2386,27 @@ public async Task<JsonElement> SendRequestAsync(string method, Dictionary<string
         public Task SendSessionEventAsync(string sessionId, string type, Dictionary<string, object?> data)
         {
             var stream = _stream ?? throw new InvalidOperationException("Client is not connected.");
+            var evt = new Dictionary<string, object?>
+            {
+                ["id"] = Guid.NewGuid().ToString(),
+                ["timestamp"] = DateTimeOffset.UtcNow.ToString("O"),
+                ["parentId"] = null,
+                ["type"] = type,
+                ["data"] = data
+            };
+            // Idle is ephemeral in the runtime and cannot be backfilled from history.
+            if (type != "session.idle")
+            {
+                _sessionEvents.Enqueue(evt);
+            }
             return WriteMessageAsync(stream, new Dictionary<string, object?>
             {
                 ["jsonrpc"] = "2.0",
                 ["method"] = "session.event",
                 ["params"] = new Dictionary<string, object?>
                 {
                     ["sessionId"] = sessionId,
-                    ["event"] = new Dictionary<string, object?>
-                    {
-                        ["id"] = Guid.NewGuid().ToString(),
-                        ["timestamp"] = DateTimeOffset.UtcNow.ToString("O"),
-                        ["parentId"] = null,
-                        ["type"] = type,
-                        ["data"] = data
-                    }
+                    ["event"] = evt
                 }
             }, _cts.Token);
         }
@@ -2437,6 +2529,11 @@ private async Task HandleRequestAsync(Stream stream, JsonElement request, Cancel
                 }, cancellationToken);
                 return;
             }
+            var requestRecord = new RpcRequestRecord(method!, paramsElement);
+            if (BeforeResponseAsync is { } beforeResponse)
+            {
+                await beforeResponse(requestRecord, cancellationToken);
+            }
             object? result = method switch
             {
                 "connect" => new Dictionary<string, object?>
@@ -2455,6 +2552,11 @@ private async Task HandleRequestAsync(Stream stream, JsonElement request, Cancel
                 {
                     ["messageId"] = "message-1"
                 },
+                "session.abort" => new Dictionary<string, object?>(),
+                "session.getMessages" => new Dictionary<string, object?>
+                {
+                    ["events"] = _sessionEvents.ToArray()
+                },
                 "session.options.update" => new Dictionary<string, object?>
                 {
                     ["success"] = true
@@ -2495,6 +2597,10 @@ private async Task HandleRequestAsync(Stream stream, JsonElement request, Cancel
                 ["id"] = id,
                 ["result"] = result
             }, cancellationToken);
+            if (AfterResponseAsync is { } afterResponse)
+            {
+                await afterResponse(requestRecord, cancellationToken);
+            }
         }
 
         private Dictionary<string, object?> CreateSessionResult(JsonElement request)