From c4af287862eaec6926980796db56c8182089f0cc Mon Sep 17 00:00:00 2001 From: Michael Jolley Date: Thu, 27 Aug 2026 13:03:25 -0500 Subject: [PATCH] CmdPal: harden JSON-RPC notification lifecycle Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 41b7617f-41e8-4a89-8331-d67c93824511 --- .../Services/JsonRpc/JsonRpcConnection.cs | 37 +++++-- .../JsonRpcConnectionTests.StreamJsonRpc.cs | 100 +++++++++++++++++- .../JsonRpcConnectionTests.cs | 29 ++++- 3 files changed, 153 insertions(+), 13 deletions(-) diff --git a/src/modules/cmdpal/Microsoft.CmdPal.UI.ViewModels/Services/JsonRpc/JsonRpcConnection.cs b/src/modules/cmdpal/Microsoft.CmdPal.UI.ViewModels/Services/JsonRpc/JsonRpcConnection.cs index ca1972cf77..429834e624 100644 --- a/src/modules/cmdpal/Microsoft.CmdPal.UI.ViewModels/Services/JsonRpc/JsonRpcConnection.cs +++ b/src/modules/cmdpal/Microsoft.CmdPal.UI.ViewModels/Services/JsonRpc/JsonRpcConnection.cs @@ -284,15 +284,12 @@ public sealed class JsonRpcConnection : IDisposable _errorPumpTask ?? Task.CompletedTask, _rpc.Completion, }; - foreach (var task in tasks) + try + { + Task.WhenAll(tasks).Wait(DisposeDrainTimeout); + } + catch (AggregateException) { - try - { - task.Wait(DisposeDrainTimeout); - } - catch (AggregateException) - { - } } _ = DisposeTokenSourcesWhenTasksCompleteAsync(tasks, _disposalCts, _connectionClosedCts); @@ -411,6 +408,11 @@ public sealed class JsonRpcConnection : IDisposable catch (ChannelClosedException) { } + catch (Exception ex) + { + Logger.LogError("The JSON-RPC notification pump ended unexpectedly.", ex); + RaiseError(ex); + } } private async Task PumpErrorStreamAsync() @@ -457,7 +459,24 @@ public sealed class JsonRpcConnection : IDisposable private void RaiseError(Exception exception) { - Error?.Invoke(this, new JsonRpcErrorEventArgs(exception)); + var handlers = Error; + if (handlers is null) + { + return; + } + + var eventArgs = new JsonRpcErrorEventArgs(exception); + foreach (EventHandler handler in handlers.GetInvocationList()) + { + try + { + handler(this, eventArgs); + } + catch (Exception ex) + { + Logger.LogError("A JSON-RPC error event handler failed.", ex); + } + } } private static int GetErrorCode(RemoteRpcException exception) diff --git a/src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.ViewModels.UnitTests/JsonRpcConnectionTests.StreamJsonRpc.cs b/src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.ViewModels.UnitTests/JsonRpcConnectionTests.StreamJsonRpc.cs index 360cdb0bee..2709995c55 100644 --- a/src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.ViewModels.UnitTests/JsonRpcConnectionTests.StreamJsonRpc.cs +++ b/src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.ViewModels.UnitTests/JsonRpcConnectionTests.StreamJsonRpc.cs @@ -3,6 +3,8 @@ // See the LICENSE file in the project root for more information. using System; +using System.Diagnostics; +using System.IO; using System.Text.Json; using System.Text.Json.Nodes; using System.Threading; @@ -160,7 +162,8 @@ public partial class JsonRpcConnectionTests public async Task Dispose_KeepsCancellationTokensAliveUntilNotificationPumpExits() { using var cts = new CancellationTokenSource(TestTimeout); - var harness = CreateHarness(); + var errorStream = new BlockingReadStream(); + var harness = CreateHarness(errorStream: errorStream); var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); @@ -174,18 +177,63 @@ public partial class JsonRpcConnectionTests await WriteFramedAsync(harness.ExtensionWrites, BuildNotification("slow", new JsonObject()), cts.Token); await entered.Task.WaitAsync(cts.Token); + await errorStream.ReadStarted.WaitAsync(cts.Token); var notificationPump = harness.Host.NotificationConsumerCompletion; - await Task.Run(harness.Host.Dispose).WaitAsync(cts.Token); + var disposeDuration = await Task.Run(() => + { + var stopwatch = Stopwatch.StartNew(); + harness.Host.Dispose(); + stopwatch.Stop(); + return stopwatch.Elapsed; + }).WaitAsync(cts.Token); + + Assert.IsTrue(disposeDuration < TimeSpan.FromSeconds(3.5), $"Dispose took {disposeDuration}."); Assert.IsFalse(notificationPump.IsCompleted); release.TrySetResult(); + errorStream.Release(); await notificationPump.WaitAsync(cts.Token); Assert.AreEqual(TaskStatus.RanToCompletion, notificationPump.Status); } finally { release.TrySetResult(); + errorStream.Release(); + harness.Host.Dispose(); + } + } + + [TestMethod] + public async Task ThrowingErrorSubscriber_DoesNotStopOtherSubscribersOrNotificationPump() + { + using var cts = new CancellationTokenSource(TestTimeout); + var harness = CreateHarness(); + var errorRaised = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var errorObserved = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var notificationReceived = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + try + { + harness.Host.Error += (_, _) => + { + errorRaised.TrySetResult(); + throw new InvalidOperationException("subscriber failed"); + }; + harness.Host.Error += (_, _) => errorObserved.TrySetResult(); + harness.Host.RegisterNotificationHandler("fail", _ => throw new InvalidOperationException("handler failed")); + harness.Host.RegisterNotificationHandler("next", _ => notificationReceived.TrySetResult()); + + await WriteFramedAsync(harness.ExtensionWrites, BuildNotification("fail", new JsonObject()), cts.Token); + await errorRaised.Task.WaitAsync(cts.Token); + await errorObserved.Task.WaitAsync(cts.Token); + await WriteFramedAsync(harness.ExtensionWrites, BuildNotification("next", new JsonObject()), cts.Token); + + await notificationReceived.Task.WaitAsync(cts.Token); + Assert.IsFalse(harness.Host.NotificationConsumerCompletion.IsCompleted); + } + finally + { harness.Host.Dispose(); } } @@ -230,4 +278,52 @@ public partial class JsonRpcConnectionTests harness.Host.Dispose(); } } + + private sealed class BlockingReadStream : Stream + { + private readonly TaskCompletionSource _readStarted = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly TaskCompletionSource _release = new(TaskCreationOptions.RunContinuationsAsynchronously); + + internal Task ReadStarted => _readStarted.Task; + + public override bool CanRead => true; + + public override bool CanSeek => false; + + public override bool CanWrite => false; + + public override long Length => throw new NotSupportedException(); + + public override long Position + { + get => throw new NotSupportedException(); + set => throw new NotSupportedException(); + } + + internal void Release() => _release.TrySetResult(0); + + public override void Flush() + { + } + + public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException(); + + public override Task ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) + { + _readStarted.TrySetResult(); + return _release.Task; + } + + public override ValueTask ReadAsync(Memory buffer, CancellationToken cancellationToken = default) + { + _readStarted.TrySetResult(); + return new ValueTask(_release.Task); + } + + public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException(); + + public override void SetLength(long value) => throw new NotSupportedException(); + + public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException(); + } } diff --git a/src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.ViewModels.UnitTests/JsonRpcConnectionTests.cs b/src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.ViewModels.UnitTests/JsonRpcConnectionTests.cs index d6db1c41db..7ab1a0dc6c 100644 --- a/src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.ViewModels.UnitTests/JsonRpcConnectionTests.cs +++ b/src/modules/cmdpal/Tests/Microsoft.CmdPal.UI.ViewModels.UnitTests/JsonRpcConnectionTests.cs @@ -214,6 +214,31 @@ public partial class JsonRpcConnectionTests } } + [TestMethod] + public async Task SendNotification_WritesNotificationWithoutRequestId() + { + using var cts = new CancellationTokenSource(TestTimeout); + var harness = CreateHarness(); + try + { + await harness.Host.SendNotificationAsync( + "statusChanged", + new JsonObject { ["status"] = "ready" }, + cts.Token); + + var (_, body) = await ReadFramedAsync(harness.ExtensionReads, cts.Token); + using var document = JsonDocument.Parse(body); + Assert.AreEqual("2.0", document.RootElement.GetProperty("jsonrpc").GetString()); + Assert.AreEqual("statusChanged", document.RootElement.GetProperty("method").GetString()); + Assert.AreEqual("ready", document.RootElement.GetProperty("params").GetProperty("status").GetString()); + Assert.IsFalse(document.RootElement.TryGetProperty("id", out _)); + } + finally + { + harness.Host.Dispose(); + } + } + [TestMethod] public async Task Notification_IsDispatchedToHandler() { @@ -415,7 +440,7 @@ public partial class JsonRpcConnectionTests await Assert.ThrowsExceptionAsync(async () => await requestTask.WaitAsync(cts.Token)); } - private static Harness CreateHarness(TimeSpan? requestTimeout = null) + private static Harness CreateHarness(TimeSpan? requestTimeout = null, Stream? errorStream = null) { var toHost = new Pipe(); var fromHost = new Pipe(); @@ -423,7 +448,7 @@ public partial class JsonRpcConnectionTests var host = new JsonRpcConnection( toHost.Reader.AsStream(), fromHost.Writer.AsStream(), - errorStream: null, + errorStream, requestTimeout: requestTimeout); host.StartListening();