Improved exception handling in handler callbacks

10 files changed, 249 insertions, 45 deletions

diff --git a/src/Components/Components/src/NavigationManager.cs b/src/Components/Components/src/NavigationManager.cs
index 3e15622672..7e70f19ff2 100644
--- a/src/Components/Components/src/NavigationManager.cs
+++ b/src/Components/Components/src/NavigationManager.cs
@@ -309,7 +309,7 @@ public abstract class NavigationManager
         {
             if (handlerCount == 1)
             {
-                var handlerTask = _locationChangingHandlers[0](context);
+                var handlerTask = InvokeLocationChangingHandlerAsync(_locationChangingHandlers[0], context);
 
                 if (context.DidPreventNavigation)
                 {
@@ -333,7 +333,13 @@ public abstract class NavigationManager
 
                     for (var i = 0; i < handlerCount; i++)
                     {
-                        var handlerTask = locationChangingHandlersCopy[i](context);
+                        var handlerTask = InvokeLocationChangingHandlerAsync(locationChangingHandlersCopy[i], context);
+
+                        if (handlerTask.IsFaulted)
+                        {
+                            await handlerTask;
+                            return false; // Unreachable because the previous line will throw.
+                        }
 
                         if (context.DidPreventNavigation)
                         {
@@ -345,11 +351,12 @@ public abstract class NavigationManager
 
                     while (locationChangingTasks.Count != 0)
                     {
-                        var completedHandler = await Task.WhenAny(locationChangingTasks).WaitAsync(cancellationToken);
+                        var completedHandlerTask = await Task.WhenAny(locationChangingTasks).WaitAsync(cancellationToken);
 
-                        if (completedHandler.Exception is { } exception)
+                        if (completedHandlerTask.IsFaulted)
                         {
-                            throw exception;
+                            await completedHandlerTask;
+                            return false; // Unreachable because the previous line will throw.
                         }
 
                         if (context.DidPreventNavigation)
@@ -357,7 +364,7 @@ public abstract class NavigationManager
                             return false;
                         }
 
-                        locationChangingTasks.Remove(completedHandler);
+                        locationChangingTasks.Remove(completedHandlerTask);
                     }
                 }
                 finally
@@ -392,6 +399,18 @@ public abstract class NavigationManager
     }
 
     /// <summary>
+    /// Invokes the provided <paramref name="handler"/>, passing it the given <paramref name="context"/>.
+    /// This method can be overridden to analyze the state of the handler task even after
+    /// <see cref="NotifyLocationChangingAsync(string, string?, bool)"/> completes. For example, this can be useful for
+    /// processing exceptions thrown from handlers that continue running after the navigation ends.
+    /// </summary>
+    /// <param name="handler">The handler to invoke.</param>
+    /// <param name="context">The context to pass to the handler.</param>
+    /// <returns></returns>
+    protected virtual ValueTask InvokeLocationChangingHandlerAsync(Func<LocationChangingContext, ValueTask> handler, LocationChangingContext context)
+        => handler(context);
+
+    /// <summary>
     /// Sets whether navigation is currently locked. If it is, then implementations should not update <see cref="Uri"/> and call
     /// <see cref="NotifyLocationChanged(bool)"/> until they have first confirmed the navigation by calling
     /// <see cref="NotifyLocationChangingAsync(string, string?, bool)"/>.
diff --git a/src/Components/Components/src/PublicAPI.Unshipped.txt b/src/Components/Components/src/PublicAPI.Unshipped.txt
index 31dbec3f86..7f129aed4f 100644
--- a/src/Components/Components/src/PublicAPI.Unshipped.txt
+++ b/src/Components/Components/src/PublicAPI.Unshipped.txt
@@ -50,4 +50,5 @@ static Microsoft.AspNetCore.Components.EventCallbackFactoryBinderExtensions.Crea
 static Microsoft.AspNetCore.Components.EventCallbackFactoryBinderExtensions.CreateBinder(this Microsoft.AspNetCore.Components.EventCallbackFactory! factory, object! receiver, Microsoft.AspNetCore.Components.EventCallback<short?> setter, short? existingValue, System.Globalization.CultureInfo? culture = null) -> Microsoft.AspNetCore.Components.EventCallback<Microsoft.AspNetCore.Components.ChangeEventArgs!>
 static Microsoft.AspNetCore.Components.EventCallbackFactoryBinderExtensions.CreateBinder(this Microsoft.AspNetCore.Components.EventCallbackFactory! factory, object! receiver, Microsoft.AspNetCore.Components.EventCallback<string?> setter, string! existingValue, System.Globalization.CultureInfo? culture = null) -> Microsoft.AspNetCore.Components.EventCallback<Microsoft.AspNetCore.Components.ChangeEventArgs!>
 static Microsoft.AspNetCore.Components.EventCallbackFactoryBinderExtensions.CreateBinder<T>(this Microsoft.AspNetCore.Components.EventCallbackFactory! factory, object! receiver, Microsoft.AspNetCore.Components.EventCallback<T> setter, T existingValue, System.Globalization.CultureInfo? culture = null) -> Microsoft.AspNetCore.Components.EventCallback<Microsoft.AspNetCore.Components.ChangeEventArgs!>
+virtual Microsoft.AspNetCore.Components.NavigationManager.InvokeLocationChangingHandlerAsync(System.Func<Microsoft.AspNetCore.Components.Routing.LocationChangingContext!, System.Threading.Tasks.ValueTask>! handler, Microsoft.AspNetCore.Components.Routing.LocationChangingContext! context) -> System.Threading.Tasks.ValueTask
 virtual Microsoft.AspNetCore.Components.NavigationManager.SetNavigationLockState(bool value) -> void
diff --git a/src/Components/Components/test/NavigationManagerTest.cs b/src/Components/Components/test/NavigationManagerTest.cs
index 6ba560197c..9bf20eb58e 100644
--- a/src/Components/Components/test/NavigationManagerTest.cs
+++ b/src/Components/Components/test/NavigationManagerTest.cs
@@ -573,9 +573,9 @@ public class NavigationManagerTest
             {
                 await Task.Delay(System.Threading.Timeout.Infinite, context.CancellationToken);
             }
-            catch (TaskCanceledException e)
+            catch (TaskCanceledException ex)
             {
-                if (e.CancellationToken == context.CancellationToken)
+                if (ex.CancellationToken == context.CancellationToken)
                 {
                     tcs.SetResult();
                 }
@@ -638,9 +638,9 @@ public class NavigationManagerTest
                 await Task.Delay(System.Threading.Timeout.Infinite, context.CancellationToken);
                 Interlocked.Increment(ref completedHandlerCount);
             }
-            catch (TaskCanceledException e)
+            catch (TaskCanceledException ex)
             {
-                if (e.CancellationToken == context.CancellationToken)
+                if (ex.CancellationToken == context.CancellationToken)
                 {
                     lock (navigationManager)
                     {
@@ -746,8 +746,7 @@ public class NavigationManagerTest
         navigationManager.AddLocationChangingHandler(HandleLocationChanging_AllowNavigation);
 
         // Act/Assert
-        var aggregateException = await Assert.ThrowsAsync<AggregateException>(() => navigationManager.RunNotifyLocationChangingAsync($"{baseUri}/subdir1", null, false));
-        var ex = Assert.Single(aggregateException.InnerExceptions);
+        var ex = await Assert.ThrowsAsync<InvalidOperationException>(() => navigationManager.RunNotifyLocationChangingAsync($"{baseUri}/subdir1", null, false));
         Assert.Equal(exceptionMessage, ex.Message);
 
         async ValueTask HandleLocationChanging_AllowNavigation(LocationChangingContext context)
@@ -763,6 +762,46 @@ public class NavigationManagerTest
     }
 
     [Fact]
+    public async Task LocationChangingHandlers_CanThrowCatchableExceptionsAsynchronously_AfterNavigationEnds()
+    {
+        // Arrange
+        var baseUri = "scheme://host/";
+        var navigationManager = new TestNavigationManagerWithLocationChangingExceptionTracking(baseUri);
+        var exceptionMessage = "Thrown from a test handler";
+        var preventNavigationTcs = new TaskCompletionSource();
+        var throwExceptionTcs = new TaskCompletionSource();
+
+        navigationManager.AddLocationChangingHandler(HandleLocationChanging_PreventNavigation);
+        navigationManager.AddLocationChangingHandler(HandleLocationChanging_ThrowException);
+
+        // Act
+        var navigation1 = navigationManager.RunNotifyLocationChangingAsync($"{baseUri}/subdir1", null, false);
+        preventNavigationTcs.SetResult();
+        var navigation1Result = await navigation1;
+
+        // Assert
+        Assert.False(navigation1Result);
+        Assert.Empty(navigationManager.ExceptionsThrownFromLocationChangingHandlers);
+
+        throwExceptionTcs.SetResult();
+
+        var ex = Assert.Single(navigationManager.ExceptionsThrownFromLocationChangingHandlers);
+        Assert.Equal(exceptionMessage, ex.Message);
+
+        async ValueTask HandleLocationChanging_PreventNavigation(LocationChangingContext context)
+        {
+            await preventNavigationTcs.Task;
+            context.PreventNavigation();
+        }
+
+        async ValueTask HandleLocationChanging_ThrowException(LocationChangingContext context)
+        {
+            await throwExceptionTcs.Task;
+            throw new InvalidOperationException(exceptionMessage);
+        }
+    }
+
+    [Fact]
     public async Task LocationChangingHandlers_CannotCancelTheNavigationAsynchronously_UntilReturning()
     {
         // Arrange
@@ -835,4 +874,32 @@ public class NavigationManagerTest
         {
         }
     }
+
+    private class TestNavigationManagerWithLocationChangingExceptionTracking : TestNavigationManager
+    {
+        private readonly List<Exception> _exceptionsThrownFromLocationChangingHandlers = new();
+
+        public IReadOnlyList<Exception> ExceptionsThrownFromLocationChangingHandlers => _exceptionsThrownFromLocationChangingHandlers;
+
+        public TestNavigationManagerWithLocationChangingExceptionTracking(string baseUri = null, string uri = null)
+            : base(baseUri, uri)
+        {
+        }
+
+        protected override async ValueTask InvokeLocationChangingHandlerAsync(Func<LocationChangingContext, ValueTask> handler, LocationChangingContext context)
+        {
+            try
+            {
+                await handler(context);
+            }
+            catch (OperationCanceledException)
+            {
+                // Ignore exceptions caused by cancellations
+            }
+            catch (Exception ex)
+            {
+                _exceptionsThrownFromLocationChangingHandlers.Add(ex);
+            }
+        }
+    }
 }
diff --git a/src/Components/Server/src/Circuits/CircuitHost.cs b/src/Components/Server/src/Circuits/CircuitHost.cs
index e6382b253c..7611caf874 100644
--- a/src/Components/Server/src/Circuits/CircuitHost.cs
+++ b/src/Components/Server/src/Circuits/CircuitHost.cs
@@ -557,9 +557,10 @@ internal partial class CircuitHost : IAsyncDisposable
         }
         catch (Exception ex)
         {
-            // Exceptions thrown by location changing handlers should be treated like unhandled exceptions in user-code.
-            Log.LocationChangeFailedInCircuit(_logger, uri, CircuitId, ex);
-            await TryNotifyClientErrorAsync(Client, GetClientErrorMessage(ex, $"Location changing to '{uri}' failed."));
+            // An exception caught at this point was probably thrown inside the NavigationManager. Treat
+            // this like bad data.
+            Log.LocationChangeFailed(_logger, uri, CircuitId, ex);
+            await TryNotifyClientErrorAsync(Client, GetClientErrorMessage(ex, $"Location change to '{uri}' failed."));
             UnhandledException?.Invoke(this, new UnhandledExceptionEventArgs(ex, isTerminating: false));
         }
     }
diff --git a/src/Components/Server/src/Circuits/RemoteNavigationManager.cs b/src/Components/Server/src/Circuits/RemoteNavigationManager.cs
index 1418a9ea9e..f71da3a382 100644
--- a/src/Components/Server/src/Circuits/RemoteNavigationManager.cs
+++ b/src/Components/Server/src/Circuits/RemoteNavigationManager.cs
@@ -99,23 +99,40 @@ internal sealed partial class RemoteNavigationManager : NavigationManager, IHost
             {
                 var shouldContinueNavigation = await NotifyLocationChangingAsync(uri, options.HistoryEntryState, false);
 
-                if (shouldContinueNavigation)
-                {
-                    await _jsRuntime.InvokeVoidAsync(Interop.NavigateTo, uri, options);
-                }
-                else
+                if (!shouldContinueNavigation)
                 {
                     Log.NavigationCanceled(_logger, uri);
+                    return;
                 }
+
+                await _jsRuntime.InvokeVoidAsync(Interop.NavigateTo, uri, options);
             }
             catch (Exception ex)
             {
+                // We shouldn't ever reach this since exceptions thrown from handlers are caught in InvokeLocationChangingHandlerAsync.
+                // But if some other exception gets thrown, we still want to know about it.
                 Log.NavigationFailed(_logger, uri, ex);
-                UnhandledException?.Invoke(this, ex);
             }
         }
     }
 
