Handle cancelled MCP session operations

src/mcpo/main.py

@@ -125,20 +125,20 @@ 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
 
         reader, writer, *_ = connection
         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):

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"}

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()}"