remove localstorage workflow on signout in cloud

Author: christian-byrne

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

@@ -7,21 +7,25 @@ import {
   mergePreservedQueryIntoQuery
 } from '@/platform/navigation/preservedQueryManager'
 import { PRESERVED_QUERY_NAMESPACES } from '@/platform/navigation/preservedQueryNamespaces'
+import { clearWorkflowPersistenceStorage } from '@/platform/workflow/persistence/workflowStorage'
 import { useSettingStore } from '@/platform/settings/settingStore'
 import { useWorkflowService } from '@/platform/workflow/core/services/workflowService'
 import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
 import { useTemplateUrlLoader } from '@/platform/workflow/templates/composables/useTemplateUrlLoader'
+import { useCurrentUser } from '@/composables/auth/useCurrentUser'
 import { api } from '@/scripts/api'
 import { app as comfyApp } from '@/scripts/app'
 import { getStorageValue, setStorageValue } from '@/scripts/utils'
 import { useCommandStore } from '@/stores/commandStore'
+import { isCloud } from '@/platform/distribution/types'
 
 export function useWorkflowPersistence() {
   const workflowStore = useWorkflowStore()
   const settingStore = useSettingStore()
   const route = useRoute()
   const router = useRouter()
   const templateUrlLoader = useTemplateUrlLoader()
+  const { onUserLogout } = useCurrentUser()
   const TEMPLATE_NAMESPACE = PRESERVED_QUERY_NAMESPACES.TEMPLATE
 
   const ensureTemplateQueryFromIntent = async () => {
@@ -183,6 +187,10 @@ export function useWorkflowPersistence() {
     }
   })
 
+  onUserLogout(() => {
+    if (isCloud) clearWorkflowPersistenceStorage()
+  })
+
   const restoreWorkflowTabsState = () => {
     if (!workflowPersistenceEnabled.value) return
     const isRestorable = storedWorkflows?.length > 0 && storedActiveIndex >= 0

src/platform/workflow/persistence/workflowStorage.ts

@@ -0,0 +1,39 @@
+const LOCAL_STORAGE_KEYS = [
+  'workflow',
+  'Comfy.PreviousWorkflow',
+  'Comfy.OpenWorkflowsPaths',
+  'Comfy.ActiveWorkflowIndex'
+]
+
+const SESSION_KEY_PREFIXES = [
+  'workflow:',
+  'Comfy.PreviousWorkflow',
+  'Comfy.OpenWorkflowsPaths',
+  'Comfy.ActiveWorkflowIndex'
+]
+
+const shouldClearSessionKey = (key: string): boolean =>
+  SESSION_KEY_PREFIXES.some((prefix) =>
+    prefix.endsWith(':')
+      ? key.startsWith(prefix)
+      : key === prefix || key.startsWith(`${prefix}:`)
+  )
+
+/**
+ * Removes any workflow state persisted in storage so that the next login
+ * starts with a clean slate.
+ */
+export const clearWorkflowPersistenceStorage = () => {
+  LOCAL_STORAGE_KEYS.forEach((key) => {
+    localStorage.removeItem(key)
+  })
+
+  // Iterate backwards because removing items mutates sessionStorage.length.
+  for (let i = sessionStorage.length - 1; i >= 0; i--) {
+    const key = sessionStorage.key(i)
+    if (!key) continue
+    if (shouldClearSessionKey(key)) {
+      sessionStorage.removeItem(key)
+    }
+  }
+}

tests-ui/tests/platform/workflow/persistence/workflowStorage.test.ts

@@ -0,0 +1,42 @@
+import { describe, expect, it, beforeEach } from 'vitest'
+
+import { clearWorkflowPersistenceStorage } from '@/platform/workflow/persistence/workflowStorage'
+
+describe('workflowStorage', () => {
+  beforeEach(() => {
+    localStorage.clear()
+    sessionStorage.clear()
+  })
+
+  it('clears all workflow persistence keys from storage', () => {
+    localStorage.setItem('workflow', 'data')
+    localStorage.setItem('Comfy.PreviousWorkflow', 'prev')
+    localStorage.setItem('Comfy.OpenWorkflowsPaths', '[]')
+    localStorage.setItem('Comfy.ActiveWorkflowIndex', '1')
+    localStorage.setItem('unrelated', 'keep')
+
+    sessionStorage.setItem('workflow:client-1', 'session-data')
+    sessionStorage.setItem('Comfy.PreviousWorkflow:client-1', 'prev')
+    sessionStorage.setItem('Comfy.ActiveWorkflowIndex:client-1', '0')
+    sessionStorage.setItem('Comfy.OpenWorkflowsPaths:client-1', '[]')
+    sessionStorage.setItem('custom', 'keep')
+
+    clearWorkflowPersistenceStorage()
+
+    expect(localStorage.getItem('workflow')).toBeNull()
+    expect(localStorage.getItem('Comfy.PreviousWorkflow')).toBeNull()
+    expect(localStorage.getItem('Comfy.OpenWorkflowsPaths')).toBeNull()
+    expect(localStorage.getItem('Comfy.ActiveWorkflowIndex')).toBeNull()
+    expect(localStorage.getItem('unrelated')).toBe('keep')
+
+    expect(sessionStorage.getItem('workflow:client-1')).toBeNull()
+    expect(sessionStorage.getItem('Comfy.PreviousWorkflow:client-1')).toBeNull()
+    expect(
+      sessionStorage.getItem('Comfy.ActiveWorkflowIndex:client-1')
+    ).toBeNull()
+    expect(
+      sessionStorage.getItem('Comfy.OpenWorkflowsPaths:client-1')
+    ).toBeNull()
+    expect(sessionStorage.getItem('custom')).toBe('keep')
+  })
+})