diff --git a/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py b/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py index 120538cf6..7d508b935 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py @@ -6,7 +6,7 @@ import os import shutil import shlex import threading -from contextlib import suppress +from contextlib import suppress, AsyncExitStack from typing import TYPE_CHECKING, Any import pydantic @@ -74,6 +74,26 @@ class MCPServerBoxConfig(pydantic.BaseModel): model_config = pydantic.ConfigDict(extra='ignore') +_HANDSHAKE_ATTEMPT_TIMEOUT_SEC = 10.0 + + +class _TransferredStack: + """Adapts an already-populated AsyncExitStack into an async context manager + so ownership of its resources can be transferred into another exit stack. + Entering is a no-op; exiting closes the wrapped stack (and thus the live WS + transport + ClientSession) when the owning session shuts down.""" + + def __init__(self, stack: AsyncExitStack): + self._stack = stack + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + await self._stack.aclose() + return False + + class BoxStdioSessionRuntime: """Encapsulate Box-backed stdio MCP session orchestration.""" @@ -204,22 +224,65 @@ class BoxStdioSessionRuntime: f'process_id={self.process_id} (transport reconnect)' ) - try: - websocket_url = workspace.get_managed_process_websocket_url(self.process_id) - 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 - raise + websocket_url = workspace.get_managed_process_websocket_url(self.process_id) - try: - await self.owner.session.initialize() - except Exception: - self.owner.error_phase = MCPSessionErrorPhase.MCP_INIT - raise + # 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' + ) + await asyncio.sleep(2) async def monitor_process_health(self) -> None: from langbot_plugin.box.models import BoxManagedProcessStatus @@ -264,6 +327,28 @@ class BoxStdioSessionRuntime: status = info.get('status', '') if isinstance(info, dict) else getattr(info, 'status', '') return status in (BoxManagedProcessStatus.RUNNING.value, BoxManagedProcessStatus.RUNNING) + async def _managed_process_has_exited(self) -> bool: + """Return True only if the process is DEFINITIVELY gone (reports EXITED). + + Distinct from ``not _managed_process_is_running()``: a process that has + just been spawned may not yet report RUNNING, and a transient query + error is not proof of exit. During the cold-start handshake retry we + must NOT treat 'not yet running' or 'query failed' as a terminal + failure, or we bail out to the outer rebuild path and churn the + process (relay then rejects the early re-attach with HTTP 400). Only a + successful query that reports EXITED stops the retry loop. + """ + from langbot_plugin.box.models import BoxManagedProcessStatus + + workspace = self._build_workspace() + try: + info = await workspace.get_managed_process(self.process_id) + except Exception: + # Unknown — treat as 'still coming up', not exited. + return False + status = info.get('status', '') if isinstance(info, dict) else getattr(info, 'status', '') + return status in (BoxManagedProcessStatus.EXITED.value, BoxManagedProcessStatus.EXITED) + async def _stage_host_path_to_shared_workspace(self, host_path: str) -> str: source_path = normalize_host_path(host_path) if not source_path: diff --git a/tests/unit_tests/provider/test_mcp_box_integration.py b/tests/unit_tests/provider/test_mcp_box_integration.py index f0e06b2e8..bb8fcbd8e 100644 --- a/tests/unit_tests/provider/test_mcp_box_integration.py +++ b/tests/unit_tests/provider/test_mcp_box_integration.py @@ -829,3 +829,94 @@ async def test_init_box_stdio_server_stages_host_path_in_shared_workspace(mcp_mo assert process_payload['command'] == 'python' assert process_payload['args'] == ['/workspace/.mcp/u1/workspace/server.py'] assert process_payload['cwd'] == '/workspace/.mcp/u1/workspace' + + +@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 + from contextlib import asynccontextmanager + + mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio'] + + calls = {'init': 0} + + class FlakyClientSession: + 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): + 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 + + @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, '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() + 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': 'slow', + 'uuid': 'slow-uuid', + 'mode': 'stdio', + 'command': 'npx', + 'args': ['-y', 'some-mcp'], + }, + ap=ap, + ) + + # The process never reports 'exited' during the cold start. + 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 + + session._box_stdio_runtime._managed_process_is_running = _is_running + + # Speed up the inter-attempt sleep. + orig_sleep = _asyncio.sleep + + async def _fast_sleep(_s): + await orig_sleep(0) + + monkeypatch.setattr(mcp_stdio_module.asyncio, 'sleep', _fast_sleep) + + await session._init_box_stdio_server() + 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