+    protected override async ValueTask InvokeLocationChangingHandlerAsync(Func<LocationChangingContext, ValueTask> handler, LocationChangingContext context)
+    {
+        try
+        {
+            await handler(context);
+        }
+        catch (OperationCanceledException)
+        {
+            // Ignore exceptions caused by cancellations.
+        }
+        catch (Exception ex)
+        {
+            Log.NavigationFailed(_logger, context.TargetLocation, ex);
+            UnhandledException?.Invoke(this, ex);
+        }
+    }
+
     protected override void SetNavigationLockState(bool value)
     {
         if (_jsRuntime is null)
diff --git a/src/Components/WebAssembly/WebAssembly/src/Services/WebAssemblyNavigationManager.cs b/src/Components/WebAssembly/WebAssembly/src/Services/WebAssemblyNavigationManager.cs
index 9af705c5e5..fdd03ff773 100644
--- a/src/Components/WebAssembly/WebAssembly/src/Services/WebAssemblyNavigationManager.cs
+++ b/src/Components/WebAssembly/WebAssembly/src/Services/WebAssemblyNavigationManager.cs
@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Diagnostics.CodeAnalysis;
+using Microsoft.AspNetCore.Components.Routing;
 using Microsoft.Extensions.Logging;
 using Microsoft.JSInterop;
 using Interop = Microsoft.AspNetCore.Components.Web.BrowserNavigationManagerInterop;
@@ -64,22 +65,39 @@ internal sealed partial class WebAssemblyNavigationManager : NavigationManager
             {
                 var shouldContinueNavigation = await NotifyLocationChangingAsync(uri, options.HistoryEntryState, false);
 
-                if (shouldContinueNavigation)
-                {
-                    DefaultWebAssemblyJSRuntime.Instance.InvokeVoid(Interop.NavigateTo, uri, options);
-                }
-                else
+                if (!shouldContinueNavigation)
                 {
                     Log.NavigationCanceled(_logger, uri);
+                    return;
                 }
+
+                DefaultWebAssemblyJSRuntime.Instance.InvokeVoid(Interop.NavigateTo, uri, options);
             }
             catch (Exception ex)
             {
+                // We shouldn't ever reach this since exceptions thrown from handlers are caught in InvokeLocationChangingHandlerAsync.
+                // But if some other exception gets thrown, we still want to know about it.
                 Log.NavigationFailed(_logger, uri, ex);
             }
         }
     }
 
