Skip to content

Commit 3fa6920

Browse files
feat: wire essentials_category for Essentials tab display
- Add essentialsNodes.ts as single source of truth for node categorization and ordering, consolidating three separate lists - Pass essentials_category through subgraphStore from blueprints and global subgraph data to node definitions - Add essentials_category to SubgraphDefinitionBase schema - Refactor toolkitNodes.ts, nodeSource.ts, nodeOrganizationService.ts to import from the centralized constants - Add tests for constants integrity and subgraph passthrough Fixes COM-15221 Amp-Thread-ID: https://ampcode.com/threads/T-019c83de-f7ab-7779-a451-0ba5940b56a9
1 parent 35f15d1 commit 3fa6920

8 files changed

Lines changed: 236 additions & 129 deletions

File tree

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
import { describe, expect, it } from 'vitest'
2+
3+
import {
4+
ESSENTIALS_CATEGORIES,
5+
ESSENTIALS_CATEGORY_MAP,
6+
ESSENTIALS_NODES,
7+
TOOLKIT_BLUEPRINT_MODULES,
8+
TOOLKIT_NOVEL_NODE_NAMES
9+
} from './essentialsNodes'
10+
11+
describe('essentialsNodes', () => {
12+
it('has no duplicate node names across categories', () => {
13+
const seen = new Map<string, string>()
14+
for (const [category, nodes] of Object.entries(ESSENTIALS_NODES)) {
15+
for (const node of nodes) {
16+
expect(
17+
seen.has(node),
18+
`"${node}" duplicated in "${category}" and "${seen.get(node)}"`
19+
).toBe(false)
20+
seen.set(node, category)
21+
}
22+
}
23+
})
24+
25+
it('ESSENTIALS_CATEGORY_MAP covers every node in ESSENTIALS_NODES', () => {
26+
for (const [category, nodes] of Object.entries(ESSENTIALS_NODES)) {
27+
for (const node of nodes) {
28+
expect(ESSENTIALS_CATEGORY_MAP[node]).toBe(category)
29+
}
30+
}
31+
})
32+
33+
it('TOOLKIT_NOVEL_NODE_NAMES excludes basics nodes', () => {
34+
for (const basicNode of ESSENTIALS_NODES.basics) {
35+
expect(TOOLKIT_NOVEL_NODE_NAMES.has(basicNode)).toBe(false)
36+
}
37+
})
38+
39+
it('TOOLKIT_NOVEL_NODE_NAMES excludes SubgraphBlueprint-prefixed nodes', () => {
40+
for (const name of TOOLKIT_NOVEL_NODE_NAMES) {
41+
expect(name.startsWith('SubgraphBlueprint.')).toBe(false)
42+
}
43+
})
44+
45+
it('ESSENTIALS_NODES keys match ESSENTIALS_CATEGORIES', () => {
46+
const nodeKeys = Object.keys(ESSENTIALS_NODES)
47+
expect(nodeKeys).toEqual([...ESSENTIALS_CATEGORIES])
48+
})
49+
50+
it('TOOLKIT_BLUEPRINT_MODULES contains comfy_essentials', () => {
51+
expect(TOOLKIT_BLUEPRINT_MODULES.has('comfy_essentials')).toBe(true)
52+
})
53+
})

