From f7949f3c9f6d82bdda11c0f702afb95d77cda537 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 10 Sep 2026 19:35:42 -0700 Subject: [PATCH] fix(knowledge): count list documents once instead of per 400-base batch #7730 switched the knowledge base lists to a KnowledgeAccessProvider, which turned the single joined count into a per-400-base loop of count plus candidate-discovery queries. The staging integ workspace holds 21k archived bases, so GET /api/knowledge?scope=archived went from ~4s to ~100s and the archive-kb integ check timed out at 30s. The stored-ACL count moves back into the list's own join. Only documents a live source (GitHub, Confluence) authorizes beyond that are counted afterwards, once for the whole list, bounded by the list filter (unpaged) or the page's ids. The provider now reports whether the reader holds any live-source credential, and the shared batch generator skips candidate discovery entirely when it does not, since the discovered predicate would be provably empty. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Lm5mdDU9bi1vb4oAvPtbjn --- apps/sim/lib/knowledge/access/scope.test.ts | 19 +++ apps/sim/lib/knowledge/access/scope.ts | 4 + apps/sim/lib/knowledge/access/types.ts | 6 + apps/sim/lib/knowledge/read-access.test.ts | 15 +++ apps/sim/lib/knowledge/read-access.ts | 1 + apps/sim/lib/knowledge/service.test.ts | 55 ++++++++- apps/sim/lib/knowledge/service.ts | 124 ++++++++++++-------- 7 files changed, 174 insertions(+), 50 deletions(-) diff --git a/apps/sim/lib/knowledge/access/scope.test.ts b/apps/sim/lib/knowledge/access/scope.test.ts index 0b1ecc25723..c91a38b0ae5 100644 --- a/apps/sim/lib/knowledge/access/scope.test.ts +++ b/apps/sim/lib/knowledge/access/scope.test.ts @@ -266,6 +266,25 @@ describe('createKnowledgeAccessProvider', () => { expect(dbChainMockFns.select).toHaveBeenCalledTimes(2) }) + it('reports live-source readers only when a member-scoped source credential exists', async () => { + queueSubjects([{ providerId: 'slack', providerTenantId: 'T1', providerSubjectId: 'U1' }]) + await expect( + createKnowledgeAccessProvider(SESSION, WORKSPACE).hasLiveSourceReaders?.() + ).resolves.toBe(false) + + queueSubjects([ + { + providerId: 'confluence', + providerTenantId: 'site-1', + providerSubjectId: 'account-1', + credentialId: 'credential-1', + }, + ]) + await expect( + createKnowledgeAccessProvider(SESSION, WORKSPACE).hasLiveSourceReaders?.() + ).resolves.toBe(true) + }) + it('retries after a failed lookup rather than caching the failure', async () => { dbChainMockFns.where.mockRejectedValueOnce(new Error('connection reset')) const provider = createKnowledgeAccessProvider(SESSION, WORKSPACE) diff --git a/apps/sim/lib/knowledge/access/scope.ts b/apps/sim/lib/knowledge/access/scope.ts index e70c62854ff..31ce1201124 100644 --- a/apps/sim/lib/knowledge/access/scope.ts +++ b/apps/sim/lib/knowledge/access/scope.ts @@ -397,6 +397,10 @@ function createAccessProvider( async get() { return (await identity()).access }, + async hasLiveSourceReaders() { + const { access, githubReaders, confluenceReaders } = await identity() + return access.kind === 'user' && (githubReaders.length > 0 || confluenceReaders.length > 0) + }, async getForConnectors(connectorIds, signal) { const ids = boundedIds(connectorIds) const cancellation = diff --git a/apps/sim/lib/knowledge/access/types.ts b/apps/sim/lib/knowledge/access/types.ts index 88995f2d3eb..028ee9efc04 100644 --- a/apps/sim/lib/knowledge/access/types.ts +++ b/apps/sim/lib/knowledge/access/types.ts @@ -93,6 +93,12 @@ export interface KnowledgeAccessProvider { documentIds: readonly string[], signal?: AbortSignal ): Promise + /** + * Whether the reader holds any credential a live source (GitHub, Confluence) could + * authorize beyond the stored ACL. Without one, candidate discovery can only re-prove + * the ordinary predicate, so readers skip it. Absent means unknown: discover. + */ + hasLiveSourceReaders?(): Promise } /** Two existing search legs each contribute at most 200 candidates to one authorization batch. */ diff --git a/apps/sim/lib/knowledge/read-access.test.ts b/apps/sim/lib/knowledge/read-access.test.ts index 9d0d1179fab..47bda229ece 100644 --- a/apps/sim/lib/knowledge/read-access.test.ts +++ b/apps/sim/lib/knowledge/read-access.test.ts @@ -49,6 +49,21 @@ describe('knowledgeReadAccessBatches', () => { expect(gt).toHaveBeenCalledWith(document.connectorId, first.at(-1)!.connectorId) }) + it('yields only the ordinary predicate for a reader without live-source credentials', async () => { + const resolve = vi.fn(async () => identity) + const provider: KnowledgeAccessProvider = { + get: async () => identity, + getForConnectors: resolve, + getForDocuments: async () => identity, + hasLiveSourceReaders: async () => false, + } + const batches = [] + for await (const predicate of knowledgeReadAccessBatches(provider, [])) batches.push(predicate) + expect(batches).toHaveLength(1) + expect(resolve).not.toHaveBeenCalled() + expect(dbChainMockFns.selectDistinct).not.toHaveBeenCalled() + }) + it('does not enumerate sources after a satisfied ordinary existence probe', async () => { const resolve = vi.fn(async () => identity) const provider: KnowledgeAccessProvider = { diff --git a/apps/sim/lib/knowledge/read-access.ts b/apps/sim/lib/knowledge/read-access.ts index e1942b698a2..8a142217859 100644 --- a/apps/sim/lib/knowledge/read-access.ts +++ b/apps/sim/lib/knowledge/read-access.ts @@ -30,6 +30,7 @@ export async function* knowledgeReadAccessBatches( const ordinary = knowledgeAccessCondition(scope) yield ordinary if (!provider || scope.kind !== 'user') return + if (provider.hasLiveSourceReaders && !(await provider.hasLiveSourceReaders())) return let cursor: string | undefined while (true) { diff --git a/apps/sim/lib/knowledge/service.test.ts b/apps/sim/lib/knowledge/service.test.ts index bd8a8780ac1..af0a8625144 100644 --- a/apps/sim/lib/knowledge/service.test.ts +++ b/apps/sim/lib/knowledge/service.test.ts @@ -551,12 +551,11 @@ describe('knowledge base counts with live source permissions', () => { id: 'kb-1', workspaceId: 'ws-1', chunkingConfig: {}, - docCount: 99, - tokenCount: 999, + docCount: 2, + tokenCount: 10, createdAt: new Date('2026-01-01'), }, ]) - queueTableRows(schemaMock.document, [{ knowledgeBaseId: 'kb-1', docCount: 2, tokenCount: 10 }]) queueTableRows(schemaMock.document, [{ connectorId: 'confluence-source' }]) queueTableRows(schemaMock.document, [{ knowledgeBaseId: 'kb-1', docCount: 3, tokenCount: 20 }]) const result = await getWorkspaceKnowledgeBases('ws-1', 'active', { access, limit: 2 }) @@ -587,6 +586,56 @@ describe('knowledge base counts with live source permissions', () => { ).toBe(true) }) + it('counts an unpaged list in one joined query and one discovery pass bounded by the list filter', async () => { + const { access, getForConnectors } = reader() + const bases = Array.from({ length: 1000 }, (_, index) => ({ + id: `kb-${index}`, + workspaceId: 'ws-1', + chunkingConfig: {}, + docCount: 1, + tokenCount: 1, + createdAt: new Date('2026-01-01'), + })) + queueTableRows(schemaMock.knowledgeBase, bases) + queueTableRows(schemaMock.document, [{ connectorId: 'confluence-source' }]) + queueTableRows(schemaMock.document, [{ knowledgeBaseId: 'kb-7', docCount: 3, tokenCount: 20 }]) + const result = await getWorkspaceKnowledgeBases('ws-1', 'archived', { access }) + expect(result.data).toHaveLength(1000) + expect(result.data[7]).toMatchObject({ docCount: 4, tokenCount: 21 }) + expect(result.data[8]).toMatchObject({ docCount: 1, tokenCount: 1 }) + expect(getForConnectors).toHaveBeenCalledOnce() + expect(dbChainMockFns.selectDistinct).toHaveBeenCalledOnce() + expect(dbChainMockFns.groupBy).toHaveBeenCalledTimes(2) + expect( + dbChainMockFns.where.mock.calls.some(([condition]) => + hasMockCondition( + condition, + (node) => node.type === 'inArray' && node.column === schemaMock.knowledgeBase.id + ) + ) + ).toBe(false) + }) + + it('never discovers live sources for a reader without live-source credentials', async () => { + const { access, getForConnectors } = reader() + access.hasLiveSourceReaders = async () => false + queueTableRows(schemaMock.knowledgeBase, [ + { + id: 'kb-1', + workspaceId: 'ws-1', + chunkingConfig: {}, + docCount: 2, + tokenCount: 10, + createdAt: new Date('2026-01-01'), + }, + ]) + const result = await getWorkspaceKnowledgeBases('ws-1', 'archived', { access }) + expect(result.data[0]).toMatchObject({ docCount: 2, tokenCount: 10 }) + expect(getForConnectors).not.toHaveBeenCalled() + expect(dbChainMockFns.selectDistinct).not.toHaveBeenCalled() + expect(dbChainMockFns.groupBy).toHaveBeenCalledOnce() + }) + it('does not retain stale totals when a live source no longer authorizes its documents', async () => { const { access, getForConnectors } = reader() queueTableRows(schemaMock.document, []) diff --git a/apps/sim/lib/knowledge/service.ts b/apps/sim/lib/knowledge/service.ts index 0cff654809d..91352dbdef0 100644 --- a/apps/sim/lib/knowledge/service.ts +++ b/apps/sim/lib/knowledge/service.ts @@ -34,7 +34,7 @@ import { generateRestoreName } from '@/lib/core/utils/restore-name' import { findActiveFolder, resolveRestoredFolderId } from '@/lib/folders/queries' import { isKnowledgeMemberAccessAvailable } from '@/lib/knowledge/access/availability' import { knowledgeAccessCondition } from '@/lib/knowledge/access/predicate' -import { MAX_KNOWLEDGE_ACCESS_CANDIDATES } from '@/lib/knowledge/access/types' +import type { KnowledgeAccessProvider } from '@/lib/knowledge/access/types' import { mirrorsSourceAcls } from '@/lib/knowledge/connectors/access-modes' import { type KnowledgeReadAccess, knowledgeReadAccessBatches } from '@/lib/knowledge/read-access' import type { @@ -174,6 +174,7 @@ async function readKnowledgeBaseRows( ): Promise< Array> > { + const scope = access && 'get' in access ? await access.get() : access const query = db .select({ id: knowledgeBase.id, @@ -201,7 +202,7 @@ async function readKnowledgeBaseRows( eq(document.userExcluded, false), isNull(document.archivedAt), isNull(document.deletedAt), - access ? ('get' in access ? sql`false` : knowledgeAccessCondition(access)) : undefined + scope ? knowledgeAccessCondition(scope) : undefined ) ) .where(where) @@ -210,59 +211,85 @@ async function readKnowledgeBaseRows( const rows = limit === undefined ? await query : await query.limit(limit) - const counts = - access && 'get' in access - ? await readKnowledgeBaseDocumentCounts( - rows.map((kb) => kb.id), + /** + * The join above already counted everything the reader's stored ACL admits. Only a + * provider can add documents a live source (GitHub, Confluence) authorizes beyond that, + * and that supplement is resolved once for the whole list: an unpaged list is bounded by + * its own filter, a page by its row IDs, so a workspace with tens of thousands of bases + * never turns into hundreds of per-batch round trips. + */ + const liveCounts = + access && 'get' in access && rows.length > 0 + ? await readLiveSourceDocumentCounts( + limit === undefined && where + ? where + : inArray( + knowledgeBase.id, + rows.map((kb) => kb.id) + ), access ) : undefined return rows.map((kb) => ({ ...kb, chunkingConfig: kb.chunkingConfig as ChunkingConfig, - docCount: counts ? (counts.get(kb.id)?.docCount ?? 0) : Number(kb.docCount), - tokenCount: counts ? (counts.get(kb.id)?.tokenCount ?? 0) : kb.tokenCount, + docCount: Number(kb.docCount) + (liveCounts?.get(kb.id)?.docCount ?? 0), + tokenCount: kb.tokenCount + (liveCounts?.get(kb.id)?.tokenCount ?? 0), })) } -/** Counts only hydrated access batches, keeping candidate discovery free of document metadata. */ -async function readKnowledgeBaseDocumentCounts( - knowledgeBaseIds: readonly string[], - access: KnowledgeReadAccess +const ACTIVE_DOCUMENT_CONDITIONS = [ + eq(document.userExcluded, false), + isNull(document.archivedAt), + isNull(document.deletedAt), +] as const + +/** + * Document totals per knowledge base for one access predicate, restricted to the bases + * `subject` selects. `subject` may reference `knowledge_base` columns. + */ +async function countDocumentsByKnowledgeBase( + subject: SQL, + accessCondition: SQL +): Promise> { + return db + .select({ + knowledgeBaseId: document.knowledgeBaseId, + docCount: count(), + tokenCount: sql`COALESCE(SUM(${document.tokenCount}), 0)`.mapWith(Number), + }) + .from(document) + .innerJoin(knowledgeBase, eq(document.knowledgeBaseId, knowledgeBase.id)) + .where(and(subject, ...ACTIVE_DOCUMENT_CONDITIONS, accessCondition)) + .groupBy(document.knowledgeBaseId) +} + +/** + * Totals for documents only a live source authorizes, on top of the reader's stored ACL. + * The ordinary predicate is skipped because every caller has already counted it; candidate + * discovery stays free of document metadata and returns nothing for a reader without + * live-source credentials. + */ +async function readLiveSourceDocumentCounts( + subject: SQL, + access: KnowledgeAccessProvider ): Promise> { const counts = new Map() - for ( - let offset = 0; - offset < knowledgeBaseIds.length; - offset += MAX_KNOWLEDGE_ACCESS_CANDIDATES - ) { - const conditions = [ - inArray( - knowledgeBase.id, - knowledgeBaseIds.slice(offset, offset + MAX_KNOWLEDGE_ACCESS_CANDIDATES) - ), - eq(document.userExcluded, false), - isNull(document.archivedAt), - isNull(document.deletedAt), - ] - for await (const accessCondition of knowledgeReadAccessBatches(access, conditions)) { - const rows = await db - .select({ - knowledgeBaseId: document.knowledgeBaseId, - docCount: count(), - tokenCount: sql`COALESCE(SUM(${document.tokenCount}), 0)`.mapWith(Number), - }) - .from(document) - .innerJoin(knowledgeBase, eq(document.knowledgeBaseId, knowledgeBase.id)) - .where(and(...conditions, accessCondition)) - .groupBy(document.knowledgeBaseId) - for (const row of rows) { - const previous = counts.get(row.knowledgeBaseId) - counts.set(row.knowledgeBaseId, { - docCount: (previous?.docCount ?? 0) + Number(row.docCount), - tokenCount: (previous?.tokenCount ?? 0) + Number(row.tokenCount), - }) - } + let ordinary = true + for await (const accessCondition of knowledgeReadAccessBatches(access, [ + subject, + ...ACTIVE_DOCUMENT_CONDITIONS, + ])) { + if (ordinary) { + ordinary = false + continue + } + for (const row of await countDocumentsByKnowledgeBase(subject, accessCondition)) { + const previous = counts.get(row.knowledgeBaseId) + counts.set(row.knowledgeBaseId, { + docCount: (previous?.docCount ?? 0) + Number(row.docCount), + tokenCount: (previous?.tokenCount ?? 0) + Number(row.tokenCount), + }) } } return counts @@ -1019,11 +1046,14 @@ export async function attachKnowledgeBaseConnectors( ): Promise { let visible = knowledgeBase if (access) { - const counts = await readKnowledgeBaseDocumentCounts([knowledgeBase.id], access) + const subject = eq(document.knowledgeBaseId, knowledgeBase.id) + const scope = 'get' in access ? await access.get() : access + const [ordinary] = await countDocumentsByKnowledgeBase(subject, knowledgeAccessCondition(scope)) + const live = 'get' in access ? await readLiveSourceDocumentCounts(subject, access) : undefined visible = { ...knowledgeBase, - docCount: counts.get(knowledgeBase.id)?.docCount ?? 0, - tokenCount: counts.get(knowledgeBase.id)?.tokenCount ?? 0, + docCount: Number(ordinary?.docCount ?? 0) + (live?.get(knowledgeBase.id)?.docCount ?? 0), + tokenCount: (ordinary?.tokenCount ?? 0) + (live?.get(knowledgeBase.id)?.tokenCount ?? 0), } } const [withConnectors] = await attachConnectorTypes([visible])