fix(search): preserve partial results and reuse document access checks (#7808)

Author: icecrasher321

apps/sim/app/api/knowledge/search/route.test.ts

@@ -19,7 +19,11 @@ describe('workspace search route', () => {
       user: { id: 'user-1', email: 'reader@fixture.test', name: 'Reader' },
       session: { id: 'session-1' },
     })
-    mocks.search.mockResolvedValue({ results: [], knowledgeBases: [] })
+    mocks.search.mockResolvedValue({
+      results: [],
+      knowledgeBases: [],
+      retrieval: { status: 'complete', timedOutLegs: [] },
+    })
   })
 
   it('passes the authenticated request cancellation signal through the existing operation', async () => {
@@ -41,11 +45,48 @@ describe('workspace search route', () => {
     expect(call.input).not.toHaveProperty('knowledgeBaseIds')
     expect(call.input.filters).toEqual({ source: 'slack', documentIds: ['doc-1'] })
     expect(call.input.signal).toBe(request.signal)
+    expect(call.input.allowPartialResults).toBe(true)
     controller.abort()
     expect(call.input.signal.aborted).toBe(true)
     await expect(response.json()).resolves.toEqual({
       success: true,
-      data: { query: 'Orion', results: [] },
+      data: { query: 'Orion', results: [], retrieval: { status: 'complete', timedOutLegs: [] } },
+    })
+  })
+
+  it('preserves usable matches and incomplete coverage when one retrieval leg times out', async () => {
+    mocks.search.mockResolvedValueOnce({
+      knowledgeBases: [{ id: 'knowledge-1', name: 'Search index' }],
+      retrieval: { status: 'partial', timedOutLegs: ['vector'] },
+      results: [
+        {
+          documentId: 'document-1',
+          knowledgeBaseId: 'knowledge-1',
+          documentName: 'Release plan',
+          sourceUrl: null,
+          connectorType: null,
+          sourceModifiedAt: null,
+          metadata: {},
+          content: 'Orion release',
+          chunkIndex: 0,
+          similarity: 0.9,
+        },
+      ],
+    })
+    const response = await POST(
+      new NextRequest('http://localhost/api/knowledge/search', {
+        method: 'POST',
+        headers: { 'content-type': 'application/json' },
+        body: JSON.stringify({ organizationId: 'organization-1', query: 'Orion' }),
+      })
+    )
+    expect(response.status).toBe(200)
+    await expect(response.json()).resolves.toMatchObject({
+      success: true,
+      data: {
+        results: [{ documentId: 'document-1', content: 'Orion release' }],
+        retrieval: { status: 'partial', timedOutLegs: ['vector'] },
+      },
     })
   })
 

apps/sim/app/api/knowledge/search/route.ts

@@ -24,16 +24,18 @@ export const POST = defineInternalJsonRoute({
     filters: body.filters,
     query: body.query,
     topK: body.topK,
+    allowPartialResults: true,
     surface: 'dashboard' as const,
     signal: request.signal,
   }),
   useCase: searchScopedKnowledge,
-  present: ({ results, knowledgeBases }, { input }) => {
+  present: ({ results, knowledgeBases, retrieval }, { input }) => {
     const knowledgeBaseNames = new Map(knowledgeBases.map((kb) => [kb.id, kb.name]))
     return {
       success: true as const,
       data: {
         query: input.query ?? '',
+        retrieval,
         results: results.map((result) => ({
           documentId: result.documentId,
           knowledgeBaseId: result.knowledgeBaseId,

apps/sim/app/o/[organizationId]/search/search.test.tsx

@@ -90,7 +90,12 @@ beforeEach(() => {
       chunkIndex: 0,
       similarity: 1,
     }
-    return { data: [result], isPending: false, isFetching: false, isError: false }
+    return {
+      data: { query, results: [result], retrieval: { status: 'complete', timedOutLegs: [] } },
+      isPending: false,
+      isFetching: false,
+      isError: false,
+    }
   })
   container = document.createElement('div')
   document.body.appendChild(container)

apps/sim/app/workspace/[workspaceId]/home/components/knowledge-search-results/knowledge-search-results.test.tsx

@@ -4,22 +4,17 @@ import { NuqsTestingAdapter } from 'nuqs/adapters/testing'
 import { createRoot, type Root } from 'react-dom/client'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
-const mocks = vi.hoisted(() => ({ overview: vi.fn() }))
+const mocks = vi.hoisted(() => ({ overview: vi.fn(), search: vi.fn(), retry: vi.fn() }))
 vi.mock('@/hooks/queries/kb/connectors', () => ({
   useSearchIndex: () => ({ data: { knowledgeBaseId: 'index' }, isPending: false }),
   useSearchSourceOverview: mocks.overview,
 }))
 vi.mock('@/hooks/queries/kb/knowledge', () => ({
-  useWorkspaceKnowledgeSearch: () => ({
-    data: [],
-    isPending: false,
-    isFetching: false,
-    isError: false,
-  }),
+  useWorkspaceKnowledgeSearch: mocks.search,
 }))
 vi.mock(
   '@/app/workspace/[workspaceId]/home/components/message-content/components/source-card',
-  () => ({ SourceCard: () => null })
+  () => ({ SourceCard: ({ source }: { source: { title: string } }) => <span>{source.title}</span> })
 )
 
 import { KnowledgeSearchResults } from '@/app/workspace/[workspaceId]/home/components/knowledge-search-results/knowledge-search-results'
@@ -28,6 +23,13 @@ let root: Root
 let container: HTMLDivElement
 beforeEach(() => {
   vi.clearAllMocks()
+  mocks.search.mockReturnValue({
+    data: { query: 'launch', results: [], retrieval: { status: 'complete', timedOutLegs: [] } },
+    isPending: false,
+    isFetching: false,
+    isError: false,
+    refetch: mocks.retry,
+  })
   vi.stubGlobal('IS_REACT_ACT_ENVIRONMENT', true)
   container = document.createElement('div')
   root = createRoot(container)
@@ -69,3 +71,52 @@ describe('source indexing context in search results', () => {
     expect(container.textContent).toContain('No documents you can read match')
   })
 })
+
+describe('incomplete search coverage', () => {
+  it.each([false, true])(
+    'keeps matches and offers a retry without claiming absence (hasResults=%s)',
+    async (hasResults) => {
+      mocks.search.mockReturnValue({
+        data: {
+          query: 'launch',
+          results: hasResults
+            ? [
+                {
+                  documentId: 'document-1',
+                  knowledgeBaseId: 'index',
+                  knowledgeBaseName: 'Search index',
+                  documentName: 'Release plan',
+                  sourceUrl: 'https://fixture.test/release',
+                  connectorType: null,
+                  sourceModifiedAt: null,
+                  author: null,
+                  content: 'launch details',
+                  chunkIndex: 0,
+                  similarity: 0.9,
+                },
+              ]
+            : [],
+          retrieval: { status: 'partial', timedOutLegs: ['vector'] },
+        },
+        isPending: false,
+        isFetching: false,
+        isError: false,
+        refetch: mocks.retry,
+      })
+      await render()
+      expect(container.textContent).not.toContain('Search couldn’t run')
+      expect(container.textContent).not.toContain('No documents')
+      expect(container.textContent).not.toContain('0 documents')
+      expect(container.textContent).toContain(
+        hasResults ? 'Some results may be missing.' : 'Search is incomplete.'
+      )
+      if (hasResults) expect(container.textContent).toContain('Release plan')
+      const retry = [...container.querySelectorAll('button')].find(
+        (button) => button.textContent === 'Try again'
+      )
+      expect(retry).toBeDefined()
+      await act(async () => retry!.click())
+      expect(mocks.retry).toHaveBeenCalledOnce()
+    }
+  )
+})

apps/sim/app/workspace/[workspaceId]/home/components/knowledge-search-results/knowledge-search-results.tsx

@@ -133,7 +133,7 @@ export function KnowledgeSearchResults({
     }
   }, [filters.source, filters.updated])
   const {
-    data: results,
+    data: search,
     isPending,
     isFetching,
     isError: searchFailed,
@@ -143,7 +143,8 @@ export function KnowledgeSearchResults({
   const indexing = (overview?.providers ?? [])
     .filter((provider) => provider.isSyncing)
     .map((provider) => connectorDisplayName(provider.connectorType))
-  const documents = useMemo(() => groupResultsByDocument(results ?? []), [results])
+  const documents = useMemo(() => groupResultsByDocument(search?.results ?? []), [search?.results])
+  const incomplete = search?.retrieval.status === 'partial'
   const sourceTypes = [
     ...new Set([
       ...(filters.source ? [filters.source] : []),
@@ -155,8 +156,7 @@ export function KnowledgeSearchResults({
   const showFilters =
     filtersActive || (documents.length >= FILTERS_MIN_RESULTS && sourceTypes.length > 1)
 
-  /* A failed search says so in one quiet line and offers to run again; the cause is
-     the server's to log, never the reader's to parse. */
+  /** A failed search offers a retry; server diagnostics carry the cause. */
   if (basesFailed || searchFailed) {
     const retrying = basesFetching || isFetching
     return (
@@ -188,7 +188,7 @@ export function KnowledgeSearchResults({
       </div>
     )
   }
-  if (isPending || (isFetching && !results)) {
+  if (isPending || (isFetching && !search)) {
     return (
       <div className='px-2 py-2'>
         <ActivityStatus label='Searching…' isActive />
@@ -205,12 +205,24 @@ export function KnowledgeSearchResults({
     <div className='flex flex-col'>
       <div className='flex items-center gap-2 px-2 py-2'>
         <span className='min-w-0 flex-1 text-[var(--text-muted)] text-caption'>
-          <span className='tabular-nums'>
-            {documents.length === 1 ? '1 document' : `${documents.length} documents`}
-          </span>
-          {' · searched as you'}
+          {incomplete && documents.length === 0 ? (
+            'Search is incomplete.'
+          ) : (
+            <>
+              <span className='tabular-nums'>
+                {documents.length === 1 ? '1 document' : `${documents.length} documents`}
+              </span>
+              {' · searched as you'}
+              {incomplete && <span className='block'>Some results may be missing.</span>}
+            </>
+          )}
           {indexingNote && <span className='block'>{indexingNote}</span>}
         </span>
+        {incomplete && (
+          <Chip variant='border' disabled={isFetching} onClick={() => void refetchSearch()}>
+            {isFetching ? 'Retrying…' : 'Try again'}
+          </Chip>
+        )}
       </div>
       {showFilters && (
         <div className='flex flex-wrap items-center gap-1.5 px-2 pb-2'>
@@ -245,11 +257,13 @@ export function KnowledgeSearchResults({
         </div>
       )}
       {documents.length === 0 ? (
-        <p className='px-2 py-2 text-[var(--text-muted)] text-caption'>
-          {filtersActive
-            ? 'No documents match these filters.'
-            : `No documents you can read match “${query}”.`}
-        </p>
+        !incomplete && (
+          <p className='px-2 py-2 text-[var(--text-muted)] text-caption'>
+            {filtersActive
+              ? 'No documents match these filters.'
+              : `No documents you can read match “${query}”.`}
+          </p>
+        )
       ) : (
         <div className='flex flex-col' onKeyDown={handleResultsKeyDown}>
           {documents.map((result) => {

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

@@ -132,12 +132,17 @@ describe('knowledge query placeholder scope', () => {
   })
 
   it('forwards search cancellation and leaves provider retries to the server', async () => {
-    mocks.requestJson.mockResolvedValueOnce({ data: { results: [] } })
+    const data = {
+      query: 'query',
+      results: [],
+      retrieval: { status: 'partial', timedOutLegs: ['vector'] },
+    }
+    mocks.requestJson.mockResolvedValueOnce({ data })
     const query = captureQuery(() =>
       useWorkspaceKnowledgeSearch('workspace-1', ' query ', { source: 'slack' })
     )
     const controller = new AbortController()
-    await query.queryFn({ signal: controller.signal })
+    await expect(query.queryFn({ signal: controller.signal })).resolves.toEqual(data)
 
     expect(query.retry).toBe(false)
     expect(mocks.requestJson).toHaveBeenLastCalledWith(

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

@@ -52,7 +52,7 @@ import {
   updateKnowledgeDocumentContract,
   updateKnowledgeDocumentTagsContract,
   type WorkspaceKnowledgeSearchBody,
-  type WorkspaceKnowledgeSearchResult,
+  type WorkspaceKnowledgeSearchData,
 } from '@/lib/api/contracts/knowledge'
 import type { WorkspaceSearchFilters } from '@/lib/api/contracts/knowledge/search'
 import type { ChunkingStrategy, StrategyOptions } from '@/lib/chunkers/types'
@@ -1197,9 +1197,9 @@ export function useBulkDeleteKnowledgeBases(workspaceId: string) {
 async function searchWorkspaceKnowledge(
   body: WorkspaceKnowledgeSearchBody,
   signal?: AbortSignal
-): Promise<WorkspaceKnowledgeSearchResult[]> {
+): Promise<WorkspaceKnowledgeSearchData> {
   const data = await requestJson(searchWorkspaceKnowledgeContract, { body, signal })
-  return data.data.results
+  return data.data
 }
 
 /** Searches the canonical index under the signed-in person's ACLs. */

apps/sim/lib/api/contracts/knowledge/search.ts

@@ -186,6 +186,16 @@ export const workspaceKnowledgeSearchBodySchema = resourceOwnerSchema.safeExtend
 })
 export type WorkspaceKnowledgeSearchBody = z.input<typeof workspaceKnowledgeSearchBodySchema>
 
+export const workspaceKnowledgeSearchDataSchema = z.object({
+  query: z.string(),
+  results: z.array(workspaceKnowledgeSearchResultSchema),
+  retrieval: z.object({
+    status: z.enum(['complete', 'partial']),
+    timedOutLegs: z.array(z.enum(['vector', 'keyword', 'tags'])).max(3),
+  }),
+})
+export type WorkspaceKnowledgeSearchData = z.output<typeof workspaceKnowledgeSearchDataSchema>
+
 /**
  * The search a signed-in person runs from the composer: what their own
  * account may read across the workspace's knowledge bases, presented as
@@ -199,10 +209,7 @@ export const searchWorkspaceKnowledgeContract = defineRouteContract({
     mode: 'json',
     schema: z.object({
       success: z.literal(true),
-      data: z.object({
-        query: z.string(),
-        results: z.array(workspaceKnowledgeSearchResultSchema),
-      }),
+      data: workspaceKnowledgeSearchDataSchema,
     }),
   },
 })