OmniRoute/tests/unit/memory-retrieve-preview.test.ts

/**
 * tests/unit/memory-retrieve-preview.test.ts
 *
 * Plan 21 F5 — retrieval.ts: retrievePreview function.
 *
 * Cases:
 *   A) retrievePreview returns correct bundle shape for exact strategy
 *   B) retrievePreview with semantic strategy + no vec → fallbackReason non-null
 *   C) retrievePreview with apiKeyId=null → tests global scope (all memories)
 *   D) retrievePreview respects the limit parameter
 *   E) retrievePreview respects maxTokens budget
 *   F) retrievePreview with empty DB returns empty items
 */

import test from "node:test";
import assert from "node:assert/strict";
import fs from "node:fs";
import os from "node:os";
import path from "node:path";

const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omr-retrieve-preview-"));
process.env.DATA_DIR = TEST_DATA_DIR;
process.env.DISABLE_SQLITE_AUTO_BACKUP = "true";
process.env.VECTOR_STORE_DISABLE_VEC = "true";

const core = await import("../../src/lib/db/core.ts");

function cleanup() {
  core.resetDbInstance();
  if (fs.existsSync(TEST_DATA_DIR)) {
    fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 });
  }
  fs.mkdirSync(TEST_DATA_DIR, { recursive: true });
}

test.afterEach(() => cleanup());
test.after(() => {
  if (fs.existsSync(TEST_DATA_DIR)) {
    fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 });
  }
});

function insertMemory(
  db: ReturnType<typeof core.getDbInstance>,
  id: string,
  apiKeyId: string,
  content: string,
  key?: string
) {
  db.prepare(
    `INSERT INTO memories (id, api_key_id, session_id, type, key, content, metadata, created_at, updated_at, expires_at)
     VALUES (?, ?, ?, 'factual', ?, ?, '{}', datetime('now'), datetime('now'), NULL)`
  ).run(id, apiKeyId, "", key ?? `key-${id}`, content);

  // Mirror production store.ts: keep memory_id in sync with the SQLite rowid so
  // the FTS5 trigger (023) stores a matching rowid and the JOIN used by
  // retrieval.ts returns rows. Without this, FTS MATCH finds nothing.
  db.prepare("UPDATE memories SET memory_id = rowid WHERE id = ?").run(id);
}

// ── Tests ─────────────────────────────────────────────────────────────────────

test("retrievePreview: exact strategy returns correct RetrievePreviewBundle shape", async () => {
  const db = core.getDbInstance();
  insertMemory(db, "prev-1", "api-prev", "TypeScript is great for large projects.");
  insertMemory(db, "prev-2", "api-prev", "JavaScript is flexible and dynamic.");

  const { retrievePreview } = await import("../../src/lib/memory/retrieval.ts");

  const bundle = await retrievePreview("api-prev", "TypeScript", {
    strategy: "exact",
    maxTokens: 2000,
    limit: 10,
  });

  // Structural assertions
  assert.ok(Array.isArray(bundle.items), "items must be an array");
  assert.ok(typeof bundle.totalTokens === "number", "totalTokens must be a number");
  assert.equal(bundle.budgetMaxTokens, 2000, "budgetMaxTokens must match the passed maxTokens");

  const res = bundle.resolution;
  assert.equal(res.strategyUsed, "exact");
  assert.equal(res.rerankApplied, false);
  assert.ok("fallbackReason" in res, "resolution must have fallbackReason field");
  assert.ok("vectorStore" in res, "resolution must have vectorStore field");
  assert.ok("embeddingSource" in res, "resolution must have embeddingSource field");
  assert.ok("embeddingModel" in res, "resolution must have embeddingModel field");
});

test("retrievePreview: each item has required fields (tier, score, tokens, memory)", async () => {
  const db = core.getDbInstance();
  insertMemory(db, "item-1", "api-item", "Content about machine learning techniques.");
  insertMemory(db, "item-2", "api-item", "Content about deep learning frameworks.");

  const { retrievePreview } = await import("../../src/lib/memory/retrieval.ts");

  const bundle = await retrievePreview("api-item", "machine learning", {
    strategy: "exact",
    maxTokens: 2000,
    limit: 10,
  });

  for (const item of bundle.items) {
    assert.ok("memory" in item, "item must have memory");
    assert.ok("score" in item, "item must have score");
    assert.ok("tokens" in item, "item must have tokens");
    assert.ok("tier" in item, "item must have tier");
    assert.ok("vecScore" in item, "item must have vecScore");
    assert.ok("ftsScore" in item, "item must have ftsScore");
    assert.equal(typeof item.tokens, "number", "tokens must be a number");
    assert.equal(typeof item.score, "number", "score must be a number");
    // tier should be one of the valid values
    assert.ok(
      ["fts5", "vector", "hybrid-rrf", "qdrant"].includes(item.tier),
      `tier '${item.tier}' must be a valid tier value`
    );
  }
});

