fix(mcp): apply timeouts to prompts and resources (#31612)

This commit is contained in:
Aiden Cline 2026-06-09 23:58:40 -05:00 committed by GitHub
commit 174ab58343
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 52 additions and 8 deletions

View file

@ -740,7 +740,7 @@ export const layer = Layer.effect(
const withClient = Effect.fnUntraced(function* <A>( const withClient = Effect.fnUntraced(function* <A>(
clientName: string, clientName: string,
fn: (client: MCPClient) => Promise<A>, fn: (client: MCPClient, timeout?: number) => Promise<A>,
label: string, label: string,
meta?: Record<string, unknown>, meta?: Record<string, unknown>,
) { ) {
@ -750,8 +750,11 @@ export const layer = Layer.effect(
yield* Effect.logWarning(`client not found for ${label}`, { clientName }) yield* Effect.logWarning(`client not found for ${label}`, { clientName })
return undefined return undefined
} }
const cfg = yield* cfgSvc.get()
const configured = cfg.mcp?.[clientName]
const staticTimeout = configured && isMcpConfigured(configured) ? configured.timeout : undefined
return yield* Effect.tryPromise({ return yield* Effect.tryPromise({
try: () => fn(client), try: () => fn(client, s.config[clientName]?.timeout ?? staticTimeout ?? cfg.experimental?.mcp_timeout),
catch: (error) => error, catch: (error) => error,
}).pipe( }).pipe(
Effect.tapError((error) => Effect.tapError((error) =>
@ -770,15 +773,21 @@ export const layer = Layer.effect(
name: string, name: string,
args?: Record<string, string>, args?: Record<string, string>,
) { ) {
return yield* withClient(clientName, (client) => client.getPrompt({ name, arguments: args }), "getPrompt", { return yield* withClient(
promptName: name, clientName,
}) (client, timeout) => client.getPrompt({ name, arguments: args }, { timeout }),
"getPrompt",
{ promptName: name },
)
}) })
const readResource = Effect.fn("MCP.readResource")(function* (clientName: string, resourceUri: string) { const readResource = Effect.fn("MCP.readResource")(function* (clientName: string, resourceUri: string) {
return yield* withClient(clientName, (client) => client.readResource({ uri: resourceUri }), "readResource", { return yield* withClient(
resourceUri, clientName,
}) (client, timeout) => client.readResource({ uri: resourceUri }, { timeout }),
"readResource",
{ resourceUri },
)
}) })
const getMcpConfig = Effect.fnUntraced(function* (mcpName: string) { const getMcpConfig = Effect.fnUntraced(function* (mcpName: string) {

View file

@ -13,6 +13,8 @@ interface MockClientState {
listToolsCalls: number listToolsCalls: number
listPromptsCalls: number listPromptsCalls: number
listResourcesCalls: number listResourcesCalls: number
getPromptTimeout?: number
readResourceTimeout?: number
requestCalls: number requestCalls: number
listToolsShouldFail: boolean listToolsShouldFail: boolean
listToolsError: string listToolsError: string
@ -206,6 +208,16 @@ void mock.module("@modelcontextprotocol/sdk/client/index.js", () => ({
return { resources: this._state?.resources ?? [] } return { resources: this._state?.resources ?? [] }
} }
async getPrompt(_params: unknown, options?: { timeout?: number }) {
if (this._state) this._state.getPromptTimeout = options?.timeout
return { messages: [] }
}
async readResource(params: { uri: string }, options?: { timeout?: number }) {
if (this._state) this._state.readResourceTimeout = options?.timeout
return { contents: [{ uri: params.uri, text: "test" }] }
}
async close() { async close() {
if (this._state) this._state.closed = true if (this._state) this._state.closed = true
} }
@ -758,6 +770,29 @@ it.instance(
}, },
) )
it.instance(
"uses per-server timeouts for prompt and resource requests",
() =>
MCP.Service.use((mcp: MCPNS.Interface) =>
Effect.gen(function* () {
lastCreatedClientName = "timeout-server"
const serverState = getOrCreateClientState("timeout-server")
yield* mcp.add("timeout-server", {
type: "local",
command: ["echo", "test"],
timeout: 2500,
})
yield* mcp.getPrompt("timeout-server", "test")
yield* mcp.readResource("timeout-server", "test://resource")
expect(serverState.getPromptTimeout).toBe(2500)
expect(serverState.readResourceTimeout).toBe(2500)
}),
),
{ config: { mcp: {}, experimental: { mcp_timeout: 5000 } } },
)
it.instance( it.instance(
"resource-only servers connect without listing tools", "resource-only servers connect without listing tools",
() => () =>