diff --git a/services/ai/localEmbeddingService.ts b/services/ai/localEmbeddingService.ts index b3bb9782..990ae993 100644 --- a/services/ai/localEmbeddingService.ts +++ b/services/ai/localEmbeddingService.ts @@ -1,9 +1,6 @@ -// QNBS-v3: Semantic embedding service — routes to inference.worker.ts via WorkerBus channels. -// Uses Xenova/all-MiniLM-L6-v2 (384-dim, L2-normalized) for semantic RAG and cross-project search. -// Adapted from CannaGuide-2025 embeddingService.ts patterns. +// QNBS-v3: [Semantic embedding service — routes to workers/v2/inference.worker.ts via WorkerBus v2 (docs/adr/0014 migration). Xenova/all-MiniLM-L6-v2, 384-dim, L2-normalized.] -// QNBS-v3: logger import was missing — restartWorker() references logger.error on restart-limit. -import { logger } from '../logger'; +import { ensureInferencePool } from '../workerBusManager'; const EMBEDDING_MODEL = 'Xenova/all-MiniLM-L6-v2'; const MAX_INPUT_CHARS = 512; @@ -22,95 +19,6 @@ function makeCacheKey(text: string): string { export type EmbeddingVector = Float32Array; -// QNBS-v3: The inference worker is loaded lazily to avoid a 2 MB bundle on app start. -let workerInstance: Worker | null = null; -// QNBS-v3: Health check — 30s ping interval; worker restarts if no pong within PONG_TIMEOUT_MS. -const PING_INTERVAL_MS = 30_000; -const PONG_TIMEOUT_MS = 5_000; -let pingTimer: ReturnType | null = null; -let pongTimeoutTimer: ReturnType | null = null; - -// QNBS-v3: Exponential backoff for worker restart to prevent infinite spin on permanent failure. -let restartAttemptCount = 0; -const MAX_RESTART_ATTEMPTS = 5; -const MAX_RESTART_BACKOFF_MS = 60_000; - -function getRestartBackoffMs(): number { - return Math.min(2 ** restartAttemptCount * 1000, MAX_RESTART_BACKOFF_MS); -} - -function clearWorkerHealthTimers(): void { - if (pingTimer) { - clearInterval(pingTimer); - pingTimer = null; - } - if (pongTimeoutTimer) { - clearTimeout(pongTimeoutTimer); - pongTimeoutTimer = null; - } -} - -function restartWorker(): void { - clearWorkerHealthTimers(); - if (workerInstance) { - workerInstance.terminate(); - workerInstance = null; - } - if (import.meta.env?.DEV) { - console.warn('[localEmbeddingService] Inference worker restarted (missed health check pong)'); - } - restartAttemptCount++; - if (restartAttemptCount > MAX_RESTART_ATTEMPTS) { - logger.error( - `[localEmbeddingService] Worker restart limit (${MAX_RESTART_ATTEMPTS}) reached. ` + - 'Embedding service is offline. Reload the app to retry.', - ); - return; - } - const backoff = getRestartBackoffMs(); - if (backoff > 0) { - setTimeout(() => startWorkerHealthCheck(), backoff); - } else { - startWorkerHealthCheck(); - } -} - -function startWorkerHealthCheck(): void { - if (typeof Worker === 'undefined') return; - pingTimer = setInterval(() => { - const w = workerInstance; - if (!w) return; - w.postMessage({ type: 'WORKER_PING' }); - pongTimeoutTimer = setTimeout(() => { - // QNBS-v3: No pong received — worker is dead or hung; restart with backoff. - restartWorker(); - }, PONG_TIMEOUT_MS); - - const pongHandler = (ev: MessageEvent<{ type?: string }>) => { - if (ev.data?.type === 'WORKER_PONG') { - // QNBS-v3: Successful pong resets the restart counter. - restartAttemptCount = 0; - if (pongTimeoutTimer) { - clearTimeout(pongTimeoutTimer); - pongTimeoutTimer = null; - } - w.removeEventListener('message', pongHandler); - } - }; - w.addEventListener('message', pongHandler); - }, PING_INTERVAL_MS); -} - -function getWorker(): Worker { - if (!workerInstance) { - workerInstance = new Worker(new URL('../../workers/inference.worker.ts', import.meta.url), { - type: 'module', - }); - startWorkerHealthCheck(); - } - return workerInstance; -} - function truncate(text: string): string { if (text.length <= MAX_INPUT_CHARS) return text; // QNBS-v3: Silent truncation — warn in dev builds only. @@ -129,30 +37,15 @@ function l2Normalize(vec: number[]): EmbeddingVector { return new Float32Array(vec.map((v) => v / magnitude)); } -function postToWorker( - task: string, - modelId: string, - input: string, -): Promise<{ ok: boolean; result?: number[]; error?: string }> { - return new Promise((resolve) => { - const messageId = `emb-${Date.now()}-${Math.random().toString(36).slice(2, 7)}`; - const worker = getWorker(); - - const handler = (event: MessageEvent) => { - const data = event.data as { - messageId: string; - ok: boolean; - result?: number[]; - error?: string; - }; - if (data.messageId !== messageId) return; - worker.removeEventListener('message', handler); - resolve(data); - }; - - worker.addEventListener('message', handler); - worker.postMessage({ messageId, task, modelId, input }); - }); +async function requestEmbedding(task: string, modelId: string, input: string): Promise { + const bus = await ensureInferencePool(); + if (!bus) throw new Error('WorkerBus v2 unavailable'); + const handle = bus.enqueue<{ task: string; modelId: string; input: string }, number[]>( + 'inference.embed', + { task, modelId, input }, + { capabilities: ['inference.embed'] }, + ); + return handle.result; } export async function embedText(text: string): Promise { @@ -167,11 +60,8 @@ export async function embedText(text: string): Promise { return cached; } - const response = await postToWorker('feature-extraction', EMBEDDING_MODEL, truncated); - if (!response.ok || !response.result) { - throw new Error(response.error ?? 'embedding failed'); - } - const vector = l2Normalize(response.result); + const raw = await requestEmbedding('feature-extraction', EMBEDDING_MODEL, truncated); + const vector = l2Normalize(raw); // QNBS-v3: Evict the oldest (first) entry when at capacity before inserting. if (embeddingCache.size >= EMBEDDING_CACHE_MAX) { @@ -206,12 +96,3 @@ export function cosineSimilarity(a: EmbeddingVector, b: EmbeddingVector): number // Both vectors are already L2-normalised, so cosine = dot product return Math.max(-1, Math.min(1, dot)); } - -// QNBS-v3: Used in testing to reset the worker instance without affecting production code. -export function _resetWorkerForTest(): void { - clearWorkerHealthTimers(); - if (workerInstance) { - workerInstance.terminate(); - workerInstance = null; - } -} diff --git a/services/workerBusManager.ts b/services/workerBusManager.ts index a8004663..caf9a97a 100644 --- a/services/workerBusManager.ts +++ b/services/workerBusManager.ts @@ -67,13 +67,35 @@ async function reRegisterDuckDbPool(bus: WorkerBus): Promise { bus.registerPool('duckdb', options.capabilities, options); } +// QNBS-v3: [Shared with reRegisterInferencePool() below, same rationale as duckDbPoolOptions().] +async function inferencePoolOptions() { + const { MAX_WORKERS_INFERENCE, MIN_WORKERS, WORKER_IDLE_TIMEOUT_MS } = await import( + '@domain/worker-bus' + ); + return { + // QNBS-v3: [Capped below MAX_WORKERS_INFERENCE — each replica loads its own transformers.js pipeline (no cross-replica cache sharing), so 4 concurrent workers could mean 4x the model memory footprint under a burst.] + maxWorkers: Math.min(2, MAX_WORKERS_INFERENCE), + minWorkers: MIN_WORKERS, + idleTimeoutMs: WORKER_IDLE_TIMEOUT_MS, + workerScript: new URL('../workers/v2/inference.worker.ts', import.meta.url).href, + capabilities: ['inference.text', 'inference.embed'] as const, + labels: { pool: 'inference', version: 'v2' }, + }; +} + +/** Re-register the 'inference' pool if it was removed via terminatePool() — a no-op if already present. */ +async function reRegisterInferencePool(bus: WorkerBus): Promise { + if (bus.hasPool('inference')) return; + const options = await inferencePoolOptions(); + bus.registerPool('inference', options.capabilities, options); +} + async function doInitWorkerBus(): Promise { try { const { WorkerBus, WorkerRegistry, MAX_WORKERS_INFERENCE, - MIN_WORKERS, WORKER_IDLE_TIMEOUT_MS, CIRCUIT_BREAKER_THRESHOLD, CIRCUIT_BREAKER_RECOVERY_MS, @@ -96,24 +118,15 @@ async function doInitWorkerBus(): Promise { const registry = new WorkerRegistry(); - // QNBS-v3: new URL(path, import.meta.url) lets Vite emit the worker script as a proper - // asset URL. The .ts extension is allowed — Vite transforms it during build. - const inferenceUrl = new URL('../workers/v2/inference.worker.ts', import.meta.url).href; // QNBS-v3: P1-1 — dedicated WebLLM (WebGPU) worker. Separate pool keeps @mlc-ai/web-llm out // of the transformers.js worker bundle and isolates the GPU lifecycle. const webllmUrl = new URL('../workers/v2/webllm.worker.ts', import.meta.url).href; + const inferenceOptions = await inferencePoolOptions(); registry.register({ poolId: 'inference', - capabilities: ['inference.text', 'inference.embed'], - options: { - maxWorkers: MAX_WORKERS_INFERENCE, - minWorkers: MIN_WORKERS, - idleTimeoutMs: WORKER_IDLE_TIMEOUT_MS, - workerScript: inferenceUrl, - capabilities: ['inference.text', 'inference.embed'], - labels: { pool: 'inference', version: 'v2' }, - }, + capabilities: inferenceOptions.capabilities, + options: inferenceOptions, }); const duckdbOptions = await duckDbPoolOptions(); @@ -189,8 +202,39 @@ export async function ensureDuckDbPool(): Promise { await initWorkerBus(); return _bus; } - // QNBS-v3: [terminatePool('duckdb') can remove the pool while the bus itself stays alive — re-register it here instead of assuming a non-null bus always has every pool.] - await reRegisterDuckDbPool(_bus); + // QNBS-v3: [terminatePool('duckdb') can remove the pool while the bus itself stays alive — re-register it here instead of assuming a non-null bus always has every pool. Catch so a re-registration failure logs instead of breaking the documented "null only if init failed" contract.] + try { + await reRegisterDuckDbPool(_bus); + } catch (err) { + log.error( + 'Failed to re-register duckdb pool', + err instanceof Error ? err : new Error(String(err)), + ); + } + return _bus; +} + +/** + * Ensure the shared local-inference worker pool is available, initializing the WorkerBus on + * demand. QNBS-v3: mirrors ensureDuckDbPool()/ensureWebLlmPool()'s decoupling from + * `enableWorkerBusV2` — embeddings/NLP were never gated by that flag in the v1 worker they + * replace, so toggling an experimental infra flag off must not silently break RAG/cross-project + * search. Returns null only if init failed. + */ +export async function ensureInferencePool(): Promise { + if (_bus === null) { + await initWorkerBus(); + return _bus; + } + // QNBS-v3: [catch so a re-registration failure logs instead of breaking the documented "null only if init failed" contract, mirroring ensureDuckDbPool().] + try { + await reRegisterInferencePool(_bus); + } catch (err) { + log.error( + 'Failed to re-register inference pool', + err instanceof Error ? err : new Error(String(err)), + ); + } return _bus; } diff --git a/tests/unit/localEmbeddingService.test.ts b/tests/unit/localEmbeddingService.test.ts index 314b1f06..45310d42 100644 --- a/tests/unit/localEmbeddingService.test.ts +++ b/tests/unit/localEmbeddingService.test.ts @@ -1,99 +1,96 @@ -import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; - -// QNBS-v3: Self-resolving MockWorker — when postMessage is called it immediately fires -// all registered 'message' handlers with a configurable response. This avoids -// async timing issues with concurrent Promise.all batches in embedBatch. -let mockResponse: { ok: boolean; result?: number[]; error?: string } = { - ok: true, - result: [0.5, 0.5], -}; - -let postMessageCalls: Array<{ messageId: string; input: string; task: string }> = []; - -class MockWorker { - private handlers: Array<(e: MessageEvent) => void> = []; - - addEventListener(_type: string, handler: (e: MessageEvent) => void) { - this.handlers.push(handler); - } - - removeEventListener(_type: string, handler: (e: MessageEvent) => void) { - const idx = this.handlers.indexOf(handler); - if (idx >= 0) this.handlers.splice(idx, 1); - } - - postMessage(msg: { messageId: string; input: string; task: string }) { - postMessageCalls.push(msg); - // QNBS-v3: Respond synchronously so Promise.all batches resolve without microtask tricks. - const response = { ...mockResponse, messageId: msg.messageId }; - const event = { data: response } as MessageEvent; - for (const handler of [...this.handlers]) { - handler(event); - } - } +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +// QNBS-v3: [localEmbeddingService now routes through WorkerBus v2 — mock ensureInferencePool() instead of the global Worker constructor.] + +const { mockEnqueue, mockEnsureInferencePool } = vi.hoisted(() => ({ + mockEnqueue: vi.fn(), + mockEnsureInferencePool: vi.fn(), +})); - terminate() {} +vi.mock('../../services/workerBusManager', () => ({ + ensureInferencePool: mockEnsureInferencePool, +})); + +const { clearEmbeddingCache, cosineSimilarity, embedBatch, embedText } = await import( + '../../services/ai/localEmbeddingService' +); + +function makeHandle(result: Promise) { + return { taskId: 't1', result, progress: (async function* () {})(), cancel: vi.fn() }; } -vi.stubGlobal('Worker', MockWorker); +function makeBus() { + return { enqueue: mockEnqueue }; +} -import { - _resetWorkerForTest, - cosineSimilarity, - embedBatch, - embedText, -} from '../../services/ai/localEmbeddingService'; +let requestCalls: Array<{ task: string; modelId: string; input: string }> = []; +let nextResult: number[] = [0.5, 0.5]; beforeEach(() => { - _resetWorkerForTest(); - mockResponse = { ok: true, result: [0.5, 0.5] }; - postMessageCalls = []; -}); - -afterEach(() => { - _resetWorkerForTest(); + vi.clearAllMocks(); + clearEmbeddingCache(); + requestCalls = []; + nextResult = [0.5, 0.5]; + mockEnsureInferencePool.mockResolvedValue(makeBus()); + mockEnqueue.mockImplementation((_taskType: string, payload: unknown) => { + requestCalls.push(payload as { task: string; modelId: string; input: string }); + return makeHandle(Promise.resolve(nextResult)); + }); }); // ─── embedText ────────────────────────────────────────────────────────────── describe('embedText', () => { - it('returns a Float32Array for a successful worker response', async () => { - mockResponse = { ok: true, result: [0.5, 0.5] }; + it('returns a Float32Array for a successful task result', async () => { + nextResult = [0.5, 0.5]; const vec = await embedText('Hello world'); expect(vec).toBeInstanceOf(Float32Array); expect(vec.length).toBe(2); }); it('L2-normalises the returned vector (magnitude ≈ 1)', async () => { - mockResponse = { ok: true, result: [3, 4] }; // magnitude = 5 → normalised to [0.6, 0.8] + nextResult = [3, 4]; // magnitude = 5 → normalised to [0.6, 0.8] const vec = await embedText('Normalise me'); const magnitude = Math.sqrt(vec[0]! ** 2 + vec[1]! ** 2); expect(magnitude).toBeCloseTo(1, 5); }); - it('throws when worker returns ok:false with error message', async () => { - mockResponse = { ok: false, error: 'OOM' }; + it('propagates the error message when the task rejects', async () => { + mockEnqueue.mockReturnValue(makeHandle(Promise.reject(new Error('OOM')))); await expect(embedText('fail case')).rejects.toThrow('OOM'); }); - it('throws with generic message when error field is absent', async () => { - mockResponse = { ok: false }; - await expect(embedText('fail case 2')).rejects.toThrow('embedding failed'); + it('throws WorkerBus v2 unavailable without enqueuing when the pool is unavailable', async () => { + mockEnsureInferencePool.mockResolvedValue(null); + await expect(embedText('fail case 2')).rejects.toThrow('WorkerBus v2 unavailable'); + expect(mockEnqueue).not.toHaveBeenCalled(); }); it('truncates input silently when text exceeds 512 chars', async () => { await embedText('a'.repeat(600)); - expect(postMessageCalls[0]?.input.length).toBe(512); + expect(requestCalls[0]?.input.length).toBe(512); }); it('does not truncate input at exactly 512 chars', async () => { await embedText('b'.repeat(512)); - expect(postMessageCalls[0]?.input.length).toBe(512); + expect(requestCalls[0]?.input.length).toBe(512); }); - it('sends feature-extraction task to worker', async () => { + it('enqueues inference.embed with the feature-extraction task and inference.embed capability', async () => { await embedText('test'); - expect(postMessageCalls[0]?.task).toBe('feature-extraction'); + expect(requestCalls[0]?.task).toBe('feature-extraction'); + expect(mockEnqueue).toHaveBeenCalledWith( + 'inference.embed', + expect.anything(), + expect.objectContaining({ capabilities: ['inference.embed'] }), + ); + }); + + it('returns the cached vector on a second call for the same text without re-enqueuing', async () => { + const first = await embedText('cache me'); + const second = await embedText('cache me'); + expect(second).toBe(first); + expect(mockEnqueue).toHaveBeenCalledTimes(1); }); }); @@ -101,7 +98,7 @@ describe('embedText', () => { describe('embedBatch', () => { it('returns an array of Float32Arrays, one per input', async () => { - mockResponse = { ok: true, result: [0.5, 0.5] }; + nextResult = [0.5, 0.5]; const results = await embedBatch(['a', 'b', 'c']); expect(results).toHaveLength(3); for (const vec of results) { @@ -110,16 +107,16 @@ describe('embedBatch', () => { }); it('returns correct count for 9 texts (spans two micro-batches of 8+1)', async () => { - mockResponse = { ok: true, result: [0.1, 0.2] }; + nextResult = [0.1, 0.2]; const texts = Array.from({ length: 9 }, (_, i) => `text-${i}`); const results = await embedBatch(texts); expect(results).toHaveLength(9); }); - it('sends exactly N worker messages for N input texts', async () => { - mockResponse = { ok: true, result: [0.1, 0.2] }; + it('sends exactly N enqueue calls for N input texts', async () => { + nextResult = [0.1, 0.2]; await embedBatch(['x', 'y', 'z']); - expect(postMessageCalls).toHaveLength(3); + expect(requestCalls).toHaveLength(3); }); it('returns empty array for empty input', async () => { diff --git a/tests/unit/workerBusManager.test.ts b/tests/unit/workerBusManager.test.ts index 151f3717..2a17c768 100644 --- a/tests/unit/workerBusManager.test.ts +++ b/tests/unit/workerBusManager.test.ts @@ -66,6 +66,19 @@ vi.mock('../../services/legacyWorkerBusAdapter', () => ({ }), })); +const { mockLogError } = vi.hoisted(() => ({ mockLogError: vi.fn() })); + +// QNBS-v3: mocked separately so re-registration-failure tests can assert on log.error without asserting on the real StructuredLogger's console/IDB side effects +vi.mock('../../services/logger', () => ({ + createLogger: () => ({ + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: mockLogError, + withContext: vi.fn(), + }), +})); + describe('workerBusManager', () => { beforeEach(() => { vi.resetModules(); @@ -76,6 +89,7 @@ describe('workerBusManager', () => { mockHasPool.mockClear(); mockHasPool.mockReturnValue(true); mockInstall.mockClear(); + mockLogError.mockClear(); }); afterEach(() => { @@ -189,5 +203,96 @@ describe('workerBusManager', () => { expect.objectContaining({ workerScript: expect.stringContaining('duckdb.worker') }), ); }); + + it('logs and returns the live bus instead of rejecting when re-registration fails', async () => { + // QNBS-v3: [Regression guard for the documented "null only if init failed" contract — a + // re-registration failure must not propagate out of ensureDuckDbPool.] + const { initWorkerBus, ensureDuckDbPool } = await import('../../services/workerBusManager'); + await initWorkerBus(); + mockHasPool.mockReturnValue(false); + mockRegisterPool.mockImplementationOnce(() => { + throw new Error('registerPool boom'); + }); + + const bus = await ensureDuckDbPool(); + + expect(bus).not.toBeNull(); + expect(mockLogError).toHaveBeenCalledWith( + 'Failed to re-register duckdb pool', + expect.any(Error), + ); + }); + }); + + describe('ensureInferencePool', () => { + // QNBS-v3: [Imports the real module (only @domain/worker-bus is mocked) — unlike + // localEmbeddingService.test.ts's fully-mocked workerBusManager, this actually + // exercises the export surface. This exact suite would have caught PR #288's + // missing-export bug (ensureInferencePool was imported but never defined).] + it('initializes the bus when not yet running', async () => { + const { ensureInferencePool, isWorkerBusReady } = await import( + '../../services/workerBusManager' + ); + const bus = await ensureInferencePool(); + expect(bus).not.toBeNull(); + expect(isWorkerBusReady()).toBe(true); + }); + + it('does not re-register the pool when it is already present', async () => { + const { initWorkerBus, ensureInferencePool } = await import( + '../../services/workerBusManager' + ); + await initWorkerBus(); + mockRegisterPool.mockClear(); + mockHasPool.mockReturnValue(true); + + await ensureInferencePool(); + + expect(mockHasPool).toHaveBeenCalledWith('inference'); + expect(mockRegisterPool).not.toHaveBeenCalled(); + }); + + it('re-registers the inference pool when the bus is alive but the pool was removed', async () => { + const { initWorkerBus, ensureInferencePool } = await import( + '../../services/workerBusManager' + ); + await initWorkerBus(); + mockRegisterPool.mockClear(); + mockHasPool.mockReturnValue(false); + + const bus = await ensureInferencePool(); + + expect(bus).not.toBeNull(); + // QNBS-v3: asserts the memory-safety cap (below MAX_WORKERS_INFERENCE) survives re-registration, not just initial registration + expect(mockRegisterPool).toHaveBeenCalledWith( + 'inference', + expect.arrayContaining(['inference.text', 'inference.embed']), + expect.objectContaining({ + maxWorkers: 2, + workerScript: expect.stringContaining('inference.worker'), + }), + ); + }); + + it('logs and returns the live bus instead of rejecting when re-registration fails', async () => { + // QNBS-v3: [Regression guard for the documented "null only if init failed" contract — a + // re-registration failure must not propagate out of ensureInferencePool.] + const { initWorkerBus, ensureInferencePool } = await import( + '../../services/workerBusManager' + ); + await initWorkerBus(); + mockHasPool.mockReturnValue(false); + mockRegisterPool.mockImplementationOnce(() => { + throw new Error('registerPool boom'); + }); + + const bus = await ensureInferencePool(); + + expect(bus).not.toBeNull(); + expect(mockLogError).toHaveBeenCalledWith( + 'Failed to re-register inference pool', + expect.any(Error), + ); + }); }); });