fix(knowledge,files): per-base kb counts and a modified-order file index

apps/sim/lib/billing/credits/weekly-refresh.test.ts

@@ -14,7 +14,6 @@ import {
 const FROZEN_NOW = new Date('2026-08-15T00:00:00.000Z')
 
 beforeEach(() => {
-  drizzleOrmMock.sum.mockImplementation(() => ({ as: () => 'sum' }))
   vi.useFakeTimers()
   vi.setSystemTime(FROZEN_NOW)
 })

apps/sim/lib/knowledge/__integration__/knowledge-base-list.integration.ts

@@ -1,6 +1,7 @@
 import { writeFileSync } from 'node:fs'
 import { db } from '@sim/db'
 import { document, knowledgeBase, organization, user, workspace } from '@sim/db/schema'
+import { explainCaptured } from '@sim/db/testing/explain'
 import { deleteWorkspaceFixture } from '@sim/db/testing/workspace-fixtures'
 import { generateId } from '@sim/utils/id'
 import { eq, inArray, sql } from 'drizzle-orm'
@@ -14,12 +15,14 @@ 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.
+ * its limit, and many bases must not read the documents other tenants share a token with. 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 archivedBases = createKnowledgeAclFixtureIds()
+const ARCHIVED_BASE_COUNT = 600
 const offPageId = generateId()
 const emptyId = generateId()
 const archivedId = generateId()
@@ -118,6 +121,16 @@ beforeAll(async () => {
       '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 seedKnowledgeAclFixture(archivedBases, { connectorType: 'google_drive' })
+  await db.execute(sql`INSERT INTO knowledge_base (id, workspace_id, user_id, name, deleted_at)
+    SELECT ${archivedBases.knowledgeBaseId} || '-' || n, ${archivedBases.workspaceId},
+      ${archivedBases.aliceId}, 'Archived fixture ' || n, now()
+    FROM generate_series(1, ${ARCHIVED_BASE_COUNT}) AS n`)
+  await db.execute(sql`INSERT INTO document
+    (id, knowledge_base_id, filename, file_url, file_size, mime_type, acl, token_count, archived_at)
+    SELECT ${archivedBases.knowledgeBaseId} || '-doc-' || n, ${archivedBases.knowledgeBaseId} || '-' || n,
+      'archived.txt', 'https://fixture.invalid/archived', 1, 'text/plain', ARRAY['ws'], 1, now()
+    FROM generate_series(1, ${ARCHIVED_BASE_COUNT}) AS n`)
   await db.execute(sql`ANALYZE knowledge_base`)
   await db.execute(sql`ANALYZE document`)
 }, 60_000)
@@ -126,7 +139,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, manyBases]) {
+    for (const fixture of [ids, foreign, manyBases, archivedBases]) {
       await deleteWorkspaceFixture(db, 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]))
@@ -172,6 +185,34 @@ describe('knowledge base list counts on real Postgres', () => {
     }
   )
 
