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

@@ -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 <pkg>` 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():

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 <pkg>` 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 <pkg>` 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 <pkg>`) 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

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()