src/constants/essentialsNodes.ts

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
/**
2+
* Single source of truth for Essentials tab node categorization and ordering.
3+
*
4+
* Adding a new node to the Essentials tab? Add it here and nowhere else.
5+
*
6+
* Source: https://www.notion.so/comfy-org/2fe6d73d365080d0a951d14cdf540778
7+
*/
8+
9+
export const ESSENTIALS_CATEGORIES = [
10+
'basics',
11+
'text generation',
12+
'image generation',
13+
'video generation',
14+
'image tools',
15+
'video tools',
16+
'audio',
17+
'3D'
18+
] as const
19+
20+
export type EssentialsCategory = (typeof ESSENTIALS_CATEGORIES)[number]
21+
22+
/**
23+
* Ordered list of nodes per category.
24+
* Array order = display order in the Essentials tab.
25+
* Presence in a category = the node's essentials_category mock fallback.
26+
*/
27+
export const ESSENTIALS_NODES: Record<EssentialsCategory, readonly string[]> = {
28+
basics: [
29+
'LoadImage',
30+
'LoadVideo',
31+
'Load3D',
32+
'SaveImage',
33+
'SaveVideo',
34+
'SaveGLB',
35+
'PrimitiveStringMultiline',
36+
'PreviewImage'
37+
],
38+
'text generation': ['OpenAIChatNode'],
39+
'image generation': [
40+
'LoraLoader',
41+
'LoraLoaderModelOnly',
42+
'ConditioningCombine'
43+
],
44+
'video generation': [
45+
'SubgraphBlueprint.pose_to_video_ltx_2_0',
46+
'SubgraphBlueprint.canny_to_video_ltx_2_0',
47+
'KlingLipSyncAudioToVideoNode',
48+
'KlingOmniProEditVideoNode'
49+
],
50+
'image tools': [
51+
'ImageBatch',
52+
'ImageCrop',
53+
'ImageCropV2',
54+
'ImageScale',
55+
'ImageScaleBy',
56+
'ImageRotate',
57+
'ImageBlur',
58+
'ImageBlend',
59+
'ImageInvert',
60+
'ImageCompare',
61+
'Canny',
62+
'RecraftRemoveBackgroundNode',
63+
'RecraftVectorizeImageNode',
64+
'LoadImageMask'
65+
],
66+
'video tools': ['GetVideoComponents', 'CreateVideo', 'Video Slice'],
67+
audio: [
68+
'LoadAudio',
69+
'SaveAudio',
70+
'SaveAudioMP3',
71+
'StabilityTextToAudio',
72+
'EmptyLatentAudio'
73+
],
74+
'3D': ['TencentTextToModelNode', 'TencentImageToModelNode']
75+
}
76+
77+
/**
78+
* Flat map: node name → category (derived from ESSENTIALS_NODES).
79+
* Used as mock/fallback when backend doesn't provide essentials_category.
80+
*/
81+
export const ESSENTIALS_CATEGORY_MAP: Record<string, string> =
82+
Object.fromEntries(
83+
Object.entries(ESSENTIALS_NODES).flatMap(([category, nodes]) =>
84+
nodes.map((node) => [node, category])
85+
)
86+
)
87+
88+
/**
89+
* "Novel" toolkit nodes for telemetry — basics excluded.
90+
* Derived from ESSENTIALS_NODES minus the 'basics' category.
91+
*/
92+
export const TOOLKIT_NOVEL_NODE_NAMES: ReadonlySet<string> = new Set(
93+
Object.entries(ESSENTIALS_NODES)
94+
.filter(([cat]) => cat !== 'basics')
95+
.flatMap(([, nodes]) => nodes)
96+
.filter((n) => !n.startsWith('SubgraphBlueprint.'))
97+
)
98+
99+
/**
100+
* python_module values that identify toolkit blueprint nodes.
101+
*/
102+
export const TOOLKIT_BLUEPRINT_MODULES: ReadonlySet<string> = new Set([
103+
'comfy_essentials'
104+
])

src/constants/toolkitNodes.ts

