scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-13 03:59:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 66

Merge 21efe15ba2 into 87e961fad0

Browse source
This commit is contained in:
ben7am1n 2026-10-08 01:02:19 +00:00 • committed by GitHub
commit 159b62eaac
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 254 additions and 88 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py
View file

@ -28,6 +28,7 @@ from litellm.proxy._experimental.mcp_server.exceptions import (
# which include '/'. Sanitize here so the same regex passes everywhere downstream.
_OPENAPI_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]")
_OPENAPI_TOOL_NAME_MAX_LEN: Final = 128
_MAX_YAML_INT_LENGTH: Final = 1024
def sanitize_openapi_tool_name(raw_name: str) -> str:
@ -147,7 +148,7 @@ def _sanitize_path_parameter_value(param_value: object, param_name: str) -> str:
return quote(value_str, safe="")
def load_openapi_spec(filepath: str) -> dict[str, Any]:
def load_openapi_spec(filepath: str) -> Mapping[str, Any]:
"""
Sync wrapper. For URL specs, use the shared/custom MCP httpx client.
"""
@ -165,6 +166,43 @@ def load_openapi_spec(filepath: str) -> dict[str, Any]:
return asyncio.run(load_openapi_spec_async(filepath))
def _parse_openapi_spec(text: str) -> Mapping[str, Any]:
"""Parse JSON or YAML OpenAPI documents into a mapping."""
try:
parsed: Any = json.loads(text)
except json.JSONDecodeError:
import yaml
from yaml.nodes import ScalarNode
class _NoMergeSafeLoader(yaml.SafeLoader):
def compose_node(self, parent: object, index: object) -> object:
if self.check_event(yaml.events.AliasEvent):
raise yaml.YAMLError("YAML aliases are not supported")
return super().compose_node(parent, index)
def flatten_mapping(self, node: object) -> None:
node_values = getattr(node, "value", ())
if any(key_node.tag == "tag:yaml.org,2002:merge" for key_node, _ in node_values):
raise yaml.YAMLError("YAML merge keys are not supported")
super().flatten_mapping(node)
def construct_yaml_int(self, node: ScalarNode) -> int:
if len(node.value) > _MAX_YAML_INT_LENGTH:
raise yaml.YAMLError("YAML integer is too long")
return super().construct_yaml_int(node)
_NoMergeSafeLoader.add_constructor(
"tag:yaml.org,2002:int",
_NoMergeSafeLoader.construct_yaml_int,
)
parsed = yaml.load(text, Loader=_NoMergeSafeLoader)
if not isinstance(parsed, dict):
raise TypeError("OpenAPI spec must be a JSON or YAML object")
return parsed
async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None) -> dict[str, Any]:
if filepath.startswith("http://") or filepath.startswith("https://"):
client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
@ -174,14 +212,14 @@ async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None
else await async_safe_get(client, filepath, max_response_bytes=max_bytes)
)
r.raise_for_status()
return r.json()
return _parse_openapi_spec(r.text)
# fallback: local file
# Local filesystem path
if not os.path.exists(filepath):
raise FileNotFoundError(f"OpenAPI spec not found at {filepath}")
with open(filepath, "r", encoding="utf-8") as f:
return json.load(f)
return _parse_openapi_spec(f.read())
def get_base_url(spec: Mapping[str, Any], spec_path: str | None = None) -> str:

2
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -4936,7 +4936,7 @@ class TestMCPServerManager:
(
httpx.Response(200, text="secret invalid JSON body"),
"unhealthy",
"OpenAPI specification could not be loaded (JSONDecodeError)",
"OpenAPI specification could not be loaded (TypeError)",
),
],
)

296
tests/unit/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py
View file

