From ae69bce729948cfeb9ab7d7a401afc4340801ac1 Mon Sep 17 00:00:00 2001 From: Alexander Piskun Date: Fri, 17 Jul 2026 13:25:09 +0300 Subject: [PATCH] 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 | 199 ++++++++++++++++++++++++- src/services/priceBadgeService.ts | 79 +++++++--- 2 files changed, 259 insertions(+), 19 deletions(-) diff --git a/src/services/priceBadgeService.test.ts b/src/services/priceBadgeService.test.ts index 0e31a1fa0b..3d16405220 100644 --- a/src/services/priceBadgeService.test.ts +++ b/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 } + } + ).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 } + } + ).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) diff --git a/src/services/priceBadgeService.ts b/src/services/priceBadgeService.ts index 0826a94168..55ae0de3ff 100644 --- a/src/services/priceBadgeService.ts +++ b/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 @@ -46,12 +51,12 @@ async function resolveComfyUIVersion(): Promise { async function fetchPriceBadgeMap(): Promise { 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 { 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]