diff --git a/src/mcpo/main.py b/src/mcpo/main.py index b0196282..a2cf2539 100644 --- a/src/mcpo/main.py +++ b/src/mcpo/main.py @@ -125,9 +125,9 @@ async def _open_session_locked(self): client_context = self._create_client_context() try: connection = await client_context.__aenter__() - except Exception: + except (Exception, asyncio.CancelledError): # Ensure the context is closed if entering fails - with contextlib.suppress(Exception): + with contextlib.suppress(Exception, asyncio.CancelledError): await client_context.__aexit__(None, None, None) raise @@ -135,10 +135,10 @@ async def _open_session_locked(self): session = ClientSession(reader, writer) try: await session.__aenter__() - except Exception: - with contextlib.suppress(Exception): + except (Exception, asyncio.CancelledError): + with contextlib.suppress(Exception, asyncio.CancelledError): await session.__aexit__(None, None, None) - with contextlib.suppress(Exception): + with contextlib.suppress(Exception, asyncio.CancelledError): await client_context.__aexit__(None, None, None) raise @@ -155,11 +155,11 @@ async def _close_session_locked(self): self._initialize_result = None if session is not None: - with contextlib.suppress(Exception): + with contextlib.suppress(Exception, asyncio.CancelledError): await session.__aexit__(None, None, None) if client_context is not None: - with contextlib.suppress(Exception): + with contextlib.suppress(Exception, asyncio.CancelledError): await client_context.__aexit__(None, None, None) def _create_client_context(self): diff --git a/src/mcpo/tests/test_main.py b/src/mcpo/tests/test_main.py index 0170fd93..3cb250eb 100644 --- a/src/mcpo/tests/test_main.py +++ b/src/mcpo/tests/test_main.py @@ -1,8 +1,15 @@ +import asyncio + import pytest +from fastapi import FastAPI, HTTPException +from starlette.requests import Request from pydantic import BaseModel, Field from typing import Any, List, Dict, Union -from mcpo.utils.main import _process_schema_property +from mcp import types +from mcp.types import CallToolResult + +from mcpo.utils.main import _process_schema_property, get_tool_handler _model_cache = {} @@ -325,3 +332,66 @@ def test_ref_to_parent_node(): assert result_type == Any assert result_field.description == "" + + +def make_request(app: FastAPI) -> Request: + return Request( + { + "type": "http", + "app": app, + "method": "POST", + "path": "/test_tool", + "headers": [], + } + ) + + +def test_tool_handler_closes_session_manager_on_cancelled_error(): + class CancelledSessionManager: + def __init__(self): + self.closed = False + + async def ensure_initialized(self): + raise asyncio.CancelledError("sse reconnect cancelled") + + async def close(self): + self.closed = True + + app = FastAPI() + session_manager = CancelledSessionManager() + app.state.session_manager = session_manager + app.state.session = object() + + handler = get_tool_handler("test_tool", {}) + + with pytest.raises(HTTPException) as exc: + asyncio.run(handler(make_request(app))) + + assert exc.value.status_code == 500 + assert exc.value.detail["message"] == "MCP session operation was cancelled" + assert session_manager.closed is True + assert app.state.session is None + + +def test_tool_handler_preserves_http_exception_from_tool_error(): + class ErrorSession: + async def call_tool(self, endpoint_name, arguments): + return CallToolResult( + content=[types.TextContent(type="text", text="tool failed")], + isError=True, + ) + + class SessionManager: + async def ensure_initialized(self): + return ErrorSession(), object() + + app = FastAPI() + app.state.session_manager = SessionManager() + + handler = get_tool_handler("test_tool", {}) + + with pytest.raises(HTTPException) as exc: + asyncio.run(handler(make_request(app))) + + assert exc.value.status_code == 500 + assert exc.value.detail == {"message": "tool failed"} diff --git a/src/mcpo/utils/main.py b/src/mcpo/utils/main.py index d45746ce..15a547de 100644 --- a/src/mcpo/utils/main.py +++ b/src/mcpo/utils/main.py @@ -1,3 +1,4 @@ +import asyncio import logging import json import traceback @@ -34,6 +35,31 @@ logger = logging.getLogger(__name__) +async def _close_session_manager_after_cancel( + request: Request, session_manager: Any, endpoint_name: str +) -> None: + logger.warning( + "MCP session operation for '%s' was cancelled; closing stale session", + endpoint_name, + exc_info=True, + ) + request.app.state.session = None + try: + await session_manager.close() + except asyncio.CancelledError: + logger.warning( + "Cancelled while closing MCP session for '%s'", + endpoint_name, + exc_info=True, + ) + except Exception: + logger.warning( + "Failed to close MCP session after cancellation for '%s'", + endpoint_name, + exc_info=True, + ) + + def normalize_server_type(server_type: str) -> str: """Normalize server_type to a standard value.""" if server_type in ["streamable_http", "streamablehttp", "streamable-http"]: @@ -299,6 +325,17 @@ async def _invoke(session): endpoint_name, ) session, _ = await session_manager.reconnect() + except asyncio.CancelledError as e: + await _close_session_manager_after_cancel( + request, session_manager, endpoint_name + ) + raise HTTPException( + status_code=500, + detail={ + "message": "MCP session operation was cancelled", + "error": str(e), + }, + ) request.app.state.session = session try: @@ -311,6 +348,17 @@ async def _invoke(session): session, _ = await session_manager.reconnect() request.app.state.session = session return await _invoke(session) + except asyncio.CancelledError as e: + await _close_session_manager_after_cancel( + request, session_manager, endpoint_name + ) + raise HTTPException( + status_code=500, + detail={ + "message": "MCP session operation was cancelled", + "error": str(e), + }, + ) session = getattr(request.app.state, "session", None) if not session: @@ -369,6 +417,8 @@ async def tool( ) return final_response + except HTTPException: + raise except McpError as e: logger.info( f"MCP Error calling {endpoint_name}: {traceback.format_exc()}" @@ -424,6 +474,8 @@ async def tool(request: Request): ) return final_response + except HTTPException: + raise except McpError as e: logger.info( f"MCP Error calling {endpoint_name}: {traceback.format_exc()}"