fix(knowledge): page connector documents in editor

Author: TheodoreSpeaks

apps/sim/app/workspace/[workspaceId]/knowledge/[id]/components/edit-connector-modal/edit-connector-modal.tsx

@@ -402,19 +402,27 @@ interface DocumentsTabProps {
 function DocumentsTab({ knowledgeBaseId, connectorId }: DocumentsTabProps) {
   const [filter, setFilter] = useState<'active' | 'excluded'>('active')
 
-  const { data, isLoading } = useConnectorDocuments(knowledgeBaseId, connectorId, {
-    includeExcluded: true,
-  })
+  const { data, isLoading, hasNextPage, isFetchingNextPage, fetchNextPage } = useConnectorDocuments(
+    knowledgeBaseId,
+    connectorId,
+    {
+      includeExcluded: true,
+    }
+  )
 
   const { mutate: excludeDoc, isPending: isExcluding } = useExcludeConnectorDocument()
   const { mutate: restoreDoc, isPending: isRestoring } = useRestoreConnectorDocument()
 
   const documents = useMemo(() => {
-    if (!data?.documents) return []
-    return data.documents.filter((d) => (filter === 'excluded' ? d.userExcluded : !d.userExcluded))
-  }, [data?.documents, filter])
+    const loadedDocuments = data?.pages.flatMap((page) => page.documents) ?? []
+    return loadedDocuments.filter((document) =>
+      filter === 'excluded' ? document.userExcluded : !document.userExcluded
+    )
+  }, [data?.pages, filter])
 
