fix(search): reuse managed MCP sessions within operations (#8567)

* fix(search): reuse managed MCP sessions within operations

* fix(search): close acceptance sessions and verify concurrency

Author: waleedlatif1

apps/sim/lib/mcp/service.ts

@@ -492,6 +492,18 @@ class McpService {
     return client
   }
 
+  /** An operation owns this unpooled client and must disconnect it in its finalizer. */
+  async openManagedMcpSession(
+    serverId: string,
+    scope: ResourceScope,
+    auth: McpOauthCredentials,
+    signal: AbortSignal
+  ): Promise<Pick<McpClient, 'listTools' | 'callTool' | 'disconnect'>> {
+    const config = await this.getServerConfig(serverId, scope)
+    if (!config) throw new Error('Managed MCP server is unavailable')
+    return this.createManagedOauthClient(config, auth, signal)
+  }
+
   async discoverManagedMcpTools(
     serverId: string,
     scope: string | ResourceScope,

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

@@ -3,7 +3,7 @@ import type { ResourceOwner } from '@/lib/core/resource-scope'
 import { resourceScopeFromOwner } from '@/lib/core/resource-scope'
 import type { PinnedConnectionPool } from '@/lib/core/security/input-validation.server'
 import type { ResolvedLiveAccount } from '@/lib/sim-search/live/accounts'
-import { createCodaMcpClient, readCodaMcp, searchCodaMcp } from '@/lib/sim-search/live/coda-mcp'
+import { readCodaMcp, searchCodaMcp } from '@/lib/sim-search/live/coda-mcp'
 import { readFirefliesMcp, searchFirefliesMcp } from '@/lib/sim-search/live/fireflies-mcp'
 import { createAdminGitLabSession } from '@/lib/sim-search/live/gitlab-admin'
 import { readGranolaMcp, searchGranolaMcp } from '@/lib/sim-search/live/granola-mcp'
@@ -39,6 +39,8 @@ export interface LiveAccountSession {
   policy: LiveSearchPolicy
   /** Service verification covered a bounded subset of the source's configured users. */
   servicePartial: boolean
+  /** Releases the operation-owned provider transport after all reads and checks settle. */
+  close(): Promise<void>
   search(input: NativeSearchInput): Promise<NativePage>
   /** True only when the document is inside the source boundary and the member may read it. */
   verify(document: Reference): Promise<boolean>
@@ -104,10 +106,9 @@ export async function openLiveAccountSession(
         'unavailable',
         'This provider does not support managed MCP Search.'
       )
-    return provider === 'coda'
-      ? createCodaMcpClient(owner, userId, account.id, signal, input.searches)
-      : createManagedSearchMcpClient(owner, userId, account.id, provider, signal, input.searches)
+    return createManagedSearchMcpClient(owner, userId, account.id, provider, signal, input.searches)
   }
+  const memberPolicy = livePolicyFor(input.policies, provider)
   const mcp = await openMcp()
   const searchMcp = (search: NativeSearchInput) => {
     if (!mcp) throw new NativeSearchError('unavailable', 'Managed MCP connection unavailable.')
@@ -176,9 +177,15 @@ export async function openLiveAccountSession(
       verify: (document: Reference) => verifyPolicy(document, document.accessMetadata),
     }
   }
-  const boundary = await sourceBoundary(livePolicyFor(input.policies, provider))
+  const boundary = await sourceBoundary(memberPolicy).catch(async (error: unknown) => {
+    await mcp?.close()
+    throw error
+  })
 
   return {
+    async close() {
+      await mcp?.close()
+    },
     policy: boundary.policy,
     servicePartial: boundary.partial,
     async search(search) {

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

@@ -44,10 +44,9 @@ vi.mock('@/lib/sim-search/live/policy-store', () => ({
   livePolicyFor: vi.fn(() => defaultLiveSearchPolicy()),
 }))
 vi.mock('@/lib/sim-search/live/managed-mcp', () => ({
-  createManagedSearchMcpClient: async () => ({ call: mocks.mcpCall }),
+  createManagedSearchMcpClient: async () => ({ call: mocks.mcpCall, close: async () => {} }),
 }))
 vi.mock('@/lib/sim-search/live/coda-mcp', () => ({
-  createCodaMcpClient: vi.fn(),
   searchCodaMcp: vi.fn(),
   readCodaMcp: vi.fn(),
 }))

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

@@ -607,12 +607,13 @@ export const searchLiveKnowledge = defineAuthorizedKnowledgeUseCase({
           results: [],
         }
       }
+      let session: LiveAccountSession | undefined
       try {
         signal.throwIfAborted()
         const resolved = await measureSearchStage('live.resolve', () =>
           resolveListedLiveAccount(input, userId, account)
         )
-        const session = await measureSearchStage('live.session', () =>
+        session = await measureSearchStage('live.session', () =>
           openLiveAccountSession({
             owner: input,
             userId,
@@ -623,9 +624,10 @@ export const searchLiveKnowledge = defineAuthorizedKnowledgeUseCase({
             searches: natives.length,
           })
         )
+        const currentSession = session
         return await Promise.all(
           natives.map((target) =>
-            searchQuery(account, resolved, session, target.native, statusFor(target)).catch(
+            searchQuery(account, resolved, currentSession, target.native, statusFor(target)).catch(
               (error) => failed(error, target)
             )
           )
@@ -634,6 +636,7 @@ export const searchLiveKnowledge = defineAuthorizedKnowledgeUseCase({
         return natives.map((target) => failed(error, target))
       } finally {
         settled.abort()
+        await session?.close()
       }
     }
     let searched: SearchedQuery[]
@@ -760,8 +763,9 @@ export const readLiveDocument = defineAuthorizedKnowledgeUseCase({
       : AbortSignal.timeout(15_000)
     const pool = createPinnedConnectionPool()
     let document: NativeDocument
+    let session: LiveAccountSession | undefined
     try {
-      const session = await openLiveAccountSession({
+      session = await openLiveAccountSession({
         owner: input,
         userId,
         resolved,
@@ -774,7 +778,10 @@ export const readLiveDocument = defineAuthorizedKnowledgeUseCase({
           'not_found',
           'Document is outside your organization’s search scope'
         )
-      document = await measureSearchStage('live.read', () => session.read(reference, input.filters))
+      const currentSession = session
+      document = await measureSearchStage('live.read', () =>
+        currentSession.read(reference, input.filters)
+      )
       /** Readers degrade section failures to warnings, so the signal decides cancellation. */
       signal.throwIfAborted()
       const current = await session.verifyCurrent(document)
@@ -786,7 +793,11 @@ export const readLiveDocument = defineAuthorizedKnowledgeUseCase({
           'Document is outside your organization’s search scope'
         )
     } finally {
-      pool.destroy()
+      try {
+        await session?.close()
+      } finally {
+        pool.destroy()
+      }
     }
     if (!matchesLiveFilters(document, input.documentId, reference.provider, input.filters))
       throw new OrchestrationError('not_found', 'Document is outside the selected search filters')

apps/sim/lib/sim-search/live/coda-mcp.ts

@@ -1,28 +1,14 @@
 import { createLogger } from '@sim/logger'
-import type { ResourceOwner } from '@/lib/core/resource-scope'
 import { parseCodaResourceUri } from '@/lib/sim-search/live/coda-uri'
 import { hasDateBounds, nativeText } from '@/lib/sim-search/live/dates'
 import { array, NativeSearchError, object, string } from '@/lib/sim-search/live/http'
-import {
-  createManagedSearchMcpClient,
-  type ManagedSearchMcpClient,
-} from '@/lib/sim-search/live/managed-mcp'
+import type { ManagedSearchMcpClient } from '@/lib/sim-search/live/managed-mcp'
 import type { NativeDocument, NativePage, NativeSearchInput } from '@/lib/sim-search/live/types'
 
 const logger = createLogger('CodaMcpSearch')
 
 export interface CodaMcpClient extends ManagedSearchMcpClient {}
 
-export function createCodaMcpClient(
-  owner: ResourceOwner,
-  userId: string,
-  credentialId: string,
-  signal: AbortSignal,
-  searches = 1
-): Promise<CodaMcpClient> {
-  return createManagedSearchMcpClient(owner, userId, credentialId, 'coda', signal, searches)
-}
-
 function requireCodaUri(uri: string): string {
   const parsed = parseCodaResourceUri(uri)
   if (!parsed)