xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-07-24 21:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mcp): make tool call timeout configurable

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit_tests/provider/test_mcp_remote_transport.py
+51 -3
View File
@@ -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):
tests/unit_tests/provider/test_mcp_resources.py
+109
View File
@@ -2,20 +2,24 @@ from __future__ import annotations
import asyncio
import base64
from datetime import timedelta
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import httpx
import pytest
from mcp import types as mcp_types
from mcp.shared.exceptions import McpError
from langbot.pkg.provider.tools.loaders.mcp import (
MCP_RESOURCE_CONTEXT_QUERY_KEY,
MCP_RESOURCE_TRACE_QUERY_KEY,
MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS,
MCP_TOOL_LIST_RESOURCES,
MCP_TOOL_READ_RESOURCE,
MCPLoader,
MCPSessionStatus,
MCPToolCallTimeoutError,
RuntimeMCPSession,
)
from langbot.pkg.telemetry import features as telemetry_features
@@ -62,6 +66,111 @@ def _http_status_error(status_code: int) -> httpx.HTTPStatusError:
return httpx.HTTPStatusError(f'HTTP {status_code}', request=request, response=response)
def _tool_result(text: str = 'ok') -> mcp_types.CallToolResult:
return mcp_types.CallToolResult(
content=[mcp_types.TextContent(type='text', text=text)],
isError=False,
)
@pytest.mark.asyncio
async def test_invoke_mcp_tool_uses_configurable_request_timeout():
session = RuntimeMCPSession(
'slow-tools',
{
'uuid': 'srv-1',
'mode': 'remote',
'tool_call_timeout_sec': 900,
},
True,
_app(),
)
session.session = SimpleNamespace(call_tool=AsyncMock(return_value=_tool_result()))
result = await session.invoke_mcp_tool('render_video', {'quality': 'high'})
assert result[0].text == 'ok'
session.session.call_tool.assert_awaited_once_with(
'render_video',
{'quality': 'high'},
read_timeout_seconds=timedelta(seconds=900),
)
@pytest.mark.asyncio
async def test_invoke_mcp_tool_zero_timeout_disables_request_deadline():
session = RuntimeMCPSession(
'unbounded-tools',
{
'uuid': 'srv-1',
'mode': 'remote',
'tool_call_timeout_sec': 0,
},
True,
_app(),
)
session.session = SimpleNamespace(call_tool=AsyncMock(return_value=_tool_result()))
await session.invoke_mcp_tool('long_job', {})
session.session.call_tool.assert_awaited_once_with(
'long_job',
{},
read_timeout_seconds=None,
)
@pytest.mark.asyncio
async def test_invoke_mcp_tool_timeout_is_not_retried_and_session_remains_usable():
session = RuntimeMCPSession(
'recoverable-tools',
{
'uuid': 'srv-1',
'mode': 'remote',
'tool_call_timeout_sec': 5,
},
True,
_app(),
)
timeout = McpError(
mcp_types.ErrorData(
code=httpx.codes.REQUEST_TIMEOUT,
message='Timed out while waiting for response to ClientRequest. Waited 5 seconds.',
)
)
call_tool = AsyncMock(side_effect=[timeout, _tool_result('recovered')])
session.session = SimpleNamespace(call_tool=call_tool)
with pytest.raises(
MCPToolCallTimeoutError,
match="MCP tool 'long_job' on server 'recoverable-tools' timed out after 5 seconds",
):
await session.invoke_mcp_tool('long_job', {})
assert call_tool.await_count == 1
second_result = await session.invoke_mcp_tool('health_check', {})
assert second_result[0].text == 'recovered'
assert call_tool.await_count == 2
@pytest.mark.parametrize('invalid_timeout', [-1, float('inf'), 'not-a-number'])
def test_invalid_tool_call_timeout_falls_back_to_default(invalid_timeout):
ap = _app()
session = RuntimeMCPSession(
'invalid-timeout',
{
'uuid': 'srv-1',
'mode': 'remote',
'tool_call_timeout_sec': invalid_timeout,
},
True,
ap,
)
assert session.tool_call_timeout_sec == MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS
ap.logger.warning.assert_called_once()
@pytest.mark.asyncio
async def test_remote_transport_falls_back_to_sse_for_compatible_http_status_in_exception_group():
session = RuntimeMCPSession(