@ -17,8 +17,11 @@ import pytest
from fastapi import HTTPException
from respx import MockRouter
from litellm.types.mcp import MCPAuth, MCPAuthType
from litellm.proxy._experimental.mcp_server import openapi_to_mcp_generator as gen
from litellm.proxy._experimental.mcp_server.exceptions import (
MCPOpenApiUpstreamError,
MCPUpstreamAuthError,
)
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
_request_auth_header,
_request_extra_headers,
@ -33,18 +36,15 @@ from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
resolve_operation_params,
)
from litellm.proxy._experimental.mcp_server.tool_outcome import JsonResult, TextResult
from litellm.proxy._experimental.mcp_server.exceptions import (
MCPOpenApiUpstreamError,
MCPUpstreamAuthError,
)
from litellm.types.mcp import MCPAuth, MCPAuthType
GET_ASYNC_CLIENT_TARGET = "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.get_async_httpx_client"
@pytest.mark.asyncio
async def test_unsupported_http_method_returns_text_without_sending_request(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch,
respx_mock: MockRouter,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
tool: Final = create_tool_function("/echo", "HEAD", {}, "https://upstream.example")
@ -58,19 +58,24 @@ async def test_unsupported_http_method_returns_text_without_sending_request(
@pytest.mark.asyncio
@pytest.mark.parametrize("body,expected", [
(' { "ok": true }\n', JsonResult({"ok": True}, ' { "ok": true }\n')),
(' [1, 2]\n', JsonResult([1, 2], ' [1, 2]\n')),
('false', JsonResult(False, 'false')),
('0', JsonResult(0, '0')),
('""', JsonResult("", '""')),
('null', TextResult('null')),
('{"unfinished":', TextResult('{"unfinished":')),
('', TextResult('')),
])
@pytest.mark.parametrize(
"body,expected",
[
(' { "ok": true }\n', JsonResult({"ok": True}, ' { "ok": true }\n')),
(" [1, 2]\n", JsonResult([1, 2], " [1, 2]\n")),
("false", JsonResult(False, "false")),
("0", JsonResult(0, "0")),
('""', JsonResult("", '""')),
("null", TextResult("null")),
('{"unfinished":', TextResult('{"unfinished":')),
("", TextResult("")),
],
)
async def test_http_response_preserves_body_and_classifies_json(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch,
body: str, expected: TextResult | JsonResult,
respx_mock: MockRouter,
monkeypatch: pytest.MonkeyPatch,
body: str,
expected: TextResult | JsonResult,
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
tool: Final = create_tool_function("/echo", "get", {}, "https://upstream.example")
@ -80,23 +85,42 @@ async def test_http_response_preserves_body_and_classifies_json(
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type,value,accepted", [
(MCPAuth.api_key, "Bearer Bearer", False), (MCPAuth.api_key, "ApiKey ApiKey", False),
(MCPAuth.api_key, "token token", False), (MCPAuth.api_key, "bEaReR BEARER", False),
(MCPAuth.api_key, "aPiKeY\tAPIKEY", False), (MCPAuth.api_key, "Bearer fixture-key", True),
(MCPAuth.api_key, "ApiKey fixture-key", True), (MCPAuth.api_key, "token fixture-key", True),
(MCPAuth.authorization, "Bearer", False), (MCPAuth.authorization, "basic", False),
(MCPAuth.authorization, "token", False), (MCPAuth.authorization, "ApiKey", False),
(MCPAuth.authorization, " bEaReR ", False), (MCPAuth.authorization, "\tTOKEN\t", False),
(MCPAuth.authorization, "opaque-secret-value", True), (MCPAuth.authorization, "Bearer abc", True),
(MCPAuth.authorization, "Custom abc", True),
])
@pytest.mark.parametrize(
"auth_type,value,accepted",
[
(MCPAuth.api_key, "Bearer Bearer", False),
(MCPAuth.api_key, "ApiKey ApiKey", False),
(MCPAuth.api_key, "token token", False),
(MCPAuth.api_key, "bEaReR BEARER", False),
(MCPAuth.api_key, "aPiKeY\tAPIKEY", False),
(MCPAuth.api_key, "Bearer fixture-key", True),
(MCPAuth.api_key, "ApiKey fixture-key", True),
(MCPAuth.api_key, "token fixture-key", True),
(MCPAuth.authorization, "Bearer", False),
(MCPAuth.authorization, "basic", False),
(MCPAuth.authorization, "token", False),
(MCPAuth.authorization, "ApiKey", False),
(MCPAuth.authorization, " bEaReR ", False),
(MCPAuth.authorization, "\tTOKEN\t", False),
(MCPAuth.authorization, "opaque-secret-value", True),
(MCPAuth.authorization, "Bearer abc", True),
(MCPAuth.authorization, "Custom abc", True),
],
)
async def test_authorization_validates_credentials_before_http(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, auth_type: MCPAuthType, value: str, accepted: bool,
respx_mock: MockRouter,
monkeypatch: pytest.MonkeyPatch,
auth_type: MCPAuthType,
value: str,
accepted: bool,
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
tool: Final = create_tool_function(
"/echo", "get", {}, "https://upstream.example", auth_type=auth_type,
"/echo",
"get",
{},
"https://upstream.example",
auth_type=auth_type,
)
destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated")
caller_token: Final = _request_auth_header.set(value)
@ -115,20 +139,44 @@ async def test_authorization_validates_credentials_before_http(
@pytest.mark.asyncio
@pytest.mark.parametrize("static,forwarded,caller,resolved,expected", [
({"Authorization": "Bearer configured"}, {"authorization": "Bearer forwarded"}, None, None, "Bearer configured"),
({"Authorization": "Bearer configured"}, None, "Bearer caller", None, "Bearer caller"),
({"Authorization": "Bearer configured"}, None, "Bearer", None, None),
({"Authorization": "Bearer configured"}, None, "Bearer caller", {"authorization": " "}, None),
({"Authorization": "Bearer configured"}, None, "Bearer", {"authorization": "Bearer resolved"}, "Bearer resolved"),
])
@pytest.mark.parametrize(
"static,forwarded,caller,resolved,expected",
[
(
{"Authorization": "Bearer configured"},
{"authorization": "Bearer forwarded"},
None,
None,
"Bearer configured",
),
({"Authorization": "Bearer configured"}, None, "Bearer caller", None, "Bearer caller"),
({"Authorization": "Bearer configured"}, None, "Bearer", None, None),
({"Authorization": "Bearer configured"}, None, "Bearer caller", {"authorization": " "}, None),
(
{"Authorization": "Bearer configured"},
None,
"Bearer",
{"authorization": "Bearer resolved"},
"Bearer resolved",
),
],
)
async def test_static_auth_validates_headers_after_existing_precedence(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch,
static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None,
resolved: dict[str, str] | None, expected: str | None,
respx_mock: MockRouter,
monkeypatch: pytest.MonkeyPatch,
static: dict[str, str],
forwarded: dict[str, str] | None,
caller: str | None,
resolved: dict[str, str] | None,
expected: str | None,
) -> None:
tool: Final = create_tool_function(
"/echo", "get", {}, "https://upstream.example", headers=static, auth_type=MCPAuth.bearer_token,
"/echo",
"get",
{},
"https://upstream.example",
headers=static,
auth_type=MCPAuth.bearer_token,
)
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated")
@ -154,12 +202,19 @@ async def test_static_auth_validates_headers_after_existing_precedence(
@pytest.mark.asyncio
@pytest.mark.parametrize("credential", ["custom-key", ""])
async def test_static_auth_uses_configured_custom_header(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, credential: str,
respx_mock: MockRouter,
monkeypatch: pytest.MonkeyPatch,
credential: str,
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
tool: Final = create_tool_function(
"/echo", "get", {}, "https://upstream.example", headers={"x-custom": credential},
auth_type=MCPAuth.api_key, upstream_token_header="X-Custom",
"/echo",
"get",
{},
"https://upstream.example",
headers={"x-custom": credential},
auth_type=MCPAuth.api_key,
upstream_token_header="X-Custom",
)
destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated")
if credential:
@ -175,11 +230,18 @@ async def test_static_auth_uses_configured_custom_header(
@pytest.mark.asyncio
@pytest.mark.parametrize("credential", ["static-key", ""])
async def test_static_auth_accepts_api_key_carried_by_static_header(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, credential: str,
respx_mock: MockRouter,
monkeypatch: pytest.MonkeyPatch,
credential: str,
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
tool: Final = create_tool_function(
"/echo", "get", {}, "https://upstream.example", headers={"apikey": credential}, auth_type=MCPAuth.api_key,
"/echo",
"get",
{},
"https://upstream.example",
headers={"apikey": credential},
auth_type=MCPAuth.api_key,
)
destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated")
if credential:
@ -193,13 +255,18 @@ async def test_static_auth_accepts_api_key_carried_by_static_header(
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type,resolved", [
(MCPAuth.none, None),
(MCPAuth.oauth2, {"Authorization": "Bearer user-oauth"}),
])
@pytest.mark.parametrize(
"auth_type,resolved",
[
(MCPAuth.none, None),
(MCPAuth.oauth2, {"Authorization": "Bearer user-oauth"}),
],
)
async def test_static_validation_preserves_no_auth_and_resolved_oauth(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch,
auth_type: MCPAuthType, resolved: dict[str, str] | None,
respx_mock: MockRouter,
monkeypatch: pytest.MonkeyPatch,
auth_type: MCPAuthType,
resolved: dict[str, str] | None,
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
tool: Final = create_tool_function("/echo", "get", {}, "https://upstream.example", auth_type=auth_type)
@ -213,6 +280,81 @@ async def test_static_validation_preserves_no_auth_and_resolved_oauth(
_request_resolved_auth_headers.reset(token)
def test_load_openapi_spec_supports_yaml_http_url(monkeypatch: pytest.MonkeyPatch) -> None:
url = "http://example.local/openapi.yaml"
response = SimpleNamespace(
text="""
openapi: 3.0.0
info:
title: YAML API
version: 1.0.0
paths: {}
""",
status_code=200,
headers={},
raise_for_status=lambda: None,
)
client = AsyncMock()
client.get = AsyncMock(return_value=response)
monkeypatch.setattr(gen, "get_async_httpx_client", lambda **kwargs: client)
monkeypatch.setattr(gen, "async_safe_get", lambda client, url, **kwargs: client.get(url))
assert gen.load_openapi_spec(url) == {
"openapi": "3.0.0",
"info": {"title": "YAML API", "version": "1.0.0"},
"paths": {},
}
def test_parse_openapi_spec_rejects_non_object_document() -> None:
with pytest.raises(TypeError, match="OpenAPI spec must be a JSON or YAML object"):
gen._parse_openapi_spec("- not an OpenAPI object\n")
def test_parse_openapi_spec_rejects_yaml_merge_aliases() -> None:
import yaml
with pytest.raises(yaml.YAMLError, match="YAML aliases are not supported"):
gen._parse_openapi_spec(
"""
base: &base
openapi: 3.0.0
info:
title: Base
version: '1.0.0'
paths: {}
merged:
<<: *base
"""
)
def test_parse_openapi_spec_rejects_yaml_aliases() -> None:
import yaml
with pytest.raises(yaml.YAMLError, match="YAML aliases are not supported"):
gen._parse_openapi_spec(
"""
shared: &shared
type: string
components:
schemas:
Request:
properties:
first: *shared
"""
)
def test_parse_openapi_spec_rejects_oversized_yaml_integer() -> None:
import yaml
oversized_integer: Final = "1" + ":11" * gen._MAX_YAML_INT_LENGTH
with pytest.raises(yaml.YAMLError, match="YAML integer is too long"):
gen._parse_openapi_spec(f"openapi: 3.0.0\nx-expensive: {oversized_integer}\n")
def _create_mock_client(method: str, response_text: str, status_code: int = 200) -> AsyncMock:
"""Utility to create a mocked async httpx client for the given method.
@ -263,9 +405,7 @@ class TestCreateToolFunction:
# Verify URL was constructed correctly
call_args = async_client.get.call_args
assert "repository-id" in str(call_args[0][0]) or "test-repo" in str(
call_args[0][0]
)
assert "repository-id" in str(call_args[0][0]) or "test-repo" in str(call_args[0][0])
@pytest.mark.asyncio
async def test_leading_digit_parameter(self):
@ -610,9 +750,7 @@ class TestExtractParameters:
{"name": "filter", "in": "query"},
{"name": "data", "in": "body"},
],
"requestBody": {
"content": {"application/json": {"schema": {"type": "object"}}}
},
"requestBody": {"content": {"application/json": {"schema": {"type": "object"}}}},
}
path_params, query_params, body_params = extract_parameters(operation)
@ -814,7 +952,6 @@ class TestGetBaseUrl:
base_url = get_base_url(spec, spec_path)
assert base_url == "https://production.example.com"
def test_fallback_with_nested_path(self):
"""Test fallback with deeply nested spec path."""
spec = {"openapi": "3.0.0", "paths": {}}
@ -1042,9 +1179,7 @@ class TestResolveOperationParams:
],
}
path_item = {"parameters": path_level_params, "get": operation}
result = resolve_operation_params(
operation, path_item, {"parameters": component_params}
)
result = resolve_operation_params(operation, path_item, {"parameters": component_params})
names = [p["name"] for p in result["parameters"]]
assert "owner" in names
assert "repo" in names
@ -1147,22 +1282,16 @@ class TestRegisterToolsFromOpenAPI:
}
}
openapi_to_mcp_generator.register_tools_from_openapi(
spec, base_url="https://api.example.com"
)
openapi_to_mcp_generator.register_tools_from_openapi(spec, base_url="https://api.example.com")
assert registered, "expected at least one registered tool"
anthropic_re = re.compile(r"^[a-zA-Z0-9_-]{1,128}$")
for name in registered:
assert anthropic_re.match(
name
), f"tool name {name!r} violates ^[a-zA-Z0-9_-]+$"
assert anthropic_re.match(name), f"tool name {name!r} violates ^[a-zA-Z0-9_-]+$"
assert "actions_download-job-logs-for-workflow-run" in registered
assert "pulls_list-files" in registered
def test_missing_operation_id_uses_sanitized_method_path_fallback(
self, monkeypatch
):
def test_missing_operation_id_uses_sanitized_method_path_fallback(self, monkeypatch):
import re
from litellm.proxy._experimental.mcp_server import openapi_to_mcp_generator
@ -1185,15 +1314,11 @@ class TestRegisterToolsFromOpenAPI:
}
}
}
openapi_to_mcp_generator.register_tools_from_openapi(
spec, base_url="https://api.example.com"
)
openapi_to_mcp_generator.register_tools_from_openapi(spec, base_url="https://api.example.com")
assert registered
for name in registered:
assert re.match(
r"^[a-zA-Z0-9_-]+$", name
), f"fallback tool name {name!r} not sanitized"
assert re.match(r"^[a-zA-Z0-9_-]+$", name), f"fallback tool name {name!r} not sanitized"
class TestRequestExtraHeaders:
@ -1349,9 +1474,7 @@ class TestRequestExtraHeaders:
async_client = _create_mock_client("get", "secure-data")
mock_client.return_value = async_client
extra_token = _request_extra_headers.set(
{"Authorization": "Bearer extra", "X-TOKEN": "token-value"}
)
extra_token = _request_extra_headers.set({"Authorization": "Bearer extra", "X-TOKEN": "token-value"})
auth_token = _request_auth_header.set("Bearer byok-credential")
try:
result = await func()
@ -1479,7 +1602,9 @@ class TestUpstreamStatusIsClassified:
@pytest.mark.asyncio
async def test_401_raises_the_reauth_signal_carrying_the_challenge(self):
tool, client = self._tool(401, text='{"error":"invalid_token"}', headers={"www-authenticate": 'Bearer realm="x"'})
tool, client = self._tool(
401, text='{"error":"invalid_token"}', headers={"www-authenticate": 'Bearer realm="x"'}
)
with patch(GET_ASYNC_CLIENT_TARGET, return_value=client):
with pytest.raises(MCPUpstreamAuthError) as exc:
await tool()
@ -1557,6 +1682,7 @@ class TestUpstreamStatusIsClassified:
assert secret_body not in str(exc.value)
assert str(exc.value) == f"upstream returned HTTP {status_code}"
class TestBoundedOpenAPISpecLoading:
@pytest.mark.asyncio
@pytest.mark.parametrize("max_bytes", [12, 13])
@ -1572,6 +1698,7 @@ class TestBoundedOpenAPISpecLoading:
@pytest.mark.parametrize("headers", [{"content-length": "1000000"}, {"content-encoding": "gzip"}])
async def test_unsafe_response_headers_reject_before_reading(self, respx_mock, monkeypatch, headers):
import httpx
from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
load_openapi_spec_async,
@ -1596,6 +1723,7 @@ class TestBoundedOpenAPISpecLoading:
@pytest.mark.asyncio
async def test_chunked_response_is_bounded_and_closed(self, respx_mock, monkeypatch):
import httpx
from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
load_openapi_spec_async,