LangBot/tests/unit_tests/provider/test_mcp_remote_transport.py

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.api.http.context import ExecutionContext
from langbot.pkg.provider.tools.loaders.mcp import MCPToolCallTimeoutError, RuntimeMCPSession


TEST_EXECUTION_CONTEXT = ExecutionContext(
    instance_uuid='instance-a',
    workspace_uuid='workspace-a',
    placement_generation=1,
)


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'},
                                    }
                                ]
                            },
                        }
                    )
                if method == 'tools/call':
                    tool_name = message.get('params', {}).get('name')
                    if tool_name == 'hang':
                        return web.Response(status=202)
                    return web.json_response(
                        {
                            'jsonrpc': '2.0',
                            'id': message['id'],
                            'result': {
                                'content': [
                                    {
                                        'type': 'text',
                                        'text': 'healthy',
                                    }
                                ],
                                'isError': False,
                            },
                        }
                    )
                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,
    tool_call_timeout_sec: float = 300,
) -> RuntimeMCPSession:
    app = cast(
        Any,
        SimpleNamespace(
            logger=Mock(),
            workspace_service=SimpleNamespace(
                get_execution_binding=AsyncMock(
                    return_value=SimpleNamespace(
                        instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid,
                        workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid,
                        placement_generation=TEST_EXECUTION_CONTEXT.placement_generation,
                    )
                )
            ),
        ),
    )
    return RuntimeMCPSession(
        'remote-transport-test',
        {
            'uuid': 'srv-1',
            'mode': 'remote',
            'url': url,
            'timeout': timeout,
            'tool_call_timeout_sec': tool_call_timeout_sec,
        },
        True,
        app,
        TEST_EXECUTION_CONTEXT,
    )


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
async def test_remote_transport_tool_timeout_does_not_poison_session():
    async with _transport_server(200) as (probe, url):
        session = _session(url, tool_call_timeout_sec=0.05)
        try:
            await session._init_remote_server()

            with pytest.raises(MCPToolCallTimeoutError, match='timed out after 0.05 seconds'):
                await session.invoke_mcp_tool('hang', {})

            result = await session.invoke_mcp_tool('health_check', {})

            assert result[0].text == 'healthy'
            assert probe.streamable_messages.count('tools/call') == 2
        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)