Skip to content

Commit 69f66b6

Browse files
roomoteroomote
committed
MCP: wait for server readiness before requests and retry once on transient -32000 'Connection closed'; add tests
1 parent 3ee6072 commit 69f66b6

File tree

2 files changed

+267
-31
lines changed

2 files changed

+267
-31
lines changed

‎src/services/mcp/McpHub.ts‎

Lines changed: 87 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -903,6 +903,48 @@ export class McpHub {
903903
(conn) => conn.server.name === serverName && (conn.server.source === "global" || !conn.server.source),
904904
)
905905
}
906+
/**
907+
* Wait until a server connection is fully ready (client connected and status updated).
908+
* Guards against race conditions where a tool/resource call arrives before initial connect completes.
909+
*/
910+
private async waitForServerReady(
911+
serverName: string,
912+
source?: "global" | "project",
913+
timeoutMs: number = 10000,
914+
): Promise<ConnectedMcpConnection> {
915+
// Immediately error if there's no connected-type connection to wait on (preserves legacy error for tests/UX)
916+
const initial = this.findConnection(serverName, source)
917+
const notFoundMsg = `No connection found for server: ${serverName}${source ? ` with source ${source}` : ""}`
918+
if (!initial || initial.type !== "connected") {
919+
throw new Error(notFoundMsg)
920+
}
921+
922+
// If we have a connected-type connection but it's still "connecting", wait briefly for readiness.
923+
const start = Date.now()
924+
while (Date.now() - start < timeoutMs) {
925+
const connection = this.findConnection(serverName, source)
926+
if (connection && connection.type === "connected" && connection.server.status === "connected") {
927+
return connection
928+
}
929+
await delay(50)
930+
}
931+
const conn = this.findConnection(serverName, source)
932+
const status = conn?.type === "connected" ? conn.server.status : "disconnected"
933+
throw new Error(`MCP server "${serverName}" is not ready (status: ${status}). Try again in a moment.`)
934+
}
935+
936+
/**
937+
* Detect transient connection-closed errors from MCP transport to allow a single retry.
938+
*/
939+
private isTransientClosedError(error: unknown): boolean {
940+
const code = (error as any)?.code
941+
const msg = error instanceof Error ? error.message : String(error)
942+
if (typeof code === "number" && code === -32000) return true
943+
if (typeof code === "string" && /-?32000/.test(code)) return true
944+
return /connection closed|closed before|socket hang up|ECONNRESET|EPIPE|ERR_STREAM_PREMATURE_CLOSE/i.test(
945+
msg || "",
946+
)
947+
}
906948

