From bf3c96026b8b508cfa2bf2cd232e37624b3f96fe Mon Sep 17 00:00:00 2001 From: Hyu Date: Fri, 3 Jul 2026 14:58:40 +0800 Subject: [PATCH] fix(mcp): fix cold-start connection drop by using lexically-intact owner exit stack (#2307) The previous _TransferredStack approach broke anyio lexical context: websocket_client/ClientSession use anyio task groups whose cancel scope is bound to the frame that entered them. Deferring their aclose via a transferred exit stack left the underlying memory streams closed once initialize() returned, so the very next request (refresh -> list_tools) failed with Connection closed. New design: - Attach on the owner exit stack (same task as the serve loop, lexically intact) - A cold-starting process makes initialize() fail; signal _ColdStartRetry up to the outer retry loop, which reuses the live process without consuming retry budget - _lifecycle_loop_with_retry handles _ColdStartRetry like _TransportReconnect: preserves process, no fatal budget, backs off 2s and retries - Two new unit tests: cold-start raises _ColdStartRetry (not fatal) when process is alive; raises fatal error when process has actually exited Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com> --- src/langbot/pkg/provider/tools/loaders/mcp.py | 20 +++- .../pkg/provider/tools/loaders/mcp_stdio.py | 101 +++++++--------- .../provider/test_mcp_box_integration.py | 111 ++++++++++++------ 3 files changed, 137 insertions(+), 95 deletions(-) diff --git a/src/langbot/pkg/provider/tools/loaders/mcp.py b/src/langbot/pkg/provider/tools/loaders/mcp.py index 539910ac7..499a00678 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp.py @@ -25,7 +25,7 @@ from ....core import app import langbot_plugin.api.entities.builtin.resource.tool as resource_tool import langbot_plugin.api.entities.builtin.provider.message as provider_message from ....entity.persistence import mcp as persistence_mcp -from .mcp_stdio import BoxStdioSessionRuntime, MCPServerBoxConfig, MCPSessionErrorPhase # noqa: F401 +from .mcp_stdio import BoxStdioSessionRuntime, MCPServerBoxConfig, MCPSessionErrorPhase, _ColdStartRetry # noqa: F401 # Synthesized LLM tools for MCP resources (not from server tools/list). # Dispatched in MCPLoader.invoke_tool; placeholder func on LLMTool is never used. @@ -490,6 +490,24 @@ class RuntimeMCPSession: self.error_phase = None await asyncio.sleep(1) continue + except _ColdStartRetry as e: + # The managed process is alive but still cold-starting (e.g. + # `npx -y ` is still installing) and cannot yet answer the + # handshake. Reuse the live process and retry the attach WITHOUT + # consuming the fatal retry budget or stopping the process, so a + # slow cold start is waited out instead of failing. Preserve the + # process across the finally-block cleanup. + if self._shutdown_event.is_set(): + return + self._preserve_managed_process = True + self.ap.logger.debug( + f'MCP session {self.server_name}: waiting for cold start ({self._describe_exception(e)})' + ) + self.status = MCPSessionStatus.CONNECTING + self.error_message = None + self.error_phase = None + await asyncio.sleep(2) + continue except Exception as e: self.retry_count = attempt + 1 if self._shutdown_event.is_set(): diff --git a/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py b/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py index 7d508b935..314023845 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py @@ -94,6 +94,15 @@ class _TransferredStack: return False +class _ColdStartRetry(Exception): + """Signal: the managed process is alive but not yet answering the MCP + handshake because it is still cold-starting (e.g. `npx -y ` is still + installing). The outer lifecycle retry treats this like a transient + reconnect: it reuses the live process and does not count toward the fatal + retry budget, so a slow cold start is waited out rather than failing. + """ + + class BoxStdioSessionRuntime: """Encapsulate Box-backed stdio MCP session orchestration.""" @@ -226,63 +235,43 @@ class BoxStdioSessionRuntime: websocket_url = workspace.get_managed_process_websocket_url(self.process_id) - # A freshly started stdio process may be slow to become ready — an - # `npx -y ` cold start downloads+installs the package before the - # server can answer the MCP handshake (measured ~27s for a simple - # server; more for heavier ones). Attaching the WS and calling - # initialize() immediately therefore fails with "Connection closed" - # even though the process is fine. The WS relay tolerates the silent - # cold-start window, so we retry the handshake (attach -> ClientSession - # -> initialize) within the startup budget, tearing down the failed - # transport between attempts, until the process is ready or we time out. - deadline = asyncio.get_running_loop().time() + max(float(self.config.startup_timeout_sec or 120), 1.0) - attempt = 0 - while True: - attempt += 1 - handshake_stack = AsyncExitStack() - try: - transport = await handshake_stack.enter_async_context(websocket_client(websocket_url)) - read_stream, write_stream = transport - session = await handshake_stack.enter_async_context(ClientSession(read_stream, write_stream)) - # Bound each handshake attempt: a process still in cold start - # will not answer, so time out and retry rather than hang until - # the transport eventually drops with 'Connection closed'. - await asyncio.wait_for(session.initialize(), timeout=_HANDSHAKE_ATTEMPT_TIMEOUT_SEC) - # Success: hand the live transport/session over to the owner's - # exit stack so it lives for the session's lifetime. - self.owner.session = session - await self.owner.exit_stack.enter_async_context(_TransferredStack(handshake_stack)) - if attempt > 1: - self.ap.logger.info( - f'MCP server {self.server_name}: handshake succeeded on attempt {attempt} ' - f'(process finished cold start)' - ) - break - except Exception as exc: - # Tear down this attempt's half-open transport before retrying. - try: - await handshake_stack.aclose() - except Exception: - pass - # Stop retrying ONLY if the process has definitively exited — - # a just-spawned process that is not yet RUNNING, or a transient - # box query error, must NOT abort the cold-start wait. - if await self._managed_process_has_exited(): - self.owner.error_phase = MCPSessionErrorPhase.MCP_INIT - raise - if asyncio.get_running_loop().time() >= deadline: - self.owner.error_phase = MCPSessionErrorPhase.MCP_INIT - raise Exception( - f'MCP handshake did not complete within ' - f'{int(self.config.startup_timeout_sec or 120)}s ' - f'({attempt} attempts): {type(exc).__name__}: {exc}' - ) - # Back off briefly and retry — the process is still coming up. - self.ap.logger.debug( - f'MCP server {self.server_name}: handshake attempt {attempt} not ready yet ' - f'({type(exc).__name__}); retrying while process starts' + # Attach the WS transport + MCP session ONCE, on the owner's exit stack, + # in the same task as the serve loop that follows. websocket_client and + # ClientSession use anyio task groups whose cancel scope is bound to the + # frame/stack that entered them, so they must live on the owner exit + # stack (not a deferred/transferred one) or the streams close the moment + # initialize() returns and the next request fails with "Connection + # closed". + # + # A slow (`npx -y `) cold start makes this single attempt fail + # while the process is still alive — the package is still installing and + # cannot answer the handshake. We surface that to the outer retry loop + # as a _ColdStartRetry: it must NOT stop the process (it is healthy and + # will be reused) and must NOT consume the fatal retry budget. The next + # attempt re-attaches to the same live process; once it has finished + # cold start the handshake succeeds and stays healthy. + try: + transport = await self.owner.exit_stack.enter_async_context(websocket_client(websocket_url)) + read_stream, write_stream = transport + self.owner.session = await self.owner.exit_stack.enter_async_context( + ClientSession(read_stream, write_stream) + ) + except Exception: + self.owner.error_phase = MCPSessionErrorPhase.RELAY_CONNECT + if not await self._managed_process_has_exited(): + # Process is alive but not yet serving (cold start) — reconnect. + raise _ColdStartRetry(f'{self.server_name}: transport not ready during cold start') + raise + + try: + await asyncio.wait_for(self.owner.session.initialize(), timeout=_HANDSHAKE_ATTEMPT_TIMEOUT_SEC) + except Exception as exc: + self.owner.error_phase = MCPSessionErrorPhase.MCP_INIT + if not await self._managed_process_has_exited(): + raise _ColdStartRetry( + f'{self.server_name}: handshake not ready during cold start ({type(exc).__name__})' ) - await asyncio.sleep(2) + raise async def monitor_process_health(self) -> None: from langbot_plugin.box.models import BoxManagedProcessStatus diff --git a/tests/unit_tests/provider/test_mcp_box_integration.py b/tests/unit_tests/provider/test_mcp_box_integration.py index bb8fcbd8e..bd0ab8a3d 100644 --- a/tests/unit_tests/provider/test_mcp_box_integration.py +++ b/tests/unit_tests/provider/test_mcp_box_integration.py @@ -832,19 +832,16 @@ async def test_init_box_stdio_server_stages_host_path_in_shared_workspace(mcp_mo @pytest.mark.asyncio -async def test_stdio_handshake_retries_during_cold_start(mcp_module, tmp_path, monkeypatch): - """A slow (npx-style) cold start: the first handshake attempts fail while - the process is still starting, but the managed process is NOT exited, so - the handshake retry loop keeps trying and succeeds once the process is - ready — rather than raising and letting the outer loop churn the process.""" - import asyncio as _asyncio +async def test_stdio_handshake_raises_coldstart_retry_while_process_alive(mcp_module, tmp_path, monkeypatch): + """During a slow (npx) cold start the handshake fails while the managed + process is still alive. initialize() must raise _ColdStartRetry (so the + outer lifecycle loop reuses the live process and retries without stopping it + or consuming the fatal budget), NOT a fatal error.""" from contextlib import asynccontextmanager mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio'] - calls = {'init': 0} - - class FlakyClientSession: + class ColdClientSession: def __init__(self, *_args): pass @@ -855,20 +852,15 @@ async def test_stdio_handshake_retries_during_cold_start(mcp_module, tmp_path, m return False async def initialize(self): - calls['init'] += 1 - # Fail the first two attempts (process still cold-starting), - # succeed on the third. - if calls['init'] < 3: - raise Exception('Connection closed') - return None + # Process still cold-starting: handshake fails. + raise Exception('Connection closed') @asynccontextmanager async def fake_websocket_client(_url: str): yield ('read-stream', 'write-stream') - monkeypatch.setattr(mcp_stdio_module, 'ClientSession', FlakyClientSession) + monkeypatch.setattr(mcp_stdio_module, 'ClientSession', ColdClientSession) monkeypatch.setattr(mcp_stdio_module, 'websocket_client', fake_websocket_client) - # Make the retry loop fast. monkeypatch.setattr(mcp_stdio_module, '_HANDSHAKE_ATTEMPT_TIMEOUT_SEC', 1.0, raising=False) ap = _make_ap() @@ -889,34 +881,77 @@ async def test_stdio_handshake_retries_during_cold_start(mcp_module, tmp_path, m ap=ap, ) - # The process never reports 'exited' during the cold start. + # Process is NOT exited (still cold-starting) and not yet running for reuse. async def _not_exited(): return False session._box_stdio_runtime._managed_process_has_exited = _not_exited - # First check (is it already running for reuse?) -> False so we start it. - running_flags = iter([False]) - async def _is_running(): - try: - return next(running_flags) - except StopIteration: - return True + async def _not_running(): + return False - session._box_stdio_runtime._managed_process_is_running = _is_running + session._box_stdio_runtime._managed_process_is_running = _not_running - # Speed up the inter-attempt sleep. - orig_sleep = _asyncio.sleep + with pytest.raises(mcp_stdio_module._ColdStartRetry): + await session._init_box_stdio_server() - async def _fast_sleep(_s): - await orig_sleep(0) - - monkeypatch.setattr(mcp_stdio_module.asyncio, 'sleep', _fast_sleep) - - await session._init_box_stdio_server() + # Process was started exactly once (the retry will reuse it, not rebuild). + assert ap.box_service.start_managed_process.await_count == 1 await session.exit_stack.aclose() - # Handshake was retried until success (3 initialize calls) and the process - # was started exactly once (never rebuilt). - assert calls['init'] == 3 - assert ap.box_service.start_managed_process.await_count == 1 + +@pytest.mark.asyncio +async def test_stdio_handshake_raises_fatal_when_process_exited(mcp_module, tmp_path, monkeypatch): + """If the handshake fails AND the process has definitively exited, that is a + real failure — initialize() must NOT swallow it as a cold-start retry.""" + from contextlib import asynccontextmanager + + mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio'] + + class DeadClientSession: + def __init__(self, *_args): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def initialize(self): + raise Exception('Connection closed') + + @asynccontextmanager + async def fake_websocket_client(_url: str): + yield ('read-stream', 'write-stream') + + monkeypatch.setattr(mcp_stdio_module, 'ClientSession', DeadClientSession) + monkeypatch.setattr(mcp_stdio_module, 'websocket_client', fake_websocket_client) + monkeypatch.setattr(mcp_stdio_module, '_HANDSHAKE_ATTEMPT_TIMEOUT_SEC', 1.0, raising=False) + + ap = _make_ap() + ap.box_service.available = True + ap.box_service.create_session = AsyncMock(return_value={}) + ap.box_service.start_managed_process = AsyncMock(return_value={}) + ap.box_service.get_managed_process_websocket_url = Mock(return_value='ws://box/p') + + session = _make_session( + mcp_module, + {'name': 'dead', 'uuid': 'dead-uuid', 'mode': 'stdio', 'command': 'npx', 'args': ['-y', 'x']}, + ap=ap, + ) + + async def _exited(): + return True + + session._box_stdio_runtime._managed_process_has_exited = _exited + + async def _not_running(): + return False + + session._box_stdio_runtime._managed_process_is_running = _not_running + + with pytest.raises(Exception) as ei: + await session._init_box_stdio_server() + assert not isinstance(ei.value, mcp_stdio_module._ColdStartRetry) + await session.exit_stack.aclose()