fix(mcp): harden remote SSE fallback (#2324)

* fix(mcp): harden remote SSE fallback

* fix(mcp): preserve transport errors during fallback

---------

Co-authored-by: RockChinQ <45992437+RockChinQ@users.noreply.github.com>

src/langbot/pkg/provider/tools/loaders/mcp.py

@@ -6,7 +6,7 @@ import json
 import re
 import time
 import typing
-from contextlib import AsyncExitStack
+from contextlib import AsyncExitStack, asynccontextmanager
 import traceback
 from langbot_plugin.api.entities.events import pipeline_query
 import sqlalchemy
@@ -18,6 +18,7 @@ from mcp import ClientSession, StdioServerParameters, types as mcp_types
 from mcp.client.stdio import stdio_client
 from mcp.client.sse import sse_client
 from mcp.client.streamable_http import streamable_http_client
+from mcp.shared.exceptions import McpError
 from pydantic import AnyUrl
 
 from .. import loader
@@ -335,23 +336,34 @@ class RuntimeMCPSession:
 
         await self.session.initialize()
 
-    async def _init_streamable_http_server(self):
-        transport = await self.exit_stack.enter_async_context(
-            streamable_http_client(
+    @asynccontextmanager
+    async def _streamable_http_session(self) -> typing.AsyncIterator[ClientSession]:
+        """Enter a fully initialized Streamable HTTP session as one context.
+
+        Initialization must happen inside the same context manager that owns the
+        MCP transport. The SDK reports request failures by cancelling the host
+        task and raises the real HTTP error from its TaskGroup during context
+        exit. Keeping these nested contexts together guarantees a failed
+        ``__aenter__`` unwinds immediately, so callers see the HTTPStatusError
+        instead of a detached CancelledError. It also owns the injected HTTPX
+        client, which the MCP SDK deliberately does not close for callers.
+        """
+        async with httpx.AsyncClient(
+            headers=self.server_config.get('headers', {}),
+            timeout=self.server_config.get('timeout', 10),
+            follow_redirects=True,
+        ) as http_client:
+            async with streamable_http_client(
                 self.server_config['url'],
-                http_client=httpx.AsyncClient(
-                    headers=self.server_config.get('headers', {}),
-                    timeout=self.server_config.get('timeout', 10),
-                    follow_redirects=True,
-                ),
-            )
-        )
+                http_client=http_client,
+            ) as transport:
+                read, write, _ = transport
+                async with ClientSession(read, write) as session:
+                    await session.initialize()
+                    yield session
 
-        read, write, _ = transport
-
-        self.session = await self.exit_stack.enter_async_context(ClientSession(read, write))
-
-        await self.session.initialize()
+    async def _init_streamable_http_server(self):
+        self.session = await self.exit_stack.enter_async_context(self._streamable_http_session())
 
     async def _init_remote_server(self):
         """Connect to a remote MCP server, auto-detecting the transport.
@@ -366,9 +378,15 @@ class RuntimeMCPSession:
             await self._init_streamable_http_server()
             return
         except Exception as e:
+            if not self._should_fallback_to_sse(e):
+                self.ap.logger.info(
+                    f'MCP server {self.server_name}: Streamable HTTP transport failed '
+                    f'({self._describe_exception(e)}); not falling back to SSE'
+                )
+                raise
             self.ap.logger.info(
-                f'MCP server {self.server_name}: Streamable HTTP transport failed '
-                f'({self._describe_exception(e)}), falling back to SSE'
+                f'MCP server {self.server_name}: Streamable HTTP initialize failed with a compatible HTTP status '
+                f'({self._describe_exception(e)}), falling back to legacy SSE'
             )
 
         # The Streamable HTTP attempt may have partially entered the transport /
@@ -575,6 +593,36 @@ class RuntimeMCPSession:
         unique = [m for m in leaves if not (m in seen or seen.add(m))]
         return '; '.join(unique) if unique else f'{type(exc).__name__}: {exc}'
 
+    @staticmethod
+    def _iter_exception_leaves(exc: BaseException) -> typing.Iterator[BaseException]:
+        sub = getattr(exc, 'exceptions', None)
+        if sub:  # ExceptionGroup / BaseExceptionGroup
+            for child in sub:
+                yield from RuntimeMCPSession._iter_exception_leaves(child)
+        else:
+            yield exc
+
+    @staticmethod
+    def _should_fallback_to_sse(exc: BaseException) -> bool:
+        """Whether a Streamable HTTP failure matches legacy-SSE fallback.
+
+        Only protocol-compatibility responses trigger fallback. Authentication,
+        authorization, throttling, and server failures must remain visible
+        instead of being retried against a different transport.
+
+        MCP SDK 1.26 translates an HTTP 404 initialize response into a synthetic
+        ``McpError(32600, 'Session terminated')`` rather than preserving the
+        HTTPStatusError, so recognize that exact SDK sentinel as 404-compatible.
+        """
+        fallback_statuses = {400, 404, 405}
+        for leaf in RuntimeMCPSession._iter_exception_leaves(exc):
+            if isinstance(leaf, httpx.HTTPStatusError):
+                if leaf.response.status_code in fallback_statuses:
+                    return True
+            elif isinstance(leaf, McpError) and leaf.error.code == 32600 and leaf.error.message == 'Session terminated':
+                return True
+        return False
+
     _MONITOR_POLL_INTERVAL = 5
     _MONITOR_MAX_CONSECUTIVE_ERRORS = 3
 

tests/unit_tests/provider/test_mcp_remote_transport.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+from typing import Any, cast
+from unittest.mock import AsyncMock, Mock
+
+import httpx
+import pytest
+from aiohttp import web
+from mcp import types as mcp_types
+
+from langbot.pkg.provider.tools.loaders.mcp import RuntimeMCPSession
+
+
+class _TransportProbe:
+    def __init__(self, streamable_status: int | None) -> None:
+        self.streamable_status = streamable_status
+        self.streamable_posts = 0
+        self.streamable_messages: list[str] = []
+        self.sse_gets = 0
+        self.sse_messages: list[str] = []
+        self.streamable_request_started = asyncio.Event()
+        self.release_streamable_request = asyncio.Event()
+        self._sse_response: web.StreamResponse | None = None
+
+    async def handle_mcp_endpoint(self, request: web.Request) -> web.StreamResponse:
+        if request.method == 'POST':
+            self.streamable_posts += 1
+            self.streamable_request_started.set()
+            if self.streamable_status is None:
+                await self.release_streamable_request.wait()
+                return web.Response(status=204)
+            if self.streamable_status == 200:
+                message = await request.json()
+                method = message.get('method', '')
+                self.streamable_messages.append(method)
+                if method == 'initialize':
+                    return web.json_response(
+                        {
+                            'jsonrpc': '2.0',
+                            'id': message['id'],
+                            'result': {
+                                'protocolVersion': mcp_types.LATEST_PROTOCOL_VERSION,
+                                'capabilities': {'tools': {}},
+                                'serverInfo': {'name': 'streamable-test', 'version': '1.0.0'},
+                            },
+                        }
+                    )
+                if method == 'tools/list':
+                    return web.json_response(
+                        {
+                            'jsonrpc': '2.0',
+                            'id': message['id'],
+                            'result': {
+                                'tools': [
+                                    {
+                                        'name': 'echo',
+                                        'description': 'Echo test input',
+                                        'inputSchema': {'type': 'object'},
+                                    }
+                                ]
+                            },
+                        }
+                    )
+                return web.Response(status=202)
+            return web.Response(status=self.streamable_status)
+
+        self.sse_gets += 1
+        response = web.StreamResponse(
+            status=200,
+            headers={
+                'Content-Type': 'text/event-stream',
+                'Cache-Control': 'no-cache',
+            },
+        )
+        await response.prepare(request)
+        self._sse_response = response
+        await response.write(b'event: endpoint\ndata: /messages?session_id=test-session\n\n')
+        try:
+            while request.transport is not None and not request.transport.is_closing():
+                await asyncio.sleep(0.05)
+        except asyncio.CancelledError:
+            raise
+        return response
+
+    async def handle_sse_message(self, request: web.Request) -> web.Response:
+        message = await request.json()
+        method = message.get('method', '')
+        self.sse_messages.append(method)
+
+        if method == 'initialize':
+            response_message = {
+                'jsonrpc': '2.0',
+                'id': message['id'],
+                'result': {
+                    'protocolVersion': mcp_types.LATEST_PROTOCOL_VERSION,
+                    'capabilities': {},
+                    'serverInfo': {'name': 'legacy-sse-test', 'version': '1.0.0'},
+                },
+            }
+            assert self._sse_response is not None
+            payload = json.dumps(response_message, separators=(',', ':'))
+            await self._sse_response.write(f'event: message\ndata: {payload}\n\n'.encode())
+
+        return web.Response(status=202)
+
+
+@asynccontextmanager
+async def _transport_server(streamable_status: int | None):
+    probe = _TransportProbe(streamable_status)
+    application = web.Application()
+    application.router.add_route('*', '/mcp', probe.handle_mcp_endpoint)
+    application.router.add_post('/messages', probe.handle_sse_message)
+    runner = web.AppRunner(application, shutdown_timeout=0.1)
+    await runner.setup()
+    site = web.TCPSite(runner, '127.0.0.1', 0)
+    await site.start()
+    server = cast(asyncio.Server, site._server)
+    port = server.sockets[0].getsockname()[1]
+    try:
+        yield probe, f'http://127.0.0.1:{port}/mcp'
+    finally:
+        await runner.cleanup()
+
+
+def _session(url: str, *, timeout: float = 2) -> RuntimeMCPSession:
+    app = cast(Any, SimpleNamespace(logger=Mock()))
+    return RuntimeMCPSession(
+        'remote-transport-test',
+        {'uuid': 'srv-1', 'mode': 'remote', 'url': url, 'timeout': timeout},
+        True,
+        app,
+    )
+
+
+def _contains_http_status(exc: BaseException, status_code: int) -> bool:
+    return any(
+        isinstance(leaf, httpx.HTTPStatusError) and leaf.response.status_code == status_code
+        for leaf in RuntimeMCPSession._iter_exception_leaves(exc)
+    )
+
+
+async def _close_session(session: RuntimeMCPSession) -> None:
+    await session.exit_stack.aclose()
+
+
+@pytest.mark.asyncio
+async def test_remote_transport_real_streamable_http_success_keeps_session_usable():
+    async with _transport_server(200) as (probe, url):
+        session = _session(url)
+        try:
+            await session._init_remote_server()
+            assert session.session is not None
+            tools = await session.session.list_tools()
+            assert [tool.name for tool in tools.tools] == ['echo']
+            assert probe.streamable_posts >= 2
+            assert probe.streamable_messages[:2] == ['initialize', 'notifications/initialized']
+            assert 'tools/list' in probe.streamable_messages
+            assert probe.sse_gets == 0
+        finally:
+            await _close_session(session)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('status_code', [400, 404, 405])
+async def test_remote_transport_real_streamable_http_error_falls_back_to_legacy_sse(status_code: int):
+    async with _transport_server(status_code) as (probe, url):
+        session = _session(url)
+        try:
+            await session._init_remote_server()
+            assert session.session is not None
+            assert probe.streamable_posts == 1
+            assert probe.sse_gets == 1
+            assert 'initialize' in probe.sse_messages
+        finally:
+            await _close_session(session)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('status_code', [401, 403, 406, 415, 429, 500])
+async def test_remote_transport_real_non_compatibility_error_does_not_fallback(status_code: int):
+    async with _transport_server(status_code) as (probe, url):
+        session = _session(url)
+        try:
+            with pytest.raises(BaseException) as exc_info:
+                await session._init_remote_server()
+            assert _contains_http_status(exc_info.value, status_code)
+            assert probe.streamable_posts == 1
+            assert probe.sse_gets == 0
+        finally:
+            await _close_session(session)
+
+
+@pytest.mark.asyncio
+async def test_remote_transport_real_timeout_does_not_fallback():
+    async with _transport_server(None) as (probe, url):
+        session = _session(url, timeout=0.05)
+        try:
+            with pytest.raises(BaseException) as exc_info:
+                await session._init_remote_server()
+            assert any(
+                isinstance(leaf, httpx.TimeoutException)
+                for leaf in RuntimeMCPSession._iter_exception_leaves(exc_info.value)
+            )
+            assert probe.streamable_posts == 1
+            assert probe.sse_gets == 0
+        finally:
+            probe.release_streamable_request.set()
+            await _close_session(session)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('error_type', [httpx.ConnectError, httpx.ConnectTimeout])
+async def test_remote_transport_connection_errors_do_not_fallback(error_type: type[httpx.RequestError]):
+    request = httpx.Request('POST', 'https://unreachable.invalid/mcp')
+    error = error_type('connection failed', request=request)
+    session = _session(str(request.url))
+    session._init_streamable_http_server = AsyncMock(side_effect=error)
+    session._init_sse_server = AsyncMock()
+
+    with pytest.raises(type(error)) as exc_info:
+        await session._init_remote_server()
+
+    assert exc_info.value is error
+    session._init_sse_server.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_remote_transport_external_cancellation_is_not_converted_to_sse_fallback():
+    async with _transport_server(None) as (probe, url):
+        session = _session(url)
+        task = asyncio.create_task(session._init_remote_server())
+        await asyncio.wait_for(probe.streamable_request_started.wait(), timeout=2)
+        task.cancel()
+        try:
+            with pytest.raises(asyncio.CancelledError):
+                await task
+            assert probe.sse_gets == 0
+        finally:
+            probe.release_streamable_request.set()
+            await _close_session(session)

tests/unit_tests/provider/test_mcp_resources.py

@@ -4,6 +4,7 @@ import base64
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, Mock
 
+import httpx
 import pytest
 from mcp import types as mcp_types
 
@@ -54,6 +55,50 @@ def _query() -> SimpleNamespace:
     return SimpleNamespace(variables={})
 
 
+def _http_status_error(status_code: int) -> httpx.HTTPStatusError:
+    request = httpx.Request('POST', 'https://example.com/mcp')
+    response = httpx.Response(status_code, request=request)
+    return httpx.HTTPStatusError(f'HTTP {status_code}', request=request, response=response)
+
+
+@pytest.mark.asyncio
+async def test_remote_transport_falls_back_to_sse_for_compatible_http_status_in_exception_group():
+    session = RuntimeMCPSession(
+        'remote',
+        {'uuid': 'srv-1', 'mode': 'remote', 'url': 'https://example.com/mcp'},
+        True,
+        _app(),
+    )
+    session._init_streamable_http_server = AsyncMock(
+        side_effect=ExceptionGroup('transport failed', [_http_status_error(405)])
+    )
+    session._init_sse_server = AsyncMock()
+
+    await session._init_remote_server()
+
+    session._init_streamable_http_server.assert_awaited_once()
+    session._init_sse_server.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_remote_transport_does_not_fallback_for_auth_http_status():
+    session = RuntimeMCPSession(
+        'remote',
+        {'uuid': 'srv-1', 'mode': 'remote', 'url': 'https://example.com/mcp'},
+        True,
+        _app(),
+    )
+    error = _http_status_error(403)
+    session._init_streamable_http_server = AsyncMock(side_effect=error)
+    session._init_sse_server = AsyncMock()
+
+    with pytest.raises(httpx.HTTPStatusError):
+        await session._init_remote_server()
+
+    session._init_streamable_http_server.assert_awaited_once()
+    session._init_sse_server.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_read_resource_envelope_truncates_caches_and_records_trace():
     session = _connected_session()