gbrain/test/query-cache.test.ts

/**
 * v0.32.x search-lite \u2014 semantic query cache.
 *
 * PGLite-backed test. Confirms:
 *   - migration v51 creates the query_cache table
 *   - store + lookup roundtrip with EXACT same embedding \u2192 hit
 *   - lookup with a similar embedding (cosine > 0.92) \u2192 hit
 *   - lookup with a far embedding \u2192 miss
 *   - TTL expiration: a stale row is skipped at read time
 *   - clear / prune / stats work as advertised
 *   - source_id isolation: brain A's cache doesn't leak to brain B
 *   - disabled cache is a pure no-op
 *
 * Uses synthetic Float32Array embeddings so the test doesn't depend on
 * any external embedding provider.
 */

import { describe, test, expect, beforeAll, afterAll, beforeEach } from 'bun:test';
import { PGLiteEngine } from '../src/core/pglite-engine.ts';
import { SemanticQueryCache, cacheRowId } from '../src/core/search/query-cache.ts';
import { configureGateway, resetGateway } from '../src/core/ai/gateway.ts';
import type { SearchResult, HybridSearchMeta } from '../src/core/types.ts';

let engine: PGLiteEngine;
let visiblePageId: number;

// Build a stable, normalized embedding. PGLite ships pgvector with 1536-dim
// support (the default); a smaller test dim won't match the column. We
// truncate / pad to 1536 to match the migration's resolved dim.
const DIM = 1536;

function makeEmbedding(seed: number, dim = DIM): Float32Array {
  const e = new Float32Array(dim);
  // Simple deterministic generator with a unique fingerprint per seed
  // so similar seeds produce similar (cosine > 0.95) vectors and distinct
  // seeds produce orthogonal-ish ones.
  for (let i = 0; i < dim; i++) {
    e[i] = Math.sin(seed * 0.001 + i * 0.01);
  }
  // L2-normalize so cosine = dot product.
  let mag = 0;
  for (let i = 0; i < dim; i++) mag += e[i] * e[i];
  mag = Math.sqrt(mag);
  if (mag > 0) for (let i = 0; i < dim; i++) e[i] /= mag;
  return e;
}

function makeOrthogonalEmbedding(seed: number, dim = DIM): Float32Array {
  // Use a totally different basis so cosine is near-zero.
  const e = new Float32Array(dim);
  for (let i = 0; i < dim; i++) {
    e[i] = Math.cos(seed * 13.7 + i * 0.97);
  }
  let mag = 0;
  for (let i = 0; i < dim; i++) mag += e[i] * e[i];
  mag = Math.sqrt(mag);
  if (mag > 0) for (let i = 0; i < dim; i++) e[i] /= mag;
  return e;
}

function makeResult(slug: string): SearchResult {
  return {
    slug,
    page_id: visiblePageId,
    title: `Title for ${slug}`,
    type: 'concept',
    chunk_text: `chunk text for ${slug}`,
    chunk_source: 'compiled_truth',
    chunk_id: 1,
    chunk_index: 0,
    score: 1.0,
    stale: false,
  };
}

const META: HybridSearchMeta = {
  vector_enabled: true,
  detail_resolved: 'medium',
  expansion_applied: false,
  intent: 'general',
};

type StoreOpts = NonNullable<Parameters<SemanticQueryCache['store']>[4]>;

async function storeCurrent(
  cache: SemanticQueryCache,
  queryText: string,
  queryEmbedding: Float32Array,
  results: SearchResult[],
  meta: HybridSearchMeta,
  opts: StoreOpts = {},
): Promise<void> {
  const rows = await engine.executeRaw<{ v: number }>(
    `SELECT COALESCE((SELECT last_value FROM page_generation_clock_seq), 0)::bigint AS v`,
  );
  await cache.store(queryText, queryEmbedding, results, meta, {
    ...opts,
    maxGenerationAtSearchStart:
      opts.maxGenerationAtSearchStart ?? Number(rows[0]?.v ?? 0),
  });
}

