albex 0.1.0 → 0.6.0

package/src/pool/coordinator.ts

@@ -0,0 +1,324 @@
+/**
+ * Pool coordinator.
+ *
+ * Runs on the main thread and orchestrates N `AlbexEngineWorker`-style
+ * shards. The corpus is split round-robin by document at index time:
+ *
+ *   indexFile(doc_i) → workers[i % N]
+ *
+ * Search broadcasts the query to every shard, collects each shard's
+ * top-K results, and performs a global heap-merge — preserving the
+ * relevance ordering of the unsharded engine. Ties broken by score.
+ *
+ * **Modes:**
+ *  - `'replicated'` (default): each worker holds its own copy of its shard.
+ *    Works in every browser, no special headers needed.
+ *  - `'shared'`: backed by SharedArrayBuffer, requires cross-origin isolation
+ *    (COOP+COEP). Single source of truth, halves the memory cost at the
+ *    expense of harder deployment.
+ *
+ * Today we ship `replicated` — `shared` is wired up to detect availability
+ * (so consumers can opt in once we land the WASM-side `SharedArrayBuffer`
+ * variant in a future iteration) but defaults to off.
+ */
+
+import type {
+  AlbexOptions,
+  IndexedDocument,
+  SearchOptions,
+  SearchResult,
+  EngineStats,
+  SearchStats,
+} from '../albex.js';
+import type { Tier } from '../profile.js';
+import { detectProfile, pickWorkerCount } from '../profile.js';
+import { AlbexInitError, AlbexError } from '../errors.js';
+import type {
+  WorkerRequest,
+  WorkerResponse,
+  WorkerOp,
+} from '../worker-protocol.js';
+
+export interface AlbexPoolOptions extends AlbexOptions {
+  /** Worker runtime URL (the `worker-runtime.js` from this package). */
+  workerUrl: string | URL;
+  /**
+   * Number of worker shards. `'auto'` (default) uses half of
+   * `navigator.hardwareConcurrency`, clamped to [1, 8].
+   *
+   * Pass an explicit number to override — for example `1` to debug, or
+   * a higher fixed value if you know your corpus benefits from more
+   * shards on a specific deployment.
+   */
+  workers?: number | 'auto';
+  /** Memory-sharing strategy. */
+  mode?: 'replicated' | 'shared';
+}
+
+interface Pending {
+  resolve: (v: unknown) => void;
+  reject:  (e: unknown) => void;
+}
+
+interface Shard {
+  worker: Worker;
+  nextId: number;
+  pending: Map<number, Pending>;
+  // Number of documents this shard currently holds — used for stats and
+  // for the round-robin counter.
+  docCount: number;
+}
+
+let _poolSearchStreamWarned = false;
+
+export class AlbexPool {
+  private readonly _opts: AlbexPoolOptions;
+  private _shards: Shard[] = [];
+  private _docsCache: IndexedDocument[] = [];
+  private _rrCursor = 0;
+  private _lastSearch: SearchStats | null = null;
+  private _tier: Tier | null = null;
+
+  constructor(opts: AlbexPoolOptions) {
+    this._opts = opts;
+  }
+
+  async init(): Promise<void> {
+    const profile = await detectProfile();
+    const n =
+      typeof this._opts.workers === 'number' && this._opts.workers > 0
+        ? Math.floor(this._opts.workers)
+        : pickWorkerCount(profile);
+
+    // `shared` mode would require SharedArrayBuffer + COOP/COEP. The check
+    // here is informational — actually using SAB inside the worker requires
+    // a corresponding WASM-side rework that is not yet shipped. Falling back
+    // to replicated is the safe path.
+    if (this._opts.mode === 'shared' && !profile.coopCoep) {
+      console.warn('[albex] pool mode=shared requested but cross-origin isolation is not active; falling back to replicated');
+    }
+
+    const shardOpts: AlbexOptions = {
+      wasmUrl:     this._opts.wasmUrl,
+      wasmBaseUrl: this._opts.wasmBaseUrl,
+      pdfWasmUrl:  this._opts.pdfWasmUrl,
+      tier:        this._opts.tier,
+      simd:        this._opts.simd,
+    };
+
+    for (let i = 0; i < n; i++) {
+      const shard = this._spawnShard();
+      await this._send(shard, { kind: 'init', opts: shardOpts });
+      this._shards.push(shard);
+    }
+
+    // Tier is the same across shards — capture it from shard 0 stats.
+    const stats0 = await this._send<EngineStats>(this._shards[0]!, { kind: 'getStats' });
+    this._tier = stats0.tier;
+  }
+
+  // ── Shard plumbing ─────────────────────────────────────────────────────
+
+  private _spawnShard(): Shard {
+    const worker = new Worker(this._opts.workerUrl, { type: 'module' });
+    const shard: Shard = { worker, nextId: 1, pending: new Map(), docCount: 0 };
+    worker.onmessage = (ev: MessageEvent<WorkerResponse>) => {
+      const p = shard.pending.get(ev.data.id);
+      if (!p) return;
+      shard.pending.delete(ev.data.id);
+      if (ev.data.ok) p.resolve(ev.data.result);
+      else            p.reject(this._rehydrate(ev.data.error));
+    };
+    worker.onerror = (e) => {
+      const err = new AlbexInitError(`Pool shard crashed: ${e.message}`);
+      for (const [, p] of shard.pending) p.reject(err);
+      shard.pending.clear();
+    };
+    return shard;
+  }
+
+  private _send<T = unknown>(
+    shard: Shard,
+    op: WorkerOp,
+    transfer: Transferable[] = [],
+  ): Promise<T> {
+    const id = shard.nextId++;
+    const req: WorkerRequest = { id, op };
+    return new Promise<T>((resolve, reject) => {
+      shard.pending.set(id, { resolve: resolve as (v: unknown) => void, reject });
+      shard.worker.postMessage(req, transfer);
+    });
+  }
+
+  private _broadcast<T = unknown>(op: WorkerOp): Promise<T[]> {
+    return Promise.all(this._shards.map(s => this._send<T>(s, op)));
+  }
+
+  private _rehydrate(e: { name: string; kind?: string; message: string }): Error {
+    // Re-importing the rich subclasses here would create a circular dep
+    // with the engine wrapper; we forward the kind unchanged in a generic
+    // AlbexError instance instead.
+    if (e.kind) return new AlbexError(e.kind, e.message);
+    const err = new Error(e.message);
+    err.name = e.name;
+    return err;
+  }
+
+  // ── Public API (mirrors AlbexEngine) ──────────────────────────────────
+
+  async indexFile(file: File): Promise<IndexedDocument> {
+    if (this._shards.length === 0) throw new AlbexInitError('Pool not initialised');
+    const idx = this._rrCursor++ % this._shards.length;
+    const shard = this._shards[idx]!;
+    const buffer = await file.arrayBuffer();
+    const doc = await this._send<IndexedDocument>(
+      shard,
+      { kind: 'indexFile', name: file.name, buffer },
+      [buffer],
+    );
+    shard.docCount++;
+    this._docsCache.push(doc);
+    return doc;
+  }
+
+  /**
+   * Map-reduce search:
+   *   1. broadcast the query to every shard,
+   *   2. each shard runs its local Bloom→Bitap→top-K,
+   *   3. coordinator merges the K-best from each shard,
+   *   4. global top-K returned in descending score order.
+   *
+   * Capped to `setMaxResults` results (default 50) AFTER merge.
+   */
+  async search(query: string, opts: SearchOptions = {}): Promise<SearchResult[]> {
+    if (this._shards.length === 0) return [];
+    const t0 = performance.now();
+    const buckets = await Promise.all(
+      this._shards.map(s => this._send<SearchResult[]>(s, { kind: 'search', query, options: opts })),
+    );
+
+    // Merge: simple flatten + sort (descending). Each shard already returned
+    // up to its local cap; the global cap could be different but in practice
+    // matches the per-shard cap.
+    const merged = buckets.flat();
+    merged.sort((a, b) => b.score - a.score);
+
+    // Aggregate search stats across shards.
+    const stats = await Promise.all(
+      this._shards.map(s => this._send<SearchStats | null>(s, { kind: 'getLastSearchStats' })),
+    );
+    let bloomTested = 0, bloomPassed = 0, bitapMatched = 0;
+    for (const s of stats) {
+      if (!s) continue;
+      bloomTested  += s.bloomTested;
+      bloomPassed  += s.bloomPassed;
+      bitapMatched += s.bitapMatched;
+    }
+    this._lastSearch = {
+      query,
+      timeMs: performance.now() - t0,
+      results: merged.length,
+      bloomTested, bloomPassed, bitapMatched,
+    };
+
+    return merged;
+  }
+
+  /**
+   * Cooperative variant of `search`. The coordinator awaits every shard,
+   * merges, sorts, and then iterates out the result list — the iterator
+   * shape exists so callers can `break` early. Streaming individual
+   * results before the cross-shard merge would deliver them in arbitrary
+   * order; that's why we materialise first.
+   */
+  async *searchCooperative(query: string, opts: SearchOptions = {}): AsyncIterable<SearchResult> {
+    const results = await this.search(query, opts);
+    for (const r of results) yield r;
+  }
+
+  /**
+   * @deprecated Renamed to `searchCooperative` in 0.3.0. Alias removed in 0.4.0.
+   */
+  async *searchStream(query: string, opts: SearchOptions = {}): AsyncIterable<SearchResult> {
+    if (!_poolSearchStreamWarned) {
+      _poolSearchStreamWarned = true;
+      console.warn('[albex] `AlbexPool.searchStream` is deprecated; rename to `searchCooperative`. Alias removed in 0.4.0.');
+    }
+    yield* this.searchCooperative(query, opts);
+  }
+
+  /**
+   * Remove a document from whichever shard owns it. Tries each shard until
+   * one reports success. O(workers) — acceptable since N ≤ 8.
+   */
+  async removeDocument(id: string): Promise<boolean> {
+    for (const shard of this._shards) {
+      const ok = await this._send<boolean>(shard, { kind: 'removeDocument', id });
+      if (ok) {
+        shard.docCount = Math.max(0, shard.docCount - 1);
+        this._docsCache = this._docsCache.filter(d => d.name !== id && d.contentHash !== id);
+        return true;
+      }
+    }
+    return false;
+  }
+
+  async compact(): Promise<void> {
+    await this._broadcast({ kind: 'compact' });
+  }
+
+  async reset(): Promise<void> {
+    await this._broadcast({ kind: 'reset' });
+    this._docsCache = [];
+    for (const s of this._shards) s.docCount = 0;
+  }
+
+  /** Aggregate engine stats across all shards. */
+  async getStats(): Promise<EngineStats> {
+    const all = await this._broadcast<EngineStats>({ kind: 'getStats' });
+    let documents = 0, chunks = 0, textUsed = 0, textCapacity = 0, wasmMemoryBytes = 0;
+    let maxChunks = 0, maxDocs = 0;
+    for (const s of all) {
+      documents       += s.documents;
+      chunks          += s.chunks;
+      textUsed        += s.textUsed;
+      textCapacity    += s.textCapacity;
+      wasmMemoryBytes += s.wasmMemoryBytes;
+      maxChunks       += s.maxChunks;
+      maxDocs         += s.maxDocs;
+    }
+    return {
+      documents, chunks, textUsed, textCapacity, wasmMemoryBytes,
+      tier: this._tier, maxChunks, maxDocs,
+    };
+  }
+
+  getLastSearchStats(): SearchStats | null {
+    return this._lastSearch;
+  }
+
+  async getDocuments(): Promise<readonly IndexedDocument[]> {
+    return this._docsCache.slice();
+  }
+
+  async setMaxErrors(n: 0 | 1 | 2 | 3):       Promise<void> { await this._broadcast({ kind: 'setMaxErrors',  n }); }
+  async setThreshold(n: number):              Promise<void> { await this._broadcast({ kind: 'setThreshold', n }); }
+  async setMaxResults(n: number):             Promise<void> { await this._broadcast({ kind: 'setMaxResults', n }); }
+  async setLanguage(lang: 'off' | 'es'):      Promise<void> { await this._broadcast({ kind: 'setLanguage', lang }); }
+
+  /** Number of shards currently running. */
+  get workerCount(): number { return this._shards.length; }
+
+  /** Tier loaded by the shards (same value across all of them). */
+  get tier(): Tier | null { return this._tier; }
+
+  [Symbol.dispose](): void {
+    for (const s of this._shards) {
+      for (const [, p] of s.pending) p.reject(new AlbexError('disposed', 'Pool disposed'));
+      s.pending.clear();
+      s.worker.terminate();
+    }
+    this._shards = [];
+    this._docsCache = [];
+  }
+}

