diff --git a/src/langbot/pkg/provider/tools/loaders/mcp.py b/src/langbot/pkg/provider/tools/loaders/mcp.py index 97b88d8a5..5b0a30610 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp.py @@ -6,7 +6,7 @@ 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.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 @@ async def _init_sse_server(self): 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, - ), - ) - ) - - read, write, _ = transport + http_client=http_client, + ) as transport: + read, write, _ = transport + async with ClientSession(read, write) as session: + await session.initialize() + yield session - 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 @@ async def _init_remote_server(self): 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 @@ def visit(e: BaseException) -> None: 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 diff --git a/tests/unit_tests/provider/test_mcp_remote_transport.py b/tests/unit_tests/provider/test_mcp_remote_transport.py new file mode 100644 index 000000000..b2f1d2e14 --- /dev/null +++ b/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) diff --git a/tests/unit_tests/provider/test_mcp_resources.py b/tests/unit_tests/provider/test_mcp_resources.py index 565422c46..773efe71c 100644 --- a/tests/unit_tests/provider/test_mcp_resources.py +++ b/tests/unit_tests/provider/test_mcp_resources.py @@ -4,6 +4,7 @@ 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()