From 4e0d3987a0e96c439596826f3f5cc1857c838fdf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?F=C3=A9lix=20Robles?= <felix@sequentech.io>
Date: Wed, 2 Sep 2026 18:20:50 +0000
Subject: [PATCH] fix(tally): restore trustee status page

---
 .../resources/Tally/TallyCeremonyTrustees.tsx | 30 ++++++--------
 .../src/resources/Tally/utils.test.ts         | 39 ++++++++++++++++++-
 .../admin-portal/src/resources/Tally/utils.ts |  9 +++++
 3 files changed, 59 insertions(+), 19 deletions(-)

diff --git a/packages/admin-portal/src/resources/Tally/TallyCeremonyTrustees.tsx b/packages/admin-portal/src/resources/Tally/TallyCeremonyTrustees.tsx
index 8f94bc0e3e9..a781d48ff2c 100644
--- a/packages/admin-portal/src/resources/Tally/TallyCeremonyTrustees.tsx
+++ b/packages/admin-portal/src/resources/Tally/TallyCeremonyTrustees.tsx
@@ -17,12 +17,7 @@ import {useGetList, useGetOne, useRecordContext} from "react-admin"
 import {WizardStyles} from "@/components/styles/WizardStyles"
 import {RESTORE_PRIVATE_KEY} from "@/queries/RestorePrivateKey"
 import {useMutation} from "@apollo/client"
