@vellumai/assistant 0.10.4 → 0.10.5-dev.202607022329.d89bf2f

package/src/persistence/job-handlers/__tests__/message-lexical.test.ts

@@ -0,0 +1,259 @@
+import { beforeAll, beforeEach, describe, expect, mock, test } from "bun:test";
+
+import type { SparseEmbedding } from "../../embeddings/embedding-types.js";
+
+mock.module("../../../util/logger.js", () => ({
+  getLogger: () =>
+    new Proxy({} as Record<string, unknown>, {
+      get: () => () => {},
+    }),
+}));
+
+// Capture calls into the lexical index singleton.
+const upsertCalls: Array<{
+  messageId: string;
+  sparse: SparseEmbedding;
+  payload: { conversationId: string; createdAt: number };
+}> = [];
+const deleteByConversationCalls: string[] = [];
+
+const fakeIndex = {
+  upsertMessage: async (
+    messageId: string,
+    sparse: SparseEmbedding,
+    payload: { conversationId: string; createdAt: number },
+  ) => {
+    upsertCalls.push({ messageId, sparse, payload });
+  },
+  deleteByConversation: async (conversationId: string) => {
+    deleteByConversationCalls.push(conversationId);
+  },
+};
+
+// Model the process-local singleton so the worker scenario (init never ran in
+// this process) is exercisable: getMessagesLexicalIndex throws until
+// initMessagesLexicalIndex sets it, mirroring the real module's behavior.
+let singletonReady = true;
+const initConfigCalls: Array<{
+  url: string;
+  collection?: string;
+  onDisk?: boolean;
+}> = [];
+
+function resetLexicalSingleton(ready: boolean): void {
+  singletonReady = ready;
+  initConfigCalls.length = 0;
+}
+
+mock.module("../../embeddings/messages-lexical-index.js", () => ({
+  MESSAGES_LEXICAL_COLLECTION: "messages_lexical",
+  getMessagesLexicalIndex: () => {
+    if (!singletonReady) {
+      throw new Error("Messages lexical index not initialized.");
+    }
+    return fakeIndex;
+  },
+  initMessagesLexicalIndex: (config: {
+    url: string;
+    collection?: string;
+    onDisk?: boolean;
+  }) => {
+    initConfigCalls.push(config);
+    singletonReady = true;
+    return fakeIndex;
+  },
+}));
+
+// `generateSparseEmbedding` is a pure local TF-IDF encoder (no provider call),
+// so it runs unmocked — mocking `embedding-backend.js` wholesale would starve
+// its other named exports that the db-init import graph pulls in.
+import { resetDbForTesting } from "../../../__tests__/db-test-helpers.js";
+import { DEFAULT_CONFIG } from "../../../config/defaults.js";
+import type { AssistantConfig } from "../../../config/types.js";
+import { getDb } from "../../db-connection.js";
+import { initializeDb } from "../../db-init.js";
+import { generateSparseEmbedding } from "../../embeddings/embedding-backend.js";
+import type { MemoryJob } from "../../jobs-store.js";
+import { conversations, messages } from "../../schema/index.js";
+import {
+  indexMessageLexicalJob,
+  purgeConversationLexicalJob,
+} from "../message-lexical.js";
+
+const TEST_CONFIG: AssistantConfig = DEFAULT_CONFIG;
+
+function makeJob(
+  type: MemoryJob["type"],
+  payload: Record<string, unknown>,
+): MemoryJob {
+  return {
+    id: "job-1",
+    type,
+    payload,
+    status: "running",
+    attempts: 0,
+    deferrals: 0,
+    runAfter: 0,
+    lastError: null,
+    startedAt: Date.now(),
+    createdAt: Date.now(),
+    updatedAt: Date.now(),
+  };
+}
+
+function insertConversation(id: string): void {
+  const now = Date.now();
+  getDb()
+    .insert(conversations)
+    .values({ id, createdAt: now, updatedAt: now })
+    .run();
+}
+
+function insertMessage(opts: {
+  id: string;
+  conversationId: string;
+  content: string;
+  createdAt: number;
+}): void {
+  getDb()
+    .insert(messages)
+    .values({
+      id: opts.id,
+      conversationId: opts.conversationId,
+      role: "user",
+      content: opts.content,
+      createdAt: opts.createdAt,
+    })
+    .run();
+}
+
+describe("indexMessageLexicalJob", () => {
+  // initializeDb runs the full migration chain; under parallel CI load it can
+  // exceed bun's default 5s hook timeout, so allow more.
+  beforeAll(async () => {
+    await initializeDb();
+  }, 30_000);
+
+  beforeEach(async () => {
+    upsertCalls.length = 0;
+    deleteByConversationCalls.length = 0;
+    resetLexicalSingleton(true);
+    resetDbForTesting();
+    await initializeDb();
+  }, 30_000);
+
+  test("loads the message and upserts it into the lexical index", async () => {
+    const createdAt = 1_700_000_000_000;
+    const content = "hello lexical world";
+    insertConversation("conv-1");
+    insertMessage({
+      id: "msg-1",
+      conversationId: "conv-1",
+      content,
+      createdAt,
+    });
+
+    await indexMessageLexicalJob(
+      makeJob("index_message_lexical", { messageId: "msg-1" }),
+      TEST_CONFIG,
+    );
+
+    expect(upsertCalls).toHaveLength(1);
+    expect(upsertCalls[0].messageId).toBe("msg-1");
+    // The handler encodes the message content via the local sparse encoder.
+    expect(upsertCalls[0].sparse).toEqual(generateSparseEmbedding(content));
+    expect(upsertCalls[0].payload).toEqual({
+      conversationId: "conv-1",
+      createdAt,
+    });
+  });
+
+  test("no-ops when messageId is missing from the payload", async () => {
+    await indexMessageLexicalJob(
+      makeJob("index_message_lexical", {}),
+      TEST_CONFIG,
+    );
+    expect(upsertCalls).toHaveLength(0);
+  });
+
+  test("no-ops when the message no longer exists", async () => {
+    await indexMessageLexicalJob(
+      makeJob("index_message_lexical", { messageId: "nonexistent" }),
+      TEST_CONFIG,
+    );
+    expect(upsertCalls).toHaveLength(0);
+  });
+
+  test("lazily initializes the index from config when not yet initialized (worker process)", async () => {
+    // Simulate the out-of-process worker, which never runs runMemoryStartup —
+    // the singleton is not initialized in this process.
+    resetLexicalSingleton(false);
+
+    const createdAt = 1_700_000_000_001;
+    insertConversation("conv-w");
+    insertMessage({
+      id: "msg-w",
+      conversationId: "conv-w",
+      content: "indexed by the worker",
+      createdAt,
+    });
+
+    await indexMessageLexicalJob(
+      makeJob("index_message_lexical", { messageId: "msg-w" }),
+      TEST_CONFIG,
+    );
+
+    // The handler initialized the index from config instead of throwing.
+    expect(initConfigCalls).toHaveLength(1);
+    expect(initConfigCalls[0]).toEqual({
+      url: TEST_CONFIG.memory.qdrant.url,
+      collection: "messages_lexical",
+      onDisk: TEST_CONFIG.memory.qdrant.onDisk,
+    });
+    // ...and proceeded to index the message.
+    expect(upsertCalls).toHaveLength(1);
+    expect(upsertCalls[0].messageId).toBe("msg-w");
+  });
+});
+
+describe("purgeConversationLexicalJob", () => {
+  beforeAll(async () => {
+    await initializeDb();
+  }, 30_000);
+
+  beforeEach(async () => {
+    upsertCalls.length = 0;
+    deleteByConversationCalls.length = 0;
+    resetLexicalSingleton(true);
+    resetDbForTesting();
+    await initializeDb();
+  }, 30_000);
+
+  test("deletes the conversation's points from the lexical index", async () => {
+    await purgeConversationLexicalJob(
+      makeJob("purge_conversation_lexical", { conversationId: "conv-9" }),
+      TEST_CONFIG,
+    );
+    expect(deleteByConversationCalls).toEqual(["conv-9"]);
+  });
+
+  test("no-ops when conversationId is missing from the payload", async () => {
+    await purgeConversationLexicalJob(
+      makeJob("purge_conversation_lexical", {}),
+      TEST_CONFIG,
+    );
+    expect(deleteByConversationCalls).toHaveLength(0);
+  });
+
+  test("lazily initializes the index from config when not yet initialized (worker process)", async () => {
+    resetLexicalSingleton(false);
+
+    await purgeConversationLexicalJob(
+      makeJob("purge_conversation_lexical", { conversationId: "conv-w" }),
+      TEST_CONFIG,
+    );
+
+    expect(initConfigCalls).toHaveLength(1);
+    expect(deleteByConversationCalls).toEqual(["conv-w"]);
+  });
+});

