fix: survive transient MCP transport faults on preflight and exhausted retries

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "switchplane"
-version = "0.10.1"
+version = "0.10.2"
 description = "Python runtime control plane for agent-based task execution, LangGraph-native"
 readme = "README.md"
 license = {text = "Apache-2.0"}

src/switchplane/mcp.py

@@ -1,5 +1,6 @@
 """MCP client lifecycle management and LangChain tool integration."""
 
+import asyncio
 import datetime
 import importlib
 import inspect
@@ -14,6 +15,59 @@
 
 logger = structlog.get_logger()
 
+# Preflight connectivity probe: retry a couple times on *transient* transport
+# faults before declaring a server unreachable. A momentary "server disconnected
+# without a response" (RemoteProtocolError) or connection drop on the startup GET
+# would otherwise abort the whole task before it begins — the same transient-fault-
+# is-fatal failure mode the RetryTransport absorbs for in-session tool calls, but on
+# the startup path the RetryTransport isn't in play yet. Persistent errors (DNS, SSL,
+# auth) fail on every attempt and still surface, so the diagnostic value is kept.
+_PREFLIGHT_MAX_ATTEMPTS = 3
+_PREFLIGHT_BACKOFF_SECONDS = 1.0
+
+
+async def _run_preflight(client, url: str, server_name: str) -> None:
+    """Probe *url* with a GET, retrying transient transport faults.
+
+    A successful response (any status) or a timeout proves reachability — MCP
+    endpoints are POST-oriented, so a GET timeout is expected and benign. A
+    transient transport fault (connection drop / "server disconnected without a
+    response") is retried with capped exponential backoff; only after attempts
+    are exhausted, or on a non-transient error (DNS/SSL/etc.), does it raise
+    ``ConnectionError``. ``client`` is an already-built ``httpx.AsyncClient``;
+    its lifecycle (close) is the caller's responsibility.
+    """
+    import httpx
+
+    # Mirrors `RetryTransport._RETRYABLE_EXC` minus the timeout, which is its
+    # own benign "reachable" signal below.
+    retryable = (httpx.ConnectError, httpx.RemoteProtocolError)
+    attempt = 0
+    while True:
+        try:
+            async with client.stream("GET", url) as resp:
+                logger.info("mcp_preflight_ok", server=server_name, status=resp.status_code)
+            return
+        except httpx.TimeoutException:
+            logger.info("mcp_preflight_timeout", server=server_name)
+            return
+        except retryable as e:
+            if attempt >= _PREFLIGHT_MAX_ATTEMPTS - 1:
+                raise ConnectionError(f"Cannot reach MCP server '{server_name}' at {url}: {e}") from e
+            delay = _PREFLIGHT_BACKOFF_SECONDS * (2**attempt)
+            logger.warning(
+                "mcp_preflight_retrying",
+                server=server_name,
+                attempt=attempt + 1,
+                max_attempts=_PREFLIGHT_MAX_ATTEMPTS,
+                error=type(e).__name__,
+                delay_seconds=round(delay, 1),
+            )
+            await asyncio.sleep(delay)
+            attempt += 1
+        except Exception as e:
+            raise ConnectionError(f"Cannot reach MCP server '{server_name}' at {url}: {e}") from e
+
 
 def _import_transport_factory(dotted_path: str):
     """Import and validate an HTTP transport factory from a dotted path.
@@ -122,12 +176,7 @@ async def start(self, stack: AsyncExitStack) -> None:
                 )
 
             try:
-                async with preflight_client.stream("GET", self.config.url) as resp:
-                    logger.info("mcp_preflight_ok", server=self.config.name, status=resp.status_code)
-            except httpx.TimeoutException:
-                logger.info("mcp_preflight_timeout", server=self.config.name)
-            except Exception as e:
-                raise ConnectionError(f"Cannot reach MCP server '{self.config.name}' at {self.config.url}: {e}") from e
+                await _run_preflight(preflight_client, self.config.url, self.config.name)
             finally:
                 if self.config.oauth or http_client is None:
                     await preflight_client.aclose()

src/switchplane/oauth.py

@@ -402,6 +402,57 @@ def _retry_after_seconds(response: httpx.Response | None, attempt: int) -> float
     return min(max(0.0, backoff), _RETRY_BACKOFF_MAX_SECONDS)
 
 
+# JSON-RPC `error.code` for the synthesized exhausted-retry response. We reuse the
+# SDK's own REQUEST_TIMEOUT (HTTP 408) value — the same code the SDK's foreground
+# `anyio.fail_after` path uses for a read-timeout (see `BaseSession.send_request`)
+# — so the surfaced `McpError` reads as a timeout regardless of which path produced
+# it. JSON-RPC permits any integer code; matching the SDK keeps the two consistent.
+_EXHAUSTED_RETRY_ERROR_CODE = httpx.codes.REQUEST_TIMEOUT
+
+
+def _exhausted_retry_response(request: httpx.Request, exc: Exception) -> httpx.Response | None:
+    """Synthesize an in-band JSON-RPC error response for a request whose retries
+    are exhausted, so the MCP SDK forwards it to the waiting caller as a catchable
+    ``McpError`` instead of letting the raw transport fault crash the session.
+
+    The MCP streamable-HTTP SDK runs each id-bearing ``tools/call`` as a *background*
+    task in a session-scoped anyio task group (``post_writer`` → ``tg.start_soon``),
+    outside its own ``try/except``. A transport fault that escapes that task crashes
+    the whole group, cancels the foreground request, and detonates at session
+    teardown — taking the entire agent task down (a transient codesearch timeout
+    should never be fatal). Returning a JSON-RPC error response whose ``id`` echoes
+    the request lets the background task complete *normally*; the SDK demuxes it to
+    the foreground (``BaseSession._receive_response``) which raises ``McpError`` —
+    the tool loop's existing exception handling then feeds it back to the model.
+
+    Returns ``None`` when the request body has no usable JSON-RPC ``id`` (e.g. a
+    notification, or an unparseable body): there is no foreground waiter to route
+    a response to, so the caller must fall back to re-raising.
+    """
+    try:
+        payload = json.loads(request.content)
+    except (ValueError, TypeError):
+        return None
+    if not isinstance(payload, dict) or "id" not in payload:
+        return None
+    body = json.dumps(
+        {
+            "jsonrpc": "2.0",
+            "id": payload["id"],
+            "error": {
+                "code": _EXHAUSTED_RETRY_ERROR_CODE,
+                "message": (f"transport fault after exhausting retries: {type(exc).__name__}: {exc}"),
+            },
+        }
+    ).encode()
+    return httpx.Response(
+        200,
+        headers={"Content-Type": "application/json"},
+        content=body,
+        request=request,
+    )
+
+
 class RetryTransport(httpx.AsyncBaseTransport):
     """Wraps an httpx transport to retry transient failures transparently.
 
