fix(search): enforce retrieval budgets in integration coverage

Author: icecrasher321

.github/workflows/test-build.yml

@@ -182,6 +182,12 @@ jobs:
           lib/memory/message-provenance.postgres.test.ts
           executor/handlers/agent/memory-harness.postgres.test.ts
 
+      - name: Verify Search vector projection upgrade in PostgreSQL
+        working-directory: packages/db
+        env:
+          KNOWLEDGE_ACL_TEST_DATABASE_URL: postgresql://postgres:postgres@127.0.0.1:5432/sim_auth_scim
+        run: bunx vitest run script-migrations/0016_backfill_search_vectors.postgres.test.ts
+
       - name: Verify Search progress, pagination, and outbox scheduling in PostgreSQL
         working-directory: apps/sim
         env:

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

@@ -36,7 +36,6 @@ import {
   createKnowledgeAclFixtureIds,
   seedKnowledgeAclFixture,
 } from '@/lib/knowledge/__integration__/seed-source-access-fixture'
-import { searchScopedKnowledge } from '@/lib/knowledge/application/workspace-search'
 import {
   SearchBudget,
   SearchDeadlineError,
@@ -166,6 +165,9 @@ const diagnosticSchema = z
     surface: z.enum(['dashboard', 'copilot']),
     outcome: z.literal('success'),
     elapsedMs: z.number(),
+    vectorBudgetMs: z.number().positive(),
+    retrievalStatus: z.enum(['complete', 'partial']),
+    timedOutLegs: z.array(z.enum(['vector', 'keyword', 'tags'])),
     toolResultBytes: z.number().int().nonnegative().optional(),
     passageBytes: z.number().int().nonnegative().optional(),
     maxPassageBytes: z.number().int().nonnegative().optional(),
@@ -217,6 +219,45 @@ async function search(
   )
 }
 
+async function searchDashboard() {
+  const authenticate = vi.spyOn(internalSessionAuth, 'authenticate').mockResolvedValue({
+    kind: 'session',
+    userId: ids.aliceId,
+    sessionId: 'fixture-dashboard',
+  })
+  try {
+    const response = await searchRoute(
+      new NextRequest('http://localhost/api/knowledge/search', {
+        method: 'POST',
+        headers: { 'content-type': 'application/json' },
+        body: JSON.stringify({
+          workspaceId: ids.workspaceId,
+          query: 'Orion deployment',
+          topK: 15,
+        }),
+      })
+    )
+    expect(response.status).toBe(200)
+    return {
+      success: true as const,
+      data: workspaceKnowledgeSearchDataSchema.parse((await response.json()).data),
+    }
+  } finally {
+    authenticate.mockRestore()
+  }
+}
+
+/** Allow either index or filtered plans, but require successful retrieval within the real surface budget. */
+function expectCompleteVectorSearch(diagnostics: z.infer<typeof diagnosticSchema>) {
+  const budget = diagnostics.surface === 'dashboard' ? 3000 : 8000
+  expect(diagnostics).toMatchObject({
+    vectorBudgetMs: budget,
+    retrievalStatus: 'complete',
+    timedOutLegs: [],
+  })
+  expect(diagnostics.stages.vector.totalMs).toBeLessThanOrEqual(budget)
+}
+
 async function sample(label: string, run: () => ReturnType<typeof search>) {
   captured.length = 0
   diagnosticLog?.mockClear()
@@ -614,11 +655,6 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
     '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,
-        sessionId: 'fixture-dashboard',
-      })
       const query = SearchBudget.prototype.query
       const delayed = vi.spyOn(SearchBudget.prototype, 'query').mockImplementation(function <T>(
         this: SearchBudget,
@@ -632,26 +668,14 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
         }) as Promise<T>
       })
       try {
-        const response = await searchRoute(
-          new NextRequest('http://localhost/api/knowledge/search', {
-            method: 'POST',
-            headers: { 'content-type': 'application/json' },
-            body: JSON.stringify({
-              workspaceId: ids.workspaceId,
-              query: 'Orion deployment',
-              topK: 15,
-            }),
-          })
-        )
-        expect(response.status).toBe(200)
+        const { data } = await searchDashboard()
         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',
           timedOutLegs: delayedLegs === 'both' ? ['vector', 'keyword'] : ['vector'],
@@ -666,7 +690,6 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
         report[`dashboard.deadline.${delayedLegs}`] = { resultCount: data.results.length }
       } finally {
         delayed.mockRestore()
-        authenticate.mockRestore()
       }
     },
     30_000
