fix(search): defer previews when verification shares their budget, and bound session opening by the account deadline

Author: waleedlatif1

apps/sim/lib/sim-search/live/account-session.ts

@@ -19,7 +19,7 @@ import { createManagedSearchMcpClient } from '@/lib/sim-search/live/managed-mcp'
 import { isManagedSearchMcpProvider } from '@/lib/sim-search/live/managed-mcp-config'
 import { readNotionMcp, searchNotionMcp } from '@/lib/sim-search/live/notion-mcp'
 import { createPolicyVerifier } from '@/lib/sim-search/live/policy'
-import type { LiveSearchPolicy } from '@/lib/sim-search/live/policy-schema'
+import { type LiveSearchPolicy, requiresScopedRetrieval } from '@/lib/sim-search/live/policy-schema'
 import { livePolicyFor, loadLiveSearchPolicies } from '@/lib/sim-search/live/policy-store'
 import { LIVE_SEARCH_PROVIDER_CATALOG } from '@/lib/sim-search/live/provider-catalog'
 import { isSearchProviderEnabled } from '@/lib/sim-search/live/provider-rollout'
@@ -57,6 +57,11 @@ export interface LiveAccountSession {
    * cancels the read when its preview is no longer wanted.
    */
   previewText(document: NativeDocument, signal: AbortSignal): Promise<string | undefined>
+  /**
+   * True when `verify` spends requests from the member budget `previewText` also draws on, so a
+   * preview read before verification finishes could leave a later result unverified.
+   */
+  verifySharesPreviewBudget: boolean
   /** Checks a fetched document against the source settings saved now, not at session start. */
   verifyCurrent(document: NativeDocument): Promise<boolean>
 }
@@ -178,14 +183,15 @@ export async function openLiveAccountSession(
       signal,
       pool: input.pool,
     })
-    if (service) return service
+    if (service) return { ...service, spendsMemberBudget: service.spendsMemberBudget === true }
     const verifyPolicy = createPolicyVerifier(provider, memberPolicy, client, origin, mcp, {
       fresh,
     })
     return {
       policy: memberPolicy,
       partial: false,
       scopeSearch: undefined,
+      spendsMemberBudget: requiresScopedRetrieval(provider, memberPolicy),
       verify: (document: Reference) => verifyPolicy(document, document.accessMetadata),
     }
   }
@@ -200,6 +206,7 @@ export async function openLiveAccountSession(
     },
     policy: boundary.policy,
     servicePartial: boundary.partial,
