Skip to content

feat(ai): migrate localEmbeddingService onto WorkerBus v2 #288

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
qnbs merged 13 commits into from
Jul 30, 2026
+239 −212
Merged

feat(ai): migrate localEmbeddingService onto WorkerBus v2 #288

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
b4b1bbf
test(worker-bus): export v2 duckdb/inference handlers + parity gate t…
qnbs Jul 29, 2026
bcc41c2
fix(worker-bus): close leaked OPFS connection on ATTACH failure + QNB…
qnbs Jul 29, 2026
bc371b3
feat(duckdb): migrate duckdbClient onto WorkerBus v2, delete v1 worker
qnbs Jul 29, 2026
dbebf5c
fix(worker-bus): make OPFS-connection cleanup best-effort on ATTACH f…
qnbs Jul 29, 2026
fc5b6e1
merge: pull in PR0's OPFS-cleanup fix
qnbs Jul 29, 2026
871cfba
feat(ai): migrate localEmbeddingService onto WorkerBus v2
qnbs Jul 29, 2026
36c5f7c
fix(worker-bus): pool-termination lifecycle, idle-respawn recovery, S…
qnbs Jul 29, 2026
8e66a9f
merge: pull in PR287's correction-loop fixes
qnbs Jul 29, 2026
776ff4e
fix(ai): add the missing ensureInferencePool export that broke prod b…
qnbs Jul 29, 2026
73efc0d
Merge remote-tracking branch 'origin/main' into feat/worker-v2-embedd…
qnbs Jul 30, 2026
8a63b7c
fix(worker-bus): guard pool re-registration failures in workerBusManager
qnbs Jul 30, 2026
a8686c4
test(worker-bus): close Codecov patch-coverage gap on PR #288
qnbs Jul 30, 2026
b3e5d1a
docs(worker-bus): add QNBS-v3 rationale comments per CodeRabbit outsi…
qnbs Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
145 changes: 13 additions & 132 deletions services/ai/localEmbeddingService.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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<typeof setInterval> | null = null;
let pongTimeoutTimer: ReturnType<typeof setTimeout> | 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.
Expand All @@ -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<number[]> {
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<EmbeddingVector> {
Expand All @@ -167,11 +60,8 @@ export async function embedText(text: string): Promise<EmbeddingVector> {
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) {
Expand Down Expand Up @@ -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;
}
}
74 changes: 59 additions & 15 deletions services/workerBusManager.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,13 +67,35 @@ async function reRegisterDuckDbPool(bus: WorkerBus): Promise<void> {
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<void> {
if (bus.hasPool('inference')) return;
const options = await inferencePoolOptions();
bus.registerPool('inference', options.capabilities, options);
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

async function doInitWorkerBus(): Promise<void> {
try {
const {
WorkerBus,
WorkerRegistry,
MAX_WORKERS_INFERENCE,
MIN_WORKERS,
WORKER_IDLE_TIMEOUT_MS,
CIRCUIT_BREAKER_THRESHOLD,
CIRCUIT_BREAKER_RECOVERY_MS,
Expand All @@ -96,24 +118,15 @@ async function doInitWorkerBus(): Promise<void> {

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();
Expand Down Expand Up @@ -189,8 +202,39 @@ export async function ensureDuckDbPool(): Promise<WorkerBus | null> {
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<WorkerBus | null> {
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;
}

Expand Down
Loading
Loading