From ca5729a8e798ed6b94e4c22f3858f6b998d3f450 Mon Sep 17 00:00:00 2001 From: AustinMroz Date: Wed, 29 Oct 2025 20:41:04 -0700 Subject: [PATCH] Add pricing badge when a subgraph contains partner nodes (#6354) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit image ![api-badge-subgraph_00003](https://github.com/user-attachments/assets/067d0398-47e9-4e97-9e1d-67fac2935e55) ┆Issue is synchronized with this [Notion page](https://www.notion.so/PR-6354-Add-pricing-badge-when-a-subgraph-contains-partner-nodes-29b6d73d365081c685bec3e9446970eb) by [Unito](https://www.unito.io) --- src/composables/node/useNodeBadge.ts | 44 ++++++------ src/composables/node/usePriceBadge.ts | 69 +++++++++++++++++++ src/composables/useCoreCommands.ts | 6 +- src/core/graph/subgraph/proxyWidget.ts | 9 ++- src/lib/litegraph/src/LGraph.ts | 8 +++ .../infrastructure/LGraphCanvasEventMap.ts | 5 ++ .../composables/node/useCreditsBadge.test.ts | 64 +++++++++++++++++ 7 files changed, 175 insertions(+), 30 deletions(-) create mode 100644 src/composables/node/usePriceBadge.ts create mode 100644 tests-ui/tests/composables/node/useCreditsBadge.test.ts diff --git a/src/composables/node/useNodeBadge.ts b/src/composables/node/useNodeBadge.ts index 174ffd705..30c4487a3 100644 --- a/src/composables/node/useNodeBadge.ts +++ b/src/composables/node/useNodeBadge.ts @@ -2,6 +2,7 @@ import _ from 'es-toolkit/compat' import { computed, onMounted, watch } from 'vue' import { useNodePricing } from '@/composables/node/useNodePricing' +import { usePriceBadge } from '@/composables/node/usePriceBadge' import { useComputedWithWidgetWatch } from '@/composables/node/useWatchWidget' import { BadgePosition, LGraphBadge } from '@/lib/litegraph/src/litegraph' import type { LGraphNode } from '@/lib/litegraph/src/litegraph' @@ -12,7 +13,6 @@ import type { ComfyNodeDefImpl } from '@/stores/nodeDefStore' import { useNodeDefStore } from '@/stores/nodeDefStore' import { useColorPaletteStore } from '@/stores/workspace/colorPaletteStore' import { NodeBadgeMode } from '@/types/nodeSource' -import { adjustColor } from '@/utils/colorUtil' /** * Add LGraphBadge to LGraphNode based on settings. @@ -27,6 +27,7 @@ export const useNodeBadge = () => { const settingStore = useSettingStore() const extensionStore = useExtensionStore() const colorPaletteStore = useColorPaletteStore() + const priceBadge = usePriceBadge() const nodeSourceBadgeMode = computed( () => @@ -118,29 +119,7 @@ export const useNodeBadge = () => { let creditsBadge const createBadge = () => { const price = nodePricing.getNodeDisplayPrice(node) - - const isLightTheme = - colorPaletteStore.completedActivePalette.light_theme - return new LGraphBadge({ - text: price, - iconOptions: { - unicode: '\ue96b', - fontFamily: 'PrimeIcons', - color: isLightTheme - ? adjustColor('#FABC25', { lightness: 0.5 }) - : '#FABC25', - bgColor: isLightTheme - ? adjustColor('#654020', { lightness: 0.5 }) - : '#654020', - fontSize: 8 - }, - fgColor: - colorPaletteStore.completedActivePalette.colors.litegraph_base - .BADGE_FG_COLOR, - bgColor: isLightTheme - ? adjustColor('#8D6932', { lightness: 0.5 }) - : '#8D6932' - }) + return priceBadge.getCreditsBadge(price) } if (hasDynamicPricing) { @@ -162,6 +141,23 @@ export const useNodeBadge = () => { node.badges.push(() => creditsBadge.value) } + }, + init() { + app.canvas.canvas.addEventListener<'litegraph:set-graph'>( + 'litegraph:set-graph', + () => { + for (const node of app.canvas.graph?.nodes ?? []) + priceBadge.updateSubgraphCredits(node) + } + ) + app.canvas.canvas.addEventListener<'subgraph-converted'>( + 'subgraph-converted', + (e) => priceBadge.updateSubgraphCredits(e.detail.subgraphNode) + ) + }, + afterConfigureGraph() { + for (const node of app.canvas.graph?.nodes ?? []) + priceBadge.updateSubgraphCredits(node) } }) }) diff --git a/src/composables/node/usePriceBadge.ts b/src/composables/node/usePriceBadge.ts new file mode 100644 index 000000000..963ca8b4a --- /dev/null +++ b/src/composables/node/usePriceBadge.ts @@ -0,0 +1,69 @@ +import type { LGraph, LGraphNode } from '@/lib/litegraph/src/litegraph' +import { LGraphBadge } from '@/lib/litegraph/src/litegraph' + +import { useColorPaletteStore } from '@/stores/workspace/colorPaletteStore' +import { adjustColor } from '@/utils/colorUtil' + +export const usePriceBadge = () => { + function updateSubgraphCredits(node: LGraphNode) { + if (!node.isSubgraphNode()) return + node.badges = node.badges.filter((b) => !isCreditsBadge(b)) + const newBadges = collectCreditsBadges(node.subgraph) + if (newBadges.length > 1) { + node.badges.push(getCreditsBadge('Partner Nodes x ' + newBadges.length)) + } else { + node.badges.push(...newBadges) + } + } + function collectCreditsBadges( + graph: LGraph, + visited: Set = new Set() + ): (LGraphBadge | (() => LGraphBadge))[] { + if (visited.has(graph.id)) return [] + visited.add(graph.id) + const badges = [] + for (const node of graph.nodes) { + badges.push( + ...(node.isSubgraphNode() + ? collectCreditsBadges(node.subgraph, visited) + : node.badges.filter((b) => isCreditsBadge(b))) + ) + } + return badges + } + + function isCreditsBadge(badge: LGraphBadge | (() => LGraphBadge)): boolean { + return ( + (typeof badge === 'function' ? badge() : badge).icon?.unicode === '\ue96b' + ) + } + + const colorPaletteStore = useColorPaletteStore() + function getCreditsBadge(price: string): LGraphBadge { + const isLightTheme = colorPaletteStore.completedActivePalette.light_theme + return new LGraphBadge({ + text: price, + iconOptions: { + unicode: '\ue96b', + fontFamily: 'PrimeIcons', + color: isLightTheme + ? adjustColor('#FABC25', { lightness: 0.5 }) + : '#FABC25', + bgColor: isLightTheme + ? adjustColor('#654020', { lightness: 0.5 }) + : '#654020', + fontSize: 8 + }, + fgColor: + colorPaletteStore.completedActivePalette.colors.litegraph_base + .BADGE_FG_COLOR, + bgColor: isLightTheme + ? adjustColor('#8D6932', { lightness: 0.5 }) + : '#8D6932' + }) + } + return { + getCreditsBadge, + updateSubgraphCredits + } +} diff --git a/src/composables/useCoreCommands.ts b/src/composables/useCoreCommands.ts index 3bf3e8582..f27e50317 100644 --- a/src/composables/useCoreCommands.ts +++ b/src/composables/useCoreCommands.ts @@ -5,10 +5,7 @@ import { DEFAULT_DARK_COLOR_PALETTE, DEFAULT_LIGHT_COLOR_PALETTE } from '@/constants/coreColorPalettes' -import { - promoteRecommendedWidgets, - tryToggleWidgetPromotion -} from '@/core/graph/subgraph/proxyWidgetUtils' +import { tryToggleWidgetPromotion } from '@/core/graph/subgraph/proxyWidgetUtils' import { showSubgraphNodeDialog } from '@/core/graph/subgraph/useSubgraphNodeDialog' import { t } from '@/i18n' import { @@ -945,7 +942,6 @@ export function useCoreCommands(): ComfyCommand[] { const { node } = res canvas.select(node) - promoteRecommendedWidgets(node) canvasStore.updateSelectedItems() } }, diff --git a/src/core/graph/subgraph/proxyWidget.ts b/src/core/graph/subgraph/proxyWidget.ts index b97ae74cf..5d54eb628 100644 --- a/src/core/graph/subgraph/proxyWidget.ts +++ b/src/core/graph/subgraph/proxyWidget.ts @@ -1,4 +1,7 @@ -import { demoteWidget } from '@/core/graph/subgraph/proxyWidgetUtils' +import { + demoteWidget, + promoteRecommendedWidgets +} from '@/core/graph/subgraph/proxyWidgetUtils' import { parseProxyWidgets } from '@/core/schemas/proxyWidget' import type { NodeProperty } from '@/lib/litegraph/src/LGraphNode' import type { @@ -62,6 +65,10 @@ export function registerProxyWidgets(canvas: LGraphCanvas) { } } }) + canvas.canvas.addEventListener<'subgraph-converted'>( + 'subgraph-converted', + (e) => promoteRecommendedWidgets(e.detail.subgraphNode) + ) SubgraphNode.prototype.onConfigure = onConfigure } diff --git a/src/lib/litegraph/src/LGraph.ts b/src/lib/litegraph/src/LGraph.ts index 4c7b1ff2a..3a5ce49f4 100644 --- a/src/lib/litegraph/src/LGraph.ts +++ b/src/lib/litegraph/src/LGraph.ts @@ -1710,6 +1710,14 @@ export class LGraph subgraphNode._setConcreteSlots() subgraphNode.arrange() + this.canvasAction((c) => + c.canvas.dispatchEvent( + new CustomEvent('subgraph-converted', { + bubbles: true, + detail: { subgraphNode: subgraphNode as SubgraphNode } + }) + ) + ) return { subgraph, node: subgraphNode as SubgraphNode } } diff --git a/src/lib/litegraph/src/infrastructure/LGraphCanvasEventMap.ts b/src/lib/litegraph/src/infrastructure/LGraphCanvasEventMap.ts index 1c25e862c..ff226c4df 100644 --- a/src/lib/litegraph/src/infrastructure/LGraphCanvasEventMap.ts +++ b/src/lib/litegraph/src/infrastructure/LGraphCanvasEventMap.ts @@ -21,6 +21,11 @@ export interface LGraphCanvasEventMap { fromNode: SubgraphNode } + /** Dispatched after a group of items has been converted to a subgraph*/ + 'subgraph-converted': { + subgraphNode: SubgraphNode + } + 'litegraph:canvas': | { subType: 'before-change' | 'after-change' } | { diff --git a/tests-ui/tests/composables/node/useCreditsBadge.test.ts b/tests-ui/tests/composables/node/useCreditsBadge.test.ts new file mode 100644 index 000000000..72bb92338 --- /dev/null +++ b/tests-ui/tests/composables/node/useCreditsBadge.test.ts @@ -0,0 +1,64 @@ +import { describe, expect, vi } from 'vitest' + +import { LGraphNode } from '@/lib/litegraph/src/litegraph' +import type { LGraphBadge } from '@/lib/litegraph/src/LGraphBadge' +import type { LGraphIcon } from '@/lib/litegraph/src/LGraphIcon' + +import { subgraphTest } from '../../litegraph/subgraph/fixtures/subgraphFixtures' + +import { usePriceBadge } from '@/composables/node/usePriceBadge' + +vi.mock('@/stores/workspace/colorPaletteStore', () => ({ + useColorPaletteStore: () => ({ + completedActivePalette: { + light_theme: false, + colors: { litegraph_base: {} } + } + }) +})) + +const { updateSubgraphCredits } = usePriceBadge() + +const mockNode = new LGraphNode('mock node') +const mockIcon: Partial = { unicode: '\ue96b' } +const badge: Partial = { + icon: mockIcon as LGraphIcon, + text: '$0.05/Run' +} +mockNode.badges = [badge as LGraphBadge] + +function getBadgeText(node: LGraphNode): string { + const badge = node.badges[0] + return (typeof badge === 'function' ? badge() : badge).text +} + +describe('subgraph pricing', () => { + subgraphTest( + 'should not display badge for subgraphs without API nodes', + ({ subgraphWithNode }) => { + const { subgraphNode } = subgraphWithNode + updateSubgraphCredits(subgraphNode) + expect(subgraphNode.badges.length).toBe(0) + } + ) + subgraphTest( + 'should return the price of a single contained API node', + ({ subgraphWithNode }) => { + const { subgraphNode, subgraph } = subgraphWithNode + subgraph.add(mockNode) + updateSubgraphCredits(subgraphNode) + expect(subgraphNode.badges.length).toBe(1) + expect(getBadgeText(subgraphNode)).toBe('$0.05/Run') + } + ) + subgraphTest( + 'should return the number of api nodes if more than one exists', + ({ subgraphWithNode }) => { + const { subgraphNode, subgraph } = subgraphWithNode + for (let i = 0; i < 5; i++) subgraph.add(mockNode) + updateSubgraphCredits(subgraphNode) + expect(subgraphNode.badges.length).toBe(1) + expect(getBadgeText(subgraphNode)).toBe('Partner Nodes x 5') + } + ) +})