fix: harden remote price badge validation and race handling

Review follow-ups:
- Require the full wire contract {engine, depends_on{widgets,inputs,
  input_groups}, expr}; validate per node so one malformed badge drops
  only that node, not the whole map
- Re-check raceLostForSession after awaiting the race so a concurrent
  overlay that timed out keeps the all-or-nothing session invariant
- Honor widgetType overrides and canonical combo spec parsing when
  bucketing declared input types; treat malformed (unvalidated
  /object_info) specs as a per-node skip instead of throwing
- Trim trailing slashes from the API base URL

src/services/priceBadgeService.test.ts

@@ -2,6 +2,8 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
 import type { ComfyNodeDef } from '@/schemas/nodeDefSchema'
 
+const mockApiBase = vi.hoisted(() => ({ url: 'https://api.example.com' }))
+
 const mockSystemStatsStore = vi.hoisted(() => ({
   isInitialized: true,
   error: null as Error | null,
@@ -16,7 +18,7 @@ vi.mock('@sentry/vue', () => ({
 }))
 
 vi.mock('@/config/comfyApi', () => ({
-  getComfyApiBaseUrl: () => 'https://api.example.com'
+  getComfyApiBaseUrl: () => mockApiBase.url
 }))
 
 vi.mock('@/platform/distribution/types', () => ({
@@ -80,6 +82,7 @@ async function importService() {
 describe('priceBadgeService', () => {
   beforeEach(() => {
     vi.resetModules()
+    mockApiBase.url = 'https://api.example.com'
     mockSystemStatsStore.isInitialized = true
     mockSystemStatsStore.error = null
     mockSystemStatsStore.systemStats = {
@@ -113,6 +116,17 @@ describe('priceBadgeService', () => {
     )
   })
 
+  it('normalizes a trailing slash on the API base URL', async () => {
+    mockApiBase.url = 'https://api.example.com/'
+    const fetchMock = mockFetchResponse({})
+    vi.stubGlobal('fetch', fetchMock)
+    const { overlayPriceBadges } = await importService()
+    await overlayPriceBadges(makeDefs())
+    expect(fetchMock).toHaveBeenCalledWith(
+      'https://api.example.com/nodes/pricing/badges?comfyui_version=0.3.50&platform=local'
+    )
+  })
+
   it('falls back to nightly when the version is unavailable', async () => {
     mockSystemStatsStore.systemStats = null
     const fetchMock = mockFetchResponse({})
@@ -124,6 +138,156 @@ describe('priceBadgeService', () => {
     )
   })
 
+  it('preserves an existing baked-in badge on nodes absent from the map', async () => {
+    vi.stubGlobal('fetch', mockFetchResponse({ PartnerNode: validBadge }))
+    const { overlayPriceBadges } = await importService()
+    const defs = makeDefs()
+    const bakedIn = {
+      engine: 'jsonata',
+      depends_on: { widgets: [], inputs: [], input_groups: [] },
+      expr: '{"type":"usd","usd":9.99}'
+    }
+    defs['OtherNode'].price_badge =
+      bakedIn as unknown as (typeof defs)['OtherNode']['price_badge']
+    await overlayPriceBadges(defs)
+    // Absent from the map: untouched, same object identity.
+    expect(defs['OtherNode'].price_badge).toBe(bakedIn)
+    // Present in the map: replaced.
+    expect(defs['PartnerNode'].price_badge).toMatchObject({
+      expr: validBadge.expr
+    })
+  })
+
+  it('skips only the malformed badge, applying valid siblings', async () => {
+    vi.stubGlobal(
+      'fetch',
+      mockFetchResponse({
+        PartnerNode: validBadge,
+        OtherNode: { expr: 'missing engine and depends_on' }
+      })
+    )
+    const { overlayPriceBadges } = await importService()
+    const defs = makeDefs()
+    await overlayPriceBadges(defs)
+    expect(defs['PartnerNode'].price_badge).toBeDefined()
+    expect(defs['OtherNode'].price_badge).toBeUndefined()
+  })
+
+  it.for([
+    ['engine', () => ({ ...validBadge, engine: undefined })],
+    ['depends_on', () => ({ ...validBadge, depends_on: undefined })],
+    ['expr', () => ({ ...validBadge, expr: undefined })],
+    [
+      'depends_on.widgets',
+      () => ({
+        ...validBadge,
+        depends_on: { inputs: [], input_groups: [] }
+      })
+    ],
+    [
+      'depends_on.inputs',
+      () => ({
+        ...validBadge,
+        depends_on: { widgets: validBadge.depends_on.widgets, input_groups: [] }
+      })
+    ],
+    [
+      'depends_on.input_groups',
+      () => ({
+        ...validBadge,
+        depends_on: { widgets: validBadge.depends_on.widgets, inputs: [] }
+      })
+    ]
+  ] as const)(
+    'rejects a badge missing required field %s, applying valid siblings',
+    async ([, makeBadge]) => {
+      vi.stubGlobal(
+        'fetch',
+        mockFetchResponse({
+          PartnerNode: makeBadge(),
+          OtherNode: {
+            ...validBadge,
+            depends_on: { widgets: [], inputs: [], input_groups: [] }
+          }
+        })
+      )
+      const { overlayPriceBadges } = await importService()
+      const defs = makeDefs()
+      await overlayPriceBadges(defs)
+      expect(defs['PartnerNode'].price_badge).toBeUndefined()
+      expect(defs['OtherNode'].price_badge).toBeDefined()
+    }
+  )
+
+  it.for([
+    ['null', null],
+    ['array', [validBadge]],
+    ['string', 'nope']
+  ] as const)('fails open when the response body is %s', async ([, body]) => {
+    vi.stubGlobal('fetch', mockFetchResponse(body))
+    const { overlayPriceBadges } = await importService()
+    const defs = makeDefs()
+    await expect(overlayPriceBadges(defs)).resolves.toBeUndefined()
+    expect(defs['PartnerNode'].price_badge).toBeUndefined()
+  })
+
+  it('skips a badge whose def input spec is malformed, without throwing', async () => {
+    const badge = {
+      ...validBadge,
+      depends_on: {
+        widgets: [{ name: 'broken', type: 'INT' }],
+        inputs: [],
+        input_groups: []
+      }
+    }
+    vi.stubGlobal('fetch', mockFetchResponse({ PartnerNode: badge }))
+    const { overlayPriceBadges } = await importService()
+    const defs = makeDefs()
+    // /object_info is unvalidated; a non-string, non-array tuple head must not
+    // crash the overlay (fail open for this node only).
+    ;(
+      defs['PartnerNode'] as unknown as {
+        input: { required: Record<string, unknown> }
+      }
+    ).input.required['broken'] = [42, {}]
+    await expect(overlayPriceBadges(defs)).resolves.toBeUndefined()
+    expect(defs['PartnerNode'].price_badge).toBeUndefined()
+  })
+
+  it('skips a badge with an unknown engine', async () => {
+    vi.stubGlobal(
+      'fetch',
+      mockFetchResponse({ PartnerNode: { ...validBadge, engine: 'python' } })
+    )
+    const { overlayPriceBadges } = await importService()
+    const defs = makeDefs()
+    await overlayPriceBadges(defs)
+    expect(defs['PartnerNode'].price_badge).toBeUndefined()
+  })
+
+  it('honors a widgetType override on the def input spec', async () => {
+    const badge = {
+      ...validBadge,
+      depends_on: {
+        widgets: [{ name: 'strength', type: 'STRING' }],
+        inputs: [],
+        input_groups: []
+      }
+    }
+    vi.stubGlobal('fetch', mockFetchResponse({ PartnerNode: badge }))
+    const { overlayPriceBadges } = await importService()
+    const defs = makeDefs()
+    // Socket type INT, but the runtime widget (and core's serialized badge
+    // dependency type) is the widgetType override.
+    ;(
+      defs['PartnerNode'] as unknown as {
+        input: { required: Record<string, unknown> }
+      }
+    ).input.required['strength'] = ['INT', { widgetType: 'STRING' }]
+    await overlayPriceBadges(defs)
+    expect(defs['PartnerNode'].price_badge).toBeDefined()
+  })
+
   it('ignores badges for nodes missing from the defs', async () => {
     vi.stubGlobal('fetch', mockFetchResponse({ UnknownNode: validBadge }))
     const { overlayPriceBadges } = await importService()
@@ -236,6 +400,39 @@ describe('priceBadgeService', () => {
     expect(defs2['PartnerNode'].price_badge).toBeUndefined()
   })
 
+  it('never applies badges on a concurrent overlay once the session lost the race', async () => {
+    vi.useFakeTimers()
+    let resolveFetch!: (value: unknown) => void
+    vi.stubGlobal(
+      'fetch',
+      vi.fn().mockReturnValue(
+        new Promise((resolve) => {
+          resolveFetch = resolve
+        })
+      )
+    )
+    const { overlayPriceBadges } = await importService()
+    const defs1 = makeDefs()
+    const defs2 = makeDefs()
+    const overlay1 = overlayPriceBadges(defs1)
+    await vi.advanceTimersByTimeAsync(1000)
+    const overlay2 = overlayPriceBadges(defs2)
+    // overlay1's timer fires at 2500ms and marks the session lost;
+    // overlay2's timer would fire at 3500ms.
+    await vi.advanceTimersByTimeAsync(1600)
+    await overlay1
+    // The fetch resolves after the session is lost but before overlay2's
+    // timeout: overlay2 must still not apply.
+    resolveFetch({
+      ok: true,
+      status: 200,
+      json: () => Promise.resolve({ PartnerNode: validBadge })
+    })
+    await overlay2
+    expect(defs1['PartnerNode'].price_badge).toBeUndefined()
+    expect(defs2['PartnerNode'].price_badge).toBeUndefined()
+  })
+
   it('reuses the same fetch across overlay calls', async () => {
     const fetchMock = mockFetchResponse({ PartnerNode: validBadge })
     vi.stubGlobal('fetch', fetchMock)

src/services/priceBadgeService.ts

@@ -4,7 +4,12 @@ import { z } from 'zod'
 
 import { getComfyApiBaseUrl } from '@/config/comfyApi'
 import { isCloud } from '@/platform/distribution/types'
-import type { ComfyNodeDef, PriceBadge } from '@/schemas/nodeDefSchema'
+import type {
+  ComfyNodeDef,
+  InputSpec,
+  PriceBadge
+} from '@/schemas/nodeDefSchema'
+import { getInputSpecType } from '@/schemas/nodeDefSchema'
 import { useSystemStatsStore } from '@/stores/systemStatsStore'
 
 const zWidgetDependency = z.object({
@@ -13,20 +18,20 @@ const zWidgetDependency = z.object({
 })
 
 const zPriceBadgeDepends = z.object({
-  widgets: z.array(zWidgetDependency).optional().default([]),
-  inputs: z.array(z.string()).optional().default([]),
-  input_groups: z.array(z.string()).optional().default([])
+  widgets: z.array(zWidgetDependency),
+  inputs: z.array(z.string()),
+  input_groups: z.array(z.string())
 })
 
+// Stricter than the canonical zPriceBadge: the wire contract always carries
+// the full shape, so nothing is defaulted and unknown engines are rejected.
 const zRemotePriceBadge = z.object({
-  engine: z.literal('jsonata').optional().default('jsonata'),
-  depends_on: zPriceBadgeDepends
-    .optional()
-    .default({ widgets: [], inputs: [], input_groups: [] }),
-  expr: z.string()
+  engine: z.literal('jsonata'),
+  depends_on: zPriceBadgeDepends,
+  expr: z.string().min(1)
 })
 
-const zPriceBadgeMap = z.record(zRemotePriceBadge)
+const zPriceBadgeMap = z.record(z.unknown())
 
 type PriceBadgeMap = Record<string, PriceBadge>
 
@@ -46,12 +51,12 @@ async function resolveComfyUIVersion(): Promise<string> {
 async function fetchPriceBadgeMap(): Promise<PriceBadgeMap | null> {
   try {
     const version = await resolveComfyUIVersion()
-    const url = `${getComfyApiBaseUrl()}/nodes/pricing/badges`
+    const base = getComfyApiBaseUrl().replace(/\/+$/, '')
     const params = new URLSearchParams({
       comfyui_version: version,
       platform: isCloud ? 'cloud' : 'local'
     })
-    const response = await fetch(`${url}?${params}`)
+    const response = await fetch(`${base}/nodes/pricing/badges?${params}`)
     if (!response.ok) {
       throw new Error(`HTTP ${response.status}`)
     }
@@ -59,7 +64,19 @@ async function fetchPriceBadgeMap(): Promise<PriceBadgeMap | null> {
     if (!parsed.success) {
       throw new Error(`invalid response: ${parsed.error.message}`)
     }
-    return parsed.data
+    const map: PriceBadgeMap = {}
+    for (const [name, rawBadge] of Object.entries(parsed.data)) {
+      const badge = zRemotePriceBadge.safeParse(rawBadge)
+      if (!badge.success) {
+        console.warn(
+          `[pricing/badges] skipping malformed badge for '${name}':`,
+          badge.error.message
+        )
+        continue
+      }
+      map[name] = badge.data
+    }
+    return map
   } catch (error) {
     console.warn(
       '[pricing/badges] fetch failed; no core-node badges this session:',
@@ -94,9 +111,26 @@ function bucketOfType(type: string): WidgetTypeBucket {
   return 'string'
 }
 
-function declaredInputBucket(spec: unknown): WidgetTypeBucket {
-  if (!Array.isArray(spec) || Array.isArray(spec[0])) return 'string'
-  return bucketOfType(String(spec[0]))
+/**
+ * Bucket the declared type of a def input spec, or null when the spec is
+ * malformed. /object_info arrives unvalidated, so the tuple head must be
+ * narrowed before it reaches the canonical parser.
+ */
+function declaredInputBucket(spec: unknown): WidgetTypeBucket | null {
+  if (!Array.isArray(spec)) return null
+  // A widgetType override is what the runtime widget is constructed as.
+  const options = spec[1]
+  const widgetType =
+    options !== null &&
+    typeof options === 'object' &&
+    'widgetType' in options &&
+    typeof options.widgetType === 'string'
+      ? options.widgetType
+      : undefined
+  if (widgetType) return bucketOfType(widgetType)
+  const head: unknown = spec[0]
+  if (typeof head !== 'string' && !Array.isArray(head)) return null
+  return bucketOfType(getInputSpecType(spec as InputSpec))
 }
 
 function validateBadgeAgainstDef(
@@ -109,8 +143,14 @@ function validateBadgeAgainstDef(
     const spec = inputs[parent]
     if (!spec) return `widget '${widget.name}' has no def input '${parent}'`
     const isNested = widget.name.includes('.')
-    if (!isNested && bucketOfType(widget.type) !== declaredInputBucket(spec)) {
-      return `widget '${widget.name}' type '${widget.type}' does not match def input type`
+    if (!isNested) {
+      const declared = declaredInputBucket(spec)
+      if (declared === null) {
+        return `widget '${widget.name}' def input spec is malformed`
+      }
+      if (bucketOfType(widget.type) !== declared) {
+        return `widget '${widget.name}' type '${widget.type}' does not match def input type`
+      }
     }
   }
   for (const name of [
@@ -154,6 +194,9 @@ export async function overlayPriceBadges(
     })
     return
   }
+  // A concurrent overlay call may have timed out and marked the session lost
+  // while we were awaiting; honor the all-or-nothing session invariant.
+  if (raceLostForSession) return
   if (!result) return
   for (const [name, badge] of Object.entries(result)) {
     const def = defs[name]