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

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<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 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
+    }
+  })
+})

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

@@ -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<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)
-      )
-    )
-    .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),
   }))
 }