improvement(search): read up to 10 documents per MCP read_document call (#8981)

* improvement(search): read up to 10 documents per MCP read_document call, five at a time

* test(search): finish later batched reads first so the order assertion can fail

Author: waleedlatif1

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

@@ -18,7 +18,7 @@ export const organizationKnowledgeMcpContract = defineRouteContract({
 
 export const liveSearchMcpSchema = searchWorkspaceInputSchema.strict()
 /** Documents one read_document call may read; each counts as one read against the rate limit. */
-const MAX_MCP_READ_DOCUMENTS = 5
+const MAX_MCP_READ_DOCUMENTS = 10
 
 export const readLiveDocumentMcpSchema = readDocumentInputSchema
   .extend({

apps/sim/lib/knowledge/mcp/server.test.ts

@@ -208,6 +208,24 @@ describe('batched document reads', () => {
     expect(mocks.liveRead).toHaveBeenCalledTimes(2)
   })
 
+  it('reads at most five documents from providers at once, in request order', async () => {
+    create()
+    let inFlight = 0
+    let peak = 0
+    mocks.liveRead.mockImplementation(async ({ input }: { input: { documentId: string } }) => {
+      inFlight++
+      peak = Math.max(peak, inFlight)
+      const laterFinishesFirst = 10 - Number(input.documentId.slice('doc-'.length))
+      for (let tick = 0; tick < laterFinishesFirst; tick++) await Promise.resolve()
+      inFlight--
+      return readResult(input.documentId)
+    })
+    const ids = Array.from({ length: 10 }, (_, index) => `doc-${index}`)
+    const result = await call('read_document', { documentIds: ids })
+    expect(peak).toBe(5)
+    expect(parsedBody(result).documents.map((document) => document.documentId)).toEqual(ids)
+  })
+
   it('reads nothing when the batch exceeds the remaining rate limit', async () => {
     create()
     mocks.rateLimit.mockResolvedValueOnce(new Response(null, { status: 429 }))

apps/sim/lib/knowledge/mcp/server.ts

@@ -12,6 +12,7 @@ import {
 import type { V2ApiKeyAuthContext } from '@/lib/api/server/routes/v2-api-key-auth'
 import { v2RateLimits } from '@/lib/api/server/routes/v2-json-route'
 import { afterResponse } from '@/lib/core/utils/after-response'
+import { mapWithConcurrency } from '@/lib/core/utils/concurrency'
 import { organizationSearchChatOperation } from '@/lib/knowledge/application/chat-operations'
 import { knowledgeOperations } from '@/lib/knowledge/application/operations'
 import {
@@ -31,6 +32,8 @@ import { v2CaughtOrchestrationError } from '@/app/api/v2/lib/response'
 import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
 
 const logger = createLogger('KnowledgeMcp')
+/** Provider reads one read_document call runs at once, however many documents it names. */
+const READ_DOCUMENT_CONCURRENCY = 5
 
 /** The message a caller may see for an expected failure; undefined for an internal error. */
 async function callerErrorMessage(error: unknown): Promise<string | undefined> {
@@ -201,8 +204,10 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe
             })
           if (documentId) return projectResult(projectReadResult(await read(documentId)), registry)
           /** One unreadable document leaves the others; each failure is reported in its place. */
-          const documents = await Promise.all(
-            batch.map(async (id) => {
+          const documents = await mapWithConcurrency(
+            batch,
+            READ_DOCUMENT_CONCURRENCY,
+            async (id) => {
               try {
                 return projectReadResult(await read(id))
               } catch (error) {
@@ -218,7 +223,7 @@ export function createKnowledgeMcpServer(context: KnowledgeMcpContext): McpServe
                   error: message ?? 'Unable to read this document. Please try again.',
                 }
               }
-            })
+            }
           )
           return projectResult({ documents }, registry)
         },