Merge pull request #23 from mbrendan/fix-vsearch-hang

src/qmd.ts

@@ -1972,8 +1972,11 @@ async function vectorSearch(query: string, opts: OutputOptions, model: string =
   const perQueryLimit = opts.all ? 500 : 20;
   const allResults = new Map<string, { file: string; displayPath: string; title: string; body: string; score: number; hash: string }>();
 
-  // Use Promise.all for concurrent vector searches
-  await Promise.all(vectorQueries.map(async (q) => {
+  // IMPORTANT: Run vector searches sequentially, not with Promise.all.
+  // node-llama-cpp's embedding context hangs when multiple concurrent embed() calls
+  // are made. This is a known limitation of the LlamaEmbeddingContext.
+  // See: https://github.com/tobi/qmd/pull/23
+  for (const q of vectorQueries) {
     const vecResults = await searchVec(db, q, model, perQueryLimit, collectionName as any);
     for (const r of vecResults) {
       const existing = allResults.get(r.filepath);
@@ -1981,7 +1984,7 @@ async function vectorSearch(query: string, opts: OutputOptions, model: string =
         allResults.set(r.filepath, { file: r.filepath, displayPath: r.displayPath, title: r.title, body: r.body || "", score: r.score, hash: r.hash });
       }
     }
-  }));
+  }
 
   // Sort by max score and limit to requested count
   const results = Array.from(allResults.values())

src/store.test.ts

@@ -1855,6 +1855,42 @@ describe("LlamaCpp Integration", () => {
     await cleanupTestDb(store);
   });
 
+  // Regression test for https://github.com/tobi/qmd/pull/23
+  // sqlite-vec virtual tables hang when combined with JOINs in the same query.
+  // The fix uses a two-step approach: vector query first, then separate JOINs.
+  test("searchVec uses two-step query to avoid sqlite-vec JOIN hang", async () => {
+    const store = await createTestStore();
+    const collectionName = await createTestCollection();
+
+    const hash = "regression_test_hash";
+    await insertTestDocument(store.db, collectionName, {
+      name: "regression-doc",
+      hash,
+      body: "Test content for vector search regression",
+      filepath: "/test/regression.md",
+      displayPath: "regression.md",
+    });
+
+    // Create vector table and insert a test vector
+    store.ensureVecTable(768);
+    const embedding = Array(768).fill(0).map(() => Math.random());
+    store.db.prepare(`INSERT INTO content_vectors (hash, seq, pos, model, embedded_at) VALUES (?, 0, 0, 'test', ?)`).run(hash, new Date().toISOString());
+    store.db.prepare(`INSERT INTO vectors_vec (hash_seq, embedding) VALUES (?, ?)`).run(`${hash}_0`, new Float32Array(embedding));
+
+    // This should complete quickly (not hang) due to the two-step fix
+    // The old code with JOINs in the sqlite-vec query would hang indefinitely
+    const startTime = Date.now();
+    const results = await store.searchVec("test content", "embeddinggemma", 5);
+    const elapsed = Date.now() - startTime;
+
+    // If the query took more than 5 seconds, something is wrong
+    // (the hang bug would cause it to never return at all)
+    expect(elapsed).toBeLessThan(5000);
+    expect(results.length).toBeGreaterThan(0);
+
+    await cleanupTestDb(store);
+  });
+
   test("expandQuery returns original plus expanded queries", async () => {
     const store = await createTestStore();
 

src/store.ts

@@ -1679,48 +1679,66 @@ export async function searchVec(db: Database, query: string, model: string, limi
   const embedding = await getEmbedding(query, model, true);
   if (!embedding) return [];
 
-  // sqlite-vec requires "k = ?" for KNN queries
-  let sql = `
+  // IMPORTANT: We use a two-step query approach here because sqlite-vec virtual tables
+  // hang indefinitely when combined with JOINs in the same query. Do NOT try to
+  // "optimize" this by combining into a single query with JOINs - it will break.
+  // See: https://github.com/tobi/qmd/pull/23
+
+  // Step 1: Get vector matches from sqlite-vec (no JOINs allowed)
+  const vecResults = db.prepare(`
+    SELECT hash_seq, distance
+    FROM vectors_vec
+    WHERE embedding MATCH ? AND k = ?
+  `).all(new Float32Array(embedding), limit * 3) as { hash_seq: string; distance: number }[];
+
+  if (vecResults.length === 0) return [];
+
+  // Step 2: Get chunk info and document data
+  const hashSeqs = vecResults.map(r => r.hash_seq);
+  const distanceMap = new Map(vecResults.map(r => [r.hash_seq, r.distance]));
+
+  // Build query for document lookup
+  const placeholders = hashSeqs.map(() => '?').join(',');
+  let docSql = `
     SELECT
-      v.hash_seq,
-      v.distance,
+      cv.hash || '_' || cv.seq as hash_seq,
+      cv.hash,
+      cv.pos,
       'qmd://' || d.collection || '/' || d.path as filepath,
       d.collection || '/' || d.path as display_path,
       d.title,
-      content.doc as body,
-      cv.hash,
-      cv.pos
-    FROM vectors_vec v
-    JOIN content_vectors cv ON cv.hash || '_' || cv.seq = v.hash_seq
+      content.doc as body
+    FROM content_vectors cv
     JOIN documents d ON d.hash = cv.hash AND d.active = 1
     JOIN content ON content.hash = d.hash
-    WHERE v.embedding MATCH ? AND k = ?
+    WHERE cv.hash || '_' || cv.seq IN (${placeholders})
   `;
-
-  const params: (Float32Array | number | string)[] = [new Float32Array(embedding), limit * 3];
+  const params: string[] = [...hashSeqs];
 
   if (collectionId) {
-    // Filter by collection name
-    sql += ` AND d.collection = ?`;
+    docSql += ` AND d.collection = ?`;
     params.push(String(collectionId));
   }
 
-  sql += ` ORDER BY v.distance`;
-
-  const rows = db.prepare(sql).all(...params) as { hash_seq: string; distance: number; filepath: string; display_path: string; title: string; body: string; hash: string; pos: number }[];
+  const docRows = db.prepare(docSql).all(...params) as {
+    hash_seq: string; hash: string; pos: number; filepath: string;
+    display_path: string; title: string; body: string;
+  }[];
 
-  const seen = new Map<string, { row: typeof rows[0]; bestDist: number }>();
-  for (const row of rows) {
+  // Combine with distances and dedupe by filepath
+  const seen = new Map<string, { row: typeof docRows[0]; bestDist: number }>();
+  for (const row of docRows) {
+    const distance = distanceMap.get(row.hash_seq) ?? 1;
     const existing = seen.get(row.filepath);
-    if (!existing || row.distance < existing.bestDist) {
-      seen.set(row.filepath, { row, bestDist: row.distance });
+    if (!existing || distance < existing.bestDist) {
+      seen.set(row.filepath, { row, bestDist: distance });
     }
   }
 
   return Array.from(seen.values())
     .sort((a, b) => a.bestDist - b.bestDist)
     .slice(0, limit)
-    .map(({ row }) => {
+    .map(({ row, bestDist }) => {
       const collectionName = row.filepath.split('//')[1]?.split('/')[0] || "";
       return {
         filepath: row.filepath,
@@ -1733,7 +1751,7 @@ export async function searchVec(db: Database, query: string, model: string, limi
         bodyLength: row.body.length,
         body: row.body,
         context: getContextForFile(db, row.filepath),
-        score: 1 - row.distance,  // Cosine similarity = 1 - cosine distance
+        score: 1 - bestDist,  // Cosine similarity = 1 - cosine distance
         source: "vec" as const,
         chunkPos: row.pos,
       };