Lines changed: 5 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,38 +1,10 @@
11
/**
22
* Toolkit (Essentials) node detection constants.
33
*
4+
* Re-exported from essentialsNodes.ts — the single source of truth.
45
* Used by telemetry to track toolkit node adoption and popularity.
5-
* Only novel nodes — basic nodes (LoadImage, SaveImage, etc.) are excluded.
6-
*
7-
* Source: https://www.notion.so/comfy-org/2fe6d73d365080d0a951d14cdf540778
8-
*/
9-
10-
/**
11-
* Canonical node type names for individual toolkit nodes.
12-
*/
13-
export const TOOLKIT_NODE_NAMES: ReadonlySet<string> = new Set([
14-
// Image Tools
15-
'ImageCrop',
16-
'ImageRotate',
17-
'ImageBlur',
18-
'ImageInvert',
19-
'ImageCompare',
20-
'Canny',
21-
22-
// Video Tools
23-
'Video Slice',
24-
25-
// API Nodes
26-
'RecraftRemoveBackgroundNode',
27-
'RecraftVectorizeImageNode',
28-
'KlingOmniProEditVideoNode'
29-
])
30-
31-
/**
32-
* python_module values that identify toolkit blueprint nodes.
33-
* Essentials blueprints are registered with node_pack 'comfy_essentials',
34-
* which maps to python_module on the node def.
356
*/
36-
export const TOOLKIT_BLUEPRINT_MODULES: ReadonlySet<string> = new Set([
37-
'comfy_essentials'
38-
])
7+
export {
8+
TOOLKIT_NOVEL_NODE_NAMES as TOOLKIT_NODE_NAMES,
9+
TOOLKIT_BLUEPRINT_MODULES
10+
} from './essentialsNodes'

src/platform/workflow/validation/schemas/workflowSchema.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -399,6 +399,7 @@ interface SubgraphDefinitionBase<
399399
/** Optional description shown as tooltip when hovering over the subgraph node. */
400400
description?: string
401401
category?: string
402+
essentials_category?: string
402403
/** Custom metadata for the subgraph (description, searchAliases, etc.) */
403404
extra?: T extends ComfyWorkflow1BaseInput
404405
? z.input<typeof zExtra> | null
@@ -437,6 +438,7 @@ const zSubgraphDefinition = zComfyWorkflow1
437438
/** Optional description shown as tooltip when hovering over the subgraph node. */
438439
description: z.string().optional(),
439440
category: z.string().optional(),
441+
essentials_category: z.string().optional(),
440442
inputNode: zExportedSubgraphIONode,
441443
outputNode: zExportedSubgraphIONode,
442444

src/services/nodeOrganizationService.ts

