fix(mcp): apply timeouts to prompts and resources (#31612)
This commit is contained in:
parent
954d618790
commit
174ab58343
2 changed files with 52 additions and 8 deletions
|
|
@ -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) {
|
||||||
|
|
|
||||||
|
|
@ -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",
|
||||||
() =>
|
() =>
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue