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])