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>

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