+    protected override async ValueTask InvokeLocationChangingHandlerAsync(Func<LocationChangingContext, ValueTask> handler, LocationChangingContext context)
+    {
+        try
+        {
+            await handler(context);
+        }
+        catch (OperationCanceledException)
+        {
+            // Ignore exceptions caused by cancellations.
+        }
+        catch (Exception ex)
+        {
+            Log.NavigationFailed(_logger, context.TargetLocation, ex);
+        }
+    }
+
     protected override void SetNavigationLockState(bool value)
         => DefaultWebAssemblyJSRuntime.Instance.InvokeVoid(Interop.SetHasLocationChangingListeners, value);
 
diff --git a/src/Components/WebView/WebView/src/IpcReceiver.cs b/src/Components/WebView/WebView/src/IpcReceiver.cs
index f1e9e5c44c..a1cb04e089 100644
--- a/src/Components/WebView/WebView/src/IpcReceiver.cs
+++ b/src/Components/WebView/WebView/src/IpcReceiver.cs
@@ -107,6 +107,6 @@ internal sealed class IpcReceiver
 
     private static void OnLocationChanging(PageContext pageContext, int callId, string uri, string? state, bool intercepted)
     {
-        pageContext.NavigationManager.HandleLocationChangingAsync(callId, uri, state, intercepted).Preserve();
+        pageContext.NavigationManager.HandleLocationChanging(callId, uri, state, intercepted);
     }
 }
