mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
greptile test fix
This commit is contained in:
shivam
parent
7f46db445b
commit
d16c1c32bb
2 changed files with 131 additions and 36 deletions
|
|
@ -2366,6 +2366,22 @@ if MCP_AVAILABLE:
|
|||
]
|
||||
return False
|
||||
|
||||
async def _send_error_response(
|
||||
scope: Scope,
|
||||
receive: Receive,
|
||||
send: Send,
|
||||
status_code: int,
|
||||
content: dict,
|
||||
headers: Optional[dict] = None,
|
||||
) -> None:
|
||||
"""Send a JSON error response. Shared by handle_streamable_http_mcp and handle_sse_mcp."""
|
||||
response = JSONResponse(
|
||||
status_code=status_code,
|
||||
content=content,
|
||||
headers=headers or {},
|
||||
)
|
||||
await response(scope, receive, send)
|
||||
|
||||
async def handle_streamable_http_mcp(
|
||||
scope: Scope, receive: Receive, send: Send
|
||||
) -> None:
|
||||
|
|
@ -2460,17 +2476,15 @@ if MCP_AVAILABLE:
|
|||
except HTTPException as e:
|
||||
try:
|
||||
detail = e.detail
|
||||
if isinstance(detail, dict):
|
||||
content = {"error": detail}
|
||||
else:
|
||||
content = {"error": {"message": str(detail)}}
|
||||
headers = dict(e.headers) if e.headers else {}
|
||||
error_response = JSONResponse(
|
||||
status_code=e.status_code,
|
||||
content=content,
|
||||
headers=headers,
|
||||
content = (
|
||||
{"error": detail}
|
||||
if isinstance(detail, dict)
|
||||
else {"error": {"message": str(detail)}}
|
||||
)
|
||||
headers = dict(e.headers) if e.headers else {}
|
||||
await _send_error_response(
|
||||
scope, receive, send, e.status_code, content, headers
|
||||
)
|
||||
await error_response(scope, receive, send)
|
||||
except Exception:
|
||||
raise e
|
||||
except ProxyException as e:
|
||||
|
|
@ -2484,21 +2498,25 @@ if MCP_AVAILABLE:
|
|||
"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 _send_error_response(
|
||||
scope,
|
||||
receive,
|
||||
send,
|
||||
status_code,
|
||||
{"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:
|
||||
error_response = JSONResponse(
|
||||
status_code=500,
|
||||
content={"error": "MCP request failed", "details": str(e)},
|
||||
await _send_error_response(
|
||||
scope,
|
||||
receive,
|
||||
send,
|
||||
500,
|
||||
{"error": "MCP request failed", "details": str(e)},
|
||||
)
|
||||
await error_response(scope, receive, send)
|
||||
except Exception as response_error:
|
||||
verbose_logger.exception(
|
||||
f"Failed to send error response: {response_error}"
|
||||
|
|
@ -2571,17 +2589,15 @@ if MCP_AVAILABLE:
|
|||
except HTTPException as e:
|
||||
try:
|
||||
detail = e.detail
|
||||
if isinstance(detail, dict):
|
||||
content = {"error": detail}
|
||||
else:
|
||||
content = {"error": {"message": str(detail)}}
|
||||
headers = dict(e.headers) if e.headers else {}
|
||||
error_response = JSONResponse(
|
||||
status_code=e.status_code,
|
||||
content=content,
|
||||
headers=headers,
|
||||
content = (
|
||||
{"error": detail}
|
||||
if isinstance(detail, dict)
|
||||
else {"error": {"message": str(detail)}}
|
||||
)
|
||||
headers = dict(e.headers) if e.headers else {}
|
||||
await _send_error_response(
|
||||
scope, receive, send, e.status_code, content, headers
|
||||
)
|
||||
await error_response(scope, receive, send)
|
||||
except Exception:
|
||||
raise e
|
||||
except ProxyException as e:
|
||||
|
|
@ -2595,21 +2611,25 @@ if MCP_AVAILABLE:
|
|||
"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 _send_error_response(
|
||||
scope,
|
||||
receive,
|
||||
send,
|
||||
status_code,
|
||||
{"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:
|
||||
error_response = JSONResponse(
|
||||
status_code=500,
|
||||
content={"error": "MCP request failed", "details": str(e)},
|
||||
await _send_error_response(
|
||||
scope,
|
||||
receive,
|
||||
send,
|
||||
500,
|
||||
{"error": "MCP request failed", "details": str(e)},
|
||||
)
|
||||
await error_response(scope, receive, send)
|
||||
except Exception as response_error:
|
||||
verbose_logger.exception(
|
||||
f"Failed to send error response: {response_error}"
|
||||
|
|
|
|||
|
|
@ -276,3 +276,78 @@ class TestHandleSseMcpErrorHandling:
|
|||
message = start_call[0][0]
|
||||
assert message["type"] == "http.response.start"
|
||||
assert message["status"] == 401
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_return_404_for_nonexistent_mcp_server(self):
|
||||
"""Non-existent MCP server names in SSE handler should return 404, not 200."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
handle_sse_mcp,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
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()
|
||||
|
||||
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_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"] == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_still_500_on_unexpected_exceptions(self):
|
||||
"""Non-ProxyException errors in SSE handler should still result in a 500 response."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
handle_sse_mcp,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
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()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=RuntimeError("unexpected crash"),
|
||||
):
|
||||
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"] == 500
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue