diff --git a/apps/sim/lib/knowledge/__integration__/knowledge-base-list.integration.ts b/apps/sim/lib/knowledge/__integration__/knowledge-base-list.integration.ts new file mode 100644 index 00000000000..f2e61a91623 --- /dev/null +++ b/apps/sim/lib/knowledge/__integration__/knowledge-base-list.integration.ts @@ -0,0 +1,235 @@ +import { writeFileSync } from 'node:fs' +import { db } from '@sim/db' +import { document, knowledgeBase, organization, user, workspace } from '@sim/db/schema' +import { generateId } from '@sim/utils/id' +import { eq, inArray, sql } from 'drizzle-orm' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { + createKnowledgeAclFixtureIds, + seedKnowledgeAclFixture, +} from '@/lib/knowledge/__integration__/seed-source-access-fixture' +import { type KnowledgeAccessScope, WORKSPACE_ACCESS_TOKENS } from '@/lib/knowledge/access/types' +import { getWorkspaceKnowledgeBases } from '@/lib/knowledge/service' + +/** + * A small KB page must not count the rest of its workspace or other tenants before applying + * its limit. Real query plans catch this even when warm caches hide it from a timing test. + * The same read must retain ACL/lifecycle filtering, empty bases, and keyset continuity. + */ +const ids = createKnowledgeAclFixtureIds() +const foreign = createKnowledgeAclFixtureIds() +const manyBases = createKnowledgeAclFixtureIds() +const offPageId = generateId() +const emptyId = generateId() +const archivedId = generateId() +const access: KnowledgeAccessScope = { kind: 'workspace', tokens: WORKSPACE_ACCESS_TOKENS } +const reports: Array> = [] + +interface CapturedQuery { + query: string + parameters: NonNullable[1]> +} + +interface ExplainNode { + 'Relation Name'?: string + 'Index Name'?: string + 'Actual Rows': number + 'Actual Loops': number + 'Rows Removed by Filter'?: number + 'Rows Removed by Index Recheck'?: number + Plans?: ExplainNode[] +} + +function documentVisits(node: ExplainNode): number { + const readsDocuments = + node['Relation Name'] === 'document' || node['Index Name']?.startsWith('doc_') + const own = readsDocuments + ? (node['Actual Rows'] + + (node['Rows Removed by Filter'] ?? 0) + + (node['Rows Removed by Index Recheck'] ?? 0)) * + node['Actual Loops'] + : 0 + return own + (node.Plans ?? []).reduce((total, child) => total + documentVisits(child), 0) +} + +beforeAll(async () => { + await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' }) + await seedKnowledgeAclFixture(foreign, { connectorType: 'google_drive' }) + await db + .update(knowledgeBase) + .set({ name: 'A small', createdAt: new Date('2026-01-01') }) + .where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + await db.insert(knowledgeBase).values([ + { + id: emptyId, + workspaceId: ids.workspaceId, + userId: ids.aliceId, + name: 'B empty', + createdAt: new Date('2026-01-02'), + }, + { + id: offPageId, + workspaceId: ids.workspaceId, + userId: ids.aliceId, + name: 'C large', + createdAt: new Date('2026-01-03'), + }, + { + id: archivedId, + workspaceId: ids.workspaceId, + userId: ids.aliceId, + name: 'D archived', + deletedAt: new Date(), + }, + ]) + await db.insert(document).values( + [ + { tokenCount: 7 }, + { tokenCount: 11 }, + { tokenCount: 100, acl: ['u:hidden@fixture.test'] }, + { tokenCount: 100, archivedAt: new Date() }, + { tokenCount: 100, deletedAt: new Date() }, + { tokenCount: 100, userExcluded: true }, + ].map((row) => ({ + id: generateId(), + knowledgeBaseId: ids.knowledgeBaseId, + filename: 'fixture.txt', + fileUrl: 'https://fixture.invalid/document', + fileSize: 1, + mimeType: 'text/plain', + acl: ['ws'], + ...row, + })) + ) + for (const baseId of [offPageId, foreign.knowledgeBaseId]) { + await db.execute(sql`INSERT INTO document + (id, knowledge_base_id, filename, file_url, file_size, mime_type, acl, token_count) + SELECT ${baseId} || '-' || n, ${baseId}, 'bulk.txt', 'https://fixture.invalid/bulk', + 1, 'text/plain', ARRAY['ws'], 1 FROM generate_series(1, 10000) AS n`) + } + await seedKnowledgeAclFixture(manyBases, { connectorType: 'google_drive' }) + await db.execute(sql`INSERT INTO knowledge_base (id, workspace_id, user_id, name) + SELECT ${manyBases.knowledgeBaseId} || '-' || n, ${manyBases.workspaceId}, + ${manyBases.aliceId}, 'Scale fixture ' || n FROM generate_series(1, 10000) AS n`) + await db.execute(sql`INSERT INTO document + (id, knowledge_base_id, filename, file_url, file_size, mime_type, acl, token_count) + VALUES (${generateId()}, ${manyBases.knowledgeBaseId}, 'first.txt', + 'https://fixture.invalid/first', 1, 'text/plain', ARRAY['ws'], 13), + (${generateId()}, ${`${manyBases.knowledgeBaseId}-10000`}, 'last.txt', + 'https://fixture.invalid/last', 1, 'text/plain', ARRAY['ws'], 17)`) + await db.execute(sql`ANALYZE knowledge_base`) + await db.execute(sql`ANALYZE document`) +}, 60_000) + +afterAll(async () => { + const reportPath = process.env.KNOWLEDGE_BASE_LIST_REPORT_PATH + if (reportPath) writeFileSync(reportPath, JSON.stringify(reports, null, 2)) + try { + for (const fixture of [ids, foreign, manyBases]) { + await db.delete(workspace).where(eq(workspace.id, fixture.workspaceId)) + await db.delete(organization).where(eq(organization.id, fixture.organizationId)) + await db.delete(user).where(inArray(user.id, [fixture.aliceId, fixture.bobId])) + } + } finally { + await db.$client.end() + } +}) + +describe('knowledge base list counts on real Postgres', () => { + it.each(['name', 'createdAt'] as const)( + 'bounds document reads to a small page ordered by %s', + async (sortBy) => { + const captured: CapturedQuery[] = [] + const previousDebug = db.$client.options.debug + db.$client.options.debug = (_connection, query, parameters) => { + if (captured.length < 30) captured.push({ query, parameters: [...parameters] }) + } + try { + const page = await getWorkspaceKnowledgeBases(ids.workspaceId, 'active', { + countsFor: access, + limit: 1, + sortBy, + }) + expect( + page.data.map(({ id, docCount, tokenCount }) => ({ id, docCount, tokenCount })) + ).toEqual([{ id: ids.knowledgeBaseId, docCount: 2, tokenCount: 18 }]) + expect(page.nextCursorKeys).not.toBeNull() + } finally { + db.$client.options.debug = previousDebug + } + const plans = [] + for (const statement of captured.filter(({ query }) => query.includes('"document"'))) { + const [result] = await db.$client.unsafe< + Array<{ 'QUERY PLAN': Array<{ Plan: ExplainNode }> }> + >(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) ${statement.query}`, statement.parameters) + plans.push(...result['QUERY PLAN']) + } + const visits = plans.reduce((total, plan) => total + documentVisits(plan.Plan), 0) + reports.push({ sortBy, visits, plans }) + expect(plans.length).toBeGreaterThan(0) + expect(visits).toBeLessThan(100) + } + ) + + it('keeps empty KBs and count visibility through pagination and unpaged reads', async () => { + const first = await getWorkspaceKnowledgeBases(ids.workspaceId, 'active', { + countsFor: access, + limit: 1, + sortBy: 'name', + }) + if (!first.nextCursorKeys) throw new Error('Expected a second knowledge-base page') + const second = await getWorkspaceKnowledgeBases(ids.workspaceId, 'active', { + countsFor: access, + limit: 1, + sortBy: 'name', + cursorKeys: first.nextCursorKeys, + }) + expect( + second.data.map(({ id, docCount, tokenCount }) => ({ id, docCount, tokenCount })) + ).toEqual([{ id: emptyId, docCount: 0, tokenCount: 0 }]) + const all = await getWorkspaceKnowledgeBases(ids.workspaceId, 'active', { + countsFor: access, + sortBy: 'name', + }) + expect(all.data.map(({ id, docCount, tokenCount }) => ({ id, docCount, tokenCount }))).toEqual([ + { id: ids.knowledgeBaseId, docCount: 2, tokenCount: 18 }, + { id: emptyId, docCount: 0, tokenCount: 0 }, + { id: offPageId, docCount: 10000, tokenCount: 10000 }, + ]) + expect(all.nextCursorKeys).toBeNull() + const archived = await getWorkspaceKnowledgeBases(ids.workspaceId, 'archived', { + countsFor: access, + }) + expect(archived.data.map(({ id }) => id)).toEqual([archivedId]) + }) + + it('counts a large unpaged workspace within a fixed database round-trip budget', async () => { + let documentQueries = 0 + const previousDebug = db.$client.options.debug + db.$client.options.debug = (_connection, query) => { + if (query.includes('"document"')) documentQueries++ + } + try { + const all = await getWorkspaceKnowledgeBases(manyBases.workspaceId, 'active', { + countsFor: access, + }) + expect(all.data).toHaveLength(10001) + expect(all.nextCursorKeys).toBeNull() + expect(all.data.find((kb) => kb.id === manyBases.knowledgeBaseId)).toMatchObject({ + docCount: 1, + tokenCount: 13, + }) + expect(all.data.find((kb) => kb.id === `${manyBases.knowledgeBaseId}-10000`)).toMatchObject({ + docCount: 1, + tokenCount: 17, + }) + expect(all.data.reduce((total, kb) => total + kb.docCount, 0)).toBe(2) + expect(all.data.reduce((total, kb) => total + kb.tokenCount, 0)).toBe(30) + reports.push({ unpagedBases: all.data.length, documentQueries }) + expect(documentQueries).toBeGreaterThan(0) + expect(documentQueries).toBeLessThan(10) + } finally { + db.$client.options.debug = previousDebug + } + }) +}) diff --git a/apps/sim/lib/knowledge/service.test.ts b/apps/sim/lib/knowledge/service.test.ts index db333fd2ba2..1204da201ef 100644 --- a/apps/sim/lib/knowledge/service.test.ts +++ b/apps/sim/lib/knowledge/service.test.ts @@ -293,18 +293,16 @@ describe('knowledge base counts with live source permissions', () => { id: 'kb-1', workspaceId: 'ws-1', chunkingConfig: {}, - docCount: 2, - tokenCount: 10, createdAt: new Date('2026-01-01'), }, ]) + queueTableRows(schemaMock.document, [{ knowledgeBaseId: 'kb-1', docCount: 2, tokenCount: 10 }]) const result = await getWorkspaceKnowledgeBases('ws-1', 'archived', { countsFor: access }) expect(result.data[0]).toMatchObject({ docCount: 2, tokenCount: 10 }) expect(getForConnectors).not.toHaveBeenCalled() expect(dbChainMockFns.select).not.toHaveBeenCalledWith({ connectorId: schemaMock.knowledgeConnector.id, }) - expect(dbChainMockFns.groupBy).toHaveBeenCalledOnce() }) it('does not retain stale totals when a live source no longer authorizes its documents', async () => { diff --git a/apps/sim/lib/knowledge/service.ts b/apps/sim/lib/knowledge/service.ts index f6e43d4bc4f..423de6a68bc 100644 --- a/apps/sim/lib/knowledge/service.ts +++ b/apps/sim/lib/knowledge/service.ts @@ -33,7 +33,7 @@ import { resourceScopeCondition } from '@/lib/core/resource-scope.server' 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 { knowledgeAccessCondition, textArrayLiteral } from '@/lib/knowledge/access/predicate' import type { KnowledgeAccessProvider } from '@/lib/knowledge/access/types' import { mirrorsSourceAcls } from '@/lib/knowledge/connectors/access-modes' import { @@ -188,8 +188,8 @@ async function readKnowledgeBaseRows( } /** - * {@link readKnowledgeBaseRows} plus the live totals of the documents `access` admits. Only the - * surfaces that display totals pay for the document join, and they always count as a reader. + * Pages bases before counting the documents `access` admits. Explicit document base IDs keep + * the count selective instead of scanning a shared ACL token across tenants before the join. */ async function readCountedKnowledgeBaseRows( where: SQL | undefined, @@ -200,31 +200,18 @@ async function readCountedKnowledgeBaseRows( Array> > { const scope = 'get' in access ? await access.get() : access - const query = db - .select({ - ...ACTIVE_KNOWLEDGE_BASE_REFERENCE_FIELDS, - tokenCount: sql`COALESCE(SUM(${document.tokenCount}), 0)`.mapWith(Number), - docCount: count(document.knowledgeBaseId), - }) - .from(knowledgeBase) - .leftJoin( - document, - and( - eq(document.knowledgeBaseId, knowledgeBase.id), - eq(document.userExcluded, false), - isNull(document.archivedAt), - isNull(document.deletedAt), - knowledgeAccessCondition(scope) - ) - ) - .where(where) - .groupBy(knowledgeBase.id) - .orderBy(...orderBy) - - const rows = limit === undefined ? await query : await query.limit(limit) + const rows = await readKnowledgeBaseRows(where, orderBy, limit) + const totals = + rows.length > 0 + ? await countDocumentsByKnowledgeBase( + sql`${document.knowledgeBaseId} = ANY(${textArrayLiteral(rows.map((kb) => kb.id))})`, + knowledgeAccessCondition(scope) + ) + : [] + const counts = new Map(totals.map((total) => [total.knowledgeBaseId, total])) /** - * The join above already counted everything the reader's stored ACL admits. Only a + * The counts above already include 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 @@ -243,9 +230,9 @@ async function readCountedKnowledgeBaseRows( ) : undefined return rows.map((kb) => ({ - ...toActiveKnowledgeBaseReference(kb), - docCount: Number(kb.docCount) + (liveCounts?.get(kb.id)?.docCount ?? 0), - tokenCount: kb.tokenCount + (liveCounts?.get(kb.id)?.tokenCount ?? 0), + ...kb, + docCount: (counts.get(kb.id)?.docCount ?? 0) + (liveCounts?.get(kb.id)?.docCount ?? 0), + tokenCount: (counts.get(kb.id)?.tokenCount ?? 0) + (liveCounts?.get(kb.id)?.tokenCount ?? 0), })) }