Merge origin/staging into transport-response-cap

Author: waleedlatif1

apps/sim/app/api/mcp/oauth/callback/route.ts

@@ -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) {
@@ -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.',
@@ -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)
     }

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/tool-input/tool-input.tsx

@@ -79,7 +79,6 @@ import {
   useCreateMcpServer,
   useForceRefreshMcpTools,
   useMcpServers,
-  useMcpToolsEvents,
   useStoredMcpTools,
 } from '@/hooks/queries/mcp'
 import { useWorkflowState, useWorkflows } from '@/hooks/queries/workflows'
@@ -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 })

apps/sim/hooks/queries/mcp.test.tsx

@@ -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 () => {
@@ -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({

apps/sim/hooks/queries/mcp.ts

@@ -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
 export const MCP_STORED_TOOL_LIST_STALE_TIME = 60 * 1000
 export const MCP_ALLOWED_DOMAINS_STALE_TIME = 5 * 60 * 1000
 
@@ -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
@@ -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
@@ -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()
@@ -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
+      }
+
       source.onerror = () => {
+        if (!opened) erroredBeforeOpen = true
         logger.warn(`SSE connection error for workspace ${workspaceId}`)
       }
 

apps/sim/lib/mcp/domain-check.ts

@@ -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
  */