mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-14 08:26:07 +00:00
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>
This commit is contained in:
Hyu
@@ -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 <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'
|
||||
)
|
||||
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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user