+    verifySharesPreviewBudget: boundary.spendsMemberBudget,
     async search(search) {
       const searchOnce = async (request: NativeSearchInput): Promise<NativePage> => {
         const scoped = boundary.scopeSearch ? boundary.scopeSearch(request) : request

apps/sim/lib/sim-search/live/application.test.ts

@@ -765,6 +765,33 @@ describe('authorized live retrieval', () => {
       expect(previewedIds()).toEqual(['a', 'd', 'b'])
     })
 
+    it('reads no preview until a verification that spends the same request budget finishes', async () => {
+      const googleDoc = 'application/vnd.google-apps.document'
+      vi.mocked(livePolicyFor).mockReturnValue({
+        ...defaultLiveSearchPolicy(),
+        fileTypes: [googleDoc],
+      })
+      mocks.search.mockResolvedValue({ documents: excerpts })
+      const pendingB = createDeferred<unknown>()
+      mocks.json.mockReset().mockImplementation(async (path: string) => {
+        const id = path.split('/').at(-1)
+        return id === 'b' ? pendingB.promise : { id, mimeType: googleDoc }
+      })
+      mocks.preview.mockImplementation(
+        async (_provider, _client, { id }: { id: string }) => `Full ${id} launch text`
+      )
+      const pending = searchLiveKnowledge.execute({ principal, input })
+      await vi.waitFor(() => expect(mocks.json).toHaveBeenCalledTimes(4))
+      expect(mocks.preview).not.toHaveBeenCalled()
+      pendingB.resolve({ id: 'b', mimeType: googleDoc })
+      expect(contents(await pending)).toEqual([
+        'Full a launch text',
+        'Full b launch text',
+        'Full c launch text',
+        'Excerpt d',
+      ])
+    })
+
     it('never surfaces the preview text of a candidate verification denies', async () => {
       mocks.search.mockResolvedValue({ documents: excerpts })
       mocks.service.mockResolvedValue({
@@ -1262,6 +1289,17 @@ describe('authorized live retrieval', () => {
       await pending
     })
 
+    it('times an account out at its deadline when its session never opens', async () => {
+      mocks.service.mockReturnValue(new Promise(() => {}))
+      let result: Awaited<ReturnType<typeof searchLiveKnowledge.execute>> | undefined
+      void searchLiveKnowledge.execute({ principal, input }).then((value) => {
+        result = value
+      })
+      await vi.advanceTimersByTimeAsync(12_000)
+      expect(result?.live?.accounts[0]?.status).toBe('timeout')
+      expect(mocks.search).not.toHaveBeenCalled()
+    })
+
     it('gives each account its own deadline from its slot, and skips one the search left no time', async () => {
       mocks.accounts.mockResolvedValue(accounts(13))
       mocks.resolveAccount.mockImplementation(() =>

apps/sim/lib/sim-search/live/application.ts

@@ -345,7 +345,8 @@ const PREVIEW_SOURCE_CHARACTERS = 100_000
  * preview is a passage around the match like other providers' results rather than a prompt to
  * read. A result is read only once verification permits it and it is among the top results the
  * denials so far leave, so reads overlap the verification of other results without reading one
- * that is denied. `attach` gives text only to the verified results it is passed and cancels every
+ * that is denied. A session whose verification spends the budget reads draw on reads only after
+ * verification, so a preview never costs a result its check. `attach` gives text only to the verified results it is passed and cancels every
  * other read. A read that fails or misses its deadline keeps the provider's preview.
  */
 function streamPreviewReads(
@@ -609,7 +610,11 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({
       let previewed: LiveCandidate[]
       try {
         verification = await measureSearchStage('live.verify', () =>
-          verifyCandidates(session, eligible, previews?.verified)
+          verifyCandidates(
+            session,
+            eligible,
+            session.verifySharesPreviewBudget ? undefined : previews?.verified
+          )
         )
         matching = matchingOf(verification.permitted)
         previewed = previews
@@ -781,17 +786,22 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({
           const resolved = await measureSearchStage('live.resolve', () =>
             waitWithAbort(resolveListedLiveAccount(input, userId, account, signal), signal)
           )
+          const opening = openLiveAccountSession({
+            owner: input,
+            userId,
+            resolved,
+            policies,
+            signal,
+            pool,
+            searches: natives.length,
+          })
+          /** Opening reads the database, which takes no signal; a session that opens late is closed. */
           session = await measureSearchStage('live.session', () =>
-            openLiveAccountSession({
-              owner: input,
-              userId,
-              resolved,
-              policies,
-              signal,
-              pool,
-              searches: natives.length,
-            })
-          )
+            waitWithAbort(opening, signal)
+          ).catch((error: unknown) => {
+            opening.then((late) => late.close()).catch(() => {})
+            throw error
+          })
           return { resolved, session }
         })
         if (!opened)

apps/sim/lib/sim-search/live/github-service.ts

@@ -132,6 +132,7 @@ export function createGitHubServiceVerifier(
   return {
     policy,
     partial: false,
+    spendsMemberBudget: true,
     async verify(document: Pick<NativeDocument, 'id' | 'container' | 'kind'>) {
       const repository = document.container?.toLowerCase()
       if (!repository) return false

apps/sim/lib/sim-search/live/service-session.ts

@@ -30,6 +30,8 @@ import { CONNECTOR_META_REGISTRY } from '@/connectors/registry'
 interface LiveServiceSession {
   policy: LiveSearchPolicy
   partial: boolean
+  /** True when verification also spends requests from the member client's budget. */
+  spendsMemberBudget?: boolean
   verify(document: Pick<NativeDocument, 'id' | 'container' | 'kind'>): Promise<boolean>
   scopeSearch?(input: NativeSearchInput): NativeSearchInput | null
 }