From 09322d2d89382ed1a7d86faceeea3df622f9a284 Mon Sep 17 00:00:00 2001 From: azizur100389 Date: Wed, 26 Aug 2026 13:57:56 +0100 Subject: [PATCH] fix(storage): load VECTOR only when needed (#3045) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(storage): load VECTOR only when needed * test(storage): verify VECTOR reopen lifecycle --------- Co-authored-by: Gergő Magyar --- gitnexus/src/core/lbug/pool-adapter.ts | 177 +++++++++++++++-- gitnexus/src/mcp/local/local-backend.ts | 82 ++++---- gitnexus/test/integration/lbug-pool.test.ts | 8 +- gitnexus/test/unit/calltool-dispatch.test.ts | 6 +- gitnexus/test/unit/lbug-pool-fts-load.test.ts | 183 ++++++++++++++++-- .../local-backend-embedding-dims-warn.test.ts | 4 + .../unit/local-backend-lazy-vector.test.ts | 119 ++++++++++++ 7 files changed, 509 insertions(+), 70 deletions(-) create mode 100644 gitnexus/test/unit/local-backend-lazy-vector.test.ts diff --git a/gitnexus/src/core/lbug/pool-adapter.ts b/gitnexus/src/core/lbug/pool-adapter.ts index a9a06aeb3..7f38e48ba 100644 --- a/gitnexus/src/core/lbug/pool-adapter.ts +++ b/gitnexus/src/core/lbug/pool-adapter.ts @@ -135,6 +135,10 @@ interface SharedDB { * scan (#2623 follow-up). Optional with `?? false` semantics so the * construction sites stay minimal. */ vectorLoaded?: boolean; + /** In-flight/completed lazy VECTOR probe for this Database lifecycle. + * Retaining a false result prevents every semantic request from retrying + * the same unavailable extension; teardown clears it before a reopen. */ + vectorLoadPromise?: Promise; /** File identity at open — used to detect reuse of a shared read-only handle * whose on-disk index was rebuilt/swapped since it opened (only reachable * when a second pool consumer shares this dbPath; #2614 F2). */ @@ -368,6 +372,7 @@ function closeOne(repoId: string): void { shared.refCount = 0; shared.ftsLoaded = false; shared.vectorLoaded = false; + shared.vectorLoadPromise = undefined; } else { shared.db.close().catch(() => {}); dbCache.delete(entry.dbPath); @@ -823,14 +828,6 @@ async function doInitLbug(repoId: string, dbPath: string): Promise { if (!shared.ftsLoaded) { shared.ftsLoaded = await loadFTSExtension(available[0], { policy: 'load-only' }); } - // VECTOR too — extension load scope is per-Database, so this one load - // makes QUERY_VECTOR_INDEX legal on every pooled connection. Same - // load-only contract as FTS above; on failure the semantic-query lane - // falls back to the exact scan with its own diagnostic (#2623 follow-up). - if (!shared.vectorLoaded) { - shared.vectorLoaded = await loadVectorExtension(available[0], { policy: 'load-only' }); - } - // Register pool entry only after all connections are pre-warmed and FTS is // loaded. Concurrent executeQuery calls see either "not initialized" // (and throw cleanly) or a fully ready pool — never a half-built one. @@ -900,12 +897,6 @@ export async function initLbugWithDb( if (!shared.ftsLoaded) { shared.ftsLoaded = await loadFTSExtension(available[0], { policy: 'load-only' }); } - // VECTOR too — same per-Database scope and load-only contract as the - // doInitLbug site above (#2623 follow-up). - if (!shared.vectorLoaded) { - shared.vectorLoaded = await loadVectorExtension(available[0], { policy: 'load-only' }); - } - pool.set(repoId, { db: existingDb, available, @@ -921,6 +912,143 @@ export async function initLbugWithDb( traceRss('init', repoId); } +/** + * Lazily load VECTOR for a semantic query. + * + * Exact graph reads never call this function, so opening their read pool does + * not probe or warn about an optional extension they do not use. The promise + * lives on SharedDB because extension scope is per Database, and also joins + * concurrent first semantic requests onto one LOAD attempt. + */ +export async function ensureVectorExtension(repoId: string): Promise { + const entry = pool.get(repoId); + if (!entry) { + throw new Error(`LadybugDB not initialized for repo "${repoId}". Call initLbug first.`); + } + + const shared = dbCache.get(entry.dbPath); + if (!shared) { + throw new Error(`LadybugDB shared handle is unavailable for repo "${repoId}".`); + } + if (shared.vectorLoaded) return true; + if (shared.vectorLoadPromise) return shared.vectorLoadPromise; + + const loadAttempt = (async () => { + const conn = await checkout(entry); + try { + const loaded = await loadVectorExtension(conn, { policy: 'load-only' }); + shared.vectorLoaded = loaded; + return loaded; + } finally { + checkin(entry, conn); + } + })(); + const cachedAttempt = loadAttempt.catch((err) => { + // A transient checkout/load failure must not poison this Database for the + // rest of its lifetime. Keep resolved false cached, but let a later + // semantic request retry a rejected attempt. + if (shared.vectorLoadPromise === cachedAttempt) { + shared.vectorLoadPromise = undefined; + } + throw err; + }); + shared.vectorLoadPromise = cachedAttempt; + + return shared.vectorLoadPromise; +} + +/** + * Detect an actual VECTOR procedure call without treating source text stored in + * Cypher literals or comments as executable syntax. + */ +function callsVectorIndex(cypher: string): boolean { + if (!/QUERY_VECTOR_INDEX/i.test(cypher)) return false; + + let code = ''; + let state: 'code' | 'single' | 'double' | 'backtick' | 'line-comment' | 'block-comment' = 'code'; + let backtickIdentifier = ''; + + for (let i = 0; i < cypher.length; i++) { + const ch = cypher[i]; + const next = cypher[i + 1]; + + if (state === 'code') { + if (ch === "'" || ch === '"' || ch === '`') { + state = ch === "'" ? 'single' : ch === '"' ? 'double' : 'backtick'; + if (state === 'backtick') backtickIdentifier = ''; + code += ' '; + } else if (ch === '/' && next === '/') { + state = 'line-comment'; + code += ' '; + i++; + } else if (ch === '/' && next === '*') { + state = 'block-comment'; + code += ' '; + i++; + } else { + code += ch; + } + continue; + } + + if (state === 'line-comment') { + if (ch === '\n' || ch === '\r') { + state = 'code'; + code += ch; + } else { + code += ' '; + } + continue; + } + + if (state === 'block-comment') { + if (ch === '*' && next === '/') { + state = 'code'; + code += ' '; + i++; + } else { + code += ch === '\n' || ch === '\r' ? ch : ' '; + } + continue; + } + + if (state === 'backtick') { + if (ch === '`' && next === '`') { + backtickIdentifier += '`'; + code += ' '; + i++; + } else if (ch === '`') { + state = 'code'; + code += + backtickIdentifier.toUpperCase() === 'QUERY_VECTOR_INDEX' ? 'QUERY_VECTOR_INDEX' : ' '; + } else if (ch === '\\' && next !== undefined) { + backtickIdentifier += next; + code += ' '; + i++; + } else { + backtickIdentifier += ch; + code += ch === '\n' || ch === '\r' ? ch : ' '; + } + continue; + } + + if (ch === '\\') { + code += ' '; + if (next !== undefined) { + code += next === '\n' || next === '\r' ? next : ' '; + i++; + } + continue; + } + + const closesLiteral = (state === 'single' && ch === "'") || (state === 'double' && ch === '"'); + if (closesLiteral) state = 'code'; + code += ch === '\n' || ch === '\r' ? ch : ' '; + } + + return /\bCALL\s+QUERY_VECTOR_INDEX\s*\(/i.test(code); +} + /** * Checkout a connection from the pool. * Returns an available connection, or creates a new one if under the cap. @@ -1028,14 +1156,31 @@ export const executeParameterized = async ( poolSidecarLogger.warn(message), ); - const entry = pool.get(repoId); + let entry = pool.get(repoId); if (!entry) { throw new Error(`LadybugDB not initialized for repo "${repoId}". Call initLbug first.`); } - entry.lastUsed = Date.now(); + // Exact reads must not pay for VECTOR, but an explicit raw vector procedure + // call is a semantic read. Preflight before taking the query connection: + // ensureVectorExtension performs its own checkout, so holding one here could + // make a saturated pool wait for a connection that every caller is holding. + // A load rejection must not replace the query's own diagnostic. + if (callsVectorIndex(cypher)) { + await ensureVectorExtension(repoId).catch(() => false); + + // The preflight suspends, so close/re-init may replace the pool entry. + // Re-read it before checkout to avoid querying through a stale handle. + entry = pool.get(repoId); + if (!entry) { + throw new Error( + `LadybugDB connection pool closed for repo "${repoId}" (re-init/teardown); retry the query.`, + ); + } + } const conn = await checkout(entry); + entry.lastUsed = Date.now(); silenceStdout(); activeQueryCount++; let queryResult: lbug.QueryResult | lbug.QueryResult[] | undefined; diff --git a/gitnexus/src/mcp/local/local-backend.ts b/gitnexus/src/mcp/local/local-backend.ts index 854b7f267..9809a76e4 100644 --- a/gitnexus/src/mcp/local/local-backend.ts +++ b/gitnexus/src/mcp/local/local-backend.ts @@ -13,6 +13,7 @@ import { initLbug, executeQuery, executeParameterized, + ensureVectorExtension, closeLbug, isLbugReady, statDbIdentity, @@ -1253,12 +1254,12 @@ export class LocalBackend { private warnedSiblingDrift: Set = new Set(); /** - * One-shot stderr warning for the VECTOR-extension fallback. Without this - * guard the diagnostic would fire on every `semanticSearch()` call on - * platforms where the extension is unsupported (e.g. Windows), making MCP - * stderr noisy per DoD §2.8. + * One-shot stderr guards for distinct VECTOR load and index-query failures. + * Keeping them separate preserves both diagnostics across semanticSearch calls + * without repeating either on hot paths. */ - private warnedVectorUnsupported = false; + private warnedVectorLoadFailed = false; + private warnedVectorQueryFailed = false; /** * One-shot warning when a pruned or Node-unloadable optional embedding stack @@ -3219,21 +3220,31 @@ export class LocalBackend { this.lastQueryEmbeddingDims.set(repo.lbugPath, dims); const queryVecStr = `[${queryVec.join(',')}]`; const maxDistance = getVectorMaxDistance(DEFAULT_MCP_VECTOR_MAX_DISTANCE); + let vectorReady = false; + try { + vectorReady = await ensureVectorExtension(repo.lbugPath); + } catch (err) { + if (!this.warnedVectorLoadFailed) { + this.warnedVectorLoadFailed = true; + logger.warn( + { err }, + 'GitNexus [query:vector]: vector extension load failed; using exact scan fallback', + ); + } + } let bestChunks = new Map< string, { distance: number; chunkIndex: number; startLine: number; endLine: number } >(); - // Always TRY the vector lane — no platform gate. LadybugDB ships the - // VECTOR extension for every supported platform, Windows included - // (#2623 follow-up; the old `platform !== 'win32'` gate was stale), so - // whether the index is queryable is a per-machine runtime fact. The - // catch below is the fallback: any failure (extension unloadable, index - // absent, older DB) degrades to the exact scan with a once-per-backend - // diagnostic instead of being silently swallowed. - try { - bestChunks = await collectBestChunks(limit, async (fetchLimit) => { - const vectorQuery = ` + // Try the vector lane only after its lazy load succeeds. An unavailable + // extension is already reported by ExtensionManager; an index/query + // failure below gets its own once-per-backend diagnostic before the exact + // scan fallback. + if (vectorReady) { + try { + bestChunks = await collectBestChunks(limit, async (fetchLimit) => { + const vectorQuery = ` CALL QUERY_VECTOR_INDEX('${EMBEDDING_TABLE_NAME}', '${EMBEDDING_INDEX_NAME}', CAST(${queryVecStr} AS FLOAT[${dims}]), ${fetchLimit}) YIELD node AS emb, distance @@ -3244,26 +3255,27 @@ export class LocalBackend { ORDER BY distance `; - const embResults = await executeQuery(repo.lbugPath, vectorQuery); - return embResults.map((row) => ({ - nodeId: row.nodeId ?? row[0], - chunkIndex: row.chunkIndex ?? row[1] ?? 0, - startLine: row.startLine ?? row[2] ?? 0, - endLine: row.endLine ?? row[3] ?? 0, - distance: row.distance ?? row[4], - })); - }); - } catch (err) { - bestChunks = new Map(); - if (!this.warnedVectorUnsupported) { - // Rare diagnostic: surface why semantic search fell back to the - // exact scan. Emitted once per `LocalBackend` instance lifetime to - // avoid noisy stderr on hot semantic-search paths (DoD §2.8). - this.warnedVectorUnsupported = true; - logger.warn( - { err }, - 'GitNexus [query:vector]: vector index query failed; using exact scan fallback', - ); + const embResults = await executeQuery(repo.lbugPath, vectorQuery); + return embResults.map((row) => ({ + nodeId: row.nodeId ?? row[0], + chunkIndex: row.chunkIndex ?? row[1] ?? 0, + startLine: row.startLine ?? row[2] ?? 0, + endLine: row.endLine ?? row[3] ?? 0, + distance: row.distance ?? row[4], + })); + }); + } catch (err) { + bestChunks = new Map(); + if (!this.warnedVectorQueryFailed) { + // Rare diagnostic: surface why semantic search fell back to the + // exact scan. Emitted once per `LocalBackend` instance lifetime to + // avoid noisy stderr on hot semantic-search paths (DoD §2.8). + this.warnedVectorQueryFailed = true; + logger.warn( + { err }, + 'GitNexus [query:vector]: vector index query failed; using exact scan fallback', + ); + } } } diff --git a/gitnexus/test/integration/lbug-pool.test.ts b/gitnexus/test/integration/lbug-pool.test.ts index 083974674..063ab1dd2 100644 --- a/gitnexus/test/integration/lbug-pool.test.ts +++ b/gitnexus/test/integration/lbug-pool.test.ts @@ -341,7 +341,7 @@ withTestLbugDB( } }); - it('QUERY_VECTOR_INDEX works through the pool once the pre-warm loads VECTOR', async (ctx) => { + it('QUERY_VECTOR_INDEX works through the pool after its lazy query preflight', async (ctx) => { const core = await import('../../src/core/lbug/lbug-adapter.js'); const { batchInsertEmbeddings } = await import('../../src/core/embeddings/embedding-pipeline.js'); @@ -376,12 +376,12 @@ withTestLbugDB( // loads are per-Database, so a shared/injected Database would inherit // the VECTOR load from createVectorIndex above and pass even without // the pre-warm fix. A fresh Database has nothing loaded — only the - // pool's own pre-warm can make the vector lane legal. + // pool's own lazy query preflight can make the vector lane legal. await core.closeLbug(); // The regression: through the POOL, the vector lane must work without - // any caller loading the extension. Pre-fix this rejects with - // "Catalog exception: function QUERY_VECTOR_INDEX is not defined". + // any caller loading the extension. The query-specific preflight loads + // VECTOR here while exact reads remain untouched. await initLbug('vec-repo', handle.dbPath); const vec = `CAST([${embedding.join(',')}] AS FLOAT[${EMBEDDING_DIMS}])`; const rows = (await executeQuery( diff --git a/gitnexus/test/unit/calltool-dispatch.test.ts b/gitnexus/test/unit/calltool-dispatch.test.ts index aaf5a344d..e3f1892ce 100644 --- a/gitnexus/test/unit/calltool-dispatch.test.ts +++ b/gitnexus/test/unit/calltool-dispatch.test.ts @@ -23,6 +23,7 @@ const { lbugMocks } = vi.hoisted(() => ({ initLbug: vi.fn().mockResolvedValue(undefined), executeQuery: vi.fn().mockResolvedValue([]), executeParameterized: vi.fn().mockResolvedValue([]), + ensureVectorExtension: vi.fn().mockResolvedValue(true), closeLbug: vi.fn().mockResolvedValue(undefined), isLbugReady: vi.fn().mockReturnValue(true), }, @@ -684,9 +685,8 @@ describe('LocalBackend.callTool', () => { }); it('falls back to the exact scan with a once-per-backend warning when the vector index query fails', async () => { - // The platform gate is gone (#2623 follow-up): the vector lane is always - // ATTEMPTED, and a runtime failure (extension unloadable, index absent) is - // what routes semantic search onto the exact scan. + // Once the lazy extension preflight succeeds, a runtime index-query failure + // routes semantic search onto the exact scan. const cap = _captureLogger(); (executeQuery as any).mockImplementation(async (_repoId: string, cypher: string) => { if (cypher.includes('COUNT(*) AS cnt')) return [{ cnt: 1 }]; diff --git a/gitnexus/test/unit/lbug-pool-fts-load.test.ts b/gitnexus/test/unit/lbug-pool-fts-load.test.ts index 7f5b6ff87..920c8b994 100644 --- a/gitnexus/test/unit/lbug-pool-fts-load.test.ts +++ b/gitnexus/test/unit/lbug-pool-fts-load.test.ts @@ -1,14 +1,30 @@ import { afterEach, describe, expect, it, vi } from 'vitest'; -const { loadFTSExtensionMock, loadVectorExtensionMock } = vi.hoisted(() => ({ - loadFTSExtensionMock: vi.fn(), - loadVectorExtensionMock: vi.fn(), -})); +import { mkdtemp, rm, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import path from 'node:path'; + +const { createLbugDatabaseMock, loadFTSExtensionMock, loadVectorExtensionMock } = vi.hoisted( + () => ({ + createLbugDatabaseMock: vi.fn(), + loadFTSExtensionMock: vi.fn(), + loadVectorExtensionMock: vi.fn(), + }), +); vi.mock('@ladybugdb/core', () => ({ default: { Database: vi.fn(), Connection: vi.fn(function (this: any) { + this.query = vi.fn(async () => ({ getAll: async () => [], close: vi.fn() })); + this.prepare = vi.fn(async () => ({ + isSuccess: () => true, + getErrorMessage: async () => '', + })); + this.execute = vi.fn(async () => ({ + getAll: async () => [], + close: vi.fn().mockResolvedValue(undefined), + })); this.close = vi.fn().mockResolvedValue(undefined); }), }, @@ -21,17 +37,22 @@ vi.mock('../../src/core/lbug/lbug-adapter.js', () => ({ })); vi.mock('../../src/core/lbug/lbug-config.js', () => ({ - createLbugDatabase: vi.fn(), + createLbugDatabase: createLbugDatabaseMock, toNativeSafePath: vi.fn((p: string) => p), isWalCorruptionError: vi.fn(() => false), WAL_RECOVERY_SUGGESTION: '', })); -const { closeLbug, initLbugWithDb } = await import('../../src/core/lbug/pool-adapter.js'); +const { closeLbug, ensureVectorExtension, executeParameterized, initLbug, initLbugWithDb } = + await import('../../src/core/lbug/pool-adapter.js'); describe('read-pool FTS loading', () => { + const tempDirs: string[] = []; + afterEach(async () => { await closeLbug().catch(() => {}); + await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true }))); + createLbugDatabaseMock.mockReset(); loadFTSExtensionMock.mockReset(); loadVectorExtensionMock.mockReset(); loadVectorExtensionMock.mockResolvedValue(false); @@ -39,7 +60,6 @@ describe('read-pool FTS loading', () => { it('loads FTS with load-only policy and caches a successful load', async () => { loadFTSExtensionMock.mockResolvedValue(true); - loadVectorExtensionMock.mockResolvedValue(true); const db = {} as any; await initLbugWithDb('repo-a', db, '/tmp/shared-fts-db'); @@ -51,7 +71,6 @@ describe('read-pool FTS loading', () => { it('does not fake a successful load when FTS is unavailable', async () => { loadFTSExtensionMock.mockResolvedValue(false); - loadVectorExtensionMock.mockResolvedValue(false); const db = {} as any; await initLbugWithDb('repo-a', db, '/tmp/shared-fts-db'); @@ -66,7 +85,7 @@ describe('read-pool FTS loading', () => { }); }); - it('loads VECTOR with load-only policy and caches a successful load (#2623 follow-up)', async () => { + it('does not probe VECTOR while initializing exact-read pools (#3021)', async () => { loadFTSExtensionMock.mockResolvedValue(true); loadVectorExtensionMock.mockResolvedValue(true); const db = {} as any; @@ -74,20 +93,160 @@ describe('read-pool FTS loading', () => { await initLbugWithDb('repo-a', db, '/tmp/shared-vec-db'); await initLbugWithDb('repo-b', db, '/tmp/shared-vec-db'); + expect(loadVectorExtensionMock).not.toHaveBeenCalled(); + }); + + it('loads VECTOR lazily once for concurrent semantic reads on a shared Database', async () => { + loadFTSExtensionMock.mockResolvedValue(true); + loadVectorExtensionMock.mockResolvedValue(true); + const db = {} as any; + + await initLbugWithDb('repo-a', db, '/tmp/shared-vec-db'); + await initLbugWithDb('repo-b', db, '/tmp/shared-vec-db'); + + await expect( + Promise.all([ensureVectorExtension('repo-a'), ensureVectorExtension('repo-b')]), + ).resolves.toEqual([true, true]); + expect(loadVectorExtensionMock).toHaveBeenCalledTimes(1); expect(loadVectorExtensionMock).toHaveBeenCalledWith(expect.anything(), { policy: 'load-only', }); }); - it('retries the VECTOR load on the next open when it was unavailable', async () => { + it('caches an unavailable VECTOR result until a non-external Database is reopened', async () => { loadFTSExtensionMock.mockResolvedValue(true); loadVectorExtensionMock.mockResolvedValue(false); + const dir = await mkdtemp(path.join(tmpdir(), 'gitnexus-vector-reopen-')); + tempDirs.push(dir); + const dbPath = path.join(dir, 'index.lbug'); + await writeFile(dbPath, 'fixture'); + const firstDb = { init: vi.fn(), close: vi.fn().mockResolvedValue(undefined) }; + const secondDb = { init: vi.fn(), close: vi.fn().mockResolvedValue(undefined) }; + createLbugDatabaseMock.mockReturnValueOnce(firstDb).mockReturnValueOnce(secondDb); + + await initLbug('repo-a', dbPath); + await expect(ensureVectorExtension('repo-a')).resolves.toBe(false); + await expect(ensureVectorExtension('repo-a')).resolves.toBe(false); + + expect(loadVectorExtensionMock).toHaveBeenCalledTimes(1); + + await closeLbug('repo-a'); + await initLbug('repo-b', dbPath); + await expect(ensureVectorExtension('repo-b')).resolves.toBe(false); + + expect(createLbugDatabaseMock).toHaveBeenCalledTimes(2); + expect(firstDb.close).toHaveBeenCalledTimes(1); + expect(loadVectorExtensionMock).toHaveBeenCalledTimes(2); + }); + + it('retries VECTOR after a rejected lazy load', async () => { + loadFTSExtensionMock.mockResolvedValue(true); + loadVectorExtensionMock.mockRejectedValueOnce(new Error('transient load failure')); + loadVectorExtensionMock.mockResolvedValueOnce(true); const db = {} as any; - await initLbugWithDb('repo-a', db, '/tmp/shared-vec-db'); - await initLbugWithDb('repo-b', db, '/tmp/shared-vec-db'); + await initLbugWithDb('repo-a', db, '/tmp/shared-vec-retry-db'); + await expect(ensureVectorExtension('repo-a')).rejects.toThrow('transient load failure'); + await expect(ensureVectorExtension('repo-a')).resolves.toBe(true); expect(loadVectorExtensionMock).toHaveBeenCalledTimes(2); }); + + it('preflights VECTOR only for executable QUERY_VECTOR_INDEX calls', async () => { + loadFTSExtensionMock.mockResolvedValue(true); + loadVectorExtensionMock.mockResolvedValue(false); + await initLbugWithDb('repo-a', {} as any, '/tmp/vector-call-detection-db'); + + const exactReads = [ + "RETURN 'CALL QUERY_VECTOR_INDEX(' AS sourceText", + 'RETURN "CALL QUERY_VECTOR_INDEX(" AS sourceText', + 'RETURN `CALL QUERY_VECTOR_INDEX(` AS propertyName', + 'RETURN `QUERY_VECTOR_INDEX` AS propertyName', + "RETURN 'CALL `QUERY_VECTOR_INDEX`(' AS sourceText", + '// CALL QUERY_VECTOR_INDEX(\nRETURN 1 AS value', + '// CALL `QUERY_VECTOR_INDEX`(\nRETURN 1 AS value', + '/* CALL QUERY_VECTOR_INDEX( */ RETURN 1 AS value', + '/* CALL `QUERY_VECTOR_INDEX`( */ RETURN 1 AS value', + ]; + for (const cypher of exactReads) { + await expect(executeParameterized('repo-a', cypher, {})).resolves.toEqual([]); + } + expect(loadVectorExtensionMock).not.toHaveBeenCalled(); + + await expect( + executeParameterized( + 'repo-a', + "call query_vector_index\n('CodeEmbedding', 'embedding_idx', [0.1], 1)", + {}, + ), + ).resolves.toEqual([]); + expect(loadVectorExtensionMock).toHaveBeenCalledTimes(1); + }); + + it('preflights VECTOR for a backtick-escaped procedure identifier', async () => { + loadFTSExtensionMock.mockResolvedValue(true); + loadVectorExtensionMock.mockResolvedValue(false); + await initLbugWithDb('repo-a', {} as any, '/tmp/vector-quoted-call-detection-db'); + + await expect( + executeParameterized( + 'repo-a', + "CALL /* legal comment */ `QUERY_VECTOR_INDEX`('CodeEmbedding', 'embedding_idx', [0.1], 1)", + {}, + ), + ).resolves.toEqual([]); + expect(loadVectorExtensionMock).toHaveBeenCalledTimes(1); + }); + + it('does not hold query connections while a saturated VECTOR preflight loads', async () => { + loadFTSExtensionMock.mockResolvedValue(true); + let releaseVectorLoad: ((loaded: boolean) => void) | undefined; + loadVectorExtensionMock.mockImplementation( + () => + new Promise((resolve) => { + releaseVectorLoad = resolve; + }), + ); + await initLbugWithDb('repo-a', {} as any, '/tmp/vector-saturation-db'); + + const calls = Array.from({ length: 8 }, () => + executeParameterized( + 'repo-a', + "CALL QUERY_VECTOR_INDEX('CodeEmbedding', 'embedding_idx', [0.1], 1)", + {}, + ), + ); + + try { + await vi.waitFor(() => expect(loadVectorExtensionMock).toHaveBeenCalledTimes(1), { + timeout: 500, + }); + releaseVectorLoad?.(true); + await expect(Promise.all(calls)).resolves.toHaveLength(8); + } finally { + if (releaseVectorLoad) { + releaseVectorLoad(true); + } else { + // Allows the old hold-one/wait-for-one ordering to unwind promptly + // instead of leaving its pool waiter alive until the 30-second timeout. + await closeLbug('repo-a'); + } + await Promise.allSettled(calls); + } + }); + + it('lets a direct vector query report its own error when preflight rejects', async () => { + loadFTSExtensionMock.mockResolvedValue(true); + loadVectorExtensionMock.mockRejectedValueOnce(new Error('transient load failure')); + await initLbugWithDb('repo-a', {} as any, '/tmp/vector-preflight-rejection-db'); + + await expect( + executeParameterized( + 'repo-a', + "CALL QUERY_VECTOR_INDEX('CodeEmbedding', 'embedding_idx', [0.1], 1)", + {}, + ), + ).resolves.toEqual([]); + }); }); diff --git a/gitnexus/test/unit/local-backend-embedding-dims-warn.test.ts b/gitnexus/test/unit/local-backend-embedding-dims-warn.test.ts index 948c7cab8..b9a80f616 100644 --- a/gitnexus/test/unit/local-backend-embedding-dims-warn.test.ts +++ b/gitnexus/test/unit/local-backend-embedding-dims-warn.test.ts @@ -15,6 +15,7 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; const executeQueryMock = vi.fn(); const executeParameterizedMock = vi.fn(); +const ensureVectorExtensionMock = vi.fn(); const loadMetaMock = vi.fn(); const embedQueryMock = vi.fn(); const getEmbeddingDimsMock = vi.fn(); @@ -24,6 +25,7 @@ vi.mock('../../src/core/lbug/pool-adapter.js', async (importOriginal) => ({ initLbug: vi.fn(), executeQuery: (...args: unknown[]) => executeQueryMock(...args), executeParameterized: (...args: unknown[]) => executeParameterizedMock(...args), + ensureVectorExtension: (...args: unknown[]) => ensureVectorExtensionMock(...args), closeLbug: vi.fn(), isLbugReady: vi.fn().mockReturnValue(true), })); @@ -121,6 +123,7 @@ describe('LocalBackend.query — index/server embedding width drift (#2798)', () beforeEach(() => { vi.clearAllMocks(); executeParameterizedMock.mockResolvedValue([]); + ensureVectorExtensionMock.mockResolvedValue(true); loadMetaMock.mockResolvedValue(null); }); @@ -225,6 +228,7 @@ describe('LocalBackend.semanticSearch — recorded query-embedding width (#2798) beforeEach(() => { vi.clearAllMocks(); executeParameterizedMock.mockResolvedValue([]); + ensureVectorExtensionMock.mockResolvedValue(true); loadMetaMock.mockResolvedValue(null); }); diff --git a/gitnexus/test/unit/local-backend-lazy-vector.test.ts b/gitnexus/test/unit/local-backend-lazy-vector.test.ts new file mode 100644 index 000000000..063c5e486 --- /dev/null +++ b/gitnexus/test/unit/local-backend-lazy-vector.test.ts @@ -0,0 +1,119 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { _captureLogger } from '../../src/core/logger.js'; + +const executeQueryMock = vi.fn(); +const ensureVectorExtensionMock = vi.fn(); +const embedQueryMock = vi.fn(); + +vi.mock('../../src/core/lbug/pool-adapter.js', async (importOriginal) => ({ + ...(await importOriginal()), + executeQuery: (...args: unknown[]) => executeQueryMock(...args), + ensureVectorExtension: (...args: unknown[]) => ensureVectorExtensionMock(...args), +})); + +vi.mock('../../src/mcp/core/embedder.js', () => ({ + embedQuery: (...args: unknown[]) => embedQueryMock(...args), + getEmbeddingDims: () => 3, +})); + +import { LocalBackend } from '../../src/mcp/local/local-backend.js'; + +interface SemanticSearchable { + semanticSearch(repo: { lbugPath: string }, query: string, limit: number): Promise; +} + +const runSemanticSearch = (backend: LocalBackend): Promise => + (backend as unknown as SemanticSearchable).semanticSearch({ lbugPath: '/tmp/index' }, 'q', 5); + +describe('LocalBackend semantic search lazy VECTOR loading (#3021)', () => { + beforeEach(() => { + executeQueryMock.mockReset(); + ensureVectorExtensionMock.mockReset(); + embedQueryMock.mockReset(); + embedQueryMock.mockResolvedValue([0.1, 0.2, 0.3]); + }); + + it('does not probe VECTOR when the exact embedding count is zero', async () => { + executeQueryMock.mockResolvedValueOnce([{ cnt: 0 }]); + + await expect(runSemanticSearch(new LocalBackend())).resolves.toEqual([]); + + expect(ensureVectorExtensionMock).not.toHaveBeenCalled(); + expect(embedQueryMock).not.toHaveBeenCalled(); + }); + + it('probes VECTOR only after embeddings are found and keeps exact-scan fallback', async () => { + executeQueryMock.mockResolvedValueOnce([{ cnt: 2 }]).mockResolvedValueOnce([]); + ensureVectorExtensionMock.mockResolvedValue(false); + + await expect(runSemanticSearch(new LocalBackend())).resolves.toEqual([]); + + expect(ensureVectorExtensionMock).toHaveBeenCalledOnce(); + expect(ensureVectorExtensionMock).toHaveBeenCalledWith('/tmp/index'); + expect(executeQueryMock).toHaveBeenCalledTimes(2); + expect( + executeQueryMock.mock.calls.some(([, cypher]) => + String(cypher).includes('QUERY_VECTOR_INDEX'), + ), + ).toBe(false); + }); + + it('uses QUERY_VECTOR_INDEX only after the lazy VECTOR load succeeds', async () => { + executeQueryMock + .mockResolvedValueOnce([{ cnt: 1 }]) + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([]); + ensureVectorExtensionMock.mockResolvedValue(true); + + await expect(runSemanticSearch(new LocalBackend())).resolves.toEqual([]); + + expect(ensureVectorExtensionMock).toHaveBeenCalledOnce(); + expect( + executeQueryMock.mock.calls.some(([, cypher]) => + String(cypher).includes('QUERY_VECTOR_INDEX'), + ), + ).toBe(true); + }); + + it('falls back to the exact scan when the lazy VECTOR load rejects', async () => { + executeQueryMock.mockResolvedValueOnce([{ cnt: 2 }]).mockResolvedValueOnce([]); + ensureVectorExtensionMock.mockRejectedValue(new Error('transient load failure')); + + await expect(runSemanticSearch(new LocalBackend())).resolves.toEqual([]); + + expect(executeQueryMock).toHaveBeenCalledTimes(2); + expect( + executeQueryMock.mock.calls.some(([, cypher]) => + String(cypher).includes('QUERY_VECTOR_INDEX'), + ), + ).toBe(false); + }); + + it('reports load and index-query failures independently', async () => { + executeQueryMock.mockImplementation(async (_repoId: string, cypher: string) => { + if (cypher.includes('COUNT(*) AS cnt')) return [{ cnt: 1 }]; + if (cypher.includes('QUERY_VECTOR_INDEX')) throw new Error('stale vector index'); + return []; + }); + ensureVectorExtensionMock + .mockRejectedValueOnce(new Error('transient load failure')) + .mockResolvedValueOnce(true); + const backend = new LocalBackend(); + const cap = _captureLogger(); + + try { + await expect(runSemanticSearch(backend)).resolves.toEqual([]); + await expect(runSemanticSearch(backend)).resolves.toEqual([]); + + const messages = cap.records().map((record) => String(record.msg ?? '')); + expect( + messages.filter((message) => message.includes('vector extension load failed')), + ).toHaveLength(1); + expect( + messages.filter((message) => message.includes('vector index query failed')), + ).toHaveLength(1); + } finally { + cap.restore(); + } + }); +});