-  const counts = data?.counts ?? { active: 0, excluded: 0 }
+  const counts = data?.pages[0]?.counts ?? { active: 0, excluded: 0 }
+  const visibleDocumentCount = filter === 'excluded' ? counts.excluded : counts.active
+  const hasMoreVisibleDocuments = Boolean(hasNextPage && documents.length < visibleDocumentCount)
 
   if (isLoading) {
     return (
@@ -435,7 +443,7 @@ function DocumentsTab({ knowledgeBaseId, connectorId }: DocumentsTabProps) {
       </ButtonGroup>
 
       <div className='max-h-[320px] min-h-0 overflow-y-auto [scrollbar-gutter:stable]'>
-        {documents.length === 0 ? (
+        {visibleDocumentCount === 0 ? (
           <p className='rounded-lg bg-[var(--surface-3)] px-3 py-8 text-center text-[var(--text-muted)] text-small'>
             {filter === 'excluded' ? 'No excluded documents' : 'No documents yet'}
           </p>
@@ -488,6 +496,17 @@ function DocumentsTab({ knowledgeBaseId, connectorId }: DocumentsTabProps) {
                 </Button>
               </div>
             ))}
+            {hasMoreVisibleDocuments && (
+              <Button
+                variant='ghost-secondary'
+                size='sm'
+                className='w-full'
+                disabled={isFetchingNextPage}
+                onClick={() => fetchNextPage()}
+              >
+                {isFetchingNextPage ? 'Loading…' : 'Load more documents'}
+              </Button>
+            )}
           </div>
         )}
       </div>

apps/sim/hooks/queries/kb/connectors.test.ts

@@ -0,0 +1,89 @@
+/**
+ * @vitest-environment node
+ */
+
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const mocks = vi.hoisted(() => ({
+  requestJson: vi.fn(),
+  useInfiniteQuery: vi.fn(),
+}))
+
+vi.mock('@tanstack/react-query', () => ({
+  keepPreviousData: Symbol('keepPreviousData'),
+  useInfiniteQuery: mocks.useInfiniteQuery,
+  useMutation: vi.fn(),
+  useQuery: vi.fn(),
+  useQueryClient: vi.fn(() => ({ invalidateQueries: vi.fn() })),
+}))
+
+vi.mock('@/lib/api/client/request', () => ({
+  requestJson: mocks.requestJson,
+}))
+
+import { listKnowledgeConnectorDocumentsContract } from '@/lib/api/contracts/knowledge'
+import { MAX_KNOWLEDGE_CONNECTOR_DOCUMENT_PAGE_SIZE } from '@/lib/knowledge/constants'
+import { useConnectorDocuments } from '@/hooks/queries/kb/connectors'
+
+interface ConnectorDocumentsPage {
+  documents: Array<{ id: string }>
+  counts: { active: number; excluded: number }
+}
+
+interface ConnectorDocumentsQueryOptions {
+  initialPageParam: number
+  queryFn: (context: { signal: AbortSignal; pageParam: number }) => Promise<unknown>
+  getNextPageParam: (
+    lastPage: ConnectorDocumentsPage,
+    pages: ConnectorDocumentsPage[]
+  ) => number | undefined
+}
+
+describe('useConnectorDocuments', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('requests bounded pages and advances until the authoritative total is loaded', async () => {
+    const firstPage = {
+      documents: [{ id: 'document-1' }, { id: 'document-2' }],
+      counts: { active: 2, excluded: 1 },
+    }
+    const finalPage = {
+      documents: [{ id: 'document-3' }],
+      counts: firstPage.counts,
+    }
+    mocks.requestJson.mockResolvedValue({ data: firstPage })
+
+    useConnectorDocuments('knowledge-1', 'connector-1', { includeExcluded: true })
+
+    const options = mocks.useInfiniteQuery.mock.calls[0]?.[0] as ConnectorDocumentsQueryOptions
+    const signal = new AbortController().signal
+    await options.queryFn({ signal, pageParam: 200 })
+
+    expect(mocks.requestJson).toHaveBeenCalledWith(listKnowledgeConnectorDocumentsContract, {
+      params: { id: 'knowledge-1', connectorId: 'connector-1' },
+      query: {
+        includeExcluded: true,
+        limit: MAX_KNOWLEDGE_CONNECTOR_DOCUMENT_PAGE_SIZE,
+        offset: 200,
+      },
+      signal,
+    })
+    expect(options.initialPageParam).toBe(0)
+    expect(options.getNextPageParam(firstPage, [firstPage])).toBe(2)
+    expect(options.getNextPageParam(finalPage, [firstPage, finalPage])).toBeUndefined()
+  })
+
+  it('does not page toward excluded documents when they were not requested', () => {
+    const activePage = {
+      documents: [{ id: 'document-1' }, { id: 'document-2' }],
+      counts: { active: 2, excluded: 10 },
+    }
+
+    useConnectorDocuments('knowledge-1', 'connector-1')
+
+    const options = mocks.useInfiniteQuery.mock.calls[0]?.[0] as ConnectorDocumentsQueryOptions
+    expect(options.getNextPageParam(activePage, [activePage])).toBeUndefined()
+  })
+})

apps/sim/hooks/queries/kb/connectors.ts

@@ -1,5 +1,11 @@
 import { createLogger } from '@sim/logger'
-import { keepPreviousData, useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
+import {
+  keepPreviousData,
+  useInfiniteQuery,
+  useMutation,
+  useQuery,
+  useQueryClient,
+} from '@tanstack/react-query'
 import { requestJson } from '@/lib/api/client/request'
 import {
   type ConnectorData,
@@ -15,6 +21,7 @@ import {
   triggerKnowledgeConnectorSyncContract,
   updateKnowledgeConnectorContract,
 } from '@/lib/api/contracts/knowledge'
+import { MAX_KNOWLEDGE_CONNECTOR_DOCUMENT_PAGE_SIZE } from '@/lib/knowledge/constants'
 import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys'
 
 const logger = createLogger('KnowledgeConnectorQueries')
@@ -245,11 +252,16 @@ async function fetchConnectorDocuments(
   knowledgeBaseId: string,
   connectorId: string,
   includeExcluded: boolean,
+  offset: number,
   signal?: AbortSignal
 ): Promise<ConnectorDocumentsData> {
   const result = await requestJson(listKnowledgeConnectorDocumentsContract, {
     params: { id: knowledgeBaseId, connectorId },
-    query: { includeExcluded },
+    query: {
+      includeExcluded,
+      limit: MAX_KNOWLEDGE_CONNECTOR_DOCUMENT_PAGE_SIZE,
+      offset,
+    },
     signal,
   })
 
@@ -261,18 +273,24 @@ export function useConnectorDocuments(
   connectorId?: string,
   options?: { includeExcluded?: boolean }
 ) {
-  return useQuery({
-    queryKey: [
-      ...connectorDocumentKeys.list(knowledgeBaseId, connectorId),
-      options?.includeExcluded ?? false,
-    ],
-    queryFn: ({ signal }) =>
+  const includeExcluded = options?.includeExcluded ?? false
+  return useInfiniteQuery({
+    queryKey: [...connectorDocumentKeys.list(knowledgeBaseId, connectorId), includeExcluded],
+    queryFn: ({ signal, pageParam }) =>
       fetchConnectorDocuments(
         knowledgeBaseId as string,
         connectorId as string,
-        options?.includeExcluded ?? false,
+        includeExcluded,
+        pageParam,
         signal
       ),
+    initialPageParam: 0,
+    getNextPageParam: (lastPage, pages) => {
+      const loadedCount = pages.reduce((total, page) => total + page.documents.length, 0)
+      const totalCount = lastPage.counts.active + (includeExcluded ? lastPage.counts.excluded : 0)
+      if (lastPage.documents.length === 0 || loadedCount >= totalCount) return undefined
+      return loadedCount
+    },
     enabled: Boolean(knowledgeBaseId && connectorId),
     staleTime: CONNECTOR_DOCUMENT_LIST_STALE_TIME,
     placeholderData: keepPreviousData,