feat(stage-ui): webgpu detect improved (#1681)

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

packages/stage-shared/src/webgpu/detect.test.ts

@@ -0,0 +1,335 @@
+import { check as gpuuCheck } from 'gpuu/webgpu'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+import {
+  detectWebGPU,
+  getCachedWebGPUCapabilities,
+  getEstimatedVRAMOverride,
+  resetWebGPUCache,
+  setEstimatedVRAMOverride,
+} from './detect'
+
+// Mock gpuu/webgpu before importing detect.ts
+vi.mock('gpuu/webgpu', () => ({
+  check: vi.fn(),
+  isWebGPUSupported: vi.fn(),
+}))
+
+const mockedCheck = vi.mocked(gpuuCheck)
+
+interface MockAdapterInfo {
+  vendor?: string
+  architecture?: string
+  device?: string
+  description?: string
+}
+
+/**
+ * Build a minimal mock GPUAdapter for tests.
+ * `maxBufferSize` drives the heuristic VRAM calculation.
+ * `info` optionally populates adapterInfo.
+ */
+function makeMockAdapter(options: {
+  maxBufferSize?: number
+  info?: MockAdapterInfo
+}): any {
+  const adapter: { limits: { maxBufferSize: number }, info?: MockAdapterInfo } = {
+    limits: { maxBufferSize: options.maxBufferSize ?? 0 },
+  }
+  if (options.info)
+    adapter.info = options.info
+  return adapter
+}
+
+describe('detectWebGPU', () => {
+  beforeEach(() => {
+    resetWebGPUCache()
+    setEstimatedVRAMOverride(null)
+  })
+
+  afterEach(() => {
+    resetWebGPUCache()
+    setEstimatedVRAMOverride(null)
+    vi.clearAllMocks()
+  })
+
+  it('should derive VRAM from the maxBufferSize * 4 heuristic', async () => {
+    mockedCheck.mockResolvedValue({
+      supported: true,
+      fp16Supported: true,
+      isNode: false,
+      reason: '',
+      adapter: makeMockAdapter({ maxBufferSize: 256 * 1024 * 1024 }),
+    })
+
+    const result = await detectWebGPU()
+
+    expect(result.supported).toBe(true)
+    expect(result.fp16Supported).toBe(true)
+    expect(result.estimatedVRAM).toBe(256 * 1024 * 1024 * 4)
+    expect(result.estimatedVRAMSource).toBe('max-buffer-heuristic')
+  })
+
+  it('should report "none" when the adapter has no maxBufferSize', async () => {
+    mockedCheck.mockResolvedValue({
+      supported: true,
+      fp16Supported: false,
+      isNode: false,
+      reason: '',
+      adapter: makeMockAdapter({ maxBufferSize: 0 }),
+    })
+
+    const result = await detectWebGPU()
+
+    expect(result.estimatedVRAM).toBe(0)
+    expect(result.estimatedVRAMSource).toBe('none')
+  })
+
+  it('should report "none" when WebGPU is unsupported', async () => {
+    mockedCheck.mockResolvedValue({
+      supported: false,
+      fp16Supported: false,
+      isNode: false,
+      reason: 'not available',
+    })
+
+    const result = await detectWebGPU()
+
+    expect(result.supported).toBe(false)
+    expect(result.estimatedVRAM).toBe(0)
+    expect(result.estimatedVRAMSource).toBe('none')
+    expect(result.reason).toBe('not available')
+  })
+
+  it('should extract adapter.info when available', async () => {
+    mockedCheck.mockResolvedValue({
+      supported: true,
+      fp16Supported: true,
+      isNode: false,
+      reason: '',
+      adapter: makeMockAdapter({
+        maxBufferSize: 1024 * 1024 * 1024,
+        info: {
+          vendor: 'apple',
+          architecture: 'apple-m1',
+          device: 'Apple M1',
+          description: 'Apple GPU',
+        },
+      }),
+    })
+
+    const result = await detectWebGPU()
+
+    expect(result.adapterInfo).toEqual({
+      vendor: 'apple',
+      architecture: 'apple-m1',
+      device: 'Apple M1',
+      description: 'Apple GPU',
+    })
+  })
+
+  it('should handle adapter.info with missing fields gracefully', async () => {
+    mockedCheck.mockResolvedValue({
+      supported: true,
+      fp16Supported: true,
+      isNode: false,
+      reason: '',
+      adapter: makeMockAdapter({
+        maxBufferSize: 1024 * 1024 * 1024,
+        info: { vendor: 'nvidia' }, // only vendor set
+      }),
+    })
+
+    const result = await detectWebGPU()
+
+    expect(result.adapterInfo).toEqual({
+      vendor: 'nvidia',
+      architecture: '',
+      device: '',
+      description: '',
+    })
+  })
+
+  it('should set adapterInfo to null when adapter.info is not exposed', async () => {
+    mockedCheck.mockResolvedValue({
+      supported: true,
+      fp16Supported: true,
+      isNode: false,
+      reason: '',
+      adapter: makeMockAdapter({ maxBufferSize: 512 * 1024 * 1024 }),
+    })
+
+    const result = await detectWebGPU()
+
+    expect(result.adapterInfo).toBeNull()
+  })
+
+  it('should fall back to requestAdapterInfo() when adapter.info is absent', async () => {
+    const legacyInfo: MockAdapterInfo = {
+      vendor: 'intel',
+      architecture: 'xe',
+      device: 'Iris Xe',
+      description: 'Intel Xe Graphics',
+    }
+
+    const legacyAdapter: any = {
+      limits: { maxBufferSize: 256 * 1024 * 1024 },
+      requestAdapterInfo: vi.fn(async () => legacyInfo),
+    }
+
+    mockedCheck.mockResolvedValue({
+      supported: true,
+      fp16Supported: false,
+      isNode: false,
+      reason: '',
+      adapter: legacyAdapter,
+    })
+
+    const result = await detectWebGPU()
+
+    expect(result.adapterInfo).toEqual({
+      vendor: 'intel',
+      architecture: 'xe',
+      device: 'Iris Xe',
+      description: 'Intel Xe Graphics',
+    })
+  })
+
+  it('should cache the detection result across calls', async () => {
+    mockedCheck.mockResolvedValue({
+      supported: true,
+      fp16Supported: true,
+      isNode: false,
+      reason: '',
+      adapter: makeMockAdapter({ maxBufferSize: 1024 * 1024 }),
+    })
+
+    const first = await detectWebGPU()
+    const second = await detectWebGPU()
+
+    expect(first).toBe(second)
+    expect(mockedCheck).toHaveBeenCalledTimes(1)
+  })
+
+  it('should deduplicate concurrent calls', async () => {
+    mockedCheck.mockResolvedValue({
+      supported: true,
+      fp16Supported: true,
+      isNode: false,
+      reason: '',
+      adapter: makeMockAdapter({ maxBufferSize: 1024 * 1024 }),
+    })
+
+    const [a, b] = await Promise.all([detectWebGPU(), detectWebGPU()])
+
+    expect(a).toBe(b)
+    expect(mockedCheck).toHaveBeenCalledTimes(1)
+  })
+
+  it('should produce a safe fallback when gpuu throws', async () => {
+    mockedCheck.mockRejectedValue(new Error('internal'))
+
+    const result = await detectWebGPU()
+
+    expect(result.supported).toBe(false)
+    expect(result.estimatedVRAM).toBe(0)
+    expect(result.estimatedVRAMSource).toBe('none')
+    expect(result.adapterInfo).toBeNull()
+    expect(result.reason).toBe('Detection threw an exception')
+  })
+})
+
+describe('vRAM override', () => {
+  beforeEach(() => {
+    resetWebGPUCache()
+    setEstimatedVRAMOverride(null)
+  })
+
+  afterEach(() => {
+    resetWebGPUCache()
+    setEstimatedVRAMOverride(null)
+    vi.clearAllMocks()
+  })
+
+  it('should apply the override when detection runs', async () => {
+    mockedCheck.mockResolvedValue({
+      supported: true,
+      fp16Supported: true,
+      isNode: false,
+      reason: '',
+      adapter: makeMockAdapter({ maxBufferSize: 128 * 1024 * 1024 }),
+    })
+
+    setEstimatedVRAMOverride(8 * 1024 * 1024 * 1024) // 8 GB
+
+    const result = await detectWebGPU()
+
+    expect(result.estimatedVRAM).toBe(8 * 1024 * 1024 * 1024)
+    expect(result.estimatedVRAMSource).toBe('override')
+  })
+
+  it('should update cached result in-place when override is set after detection', async () => {
+    mockedCheck.mockResolvedValue({
+      supported: true,
+      fp16Supported: true,
+      isNode: false,
+      reason: '',
+      adapter: makeMockAdapter({ maxBufferSize: 128 * 1024 * 1024 }),
+    })
+
+    await detectWebGPU()
+    expect(getCachedWebGPUCapabilities()?.estimatedVRAMSource).toBe('max-buffer-heuristic')
+
+    setEstimatedVRAMOverride(4 * 1024 * 1024 * 1024)
+
+    expect(getCachedWebGPUCapabilities()?.estimatedVRAM).toBe(4 * 1024 * 1024 * 1024)
+    expect(getCachedWebGPUCapabilities()?.estimatedVRAMSource).toBe('override')
+  })
+
+  it('should revert to heuristic when override is cleared with null', async () => {
+    mockedCheck.mockResolvedValue({
+      supported: true,
+      fp16Supported: true,
+      isNode: false,
+      reason: '',
+      adapter: makeMockAdapter({ maxBufferSize: 128 * 1024 * 1024 }),
+    })
+
+    await detectWebGPU()
+    setEstimatedVRAMOverride(4 * 1024 * 1024 * 1024)
+    expect(getCachedWebGPUCapabilities()?.estimatedVRAMSource).toBe('override')
+
+    setEstimatedVRAMOverride(null)
+
+    expect(getCachedWebGPUCapabilities()?.estimatedVRAM).toBe(128 * 1024 * 1024 * 4)
+    expect(getCachedWebGPUCapabilities()?.estimatedVRAMSource).toBe('max-buffer-heuristic')
+  })
+
+  it('should expose the current override via getEstimatedVRAMOverride()', () => {
+    expect(getEstimatedVRAMOverride()).toBeNull()
+    setEstimatedVRAMOverride(2 * 1024 * 1024 * 1024)
+    expect(getEstimatedVRAMOverride()).toBe(2 * 1024 * 1024 * 1024)
+  })
+
+  it('should reject invalid override values', () => {
+    expect(() => setEstimatedVRAMOverride(-1)).toThrow()
+    expect(() => setEstimatedVRAMOverride(Number.NaN)).toThrow()
+    expect(() => setEstimatedVRAMOverride(Number.POSITIVE_INFINITY)).toThrow()
+  })
+
+  it('should accept zero as a no-op override (reverts to heuristic)', async () => {
+    mockedCheck.mockResolvedValue({
+      supported: true,
+      fp16Supported: true,
+      isNode: false,
+      reason: '',
+      adapter: makeMockAdapter({ maxBufferSize: 128 * 1024 * 1024 }),
+    })
+
+    setEstimatedVRAMOverride(0)
+    const result = await detectWebGPU()
+
+    // Zero is considered "no override" for the purpose of estimation
+    expect(result.estimatedVRAMSource).toBe('max-buffer-heuristic')
+  })
+})

packages/stage-shared/src/webgpu/detect.ts

@@ -3,17 +3,73 @@
  *
  * Wraps `gpuu/webgpu` and caches the result so every consumer
  * gets the same answer without redundant adapter requests.
+ *
+ * ## VRAM estimation
+ *
+ * The web platform does not expose GPU memory usage or total VRAM.
+ * We approximate via three ordered sources:
+ *
+ *   1. User override (`setEstimatedVRAMOverride(bytes)`)
+ *   2. `adapter.limits.maxBufferSize * 4` heuristic (fallback)
+ *   3. Zero (unavailable)
+ *
+ * The provenance is reported via `WebGPUCapabilities.estimatedVRAMSource`
+ * so consumers can surface it for diagnostics.
  */
 
 import { check as gpuuCheck, isWebGPUSupported as gpuuIsSupported } from 'gpuu/webgpu'
 
+// Minimal structural subset of the WebGPU types we interact with.
+// Avoids depending on `@webgpu/types` (which is shipped transitively via
+// transformers.js but not declared by this package).
+interface GPUAdapterInfoLike {
+  vendor?: string
+  architecture?: string
+  device?: string
+  description?: string
+}
+
+interface GPUAdapterLike {
+  limits?: { maxBufferSize?: number }
+  info?: GPUAdapterInfoLike
+  requestAdapterInfo?: () => Promise<GPUAdapterInfoLike>
+}
+
+/**
+ * Subset of `GPUAdapterInfo` that we surface to consumers. Values come
+ * directly from the browser's WebGPU implementation — treat them as
+ * opaque strings; vendor/architecture naming is not standardized.
+ */
+export interface WebGPUAdapterInfo {
+  /** Vendor name, e.g. "nvidia", "apple", "intel" */
+  vendor: string
+  /** Architecture name, e.g. "ada-lovelace", "apple-m1" */
+  architecture: string
+  /** Device description, e.g. "NVIDIA GeForce RTX 4090" */
+  device: string
+  /** Free-form description string from the driver */
+  description: string
+}
+
+/**
+ * Source of the VRAM estimate, reported for observability.
+ * - `override`: user-provided value via `setEstimatedVRAMOverride()`
+ * - `max-buffer-heuristic`: derived from `adapter.limits.maxBufferSize * 4`
+ * - `none`: no estimate available (WebGPU unsupported or adapter query failed)
+ */
+export type VRAMSource = 'override' | 'max-buffer-heuristic' | 'none'
+
 export interface WebGPUCapabilities {
   /** Whether WebGPU is available in this environment */
   supported: boolean
   /** Whether fp16 shader operations are supported */
   fp16Supported: boolean
-  /** Estimated VRAM in bytes (heuristic, 0 when unavailable) */
+  /** Estimated VRAM in bytes (0 when unavailable) */
   estimatedVRAM: number
+  /** Provenance of the VRAM estimate */
+  estimatedVRAMSource: VRAMSource
+  /** Adapter-reported vendor/architecture/device info, when available */
+  adapterInfo: WebGPUAdapterInfo | null
   /** Raw reason string from gpuu when unsupported */
   reason: string
 }
@@ -21,6 +77,71 @@ export interface WebGPUCapabilities {
 let cachedResult: WebGPUCapabilities | null = null
 let pendingDetection: Promise<WebGPUCapabilities> | null = null
 
+// NOTICE: User override for VRAM estimation. When set, this value takes
+// priority over all heuristics. Useful for users with known hardware where
+// the heuristic is inaccurate (e.g. discrete GPUs with small maxBufferSize).
+let vramOverride: number | null = null
+
+// Cached heuristic value so we can restore it when the override is cleared.
+// Computed during detectWebGPU() and persisted for the lifetime of the cache.
+let cachedHeuristicVRAM = 0
+
+/**
+ * Best-effort extraction of `GPUAdapterInfo` from a `GPUAdapter`. Tries
+ * the modern synchronous `adapter.info` first, then falls back to the
+ * legacy `requestAdapterInfo()` promise API. Returns null if neither works.
+ *
+ * References:
+ *   - https://www.w3.org/TR/webgpu/#gpu-adapterinfo
+ */
+async function extractAdapterInfo(adapter: GPUAdapterLike): Promise<WebGPUAdapterInfo | null> {
+  try {
+    // Modern API: synchronous `info` property (Chrome 114+, Safari 17.4+)
+    const info = adapter.info
+    if (info) {
+      return {
+        vendor: info.vendor ?? '',
+        architecture: info.architecture ?? '',
+        device: info.device ?? '',
+        description: info.description ?? '',
+      }
+    }
+
+    // Legacy API: requestAdapterInfo() returns a Promise
+    const legacy = adapter.requestAdapterInfo
+    if (typeof legacy === 'function') {
+      const legacyInfo = await legacy.call(adapter)
+      return {
+        vendor: legacyInfo.vendor ?? '',
+        architecture: legacyInfo.architecture ?? '',
+        device: legacyInfo.device ?? '',
+        description: legacyInfo.description ?? '',
+      }
+    }
+  }
+  catch {
+    // Fall through to null — adapter info is best-effort, not required
+  }
+  return null
+}
+
+/** Compute the heuristic VRAM estimate from `maxBufferSize`. */
+function computeHeuristicVRAM(adapter: GPUAdapterLike): number {
+  const maxBuffer = adapter.limits?.maxBufferSize ?? 0
+  // Typical values: 256 MB on integrated GPUs, 2-4 GB on discrete.
+  // Multiply by 4 as a conservative total VRAM heuristic.
+  return maxBuffer > 0 ? maxBuffer * 4 : 0
+}
+
+/** Decide the VRAM estimate based on override > heuristic > none. */
+function resolveVRAM(heuristic: number): { bytes: number, source: VRAMSource } {
+  if (vramOverride !== null && vramOverride > 0)
+    return { bytes: vramOverride, source: 'override' }
+  if (heuristic > 0)
+    return { bytes: heuristic, source: 'max-buffer-heuristic' }
+  return { bytes: 0, source: 'none' }
+}
+
 /**
  * Detect WebGPU capabilities. The result is cached as a singleton
  * after the first successful call -- safe to call repeatedly.
@@ -37,27 +158,33 @@ export async function detectWebGPU(): Promise<WebGPUCapabilities> {
     try {
       const result = await gpuuCheck()
 
-      let estimatedVRAM = 0
+      let adapterInfo: WebGPUAdapterInfo | null = null
+      let heuristic = 0
       if (result.supported && result.adapter) {
-        // Use maxBufferSize as a rough proxy -- typically 256 MB on
-        // integrated GPUs, 2-4 GB on discrete GPUs.
-        // Multiply by 4 as a conservative total VRAM heuristic.
-        const maxBuffer = result.adapter.limits?.maxBufferSize ?? 0
-        estimatedVRAM = maxBuffer > 0 ? maxBuffer * 4 : 0
+        heuristic = computeHeuristicVRAM(result.adapter)
+        adapterInfo = await extractAdapterInfo(result.adapter)
       }
 
+      cachedHeuristicVRAM = heuristic
+      const vram = resolveVRAM(heuristic)
+
       cachedResult = {
         supported: result.supported,
         fp16Supported: result.fp16Supported ?? false,
-        estimatedVRAM,
+        estimatedVRAM: vram.bytes,
+        estimatedVRAMSource: vram.source,
+        adapterInfo,
         reason: result.reason ?? '',
       }
     }
     catch {
+      cachedHeuristicVRAM = 0
       cachedResult = {
         supported: false,
         fp16Supported: false,
         estimatedVRAM: 0,
+        estimatedVRAMSource: 'none',
+        adapterInfo: null,
         reason: 'Detection threw an exception',
       }
     }
@@ -86,10 +213,44 @@ export async function isWebGPUSupported(): Promise<boolean> {
   return gpuuIsSupported()
 }
 
+/**
+ * Override the estimated VRAM value. Pass `null` to clear the override and
+ * revert to the heuristic. The override applies to future detections, and if
+ * a result is already cached its VRAM fields are updated immediately, so
+ * `resetWebGPUCache()` is not required.
+ *
+ * Intended for user preference UI ("I have 8 GB VRAM") and testing.
+ */
+export function setEstimatedVRAMOverride(bytes: number | null): void {
+  if (bytes !== null && (!Number.isFinite(bytes) || bytes < 0))
+    throw new Error(`Invalid VRAM override: ${bytes} (expected null or non-negative finite number)`)
+
+  vramOverride = bytes
+
+  // If we already have a cached result, update it in-place so consumers
+  // see the new value without needing to call resetWebGPUCache(). The
+  // original heuristic value is preserved in `cachedHeuristicVRAM` so we
+  // can revert when the override is cleared.
+  if (cachedResult) {
+    const vram = resolveVRAM(cachedHeuristicVRAM)
+    cachedResult = {
+      ...cachedResult,
+      estimatedVRAM: vram.bytes,
+      estimatedVRAMSource: vram.source,
+    }
+  }
+}
+
+/** Read the current VRAM override, or null if unset. */
+export function getEstimatedVRAMOverride(): number | null {
+  return vramOverride
+}
+
 /**
  * Reset the cached detection result. Intended for tests only.
  */
 export function resetWebGPUCache(): void {
   cachedResult = null
   pendingDetection = null
+  cachedHeuristicVRAM = 0
 }

packages/stage-shared/src/webgpu/index.ts

@@ -1,7 +1,9 @@
 export {
   detectWebGPU,
   getCachedWebGPUCapabilities,
+  getEstimatedVRAMOverride,
   isWebGPUSupported,
   resetWebGPUCache,
+  setEstimatedVRAMOverride,
 } from './detect'
-export type { WebGPUCapabilities } from './detect'
+export type { VRAMSource, WebGPUAdapterInfo, WebGPUCapabilities } from './detect'