diff --git a/src/Components/WebView/WebView/src/Services/WebViewNavigationManager.cs b/src/Components/WebView/WebView/src/Services/WebViewNavigationManager.cs
index ae8b347dca..b197833304 100644
--- a/src/Components/WebView/WebView/src/Services/WebViewNavigationManager.cs
+++ b/src/Components/WebView/WebView/src/Services/WebViewNavigationManager.cs
@@ -1,6 +1,7 @@
 // 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.Routing;
 using Microsoft.Extensions.Logging;
 
 namespace Microsoft.AspNetCore.Components.WebView.Services;
@@ -28,26 +29,30 @@ internal sealed partial class WebViewNavigationManager : NavigationManager
         NotifyLocationChanged(intercepted);
     }
 
-    public async ValueTask HandleLocationChangingAsync(int callId, string uri, string? state, bool intercepted)
+    public void HandleLocationChanging(int callId, string uri, string? state, bool intercepted)
     {
-        bool shouldContinueNavigation;
+        _ = HandleLocationChangingAsync();
 
-        try
+        async Task HandleLocationChangingAsync()
         {
-            shouldContinueNavigation = await NotifyLocationChangingAsync(uri, state, intercepted);
+            try
+            {
+                var shouldContinueNavigation = await NotifyLocationChangingAsync(uri, state, intercepted);
+
+                if (!shouldContinueNavigation)
+                {
+                    Log.NavigationCanceled(_logger, uri);
+                }
 
-            if (!shouldContinueNavigation)
+                _ipcSender.EndLocationChanging(callId, shouldContinueNavigation);
+            }
+            catch (Exception ex)
             {
-                Log.NavigationCanceled(_logger, uri);
+                // We shouldn't ever reach this since exceptions thrown from handlers are caught in InvokeLocationChangingHandlerAsync.
+                // But if some other exception gets thrown, we still want to know about it.
+                Log.NavigationFailed(_logger, uri, ex);
             }
         }
-        catch (Exception ex)
-        {
-            shouldContinueNavigation = false;
-            Log.NavigationFailed(_logger, uri, ex);
-        }
-
-        _ipcSender.EndLocationChanging(callId, shouldContinueNavigation);
     }
 
     protected override void NavigateToCore(string uri, NavigationOptions options)
