diff --git a/dotnet/src/Session.cs b/dotnet/src/Session.cs index 04306b7f62..92174190a3 100644 --- a/dotnet/src/Session.cs +++ b/dotnet/src/Session.cs @@ -81,13 +81,18 @@ private sealed record EventSubscription(Type EventType, Action Han private IReadOnlyList _openCanvases = Array.Empty(); private int _isDisposed; + private long _eventEnqueueVersion; + + private abstract record EventDispatchItem; + private sealed record EventItem(SessionEvent Event) : EventDispatchItem; + private sealed record EventBarrier(TaskCompletionSource Completion) : EventDispatchItem; /// /// Channel that serializes event dispatch. enqueues; /// a single background consumer () dequeues and /// invokes handlers one at a time, preserving arrival order. /// - private readonly Channel _eventChannel = Channel.CreateUnbounded( + private readonly Channel _eventChannel = Channel.CreateUnbounded( new() { SingleReader = true }); /// @@ -336,7 +341,8 @@ public async Task SendAsync(MessageOptions options, CancellationToken ca var totalTimestamp = Stopwatch.GetTimestamp(); var effectiveTimeout = timeout ?? TimeSpan.FromSeconds(60); - var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var tcs = new TaskCompletionSource<(AssistantMessageEvent? Message, string CompletedBy)>(TaskCreationOptions.RunContinuationsAsynchronously); + var completionCandidate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); AssistantMessageEvent? lastAssistantMessage = null; var firstAssistantMessageLogged = false; @@ -346,6 +352,7 @@ void Handler(SessionEvent evt) { case AssistantMessageEvent assistantMessage: lastAssistantMessage = assistantMessage; + completionCandidate.TrySetResult(true); if (!firstAssistantMessageLogged) { firstAssistantMessageLogged = true; @@ -356,12 +363,16 @@ void Handler(SessionEvent evt) } break; + case AssistantTurnEndEvent: + completionCandidate.TrySetResult(true); + break; + case SessionIdleEvent: LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null, "CopilotSession.SendAndWaitAsync idle received. Elapsed={Elapsed}, SessionId={SessionId}", totalTimestamp, SessionId); - tcs.TrySetResult(lastAssistantMessage); + tcs.TrySetResult((lastAssistantMessage, "idle")); break; case SessionErrorEvent errorEvent: @@ -371,6 +382,77 @@ void Handler(SessionEvent evt) } } + async Task MonitorRuntimeCompletionAsync(CancellationToken waitCancellationToken) + { + try + { + var candidateOrCompletion = await Task.WhenAny(completionCandidate.Task, tcs.Task).ConfigureAwait(false); + if (candidateOrCompletion != completionCandidate.Task) + { + return; + } + + // In the normal path session.idle follows the assistant events in the + // same notification burst. Give it a brief chance to arrive so the + // fallback adds no RPC traffic to healthy turns. + var graceDelay = Task.Delay(TimeSpan.FromMilliseconds(250), waitCancellationToken); + if (await Task.WhenAny(graceDelay, tcs.Task).ConfigureAwait(false) != graceDelay) + { + return; + } + + while (!tcs.Task.IsCompleted) + { + // The runtime owns the authoritative view of background tasks and + // follow-up turns. This RPC returns only after those have settled. + await Rpc.Tasks.WaitForPendingAsync(waitCancellationToken).ConfigureAwait(false); + + var activity = await Rpc.Metadata.ActivityAsync(waitCancellationToken).ConfigureAwait(false); + if (!activity.HasActiveWork) + { + // RPC responses and session.event notifications share one ordered + // connection, but user handlers run on a separate FIFO channel. + // Flush that channel before confirming the runtime is still idle. + await FlushEventDispatchAsync(waitCancellationToken).ConfigureAwait(false); + activity = await Rpc.Metadata.ActivityAsync(waitCancellationToken).ConfigureAwait(false); + if (!activity.HasActiveWork) + { + var stableEventVersion = Volatile.Read(ref _eventEnqueueVersion); + await FlushEventDispatchAsync(waitCancellationToken).ConfigureAwait(false); + if (stableEventVersion != Volatile.Read(ref _eventEnqueueVersion)) + { + continue; + } + + activity = await Rpc.Metadata.ActivityAsync(waitCancellationToken).ConfigureAwait(false); + if (!activity.HasActiveWork + && stableEventVersion == Volatile.Read(ref _eventEnqueueVersion)) + { + tcs.TrySetResult((lastAssistantMessage, "activity")); + return; + } + } + } + + await Task.Delay(TimeSpan.FromMilliseconds(100), waitCancellationToken).ConfigureAwait(false); + } + } + catch (RemoteRpcException ex) when (ex.ErrorCode == RemoteRpcException.MethodNotFoundErrorCode) + { + // Older runtimes may not expose activity/task-drain RPCs. Preserve the + // existing event-only behavior and let session.idle or the timeout win. + LogRuntimeCompletionFallbackUnavailable(ex, SessionId); + } + catch (OperationCanceledException) when (waitCancellationToken.IsCancellationRequested) + { + // The timeout/caller-cancellation registration completes tcs. + } + catch (Exception ex) when (ex is RemoteRpcException or IOException or ObjectDisposedException or JsonException) + { + tcs.TrySetException(ex); + } + } + using var subscription = On(Handler); await SendAsync(options, cancellationToken); @@ -385,16 +467,17 @@ void Handler(SessionEvent evt) else tcs.TrySetException(new TimeoutException($"SendAndWaitAsync timed out after {effectiveTimeout}")); }); + var runtimeCompletionTask = MonitorRuntimeCompletionAsync(cts.Token); try { - var result = await tcs.Task; + var completion = await tcs.Task; LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null, "CopilotSession.SendAndWaitAsync complete. Elapsed={Elapsed}, SessionId={SessionId}, CompletedBy={CompletedBy}, AssistantMessageReceived={AssistantMessageReceived}", totalTimestamp, SessionId, - "idle", - result is not null); - return result; + completion.CompletedBy, + completion.Message is not null); + return completion.Message; } catch (Exception ex) when (ex is TimeoutException) { @@ -414,6 +497,11 @@ void Handler(SessionEvent evt) "error"); throw; } + finally + { + cts.Cancel(); + await runtimeCompletionTask.ConfigureAwait(false); + } } /// @@ -485,8 +573,27 @@ internal void DispatchEvent(SessionEvent sessionEvent) // never completes (multi-client permission scenario). _ = HandleBroadcastEventAsync(sessionEvent); - // Queue the event for serial processing by user handlers. - _eventChannel.Writer.TryWrite(sessionEvent); + // Queue the event for serial processing by user handlers. Increment first so + // a completion monitor can detect an event even if this thread is preempted + // before the channel write. + Interlocked.Increment(ref _eventEnqueueVersion); + _eventChannel.Writer.TryWrite(new EventItem(sessionEvent)); + } + + /// + /// Waits until every event already queued for this session has been delivered to + /// user handlers. Events arriving after the barrier remain queued behind it. + /// + private async Task FlushEventDispatchAsync(CancellationToken cancellationToken) + { + var completion = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var queued = _eventChannel.Writer.TryWrite(new EventBarrier(completion)); + ObjectDisposedException.ThrowIf(!queued, this); + + using var registration = cancellationToken.Register( + static state => ((TaskCompletionSource)state!).TrySetCanceled(), + completion); + await completion.Task.ConfigureAwait(false); } /// @@ -495,8 +602,15 @@ internal void DispatchEvent(SessionEvent sessionEvent) /// private async Task ProcessEventsAsync() { - await foreach (var sessionEvent in _eventChannel.Reader.ReadAllAsync()) + await foreach (var item in _eventChannel.Reader.ReadAllAsync()) { + if (item is EventBarrier barrier) + { + barrier.Completion.TrySetResult(true); + continue; + } + + var sessionEvent = ((EventItem)item).Event; var dispatchTimestamp = Stopwatch.GetTimestamp(); var eventType = sessionEvent.GetType(); foreach (var subscription in _eventHandlers) @@ -1935,6 +2049,9 @@ await InvokeRpcAsync( [LoggerMessage(Level = LogLevel.Error, Message = "Unhandled exception in session event handler")] private partial void LogEventHandlerError(Exception exception); + [LoggerMessage(Level = LogLevel.Debug, Message = "Runtime completion fallback unavailable. SessionId={sessionId}")] + private partial void LogRuntimeCompletionFallbackUnavailable(Exception exception, string sessionId); + [LoggerMessage(Level = LogLevel.Debug, Message = "Failed to fetch tool metadata for {toolName}")] private partial void LogToolMetadataFetchFailed(Exception exception, string toolName); diff --git a/dotnet/test/E2E/SendAndWaitReliabilityE2ETests.cs b/dotnet/test/E2E/SendAndWaitReliabilityE2ETests.cs new file mode 100644 index 0000000000..5a482c5b33 --- /dev/null +++ b/dotnet/test/E2E/SendAndWaitReliabilityE2ETests.cs @@ -0,0 +1,84 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +using System.Collections; +using System.Reflection; +using Xunit; +using Xunit.Abstractions; + +namespace GitHub.Copilot.Test.E2E; + +/// +/// Verifies that .NET completion remains reliable when the runtime's ephemeral +/// session.idle notification does not reach the SendAndWait handler. +/// +public class SendAndWaitReliabilityE2ETests(E2ETestFixture fixture, ITestOutputHelper output) + : E2ETestBase(fixture, "send_and_wait_reliability", output) +{ + [Fact] + public async Task Should_Complete_When_Live_SessionIdle_Is_Dropped() + { + await using var session = await CreateSessionAsync(); + var userMessageDispatched = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseUserMessage = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var idleObserved = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + using var userSubscription = session.On(_ => + { + userMessageDispatched.TrySetResult(); + releaseUserMessage.Task.GetAwaiter().GetResult(); + }); + using var idleSubscription = session.On(_ => idleObserved.TrySetResult()); + + var completionTask = session.SendAndWaitAsync( + new MessageOptions { Prompt = "What is 5+5? Reply with just the number." }, + timeout: TimeSpan.FromSeconds(30)); + + try + { + await userMessageDispatched.Task.WaitAsync(TimeSpan.FromSeconds(10)); + SuppressIdleForSendAndWaitHandler(session); + } + finally + { + releaseUserMessage.TrySetResult(); + } + + var response = await completionTask; + await idleObserved.Task.WaitAsync(TimeSpan.FromSeconds(10)); + + Assert.NotNull(response); + Assert.Contains("10", response.Data.Content); + } + + private static void SuppressIdleForSendAndWaitHandler(CopilotSession session) + { + var handlersField = typeof(CopilotSession).GetField( + "_eventHandlers", + BindingFlags.Instance | BindingFlags.NonPublic) + ?? throw new InvalidOperationException("CopilotSession._eventHandlers was not found."); + var handlers = (IEnumerable)(handlersField.GetValue(session) + ?? throw new InvalidOperationException("CopilotSession._eventHandlers was null.")); + + var sendAndWaitSubscription = handlers.Cast().Last(subscription => + { + var eventType = subscription.GetType().GetProperty("EventType")?.GetValue(subscription); + return Equals(eventType, typeof(SessionEvent)); + }); + var handlerField = sendAndWaitSubscription.GetType().GetField( + "k__BackingField", + BindingFlags.Instance | BindingFlags.NonPublic) + ?? throw new InvalidOperationException("EventSubscription.Handler backing field was not found."); + var original = (Action)(handlerField.GetValue(sendAndWaitSubscription) + ?? throw new InvalidOperationException("SendAndWait event handler was null.")); + + handlerField.SetValue(sendAndWaitSubscription, (Action)(evt => + { + if (evt is not SessionIdleEvent) + { + original(evt); + } + })); + } +} diff --git a/dotnet/test/Unit/ClientSessionLifetimeTests.cs b/dotnet/test/Unit/ClientSessionLifetimeTests.cs index f736e8dee4..fe59411aa2 100644 --- a/dotnet/test/Unit/ClientSessionLifetimeTests.cs +++ b/dotnet/test/Unit/ClientSessionLifetimeTests.cs @@ -364,6 +364,107 @@ public async Task Generated_Session_Rpc_Throws_When_Session_Disposed() await Assert.ThrowsAsync(() => session.Rpc.Model.GetCurrentAsync()); } + [Fact] + public async Task SendAndWaitAsync_Completes_When_SessionIdle_Notification_Is_Dropped() + { + await using var server = await FakeCopilotServer.StartAsync(); + server.ConfigureDroppedIdleCompletion(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + await using var session = await client.CreateSessionAsync(new SessionConfig + { + OnPermissionRequest = PermissionHandler.ApproveAll + }); + + var response = await session.SendAndWaitAsync( + new MessageOptions { Prompt = "complete without idle" }, + timeout: TimeSpan.FromSeconds(2)); + + Assert.NotNull(response); + Assert.Equal("completed response", response.Data.Content); + Assert.Contains(server.Requests, request => request.Method == "session.metadata.activity"); + } + + [Fact] + public async Task SendAndWaitAsync_DroppedIdle_Fallback_Flushes_Preceding_Event_Handlers() + { + await using var server = await FakeCopilotServer.StartAsync(); + server.ConfigureDroppedIdleCompletion(delayEvents: true); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + await using var session = await client.CreateSessionAsync(new SessionConfig + { + OnPermissionRequest = PermissionHandler.ApproveAll + }); + + var completionTask = session.SendAndWaitAsync( + new MessageOptions { Prompt = "flush handlers before completion" }, + timeout: TimeSpan.FromSeconds(5)); + + var handlerStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseHandler = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var subscription = session.On(_ => + { + handlerStarted.TrySetResult(); + releaseHandler.Task.GetAwaiter().GetResult(); + }); + + server.ReleaseCompletionEvents(); + await handlerStarted.Task.WaitAsync(TimeSpan.FromSeconds(2)); + await Task.Delay(200); + + Assert.False(completionTask.IsCompleted, "Completion must wait for preceding FIFO event handlers to finish."); + + releaseHandler.TrySetResult(); + var response = await completionTask; + Assert.Equal("completed response", response?.Data.Content); + } + + [Fact] + public async Task SendAndWaitAsync_DroppedIdle_Fallback_Rechecks_Activity_After_Final_Barrier() + { + await using var server = await FakeCopilotServer.StartAsync(); + server.ConfigureActivityReactivationDuringFinalBarrier(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + await using var session = await client.CreateSessionAsync(new SessionConfig + { + OnPermissionRequest = PermissionHandler.ApproveAll + }); + + var barrierHandlerStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseBarrierHandler = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var subscription = session.On(evt => + { + if (evt.Data.TurnId == "activity-reactivation-barrier") + { + barrierHandlerStarted.TrySetResult(); + releaseBarrierHandler.Task.GetAwaiter().GetResult(); + } + }); + + var completionTask = session.SendAndWaitAsync( + new MessageOptions { Prompt = "recheck activity after final barrier" }, + timeout: TimeSpan.FromSeconds(5)); + + try + { + await barrierHandlerStarted.Task.WaitAsync(TimeSpan.FromSeconds(2)); + await server.ActivityReactivated.WaitAsync(TimeSpan.FromSeconds(2)); + Assert.False(completionTask.IsCompleted); + + releaseBarrierHandler.TrySetResult(); + await server.ReactivatedActivityObserved.WaitAsync(TimeSpan.FromSeconds(2)); + Assert.False(completionTask.IsCompleted, "Completion must not use stale idle activity after the barrier."); + + server.CompleteReactivatedActivity(); + var response = await completionTask; + Assert.Equal("completed response", response?.Data.Content); + } + finally + { + releaseBarrierHandler.TrySetResult(); + server.CompleteReactivatedActivity(); + } + } + [MethodImpl(MethodImplOptions.NoInlining)] private static async Task> CreateDroppedSessionAsync(CopilotClient client) { @@ -461,12 +562,20 @@ private sealed class FakeCopilotServer : IAsyncDisposable private readonly SemaphoreSlim _writeLock = new(1, 1); private readonly TaskCompletionSource _destroyStarted = new(TaskCreationOptions.RunContinuationsAsynchronously); private readonly TaskCompletionSource _allowDestroy = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly TaskCompletionSource _allowCompletionEvents = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly TaskCompletionSource _activityReactivated = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly TaskCompletionSource _reactivatedActivityObserved = new(TaskCreationOptions.RunContinuationsAsynchronously); private readonly Task _serverTask; private readonly List _requests = []; private readonly object _requestsLock = new(); private string? _lastSessionId; private bool _delayDestroy; private bool _failRuntimeShutdown; + private bool _emitCompletionWithoutIdle; + private bool _delayCompletionEvents; + private bool _reactivateDuringFinalBarrier; + private bool _hasActiveWork; + private int _activityRequestCount; private FakeCopilotServer(TcpListener listener) { @@ -492,6 +601,10 @@ public static Task StartAsync() public Task DestroyStarted => _destroyStarted.Task; + public Task ActivityReactivated => _activityReactivated.Task; + + public Task ReactivatedActivityObserved => _reactivatedActivityObserved.Task; + public int RuntimeShutdownCount { get; private set; } public IReadOnlyList Requests @@ -528,6 +641,32 @@ public void FailRuntimeShutdown() _failRuntimeShutdown = true; } + public void ConfigureDroppedIdleCompletion(bool delayEvents = false) + { + _emitCompletionWithoutIdle = true; + _delayCompletionEvents = delayEvents; + if (!delayEvents) + { + _allowCompletionEvents.TrySetResult(); + } + } + + public void ReleaseCompletionEvents() + { + _allowCompletionEvents.TrySetResult(); + } + + public void ConfigureActivityReactivationDuringFinalBarrier() + { + ConfigureDroppedIdleCompletion(); + _reactivateDuringFinalBarrier = true; + } + + public void CompleteReactivatedActivity() + { + Volatile.Write(ref _hasActiveWork, false); + } + public async ValueTask DisposeAsync() { _allowDestroy.TrySetResult(); @@ -595,6 +734,21 @@ private async Task HandleRequestAsync(Stream stream, JsonElement request, Cancel { _requests.Add(new RpcRequestRecord(method!, paramsElement)); } + + var activityRequestNumber = 0; + if (method == "session.metadata.activity") + { + activityRequestNumber = Interlocked.Increment(ref _activityRequestCount); + if (_reactivateDuringFinalBarrier && activityRequestNumber == 2) + { + await EmitActivityReactivationBarrierEventAsync(stream, cancellationToken); + } + else if (_reactivateDuringFinalBarrier && activityRequestNumber >= 3 && Volatile.Read(ref _hasActiveWork)) + { + _reactivatedActivityObserved.TrySetResult(); + } + } + object? result = method switch { "connect" => new Dictionary @@ -613,6 +767,12 @@ private async Task HandleRequestAsync(Stream stream, JsonElement request, Cancel { ["messageId"] = "message-1" }, + "session.metadata.activity" => new Dictionary + { + ["abortable"] = false, + ["hasActiveWork"] = Volatile.Read(ref _hasActiveWork) + }, + "session.tasks.waitForPending" => new Dictionary(), "session.mcp.oauth.handlePendingRequest" => new Dictionary { ["success"] = true @@ -632,6 +792,89 @@ private async Task HandleRequestAsync(Stream stream, JsonElement request, Cancel ["id"] = id, ["result"] = result }, cancellationToken); + + if (_reactivateDuringFinalBarrier && method == "session.metadata.activity" && activityRequestNumber == 2) + { + Volatile.Write(ref _hasActiveWork, true); + _activityReactivated.TrySetResult(); + } + + if (method == "session.send" && _emitCompletionWithoutIdle) + { + _ = EmitCompletionWithoutIdleAsync(stream, cancellationToken); + } + } + + private Task EmitActivityReactivationBarrierEventAsync(Stream stream, CancellationToken cancellationToken) + { + return WriteMessageAsync(stream, new Dictionary + { + ["jsonrpc"] = "2.0", + ["method"] = "session.event", + ["params"] = new Dictionary + { + ["sessionId"] = _lastSessionId, + ["event"] = new Dictionary + { + ["type"] = "assistant.turn_end", + ["id"] = Guid.NewGuid().ToString(), + ["timestamp"] = DateTimeOffset.UtcNow.ToString("O"), + ["data"] = new Dictionary + { + ["turnId"] = "activity-reactivation-barrier" + } + } + } + }, cancellationToken); + } + + private async Task EmitCompletionWithoutIdleAsync(Stream stream, CancellationToken cancellationToken) + { + if (_delayCompletionEvents) + { + await _allowCompletionEvents.Task.WaitAsync(cancellationToken); + } + + await WriteMessageAsync(stream, new Dictionary + { + ["jsonrpc"] = "2.0", + ["method"] = "session.event", + ["params"] = new Dictionary + { + ["sessionId"] = _lastSessionId, + ["event"] = new Dictionary + { + ["type"] = "assistant.message", + ["id"] = Guid.NewGuid().ToString(), + ["timestamp"] = DateTimeOffset.UtcNow.ToString("O"), + ["data"] = new Dictionary + { + ["content"] = "completed response", + ["messageId"] = "assistant-message-1" + } + } + } + }, cancellationToken); + + await WriteMessageAsync(stream, new Dictionary + { + ["jsonrpc"] = "2.0", + ["method"] = "session.event", + ["params"] = new Dictionary + { + ["sessionId"] = _lastSessionId, + ["event"] = new Dictionary + { + ["type"] = "assistant.turn_end", + ["id"] = Guid.NewGuid().ToString(), + ["timestamp"] = DateTimeOffset.UtcNow.ToString("O"), + ["data"] = new Dictionary + { + ["turnId"] = "turn-1" + } + } + } + }, cancellationToken); } private Dictionary CreateSessionResult(JsonElement request) diff --git a/test/snapshots/send_and_wait_reliability/should_complete_when_live_sessionidle_is_dropped.yaml b/test/snapshots/send_and_wait_reliability/should_complete_when_live_sessionidle_is_dropped.yaml new file mode 100644 index 0000000000..48667da723 --- /dev/null +++ b/test/snapshots/send_and_wait_reliability/should_complete_when_live_sessionidle_is_dropped.yaml @@ -0,0 +1,10 @@ +models: + - claude-sonnet-4.5 +conversations: + - messages: + - role: system + content: ${system} + - role: user + content: What is 5+5? Reply with just the number. + - role: assistant + content: "10"