From ba07b4e4da148399206a91d3a20a9468f2661153 Mon Sep 17 00:00:00 2001 From: Vikhyath Mondreti Date: Fri, 2 Oct 2026 10:46:56 -0700 Subject: [PATCH 1/2] fix(knowledge): scope list counts to selected knowledge bases --- .../knowledge-base-list.integration.ts | 194 ++++++++++++++++++ apps/sim/lib/knowledge/service.test.ts | 4 +- apps/sim/lib/knowledge/service.ts | 46 ++--- 3 files changed, 214 insertions(+), 30 deletions(-) create mode 100644 apps/sim/lib/knowledge/__integration__/knowledge-base-list.integration.ts 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..920bf7bc79e --- /dev/null +++ b/apps/sim/lib/knowledge/__integration__/knowledge-base-list.integration.ts @@ -0,0 +1,194 @@ +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 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 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]) { + 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]) + }) +}) 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..730c75c352b 100644 --- a/apps/sim/lib/knowledge/service.ts +++ b/apps/sim/lib/knowledge/service.ts @@ -2,6 +2,7 @@ import { db } from '@sim/db' import { document, knowledgeBase, knowledgeConnector, workspaceFiles } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { getPostgresConstraintName, getPostgresErrorCode } from '@sim/utils/errors' +import { chunkArray } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' import { filterUndefined } from '@sim/utils/object' import type { SQL } from 'drizzle-orm' @@ -52,6 +53,7 @@ import type { import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils' const logger = createLogger('KnowledgeBaseService') +const KNOWLEDGE_BASE_COUNT_BATCH_SIZE = 100 /** * Every caller-fixable knowledge-base failure is an {@link OrchestrationError}, @@ -188,8 +190,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 +202,21 @@ 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) - ) + const rows = await readKnowledgeBaseRows(where, orderBy, limit) + const counts = new Map() + for (const batch of chunkArray(rows, KNOWLEDGE_BASE_COUNT_BATCH_SIZE)) { + const totals = await countDocumentsByKnowledgeBase( + inArray( + document.knowledgeBaseId, + batch.map((kb) => kb.id) + ), + knowledgeAccessCondition(scope) ) - .where(where) - .groupBy(knowledgeBase.id) - .orderBy(...orderBy) - - const rows = limit === undefined ? await query : await query.limit(limit) + for (const total of totals) counts.set(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 +235,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), })) } From 2a16f914c3db1d35a7f53abbdb10758803bc040b Mon Sep 17 00:00:00 2001 From: Vikhyath Mondreti Date: Fri, 2 Oct 2026 11:24:53 -0700 Subject: [PATCH 2/2] fix(knowledge): avoid count-query fan-out for unpaged lists --- .../knowledge-base-list.integration.ts | 43 ++++++++++++++++++- apps/sim/lib/knowledge/service.ts | 23 ++++------ 2 files changed, 51 insertions(+), 15 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/knowledge-base-list.integration.ts b/apps/sim/lib/knowledge/__integration__/knowledge-base-list.integration.ts index 920bf7bc79e..f2e61a91623 100644 --- a/apps/sim/lib/knowledge/__integration__/knowledge-base-list.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/knowledge-base-list.integration.ts @@ -18,6 +18,7 @@ import { getWorkspaceKnowledgeBases } from '@/lib/knowledge/service' */ const ids = createKnowledgeAclFixtureIds() const foreign = createKnowledgeAclFixtureIds() +const manyBases = createKnowledgeAclFixtureIds() const offPageId = generateId() const emptyId = generateId() const archivedId = generateId() @@ -106,6 +107,16 @@ beforeAll(async () => { 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) @@ -114,7 +125,7 @@ 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]) { + 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])) @@ -191,4 +202,34 @@ describe('knowledge base list counts on real Postgres', () => { }) 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.ts b/apps/sim/lib/knowledge/service.ts index 730c75c352b..423de6a68bc 100644 --- a/apps/sim/lib/knowledge/service.ts +++ b/apps/sim/lib/knowledge/service.ts @@ -2,7 +2,6 @@ import { db } from '@sim/db' import { document, knowledgeBase, knowledgeConnector, workspaceFiles } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { getPostgresConstraintName, getPostgresErrorCode } from '@sim/utils/errors' -import { chunkArray } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' import { filterUndefined } from '@sim/utils/object' import type { SQL } from 'drizzle-orm' @@ -34,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 { @@ -53,7 +52,6 @@ import type { import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils' const logger = createLogger('KnowledgeBaseService') -const KNOWLEDGE_BASE_COUNT_BATCH_SIZE = 100 /** * Every caller-fixable knowledge-base failure is an {@link OrchestrationError}, @@ -203,17 +201,14 @@ async function readCountedKnowledgeBaseRows( > { const scope = 'get' in access ? await access.get() : access const rows = await readKnowledgeBaseRows(where, orderBy, limit) - const counts = new Map() - for (const batch of chunkArray(rows, KNOWLEDGE_BASE_COUNT_BATCH_SIZE)) { - const totals = await countDocumentsByKnowledgeBase( - inArray( - document.knowledgeBaseId, - batch.map((kb) => kb.id) - ), - knowledgeAccessCondition(scope) - ) - for (const total of totals) counts.set(total.knowledgeBaseId, total) - } + 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 counts above already include everything the reader's stored ACL admits. Only a