fix(knowledge): scope list counts to selected knowledge bases

Author: icecrasher321

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<Record<string, unknown>> = []
+
+interface CapturedQuery {
+  query: string
+  parameters: NonNullable<Parameters<typeof db.$client.unsafe>[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])
+  })
+})

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 () => {

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<ActiveKnowledgeBaseReference & Pick<KnowledgeBaseWithCounts, 'docCount' | 'tokenCount'>>
 > {
   const scope = 'get' in access ? await access.get() : access
-  const query = db
-    .select({
-      ...ACTIVE_KNOWLEDGE_BASE_REFERENCE_FIELDS,
-      tokenCount: sql<number>`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<string, { docCount: number; tokenCount: number }>()
+  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),
   }))
 }