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)