diff --git a/src/components/actionbar/ComfyRunButton/ComfyQueueButton.test.ts b/src/components/actionbar/ComfyRunButton/ComfyQueueButton.test.ts new file mode 100644 index 0000000000..41470be2b2 --- /dev/null +++ b/src/components/actionbar/ComfyRunButton/ComfyQueueButton.test.ts @@ -0,0 +1,197 @@ +import { createTestingPinia } from '@pinia/testing' +import { mount } from '@vue/test-utils' +import { describe, expect, it, vi } from 'vitest' +import { defineComponent, nextTick } from 'vue' +import { createI18n } from 'vue-i18n' + +import type { + JobListItem, + JobStatus +} from '@/platform/remote/comfyui/jobs/jobTypes' +import { useCommandStore } from '@/stores/commandStore' +import { + TaskItemImpl, + useQueueSettingsStore, + useQueueStore +} from '@/stores/queueStore' + +import ComfyQueueButton from './ComfyQueueButton.vue' + +vi.mock('@/platform/distribution/types', () => ({ + isCloud: false +})) + +vi.mock('@/platform/telemetry', () => ({ + useTelemetry: () => null +})) + +vi.mock('@/workbench/extensions/manager/utils/graphHasMissingNodes', () => ({ + graphHasMissingNodes: () => false +})) + +vi.mock('@/scripts/app', () => ({ + app: { + rootGraph: {} + } +})) + +vi.mock('@/stores/workspaceStore', () => ({ + useWorkspaceStore: () => ({ + shiftDown: false + }) +})) + +const SplitButtonStub = defineComponent({ + name: 'SplitButton', + props: { + label: { + type: String, + default: '' + }, + severity: { + type: String, + default: 'primary' + } + }, + template: ` + + ` +}) + +const i18n = createI18n({ + legacy: false, + locale: 'en', + messages: { + en: { + menu: { + run: 'Run', + disabledTooltip: 'Disabled tooltip', + onChange: 'On Change', + onChangeTooltip: 'On change tooltip', + instant: 'Instant', + instantTooltip: 'Instant tooltip', + stopRunInstant: 'Stop Run (Instant)', + stopRunInstantTooltip: 'Stop running', + runWorkflow: 'Run workflow', + runWorkflowFront: 'Run workflow front', + runWorkflowDisabled: 'Run workflow disabled' + } + } + } +}) + +function createTask(id: string, status: JobStatus): TaskItemImpl { + const job: JobListItem = { + id, + status, + create_time: Date.now(), + priority: 1 + } + + return new TaskItemImpl(job) +} + +function createWrapper() { + const pinia = createTestingPinia({ createSpy: vi.fn }) + + return mount(ComfyQueueButton, { + global: { + plugins: [pinia, i18n], + directives: { + tooltip: () => {} + }, + stubs: { + SplitButton: SplitButtonStub, + BatchCountEdit: true + } + } + }) +} + +describe('ComfyQueueButton', () => { + it('keeps the run instant presentation while idle even with active jobs', async () => { + const wrapper = createWrapper() + const queueSettingsStore = useQueueSettingsStore() + const queueStore = useQueueStore() + + queueSettingsStore.mode = 'instant-idle' + queueStore.runningTasks = [createTask('run-1', 'in_progress')] + await nextTick() + + const splitButton = wrapper.get('[data-testid="queue-button"]') + + expect(splitButton.attributes('data-label')).toBe('Run (Instant)') + expect(splitButton.attributes('data-severity')).toBe('primary') + expect(wrapper.find('.icon-\\[lucide--fast-forward\\]').exists()).toBe(true) + }) + + it('switches to stop presentation when instant mode is armed', async () => { + const wrapper = createWrapper() + const queueSettingsStore = useQueueSettingsStore() + + queueSettingsStore.mode = 'instant-running' + await nextTick() + + const splitButton = wrapper.get('[data-testid="queue-button"]') + + expect(splitButton.attributes('data-label')).toBe('Stop Run (Instant)') + expect(splitButton.attributes('data-severity')).toBe('danger') + expect(wrapper.find('.icon-\\[lucide--square\\]').exists()).toBe(true) + }) + + it('disarms instant mode without interrupting even when jobs are active', async () => { + const wrapper = createWrapper() + const queueSettingsStore = useQueueSettingsStore() + const queueStore = useQueueStore() + const commandStore = useCommandStore() + + queueSettingsStore.mode = 'instant-running' + queueStore.runningTasks = [createTask('run-1', 'in_progress')] + await nextTick() + + await wrapper.get('[data-testid="queue-button"]').trigger('click') + await nextTick() + + expect(queueSettingsStore.mode).toBe('instant-idle') + const splitButtonWhileStopping = wrapper.get('[data-testid="queue-button"]') + expect(splitButtonWhileStopping.attributes('data-label')).toBe( + 'Run (Instant)' + ) + expect(splitButtonWhileStopping.attributes('data-severity')).toBe('primary') + expect(wrapper.find('.icon-\\[lucide--fast-forward\\]').exists()).toBe(true) + + expect(commandStore.execute).not.toHaveBeenCalled() + + const splitButton = wrapper.get('[data-testid="queue-button"]') + expect(queueSettingsStore.mode).toBe('instant-idle') + expect(splitButton.attributes('data-label')).toBe('Run (Instant)') + expect(splitButton.attributes('data-severity')).toBe('primary') + expect(wrapper.find('.icon-\\[lucide--fast-forward\\]').exists()).toBe(true) + }) + + it('activates instant running mode when queueing again', async () => { + const wrapper = createWrapper() + const queueSettingsStore = useQueueSettingsStore() + const commandStore = useCommandStore() + + queueSettingsStore.mode = 'instant-idle' + await nextTick() + + await wrapper.get('[data-testid="queue-button"]').trigger('click') + await nextTick() + + expect(queueSettingsStore.mode).toBe('instant-running') + expect(commandStore.execute).toHaveBeenCalledWith('Comfy.QueuePrompt', { + metadata: { + subscribe_to_run: false, + trigger_source: 'button' + } + }) + }) +}) diff --git a/src/components/actionbar/ComfyRunButton/ComfyQueueButton.vue b/src/components/actionbar/ComfyRunButton/ComfyQueueButton.vue index 4c0ea84e43..1218b7228e 100644 --- a/src/components/actionbar/ComfyRunButton/ComfyQueueButton.vue +++ b/src/components/actionbar/ComfyRunButton/ComfyQueueButton.vue @@ -6,8 +6,8 @@ showDelay: 600 }" class="comfyui-queue-button" - :label="String(activeQueueModeMenuItem?.label ?? '')" - severity="primary" + :label="queueButtonLabel" + :severity="queueButtonSeverity" size="small" :model="queueModeMenuItems" data-testid="queue-button" @@ -22,7 +22,7 @@ value: item.tooltip, showDelay: 600 }" - :variant="item.key === queueMode ? 'primary' : 'secondary'" + :variant="item.key === selectedQueueMode ? 'primary' : 'secondary'" size="sm" class="w-full justify-start" > @@ -48,7 +48,11 @@ import { useTelemetry } from '@/platform/telemetry' import { app } from '@/scripts/app' import { useCommandStore } from '@/stores/commandStore' import { useNodeDefStore } from '@/stores/nodeDefStore' -import { useQueueSettingsStore } from '@/stores/queueStore' +import { + isInstantMode, + isInstantRunningMode, + useQueueSettingsStore +} from '@/stores/queueStore' import { useWorkspaceStore } from '@/stores/workspaceStore' import { graphHasMissingNodes } from '@/workbench/extensions/manager/utils/graphHasMissingNodes' @@ -63,6 +67,12 @@ const hasMissingNodes = computed(() => ) const { t } = useI18n() +type QueueModeMenuKey = 'disabled' | 'change' | 'instant-idle' + +const selectedQueueMode = computed(() => + isInstantMode(queueMode.value) ? 'instant-idle' : queueMode.value +) + const queueModeMenuItemLookup = computed(() => { const items: Record = { disabled: { @@ -86,15 +96,15 @@ const queueModeMenuItemLookup = computed(() => { } } if (!isCloud) { - items.instant = { - key: 'instant', + items['instant-idle'] = { + key: 'instant-idle', label: `${t('menu.run')} (${t('menu.instant')})`, tooltip: t('menu.instantTooltip'), command: () => { useTelemetry()?.trackUiButtonClicked({ button_id: 'queue_mode_option_run_instant_selected' }) - queueMode.value = 'instant' + queueMode.value = 'instant-idle' } } } @@ -104,7 +114,7 @@ const queueModeMenuItemLookup = computed(() => { const activeQueueModeMenuItem = computed(() => { // Fallback to disabled mode if current mode is not available (e.g., instant mode in cloud) return ( - queueModeMenuItemLookup.value[queueMode.value] || + queueModeMenuItemLookup.value[selectedQueueMode.value] || queueModeMenuItemLookup.value.disabled ) }) @@ -112,7 +122,24 @@ const queueModeMenuItems = computed(() => Object.values(queueModeMenuItemLookup.value) ) +const isStopInstantAction = computed(() => + isInstantRunningMode(queueMode.value) +) + +const queueButtonLabel = computed(() => + isStopInstantAction.value + ? t('menu.stopRunInstant') + : String(activeQueueModeMenuItem.value?.label ?? '') +) + +const queueButtonSeverity = computed(() => + isStopInstantAction.value ? 'danger' : 'primary' +) + const iconClass = computed(() => { + if (isStopInstantAction.value) { + return 'icon-[lucide--square]' + } if (hasMissingNodes.value) { return 'icon-[lucide--triangle-alert]' } @@ -122,7 +149,7 @@ const iconClass = computed(() => { if (queueMode.value === 'disabled') { return 'icon-[lucide--play]' } - if (queueMode.value === 'instant') { + if (isInstantMode(queueMode.value)) { return 'icon-[lucide--fast-forward]' } if (queueMode.value === 'change') { @@ -132,6 +159,9 @@ const iconClass = computed(() => { }) const queueButtonTooltip = computed(() => { + if (isStopInstantAction.value) { + return t('menu.stopRunInstantTooltip') + } if (hasMissingNodes.value) { return t('menu.runWorkflowDisabled') } @@ -143,11 +173,20 @@ const queueButtonTooltip = computed(() => { const commandStore = useCommandStore() const queuePrompt = async (e: Event) => { + if (isStopInstantAction.value) { + queueMode.value = 'instant-idle' + return + } + const isShiftPressed = 'shiftKey' in e && e.shiftKey const commandId = isShiftPressed ? 'Comfy.QueuePromptFront' : 'Comfy.QueuePrompt' + if (isInstantMode(queueMode.value)) { + queueMode.value = 'instant-running' + } + if (batchCount.value > 1) { useTelemetry()?.trackUiButtonClicked({ button_id: 'queue_run_multiple_batches_submitted' diff --git a/src/locales/en/main.json b/src/locales/en/main.json index 48c43f1035..749f0affe6 100644 --- a/src/locales/en/main.json +++ b/src/locales/en/main.json @@ -918,6 +918,8 @@ "runWorkflowFront": "Run workflow (Queue at front)", "runWorkflowDisabled": "Workflow contains unsupported nodes (highlighted red). Remove these to run the workflow.", "run": "Run", + "stopRunInstant": "Stop Run (Instant)", + "stopRunInstantTooltip": "Stop running", "execute": "Execute", "interrupt": "Cancel current run", "refresh": "Refresh node definitions", diff --git a/src/services/autoQueueService.ts b/src/services/autoQueueService.ts index ffbc5dede4..6a62a9f0ec 100644 --- a/src/services/autoQueueService.ts +++ b/src/services/autoQueueService.ts @@ -1,6 +1,7 @@ import { api } from '@/scripts/api' import { app } from '@/scripts/app' import { + isInstantRunningMode, useQueuePendingTaskCountStore, useQueueSettingsStore } from '@/stores/queueStore' @@ -29,7 +30,7 @@ export function setupAutoQueueHandler() { internalCount = queueCountStore.count if (!internalCount && !app.lastExecutionError) { if ( - queueSettingsStore.mode === 'instant' || + isInstantRunningMode(queueSettingsStore.mode) || (queueSettingsStore.mode === 'change' && graphHasChanged) ) { graphHasChanged = false diff --git a/src/stores/queueStore.ts b/src/stores/queueStore.ts index 07ffd9060a..18065b1473 100644 --- a/src/stores/queueStore.ts +++ b/src/stores/queueStore.ts @@ -627,7 +627,20 @@ export const useQueuePendingTaskCountStore = defineStore( } ) -export type AutoQueueMode = 'disabled' | 'instant' | 'change' +export type AutoQueueMode = + | 'disabled' + | 'change' + | 'instant-idle' + | 'instant-running' + +export const isInstantMode = ( + mode: AutoQueueMode +): mode is 'instant-idle' | 'instant-running' => + mode === 'instant-idle' || mode === 'instant-running' + +export const isInstantRunningMode = ( + mode: AutoQueueMode +): mode is 'instant-running' => mode === 'instant-running' export const useQueueSettingsStore = defineStore('queueSettingsStore', { state: () => ({