beforeAll(async () => {
  // v0.36.2.0: DEFAULT_EMBEDDING_DIMENSIONS flipped to 1280 (ZE Matryoshka).
  // This test hardcodes DIM=1536 in its embeddings. If another test file in
  // the same shard configured the gateway before us, initSchema() would size
  // query_cache.embedding at vector(1280) and every insert below would fail
  // with "expected 1280 dimensions, not 1536". Pin the gateway to 1536d
  // explicitly so this file is hermetic regardless of cross-file state.
  resetGateway();
  configureGateway({
    embedding_model: 'openai:text-embedding-3-large',
    embedding_dimensions: 1536,
    env: { OPENAI_API_KEY: 'sk-fake' },
  });
  engine = new PGLiteEngine();
  await engine.connect({});
  await engine.initSchema();
  const page = await engine.putPage('cache/visible-fixture', {
    type: 'note',
    title: 'Visible cache fixture',
    compiled_truth: 'visible cache fixture',
    timeline: '',
    frontmatter: {},
  });
  visiblePageId = page.id;
});

afterAll(async () => {
  try { await engine.disconnect(); } catch { /* ignore */ }
  resetGateway();
});

beforeEach(async () => {
  // Wipe the cache between tests so ordering doesn't matter.
  await engine.executeRaw(`DELETE FROM query_cache`);
});

describe('migration v51 \u2014 query_cache table exists', () => {
  test('table is present and has expected columns', async () => {
    const rows = await engine.executeRaw<{ column_name: string }>(
      `SELECT column_name FROM information_schema.columns
        WHERE table_name = 'query_cache'`,
    );
    const names = rows.map(r => r.column_name);
    expect(names).toContain('id');
    expect(names).toContain('query_text');
    expect(names).toContain('source_id');
    expect(names).toContain('embedding');
    expect(names).toContain('results');
    expect(names).toContain('meta');
    expect(names).toContain('ttl_seconds');
    expect(names).toContain('created_at');
    expect(names).toContain('hit_count');
  });
});

describe('cacheRowId', () => {
  test('is deterministic across same input', () => {
    expect(cacheRowId('hello', 'default')).toBe(cacheRowId('hello', 'default'));
  });
  test('differs across source_id', () => {
    expect(cacheRowId('hello', 'a')).not.toBe(cacheRowId('hello', 'b'));
  });
});

describe('SemanticQueryCache \u2014 store + lookup', () => {
  test('roundtrip: exact embedding match returns a hit', async () => {
    const cache = new SemanticQueryCache(engine);
    const emb = makeEmbedding(1);
    const results = [makeResult('a'), makeResult('b')];

    await storeCurrent(cache, 'what is foo', emb, results, META);
    const hit = await cache.lookup(emb);

    expect(hit.hit).toBe(true);
    expect(hit.results).toHaveLength(2);
    expect(hit.results?.[0].slug).toBe('a');
    expect(hit.similarity).toBeGreaterThan(0.99);
  });

  test('similar embedding (cosine > 0.92) is a hit', async () => {
    const cache = new SemanticQueryCache(engine);
    const base = makeEmbedding(100);

    // Construct a near-neighbor: tweak a few dims so cosine stays > 0.92.
    const near = new Float32Array(base);
    for (let i = 0; i < 10; i++) near[i] += 0.005;
    // Re-normalize.
    let mag = 0;
    for (let i = 0; i < DIM; i++) mag += near[i] * near[i];
    mag = Math.sqrt(mag);
    for (let i = 0; i < DIM; i++) near[i] /= mag;

    await storeCurrent(cache, 'what is foo', base, [makeResult('a')], META);
    const hit = await cache.lookup(near);

    expect(hit.hit).toBe(true);
    expect(hit.similarity).toBeGreaterThan(0.92);
  });

  test('orthogonal embedding is a miss', async () => {
    const cache = new SemanticQueryCache(engine);
    const a = makeEmbedding(1);
    const b = makeOrthogonalEmbedding(2);
    await storeCurrent(cache, 'q1', a, [makeResult('a')], META);
    const hit = await cache.lookup(b);
    expect(hit.hit).toBe(false);
  });
});

