Compare commits

...

3 Commits

Author SHA1 Message Date
huang47
76fb6a7204 feat: block disabled partner nodes from queue 2026-07-17 02:44:25 -07:00
huang47
6e0af560cc feat: filter disabled partner nodes from discovery 2026-07-17 02:00:29 -07:00
huang47
cbdacb3316 feat: add partner node enforcement control 2026-07-17 01:32:55 -07:00
17 changed files with 689 additions and 20 deletions

View File

@@ -0,0 +1,99 @@
import type { Page } from '@playwright/test'
import type { ListAssetsResponse } from '@comfyorg/ingest-types'
import type { PartnerNodePolicyResponse } from '@/platform/workspace/api/partnerNodePolicyApi'
import type { RemoteConfig } from '@/platform/remoteConfig/types'
import type { PromptResponse } from '@/schemas/apiSchema'
import type { ComfyNodeDef } from '@/schemas/nodeDefSchema'
import { comfyPageFixture as base } from '@e2e/fixtures/ComfyPage'
import { WORKSPACE_FEATURE_FLAG } from '@e2e/fixtures/data/cloudWorkspace'
import { CloudWorkspaceMockHelper } from '@e2e/fixtures/helpers/CloudWorkspaceMockHelper'
import { jsonRoute } from '@e2e/fixtures/utils/jsonRoute'
const APP_URL = process.env.PLAYWRIGHT_TEST_URL || 'http://localhost:8188'
const DISABLED_NODE = 'DisabledPartnerNode'
interface PartnerNodeGovernanceFixture {
promptRequestCount: () => number
}
function disabledPartnerNode(): ComfyNodeDef {
return {
name: DISABLED_NODE,
display_name: 'Disabled Partner Node',
category: 'partner/image/Acme',
python_module: 'comfy_api_nodes.acme',
description: '',
input: {},
output: [],
output_is_list: [],
output_name: [],
output_node: true,
api_node: true
}
}
async function setupGovernedWorkspace(page: Page) {
await new CloudWorkspaceMockHelper(page).setup()
await page.route('**/api/features', (route) =>
route.fulfill(
jsonRoute({
...WORKSPACE_FEATURE_FLAG,
partner_node_governance_enabled: true
} satisfies RemoteConfig)
)
)
await page.route('**/api/object_info', (route) =>
route.fulfill(jsonRoute({ [DISABLED_NODE]: disabledPartnerNode() }))
)
await page.route(/\/api\/assets(?:\?.*)?$/, (route) =>
route.fulfill(
jsonRoute({
assets: [],
total: 0,
has_more: false
} satisfies ListAssetsResponse)
)
)
await page.route('**/api/workspace/partner-node-policy', (route) =>
route.fulfill(
jsonRoute({
enforcement_enabled: true,
nodes: { [DISABLED_NODE]: false }
} satisfies PartnerNodePolicyResponse)
)
)
}
export const partnerNodeGovernanceTest = base.extend<{
partnerNodeGovernance: PartnerNodeGovernanceFixture
}>({
partnerNodeGovernance: async ({ page }, use) => {
await setupGovernedWorkspace(page)
let promptRequests = 0
await page.route('**/api/prompt', (route) => {
promptRequests++
return route.fulfill(
jsonRoute({
prompt_id: 'unexpected-prompt',
node_errors: {},
error: ''
} satisfies PromptResponse)
)
})
await page.goto(APP_URL)
await page.waitForFunction(() => !!window.app?.extensionManager, null, {
timeout: 45_000
})
await page.evaluate((nodeType) => {
const node = window.LiteGraph!.createNode(nodeType)
if (!node) throw new Error(`Failed to create ${nodeType}`)
window.app!.rootGraph.add(node)
}, DISABLED_NODE)
await use({ promptRequestCount: () => promptRequests })
}
})

View File