907949
private async fetchToolsList(serverName: string, source?: "global" | "project"): Promise<McpTool[]> {
908950
try {
@@ -1559,22 +1601,31 @@ export class McpHub {
15591601
}
15601602

15611603
async readResource(serverName: string, uri: string, source?: "global" | "project"): Promise<McpResourceResponse> {
1562-
const connection = this.findConnection(serverName, source)
1563-
if (!connection || connection.type !== "connected") {
1564-
throw new Error(`No connection found for server: ${serverName}${source ? ` with source ${source}` : ""}`)
1565-
}
1604+
let connection = await this.waitForServerReady(serverName, source)
15661605
if (connection.server.disabled) {
15671606
throw new Error(`Server "${serverName}" is disabled`)
15681607
}
1569-
return await connection.client.request(
1570-
{
1571-
method: "resources/read",
1572-
params: {
1573-
uri,
1608+
1609+
const doRequest = async () =>
1610+
connection.client.request(
1611+
{
1612+
method: "resources/read",
1613+
params: { uri },
15741614
},
1575-
},
1576-
ReadResourceResultSchema,
1577-
)
1615+
ReadResourceResultSchema,
1616+
)
1617+
1618+
try {
1619+
return await doRequest()
1620+
} catch (error) {
1621+
// Handle initial transient "Connection closed" on first invocation after install/start
1622+
if (this.isTransientClosedError(error)) {
1623+
await delay(200)
1624+
connection = await this.waitForServerReady(serverName, source)
1625+
return await doRequest()
1626+
}
1627+
throw error
1628+
}
15781629
}
15791630

15801631
async callTool(
@@ -1583,12 +1634,7 @@ export class McpHub {
15831634
toolArguments?: Record<string, unknown>,
15841635
source?: "global" | "project",
15851636
): Promise<McpToolCallResponse> {
1586-
const connection = this.findConnection(serverName, source)
1587-
if (!connection || connection.type !== "connected") {
1588-
throw new Error(
1589-
`No connection found for server: ${serverName}${source ? ` with source ${source}` : ""}. Please make sure to use MCP servers available under 'Connected MCP Servers'.`,
1590-
)
1591-
}
1637+
let connection = await this.waitForServerReady(serverName, source)
15921638
if (connection.server.disabled) {
15931639
throw new Error(`Server "${serverName}" is disabled and cannot be used`)
15941640
}
@@ -1599,23 +1645,33 @@ export class McpHub {
15991645
timeout = (parsedConfig.timeout ?? 60) * 1000
16001646
} catch (error) {
16011647
console.error("Failed to parse server config for timeout:", error)
1602-
// Default to 60 seconds if parsing fails
16031648
timeout = 60 * 1000
16041649
}
16051650

1606-
return await connection.client.request(
1607-
{
1608-
method: "tools/call",
1609-
params: {
1610-
name: toolName,
1611-
arguments: toolArguments,
1651+
const doRequest = async () =>
1652+
connection.client.request(
1653+
{
1654+
method: "tools/call",
1655+
params: {
1656+
name: toolName,
1657+
arguments: toolArguments,
1658+
},
16121659
},
1613-
},
1614-
CallToolResultSchema,
1615-
{
1616-
timeout,
1617-
},
1618-
)
1660+
CallToolResultSchema,
1661+
{ timeout },
1662+
)
1663+
1664+
try {
1665+
return await doRequest()
1666+
} catch (error) {
1667+
// Handle initial transient "Connection closed" that can happen on the very first call
1668+
if (this.isTransientClosedError(error)) {
1669+
await delay(200)
1670+
connection = await this.waitForServerReady(serverName, source)
1671+
return await doRequest()
1672+
}
1673+
throw error
1674+
}
16191675
}
16201676

16211677
/**

‎src/services/mcp/__tests__/McpHub.retry.spec.ts‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
import type { ClineProvider } from "../../../core/webview/ClineProvider"
2+
import { McpHub } from "../McpHub"
3+
import { vi, describe, it, expect, beforeEach } from "vitest"
4+
import fs from "fs/promises"
5+
6+
// Minimal VSCode and FS mocks to keep constructor side-effects inert
7+
vi.mock("vscode", () => ({
8+
workspace: {
9+
createFileSystemWatcher: vi.fn().mockReturnValue({
10+
onDidChange: vi.fn(),
11+
onDidCreate: vi.fn(),
12+
onDidDelete: vi.fn(),
13+
dispose: vi.fn(),
14+
}),
15+
onDidChangeWorkspaceFolders: vi.fn(),
16+
workspaceFolders: [],
17+
},
18+
window: {
19+
showErrorMessage: vi.fn(),
20+
showInformationMessage: vi.fn(),
21+
showWarningMessage: vi.fn(),
22+
},
23+
Disposable: {
24+
from: vi.fn(),
25+
},
26+
}))
27+
28+
vi.mock("fs/promises", () => ({
29+
default: {
30+
access: vi.fn().mockResolvedValue(undefined),
31+
writeFile: vi.fn().mockResolvedValue(undefined),
32+
readFile: vi.fn().mockResolvedValue("{}"),
33+
unlink: vi.fn().mockResolvedValue(undefined),
34+
rename: vi.fn().mockResolvedValue(undefined),
35+
lstat: vi.fn().mockResolvedValue({ isDirectory: () => true }),
36+
mkdir: vi.fn().mockResolvedValue(undefined),
37+
},
38+
access: vi.fn().mockResolvedValue(undefined),
39+
writeFile: vi.fn().mockResolvedValue(undefined),
40+
readFile: vi.fn().mockResolvedValue("{}"),
41+
unlink: vi.fn().mockResolvedValue(undefined),
42+
rename: vi.fn().mockResolvedValue(undefined),
43+
lstat: vi.fn().mockResolvedValue({ isDirectory: () => true }),
44+
mkdir: vi.fn().mockResolvedValue(undefined),
45+
}))
46+
47+
describe("McpHub transient-retry and readiness", () => {
48+
let mcpHub: McpHub
49+
let mockProvider: Partial<ClineProvider>
50+
51+
beforeEach(() => {
52+
vi.clearAllMocks()
53+
54+
// Prevent constructor from kicking off async initialization that races our injected connections
55+
vi.spyOn(McpHub.prototype as any, "initializeGlobalMcpServers").mockResolvedValue(undefined as any)
56+
vi.spyOn(McpHub.prototype as any, "initializeProjectMcpServers").mockResolvedValue(undefined as any)
57+
58+
mockProvider = {
59+
ensureSettingsDirectoryExists: vi.fn().mockResolvedValue("/mock/settings/path"),
60+
ensureMcpServersDirectoryExists: vi.fn().mockResolvedValue("/mock/settings/path"),
61+
postMessageToWebview: vi.fn().mockResolvedValue(undefined),
62+
getState: vi.fn().mockResolvedValue({ mcpEnabled: true }),
63+
context: {
64+
subscriptions: [],
65+
workspaceState: {} as any,
66+
globalState: {} as any,
67+
secrets: {} as any,
68+
extensionUri: { fsPath: "/test/path" } as any,
69+
extensionPath: "/test/path",
70+
asAbsolutePath: (p: string) => p,
71+
globalStorageUri: { fsPath: "/test/global" } as any,
72+
globalStoragePath: "/test/global",
73+
extension: {
74+
packageJSON: { version: "1.0.0" },
75+
} as any,
76+
} as any,
77+
}
78+
79+
// Keep global config empty so constructor does not attempt to connect anything
80+
vi.mocked(fs.readFile).mockResolvedValue("{}")
81+
82+
mcpHub = new McpHub(mockProvider as ClineProvider)
83+
})
84+
85+
it("retries once on transient -32000 'Connection closed' for tools/call", async () => {
86+
// Prepare a connected connection
87+
const firstError: any = new Error("Connection closed")
88+
firstError.code = -32000
89+
90+
const request = vi
91+
.fn()
92+
// First call: throw transient
93+
.mockRejectedValueOnce(firstError)
94+
// Second call: succeed
95+
.mockResolvedValueOnce({ content: [{ type: "text", text: "ok" }] })
96+
97+
// Inject a connected connection directly
98+
;(mcpHub as any).connections = [
99+
{
100+
type: "connected",
101+
server: {
102+
name: "retry-server",
103+
config: JSON.stringify({ type: "stdio", command: "node", timeout: 60 }),
104+
status: "connected",
105+
source: "global",
106+
errorHistory: [],
107+
},
108+
client: { request } as any,
109+
transport: {} as any,
110+
},
111+
]
112+
113+
const result = await mcpHub.callTool("retry-server", "do_something", { a: 1 })
114+
expect(result).toBeTruthy()
115+
expect(request).toHaveBeenCalledTimes(2)
116+
})
117+
118+
it("retries once on transient -32000 'Connection closed' for resources/read", async () => {
119+
const firstError: any = new Error("Connection closed before response")
120+
firstError.code = -32000
121+
122+
const request = vi
123+
.fn()
124+
.mockRejectedValueOnce(firstError)
125+
.mockResolvedValueOnce({ contents: [{ type: "text", text: "resource" }] })
126+
127+
;(mcpHub as any).connections = [
128+
{
129+
type: "connected",
130+
server: {
131+
name: "retry-server",
132+
config: JSON.stringify({ type: "stdio", command: "node", timeout: 60 }),
133+
status: "connected",
134+
source: "global",
135+
errorHistory: [],
136+
},
137+
client: { request } as any,
138+
transport: {} as any,
139+
},
140+
]
141+
142+
const result = await mcpHub.readResource("retry-server", "file:///tmp/x.txt")
143+
expect(result).toBeTruthy()
144+
expect(request).toHaveBeenCalledTimes(2)
145+
})
146+
147+
it("waits for server to become connected before first tool call", async () => {
148+
// Start 'connecting' and then flip to 'connected' shortly after
149+
const request = vi.fn().mockResolvedValue({ content: [{ type: "text", text: "ok" }] })
150+
151+
const connection: any = {
152+
type: "connected",
153+
server: {
154+
name: "slow-server",
155+
config: JSON.stringify({ type: "stdio", command: "node", timeout: 60 }),
156+
status: "connecting",
157+
source: "global",
158+
errorHistory: [],
159+
},
160+
client: { request },
161+
transport: {},
162+
}
163+
164+
;(mcpHub as any).connections = [connection]
165+
166+
// Flip status to 'connected' after 100ms
167+
setTimeout(() => {
168+
connection.server.status = "connected"
169+
}, 100)
170+
171+
const start = Date.now()
172+
const result = await mcpHub.callTool("slow-server", "ping", {})
173+
const elapsed = Date.now() - start
174+
175+
expect(result).toBeTruthy()
176+
expect(request).toHaveBeenCalledTimes(1)
177+
// Should have waited at least ~100ms before making the request
178+
expect(elapsed).toBeGreaterThanOrEqual(90)
179+
})
180+
})

0 commit comments

Comments
 (0)