From b6abae6abe089c43f95b8192d1c3f42ef6eed4b9 Mon Sep 17 00:00:00 2001
From: Ali Waseem <waseema393@gmail.com>
Date: Wed, 19 Aug 2026 07:51:56 -0600
Subject: [PATCH] fix(studio): only show restore completion once the restore
 has run (#48948)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Resolves
[FE-4144](https://linear.app/supabase/issue/FE-4144/restore-flow-shows-completion-before-restore-is-actually-done)

## Problem

`RestoringState` treated any `ACTIVE_HEALTHY` reading from the project
status endpoint as "restore finished". Right after a restore is
triggered the backend still reports the pre-restore status, so the first
poll could land on `ACTIVE_HEALTHY` and flip the UI to "Restoration
complete!" seconds into a restore that had barely started. `isCompleted`
was local state nothing reset and polling stopped on that first reading,
so the screen never self-corrected — "Return to project" then hung until
a manual refresh.

## Changes

- Gate completion on having observed the project leave the healthy
state, so a stale pre-restore reading is no longer mistaken for a
finished restore.
- Keep polling through an unconfirmed healthy reading instead of
stopping on it.
- `onConfirm` clears its loading flag rather than relying on the layout
to unmount the component.
- Component tests covering both the premature completion and the stuck
button.

## Needs validation

Not yet verified against a real restore — please confirm on staging
before merging. Worth checking in particular that a restore which
completes normally still reaches the completion screen.

There is one residual edge case left in place deliberately: if the
details endpoint reports `RESTORING` while the status endpoint reports
`ACTIVE_HEALTHY`, the UI now stays on "Restoration in progress" until
the details query catches up. Fixing that properly needs an
authoritative "restore initiated at" timestamp from the API, which does
not exist today.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **Bug Fixes**
- Improved project restoration tracking to prevent completion from being
reported prematurely.
- Restoration now correctly detects failures and stops polling when
appropriate.
- Restore status and saved transition information are cleared after
successful completion or failure.
- Confirmation actions now remain reliable while project details
refresh.
  - Restoring controls become usable again after the process finishes.
- Improved the restore menu trigger behavior for more consistent
interaction.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Alaister Young <10985857+alaister@users.noreply.github.com>
---
 .../ProjectLayout/RestoreFailedState.tsx      |   2 +-
 .../ProjectLayout/RestoringState.test.tsx     | 178 ++++++++++++++++++
 .../layouts/ProjectLayout/RestoringState.tsx  |  64 ++++---
 3 files changed, 218 insertions(+), 26 deletions(-)
 create mode 100644 apps/studio/components/layouts/ProjectLayout/RestoringState.test.tsx

diff --git a/apps/studio/components/layouts/ProjectLayout/RestoreFailedState.tsx b/apps/studio/components/layouts/ProjectLayout/RestoreFailedState.tsx
index 2855c96f40a..ade2b7c85b0 100644
--- a/apps/studio/components/layouts/ProjectLayout/RestoreFailedState.tsx
+++ b/apps/studio/components/layouts/ProjectLayout/RestoreFailedState.tsx
@@ -120,7 +120,7 @@ export const RestoreFailedState = () => {
               </ButtonTooltip>
 
               <DropdownMenu>
-                <DropdownMenuTrigger>
+                <DropdownMenuTrigger asChild>
                   <Button variant="default" className="w-7" icon={<MoreVertical />} />
                 </DropdownMenuTrigger>
                 <DropdownMenuContent className="w-72" align="end">
diff --git a/apps/studio/components/layouts/ProjectLayout/RestoringState.test.tsx b/apps/studio/components/layouts/ProjectLayout/RestoringState.test.tsx
new file mode 100644
index 00000000000..7829cd41a33
--- /dev/null
+++ b/apps/studio/components/layouts/ProjectLayout/RestoringState.test.tsx
@@ -0,0 +1,178 @@
+import { screen, waitFor } from '@testing-library/react'
+import userEvent from '@testing-library/user-event'
+import { platformComponents as components } from 'api-types'
+import { LOCAL_STORAGE_KEYS } from 'common'
+import { HttpResponse } from 'msw'
+import { beforeEach, describe, expect, test } from 'vitest'
+
+import { POLL_INTERVAL_MS, RestoringState } from './RestoringState'
+import { customRender } from '@/tests/lib/custom-render'
+import { addAPIMock } from '@/tests/lib/msw'
+
+type ProjectDetailResponse = components['schemas']['ProjectDetailResponse']
+type ProjectStatus = ProjectDetailResponse['status']
+type DownloadableBackupsResponse = components['schemas']['DownloadableBackupsResponse']
+
+const PROJECT_REF = 'default'
+
+const SETTLE_MS = 2000
+
+const createProject = (status: ProjectStatus): ProjectDetailResponse => ({
+  cloud_provider: 'AWS',
+  connectionString: `postgresql://postgres:password@db.${PROJECT_REF}.supabase.co:5432/postgres`,
+  db_host: `db.${PROJECT_REF}.supabase.co`,
+  high_availability: false,
+  id: 1,
+  infra_compute_size: 'micro',
+  inserted_at: '2026-01-01T00:00:00.000Z',
+  integration_source: null,
+  is_branch_enabled: false,
+  is_physical_backups_enabled: false,
+  name: 'Production',
+  organization_id: 1,
+  ref: PROJECT_REF,
+  region: 'us-east-1',
+  restUrl: `https://${PROJECT_REF}.supabase.co`,
+  status,
+  subscription_id: 'subscription-1',
+  updated_at: '2026-01-01T00:00:00.000Z',
+})
+
+const createStatusSequence = <T,>(statuses: [T, ...T[]]) => {
+  const queue = [...statuses]
+  return () => (queue.length > 1 ? queue.shift()! : queue[0])
+}
+
+const mockProjectDetail = (statuses: [ProjectStatus, ...ProjectStatus[]]) => {
+  const nextStatus = createStatusSequence(statuses)
+  addAPIMock({
+    method: 'get',
+    path: '/platform/projects/:ref',
+    response: () => HttpResponse.json<ProjectDetailResponse>(createProject(nextStatus())),
+  })
+}
+
+const mockProjectStatus = (statuses: [ProjectStatus, ...ProjectStatus[]]) => {
+  const nextStatus = createStatusSequence(statuses)
+  addAPIMock({
+    method: 'get',
+    path: '/platform/projects/:ref/status',
+    response: () => HttpResponse.json({ status: nextStatus() }),
+  })
+}
+
+const mockDownloadableBackups = () => {
+  addAPIMock({
+    method: 'get',
+    path: '/platform/database/:ref/backups/downloadable-backups',
+    response: () =>
+      HttpResponse.json<DownloadableBackupsResponse>({
+        backups: [
+          {
+            id: 1,
+            inserted_at: '2026-01-01T00:00:00.000Z',
+            isPhysicalBackup: false,
+            project_id: 1,
+            status: 'COMPLETED',
+          },
+        ],
+        status: 'ok',
+      }),
+  })
+}
+
+describe('RestoringState', () => {
+  beforeEach(() => {
+    window.localStorage.clear()
+    mockDownloadableBackups()
+  })
+
+  test('shows restoration in progress while the backend reports RESTORING', async () => {
+    mockProjectDetail(['RESTORING'])
+    mockProjectStatus(['RESTORING'])
+
+    customRender(<RestoringState />)
+
+    expect(await screen.findByText('Restoration in progress')).toBeInTheDocument()
+    expect(screen.queryByText('Restoration complete!')).not.toBeInTheDocument()
+  })
+
+  test('shows restoration complete once the project goes healthy after restoring', async () => {
+    mockProjectDetail(['RESTORING'])
+    mockProjectStatus(['RESTORING', 'ACTIVE_HEALTHY'])
+
+    customRender(<RestoringState />)
+
+    expect(await screen.findByText('Restoration in progress')).toBeInTheDocument()
+    expect(
+      await screen.findByText('Restoration complete!', {}, { timeout: POLL_INTERVAL_MS + 2000 })
+    ).toBeInTheDocument()
+  })
+
+  test('does not claim completion when the first status poll still reports the pre-restore healthy state', async () => {
+    mockProjectDetail(['RESTORING'])
+    mockProjectStatus(['ACTIVE_HEALTHY', 'RESTORING'])
+
+    customRender(<RestoringState />)
+
+    await expect(
+      screen.findByText('Restoration complete!', {}, { timeout: SETTLE_MS })
+    ).rejects.toThrow()
+    expect(screen.getByText('Restoration in progress')).toBeInTheDocument()
+  })
+
+  test(
+    'still reaches completion after an early healthy reading, once the restore actually finishes',
+    async () => {
+      mockProjectDetail(['RESTORING'])
+      mockProjectStatus(['ACTIVE_HEALTHY', 'RESTORING', 'ACTIVE_HEALTHY'])
+
+      customRender(<RestoringState />)
+
+      expect(await screen.findByText('Restoration in progress')).toBeInTheDocument()
+      expect(
+        await screen.findByText(
+          'Restoration complete!',
+          {},
+          { timeout: POLL_INTERVAL_MS * 3 + SETTLE_MS }
+        )
+      ).toBeInTheDocument()
+    },
+    POLL_INTERVAL_MS * 3 + SETTLE_MS * 2
+  )
+
+  test('clears the persisted restore start time when the restore fails', async () => {
+    const storageKey = LOCAL_STORAGE_KEYS.PROJECT_RESTORING_STARTED_AT(PROJECT_REF)
+    window.localStorage.setItem(storageKey, String(Date.now()))
+    mockProjectDetail(['RESTORING'])
+    mockProjectStatus(['RESTORE_FAILED'])
+
+    customRender(<RestoringState />)
+
+    await waitFor(() => {
+      expect(window.localStorage.getItem(storageKey)).toBeNull()
+    })
+  })
+
+  test(
+    'does not leave Return to project stuck loading when the details still report restoring',
+    async () => {
+      mockProjectDetail(['RESTORING'])
+      mockProjectStatus(['RESTORING', 'ACTIVE_HEALTHY'])
+
+      customRender(<RestoringState />)
+
+      const returnToProject = await screen.findByRole(
+        'button',
+        { name: 'Return to project' },
+        { timeout: POLL_INTERVAL_MS + SETTLE_MS }
+      )
+      await userEvent.click(returnToProject)
+
+      await waitFor(() => {
+        expect(returnToProject).not.toBeDisabled()
+      })
+    },
+    POLL_INTERVAL_MS + SETTLE_MS * 2
+  )
+})
diff --git a/apps/studio/components/layouts/ProjectLayout/RestoringState.tsx b/apps/studio/components/layouts/ProjectLayout/RestoringState.tsx
index a223931bb4b..d90039e9f04 100644
--- a/apps/studio/components/layouts/ProjectLayout/RestoringState.tsx
+++ b/apps/studio/components/layouts/ProjectLayout/RestoringState.tsx
@@ -20,12 +20,14 @@ import {
 } from '@/lib/project-transition-state'
 import { getRestoreLongRunningThresholdMinutes } from '@/lib/restore-estimate'
 
+export const POLL_INTERVAL_MS = 4000
+
 export const RestoringState = () => {
   const { ref } = useParams()
   const { data: project } = useSelectedProjectQuery()
 
-  const [loading, setLoading] = useState(false)
-  const [isCompleted, setIsCompleted] = useState(false)
+  const [isConfirming, setIsConfirming] = useState(false)
+  const [hasLeftHealthyState, setHasLeftHealthyState] = useState(false)
   const restoreStateStartStorageKey = ref
     ? LOCAL_STORAGE_KEYS.PROJECT_RESTORING_STARTED_AT(ref)
     : null
@@ -42,20 +44,35 @@ export const RestoringState = () => {
 
   const { invalidateProjectDetailsQuery } = useInvalidateProjectDetailsQuery()
 
-  const { data: projectStatusData, isSuccess: isProjectStatusSuccess } = useProjectStatusQuery(
+  const { data: projectStatusData } = useProjectStatusQuery(
     { projectRef: ref },
     {
       enabled: project?.status !== PROJECT_STATUS.ACTIVE_HEALTHY,
       refetchInterval: (query) => {
-        const data = query.state.data
-        return data?.status === PROJECT_STATUS.ACTIVE_HEALTHY ||
-          data?.status === PROJECT_STATUS.RESTORE_FAILED
-          ? false
-          : 4000
+        const status = query.state.data?.status
+        if (status === PROJECT_STATUS.RESTORE_FAILED) return false
+        if (status === PROJECT_STATUS.ACTIVE_HEALTHY && hasLeftHealthyState) return false
+        return POLL_INTERVAL_MS
       },
     }
   )
 
+  const projectStatus = projectStatusData?.status
+
+  // Right after a restore is triggered the status endpoint can still report the stale
+  // pre-restore ACTIVE_HEALTHY, so completion is only trusted once the status has been
+  // observed leaving the healthy state.
+  if (
+    !hasLeftHealthyState &&
+    projectStatus !== undefined &&
+    projectStatus !== PROJECT_STATUS.ACTIVE_HEALTHY
+  ) {
+    setHasLeftHealthyState(true)
+  }
+
+  const hasRestoreFailed = projectStatus === PROJECT_STATUS.RESTORE_FAILED
+  const isCompleted = hasLeftHealthyState && projectStatus === PROJECT_STATUS.ACTIVE_HEALTHY
+
   const { mutate: downloadBackup, isPending: isDownloading } = useBackupDownloadMutation({
     onSuccess: (res) => {
       const { fileUrl } = res
@@ -77,28 +94,25 @@ export const RestoringState = () => {
   }
 
   const onConfirm = async () => {
-    if (!project) return console.error('Project is required')
-    setLoading(true)
-    if (ref) await invalidateProjectDetailsQuery(ref)
+    if (!ref) return console.error('Project ref is required')
+    setIsConfirming(true)
+    try {
+      await invalidateProjectDetailsQuery(ref)
+    } finally {
+      setIsConfirming(false)
+    }
   }
 
   useEffect(() => {
-    if (!isProjectStatusSuccess) return
+    if (!isCompleted && !hasRestoreFailed) return
 
-    if (projectStatusData.status === PROJECT_STATUS.ACTIVE_HEALTHY) {
-      if (restoreStateStartStorageKey) {
-        clearPersistedTransitionStartTime(restoreStateStartStorageKey)
-      }
-      setIsCompleted(true)
-    } else if (projectStatusData.status === PROJECT_STATUS.RESTORE_FAILED) {
-      if (restoreStateStartStorageKey) {
-        clearPersistedTransitionStartTime(restoreStateStartStorageKey)
-      }
-      if (ref) void invalidateProjectDetailsQuery(ref)
+    if (restoreStateStartStorageKey) {
+      clearPersistedTransitionStartTime(restoreStateStartStorageKey)
     }
+    if (hasRestoreFailed && ref) void invalidateProjectDetailsQuery(ref)
   }, [
-    isProjectStatusSuccess,
-    projectStatusData,
+    isCompleted,
+    hasRestoreFailed,
     restoreStateStartStorageKey,
     ref,
     invalidateProjectDetailsQuery,
@@ -121,7 +135,7 @@ export const RestoringState = () => {
               </div>
             </div>
             <div className="border-t border-overlay flex items-center justify-end py-4 px-8">
-              <Button disabled={loading} loading={loading} onClick={onConfirm}>
+              <Button disabled={isConfirming} loading={isConfirming} onClick={onConfirm}>
                 Return to project
               </Button>
             </div>