mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-14 22:40:59 +00:00
fix(mcp): make tool call timeout configurable
This commit is contained in:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user