From ff297010da40dc04bc89807a691b9b094d6ef2a2 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 18 Sep 2026 01:45:31 -0700 Subject: [PATCH 1/3] fix(knowledge): recover vector candidates an HNSW post-filter discards pgvector's HNSW index post-filters by construction, so a visibility predicate can only discard neighbours the graph walk already committed to. When the documents a caller may read are a small share of the index, the traversal returns a handful of candidates instead of its limit. Rank the permitted set exactly when the traversal comes back underfilled and a bounded probe says that set is small enough to afford. A traversal that fills its limit is returned untouched, so a scope the graph serves well pays nothing. Bound the probe by documents examined rather than chunks accumulated, and give it its own sub-budget so deciding against a rescue can never cost the leg its results. Retune the scan settings so the tuple budget sits an order of magnitude above the beam, which is what lets the iterative scan iterate at all, and so the narrower beam lowers the leg's uninterruptible floor. --- .../app/api/knowledge/search/utils.test.ts | 35 ++- apps/sim/lib/knowledge/search/budget.ts | 12 + apps/sim/lib/knowledge/search/diagnostics.ts | 5 +- apps/sim/lib/knowledge/search/queries.test.ts | 279 +++++++++++------- apps/sim/lib/knowledge/search/queries.ts | 277 +++++++++++------ 5 files changed, 381 insertions(+), 227 deletions(-) diff --git a/apps/sim/app/api/knowledge/search/utils.test.ts b/apps/sim/app/api/knowledge/search/utils.test.ts index 7106308b9d5..83217cbd6b2 100644 --- a/apps/sim/app/api/knowledge/search/utils.test.ts +++ b/apps/sim/app/api/knowledge/search/utils.test.ts @@ -212,11 +212,14 @@ describe('Knowledge Search Utils', () => { describe('handleTagAndVectorSearch', () => { it('returns only bounded ranked rows without first materializing every matching tag ID', async () => { resetDbChainMock() - queueTableRows( - schemaMock.embedding, - Array.from({ length: 201 }, (_, index) => ({ id: `candidate-${index}` })) - ) - queueTableRows(schemaMock.embedding, [makeResult('second', 0.2), makeResult('first', 0.1)]) + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = (query as { toSQL: () => { sql: string } }).toSQL().sql + if (statement.includes('AS visible')) return [] + if (statement.includes(') + 0 LIMIT')) return [{ id: 'first' }, { id: 'second' }] + if (statement.includes('WITH scored_search_candidates')) + return [makeResult('second', 0.2), makeResult('first', 0.1)] + return [{ id: 'doc-first' }, { id: 'doc-second' }] + }) queueTableRows(schemaMock.embedding, [makeResult('second', 0.2), makeResult('first', 0.1)]) const results = await handleTagAndVectorSearch({ @@ -231,11 +234,13 @@ describe('Knowledge Search Utils', () => { }) expect(results.map((row) => row.id)).toEqual(['first', 'second']) - expect(dbChainMockFns.select).toHaveBeenCalledTimes(3) - expect(Object.keys(dbChainMockFns.select.mock.calls[0][0])).toEqual(['id']) - expect(dbChainMockFns.limit).toHaveBeenNthCalledWith(1, 400) - expect(dbChainMockFns.select.mock.calls[1][0]).toHaveProperty('distance') - expect(dbChainMockFns.limit).toHaveBeenCalledWith(20) + /** Only hydration reads through the query builder; ranking never materializes tag IDs. */ + expect(dbChainMockFns.select).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.select.mock.calls[0][0]).toHaveProperty('distance') + const exact = dbChainMockFns.execute.mock.calls + .map(([query]) => (query as { toSQL: () => { sql: string; params: unknown[] } }).toSQL()) + .find((statement) => statement.sql.includes(') + 0 LIMIT'))! + expect(exact.params).toContain(400) }) it('should throw error when no filters provided', async () => { @@ -552,13 +557,13 @@ describe('Knowledge Search Utils', () => { }) expect(results.map((r) => r.id)).toEqual(['vector-hit']) - expect(dbChainMockFns.select).toHaveBeenCalledTimes(2) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(1) }) it('runs both legs and fuses them in hybrid mode', async () => { /** - * The raw vector probe does not consume a table chain. Keyword ranking and - * hydration complete before vector exact ranking and content hydration. + * Vector ranking is raw SQL throughout and consumes no table chain. Keyword ranking and + * hydration complete before vector content hydration. */ dbChainMockFns.execute.mockResolvedValue([{ id: 'vector-hit' }]) queueTableRows(schemaMock.embedding, [{ id: 'keyword-hit', keywordRank: 0.9 }]) @@ -576,7 +581,7 @@ describe('Knowledge Search Utils', () => { }) expect(results.map((r) => r.id).sort()).toEqual(['keyword-hit', 'vector-hit']) - expect(dbChainMockFns.select).toHaveBeenCalledTimes(4) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(3) }) it('propagates unexpected keyword errors after the vector leg finishes', async () => { @@ -601,7 +606,7 @@ describe('Knowledge Search Utils', () => { queryVector: { vector: JSON.stringify(TEST_EMBEDDING), dimensions: 1536 }, }) ).rejects.toBe(failure) - expect(dbChainMockFns.select).toHaveBeenCalledTimes(3) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(2) }) it('skips both query legs when only tag filters are provided', async () => { diff --git a/apps/sim/lib/knowledge/search/budget.ts b/apps/sim/lib/knowledge/search/budget.ts index a80279ec3e1..fcb1ac9e5cd 100644 --- a/apps/sim/lib/knowledge/search/budget.ts +++ b/apps/sim/lib/knowledge/search/budget.ts @@ -39,6 +39,18 @@ export class SearchBudget { return remaining } + /** + * A shorter deadline for one step of the leg, carrying the same leg and cancellation signal. + * Spending it is the step's own business: the leg's deadline is untouched and still governs. + */ + capped(milliseconds: number): SearchBudget { + return new SearchBudget( + this.leg, + Math.min(this.deadline, performance.now() + milliseconds), + this.signal + ) + } + isTimeout(error: unknown): boolean { this.signal?.throwIfAborted() if (error instanceof SearchDeadlineError || getPostgresErrorCode(error) === '57014') { diff --git a/apps/sim/lib/knowledge/search/diagnostics.ts b/apps/sim/lib/knowledge/search/diagnostics.ts index 293d0acb0e7..8406c131027 100644 --- a/apps/sim/lib/knowledge/search/diagnostics.ts +++ b/apps/sim/lib/knowledge/search/diagnostics.ts @@ -47,6 +47,7 @@ export type SearchStage = | 'vector.settings' | 'vector.probe' | 'vector.rerank' + | 'vector.exact_candidates' | 'vector.exact' | 'vector.candidate_search' | 'source_overview' @@ -76,7 +77,7 @@ export interface SearchDiagnosticMetadata { searchMode?: 'hybrid' | 'vector' boostRecency?: boolean embeddingDimensions?: number - vectorRanking?: 'exact' | 'candidate-rerank' + vectorRanking?: 'exact' | 'exact-candidates' | 'candidate-rerank' vectorCandidateStorage?: 'stored-halfvec' /** * Whether the bounded traversal filled its candidate limit. `underfilled` means visibility @@ -86,6 +87,8 @@ export interface SearchDiagnosticMetadata { vectorCandidateScan?: 'planned' | 'underfilled' vectorBudgetMs?: number vectorCandidateLimit?: number + /** Visible documents the tractability probe enumerated, capped at its own document limit. */ + vectorProbeDocumentCount?: number vectorCandidateCount?: number vectorCandidateDimensions?: number resultCount?: number diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index a3a3360b0e4..f38447398a7 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -30,6 +30,7 @@ import { handleVectorOnlySearch, retrieveKnowledgeSearch, type SearchParams, + VECTOR_PROBE_DOCUMENT_LIMIT, } from '@/lib/knowledge/search/queries' import type { StructuredFilter } from '@/lib/knowledge/types' @@ -57,7 +58,7 @@ describe('retrieval leg budgets', () => { vi.spyOn(SearchBudget.prototype, 'remaining').mockImplementation(function ( this: SearchBudget ) { - deadlines.set(this.leg, this.deadline) + if (!deadlines.has(this.leg)) deadlines.set(this.leg, this.deadline) return remaining.call(this) }) const access: UserAccessScope = { @@ -106,6 +107,18 @@ function render(condition: unknown) { return (condition as { toSQL: () => { sql: string; params: unknown[] } }).toSQL() } +/** The document probe is the only vector statement that selects ids without ordering them. */ +function isProbeStatement(sql: string) { + return sql.includes('AS id FROM') && !sql.includes('ORDER BY') +} + +/** `+ 0` is what keeps the exact ranking off the ANN index, so it also identifies the statement. */ +function isExactRanking(sql: string) { + return sql.includes(') + 0 LIMIT') +} + +const statements = () => dbChainMockFns.execute.mock.calls.map(([query]) => render(query)) + function renderOne(filters: StructuredFilter[]) { const conditions = getStructuredTagFilters(filters, embeddingTable) expect(conditions).toHaveLength(1) @@ -354,6 +367,8 @@ describe('workspace-scoped vector retrieval', () => { }, ] let probeRows: Array<{ id: string }> + let traversedRows: Array<{ id: string; initial_count?: number }> + let exactRows: Array<{ id: string }> let failSettings: unknown let failCandidates: unknown @@ -361,20 +376,23 @@ describe('workspace-scoped vector retrieval', () => { resetDbChainMock() getForConnectors.mockReset() probeRows = probe + traversedRows = candidates + exactRows = probe failSettings = undefined failCandidates = undefined dbChainMockFns.execute.mockImplementation(async (query) => { const statement = render(query).sql - if (statement.includes('SELECT scoped_chunk.id')) return probeRows if (statement.includes('hnsw.iterative_scan')) { if (failSettings) throw failSettings return [] } if (statement.includes('AS visible')) { if (failCandidates) throw failCandidates - return candidates + return traversedRows } if (statement.includes('WITH scored_search_candidates')) return ranked + if (isExactRanking(statement)) return exactRows + if (statement.includes('AS id FROM')) return probeRows return [] }) }) @@ -383,8 +401,6 @@ describe('workspace-scoped vector retrieval', () => { vi.useRealTimers() }) - const statements = () => dbChainMockFns.execute.mock.calls.map(([query]) => render(query)) - it.each([handleVectorOnlySearch, handleTagAndVectorSearch])( 'does not acquire a connection or start SQL after the KB retrieval deadline', async (search) => { @@ -405,32 +421,84 @@ describe('workspace-scoped vector retrieval', () => { } ) - it('ranks an exhausted visible scope exactly and rechecks access before returning content', async () => { - probeRows = ranked - queueTableRows(schemaMock.embedding, ranked) + it('rescues an underfilled traversal by ranking the permitted set exactly', async () => { + traversedRows = ranked + probeRows = [{ id: 'near-doc' }, { id: 'far-doc' }] + exactRows = ranked queueTableRows(schemaMock.embedding, [...ranked].reverse()) expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['near', 'far']) - expect(statements()).toHaveLength(1) - expect(statements()[0].sql).not.toContain('<=>') - expect(render(dbChainMockFns.orderBy.mock.calls[0][0]).sql).toContain('+ 0') - for (const [condition] of dbChainMockFns.where.mock.calls) { - expect( - hasMockCondition( - condition, - (node) => - node.type === 'inArray' && - node.column === schemaMock.embedding.id && - Array.isArray(node.values) && - node.values.length === 2 && - node.values.includes('near') - ) - ).toBe(true) - expect(JSON.stringify(condition)).toContain('required_clause') - expect(JSON.stringify(condition)).toContain('aclVerifiedAt') - } + const exact = statements().find((query) => isExactRanking(query.sql))! + expect(exact.sql).not.toContain('CROSS JOIN LATERAL') + expect(JSON.stringify(exact)).toContain('near-doc') + const probeStatement = statements().find((query) => isProbeStatement(query.sql))! + expect(probeStatement.sql).not.toContain('<=>') + expect(probeStatement.params).toContain(VECTOR_PROBE_DOCUMENT_LIMIT + 1) + expect(JSON.stringify(probeStatement)).toContain('required_clause') expect(getForConnectors).not.toHaveBeenCalled() }) + it('keeps the tuple budget an order of magnitude above the beam so the scan can iterate', async () => { + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + await handleVectorOnlySearch(params) + const settings = statements().find((query) => query.sql.includes('hnsw.iterative_scan'))! + const [maxScanTuples, efSearch] = settings.params as string[] + /** + * `max_scan_tuples` excludes the first beam, so a budget at or below `ef_search` is spent + * before `ResumeScanItems` can widen anything and the scan never iterates (pgvector#912). + */ + expect(Number(maxScanTuples)).toBeGreaterThanOrEqual(Number(efSearch) * 10) + /** A beam is uninterruptible, so its width is also this leg's cancellation floor. */ + expect(Number(efSearch)).toBeLessThanOrEqual(400) + }) + + it('leaves a filled traversal alone instead of probing for an exact ranking', async () => { + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['near', 'far']) + expect(statements().filter((query) => isProbeStatement(query.sql))).toHaveLength(0) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) + }) + + it('keeps an underfilled traversal when the permitted set is too large to rank exactly', async () => { + traversedRows = ranked + probeRows = new Array(VECTOR_PROBE_DOCUMENT_LIMIT + 1).fill({ id: 'doc' }) + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['near', 'far']) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) + }) + + it('finishes a scope the probe finds nothing in without ranking anything', async () => { + traversedRows = [] + probeRows = [] + expect(await handleVectorOnlySearch(params)).toEqual([]) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) + expect( + statements().filter((query) => query.sql.includes('WITH scored_search_candidates')) + ).toHaveLength(0) + expect(getForConnectors).not.toHaveBeenCalled() + }) + + it('spends only its own share of the leg on a probe that runs long', async () => { + const budget = new SearchBudget('vector', performance.now() + 8000) + traversedRows = ranked + exactRows = ranked + const execute = dbChainMockFns.execute.getMockImplementation()! + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query).sql + if (isProbeStatement(statement)) + throw new Error('canceling statement due to statement timeout', { + cause: { code: '57014' }, + }) + return execute(query) + }) + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + expect((await handleVectorOnlySearch({ ...params, budget })).map((row) => row.id)).toEqual([ + 'near', + 'far', + ]) + expect(budget.timedOut).toBe(false) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) + }) + it('uses compact candidates for a large KB and applies full workspace access before its limit', async () => { queueTableRows(schemaMock.embedding, [...ranked].reverse()) expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['near', 'far']) @@ -504,14 +572,13 @@ describe('workspace-scoped vector retrieval', () => { }) it('does not turn a broad tag filter into exhaustive full-vector ranking', async () => { - queueTableRows(schemaMock.embedding, probe) queueTableRows(schemaMock.embedding, ranked) const rows = await handleTagAndVectorSearch({ ...params, structuredFilters: [{ tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'common' }], }) expect(rows.map((row) => row.id)).toEqual(['near', 'far']) - expect(Object.keys(dbChainMockFns.select.mock.calls[0][0])).toEqual(['id']) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) const candidate = statements().find((query) => query.sql.includes('AS visible'))! expect(JSON.stringify(candidate)).toContain('common') expect(JSON.stringify(candidate)).toContain(String(schemaMock.embedding.tag1)) @@ -534,7 +601,7 @@ describe('workspace-scoped vector retrieval', () => { expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() }) - it.each(['vector.probe', 'vector.candidate_search', 'vector.rerank', 'vector.sql'] as const)( + it.each(['vector.candidate_search', 'vector.rerank', 'vector.sql'] as const)( 'reports a %s timeout as partial, not a complete empty search', async (failedStage) => { const query = SearchBudget.prototype.query @@ -565,7 +632,6 @@ describe('workspace-scoped vector retrieval', () => { run: (executor: SearchExecutor) => PromiseLike ) { const result = await (query.bind(this) as SearchBudget['query'])(stage, run) - if (stage === 'vector.probe') vi.spyOn(performance, 'now').mockReturnValue(30) if (stage === 'vector.candidate_search') vi.spyOn(performance, 'now').mockReturnValue(60) if (stage === 'vector.rerank') vi.spyOn(performance, 'now').mockReturnValue(80) return result @@ -576,7 +642,7 @@ describe('workspace-scoped vector retrieval', () => { statements() .filter((query) => query.sql.includes('statement_timeout')) .map((query) => query.params[0]) - ).toEqual(['100', '70', '70', '40', '20']) + ).toEqual(['100', '100', '40', '20']) }) it('does not convert an unexpected candidate failure into partial retrieval', async () => { @@ -778,6 +844,7 @@ describe('live repository authorization follows ranked candidates', () => { } const probePages: Array> = [] + const exactPages: Array> = [] const candidatePages: Array> = [] const rerankPages: Array>> = [] const keywordPages: Array>> = [] @@ -791,20 +858,19 @@ describe('live repository authorization follows ranked candidates', () => { beforeEach(() => { resetDbChainMock() probePages.length = 0 + exactPages.length = 0 candidatePages.length = 0 rerankPages.length = 0 keywordPages.length = 0 - dbChainMockFns.execute.mockImplementation(async (query) => - render(query).sql.includes('SELECT scoped_chunk.id') - ? (probePages.shift() ?? []) - : render(query).sql.includes('AS visible') - ? (candidatePages.shift() ?? []) - : render(query).sql.includes('WITH scored_search_candidates') - ? (rerankPages.shift() ?? []) - : render(query).sql.includes('WITH visible_keyword_documents') - ? (keywordPages.shift() ?? []) - : [] - ) + dbChainMockFns.execute.mockImplementation(async (query) => { + const statement = render(query).sql + if (statement.includes('AS visible')) return candidatePages.shift() ?? [] + if (statement.includes('WITH scored_search_candidates')) return rerankPages.shift() ?? [] + if (statement.includes('WITH visible_keyword_documents')) return keywordPages.shift() ?? [] + if (isExactRanking(statement)) return exactPages.shift() ?? [] + if (statement.includes('AS id FROM')) return probePages.shift() ?? [] + return [] + }) getForConnectors.mockReset().mockResolvedValue(allowed) }) @@ -847,73 +913,58 @@ describe('live repository authorization follows ranked candidates', () => { ) }) - it('finishes empty scopes after the bounded probe without scanning HNSW or calling providers', async () => { + it('finishes a scope the probe finds nothing in without ranking or calling providers', async () => { probePages.push([]) expect(await handleVectorOnlySearch({ ...params, structuredFilters: undefined })).toEqual([]) - expect(dbChainMockFns.select).not.toHaveBeenCalled() - const probe = render(dbChainMockFns.execute.mock.calls[0][0]) - expect(probe.sql).toContain('CROSS JOIN LATERAL') - expect(probe.params.filter((value) => value === 400)).toHaveLength(2) - expect(dbChainMockFns.orderBy).not.toHaveBeenCalled() + const probe = statements().find((query) => isProbeStatement(query.sql))! + expect(probe.params).toContain(VECTOR_PROBE_DOCUMENT_LIMIT + 1) + expect(probe.sql).not.toContain('<=>') + expect(JSON.stringify(probe)).toContain('required_clause') + expect( + statements().filter((query) => query.sql.includes('WITH scored_search_candidates')) + ).toHaveLength(0) expect(getForConnectors).not.toHaveBeenCalled() }) - it('reads vectors only for the bounded IDs when a broad scope has few candidates', async () => { - probePages.push([candidate('selected', 'allowed-source')]) - queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')]) + it('ranks the permitted set exactly when the traversal comes back underfilled', async () => { + probePages.push([{ id: 'doc-selected' }]) + exactPages.push([{ id: 'selected' }]) + queueRerank([candidate('selected', 'allowed-source')]) queueTableRows(schemaMock.embedding, [ { id: 'selected', content: 'Verified small scope', distance: 0.1 }, ]) expect(await handleVectorOnlySearch({ ...params, structuredFilters: undefined })).toEqual([ { id: 'selected', content: 'Verified small scope', distance: 0.1 }, ]) - const probe = dbChainMockFns.execute.mock.calls[0][0] - expect(render(probe).sql).toContain('SELECT scoped_chunk.id') - expect(render(probe).sql).not.toContain('<=>') - expect(JSON.stringify(probe)).toContain('required_clause') - expect( - hasMockCondition( - dbChainMockFns.where.mock.calls[0][0], - (node) => - node.type === 'inArray' && - node.column === schemaMock.embedding.id && - Array.isArray(node.values) && - node.values.length === 1 && - node.values[0] === 'selected' - ) - ).toBe(true) + const exact = statements().find((query) => isExactRanking(query.sql))! + expect(exact.sql).not.toContain('CROSS JOIN LATERAL') + expect(JSON.stringify(exact)).toContain('doc-selected') expect(getForConnectors).toHaveBeenCalledExactlyOnceWith(['allowed-source'], undefined) }) - it.each([199, 200, 399])( - 'ranks an exhausted scope of %s chunks once without repeating candidate search', + it.each([1, 200, 399])( + 'ranks a permitted set of %s documents exactly without repeating the traversal', async (count) => { - const probe = Array.from({ length: count }, (_, index) => ({ id: `chunk-${index}` })) - probePages.push(probe) - queueTableRows(schemaMock.embedding, [candidate('chunk-0', 'allowed-source')]) + probePages.push(Array.from({ length: count }, (_, index) => ({ id: `doc-${index}` }))) + exactPages.push([{ id: 'chunk-0' }]) + queueRerank([candidate('chunk-0', 'allowed-source')]) queueTableRows(schemaMock.embedding, [ { id: 'chunk-0', content: 'Authorized passage', distance: 0.1 }, ]) const rows = await handleVectorOnlySearch({ ...params, structuredFilters: undefined }) expect(rows.map((row) => row.id)).toEqual(['chunk-0']) - expect(dbChainMockFns.execute).toHaveBeenCalledOnce() - expect( - hasMockCondition( - dbChainMockFns.where.mock.calls[0][0], - (node) => - node.type === 'inArray' && - node.column === schemaMock.embedding.id && - Array.isArray(node.values) && - node.values.length === count - ) - ).toBe(true) - expect(dbChainMockFns.orderBy).toHaveBeenCalledOnce() + expect(statements().filter((query) => query.sql.includes('AS visible'))).toHaveLength(1) + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(1) } ) - it('keeps an underfilled ANN result instead of rescoring the whole projection', async () => { - probePages.push(Array.from({ length: 400 }, (_, index) => ({ id: `probe-${index}` }))) + it('keeps an underfilled traversal when the permitted set is past the probe bound', async () => { queueCandidates([{ id: 'selected' }], 1) + probePages.push( + Array.from({ length: VECTOR_PROBE_DOCUMENT_LIMIT + 1 }, (_, index) => ({ + id: `doc-${index}`, + })) + ) queueRerank([candidate('selected', 'allowed-source')]) queueTableRows(schemaMock.embedding, [ { id: 'selected', content: 'Verified fallback', distance: 0.1 }, @@ -921,13 +972,12 @@ describe('live repository authorization follows ranked candidates', () => { expect(await handleVectorOnlySearch({ ...params, structuredFilters: undefined })).toEqual([ { id: 'selected', content: 'Verified fallback', distance: 0.1 }, ]) - const candidateQuery = dbChainMockFns.execute.mock.calls.find(([query]) => - render(query).sql.includes('AS visible') - )![0] + const candidateQuery = statements().find((query) => query.sql.includes('AS visible'))! /** Widening the scan on underfill is what made this leg exceed its budget on a large corpus. */ - expect(render(candidateQuery).sql).not.toContain('UNION ALL') - expect(render(candidateQuery).sql).not.toContain('filtered_scores') - expect(render(candidateQuery).sql).toContain('CROSS JOIN LATERAL') + expect(candidateQuery.sql).not.toContain('UNION ALL') + expect(candidateQuery.sql).not.toContain('filtered_scores') + expect(candidateQuery.sql).toContain('CROSS JOIN LATERAL') + expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(0) expect(JSON.stringify(dbChainMockFns.where.mock.calls.at(-1)![0])).toContain( 'github_read_grant' ) @@ -1010,10 +1060,13 @@ describe('live repository authorization follows ranked candidates', () => { '%s ranks identifiers before verification and loads content under the full predicate', async (mode) => { const candidates = [candidate('selected', 'allowed-source')] - if (mode === 'vector' || mode === 'tag-vector') - queueTableRows(schemaMock.embedding, candidates) + if (mode === 'vector' || mode === 'tag-vector') { + probePages.push([{ id: 'doc-selected' }]) + exactPages.push([{ id: 'selected' }]) + queueRerank(candidates) + } if (mode === 'keyword') keywordPages.push(candidates) - else queueTableRows(schemaMock.embedding, candidates) + if (mode === 'tags') queueTableRows(schemaMock.embedding, candidates) queueTableRows(schemaMock.embedding, [{ id: 'selected', content: 'verified result' }]) const rows = mode === 'vector' @@ -1035,23 +1088,19 @@ describe('live repository authorization follows ranked candidates', () => { expect(ranking).toContain('ORDER BY keyword_rank DESC, id LIMIT') expect(ranking).not.toContain('<=>') expect(ranking).not.toContain('"content"') - } else { - expect( - Object.keys(dbChainMockFns.select.mock.calls[mode === 'tags' ? 0 : 1][0]).sort() - ).toEqual( - [ - 'id', - 'documentId', - 'connectorId', - 'liveAuthorizationSource', - ...(mode === 'tags' ? [] : ['distance']), - ].sort() + } else if (mode === 'tags') { + expect(Object.keys(dbChainMockFns.select.mock.calls[0][0]).sort()).toEqual( + ['id', 'documentId', 'connectorId', 'liveAuthorizationSource'].sort() ) + } else { + const ranking = statements().find((query) => isExactRanking(query.sql))! + expect(ranking.sql).toContain('AS id FROM') + expect(ranking.sql).not.toContain('"content"') } const rankingOrder = - mode === 'keyword' - ? dbChainMockFns.execute.mock.invocationCallOrder[0] - : dbChainMockFns.select.mock.invocationCallOrder[0] + mode === 'tags' + ? dbChainMockFns.select.mock.invocationCallOrder[0] + : dbChainMockFns.execute.mock.invocationCallOrder[0] expect(rankingOrder).toBeLessThan(getForConnectors.mock.invocationCallOrder[0]) expect(getForConnectors.mock.invocationCallOrder[0]).toBeLessThan( dbChainMockFns.select.mock.invocationCallOrder.at(-1)! @@ -1080,10 +1129,13 @@ describe('live repository authorization follows ranked candidates', () => { async (mode) => { getForConnectors.mockResolvedValue(identity) const candidates = [{ ...candidate('gmail', 'gmail-source'), installationSource: false }] - if (mode === 'vector' || mode === 'tag-vector') - queueTableRows(schemaMock.embedding, candidates) + if (mode === 'vector' || mode === 'tag-vector') { + probePages.push([{ id: 'doc-gmail' }]) + exactPages.push([{ id: 'gmail' }]) + queueRerank(candidates) + } if (mode === 'keyword') keywordPages.push(candidates) - else queueTableRows(schemaMock.embedding, candidates) + if (mode === 'tags') queueTableRows(schemaMock.embedding, candidates) const hydrated = [{ id: 'gmail', content: 'current permitted content' }] queueTableRows(schemaMock.embedding, hydrated) const searchParams = { ...params, filters: { source: 'gmail' } } @@ -1111,9 +1163,8 @@ describe('live repository authorization follows ranked candidates', () => { const hydration = JSON.stringify(dbChainMockFns.where.mock.calls.at(-1)![0]) expect(hydration).toContain('acl') expect(hydration).toContain('knowledgeConnectorMember') - expect(dbChainMockFns.select).toHaveBeenCalledTimes( - mode === 'keyword' ? 1 : mode === 'tags' ? 2 : 3 - ) + /** Vector ranking is raw SQL throughout; only hydration reads through the query builder. */ + expect(dbChainMockFns.select).toHaveBeenCalledTimes(mode === 'tags' ? 2 : 1) } ) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 341c62d91d8..d64af717c31 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -12,6 +12,7 @@ import { and, eq, inArray, isNull, type SQL, sql } from 'drizzle-orm' import { knowledgeAccessCondition, knowledgeMetadataCandidateAccessCondition, + textArrayLiteral, } from '@/lib/knowledge/access/predicate' import type { KnowledgeAccessProvider, KnowledgeAccessScope } from '@/lib/knowledge/access/types' import type { KbEmbeddingDimensions } from '@/lib/knowledge/embedding-models' @@ -50,20 +51,53 @@ const UNDEFINED_OBJECT_SQLSTATE = '42704' /** Bound candidate pages retained while live permissions are checked. */ const MAX_AUTHORIZED_SEARCH_CANDIDATES = 20_000 /** - * Bounds a permission-starved graph walk, which returns fewer candidates rather than widening. - * This approximate iterative-visit threshold excludes pgvector's initial scan; it is not a row limit. + * Approximate iterative-visit threshold for a permission-starved graph walk. It excludes + * pgvector's initial beam, so it only takes effect once `ResumeScanItems` starts widening. * - * Raising it trades recall for latency far more steeply than its size suggests. Measured on a - * search index where visibility admitted a tenth of the corpus, visiting four times as many tuples - * took 5.1s and 9.7s on consecutive identical runs, against ~115ms for the bounded walk — the walk - * degrades superlinearly with depth, and unpredictably. Re-measure before changing it. + * It must therefore stay roughly an order of magnitude above `ef_search`, or the first beam + * already exhausts the tuple budget and the scan stops before it can iterate at all — pgvector's + * maintainer says as much in pgvector#912. Measured on a production-shaped corpus, the previous + * pairing of a 1,000-wide beam against a 1,000-tuple budget returned fewer candidates than a + * narrower beam allowed to iterate, and spent longer inside the one uninterruptible beam. */ -const CANDIDATE_HNSW_MAX_SCAN_TUPLES = '1000' -const CANDIDATE_HNSW_EF_SEARCH = '1000' +const CANDIDATE_HNSW_MAX_SCAN_TUPLES = '20000' +/** + * Beam width per iteration. A beam is the granularity of cancellation: pgvector calls + * `CHECK_FOR_INTERRUPTS` only while building an index, never inside `hnswgettuple`, so neither + * `statement_timeout` nor a cancellation request can interrupt one. A narrower beam that iterates + * therefore bounds the leg's uninterruptible floor as well as widening its reach. + */ +const CANDIDATE_HNSW_EF_SEARCH = '200' const CANDIDATE_HNSW_SCAN_MEM_MULTIPLIER = '2' const MIN_VECTOR_RERANK_CANDIDATES = 400 const MAX_VECTOR_RERANK_CANDIDATES = 1600 const VECTOR_RERANK_OVERSAMPLING = 32 +/** + * The probe's share of the leg. It ranks nothing, so it must never be why the leg misses its + * own deadline. + * + * Its share comes out of what the rescue can claim from the narrowest live budget, + * `DIRECT_SEARCH_VECTOR_BUDGET_MS`, before live authorization, hydration and the exact rerank + * need the rest. + */ +const VECTOR_PROBE_BUDGET_MS = 600 +/** + * What one document costs the probe, measured on a corpus shaped like a search index under + * comparable cache pressure: the access predicate, evaluated once per document. + */ +const VECTOR_PROBE_MICROSECONDS_PER_DOCUMENT = 6 +/** + * Documents the probe enumerates before it concludes the permitted set is too large to rank + * exactly. Derived so that reaching it is what spends the probe's budget, rather than a separate + * number that a change to that budget could silently invalidate. + * + * It bounds the rescue's second step too: exact ranking of the `halfvec` projection measures at + * around half the probe's per-document cost, so a permitted set within this bound is affordable + * by construction. + */ +export const VECTOR_PROBE_DOCUMENT_LIMIT = Math.round( + (VECTOR_PROBE_BUDGET_MS * 1000) / VECTOR_PROBE_MICROSECONDS_PER_DOCUMENT +) /** How long to stop trying the iterative-scan settings after the server rejected them. */ const HNSW_SETTINGS_UNSUPPORTED_RETRY_MS = 10 * 60 * 1000 @@ -781,10 +815,54 @@ export async function handleVectorOnlySearch(params: SearchParams): Promise { + const probeBudget = budget?.capped(VECTOR_PROBE_BUDGET_MS) + try { + const probed = await runSearchQuery(probeBudget, 'vector.probe', (executor) => + executor.execute<{ id: string }>(sql` + SELECT ${document.id} AS id FROM ${document} + WHERE ${and(...conditions)} + LIMIT ${VECTOR_PROBE_DOCUMENT_LIMIT + 1} + `) + ) + if (probed.length > VECTOR_PROBE_DOCUMENT_LIMIT) return null + return probed.map(({ id }) => id) + } catch (error) { + if (!budget || !probeBudget?.isTimeout(error)) throw error + /** Only the probe's share was spent; the leg's own deadline still governs. */ + budget.remaining() + return null + } +} + +/** Tags live on chunks, so a row qualifies when a chunk it joins to carries them. */ +function chunkTagCondition(join: SQL, tagConditions: SQL[]): SQL | undefined { + if (!tagConditions.length) return undefined + return sql`EXISTS ( + SELECT 1 FROM ${embedding} + WHERE ${and(join, ...tagConditions)} + )` +} + +/** + * Select a bounded candidate pool and rerank it against the original vectors. + * + * A bounded ANN traversal fills that pool. When visibility leaves the traversal short of its + * limit, a bounded probe decides whether the permitted set is small enough to rank exactly + * instead, which recovers the candidates the traversal's post-filter discarded. + * * Live source authorization and content hydration still run after candidate ranking. */ async function selectVectorResults(params: SearchParams): Promise { @@ -796,13 +874,11 @@ async function selectVectorResults(params: SearchParams): Promise } | undefined return selectAuthorizedSearchResults({ leg: 'vector', access: params.access, @@ -836,21 +922,15 @@ async function selectVectorResults(params: SearchParams): Promise { + /** Explicit document IDs are already a bounded scope, and retain exhaustive ordering. */ + const exactPage = async () => { annotateSearchDiagnostics({ vectorRanking: 'exact' }) const candidates = await runSearchQuery(params.budget, 'vector.exact', (executor) => executor .select({ ...SEARCH_READ_CANDIDATE_FIELDS, distance: distance.as('distance') }) .from(embedding) .innerJoin(document, eq(embedding.documentId, document.id)) - .where( - and( - ...conditions, - ...visibility, - candidateIds ? inArray(embedding.id, candidateIds) : undefined - ) - ) + .where(and(...conditions, ...visibility)) .orderBy(sql`(${distance}) + 0`, embedding.id) .limit(limit) .offset(offset) @@ -858,69 +938,25 @@ async function selectVectorResults(params: SearchParams): Promise - tagConditions.length - ? executor - .select({ id: embedding.id }) - .from(embedding) - .innerJoin(document, eq(document.id, embedding.documentId)) - .where( - and( - inArray(embedding.knowledgeBaseId, params.knowledgeBaseIds), - ...visibility, - ...tagConditions - ) - ) - .limit(candidateLimit) - : executor.execute<{ id: string }>(sql` - SELECT scoped_chunk.id FROM ${document} - CROSS JOIN LATERAL ( - SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} - WHERE ${and( - eq(embeddingSearch.documentId, document.id), - inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds), - eq(embeddingSearch.enabled, true) - )} - LIMIT ${candidateLimit} - ) AS scoped_chunk - WHERE ${and(...candidateDocumentVisibility)} - LIMIT ${candidateLimit} - `) - ) - if (probe.length === 0) return { candidates: [], nextOffset: offset } - if (probe.length < candidateLimit) { - return exactPage(probe.map((candidate) => candidate.id)) - } - annotateSearchDiagnostics({ - vectorRanking: 'candidate-rerank', - vectorCandidateStorage: 'stored-halfvec', - vectorCandidateLimit: candidateLimit, - vectorCandidateScan: 'planned', - vectorCandidateDimensions: embeddingCandidateDimensions( - queryVector.dimensions, - queryVector.model - ), - }) - /** - * The bounded ANN traversal is the whole candidate set. LIMIT keeps document authorization - * downstream of the traversal, with a primary-key lookup per candidate. - * - * An underfilled traversal yields fewer candidates rather than widening the search. Widening - * it has no affordable form here: rescoring the projection exhaustively is O(corpus) and a - * deeper `hnsw.max_scan_tuples` is worse still. Measured where visibility admitted a tenth of - * the corpus, the exhaustive rescan took 1.9s while visiting four times as many tuples took - * 5.1s and 9.7s on consecutive identical runs. Both exceed the retrieval budget once the - * corpus grows, and a leg that exceeds its budget returns nothing at all, so fewer candidates - * strictly beats every widening strategy available. - */ - const identities = await withVectorScanSettings( - (executor) => - executor.execute<{ id: string }>(sql` + const excludedKey = excludedSources.join('\u0000') + if (candidatePool?.excludedKey !== excludedKey) { + annotateSearchDiagnostics({ + vectorRanking: 'candidate-rerank', + vectorCandidateStorage: 'stored-halfvec', + vectorCandidateLimit: candidateLimit, + vectorCandidateScan: 'planned', + vectorCandidateDimensions: embeddingCandidateDimensions( + queryVector.dimensions, + queryVector.model + ), + }) + /** + * The bounded ANN traversal is the whole candidate set. LIMIT keeps document + * authorization downstream of the traversal, with a primary-key lookup per candidate. + */ + const traversed = await withVectorScanSettings( + (executor) => + executor.execute<{ id: string }>(sql` SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} CROSS JOIN LATERAL ( SELECT 1 FROM ${document} @@ -933,12 +969,59 @@ async function selectVectorResults(params: SearchParams): Promise = traversed + if (traversed.length < candidateLimit) { + const visibleDocumentIds = await probeVisibleDocuments( + [...candidateDocumentVisibility, documentTagCondition], + params.budget + ) + if (visibleDocumentIds) { + annotateSearchDiagnostics({ + vectorRanking: 'exact-candidates', + vectorProbeDocumentCount: visibleDocumentIds.length, + }) + /** + * `+ 0` keeps the planner off the ANN index, and the probed identities keep the scan + * on `embedding_search_document_lookup_idx`, so this reads what the permitted set + * costs rather than re-deriving permission across the whole index. Exact ranking also + * honours `statement_timeout`, which a traversal cannot. + */ + selected = visibleDocumentIds.length + ? await runSearchQuery(params.budget, 'vector.exact_candidates', (executor) => + executor.execute<{ id: string }>(sql` + SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch} + WHERE ${and( + inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds), + eq(embeddingSearch.enabled, true), + sql`${embeddingSearch.documentId} = ANY(${textArrayLiteral(visibleDocumentIds)})`, + candidateTagCondition + )} + ORDER BY (${candidateDistance}) + 0 LIMIT ${candidateLimit} + `) + ) + : [] + } + } + candidatePool = { excludedKey, identities: selected } + annotateSearchDiagnostics({ + vectorCandidateCount: selected.length, + vectorCandidateScan: selected.length < candidateLimit ? 'underfilled' : 'planned', + }) + } + const { identities } = candidatePool if (!identities.length) return { candidates: [], nextOffset: offset } /** Score each bounded candidate once; sorting the materialized scalar cannot invoke HNSW again. */ const page = await runSearchQuery(params.budget, 'vector.rerank', (executor) => From 18279e4e9b727750a41949fb727dd1901a81c0da Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 18 Sep 2026 01:48:32 -0700 Subject: [PATCH 2/3] test(knowledge): cover the capped step budget --- apps/sim/lib/knowledge/search/budget.test.ts | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/apps/sim/lib/knowledge/search/budget.test.ts b/apps/sim/lib/knowledge/search/budget.test.ts index 941fbdbd98e..40f6d7a791e 100644 --- a/apps/sim/lib/knowledge/search/budget.test.ts +++ b/apps/sim/lib/knowledge/search/budget.test.ts @@ -44,6 +44,23 @@ describe('search SQL deadline', () => { } }) + it('caps a step without letting it shorten or outlive the leg it came from', () => { + vi.spyOn(performance, 'now').mockReturnValue(0) + const controller = new AbortController() + const leg = new SearchBudget('vector', 1000, controller.signal) + const step = leg.capped(600) + expect(step.deadline).toBe(600) + expect(step.leg).toBe('vector') + expect(step.signal).toBe(controller.signal) + /** Spending the step is the step's own business; the leg keeps its deadline and its verdict. */ + step.isTimeout(new SearchDeadlineError()) + expect(step.timedOut).toBe(true) + expect(leg.timedOut).toBe(false) + expect(leg.remaining()).toBe(1000) + /** A step longer than what the leg has left cannot extend it. */ + expect(new SearchBudget('vector', 200).capped(600).deadline).toBe(200) + }) + it('preserves cancellation and unexpected errors instead of labeling them incomplete evidence', async () => { const controller = new AbortController() const budget = new SearchBudget('keyword', performance.now() - 1, controller.signal) From a32fbb06e4a74d8db1285216929035e403c05b01 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 18 Sep 2026 02:02:19 -0700 Subject: [PATCH 3/3] fix(knowledge): require an enabled chunk when a tag filter decides a document A document whose only tagged chunk is disabled could be admitted by the probe and then discarded by ranking, spending the probe's document bound on a document that can contribute no candidate. Update the KB block fan-out integration test to the strategy it now exercises: a scope too small to fill the traversal probes once and rescues once. --- .../kb-block-search.integration.ts | 25 ++++++++++++++++--- apps/sim/lib/knowledge/search/queries.test.ts | 18 +++++++++++++ apps/sim/lib/knowledge/search/queries.ts | 8 ++++-- 3 files changed, 45 insertions(+), 6 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts b/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts index b54266037d7..5abf70817cd 100644 --- a/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts @@ -91,7 +91,7 @@ describe('API-key KB block fan-out', () => { const previousDebug = db.$client.options.debug const statements: string[] = [] db.$client.options.debug = (_connection, query) => { - if (statements.length < 250) statements.push(query) + if (statements.length < 1000) statements.push(query) } try { const results = await Promise.all( @@ -124,9 +124,26 @@ describe('API-key KB block fan-out', () => { expect(result.rows[0].knowledgeBaseId).toBe(bases[index].id) expect(result.rows[0].distance).toBeCloseTo(0) } - expect(statements.filter((query) => query.includes('statement_timeout'))).toHaveLength(54) - expect(statements.filter((query) => query.includes('+ 0'))).toHaveLength(18) - expect(statements.some((query) => query.includes('hnsw.iterative_scan'))).toBe(false) + const matching = (fragment: string) => + statements.filter((query) => query.includes(fragment)) + /** + * Every statement runs under the leg's deadline: the candidate search reinstates it after + * tuning the scan, and the probe, the exact ranking, the rerank and hydration each open + * with one of their own. + */ + expect(matching('statement_timeout')).toHaveLength(bases.length * 6) + /** + * A scope this small leaves the bounded traversal short of its candidate limit, so every + * search probes once and rescues once — never a widening retry loop. + */ + expect(matching('hnsw.iterative_scan')).toHaveLength(bases.length) + expect(matching('AS visible')).toHaveLength(bases.length) + expect(matching(') + 0 LIMIT')).toHaveLength(bases.length) + expect(matching('scored_search_candidates')).toHaveLength(bases.length) + /** The probe enumerates visible documents; it never ranks them. */ + expect( + statements.filter((query) => query.includes('AS id FROM') && !query.includes('ORDER BY')) + ).toHaveLength(bases.length) } finally { db.$client.options.debug = previousDebug } diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index f38447398a7..ce1f0031c63 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -437,6 +437,24 @@ describe('workspace-scoped vector retrieval', () => { expect(getForConnectors).not.toHaveBeenCalled() }) + it('counts only chunks the search can return when a tag filter decides a document', async () => { + traversedRows = ranked + probeRows = [{ id: 'near-doc' }] + exactRows = ranked + queueTableRows(schemaMock.embedding, [...ranked].reverse()) + await handleTagAndVectorSearch({ + ...params, + structuredFilters: [{ tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'common' }], + }) + /** + * A document whose only tagged chunk is disabled contributes no candidate, so admitting it + * would spend the probe's document bound on a document the ranking then discards. + */ + const probe = JSON.stringify(statements().find((query) => isProbeStatement(query.sql))!) + expect(probe).toContain(String(schemaMock.embedding.tag1)) + expect(probe).toContain(`"left":"${schemaMock.embedding.enabled}","right":true`) + }) + it('keeps the tuple budget an order of magnitude above the beam so the scan can iterate', async () => { queueTableRows(schemaMock.embedding, [...ranked].reverse()) await handleVectorOnlySearch(params) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index d64af717c31..bb5ae17f674 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -847,12 +847,16 @@ async function probeVisibleDocuments( } } -/** Tags live on chunks, so a row qualifies when a chunk it joins to carries them. */ +/** + * Tags live on chunks, so a row qualifies when a chunk it joins to carries them — and only a + * chunk the search can actually return counts, or a document whose sole match is disabled would + * be admitted by a check that ranking then discards. + */ function chunkTagCondition(join: SQL, tagConditions: SQL[]): SQL | undefined { if (!tagConditions.length) return undefined return sql`EXISTS ( SELECT 1 FROM ${embedding} - WHERE ${and(join, ...tagConditions)} + WHERE ${and(join, eq(embedding.enabled, true), ...tagConditions)} )` }