From fd418a3194735dbf1d301b34e179eac9b912e444 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 21 Sep 2026 11:20:15 -0700 Subject: [PATCH 1/2] fix(knowledge): give the reach count the probe's share of the deadline, not the leg's Counting a caller's reach reads as many index entries as they reach, so on a large index it took the whole vector budget and left the leg no time to walk; the search then reported the leg timed out with the count's stage measured twice. The count now runs under the probe's share of the deadline, a count that runs out answers nothing for this search without marking the leg, and the reach memo holds for an hour: reach moves slowly and a stale answer costs speed, never access. --- apps/sim/lib/knowledge/search/queries.test.ts | 5 +- apps/sim/lib/knowledge/search/queries.ts | 100 ++++++++++-------- 2 files changed, 58 insertions(+), 47 deletions(-) diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index 6e76b50d677..b3735836741 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -2075,9 +2075,10 @@ describe('permitted-document planner', () => { ) expect(reach).toEqual({ kind: 'unbounded', broad: false }) expect(reachCounts()).toHaveLength(1) - /** The count is the search's own read: it runs inside the leg's deadline statement. */ + /** The count is the search's own read, under the probe's share of the deadline, not the leg's. */ const countAt = statements().findIndex((query) => query.sql.includes(') reached')) expect(statements()[countAt - 1].sql).toContain('statement_timeout') + expect(Number(statements()[countAt - 1].params[0])).toBeLessThanOrEqual(600) }) it('does not remember a reach whose count ran out of time', async () => { @@ -2102,6 +2103,8 @@ describe('permitted-document planner', () => { }) expect(plan).toEqual({ kind: 'unbounded', broad: true }) expect(reachCounts()).toHaveLength(1) + /** Only the count's share of the deadline was spent; the leg is not the one that timed out. */ + expect(budget.timedOut).toBe(false) /** The next search counts again rather than trusting an answer that never came. */ await resolveReach( ['org-index'], diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index af01100b26e..3c6c71f2ce1 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -1180,9 +1180,11 @@ export const BROAD_REACH_SHARE = 0.25 /** * 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. + * index with the full access predicate, so a stale answer costs speed, never access. Counting it + * is the one read of a search that scales with the caller's reach rather than the query, so it is + * remembered for long. */ -const SATURATED_REACH_TTL_MS = 5 * 60 * 1000 +const SATURATED_REACH_TTL_MS = 60 * 60 * 1000 /** A saturated reach, and whether it is broad enough to walk the whole graph for. */ const saturatedReach = new LRUCache({ @@ -1253,6 +1255,11 @@ async function estimateFilteredDocuments( * documents. Counted once against that bound and remembered, so the first search after the * window pays for it and the rest do not. A caller whose probe already saturated is known to * reach past the probe's limit, so a bound inside that limit is met without counting. + * + * The count reads as many index entries as the caller reaches, so on a large index it can cost + * more than the leg it serves; it gets the probe's share of the deadline, never the whole leg's. + * A count that runs out of that share answers `null`: the leg keeps its time and its deadline + * intact, and the caller decides this search alone without remembering anything. */ async function reachIsBroad( knowledgeBaseIds: string[], @@ -1260,28 +1267,36 @@ async function reachIsBroad( budget: SearchBudget | undefined, plan: SearchAccessPlan | undefined, saturated: boolean -): Promise { +): Promise { if (access.kind !== 'user') return true - const total = - (await indexDocumentCounts.fetch([...knowledgeBaseIds].sort().join(','), { - context: budget, - })) ?? 0 - const bound = Math.ceil(total * BROAD_REACH_SHARE) - if (saturated && bound <= VECTOR_PROBE_DOCUMENT_LIMIT) return true - const [row] = await runSearchQuery(budget, 'permitted_documents', (executor) => - executor.execute<{ n: number }>(sql` - SELECT count(*) AS n FROM ( - SELECT 1 FROM ${document} - WHERE ${and( - isNull(document.deletedAt), - knowledgeAclOverlapCondition(access), - inArray(document.knowledgeBaseId, knowledgeBaseIds), - planSourceCondition(plan) - )} - LIMIT ${bound} - ) reached`) - ) - return Number(row?.n ?? 0) >= bound + const countBudget = budget?.capped(VECTOR_PROBE_BUDGET_MS) + try { + const total = + (await indexDocumentCounts.fetch([...knowledgeBaseIds].sort().join(','), { + context: countBudget, + })) ?? 0 + const bound = Math.ceil(total * BROAD_REACH_SHARE) + if (saturated && bound <= VECTOR_PROBE_DOCUMENT_LIMIT) return true + const [row] = await runSearchQuery(countBudget, 'permitted_documents', (executor) => + executor.execute<{ n: number }>(sql` + SELECT count(*) AS n FROM ( + SELECT 1 FROM ${document} + WHERE ${and( + isNull(document.deletedAt), + knowledgeAclOverlapCondition(access), + inArray(document.knowledgeBaseId, knowledgeBaseIds), + planSourceCondition(plan) + )} + LIMIT ${bound} + ) reached`) + ) + return Number(row?.n ?? 0) >= bound + } catch (error) { + if (!budget || !countBudget?.isTimeout(error)) throw error + /** Only the count's share was spent; the leg's own deadline still governs. */ + budget.remaining() + return null + } } /** @@ -1331,15 +1346,11 @@ export async function resolveReach( const key = reachKey(knowledgeBaseIds, access, plan) const remembered = key ? saturatedReach.get(key) : undefined if (remembered) return { kind: 'unbounded', broad: remembered.broad } - try { - const broad = await reachIsBroad(knowledgeBaseIds, access, budget, plan, false) - if (key) saturatedReach.set(key, { broad }) - return { kind: 'unbounded', broad } - } catch (error) { - if (!budget?.isTimeout(error)) throw error - /** A count that ran out of time decides this search only; the next one counts again. */ - return { kind: 'unbounded', broad: true } - } + const broad = await reachIsBroad(knowledgeBaseIds, access, budget, plan, false) + /** A count that ran out of time decides this search only; the next one counts again. */ + if (broad === null) return { kind: 'unbounded', broad: true } + if (key) saturatedReach.set(key, { broad }) + return { kind: 'unbounded', broad } } /** @@ -1392,18 +1403,17 @@ export async function resolvePermittedDocuments(params: { probe = { kind: 'timed_out' } } if (probe.kind === 'saturated') { - try { - broad = await reachIsBroad( - params.knowledgeBaseIds, - params.access, - params.budget, - params.accessPlan, - true - ) + const counted = await reachIsBroad( + params.knowledgeBaseIds, + params.access, + params.budget, + params.accessPlan, + true + ) + /** A count that ran out of time decides this search only; the next one counts again. */ + if (counted !== null) { + broad = counted if (key) saturatedReach.set(key, { broad }) - } catch (error) { - if (!params.budget?.isTimeout(error)) throw error - /** A count that ran out of time decides this search only; the next one counts again. */ } } } @@ -2658,9 +2668,7 @@ export async function retrieveKnowledgeSearch( * readable documents enumerated ahead of ranking: its reach alone chooses between one * walk over the whole graph and a search of each source. */ - await measureSearchStage('permitted_documents', () => - resolveReach(knowledgeBaseIds, access, budgets.vector, accessPlan) - ) + await resolveReach(knowledgeBaseIds, access, budgets.vector, accessPlan) : await resolvePermittedDocuments({ knowledgeBaseIds, access, From 1ea4c24d8f17e5a138b2fcfb5b45c9a2c0f39af9 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 21 Sep 2026 11:26:52 -0700 Subject: [PATCH 2/2] fix(knowledge): keep the reach memo at five minutes and treat a spent leg during the count as short A stale strategy costs recall at the margin, so the memo stays at five minutes now that a count costs at most the probe's share. A leg whose own deadline passes during the count reports short, as it did before, instead of failing the search. --- apps/sim/lib/knowledge/search/queries.test.ts | 14 ++++++ apps/sim/lib/knowledge/search/queries.ts | 49 +++++++++++-------- 2 files changed, 43 insertions(+), 20 deletions(-) diff --git a/apps/sim/lib/knowledge/search/queries.test.ts b/apps/sim/lib/knowledge/search/queries.test.ts index b3735836741..8dc22757afb 100644 --- a/apps/sim/lib/knowledge/search/queries.test.ts +++ b/apps/sim/lib/knowledge/search/queries.test.ts @@ -2051,6 +2051,20 @@ describe('permitted-document planner', () => { expect(reachCounts()).toHaveLength(2) }) + it('reports a leg whose own deadline passed during the count as short, not failed', async () => { + const budget = new SearchBudget('vector', performance.now() - 1) + await expect( + resolveReach(['org-index'], scope('spent-leg'), budget, { + connectors: { workspace: [], admin: [], members: [], liveProofRequired: [] }, + observers: { confirmed: [], observed: [] }, + memberSources: [], + connectorTypes: new Map(), + uploads: true, + }) + ).resolves.toEqual({ kind: 'unbounded', broad: true }) + expect(budget.timedOut).toBe(true) + }) + it('counts a resolved reach against a small index instead of assuming it broad', async () => { /** A bound inside the probe limit proves nothing without a saturated probe. */ dbChainMockFns.execute.mockImplementation(async (query) => { diff --git a/apps/sim/lib/knowledge/search/queries.ts b/apps/sim/lib/knowledge/search/queries.ts index 3c6c71f2ce1..4ae24ac79ad 100644 --- a/apps/sim/lib/knowledge/search/queries.ts +++ b/apps/sim/lib/knowledge/search/queries.ts @@ -1180,11 +1180,9 @@ export const BROAD_REACH_SHARE = 0.25 /** * 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. Counting it - * is the one read of a search that scales with the caller's reach rather than the query, so it is - * remembered for long. + * index with the full access predicate, so a stale answer costs speed, never access. */ -const SATURATED_REACH_TTL_MS = 60 * 60 * 1000 +const SATURATED_REACH_TTL_MS = 5 * 60 * 1000 /** A saturated reach, and whether it is broad enough to walk the whole graph for. */ const saturatedReach = new LRUCache({ @@ -1346,11 +1344,17 @@ export async function resolveReach( const key = reachKey(knowledgeBaseIds, access, plan) const remembered = key ? saturatedReach.get(key) : undefined if (remembered) return { kind: 'unbounded', broad: remembered.broad } - const broad = await reachIsBroad(knowledgeBaseIds, access, budget, plan, false) - /** A count that ran out of time decides this search only; the next one counts again. */ - if (broad === null) return { kind: 'unbounded', broad: true } - if (key) saturatedReach.set(key, { broad }) - return { kind: 'unbounded', broad } + try { + const broad = await reachIsBroad(knowledgeBaseIds, access, budget, plan, false) + /** A count that ran out of time decides this search only; the next one counts again. */ + if (broad === null) return { kind: 'unbounded', broad: true } + if (key) saturatedReach.set(key, { broad }) + return { kind: 'unbounded', broad } + } catch (error) { + /** The leg's own deadline passed during the count: the leg is short, the search is not failed. */ + if (!budget?.isTimeout(error)) throw error + return { kind: 'unbounded', broad: true } + } } /** @@ -1403,17 +1407,22 @@ export async function resolvePermittedDocuments(params: { probe = { kind: 'timed_out' } } if (probe.kind === 'saturated') { - const counted = await reachIsBroad( - params.knowledgeBaseIds, - params.access, - params.budget, - params.accessPlan, - true - ) - /** A count that ran out of time decides this search only; the next one counts again. */ - if (counted !== null) { - broad = counted - if (key) saturatedReach.set(key, { broad }) + try { + const counted = await reachIsBroad( + params.knowledgeBaseIds, + params.access, + params.budget, + params.accessPlan, + true + ) + /** A count that ran out of time decides this search only; the next one counts again. */ + if (counted !== null) { + broad = counted + if (key) saturatedReach.set(key, { broad }) + } + } catch (error) { + /** The leg's own deadline passed during the count: the leg is short, the search is not failed. */ + if (!params.budget?.isTimeout(error)) throw error } } }