fix(search): budget batched kinds like separate searches and hide only retried search failures

Author: waleedlatif1

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/search-activity.test.tsx

@@ -236,6 +236,21 @@ describe('search in shared tool activity', () => {
     )
   })
 
+  it('keeps a failed search that no later search retried', () => {
+    const read: ToolCallData = {
+      id: 'read',
+      toolName: 'read_document',
+      displayTitle: 'Reading document',
+      activityDescription: 'Reading the launch plan',
+      status: 'success',
+    }
+    render([read])
+    expect(header()).toBeNull()
+    render([{ ...completedSearch('one', 'First query'), status: 'error' }, read])
+    expect(header()).not.toBeNull()
+    expect(container.textContent).not.toMatch(/failed/i)
+  })
+
   it('keeps the last search visible when every search failed', () => {
     render([
       { ...completedSearch('one', 'First query'), status: 'error' },

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/tool-activity-group.tsx

@@ -36,14 +36,19 @@ function isFailedTool(tool: ToolCallData): boolean {
 }
 
 /**
- * Drops a search that errored when another call follows it: the model corrected the query
- * and moved on, so the failed attempt is not part of what the user reads.
- * The last call always stays, so a run whose searches all failed still shows its outcome.
+ * Drops a search that errored when a later search follows it: the model corrected the query
+ * and searched again, so the failed attempt is not part of what the user reads. The last search
+ * always stays, so a failure that no search retried, or a run whose searches all failed, still
+ * shows its outcome.
  */
-function withoutRecoveredSearchFailures(tools: ToolCallData[]): ToolCallData[] {
+function withoutRetriedSearchFailures(tools: ToolCallData[]): ToolCallData[] {
+  let lastSearch = -1
+  for (const [index, tool] of tools.entries()) {
+    if (tool.toolName === SearchWorkspace.id) lastSearch = index
+  }
   return tools.filter(
     (tool, index) =>
-      index === tools.length - 1 ||
+      index >= lastSearch ||
       tool.toolName !== SearchWorkspace.id ||
       tool.status !== ToolCallStatus.error
   )
@@ -216,7 +221,7 @@ export function ToolActivityGroup({
   isLive = false,
 }: ToolActivityGroupProps) {
   const [expanded, setExpanded] = useState(false)
-  const tools = withoutRecoveredSearchFailures(calls)
+  const tools = withoutRetriedSearchFailures(calls)
   const statusTool = getActivityStatusTool(tools)
   if (!statusTool) return null
   const headerTool = getActivityHeaderTool(tools, statusTool)

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

@@ -4,7 +4,7 @@ import type { PinnedConnectionPool } from '@/lib/core/security/input-validation.
 import type { ResolvedLiveAccount } from '@/lib/sim-search/live/accounts'
 import { createCodaMcpClient, readCodaMcp, searchCodaMcp } from '@/lib/sim-search/live/coda-mcp'
 import { createAdminGitLabSession } from '@/lib/sim-search/live/gitlab-admin'
-import { createNativeClient } from '@/lib/sim-search/live/http'
+import { createNativeClient, NATIVE_SEARCH_REQUEST_BUDGET } from '@/lib/sim-search/live/http'
 import { createPolicyVerifier } from '@/lib/sim-search/live/policy'
 import type { LiveSearchPolicy } from '@/lib/sim-search/live/policy-schema'
 import { livePolicyFor, loadLiveSearchPolicies } from '@/lib/sim-search/live/policy-store'
@@ -40,6 +40,8 @@ interface OpenLiveAccountSessionInput {
   policies: Record<string, unknown>
   signal: AbortSignal
   pool?: PinnedConnectionPool
+  /** Native searches this session serves; each gets the budget a separate call would have. */
+  searches?: number
 }
 
 /**
@@ -58,7 +60,13 @@ export async function openLiveAccountSession(
   const client =
     account.type === 'managed_mcp'
       ? null
-      : createNativeClient({ origin, accessToken: resolved.accessToken, signal, pool: input.pool })
+      : createNativeClient({
+          origin,
+          accessToken: resolved.accessToken,
+          signal,
+          pool: input.pool,
+          requestBudget: NATIVE_SEARCH_REQUEST_BUDGET * (input.searches ?? 1),
+        })
   const admin =
     'adminSource' in resolved && client
       ? await createAdminGitLabSession({

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

@@ -474,7 +474,15 @@ export const searchLiveKnowledge = defineAuthorizedKnowledgeUseCase({
           resolveListedLiveAccount(input, userId, account)
         )
         const session = await measureSearchStage('live.session', () =>
-          openLiveAccountSession({ owner: input, userId, resolved, policies, signal, pool })
+          openLiveAccountSession({
+            owner: input,
+            userId,
+            resolved,
+            policies,
+            signal,
+            pool,
+            searches: natives.length,
+          })
         )
         return await Promise.all(
           natives.map((native) =>

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

@@ -30,6 +30,26 @@ describe('native search network boundary', () => {
       })
     )
   })
+  it('stops at its request budget, which defaults to one search', async () => {
+    mocks.fetch.mockImplementation(async () => new Response('{}', { status: 200 }))
+    const clientWith = (requestBudget?: number) =>
+      createNativeClient({
+        origin: 'https://api.github.com',
+        accessToken: 'private',
+        signal: new AbortController().signal,
+        requestBudget,
+      })
+    const exhaust = async (client: ReturnType<typeof clientWith>, requests: number) => {
+      for (let request = 0; request < requests; request++) await client.json('/user')
+    }
+    await exhaust(clientWith(), 30)
+    const single = clientWith()
+    await exhaust(single, 30)
+    await expect(single.json('/user')).rejects.toThrow('Request budget reached')
+    const shared = clientWith(60)
+    await exhaust(shared, 60)
+    await expect(shared.json('/user')).rejects.toThrow('Request budget reached')
+  })
   it('does not send a token to an absolute or protocol-relative model URL', async () => {
     const client = createNativeClient({
       origin: 'https://api.github.com',

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

@@ -15,13 +15,18 @@ export class NativeSearchError extends Error {
   }
 }
 
+/** Provider requests one native search may make, including discovery and verification. */
+export const NATIVE_SEARCH_REQUEST_BUDGET = 30
+
 /** Tokens only go to a code-selected provider origin; redirects never carry credentials. */
 export function createNativeClient(input: {
   origin: string
   accessToken: string
   signal: AbortSignal
   /** Reuses connections across this client's requests; the caller owns its lifetime. */
   pool?: PinnedConnectionPool
+  /** Requests this client may make; defaults to one search's budget. */
+  requestBudget?: number
 }): NativeClient {
   let requests = 0
   async function request(
@@ -33,7 +38,7 @@ export function createNativeClient(input: {
     }
   ) {
     input.signal.throwIfAborted()
-    if (++requests > 30)
+    if (++requests > (input.requestBudget ?? NATIVE_SEARCH_REQUEST_BUDGET))
       throw new NativeSearchError('unavailable', 'Request budget reached. Narrow the query.')
     const url = new URL(input.origin)
     if (options?.googleService === 'sheets') {