diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 7d5cbfbab88..acb205cda09 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -1529,13 +1529,57 @@ describe('permitted-document planner', () => { probeRows = [...rows] const permitted = await resolvePermittedDocuments({ knowledgeBaseIds: ['org-index'], - access: reader, + access: { ...reader, tokens: [`u:resolves-${kind}@example.com`] }, }) expect(permitted.kind).toBe(kind) if (permitted.kind === 'bounded') expect(permitted.documents).toEqual([{ id: 'doc-a', connectorId: null }]) }) + describe('saturated reach', () => { + const scope = (name: string): UserAccessScope => ({ + ...reader, + tokens: [`u:${name}@example.com`], + }) + const resolve = (access: UserAccessScope, knowledgeBaseIds = ['org-index']) => + resolvePermittedDocuments({ knowledgeBaseIds, access }) + const probes = () => statements().filter((query) => isProbeStatement(query.sql)).length + + it('is remembered, so a broad caller skips the probe on the next search', async () => { + probeRows = [{ id: null, connectorId: null, saturated: true }] + const broad = scope('broad') + expect((await resolve(broad)).kind).toBe('unbounded') + expect((await resolve({ ...broad, tokens: [...broad.tokens].reverse() })).kind).toBe( + 'unbounded' + ) + expect(probes()).toBe(1) + }) + + it('is remembered per set of bases and tokens', async () => { + probeRows = [{ id: null, connectorId: null, saturated: true }] + await resolve(scope('per-key')) + probeRows = [{ id: 'doc-a', connectorId: null, saturated: false }] + expect((await resolve(scope('per-key'), ['other-index'])).kind).toBe('bounded') + expect((await resolve(scope('per-key-other'))).kind).toBe('bounded') + expect(probes()).toBe(3) + }) + + it('is not inferred from a bounded set or a probe that ran out of time', async () => { + probeRows = [{ id: 'doc-a', connectorId: null, saturated: false }] + await resolve(scope('bounded')) + await resolve(scope('bounded')) + expect(probes()).toBe(2) + const budget = new SearchBudget('vector', performance.now() - 1) + await resolvePermittedDocuments({ + knowledgeBaseIds: ['org-index'], + access: scope('timed-out'), + budget, + }) + probeRows = [{ id: 'doc-a', connectorId: null, saturated: false }] + expect((await resolve(scope('timed-out'))).kind).toBe('bounded') + }) + }) + it('reports an exhausted vector budget as unbounded instead of failing both legs', async () => { const budget = new SearchBudget('vector', performance.now() - 1) const permitted = await resolvePermittedDocuments({ diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 9152057475c..edb7038ef47 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -7,8 +7,10 @@ import { knowledgeConnector, } from '@sim/db/schema' import { createLogger } from '@sim/logger' +import { sha256Hex } from '@sim/security/hash' import { getErrorMessage, getPostgresErrorCode } from '@sim/utils/errors' import { and, eq, inArray, isNull, type SQL, sql } from 'drizzle-orm' +import { LRUCache } from 'lru-cache' import { knowledgeAccessCondition, knowledgeAclOverlapCondition, @@ -833,15 +835,21 @@ export async function handleVectorOnlySearch(params: SearchParams): Promise { +): Promise { const probeBudget = budget?.capped(VECTOR_PROBE_BUDGET_MS) try { const probed = await runSearchQuery(probeBudget, stage, (executor) => @@ -858,13 +866,18 @@ async function probeVisibleDocuments( ) ) /** The saturation sentinel is only ever emitted alone. */ - if (probed.length > VECTOR_PROBE_DOCUMENT_LIMIT || probed[0]?.saturated) return null - return probed.map(({ id, connectorId }) => ({ id, connectorId })) + if (probed.length > VECTOR_PROBE_DOCUMENT_LIMIT || probed[0]?.saturated) { + return { kind: 'saturated' } + } + return { + kind: 'documents', + documents: probed.map(({ id, connectorId }) => ({ id, connectorId })), + } } 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 + return { kind: 'timed_out' } } } @@ -949,6 +962,24 @@ export type PermittedDocuments = | { kind: 'bounded'; documents: readonly PermittedDocument[] } | { kind: 'unbounded' } +/** + * How long a caller's saturated reach is remembered. Reach counts the documents a caller's tokens + * touch in the bases, which moves slowly, and an unbounded set only means the legs search the + * index with the full access predicate, so a stale answer costs speed, never access. + */ +const SATURATED_REACH_TTL_MS = 5 * 60 * 1000 + +const saturatedReach = new LRUCache({ max: 10_000, ttl: SATURATED_REACH_TTL_MS }) + +/** Reach depends only on the bases and the caller's tokens; filters narrow the set, not the reach. */ +function reachKey( + knowledgeBaseIds: readonly string[], + access: KnowledgeAccessScope +): string | null { + if (access.kind !== 'user') return null + return `${[...knowledgeBaseIds].sort().join(',')}:${sha256Hex([...access.tokens].sort().join('\n'))}` +} + /** * Resolve the permitted set with the candidate predicate both legs apply, so restricting a leg * to it never admits a document the leg would otherwise refuse. Tag filters stay chunk-level in @@ -963,30 +994,37 @@ export async function resolvePermittedDocuments(params: { filters?: WorkspaceSearchFilters budget?: SearchBudget }): Promise { - let documents: PermittedDocument[] | null - try { - documents = await probeVisibleDocuments( - params.knowledgeBaseIds, - candidateDocumentConditions( + const key = reachKey(params.knowledgeBaseIds, params.access) + let probe: ProbeOutcome + if (key && saturatedReach.get(key)) { + probe = { kind: 'saturated' } + } else { + try { + probe = await probeVisibleDocuments( params.knowledgeBaseIds, + candidateDocumentConditions( + params.knowledgeBaseIds, + params.access, + params.filters, + knowledgeMetadataCandidateAccessCondition(params.access) + ), params.access, - params.filters, - knowledgeMetadataCandidateAccessCondition(params.access) - ), - params.access, - params.budget, - 'permitted_documents' - ) - } catch (error) { - if (!params.budget?.isTimeout(error)) throw error - documents = null + params.budget, + 'permitted_documents' + ) + } catch (error) { + if (!params.budget?.isTimeout(error)) throw error + probe = { kind: 'timed_out' } + } + if (key && probe.kind === 'saturated') saturatedReach.set(key, true) } - const permitted: PermittedDocuments = documents - ? { kind: 'bounded', documents } - : { kind: 'unbounded' } + const permitted: PermittedDocuments = + probe.kind === 'documents' + ? { kind: 'bounded', documents: probe.documents } + : { kind: 'unbounded' } annotateSearchDiagnostics({ permittedDocuments: permitted.kind, - ...(documents ? { permittedDocumentCount: documents.length } : {}), + ...(probe.kind === 'documents' ? { permittedDocumentCount: probe.documents.length } : {}), }) return permitted } @@ -1180,16 +1218,16 @@ async function selectVectorResults(params: SearchParams): Promise id)) + if (probe.kind === 'documents') { + annotateSearchDiagnostics({ vectorProbeDocumentCount: probe.documents.length }) + selected = await rankPermittedExactly(probe.documents.map(({ id }) => id)) } } }