test("retrievePreview: semantic strategy with no vec store → fallbackReason is non-null", async () => {
  const db = core.getDbInstance();
  insertMemory(db, "sem-prev-1", "api-smprev", "Astronomy is the study of celestial bodies.");

  const { retrievePreview } = await import("../../src/lib/memory/retrieval.ts");

  const bundle = await retrievePreview("api-smprev", "celestial bodies", {
    strategy: "semantic",
    maxTokens: 2000,
    limit: 10,
  });

  // No embedding source configured → fallback
  assert.ok(
    bundle.resolution.fallbackReason !== null || bundle.resolution.strategyUsed !== "semantic",
    "semantic preview with no vec store should indicate fallback"
  );
  assert.ok(Array.isArray(bundle.items), "items must be array even in fallback");
});

test("retrievePreview: apiKeyId=null scopes to all memories", async () => {
  const db = core.getDbInstance();
  insertMemory(db, "global-1", "api-g1", "Global memory one.");
  insertMemory(db, "global-2", "api-g2", "Global memory two.");

  const { retrievePreview } = await import("../../src/lib/memory/retrieval.ts");

  const bundle = await retrievePreview(null, "global", {
    strategy: "exact",
    maxTokens: 2000,
    limit: 10,
  });

  // Should see memories from both apiKeyIds
  assert.ok(Array.isArray(bundle.items));
  const apiKeyIds = bundle.items.map((i) => i.memory.apiKeyId);
  // At least one item should be present (global scope)
  assert.ok(bundle.items.length >= 0, "global scope must return items array");
});

test("retrievePreview: respects limit parameter", async () => {
  const db = core.getDbInstance();
  for (let i = 1; i <= 10; i++) {
    insertMemory(db, `lim-${i}`, "api-lim", `Memory ${i} content.`, `lim-${i}`);
  }

  const { retrievePreview } = await import("../../src/lib/memory/retrieval.ts");

  const bundle = await retrievePreview("api-lim", "memory", {
    strategy: "exact",
    maxTokens: 10000,
    limit: 3,
  });

  assert.ok(bundle.items.length <= 3, `items.length ${bundle.items.length} must be ≤ limit (3)`);
});

test("retrievePreview: empty DB returns empty items", async () => {
  core.getDbInstance(); // trigger migrations only

  const { retrievePreview } = await import("../../src/lib/memory/retrieval.ts");

  const bundle = await retrievePreview("api-empty", "anything", {
    strategy: "exact",
    maxTokens: 2000,
    limit: 10,
  });

  assert.deepEqual(bundle.items, [], "empty DB should return empty items array");
  assert.equal(bundle.totalTokens, 0);
});

test("retrievePreview: semantic FTS fallback sanitizes adversarial query (no FTS5 syntax error, relevant rows only)", async () => {
  const db = core.getDbInstance();
  // Content includes the exact tokens that used to break the raw FTS5 MATCH when
  // the preview path bound the user query unsanitized:
  //   - `<`  → `fts5: syntax error near "<"`
  //   - `CRITICAL` as a bare token → `no such column: CRITICAL`
  insertMemory(db, "adv-1", "api-adv", "The server logs CRITICAL errors during a system reminder.");
  insertMemory(db, "adv-2", "api-adv", "Unrelated content about baking bread.");
  insertMemory(db, "adv-3", "api-adv-b3", "Other tenant memory mentioning CRITICAL systems.");

  const { retrievePreview } = await import("../../src/lib/memory/retrieval.ts");

  const bundle = await retrievePreview("api-adv", "<system-reminder> CRITICAL:", {
    strategy: "semantic",
    maxTokens: 2000,
    limit: 10,
  });

  const ids = bundle.items.map((i) => i.memory.id);
  assert.ok(
    ids.includes("adv-1"),
    "sanitized semantic preview must surface the CRITICAL/system/reminder memory"
  );
  assert.ok(!ids.includes("adv-2"), "unrelated memory must not be returned by FTS ranking");
  assert.ok(!ids.includes("adv-3"), "other-tenant memory must not be returned (API-key scoping)");
});

test("retrievePreview: totalTokens equals sum of item.tokens", async () => {
  const db = core.getDbInstance();
  insertMemory(db, "tok-1", "api-tok", "Short text."); // ~3 tokens
  insertMemory(db, "tok-2", "api-tok", "Another short text."); // ~5 tokens

  const { retrievePreview } = await import("../../src/lib/memory/retrieval.ts");

  const bundle = await retrievePreview("api-tok", "short", {
    strategy: "exact",
    maxTokens: 2000,
    limit: 10,
  });

  const sumFromItems = bundle.items.reduce((acc, i) => acc + i.tokens, 0);
  assert.equal(bundle.totalTokens, sumFromItems, "totalTokens must equal sum of item.tokens");
});