diff --git a/src/langbot/pkg/provider/tools/loaders/mcp.py b/src/langbot/pkg/provider/tools/loaders/mcp.py index 87f04b997..cdb6a0a2e 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp.py @@ -3,10 +3,12 @@ from __future__ import annotations import base64 import enum import json +import math import re import time import typing from contextlib import AsyncExitStack, asynccontextmanager +from datetime import timedelta import traceback from langbot_plugin.api.entities.events import pipeline_query import sqlalchemy @@ -51,6 +53,7 @@ MCP_RESOURCE_CONTEXT_MAX_BYTES = 96 * 1024 MCP_RESOURCE_TRACE_QUERY_KEY = '_mcp_resource_reads' MCP_RESOURCE_LINKS_QUERY_KEY = '_mcp_resource_links' MCP_RESOURCE_CONTEXT_QUERY_KEY = '_mcp_resource_context' +MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS = 300.0 TEXT_LIKE_MIME_TYPES = { 'application/json', @@ -213,6 +216,10 @@ class _CallerReconnect(Exception): """ +class MCPToolCallTimeoutError(TimeoutError): + """An MCP tool call exceeded its configured per-server deadline.""" + + class RuntimeMCPSession: """运行时 MCP 会话""" @@ -262,6 +269,9 @@ class RuntimeMCPSession: self.ap = ap self.enable = enable self.session = None + self.tool_call_timeout_sec = self._parse_tool_call_timeout( + server_config.get('tool_call_timeout_sec', MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS) + ) # Transient test sessions (created from the config page "test" button, # which carry no persisted server UUID) must NOT share the live @@ -302,6 +312,21 @@ class RuntimeMCPSession: self._box_stdio_runtime = BoxStdioSessionRuntime(self) self.box_config = self._box_stdio_runtime.config + def _parse_tool_call_timeout(self, value: typing.Any) -> float: + """Return a safe tool-call timeout; zero explicitly disables it.""" + try: + timeout = float(value) + except (TypeError, ValueError): + timeout = -1 + + if not math.isfinite(timeout) or timeout < 0: + self.ap.logger.warning( + f'Invalid MCP tool call timeout {value!r} for {self.server_name}; ' + f'using {MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS:g} seconds' + ) + return MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS + return timeout + async def _init_stdio_python_server(self): if self._uses_box_stdio(): await self._box_stdio_runtime.initialize() @@ -968,13 +993,22 @@ class RuntimeMCPSession: raise Exception('MCP session is not connected') try: - async with asyncio.timeout(30): - result = await self.session.call_tool(tool_name, arguments) - except TimeoutError as e: - raise Exception( - f"MCP tool '{tool_name}' on server '{self.server_name}' timed out after 30 seconds" - ) from e + read_timeout = timedelta(seconds=self.tool_call_timeout_sec) if self.tool_call_timeout_sec > 0 else None + result = await self.session.call_tool( + tool_name, + arguments, + read_timeout_seconds=read_timeout, + ) except Exception as e: + if self._is_tool_call_timeout(e): + self.ap.logger.warning( + f'MCP tool {tool_name} on {self.server_name} timed out after ' + f'{self.tool_call_timeout_sec:g} seconds' + ) + raise MCPToolCallTimeoutError( + f"MCP tool '{tool_name}' on server '{self.server_name}' timed out after " + f'{self.tool_call_timeout_sec:g} seconds' + ) from e if attempt == 0 and self._is_session_terminated(e): self.ap.logger.warning( f'MCP tool {tool_name} on {self.server_name} got session terminated, triggering reconnect...' @@ -997,6 +1031,14 @@ class RuntimeMCPSession: raise Exception('MCP session is not connected') + @staticmethod + def _is_tool_call_timeout(exc: BaseException) -> bool: + """Recognize the MCP SDK's per-request timeout without retrying it.""" + return any( + isinstance(leaf, McpError) and leaf.error.code == httpx.codes.REQUEST_TIMEOUT + for leaf in RuntimeMCPSession._iter_exception_leaves(exc) + ) + def get_tools(self) -> list[resource_tool.LLMTool]: return self.functions diff --git a/tests/unit_tests/provider/test_mcp_remote_transport.py b/tests/unit_tests/provider/test_mcp_remote_transport.py index b2f1d2e14..b7761ec8a 100644 --- a/tests/unit_tests/provider/test_mcp_remote_transport.py +++ b/tests/unit_tests/provider/test_mcp_remote_transport.py @@ -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): diff --git a/tests/unit_tests/provider/test_mcp_resources.py b/tests/unit_tests/provider/test_mcp_resources.py index 05f74b0f0..a0ae37a16 100644 --- a/tests/unit_tests/provider/test_mcp_resources.py +++ b/tests/unit_tests/provider/test_mcp_resources.py @@ -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( diff --git a/web/src/app/home/mcp/components/mcp-form/MCPForm.tsx b/web/src/app/home/mcp/components/mcp-form/MCPForm.tsx index 753e7a73c..058946395 100644 --- a/web/src/app/home/mcp/components/mcp-form/MCPForm.tsx +++ b/web/src/app/home/mcp/components/mcp-form/MCPForm.tsx @@ -435,6 +435,10 @@ const getFormSchema = (t: TFunction) => .number({ invalid_type_error: t('mcp.timeoutMustBeNumber') }) .positive({ message: t('mcp.timeoutMustBePositive') }) .default(30), + tool_call_timeout_sec: z + .number({ invalid_type_error: t('mcp.timeoutMustBeNumber') }) + .nonnegative({ message: t('mcp.timeoutNonNegative') }) + .default(300), ssereadtimeout: z .number({ invalid_type_error: t('mcp.sseTimeoutMustBeNumber') }) .positive({ message: t('mcp.timeoutMustBePositive') }) @@ -474,6 +478,7 @@ const getFormSchema = (t: TFunction) => type FormValues = z.infer> & { timeout: number; + tool_call_timeout_sec: number; ssereadtimeout: number; }; @@ -535,6 +540,7 @@ const MCPForm = forwardRef(function MCPForm( command: '', args: [], timeout: 30, + tool_call_timeout_sec: 300, ssereadtimeout: 300, extra_args: [], ...initialDraftRef.current, @@ -609,6 +615,7 @@ const MCPForm = forwardRef(function MCPForm( command: '', args: [], timeout: 30, + tool_call_timeout_sec: 300, ssereadtimeout: 300, extra_args: [], ...initialDraftRef.current, @@ -687,10 +694,16 @@ const MCPForm = forwardRef(function MCPForm( command: '', args: [], timeout: 30, + tool_call_timeout_sec: 300, ssereadtimeout: 300, extra_args: [], }; + if (typeof server.extra_args.tool_call_timeout_sec === 'number') { + formValues.tool_call_timeout_sec = + server.extra_args.tool_call_timeout_sec; + } + let newExtraArgs: { key: string; type: 'string' | 'number' | 'boolean'; @@ -770,6 +783,7 @@ const MCPForm = forwardRef(function MCPForm( url: value.url!, headers, timeout: value.timeout, + tool_call_timeout_sec: value.tool_call_timeout_sec, }, }; } else { @@ -786,6 +800,7 @@ const MCPForm = forwardRef(function MCPForm( command: value.command!, args: value.args?.map((arg) => arg.value) || [], env, + tool_call_timeout_sec: value.tool_call_timeout_sec, }, }; } @@ -835,6 +850,7 @@ const MCPForm = forwardRef(function MCPForm( extraArgsData = { url: form.getValues('url')!, timeout: form.getValues('timeout'), + tool_call_timeout_sec: form.getValues('tool_call_timeout_sec'), headers: Object.fromEntries( formExtraArgs.map((arg) => [arg.key, arg.value]), ), @@ -846,6 +862,7 @@ const MCPForm = forwardRef(function MCPForm( env: Object.fromEntries( formExtraArgs.map((arg) => [arg.key, arg.value]), ), + tool_call_timeout_sec: form.getValues('tool_call_timeout_sec'), }; } @@ -1047,6 +1064,30 @@ const MCPForm = forwardRef(function MCPForm( )} /> + ( + + {t('mcp.toolCallTimeout')} + + field.onChange(Number(e.target.value))} + /> + + + {t('mcp.toolCallTimeoutDescription')} + + + + )} + /> + {watchMode === 'remote' && ( <> ; timeout: number; ssereadtimeout: number; + tool_call_timeout_sec?: number; } export interface MCPServerExtraArgsStdio { command: string; args: string[]; env: Record; + tool_call_timeout_sec?: number; } export interface MCPServerExtraArgsHttp { url: string; headers: Record; timeout: number; + tool_call_timeout_sec?: number; } // "remote" mode: the user only supplies a URL; the backend auto-detects the @@ -540,6 +543,7 @@ export interface MCPServerExtraArgsRemote { url: string; headers?: Record; timeout?: number; + tool_call_timeout_sec?: number; } export enum MCPSessionStatus { diff --git a/web/src/i18n/locales/en-US.ts b/web/src/i18n/locales/en-US.ts index 2f9a71077..dbd1aea14 100644 --- a/web/src/i18n/locales/en-US.ts +++ b/web/src/i18n/locales/en-US.ts @@ -797,6 +797,9 @@ const enUS = { url: 'URL', headers: 'Headers', timeout: 'Timeout', + toolCallTimeout: 'Tool call timeout (seconds)', + toolCallTimeoutDescription: + 'Maximum wait for one tool call. Set to 0 for no timeout. Defaults to 300 seconds.', addArgument: 'Add Argument', addEnvVar: 'Add Environment Variable', addHeader: 'Add Header', diff --git a/web/src/i18n/locales/es-ES.ts b/web/src/i18n/locales/es-ES.ts index 6d25d0112..74f1985f1 100644 --- a/web/src/i18n/locales/es-ES.ts +++ b/web/src/i18n/locales/es-ES.ts @@ -811,6 +811,9 @@ const esES = { url: 'URL', headers: 'Encabezados', timeout: 'Tiempo de espera', + toolCallTimeout: 'Tiempo de espera de herramienta (segundos)', + toolCallTimeoutDescription: + 'Espera máxima para una llamada de herramienta. Use 0 para no limitar. El valor predeterminado es 300 segundos.', addArgument: 'Añadir argumento', addEnvVar: 'Añadir variable de entorno', addHeader: 'Añadir encabezado', diff --git a/web/src/i18n/locales/ja-JP.ts b/web/src/i18n/locales/ja-JP.ts index e3d59c55f..614891d98 100644 --- a/web/src/i18n/locales/ja-JP.ts +++ b/web/src/i18n/locales/ja-JP.ts @@ -803,6 +803,9 @@ const jaJP = { url: 'URL', headers: 'ヘッダー', timeout: 'タイムアウト', + toolCallTimeout: 'ツール呼び出しタイムアウト(秒)', + toolCallTimeoutDescription: + '1 回のツール呼び出しの最大待機時間です。0 で無制限、既定値は 300 秒です。', addArgument: '引数を追加', addEnvVar: '環境変数を追加', addHeader: 'ヘッダーを追加', diff --git a/web/src/i18n/locales/ru-RU.ts b/web/src/i18n/locales/ru-RU.ts index 0f6072c49..eef09cdee 100644 --- a/web/src/i18n/locales/ru-RU.ts +++ b/web/src/i18n/locales/ru-RU.ts @@ -808,6 +808,9 @@ const ruRU = { url: 'URL', headers: 'Заголовки', timeout: 'Таймаут', + toolCallTimeout: 'Таймаут вызова инструмента (секунды)', + toolCallTimeoutDescription: + 'Максимальное ожидание одного вызова. 0 отключает ограничение. По умолчанию 300 секунд.', addArgument: 'Добавить аргумент', addEnvVar: 'Добавить переменную окружения', addHeader: 'Добавить заголовок', diff --git a/web/src/i18n/locales/th-TH.ts b/web/src/i18n/locales/th-TH.ts index 704daf9cd..4ced7ace4 100644 --- a/web/src/i18n/locales/th-TH.ts +++ b/web/src/i18n/locales/th-TH.ts @@ -786,6 +786,9 @@ const thTH = { url: 'URL', headers: 'ส่วนหัว', timeout: 'หมดเวลา', + toolCallTimeout: 'หมดเวลาการเรียกเครื่องมือ (วินาที)', + toolCallTimeoutDescription: + 'เวลารอสูงสุดต่อการเรียกเครื่องมือหนึ่งครั้ง ตั้งเป็น 0 เพื่อไม่จำกัด ค่าเริ่มต้นคือ 300 วินาที', addArgument: 'เพิ่มอาร์กิวเมนต์', addEnvVar: 'เพิ่มตัวแปรสภาพแวดล้อม', addHeader: 'เพิ่มส่วนหัว', diff --git a/web/src/i18n/locales/vi-VN.ts b/web/src/i18n/locales/vi-VN.ts index b9770314e..26c415688 100644 --- a/web/src/i18n/locales/vi-VN.ts +++ b/web/src/i18n/locales/vi-VN.ts @@ -801,6 +801,9 @@ const viVN = { url: 'URL', headers: 'Tiêu đề', timeout: 'Thời gian chờ', + toolCallTimeout: 'Thời gian chờ gọi công cụ (giây)', + toolCallTimeoutDescription: + 'Thời gian chờ tối đa cho một lần gọi công cụ. Đặt 0 để không giới hạn. Mặc định là 300 giây.', addArgument: 'Thêm tham số', addEnvVar: 'Thêm biến môi trường', addHeader: 'Thêm tiêu đề', diff --git a/web/src/i18n/locales/zh-Hans.ts b/web/src/i18n/locales/zh-Hans.ts index a4a04e0de..96e24759d 100644 --- a/web/src/i18n/locales/zh-Hans.ts +++ b/web/src/i18n/locales/zh-Hans.ts @@ -763,6 +763,9 @@ const zhHans = { url: 'URL地址', headers: '请求头', timeout: '超时时间', + toolCallTimeout: '工具调用超时(秒)', + toolCallTimeoutDescription: + '单次工具调用的最长等待时间。设为 0 表示不限制,默认 300 秒。', addArgument: '添加参数', addEnvVar: '添加环境变量', addHeader: '添加请求头', diff --git a/web/src/i18n/locales/zh-Hant.ts b/web/src/i18n/locales/zh-Hant.ts index 1ae98f758..57f0bc58e 100644 --- a/web/src/i18n/locales/zh-Hant.ts +++ b/web/src/i18n/locales/zh-Hant.ts @@ -762,6 +762,9 @@ const zhHant = { url: 'URL位址', headers: '請求標頭', timeout: '逾時時間', + toolCallTimeout: '工具呼叫逾時(秒)', + toolCallTimeoutDescription: + '單次工具呼叫的最長等待時間。設為 0 表示不限制,預設 300 秒。', addArgument: '新增參數', addEnvVar: '新增環境變數', addHeader: '新增請求標頭',