[backport rh-test] fix: improve template URL loading UX and prevent re-triggering (#6654)

Backport of #6593 to `rh-test`

Automatically created by backport workflow.

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-6654-backport-rh-test-fix-improve-template-URL-loading-UX-and-prevent-re-triggering-2a86d73d36508163834fdea17353ed37)
by [Unito](https://www.unito.io)

Co-authored-by: Christian Byrne <cbyrne@comfy.org>
Co-authored-by: Christian Byrne <c.byrne@comfy.org>

src/components/graph/GraphCanvas.vue

@@ -443,10 +443,13 @@ onMounted(async () => {
     'Comfy.CustomColorPalettes'
   )
 
-  // Restore workflow and workflow tabs state from storage
-  await workflowPersistence.restorePreviousWorkflow()
+  // Restore saved workflow and workflow tabs state
+  await workflowPersistence.initializeWorkflow()
   workflowPersistence.restoreWorkflowTabsState()
 
+  // Load template from URL if present
+  await workflowPersistence.loadTemplateFromUrlIfPresent()
+
   // Initialize release store to fetch releases from comfy-api (fire-and-forget)
   const { useReleaseStore } = await import(
     '@/platform/updates/common/releaseStore'

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

@@ -1,9 +1,11 @@
 import { tryOnScopeDispose } from '@vueuse/core'
 import { computed, watch } from 'vue'
+import { useRoute } from 'vue-router'
 
 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 { api } from '@/scripts/api'
 import { app as comfyApp } from '@/scripts/app'
 import { getStorageValue, setStorageValue } from '@/scripts/utils'
@@ -12,6 +14,8 @@ import { useCommandStore } from '@/stores/commandStore'
 export function useWorkflowPersistence() {
   const workflowStore = useWorkflowStore()
   const settingStore = useSettingStore()
+  const route = useRoute()
+  const templateUrlLoader = useTemplateUrlLoader()
 
   const workflowPersistenceEnabled = computed(() =>
     settingStore.get('Comfy.Workflow.Persist')
@@ -63,8 +67,9 @@ export function useWorkflowPersistence() {
     }
   }
 
-  const restorePreviousWorkflow = async () => {
+  const initializeWorkflow = async () => {
     if (!workflowPersistenceEnabled.value) return
+
     try {
       const restored = await loadPreviousWorkflowFromStorage()
       if (!restored) {
@@ -76,6 +81,15 @@ export function useWorkflowPersistence() {
     }
   }
 
+  const loadTemplateFromUrlIfPresent = async () => {
+    const hasTemplateUrl =
+      route.query.template && typeof route.query.template === 'string'
+
+    if (hasTemplateUrl) {
+      await templateUrlLoader.loadTemplateFromUrl()
+    }
+  }
+
   // Setup watchers
   watch(
     () => workflowStore.activeWorkflow?.key,
@@ -141,7 +155,8 @@ export function useWorkflowPersistence() {
   }
 
   return {
-    restorePreviousWorkflow,
+    initializeWorkflow,
+    loadTemplateFromUrlIfPresent,
     restoreWorkflowTabsState
   }
 }

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

@@ -1,6 +1,6 @@
 import { useToast } from 'primevue/usetoast'
 import { useI18n } from 'vue-i18n'
-import { useRoute } from 'vue-router'
+import { useRoute, useRouter } from 'vue-router'
 
 import { useTemplateWorkflows } from './useTemplateWorkflows'
 
@@ -17,6 +17,7 @@ import { useTemplateWorkflows } from './useTemplateWorkflows'
  */
 export function useTemplateUrlLoader() {
   const route = useRoute()
+  const router = useRouter()
   const { t } = useI18n()
   const toast = useToast()
   const templateWorkflows = useTemplateWorkflows()
@@ -28,6 +29,16 @@ export function useTemplateUrlLoader() {
     return /^[a-zA-Z0-9_-]+$/.test(param)
   }
 
+  /**
+   * Removes template and source parameters from URL
+   */
+  const cleanupUrlParams = () => {
+    const newQuery = { ...route.query }
+    delete newQuery.template
+    delete newQuery.source
+    void router.replace({ query: newQuery })
+  }
+
   /**
    * Loads template from URL query parameters if present
    * Handles errors internally and shows appropriate user feedback
@@ -39,7 +50,6 @@ export function useTemplateUrlLoader() {
       return
     }
 
-    // Validate template name format
     if (!isValidParameter(templateParam)) {
       console.warn(
         `[useTemplateUrlLoader] Invalid template parameter format: ${templateParam}`
@@ -49,7 +59,6 @@ export function useTemplateUrlLoader() {
 
     const sourceParam = (route.query.source as string | undefined) || 'default'
 
-    // Validate source parameter format
     if (!isValidParameter(sourceParam)) {
       console.warn(
         `[useTemplateUrlLoader] Invalid source parameter format: ${sourceParam}`
@@ -57,7 +66,6 @@ export function useTemplateUrlLoader() {
       return
     }
 
-    // Load template with error handling
     try {
       await templateWorkflows.loadTemplates()
 
@@ -87,6 +95,8 @@ export function useTemplateUrlLoader() {
         detail: t('g.errorLoadingTemplate'),
         life: 3000
       })
+    } finally {
+      cleanupUrlParams()
     }
   }
 

src/views/GraphView.vue

@@ -51,7 +51,6 @@ import { useSettingStore } from '@/platform/settings/settingStore'
 import { useTelemetry } from '@/platform/telemetry'
 import { useFrontendVersionMismatchWarning } from '@/platform/updates/common/useFrontendVersionMismatchWarning'
 import { useVersionCompatibilityStore } from '@/platform/updates/common/versionCompatibilityStore'
-import { useTemplateUrlLoader } from '@/platform/workflow/templates/composables/useTemplateUrlLoader'
 import type { StatusWsMessageStatus } from '@/schemas/apiSchema'
 import { api } from '@/scripts/api'
 import { app } from '@/scripts/app'
@@ -78,9 +77,6 @@ setupAutoQueueHandler()
 useProgressFavicon()
 useBrowserTabTitle()
 
-// Template URL loading
-const { loadTemplateFromUrl } = useTemplateUrlLoader()
-
 const { t } = useI18n()
 const toast = useToast()
 const settingStore = useSettingStore()
@@ -347,9 +343,6 @@ const onGraphReady = () => {
       tabCountChannel.postMessage({ type: 'heartbeat', tabId: currentTabId })
     }
 
-    // Load template from URL if present
-    void loadTemplateFromUrl()
-
     // Setting values now available after comfyApp.setup.
     // Load keybindings.
     wrapWithErrorHandling(useKeybindingService().registerUserKeybindings)()

tests-ui/tests/platform/workflow/templates/composables/useTemplateUrlLoader.test.ts

@@ -14,10 +14,14 @@ import { useTemplateUrlLoader } from '@/platform/workflow/templates/composables/
 
 // Mock vue-router
 let mockQueryParams: Record<string, string | undefined> = {}
+const mockRouterReplace = vi.fn()
 
 vi.mock('vue-router', () => ({
   useRoute: vi.fn(() => ({
     query: mockQueryParams
+  })),
+  useRouter: vi.fn(() => ({
+    replace: mockRouterReplace
   }))
 }))
 
@@ -217,4 +221,43 @@ describe('useTemplateUrlLoader', () => {
       life: 3000
     })
   })
+
+  it('removes template params from URL after successful load', async () => {
+    mockQueryParams = {
+      template: 'flux_simple',
+      source: 'custom',
+      other: 'param'
+    }
+
+    const { loadTemplateFromUrl } = useTemplateUrlLoader()
+    await loadTemplateFromUrl()
+
+    expect(mockRouterReplace).toHaveBeenCalledWith({
+      query: { other: 'param' }
+    })
+  })
+
+  it('removes template params from URL even on error', async () => {
+    mockQueryParams = { template: 'invalid', source: 'custom', other: 'param' }
+    mockLoadWorkflowTemplate.mockResolvedValueOnce(false)
+
+    const { loadTemplateFromUrl } = useTemplateUrlLoader()
+    await loadTemplateFromUrl()
+
+    expect(mockRouterReplace).toHaveBeenCalledWith({
+      query: { other: 'param' }
+    })
+  })
+
+  it('removes template params from URL even on exception', async () => {
+    mockQueryParams = { template: 'flux_simple', other: 'param' }
+    mockLoadTemplates.mockRejectedValueOnce(new Error('Network error'))
+
+    const { loadTemplateFromUrl } = useTemplateUrlLoader()
+    await loadTemplateFromUrl()
+
+    expect(mockRouterReplace).toHaveBeenCalledWith({
+      query: { other: 'param' }
+    })
+  })
 })