@@ -59,9 +59,7 @@ async function openAllowlist(page: Page) {
test.describe('Partner node allowlist', { tag: '@cloud' }, () => {
test.describe.configure({ timeout: 60_000 })
test('saves a whole policy from the workspace settings tab', async ({
page
}) => {
test('saves enforcement and allowlist as one policy', async ({ page }) => {
await new CloudWorkspaceMockHelper(page).setup()
await page.route('**/api/features', (route) =>
@@ -111,13 +109,16 @@ test.describe('Partner node allowlist', { tag: '@cloud' }, () => {
content.getByRole('switch', { name: 'Allow Flux Expand' })
).toBeChecked()
await content
.getByRole('switch', { name: 'Enforce partner node allowlist' })
.click()
await fluxFill.click()
await content.getByRole('button', { name: 'Save' }).click()
await expect(page.getByText('Allowlist saved')).toBeVisible()
await expect(page.getByText('Partner node policy saved')).toBeVisible()
expect(updates).toEqual([
{
enforcement_enabled: false,
enforcement_enabled: true,
nodes: { FluxFill: true, FluxExpand: true, VeoVideo: true }
}
] satisfies PartnerNodePolicyResponse[])

View File

@@ -0,0 +1,97 @@
import { expect } from '@playwright/test'
import type { Page } from '@playwright/test'
import type { ListAssetsResponse } from '@comfyorg/ingest-types'
import type { PartnerNodePolicyResponse } from '@/platform/workspace/api/partnerNodePolicyApi'
import type { RemoteConfig } from '@/platform/remoteConfig/types'
import type { ComfyNodeDef } from '@/schemas/nodeDefSchema'
import { comfyPageFixture as test } from '@e2e/fixtures/ComfyPage'
import { WORKSPACE_FEATURE_FLAG } from '@e2e/fixtures/data/cloudWorkspace'
import { CloudWorkspaceMockHelper } from '@e2e/fixtures/helpers/CloudWorkspaceMockHelper'
import { jsonRoute } from '@e2e/fixtures/utils/jsonRoute'
const APP_URL = process.env.PLAYWRIGHT_TEST_URL || 'http://localhost:8188'
function partnerNode(name: string): ComfyNodeDef {
return {
name,
display_name: name,
category: 'partner/image/Acme',
python_module: 'comfy_api_nodes.acme',
description: '',
input: {},
output: [],
output_is_list: [],
output_name: [],
output_node: false,
api_node: true
}
}
async function setupGovernedWorkspace(page: Page) {
await new CloudWorkspaceMockHelper(page).setup()
await page.route('**/api/features', (route) =>
route.fulfill(
jsonRoute({
...WORKSPACE_FEATURE_FLAG,
partner_node_governance_enabled: true
} satisfies RemoteConfig)
)
)
await page.route('**/api/object_info', (route) =>
route.fulfill(
jsonRoute({
AllowedPartnerNode: partnerNode('AllowedPartnerNode'),
DisabledPartnerNode: partnerNode('DisabledPartnerNode')
})
)
)
await page.route(/\/api\/assets(?:\?.*)?$/, (route) =>
route.fulfill(
jsonRoute({
assets: [],
total: 0,
has_more: false
} satisfies ListAssetsResponse)
)
)
await page.route('**/api/workspace/partner-node-policy', (route) =>
route.fulfill(
jsonRoute({
enforcement_enabled: true,
nodes: {
AllowedPartnerNode: true,
DisabledPartnerNode: false
}
} satisfies PartnerNodePolicyResponse)
)
)
}
test.describe('Partner node governance discovery', { tag: '@cloud' }, () => {
test('hides disabled nodes from search', async ({ page }) => {
await setupGovernedWorkspace(page)
await page.goto(APP_URL)
await page.waitForFunction(() => !!window.app?.extensionManager, null, {
timeout: 45_000
})
await page.evaluate(async () => {
await window.app!.extensionManager.setting.set(
'Comfy.NodeSearchBoxImpl',
'default'
)
await window.app!.extensionManager.command.execute(
'Workspace.SearchBox.Toggle'
)
})
const search = page.getByRole('search')
await expect(search).toBeVisible()
await search.getByRole('combobox').fill('PartnerNode')
await expect(search.getByText('AllowedPartnerNode')).toBeVisible()
await expect(search.getByText('DisabledPartnerNode')).toHaveCount(0)
})
})

View File

@@ -0,0 +1,18 @@
import { expect } from '@playwright/test'
import { partnerNodeGovernanceTest as test } from '@e2e/fixtures/partnerNodeGovernanceFixture'
import { TestIds } from '@e2e/fixtures/selectors'
test.describe('Partner node governance workbench', { tag: '@cloud' }, () => {
test('blocks queueing a disabled partner node', async ({
page,
partnerNodeGovernance
}) => {
await page.getByTestId(TestIds.topbar.queueButton).click()
await expect(page.getByRole('alert')).toContainText(
'Workflow blocked by workspace policy'
)
expect(partnerNodeGovernance.promptRequestCount()).toBe(0)
})
})

View File

@@ -0,0 +1,91 @@
import { fromPartial } from '@total-typescript/shoehorn'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { LGraph, LGraphNode } from '@/lib/litegraph/src/litegraph'
import type { ComfyApp } from '@/scripts/app'
import type { QueuePromptGuard } from '@/services/queuePromptGuardService'
import type { ComfyExtension } from '@/types/comfy'
const {
addToast,
isNodeDisabled,
registerQueuePromptGuard,
registerExtension,
usePartnerNodeGovernanceStore
} = vi.hoisted(() => {
const isNodeDisabled = vi.fn()
return {
addToast: vi.fn(),
isNodeDisabled,
registerQueuePromptGuard: vi.fn<
(id: string, guard: QueuePromptGuard) => () => void
>(() => () => {}),
registerExtension: vi.fn<(extension: ComfyExtension) => void>(),
usePartnerNodeGovernanceStore: vi.fn(() => ({ isNodeDisabled }))
}
})
vi.mock('@/platform/workspace/stores/partnerNodeGovernanceStore', () => ({
usePartnerNodeGovernanceStore
}))
vi.mock('@/platform/updates/common/toastStore', () => ({
useToastStore: () => ({ add: addToast })
}))
vi.mock('@/services/queuePromptGuardService', () => ({
registerQueuePromptGuard
}))
vi.mock('@/services/extensionService', () => ({
useExtensionService: () => ({ registerExtension })
}))
describe('cloudPartnerNodeGovernance', () => {
beforeEach(() => {
vi.resetModules()
vi.clearAllMocks()
})
async function loadExtension(): Promise<ComfyExtension> {
await import('./cloudPartnerNodeGovernance')
const extension = registerExtension.mock.calls[0]?.[0]
if (!extension)
throw new Error('Expected governance extension registration')
return extension
}
it('initializes governance during Cloud setup', async () => {
const extension = await loadExtension()
expect(extension.name).toBe('Comfy.Cloud.PartnerNodeGovernance')
extension.setup?.(fromPartial<ComfyApp>({}))
expect(usePartnerNodeGovernanceStore).toHaveBeenCalledOnce()
expect(registerQueuePromptGuard).toHaveBeenCalledOnce()
})
it('blocks queueing when the graph contains a disabled partner node', async () => {
const graph = new LGraph()
graph.add(new LGraphNode('DisabledPartnerNode', 'DisabledPartnerNode'))
isNodeDisabled.mockImplementation(
(nodeType) => nodeType === 'DisabledPartnerNode'
)
const extension = await loadExtension()
extension.setup?.(fromPartial<ComfyApp>({}))
const guard = registerQueuePromptGuard.mock.calls[0]?.[1]
if (!guard) throw new Error('Expected queue guard registration')
const result = await guard({ rootGraph: graph })
expect(result).toBe(false)
expect(isNodeDisabled).toHaveBeenCalledWith('DisabledPartnerNode')
expect(addToast).toHaveBeenCalledWith({
severity: 'error',
summary: 'Workflow blocked by workspace policy',
detail:
'Remove disabled partner nodes from this workflow or ask a workspace owner to update the policy.',
life: 6000
})
})
})

View File

@@ -0,0 +1,35 @@
import { t } from '@/i18n'
import { usePartnerNodeGovernanceStore } from '@/platform/workspace/stores/partnerNodeGovernanceStore'
import { useToastStore } from '@/platform/updates/common/toastStore'
import { useExtensionService } from '@/services/extensionService'
import { registerQueuePromptGuard } from '@/services/queuePromptGuardService'
import { forEachNode } from '@/utils/graphTraversalUtil'
const QUEUE_GUARD_ID = 'workspace.partner-node-governance'
useExtensionService().registerExtension({
name: 'Comfy.Cloud.PartnerNodeGovernance',
setup: () => {
usePartnerNodeGovernanceStore()
registerQueuePromptGuard(QUEUE_GUARD_ID, (context) => {
const governanceStore = usePartnerNodeGovernanceStore()
let hasDisabledPartnerNode = false
forEachNode(context.rootGraph, (node) => {
if (node.type && governanceStore.isNodeDisabled(node.type)) {
hasDisabledPartnerNode = true
}
})
if (!hasDisabledPartnerNode) return true
useToastStore().add({
severity: 'error',
summary: t('workspacePanel.allowlist.runBlocked.summary'),
detail: t('workspacePanel.allowlist.runBlocked.detail'),
life: 6000
})
return false
})
}
})

View File

@@ -38,6 +38,7 @@ if (isCloud) {
await import('./cloudRemoteConfig')
await import('./cloudBadges')
await import('./cloudSessionCookie')
await import('./cloudPartnerNodeGovernance')
}
// Feedback button for cloud and nightly builds

View File

@@ -2933,8 +2933,17 @@
"noMatches": "No partner nodes match your search.",
"allowedCount": "{enabled} of {total} allowed",
"nodeToggle": "Allow {name}",
"saved": "Allowlist saved",
"saveError": "Allowlist could not be saved."
"saved": "Partner node policy saved",
"saveError": "Partner node policy could not be saved.",
"runBlocked": {
"summary": "Workflow blocked by workspace policy",
"detail": "Remove disabled partner nodes from this workflow or ask a workspace owner to update the policy."
},
"enforcement": {
"title": "Enforce partner node allowlist",
"description": "Block partner nodes that are not explicitly allowed.",
"toggle": "Enforce partner node allowlist"
}
},
"dashboard": {
"placeholder": "Dashboard workspace settings"

View File

@@ -139,6 +139,11 @@ describe('PartnerNodeAllowlistPanel', () => {
expect(screen.getByText('BFL')).toBeInTheDocument()
expect(screen.getByText('Google')).toBeInTheDocument()
expect(
screen.getByRole('switch', {
name: 'Enforce partner node allowlist'
})
).toBeChecked()
expect(
screen.getByRole('switch', { name: 'Allow Flux Fill' })
).not.toBeChecked()
@@ -166,19 +171,56 @@ describe('PartnerNodeAllowlistPanel', () => {
})
expect(mockToastAdd).toHaveBeenCalledWith({
severity: 'success',
summary: 'Allowlist saved',
summary: 'Partner node policy saved',
life: 2000
})
})
it('saves enforcement with the existing allowlist', async () => {
const user = userEvent.setup()
state.policy.value = {
enforcementEnabled: false,
nodes: {
FluxFill: false,
FluxExpand: true,
VeoVideo: true,
RetiredNode: false
}
}
renderPanel()
await user.click(
screen.getByRole('switch', {
name: 'Enforce partner node allowlist'
})
)
await user.click(screen.getByRole('button', { name: 'Save' }))
expect(mockSavePolicy).toHaveBeenCalledWith({
enforcementEnabled: true,
nodes: {
FluxFill: false,
FluxExpand: true,
VeoVideo: true,
RetiredNode: false
}
})
})
it('treats an unconfigured workspace as allow-all draft', async () => {
state.policy.value = null
state.status.value = 'unconfigured'
renderPanel()
await nextTick()
expect(screen.getAllByRole('switch')).toHaveLength(3)
for (const toggle of screen.getAllByRole('switch')) {
expect(
screen.getByRole('switch', {
name: 'Enforce partner node allowlist'
})
).not.toBeChecked()
const nodeToggles = screen.getAllByRole('switch', { name: /^Allow / })
expect(nodeToggles).toHaveLength(3)
for (const toggle of nodeToggles) {
expect(toggle).toBeChecked()
}
expect(screen.getByRole('button', { name: 'Save' })).toBeDisabled()
@@ -222,7 +264,7 @@ describe('PartnerNodeAllowlistPanel', () => {
expect(toggle).toBeChecked()
expect(screen.getByRole('alert')).toHaveTextContent(
'Allowlist could not be saved.'
'Partner node policy could not be saved.'
)
})
})

View File

@@ -21,6 +21,26 @@
/>
</div>
<div
v-if="isReady"
class="flex min-h-14 items-center gap-4 rounded-xl border border-interface-stroke/60 px-4 py-3"
>
<div class="min-w-0 flex-1">
<div class="text-sm text-base-foreground">
{{ $t('workspacePanel.allowlist.enforcement.title') }}
</div>
<div class="mt-0.5 text-xs text-muted-foreground">
{{ $t('workspacePanel.allowlist.enforcement.description') }}
</div>
</div>
<ToggleSwitch
:model-value="draftEnforcementEnabled"
:disabled="isSaving"
:aria-label="$t('workspacePanel.allowlist.enforcement.toggle')"
@update:model-value="setEnforcementEnabled"
/>
</div>
<div
v-if="status === 'loading'"
role="status"
@@ -157,6 +177,7 @@ const toastStore = useToastStore()
const searchQuery = ref('')
const draftNodes = ref<Record<string, boolean>>({})
const draftEnforcementEnabled = ref(false)
const isSaving = ref(false)
const saveError = ref(false)
@@ -169,6 +190,7 @@ function originalNodeValue(nodeId: string): boolean {
}
function resetDraft(): void {
draftEnforcementEnabled.value = policy.value?.enforcementEnabled ?? false
draftNodes.value = Object.fromEntries(
partnerNodes.value.map((node) => [node.id, originalNodeValue(node.id)])
)
@@ -179,10 +201,13 @@ watch([governedWorkspaceId, partnerNodes, policy], resetDraft, {
immediate: true
})
const hasChanges = computed(() =>
partnerNodes.value.some(
(node) => draftNodes.value[node.id] !== originalNodeValue(node.id)
)
const hasChanges = computed(
() =>
draftEnforcementEnabled.value !==
(policy.value?.enforcementEnabled ?? false) ||
partnerNodes.value.some(
(node) => draftNodes.value[node.id] !== originalNodeValue(node.id)
)
)
const groups = computed<PartnerNodeGroup[]>(() => {
@@ -224,6 +249,11 @@ function setNodeEnabled(nodeId: string, enabled: boolean): void {
saveError.value = false
}
function setEnforcementEnabled(enabled: boolean): void {
draftEnforcementEnabled.value = enabled
saveError.value = false
}
async function save(): Promise<void> {
if (!hasChanges.value || isSaving.value) return
@@ -231,7 +261,7 @@ async function save(): Promise<void> {
saveError.value = false
try {
const applied = await governanceStore.savePolicy({
enforcementEnabled: policy.value?.enforcementEnabled ?? false,
enforcementEnabled: draftEnforcementEnabled.value,
nodes: {
...(policy.value?.nodes ?? {}),
...draftNodes.value

View File

@@ -3,6 +3,7 @@ import { setActivePinia } from 'pinia'
import { nextTick } from 'vue'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { LGraphNode, LiteGraph } from '@/lib/litegraph/src/litegraph'
import type * as PartnerNodePolicyApi from '@/platform/workspace/api/partnerNodePolicyApi'
import type { PartnerNodePolicy } from '@/platform/workspace/api/partnerNodePolicyApi'
import { PartnerNodePolicyApiError } from '@/platform/workspace/api/partnerNodePolicyApi'
@@ -106,6 +107,7 @@ describe('partnerNodeGovernanceStore', () => {
})
afterEach(() => {
delete LiteGraph.registered_node_types.LegacyGovernedNode
store?.$dispose()
store = undefined
})
@@ -160,6 +162,47 @@ describe('partnerNodeGovernanceStore', () => {
expect(store.isNodeDisabled('CoreNode')).toBe(false)
})
it('filters disabled nodes out of discovery', async () => {
mockGetPartnerNodePolicy.mockResolvedValue({
enforcementEnabled: true,
nodes: { AllowedNode: true, DisabledNode: false }
} satisfies PartnerNodePolicy)
store = await createLoadedStore()
const nodeDefStore = useNodeDefStore()
expect(nodeDefStore.visibleNodeDefs.map((node) => node.name)).toEqual([
'AllowedNode',
'CoreNode'
])
expect(nodeDefStore.nodeSearchService.searchNode('DisabledNode')).toEqual(
[]
)
})
it('keeps the litegraph legacy menu in sync with policy changes', async () => {
class LegacyGovernedNode extends LGraphNode {}
LiteGraph.registered_node_types.LegacyGovernedNode = LegacyGovernedNode
useNodeDefStore().addNodeDef(nodeDef('LegacyGovernedNode'))
mockGetPartnerNodePolicy.mockResolvedValue({
enforcementEnabled: true,
nodes: { LegacyGovernedNode: false }
} satisfies PartnerNodePolicy)
store = await createLoadedStore()
expect(LegacyGovernedNode.skip_list).toBe(true)
mockGetPartnerNodePolicy.mockResolvedValue({
enforcementEnabled: false,
nodes: { LegacyGovernedNode: false }
} satisfies PartnerNodePolicy)
await store.loadPolicy()
await nextTick()
expect(LegacyGovernedNode.skip_list).toBe(false)
})
it('fails closed for a 503 from an enforcing workspace', async () => {
mockGetPartnerNodePolicy.mockRejectedValue(
new PartnerNodePolicyApiError(503, 'Service Unavailable')

View File

@@ -1,7 +1,8 @@
import { defineStore } from 'pinia'
import { computed, ref, shallowRef, watch } from 'vue'
import { computed, ref, shallowRef, watch, watchEffect } from 'vue'
import { useFeatureFlags } from '@/composables/useFeatureFlags'
import { LiteGraph } from '@/lib/litegraph/src/litegraph'
import {
getPartnerNodePolicy,
PartnerNodePolicyApiError,
@@ -25,6 +26,8 @@ export type PartnerNodePolicyStatus =
| 'unavailable'
| 'error'
const DISCOVERY_FILTER_ID = 'workspace.partner-node-governance'
export const usePartnerNodeGovernanceStore = defineStore(
'partnerNodeGovernance',
() => {
@@ -72,6 +75,34 @@ export const usePartnerNodeGovernanceStore = defineStore(
)
}
nodeDefStore.registerNodeDefFilter({
id: DISCOVERY_FILTER_ID,
name: 'Workspace partner-node governance',
predicate: (nodeDef) => !isNodeDisabled(nodeDef.name)
})
const legacyHiddenNodeTypes = new Set<string>()
watchEffect(() => {
const showDevOnly = nodeDefStore.showDevOnly
for (const nodeDef of Object.values(nodeDefStore.nodeDefsByName)) {
if (!nodeDef.api_node) continue
const nodeType = LiteGraph.registered_node_types[nodeDef.name]
if (!nodeType) continue
if (isNodeDisabled(nodeDef.name)) {
if (!nodeType.skip_list) {
nodeType.skip_list = true
legacyHiddenNodeTypes.add(nodeDef.name)
}
continue
}
if (legacyHiddenNodeTypes.delete(nodeDef.name)) {
nodeType.skip_list = nodeDef.dev_only && !showDevOnly
}
}
})
async function loadPolicy(): Promise<void> {
const workspaceId = governedWorkspaceId.value
const version = ++loadVersion

View File

@@ -32,6 +32,7 @@ import {
createTestSubgraphNode
} from '@/lib/litegraph/src/subgraph/__fixtures__/subgraphHelpers'
import { useWidgetValueStore } from '@/stores/widgetValueStore'
import { registerQueuePromptGuard } from '@/services/queuePromptGuardService'
const {
mockApiKeyAuthStore,
@@ -189,6 +190,93 @@ describe('ComfyApp', () => {
})
describe('queuePrompt', () => {
it('blocks before adding a disallowed prompt to the client queue', async () => {
const unregisterGuard = registerQueuePromptGuard(
'app.test.before-queue',
() => false
)
const dispatchEvent = vi
.spyOn(api, 'dispatchCustomEvent')
.mockImplementation(() => true)
const queuePrompt = vi.spyOn(api, 'queuePrompt')
try {
await expect(app.queuePrompt(0)).resolves.toBe(false)
} finally {
unregisterGuard()
}
expect(mockAuthStore.getAuthToken).not.toHaveBeenCalled()
expect(dispatchEvent).not.toHaveBeenCalled()
expect(queuePrompt).not.toHaveBeenCalled()
})
it('checks again after deferred authentication resolves', async () => {
const graph = new LGraph()
Reflect.set(app, 'rootGraphInternal', graph)
let resolveAuth: (token: string | undefined) => void = () => {}
mockAuthStore.getAuthToken.mockReturnValue(
new Promise((resolve) => {
resolveAuth = resolve
})
)
let guardCalls = 0
const unregisterGuard = registerQueuePromptGuard(
'app.test.after-auth',
() => ++guardCalls === 1
)
const queuePrompt = vi.spyOn(api, 'queuePrompt')
try {
const result = app.queuePrompt(0)
await vi.waitFor(() => {
expect(mockAuthStore.getAuthToken).toHaveBeenCalledOnce()
})
resolveAuth(undefined)
await expect(result).resolves.toBe(false)
} finally {
unregisterGuard()
}
expect(guardCalls).toBe(2)
expect(queuePrompt).not.toHaveBeenCalled()
})
it('dispatches queued progress when a later batch item is blocked', async () => {
const graph = new LGraph()
Reflect.set(app, 'rootGraphInternal', graph)
vi.spyOn(app, 'graphToPrompt').mockResolvedValue({
output: {},
workflow: createWorkflowGraphData()
})
vi.spyOn(api, 'queuePrompt').mockResolvedValue({
prompt_id: 'job-1',
node_errors: {},
error: ''
})
const dispatchEvent = vi
.spyOn(api, 'dispatchCustomEvent')
.mockImplementation(() => true)
let guardCalls = 0
const unregisterGuard = registerQueuePromptGuard(
'app.test.partial-batch',
() => ++guardCalls < 3
)
try {
await expect(app.queuePrompt(0, 2)).resolves.toBe(false)
} finally {
unregisterGuard()
}
expect(api.queuePrompt).toHaveBeenCalledOnce()
expect(dispatchEvent).toHaveBeenCalledWith('promptQueued', {
number: 0,
batchCount: 1,
requestId: 1
})
})
it('shows the error overlay for successful prompt responses with node errors', async () => {
const graph = new LGraph()
const workflow = new ComfyWorkflow({

View File

@@ -68,6 +68,7 @@ import { resolveAccountPrecondition } from '@/platform/errorCatalog/accountPreco
import { useDialogService } from '@/services/dialogService'
import { useExtensionService } from '@/services/extensionService'
import { useLitegraphService } from '@/services/litegraphService'
import { runQueuePromptGuards } from '@/services/queuePromptGuardService'
import { useSubgraphService } from '@/services/subgraphService'
import { useApiKeyAuthStore } from '@/stores/apiKeyAuthStore'
import { useCommandStore } from '@/stores/commandStore'
@@ -1611,11 +1612,17 @@ export class ComfyApp {
})
}
private async isQueuePromptAllowed(): Promise<boolean> {
return await runQueuePromptGuards(this)
}
async queuePrompt(
number: number,
batchCount: number = 1,
queueNodeIds?: NodeExecutionId[]
): Promise<boolean> {
if (!(await this.isQueuePromptAllowed())) return false
const requestId = this.nextQueueRequestId++
this.queueItems.push({ number, batchCount, queueNodeIds, requestId })
api.dispatchCustomEvent('promptQueueing', {
@@ -1648,6 +1655,18 @@ export class ComfyApp {
const isPartialExecution = !!queueNodeIds?.length
for (let i = 0; i < batchCount; i++) {
if (!(await this.isQueuePromptAllowed())) {
this.queueItems.length = 0
if (queuedCount > 0) {
api.dispatchCustomEvent('promptQueued', {
number,
batchCount: queuedCount,
requestId
})
}
return false
}
// Allow widgets to run callbacks before a prompt has been queued
// e.g. random seed before every gen
forEachNode(this.rootGraph, (node) => {

View File

@@ -0,0 +1,30 @@
import type { LGraph } from '@/lib/litegraph/src/litegraph'
export interface QueuePromptGuardContext {
readonly rootGraph: LGraph
}
export type QueuePromptGuard = (
context: QueuePromptGuardContext
) => boolean | Promise<boolean>
const guards = new Map<string, QueuePromptGuard>()
export function registerQueuePromptGuard(
id: string,
guard: QueuePromptGuard
): () => void {
guards.set(id, guard)
return () => {
if (guards.get(id) === guard) guards.delete(id)
}
}
export async function runQueuePromptGuards(
context: QueuePromptGuardContext
): Promise<boolean> {
const results = await Promise.all(
[...guards.values()].map((guard) => guard(context))
)
return results.every((result) => result !== false)
}

View File

@@ -1,6 +1,7 @@
import axios from 'axios'
import { createTestingPinia } from '@pinia/testing'
import { setActivePinia } from 'pinia'
import { beforeEach, describe, expect, it } from 'vitest'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { promoteValueWidgetViaSubgraphInput } from '@/core/graph/subgraph/promotionUtils'
import { LGraphNode } from '@/lib/litegraph/src/litegraph'
@@ -10,7 +11,7 @@ import {
createTestSubgraphNode
} from '@/lib/litegraph/src/subgraph/__fixtures__/subgraphHelpers'
import type { ComfyNodeDef } from '@/schemas/nodeDefSchema'
import { useNodeDefStore } from '@/stores/nodeDefStore'
import { useNodeDefStore, useNodeFrequencyStore } from '@/stores/nodeDefStore'
import type { NodeDefFilter } from '@/stores/nodeDefStore'
describe('useNodeDefStore', () => {
@@ -21,6 +22,10 @@ describe('useNodeDefStore', () => {
store = useNodeDefStore()
})
afterEach(() => {
vi.restoreAllMocks()
})
const createMockNodeDef = (
overrides: Partial<ComfyNodeDef> = {}
): ComfyNodeDef => ({
@@ -286,6 +291,30 @@ describe('useNodeDefStore', () => {
})
})
describe('node frequency visibility', () => {
it('excludes node definitions hidden by registered filters', async () => {
vi.spyOn(axios, 'get').mockResolvedValue({
data: { HiddenNode: 2, VisibleNode: 1 }
})
store.updateNodeDefs([
createMockNodeDef({ name: 'HiddenNode' }),
createMockNodeDef({ name: 'VisibleNode' })
])
store.registerNodeDefFilter({
id: 'test.hide-node',
name: 'Hide node',
predicate: (nodeDef) => nodeDef.name !== 'HiddenNode'
})
const frequencyStore = useNodeFrequencyStore()
await frequencyStore.loadNodeFrequencies()
expect(frequencyStore.topNodeDefs.map((nodeDef) => nodeDef.name)).toEqual(
['VisibleNode']
)
})
})
describe('allNodeDefsByName', () => {
it('should include all node defs by name', () => {
const node1 = createMockNodeDef({ name: 'Node1' })

View File

@@ -562,9 +562,15 @@ export const useNodeFrequencyStore = defineStore('nodeFrequency', () => {
const nodeDefStore = useNodeDefStore()
const topNodeDefs = computed<ComfyNodeDefImpl[]>(() => {
const visibleNodeNames = new Set(
nodeDefStore.visibleNodeDefs.map((nodeDef) => nodeDef.name)
)
return nodeNamesByFrequency.value
.map((nodeName: string) => nodeDefStore.nodeDefsByName[nodeName])
.filter((nodeDef: ComfyNodeDefImpl) => nodeDef !== undefined)
.filter(
(nodeDef): nodeDef is ComfyNodeDefImpl =>
nodeDef !== undefined && visibleNodeNames.has(nodeDef.name)
)
.slice(0, topNodeDefLimit.value)
})