@@ -474,6 +525,24 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
                 response = await self._wrapped.handle_async_request(request)
             except self._RETRYABLE_EXC as exc:
                 if attempt >= self._max_retries:
+                    # Retries exhausted. Re-raising here escapes back into the MCP
+                    # SDK's session-scoped task group (the request runs as a
+                    # `tg.start_soon` background task) and detonates at session
+                    # teardown, taking the whole agent task down. Instead, hand the
+                    # SDK an in-band JSON-RPC error response echoing the request id:
+                    # the background task completes normally and the fault surfaces
+                    # in the foreground as a catchable `McpError`. Only re-raise when
+                    # we can't synthesize one (no JSON-RPC id to route a response to).
+                    synthesized = _exhausted_retry_response(request, exc)
+                    logger.warning(
+                        "mcp_transport_retries_exhausted",
+                        server=self._server_name,
+                        max_retries=self._max_retries,
+                        error=type(exc).__name__,
+                        surfaced="in_band_error" if synthesized is not None else "reraised",
+                    )
+                    if synthesized is not None:
+                        return synthesized
                     raise
                 delay = _retry_after_seconds(None, attempt)
                 logger.warning(

tests/test_mcp.py

@@ -379,6 +379,109 @@ def test_no_input_schema(self):
         assert tool.name == "simple"
 
 
+class _FakeStreamCM:
+    """Async context manager standing in for `client.stream(...)`: raises the
+    queued outcome on enter, or yields a response stub with `status_code`."""
+
+    def __init__(self, outcome):
+        self._outcome = outcome
+
+    async def __aenter__(self):
+        if isinstance(self._outcome, Exception):
+            raise self._outcome
+        resp = MagicMock()
+        resp.status_code = self._outcome
+        return resp
+
+    async def __aexit__(self, *exc):
+        return False
+
+
+class _FakePreflightClient:
+    """Fake httpx.AsyncClient that replays a queue of `stream()` outcomes."""
+
+    def __init__(self, outcomes):
+        self._outcomes = list(outcomes)
+        self.stream_calls = 0
+
+    def stream(self, method, url):
+        self.stream_calls += 1
+        return _FakeStreamCM(self._outcomes.pop(0))
+
+    async def aclose(self):
+        pass
+
+
+class TestPreflightRetry:
+    """Unit tests for the `_run_preflight` connectivity retry (no live endpoint)."""
+
+    @pytest.mark.asyncio
+    async def test_succeeds_first_try(self, monkeypatch):
+        from switchplane import mcp as mcp_mod
+
+        client = _FakePreflightClient([200])
+        sleep = AsyncMock()
+        monkeypatch.setattr(mcp_mod.asyncio, "sleep", sleep)
+        await mcp_mod._run_preflight(client, "https://x/mcp", "cs")
+        assert client.stream_calls == 1
+        assert sleep.await_count == 0
+
+    @pytest.mark.asyncio
+    async def test_timeout_is_benign(self, monkeypatch):
+        import httpx
+
+        from switchplane import mcp as mcp_mod
+
+        client = _FakePreflightClient([httpx.ConnectTimeout("slow")])
+        sleep = AsyncMock()
+        monkeypatch.setattr(mcp_mod.asyncio, "sleep", sleep)
+        # Must NOT raise — a GET timeout proves reachability for POST-oriented MCP.
+        await mcp_mod._run_preflight(client, "https://x/mcp", "cs")
+        assert client.stream_calls == 1
+        assert sleep.await_count == 0
+
+    @pytest.mark.asyncio
+    async def test_retries_transient_fault_then_raises(self, monkeypatch):
+        import httpx
+
+        from switchplane import mcp as mcp_mod
+
+        client = _FakePreflightClient([httpx.RemoteProtocolError("disconnected")] * 5)
+        sleep = AsyncMock()
+        monkeypatch.setattr(mcp_mod.asyncio, "sleep", sleep)
+        with pytest.raises(ConnectionError, match="Cannot reach MCP server 'cs'"):
+            await mcp_mod._run_preflight(client, "https://x/mcp", "cs")
+        # initial + 2 retries = 3 attempts, 2 backoff sleeps
+        assert client.stream_calls == 3
+        assert sleep.await_count == 2
+
+    @pytest.mark.asyncio
+    async def test_retries_transient_fault_then_succeeds(self, monkeypatch):
+        import httpx
+
+        from switchplane import mcp as mcp_mod
+
+        client = _FakePreflightClient([httpx.RemoteProtocolError("disconnected"), 200])
+        sleep = AsyncMock()
+        monkeypatch.setattr(mcp_mod.asyncio, "sleep", sleep)
+        await mcp_mod._run_preflight(client, "https://x/mcp", "cs")
+        assert client.stream_calls == 2  # one fault, then success
+        assert sleep.await_count == 1
+
+    @pytest.mark.asyncio
+    async def test_non_transient_error_raises_without_retry(self, monkeypatch):
+        from switchplane import mcp as mcp_mod
+
+        # A non-retryable fault (e.g. SSL) surfaces immediately, no retries.
+        client = _FakePreflightClient([ValueError("bad cert")])
+        sleep = AsyncMock()
+        monkeypatch.setattr(mcp_mod.asyncio, "sleep", sleep)
+        with pytest.raises(ConnectionError, match="Cannot reach MCP server 'cs'"):
+            await mcp_mod._run_preflight(client, "https://x/mcp", "cs")
+        assert client.stream_calls == 1
+        assert sleep.await_count == 0
+
+
 class TestPreflightIntegration:
     """Integration tests against a live MCP endpoint. Only run with ITEST=1."""
 

tests/test_oauth.py

@@ -5,6 +5,7 @@
 elsewhere; here we cover the rate-limit retry that wraps the MCP HTTP client.
 """
 
+import json
 from unittest.mock import AsyncMock, patch
 
 import httpx
@@ -195,14 +196,49 @@ async def test_retry_transport_replays_body_unchanged_across_retries():
     assert inner.seen_bodies == [b'{"jsonrpc":"2.0"}', b'{"jsonrpc":"2.0"}']
 
 
-async def test_retry_transport_reraises_timeout_after_max_retries():
+async def test_retry_transport_reraises_timeout_after_max_retries_without_id():
+    # `_req()`'s body carries no JSON-RPC `id`, so there is no foreground waiter
+    # to route a synthesized error response to — the transport must re-raise.
     inner = _FaultTransport([httpx.ReadTimeout("hung")] * 5)
     rt = RetryTransport(inner, max_retries=2, server_name="codesearch")
     with patch("switchplane.oauth.asyncio.sleep", new=AsyncMock()), pytest.raises(httpx.ReadTimeout):
         await rt.handle_async_request(_req())
     assert inner.request_count == 3  # initial + 2 retries
 
 
+def _req_with_id(request_id=7):
+    # An id-bearing JSON-RPC request, as the MCP SDK sends for `tools/call`.
+    body = f'{{"jsonrpc":"2.0","id":{request_id},"method":"tools/call"}}'.encode()
+    return httpx.Request("POST", "https://mcp.codesearch/mcp", content=body)
+
+
+async def test_retry_transport_synthesizes_in_band_error_on_exhaustion_with_id():
+    # An id-bearing request whose retries are exhausted must NOT re-raise (that
+    # would crash the SDK's session task group); it returns a JSON-RPC error
+    # response echoing the request id so the SDK surfaces a catchable McpError.
+    inner = _FaultTransport([httpx.ReadTimeout("hung")] * 5)
+    rt = RetryTransport(inner, max_retries=2, server_name="codesearch")
+    with patch("switchplane.oauth.asyncio.sleep", new=AsyncMock()):
+        resp = await rt.handle_async_request(_req_with_id(7))
+    assert inner.request_count == 3  # initial + 2 retries
+    assert resp.status_code == 200
+    assert resp.headers["Content-Type"] == "application/json"
+    payload = json.loads(resp.content)
+    assert payload["id"] == 7
+    assert payload["error"]["code"] == int(httpx.codes.REQUEST_TIMEOUT)
+    assert "ReadTimeout" in payload["error"]["message"]
+
+
+async def test_retry_transport_connect_error_exhaustion_also_synthesizes():
+    inner = _FaultTransport([httpx.ConnectError("refused")] * 5)
+    rt = RetryTransport(inner, max_retries=1, server_name="codesearch")
+    with patch("switchplane.oauth.asyncio.sleep", new=AsyncMock()):
+        resp = await rt.handle_async_request(_req_with_id(3))
+    payload = json.loads(resp.content)
+    assert payload["id"] == 3
+    assert "ConnectError" in payload["error"]["message"]
+
+
 async def test_retry_transport_retries_connect_error():
     inner = _FaultTransport([httpx.ConnectError("refused"), 200])
     rt = RetryTransport(inner, max_retries=3, server_name="codesearch")