@@ -60,22 +65,39 @@ internal sealed partial class WebViewNavigationManager : NavigationManager
             {
                 var shouldContinueNavigation = await NotifyLocationChangingAsync(uri, options.HistoryEntryState, false);
 
-                if (shouldContinueNavigation)
-                {
-                    _ipcSender.Navigate(uri, options);
-                }
-                else
+                if (!shouldContinueNavigation)
                 {
                     Log.NavigationCanceled(_logger, uri);
+                    return;
                 }
+
+                _ipcSender.Navigate(uri, options);
             }
             catch (Exception ex)
             {
+                // We shouldn't ever reach this since exceptions thrown from handlers are caught in InvokeLocationChangingHandlerAsync.
+                // But if some other exception gets thrown, we still want to know about it.
                 Log.NavigationFailed(_logger, uri, ex);
             }
         }
     }
 
+    protected override async ValueTask InvokeLocationChangingHandlerAsync(Func<LocationChangingContext, ValueTask> handler, LocationChangingContext context)
+    {
+        try
+        {
+            await handler(context);
+        }
+        catch (OperationCanceledException)
+        {
+            // Ignore exceptions caused by cancellations.
+        }
+        catch (Exception ex)
+        {
+            Log.NavigationFailed(_logger, context.TargetLocation, ex);
+        }
+    }
+
     private static partial class Log
     {
         [LoggerMessage(1, LogLevel.Debug, "Navigation canceled when changing the location to {Uri}", EventName = "NavigationCanceled")]
diff --git a/src/Components/test/E2ETest/Tests/RoutingTest.cs b/src/Components/test/E2ETest/Tests/RoutingTest.cs
index c71a6738d4..1ac1946385 100644
--- a/src/Components/test/E2ETest/Tests/RoutingTest.cs
+++ b/src/Components/test/E2ETest/Tests/RoutingTest.cs
@@ -1159,6 +1159,64 @@ public class RoutingTest : ServerTestBase<ToggleExecutionModeServerFixture<Progr
     }
 
     [Fact]
