fix: template query param stripped during login views (#6677)

Fixes issue where query params from
https://github.com/Comfy-Org/ComfyUI_frontend/pull/6593 are stripped
during the login/signup views/flow by storing initial params in session
storage via router plugin.



https://github.com/user-attachments/assets/51642e8c-af5c-43ef-ab7d-133bc7e511aa




┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-6677-fix-template-query-param-stripped-during-login-views-2aa6d73d365081a1bdc7d22b35f72a77)
by [Unito](https://www.unito.io)

src/platform/navigation/preservedQueryManager.ts

@@ -0,0 +1,109 @@
+import type { LocationQuery, LocationQueryRaw } from 'vue-router'
+
+const STORAGE_PREFIX = 'Comfy.PreservedQuery.'
+const preservedQueries = new Map<string, Record<string, string>>()
+
+const readQueryParam = (value: unknown): string | undefined => {
+  return typeof value === 'string' ? value : undefined
+}
+
+const getStorageKey = (namespace: string) => `${STORAGE_PREFIX}${namespace}`
+
+const isValidQueryRecord = (
+  value: unknown
+): value is Record<string, string> => {
+  if (typeof value !== 'object' || value === null || Array.isArray(value)) {
+    return false
+  }
+  return Object.values(value).every((v) => typeof v === 'string')
+}
+
+const readFromStorage = (namespace: string): Record<string, string> | null => {
+  try {
+    const raw = sessionStorage.getItem(getStorageKey(namespace))
+    if (!raw) return null
+
+    const parsed = JSON.parse(raw)
+    if (!isValidQueryRecord(parsed)) {
+      console.warn('[preservedQuery] invalid storage format')
+      sessionStorage.removeItem(getStorageKey(namespace))
+      return null
+    }
+    return parsed
+  } catch (error) {
+    console.warn('[preservedQuery] storage operation failed')
+    sessionStorage.removeItem(getStorageKey(namespace))
+    return null
+  }
+}
+
+const writeToStorage = (
+  namespace: string,
+  payload: Record<string, string> | null
+) => {
+  try {
+    if (!payload || Object.keys(payload).length === 0) {
+      sessionStorage.removeItem(getStorageKey(namespace))
+      return
+    }
+    sessionStorage.setItem(getStorageKey(namespace), JSON.stringify(payload))
+  } catch (error) {
+    console.warn('[preservedQuery] failed to write storage', {
+      namespace,
+      error
+    })
+  }
+}
+
+export const hydratePreservedQuery = (namespace: string) => {
+  if (preservedQueries.has(namespace)) return
+  const payload = readFromStorage(namespace)
+  if (payload) {
+    preservedQueries.set(namespace, payload)
+  }
+}
+
+export const capturePreservedQuery = (
+  namespace: string,
+  query: LocationQuery,
+  keys: string[]
+) => {
+  const payload: Record<string, string> = {}
+
+  keys.forEach((key) => {
+    const value = readQueryParam(query[key])
+    if (value) {
+      payload[key] = value
+    }
+  })
+
+  if (Object.keys(payload).length === 0) return
+
+  preservedQueries.set(namespace, payload)
+  writeToStorage(namespace, payload)
+}
+
+export const mergePreservedQueryIntoQuery = (
+  namespace: string,
+  query?: LocationQueryRaw
+): LocationQueryRaw | undefined => {
+  const payload = preservedQueries.get(namespace)
+  if (!payload) return undefined
+
+  const nextQuery: LocationQueryRaw = { ...(query || {}) }
+  let changed = false
+
+  for (const [key, value] of Object.entries(payload)) {
+    if (typeof nextQuery[key] === 'string') continue
+    nextQuery[key] = value
+    changed = true
+  }
+
+  return changed ? nextQuery : undefined
+}
+
+export const clearPreservedQuery = (namespace: string) => {
+  if (!preservedQueries.has(namespace)) return
+  preservedQueries.delete(namespace)
+  writeToStorage(namespace, null)
+}

src/platform/navigation/preservedQueryNamespaces.ts

@@ -0,0 +1,3 @@
+export const PRESERVED_QUERY_NAMESPACES = {
+  TEMPLATE: 'template'
+} as const

src/platform/navigation/preservedQueryTracker.ts

@@ -0,0 +1,29 @@
+import type { Router } from 'vue-router'
+
+import {
+  capturePreservedQuery,
+  hydratePreservedQuery
+} from '@/platform/navigation/preservedQueryManager'
+
+export const installPreservedQueryTracker = (
+  router: Router,
+  definitions: Array<{ namespace: string; keys: string[] }>
+) => {
+  const trackedDefinitions = definitions.map((definition) => ({
+    ...definition
+  }))
+
+  router.beforeEach((to, _from, next) => {
+    const queryKeys = new Set(Object.keys(to.query))
+
+    trackedDefinitions.forEach(({ namespace, keys }) => {
+      hydratePreservedQuery(namespace)
+      const shouldCapture = keys.some((key) => queryKeys.has(key))
+      if (shouldCapture) {
+        capturePreservedQuery(namespace, to.query, keys)
+      }
+    })
+
+    next()
+  })
+}

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

@@ -1,7 +1,12 @@
 import { tryOnScopeDispose } from '@vueuse/core'
 import { computed, watch } from 'vue'
-import { useRoute } from 'vue-router'
+import { useRoute, useRouter } from 'vue-router'
 
+import {
+  hydratePreservedQuery,
+  mergePreservedQueryIntoQuery
+} from '@/platform/navigation/preservedQueryManager'
+import { PRESERVED_QUERY_NAMESPACES } from '@/platform/navigation/preservedQueryNamespaces'
 import { useSettingStore } from '@/platform/settings/settingStore'
 import { useWorkflowService } from '@/platform/workflow/core/services/workflowService'
 import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
@@ -15,7 +20,23 @@ export function useWorkflowPersistence() {
   const workflowStore = useWorkflowStore()
   const settingStore = useSettingStore()
   const route = useRoute()
+  const router = useRouter()
   const templateUrlLoader = useTemplateUrlLoader()
+  const TEMPLATE_NAMESPACE = PRESERVED_QUERY_NAMESPACES.TEMPLATE
+
+  const ensureTemplateQueryFromIntent = async () => {
+    hydratePreservedQuery(TEMPLATE_NAMESPACE)
+    const mergedQuery = mergePreservedQueryIntoQuery(
+      TEMPLATE_NAMESPACE,
+      route.query
+    )
+
+    if (mergedQuery) {
+      await router.replace({ query: mergedQuery })
+    }
+
+    return mergedQuery ?? route.query
+  }
 
   const workflowPersistenceEnabled = computed(() =>
     settingStore.get('Comfy.Workflow.Persist')
@@ -101,8 +122,8 @@ export function useWorkflowPersistence() {
   }
 
   const loadTemplateFromUrlIfPresent = async () => {
-    const hasTemplateUrl =
-      route.query.template && typeof route.query.template === 'string'
+    const query = await ensureTemplateQueryFromIntent()
+    const hasTemplateUrl = query.template && typeof query.template === 'string'
 
     if (hasTemplateUrl) {
       await templateUrlLoader.loadTemplateFromUrl()

src/platform/workflow/templates/composables/useTemplateUrlLoader.ts

@@ -2,6 +2,9 @@ import { useToast } from 'primevue/usetoast'
 import { useI18n } from 'vue-i18n'
 import { useRoute, useRouter } from 'vue-router'
 
+import { clearPreservedQuery } from '@/platform/navigation/preservedQueryManager'
+import { PRESERVED_QUERY_NAMESPACES } from '@/platform/navigation/preservedQueryNamespaces'
+
 import { useTemplateWorkflows } from './useTemplateWorkflows'
 
 /**
@@ -21,6 +24,7 @@ export function useTemplateUrlLoader() {
   const { t } = useI18n()
   const toast = useToast()
   const templateWorkflows = useTemplateWorkflows()
+  const TEMPLATE_NAMESPACE = PRESERVED_QUERY_NAMESPACES.TEMPLATE
 
   /**
    * Validates parameter format to prevent path traversal and injection attacks
@@ -97,6 +101,7 @@ export function useTemplateUrlLoader() {
       })
     } finally {
       cleanupUrlParams()
+      clearPreservedQuery(TEMPLATE_NAMESPACE)
     }
   }