package/src/profile.ts

@@ -0,0 +1,280 @@
+/**
+ * Capability profile detection.
+ *
+ * Probes the host environment ONCE per session and returns a structured
+ * snapshot. Used by `Albex.create()` to choose the optimal binary tier
+ * (mini/std/pro), enable SIMD when available, decide whether to spin up a
+ * worker pool, allocate WebGPU resources, etc.
+ *
+ * The function is intentionally side-effect-free except for the optional
+ * `sessionStorage` cache — every property is a synchronous-or-async probe
+ * with a sensible default if it throws.
+ */
+
+const CACHE_KEY = 'albex.profile.v1';
+
+export interface DeviceProfile {
+  /** Number of logical cores reported by the browser (1-128 typical). */
+  cores: number;
+
+  /**
+   * Approximate device memory in GB, capped at 8 by the spec for privacy.
+   * `null` if `navigator.deviceMemory` is unavailable (Safari, Firefox).
+   */
+  memoryGB: number | null;
+
+  wasm: {
+    /** WebAssembly SIMD (v128) instructions supported. */
+    simd: boolean;
+    /** Bulk memory ops (memory.copy, memory.fill) supported. */
+    bulkMemory: boolean;
+    /** Threads (Atomics + SharedArrayBuffer) supported AND page is cross-origin isolated. */
+    threads: boolean;
+  };
+
+  /** `true` if `navigator.gpu` is available. Does NOT mean a device was acquired yet. */
+  webgpu: boolean;
+
+  /** Cross-Origin Isolation (required for SAB-based shared mode). */
+  coopCoep: boolean;
+
+  /**
+   * OPFS / IndexedDB free space, if reported. Some browsers gate this behind
+   * permissions or return overly conservative values. `null` if unknown.
+   */
+  storage: {
+    quotaBytes: number | null;
+    usageBytes: number | null;
+  };
+
+  /**
+   * Network conditions if reported via Network Information API (Chrome only
+   * at time of writing). Used to throttle PDF wasm prefetch on slow links.
+   */
+  net: {
+    effectiveType: 'slow-2g' | '2g' | '3g' | '4g' | 'unknown';
+    saveData: boolean;
+  };
+
+  /**
+   * Battery state if Battery Status API is available (Chrome on mobile).
+   * `null` everywhere else. Used to opt into low-power mode.
+   */
+  battery: {
+    level: number | null;     // 0-1
+    charging: boolean | null;
+  } | null;
+
+  /** Page visibility at probe time. Engines should re-listen after this. */
+  visible: boolean;
+}
+
+// ── WASM feature probes ──────────────────────────────────────────────────────
+//
+// Each probe is a minimal valid module that uses exactly one feature. If
+// `WebAssembly.validate` returns true, the host supports it. We hand-coded
+// the byte sequences instead of pulling in a builder; the modules are
+// tiny and stable across spec revisions.
+
+const PROBE_SIMD = new Uint8Array([
+  // magic + version
+  0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00,
+  // type section: (func (result v128))
+  0x01, 0x05, 0x01, 0x60, 0x00, 0x01, 0x7b,
+  // function section
+  0x03, 0x02, 0x01, 0x00,
+  // code section: v128.const 0
+  0x0a, 0x16, 0x01, 0x14, 0x00,
+  0xfd, 0x0c,
+  0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0,
+  0x0b,
+]);
+
+const PROBE_BULK_MEMORY = new Uint8Array([
+  // magic + version
+  0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00,
+  // type section: (func)
+  0x01, 0x04, 0x01, 0x60, 0x00, 0x00,
+  // function section: one function of type 0
+  0x03, 0x02, 0x01, 0x00,
+  // memory section: min=1 page
+  0x05, 0x03, 0x01, 0x00, 0x01,
+  // code section
+  //   section id 0x0a, section size 13
+  //   func count = 1
+  //   body size = 11
+  //   locals = 0
+  //   i32.const 0, i32.const 0, i32.const 0
+  //   memory.fill (0xfc 0x0b 0x00)
+  //   end (0x0b)
+  0x0a, 0x0d, 0x01, 0x0b, 0x00,
+  0x41, 0x00, 0x41, 0x00, 0x41, 0x00,
+  0xfc, 0x0b, 0x00,
+  0x0b,
+]);
+
+const PROBE_THREADS = new Uint8Array([
+  // magic + version
+  0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00,
+  // memory: shared, min=1, max=1
+  0x05, 0x04, 0x01, 0x03, 0x01, 0x01,
+]);
+
+function safeValidate(bytes: Uint8Array): boolean {
+  try {
+    if (typeof WebAssembly === 'undefined') return false;
+    // WebAssembly.validate types insist on ArrayBuffer-backed BufferSource;
+    // copying through a fresh Uint8Array satisfies that.
+    const copy = new Uint8Array(bytes.byteLength);
+    copy.set(bytes);
+    return WebAssembly.validate(copy);
+  } catch {
+    return false;
+  }
+}
+
+// ── Storage probe (async) ────────────────────────────────────────────────────
+
+async function probeStorage(): Promise<{ quotaBytes: number | null; usageBytes: number | null }> {
+  try {
+    if (typeof navigator !== 'undefined' && navigator.storage?.estimate) {
+      const est = await navigator.storage.estimate();
+      return {
+        quotaBytes: typeof est.quota === 'number' ? est.quota : null,
+        usageBytes: typeof est.usage === 'number' ? est.usage : null,
+      };
+    }
+  } catch { /* fall through */ }
+  return { quotaBytes: null, usageBytes: null };
+}
+
+// ── Battery probe (async) ────────────────────────────────────────────────────
+
+interface BatteryManagerLike {
+  level: number;
+  charging: boolean;
+}
+
+async function probeBattery(): Promise<DeviceProfile['battery']> {
+  try {
+    // @ts-expect-error Battery API typing is sparse
+    const getBat: (() => Promise<BatteryManagerLike>) | undefined = navigator?.getBattery?.bind(navigator);
+    if (!getBat) return null;
+    const b = await getBat();
+    return { level: b.level, charging: b.charging };
+  } catch {
+    return null;
+  }
+}
+
+// ── Public API ───────────────────────────────────────────────────────────────
+
+/**
+ * Build a fresh `DeviceProfile`. The result is cached in `sessionStorage`
+ * (key `albex.profile.v1`) so subsequent calls in the same tab are free.
+ *
+ * Pass `{ fresh: true }` to bypass the cache (useful in tests or after a
+ * known capability change such as connecting to power).
+ */
+export async function detectProfile(opts: { fresh?: boolean } = {}): Promise<DeviceProfile> {
+  if (!opts.fresh) {
+    try {
+      const cached = sessionStorage.getItem(CACHE_KEY);
+      if (cached) return JSON.parse(cached) as DeviceProfile;
+    } catch { /* ignore */ }
+  }
+
+  const nav: Navigator | undefined = typeof navigator !== 'undefined' ? navigator : undefined;
+
+  const simd       = safeValidate(PROBE_SIMD);
+  const bulkMemory = safeValidate(PROBE_BULK_MEMORY);
+  // Threads need both validate (shared memory) AND cross-origin isolation.
+  const sabAvailable = typeof SharedArrayBuffer !== 'undefined';
+  const coopCoep =
+    typeof self !== 'undefined' &&
+    typeof (self as unknown as { crossOriginIsolated?: boolean }).crossOriginIsolated === 'boolean' &&
+    (self as unknown as { crossOriginIsolated: boolean }).crossOriginIsolated === true;
+  const threadsValidated = safeValidate(PROBE_THREADS);
+  const threads = sabAvailable && coopCoep && threadsValidated;
+
+  const webgpu = !!(nav && 'gpu' in nav);
+
+  const [storage, battery] = await Promise.all([
+    probeStorage(),
+    probeBattery(),
+  ]);
+
+  const connection = (nav as unknown as { connection?: { effectiveType?: string; saveData?: boolean } } | undefined)?.connection;
+
+  const profile: DeviceProfile = {
+    cores: nav?.hardwareConcurrency ?? 1,
+    memoryGB:
+      (nav as unknown as { deviceMemory?: number } | undefined)?.deviceMemory ?? null,
+    wasm: { simd, bulkMemory, threads },
+    webgpu,
+    coopCoep,
+    storage,
+    net: {
+      effectiveType: (connection?.effectiveType as DeviceProfile['net']['effectiveType']) ?? 'unknown',
+      saveData: connection?.saveData ?? false,
+    },
+    battery,
+    visible: typeof document !== 'undefined'
+      ? document.visibilityState === 'visible'
+      : true,
+  };
+
+  try {
+    sessionStorage.setItem(CACHE_KEY, JSON.stringify(profile));
+  } catch { /* sessionStorage may be full or disabled */ }
+
+  return profile;
+}
+
+// ── Tier selection ───────────────────────────────────────────────────────────
+
+/**
+ * @deprecated The tier system was removed in 0.5.0 (audit 4.1: "6
+ * binaries × no proven benefit"). The type remains exported as `'std'`
+ * for backwards compatibility with code that read `engine.getStats().tier`.
+ */
+export type Tier = 'std';
+
+/**
+ * @deprecated Always returns `'std'` as of 0.5.0. Albex ships exactly
+ * two main binaries (baseline + SIMD); the only runtime variant is the
+ * SIMD probe, not a capacity tier. Kept callable so existing integrators
+ * don't break, but the value has no operational meaning anymore.
+ */
+export function pickTier(_profile: DeviceProfile): Tier {
+  return 'std';
+}
+
+/**
+ * Suggested number of search/index workers given a profile.
+ *
+ * Heuristic: half of logical cores, clamped to [1, 8]. The clamp is
+ * deliberate — beyond 8 workers the postMessage coordination overhead
+ * starts eating the parallelism gain for typical Albex query shapes.
+ *
+ * Battery awareness: when battery is reported and below 20% and discharging,
+ * we shrink the pool to 1 to preserve the user's session.
+ */
+export function pickWorkerCount(profile: DeviceProfile): number {
+  if (profile.battery && profile.battery.level !== null
+      && profile.battery.level < 0.2
+      && profile.battery.charging === false) {
+    return 1;
+  }
+  const half = Math.floor(profile.cores / 2);
+  return Math.max(1, Math.min(8, half));
+}
+
+/**
+ * Whether WebGPU is worth using for the current workload.
+ *
+ * Threshold is exposed so callers can tweak per index size.
+ */
+export function shouldUseGpu(profile: DeviceProfile, chunkCount: number, threshold = 20_000): boolean {
+  return profile.webgpu && chunkCount > threshold;
+}