From a6adb0dd970f9945bea4f7eebdea337a3d99a517 Mon Sep 17 00:00:00 2001 From: alichherawalla Date: Sat, 4 Jul 2026 11:59:20 +0530 Subject: [PATCH] fix(rag): roll back + surface a document that fails to embed (no silent dead entries) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit indexDocument caught embedding failure non-fatally and still returned success, so a document with ZERO embeddings showed as 'added' in the Knowledge Base but was not searchable (semantic search skips it) — 'looks indexed but doesn't actually work'. backfillEmbeddings has no auto-trigger, so the swallowed failure stranded the doc permanently. Now: on any embedding failure (load throw, embedBatch throw/OOM, or an embedding-count mismatch), roll back the just-inserted document + chunks and throw. KnowledgeBaseScreen already alerts on an index throw, so the user sees the real error and can retry instead of a silent non-functional entry. Tests: replace the 'continues without embeddings' test (which encoded the lie) with three fails-before/passes-after cases — load failure, embedBatch failure, and embedding-count mismatch all roll back (deleteDocument called) and throw, with insertEmbeddingsBatch never called. Co-Authored-By: Claude Opus 4.8 (1M context) --- __tests__/unit/services/rag/index.test.ts | 33 ++++++++++++++++++++--- src/services/rag/index.ts | 13 ++++++++- 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/__tests__/unit/services/rag/index.test.ts b/__tests__/unit/services/rag/index.test.ts index 0aa72c24a..b630e934e 100644 --- a/__tests__/unit/services/rag/index.test.ts +++ b/__tests__/unit/services/rag/index.test.ts @@ -113,10 +113,37 @@ describe('RagService', () => { .rejects.toThrow('already in the knowledge base'); }); - it('continues without embeddings if embedding fails', async () => { + // HONESTY: a doc with no embeddings is not searchable, so it must NOT be left in + // the KB looking "added". On any embedding failure, indexDocument rolls back the + // inserted doc + chunks and throws, so the KB add flow surfaces the real error + // (instead of a silent dead entry). backfillEmbeddings has no auto-trigger, so the + // old "continue without embeddings" swallow stranded the doc permanently. + it('rolls back the inserted document and throws when embedding LOAD fails', async () => { mockEmbedding.load.mockRejectedValueOnce(new Error('model not found')); - const docId = await ragService.indexDocument({ projectId: 'proj1', filePath: '/p', fileName: 'test.txt', fileSize: 100 }); - expect(docId).toBe(1); // Still returns docId + await expect( + ragService.indexDocument({ projectId: 'proj1', filePath: '/p', fileName: 'test.txt', fileSize: 100 }), + ).rejects.toThrow(/Could not index "test.txt".*embeddings failed/); + expect(mockDb.deleteDocument).toHaveBeenCalledWith(1); // the just-inserted docId + expect(mockDb.insertEmbeddingsBatch).not.toHaveBeenCalled(); + }); + + it('rolls back the inserted document and throws when embedBatch fails (e.g. OOM)', async () => { + mockEmbedding.embedBatch.mockRejectedValueOnce(new Error('Embedding failed (native error). OOM')); + await expect( + ragService.indexDocument({ projectId: 'proj1', filePath: '/p', fileName: 'test.txt', fileSize: 100 }), + ).rejects.toThrow(/embeddings failed to generate/); + expect(mockDb.deleteDocument).toHaveBeenCalledWith(1); + expect(mockDb.insertEmbeddingsBatch).not.toHaveBeenCalled(); + }); + + it('rolls back and throws when embedBatch returns fewer embeddings than chunks', async () => { + // insertChunks mock returns [1,2] (2 chunks); return only 1 embedding → mismatch. + mockEmbedding.embedBatch.mockResolvedValueOnce([[0.1, 0.2]]); + await expect( + ragService.indexDocument({ projectId: 'proj1', filePath: '/p', fileName: 'test.txt', fileSize: 100 }), + ).rejects.toThrow(/embeddings failed to generate/); + expect(mockDb.deleteDocument).toHaveBeenCalledWith(1); + expect(mockDb.insertEmbeddingsBatch).not.toHaveBeenCalled(); }); }); diff --git a/src/services/rag/index.ts b/src/services/rag/index.ts index d9fdb30ce..66bf9960b 100644 --- a/src/services/rag/index.ts +++ b/src/services/rag/index.ts @@ -63,6 +63,9 @@ class RagService { await embeddingService.load(); const texts = chunks.map(c => c.content); const embeddings = await embeddingService.embedBatch(texts); + if (embeddings.length !== rowIds.length) { + throw new Error(`embedding count ${embeddings.length} does not match ${rowIds.length} chunks`); + } const entries = rowIds.map((rowId, i) => ({ chunkRowid: rowId, docId, @@ -71,7 +74,15 @@ class RagService { ragDatabase.insertEmbeddingsBatch(entries); logger.log(`[RAG] Generated ${embeddings.length} embeddings for ${fileName}`); } catch (err) { - logger.error('[RAG] Embedding generation failed (non-fatal):', err); + // A document with no embeddings is NOT searchable — semantic search skips it — + // yet the KB would show it as "added". That is the "looks indexed but doesn't + // work" lie. backfillEmbeddings has no auto-trigger, so a swallowed failure + // strands the doc permanently. Roll back the just-inserted doc + chunks and + // surface the failure so the user sees the real error (KnowledgeBaseScreen + // already alerts on an index throw) and can retry — never a silent dead entry. + ragDatabase.deleteDocument(docId); + logger.error('[RAG] Embedding generation failed — rolled back document:', err); + throw new Error(`Could not index "${fileName}": embeddings failed to generate. ${err instanceof Error ? err.message : String(err)}`); } onProgress?.({ stage: 'done', message: 'Done' });