Partner node tracking cleanup

src/composables/node/usePriceBadge.ts

@@ -1,8 +1,10 @@
-import { computed, toValue } from 'vue'
+import { createSharedComposable } from '@vueuse/core'
+import { computed, ref, toValue } from 'vue'
 
 import type { LGraph, LGraphNode } from '@/lib/litegraph/src/litegraph'
 import { LGraphBadge } from '@/lib/litegraph/src/litegraph'
 
+import { useChainCallback } from '@/composables/functional/useChainCallback'
 import { useNodePricing } from '@/composables/node/useNodePricing'
 import type { INodeInputSlot } from '@/lib/litegraph/src/interfaces'
 import type { SubgraphInput } from '@/lib/litegraph/src/subgraph/SubgraphInput'
@@ -149,8 +151,26 @@ export const usePriceBadge = () => {
         : '#8D6932'
     })
   }
-  const creditsBadges = computed(() =>
-    mapAllNodes(app.graph, (node) => {
+  return {
+    getCreditsBadge,
+    isCreditsBadge,
+    updateSubgraphCredits
+  }
+}
+export const useCreditsBadgesInGraph = createSharedComposable(() => {
+  const { isCreditsBadge } = usePriceBadge()
+  const nodeTrigger = ref(0)
+  app.graph.onNodeAdded = useChainCallback(
+    app.graph.onNodeAdded,
+    () => nodeTrigger.value++
+  )
+  app.graph.onNodeRemoved = useChainCallback(
+    app.graph.onNodeRemoved,
+    () => nodeTrigger.value++
+  )
+  return computed(() => {
+    void nodeTrigger.value
+    return mapAllNodes(app.graph, (node) => {
       if (node.isSubgraphNode()) return
 
       const priceBadge = node.badges.find(isCreditsBadge)
@@ -159,11 +179,5 @@ export const usePriceBadge = () => {
       trackNodePrice(node)
       return [node.title, toValue(priceBadge).text, node.id] as const
     })
-  )
-  return {
-    creditsBadges,
-    getCreditsBadge,
-    isCreditsBadge,
-    updateSubgraphCredits
-  }
-}
+  })
+})

src/platform/cloud/subscription/components/FreeTierQuota.vue

@@ -3,7 +3,6 @@ import { computed } from 'vue'
 import { useI18n } from 'vue-i18n'
 
 import { cn } from '@comfyorg/tailwind-utils'
-import { usePriceBadge } from '@/composables/node/usePriceBadge'
 import { useFreeTierQuota } from '@/platform/cloud/subscription/composables/useFreeTierQuota'
 
 const DOT_COLORS = [
@@ -13,8 +12,8 @@ const DOT_COLORS = [
 ]
 
 const { t } = useI18n()
-const { creditsBadges } = usePriceBadge()
-const { available, maxAvailable, quotaEnabled } = useFreeTierQuota()
+const { available, hasInvalidNodes, maxAvailable, quotaEnabled } =
+  useFreeTierQuota()
 
 const dotColor = computed(() => {
   const ratio = maxAvailable.value ? available.value / maxAvailable.value : 0
@@ -37,7 +36,7 @@ const label = computed(() =>
     class="mt-2 w-full border-t border-border-subtle bg-comfy-menu-bg px-4 pt-2"
   >
     <div
-      v-if="creditsBadges.length"
+      v-if="hasInvalidNodes"
       class="flex w-full items-center justify-center gap-2"
     >
       <i class="icon-[comfy--credits] bg-amber-400" />

src/platform/cloud/subscription/composables/useFreeTierQuota.ts

@@ -1,6 +1,7 @@
 import { createSharedComposable } from '@vueuse/core'
 import { computed, ref, watch } from 'vue'
 
+import { useCreditsBadgesInGraph } from '@/composables/node/usePriceBadge'
 import { useFeatureFlags } from '@/composables/useFeatureFlags'
 import { remoteConfig } from '@/platform/remoteConfig/remoteConfig'
 
@@ -12,6 +13,8 @@ type FreeTierBalance = {
 
 export const useFreeTierQuota = createSharedComposable(function () {
   const { flags } = useFeatureFlags()
+  const creditsBadges = useCreditsBadgesInGraph()
+
   const freeTierBalance = computed(() => {
     return (remoteConfig.value as { free_tier_balance?: FreeTierBalance })
       ?.free_tier_balance
@@ -31,10 +34,21 @@ export const useFreeTierQuota = createSharedComposable(function () {
   const quotaEnabled = computed(
     () => flags.freeTierJobAllowanceEnabled && maxAvailable.value > 0
   )
+  const hasInvalidNodes = computed(() => creditsBadges.value.length > 0)
+  const freeTierExecutionPermitted = computed(
+    () => !hasInvalidNodes.value && quotaEnabled.value && available.value > 0
+  )
 
   function trackRun() {
     if (available.value > 0) available.value--
   }
 
-  return { available, maxAvailable, quotaEnabled, trackRun }
+  return {
+    available,
+    freeTierExecutionPermitted,
+    hasInvalidNodes,
+    maxAvailable,
+    quotaEnabled,
+    trackRun
+  }
 })

src/renderer/extensions/linearMode/PartnerNodesList.vue

@@ -8,12 +8,12 @@ import { useI18n } from 'vue-i18n'
 
 import Button from '@/components/ui/button/Button.vue'
 import Popover from '@/components/ui/Popover.vue'
-import { usePriceBadge } from '@/composables/node/usePriceBadge'
+import { useCreditsBadgesInGraph } from '@/composables/node/usePriceBadge'
 import PartnerNodeItem from '@/renderer/extensions/linearMode/PartnerNodeItem.vue'
 
 defineProps<{ mobile?: boolean }>()
 
-const { creditsBadges } = usePriceBadge()
+const creditsBadges = useCreditsBadgesInGraph()
 const { t } = useI18n()
 </script>
 <template>