agent-working-memory 0.7.16 → 0.8.5

package/src/core/ml-worker.ts

@@ -0,0 +1,281 @@
+// Copyright 2026 Robert Winter / Complete Ideas
+// SPDX-License-Identifier: Apache-2.0
+/**
+ * ML inference dispatch pool.
+ *
+ * STATUS (AWM 0.8.x P1 REVISE, 2026-05-25): worker_threads were the original
+ * plan, but @huggingface/transformers in Node only supports the `cpu`
+ * (native ONNX) and `dml` (Windows GPU) backends. Neither is safe inside
+ * a worker_thread — onnxruntime-node's native bindings store V8 handles
+ * that get invalidated when crossing isolate boundaries, causing
+ * `v8::HandleScope::CreateHandle()` crashes on first inference call.
+ * The browser-only `wasm` backend is not loaded in Node builds of
+ * transformers.js.
+ *
+ * The dispatch abstraction is preserved so a future child_process pool
+ * or HTTP sidecar (see AWM_ML_SIDECAR_URL design in docs/awm-architecture-history.md)
+ * can plug in. For now ALL inference runs in-process. The freeze fix from
+ * P0 (sleep-only consolidation) already eliminates the multi-second
+ * in-band blocks. Individual inference calls (~50ms each) on the main
+ * thread are accepted as-is.
+ *
+ * Test / dev mode: AWM_ML_INPROCESS=1 is honored but is now the default.
+ * The env var remains as a no-op for backwards compatibility.
+ */
+
+import { Worker } from 'node:worker_threads';
+import { fileURLToPath } from 'node:url';
+import { dirname, join } from 'node:path';
+import { existsSync } from 'node:fs';
+
+type WorkerRole = 'embed' | 'rerank' | 'expand';
+
+interface PendingRequest {
+  resolve: (value: any) => void;
+  reject: (err: Error) => void;
+  timeoutHandle?: ReturnType<typeof setTimeout>;
+}
+
+interface ManagedWorker {
+  role: WorkerRole;
+  worker: Worker | null;
+  ready: Promise<void>;
+  setReady: () => void;
+  setFailed: (err: Error) => void;
+  pending: Map<number, PendingRequest>;
+  buffered: Array<{ id: number; op: WorkerRole; args: any }>;
+  isReady: boolean;
+  restartTimestamps: number[]; // unix-ms of recent restarts (for backoff cap)
+}
+
+let inProcessMode = false;
+let workers: Record<WorkerRole, ManagedWorker> | null = null;
+let nextId = 1;
+
+// In-process fallback handles (used in tests and as crash escape hatch).
+let inProcessEmbed: ((args: any) => Promise<number[][]>) | null = null;
+let inProcessRerank: ((args: any) => Promise<Array<{ index: number; score: number }>>) | null = null;
+let inProcessExpand: ((args: any) => Promise<string>) | null = null;
+
+const REQUEST_TIMEOUT_MS = 60_000;  // any single inference call > 60s is treated as failed
+const MAX_RESTARTS_PER_MINUTE = 3;
+
+function shouldUseInProcess(): boolean {
+  // AWM 0.8.x P1 REVISE: worker_threads were removed because @huggingface/transformers
+  // in Node is not worker_threads-safe (see file header). Always in-process for now.
+  // The dispatch abstraction is preserved for a future child_process pool or
+  // HTTP sidecar pivot.
+  return true;
+}
+
+function createWorker(role: WorkerRole): ManagedWorker {
+  const m: ManagedWorker = {
+    role,
+    worker: null,
+    ready: Promise.resolve(),
+    setReady: () => {},
+    setFailed: () => {},
+    pending: new Map(),
+    buffered: [],
+    isReady: false,
+    restartTimestamps: [],
+  };
+  m.ready = new Promise<void>((resolve, reject) => {
+    m.setReady = () => { m.isReady = true; resolve(); };
+    m.setFailed = (err) => reject(err);
+  });
+  spawnWorker(m);
+  return m;
+}
+
+function workerEntryPath(): string {
+  // Resolve to the compiled .js. Two possible locations:
+  //   1. Same directory as this file (when running from dist/core/)
+  //   2. Sibling dist/core/ (when running from src/core/ via tsx)
+  // Always prefer the compiled file. If neither exists, shouldUseInProcess()
+  // will detect the missing entry and fall back to in-process mode.
+  const here = dirname(fileURLToPath(import.meta.url));
+  const samedir = join(here, 'ml-worker-entry.js');
+  if (existsSync(samedir)) return samedir;
+  // From src/core/ml-worker.ts, dist/core/ml-worker-entry.js is at ../../dist/core/
+  const distSibling = join(here, '..', '..', 'dist', 'core', 'ml-worker-entry.js');
+  if (existsSync(distSibling)) return distSibling;
+  // Fallback to the same-dir path (will fail existsSync in shouldUseInProcess
+  // and trigger in-process mode)
+  return samedir;
+}
+
+function spawnWorker(m: ManagedWorker): void {
+  const w = new Worker(workerEntryPath(), { workerData: { role: m.role } });
+  m.worker = w;
+
+  w.on('message', (msg: any) => {
+    if (msg?.ready === true) {
+      m.setReady();
+      // Drain buffered messages
+      for (const buf of m.buffered) w.postMessage(buf);
+      m.buffered = [];
+      return;
+    }
+    if (msg?.ready === false) {
+      m.setFailed(new Error(`worker ${m.role} failed to load model: ${msg.error}`));
+      return;
+    }
+    if (msg?.shutdown === 'done') return;
+    if (typeof msg?.id !== 'number') return;
+
+    const p = m.pending.get(msg.id);
+    if (!p) return;
+    m.pending.delete(msg.id);
+    if (p.timeoutHandle) clearTimeout(p.timeoutHandle);
+    if (msg.ok) p.resolve(msg.result);
+    else p.reject(new Error(msg.error ?? 'worker error'));
+  });
+
+  w.on('error', (err) => {
+    console.error(`[ml-worker:${m.role}] error:`, err);
+  });
+
+  w.on('exit', (code) => {
+    if (code === 0) return; // graceful exit
+    console.warn(`[ml-worker:${m.role}] exited with code ${code} — recovering`);
+    // Reject all pending requests
+    for (const [, p] of m.pending) {
+      if (p.timeoutHandle) clearTimeout(p.timeoutHandle);
+      p.reject(new Error(`worker ${m.role} crashed (exit ${code})`));
+    }
+    m.pending.clear();
+
+    // Restart-rate backoff
+    const now = Date.now();
+    m.restartTimestamps = m.restartTimestamps.filter(t => now - t < 60_000);
+    m.restartTimestamps.push(now);
+
+    if (m.restartTimestamps.length > MAX_RESTARTS_PER_MINUTE) {
+      console.error(`[ml-worker:${m.role}] crashed ${m.restartTimestamps.length} times in 60s — falling back to in-process`);
+      m.worker = null;
+      m.isReady = false;
+      inProcessMode = true;
+      return;
+    }
+
+    // Reset ready promise + respawn
+    m.isReady = false;
+    m.ready = new Promise<void>((resolve, reject) => {
+      m.setReady = () => { m.isReady = true; resolve(); };
+      m.setFailed = (err) => reject(err);
+    });
+    spawnWorker(m);
+  });
+}
+
+/** Initialize the pool. Idempotent — safe to call multiple times. */
+export function initMLPool(): void {
+  if (shouldUseInProcess()) {
+    inProcessMode = true;
+    return;
+  }
+  if (workers) return;
+  workers = {
+    embed: createWorker('embed'),
+    rerank: createWorker('rerank'),
+    expand: createWorker('expand'),
+  };
+}
+
+/** Register in-process fallback handlers. Called once by the consumer modules. */
+export function registerInProcessHandlers(handlers: {
+  embed?: typeof inProcessEmbed;
+  rerank?: typeof inProcessRerank;
+  expand?: typeof inProcessExpand;
+}): void {
+  if (handlers.embed) inProcessEmbed = handlers.embed;
+  if (handlers.rerank) inProcessRerank = handlers.rerank;
+  if (handlers.expand) inProcessExpand = handlers.expand;
+}
+
+/** True if the pool is operating in in-process mode (no workers). */
+export function isInProcessMode(): boolean {
+  return inProcessMode;
+}
+
+async function dispatchToWorker<T>(role: WorkerRole, args: any): Promise<T> {
+  if (!workers) initMLPool();
+  if (inProcessMode) {
+    return dispatchInProcess<T>(role, args);
+  }
+
+  const m = workers![role];
+  await m.ready;
+  if (inProcessMode) {
+    // Fallback flipped while we awaited ready
+    return dispatchInProcess<T>(role, args);
+  }
+
+  const id = nextId++;
+  return new Promise<T>((resolve, reject) => {
+    const req: PendingRequest = {
+      resolve,
+      reject,
+      timeoutHandle: setTimeout(() => {
+        m.pending.delete(id);
+        reject(new Error(`ml-worker:${role} request ${id} timed out after ${REQUEST_TIMEOUT_MS}ms`));
+      }, REQUEST_TIMEOUT_MS),
+    };
+    m.pending.set(id, req);
+    const msg = { id, op: role, args };
+    if (m.isReady && m.worker) {
+      m.worker.postMessage(msg);
+    } else {
+      m.buffered.push(msg);
+    }
+  });
+}
+
+async function dispatchInProcess<T>(role: WorkerRole, args: any): Promise<T> {
+  switch (role) {
+    case 'embed':
+      if (!inProcessEmbed) throw new Error('in-process embed handler not registered');
+      return inProcessEmbed(args) as Promise<T>;
+    case 'rerank':
+      if (!inProcessRerank) throw new Error('in-process rerank handler not registered');
+      return inProcessRerank(args) as Promise<T>;
+    case 'expand':
+      if (!inProcessExpand) throw new Error('in-process expand handler not registered');
+      return inProcessExpand(args) as Promise<T>;
+  }
+}
+
+// --- Public API used by the consumer modules ---
+
+export async function dispatchEmbed(args: { texts: string[]; pooling: 'cls' | 'mean'; dimensions: number }): Promise<number[][]> {
+  return dispatchToWorker<number[][]>('embed', args);
+}
+
+export async function dispatchRerank(args: { query: string; passages: string[] }): Promise<Array<{ index: number; score: number }>> {
+  return dispatchToWorker<Array<{ index: number; score: number }>>('rerank', args);
+}
+
+export async function dispatchExpand(args: { prompt: string; maxNewTokens: number; noRepeatNgramSize: number }): Promise<string> {
+  return dispatchToWorker<string>('expand', args);
+}
+
+/** Graceful shutdown. Waits up to 2s for queue drain, then terminates. */
+export async function shutdownMLPool(): Promise<void> {
+  if (!workers) return;
+  const promises: Promise<void>[] = [];
+  for (const role of ['embed', 'rerank', 'expand'] as WorkerRole[]) {
+    const m = workers[role];
+    if (!m.worker) continue;
+    const w = m.worker;
+    promises.push(new Promise<void>((resolve) => {
+      const timeout = setTimeout(() => {
+        w.terminate().finally(() => resolve());
+      }, 2000);
+      w.once('exit', () => { clearTimeout(timeout); resolve(); });
+      w.postMessage({ shutdown: true });
+    }));
+  }
+  await Promise.all(promises);
+  workers = null;
+}