+  /**
+   * Bases whose documents are all archived, beside thousands of other tenants' live `ws`
+   * documents. No small fixture reproduces the production estimates that once turned this count
+   * into a scan of every tenant's documents (see `countDocumentsByKnowledgeBase`), so this pins
+   * the shape that rules that plan out: each index read of `document` seeks one base.
+   */
+  it('counts many archived bases one base at a time', async () => {
+    const list = await getWorkspaceKnowledgeBases(archivedBases.workspaceId, 'archived', {
+      countsFor: access,
+    })
+    expect(list.data).toHaveLength(ARCHIVED_BASE_COUNT)
+    expect(list.data.every((kb) => kb.docCount === 0 && kb.tokenCount === 0)).toBe(true)
+    const nodes = await explainCaptured(
+      db,
+      () =>
+        getWorkspaceKnowledgeBases(archivedBases.workspaceId, 'archived', { countsFor: access }),
+      { match: /"document"/ }
+    )
+    expect(
+      nodes.some((node) => node['Node Type'] === 'Seq Scan' && node['Relation Name'] === 'document')
+    ).toBe(false)
+    const documentIndexReads = nodes.filter((node) => node['Index Name']?.startsWith('doc_'))
+    expect(documentIndexReads.length).toBeGreaterThan(0)
+    expect(documentIndexReads.map((node) => node['Index Cond'])).toEqual(
+      documentIndexReads.map(() => '(knowledge_base_id = knowledge_base.id)')
+    )
+  })
+
   it('keeps empty KBs and count visibility through pagination and unpaged reads', async () => {
     const first = await getWorkspaceKnowledgeBases(ids.workspaceId, 'active', {
       countsFor: access,

apps/sim/lib/knowledge/service.test.ts

@@ -296,7 +296,9 @@ describe('knowledge base counts with live source permissions', () => {
         createdAt: new Date('2026-01-01'),
       },
     ])
-    queueTableRows(schemaMock.document, [{ knowledgeBaseId: 'kb-1', docCount: 2, tokenCount: 10 }])
+    queueTableRows(schemaMock.knowledgeBase, [
+      { 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()
@@ -307,9 +309,9 @@ describe('knowledge base counts with live source permissions', () => {
 
   it('does not retain stale totals when a live source no longer authorizes its documents', async () => {
     const { access, getForConnectors } = reader()
-    queueTableRows(schemaMock.document, [])
+    queueTableRows(schemaMock.knowledgeBase, [])
     queueTableRows(schemaMock.knowledgeConnector, [{ connectorId: 'confluence-source' }])
-    queueTableRows(schemaMock.document, [])
+    queueTableRows(schemaMock.knowledgeBase, [])
     const base = { id: 'kb-1', docCount: 5, tokenCount: 50 } as KnowledgeBaseWithCounts
     await expect(attachKnowledgeBaseConnectors(base, access)).resolves.toMatchObject({
       docCount: 0,

apps/sim/lib/knowledge/service.ts

@@ -189,8 +189,8 @@ async function readKnowledgeBaseRows(
 }
 
 /**
- * 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.
+ * Pages bases before counting the documents `access` admits, so a small page never counts the
+ * rest of its workspace or another tenant's documents.
  */
 async function readCountedKnowledgeBaseRows(
   where: SQL | undefined,
@@ -202,34 +202,24 @@ async function readCountedKnowledgeBaseRows(
 > {
   const scope = 'get' in access ? await access.get() : access
   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)
-        )
-      : []
+  if (rows.length === 0) return []
+  /**
+   * The bases to count: an unpaged list is bounded by its own filter, a page by its row IDs.
+   * Every count below runs once for the whole list, so a workspace with tens of thousands of
+   * bases never turns into hundreds of per-batch round trips.
+   */
+  const subject =
+    limit === undefined && where
+      ? where
+      : sql`${knowledgeBase.id} = ANY(${textArrayLiteral(rows.map((kb) => kb.id))})`
+  const totals = await countDocumentsByKnowledgeBase(subject, 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
-   * 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.
+   * provider can add documents a live source (GitHub, Confluence) authorizes beyond that.
    */
   const liveCounts =
-    'get' in access && rows.length > 0
-      ? await readLiveSourceDocumentCounts(
-          limit === undefined && where
-            ? where
-            : inArray(
-                knowledgeBase.id,
-                rows.map((kb) => kb.id)
-              ),
-          access
-        )
-      : undefined
+    'get' in access ? await readLiveSourceDocumentCounts(subject, access) : undefined
   return rows.map((kb) => ({
     ...kb,
     docCount: (counts.get(kb.id)?.docCount ?? 0) + (liveCounts?.get(kb.id)?.docCount ?? 0),
@@ -244,23 +234,43 @@ const ACTIVE_DOCUMENT_CONDITIONS = [
 ] as const
 
 /**
- * Document totals per knowledge base for one access predicate, restricted to the bases
- * `subject` selects. `subject` may reference `knowledge_base` columns.
+ * Document totals per knowledge base for one access predicate, one row per base `subject`
+ * selects from `knowledge_base` (zero when it has no readable documents).
+ *
+ * Each base is counted on its own. One grouped scan over many bases lets the planner
+ * over-estimate their documents and AND the base index with bitmaps of every tenant's documents
+ * (the shared-token ACL index, a partial index of unarchived documents) before it reads a row;
+ * correlated on one base ID, the count is a plain scan of that base's active documents.
  */
 async function countDocumentsByKnowledgeBase(
   subject: SQL,
   accessCondition: SQL
 ): Promise<Array<{ knowledgeBaseId: string; docCount: number; tokenCount: number }>> {
-  return db
+  const totals = db
     .select({
-      knowledgeBaseId: document.knowledgeBaseId,
-      docCount: count(),
-      tokenCount: sql<number>`COALESCE(SUM(${document.tokenCount}), 0)`.mapWith(Number),
+      docCount: count().as('readable_doc_count'),
+      tokenCount: sql<number>`COALESCE(SUM(${document.tokenCount}), 0)`
+        .mapWith(Number)
+        .as('readable_token_count'),
     })
     .from(document)
-    .innerJoin(knowledgeBase, eq(document.knowledgeBaseId, knowledgeBase.id))
-    .where(and(subject, ...ACTIVE_DOCUMENT_CONDITIONS, accessCondition))
-    .groupBy(document.knowledgeBaseId)
+    .where(
+      and(
+        eq(document.knowledgeBaseId, knowledgeBase.id),
+        ...ACTIVE_DOCUMENT_CONDITIONS,
+        accessCondition
+      )
+    )
+    .as('totals')
+  return db
+    .select({
+      knowledgeBaseId: knowledgeBase.id,
+      docCount: totals.docCount,
+      tokenCount: totals.tokenCount,
+    })
+    .from(knowledgeBase)
+    .crossJoinLateral(totals)
+    .where(subject)
 }
 
 /**
@@ -952,18 +962,18 @@ export async function getActiveKnowledgeBaseReferences(
  * operation that only resolves its context pays for neither the count nor the connector read.
  */
 export async function attachKnowledgeBaseConnectors(
-  knowledgeBase: ActiveKnowledgeBaseReference,
+  base: ActiveKnowledgeBaseReference,
   access: KnowledgeReadAccess
 ): Promise<KnowledgeBaseWithCounts> {
-  const subject = eq(document.knowledgeBaseId, knowledgeBase.id)
+  const subject = eq(knowledgeBase.id, base.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
   const [withConnectors] = await attachConnectorTypes([
     {
-      ...knowledgeBase,
-      docCount: Number(ordinary?.docCount ?? 0) + (live?.get(knowledgeBase.id)?.docCount ?? 0),
-      tokenCount: (ordinary?.tokenCount ?? 0) + (live?.get(knowledgeBase.id)?.tokenCount ?? 0),
+      ...base,
+      docCount: Number(ordinary?.docCount ?? 0) + (live?.get(base.id)?.docCount ?? 0),
+      tokenCount: (ordinary?.tokenCount ?? 0) + (live?.get(base.id)?.tokenCount ?? 0),
     },
   ])
   return withConnectors

apps/sim/lib/workspace-files/__integration__/file-listing.integration.ts

@@ -679,6 +679,27 @@ describe('bounded workspace file reads in PostgreSQL', () => {
         )
       }
     })
+
+    /**
+     * The file browser's default: every live file, most recently modified first. With sorting
+     * ruled out, an ordered plan exists only through an index keyed on that order.
+     */
+    it('reads the whole workspace in modified order from its workspace index', async () => {
+      const nodes = await explainCaptured(
+        db,
+        () =>
+          queryWorkspaceFiles(sortWorkspaceId, {
+            sortBy: 'updatedAt',
+            sortOrder: 'desc',
+            limit: 50,
+          }),
+        { match: /order by/, gucs: { enable_sort: 'off' } }
+      )
+      expect(nodes.map((node) => node['Node Type'])).not.toContain('Sort')
+      expect(nodes.map((node) => node['Index Name'])).toContain(
+        'workspace_files_workspace_active_updated_idx'
+      )
+    })
   })
 
   describe('SQL type labels and kinds', () => {

packages/db/migrations/0458_workspace_files_active_updated_idx.sql

@@ -0,0 +1,30 @@
+-- A workspace's live file list, most recently modified first: the file browser's default sort.
+-- 0442 indexed the upload order, but no index served this one, so each list read every live file
+-- of the workspace and sorted them, which for a large workspace runs from disk and reaches the
+-- statement timeout.
+--
+-- The predicate and the heap-only owner checks are 0442's, for the same reasons; the key is the
+-- (updated_at, id) keyset the list orders and resumes on, kept as narrow.
+--
+-- Replay-safety: as 0442. The build is CONCURRENTLY below an embedded COMMIT, and a replay moves
+-- only an INVALID index left by a failed build aside to be dropped, keeping a valid one.
+DO $$
+BEGIN
+  IF EXISTS (
+    SELECT 1 FROM pg_index WHERE indexrelid = to_regclass('workspace_files_workspace_active_updated_idx_failed')
+      AND (indisvalid OR indrelid <> to_regclass('workspace_files'))
+  ) THEN
+    RAISE EXCEPTION 'Refusing to drop unexpected recovery index workspace_files_workspace_active_updated_idx_failed';
+  END IF;
+  IF EXISTS (
+    SELECT 1 FROM pg_index WHERE indexrelid = to_regclass('workspace_files_workspace_active_updated_idx')
+      AND indrelid = to_regclass('workspace_files') AND NOT indisvalid
+  ) THEN
+    ALTER INDEX "workspace_files_workspace_active_updated_idx" RENAME TO "workspace_files_workspace_active_updated_idx_failed";
+  END IF;
+END $$;--> statement-breakpoint
+COMMIT;--> statement-breakpoint
+SET lock_timeout = 0;--> statement-breakpoint
+DROP INDEX CONCURRENTLY IF EXISTS "workspace_files_workspace_active_updated_idx_failed";--> statement-breakpoint
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "workspace_files_workspace_active_updated_idx" ON "workspace_files" USING btree ("workspace_id","updated_at","id") WHERE "workspace_files"."deleted_at" IS NULL AND "workspace_files"."context" = 'workspace' AND "workspace_files"."workspace_id" IS NOT NULL;--> statement-breakpoint
+SET lock_timeout = '5s';