fix(search): shorten direct retrieval and simplify result states

Author: icecrasher321

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

@@ -46,6 +46,7 @@ describe('workspace search route', () => {
     expect(call.input.filters).toEqual({ source: 'slack', documentIds: ['doc-1'] })
     expect(call.input.signal).toBe(request.signal)
     expect(call.input.allowPartialResults).toBe(true)
+    expect(call.input.vectorBudgetMs).toBe(3000)
     controller.abort()
     expect(call.input.signal.aborted).toBe(true)
     await expect(response.json()).resolves.toEqual({

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

@@ -9,6 +9,8 @@ import { knowledgeOperations } from '@/lib/knowledge/application/operations'
 import { searchScopedKnowledge } from '@/lib/knowledge/application/workspace-search'
 import { sourceAuthor } from '@/lib/knowledge/search/author'
 
+const DIRECT_SEARCH_VECTOR_BUDGET_MS = 3000
+
 export const POST = defineInternalJsonRoute({
   contract: searchWorkspaceKnowledgeContract,
   auth: internalSessionAuth,
@@ -25,6 +27,7 @@ export const POST = defineInternalJsonRoute({
     query: body.query,
     topK: body.topK,
     allowPartialResults: true,
+    vectorBudgetMs: DIRECT_SEARCH_VECTOR_BUDGET_MS,
     surface: 'dashboard' as const,
     signal: request.signal,
   }),

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

@@ -68,13 +68,13 @@ describe('source indexing context in search results', () => {
     mocks.overview.mockReturnValue({ data: undefined })
     await render()
     expect(container.textContent).not.toContain('Still indexing')
-    expect(container.textContent).toContain('No documents you can read match')
+    expect(container.textContent).toContain('Search found no results.')
   })
 })
 
 describe('incomplete search coverage', () => {
   it.each([false, true])(
-    'keeps matches and offers a retry without claiming absence (hasResults=%s)',
+    'shows matches without timeout copy or retry controls (hasResults=%s)',
     async (hasResults) => {
       mocks.search.mockReturnValue({
         data: {
@@ -106,17 +106,17 @@ describe('incomplete search coverage', () => {
       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).not.toContain('Some results may be missing.')
+      expect(container.textContent).not.toContain('Search is incomplete.')
       expect(container.textContent).toContain(
-        hasResults ? 'Some results may be missing.' : 'Search is incomplete.'
+        hasResults ? '1 document' : 'Search found no results.'
       )
       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()
+      expect(retry).toBeUndefined()
+      expect(mocks.retry).not.toHaveBeenCalled()
     }
   )
 })

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

@@ -144,7 +144,6 @@ export function KnowledgeSearchResults({
     .filter((provider) => provider.isSyncing)
     .map((provider) => connectorDisplayName(provider.connectorType))
   const documents = useMemo(() => groupResultsByDocument(search?.results ?? []), [search?.results])
-  const incomplete = search?.retrieval.status === 'partial'
   const sourceTypes = [
     ...new Set([
       ...(filters.source ? [filters.source] : []),
@@ -205,24 +204,18 @@ 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'>
-          {incomplete && documents.length === 0 ? (
-            'Search is incomplete.'
+          {documents.length === 0 ? (
+            'Search found no results.'
           ) : (
             <>
               <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'>
@@ -256,15 +249,7 @@ export function KnowledgeSearchResults({
           ))}
         </div>
       )}
-      {documents.length === 0 ? (
-        !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>
-        )
-      ) : (
+      {documents.length > 0 && (
         <div className='flex flex-col' onKeyDown={handleResultsKeyDown}>
           {documents.map((result) => {
             const source = toSource(result, query, scope)

apps/sim/lib/knowledge/__integration__/search-latency.integration.ts

@@ -480,17 +480,25 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
     expect(run).not.toHaveBeenCalled()
   })
 
-  it.each(['keyword', 'both'] as const)(
-    'handles %s SQL branches exceeding the deadline explicitly',
+  it.each(['vector', 'keyword', 'both'] as const)(
+    'keeps the Assistant budget when %s SQL branches are delayed',
     async (delayedLegs) => {
+      diagnosticLog?.mockClear()
+      let vectorDelayed = false
       const query = SearchBudget.prototype.query
       const delayed = vi.spyOn(SearchBudget.prototype, 'query').mockImplementation(function <T>(
         this: SearchBudget,
         stage: SearchStage,
         run: (executor: SearchExecutor) => PromiseLike<T>
       ): Promise<T> {
         return query.call(this, stage, async (tx) => {
-          if (delayedLegs === 'both' || this.leg === 'keyword')
+          if (delayedLegs === 'vector' && this.leg === 'vector' && !vectorDelayed) {
+            vectorDelayed = true
+            await tx.execute(sql`SELECT pg_sleep(4)`)
+          } else if (
+            delayedLegs === 'both' ||
+            (delayedLegs === 'keyword' && this.leg === 'keyword')
+          )
             await tx.execute(sql`SELECT pg_sleep(9)`)
           return run(tx)
         }) as Promise<T>
@@ -509,6 +517,10 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
             }),
           }
         )
+        const completed = diagnosticLog?.mock.calls.find(
+          ([message]) => message === 'Knowledge search completed'
+        )
+        expect(diagnosticSchema.parse(completed?.[1]).vectorBudgetMs).toBe(8000)
         if (delayedLegs === 'both') {
           expect(result).toMatchObject({
             success: true,
@@ -521,14 +533,19 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
         }
         expect(result).toMatchObject({
           success: true,
-          data: { retrieval: { status: 'partial', timedOutLegs: ['keyword'] } },
+          data: {
+            retrieval:
+              delayedLegs === 'vector'
+                ? { status: 'complete', timedOutLegs: [] }
+                : { status: 'partial', timedOutLegs: ['keyword'] },
+          },
         })
         const parsed = resultSchema.parse(result)
         expect(parsed.data.results.length).toBeGreaterThan(0)
         expect(
           parsed.data.results.every((row) => row.knowledgeBaseId === ids.knowledgeBaseId)
         ).toBe(true)
-        report['deadline.partial'] = { resultCount: parsed.data.results.length }
+        report[`assistant.deadline.${delayedLegs}`] = { resultCount: parsed.data.results.length }
       } finally {
         delayed.mockRestore()
       }
@@ -539,6 +556,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
   it.each(['vector', 'both'] as const)(
     'returns incomplete dashboard coverage when %s SQL branches exceed their deadline',
     async (delayedLegs) => {
+      diagnosticLog?.mockClear()
       const authenticate = vi.spyOn(internalSessionAuth, 'authenticate').mockResolvedValue({
         kind: 'session',
         userId: ids.aliceId,
@@ -552,7 +570,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
       ): Promise<T> {
         return query.call(this, stage, async (tx) => {
           if (delayedLegs === 'both' || this.leg === 'vector')
-            await tx.execute(sql`SELECT pg_sleep(9)`)
+            await tx.execute(sql`SELECT pg_sleep(${delayedLegs === 'both' ? 9 : 4})`)
           return run(tx)
         }) as Promise<T>
       })
@@ -569,6 +587,13 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
           })
         )
         expect(response.status).toBe(200)
+        const completed = diagnosticLog?.mock.calls.find(
+          ([message]) => message === 'Knowledge search completed'
+        )
+        const diagnostics = diagnosticSchema.parse(completed?.[1])
+        expect(diagnostics.vectorBudgetMs).toBe(3000)
+        expect(diagnostics.stages.vector.totalMs).toBeGreaterThan(2500)
+        expect(diagnostics.stages.vector.totalMs).toBeLessThan(4000)
         const data = workspaceKnowledgeSearchDataSchema.parse((await response.json()).data)
         expect(data.retrieval).toEqual({
           status: 'partial',

apps/sim/lib/knowledge/application/search.test.ts

@@ -234,6 +234,25 @@ describe('knowledge search application use case', () => {
     }
   )
 
+  it.each([
+    { surface: 'dashboard' as const, vectorBudgetMs: 3000 },
+    { surface: 'copilot' as const, vectorBudgetMs: undefined },
+    { surface: 'workflow' as const, vectorBudgetMs: undefined },
+  ])('forwards only the configured vector budget for $surface', async (options) => {
+    await searchKnowledge.execute({
+      principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' },
+      input: {
+        knowledgeBaseIds: ['knowledge-1'],
+        query: 'release',
+        topK: 10,
+        ...options,
+      },
+    })
+    expect(mocks.executeSearch).toHaveBeenCalledWith(
+      expect.objectContaining({ vectorBudgetMs: options.vectorBudgetMs })
+    )
+  })
+
   describe.each(['workspace', 'organization'] as const)('%s ranking policy', (scope) => {
     beforeEach(() => {
       if (scope === 'organization') {

apps/sim/lib/knowledge/application/search.ts

@@ -90,8 +90,10 @@ export class KnowledgeSearchProvenanceUnavailableError extends Error {
 export type KnowledgeSearchTagFilter = KnowledgeTagNameFilter
 
 export interface SearchKnowledgeInput {
-  /** Only surfaces displaying retrieval status may accept incomplete evidence. */
+  /** Allows returning available results when a retrieval leg times out. */
   allowPartialResults?: boolean
+  /** Trusted adapter's vector retrieval budget; omitted callers use the shared default. */
+  vectorBudgetMs?: number
   /** Optional assertion from a trusted adapter or public contract. */
   workspaceId?: string
   organizationId?: string
@@ -404,6 +406,7 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({
       : input.topK
     const retrieved = await measureSearchStage('retrieval', () =>
       retrieveKnowledgeSearch({
+        vectorBudgetMs: input.vectorBudgetMs,
         knowledgeBaseIds,
         topK: candidateTopK,
         filters: input.filters,

apps/sim/lib/knowledge/search/budget.ts

@@ -19,7 +19,7 @@ export class SearchDeadlineError extends Error {
   }
 }
 
-/** One leg's state, using the hybrid request's shared deadline across every refill and fallback. */
+/** One leg's absolute deadline, reused across every refill and fallback. */
 export class SearchBudget {
   timedOut = false
 

apps/sim/lib/knowledge/search/diagnostics.ts

@@ -70,6 +70,7 @@ export interface SearchDiagnosticMetadata {
   boostRecency?: boolean
   embeddingDimensions?: number
   vectorRanking?: 'exact' | 'binary-rerank'
+  vectorBudgetMs?: number
   vectorCandidateLimit?: number
   vectorCandidateCount?: number
   resultCount?: number

apps/sim/lib/knowledge/search/queries.test.ts

@@ -15,12 +15,14 @@ import {
   WORKSPACE_ACCESS_TOKENS,
 } from '@/lib/knowledge/access/types'
 import { buildTagFilterCondition } from '@/lib/knowledge/documents/tag-filter'
+import { SearchBudget } from '@/lib/knowledge/search/budget'
 import {
   executeKeywordSearch,
   getStructuredTagFilters,
   handleTagAndVectorSearch,
   handleTagOnlySearch,
   handleVectorOnlySearch,
+  retrieveKnowledgeSearch,
   type SearchParams,
 } from '@/lib/knowledge/search/queries'
 import type { StructuredFilter } from '@/lib/knowledge/types'
@@ -36,6 +38,59 @@ const embeddingTable = {
   boolean1: 'boolean1',
 }
 
+describe('retrieval leg budgets', () => {
+  afterEach(() => vi.restoreAllMocks())
+
+  it.each([undefined, 3000])(
+    'applies vector budget %s without shortening keyword or tag retrieval',
+    async (vectorBudgetMs) => {
+      resetDbChainMock()
+      vi.spyOn(performance, 'now').mockReturnValue(1000)
+      const remaining = SearchBudget.prototype.remaining
+      const deadlines = new Map<string, number>()
+      vi.spyOn(SearchBudget.prototype, 'remaining').mockImplementation(function (
+        this: SearchBudget
+      ) {
+        deadlines.set(this.leg, this.deadline)
+        return remaining.call(this)
+      })
+      const access: UserAccessScope = {
+        kind: 'user',
+        userId: 'user-1',
+        tokens: WORKSPACE_ACCESS_TOKENS,
+      }
+      const params = {
+        knowledgeBaseIds: ['knowledge-1'],
+        topK: 10,
+        access,
+        accessProvider: {
+          get: async () => access,
+          getForConnectors: async () => access,
+          getForDocuments: async () => access,
+        },
+        searchMode: 'hybrid' as const,
+        vectorBudgetMs,
+      }
+      await retrieveKnowledgeSearch({
+        ...params,
+        query: 'release',
+        queryVector: { vector: '[1,0]', dimensions: 1536 },
+      })
+      await retrieveKnowledgeSearch({
+        ...params,
+        structuredFilters: [
+          { tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'release' },
+        ],
+      })
+      expect(Object.fromEntries(deadlines)).toEqual({
+        vector: 1000 + (vectorBudgetMs ?? 8000),
+        keyword: 9000,
+        tags: 9000,
+      })
+    }
+  )
+})
+
 /**
  * The global `drizzle-orm` mock renders `sql` fragments to a `?`-placeholder
  * string via `toSQL()`, so we can assert the exact predicate each filter builds.