From ab497ad6a77e5a5865e430cdba88147c7be063a3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89mile=20R=C3=A9?= <emile@getprobo.com>
Date: Tue, 10 Feb 2026 13:06:56 +0400
Subject: [PATCH] Add error boundaries for auth and assume
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Émile Ré <emile@getprobo.com>
---
 .../components/OrganizationErrorBoundary.tsx  | 22 ++++++++++++
 .../src/components/RootErrorBoundary.tsx      | 15 ++++++++
 apps/console/src/routes.tsx                   | 34 ++++---------------
 3 files changed, 43 insertions(+), 28 deletions(-)
 create mode 100644 apps/console/src/components/OrganizationErrorBoundary.tsx
 create mode 100644 apps/console/src/components/RootErrorBoundary.tsx

diff --git a/apps/console/src/components/OrganizationErrorBoundary.tsx b/apps/console/src/components/OrganizationErrorBoundary.tsx
new file mode 100644
index 000000000..70f7bfe97
--- /dev/null
+++ b/apps/console/src/components/OrganizationErrorBoundary.tsx
@@ -0,0 +1,22 @@
+import { AssumptionRequiredError, UnAuthenticatedError } from "@probo/relay";
+import { Navigate, useLocation, useRouteError } from "react-router";
+
+import { useOrganizationId } from "#/hooks/useOrganizationId";
+
+import { PageError } from "./PageError";
+
+export function OrganizationErrorBoundary() {
+  const error = useRouteError();
+  const location = useLocation();
+  const organizationId = useOrganizationId();
+
+  if (error instanceof UnAuthenticatedError) {
+    return <Navigate to="/auth/login" state={{ from: location.pathname }} />;
+  }
+
+  if (error instanceof AssumptionRequiredError) {
+    return <Navigate to="/auth/login" state={{ from: location.pathname, organizationId }} />;
+  }
+
+  return <PageError error={error instanceof Error ? error : new Error("unknown error")} />;
+}
diff --git a/apps/console/src/components/RootErrorBoundary.tsx b/apps/console/src/components/RootErrorBoundary.tsx
new file mode 100644
index 000000000..a6f698230
--- /dev/null
+++ b/apps/console/src/components/RootErrorBoundary.tsx
@@ -0,0 +1,15 @@
+import { UnAuthenticatedError } from "@probo/relay";
+import { Navigate, useLocation, useRouteError } from "react-router";
+
+import { PageError } from "./PageError";
+
+export function RootErrorBoundary() {
+  const error = useRouteError();
+  const location = useLocation();
+
+  if (error instanceof UnAuthenticatedError) {
+    return <Navigate to="/auth/login" state={{ from: location.pathname }} />;
+  }
+
+  return <PageError error={error instanceof Error ? error : new Error("unknown error")} />;
+}
diff --git a/apps/console/src/routes.tsx b/apps/console/src/routes.tsx
index cfcbdea6b..51589f5cb 100644
--- a/apps/console/src/routes.tsx
+++ b/apps/console/src/routes.tsx
@@ -1,9 +1,5 @@
 import { Role } from "@probo/helpers";
 import { lazy } from "@probo/react-lazy";
-import {
-  AssumptionRequiredError,
-  UnAuthenticatedError,
-} from "@probo/relay";
 import { type AppRoute, routeFromAppRoute } from "@probo/routes";
 import { CenteredLayout } from "@probo/ui";
 import { use } from "react";
@@ -11,10 +7,11 @@ import {
   createBrowserRouter,
   Navigate,
   redirect,
-  useRouteError,
 } from "react-router";
 
+import { OrganizationErrorBoundary } from "./components/OrganizationErrorBoundary";
 import { PageError } from "./components/PageError";
+import { RootErrorBoundary } from "./components/RootErrorBoundary";
 import { PageSkeleton } from "./components/skeletons/PageSkeleton";
 import { ViewerLayoutLoading } from "./pages/iam/memberships/ViewerLayoutLoading";
 import { compliancePageRoutes } from "./pages/organizations/compliance-page/routes";
@@ -38,24 +35,6 @@ import { statesOfApplicabilityRoutes } from "./routes/statesOfApplicabilityRoute
 import { taskRoutes } from "./routes/taskRoutes";
 import { vendorRoutes } from "./routes/vendorRoutes";
 
-/**
- * Top level error boundary
- */
-function ErrorBoundary() {
-  const error = useRouteError();
-
-  if (error instanceof UnAuthenticatedError) {
-    return <Navigate to="/auth/login" />;
-  }
-
-  if (error instanceof AssumptionRequiredError) {
-    // TODO redirect to right URL
-    return <Navigate to="/" />;
-  }
-
-  return <PageError error={error instanceof Error ? error : new Error("unknown error")} />;
-}
-
 const routes = [
   {
     path: "/auth",
@@ -101,7 +80,7 @@ const routes = [
   },
   {
     path: "/",
-    ErrorBoundary: ErrorBoundary,
+    ErrorBoundary: RootErrorBoundary,
     children: [
       {
         Component: lazy(() => import("./pages/iam/memberships/ViewerLayoutLoader")),
@@ -136,7 +115,7 @@ const routes = [
   },
   {
     path: "documents/signing-requests",
-    ErrorBoundary: ErrorBoundary,
+    ErrorBoundary: RootErrorBoundary,
     Component: lazy(
       () => import("./pages/DocumentSigningRequestsPage"),
     ),
@@ -146,7 +125,7 @@ const routes = [
     Component: lazy(
       () => import("./pages/organizations/employee/EmployeeLayoutLoader"),
     ),
-    ErrorBoundary: ErrorBoundary,
+    ErrorBoundary: OrganizationErrorBoundary,
     children: [
       {
         index: true,
@@ -157,7 +136,6 @@ const routes = [
       },
       {
         path: ":documentId",
-        ErrorBoundary: ErrorBoundary,
         Component: lazy(
           () =>
             import("./pages/organizations/employee/EmployeeDocumentSignaturePageLoader"),
@@ -170,7 +148,7 @@ const routes = [
     Component: lazy(
       () => import("./pages/iam/organizations/ViewerMembershipLayoutLoader"),
     ),
-    ErrorBoundary: ErrorBoundary,
+    ErrorBoundary: OrganizationErrorBoundary,
     children: [
       {
         path: "",