mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
right error codes for mcp
This commit is contained in:
shivam
parent
229d2008a3
commit
7b234ebfd6
5 changed files with 406 additions and 23 deletions
|
|
@ -77,8 +77,7 @@ def get_proxy_server_request_headers(litellm_params: Optional[dict]) -> dict:
|
|||
if litellm_params is None:
|
||||
return {}
|
||||
|
||||
proxy_request_headers = (
|
||||
litellm_params.get("proxy_server_request", {}).get("headers", {}) or {}
|
||||
)
|
||||
proxy_server_request = litellm_params.get("proxy_server_request") or {}
|
||||
proxy_request_headers = proxy_server_request.get("headers", {}) or {}
|
||||
|
||||
return proxy_request_headers
|
||||
|
|
|
|||
|
|
@ -2192,7 +2192,9 @@ if MCP_AVAILABLE:
|
|||
|
||||
def _get_mcp_servers_in_path(path: str) -> Optional[List[str]]:
|
||||
"""
|
||||
Get the MCP servers from the path
|
||||
Get the MCP servers from the path.
|
||||
Handles both /mcp/<server> (full path) and /<server> (child app receives
|
||||
stripped path when mounted at /mcp).
|
||||
"""
|
||||
import re
|
||||
|
||||
|
|
@ -2234,6 +2236,12 @@ if MCP_AVAILABLE:
|
|||
mcp_servers_from_path = [server_name]
|
||||
else:
|
||||
mcp_servers_from_path = [servers_and_path]
|
||||
else:
|
||||
# Child app receives path like /undefined when mounted at /mcp
|
||||
# Extract first path segment as server name
|
||||
segments = path.strip("/").split("/")
|
||||
if segments and segments[0] and "?" not in segments[0]:
|
||||
mcp_servers_from_path = [segments[0]]
|
||||
return mcp_servers_from_path
|
||||
|
||||
async def extract_mcp_auth_context(scope, path):
|
||||
|
|
@ -2362,6 +2370,8 @@ if MCP_AVAILABLE:
|
|||
scope: Scope, receive: Receive, send: Send
|
||||
) -> None:
|
||||
"""Handle MCP requests through StreamableHTTP."""
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
try:
|
||||
path = scope.get("path", "")
|
||||
(
|
||||
|
|
@ -2382,12 +2392,17 @@ if MCP_AVAILABLE:
|
|||
verbose_logger.debug(
|
||||
f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
|
||||
)
|
||||
# https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response
|
||||
# Validate MCP servers exist and check OAuth requirements
|
||||
for server_name in mcp_servers or []:
|
||||
server = global_mcp_server_manager.get_mcp_server_by_name(
|
||||
server_name, client_ip=_client_ip
|
||||
)
|
||||
if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers:
|
||||
if server is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=f"MCP server '{server_name}' not found",
|
||||
)
|
||||
if server.auth_type == MCPAuth.oauth2 and not oauth2_headers:
|
||||
request = StarletteRequest(scope)
|
||||
base_url = get_request_base_url(request)
|
||||
|
||||
|
|
@ -2442,18 +2457,38 @@ if MCP_AVAILABLE:
|
|||
return
|
||||
|
||||
await session_manager.handle_request(scope, receive, send)
|
||||
except HTTPException:
|
||||
# Re-raise HTTP exceptions to preserve status codes and details
|
||||
raise
|
||||
except HTTPException as e:
|
||||
try:
|
||||
detail = e.detail
|
||||
if isinstance(detail, dict):
|
||||
content = {"error": detail}
|
||||
else:
|
||||
content = {"error": {"message": str(detail)}}
|
||||
error_response = JSONResponse(
|
||||
status_code=e.status_code,
|
||||
content=content,
|
||||
)
|
||||
await error_response(scope, receive, send)
|
||||
except Exception:
|
||||
raise e
|
||||
except ProxyException as e:
|
||||
status_code = int(e.code) if e.code else 500
|
||||
verbose_logger.warning(
|
||||
"MCP auth error (status=%s): %s", status_code, e.message
|
||||
)
|
||||
try:
|
||||
error_response = JSONResponse(
|
||||
status_code=status_code,
|
||||
content={"error": {"message": e.message, "type": e.type}},
|
||||
)
|
||||
await error_response(scope, receive, send)
|
||||
except Exception:
|
||||
raise e
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error handling MCP request: {e}")
|
||||
# Try to send a graceful error response for non-HTTP exceptions
|
||||
try:
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
|
||||
|
||||
error_response = JSONResponse(
|
||||
status_code=HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
status_code=500,
|
||||
content={"error": "MCP request failed", "details": str(e)},
|
||||
)
|
||||
await error_response(scope, receive, send)
|
||||
|
|
@ -2461,11 +2496,12 @@ if MCP_AVAILABLE:
|
|||
verbose_logger.exception(
|
||||
f"Failed to send error response: {response_error}"
|
||||
)
|
||||
# If we can't send a proper response, re-raise the original error
|
||||
raise e
|
||||
|
||||
async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
"""Handle MCP requests through SSE."""
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
try:
|
||||
path = scope.get("path", "")
|
||||
(
|
||||
|
|
@ -2486,6 +2522,16 @@ if MCP_AVAILABLE:
|
|||
verbose_logger.debug(
|
||||
f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
|
||||
)
|
||||
# Validate MCP servers exist (same as handle_streamable_http_mcp)
|
||||
for server_name in mcp_servers or []:
|
||||
server = global_mcp_server_manager.get_mcp_server_by_name(
|
||||
server_name, client_ip=_sse_client_ip
|
||||
)
|
||||
if server is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=f"MCP server '{server_name}' not found",
|
||||
)
|
||||
set_auth_context(
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
mcp_auth_header=mcp_auth_header,
|
||||
|
|
@ -2501,16 +2547,38 @@ if MCP_AVAILABLE:
|
|||
await asyncio.sleep(0.1)
|
||||
|
||||
await sse_session_manager.handle_request(scope, receive, send)
|
||||
except HTTPException as e:
|
||||
try:
|
||||
detail = e.detail
|
||||
if isinstance(detail, dict):
|
||||
content = {"error": detail}
|
||||
else:
|
||||
content = {"error": {"message": str(detail)}}
|
||||
error_response = JSONResponse(
|
||||
status_code=e.status_code,
|
||||
content=content,
|
||||
)
|
||||
await error_response(scope, receive, send)
|
||||
except Exception:
|
||||
raise e
|
||||
except ProxyException as e:
|
||||
status_code = int(e.code) if e.code else 500
|
||||
verbose_logger.warning(
|
||||
"MCP SSE auth error (status=%s): %s", status_code, e.message
|
||||
)
|
||||
try:
|
||||
error_response = JSONResponse(
|
||||
status_code=status_code,
|
||||
content={"error": {"message": e.message, "type": e.type}},
|
||||
)
|
||||
await error_response(scope, receive, send)
|
||||
except Exception:
|
||||
raise e
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error handling MCP request: {e}")
|
||||
# Instead of re-raising, try to send a graceful error response
|
||||
try:
|
||||
# Send a proper HTTP error response instead of letting the exception bubble up
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
|
||||
|
||||
error_response = JSONResponse(
|
||||
status_code=HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
status_code=500,
|
||||
content={"error": "MCP request failed", "details": str(e)},
|
||||
)
|
||||
await error_response(scope, receive, send)
|
||||
|
|
@ -2518,7 +2586,6 @@ if MCP_AVAILABLE:
|
|||
verbose_logger.exception(
|
||||
f"Failed to send error response: {response_error}"
|
||||
)
|
||||
# If we can't send a proper response, re-raise the original error
|
||||
raise e
|
||||
|
||||
app = FastAPI(
|
||||
|
|
|
|||
|
|
@ -75,7 +75,7 @@ class UserAPIKeyAuthExceptionHandler:
|
|||
request=request,
|
||||
use_x_forwarded_for=general_settings.get("use_x_forwarded_for", False),
|
||||
)
|
||||
verbose_proxy_logger.exception(
|
||||
verbose_proxy_logger.warning(
|
||||
"litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - {}\nRequester IP Address:{}".format(
|
||||
str(e),
|
||||
requester_ip,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,39 @@
|
|||
"""
|
||||
Tests for llm_request_utils helper functions.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.litellm_core_utils.llm_request_utils import (
|
||||
get_proxy_server_request_headers,
|
||||
)
|
||||
|
||||
|
||||
class TestGetProxyServerRequestHeaders:
|
||||
"""Tests for get_proxy_server_request_headers."""
|
||||
|
||||
def test_should_return_empty_dict_when_litellm_params_is_none(self):
|
||||
assert get_proxy_server_request_headers(None) == {}
|
||||
|
||||
def test_should_return_empty_dict_when_proxy_server_request_is_none(self):
|
||||
"""When proxy_server_request is explicitly None (e.g. MCP tool calls),
|
||||
should return {} instead of raising 'NoneType has no attribute get'."""
|
||||
litellm_params = {"proxy_server_request": None}
|
||||
assert get_proxy_server_request_headers(litellm_params) == {}
|
||||
|
||||
def test_should_return_empty_dict_when_proxy_server_request_missing(self):
|
||||
litellm_params = {}
|
||||
assert get_proxy_server_request_headers(litellm_params) == {}
|
||||
|
||||
def test_should_return_empty_dict_when_headers_is_none(self):
|
||||
litellm_params = {"proxy_server_request": {"headers": None}}
|
||||
assert get_proxy_server_request_headers(litellm_params) == {}
|
||||
|
||||
def test_should_return_headers_when_present(self):
|
||||
expected = {"Authorization": "Bearer sk-123", "Content-Type": "application/json"}
|
||||
litellm_params = {"proxy_server_request": {"headers": expected}}
|
||||
assert get_proxy_server_request_headers(litellm_params) == expected
|
||||
|
||||
def test_should_return_empty_dict_when_proxy_server_request_has_no_headers_key(self):
|
||||
litellm_params = {"proxy_server_request": {"body": "{}"}}
|
||||
assert get_proxy_server_request_headers(litellm_params) == {}
|
||||
|
|
@ -0,0 +1,278 @@
|
|||
"""
|
||||
Tests for MCP endpoint error handling.
|
||||
|
||||
Verifies that auth failures and other errors in MCP endpoints return proper
|
||||
HTTP status codes (401, 403) instead of 500, and that ProxyException is
|
||||
caught and converted to a JSON error response.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
||||
class TestHandleStreamableHttpMcpErrorHandling:
|
||||
"""Tests that handle_streamable_http_mcp properly handles ProxyException."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_return_401_on_auth_failure(self):
|
||||
"""Auth failures (ProxyException with code=401) should return 401, not 500."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
handle_streamable_http_mcp,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/mcp/undefined",
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
"query_string": b"",
|
||||
"server": ("localhost", 8000),
|
||||
"scheme": "http",
|
||||
}
|
||||
receive = AsyncMock()
|
||||
send = AsyncMock()
|
||||
|
||||
auth_error = ProxyException(
|
||||
message="Malformed API Key passed in. Ensure Key has `Bearer ` prefix.",
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="None",
|
||||
code=401,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=auth_error,
|
||||
):
|
||||
await handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert send.called
|
||||
# First call is http.response.start with status code
|
||||
start_call = send.call_args_list[0]
|
||||
message = start_call[0][0]
|
||||
assert message["type"] == "http.response.start"
|
||||
assert message["status"] == 401
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_return_403_on_forbidden(self):
|
||||
"""ProxyException with code=403 should return 403, not 500."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
handle_streamable_http_mcp,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/mcp/some_server",
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
"query_string": b"",
|
||||
"server": ("localhost", 8000),
|
||||
"scheme": "http",
|
||||
}
|
||||
receive = AsyncMock()
|
||||
send = AsyncMock()
|
||||
|
||||
forbidden_error = ProxyException(
|
||||
message="Access denied",
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="None",
|
||||
code=403,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=forbidden_error,
|
||||
):
|
||||
await handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert send.called
|
||||
start_call = send.call_args_list[0]
|
||||
message = start_call[0][0]
|
||||
assert message["type"] == "http.response.start"
|
||||
assert message["status"] == 403
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_return_error_body_as_json(self):
|
||||
"""ProxyException responses should include error message and type in JSON body."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
handle_streamable_http_mcp,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
import json
|
||||
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/mcp/undefined",
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
"query_string": b"",
|
||||
"server": ("localhost", 8000),
|
||||
"scheme": "http",
|
||||
}
|
||||
receive = AsyncMock()
|
||||
send = AsyncMock()
|
||||
|
||||
auth_error = ProxyException(
|
||||
message="Malformed API Key",
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="None",
|
||||
code=401,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=auth_error,
|
||||
):
|
||||
await handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
# Second call is http.response.body
|
||||
body_call = send.call_args_list[1]
|
||||
body_message = body_call[0][0]
|
||||
assert body_message["type"] == "http.response.body"
|
||||
body = json.loads(body_message["body"])
|
||||
assert body["error"]["message"] == "Malformed API Key"
|
||||
assert body["error"]["type"] == ProxyErrorTypes.auth_error
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_return_404_for_nonexistent_mcp_server(self):
|
||||
"""Non-existent MCP server names should return 404, not 200."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
handle_streamable_http_mcp,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/undefined",
|
||||
"headers": [
|
||||
(b"content-type", b"application/json"),
|
||||
(b"accept", b"application/json, text/event-stream"),
|
||||
(b"authorization", b"Bearer sk-1620"),
|
||||
],
|
||||
"query_string": b"",
|
||||
"server": ("localhost", 8000),
|
||||
"scheme": "http",
|
||||
}
|
||||
receive = AsyncMock()
|
||||
send = AsyncMock()
|
||||
|
||||
mock_auth = MagicMock()
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(mock_auth, None, ["undefined"], {}, None, {}),
|
||||
), patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager"
|
||||
) as mock_mgr:
|
||||
mock_mgr.get_mcp_server_by_name.return_value = None
|
||||
|
||||
await handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert send.called
|
||||
start_call = send.call_args_list[0]
|
||||
message = start_call[0][0]
|
||||
assert message["type"] == "http.response.start"
|
||||
assert message["status"] == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_still_500_on_unexpected_exceptions(self):
|
||||
"""Non-ProxyException errors should still result in a 500 response."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
handle_streamable_http_mcp,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/mcp/test",
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
"query_string": b"",
|
||||
"server": ("localhost", 8000),
|
||||
"scheme": "http",
|
||||
}
|
||||
receive = AsyncMock()
|
||||
send = AsyncMock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=RuntimeError("unexpected crash"),
|
||||
):
|
||||
await handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert send.called
|
||||
start_call = send.call_args_list[0]
|
||||
message = start_call[0][0]
|
||||
assert message["type"] == "http.response.start"
|
||||
assert message["status"] == 500
|
||||
|
||||
|
||||
class TestHandleSseMcpErrorHandling:
|
||||
"""Tests that handle_sse_mcp properly handles ProxyException."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_return_401_on_auth_failure(self):
|
||||
"""Auth failures in SSE handler should return 401, not 500."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
handle_sse_mcp,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "GET",
|
||||
"path": "/sse",
|
||||
"headers": [(b"accept", b"text/event-stream")],
|
||||
"query_string": b"",
|
||||
"server": ("localhost", 8000),
|
||||
"scheme": "http",
|
||||
}
|
||||
receive = AsyncMock()
|
||||
send = AsyncMock()
|
||||
|
||||
auth_error = ProxyException(
|
||||
message="Malformed API Key passed in.",
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="None",
|
||||
code=401,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=auth_error,
|
||||
):
|
||||
await handle_sse_mcp(scope, receive, send)
|
||||
|
||||
assert send.called
|
||||
start_call = send.call_args_list[0]
|
||||
message = start_call[0][0]
|
||||
assert message["type"] == "http.response.start"
|
||||
assert message["status"] == 401
|
||||
Loading…
Add table
Reference in a new issue