+    public void NavigationLock_CanRenderUIForExceptions_ProgrammaticNavigation()
+    {
+        SetUrlViaPushState("/");
+
+        var app = Browser.MountTestComponent<NavigationManagerComponent>();
+
+        // Add a navigation lock that blocks internal navigations
+        Browser.FindElement(By.Id("add-navigation-lock")).Click();
+        Browser.FindElement(By.CssSelector("#navigation-lock-0 > input.block-internal-navigation")).Click();
+        
+        var uriBeforeBlockedNavigation = Browser.FindElement(By.Id("test-info")).Text;
+
+        Browser.FindElement(By.Id("programmatic-navigation")).Click();
+
+        // The navigation lock has initiated its "location changing" handler and is displaying navigation controls
+        Browser.Exists(By.CssSelector("#navigation-lock-0 > div.blocking-controls"));
+
+        // The location was reverted to what it was before the navigation started
+        Browser.Equal(uriBeforeBlockedNavigation, () => app.FindElement(By.Id("test-info")).Text);
+
+        // Throw an exception for the current navigation
+        Browser.FindElement(By.CssSelector("#navigation-lock-0 > div.blocking-controls > button.navigation-throw-exception")).Click();
+
+        // The exception shows up in the UI
+        var errorUiElem = Browser.Exists(By.Id("blazor-error-ui"), TimeSpan.FromSeconds(10));
+        Assert.NotNull(errorUiElem);
+    }
+
+    [Fact]
+    public void NavigationLock_CanRenderUIForExceptions_InternalLinkNavigation()
+    {
+        SetUrlViaPushState("/");
+
+        var app = Browser.MountTestComponent<NavigationManagerComponent>();
+
+        // Add a navigation lock that blocks internal navigations
+        Browser.FindElement(By.Id("add-navigation-lock")).Click();
+        Browser.FindElement(By.CssSelector("#navigation-lock-0 > input.block-internal-navigation")).Click();
+        
+        var uriBeforeBlockedNavigation = Browser.FindElement(By.Id("test-info")).Text;
+
+        Browser.FindElement(By.Id("internal-link-navigation")).Click();
+
+        // The navigation lock has initiated its "location changing" handler and is displaying navigation controls
+        Browser.Exists(By.CssSelector("#navigation-lock-0 > div.blocking-controls"));
+
+        // The location was reverted to what it was before the navigation started
+        Browser.Equal(uriBeforeBlockedNavigation, () => app.FindElement(By.Id("test-info")).Text);
+
+        // Throw an exception for the current navigation
+        Browser.FindElement(By.CssSelector("#navigation-lock-0 > div.blocking-controls > button.navigation-throw-exception")).Click();
+
+        // The exception shows up in the UI
+        var errorUiElem = Browser.Exists(By.Id("blazor-error-ui"), TimeSpan.FromSeconds(10));
+        Assert.NotNull(errorUiElem);
+    }
+
+    [Fact]
     public void CanArriveAtRouteWithExtension()
     {
         // This is an odd test, but it's primarily here to verify routing for routeablecomponentfrompackage isn't available due to
diff --git a/src/Components/test/testassets/BasicTestApp/RouterTest/ConfigurableNavigationLock.razor b/src/Components/test/testassets/BasicTestApp/RouterTest/ConfigurableNavigationLock.razor
index 856ba6f15d..8cc31ef793 100644
--- a/src/Components/test/testassets/BasicTestApp/RouterTest/ConfigurableNavigationLock.razor
+++ b/src/Components/test/testassets/BasicTestApp/RouterTest/ConfigurableNavigationLock.razor
@@ -68,6 +68,7 @@
         catch (TaskCanceledException)
         {
             Log($"Canceling '{context.TargetLocation}'");
+            throw;
         }
         catch
         {