Preserve JsonRpc proxy lifetimes on refresh

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 813c51f6-19e2-49ee-bf03-10e3187d3728
2026-09-01 19:51:34 +02:00 · 2026-08-27 16:33:41 -05:00
parent 0358bfbaff
commit 8aadbe1bd0
6 changed files with 230 additions and 154 deletions
--- a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSContentPageProxy.cs
+++ b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSContentPageProxy.cs
@@ -3,7 +3,6 @@
 // See the LICENSE file in the project root for more information.

 using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Runtime.CompilerServices;
 using System.Text.Json;
@@ -46,11 +45,7 @@ internal sealed partial class JSContentPageProxy : JSObservableProxyBase, IConte
        _registry = Registries.GetValue(Connection, static _ => new PageRegistry());
        _registry.EnsureSubscribed(Connection);

-        var pages = _registry.Pages.GetOrAdd(_pageId, static _ => new List<WeakReference<JSContentPageProxy>>());
-        lock (pages)
-        {
-            pages.Add(new WeakReference<JSContentPageProxy>(this));
-        }
+        _registry.Pages.Register(_pageId, this);
    }

    public event TypedEventHandler<object, IItemsChangedEventArgs>? ItemsChanged;
@@ -99,17 +94,7 @@ internal sealed partial class JSContentPageProxy : JSObservableProxyBase, IConte
    {
        _details.Dispose();
        _commands.Dispose();
-        if (_registry.Pages.TryGetValue(_pageId, out var pages))
-        {
-            lock (pages)
-            {
-                pages.RemoveAll(weak => !weak.TryGetTarget(out var target) || ReferenceEquals(target, this));
-                if (pages.Count == 0)
-                {
-                    _registry.Pages.TryRemove(_pageId, out _);
-                }
-            }
-        }
+        _registry.Pages.Unregister(_pageId, this);

        base.Dispose();
    }
@@ -163,30 +148,12 @@ internal sealed partial class JSContentPageProxy : JSObservableProxyBase, IConte
            }

            var pageId = pageProperty.GetString();
-            if (pageId is null || !registry.Pages.TryGetValue(pageId, out var pageReferences))
+            if (pageId is null)
            {
                return;
            }

-            List<JSContentPageProxy> targets = [];
-            lock (pageReferences)
-            {
-                pageReferences.RemoveAll(weak => !weak.TryGetTarget(out _));
-                foreach (var weak in pageReferences)
-                {
-                    if (weak.TryGetTarget(out var target))
-                    {
-                        targets.Add(target);
-                    }
-                }
-
-                if (pageReferences.Count == 0)
-                {
-                    registry.Pages.TryRemove(pageId, out _);
-                }
-            }
-
-            foreach (var target in targets)
+            foreach (var target in registry.Pages.GetLiveTargets(pageId))
            {
                var handler = target.ItemsChanged;
                if (handler is not null)
@@ -206,7 +173,7 @@ internal sealed partial class JSContentPageProxy : JSObservableProxyBase, IConte
        private readonly object _subscribeLock = new();
        private bool _subscribed;

-        public ConcurrentDictionary<string, List<WeakReference<JSContentPageProxy>>> Pages { get; } = new();
+        public JSWeakReferenceRegistry<string, JSContentPageProxy> Pages { get; } = new();

        public void EnsureSubscribed(JsonRpcConnection connection)
        {
--- a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSLazyCache.cs
+++ b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSLazyCache.cs
@@ -63,7 +63,7 @@ internal sealed partial class JSLazyCache<T> : IDisposable
                return;
            }

-            DisposeCreatedValue();
+            // The host may still own the previous proxy, so invalidation only releases the cache reference.
            _value = default!;
            _hasValue = false;
        }
--- a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSListPageProxy.cs
+++ b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSListPageProxy.cs
@@ -3,7 +3,6 @@
 // See the LICENSE file in the project root for more information.

 using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 using System.Runtime.CompilerServices;
@@ -56,11 +55,7 @@ internal sealed partial class JSListPageProxy : JSObservableProxyBase, IListPage
        _registry = Registries.GetValue(Connection, static _ => new PageRegistry());
        _registry.EnsureSubscribed(Connection);

-        var list = _registry.Pages.GetOrAdd(_pageId, static _ => new List<WeakReference<JSListPageProxy>>());
-        lock (list)
-        {
-            list.Add(new WeakReference<JSListPageProxy>(this));
-        }
+        _registry.Pages.Register(_pageId, this);
    }

    public event TypedEventHandler<object, IItemsChangedEventArgs>? ItemsChanged;
@@ -295,17 +290,7 @@ internal sealed partial class JSListPageProxy : JSObservableProxyBase, IListPage
        _emptyContent.Dispose();
        base.Dispose();

-        if (_registry.Pages.TryGetValue(_pageId, out var list))
-        {
-            lock (list)
-            {
-                list.RemoveAll(weak => !weak.TryGetTarget(out var target) || ReferenceEquals(target, this));
-                if (list.Count == 0)
-                {
-                    _registry.Pages.TryRemove(_pageId, out _);
-                }
-            }
-        }
+        _registry.Pages.Unregister(_pageId, this);
    }

    private static void DispatchItemsChanged(PageRegistry registry, JsonElement paramsElement)
@@ -319,7 +304,7 @@ internal sealed partial class JSListPageProxy : JSObservableProxyBase, IListPage
            }

            var pageId = pageProp.GetString();
-            if (pageId == null || !registry.Pages.TryGetValue(pageId, out var proxyRefs))
+            if (pageId == null)
            {
                return;
            }
@@ -331,27 +316,7 @@ internal sealed partial class JSListPageProxy : JSObservableProxyBase, IListPage
                totalItems = totalItemsProp.GetInt32();
            }

-            // Snapshot live proxies and prune collected ones so the registry does
-            // not grow as pages come and go.
-            List<JSListPageProxy> targets = new();
-            lock (proxyRefs)
-            {
-                proxyRefs.RemoveAll(weak => !weak.TryGetTarget(out _));
-                foreach (var weak in proxyRefs)
-                {
-                    if (weak.TryGetTarget(out var proxy))
-                    {
-                        targets.Add(proxy);
-                    }
-                }
-
-                if (proxyRefs.Count == 0)
-                {
-                    registry.Pages.TryRemove(pageId, out _);
-                }
-            }
-
-            foreach (var proxy in targets)
+            foreach (var proxy in registry.Pages.GetLiveTargets(pageId))
            {
                proxy.UpdatePageState(paramsElement);

@@ -458,7 +423,7 @@ internal sealed partial class JSListPageProxy : JSObservableProxyBase, IListPage
        private readonly object _subscribeLock = new();
        private bool _subscribed;

-        public ConcurrentDictionary<string, List<WeakReference<JSListPageProxy>>> Pages { get; } = new();
+        public JSWeakReferenceRegistry<string, JSListPageProxy> Pages { get; } = new();

        // Binds the itemsChanged handler to the retained registry once. Binding here,
        // instead of inside the ConditionalWeakTable factory, keeps the handler from
--- a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSPropertyChangeRegistry.cs
+++ b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSPropertyChangeRegistry.cs
@@ -3,8 +3,6 @@
 // See the LICENSE file in the project root for more information.

 using System;
-using System.Collections.Concurrent;
-using System.Collections.Generic;
 using System.Runtime.CompilerServices;
 using System.Text.Json;
 using Microsoft.CmdPal.JsonRpc;
@@ -18,46 +16,17 @@ internal static class JSPropertyChangeRegistry
    internal static void Register(JsonRpcConnection connection, string commandId, IJSPropertyChangeTarget target)
    {
        var registry = Registries.GetValue(connection, static _ => new Registry());
-        while (true)
-        {
-            var targets = registry.Targets.GetOrAdd(commandId, static _ => []);
-            lock (targets)
-            {
-                if (!registry.Targets.TryGetValue(commandId, out var currentTargets) ||
-                    !ReferenceEquals(targets, currentTargets))
-                {
-                    continue;
-                }
-
-                targets.RemoveAll(reference => !reference.TryGetTarget(out _));
-                if (!targets.Exists(reference =>
-                    reference.TryGetTarget(out var current) && ReferenceEquals(current, target)))
-                {
-                    targets.Add(new WeakReference<IJSPropertyChangeTarget>(target));
-                }
-
-                return;
-            }
-        }
+        registry.Targets.Register(commandId, target);
    }

    internal static void Unregister(JsonRpcConnection connection, string commandId, IJSPropertyChangeTarget target)
    {
-        if (!Registries.TryGetValue(connection, out var registry) ||
-            !registry.Targets.TryGetValue(commandId, out var targets))
+        if (!Registries.TryGetValue(connection, out var registry))
        {
            return;
        }

-        lock (targets)
-        {
-            targets.RemoveAll(reference =>
-                !reference.TryGetTarget(out var current) || ReferenceEquals(current, target));
-            if (targets.Count == 0)
-            {
-                RemoveTargets(registry, commandId, targets);
-            }
-        }
+        registry.Targets.Unregister(commandId, target);
    }

    internal static void Dispatch(JsonRpcConnection connection, JsonElement paramsElement)
@@ -73,30 +42,12 @@ internal static class JSPropertyChangeRegistry
        }

        var commandId = commandIdProperty.GetString();
-        if (commandId is null || !registry.Targets.TryGetValue(commandId, out var targets))
+        if (commandId is null)
        {
            return;
        }

-        List<IJSPropertyChangeTarget> liveTargets = [];
-        lock (targets)
-        {
-            targets.RemoveAll(reference => !reference.TryGetTarget(out _));
-            foreach (var reference in targets)
-            {
-                if (reference.TryGetTarget(out var target))
-                {
-                    liveTargets.Add(target);
-                }
-            }
-
-            if (targets.Count == 0)
-            {
-                RemoveTargets(registry, commandId, targets);
-            }
-        }
-
-        foreach (var target in liveTargets)
+        foreach (var target in registry.Targets.GetLiveTargets(commandId))
        {
            target.ApplyPropertyChanges(properties);
        }
@@ -104,29 +55,13 @@ internal static class JSPropertyChangeRegistry

    internal static int GetRegistrationCount(JsonRpcConnection connection, string commandId)
    {
-        if (!Registries.TryGetValue(connection, out var registry) ||
-            !registry.Targets.TryGetValue(commandId, out var targets))
-        {
-            return 0;
-        }
-
-        lock (targets)
-        {
-            return targets.Count;
-        }
-    }
-
-    private static void RemoveTargets(
-        Registry registry,
-        string commandId,
-        List<WeakReference<IJSPropertyChangeTarget>> targets)
-    {
-        ((ICollection<KeyValuePair<string, List<WeakReference<IJSPropertyChangeTarget>>>>)registry.Targets)
-            .Remove(new KeyValuePair<string, List<WeakReference<IJSPropertyChangeTarget>>>(commandId, targets));
+        return Registries.TryGetValue(connection, out var registry)
+            ? registry.Targets.GetRegistrationCount(commandId)
+            : 0;
    }

    private sealed class Registry
    {
-        internal ConcurrentDictionary<string, List<WeakReference<IJSPropertyChangeTarget>>> Targets { get; } = new();
+        internal JSWeakReferenceRegistry<string, IJSPropertyChangeTarget> Targets { get; } = new();
    }
 }
--- a/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSWeakReferenceRegistry.cs
+++ b/src/modules/cmdpal/Microsoft.CmdPal.JsonRpc/Models/JSWeakReferenceRegistry.cs
@@ -0,0 +1,122 @@
+// Copyright (c) Microsoft Corporation
+// The Microsoft Corporation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+
+namespace Microsoft.CmdPal.JsonRpc.Models;
+
+internal sealed partial class JSWeakReferenceRegistry<TKey, TTarget>
+    where TKey : notnull
+    where TTarget : class
+{
+    private readonly ConcurrentDictionary<TKey, List<WeakReference<TTarget>>> _targets = new();
+
+    internal void Register(TKey key, TTarget target)
+    {
+        while (true)
+        {
+            var targets = _targets.GetOrAdd(key, static _ => []);
+            lock (targets)
+            {
+                if (!IsCurrent(key, targets))
+                {
+                    continue;
+                }
+
+                targets.RemoveAll(reference => !reference.TryGetTarget(out _));
+                if (!targets.Exists(reference =>
+                    reference.TryGetTarget(out var current) && ReferenceEquals(current, target)))
+                {
+                    targets.Add(new WeakReference<TTarget>(target));
+                }
+
+                return;
+            }
+        }
+    }
+
+    internal void Unregister(TKey key, TTarget target)
+    {
+        if (!_targets.TryGetValue(key, out var targets))
+        {
+            return;
+        }
+
+        lock (targets)
+        {
+            if (!IsCurrent(key, targets))
+            {
+                return;
+            }
+
+            targets.RemoveAll(reference =>
+                !reference.TryGetTarget(out var current) || ReferenceEquals(current, target));
+            RemoveIfEmpty(key, targets);
+        }
+    }
+
+    internal List<TTarget> GetLiveTargets(TKey key)
+    {
+        while (_targets.TryGetValue(key, out var targets))
+        {
+            lock (targets)
+            {
+                if (!IsCurrent(key, targets))
+                {
+                    continue;
+                }
+
+                var liveTargets = new List<TTarget>(targets.Count);
+                targets.RemoveAll(reference => !reference.TryGetTarget(out _));
+                foreach (var reference in targets)
+                {
+                    if (reference.TryGetTarget(out var target))
+                    {
+                        liveTargets.Add(target);
+                    }
+                }
+
+                RemoveIfEmpty(key, targets);
+                return liveTargets;
+            }
+        }
+
+        return [];
+    }
+
+    internal int GetRegistrationCount(TKey key)
+    {
+        while (_targets.TryGetValue(key, out var targets))
+        {
+            lock (targets)
+            {
+                if (!IsCurrent(key, targets))
+                {
+                    continue;
+                }
+
+                return targets.Count;
+            }
+        }
+
+        return 0;
+    }
+
+    private bool IsCurrent(TKey key, List<WeakReference<TTarget>> targets)
+    {
+        return _targets.TryGetValue(key, out var currentTargets) &&
+            ReferenceEquals(targets, currentTargets);
+    }
+
+    private void RemoveIfEmpty(TKey key, List<WeakReference<TTarget>> targets)
+    {
+        if (targets.Count == 0)
+        {
+            ((ICollection<KeyValuePair<TKey, List<WeakReference<TTarget>>>>)_targets)
+                .Remove(new KeyValuePair<TKey, List<WeakReference<TTarget>>>(key, targets));
+        }
+    }
+}
--- a/src/modules/cmdpal/Tests/Microsoft.CmdPal.JsonRpc.UnitTests/JSAdapterLifecycleTests.cs
+++ b/src/modules/cmdpal/Tests/Microsoft.CmdPal.JsonRpc.UnitTests/JSAdapterLifecycleTests.cs
@@ -51,6 +51,45 @@ public partial class JSAdapterTests
        JSPropertyChangeRegistry.Unregister(fake.Connection, commandId, liveTarget);
    }

+    [TestMethod]
+    public void WeakReferenceRegistry_PrunesDeadTargetsDuringRegistration()
+    {
+        var registry = new JSWeakReferenceRegistry<string, object>();
+        var deadTarget = RegisterTemporaryTarget(registry, "page");
+        GC.Collect();
+        GC.WaitForPendingFinalizers();
+        GC.Collect();
+
+        Assert.IsFalse(deadTarget.TryGetTarget(out _));
+        Assert.AreEqual(1, registry.GetRegistrationCount("page"));
+
+        var liveTarget = new object();
+        registry.Register("page", liveTarget);
+
+        Assert.AreEqual(1, registry.GetRegistrationCount("page"));
+        Assert.AreSame(liveTarget, registry.GetLiveTargets("page")[0]);
+    }
+
+    [TestMethod]
+    public void WeakReferenceRegistry_ConcurrentReplacementKeepsNewestTarget()
+    {
+        for (var i = 0; i < 128; i++)
+        {
+            var registry = new JSWeakReferenceRegistry<string, object>();
+            var oldTarget = new object();
+            var newTarget = new object();
+            registry.Register("page", oldTarget);
+
+            Parallel.Invoke(
+                () => registry.Unregister("page", oldTarget),
+                () => registry.Register("page", newTarget));
+
+            var liveTargets = registry.GetLiveTargets("page");
+            Assert.AreEqual(1, liveTargets.Count);
+            Assert.AreSame(newTarget, liveTargets[0]);
+        }
+    }
+
    [TestMethod]
    public void NestedProxyGetters_CacheIdentityAndRegistration()
    {
@@ -139,6 +178,7 @@ public partial class JSAdapterTests
            fake.Connection);
        var originalDetails = item.Details;
        var originalMoreCommands = item.MoreCommands;
+        var originalCommand = ((ICommandContextItem)originalMoreCommands[0]).Command;

        JSPropertyChangeRegistry.Dispatch(
            fake.Connection,
@@ -167,8 +207,45 @@ public partial class JSAdapterTests
        Assert.AreEqual("New details", item.Details?.Title);
        Assert.AreNotSame(originalDetails, item.Details);
        Assert.AreNotSame(originalMoreCommands, item.MoreCommands);
-        Assert.AreEqual(0, JSPropertyChangeRegistry.GetRegistrationCount(fake.Connection, "old-more"));
+        Assert.AreSame(item.MoreCommands, item.MoreCommands);
+        Assert.AreEqual(1, JSPropertyChangeRegistry.GetRegistrationCount(fake.Connection, "old-more"));
        Assert.AreEqual(1, JSPropertyChangeRegistry.GetRegistrationCount(fake.Connection, "new-more"));
+
+        JSPropertyChangeRegistry.Dispatch(
+            fake.Connection,
+            ParseElement(new JsonObject
+            {
+                ["commandId"] = "old-more",
+                ["properties"] = new JsonObject { ["name"] = "Still alive" },
+            }));
+
+        Assert.AreEqual("Still alive", originalCommand.Name);
+        (originalCommand as IDisposable)?.Dispose();
+    }
+
+    [TestMethod]
+    public async Task PageRegistries_RouteToNewestProxyAfterOlderProxyIsRemoved()
+    {
+        using var fake = new JSFakeExtension();
+        using var oldListPage = new JSListPageProxy("list-page", fake.Connection);
+        using var newListPage = new JSListPageProxy("list-page", fake.Connection);
+        using var oldContentPage = new JSContentPageProxy("content-page", fake.Connection);
+        using var newContentPage = new JSContentPageProxy("content-page", fake.Connection);
+        var listChanged = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        var contentChanged = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        newListPage.ItemsChanged += (_, _) => listChanged.TrySetResult();
+        newContentPage.ItemsChanged += (_, _) => contentChanged.TrySetResult();
+
+        oldListPage.Dispose();
+        oldContentPage.Dispose();
+        await fake.PushNotificationAsync(
+            "listPage/itemsChanged",
+            new JsonObject { ["pageId"] = "list-page" });
+        await fake.PushNotificationAsync(
+            "contentPage/itemsChanged",
+            new JsonObject { ["pageId"] = "content-page" });
+
+        await Task.WhenAll(listChanged.Task, contentChanged.Task).WaitAsync(Timeout);
    }

    [TestMethod]
@@ -428,6 +505,16 @@ public partial class JSAdapterTests
        return new WeakReference<RecordingPropertyChangeTarget>(target);
    }

+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static WeakReference<object> RegisterTemporaryTarget(
+        JSWeakReferenceRegistry<string, object> registry,
+        string key)
+    {
+        var target = new object();
+        registry.Register(key, target);
+        return new WeakReference<object>(target);
+    }
+
    private static JsonObject Command(string id) => new()
    {
        ["id"] = id,