fix(mcp): make tool call timeout configurable

This commit is contained in:
Junyan Qin
2026-07-23 18:15:57 +08:00
parent 7677d1a288
commit 0dfae76e39
13 changed files with 277 additions and 9 deletions
@@ -12,7 +12,7 @@ import pytest
from aiohttp import web
from mcp import types as mcp_types
from langbot.pkg.provider.tools.loaders.mcp import RuntimeMCPSession
from langbot.pkg.provider.tools.loaders.mcp import MCPToolCallTimeoutError, RuntimeMCPSession
class _TransportProbe:
@@ -65,6 +65,25 @@ class _TransportProbe:
},
}
)
if method == 'tools/call':
tool_name = message.get('params', {}).get('name')
if tool_name == 'hang':
return web.Response(status=202)
return web.json_response(
{
'jsonrpc': '2.0',
'id': message['id'],
'result': {
'content': [
{
'type': 'text',
'text': 'healthy',
}
],
'isError': False,
},
}
)
return web.Response(status=202)
return web.Response(status=self.streamable_status)
@@ -126,11 +145,22 @@ async def _transport_server(streamable_status: int | None):
await runner.cleanup()
def _session(url: str, *, timeout: float = 2) -> RuntimeMCPSession:
def _session(
url: str,
*,
timeout: float = 2,
tool_call_timeout_sec: float = 300,
) -> RuntimeMCPSession:
app = cast(Any, SimpleNamespace(logger=Mock()))
return RuntimeMCPSession(
'remote-transport-test',
{'uuid': 'srv-1', 'mode': 'remote', 'url': url, 'timeout': timeout},
{
'uuid': 'srv-1',
'mode': 'remote',
'url': url,
'timeout': timeout,
'tool_call_timeout_sec': tool_call_timeout_sec,
},
True,
app,
)
@@ -164,6 +194,24 @@ async def test_remote_transport_real_streamable_http_success_keeps_session_usabl
await _close_session(session)
@pytest.mark.asyncio
async def test_remote_transport_tool_timeout_does_not_poison_session():
async with _transport_server(200) as (probe, url):
session = _session(url, tool_call_timeout_sec=0.05)
try:
await session._init_remote_server()
with pytest.raises(MCPToolCallTimeoutError, match='timed out after 0.05 seconds'):
await session.invoke_mcp_tool('hang', {})
result = await session.invoke_mcp_tool('health_check', {})
assert result[0].text == 'healthy'
assert probe.streamable_messages.count('tools/call') == 2
finally:
await _close_session(session)
@pytest.mark.asyncio
@pytest.mark.parametrize('status_code', [400, 404, 405])
async def test_remote_transport_real_streamable_http_error_falls_back_to_legacy_sse(status_code: int):