Fix Blazor per-page activity links

src/Components/Components/src/ComponentsActivityState.cs

@@ -0,0 +1,73 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.AspNetCore.Components.Infrastructure;
+
+namespace Microsoft.AspNetCore.Components;
+
+internal sealed class ComponentsActivityState
+{
+    private ComponentsActivityLinkStore? _activityLinkStore;
+    private ComponentsActivityPersistentStateUpdate? _pendingState;
+    private bool _hasPendingState;
+
+    public ComponentsActivityPersistentStateUpdate Capture()
+    {
+        ComponentsActivityPersistentState? routeState = null;
+        _activityLinkStore?.TryCreatePersistentRouteState(out routeState);
+        return new ComponentsActivityPersistentStateUpdate(routeState);
+    }
+
+    public void Apply(ComponentsActivityPersistentStateUpdate state)
+    {
+        _pendingState = state;
+        _hasPendingState = true;
+        ApplyPendingState();
+    }
+
+    public void Initialize(ComponentsActivityLinkStore activityLinkStore)
+    {
+        _activityLinkStore = activityLinkStore;
+        ApplyPendingState();
+    }
+
+    private void ApplyPendingState()
+    {
+        if (_activityLinkStore is null || !_hasPendingState)
+        {
+            return;
+        }
+
+        if (_pendingState?.Route is { } routeState)
+        {
+            _activityLinkStore.RestorePersistentRouteState(routeState);
+        }
+        else
+        {
+            _activityLinkStore.RemoveActivityContext(ComponentsActivityLinkStore.Route);
+        }
+
+        _pendingState = null;
+        _hasPendingState = false;
+    }
+}
+
+internal sealed class ServerComponentsActivityState(ComponentsActivityState activityState)
+{
+    [PersistentState(AllowUpdates = true)]
+    public ComponentsActivityPersistentStateUpdate ActivityState
+    {
+        get => activityState.Capture();
+        set => activityState.Apply(value);
+    }
+}
+
+internal sealed class WebAssemblyComponentsActivityState(ComponentsActivityState activityState)
+{
+    [PersistentState(AllowUpdates = true)]
+    public ComponentsActivityPersistentStateUpdate ActivityState
+    {
+        get => activityState.Capture();
+        set => activityState.Apply(value);
+    }
+}

src/Components/Components/src/RegisterRenderingMetricsServiceCollectionExtensions.cs

@@ -1,9 +1,12 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Diagnostics.CodeAnalysis;
 using System.Diagnostics.Metrics;
+using System.Linq;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.DependencyInjection.Extensions;
+using static Microsoft.AspNetCore.Internal.LinkerFlags;
 
 namespace Microsoft.AspNetCore.Components.Infrastructure;
 
@@ -39,10 +42,66 @@ public static IServiceCollection AddComponentsTracing(
         IServiceCollection services)
     {
         services.TryAddScoped<ComponentsActivitySource>();
+        if (services.Any(descriptor =>
+            descriptor.ServiceType == typeof(IComponentRenderMode) &&
+            Equals(descriptor.ServiceKey, "Microsoft.AspNetCore.Components.ActivityState.WebAssembly")))
+        {
+            services.TryAddSingleton<ComponentsActivityState>();
+        }
+        else
+        {
+            services.TryAddScoped<ComponentsActivityState>();
+        }
+        services.TryAddScoped<ServerComponentsActivityState>();
+        services.TryAddScoped<WebAssemblyComponentsActivityState>();
+        services.TryAddKeyedSingleton<IComponentRenderMode, UnsupportedComponentsActivityStateRenderMode>(
+            "Microsoft.AspNetCore.Components.ActivityState.Server");
+        services.TryAddKeyedSingleton<IComponentRenderMode, UnsupportedComponentsActivityStateRenderMode>(
+            "Microsoft.AspNetCore.Components.ActivityState.WebAssembly");
+        services.TryAddEnumerable(
+            ServiceDescriptor.Singleton<IPersistentServiceRegistration, ServerComponentsActivityStatePersistentServiceRegistration>());
+        services.TryAddEnumerable(
+            ServiceDescriptor.Singleton<IPersistentServiceRegistration, WebAssemblyComponentsActivityStatePersistentServiceRegistration>());
 
         return services;
     }
 
+    private sealed class UnsupportedComponentsActivityStateRenderMode : IComponentRenderMode;
+
+    private sealed class ServerComponentsActivityStatePersistentServiceRegistration(IServiceProvider serviceProvider)
+        : IPersistentServiceRegistration
+    {
+        private readonly PersistentServiceRegistration<ServerComponentsActivityState> _registration = new(
+            serviceProvider.GetRequiredKeyedService<IComponentRenderMode>(
+                "Microsoft.AspNetCore.Components.ActivityState.Server"));
+
+        public string Assembly => _registration.Assembly;
+
+        public string FullTypeName => _registration.FullTypeName;
+
+        public IComponentRenderMode? GetRenderModeOrDefault() => _registration.GetRenderModeOrDefault();
+
+        [DynamicDependency(JsonSerialized, typeof(ServerComponentsActivityState))]
+        public Type? GetResolvedTypeOrNull() => _registration.GetResolvedTypeOrNull();
+    }
+
+    private sealed class WebAssemblyComponentsActivityStatePersistentServiceRegistration(IServiceProvider serviceProvider)
+        : IPersistentServiceRegistration
+    {
+        private readonly PersistentServiceRegistration<WebAssemblyComponentsActivityState> _registration = new(
+            serviceProvider.GetRequiredKeyedService<IComponentRenderMode>(
+                "Microsoft.AspNetCore.Components.ActivityState.WebAssembly"));
+
+        public string Assembly => _registration.Assembly;
+
+        public string FullTypeName => _registration.FullTypeName;
+
+        public IComponentRenderMode? GetRenderModeOrDefault() => _registration.GetRenderModeOrDefault();
+
+        [DynamicDependency(JsonSerialized, typeof(WebAssemblyComponentsActivityState))]
+        public Type? GetResolvedTypeOrNull() => _registration.GetResolvedTypeOrNull();
+    }
+
     private static bool IsMeterFactoryRegistered(IServiceCollection services)
     {
         foreach (var service in services)

src/Components/Components/src/RenderTree/Renderer.cs

@@ -112,7 +112,12 @@ public Renderer(IServiceProvider serviceProvider, ILoggerFactory loggerFactory,
         if (ComponentsActivitySource.IsSupported)
         {
             _componentsActivitySource = serviceProvider.GetService<ComponentsActivitySource>();
-            _componentsActivitySource?.Init(new ComponentsActivityLinkStore(this));
+            if (_componentsActivitySource is not null)
+            {
+                var activityLinkStore = new ComponentsActivityLinkStore(this);
+                _componentsActivitySource.Init(activityLinkStore);
+                serviceProvider.GetService<ComponentsActivityState>()?.Initialize(activityLinkStore);
+            }
         }
 
         ServiceProviderCascadingValueSuppliers = serviceProvider.GetService<ICascadingValueSupplier>() is null

src/Components/Components/test/ComponentsActivitySourceTest.cs

@@ -71,25 +71,22 @@ public void StartRouteActivity_CreatesAndStartsActivity()
     }
 
     [Fact]
-    public void StartEventActivity_CreatesAndStartsActivity()
+    public void StartEventActivity_WithSharedLinkStore_LinksToRoute()
     {
-        // Arrange
         var componentsActivitySource = new ComponentsActivitySource();
         var linkstore = new ComponentsActivityLinkStore(null);
         componentsActivitySource.Init(linkstore);
         var componentType = "TestComponent";
         var methodName = "OnClick";
         var attributeName = "onclick";
 
-        // First set up a circuit and route context
         linkstore.SetActivityContext(ComponentsActivityLinkStore.Circuit, default, new KeyValuePair<string, object>("aspnetcore.components.circuit.id", "test-circuit-id"));
-        componentsActivitySource.StartNavigateActivity("ParentComponent", "/parent");
+        var routeActivityHandle = componentsActivitySource.StartNavigateActivity("ParentComponent", "/parent");
+        componentsActivitySource.StopNavigateActivity(routeActivityHandle, null);
 
-        // Act
         var activityHandle = ComponentsActivitySource.StartHandleEventActivity(componentType, methodName, attributeName);
         var activity = activityHandle.Activity;
 
-        // Assert
         Assert.NotNull(activity);
         Assert.Equal(ComponentsActivitySource.OnEventName, activity.OperationName);
         Assert.Equal($"Event {attributeName} -> {componentType}.{methodName}", activity.DisplayName);
@@ -100,10 +97,33 @@ public void StartEventActivity_CreatesAndStartsActivity()
         Assert.Equal(attributeName, activity.GetTagItem("aspnetcore.components.attribute.name"));
         Assert.False(activity.IsStopped);
 
-        componentsActivitySource.StopNavigateActivity(activityHandle, null);
+        componentsActivitySource.StopHandleEventActivity(activityHandle, null);
         Assert.True(activity.IsStopped);
         Assert.Equal("test-circuit-id", activity.GetTagItem("aspnetcore.components.circuit.id"));
-        Assert.Empty(activity.Links);
+        Assert.Collection(activity.Links, link => Assert.Equal(routeActivityHandle.Activity.Context, link.Context));
+    }
+
+    [Fact]
+    public void StartEventActivity_WithSeparateLinkStores_DoesNotLinkToRoute()
+    {
+        var endpointActivitySource = new ComponentsActivitySource();
+        var endpointLinkStore = new ComponentsActivityLinkStore(null);
+        endpointActivitySource.Init(endpointLinkStore);
+
+        var routeActivityHandle = endpointActivitySource.StartNavigateActivity("ParentComponent", "/parent");
+        endpointActivitySource.StopNavigateActivity(routeActivityHandle, null);
+        Assert.NotNull(routeActivityHandle.Activity);
+
+        var circuitActivitySource = new ComponentsActivitySource();
+        var circuitLinkStore = new ComponentsActivityLinkStore(null);
+        circuitActivitySource.Init(circuitLinkStore);
+
+        var eventActivityHandle = ComponentsActivitySource.StartHandleEventActivity("TestComponent", "OnClick", "onclick");
+        circuitActivitySource.StopHandleEventActivity(eventActivityHandle, null);
+        Assert.NotNull(eventActivityHandle.Activity);
+
+        Assert.Empty(eventActivityHandle.Activity.Links);
+        Assert.DoesNotContain(eventActivityHandle.Activity.Links, link => link.Context == routeActivityHandle.Activity.Context);
     }
 
     [Fact]

src/Components/Components/test/ComponentsActivityStateTest.cs

@@ -0,0 +1,92 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Diagnostics;
+using Microsoft.AspNetCore.Components.Infrastructure;
+
+namespace Microsoft.AspNetCore.Components;
+
+public class ComponentsActivityStateTest
+{
+    [Fact]
+    public void CircuitInitialization_AppliesStateRestoredBeforeRendererInitialization()
+    {
+        var routeContext = CreateActivityContext();
+        var activityState = new ComponentsActivityState();
+        activityState.Apply(CreateState(routeContext, "/counter"));
+        var linkStore = new ComponentsActivityLinkStore(null);
+
+        activityState.Initialize(linkStore);
+
+        AssertRoute(linkStore, routeContext, "/counter");
+    }
+
+    [Fact]
+    public void NavigationUpdate_ReplacesRouteWhenInteractiveTreeHasNoRouter()
+    {
+        var initialContext = CreateActivityContext();
+        var updatedContext = CreateActivityContext();
+        var linkStore = new ComponentsActivityLinkStore(null);
+        var activityState = new ComponentsActivityState();
+        activityState.Initialize(linkStore);
+        activityState.Apply(CreateState(initialContext, "/initial"));
+
+        activityState.Apply(CreateState(updatedContext, "/updated"));
+
+        AssertRoute(linkStore, updatedContext, "/updated");
+    }
+
+    [Fact]
+    public void NavigationUpdate_ClearsRouteWhenStateHasNoRoute()
+    {
+        var linkStore = new ComponentsActivityLinkStore(null);
+        var activityState = new ComponentsActivityState();
+        activityState.Initialize(linkStore);
+        activityState.Apply(CreateState(CreateActivityContext(), "/counter"));
+
+        activityState.Apply(new ComponentsActivityPersistentStateUpdate(null));
+
+        Assert.False(linkStore.TryGetActivityContext(ComponentsActivityLinkStore.Route, out _, out _));
+    }
+
+    [Fact]
+    public void CircuitInitialization_IgnoresMalformedActivityContext()
+    {
+        var activityState = new ComponentsActivityState();
+        activityState.Apply(new ComponentsActivityPersistentStateUpdate(
+            new ComponentsActivityPersistentState("invalid", null, false, "/counter")));
+        var linkStore = new ComponentsActivityLinkStore(null);
+
+        activityState.Initialize(linkStore);
+
+        Assert.False(linkStore.TryGetActivityContext(ComponentsActivityLinkStore.Route, out _, out _));
+    }
+
+    private static ComponentsActivityPersistentStateUpdate CreateState(ActivityContext context, string route)
+        => new(new ComponentsActivityPersistentState(
+            $"00-{context.TraceId}-{context.SpanId}-{(byte)context.TraceFlags:x2}",
+            context.TraceState,
+            context.IsRemote,
+            route));
+
+    private static ActivityContext CreateActivityContext()
+        => new(
+            ActivityTraceId.CreateRandom(),
+            ActivitySpanId.CreateRandom(),
+            ActivityTraceFlags.Recorded,
+            traceState: "vendor=value",
+            isRemote: false);
+
+    private static void AssertRoute(
+        ComponentsActivityLinkStore linkStore,
+        ActivityContext expectedContext,
+        string expectedRoute)
+    {
+        Assert.True(linkStore.TryGetActivityContext(
+            ComponentsActivityLinkStore.Route,
+            out var context,
+            out var tag));
+        Assert.Equal(expectedContext, context);
+        Assert.Equal(new KeyValuePair<string, object>("aspnetcore.components.route", expectedRoute), tag);
+    }
+}

src/Components/Server/src/DependencyInjection/ServerRazorComponentsBuilderExtensions.cs

@@ -35,6 +35,9 @@ public static IServerSideBlazorBuilder AddInteractiveServerComponents(this IRazo
         ArgumentNullException.ThrowIfNull(builder, nameof(builder));
 
         builder.Services.AddServerSideBlazor(configure);
+        builder.Services.AddKeyedSingleton<IComponentRenderMode>(
+            "Microsoft.AspNetCore.Components.ActivityState.Server",
+            RenderMode.InteractiveServer);
         builder.Services.TryAddEnumerable(ServiceDescriptor.Singleton<RenderModeEndpointProvider, CircuitEndpointProvider>());
 
         return new DefaultServerSideBlazorBuilder(builder.Services);

src/Components/WebAssembly/Server/src/WebAssemblyRazorComponentsBuilderExtensions.cs

@@ -27,6 +27,9 @@ public static IRazorComponentsBuilder AddInteractiveWebAssemblyComponents(this I
     {
         ArgumentNullException.ThrowIfNull(builder, nameof(builder));
 
+        builder.Services.AddKeyedSingleton<IComponentRenderMode>(
+            "Microsoft.AspNetCore.Components.ActivityState.WebAssembly",
+            RenderMode.InteractiveWebAssembly);
         builder.Services.TryAddEnumerable(ServiceDescriptor.Singleton<RenderModeEndpointProvider, WebAssemblyEndpointProvider>());
         builder.Services.TryAddScoped<LazyAssemblyLoader>();
 

src/Components/WebAssembly/WebAssembly/src/Hosting/WebAssemblyHostBuilder.cs

@@ -358,6 +358,9 @@ internal void InitializeDefaultServices()
         // Register metrics and tracing when supported and not disabled by the feature switch.
         if (IsMeterSupported && IsMeterEnabled)
         {
+            Services.AddKeyedSingleton<IComponentRenderMode>(
+                "Microsoft.AspNetCore.Components.ActivityState.WebAssembly",
+                RenderMode.InteractiveWebAssembly);
             ComponentsMetricsServiceCollectionExtensions.AddComponentsMetrics(Services);
             ComponentsMetricsServiceCollectionExtensions.AddComponentsTracing(Services);
         }