fix(mcp): retry the stdio handshake during slow (npx) cold starts (#2306)

A node/npx stdio MCP server (e.g. firecrawl-mcp via npx -y) failed on first
connect with Connection closed / Failed after 4 attempts, even though the
process was fine. An npx cold start downloads+installs the package before the
server can answer the MCP handshake (measured ~27s for a simple official
server; longer for heavier ones). The old code attached the WS and called
session.initialize() the instant the process was started, so the handshake ran
before the process could answer and failed; the outer lifecycle retry then
rebuilt the process, churning it in a loop.

Verified decisively: attaching + initialize() against a mid-cold-start process
times out on attempt 1 (process still installing) but SUCCEEDS at t+0.6s on
attempt 2 once the process is ready. So the fix is to retry the handshake in
place, not to rebuild the process.

Changes (mcp_stdio.initialize):
- Start the managed process ONCE, then loop attach WS -> ClientSession ->
  initialize() within the startup_timeout budget, tearing down each failed
  attempt cleanly, until the handshake succeeds or the budget elapses. A
  successful transport/session is transferred into the owner exit stack via a
  small _TransferredStack adapter.
- Bound each attempt with asyncio.wait_for(initialize, _HANDSHAKE_ATTEMPT_TIMEOUT_SEC=10s)
  so a cold-starting process fails fast and retries instead of hanging until
  the transport drops.
- Stop retrying ONLY when the process has DEFINITIVELY exited: new
  _managed_process_has_exited() (checks EXITED status) replaces the previous
  not-_managed_process_is_running() test, which false-negatived on a
  just-spawned process that had not yet reported RUNNING and made the loop bail
  to the outer rebuild path (relay then rejected the early re-attach with HTTP 400).

Adds a unit test that fails the first two handshakes with the process alive and
asserts the loop retries to success while starting the process exactly once.

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>

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)
-            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
 
-        try:
+        # A freshly started stdio process may be slow to become ready — an
-            await self.owner.session.initialize()
+        # `npx -y <pkg>` cold start downloads+installs the package before the
-        except Exception:
+        # server can answer the MCP handshake (measured ~27s for a simple
-            self.owner.error_phase = MCPSessionErrorPhase.MCP_INIT
+        # server; more for heavier ones). Attaching the WS and calling
-            raise
+        # 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:

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