package/src/persistence/job-handlers/cleanup.ts

@@ -1,4 +1,5 @@
 import type { AssistantConfig } from "../../config/types.js";
+import { rotateToolInvocations } from "../../telemetry/tool-usage-store.js";
 import { getLogger } from "../../util/logger.js";
 import { getLogsDbPath } from "../../util/logs-db-path.js";
 import { runAsyncSqlite } from "../db-async-query.js";
@@ -139,6 +140,32 @@ SELECT changes();`,
   );
 }
 
+/**
+ * Delete audit-log (`tool_invocations`) entries older than the configured
+ * retention window. Retention comes from `auditLog.retentionDays` (0 = keep
+ * forever). The DELETE itself lives in {@link rotateToolInvocations}; this
+ * handler just resolves the window and dispatches, mirroring the other
+ * scheduled cleanup jobs.
+ */
+export async function pruneOldToolInvocationsJob(
+  job: MemoryJob,
+  config: AssistantConfig,
+): Promise<void> {
+  const rawRetention = job.payload.retentionDays;
+  const retentionDays =
+    typeof rawRetention === "number" &&
+    Number.isFinite(rawRetention) &&
+    rawRetention >= 0
+      ? rawRetention
+      : config.auditLog.retentionDays;
+
+  // 0 means disabled. rotateToolInvocations no-ops on this too, but
+  // short-circuit so we don't dispatch a pointless async query.
+  if (!retentionDays) return;
+
+  await rotateToolInvocations(retentionDays);
+}
+
 /**
  * Parse the `SELECT changes()` result emitted by the sqlite3 CLI after
  * the prune DELETE. Returns 0 if stdout is missing or unparseable —
@@ -170,10 +197,9 @@ function parseDeletedCount(stdout: string | undefined): number {
  *
  * Tables with onDelete cascade on conversation FK (memory_segments,
  * conversation_keys, channel_inbound_events, message_runs, call_sessions,
- * external_conversation_bindings, assistant_inbox_conversation_state) are handled
- * automatically. Tables without cascade (messages, tool_invocations,
- * llm_request_logs, skill_loaded_events) are deleted explicitly before
- * removing the conversation row.
+ * external_conversation_bindings) are handled automatically. Tables without
+ * cascade (messages, tool_invocations, llm_request_logs, skill_loaded_events)
+ * are deleted explicitly before removing the conversation row.
  */
 export function pruneOldConversationsJob(
   job: MemoryJob,
@@ -192,6 +218,7 @@ export function pruneOldConversationsJob(
   const cutoffMs = Date.now() - retentionDays * 86_400_000;
 
   const stale = rawAll<{ id: string }>(
+    "cleanup:pruneOldConversations:stale",
     `SELECT id FROM conversations WHERE updated_at < ? ORDER BY updated_at ASC LIMIT ?`,
     cutoffMs,
     PRUNE_BATCH_LIMIT,
@@ -205,6 +232,7 @@ export function pruneOldConversationsJob(
       // Re-check staleness inside the transaction to avoid racing with a conversation
       // that became active again between the initial SELECT and this DELETE.
       const still = rawAll<{ id: string }>(
+        "cleanup:pruneOldConversations:recheck",
         `SELECT id FROM conversations WHERE id = ? AND updated_at < ?`,
         id,
         cutoffMs,
@@ -213,12 +241,32 @@ export function pruneOldConversationsJob(
 
       // Non-cascading tables. llm_request_logs lives in the dedicated logs
       // connection, so it is deleted there (outside this main-DB transaction).
-      rawLogsRun(`DELETE FROM llm_request_logs WHERE conversation_id = ?`, id);
-      rawRun(`DELETE FROM tool_invocations WHERE conversation_id = ?`, id);
-      rawRun(`DELETE FROM messages WHERE conversation_id = ?`, id);
-      rawRun(`DELETE FROM skill_loaded_events WHERE conversation_id = ?`, id);
+      rawLogsRun(
+        "cleanup:pruneOldConversations:logs",
+        `DELETE FROM llm_request_logs WHERE conversation_id = ?`,
+        id,
+      );
+      rawRun(
+        "cleanup:pruneOldConversations:toolInv",
+        `DELETE FROM tool_invocations WHERE conversation_id = ?`,
+        id,
+      );
+      rawRun(
+        "cleanup:pruneOldConversations:messages",
+        `DELETE FROM messages WHERE conversation_id = ?`,
+        id,
+      );
+      rawRun(
+        "cleanup:pruneOldConversations:skills",
+        `DELETE FROM skill_loaded_events WHERE conversation_id = ?`,
+        id,
+      );
       // Conversation row deletion cascades to remaining dependent tables
-      rawRun(`DELETE FROM conversations WHERE id = ?`, id);
+      rawRun(
+        "cleanup:pruneOldConversations:conv",
+        `DELETE FROM conversations WHERE id = ?`,
+        id,
+      );
       pruned++;
     });
   }

package/src/persistence/job-handlers/message-lexical-backfill.ts

@@ -0,0 +1,187 @@
+import { and, asc, eq, gt, or } from "drizzle-orm";
+
+import type { AssistantConfig } from "../../config/types.js";
+import { getLogger } from "../../util/logger.js";
+import {
+  clearLexicalBackfillComplete,
+  isLexicalBackfillComplete,
+  LEXICAL_BACKFILL_COMPLETE_KEY,
+  readMessageCursorCheckpoint,
+  resetMessageCursorCheckpoint,
+  setMemoryCheckpoint,
+  writeMessageCursorCheckpoint,
+} from "../checkpoints.js";
+import { getDb } from "../db-connection.js";
+import { generateSparseEmbedding } from "../embeddings/embedding-backend.js";
+import { withQdrantBreaker } from "../embeddings/qdrant-circuit-breaker.js";
+import {
+  enqueueMemoryJob,
+  hasActiveJobOfType,
+  type MemoryJob,
+} from "../jobs-store.js";
+import { messages } from "../schema/index.js";
+import { resolveLexicalIndex } from "./message-lexical.js";
+
+const log = getLogger("lexical-backfill");
+
+/**
+ * Cursor checkpoint keys for the lexical-index backfill. Kept distinct from the
+ * dense embedding backfill's keys (`memory:backfill:*`) so the two backfills
+ * advance independently — enabling lexical without re-running the dense one, and
+ * vice versa.
+ */
+const LEXICAL_BACKFILL_CHECKPOINT_KEY = "lexical:messages:last_created_at";
+const LEXICAL_BACKFILL_CHECKPOINT_ID_KEY = "lexical:messages:last_id";
+
+/**
+ * One-time, self-healing auto-enqueue of the messages lexical-index backfill,
+ * invoked once from daemon startup (`daemon/lifecycle.ts`). Ensures each
+ * instance populates its Qdrant lexical index (`messages_lexical`) exactly once
+ * on upgrade — in the background — so message-content search does not open
+ * onto an empty index.
+ *
+ * Deliberate, narrow exception to the "not run at daemon startup" note carried by
+ * the manual backfill route (`messages-lexical-routes.ts`, added in the FTS→Qdrant
+ * PR 5): this call only *enqueues* a memory job — it does no indexing, no LLM
+ * work, and no other heavy lifting on the boot thread. The backfill itself runs
+ * off the event loop via the existing memory job worker, in small
+ * cursor-checkpointed batches. It is a one-time, checkpoint-guarded, no-LLM data
+ * backfill (migration-like), so it is consistent with the daemon-startup
+ * philosophy of never blocking or doing expensive work at boot.
+ *
+ * Guards (all must pass to enqueue):
+ * - The completion sentinel must be unset. Once the backfill has fully drained
+ *   on this instance there is nothing to do; the marker makes this idempotent
+ *   across restarts.
+ * - No backfill job may already be pending or running. Without this, every
+ *   restart during a long backfill would pile up a duplicate job. (The batches
+ *   are idempotent, but redundant jobs waste worker cycles.)
+ *
+ * Unconditional beyond those guards: message-content search is host
+ * infrastructure, so the backfill runs regardless of the memory feature's or
+ * memory plugin's state (the job worker drains the lexical types even while
+ * memory is disabled).
+ *
+ * An instance that was already manually backfilled via the route but predates
+ * the completion sentinel re-enqueues once here and re-runs the idempotent
+ * backfill (no duplicate points; ~minutes) — an acceptable one-time cost.
+ *
+ * Best-effort and lightweight: any failure (e.g. the memory database is briefly
+ * unavailable at boot) is swallowed and logged so a memory hiccup never blocks
+ * daemon startup, matching the surrounding startup steps.
+ */
+export function maybeEnqueueLexicalBackfillOnUpgrade(): void {
+  try {
+    if (isLexicalBackfillComplete()) {
+      return;
+    }
+    if (hasActiveJobOfType("backfill_lexical_index")) {
+      return;
+    }
+    const jobId = enqueueMemoryJob("backfill_lexical_index", {});
+    log.info(
+      { jobId },
+      "Enqueued one-time messages lexical-index backfill on startup",
+    );
+  } catch (err) {
+    log.warn(
+      { err },
+      "Failed to enqueue startup lexical-index backfill (non-fatal)",
+    );
+  }
+}
+
+/** Messages indexed per job invocation before the job re-enqueues itself. */
+const BATCH_SIZE = 200;
+
+/**
+ * Resumable, cursor-checkpointed backfill of existing messages into the Qdrant
+ * lexical index. Selects the next {@link BATCH_SIZE} messages ordered by
+ * `(createdAt asc, id asc)` after the persisted cursor, encodes each with the
+ * local sparse encoder, batch-upserts them, advances the cursor, and
+ * re-enqueues itself when the batch is full — draining the whole `messages`
+ * table across many short invocations without blocking the event loop.
+ *
+ * Idempotent: the upsert keys on a deterministic point id derived from the
+ * message id, so re-indexing the same message overwrites the same point. Pass
+ * `payload.force` to reset the cursor and re-index from the beginning.
+ */
+export async function backfillLexicalIndexJob(
+  job: MemoryJob,
+  config: AssistantConfig,
+): Promise<void> {
+  const db = getDb();
+  const force = job.payload.force === true;
+  if (force) {
+    resetMessageCursorCheckpoint(
+      LEXICAL_BACKFILL_CHECKPOINT_KEY,
+      LEXICAL_BACKFILL_CHECKPOINT_ID_KEY,
+    );
+    // Clearing the completion sentinel lets a forced re-run drain and re-mark
+    // as complete, and re-arms the one-time startup auto-enqueue for this run.
+    // The route also clears this at enqueue time (so reads fall back to FTS
+    // before the worker claims the job); clearing again here is idempotent and
+    // keeps the forced-run semantics self-contained in the handler.
+    clearLexicalBackfillComplete();
+  }
+
+  const cursor = readMessageCursorCheckpoint(
+    LEXICAL_BACKFILL_CHECKPOINT_KEY,
+    LEXICAL_BACKFILL_CHECKPOINT_ID_KEY,
+  );
+  const batch = db
+    .select({
+      id: messages.id,
+      conversationId: messages.conversationId,
+      content: messages.content,
+      createdAt: messages.createdAt,
+    })
+    .from(messages)
+    .where(
+      or(
+        gt(messages.createdAt, cursor.createdAt),
+        and(
+          eq(messages.createdAt, cursor.createdAt),
+          gt(messages.id, cursor.messageId),
+        ),
+      ),
+    )
+    .orderBy(asc(messages.createdAt), asc(messages.id))
+    .limit(BATCH_SIZE)
+    .all();
+
+  if (batch.length > 0) {
+    const index = resolveLexicalIndex(config);
+    const points = batch.map((message) => ({
+      messageId: message.id,
+      sparse: generateSparseEmbedding(message.content),
+      conversationId: message.conversationId,
+      createdAt: message.createdAt,
+    }));
+    await withQdrantBreaker(() => index.upsertMessagesBatch(points));
+
+    const lastMessage = batch[batch.length - 1];
+    writeMessageCursorCheckpoint(
+      LEXICAL_BACKFILL_CHECKPOINT_KEY,
+      LEXICAL_BACKFILL_CHECKPOINT_ID_KEY,
+      {
+        createdAt: lastMessage.createdAt,
+        messageId: lastMessage.id,
+      },
+    );
+  }
+
+  // A full batch means there may be more rows past the new cursor — re-enqueue
+  // to continue draining. Never carry `force` forward: the reset already
+  // happened this invocation, and re-resetting would loop over the same head.
+  if (batch.length === BATCH_SIZE) {
+    enqueueMemoryJob("backfill_lexical_index", {});
+  } else {
+    // A short batch (including an empty one) means the cursor has reached the
+    // end of the `messages` table — the backfill has fully drained. Record the
+    // completion sentinel so the one-time startup auto-enqueue never schedules
+    // it again on this instance. Idempotent: re-running after completion draws
+    // an empty batch and simply re-sets the same marker.
+    setMemoryCheckpoint(LEXICAL_BACKFILL_COMPLETE_KEY, "1");
+  }
+}