diff --git a/dotnet/src/Session.cs b/dotnet/src/Session.cs index 04306b7f6..011ae3099 100644 --- a/dotnet/src/Session.cs +++ b/dotnet/src/Session.cs @@ -81,13 +81,19 @@ private sealed record EventSubscription(Type EventType, Action Han private IReadOnlyList _openCanvases = Array.Empty(); private int _isDisposed; + private readonly object _eventDispatchGate = new(); + 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 +342,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 +353,7 @@ void Handler(SessionEvent evt) { case AssistantMessageEvent assistantMessage: lastAssistantMessage = assistantMessage; + completionCandidate.TrySetResult(true); if (!firstAssistantMessageLogged) { firstAssistantMessageLogged = true; @@ -356,12 +364,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 +383,86 @@ 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 = 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 (IOException ex) when (ex.InnerException is RemoteRpcException + { + ErrorCode: RemoteRpcException.MethodNotFoundErrorCode + }) + { + // Generated RPC methods surface remote errors through CopilotClient, + // which wraps them in IOException. Treat an older runtime's missing + // fallback methods the same as a direct method-not-found response. + 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 +477,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 +507,11 @@ void Handler(SessionEvent evt) "error"); throw; } + finally + { + cts.Cancel(); + await runtimeCompletionTask.ConfigureAwait(false); + } } /// @@ -485,8 +583,39 @@ 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); + // Queueing and publishing the new version share the same gate used to capture + // a completion barrier's stable version. Publish first so a monitor whose + // barrier is already queued cannot observe the old version after this event + // is written behind it; the gate prevents a new barrier from splitting the + // version update from the channel write. + lock (_eventDispatchGate) + { + Interlocked.Increment(ref _eventEnqueueVersion); + var queued = _eventChannel.Writer.TryWrite(new EventItem(sessionEvent)); + ObjectDisposedException.ThrowIf(!queued, this); + } + } + + /// + /// 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); + long stableEventVersion; + lock (_eventDispatchGate) + { + stableEventVersion = _eventEnqueueVersion; + 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); + return stableEventVersion; } /// @@ -495,8 +624,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 +2071,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/Unit/ClientSessionLifetimeTests.cs b/dotnet/test/Unit/ClientSessionLifetimeTests.cs index e1143db17..5389c1848 100644 --- a/dotnet/test/Unit/ClientSessionLifetimeTests.cs +++ b/dotnet/test/Unit/ClientSessionLifetimeTests.cs @@ -415,6 +415,255 @@ 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_Preserves_Event_Completion_When_Legacy_Runtime_Lacks_Fallback_Rpcs() + { + await using var server = await FakeCopilotServer.StartAsync(); + server.ConfigureLegacyRuntimeCompletion(); + 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 from delayed legacy idle" }, + timeout: TimeSpan.FromSeconds(3)); + + Assert.NotNull(response); + Assert.Equal("completed response", response.Data.Content); + Assert.Contains(server.Requests, request => request.Method == "session.tasks.waitForPending"); + } + + [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(); + } + } + + [Fact] + public async Task SendAndWaitAsync_DroppedIdle_Fallback_Flushes_Event_Enqueued_During_Final_Barrier() + { + await using var server = await FakeCopilotServer.StartAsync(); + server.ConfigureEventEnqueueDuringFinalBarrier(); + 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 queuedHandlerStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseQueuedHandler = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var subscription = session.On(evt => + { + if (evt.Data.TurnId == "activity-reactivation-barrier") + { + DispatchEvent(session, new AssistantTurnEndEvent + { + Data = new AssistantTurnEndData { TurnId = "queued-during-final-barrier" } + }); + } + else if (evt.Data.TurnId == "queued-during-final-barrier") + { + queuedHandlerStarted.TrySetResult(); + releaseQueuedHandler.Task.GetAwaiter().GetResult(); + } + }); + + var completionTask = session.SendAndWaitAsync( + new MessageOptions { Prompt = "flush an event queued during the final barrier" }, + timeout: TimeSpan.FromSeconds(5)); + + try + { + var first = await Task.WhenAny(queuedHandlerStarted.Task, completionTask) + .WaitAsync(TimeSpan.FromSeconds(2)); + Assert.Same(queuedHandlerStarted.Task, first); + await Task.Delay(200); + Assert.False(completionTask.IsCompleted, "Completion must wait for events enqueued after the stable snapshot."); + + releaseQueuedHandler.TrySetResult(); + var response = await completionTask; + Assert.Equal("completed response", response?.Data.Content); + } + finally + { + releaseQueuedHandler.TrySetResult(); + } + } + + [Fact] + public async Task SendAndWaitAsync_DroppedIdle_Fallback_Serializes_Concurrent_Enqueue_With_Final_Barrier() + { + await using var server = await FakeCopilotServer.StartAsync(); + server.ConfigureEventEnqueueDuringFinalBarrier(); + 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 dispatchGate = typeof(CopilotSession).GetField("_eventDispatchGate", BindingFlags.Instance | BindingFlags.NonPublic) + ?.GetValue(session) + ?? throw new InvalidOperationException("Event dispatch synchronization gate was not found."); + var enqueueStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseEnqueue = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var queuedHandlerStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseQueuedHandler = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + Task? enqueueTask = null; + using var subscription = session.On(evt => + { + if (evt.Data.TurnId == "activity-reactivation-barrier") + { + lock (dispatchGate) + { + var dispatchAttempted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var task = Task.Run(() => + { + dispatchAttempted.TrySetResult(); + DispatchEvent(session, new AssistantTurnEndEvent + { + Data = new AssistantTurnEndData { TurnId = "queued-after-concurrent-enqueue" } + }); + }); + enqueueTask = task; + dispatchAttempted.Task.GetAwaiter().GetResult(); + enqueueStarted.TrySetResult(); + releaseEnqueue.Task.GetAwaiter().GetResult(); + } + } + else if (evt.Data.TurnId == "queued-after-concurrent-enqueue") + { + queuedHandlerStarted.TrySetResult(); + releaseQueuedHandler.Task.GetAwaiter().GetResult(); + } + }); + + var completionTask = session.SendAndWaitAsync( + new MessageOptions { Prompt = "serialize a concurrent enqueue with the final barrier" }, + timeout: TimeSpan.FromSeconds(5)); + + try + { + await enqueueStarted.Task.WaitAsync(TimeSpan.FromSeconds(2)); + await Task.Delay(200); + Assert.False(enqueueTask!.IsCompleted, "DispatchEvent must be blocked by the shared enqueue/barrier gate."); + Assert.False(completionTask.IsCompleted, "Completion must not cross an event enqueue that is contending with the final barrier."); + + releaseEnqueue.TrySetResult(); + await enqueueTask!.WaitAsync(TimeSpan.FromSeconds(2)); + var first = await Task.WhenAny(queuedHandlerStarted.Task, completionTask) + .WaitAsync(TimeSpan.FromSeconds(2)); + Assert.Same(queuedHandlerStarted.Task, first); + Assert.False(completionTask.IsCompleted, "Completion must wait for the in-flight event's handler to cross the FIFO barrier."); + + releaseQueuedHandler.TrySetResult(); + var response = await completionTask; + Assert.Equal("completed response", response?.Data.Content); + } + finally + { + releaseEnqueue.TrySetResult(); + releaseQueuedHandler.TrySetResult(); + } + } + + [MethodImpl(MethodImplOptions.NoInlining)] private static async Task> CreateDroppedSessionAsync(CopilotClient client) { @@ -512,12 +761,23 @@ 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 _emitDelayedIdle; + private bool _fallbackMethodsUnavailable; + private bool _reactivateDuringFinalBarrier; + private bool _enqueueDuringFinalBarrier; + private bool _hasActiveWork; + private int _activityRequestCount; private FakeCopilotServer(TcpListener listener) { @@ -543,6 +803,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 @@ -579,6 +843,46 @@ public void FailRuntimeShutdown() _failRuntimeShutdown = true; } + public void ConfigureDroppedIdleCompletion(bool delayEvents = false) + { + _emitCompletionWithoutIdle = true; + _delayCompletionEvents = delayEvents; + if (!delayEvents) + { + _allowCompletionEvents.TrySetResult(); + } + } + + public void ConfigureLegacyRuntimeCompletion() + { + ConfigureDroppedIdleCompletion(); + _emitDelayedIdle = true; + _fallbackMethodsUnavailable = true; + } + + public void ReleaseCompletionEvents() + { + _allowCompletionEvents.TrySetResult(); + } + + public void ConfigureActivityReactivationDuringFinalBarrier() + { + ConfigureDroppedIdleCompletion(); + _reactivateDuringFinalBarrier = true; + } + + public void ConfigureEventEnqueueDuringFinalBarrier() + { + ConfigureDroppedIdleCompletion(); + _enqueueDuringFinalBarrier = true; + } + + + public void CompleteReactivatedActivity() + { + Volatile.Write(ref _hasActiveWork, false); + } + public async ValueTask DisposeAsync() { _allowDestroy.TrySetResult(); @@ -646,6 +950,37 @@ private async Task HandleRequestAsync(Stream stream, JsonElement request, Cancel { _requests.Add(new RpcRequestRecord(method!, paramsElement)); } + + if (_fallbackMethodsUnavailable + && method is "session.tasks.waitForPending" or "session.metadata.activity") + { + await WriteMessageAsync(stream, new Dictionary + { + ["jsonrpc"] = "2.0", + ["id"] = id, + ["error"] = new Dictionary + { + ["code"] = -32601, + ["message"] = $"Method not found: {method}" + } + }, cancellationToken); + return; + } + + var activityRequestNumber = 0; + if (method == "session.metadata.activity") + { + activityRequestNumber = Interlocked.Increment(ref _activityRequestCount); + if ((_reactivateDuringFinalBarrier || _enqueueDuringFinalBarrier) && activityRequestNumber == 2) + { + await EmitActivityReactivationBarrierEventAsync(stream, cancellationToken); + } + else if (_reactivateDuringFinalBarrier && activityRequestNumber >= 3 && Volatile.Read(ref _hasActiveWork)) + { + _reactivatedActivityObserved.TrySetResult(); + } + } + object? result = method switch { "connect" => new Dictionary @@ -664,6 +999,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 @@ -683,6 +1024,110 @@ 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); + + if (_emitDelayedIdle) + { + await Task.Delay(TimeSpan.FromMilliseconds(500), cancellationToken); + await WriteMessageAsync(stream, new Dictionary + { + ["jsonrpc"] = "2.0", + ["method"] = "session.event", + ["params"] = new Dictionary + { + ["sessionId"] = _lastSessionId, + ["event"] = new Dictionary + { + ["type"] = "session.idle", + ["id"] = Guid.NewGuid().ToString(), + ["timestamp"] = DateTimeOffset.UtcNow.ToString("O"), + ["data"] = new Dictionary() + } + } + }, cancellationToken); + } } private Dictionary CreateSessionResult(JsonElement request)