Add pricing badge when a subgraph contains partner nodes (#6354)

<img width="596" height="213" alt="image"
src="https://github.com/user-attachments/assets/174c5461-f638-42de-b3ad-0e108dee3983"
/>


![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)

Author: AustinMroz

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)
       }
     })
   })

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<string> = 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
+  }
+}

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()
       }
     },

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
 }
 

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 }
   }
 

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' }
     | {

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<LGraphIcon> = { unicode: '\ue96b' }
+const badge: Partial<LGraphBadge> = {
+  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')
+    }
+  )
+})