diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue index 20f73ad88..a1f893922 100644 --- a/src/components/graph/GraphCanvas.vue +++ b/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' diff --git a/src/platform/workflow/persistence/composables/useWorkflowPersistence.ts b/src/platform/workflow/persistence/composables/useWorkflowPersistence.ts index f6f5f1828..2be36ca9c 100644 --- a/src/platform/workflow/persistence/composables/useWorkflowPersistence.ts +++ b/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 } } diff --git a/src/platform/workflow/templates/composables/useTemplateUrlLoader.ts b/src/platform/workflow/templates/composables/useTemplateUrlLoader.ts index c23e570e6..cccadfe07 100644 --- a/src/platform/workflow/templates/composables/useTemplateUrlLoader.ts +++ b/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() } } diff --git a/src/views/GraphView.vue b/src/views/GraphView.vue index d05446125..71d3146a1 100644 --- a/src/views/GraphView.vue +++ b/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)() diff --git a/tests-ui/tests/platform/workflow/templates/composables/useTemplateUrlLoader.test.ts b/tests-ui/tests/platform/workflow/templates/composables/useTemplateUrlLoader.test.ts index 5ed3b94ce..10127a47f 100644 --- a/tests-ui/tests/platform/workflow/templates/composables/useTemplateUrlLoader.test.ts +++ b/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 = {} +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' } + }) + }) })