From 57fdf637fba30d69500a58693d05b972b96671b4 Mon Sep 17 00:00:00 2001 From: felirami Date: Sun, 16 Aug 2026 19:08:28 -0400 Subject: [PATCH] Add WebGPU wedge watchdog with one-way WASM fallback Live smoke of v1.3.1 (M3 Pro under GPU contention): a WebGPU session.run stalled mid-page on Wikipedia and Amazon search, the offscreen document failed its 30 second init health check, and 104 badges stayed pending with no recovery path. src/webgpu-watchdog.js (new, pure, node:test-covered): - MODEL_RUN_WEDGE_TIMEOUT_MS 20s per single session.run, well under the 45s inference budget so the wedge verdict plus WASM retry fit inside it - isWedgeError / isWebGpuDeviceError / shouldFallBackToWasm decision logic, watchModelRun race, wrapSessionWithWatchdog session wrapper, createWebGpuHealth one-way latch src/offscreen.js: - WebGPU CF and DINO sessions are wrapped with the watchdog - GPUDevice.lost marks the health latch - On wedge or device loss: dispose both sessions, rebuild both on WASM, log the switch once, rescore the in-flight image, and let everything behind the inference lock and in the page queue run on WASM - One-way for the browser session; no flap back until the offscreen document restarts - An error after the WASM retry propagates so the badge gets its normal error state instead of staying pending tests/webgpu-watchdog.test.mjs: 20 new tests. Full suite 248 passing (was 228). No fusion, threshold, model, or probe changes. --- package.json | 2 +- src/offscreen.js | 195 +++++++++++++++++++++++++----- src/webgpu-watchdog.js | 145 ++++++++++++++++++++++ tests/webgpu-watchdog.test.mjs | 214 +++++++++++++++++++++++++++++++++ 4 files changed, 525 insertions(+), 31 deletions(-) create mode 100644 src/webgpu-watchdog.js create mode 100644 tests/webgpu-watchdog.test.mjs diff --git a/package.json b/package.json index 1e86d74..06a0aa5 100644 --- a/package.json +++ b/package.json @@ -12,7 +12,7 @@ "fetch-model": "node scripts/fetch-model.mjs", "icons": "node scripts/generate-icons.mjs", "build": "node scripts/generate-icons.mjs && node scripts/build.mjs", - "test": "node --test tests/benchmark-docs.test.mjs tests/pixel-resize.test.mjs tests/fuse.test.mjs tests/scoring.test.mjs tests/inference-policy.test.mjs tests/graphic-gate.test.mjs tests/eval-policy.test.mjs tests/calibration.test.mjs tests/heuristics.test.mjs tests/clip-preprocess.test.mjs tests/community-forensics.test.mjs tests/bytes.test.mjs tests/image-limits.test.mjs tests/analyze-retry.test.mjs tests/analyze-queue.test.mjs tests/store-listing.test.mjs tests/release-version.test.mjs tests/release-workflow.test.mjs tests/site-astro.test.mjs tests/site-worker.test.mjs tests/model-fetch.test.mjs tests/image-url.test.mjs tests/image-job.test.mjs tests/messages.test.mjs tests/overlay.test.mjs tests/update-checker.test.mjs tests/element-pixels.test.mjs tests/dino.test.mjs tests/walkthrough.test.mjs", + "test": "node --test tests/benchmark-docs.test.mjs tests/pixel-resize.test.mjs tests/fuse.test.mjs tests/scoring.test.mjs tests/inference-policy.test.mjs tests/graphic-gate.test.mjs tests/eval-policy.test.mjs tests/calibration.test.mjs tests/heuristics.test.mjs tests/clip-preprocess.test.mjs tests/community-forensics.test.mjs tests/bytes.test.mjs tests/image-limits.test.mjs tests/analyze-retry.test.mjs tests/analyze-queue.test.mjs tests/webgpu-watchdog.test.mjs tests/store-listing.test.mjs tests/release-version.test.mjs tests/release-workflow.test.mjs tests/site-astro.test.mjs tests/site-worker.test.mjs tests/model-fetch.test.mjs tests/image-url.test.mjs tests/image-job.test.mjs tests/messages.test.mjs tests/overlay.test.mjs tests/update-checker.test.mjs tests/element-pixels.test.mjs tests/dino.test.mjs tests/walkthrough.test.mjs", "eval": "node eval/harness.mjs", "package": "node scripts/package.mjs", "release:check-version": "node scripts/check-release-version.mjs", diff --git a/src/offscreen.js b/src/offscreen.js index 0652929..7ec5366 100644 --- a/src/offscreen.js +++ b/src/offscreen.js @@ -35,6 +35,13 @@ import { normalizeTtaMode, runExclusiveAfterStart, } from './analyze-queue.js'; +import { + createWebGpuHealth, + isWebGpuDeviceError, + isWedgeError, + shouldFallBackToWasm, + wrapSessionWithWatchdog, +} from './webgpu-watchdog.js'; const HF_FILES = [MODEL_ONNX_PATH, 'preprocessor_config.json', 'config.json']; const PROBE_URL_PATH = 'models/probe/dino-probe.json'; @@ -48,6 +55,15 @@ let sessionPromise = null; let threshold = DEFAULT_THRESHOLD; const inferLock = createExclusiveLock(); +// Runtime WebGPU recovery (live smoke of v1.3.1: a wedged WebGPU run under +// GPU contention left 104 badges pending). One-way for the lifetime of this +// offscreen document: once WebGPU wedges or the device is lost, both +// sessions are rebuilt on WASM and never flap back. +const webGpuHealth = createWebGpuHealth(); +let rawCfSession = null; +let rawDinoSession = null; +let wasmFallbackPromise = null; + /** * Multithreaded WASM needs SharedArrayBuffer, which needs the COOP/COEP * headers declared in manifest.json. Feature-detect so a Chrome that @@ -93,16 +109,39 @@ async function tryCreateDinoHead() { }) ); - dinoSession = created; + rawDinoSession = created; + dinoSession = sessionDevice === 'webgpu' ? wrapSessionWithWatchdog(created) : created; dinoProbe = probe; console.log('DINOv2 probe head ready'); } catch (err) { console.debug('DINOv2 head unavailable, running CommunityForensics only:', err?.message || err); + rawDinoSession = null; dinoSession = null; dinoProbe = null; } } +/** + * GPUDevice.lost resolves (it never rejects) when the device goes away. + * onnxruntime-web keeps one shared device on ort.env.webgpu after a WebGPU + * session exists. Marking the latch here means the next analyze rebuilds on + * WASM before it ever starts another WebGPU run. + */ +function watchGpuDeviceLoss() { + let lost; + try { + lost = ort.env?.webgpu?.device?.lost; + } catch { + return; + } + if (typeof lost?.then !== 'function') return; + lost + .then((info) => { + webGpuHealth.markUnhealthy(`GPUDevice lost (${info?.reason || 'unknown'})`); + }) + .catch(() => {}); +} + async function ensureSession() { if (session) return { session, device: sessionDevice }; if (sessionPromise) return sessionPromise; @@ -117,20 +156,23 @@ async function ensureSession() { } const modelUrl = chrome.runtime.getURL(`models/${MODEL_ID}/${MODEL_ONNX_PATH}`); - const adapter = await probeWebGpuAdapter(); + const adapter = webGpuHealth.isUnhealthy() ? null : await probeWebGpuAdapter(); const created = await createCommunityForensicsSession({ modelUrl, wasmPaths: chrome.runtime.getURL('lib/'), - preferWebGpu: Boolean(adapter), + preferWebGpu: Boolean(adapter) && !webGpuHealth.isUnhealthy(), numThreads: wasmThreadCount(), }); - session = created.session; + rawCfSession = created.session; sessionDevice = created.device; + session = + created.device === 'webgpu' ? wrapSessionWithWatchdog(created.session) : created.session; + if (created.device === 'webgpu') watchGpuDeviceLoss(); await tryCreateDinoHead(); - return created; + return { session, device: sessionDevice }; })(); try { @@ -142,6 +184,52 @@ async function ensureSession() { } } +/** + * One-way WebGPU to WASM fallback. Disposes both (possibly wedged) sessions + * without awaiting them, rebuilds both on the WASM execution provider, and + * logs the switch exactly once. The caller retries the in-flight image; + * everything still waiting behind the inference lock or in the page queue + * simply runs on the rebuilt WASM sessions. There is no path back to WebGPU + * until the offscreen document restarts. + * @param {string} reason + */ +function fallbackToWasm(reason) { + if (wasmFallbackPromise) return wasmFallbackPromise; + + webGpuHealth.markUnhealthy(reason); + console.warn( + `Clueside: WebGPU backend unhealthy (${webGpuHealth.reason()}). ` + + 'Rebuilding both model sessions on WASM for the rest of this browser session.' + ); + + wasmFallbackPromise = (async () => { + // Fire-and-forget disposal: release on a wedged device can itself hang. + try { + Promise.resolve(rawCfSession?.release?.()).catch(() => {}); + } catch { + // Session already unusable. + } + try { + Promise.resolve(rawDinoSession?.release?.()).catch(() => {}); + } catch { + // Session already unusable. + } + + session = null; + rawCfSession = null; + dinoSession = null; + rawDinoSession = null; + dinoProbe = null; + sessionDevice = 'wasm'; + sessionPromise = null; + initError = null; + + await ensureSession(); + })(); + + return wasmFallbackPromise; +} + /** * @param {ImageBitmap} bitmap * @returns {Promise} p(AI) from the DINOv2 probe, or null @@ -154,6 +242,11 @@ async function dinoScore(bitmap) { const outputs = await dinoSession.run({ [DINO_INPUT_NAME]: input }); return dinoScoreHiddenState(outputs.last_hidden_state, dinoProbe); } catch (err) { + if (sessionDevice === 'webgpu' && (isWedgeError(err) || isWebGpuDeviceError(err))) { + // Do not silently degrade to CF-only on a wedged backend: let the + // caller run the WASM fallback and rescore with both heads. + throw err; + } console.debug('DINO head inference failed:', err?.message || err); return null; } @@ -218,6 +311,67 @@ async function resolveImageBytes(message) { throw new Error('No image bytes received (missing bufferB64 and non-http URL)'); } +/** + * One neural pass (graphic gate + DINO + CF) on whatever backend is current. + * @param {ArrayBuffer} bytes + * @param {'adaptive'|'always'|'center'} mode + * @returns {Promise<{ cfPAi: number, dinoPAi: number|null, graphicGate: boolean }>} + */ +async function scoreImageOnce(bytes, mode) { + const { session: activeSession } = await ensureSession(); + const mime = sniffMime(bytes); + const blob = new Blob([bytes], { type: mime }); + const bitmap = await createImageBitmap(blob); + try { + const graphicGate = analyzeGraphicGate(bitmap).isGraphic; + const dinoPAi = await dinoScore(bitmap); + + let cfPAi; + if (mode === 'center') { + const chw = await preprocessBitmap(bitmap); + cfPAi = await predictCHW(activeSession, chw); + } else { + const views = await preprocessBitmapViews(bitmap); + const effectiveMode = effectiveTtaMode(mode, dinoPAi); + cfPAi = (await predictAdaptiveViews(activeSession, views, { mode: effectiveMode })).neuralPAi; + } + + return { cfPAi, dinoPAi, graphicGate }; + } finally { + bitmap.close(); + } +} + +/** + * Neural scoring with the runtime WebGPU escape hatch. If a run wedges past + * MODEL_RUN_WEDGE_TIMEOUT_MS or throws a device-level WebGPU error, fall back + * to WASM once and rescore this same image there. An error after the WASM + * retry propagates to the caller so the badge gets its normal error state + * instead of staying pending forever. + * @param {ArrayBuffer} bytes + * @param {'adaptive'|'always'|'center'} mode + */ +async function scoreImageWithFallback(bytes, mode) { + if (sessionDevice === 'webgpu' && webGpuHealth.isUnhealthy()) { + // Device loss reported between runs: rebuild before starting another + // WebGPU run. + await fallbackToWasm(webGpuHealth.reason()); + } + + try { + return await scoreImageOnce(bytes, mode); + } catch (err) { + const fallBack = shouldFallBackToWasm({ + device: sessionDevice, + error: err, + alreadyFellBack: Boolean(wasmFallbackPromise), + }); + if (!fallBack) throw err; + await fallbackToWasm(err?.message || String(err)); + return scoreImageOnce(bytes, mode); + } +} + async function classifyImage(rawBytes, url, customThreshold, ttaMode) { const bytes = toArrayBuffer(rawBytes); const heuristics = await analyzeHeuristics(bytes, url); @@ -228,32 +382,13 @@ async function classifyImage(rawBytes, url, customThreshold, ttaMode) { let modelError = null; try { - const { session: activeSession } = await ensureSession(); - const mime = sniffMime(bytes); - const blob = new Blob([bytes], { type: mime }); - const bitmap = await createImageBitmap(blob); - try { - const graphicGate = analyzeGraphicGate(bitmap).isGraphic; - const dinoPAi = await dinoScore(bitmap); - - let cfPAi; - if (mode === 'center') { - const chw = await preprocessBitmap(bitmap); - cfPAi = await predictCHW(activeSession, chw); - } else { - const views = await preprocessBitmapViews(bitmap); - const effectiveMode = effectiveTtaMode(mode, dinoPAi); - cfPAi = (await predictAdaptiveViews(activeSession, views, { mode: effectiveMode })).neuralPAi; - } + const { cfPAi, dinoPAi, graphicGate } = await scoreImageWithFallback(bytes, mode); - fused = fuseInferenceScores(cfPAi, dinoPAi, heuristics, activeThreshold, { - graphicGate, - }); - if (graphicGate && cfPAi < DEFAULT_THRESHOLD) { - fused.reasons.push('Flat graphic gate: DINO lift suppressed'); - } - } finally { - bitmap.close(); + fused = fuseInferenceScores(cfPAi, dinoPAi, heuristics, activeThreshold, { + graphicGate, + }); + if (graphicGate && cfPAi < DEFAULT_THRESHOLD) { + fused.reasons.push('Flat graphic gate: DINO lift suppressed'); } } catch (err) { modelError = err.message || String(err); diff --git a/src/webgpu-watchdog.js b/src/webgpu-watchdog.js new file mode 100644 index 0000000..cd01cad --- /dev/null +++ b/src/webgpu-watchdog.js @@ -0,0 +1,145 @@ +/** + * WebGPU wedge watchdog and one-way WASM fallback decisions. + * + * Live smoke of v1.3.1 (M3 Pro under GPU contention): a WebGPU session.run + * stalled mid-page on Wikipedia and Amazon, the offscreen document failed its + * 30 second init health check, and 104 badges stayed pending with no recovery. + * AGENTS.md already forbids latching a WebGPU error at init ("probe the + * adapter"), but a backend that hangs mid-run needs a runtime escape hatch. + * + * Everything here is pure so node:test can exercise the decision logic + * without a browser. The offscreen document owns the actual session disposal + * and recreation (src/offscreen.js). + */ + +/** + * A single model run (one ORT session.run call) that exceeds this is treated + * as a wedged backend. Generous on purpose: one CommunityForensics 384 pass + * finishes in well under 5 s on an M3 Pro even under GPU contention, and the + * DINO 224 pass is cheaper still. Anything past 20 s means the exclusive + * inference lock is starving the whole queue. Must stay well below + * INFERENCE_TIMEOUT_MS (45 s, analyze-retry.js) so the wedge verdict, the + * WASM session rebuild, and the retry of the in-flight image all fit inside + * one inference budget. + */ +export const MODEL_RUN_WEDGE_TIMEOUT_MS = 20_000; + +/** Marker prefix used to recognize watchdog timeouts in error text. */ +export const MODEL_RUN_WEDGE_MESSAGE = 'Model run wedged'; + +function errorText(error) { + return String(error?.message ?? error ?? ''); +} + +/** + * @param {unknown} error + * @returns {boolean} true when the error came from watchModelRun timing out + */ +export function isWedgeError(error) { + return errorText(error).includes(MODEL_RUN_WEDGE_MESSAGE); +} + +/** + * Device-level WebGPU failures (Dawn / Chrome / onnxruntime-web wording). + * Fetch, decode, and size-cap errors must never match: they have their own + * skip and error paths and must not tear down a healthy backend. + * @param {unknown} error + */ +export function isWebGpuDeviceError(error) { + const msg = errorText(error).toLowerCase(); + if (!msg) return false; + return ( + msg.includes('device lost') || + msg.includes('device is lost') || + msg.includes('device was lost') || + msg.includes('device destroyed') || + msg.includes('gpu device') || + msg.includes('gpudevice') || + msg.includes('gpu connection') || + msg.includes('gpu process') || + msg.includes('webgpu') + ); +} + +/** + * Decide whether an inference error must trigger the one-way WebGPU to WASM + * fallback. Only a wedge or a device-level error on an active WebGPU backend + * qualifies. WASM has nothing to fall back to, and a backend that already + * fell back must not flap. + * + * @param {{ device: 'webgpu' | 'wasm', error: unknown, alreadyFellBack?: boolean }} input + * @returns {boolean} + */ +export function shouldFallBackToWasm({ device, error, alreadyFellBack = false }) { + if (alreadyFellBack) return false; + if (device !== 'webgpu') return false; + return isWedgeError(error) || isWebGpuDeviceError(error); +} + +/** + * One-way health latch for the lifetime of the offscreen document. + * markUnhealthy returns true only on the first call so the caller can log + * the fallback exactly once. There is deliberately no way to reset it: + * WebGPU comes back only when the offscreen document restarts. + */ +export function createWebGpuHealth() { + let unhealthy = false; + let reason = null; + + return { + markUnhealthy(newReason) { + if (unhealthy) return false; + unhealthy = true; + reason = String(newReason || 'unknown'); + return true; + }, + isUnhealthy() { + return unhealthy; + }, + reason() { + return reason; + }, + }; +} + +/** + * Race one started model run against the wedge timeout. WebGPU has no abort: + * the underlying promise may never settle, so on timeout the caller must stop + * using the wedged session and dispose it instead of awaiting it again. A + * late settle of the abandoned run is swallowed so it cannot surface as an + * unhandled rejection. + * + * @param {Promise | (() => Promise)} work + * @param {{ timeoutMs?: number }} [options] + */ +export function watchModelRun(work, options = {}) { + const timeoutMs = options.timeoutMs ?? MODEL_RUN_WEDGE_TIMEOUT_MS; + const promise = + typeof work === 'function' ? Promise.resolve().then(work) : Promise.resolve(work); + promise.catch(() => {}); + + let timer; + return Promise.race([ + promise, + new Promise((_, reject) => { + timer = setTimeout(() => { + reject(new Error(`${MODEL_RUN_WEDGE_MESSAGE}: run exceeded ${timeoutMs} ms`)); + }, timeoutMs); + }), + ]).finally(() => clearTimeout(timer)); +} + +/** + * Wrap an ORT session so every run is watched by the wedge timer. predictCHW, + * predictAdaptiveViews, and the DINO head only call session.run, so a run and + * release surface is enough. + * + * @param {{ run: (...args: unknown[]) => Promise, release?: () => Promise }} session + * @param {{ timeoutMs?: number }} [options] + */ +export function wrapSessionWithWatchdog(session, options = {}) { + return { + run: (...args) => watchModelRun(() => session.run(...args), options), + release: () => session.release?.(), + }; +} diff --git a/tests/webgpu-watchdog.test.mjs b/tests/webgpu-watchdog.test.mjs new file mode 100644 index 0000000..fcf12d1 --- /dev/null +++ b/tests/webgpu-watchdog.test.mjs @@ -0,0 +1,214 @@ +import { describe, it } from 'node:test'; +import assert from 'node:assert/strict'; +import { readFile } from 'node:fs/promises'; +import { dirname, join } from 'node:path'; +import { fileURLToPath } from 'node:url'; +import { + MODEL_RUN_WEDGE_TIMEOUT_MS, + MODEL_RUN_WEDGE_MESSAGE, + isWedgeError, + isWebGpuDeviceError, + shouldFallBackToWasm, + createWebGpuHealth, + watchModelRun, + wrapSessionWithWatchdog, +} from '../src/webgpu-watchdog.js'; +import { INFERENCE_TIMEOUT_MS } from '../src/analyze-retry.js'; + +const ROOT = join(dirname(fileURLToPath(import.meta.url)), '..'); + +function sleep(ms) { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +describe('wedge timeout budget', () => { + it('is generous but leaves room for the WASM rebuild and retry', () => { + assert.equal(MODEL_RUN_WEDGE_TIMEOUT_MS, 20_000); + assert.ok(MODEL_RUN_WEDGE_TIMEOUT_MS < INFERENCE_TIMEOUT_MS / 2); + }); +}); + +describe('isWedgeError', () => { + it('recognizes watchdog timeouts by message', () => { + assert.equal(isWedgeError(new Error(`${MODEL_RUN_WEDGE_MESSAGE}: run exceeded 20000 ms`)), true); + assert.equal(isWedgeError(`${MODEL_RUN_WEDGE_MESSAGE}: run exceeded 20000 ms`), true); + }); + + it('does not match unrelated errors', () => { + assert.equal(isWedgeError(new Error('Image fetch timed out')), false); + assert.equal(isWedgeError(new Error('Inference timed out')), false); + assert.equal(isWedgeError(null), false); + assert.equal(isWedgeError(undefined), false); + }); +}); + +describe('isWebGpuDeviceError', () => { + it('matches device-level WebGPU failures', () => { + assert.equal(isWebGpuDeviceError(new Error('Device lost: destroyed')), true); + assert.equal(isWebGpuDeviceError(new Error('GPUDevice is invalid')), true); + assert.equal(isWebGpuDeviceError(new Error('WebGPU validation error')), true); + assert.equal(isWebGpuDeviceError(new Error('The GPU process crashed')), true); + assert.equal(isWebGpuDeviceError('gpu connection lost'), true); + }); + + it('never matches fetch, decode, or size-cap errors', () => { + assert.equal(isWebGpuDeviceError(new Error('Image fetch failed (404)')), false); + assert.equal(isWebGpuDeviceError(new Error('Image fetch timed out')), false); + assert.equal(isWebGpuDeviceError(new Error('Image exceeds size cap')), false); + assert.equal(isWebGpuDeviceError(new Error('Not an image content-type (text/html)')), false); + assert.equal(isWebGpuDeviceError(new Error('The source image could not be decoded')), false); + assert.equal(isWebGpuDeviceError(''), false); + assert.equal(isWebGpuDeviceError(null), false); + }); +}); + +describe('shouldFallBackToWasm', () => { + const wedge = new Error(`${MODEL_RUN_WEDGE_MESSAGE}: run exceeded 20000 ms`); + + it('falls back on a wedged WebGPU run', () => { + assert.equal(shouldFallBackToWasm({ device: 'webgpu', error: wedge }), true); + }); + + it('falls back on a WebGPU device error', () => { + assert.equal( + shouldFallBackToWasm({ device: 'webgpu', error: new Error('Device lost') }), + true + ); + }); + + it('never falls back from WASM (nothing below it)', () => { + assert.equal(shouldFallBackToWasm({ device: 'wasm', error: wedge }), false); + assert.equal( + shouldFallBackToWasm({ device: 'wasm', error: new Error('Device lost') }), + false + ); + }); + + it('is one-way: no second fallback once fallen back', () => { + assert.equal( + shouldFallBackToWasm({ device: 'webgpu', error: wedge, alreadyFellBack: true }), + false + ); + }); + + it('ignores ordinary inference errors on WebGPU', () => { + assert.equal( + shouldFallBackToWasm({ device: 'webgpu', error: new Error('invalid dims') }), + false + ); + assert.equal( + shouldFallBackToWasm({ device: 'webgpu', error: new Error('Image fetch timed out') }), + false + ); + }); +}); + +describe('createWebGpuHealth', () => { + it('latches on the first mark and reports it once', () => { + const health = createWebGpuHealth(); + assert.equal(health.isUnhealthy(), false); + assert.equal(health.reason(), null); + + assert.equal(health.markUnhealthy('Device lost'), true); + assert.equal(health.isUnhealthy(), true); + assert.equal(health.reason(), 'Device lost'); + + assert.equal(health.markUnhealthy('later error'), false); + assert.equal(health.reason(), 'Device lost'); + assert.equal(health.isUnhealthy(), true); + }); + + it('defaults an empty reason to unknown', () => { + const health = createWebGpuHealth(); + health.markUnhealthy(''); + assert.equal(health.reason(), 'unknown'); + }); +}); + +describe('watchModelRun', () => { + it('passes through a run that finishes in time', async () => { + assert.equal(await watchModelRun(async () => 'ok', { timeoutMs: 100 }), 'ok'); + }); + + it('rejects a wedged run with the wedge message', async () => { + let hungResolve; + const hung = new Promise((resolve) => { + hungResolve = resolve; + }); + await assert.rejects( + watchModelRun(() => hung, { timeoutMs: 20 }), + new RegExp(MODEL_RUN_WEDGE_MESSAGE) + ); + hungResolve('late'); + }); + + it('propagates the run error unchanged when the run fails fast', async () => { + await assert.rejects( + watchModelRun(() => Promise.reject(new Error('Device lost')), { timeoutMs: 100 }), + /Device lost/ + ); + }); + + it('swallows a late rejection of the abandoned run', async () => { + let hungReject; + const hung = new Promise((_, reject) => { + hungReject = reject; + }); + await assert.rejects(watchModelRun(hung, { timeoutMs: 10 }), new RegExp(MODEL_RUN_WEDGE_MESSAGE)); + hungReject(new Error('late device loss')); + await sleep(10); + }); +}); + +describe('wrapSessionWithWatchdog', () => { + it('forwards feeds and results for healthy runs', async () => { + const seen = []; + const wrapped = wrapSessionWithWatchdog( + { + run: async (feeds) => { + seen.push(feeds); + return { logits: { data: [1.5] } }; + }, + }, + { timeoutMs: 100 } + ); + const out = await wrapped.run({ input: 'tensor' }); + assert.deepEqual(seen, [{ input: 'tensor' }]); + assert.equal(out.logits.data[0], 1.5); + }); + + it('turns a hung session.run into a wedge rejection', async () => { + const wrapped = wrapSessionWithWatchdog( + { run: () => new Promise(() => {}) }, + { timeoutMs: 15 } + ); + await assert.rejects(wrapped.run({}), new RegExp(MODEL_RUN_WEDGE_MESSAGE)); + }); + + it('delegates release to the underlying session', async () => { + let released = false; + const wrapped = wrapSessionWithWatchdog({ + run: async () => ({}), + release: async () => { + released = true; + }, + }); + await wrapped.release(); + assert.equal(released, true); + }); +}); + +describe('offscreen wiring', () => { + it('wraps WebGPU sessions with the watchdog and keeps the fallback one-way', async () => { + const source = await readFile(join(ROOT, 'src/offscreen.js'), 'utf8'); + assert.match(source, /wrapSessionWithWatchdog/); + assert.match(source, /watchGpuDeviceLoss/); + assert.match(source, /fallbackToWasm/); + assert.match(source, /shouldFallBackToWasm/); + // Logged once: the console.warn lives behind the single-shot promise guard. + assert.match(source, /if \(wasmFallbackPromise\) return wasmFallbackPromise;/); + assert.match(source, /WebGPU backend unhealthy/); + // One-way: nothing in the offscreen document resets the health latch. + assert.doesNotMatch(source, /markHealthy|reset\(\)/); + }); +});