From daedce3ca907ddb461a40d193a6412f70c52fa7f Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Fri, 11 Sep 2026 21:04:50 -0400 Subject: [PATCH] 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> --- dotnet/test/E2E/SessionE2ETests.cs | 14 +- .../test/Unit/ClientSessionLifetimeTests.cs | 122 ++++++++++++++++-- 2 files changed, 123 insertions(+), 13 deletions(-) diff --git a/dotnet/test/E2E/SessionE2ETests.cs b/dotnet/test/E2E/SessionE2ETests.cs index 288ba1e65f..08781ce209 100644 --- a/dotnet/test/E2E/SessionE2ETests.cs +++ b/dotnet/test/E2E/SessionE2ETests.cs @@ -349,11 +349,15 @@ await Assert.ThrowsAsync(() => [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(session); - var sessionIdleTask = TestHelper.GetNextEventOfTypeAsync(session); + var toolStartTask = TestHelper.GetNextEventOfTypeAsync(session, timeout); + var sessionIdleTask = TestHelper.GetNextEventOfTypeAsync(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); } diff --git a/dotnet/test/Unit/ClientSessionLifetimeTests.cs b/dotnet/test/Unit/ClientSessionLifetimeTests.cs index 173569788c..dd7fdc2bbb 100644 --- a/dotnet/test/Unit/ClientSessionLifetimeTests.cs +++ b/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()).Data.Content); + + async Task SendAndDrainAsync(string type, Dictionary data, CancellationToken cancellationToken) + { + var drained = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var subscription = session.On(_ => 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 _requests = []; + private readonly ConcurrentQueue _sessionEvents = new(); private readonly object _requestsLock = new(); private readonly ConcurrentDictionary> _pendingRequests = new(); private NetworkStream? _stream; @@ -2228,6 +2310,10 @@ public static Task StartAsync() public int RuntimeShutdownCount { get; private set; } + public Func? BeforeResponseAsync { get; set; } + + public Func? AfterResponseAsync { get; set; } + public IReadOnlyList Requests { get @@ -2300,6 +2386,19 @@ public async Task SendRequestAsync(string method, Dictionary data) { var stream = _stream ?? throw new InvalidOperationException("Client is not connected."); + var evt = new Dictionary + { + ["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 { ["jsonrpc"] = "2.0", @@ -2307,14 +2406,7 @@ public Task SendSessionEventAsync(string sessionId, string type, Dictionary { ["sessionId"] = sessionId, - ["event"] = new Dictionary - { - ["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 @@ -2455,6 +2552,11 @@ private async Task HandleRequestAsync(Stream stream, JsonElement request, Cancel { ["messageId"] = "message-1" }, + "session.abort" => new Dictionary(), + "session.getMessages" => new Dictionary + { + ["events"] = _sessionEvents.ToArray() + }, "session.options.update" => new Dictionary { ["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 CreateSessionResult(JsonElement request)