From 020dc2ae1105f4f2a9ab39fcdc041fbba464ebd6 Mon Sep 17 00:00:00 2001 From: Arnas Donauskas Date: Thu, 2 Jul 2026 17:20:38 +0300 Subject: [PATCH] fix(ai-kb): address code-review findings - Security: REVOKE ALL FROM PUBLIC + GRANT EXECUTE (authenticated, service_role) on the match RPCs so anon can't read a KB by supplying an account_id (mirrors migrations 018/025). - Ingest: store chunks even when embedding fails (insert first, rethrow after) so a doc stays lexically searchable; the "lexical still works" warning is now truthful. - Retrieval: short-circuit on an empty KB (cheap COUNT) to avoid a paid query embedding + 2 RPCs on every draft/auto-reply. - Migration: HNSW instead of IVFFlat (robust on small/growing tables); language-neutral 'simple' FTS config (was 'english'); semantic RPC takes text and casts ::vector(1536) to remove PostgREST bind ambiguity. - Embeddings key: distinguish corrupt-from-absent (loadEmbeddingsKey returns {key,corrupt}); log on decrypt failure; reindex now reports failure instead of false success; POST/PATCH warn on keyword-only save. - embedTexts: fail loud when the provider omits result indices (was silently misaligning chunks with vectors). Tests: +3 (empty-KB guard, embed-failure-still-indexes, missing-index). 587 pass; typecheck/lint/build green. Co-Authored-By: Claude Opus 4.8 --- src/app/api/ai/knowledge/[id]/route.ts | 12 +++++- src/app/api/ai/knowledge/reindex/route.ts | 19 ++++++++- src/app/api/ai/knowledge/route.ts | 24 +++++++++-- src/lib/ai/config.ts | 24 ++++++++--- src/lib/ai/embeddings.test.ts | 12 ++++++ src/lib/ai/embeddings.ts | 11 ++++- src/lib/ai/knowledge.test.ts | 26 ++++++++++++ src/lib/ai/knowledge.ts | 37 +++++++++++++--- supabase/migrations/030_ai_knowledge.sql | 52 ++++++++++++++++++----- 9 files changed, 187 insertions(+), 30 deletions(-) diff --git a/src/app/api/ai/knowledge/[id]/route.ts b/src/app/api/ai/knowledge/[id]/route.ts index 5fc6a31..6b7d151 100644 --- a/src/app/api/ai/knowledge/[id]/route.ts +++ b/src/app/api/ai/knowledge/[id]/route.ts @@ -77,7 +77,10 @@ export async function PATCH(request: Request, { params }: Params) { if (!updated) return NextResponse.json({ error: 'Not found' }, { status: 404 }) if (content !== undefined) { - const embeddingsApiKey = await loadEmbeddingsKey(supabase, accountId) + const { key: embeddingsApiKey, corrupt } = await loadEmbeddingsKey( + supabase, + accountId, + ) try { await ingestDocument(supabase, accountId, { embeddingsApiKey }, id, content) } catch (err) { @@ -91,6 +94,13 @@ export async function PATCH(request: Request, { params }: Params) { { status: 200 }, ) } + if (corrupt) { + return NextResponse.json({ + success: true, + warning: + 'Updated with keyword search only — your embeddings key could not be decrypted (check ENCRYPTION_KEY, then re-enter the key).', + }) + } } return NextResponse.json({ success: true }) diff --git a/src/app/api/ai/knowledge/reindex/route.ts b/src/app/api/ai/knowledge/reindex/route.ts index f36d824..875ddf8 100644 --- a/src/app/api/ai/knowledge/reindex/route.ts +++ b/src/app/api/ai/knowledge/reindex/route.ts @@ -31,7 +31,24 @@ export async function POST() { ) } - const embeddingsApiKey = await loadEmbeddingsKey(supabase, accountId) + const { key: embeddingsApiKey, corrupt } = await loadEmbeddingsKey( + supabase, + accountId, + ) + // The whole point of Reindex is usually to backfill embeddings — so + // if a key is configured but can't be decrypted, don't quietly do a + // lexical-only pass and report success. Stop and tell the admin. + if (corrupt) { + return NextResponse.json( + { + success: false, + reindexed: 0, + error: + 'Your embeddings key could not be decrypted (check ENCRYPTION_KEY, then re-enter the key in Settings → AI Assistant). Nothing was reindexed.', + }, + { status: 200 }, + ) + } let reindexed = 0 for (const doc of docs ?? []) { diff --git a/src/app/api/ai/knowledge/route.ts b/src/app/api/ai/knowledge/route.ts index affcaf2..8cede6b 100644 --- a/src/app/api/ai/knowledge/route.ts +++ b/src/app/api/ai/knowledge/route.ts @@ -70,12 +70,20 @@ export async function POST(request: Request) { ) } - const embeddingsApiKey = await loadEmbeddingsKey(supabase, accountId) + const { key: embeddingsApiKey, corrupt } = await loadEmbeddingsKey( + supabase, + accountId, + ) try { - await ingestDocument(supabase, accountId, { embeddingsApiKey }, doc.id, content) + await ingestDocument( + supabase, + accountId, + { embeddingsApiKey }, + doc.id, + content, + ) } catch (err) { - const message = - err instanceof AiError ? err.message : 'indexing failed' + const message = err instanceof AiError ? err.message : 'indexing failed' console.error('[ai/knowledge POST] ingest error:', err) return NextResponse.json( { @@ -87,6 +95,14 @@ export async function POST(request: Request) { ) } + if (corrupt) { + return NextResponse.json({ + success: true, + id: doc.id, + warning: + 'Saved with keyword search only — your embeddings key could not be decrypted (check ENCRYPTION_KEY, then re-enter the key).', + }) + } return NextResponse.json({ success: true, id: doc.id }) } catch (err) { return toErrorResponse(err) diff --git a/src/lib/ai/config.ts b/src/lib/ai/config.ts index ee7effa..df5c3de 100644 --- a/src/lib/ai/config.ts +++ b/src/lib/ai/config.ts @@ -55,6 +55,11 @@ export async function loadAiConfig( try { embeddingsApiKey = decrypt(row.embeddings_api_key) } catch { + // Not silent — a rotated/mismatched ENCRYPTION_KEY here means + // semantic search quietly stops working, so leave a breadcrumb. + console.error( + `[ai config] embeddings key for account ${accountId} could not be decrypted — check ENCRYPTION_KEY; semantic search is disabled until it is re-entered.`, + ) embeddingsApiKey = null } } @@ -75,22 +80,29 @@ export async function loadAiConfig( * Load + decrypt just the embeddings key, independent of `is_active`. * Used by the knowledge-base ingest routes so the KB gets embedded (and * semantic search works) whenever an embeddings key is present, even if - * the assistant's master switch is currently off. Returns null when - * there's no key or it can't be decrypted. + * the assistant's master switch is currently off. + * + * Returns `{ key, corrupt }`: `key` is null when there's no key OR it + * can't be decrypted; `corrupt` distinguishes those cases so callers can + * warn ("a key is set but unusable") rather than silently indexing + * lexical-only and reporting success. */ export async function loadEmbeddingsKey( db: SupabaseClient, accountId: string, -): Promise { +): Promise<{ key: string | null; corrupt: boolean }> { const { data, error } = await db .from('ai_configs') .select('embeddings_api_key') .eq('account_id', accountId) .maybeSingle() - if (error || !data?.embeddings_api_key) return null + if (error || !data?.embeddings_api_key) return { key: null, corrupt: false } try { - return decrypt(data.embeddings_api_key) + return { key: decrypt(data.embeddings_api_key), corrupt: false } } catch { - return null + console.error( + `[ai config] embeddings key for account ${accountId} could not be decrypted — check ENCRYPTION_KEY.`, + ) + return { key: null, corrupt: true } } } diff --git a/src/lib/ai/embeddings.test.ts b/src/lib/ai/embeddings.test.ts index 8ee934d..68af460 100644 --- a/src/lib/ai/embeddings.test.ts +++ b/src/lib/ai/embeddings.test.ts @@ -85,6 +85,18 @@ describe('embedTexts', () => { }) }) + it('throws when the provider omits result indices', async () => { + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue({ + ok: true, + status: 200, + json: async () => ({ data: [{ embedding: [0.1] }, { embedding: [0.2] }] }), + } as unknown as Response), + ) + await expect(embedTexts('sk-x', ['a', 'b'])).rejects.toBeInstanceOf(AiError) + }) + it('throws on a malformed response (count mismatch)', async () => { vi.stubGlobal( 'fetch', diff --git a/src/lib/ai/embeddings.ts b/src/lib/ai/embeddings.ts index d14b718..75aa96a 100644 --- a/src/lib/ai/embeddings.ts +++ b/src/lib/ai/embeddings.ts @@ -77,8 +77,15 @@ export async function embedTexts( } // Sort by index so order matches the input batch regardless of how - // the provider returns them. - const ordered = [...rows].sort((a, b) => (a.index ?? 0) - (b.index ?? 0)) + // the provider returns them. Require a real numeric index — defaulting + // a missing one to 0 would silently misalign chunks with their + // vectors (chunk N gets chunk M's embedding), so fail loud instead. + if (rows.some((r) => typeof r.index !== 'number')) { + throw new AiError('Embeddings response was missing result indices.', { + code: 'embeddings_malformed', + }) + } + const ordered = [...rows].sort((a, b) => a.index! - b.index!) for (const r of ordered) { if (!Array.isArray(r.embedding)) { throw new AiError('Embeddings response missing a vector.', { diff --git a/src/lib/ai/knowledge.test.ts b/src/lib/ai/knowledge.test.ts index 2eead6c..c01fd53 100644 --- a/src/lib/ai/knowledge.test.ts +++ b/src/lib/ai/knowledge.test.ts @@ -12,6 +12,7 @@ import { retrieveKnowledge, ingestDocument } from './knowledge' interface FakeState { semantic: { id: string; content: string }[] fts: { id: string; content: string }[] + chunkCount: number rpcCalls: string[] inserted: Record[] | null deletedFor: string | null @@ -21,6 +22,7 @@ function makeDb() { const state: FakeState = { semantic: [], fts: [], + chunkCount: 5, // account has a non-empty KB by default rpcCalls: [], inserted: null, deletedFor: null, @@ -35,6 +37,10 @@ function makeDb() { return Promise.resolve({ data: null, error: null }) }, from: () => ({ + // retrieveKnowledge's empty-KB count guard. + select: () => ({ + eq: () => Promise.resolve({ count: state.chunkCount, error: null }), + }), delete: () => ({ eq: (_col: string, val: string) => { state.deletedFor = val @@ -64,6 +70,15 @@ describe('retrieveKnowledge', () => { expect(state.rpcCalls).toEqual([]) }) + it('short-circuits (no embed, no RPC) when the KB is empty', async () => { + const { db, state } = makeDb() + state.chunkCount = 0 + const out = await retrieveKnowledge(db, 'acct', { embeddingsApiKey: 'sk-x' }, 'q') + expect(out).toEqual([]) + expect(h.embedTexts).not.toHaveBeenCalled() + expect(state.rpcCalls).toEqual([]) + }) + it('uses lexical FTS only when there is no embeddings key', async () => { const { db, state } = makeDb() state.fts = [{ id: 'f1', content: 'F1' }] @@ -131,4 +146,15 @@ describe('ingestDocument', () => { expect(state.inserted).toBeNull() expect(h.embedTexts).not.toHaveBeenCalled() }) + + it('still stores lexical chunks when embedding fails, then rethrows', async () => { + const { db, state } = makeDb() + h.embedTexts.mockRejectedValueOnce(new Error('rate limited')) + await expect( + ingestDocument(db, 'acct', { embeddingsApiKey: 'sk-x' }, 'doc-1', 'hello world'), + ).rejects.toThrow('rate limited') + // Chunks were inserted (lexical search works) despite the embed failure… + expect(state.inserted).toHaveLength(1) + expect(state.inserted![0].embedding).toBeNull() + }) }) diff --git a/src/lib/ai/knowledge.ts b/src/lib/ai/knowledge.ts index cdd95a8..3a3842c 100644 --- a/src/lib/ai/knowledge.ts +++ b/src/lib/ai/knowledge.ts @@ -33,11 +33,6 @@ export async function ingestDocument( ): Promise { const chunks = chunkText(content) - let embeddings: number[][] | null = null - if (chunks.length > 0 && config.embeddingsApiKey) { - embeddings = await embedTexts(config.embeddingsApiKey, chunks) - } - // Replace, don't append — re-ingest must be idempotent. const { error: delErr } = await db .from('ai_knowledge_chunks') @@ -47,6 +42,22 @@ export async function ingestDocument( if (chunks.length === 0) return + // Embed if a key is set, but DON'T let an embedding failure stop the + // chunks from being stored: a failed embed must still leave the + // document searchable lexically. We record the error and rethrow it + // AFTER inserting (embedding-less) rows, so the route can warn + // "semantic indexing failed" — which is now truthful, because lexical + // search really does still work. + let embeddings: number[][] | null = null + let embedError: unknown = null + if (config.embeddingsApiKey) { + try { + embeddings = await embedTexts(config.embeddingsApiKey, chunks) + } catch (err) { + embedError = err + } + } + const rows = chunks.map((content, i) => ({ document_id: documentId, account_id: accountId, @@ -57,6 +68,8 @@ export async function ingestDocument( const { error: insErr } = await db.from('ai_knowledge_chunks').insert(rows) if (insErr) throw insErr + + if (embedError) throw embedError } /** @@ -78,6 +91,20 @@ export async function retrieveKnowledge( const query = queryText.trim() if (!query || k <= 0) return [] + // Skip everything when the account has no knowledge base — otherwise + // every draft / auto-reply would pay for a query embedding + two RPCs + // just to get []. One cheap indexed COUNT (head, no rows) instead of a + // paid embeddings call on the hot path. + try { + const { count, error } = await db + .from('ai_knowledge_chunks') + .select('id', { count: 'exact', head: true }) + .eq('account_id', accountId) + if (error || !count) return [] + } catch { + return [] + } + const picked = new Map() // id → content, preserves order // Semantic path. diff --git a/supabase/migrations/030_ai_knowledge.sql b/supabase/migrations/030_ai_knowledge.sql index 25bbd5d..4b88a24 100644 --- a/supabase/migrations/030_ai_knowledge.sql +++ b/supabase/migrations/030_ai_knowledge.sql @@ -96,7 +96,14 @@ CREATE TABLE IF NOT EXISTS ai_knowledge_chunks ( account_id uuid NOT NULL REFERENCES accounts(id) ON DELETE CASCADE, chunk_index integer NOT NULL DEFAULT 0, content text NOT NULL, - fts tsvector GENERATED ALWAYS AS (to_tsvector('english', content)) STORED, + -- Language-neutral FTS config: wacrm is used in many languages + -- (its markets include BR / LATAM / IN), and this lexical path is the + -- fallback for accounts without an embeddings key. `'simple'` tokenizes + -- + lowercases without English-only stemming/stopwords, so it degrades + -- gracefully in any language. (Per-account language config is a + -- follow-up; accounts wanting paraphrase/morphology matching add an + -- embeddings key for the semantic path.) + fts tsvector GENERATED ALWAYS AS (to_tsvector('simple', content)) STORED, embedding vector(1536), created_at timestamptz NOT NULL DEFAULT now() ); @@ -109,9 +116,14 @@ CREATE INDEX IF NOT EXISTS ai_knowledge_chunks_fts_idx ON ai_knowledge_chunks USING gin (fts); -- Cosine-distance ANN index for the semantic path. Rows with a NULL -- embedding (lexical-only accounts) are simply absent from it. +-- +-- HNSW (not IVFFlat): per-account knowledge bases start empty and grow +-- incrementally, and IVFFlat must be trained on existing rows — built +-- against an empty/tiny table its centroids are meaningless and recall +-- is poor until it's large and REINDEXed. HNSW needs no training and is +-- accurate from the first row. CREATE INDEX IF NOT EXISTS ai_knowledge_chunks_embedding_idx - ON ai_knowledge_chunks USING ivfflat (embedding vector_cosine_ops) - WITH (lists = 100); + ON ai_knowledge_chunks USING hnsw (embedding vector_cosine_ops); ALTER TABLE ai_knowledge_chunks ENABLE ROW LEVEL SECURITY; @@ -138,7 +150,8 @@ CREATE POLICY ai_knowledge_chunks_delete ON ai_knowledge_chunks FOR DELETE -- ============================================================ -- Lexical: full-text rank. `plainto_tsquery` turns a raw customer --- message into a query safely (no operator injection). +-- message into a query safely (no operator injection). Uses the same +-- language-neutral `'simple'` config as the stored `fts` column. CREATE OR REPLACE FUNCTION public.match_ai_knowledge_fts( p_account_id uuid, p_query text, @@ -147,28 +160,45 @@ CREATE OR REPLACE FUNCTION public.match_ai_knowledge_fts( RETURNS TABLE (id uuid, content text, rank real) AS $$ SELECT c.id, c.content, - ts_rank(c.fts, plainto_tsquery('english', p_query)) AS rank + ts_rank(c.fts, plainto_tsquery('simple', p_query)) AS rank FROM ai_knowledge_chunks c WHERE c.account_id = p_account_id - AND c.fts @@ plainto_tsquery('english', p_query) + AND c.fts @@ plainto_tsquery('simple', p_query) ORDER BY rank DESC LIMIT GREATEST(p_match_count, 0); $$ LANGUAGE sql STABLE SECURITY DEFINER SET search_path = public; -- Semantic: cosine distance against the query embedding. Only rows -- that actually have an embedding participate. +-- +-- `p_query_embedding` is declared `text` (not `vector`) and cast inside: +-- the caller sends the canonical pgvector literal `[0.1,0.2,...]` as a +-- plain string, so there's no ambiguity in how PostgREST binds a JSON +-- value to a `vector` parameter. Casting a literal to a constant vector +-- still lets the HNSW index serve the `<=>` order-by. CREATE OR REPLACE FUNCTION public.match_ai_knowledge_semantic( - p_account_id uuid, - p_query_embedding vector(1536), - p_match_count integer + p_account_id uuid, + p_query_embedding text, + p_match_count integer ) RETURNS TABLE (id uuid, content text, distance real) AS $$ SELECT c.id, c.content, - (c.embedding <=> p_query_embedding) AS distance + (c.embedding <=> p_query_embedding::vector(1536)) AS distance FROM ai_knowledge_chunks c WHERE c.account_id = p_account_id AND c.embedding IS NOT NULL - ORDER BY c.embedding <=> p_query_embedding + ORDER BY c.embedding <=> p_query_embedding::vector(1536) LIMIT GREATEST(p_match_count, 0); $$ LANGUAGE sql STABLE SECURITY DEFINER SET search_path = public; + +-- Lock down EXECUTE (mirrors migrations 018 / 025). These are +-- SECURITY DEFINER and would otherwise default to PUBLIC — i.e. the +-- anon role — which, since the function bypasses RLS and only gates on +-- the passed account_id, would let an unauthenticated caller read any +-- account's knowledge base. The draft path calls them as `authenticated` +-- and the auto-reply bot as `service_role`. +REVOKE ALL ON FUNCTION public.match_ai_knowledge_fts(uuid, text, integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION public.match_ai_knowledge_fts(uuid, text, integer) TO authenticated, service_role; +REVOKE ALL ON FUNCTION public.match_ai_knowledge_semantic(uuid, text, integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION public.match_ai_knowledge_semantic(uuid, text, integer) TO authenticated, service_role;