fix: close workflow transition review gaps

Co-authored-by: Amp <amp@ampcode.com>
Amp-Thread-ID: https://ampcode.com/threads/T-019fb985-aceb-70a8-96f3-776224f1bf38

Author: dante01yoon

src/platform/workflow/persistence/composables/useWorkflowPersistenceV2.test.ts

@@ -8,10 +8,7 @@ import { WORKSPACE_STORAGE_KEYS } from '@/platform/workspace/workspaceConstants'
 import { useTeamWorkspaceStore } from '@/platform/workspace/stores/teamWorkspaceStore'
 import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
 import { StorageKeys } from '../base/storageKeys'
-import {
-  prepareWorkflowWorkspaceTransition,
-  writePayload
-} from '../base/storageIO'
+import * as storageIO from '../base/storageIO'
 import { useWorkflowDraftStoreV2 } from '../stores/workflowDraftStoreV2'
 import { useWorkflowPersistenceV2 } from './useWorkflowPersistenceV2'
 
@@ -680,7 +677,7 @@ describe('useWorkflowPersistenceV2', () => {
     )
     expect(localStorage.getItem(sourcePayloadKey)).toBeNull()
 
-    const cancelTransition = prepareWorkflowWorkspaceTransition()
+    const cancelTransition = storageIO.prepareWorkflowWorkspaceTransition()
     sessionStorage.setItem(
       WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE,
       JSON.stringify({ id: destinationWorkspaceId, type: 'team' })
@@ -720,16 +717,46 @@ describe('useWorkflowPersistenceV2', () => {
     expect(localStorage).toHaveLength(0)
     expect(sessionStorage).toHaveLength(0)
     expect(
-      writePayload('workspace-a', 'blocked', { data: '{}', updatedAt: 1 })
+      storageIO.writePayload('workspace-a', 'blocked', {
+        data: '{}',
+        updatedAt: 1
+      })
     ).toBe(false)
 
     onUserResolved({ id: 'user-a' })
 
     expect(
-      writePayload('workspace-a', 'resumed', { data: '{}', updatedAt: 2 })
+      storageIO.writePayload('workspace-a', 'resumed', {
+        data: '{}',
+        updatedAt: 2
+      })
     ).toBe(true)
   })
 
+  it('cancels stale workspace-readiness watchers across authentication episodes', async () => {
+    distributionMocks.isCloud = true
+    featureFlagMocks.teamWorkspacesEnabled = true
+    const completeTransitionSpy = vi.spyOn(
+      storageIO,
+      'completeWorkflowLogoutTransition'
+    )
+    mountWorkflowPersistence()
+
+    const onLogout = currentUserMocks.onUserLogout.mock.calls[0][0]
+    const onUserResolved = currentUserMocks.onUserResolved.mock.calls[0][0]
+    onLogout()
+    onUserResolved({ id: 'user-b' })
+    onLogout()
+    onUserResolved({ id: 'user-c' })
+
+    const teamWorkspaceStore = useTeamWorkspaceStore()
+    teamWorkspaceStore.activeWorkspaceId = 'workspace-c'
+    teamWorkspaceStore.initState = 'ready'
+    await nextTick()
+
+    expect(completeTransitionSpy).toHaveBeenCalledOnce()
+  })
+
   it('waits for workspace readiness and drops pending pre-logout edits', async () => {
     distributionMocks.isCloud = true
     featureFlagMocks.teamWorkspacesEnabled = true

src/platform/workflow/persistence/composables/useWorkflowPersistenceV2.ts

@@ -62,6 +62,12 @@ export function useWorkflowPersistenceV2() {
   const { onUserLogout, onUserResolved } = useCurrentUser()
   const { flags } = useFeatureFlags()
   const teamWorkspaceStore = useTeamWorkspaceStore()
+  let stopWorkspaceReadinessWatcher: (() => void) | undefined
+
+  function stopPendingWorkspaceReadinessWatcher(): void {
+    stopWorkspaceReadinessWatcher?.()
+    stopWorkspaceReadinessWatcher = undefined
+  }
 
   // Run migration on module load, passing clientId for tab state migration
   migrateV1toV2(undefined, api.clientId ?? api.initialClientId ?? undefined)
@@ -139,23 +145,34 @@ export function useWorkflowPersistenceV2() {
 
   onUserLogout(() => {
     if (!isCloud) return
+    stopPendingWorkspaceReadinessWatcher()
     debouncedPersist.cancel()
     prepareWorkflowLogoutTransition()
     clearAllWorkflowStorage()
   })
   onUserResolved(() => {
     if (!isCloud) return
+    stopPendingWorkspaceReadinessWatcher()
     if (!flags.teamWorkspacesEnabled) {
       completeWorkflowLogoutTransition()
       return
     }
 
-    whenever(
-      () =>
-        teamWorkspaceStore.initState === 'ready' &&
-        teamWorkspaceStore.activeWorkspaceId !== null,
-      completeWorkflowLogoutTransition,
-      { immediate: true, once: true }
+    const isWorkspaceReady = () =>
+      teamWorkspaceStore.initState === 'ready' &&
+      teamWorkspaceStore.activeWorkspaceId !== null
+    if (isWorkspaceReady()) {
+      completeWorkflowLogoutTransition()
+      return
+    }
+
+    stopWorkspaceReadinessWatcher = whenever(
+      isWorkspaceReady,
+      () => {
+        stopWorkspaceReadinessWatcher = undefined
+        completeWorkflowLogoutTransition()
+      },
+      { once: true }
     )
   })
 
@@ -282,6 +299,7 @@ export function useWorkflowPersistenceV2() {
     api.removeEventListener('graphChanged', debouncedPersist)
     unregisterPersistenceFlush()
     debouncedPersist.cancel()
+    stopPendingWorkspaceReadinessWatcher()
   })
 
   // Restore workflow tabs states

src/platform/workspace/stores/useWorkspaceAuth.test.ts

@@ -1200,6 +1200,10 @@ describe('useWorkspaceAuthStore', () => {
       const store = useWorkspaceAuthStore()
       const { currentWorkspace, workspaceToken } = storeToRefs(store)
       let workspaceWhenRevocationHandled: string | null = null
+      const cancelWorkflowTransition = vi.fn()
+      mockPrepareWorkflowWorkspaceTransition.mockReturnValue(
+        cancelWorkflowTransition
+      )
       mockForgetRevokedActiveWorkspace.mockImplementation(() => {
         workspaceWhenRevocationHandled = sessionStorage.getItem(
           WORKSPACE_STORAGE_KEYS.CURRENT_WORKSPACE
@@ -1222,6 +1226,8 @@ describe('useWorkspaceAuthStore', () => {
       expect(workspaceWhenRevocationHandled).toBe(
         JSON.stringify(mockWorkspaceWithRole)
       )
+      expect(cancelWorkflowTransition).toHaveBeenCalledOnce()
+      expect(mockReload).not.toHaveBeenCalled()
     })
   })
 

src/platform/workspace/stores/workspaceAuthStore.ts

@@ -981,7 +981,9 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
 
   function endWorkspaceSession(revokedWorkspaceId?: string): void {
     const hadContext = currentWorkspace.value !== null
-    if (hadContext) prepareWorkflowWorkspaceTransition()
+    const cancelWorkflowTransition = hadContext
+      ? prepareWorkflowWorkspaceTransition()
+      : undefined
     const revokedWorkspaceHandled = revokedWorkspaceId
       ? useTeamWorkspaceStore().forgetRevokedActiveWorkspace(revokedWorkspaceId)
       : false
@@ -991,7 +993,9 @@ export const useWorkspaceAuthStore = defineStore('workspaceAuth', () => {
       : hadContext
     if (shouldReload) {
       window.location.reload()
+      return
     }
+    cancelWorkflowTransition?.()
   }
 
   return {