-import {
-    ICeremonyStatus,
-    ITallyExecutionStatus,
-    ITallyTrusteeStatus,
-    ITrusteeStatus,
-} from "@/types/ceremonies"
+import {ITallyCeremonyStatus, ITallyTrusteeStatus} from "@/types/ceremonies"
 import {Box} from "@mui/material"
 import {
     RestorePrivateKeyMutation,
@@ -34,6 +29,7 @@ import {
 import {AuthContext} from "@/providers/AuthContextProvider"
 import {useTenantStore} from "@/providers/TenantContextProvider"
 import {SettingsContext} from "@/providers/SettingsContextProvider"
+import {canTrusteeRestorePrivateKey} from "./utils"
 
 const WizardSteps = {
     Start: 0,
@@ -48,14 +44,14 @@ export const TallyCeremonyTrustees: React.FC = () => {
     const [tenantId] = useTenantStore()
     const authContext = useContext(AuthContext)
 
-    const [page, setPage] = useState<number>(WizardSteps.Start)
+    const [page, setPage] = useState<number>(WizardSteps.Status)
     const [selectedElections, setSelectedElections] = useState<string[]>([])
     const [selectedTrustees, setSelectedTrustees] = useState<boolean>(false)
     const [tally, setTally] = useState<Sequent_Backend_Tally_Session>()
     const [verified, setVerified] = useState<boolean>(false)
     const [uploading, setUploading] = useState<boolean>(false)
     const [errors, setErrors] = useState<String | null>(null)
-    const [trusteeStatus, setTrusteeStatus] = useState<ITrusteeStatus | null>(null)
+    const [trusteeStatus, setTrusteeStatus] = useState<ITallyTrusteeStatus | null>(null)
     const {globalSettings} = useContext(SettingsContext)
     const [isTallyCompleted, setIsTallyCompleted] = useState<boolean>(false)
 
@@ -122,25 +118,23 @@ export const TallyCeremonyTrustees: React.FC = () => {
 
     useEffect(() => {
         if (tallySessionExecutions) {
-            const username = authContext?.username
-            const ceremonyStatus: ICeremonyStatus | undefined = tallySessionExecutions?.[0]?.status
+            const trusteeName = authContext?.trustee
+            const ceremonyStatus: ITallyCeremonyStatus | undefined =
+                tallySessionExecutions?.[0]?.status
             const trusteeStatus = ceremonyStatus?.trustees.find(
-                (item) => item.name === username
+                (item) => item.name === trusteeName
             )?.status
             setTrusteeStatus(trusteeStatus ?? null)
         }
-    }, [tallySessionExecutions])
+    }, [authContext?.trustee, tallySessionExecutions])
 
     useEffect(() => {
         setPage(
-            !trusteeStatus && tally?.execution_status !== ITallyExecutionStatus.CANCELLED
+            canTrusteeRestorePrivateKey(trusteeStatus, tally?.execution_status)
                 ? WizardSteps.Start
-                : trusteeStatus === ITrusteeStatus.WAITING &&
-                    tally?.execution_status !== ITallyExecutionStatus.CANCELLED
-                  ? WizardSteps.Start
-                  : WizardSteps.Status
+                : WizardSteps.Status
         )
-    }, [trusteeStatus])
+    }, [tally?.execution_status, trusteeStatus])
 
     const CancelButton = styled(Button)`
         background-color: ${({theme}) => theme.palette.white};
diff --git a/packages/admin-portal/src/resources/Tally/utils.test.ts b/packages/admin-portal/src/resources/Tally/utils.test.ts
index 5a1c9d0190d..7ac096a995f 100644
--- a/packages/admin-portal/src/resources/Tally/utils.test.ts
+++ b/packages/admin-portal/src/resources/Tally/utils.test.ts
@@ -2,7 +2,44 @@
 //
 // SPDX-License-Identifier: AGPL-3.0-only
 
-import {orderItemsByIds} from "./utils"
+import {ITallyExecutionStatus, ITallyTrusteeStatus} from "@/types/ceremonies"
+import {canTrusteeRestorePrivateKey, orderItemsByIds} from "./utils"
+
+describe("canTrusteeRestorePrivateKey", () => {
+    it.each([ITallyExecutionStatus.STARTED, ITallyExecutionStatus.CONNECTED])(
+        "allows a waiting trustee while the tally accepts keys (%s)",
+        (executionStatus) => {
+            expect(canTrusteeRestorePrivateKey(ITallyTrusteeStatus.WAITING, executionStatus)).toBe(
+                true
+            )
+        }
+    )
+
+    it.each([
+        ITallyExecutionStatus.NOT_STARTED,
+        ITallyExecutionStatus.IN_PROGRESS,
+        ITallyExecutionStatus.AWAITING_INPUT,
+        ITallyExecutionStatus.SUCCESS,
+        ITallyExecutionStatus.CANCELLED,
+    ])("does not allow key restoration while the tally is %s", (executionStatus) => {
+        expect(canTrusteeRestorePrivateKey(ITallyTrusteeStatus.WAITING, executionStatus)).toBe(
+            false
+        )
+    })
+
+    it("does not allow a trustee whose key is already restored", () => {
+        expect(
+            canTrusteeRestorePrivateKey(
+                ITallyTrusteeStatus.KEY_RESTORED,
+                ITallyExecutionStatus.STARTED
+            )
+        ).toBe(false)
+    })
+
+    it("does not allow a user who is absent from the tally ceremony", () => {
+        expect(canTrusteeRestorePrivateKey(null, ITallyExecutionStatus.STARTED)).toBe(false)
+    })
+})
 
 describe("orderItemsByIds", () => {
     const items = [
diff --git a/packages/admin-portal/src/resources/Tally/utils.ts b/packages/admin-portal/src/resources/Tally/utils.ts
index 594e38389b9..41dee132d2a 100644
--- a/packages/admin-portal/src/resources/Tally/utils.ts
+++ b/packages/admin-portal/src/resources/Tally/utils.ts
@@ -5,6 +5,15 @@
 import {ParsedAnnotations, RunoffStatus} from "./types"
 import {Sequent_Backend_Candidate, Sequent_Backend_Contest} from "@/gql/graphql"
 import {ICandidate, IContest, ICountingAlgorithm} from "@sequentech/ui-core"
+import {ITallyExecutionStatus, ITallyTrusteeStatus} from "@/types/ceremonies"
+
+export const canTrusteeRestorePrivateKey = (
+    trusteeStatus: ITallyTrusteeStatus | null,
+    tallyExecutionStatus: string | null | undefined
+): boolean =>
+    trusteeStatus === ITallyTrusteeStatus.WAITING &&
+    (tallyExecutionStatus === ITallyExecutionStatus.STARTED ||
+        tallyExecutionStatus === ITallyExecutionStatus.CONNECTED)
 
 /**
  * Safely extracts the value from a GraphQL 'Maybe<T>' type.