Lines changed: 13 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,8 @@
1+
import type { EssentialsCategory } from '@/constants/essentialsNodes'
2+
import {
3+
ESSENTIALS_CATEGORIES,
4+
ESSENTIALS_NODES
5+
} from '@/constants/essentialsNodes'
16
import type { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
27
import { buildNodeDefTree } from '@/stores/nodeDefStore'
38
import type {
@@ -14,46 +19,6 @@ import { upperCase } from 'es-toolkit/string'
1419

1520
const DEFAULT_ICON = 'pi pi-sort'
1621

17-
const NODE_ORDER_BY_FOLDER = {
18-
basics: [
19-
'LoadImage',
20-
'LoadVideo',
21-
'Load3D',
22-
'SaveImage',
23-
'SaveVideo',
24-
'SaveGLB',
25-
'PrimitiveStringMultiline',
26-
'PreviewImage'
27-
],
28-
'image tools': [
29-
'ImageBatch',
30-
'ImageCrop',
31-
'ImageCropV2',
32-
'ImageScale',
33-
'ImageScaleBy',
34-
'ImageRotate',
35-
'ImageBlur',
36-
'ImageBlend',
37-
'ImageInvert',
38-
'Canny',
39-
'RecraftRemoveBackgroundNode',
40-
'LoadImageMask'
41-
],
42-
'video tools': ['GetVideoComponents', 'CreateVideo'],
43-
'image generation': [
44-
'LoraLoader',
45-
'LoraLoaderModelOnly',
46-
'ConditioningCombine'
47-
],
48-
audio: [
49-
'LoadAudio',
50-
'SaveAudio',
51-
'SaveAudioMP3',
52-
'StabilityTextToAudio',
53-
'EmptyLatentAudio'
54-
]
55-
} as const satisfies Record<string, readonly string[]>
56-
5722
export const DEFAULT_GROUPING_ID = 'category' as const
5823
export const DEFAULT_SORTING_ID = 'original' as const
5924
export const DEFAULT_TAB_ID = 'all' as const
@@ -178,34 +143,25 @@ class NodeOrganizationService {
178143
const tree = buildNodeDefTree(essentialNodes, {
179144
pathExtractor: essentialsPathExtractor
180145
})
181-
const folderOrder = [
182-
'basics',
183-
'text generation',
184-
'image generation',
185-
'video generation',
186-
'image tools',
187-
'video tools',
188-
'audio',
189-
'3D'
190-
]
191146
if (tree.children) {
192-
const len = folderOrder.length
147+
const len = ESSENTIALS_CATEGORIES.length
193148
const originalIndex = new Map(
194149
tree.children.map((child, i) => [child, i])
195150
)
196151
tree.children.sort((a, b) => {
197-
const ai = folderOrder.indexOf(a.label ?? '')
198-
const bi = folderOrder.indexOf(b.label ?? '')
152+
const ai = ESSENTIALS_CATEGORIES.indexOf(
153+
a.label as EssentialsCategory
154+
)
155+
const bi = ESSENTIALS_CATEGORIES.indexOf(
156+
b.label as EssentialsCategory
157+
)
199158
const orderA = ai === -1 ? len + originalIndex.get(a)! : ai
200159
const orderB = bi === -1 ? len + originalIndex.get(b)! : bi
201160
return orderA - orderB
202161
})
203162
for (const folder of tree.children) {
204163
if (!folder.children) continue
205-
const order =
206-
NODE_ORDER_BY_FOLDER[
207-
folder.label as keyof typeof NODE_ORDER_BY_FOLDER
208-
]
164+
const order = ESSENTIALS_NODES[folder.label as EssentialsCategory]
209165
if (!order) continue
210166
const nodeOrder: readonly string[] = order
211167
const orderLen = nodeOrder.length

src/stores/subgraphStore.test.ts

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -515,6 +515,56 @@ describe('useSubgraphStore', () => {
515515
})
516516
})
517517

518+
describe('essentials_category passthrough', () => {
519+
it('should pass essentials_category from GlobalSubgraphData to node def', async () => {
520+
await mockFetch(
521+
{},
522+
{
523+
bp_essentials: {
524+
name: 'Test Essentials Blueprint',
525+
info: { node_pack: 'test_pack', category: 'Test Category' },
526+
data: JSON.stringify(mockGraph),
527+
essentials_category: 'Image Generation'
528+
}
529+
}
530+
)
531+
const nodeDef = useNodeDefStore().nodeDefs.find(
532+
(d) => d.name === 'SubgraphBlueprint.bp_essentials'
533+
)
534+
expect(nodeDef).toBeDefined()
535+
expect(nodeDef?.essentials_category).toBe('Image Generation')
536+
})
537+
538+
it('should extract essentials_category from subgraph definition as fallback', async () => {
539+
const graphWithEssentials = {
540+
...mockGraph,
541+
definitions: {
542+
subgraphs: [
543+
{
544+
...mockGraph.definitions?.subgraphs?.[0],
545+
essentials_category: 'Image Tools'
546+
}
547+
]
548+
}
549+
}
550+
await mockFetch(
551+
{},
552+
{
553+
bp_fallback: {
554+
name: 'Fallback Blueprint',
555+
info: { node_pack: 'test_pack' },
556+
data: JSON.stringify(graphWithEssentials)
557+
}
558+
}
559+
)
560+
const nodeDef = useNodeDefStore().nodeDefs.find(
561+
(d) => d.name === 'SubgraphBlueprint.bp_fallback'
562+
)
563+
expect(nodeDef).toBeDefined()
564+
expect(nodeDef?.essentials_category).toBe('Image Tools')
565+
})
566+
})
567+
518568
describe('global blueprint filtering', () => {
519569
function globalBlueprint(
520570
overrides: Partial<GlobalSubgraphData['info']> = {}

0 commit comments

Comments
 (0)