Skip to content

fix(dotnet): recover from dropped session idle events #2016

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
hydraxman wants to merge 1 commit into
base: main
Choose a base branch
from
+464 −10
Draft

fix(dotnet): recover from dropped session idle events #2016

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
137 changes: 127 additions & 10 deletions dotnet/src/Session.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,13 +81,18 @@ private sealed record EventSubscription(Type EventType, Action<SessionEvent> Han
private IReadOnlyList<OpenCanvasInstance> _openCanvases = Array.Empty<OpenCanvasInstance>();

private int _isDisposed;
private long _eventEnqueueVersion;

private abstract record EventDispatchItem;
private sealed record EventItem(SessionEvent Event) : EventDispatchItem;
private sealed record EventBarrier(TaskCompletionSource<bool> Completion) : EventDispatchItem;

/// <summary>
/// Channel that serializes event dispatch. <see cref="DispatchEvent"/> enqueues;
/// a single background consumer (<see cref="ProcessEventsAsync"/>) dequeues and
/// invokes handlers one at a time, preserving arrival order.
/// </summary>
private readonly Channel<SessionEvent> _eventChannel = Channel.CreateUnbounded<SessionEvent>(
private readonly Channel<EventDispatchItem> _eventChannel = Channel.CreateUnbounded<EventDispatchItem>(
new() { SingleReader = true });

/// <summary>
Expand Down Expand Up @@ -336,7 +341,8 @@ public async Task<string> SendAsync(MessageOptions options, CancellationToken ca

var totalTimestamp = Stopwatch.GetTimestamp();
var effectiveTimeout = timeout ?? TimeSpan.FromSeconds(60);
var tcs = new TaskCompletionSource<AssistantMessageEvent?>(TaskCreationOptions.RunContinuationsAsynchronously);
var tcs = new TaskCompletionSource<(AssistantMessageEvent? Message, string CompletedBy)>(TaskCreationOptions.RunContinuationsAsynchronously);
var completionCandidate = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
AssistantMessageEvent? lastAssistantMessage = null;
var firstAssistantMessageLogged = false;

Expand All @@ -346,6 +352,7 @@ void Handler(SessionEvent evt)
{
case AssistantMessageEvent assistantMessage:
lastAssistantMessage = assistantMessage;
completionCandidate.TrySetResult(true);
if (!firstAssistantMessageLogged)
{
firstAssistantMessageLogged = true;
Expand All @@ -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:
Expand All @@ -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<SessionEvent>(Handler);

await SendAsync(options, cancellationToken);
Expand All @@ -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)
{
Expand All @@ -414,6 +497,11 @@ void Handler(SessionEvent evt)
"error");
throw;
}
finally
{
cts.Cancel();
await runtimeCompletionTask.ConfigureAwait(false);
}
}

/// <summary>
Expand Down Expand Up @@ -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));
}

/// <summary>
/// Waits until every event already queued for this session has been delivered to
/// user handlers. Events arriving after the barrier remain queued behind it.
/// </summary>
private async Task FlushEventDispatchAsync(CancellationToken cancellationToken)
{
var completion = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
var queued = _eventChannel.Writer.TryWrite(new EventBarrier(completion));
ObjectDisposedException.ThrowIf(!queued, this);

using var registration = cancellationToken.Register(
static state => ((TaskCompletionSource<bool>)state!).TrySetCanceled(),
completion);
await completion.Task.ConfigureAwait(false);
}

/// <summary>
Expand All @@ -495,8 +602,15 @@ internal void DispatchEvent(SessionEvent sessionEvent)
/// </summary>
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)
Expand Down Expand Up @@ -1935,6 +2049,9 @@ await InvokeRpcAsync<object>(
[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);

Expand Down
84 changes: 84 additions & 0 deletions dotnet/test/E2E/SendAndWaitReliabilityE2ETests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;

/// <summary>
/// Verifies that .NET completion remains reliable when the runtime's ephemeral
/// <c>session.idle</c> notification does not reach the SendAndWait handler.
/// </summary>
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<UserMessageEvent>(_ =>
{
userMessageDispatched.TrySetResult();
releaseUserMessage.Task.GetAwaiter().GetResult();
});
using var idleSubscription = session.On<SessionIdleEvent>(_ => 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<object>().Last(subscription =>
{
var eventType = subscription.GetType().GetProperty("EventType")?.GetValue(subscription);
return Equals(eventType, typeof(SessionEvent));
});
var handlerField = sendAndWaitSubscription.GetType().GetField(
"<Handler>k__BackingField",
BindingFlags.Instance | BindingFlags.NonPublic)
?? throw new InvalidOperationException("EventSubscription.Handler backing field was not found.");
var original = (Action<SessionEvent>)(handlerField.GetValue(sendAndWaitSubscription)
?? throw new InvalidOperationException("SendAndWait event handler was null."));

handlerField.SetValue(sendAndWaitSubscription, (Action<SessionEvent>)(evt =>
{
if (evt is not SessionIdleEvent)
{
original(evt);
}
}));
}
}
Loading