Skip to content

Commit c9baae7

Browse files
authored
fix(mcp): recover cleanly from OAuth failures (#5595)
* fix(mcp): improve OAuth failure recovery * fix(mcp): address OAuth recovery review findings * fix(mcp): prefer static bearer auth in connection tests * fix(mcp): preserve discovery failure state Treat static-header 401s as credential failures, keep OAuth failures pending, and prevent failed refreshes or reflected upstream errors from masquerading as connected state.
1 parent 7e975e7 commit c9baae7

16 files changed

Lines changed: 1448 additions & 195 deletions

File tree

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
/**
2+
* @vitest-environment node
3+
*/
4+
import type { NextRequest } from 'next/server'
5+
import { beforeEach, describe, expect, it, vi } from 'vitest'
6+
7+
const { mockClearCache, mockDiscoverServerTools, mockSelect, mockUpdateSet } = vi.hoisted(() => ({
8+
mockClearCache: vi.fn(),
9+
mockDiscoverServerTools: vi.fn(),
10+
mockSelect: vi.fn(),
11+
mockUpdateSet: vi.fn(),
12+
}))
13+
14+
vi.mock('@sim/db', () => ({
15+
db: {
16+
select: mockSelect,
17+
update: vi.fn().mockReturnValue({ set: mockUpdateSet }),
18+
},
19+
}))
20+
21+
vi.mock('@/lib/core/utils/with-route-handler', () => ({
22+
withRouteHandler: (handler: unknown) => handler,
23+
}))
24+
25+
vi.mock('@/lib/mcp/middleware', () => ({
26+
withMcpAuth:
27+
() =>
28+
(
29+
handler: (
30+
request: NextRequest,
31+
context: { userId: string; workspaceId: string; requestId: string },
32+
routeContext: { params: Promise<{ id: string }> }
33+
) => Promise<Response>
34+
) =>
35+
(request: NextRequest, routeContext: { params: Promise<{ id: string }> }) =>
36+
handler(
37+
request,
38+
{ userId: 'user-1', workspaceId: 'workspace-1', requestId: 'request-1' },
39+
routeContext
40+
),
41+
}))
42+
43+
vi.mock('@/lib/mcp/service', () => ({
44+
mcpService: {
45+
clearCache: mockClearCache,
46+
discoverServerTools: mockDiscoverServerTools,
47+
},
48+
}))
49+
50+
import { POST } from '@/app/api/mcp/servers/[id]/refresh/route'
51+
52+
const initialServer = {
53+
id: 'server-1',
54+
workspaceId: 'workspace-1',
55+
name: 'OAuth Server',
56+
url: 'https://example.com/mcp',
57+
connectionStatus: 'connected',
58+
lastError: null,
59+
lastConnected: new Date('2026-01-01T00:00:00.000Z'),
60+
toolCount: 4,
61+
statusConfig: { consecutiveFailures: 0, lastSuccessfulDiscovery: null },
62+
}
63+
64+
const persistedServer = {
65+
...initialServer,
66+
connectionStatus: 'disconnected',
67+
lastError: null,
68+
toolCount: 0,
69+
}
70+
71+
function selectRows(rows: unknown[]) {
72+
return {
73+
from: vi.fn().mockReturnValue({
74+
where: vi.fn().mockReturnValue({
75+
limit: vi.fn().mockResolvedValue(rows),
76+
}),
77+
}),
78+
}
79+
}
80+
81+
describe('MCP server refresh route', () => {
82+
beforeEach(() => {
83+
vi.clearAllMocks()
84+
mockSelect.mockReturnValueOnce(selectRows([initialServer]))
85+
mockUpdateSet.mockReturnValue({
86+
where: vi.fn().mockReturnValue({ returning: vi.fn().mockResolvedValue([persistedServer]) }),
87+
})
88+
})
89+
90+
it('preserves the service-persisted OAuth pending status', async () => {
91+
mockDiscoverServerTools.mockRejectedValueOnce(new Error('OAuth authorization required'))
92+
93+
const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', {
94+
method: 'POST',
95+
}) as NextRequest
96+
const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) })
97+
const body = await response.json()
98+
99+
expect(body.data).toEqual(
100+
expect.objectContaining({
101+
status: 'disconnected',
102+
error: null,
103+
})
104+
)
105+
expect(mockUpdateSet).not.toHaveBeenCalledWith(
106+
expect.objectContaining({ connectionStatus: expect.anything() })
107+
)
108+
})
109+
110+
it('reports the discovery failure when status persistence leaves a stale connected row', async () => {
111+
const reflectedSecret = 'Bearer reflected-static-token'
112+
mockDiscoverServerTools.mockRejectedValueOnce(
113+
new Error(`Upstream reflected ${reflectedSecret}`)
114+
)
115+
mockUpdateSet.mockReturnValueOnce({
116+
where: vi.fn().mockReturnValue({
117+
returning: vi.fn().mockResolvedValue([initialServer]),
118+
}),
119+
})
120+
121+
const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', {
122+
method: 'POST',
123+
}) as NextRequest
124+
const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) })
125+
const body = await response.json()
126+
127+
expect(body.data).toEqual(
128+
expect.objectContaining({
129+
status: 'disconnected',
130+
error: 'Internal server error',
131+
workflowsUpdated: 0,
132+
})
133+
)
134+
expect(JSON.stringify(body)).not.toContain(reflectedSecret)
135+
expect(mockClearCache).not.toHaveBeenCalled()
136+
})
137+
138+
it('preserves a connected status from a newer successful discovery', async () => {
139+
mockDiscoverServerTools.mockRejectedValueOnce(new Error('Connection failed'))
140+
const newerSuccessfulServer = {
141+
...initialServer,
142+
lastConnected: new Date(Date.now() + 60_000),
143+
toolCount: 7,
144+
}
145+
mockUpdateSet.mockReturnValueOnce({
146+
where: vi.fn().mockReturnValue({
147+
returning: vi.fn().mockResolvedValue([newerSuccessfulServer]),
148+
}),
149+
})
150+
151+
const request = new Request('http://localhost/api/mcp/servers/server-1/refresh', {
152+
method: 'POST',
153+
}) as NextRequest
154+
const response = await POST(request, { params: Promise.resolve({ id: 'server-1' }) })
155+
const body = await response.json()
156+
157+
expect(body.data).toEqual(
158+
expect.objectContaining({
159+
status: 'connected',
160+
error: null,
161+
toolCount: 7,
162+
workflowsUpdated: 0,
163+
})
164+
)
165+
expect(mockClearCache).toHaveBeenCalledWith('workspace-1')
166+
})
167+
})

