From 8f8cd406dca22f97ec0b74ddc388b8f43dba5fdf Mon Sep 17 00:00:00 2001
From: bmen25124 <bmen25124@gmail.com>
Date: Thu, 10 Apr 2025 21:29:29 +0300
Subject: [PATCH] Added proper HTTP exceptions

---
 src/mcpo/utils/main.py | 72 ++++++++++++++++++++++++++++++++----------
 1 file changed, 55 insertions(+), 17 deletions(-)

diff --git a/src/mcpo/utils/main.py b/src/mcpo/utils/main.py
index 5992f98..ea53e35 100644
--- a/src/mcpo/utils/main.py
+++ b/src/mcpo/utils/main.py
@@ -2,11 +2,20 @@ from typing import Any, Dict, List, Type, Union, ForwardRef, Optional
 from pydantic import BaseModel, create_model, Field
 from pydantic.fields import FieldInfo
 from mcp import ClientSession, types
-from mcp.types import CallToolResult
+from fastapi import HTTPException
+from mcp.types import CallToolResult, PARSE_ERROR, INVALID_REQUEST, METHOD_NOT_FOUND, INVALID_PARAMS, INTERNAL_ERROR
 from mcp.shared.exceptions import McpError
 
 import json
 
+MCP_ERROR_TO_HTTP_STATUS = {
+    PARSE_ERROR: 400,
+    INVALID_REQUEST: 400,
+    METHOD_NOT_FOUND: 404,
+    INVALID_PARAMS: 422,
+    INTERNAL_ERROR: 500,
+}
+
 
 class ToolResponse(BaseModel):
     response: Optional[Any] = None
@@ -140,10 +149,18 @@ def get_tool_handler(session, endpoint_name, form_model_name, model_fields):
                     result = await session.call_tool(endpoint_name, arguments=args)
 
                     if result.isError:
-                        errorMessage = "Unknown tool execution error"
-                        if result.content and isinstance(result.content[0], types.TextContent):
-                            errorMessage = result.content[0].text
-                        return ToolResponse(errorMessage=errorMessage)
+                        error_message = "Unknown tool execution error"
+                        error_data = None  # Initialize error_data
+                        if result.content:
+                            if isinstance(result.content[0], types.TextContent):
+                                error_message = result.content[0].text
+                        detail = {"message": error_message}
+                        if error_data is not None:
+                            detail["data"] = error_data
+                        raise HTTPException(
+                            status_code=500,
+                            detail=detail,
+                        )
 
                     response_data = process_tool_response(result)
                     final_response = response_data[0] if len(response_data) == 1 else response_data
@@ -151,13 +168,21 @@ def get_tool_handler(session, endpoint_name, form_model_name, model_fields):
 
                 except McpError as e:
                     print(f"MCP Error calling {endpoint_name}: {e.error}")
-                    return ToolResponse(
-                        errorMessage=e.error.message,
-                        errorData=e.error.data,
+                    status_code = MCP_ERROR_TO_HTTP_STATUS.get(e.error.code, 500)
+                    raise HTTPException(
+                        status_code=status_code,
+                        detail=(
+                            {"message": e.error.message, "data": e.error.data}
+                            if e.error.data is not None
+                            else {"message": e.error.message}
+                        ),
                     )
                 except Exception as e:
                     print(f"Unexpected error calling {endpoint_name}: {e}")
-                    return ToolResponse(errorMessage=f"An unexpected internal error occurred: {e}")
+                    raise HTTPException(
+                        status_code=500,
+                        detail={"message": "Unexpected error", "error": str(e)},
+                    )
 
             return tool
 
@@ -176,9 +201,14 @@ def get_tool_handler(session, endpoint_name, form_model_name, model_fields):
 
                     if result.isError:
                         error_message = "Unknown tool execution error"
-                        if result.content and isinstance(result.content[0], types.TextContent):
-                            error_message = result.content[0].text
-                        return ToolResponse(errorMessage=error_message)
+                        if result.content:
+                            if isinstance(result.content[0], types.TextContent):
+                                error_message = result.content[0].text
+                        detail = {"message": error_message}
+                        raise HTTPException(
+                            status_code=500,
+                            detail=detail,
+                        )
 
                     response_data = process_tool_response(result)
                     final_response = response_data[0] if len(response_data) == 1 else response_data
@@ -186,14 +216,22 @@ def get_tool_handler(session, endpoint_name, form_model_name, model_fields):
 
                 except McpError as e:
                     print(f"MCP Error calling {endpoint_name}: {e.error}")
-                    # Propagate the error received from MCP
-                    return ToolResponse(
-                        errorMessage=e.error.message,
-                        errorData=e.error.data,
+                    status_code = MCP_ERROR_TO_HTTP_STATUS.get(e.error.code, 500)
+                    # Propagate the error received from MCP as an HTTP exception
+                    raise HTTPException(
+                        status_code=status_code,
+                        detail=(
+                            {"message": e.error.message, "data": e.error.data}
+                            if e.error.data is not None
+                            else {"message": e.error.message}
+                        ),
                     )
                 except Exception as e:
                     print(f"Unexpected error calling {endpoint_name}: {e}")
-                    return ToolResponse(errorMessage=f"An unexpected internal error occurred: {e}")
+                    raise HTTPException(
+                        status_code=500,
+                        detail={"message": "Unexpected error", "error": str(e)},
+                    )
 
             return tool