@vellumai/assistant 0.7.2 → 0.8.0

package/src/memory/rerank-local.ts

@@ -0,0 +1,374 @@
+/** Local cross-encoder rerank backend — drives the rerank-worker subprocess. */
+import { existsSync } from "node:fs";
+
+import type { RerankDtype } from "../config/schemas/memory-v2.js";
+import { getLogger } from "../util/logger.js";
+import { getEmbeddingModelsDir } from "../util/platform.js";
+import { PromiseGuard } from "../util/promise-guard.js";
+import { EmbeddingRuntimeManager } from "./embedding-runtime-manager.js";
+
+const log = getLogger("memory-rerank-local");
+
+interface WorkerResponse {
+  id?: number;
+  type?: string;
+  scores?: number[];
+  error?: string;
+}
+
+export class LocalRerankBackend {
+  readonly model: string;
+  readonly dtype: RerankDtype;
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  private workerProc: any = null;
+  private stdoutBuffer = "";
+  private requestCounter = 0;
+  private pendingRequests = new Map<
+    number,
+    { resolve: (response: WorkerResponse) => void }
+  >();
+  private stdoutReaderActive = false;
+  private activeRequests = 0;
+  private disposeRequested = false;
+
+  private readyResolve: (() => void) | null = null;
+  private readyReject: ((err: Error) => void) | null = null;
+
+  private readonly initGuard = new PromiseGuard<void>();
+
+  constructor(model: string, dtype: RerankDtype) {
+    this.model = model;
+    this.dtype = dtype;
+  }
+
+  /**
+   * Score paired `(queries[i], passages[i])` tuples in one batched ONNX
+   * inference call. Multiple distinct queries can ride in a single batch
+   * so callers can score the user-channel and assistant-channel queries
+   * against a shared candidate set in one tokenizer + forward pass.
+   */
+  async score(queries: string[], passages: string[]): Promise<number[]> {
+    if (this.disposeRequested) {
+      throw new Error("Local rerank backend is shutting down");
+    }
+    if (passages.length === 0) return [];
+    if (queries.length !== passages.length) {
+      throw new Error(
+        `Rerank backend got ${queries.length} queries for ${passages.length} passages`,
+      );
+    }
+
+    this.activeRequests++;
+    try {
+      await this.ensureInitialized();
+      const response = await this.sendRequest({ queries, passages });
+      if (response.error) {
+        throw new Error(`Rerank worker error: ${response.error}`);
+      }
+      if (!response.scores) {
+        throw new Error("Rerank worker returned no scores");
+      }
+      if (response.scores.length !== passages.length) {
+        throw new Error(
+          `Rerank worker returned ${response.scores.length} scores for ${passages.length} passages`,
+        );
+      }
+      return response.scores;
+    } finally {
+      this.activeRequests--;
+      this.disposeIfIdle();
+    }
+  }
+
+  dispose(): void {
+    this.disposeRequested = true;
+    this.disposeIfIdle();
+  }
+
+  private sendRequest(payload: {
+    queries: string[];
+    passages: string[];
+  }): Promise<WorkerResponse> {
+    const id = ++this.requestCounter;
+    return new Promise((resolve) => {
+      if (!this.workerProc) {
+        resolve({ error: "Worker not initialized" });
+        return;
+      }
+      this.pendingRequests.set(id, { resolve });
+      this.workerProc.stdin.write(JSON.stringify({ id, ...payload }) + "\n");
+      try {
+        this.workerProc.stdin.flush();
+      } catch {
+        // Worker may have exited — stdout reader cleanup resolves pending requests.
+      }
+    });
+  }
+
+  private async ensureInitialized(): Promise<void> {
+    if (this.workerProc) return;
+    await this.initGuard.run(() => this.initialize());
+  }
+
+  private async initialize(): Promise<void> {
+    log.info({ model: this.model }, "Initializing local rerank backend");
+
+    const runtimeManager = new EmbeddingRuntimeManager();
+    if (!runtimeManager.isReady()) {
+      log.info("Embedding runtime not yet available, waiting for download...");
+      await runtimeManager.ensureInstalled();
+    }
+
+    const bunPath = runtimeManager.getBunPath();
+    const workerPath = runtimeManager.getRerankWorkerPath();
+
+    if (!bunPath) {
+      throw new Error("Local rerank backend unavailable: no bun binary found");
+    }
+    if (!existsSync(workerPath)) {
+      throw new Error(
+        `Local rerank backend unavailable: worker script not found at ${workerPath}`,
+      );
+    }
+
+    await this.startWorker(bunPath, workerPath);
+  }
+
+  private async startWorker(
+    bunPath: string,
+    workerPath: string,
+  ): Promise<void> {
+    const embeddingModelsDir = getEmbeddingModelsDir();
+    const modelCacheDir = `${embeddingModelsDir}/model-cache`;
+
+    log.info(
+      { bunPath, workerPath, model: this.model, dtype: this.dtype },
+      "Spawning rerank worker process",
+    );
+
+    const proc = Bun.spawn({
+      cmd: [
+        bunPath,
+        "--smol",
+        workerPath,
+        this.model,
+        modelCacheDir,
+        this.dtype,
+      ],
+      stdin: "pipe",
+      stdout: "pipe",
+      stderr: "pipe",
+      cwd: embeddingModelsDir,
+    });
+
+    this.workerProc = proc;
+    this.startStdoutReader();
+
+    try {
+      await this.waitForReady();
+    } catch (err) {
+      this.workerProc = null;
+      this.stdoutReaderActive = false;
+      try {
+        proc.kill();
+      } catch {
+        /* may already be dead */
+      }
+      const exitCode = await proc.exited.catch(() => undefined);
+      const stderr = await new Response(proc.stderr).text().catch(() => "");
+      if (stderr.trim()) {
+        log.warn({ stderr: stderr.trim(), exitCode }, "Rerank worker stderr");
+      }
+      throw new Error(
+        `Rerank worker exited (code ${exitCode ?? "unknown"}): ${
+          stderr.trim() || (err instanceof Error ? err.message : String(err))
+        }`,
+      );
+    }
+
+    this.drainStderr(proc.stderr);
+    log.info(
+      { pid: proc.pid, model: this.model },
+      "Rerank worker process started",
+    );
+    this.disposeIfIdle();
+  }
+
+  private drainStderr(stderr: ReadableStream<Uint8Array>): void {
+    const reader = stderr.getReader();
+    const decoder = new TextDecoder();
+    (async () => {
+      try {
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          const text = decoder.decode(value, { stream: true }).trim();
+          if (text) log.debug({ workerStderr: text }, "Rerank worker stderr");
+        }
+      } catch {
+        /* expected on shutdown */
+      }
+    })();
+  }
+
+  private startStdoutReader(): void {
+    if (this.stdoutReaderActive || !this.workerProc) return;
+    this.stdoutReaderActive = true;
+
+    const proc = this.workerProc;
+    const reader = proc.stdout.getReader();
+    const decoder = new TextDecoder();
+
+    (async () => {
+      try {
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          this.stdoutBuffer += decoder.decode(value, { stream: true });
+          this.processStdoutBuffer();
+        }
+      } catch {
+        /* reader cancelled or stream errored */
+      }
+
+      if (this.workerProc === proc) {
+        for (const pending of this.pendingRequests.values()) {
+          pending.resolve({
+            error: "Rerank worker process exited unexpectedly",
+          });
+        }
+        this.pendingRequests.clear();
+        this.workerProc = null;
+        this.stdoutReaderActive = false;
+        this.stdoutBuffer = "";
+        this.initGuard.reset();
+      }
+    })();
+  }
+
+  private processStdoutBuffer(): void {
+    let idx: number;
+    while ((idx = this.stdoutBuffer.indexOf("\n")) !== -1) {
+      const line = this.stdoutBuffer.slice(0, idx);
+      this.stdoutBuffer = this.stdoutBuffer.slice(idx + 1);
+      if (!line.trim()) continue;
+
+      let msg: WorkerResponse;
+      try {
+        msg = JSON.parse(line);
+      } catch {
+        continue;
+      }
+
+      if (msg.type === "ready") {
+        this.readyResolve?.();
+        this.readyResolve = null;
+        this.readyReject = null;
+        continue;
+      }
+      if (msg.type === "error" && this.readyReject) {
+        this.readyReject(
+          new Error(msg.error ?? "Worker initialization failed"),
+        );
+        this.readyResolve = null;
+        this.readyReject = null;
+        continue;
+      }
+
+      if (msg.id !== undefined) {
+        const pending = this.pendingRequests.get(msg.id);
+        if (pending) {
+          this.pendingRequests.delete(msg.id);
+          pending.resolve(msg);
+          this.disposeIfIdle();
+        }
+      }
+    }
+  }
+
+  private waitForReady(): Promise<void> {
+    return new Promise<void>((resolve, reject) => {
+      // First-call timeout. Generous because the first run downloads the
+      // ONNX weights (~280 MB to ~1 GB depending on model) before loading.
+      const timeout = setTimeout(() => {
+        this.readyResolve = null;
+        this.readyReject = null;
+        reject(new Error("Rerank worker timed out waiting for model to load"));
+      }, 120_000);
+
+      this.readyResolve = () => {
+        clearTimeout(timeout);
+        resolve();
+      };
+      this.readyReject = (err: Error) => {
+        clearTimeout(timeout);
+        reject(err);
+      };
+
+      this.workerProc?.exited.then(() => {
+        if (this.readyResolve) {
+          clearTimeout(timeout);
+          this.readyResolve = null;
+          this.readyReject = null;
+          reject(
+            new Error("Rerank worker process exited before becoming ready"),
+          );
+        }
+      });
+    });
+  }
+
+  private disposeIfIdle(): void {
+    if (!this.disposeRequested) return;
+    if (this.activeRequests > 0) return;
+    if (this.pendingRequests.size > 0) return;
+    if (this.readyResolve || this.readyReject) return;
+
+    const proc = this.workerProc;
+    this.workerProc = null;
+    this.stdoutReaderActive = false;
+    this.stdoutBuffer = "";
+    this.initGuard.reset();
+
+    if (!proc) return;
+
+    try {
+      proc.kill();
+    } catch {
+      /* may already be exiting */
+    }
+  }
+}
+
+// ── Module-level singleton management ─────────────────────────────────
+
+let _backend: LocalRerankBackend | null = null;
+
+export function getOrCreateRerankBackend(
+  model: string,
+  dtype: RerankDtype,
+): LocalRerankBackend {
+  if (_backend?.model === model && _backend.dtype === dtype) return _backend;
+  if (_backend) {
+    try {
+      _backend.dispose();
+    } catch {
+      /* best effort */
+    }
+  }
+  _backend = new LocalRerankBackend(model, dtype);
+  return _backend;
+}
+
+/** @internal Test-only: reset the cached backend. */
+export function _resetRerankBackendForTests(): void {
+  if (_backend) {
+    try {
+      _backend.dispose();
+    } catch {
+      /* best effort */
+    }
+  }
+  _backend = null;
+}

