Skip to content

Commit ae7f8b6

Browse files
waleedlatif1waleedlatif1
committed
fix(mcp): reset connection on any auth-type flip, scope h2 to live transport
1 parent eea4aa1 commit ae7f8b6

7 files changed

Lines changed: 65 additions & 26 deletions

File tree

‎apps/sim/lib/mcp/client.ts‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,6 @@ export class McpClient {
111111
}
112112
)
113113

114-
// Observe out-of-band transport errors the SDK would otherwise drop silently.
115114
this.client.onerror = (error) => {
116115
logger.warn(`MCP transport error for ${this.config.name}`, {
117116
serverId: this.config.id,
@@ -230,8 +229,6 @@ export class McpClient {
230229
}
231230

232231
const configuredTimeout = this.config.timeout
233-
// Idle timeout honors the per-server config but never exceeds the absolute
234-
// discovery ceiling, so tools/list can't hang the UI past that cap.
235232
const idleTimeoutMs = Math.min(
236233
configuredTimeout !== undefined && Number.isFinite(configuredTimeout) && configuredTimeout > 0
237234
? Math.floor(configuredTimeout)

‎apps/sim/lib/mcp/oauth/probe.test.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,10 @@ const { mockCreatePinnedFetch, mockCreateSsrfGuardedMcpFetch, mockPinnedFetch, m
1515
}
1616
})
1717

18+
vi.mock('@/lib/core/security/input-validation.server', () => ({
19+
createPinnedFetch: mockCreatePinnedFetch,
20+
}))
1821
vi.mock('@/lib/mcp/pinned-fetch', () => ({
19-
createPinnedMcpFetch: mockCreatePinnedFetch,
2022
createSsrfGuardedMcpFetch: mockCreateSsrfGuardedMcpFetch,
2123
}))
2224

‎apps/sim/lib/mcp/oauth/probe.ts‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
import { extractWWWAuthenticateParams } from '@modelcontextprotocol/sdk/client/auth.js'
22
import type { FetchLike } from '@modelcontextprotocol/sdk/shared/transport.js'
33
import { createLogger } from '@sim/logger'
4+
import { createPinnedFetch } from '@/lib/core/security/input-validation.server'
45
import { isLoopbackHostname } from '@/lib/core/utils/urls'
5-
import { createPinnedMcpFetch, createSsrfGuardedMcpFetch } from '@/lib/mcp/pinned-fetch'
6+
import { createSsrfGuardedMcpFetch } from '@/lib/mcp/pinned-fetch'
67
import type { McpAuthType } from '@/lib/mcp/types'
78

89
const logger = createLogger('McpOauthProbe')
@@ -33,7 +34,7 @@ export async function detectMcpAuthType(
3334
}
3435

3536
const probeFetch: FetchLike = resolvedIP
36-
? createPinnedMcpFetch(resolvedIP)
37+
? createPinnedFetch(resolvedIP)
3738
: createSsrfGuardedMcpFetch()
3839

3940
const controller = new AbortController()
@@ -67,10 +68,7 @@ export async function detectMcpAuthType(
6768

6869
if (res.status === 401) {
6970
const params = extractWWWAuthenticateParams(res)
70-
// Per RFC 9728, an OAuth-protected resource signals OAuth via
71-
// `resource_metadata=...` in WWW-Authenticate. `scope=...` is also an
72-
// OAuth-specific hint. A bare `error="invalid_token"` is generic Bearer
73-
// and used by plain API-key servers too, so it must not classify as OAuth.
71+
// RFC 9728: resource_metadata / scope signal OAuth; a bare invalid_token is generic Bearer (API-key servers use it too).
7472
if (params.resourceMetadataUrl || params.scope) {
7573
return 'oauth'
7674
}

‎apps/sim/lib/mcp/orchestration/server-lifecycle.test.ts‎

Lines changed: 45 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -88,11 +88,53 @@ describe('MCP server lifecycle orchestration', () => {
8888

8989
expect(result.success).toBe(true)
9090
expect(dbChainMockFns.set).toHaveBeenCalledWith(expect.objectContaining({ authType: 'oauth' }))
91-
// Flipping to OAuth must reset the server to disconnected — it hasn't
92-
// completed an auth flow, so it can't remain 'connected'.
91+
// Flipping to OAuth must reset to disconnected — it hasn't completed an auth flow.
9392
expect(dbChainMockFns.set).toHaveBeenCalledWith(
94-
expect.objectContaining({ connectionStatus: 'disconnected', lastConnected: null })
93+
expect.objectContaining({
94+
connectionStatus: 'disconnected',
95+
lastConnected: null,
96+
lastError: null,
97+
})
9598
)
9699
expect(mockClearCache).toHaveBeenCalledWith('workspace-1')
97100
})
101+
102+
it('resets an OAuth server to disconnected when its auth type flips to headers', async () => {
103+
dbChainMockFns.limit.mockResolvedValueOnce([
104+
{
105+
url: 'https://example.com/mcp',
106+
authType: 'oauth',
107+
oauthClientId: 'client-1',
108+
oauthClientSecret: 'secret-1',
109+
},
110+
])
111+
dbChainMockFns.returning.mockResolvedValueOnce([
112+
{
113+
id: 'server-1',
114+
workspaceId: 'workspace-1',
115+
name: 'Example',
116+
transport: 'streamable-http',
117+
url: 'https://example.com/mcp',
118+
authType: 'headers',
119+
},
120+
])
121+
122+
const result = await performUpdateMcpServer({
123+
workspaceId: 'workspace-1',
124+
userId: 'user-1',
125+
serverId: 'server-1',
126+
authType: 'headers',
127+
})
128+
129+
expect(result.success).toBe(true)
130+
// Flipping away from OAuth must reset too — no stale 'connected'/lastError until re-discovery.
131+
expect(dbChainMockFns.set).toHaveBeenCalledWith(
132+
expect.objectContaining({
133+
authType: 'headers',
134+
connectionStatus: 'disconnected',
135+
lastConnected: null,
136+
lastError: null,
137+
})
138+
)
139+
})
98140
})

‎apps/sim/lib/mcp/orchestration/server-lifecycle.ts‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -352,12 +352,11 @@ export async function performUpdateMcpServer(
352352
const shouldClearOauth = urlChanged || credsChanged
353353
const resolvedAuthType = (updateData.authType ?? currentServer.authType) as McpAuthType
354354
const authTypeChanged = resolvedAuthType !== currentServer.authType
355-
// Switching a server to OAuth (via creds/URL change or a plain authType flip)
356-
// must reset it to disconnected: an OAuth server has not completed its auth
357-
// flow, so it can't legitimately remain 'connected' (matches the create path).
358-
if ((shouldClearOauth || authTypeChanged) && resolvedAuthType === 'oauth') {
355+
// An auth-type flip (either direction) or OAuth creds/URL change invalidates the connection: reset and clear stale state.
356+
if (authTypeChanged || (shouldClearOauth && resolvedAuthType === 'oauth')) {
359357
updateData.connectionStatus = 'disconnected'
360358
updateData.lastConnected = null
359+
updateData.lastError = null
361360
}
362361

363362
if (shouldClearOauth) await revokeMcpOauthTokens(params.serverId)

‎apps/sim/lib/mcp/pinned-fetch.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ describe('createSsrfGuardedMcpFetch', () => {
3131
await fetchLike('https://attacker.example/revoke', { method: 'POST' })
3232

3333
expect(mockValidateMcpServerSsrf).toHaveBeenCalledWith('https://attacker.example/revoke')
34-
expect(mockCreatePinnedFetch).toHaveBeenCalledWith('203.0.113.10', { allowH2: true })
34+
expect(mockCreatePinnedFetch).toHaveBeenCalledWith('203.0.113.10')
3535
expect(sentinelFetch).toHaveBeenCalledWith('https://attacker.example/revoke', {
3636
method: 'POST',
3737
})
@@ -54,7 +54,7 @@ describe('createSsrfGuardedMcpFetch', () => {
5454
await fetchLike(new URL('https://attacker.example/discover'))
5555

5656
expect(mockValidateMcpServerSsrf).toHaveBeenCalledWith('https://attacker.example/discover')
57-
expect(mockCreatePinnedFetch).toHaveBeenCalledWith('203.0.113.10', { allowH2: true })
57+
expect(mockCreatePinnedFetch).toHaveBeenCalledWith('203.0.113.10')
5858
})
5959

6060
it('falls back to global fetch when validation returns no IP', async () => {

‎apps/sim/lib/mcp/pinned-fetch.ts‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,13 @@ import { createPinnedFetch } from '@/lib/core/security/input-validation.server'
33
import { validateMcpServerSsrf } from '@/lib/mcp/domain-check'
44

55
/**
6-
* Pinned fetch for all MCP traffic. MCP servers are commonly deployed behind
7-
* HTTP/2 fronts (CDNs, cloud LBs), and undici's Agent is h1.1-only unless opted
8-
* into h2 via ALPN — so every MCP connection enables it. This is the single
9-
* source of the "MCP implies h2" decision; the base `createPinnedFetch` keeps
10-
* its h1.1 default for non-MCP consumers. Pinning is unaffected: the pinned
11-
* lookup forces the socket to `resolvedIP` regardless of negotiated protocol.
6+
* Pinned fetch for the long-lived MCP transport, which reuses one Agent across
7+
* a connection's requests. MCP servers are commonly behind HTTP/2 fronts (CDNs,
8+
* cloud LBs), and undici's Agent is h1.1-only unless opted into h2 via ALPN, so
9+
* the transport enables it. h2 is *not* used for one-shot flows (OAuth discovery,
10+
* auth-type probe), where a per-request Agent would leave idle h2 sessions with
11+
* no reuse benefit. Pinning is unaffected: the pinned lookup forces the socket to
12+
* `resolvedIP` regardless of negotiated protocol.
1213
*/
1314
export function createPinnedMcpFetch(resolvedIP: string): typeof fetch {
1415
return createPinnedFetch(resolvedIP, { allowH2: true })
@@ -37,7 +38,7 @@ export function createSsrfGuardedMcpFetch(): FetchLike {
3738
return (async (url, init) => {
3839
const target = typeof url === 'string' ? url : url.href
3940
const resolvedIP = await validateMcpServerSsrf(target)
40-
const pinnedFetch: FetchLike = resolvedIP ? createPinnedMcpFetch(resolvedIP) : globalThis.fetch
41+
const pinnedFetch: FetchLike = resolvedIP ? createPinnedFetch(resolvedIP) : globalThis.fetch
4142
return pinnedFetch(url, init)
4243
}) satisfies FetchLike
4344
}

0 commit comments

Comments
 (0)