Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions apps/sim/lib/mcp/service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
17 changes: 12 additions & 5 deletions apps/sim/lib/sim-search/live/account-session.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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>
Expand Down Expand Up @@ -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.')
Expand Down Expand Up @@ -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) {
Expand Down
3 changes: 1 addition & 2 deletions apps/sim/lib/sim-search/live/application.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
}))
Expand Down
21 changes: 16 additions & 5 deletions apps/sim/lib/sim-search/live/application.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)
)
)
Expand All @@ -634,6 +636,7 @@ export const searchLiveKnowledge = defineAuthorizedKnowledgeUseCase({
return natives.map((target) => failed(error, target))
} finally {
settled.abort()
await session?.close()
}
}
let searched: SearchedQuery[]
Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand All @@ -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')
Expand Down
16 changes: 1 addition & 15 deletions apps/sim/lib/sim-search/live/coda-mcp.ts
Original file line number Diff line number Diff line change
@@ -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)
Expand Down
Loading
Loading