describe('SemanticQueryCache \u2014 TTL', () => {
  test('stale row (past TTL) is not returned', async () => {
    const cache = new SemanticQueryCache(engine, { ttlSeconds: 1 });
    const emb = makeEmbedding(42);
    await storeCurrent(cache, 'q', emb, [makeResult('a')], META, { ttlSeconds: 1 });

    // Manually rewind created_at to simulate expiration.
    await engine.executeRaw(
      `UPDATE query_cache SET created_at = now() - interval '10 seconds'`,
    );
    const hit = await cache.lookup(emb);
    expect(hit.hit).toBe(false);
  });
});

describe('SemanticQueryCache — source isolation', () => {
  test('different source_id cannot read each other’s rows', async () => {
    const cache = new SemanticQueryCache(engine);
    const emb = makeEmbedding(7);
    await storeCurrent(cache, 'q', emb, [makeResult('a')], META, { sourceId: 'src-A' });
    const hitB = await cache.lookup(emb, { sourceId: 'src-B' });
    expect(hitB.hit).toBe(false);
    const hitA = await cache.lookup(emb, { sourceId: 'src-A' });
    expect(hitA.hit).toBe(true);
  });

  test('archiving a source invalidates its cached result without a page write', async () => {
    const sourceId = 'cache-archive-src';
    await engine.executeRaw(
      `INSERT INTO sources (id, name, archived)
       VALUES ($1, $1, false)
       ON CONFLICT (id) DO UPDATE SET archived = false`,
      [sourceId],
    );
    const page = await engine.putPage(
      'cache/archive-visibility',
      {
        type: 'note',
        title: 'Archive visibility',
        compiled_truth: 'archive visibility cache canary',
        timeline: '',
        frontmatter: {},
      },
      { sourceId },
    );
    const result = { ...makeResult(page.slug), page_id: page.id, source_id: sourceId };
    const cache = new SemanticQueryCache(engine);
    const emb = makeEmbedding(71);

    await storeCurrent(cache, 'archive visibility', emb, [result], META, { sourceId });
    expect((await cache.lookup(emb, { sourceId })).hit).toBe(true);

    await engine.executeRaw(`UPDATE sources SET archived = true WHERE id = $1`, [sourceId]);
    try {
      expect((await cache.lookup(emb, { sourceId })).hit).toBe(false);
    } finally {
      await engine.executeRaw(`UPDATE sources SET archived = false WHERE id = $1`, [sourceId]);
    }
  });

  test('a result hard-deleted before cache snapshot cannot be served', async () => {
    const page = await engine.putPage('cache/deleted-before-store', {
      type: 'note',
      title: 'Deleted before store',
      compiled_truth: 'private cache race canary',
      timeline: '',
      frontmatter: {},
    });
    const result = { ...makeResult(page.slug), page_id: page.id, source_id: 'default' };
    await engine.executeRaw(`DELETE FROM pages WHERE id = $1`, [page.id]);

    const cache = new SemanticQueryCache(engine);
    const emb = makeEmbedding(72);
    await storeCurrent(cache, 'deleted before store', emb, [result], META);

    expect((await cache.lookup(emb)).hit).toBe(false);
  });

  test('writeback without a pre-search generation fails closed', async () => {
    const cache = new SemanticQueryCache(engine);
    const emb = makeEmbedding(77);

    await cache.store('missing generation', emb, [makeResult('a')], META);

    expect((await cache.lookup(emb)).hit).toBe(false);
    expect((await cache.stats()).total_rows).toBe(0);
  });

  test('a concurrent page mutation prevents stale result writeback', async () => {
    const clock = await engine.executeRaw<{ v: number }>(
      `SELECT COALESCE((SELECT last_value FROM page_generation_clock_seq), 0)::bigint AS v`,
    );
    const generationAtSearchStart = Number(clock[0]?.v ?? 0);
    const staleResult = { ...makeResult('cache/visible-fixture'), page_id: visiblePageId };
    await engine.executeRaw(
      `UPDATE pages SET compiled_truth = 'redacted after search' WHERE id = $1`,
      [visiblePageId],
    );

    const cache = new SemanticQueryCache(engine);
    const emb = makeEmbedding(73);
    await storeCurrent(cache, 'concurrent mutation', emb, [staleResult], META, {
      maxGenerationAtSearchStart: generationAtSearchStart,
    });

    expect((await cache.lookup(emb)).hit).toBe(false);
  });
});

