From 8b1a096b690cfc9317ad19ffdd82349c4ce2ea77 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 18 Sep 2026 23:57:00 -0700 Subject: [PATCH 1/2] improvement(knowledge): remember a caller's saturated search reach A caller whose tokens reach more documents than the permitted-set limit paid the reach count on every search only to learn again that the set is unbounded. That answer is now remembered per bases and token set for five minutes. The probe now reports saturation apart from a timeout, and only saturation is remembered; an unbounded set only means the legs apply the full access predicate per candidate, so a stale answer costs speed, never access. --- apps/sim/lib/knowledge/search/queries.test.ts | 46 ++++++++- apps/sim/lib/knowledge/search/queries.ts | 98 +++++++++++++------ 2 files changed, 113 insertions(+), 31 deletions(-) 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..9c8bc37b049 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' } } } @@ -957,36 +970,61 @@ export type PermittedDocuments = * It runs ahead of both legs on the vector leg's budget, so exhausting that budget here reports * `unbounded` and marks the vector leg timed out rather than failing the keyword leg with it. */ +/** + * 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'))}` +} + export async function resolvePermittedDocuments(params: { knowledgeBaseIds: string[] access: KnowledgeAccessScope 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)) } } } From 3cf9dd1248040bae3f5bb07bb83222601939661b Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sat, 19 Sep 2026 00:01:26 -0700 Subject: [PATCH 2/2] improvement(knowledge): keep the permitted-set resolver's TSDoc on its function --- apps/sim/lib/knowledge/search/queries.ts | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 9c8bc37b049..edb7038ef47 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -962,14 +962,6 @@ export type PermittedDocuments = | { kind: 'bounded'; documents: readonly PermittedDocument[] } | { kind: 'unbounded' } -/** - * 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 - * each leg; the set is the document-level superset they narrow. - * - * It runs ahead of both legs on the vector leg's budget, so exhausting that budget here reports - * `unbounded` and marks the vector leg timed out rather than failing the keyword leg with it. - */ /** * 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 @@ -988,6 +980,14 @@ function reachKey( 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 + * each leg; the set is the document-level superset they narrow. + * + * It runs ahead of both legs on the vector leg's budget, so exhausting that budget here reports + * `unbounded` and marks the vector leg timed out rather than failing the keyword leg with it. + */ export async function resolvePermittedDocuments(params: { knowledgeBaseIds: string[] access: KnowledgeAccessScope