Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions litellm/proxy/_experimental/mcp_server/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,13 @@
_MCP_PROTOCOL_VERSION_HEADER: Final = b"mcp-protocol-version"


def reject_disallowed_mcp_origin(request: StarletteRequest) -> None:
from litellm.proxy.proxy_server import origins # noqa: PLC0415 # proxy imports this module during startup

if "*" not in origins and any(origin not in origins for origin in request.headers.getlist("origin")):
raise HTTPException(status_code=403, detail="Invalid Origin header")


def unsupported_protocol_version(scope: Scope) -> str | None:
"""Return the unsupported ``MCP-Protocol-Version`` header value, if any.

Expand Down Expand Up @@ -1931,6 +1938,7 @@ async def _check_passthrough_upstream_auth(
async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None:
"""Handle MCP requests through StreamableHTTP."""
try:
reject_disallowed_mcp_origin(StarletteRequest(scope))
bad_version: Final = unsupported_protocol_version(scope)
if bad_version is not None:
supported: Final = ", ".join(sorted(HANDSHAKE_PROTOCOL_VERSIONS))
Expand Down Expand Up @@ -2275,6 +2283,7 @@ async def _dispatch() -> None:
async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None:
"""Handle MCP requests through SSE."""
try:
reject_disallowed_mcp_origin(StarletteRequest(scope))
bad_version: Final = unsupported_protocol_version(scope)
if bad_version is not None:
supported: Final = ", ".join(sorted(HANDSHAKE_PROTOCOL_VERSIONS))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10201,6 +10201,64 @@ async def test_active_request_ctx_var_feeds_get_current_session(_mcp_request_ctx
assert _get_current_session() is None


@pytest.mark.asyncio
@pytest.mark.parametrize(
("method", "path", "session_headers"),
(
("POST", "/mcp", ()),
("GET", "/mcp", (("mcp-session-id", "existing-session"),)),
("DELETE", "/mcp", (("mcp-session-id", "existing-session"),)),
("POST", "/server/mcp", ()),
("GET", "/sse", ()),
("POST", "/sse/messages", ()),
),
)
@pytest.mark.parametrize(
("allowed_origins", "origin_headers", "expected_status"),
(
(("https://allowed.example",), (("origin", "https://evil.example"),), 403),
(("https://allowed.example",), (("origin", "https://allowed.example.evil.example"),), 403),
(("https://allowed.example",), (("origin", "null"),), 403),
(("https://allowed.example",), (("origin", ""),), 403),
(
("https://allowed.example",),
(("origin", "https://allowed.example"), ("origin", "https://evil.example")),
403,
),
(("https://allowed.example",), (("origin", "https://allowed.example"),), 401),
(("https://allowed.example",), (), 401),
(("*",), (("origin", "https://another.example"),), 401),
),
)
async def test_mcp_origin_admission_precedes_authentication(
method: str,
path: str,
session_headers: tuple[tuple[str, str], ...],
allowed_origins: tuple[str, ...],
origin_headers: tuple[tuple[str, str], ...],
expected_status: int,
) -> None:
import httpx

from litellm.proxy._experimental.mcp_server import server

authenticate: Final = AsyncMock(side_effect=HTTPException(status_code=401, detail="authentication required"))
with (
patch("litellm.proxy.proxy_server.origins", allowed_origins),
patch.object(server, "extract_mcp_auth_context", authenticate),
):
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server.app), base_url="http://gateway") as client:
response: Final = await client.request(method, path, headers=(*session_headers, *origin_headers))

assert response.status_code == expected_status
if expected_status == 403:
assert response.json() == {"detail": "Invalid Origin header"}
authenticate.assert_not_awaited()
else:
assert response.json() == {"detail": "authentication required"}
authenticate.assert_awaited_once()


@pytest.mark.asyncio
async def test_active_request_ctx_var_feeds_auth_resolution_recording(_mcp_request_ctx) -> None:
from starlette.requests import Request
Expand Down
Loading