@@ -675,7 +698,8 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
   it('records first and repeated application searches with the actual SQL plans', async () => {
     const before = embeddingCalls
     for (let iteration = 0; iteration < 2; iteration++) {
-      const { result, plans } = await sample(`broad.${iteration}`, () => search())
+      const { result, plans, diagnostics } = await sample(`broad.${iteration}`, () => search())
+      expectCompleteVectorSearch(diagnostics)
       expect(result.data.results).toHaveLength(15)
       expect(result.data.results.every((row) => row.knowledgeBaseId === ids.knowledgeBaseId)).toBe(
         true
@@ -703,9 +727,10 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
 
   it('preserves exact-neighbor recall across different query vectors', async () => {
     for (const topic of [3, 11, 23]) {
-      const { plans } = await sample(`topic.${topic}`, () =>
+      const { plans, diagnostics } = await sample(`topic.${topic}`, () =>
         search(ids.aliceId, `Topic ${topic} deployment`)
       )
+      expectCompleteVectorSearch(diagnostics)
       const candidates = plans.find((plan) => plan.kind === 'vector')!
       assertCompactCandidates(candidates.plan[0].Plan)
       const rerank = plans.find((plan) => plan.kind === 'rerank')!
@@ -723,19 +748,10 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
   }, 180_000)
 
   it('compares the Search tab and Assistant with the same person, query and index', async () => {
-    const dashboard = await sample('dashboard', async () => {
-      const result = await searchScopedKnowledge.execute({
-        principal: { kind: 'session', userId: ids.aliceId, sessionId: 'fixture-dashboard' },
-        input: {
-          workspaceId: ids.workspaceId,
-          query: 'Orion deployment',
-          topK: 15,
-          surface: 'dashboard',
-        },
-      })
-      return resultSchema.parse({ success: true, data: result })
-    })
+    const dashboard = await sample('dashboard', searchDashboard)
     const assistant = await sample('assistant.comparison', () => search())
+    expectCompleteVectorSearch(dashboard.diagnostics)
+    expectCompleteVectorSearch(assistant.diagnostics)
     expect(dashboard.diagnostics.surface).toBe('dashboard')
     expect(assistant.diagnostics.surface).toBe('copilot')
     expect(dashboard.diagnostics.stages.result_provenance).toBeUndefined()
@@ -767,21 +783,9 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
       for (const surface of ['copilot', 'dashboard'] as const) {
         const { result, plans, diagnostics } = await sample(
           `filtered-neighborhood.${surface}`,
-          async () => {
-            if (surface === 'copilot') return search()
-            const data = await searchScopedKnowledge.execute({
-              principal: { kind: 'session', userId: ids.aliceId, sessionId: 'fixture-dashboard' },
-              input: {
-                workspaceId: ids.workspaceId,
-                query: 'Orion deployment',
-                topK: 15,
-                surface,
-              },
-            })
-            return resultSchema.parse({ success: true, data })
-          }
+          surface === 'copilot' ? () => search() : searchDashboard
         )
-        expect(diagnostics).toMatchObject({ retrievalStatus: 'complete', timedOutLegs: [] })
+        expectCompleteVectorSearch(diagnostics)
         expect(result.data.results).toHaveLength(15)
         const rerank = plans.find((plan) => plan.kind === 'rerank')!
         expect(rerank).toBeDefined()
@@ -852,6 +856,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
   }, 180_000)
 
   it('runs two independent Assistant searches concurrently', async () => {
+    diagnosticLog?.mockClear()
     const start = performance.now()
     const results = await Promise.all([search(), search(ids.aliceId, 'Engineering operations')])
     report.concurrent = {
@@ -860,6 +865,12 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
     }
     saveReport()
     for (const result of results) expect(result.data.results).toHaveLength(15)
+    const completed = diagnosticLog!.mock.calls.filter(
+      ([message]) => message === 'Knowledge search completed'
+    )
+    expect(completed).toHaveLength(2)
+    for (const [, metadata] of completed)
+      expectCompleteVectorSearch(diagnosticSchema.parse(metadata))
   }, 180_000)
 
   it('checks live reader access on every search, including after revocation', async () => {

packages/db/script-migrations-paused-billing-attribution.test.ts

@@ -452,6 +452,7 @@ describe('script migration registry', () => {
       '0013_backfill_legacy_knowledge_base_workspaces',
       '0014_require_knowledge_base_owner',
       '0015_backfill_embedding_search',
+      '0016_backfill_search_vectors',
     ])
   })
 })

packages/db/script-migrations/0016_backfill_search_vectors.postgres.test.ts

@@ -16,9 +16,9 @@ describe.runIf(Boolean(databaseUrl))('search projection upgrade in PostgreSQL',
     const url = new URL(databaseUrl!)
     if (
       !['localhost', '127.0.0.1'].includes(url.hostname) ||
-      !url.pathname.startsWith('/sim_acl_test')
+      (!url.pathname.startsWith('/sim_acl_test') && url.pathname !== '/sim_auth_scim')
     ) {
-      throw new Error('Projection tests require a disposable local sim_acl_test database')
+      throw new Error('Projection tests require a disposable local integration database')
     }
     admin = postgres(url.toString(), { max: 1, onnotice: () => undefined })
     await admin.unsafe(`CREATE SCHEMA "${schemaName}"`)

scripts/test-knowledge-acls.ts

@@ -180,7 +180,7 @@ try {
       env: environment,
     }
   )
-  if (!scale && testFilters.length === 0)
+  if (!scale && testFilters.length === 0) {
     run(
       'bunx',
       [
@@ -195,6 +195,15 @@ try {
         env: environment,
       }
     )
+    run(
+      'bunx',
+      ['vitest', 'run', 'script-migrations/0016_backfill_search_vectors.postgres.test.ts'],
+      {
+        cwd: path.join(root, 'packages/db'),
+        env: environment,
+      }
+    )
+  }
   logger.info(
     scale
       ? 'Opt-in knowledge scale measurements passed'