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
25 changes: 17 additions & 8 deletions apps/sim/app/api/mcp/oauth/callback/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -125,12 +125,19 @@ export const GET = withRouteHandler(async (request: NextRequest) => {
serverId?: string
) => htmlClose(message, ok, reason, serverId, state)

const initialRow = state ? await loadOauthRowByState(state).catch(() => null) : null
const initialRow = state
? await timedStep('loadOauthRowByState', 15_000, () => loadOauthRowByState(state)).catch(
() => null
)
: null
const stateRowServerId = initialRow?.mcpServerId

if (errorParam) {
logger.warn(`MCP OAuth callback received error: ${errorParam}`)
if (initialRow) await clearState(initialRow.id, 'callback:provider_error').catch(() => {})
if (initialRow)
await timedStep('clearState(provider_error)', 10_000, () =>
clearState(initialRow.id, 'callback:provider_error')
).catch(() => {})
return respond(`Authorization failed: ${errorParam}`, false, 'provider_error', stateRowServerId)
}
if (!state || !code) {
Expand All @@ -144,7 +151,7 @@ export const GET = withRouteHandler(async (request: NextRequest) => {

let serverId: string | undefined
try {
const session = await getSession()
const session = await timedStep('getSession', 15_000, () => getSession())
if (!session?.user?.id) {
return respond(
'You must be signed in to complete authorization.',
Expand All @@ -169,11 +176,13 @@ export const GET = withRouteHandler(async (request: NextRequest) => {
)
}

const [server] = await db
.select({ id: mcpServers.id, url: mcpServers.url, workspaceId: mcpServers.workspaceId })
.from(mcpServers)
.where(and(eq(mcpServers.id, row.mcpServerId), isNull(mcpServers.deletedAt)))
.limit(1)
const [server] = await timedStep('loadServer', 15_000, () =>
db
.select({ id: mcpServers.id, url: mcpServers.url, workspaceId: mcpServers.workspaceId })
.from(mcpServers)
.where(and(eq(mcpServers.id, row.mcpServerId), isNull(mcpServers.deletedAt)))
.limit(1)
)
if (!server || !server.url) {
return respond('Server no longer exists.', false, 'server_gone', serverId)
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,6 @@ import {
useCreateMcpServer,
useForceRefreshMcpTools,
useMcpServers,
useMcpToolsEvents,
useStoredMcpTools,
} from '@/hooks/queries/mcp'
import { useWorkflowState, useWorkflows } from '@/hooks/queries/workflows'
Expand Down Expand Up @@ -570,7 +569,6 @@ export const ToolInput = memo(function ToolInput({
const { data: mcpServers = [], isLoading: mcpServersLoading } = useMcpServers(workspaceId)
const { data: storedMcpTools = [] } = useStoredMcpTools(workspaceId)
const forceRefreshMcpTools = useForceRefreshMcpTools().mutate
useMcpToolsEvents(workspaceId)
const { navigateToSettings } = useSettingsNavigation()
const createMcpServer = useCreateMcpServer()
const { startOauthForServer } = useMcpOauthPopup({ workspaceId })
Expand Down
22 changes: 21 additions & 1 deletion apps/sim/hooks/queries/mcp.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -103,13 +103,29 @@ function mockServers(servers: McpServer[]) {
})
}

// jsdom has no EventSource; useMcpToolsQuery mounts the shared SSE subscription.
class FakeEventSource {
onopen: (() => void) | null = null
onerror: (() => void) | null = null
constructor(public url: string) {}
addEventListener(): void {}
close(): void {}
}

describe('useMcpToolsQuery', () => {
beforeEach(() => {
vi.clearAllMocks()
;(globalThis as unknown as { EventSource: unknown }).EventSource = FakeEventSource
})

afterEach(() => {
vi.restoreAllMocks()
// mcp.ts captured these Map/Set instances in module consts at import, so reassigning the
// globalThis property wouldn't reset what the module uses — clear the shared instances.
;(
globalThis as unknown as { __mcp_sse_connections?: Map<string, unknown> }
).__mcp_sse_connections?.clear()
;(globalThis as unknown as { __mcp_sse_subscribed?: Set<string> }).__mcp_sse_subscribed?.clear()
})

it('does not auto-discover disconnected or errored OAuth servers', async () => {
Expand Down Expand Up @@ -139,7 +155,11 @@ describe('useMcpToolsQuery', () => {
const { unmount } = renderHookWithClient(() => useMcpToolsQuery(WORKSPACE_ID))
await flush()

expect(mockRequestJson).toHaveBeenCalledTimes(3)
// Both eligible servers get discovered.
const discoveryCalls = mockRequestJson.mock.calls.filter(
([contract]) => contract === discoverMcpToolsContract
)
expect(discoveryCalls).toHaveLength(2)
expect(mockRequestJson).toHaveBeenCalledWith(
discoverMcpToolsContract,
expect.objectContaining({
Expand Down
43 changes: 38 additions & 5 deletions apps/sim/hooks/queries/mcp.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,13 @@ const logger = createLogger('McpQueries')
export type { McpServerStatusConfig, McpTool, StoredMcpTool }

export const MCP_SERVER_LIST_STALE_TIME = 60 * 1000
export const MCP_SERVER_TOOLS_STALE_TIME = 30 * 1000
/**
* Tool discovery is kept fresh by the `list_changed` → SSE push (see `useMcpToolsEvents`),
* so the query only needs a re-probe-on-visit fallback for servers without push. Matches the
* server-side cache TTL (`MCP_CONSTANTS.CACHE_TIMEOUT`) — no reference MCP client re-probes
* more often than its cache; real changes arrive via push regardless of this value.
*/
export const MCP_SERVER_TOOLS_STALE_TIME = 5 * 60 * 1000
Comment thread
waleedlatif1 marked this conversation as resolved.
Comment thread
waleedlatif1 marked this conversation as resolved.
export const MCP_STORED_TOOL_LIST_STALE_TIME = 60 * 1000
export const MCP_ALLOWED_DOMAINS_STALE_TIME = 5 * 60 * 1000

Expand Down Expand Up @@ -140,6 +146,11 @@ function isServerEligibleForDiscovery(server: McpServer, workspaceId: string): b
export function useMcpToolsQuery(workspaceId: string) {
const queryClient = useQueryClient()
const { data: servers, isLoading: serversLoading } = useMcpServers(workspaceId)
// Push is intrinsic to consuming the tools query: every surface that reads tools (settings,
// tool picker, dynamic args, tool selector, canvas block) gets real-time `list_changed`
// refresh via the shared, reference-counted subscription — so the 5-min stale time is always
// push-backed and no consumer is left re-probing.
useMcpToolsEvents(workspaceId)

/**
* Skip disabled rows, rows retained from a previous workspace, and OAuth rows
Expand Down Expand Up @@ -192,10 +203,10 @@ export function useMcpToolsQuery(workspaceId: string) {
const serverId = serverIds[index]
const status = serverId ? statusById.get(serverId) : undefined
const persistentlyFailed = status === 'error' || status === 'disconnected'
// Keep last-known-good tools through a transient failure (React Query retains `data`, the
// stored status is still healthy) so a populated server doesn't blank — but drop them once
// the stored status crosses its failure threshold, so the workflow editor stops offering a
// dead server's stale tools. Matches how reference MCP clients treat transient vs. closed.
// Keep last-known-good tools while the stored status is still `connected` (React Query
// retains `data` across a failed refetch, so a populated server doesn't blank on a
// transient probe error) — but drop them once the stored status leaves `connected`
// (disconnected/error), so the workflow editor stops offering a dead server's stale tools.
if (result.data && (!result.isError || !persistentlyFailed)) {
tools.push(...result.data)
hasData = true
Expand Down Expand Up @@ -527,6 +538,12 @@ const sseConnections: Map<string, SseEntry> =
((globalThis as Record<string, unknown>)[SSE_KEY] as Map<string, SseEntry>) ??
((globalThis as Record<string, unknown>)[SSE_KEY] = new Map<string, SseEntry>())

/** Per-workspace flag: has this session ever held a live SSE subscription for it? */
const SSE_SUBSCRIBED_KEY = '__mcp_sse_subscribed' as const
const sseEverSubscribed: Set<string> =
((globalThis as Record<string, unknown>)[SSE_SUBSCRIBED_KEY] as Set<string>) ??
((globalThis as Record<string, unknown>)[SSE_SUBSCRIBED_KEY] = new Set<string>())

/** Subscribes to `tools_changed` SSE events and invalidates the affected query keys. */
export function useMcpToolsEvents(workspaceId: string) {
const queryClient = useQueryClient()
Expand Down Expand Up @@ -563,7 +580,23 @@ export function useMcpToolsEvents(workspaceId: string) {
invalidate(serverId)
})

// EventSource fires `onopen` on the initial connect and on every auto-reconnect. Re-sync
// the workspace whenever we could have missed a `tools_changed` event: on any reconnect,
// on the first open of a RE-subscription (leaving the tab tears the connection down), and
// on a first open that only succeeded after an earlier connection error (the initial tools
// query may have failed during that gap and won't retry itself). Skip only a clean first
// subscription — the queries fetch fresh on their own initial mount.
const isResubscribe = sseEverSubscribed.has(workspaceId)
sseEverSubscribed.add(workspaceId)
let opened = false
let erroredBeforeOpen = false
source.onopen = () => {
if (opened || isResubscribe || erroredBeforeOpen) invalidate()
opened = true
}
Comment thread
waleedlatif1 marked this conversation as resolved.

source.onerror = () => {
if (!opened) erroredBeforeOpen = true
logger.warn(`SSE connection error for workspace ${workspaceId}`)
}
Comment thread
waleedlatif1 marked this conversation as resolved.

Expand Down
19 changes: 11 additions & 8 deletions apps/sim/lib/mcp/domain-check.ts
Original file line number Diff line number Diff line change
Expand Up @@ -139,14 +139,17 @@ function isLocalhostHostname(hostname: string): boolean {
* URLs with env var references in the hostname are skipped — they will be
* validated after resolution at execution time.
*
* Returns the IP address to pin subsequent connections to (the resolved IP for
* hostnames, or the literal itself for public IP-literal URLs) so the caller can
* prevent DNS-rebinding TOCTOU attacks and stop redirects from escaping to
* internal hosts. Pinning matters for IP literals too: without it the transport
* uses the default fetch, which follows an attacker-controlled 3xx redirect to a
* private/metadata address. Returns null only when pinning is unnecessary or
* impossible: no URL, allowlist-only mode, env-var hostnames (validated later),
* and localhost on self-hosted (no rebinding risk against a fixed loopback).
* Returns the resolved IP (or the literal itself for IP-literal URLs) as a
* non-null **policy signal**: the SSRF guard is active for this server. A public
* resolution selects the validate-at-connect guarded fetch — DNS-rebinding TOCTOU
* and redirect escapes are prevented by re-validating every socket connect and
* following redirects under per-hop validation (see `createSsrfGuardedMcpFetch` /
* `followRedirectsGuarded`), NOT by pinning to this address. The value is literally
* pinned only for the self-hosted private/loopback carve-out (a policy-permitted
* DNS alias the guarded lookup would otherwise filter). Returns null when the guard
* is unnecessary or impossible: no URL, allowlist-only mode, env-var hostnames
* (validated later), and localhost on self-hosted (no rebinding risk against a
* fixed loopback).
*
* @throws McpSsrfError if the URL resolves to a blocked IP address
*/
Expand Down
Loading