describe('SemanticQueryCache \u2014 management', () => {
  test('clear() wipes all rows', async () => {
    const cache = new SemanticQueryCache(engine);
    const emb = makeEmbedding(9);
    await storeCurrent(cache, 'q1', emb, [makeResult('a')], META);
    await storeCurrent(cache, 'q2', makeEmbedding(10), [makeResult('b')], META);
    const removed = await cache.clear();
    expect(removed).toBeGreaterThanOrEqual(2);
    const stats = await cache.stats();
    expect(stats.total_rows).toBe(0);
  });

  test('source-scoped clear removes typed scalar and federated scopes containing the source', async () => {
    const cache = new SemanticQueryCache(engine);
    const emb = makeEmbedding(74);
    await storeCurrent(cache, 'scalar-a', emb, [makeResult('a')], META, { sourceId: '["scalar","source-a"]' });
    await storeCurrent(cache, 'set-ab', makeEmbedding(75), [makeResult('ab')], META, { sourceId: '["set","source-a","source-b"]' });
    await storeCurrent(cache, 'scalar-b', makeEmbedding(76), [makeResult('b')], META, { sourceId: '["scalar","source-b"]' });

    expect(await cache.clear({ sourceId: 'source-a' })).toBe(2);
    const rows = await engine.executeRaw<{ source_id: string }>(`SELECT source_id FROM query_cache ORDER BY source_id`);
    expect(rows.map(r => r.source_id)).toEqual(['["scalar","source-b"]']);
  });

  test('prune() deletes only stale rows', async () => {
    const cache = new SemanticQueryCache(engine);
    await storeCurrent(cache, 'fresh', makeEmbedding(11), [makeResult('a')], META);
    await storeCurrent(cache, 'stale', makeEmbedding(12), [makeResult('b')], META, { ttlSeconds: 1 });
    await engine.executeRaw(
      `UPDATE query_cache SET created_at = now() - interval '10 seconds' WHERE query_text = 'stale'`,
    );
    const removed = await cache.prune();
    expect(removed).toBe(1);
    const stats = await cache.stats();
    expect(stats.total_rows).toBe(1);
    expect(stats.fresh_rows).toBe(1);
  });

  test('stats() reports fresh / stale / total / hit counters', async () => {
    const cache = new SemanticQueryCache(engine);
    const emb = makeEmbedding(13);
    await storeCurrent(cache, 'q', emb, [makeResult('a')], META);
    await cache.lookup(emb);  // bump hit
    // Hit bump is async/fire-and-forget; give it a moment to land.
    await new Promise(r => setTimeout(r, 50));
    const stats = await cache.stats();
    expect(stats.total_rows).toBe(1);
    expect(stats.fresh_rows).toBe(1);
    expect(stats.stale_rows).toBe(0);
    expect(stats.total_hits).toBeGreaterThanOrEqual(1);
  });
});

describe('SemanticQueryCache \u2014 disabled', () => {
  test('disabled cache is a pure no-op on lookup', async () => {
    const cache = new SemanticQueryCache(engine, { enabled: false });
    const emb = makeEmbedding(99);
    await storeCurrent(cache, 'q', emb, [makeResult('a')], META);
    // Even after a store call, lookup must miss because enabled=false.
    const hit = await cache.lookup(emb);
    expect(hit.hit).toBe(false);
  });
});