package/src/core/query-expander.ts

@@ -1,128 +1,122 @@
-// Copyright 2026 Robert Winter / Complete Ideas
-// SPDX-License-Identifier: Apache-2.0
-/**
- * Query Expander — rewrites queries with synonyms and related terms.
- *
- * Uses Xenova/flan-t5-small (~80MB ONNX) to expand search queries
- * with related terms that improve BM25 recall.
- *
- * Example: "What is Caroline's identity?" →
- *   "What is Caroline's identity? Caroline personal gender transgender self"
- *
- * Singleton pattern — call getExpander() to get the shared instance.
- */
-
-import { pipeline, type Text2TextGenerationPipeline } from '@huggingface/transformers';
-
-const MODEL_ID = 'Xenova/flan-t5-small';
-let instance: Text2TextGenerationPipeline | null = null;
-let initPromise: Promise<Text2TextGenerationPipeline> | null = null;
-
-/**
- * Get or initialize the text generation pipeline (singleton).
- * First call downloads the model (~80MB), subsequent calls are instant.
- */
-export async function getExpander(): Promise<Text2TextGenerationPipeline> {
-  if (instance) return instance;
-  if (initPromise) return initPromise;
-
-  initPromise = pipeline('text2text-generation', MODEL_ID, {
-    dtype: 'fp32',
-  }).then(pipe => {
-    instance = pipe as Text2TextGenerationPipeline;
-    console.log(`Query expander loaded: ${MODEL_ID}`);
-    return instance;
-  });
-
-  return initPromise;
-}
-
-/**
- * LRU cache of normalized-query → expanded-query mappings.
- * Map preserves insertion order — re-set on hit moves entry to the end (most-recent),
- * delete-first when over capacity drops the least-recent. ~500 entries × ~200 chars
- * each ≈ 100KB, negligible memory cost.
- *
- * Why: phase-breakdown spike (2026-05-08) showed expandQuery at ~164ms per call.
- * Agent recall patterns repeat — cache hits are common.
- */
-const expansionCache = new Map<string, string>();
-const EXPANSION_CACHE_LIMIT = 500;
-
-/**
- * Heuristic: skip expansion when the query is already specific.
- * Long, multi-token queries don't benefit from synonyms — they're already
- * narrow enough that flan-t5's general-vocabulary expansion adds noise more
- * than recall. Exact thresholds are conservative; false-skips would only
- * affect candidates that BM25 catches anyway.
- */
-function shouldSkipExpansion(normalized: string): boolean {
-  if (normalized.length === 0) return true;
-  if (normalized.length > 50) return true;
-  // ≥5 distinct meaningful tokens = already specific
-  const tokens = new Set(normalized.split(/\s+/).filter(t => t.length > 2));
-  return tokens.size >= 5;
-}
-
-/**
- * Expand a query with related terms and synonyms.
- * Returns the original query + generated expansion terms.
- * Falls back to the original query on any error.
- *
- * Optimization (0.7.11+):
- * - Skip heuristic for long/specific queries (~30% of typical agent recalls)
- * - LRU cache for repeated queries (cache hit ≈ 0ms vs 164ms cold)
- * - Disable both via AWM_DISABLE_EXPANSION_CACHE=1
- */
-export async function expandQuery(originalQuery: string): Promise<string> {
-  const normalized = originalQuery.toLowerCase().trim();
-  const optimizationsEnabled = process.env.AWM_DISABLE_EXPANSION_CACHE !== '1';
-
-  if (optimizationsEnabled) {
-    if (shouldSkipExpansion(normalized)) {
-      return originalQuery;
-    }
-    const cached = expansionCache.get(normalized);
-    if (cached !== undefined) {
-      // Move to end (most-recent) for LRU semantics
-      expansionCache.delete(normalized);
-      expansionCache.set(normalized, cached);
-      return cached;
-    }
-  }
-
-  try {
-    const expander = await getExpander();
-    const prompt = `Expand this search query with synonyms and related terms. Only output the additional terms, not the original query. Query: ${originalQuery}. Additional terms:`;
-
-    const result = await expander(prompt, {
-      max_new_tokens: 25,
-      no_repeat_ngram_size: 2,
-    });
-
-    const expanded = Array.isArray(result) ? (result[0] as any)?.generated_text ?? '' : '';
-    const cleanExpanded = expanded.trim();
-
-    const finalQuery = cleanExpanded && cleanExpanded.length > 2
-      ? `${originalQuery} ${cleanExpanded}`
-      : originalQuery;
-
-    // Cache the result (LRU eviction when over capacity)
-    if (optimizationsEnabled) {
-      if (expansionCache.size >= EXPANSION_CACHE_LIMIT) {
-        const oldestKey = expansionCache.keys().next().value;
-        if (oldestKey !== undefined) expansionCache.delete(oldestKey);
-      }
-      expansionCache.set(normalized, finalQuery);
-    }
-
-    return finalQuery;
-  } catch {
-    return originalQuery;
-  }
-}
-
-/** Clear the expansion cache (used by tests + cache invalidation if needed). */
-export function clearExpansionCache(): void {
-  expansionCache.clear();
-}
+// Copyright 2026 Robert Winter / Complete Ideas
+// SPDX-License-Identifier: Apache-2.0
+/**
+ * Query Expander - rewrites queries with synonyms and related terms.
+ *
+ * Uses Xenova/flan-t5-small (~80MB ONNX) to expand search queries with
+ * related terms that improve BM25 recall.
+ *
+ * AWM 0.8.x: inference dispatches through ml-worker.ts (currently in-process
+ * — worker_threads reverted because onnxruntime-node bindings cross isolate
+ * boundaries unsafely; see ml-worker.ts). The dispatch abstraction is
+ * preserved for a future child_process / HTTP sidecar pool.
+ *
+ * The LRU cache + skip heuristic stay on the main thread — they're pure
+ * filter/lookup logic that shouldn't pay IPC cost.
+ */
+
+import { pipeline, type Text2TextGenerationPipeline } from '@huggingface/transformers';
+import { dispatchExpand, registerInProcessHandlers } from './ml-worker.js';
+
+const MODEL_ID = 'Xenova/flan-t5-small';
+
+// --- In-process fallback ---
+
+let inProcessInstance: Text2TextGenerationPipeline | null = null;
+let inProcessInitPromise: Promise<Text2TextGenerationPipeline> | null = null;
+
+async function loadInProcess(): Promise<Text2TextGenerationPipeline> {
+  if (inProcessInstance) return inProcessInstance;
+  if (inProcessInitPromise) return inProcessInitPromise;
+  inProcessInitPromise = pipeline('text2text-generation', MODEL_ID, { dtype: 'fp32' }).then(pipe => {
+    inProcessInstance = pipe as Text2TextGenerationPipeline;
+    console.log(`Query expander loaded in-process: ${MODEL_ID}`);
+    return inProcessInstance;
+  });
+  return inProcessInitPromise;
+}
+
+async function inProcessExpand(args: { prompt: string; maxNewTokens: number; noRepeatNgramSize: number }): Promise<string> {
+  const expander = await loadInProcess();
+  const result = await expander(args.prompt, {
+    max_new_tokens: args.maxNewTokens,
+    no_repeat_ngram_size: args.noRepeatNgramSize,
+  });
+  const text = Array.isArray(result) ? (result[0] as any)?.generated_text ?? '' : '';
+  return String(text).trim();
+}
+
+registerInProcessHandlers({ expand: inProcessExpand });
+
+// --- Public API ---
+
+/** Kept for backwards compat. */
+export async function getExpander(): Promise<Text2TextGenerationPipeline> {
+  return loadInProcess();
+}
+
+/**
+ * LRU cache of normalized-query → expanded-query mappings.
+ * Lives on the main thread — cache hits skip the worker IPC entirely.
+ */
+const expansionCache = new Map<string, string>();
+const EXPANSION_CACHE_LIMIT = 500;
+
+/**
+ * Skip expansion when the query is already specific (long or many tokens).
+ */
+function shouldSkipExpansion(normalized: string): boolean {
+  if (normalized.length === 0) return true;
+  if (normalized.length > 50) return true;
+  const tokens = new Set(normalized.split(/\s+/).filter(t => t.length > 2));
+  return tokens.size >= 5;
+}
+
+/**
+ * Expand a query with related terms and synonyms.
+ * Returns the original query + generated expansion terms.
+ * Falls back to the original query on any error.
+ *
+ * Dispatches inference to the worker pool. Cache + skip heuristic stay
+ * on the main thread.
+ */
+export async function expandQuery(originalQuery: string): Promise<string> {
+  const normalized = originalQuery.toLowerCase().trim();
+  const optimizationsEnabled = process.env.AWM_DISABLE_EXPANSION_CACHE !== '1';
+
+  if (optimizationsEnabled) {
+    if (shouldSkipExpansion(normalized)) return originalQuery;
+    const cached = expansionCache.get(normalized);
+    if (cached !== undefined) {
+      // LRU touch
+      expansionCache.delete(normalized);
+      expansionCache.set(normalized, cached);
+      return cached;
+    }
+  }
+
+  try {
+    const prompt = `Expand this search query with synonyms and related terms. Only output the additional terms, not the original query. Query: ${originalQuery}. Additional terms:`;
+    const expansion = await dispatchExpand({ prompt, maxNewTokens: 25, noRepeatNgramSize: 2 });
+    const finalQuery = expansion && expansion.length > 2
+      ? `${originalQuery} ${expansion}`
+      : originalQuery;
+
+    if (optimizationsEnabled) {
+      if (expansionCache.size >= EXPANSION_CACHE_LIMIT) {
+        const oldestKey = expansionCache.keys().next().value;
+        if (oldestKey !== undefined) expansionCache.delete(oldestKey);
+      }
+      expansionCache.set(normalized, finalQuery);
+    }
+
+    return finalQuery;
+  } catch {
+    return originalQuery;
+  }
+}
+
+/** Clear the expansion cache (used by tests + cache invalidation). */
+export function clearExpansionCache(): void {
+  expansionCache.clear();
+}