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)

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