apps/sim/app/api/mcp/servers/[id]/refresh/route.ts

Lines changed: 43 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,17 @@ import { db } from '@sim/db'
22
import { mcpServers, workflow, workflowBlocks } from '@sim/db/schema'
33
import { createLogger } from '@sim/logger'
44
import { toError } from '@sim/utils/errors'
5+
import { truncate } from '@sim/utils/string'
56
import { and, eq, inArray, isNull } from 'drizzle-orm'
67
import type { NextRequest } from 'next/server'
78
import { mcpServerIdParamsSchema } from '@/lib/api/contracts/mcp'
89
import { validationErrorResponse } from '@/lib/api/server'
910
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
1011
import { withMcpAuth } from '@/lib/mcp/middleware'
1112
import { mcpService } from '@/lib/mcp/service'
12-
import type { McpServerStatusConfig, McpTool, McpToolSchema } from '@/lib/mcp/types'
13+
import type { McpTool, McpToolSchema } from '@/lib/mcp/types'
1314
import {
15+
categorizeError,
1416
createMcpErrorResponse,
1517
createMcpSuccessResponse,
1618
MCP_TOOL_CORE_PARAMS,
@@ -184,17 +186,10 @@ export const POST = withRouteHandler(
184186
)
185187
}
186188

187-
let connectionStatus: 'connected' | 'disconnected' | 'error' = 'error'
188-
let toolCount = 0
189-
let lastError: string | null = null
190189
let syncResult: SyncResult = { updatedCount: 0, updatedWorkflowIds: [] }
191190
let discoveredTools: McpTool[] = []
192-
193-
const currentStatusConfig: McpServerStatusConfig =
194-
(server.statusConfig as McpServerStatusConfig | null) ?? {
195-
consecutiveFailures: 0,
196-
lastSuccessfulDiscovery: null,
197-
}
191+
let discoveryError: string | null = null
192+
const discoveryStartedAt = new Date()
198193

