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
11 changes: 11 additions & 0 deletions apps/sim/lib/core/network/context.server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,3 +30,14 @@ export function captureOutboundScope(): <T>(run: () => T) => T {
const scope = storage.getStore()
return <T>(run: () => T): T => (scope ? storage.run(scope, run) : storage.exit(run))
}

/**
* Captures the outbound scope for a long-lived object whose callbacks may run from a context
* with no scope (e.g. an abort listener fired by a timer). A scope current at call time wins,
* so a pooled object never pins its creator's organization; with neither, the call fails
* closed as usual.
*/
export function captureOutboundScopeFallback(): <T>(run: () => T) => T {
const scope = storage.getStore()
return <T>(run: () => T): T => (scope && !storage.getStore() ? storage.run(scope, run) : run())
}
80 changes: 72 additions & 8 deletions apps/sim/lib/mcp/client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import {
executionLimitsMockFns,
} from '@sim/testing/mocks/execution-limits.mock'
import { getMockLogger } from '@sim/testing/mocks/logger.mock'
import { networkConfigMock, networkConfigMockFns } from '@sim/testing/mocks/network-config.mock'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'

const { mockSdkConnect, mockSdkListTools, mockPinnedClose } = vi.hoisted(() => ({
Expand Down Expand Up @@ -48,9 +49,15 @@ vi.mock('@modelcontextprotocol/sdk/types.js', () => ({
}))

vi.mock('@/lib/core/execution-limits', () => executionLimitsMock)
vi.mock('@/lib/core/network/config.server', () => networkConfigMock)

import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
import { getMaxExecutionTimeout } from '@/lib/core/execution-limits'
import {
resolveCurrentOutboundRoute,
runWithOutboundOrganization,
} from '@/lib/core/network/context.server'
import { OutboundRoutingError } from '@/lib/core/network/routing'
import { McpClient } from '@/lib/mcp/client'
import { createGuardedMcpFetch, createPinnedPrivateMcpFetch } from '@/lib/mcp/pinned-fetch'
import {
Expand Down Expand Up @@ -186,7 +193,7 @@ describe('McpClient notification handler', () => {
expect(JSON.stringify(mockLogger.info.mock.calls)).not.toContain('also-secret')
})

it('keeps the transport on the SSRF guard when no validated address is supplied', () => {
it('keeps the transport on the SSRF guard when no validated address is supplied', async () => {
new McpClient({
config: createConfig(),
securityPolicy: { requireConsent: false, auditLevel: 'basic' },
Expand All @@ -195,12 +202,14 @@ describe('McpClient notification handler', () => {
const guarded = vi.mocked(createGuardedMcpFetch).mock.results.at(-1)?.value
expect(createGuardedMcpFetch).toHaveBeenCalledWith('https://test.example.com/mcp')
expect(createPinnedPrivateMcpFetch).not.toHaveBeenCalled()
expect(vi.mocked(StreamableHTTPClientTransport).mock.calls.at(-1)?.[1]?.fetch).toBe(
guarded.fetch
)
await vi
.mocked(StreamableHTTPClientTransport)
.mock.calls.at(-1)?.[1]
?.fetch?.('https://test.example.com/mcp')
expect(guarded.fetch).toHaveBeenCalledWith('https://test.example.com/mcp', undefined)
})

it('pins the transport to a validated private address', () => {
it('pins the transport to a validated private address', async () => {
new McpClient({
config: createConfig(),
securityPolicy: { requireConsent: false, auditLevel: 'basic' },
Expand All @@ -213,9 +222,11 @@ describe('McpClient notification handler', () => {
'https://test.example.com/mcp'
)
expect(createGuardedMcpFetch).not.toHaveBeenCalled()
expect(vi.mocked(StreamableHTTPClientTransport).mock.calls.at(-1)?.[1]?.fetch).toBe(
pinned.fetch
)
await vi
.mocked(StreamableHTTPClientTransport)
.mock.calls.at(-1)?.[1]
?.fetch?.('https://test.example.com/mcp')
expect(pinned.fetch).toHaveBeenCalledWith('https://test.example.com/mcp', undefined)
})

it('closes the pinned transport Agent when connect fails', async () => {
Expand Down Expand Up @@ -322,3 +333,56 @@ describe('McpClient notification handler', () => {
)
})
})

describe('McpClient outbound scope', () => {
beforeEach(() => {
networkConfigMockFns.mockIsOutboundRoutingEnabled.mockReturnValue(true)
networkConfigMockFns.mockResolveOutboundRoute.mockImplementation(async (organizationId) => {
if (organizationId === undefined) throw new OutboundRoutingError('MISSING_SCOPE')
return { kind: 'direct', organizationId }
})
vi.mocked(createGuardedMcpFetch).mockReturnValue({
fetch: vi.fn(async () => Response.json(await resolveCurrentOutboundRoute())),
close: mockPinnedClose,
})
})

afterEach(() => {
networkConfigMockFns.mockIsOutboundRoutingEnabled.mockReturnValue(false)
networkConfigMockFns.mockResolveOutboundRoute.mockImplementation(async () => ({
kind: 'direct',
}))
vi.mocked(createGuardedMcpFetch).mockImplementation(() => ({
fetch: vi.fn(),
close: mockPinnedClose,
}))
})

function constructedTransportFetch() {
const fetch = vi.mocked(StreamableHTTPClientTransport).mock.calls.at(-1)?.[1]?.fetch
if (!fetch) throw new Error('transport was constructed without a fetch')
return async () => (await fetch('https://test.example.com/mcp')).json()
}

it('routes a request the SDK sends from outside any scope under the constructing organization', async () => {
runWithOutboundOrganization('org_a', () => new McpClient({ config: createConfig() }))
const transportFetch = constructedTransportFetch()

await expect(transportFetch()).resolves.toEqual({ kind: 'direct', organizationId: 'org_a' })
await expect(runWithOutboundOrganization('org_b', transportFetch)).resolves.toEqual({
kind: 'direct',
organizationId: 'org_b',
})
})

it('still fails closed when the client was constructed outside any scope', async () => {
new McpClient({ config: createConfig() })
const transportFetch = constructedTransportFetch()

await expect(transportFetch()).rejects.toMatchObject({ code: 'MISSING_SCOPE' })
await expect(runWithOutboundOrganization('org_b', transportFetch)).resolves.toEqual({
kind: 'direct',
organizationId: 'org_b',
})
})
})
6 changes: 5 additions & 1 deletion apps/sim/lib/mcp/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import { createLogger } from '@sim/logger'
import { isPrivateIp } from '@sim/security/ssrf'
import { getErrorMessage } from '@sim/utils/errors'
import { getMaxExecutionTimeout } from '@/lib/core/execution-limits'
import { captureOutboundScopeFallback } from '@/lib/core/network/context.server'
import { getMcpSafeErrorDiagnostics } from '@/lib/mcp/error-diagnostics'
import { McpOauthRedirectRequired } from '@/lib/mcp/oauth'
import {
Expand Down Expand Up @@ -142,10 +143,13 @@ export class McpClient {
fetch: oauthFetch,
})
: (oauthFetch ?? guarded.fetch)
// The SDK sends some requests (e.g. `notifications/cancelled` from an abort listener)
// from whatever context fired the callback, which may carry no outbound scope.
const runInOutboundScope = captureOutboundScopeFallback()
this.transport = new StreamableHTTPClientTransport(new URL(this.config.url), {
authProvider: useOauth ? this.authProvider : undefined,
...(useOauth ? {} : { requestInit: { headers: this.config.headers } }),
fetch: transportFetch,
fetch: (input, init) => runInOutboundScope(() => transportFetch(input, init)),
})

this.client = new Client(
Expand Down
Loading