package/src/memory/search/semantic.ts

@@ -1,7 +1,6 @@
 import { inArray } from "drizzle-orm";
 
 import { getConfig } from "../../config/loader.js";
-import { isMemoryV2ReadActive } from "../context-search/sources/memory-v2.js";
 import { getDb } from "../db-connection.js";
 import { withQdrantBreaker } from "../qdrant-circuit-breaker.js";
 import type {
@@ -9,7 +8,7 @@ import type {
   QdrantSparseVector,
 } from "../qdrant-client.js";
 import { getQdrantClient } from "../qdrant-client.js";
-import { conversations, memorySegments, memorySummaries } from "../schema.js";
+import { memorySegments, memorySummaries } from "../schema.js";
 // ── Types (inlined from deleted types.ts) ──────────────────────────
 
 type CandidateType = "segment" | "item" | "summary" | "media";
@@ -52,15 +51,14 @@ export async function semanticSearch(
   _model: string,
   limit: number,
   excludedMessageIds: string[] = [],
-  scopeIds?: string[],
   sparseVector?: QdrantSparseVector,
 ): Promise<Candidate[]> {
   if (limit <= 0) return [];
 
-  // v2 owns the read path when both gates are on; the v1 `memory` collection
-  // is in active retirement, and routing semantic recall there would re-enter
-  // the same corrupted sparse segments that can OOM-crash Qdrant.
-  if (isMemoryV2ReadActive(getConfig())) return [];
+  // v2 owns the read path when enabled; the v1 `memory` collection is in
+  // active retirement, and routing semantic recall there would re-enter the
+  // same corrupted sparse segments that can OOM-crash Qdrant.
+  if (getConfig().memory.v2.enabled) return [];
 
   const qdrant = getQdrantClient();
 
@@ -75,7 +73,7 @@ export async function semanticSearch(
   let isHybrid = false;
   if (sparseVector && sparseVector.indices.length > 0) {
     isHybrid = true;
-    const filter = buildHybridFilter(excludedMessageIds, scopeIds);
+    const filter = buildHybridFilter(excludedMessageIds);
     results = await withQdrantBreaker(() =>
       qdrant.hybridSearch({
         denseVector: queryVector,
@@ -92,7 +90,6 @@ export async function semanticSearch(
         fetchLimit,
         ["summary", "segment", "media"],
         excludedMessageIds,
-        scopeIds,
       ),
     );
   }
@@ -102,18 +99,15 @@ export async function semanticSearch(
   // Batch-fetch all backing records upfront to avoid N+1 queries per result
   const summaryTargetIds: string[] = [];
   const segmentTargetIds: string[] = [];
-  const mediaConversationIds: string[] = [];
   for (const r of results) {
     if (r.payload.target_type === "summary")
       summaryTargetIds.push(r.payload.target_id);
     else if (r.payload.target_type === "segment")
       segmentTargetIds.push(r.payload.target_id);
-    else if (r.payload.target_type === "media" && r.payload.conversation_id)
-      mediaConversationIds.push(r.payload.conversation_id);
   }
 
   const summariesMap = new Map<string, typeof memorySummaries.$inferSelect>();
-  if (scopeIds && summaryTargetIds.length > 0) {
+  if (summaryTargetIds.length > 0) {
     const allSummaries = db
       .select()
       .from(memorySummaries)
@@ -123,7 +117,7 @@ export async function semanticSearch(
   }
 
   const segmentsMap = new Map<string, typeof memorySegments.$inferSelect>();
-  if (scopeIds && segmentTargetIds.length > 0) {
+  if (segmentTargetIds.length > 0) {
     const allSegments = db
       .select()
       .from(memorySegments)
@@ -132,23 +126,6 @@ export async function semanticSearch(
     for (const seg of allSegments) segmentsMap.set(seg.id, seg);
   }
 
-  // Batch-fetch conversation scope IDs for media results to avoid N+1 queries.
-  // When a conversation is not found (deleted), its media is excluded rather than
-  // falling back to "default" scope, which would leak private media.
-  const mediaScopeMap = new Map<string, string>();
-  if (scopeIds && mediaConversationIds.length > 0) {
-    const unique = [...new Set(mediaConversationIds)];
-    const rows = db
-      .select({
-        id: conversations.id,
-        memoryScopeId: conversations.memoryScopeId,
-      })
-      .from(conversations)
-      .where(inArray(conversations.id, unique))
-      .all();
-    for (const row of rows) mediaScopeMap.set(row.id, row.memoryScopeId);
-  }
-
   const candidates: Candidate[] = [];
   for (const result of results) {
     const { payload, score } = result;
@@ -160,10 +137,7 @@ export async function semanticSearch(
       // Legacy item vectors — skip (table dropped, Qdrant cleanup pending)
       continue;
     } else if (payload.target_type === "summary") {
-      if (scopeIds) {
-        const summary = summariesMap.get(payload.target_id);
-        if (!summary || !scopeIds.includes(summary.scopeId)) continue;
-      }
+      if (!summariesMap.has(payload.target_id)) continue;
       candidates.push({
         key: `summary:${payload.target_id}`,
         type: "summary",
@@ -180,22 +154,6 @@ export async function semanticSearch(
         finalScore: 0,
       });
     } else if (payload.target_type === "media") {
-      // Use stored memory_scope_id when available; fall back to deriving
-      // scope from conversation_id for legacy media points.
-      // If the conversation was deleted, skip the media to avoid leaking
-      // private media into the default scope.
-      if (scopeIds) {
-        let mediaScopeId: string | undefined;
-        if (payload.memory_scope_id) {
-          mediaScopeId = payload.memory_scope_id;
-        } else if (payload.conversation_id) {
-          mediaScopeId = mediaScopeMap.get(payload.conversation_id);
-          if (!mediaScopeId) continue; // conversation deleted — skip
-        } else {
-          mediaScopeId = "default";
-        }
-        if (!scopeIds.includes(mediaScopeId)) continue;
-      }
       candidates.push({
         key: `media:${payload.target_id}`,
         type: "media",
@@ -212,10 +170,7 @@ export async function semanticSearch(
         finalScore: 0,
       });
     } else {
-      if (scopeIds) {
-        const segment = segmentsMap.get(payload.target_id);
-        if (!segment || !scopeIds.includes(segment.scopeId)) continue;
-      }
+      if (!segmentsMap.has(payload.target_id)) continue;
       candidates.push({
         key: `segment:${payload.target_id}`,
         type: "segment",
@@ -254,14 +209,9 @@ export async function semanticSearch(
 /**
  * Build a Qdrant filter for hybrid search. Mirrors the logic in
  * `searchWithFilter` but as a standalone object for the query API.
- *
- * Scope filtering: points with a `memory_scope_id` payload field are
- * filtered at the Qdrant level. Legacy points without the field pass
- * through and are caught by post-query DB filtering.
  */
 function buildHybridFilter(
   excludeMessageIds: string[],
-  scopeIds?: string[],
 ): Record<string, unknown> {
   const mustConditions: Array<Record<string, unknown>> = [
     {
@@ -270,18 +220,6 @@ function buildHybridFilter(
     },
   ];
 
-  // Scope filtering: accept points whose memory_scope_id matches one of the
-  // allowed scopes, OR points that lack the field entirely (legacy data).
-  // Post-query DB filtering remains as defense-in-depth for legacy points.
-  if (scopeIds && scopeIds.length > 0) {
-    mustConditions.push({
-      should: [
-        { key: "memory_scope_id", match: { any: scopeIds } },
-        { is_empty: { key: "memory_scope_id" } },
-      ],
-    });
-  }
-
   const mustNotConditions: Array<Record<string, unknown>> = [
     { key: "_meta", match: { value: true } },
   ];

package/src/memory/trace-event-store.ts

@@ -1,11 +1,10 @@
-import { and, asc, eq, gt, lt, sql } from "drizzle-orm";
+import { and, asc, eq, gt, sql } from "drizzle-orm";
 
 import type {
   TraceEvent,
   TraceEventKind,
 } from "../daemon/message-types/messages.js";
 import { getDb } from "./db-connection.js";
-import { rawChanges } from "./raw-query.js";
 import { traceEvents } from "./schema.js";
 
 // ---------------------------------------------------------------------------
@@ -115,21 +114,6 @@ export function getTraceEvents(
   return rows.map(rowToTraceEventRow);
 }
 
-// ---------------------------------------------------------------------------
-// Cleanup
-// ---------------------------------------------------------------------------
-
-/**
- * Delete trace events older than `maxAgeDays` based on `created_at`.
- * Returns the count of deleted rows.
- */
-export function deleteOldTraceEvents(maxAgeDays: number): number {
-  const db = getDb();
-  const cutoff = Date.now() - maxAgeDays * 24 * 60 * 60 * 1000;
-  db.delete(traceEvents).where(lt(traceEvents.createdAt, cutoff)).run();
-  return rawChanges();
-}
-
 // ---------------------------------------------------------------------------
 // Sequence
 // ---------------------------------------------------------------------------