199194
try {
200195
discoveredTools = await mcpService.discoverServerTools(
@@ -203,48 +198,62 @@ export const POST = withRouteHandler(
203198
workspaceId,
204199
true
205200
)
206-
connectionStatus = 'connected'
207-
toolCount = discoveredTools.length
208-
logger.info(`[${requestId}] Discovered ${toolCount} tools from server ${serverId}`)
201+
logger.info(
202+
`[${requestId}] Discovered ${discoveredTools.length} tools from server ${serverId}`
203+
)
204+
} catch (error) {
205+
discoveryError = truncate(categorizeError(error).message, 200, '')
206+
logger.warn(`[${requestId}] Failed to connect to server ${serverId}`, {
207+
error: discoveryError,
208+
})
209+
}
209210

211+
if (discoveryError === null) {
210212
syncResult = await syncToolSchemasToWorkflows(
211213
workspaceId,
212214
serverId,
213215
discoveredTools,
214216
requestId,
215217
{ url: server.url ?? undefined, name: server.name ?? undefined }
216218
)
217-
} catch (error) {
218-
connectionStatus = 'error'
219-
lastError =
220-
error instanceof Error
221-
? error.message.split('\n')[0].slice(0, 200)
222-
: 'Connection failed'
223-
logger.warn(`[${requestId}] Failed to connect to server ${serverId}:`, error)
224219
}
225220

226221
const now = new Date()
227-
const newStatusConfig =
228-
connectionStatus === 'connected'
229-
? { consecutiveFailures: 0, lastSuccessfulDiscovery: now.toISOString() }
230-
: {
231-
consecutiveFailures: currentStatusConfig.consecutiveFailures + 1,
232-
lastSuccessfulDiscovery: currentStatusConfig.lastSuccessfulDiscovery,
233-
}
234222

235223
const [refreshedServer] = await db
236224
.update(mcpServers)
237225
.set({
238226
lastToolsRefresh: now,
239-
connectionStatus,
240-
lastError,
241-
lastConnected: connectionStatus === 'connected' ? now : server.lastConnected,
242-
toolCount,
243-
statusConfig: newStatusConfig,
244227
updatedAt: now,
245228
})
246-
.where(eq(mcpServers.id, serverId))
247-
.returning()
229+
.where(
230+
and(
231+
eq(mcpServers.id, serverId),
232+
eq(mcpServers.workspaceId, workspaceId),
233+
isNull(mcpServers.deletedAt)
234+
)
235+
)
236+
.returning({
237+
connectionStatus: mcpServers.connectionStatus,
238+
lastConnected: mcpServers.lastConnected,
239+
lastError: mcpServers.lastError,
240+
toolCount: mcpServers.toolCount,
241+
})
242+
243+
let connectionStatus = refreshedServer?.connectionStatus ?? 'error'
244+
let lastError = refreshedServer ? refreshedServer.lastError : discoveryError
245+
const toolCount = refreshedServer?.toolCount ?? discoveredTools.length
246+
247+
if (discoveryError !== null && connectionStatus === 'connected') {
248+
const newerSuccessWonRace =
249+
refreshedServer?.lastConnected != null &&
250+
refreshedServer.lastConnected > discoveryStartedAt
251+
252+
if (!newerSuccessWonRace) {
253+
connectionStatus = 'disconnected'
254+
lastError = discoveryError
255+
}
256+
}
248257

249258
if (connectionStatus === 'connected') {
250259
await mcpService.clearCache(workspaceId)

0 commit comments

Comments
 (0)