fix: model workflow storage transitions explicitly

src/composables/auth/useAuthActions.test.ts

@@ -43,6 +43,7 @@ const mockBillingState = vi.hoisted(() => ({
   canAccessSubscriptionFeatures: false
 }))
 const mockClearAllWorkflowStorage = vi.hoisted(() => vi.fn())
+const mockPrepareWorkflowLogoutTransition = vi.hoisted(() => vi.fn())
 
 const knownAuthErrorCodes = new Set([
   'auth/invalid-credential',
@@ -77,7 +78,8 @@ vi.mock('@/platform/updates/common/toastStore', () => ({
 }))
 
 vi.mock('@/platform/workflow/persistence/base/storageIO', () => ({
-  clearAllWorkflowStorage: mockClearAllWorkflowStorage
+  clearAllWorkflowStorage: mockClearAllWorkflowStorage,
+  prepareWorkflowLogoutTransition: mockPrepareWorkflowLogoutTransition
 }))
 
 vi.mock('@/platform/workflow/management/stores/workflowStore', () => ({
@@ -214,12 +216,14 @@ describe('useAuthActions.logout', () => {
 
     await logout()
 
-    expect(mockClearAllWorkflowStorage).toHaveBeenCalledExactlyOnceWith({
-      blockWrites: true
-    })
+    expect(mockPrepareWorkflowLogoutTransition).toHaveBeenCalledOnce()
+    expect(mockClearAllWorkflowStorage).toHaveBeenCalledExactlyOnceWith()
     expect(mockAuthStore.logout.mock.invocationCallOrder[0]).toBeLessThan(
-      mockClearAllWorkflowStorage.mock.invocationCallOrder[0]
+      mockPrepareWorkflowLogoutTransition.mock.invocationCallOrder[0]
     )
+    expect(
+      mockPrepareWorkflowLogoutTransition.mock.invocationCallOrder[0]
+    ).toBeLessThan(mockClearAllWorkflowStorage.mock.invocationCallOrder[0])
     expect(
       mockClearAllWorkflowStorage.mock.invocationCallOrder[0]
     ).toBeLessThan(navigationSpy.mock.invocationCallOrder[0])
@@ -231,6 +235,7 @@ describe('useAuthActions.logout', () => {
 
     await logout()
 
+    expect(mockPrepareWorkflowLogoutTransition).not.toHaveBeenCalled()
     expect(mockClearAllWorkflowStorage).not.toHaveBeenCalled()
   })
 

src/composables/auth/useAuthActions.ts

@@ -10,7 +10,10 @@ import { isCloud } from '@/platform/distribution/types'
 import { useTelemetry } from '@/platform/telemetry'
 import type { AuthFlowAction } from '@/platform/telemetry/types'
 import { useToastStore } from '@/platform/updates/common/toastStore'
-import { clearAllWorkflowStorage } from '@/platform/workflow/persistence/base/storageIO'
+import {
+  clearAllWorkflowStorage,
+  prepareWorkflowLogoutTransition
+} from '@/platform/workflow/persistence/base/storageIO'
 import { useWorkflowService } from '@/platform/workflow/core/services/workflowService'
 import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
 import { useDialogService } from '@/services/dialogService'
@@ -113,7 +116,10 @@ export const useAuthActions = () => {
     }
 
     await authStore.logout()
-    if (isCloud) clearAllWorkflowStorage({ blockWrites: true })
+    if (isCloud) {
+      prepareWorkflowLogoutTransition()
+      clearAllWorkflowStorage()
+    }
 
     toastStore.add({
       severity: 'success',

src/platform/workflow/persistence/base/storageIO.test.ts

@@ -337,10 +337,11 @@ describe('storageIO', () => {
       expect(sessionStorage.getItem('unrelated')).toBe('keep')
     })
 
-    it('blocks workflow writes after cleanup starts', async () => {
+    it('blocks workflow writes during logout cleanup', async () => {
       const isolatedStorageIO = await import('./storageIO')
 
-      isolatedStorageIO.clearAllWorkflowStorage({ blockWrites: true })
+      isolatedStorageIO.prepareWorkflowLogoutTransition()
+      isolatedStorageIO.clearAllWorkflowStorage()
 
       expect(isolatedStorageIO.isStorageAvailable()).toBe(false)
       expect(
@@ -387,7 +388,7 @@ describe('storageIO', () => {
     })
   })
 
-  describe('clearWorkflowRestoreState', () => {
+  describe('workflow storage transitions', () => {
     it('blocks writes and clears restore state when a persistence flush fails', async () => {
       const isolatedStorageIO = await import('./storageIO')
       localStorage.setItem('workflow', '{}')
@@ -421,6 +422,98 @@ describe('storageIO', () => {
       consoleWarnSpy.mockRestore()
     })
 
+    it('resumes writes when a workspace transition is cancelled', async () => {
+      const isolatedStorageIO = await import('./storageIO')
+
+      const cancelTransition =
+        isolatedStorageIO.prepareWorkflowWorkspaceTransition()
+      expect(
+        isolatedStorageIO.writePayload('ws-1', 'blocked', {
+          data: '{}',
+          updatedAt: 1
+        })
+      ).toBe(false)
+
+      cancelTransition()
+
+      expect(
+        isolatedStorageIO.writePayload('ws-1', 'resumed', {
+          data: '{}',
+          updatedAt: 2
+        })
+      ).toBe(true)
+    })
+
+    it('preserves unavailable storage when a workspace transition is cancelled', async () => {
+      const isolatedStorageIO = await import('./storageIO')
+
+      isolatedStorageIO.markStorageUnavailable()
+      const cancelTransition =
+        isolatedStorageIO.prepareWorkflowWorkspaceTransition()
+      cancelTransition()
+
+      expect(isolatedStorageIO.isStorageAvailable()).toBe(false)
+      expect(
+        isolatedStorageIO.writePayload('ws-1', 'draft', {
+          data: '{}',
+          updatedAt: 1
+        })
+      ).toBe(false)
+    })
+
+    it('only resumes logout transitions after authentication recovers', async () => {
+      const isolatedStorageIO = await import('./storageIO')
+
+      isolatedStorageIO.prepareWorkflowWorkspaceTransition()
+      isolatedStorageIO.completeWorkflowLogoutTransition()
+      expect(isolatedStorageIO.isStorageAvailable()).toBe(false)
+
+      isolatedStorageIO.prepareWorkflowLogoutTransition()
+      isolatedStorageIO.completeWorkflowLogoutTransition()
+
+      expect(isolatedStorageIO.isStorageAvailable()).toBe(true)
+    })
+
+    it('preserves unavailable storage after logout recovery', async () => {
+      const isolatedStorageIO = await import('./storageIO')
+
+      isolatedStorageIO.markStorageUnavailable()
+      isolatedStorageIO.prepareWorkflowLogoutTransition()
+      isolatedStorageIO.completeWorkflowLogoutTransition()
+
+      expect(isolatedStorageIO.isStorageAvailable()).toBe(false)
+    })
+
+    it('does not let a duplicate transition cancel the owner transition', async () => {
+      const isolatedStorageIO = await import('./storageIO')
+
+      const cancelOwner = isolatedStorageIO.prepareWorkflowWorkspaceTransition()
+      const cancelDuplicate =
+        isolatedStorageIO.prepareWorkflowWorkspaceTransition()
+      cancelDuplicate()
+      expect(isolatedStorageIO.isStorageAvailable()).toBe(false)
+
+      cancelOwner()
+      expect(isolatedStorageIO.isStorageAvailable()).toBe(true)
+    })
+
+    it('keeps reads available while a workspace transition blocks writes', async () => {
+      const isolatedStorageIO = await import('./storageIO')
+
+      isolatedStorageIO.writePayload('ws-1', 'draft', {
+        data: '{}',
+        updatedAt: 1
+      })
+      const cancelTransition =
+        isolatedStorageIO.prepareWorkflowWorkspaceTransition()
+
+      expect(isolatedStorageIO.readPayload('ws-1', 'draft')).not.toBeNull()
+      expect(isolatedStorageIO.getPayloadKeys('ws-1')).toContain('draft')
+      cancelTransition()
+    })
+  })
+
+  describe('clearWorkflowRestoreState', () => {
     it('clears cross-workspace restore state without deleting scoped drafts', () => {
       localStorage.setItem('Comfy.Workflow.DraftIndex.v2:ws-1', '{}')
       localStorage.setItem('Comfy.Workflow.Draft.v2:ws-1:abc', '{}')

src/platform/workflow/persistence/base/storageIO.ts

@@ -12,9 +12,25 @@ import type {
 } from './draftTypes'
 import { StorageKeys } from './storageKeys'
 
-/** Flag indicating if storage is available */
-let storageAvailable = true
-let workflowWritesBlocked = false
+type StorageAvailability = 'available' | 'unavailable'
+type WorkflowStorageState =
+  | { status: 'ready'; availability: StorageAvailability }
+  | {
+      status: 'transitioning'
+      reason: 'workspace'
+      resumeAvailability: StorageAvailability
+      ownerId: symbol
+    }
+  | {
+      status: 'transitioning'
+      reason: 'logout'
+      resumeAvailability: StorageAvailability
+    }
+
+let workflowStorageState: WorkflowStorageState = {
+  status: 'ready',
+  availability: 'available'
+}
 const pendingPersistenceFlushes = new Set<() => void>()
 
 export function registerWorkflowPersistenceFlush(
@@ -35,11 +51,23 @@ function flushPendingWorkflowPersistence(): void {
 }
 
 export function isStorageAvailable(): boolean {
-  return storageAvailable && !workflowWritesBlocked
+  return (
+    workflowStorageState.status === 'ready' &&
+    workflowStorageState.availability === 'available'
+  )
 }
 
 export function markStorageUnavailable(): void {
-  storageAvailable = false
+  workflowStorageState =
+    workflowStorageState.status === 'transitioning'
+      ? { ...workflowStorageState, resumeAvailability: 'unavailable' }
+      : { status: 'ready', availability: 'unavailable' }
+}
+
+function isStorageReadable(): boolean {
+  return workflowStorageState.status === 'transitioning'
+    ? workflowStorageState.resumeAvailability === 'available'
+    : workflowStorageState.availability === 'available'
 }
 
 function isQuotaExceeded(error: unknown): boolean {
@@ -68,7 +96,7 @@ function isValidIndex(value: unknown): value is DraftIndexV2 {
  * Reads and parses the draft index from localStorage.
  */
 export function readIndex(workspaceId: string): DraftIndexV2 | null {
-  if (!storageAvailable) return null
+  if (!isStorageReadable()) return null
 
   try {
     const key = StorageKeys.draftIndex(workspaceId)
@@ -88,7 +116,7 @@ export function readIndex(workspaceId: string): DraftIndexV2 | null {
  * Writes the draft index to localStorage.
  */
 export function writeIndex(workspaceId: string, index: DraftIndexV2): boolean {
-  if (!storageAvailable || workflowWritesBlocked) return false
+  if (!isStorageAvailable()) return false
 
   try {
     const key = StorageKeys.draftIndex(workspaceId)
@@ -107,7 +135,7 @@ export function readPayload(
   workspaceId: string,
   draftKey: string
 ): DraftPayloadV2 | null {
-  if (!storageAvailable) return null
+  if (!isStorageReadable()) return null
 
   try {
     const key = `${StorageKeys.prefixes.draftPayload}${workspaceId}:${draftKey}`
@@ -128,7 +156,7 @@ export function writePayload(
   draftKey: string,
   payload: DraftPayloadV2
 ): boolean {
-  if (!storageAvailable || workflowWritesBlocked) return false
+  if (!isStorageAvailable()) return false
 
   try {
     const key = `${StorageKeys.prefixes.draftPayload}${workspaceId}:${draftKey}`
@@ -165,7 +193,7 @@ export function deletePayloads(workspaceId: string, draftKeys: string[]): void {
  * Gets all draft payload keys for a workspace from localStorage.
  */
 export function getPayloadKeys(workspaceId: string): string[] {
-  if (!storageAvailable) return []
+  if (!isStorageReadable()) return []
 
   const prefix = `${StorageKeys.prefixes.draftPayload}${workspaceId}:`
   const keys: string[] = []
@@ -394,7 +422,7 @@ function readLocalPointer<T>(
 }
 
 function writeStorage(storage: Storage, key: string, value: string): void {
-  if (!storageAvailable || workflowWritesBlocked) return
+  if (!isStorageAvailable()) return
 
   try {
     storage.setItem(key, value)
@@ -452,29 +480,65 @@ function removeStorageKeys(
   }
 }
 
-export function clearWorkflowRestoreState(
-  options: { blockWrites?: boolean } = {}
-): void {
-  if (options.blockWrites) {
-    prepareWorkflowWorkspaceTransition()
-    return
-  }
-
+export function clearWorkflowRestoreState(): void {
   removeStorageKeys(localStorage, legacyLocalRestoreKeys)
   removeStorageKeys(sessionStorage, sessionRestoreKeys, sessionRestorePrefixes)
 }
 
-export function prepareWorkflowWorkspaceTransition(): void {
-  if (!workflowWritesBlocked) flushPendingWorkflowPersistence()
-  workflowWritesBlocked = true
+export function prepareWorkflowWorkspaceTransition(): () => void {
+  let ownerId: symbol | undefined
+  if (workflowStorageState.status === 'ready') {
+    flushPendingWorkflowPersistence()
+    ownerId = Symbol('workflow-storage-transition')
+    workflowStorageState = {
+      status: 'transitioning',
+      reason: 'workspace',
+      resumeAvailability: workflowStorageState.availability,
+      ownerId
+    }
+  }
   clearWorkflowRestoreState()
+
+  return () => {
+    if (
+      workflowStorageState.status !== 'transitioning' ||
+      workflowStorageState.reason !== 'workspace' ||
+      workflowStorageState.ownerId !== ownerId
+    )
+      return
+
+    workflowStorageState = {
+      status: 'ready',
+      availability: workflowStorageState.resumeAvailability
+    }
+  }
 }
 
-export function clearAllWorkflowStorage(
-  options: { blockWrites?: boolean } = {}
-): void {
-  if (options.blockWrites) workflowWritesBlocked = true
+export function prepareWorkflowLogoutTransition(): void {
+  workflowStorageState = {
+    status: 'transitioning',
+    reason: 'logout',
+    resumeAvailability:
+      workflowStorageState.status === 'transitioning'
+        ? workflowStorageState.resumeAvailability
+        : workflowStorageState.availability
+  }
+}
+
+export function completeWorkflowLogoutTransition(): void {
+  if (
+    workflowStorageState.status !== 'transitioning' ||
+    workflowStorageState.reason !== 'logout'
+  )
+    return
+
+  workflowStorageState = {
+    status: 'ready',
+    availability: workflowStorageState.resumeAvailability
+  }
+}
 
+export function clearAllWorkflowStorage(): void {
   const localPrefixes = [
     StorageKeys.prefixes.draftIndex,
     StorageKeys.prefixes.draftPayload,