mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
feat(mcp): allowlist client applications for MCP gateway access
Adds the mcp_allowed_clients general setting, enforced against the clientInfo.name each MCP client sends in its initialize request. A client not on the list, or one that does not identify itself, is rejected with 403 before any stateful session is created. The setting is configurable from config.yaml and from the Admin UI MCP network settings page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
d5b8400aa9
commit
47be6c8aeb
9 changed files with 745 additions and 53 deletions
67
litellm/proxy/_experimental/mcp_server/client_allowlist.py
Normal file
67
litellm/proxy/_experimental/mcp_server/client_allowlist.py
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
"""
|
||||
Gateway-level allowlist of MCP client applications, matched against the
|
||||
``clientInfo.name`` a client sends in its JSON-RPC ``initialize`` request. The
|
||||
name is client-supplied, so this is a policy control and not a security boundary.
|
||||
"""
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
MCP_ALLOWED_CLIENTS_SETTING: Final = "mcp_allowed_clients"
|
||||
|
||||
_ALLOWED_CLIENTS_ADAPTER: Final = TypeAdapter(list[str])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MCPClientRejection:
|
||||
client_name: str | None
|
||||
|
||||
@property
|
||||
def details(self) -> str:
|
||||
if self.client_name is None:
|
||||
return (
|
||||
"MCP initialize request did not identify the client application (clientInfo.name). "
|
||||
f"This gateway only admits clients listed in {MCP_ALLOWED_CLIENTS_SETTING}."
|
||||
)
|
||||
return f"MCP client '{self.client_name}' is not listed in this gateway's {MCP_ALLOWED_CLIENTS_SETTING}."
|
||||
|
||||
|
||||
def parse_allowed_mcp_clients(raw_setting: object) -> frozenset[str] | None:
|
||||
"""None when the setting is absent (not enforced). A malformed setting admits nobody."""
|
||||
if raw_setting is None:
|
||||
return None
|
||||
try:
|
||||
return frozenset(_ALLOWED_CLIENTS_ADAPTER.validate_python(raw_setting))
|
||||
except ValidationError:
|
||||
verbose_logger.warning(
|
||||
"%s is not a list of client names (%r); rejecting every MCP client until it is fixed",
|
||||
MCP_ALLOWED_CLIENTS_SETTING,
|
||||
raw_setting,
|
||||
)
|
||||
return frozenset()
|
||||
|
||||
|
||||
def extract_mcp_client_name(body: bytes) -> str | None:
|
||||
try:
|
||||
data: Final = json.loads(body)
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
return None
|
||||
params: Final = data.get("params") if isinstance(data, dict) else None
|
||||
client_info: Final = params.get("clientInfo") if isinstance(params, dict) else None
|
||||
name: Final = client_info.get("name") if isinstance(client_info, dict) else None
|
||||
return name if isinstance(name, str) and name else None
|
||||
|
||||
|
||||
def check_mcp_client_allowed(body: bytes, allowed_clients: frozenset[str] | None) -> MCPClientRejection | None:
|
||||
"""None when the initialize is admitted, otherwise the rejection to send back as a 403."""
|
||||
if allowed_clients is None:
|
||||
return None
|
||||
client_name: Final = extract_mcp_client_name(body)
|
||||
if client_name is not None and client_name in allowed_clients:
|
||||
return None
|
||||
return MCPClientRejection(client_name=client_name)
|
||||
|
|
@ -463,6 +463,11 @@ if MCP_AVAILABLE:
|
|||
from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import (
|
||||
MCPAuthenticatedUser,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.client_allowlist import (
|
||||
MCP_ALLOWED_CLIENTS_SETTING,
|
||||
check_mcp_client_allowed,
|
||||
parse_allowed_mcp_clients,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
|
||||
SERVER_OUTCOMES_META_KEY,
|
||||
AggregateToolListing,
|
||||
|
|
@ -3810,6 +3815,43 @@ if MCP_AVAILABLE:
|
|||
except (json.JSONDecodeError, TypeError):
|
||||
return False
|
||||
|
||||
def _load_allowed_mcp_clients() -> frozenset[str] | None:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
return parse_allowed_mcp_clients(general_settings.get(MCP_ALLOWED_CLIENTS_SETTING))
|
||||
|
||||
async def _reject_initialize_from_disallowed_client(
|
||||
scope: Scope,
|
||||
receive: Receive,
|
||||
send: Send,
|
||||
body: bytes,
|
||||
client_ip: str | None,
|
||||
) -> bool:
|
||||
"""Send a 403 and return True when the initialize body names a client the gateway does not admit."""
|
||||
rejection: Final = check_mcp_client_allowed(body, _load_allowed_mcp_clients())
|
||||
if rejection is None:
|
||||
return False
|
||||
verbose_logger.warning(
|
||||
"Rejecting MCP initialize from client %r (ip=%s): not listed in %s",
|
||||
rejection.client_name,
|
||||
client_ip,
|
||||
MCP_ALLOWED_CLIENTS_SETTING,
|
||||
)
|
||||
forbidden: Final = JSONResponse(
|
||||
status_code=403,
|
||||
content={"error": "Forbidden", "details": rejection.details},
|
||||
)
|
||||
await forbidden(scope, receive, send)
|
||||
return True
|
||||
|
||||
def _replay_consumed_messages(consumed_messages: list[Message], receive: Receive) -> Receive:
|
||||
async def wrapped_receive() -> Message:
|
||||
if consumed_messages:
|
||||
return consumed_messages.pop(0)
|
||||
return await receive()
|
||||
|
||||
return wrapped_receive
|
||||
|
||||
async def _read_request_body_for_routing(
|
||||
receive: Receive,
|
||||
) -> tuple[list[Message], bytes]:
|
||||
|
|
@ -4510,6 +4552,10 @@ if MCP_AVAILABLE:
|
|||
if scope.get("method") == "POST":
|
||||
consumed_messages, body = await _read_request_body_for_routing(receive)
|
||||
is_initialize = _is_initialize_request(body)
|
||||
if is_initialize and await _reject_initialize_from_disallowed_client(
|
||||
scope, receive, send, body, _client_ip
|
||||
):
|
||||
return
|
||||
|
||||
use_stateful: Final = bool(session_id or is_initialize)
|
||||
target_manager: Final = session_manager_stateful if use_stateful else session_manager_stateless
|
||||
|
|
@ -4540,15 +4586,8 @@ if MCP_AVAILABLE:
|
|||
return
|
||||
|
||||
# Replay body messages if we consumed them for peeking
|
||||
original_receive: Final = receive
|
||||
if consumed_messages:
|
||||
|
||||
async def wrapped_receive():
|
||||
if consumed_messages:
|
||||
return consumed_messages.pop(0)
|
||||
return await original_receive()
|
||||
|
||||
receive = wrapped_receive
|
||||
receive = _replay_consumed_messages(consumed_messages, receive)
|
||||
|
||||
# Serialize requests on the same stateful session so concurrent
|
||||
# callers don't clobber each other's auth context mid-flight.
|
||||
|
|
@ -4785,6 +4824,15 @@ if MCP_AVAILABLE:
|
|||
await initialize_session_managers()
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
sse_consumed_messages, sse_body = (
|
||||
await _read_request_body_for_routing(receive) if scope.get("method") == "POST" else ([], b"")
|
||||
)
|
||||
if _is_initialize_request(sse_body) and await _reject_initialize_from_disallowed_client(
|
||||
scope, receive, send, sse_body, _sse_client_ip
|
||||
):
|
||||
return
|
||||
sse_receive: Final = _replay_consumed_messages(sse_consumed_messages, receive)
|
||||
|
||||
async with _gateway_initialize_instructions_request_scope(
|
||||
user_api_key_auth,
|
||||
mcp_servers,
|
||||
|
|
@ -4792,7 +4840,7 @@ if MCP_AVAILABLE:
|
|||
scoped_server_endpoint=scoped_server_endpoint,
|
||||
is_initialize=scope.get("method") == "GET",
|
||||
):
|
||||
await sse_session_manager.handle_request(scope, receive, send)
|
||||
await sse_session_manager.handle_request(scope, sse_receive, send)
|
||||
except MCPUpstreamAuthError as e:
|
||||
# Upstream delegated auth returned 401; surface it to the client so
|
||||
# standards-compliant MCP clients trigger the upstream OAuth flow.
|
||||
|
|
|
|||
|
|
@ -2853,6 +2853,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
None,
|
||||
description="Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).",
|
||||
)
|
||||
mcp_allowed_clients: list[str] | None = Field(
|
||||
None,
|
||||
description="MCP client applications admitted by the gateway, matched exactly against the clientInfo.name the client sends in its initialize request (for example 'claude-code'). When set, an initialize from any other client, or one that does not identify itself, is rejected with 403. Unset means every client is admitted. The name is client-supplied, so this is a policy control rather than a security boundary.",
|
||||
)
|
||||
mcp_trusted_proxy_ranges: list[str] | None = Field(
|
||||
None,
|
||||
description="CIDR ranges of trusted reverse proxies. When set, X-Forwarded-For and X-Forwarded-* origin headers are only trusted from these IPs.",
|
||||
|
|
|
|||
|
|
@ -7179,6 +7179,9 @@ class ProxyConfig:
|
|||
"enable_openai_websocket_passthrough"
|
||||
)
|
||||
|
||||
if "mcp_allowed_clients" not in self._yaml_general_settings_keys:
|
||||
general_settings["mcp_allowed_clients"] = _general_settings.get("mcp_allowed_clients")
|
||||
|
||||
if "user_api_key_cache_max_size" not in self._yaml_general_settings_keys:
|
||||
db_cache_max_size: Final = _general_settings.get("user_api_key_cache_max_size")
|
||||
try:
|
||||
|
|
@ -17137,6 +17140,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
|
|||
"maximum_spend_logs_cleanup_run_budget": "String",
|
||||
"maximum_spend_logs_cleanup_batch_timeout": "String",
|
||||
"mcp_internal_ip_ranges": "List",
|
||||
"mcp_allowed_clients": "List",
|
||||
"mcp_trusted_proxy_ranges": "List",
|
||||
"mcp_xff_num_trusted_hops": "Integer",
|
||||
"always_include_stream_usage": "Boolean",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,105 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.client_allowlist import (
|
||||
MCP_ALLOWED_CLIENTS_SETTING,
|
||||
MCPClientRejection,
|
||||
check_mcp_client_allowed,
|
||||
extract_mcp_client_name,
|
||||
parse_allowed_mcp_clients,
|
||||
)
|
||||
|
||||
|
||||
def _initialize_body(client_info: object) -> bytes:
|
||||
return json.dumps(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 0,
|
||||
"method": "initialize",
|
||||
"params": {"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": client_info},
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
CLAUDE_CODE: Final = _initialize_body({"name": "claude-code", "version": "2.1.274"})
|
||||
ANTIGRAVITY: Final = _initialize_body({"name": "antigravity-cli", "version": "1.0.0"})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("raw_setting", "expected"),
|
||||
(
|
||||
(None, None),
|
||||
([], frozenset()),
|
||||
(["antigravity-cli"], frozenset({"antigravity-cli"})),
|
||||
(["antigravity-cli", "codex-mcp-client"], frozenset({"antigravity-cli", "codex-mcp-client"})),
|
||||
("antigravity-cli", frozenset()),
|
||||
([1, "antigravity-cli"], frozenset()),
|
||||
({"name": "antigravity-cli"}, frozenset()),
|
||||
),
|
||||
)
|
||||
def test_parse_allowed_mcp_clients(raw_setting: object, expected: frozenset[str] | None) -> None:
|
||||
assert parse_allowed_mcp_clients(raw_setting) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("body", "expected"),
|
||||
(
|
||||
(CLAUDE_CODE, "claude-code"),
|
||||
(_initialize_body({"name": "", "version": "1"}), None),
|
||||
(_initialize_body({"version": "1"}), None),
|
||||
(_initialize_body({"name": 7}), None),
|
||||
(_initialize_body("claude-code"), None),
|
||||
(b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', None),
|
||||
(b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":[]}', None),
|
||||
(b'["not", "an", "object"]', None),
|
||||
(b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"clientInfo":{"name":"clau', None),
|
||||
(b"\xff\xfe", None),
|
||||
(b"", None),
|
||||
),
|
||||
)
|
||||
def test_extract_mcp_client_name(body: bytes, expected: str | None) -> None:
|
||||
assert extract_mcp_client_name(body) == expected
|
||||
|
||||
|
||||
def test_unconfigured_allowlist_admits_every_client_including_unidentified_ones() -> None:
|
||||
assert check_mcp_client_allowed(CLAUDE_CODE, None) is None
|
||||
assert check_mcp_client_allowed(b'{"method":"initialize","params":{}}', None) is None
|
||||
assert check_mcp_client_allowed(b"garbage", None) is None
|
||||
|
||||
|
||||
def test_listed_client_is_admitted_and_unlisted_client_is_rejected_by_name() -> None:
|
||||
allowed: Final = frozenset({"antigravity-cli"})
|
||||
assert check_mcp_client_allowed(ANTIGRAVITY, allowed) is None
|
||||
assert check_mcp_client_allowed(CLAUDE_CODE, allowed) == MCPClientRejection(client_name="claude-code")
|
||||
|
||||
|
||||
def test_matching_is_exact_not_prefix_or_case_insensitive() -> None:
|
||||
allowed: Final = frozenset({"claude-code"})
|
||||
assert check_mcp_client_allowed(_initialize_body({"name": "Claude-Code"}), allowed) is not None
|
||||
assert check_mcp_client_allowed(_initialize_body({"name": "claude-code-sdk"}), allowed) is not None
|
||||
assert check_mcp_client_allowed(_initialize_body({"name": " claude-code"}), allowed) is not None
|
||||
|
||||
|
||||
def test_empty_allowlist_rejects_every_client() -> None:
|
||||
assert check_mcp_client_allowed(ANTIGRAVITY, frozenset()) == MCPClientRejection(client_name="antigravity-cli")
|
||||
assert check_mcp_client_allowed(CLAUDE_CODE, frozenset()) == MCPClientRejection(client_name="claude-code")
|
||||
|
||||
|
||||
def test_missing_or_malformed_client_metadata_is_rejected_when_allowlist_is_set() -> None:
|
||||
allowed: Final = frozenset({"antigravity-cli"})
|
||||
assert check_mcp_client_allowed(_initialize_body({"version": "1"}), allowed) == MCPClientRejection(None)
|
||||
assert check_mcp_client_allowed(b'{"method":"initialize","params":{}}', allowed) == MCPClientRejection(None)
|
||||
assert check_mcp_client_allowed(b"{not json", allowed) == MCPClientRejection(None)
|
||||
|
||||
|
||||
def test_rejection_details_name_the_setting_and_the_offending_client() -> None:
|
||||
named: Final = MCPClientRejection(client_name="claude-code").details
|
||||
assert "claude-code" in named
|
||||
assert MCP_ALLOWED_CLIENTS_SETTING in named
|
||||
|
||||
anonymous: Final = MCPClientRejection(client_name=None).details
|
||||
assert "clientInfo.name" in anonymous
|
||||
assert MCP_ALLOWED_CLIENTS_SETTING in anonymous
|
||||
assert "None" not in anonymous
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import contextvars
|
||||
import os
|
||||
from datetime import datetime, timedelta
|
||||
|
|
@ -2035,6 +2036,249 @@ async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless(
|
|||
assert not any(name.startswith(b"x-mcp-debug") for name in headers)
|
||||
|
||||
|
||||
_CLAUDE_CODE_INITIALIZE: Final = (
|
||||
b'{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-06-18",'
|
||||
b'"capabilities":{},"clientInfo":{"name":"claude-code","version":"2.1.274"}}}'
|
||||
)
|
||||
_ANTIGRAVITY_INITIALIZE: Final = (
|
||||
b'{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-06-18",'
|
||||
b'"capabilities":{},"clientInfo":{"name":"antigravity-cli","version":"1.0.0"}}}'
|
||||
)
|
||||
_ANONYMOUS_INITIALIZE: Final = b'{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"capabilities":{}}}'
|
||||
_TOOLS_LIST: Final = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}'
|
||||
|
||||
|
||||
async def _drain_body(receive) -> bytes:
|
||||
chunks: list[bytes] = []
|
||||
while True:
|
||||
message = await receive()
|
||||
chunks.append(message.get("body", b""))
|
||||
if not message.get("more_body", False):
|
||||
return b"".join(chunks)
|
||||
|
||||
|
||||
def _forbidden_client_response(send: AsyncMock) -> tuple[int, dict[str, str]]:
|
||||
import json as _json
|
||||
|
||||
start: Final = send.call_args_list[0].args[0]
|
||||
body: Final = b"".join(call.args[0].get("body", b"") for call in send.call_args_list[1:])
|
||||
return start["status"], _json.loads(body)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _client_allowlist_patches(allowed_clients: object):
|
||||
settings: Final = {} if allowed_clients is None else {"mcp_allowed_clients": allowed_clients}
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(UserAPIKeyAuth(user_id="allowlist-user"), None, None, None, None, {}),
|
||||
),
|
||||
patch("litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True),
|
||||
patch("litellm.proxy.proxy_server.general_settings", settings),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("request_body", "expected_details"),
|
||||
(
|
||||
(
|
||||
_CLAUDE_CODE_INITIALIZE,
|
||||
"MCP client 'claude-code' is not listed in this gateway's mcp_allowed_clients.",
|
||||
),
|
||||
(
|
||||
_ANONYMOUS_INITIALIZE,
|
||||
"MCP initialize request did not identify the client application (clientInfo.name). "
|
||||
"This gateway only admits clients listed in mcp_allowed_clients.",
|
||||
),
|
||||
),
|
||||
)
|
||||
async def test_streamable_http_rejects_initialize_from_unlisted_client_before_session_creation(
|
||||
request_body: bytes, expected_details: str
|
||||
) -> None:
|
||||
from starlette.types import Scope
|
||||
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
|
||||
scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
|
||||
receive: Final = AsyncMock(return_value={"type": "http.request", "body": request_body, "more_body": False})
|
||||
send: Final = AsyncMock()
|
||||
stateful_handle: Final = AsyncMock()
|
||||
stateless_handle: Final = AsyncMock()
|
||||
session_cap: Final = AsyncMock(return_value=True)
|
||||
|
||||
with (
|
||||
_client_allowlist_patches(["antigravity-cli"]),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.session_manager_stateful",
|
||||
SimpleNamespace(handle_request=stateful_handle),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.session_manager_stateless",
|
||||
SimpleNamespace(handle_request=stateless_handle),
|
||||
),
|
||||
patch("litellm.proxy._experimental.mcp_server.server._enforce_stateful_session_cap_for_owner", session_cap),
|
||||
):
|
||||
await mcp_module.handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert _forbidden_client_response(send) == (403, {"error": "Forbidden", "details": expected_details})
|
||||
stateful_handle.assert_not_awaited()
|
||||
stateless_handle.assert_not_awaited()
|
||||
session_cap.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("allowed_clients", "request_body"),
|
||||
(
|
||||
(["antigravity-cli"], _ANTIGRAVITY_INITIALIZE),
|
||||
(["claude-code", "antigravity-cli"], _CLAUDE_CODE_INITIALIZE),
|
||||
(None, _CLAUDE_CODE_INITIALIZE),
|
||||
(None, _ANONYMOUS_INITIALIZE),
|
||||
),
|
||||
)
|
||||
async def test_streamable_http_admits_listed_or_unrestricted_initialize_and_replays_body(
|
||||
allowed_clients: list[str] | None, request_body: bytes
|
||||
) -> None:
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
|
||||
scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
|
||||
receive: Final = AsyncMock(
|
||||
side_effect=[
|
||||
{"type": "http.request", "body": request_body[:20], "more_body": True},
|
||||
{"type": "http.request", "body": request_body[20:], "more_body": False},
|
||||
]
|
||||
)
|
||||
send: Final = AsyncMock()
|
||||
downstream_bodies: Final[list[bytes]] = []
|
||||
|
||||
async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None:
|
||||
downstream_bodies.append(await _drain_body(downstream_receive))
|
||||
|
||||
stateful_handle: Final = AsyncMock(side_effect=handle_request)
|
||||
stateless_handle: Final = AsyncMock()
|
||||
|
||||
with (
|
||||
_client_allowlist_patches(allowed_clients),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.session_manager_stateful",
|
||||
SimpleNamespace(handle_request=stateful_handle),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.session_manager_stateless",
|
||||
SimpleNamespace(handle_request=stateless_handle),
|
||||
),
|
||||
):
|
||||
await mcp_module.handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert downstream_bodies == [request_body]
|
||||
stateless_handle.assert_not_awaited()
|
||||
send.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("allowed_clients", ([], "claude-code", [{"name": "claude-code"}]))
|
||||
async def test_streamable_http_empty_or_malformed_allowlist_admits_nobody(allowed_clients: object) -> None:
|
||||
from starlette.types import Scope
|
||||
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
|
||||
scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
|
||||
receive: Final = AsyncMock(
|
||||
return_value={"type": "http.request", "body": _CLAUDE_CODE_INITIALIZE, "more_body": False}
|
||||
)
|
||||
send: Final = AsyncMock()
|
||||
stateful_handle: Final = AsyncMock()
|
||||
|
||||
with (
|
||||
_client_allowlist_patches(allowed_clients),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.session_manager_stateful",
|
||||
SimpleNamespace(handle_request=stateful_handle),
|
||||
),
|
||||
):
|
||||
await mcp_module.handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
status, body = _forbidden_client_response(send)
|
||||
assert status == 403
|
||||
assert body["details"] == "MCP client 'claude-code' is not listed in this gateway's mcp_allowed_clients."
|
||||
stateful_handle.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streamable_http_allowlist_only_inspects_initialize_requests() -> None:
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
|
||||
scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
|
||||
receive: Final = AsyncMock(side_effect=[{"type": "http.request", "body": _TOOLS_LIST, "more_body": False}])
|
||||
send: Final = AsyncMock()
|
||||
downstream_bodies: Final[list[bytes]] = []
|
||||
|
||||
async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None:
|
||||
downstream_bodies.append(await _drain_body(downstream_receive))
|
||||
|
||||
with (
|
||||
_client_allowlist_patches(["antigravity-cli"]),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.session_manager_stateless",
|
||||
SimpleNamespace(handle_request=AsyncMock(side_effect=handle_request)),
|
||||
),
|
||||
):
|
||||
await mcp_module.handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert downstream_bodies == [_TOOLS_LIST]
|
||||
send.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("request_body", "admitted"),
|
||||
((_ANTIGRAVITY_INITIALIZE, True), (_CLAUDE_CODE_INITIALIZE, False), (_ANONYMOUS_INITIALIZE, False)),
|
||||
)
|
||||
async def test_sse_endpoint_applies_the_same_client_allowlist(request_body: bytes, admitted: bool) -> None:
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_module
|
||||
|
||||
scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp/sse", "headers": []}
|
||||
receive: Final = AsyncMock(side_effect=[{"type": "http.request", "body": request_body, "more_body": False}])
|
||||
send: Final = AsyncMock()
|
||||
downstream_bodies: Final[list[bytes]] = []
|
||||
|
||||
async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None:
|
||||
downstream_bodies.append(await _drain_body(downstream_receive))
|
||||
|
||||
with (
|
||||
_client_allowlist_patches(["antigravity-cli"]),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._raise_preemptive_401_for_unauthenticated_servers",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._check_passthrough_upstream_auth",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch.object(mcp_module.sse_session_manager, "handle_request", side_effect=handle_request),
|
||||
):
|
||||
await mcp_module.handle_sse_mcp(scope, receive, send)
|
||||
|
||||
if admitted:
|
||||
assert downstream_bodies == [request_body]
|
||||
send.assert_not_awaited()
|
||||
return
|
||||
assert downstream_bodies == []
|
||||
status, body = _forbidden_client_response(send)
|
||||
assert status == 403
|
||||
assert body["error"] == "Forbidden"
|
||||
assert "mcp_allowed_clients" in body["details"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_routing_chunked_initialize_to_stateful():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -7428,10 +7428,18 @@ async def test_update_general_settings_keeps_yaml_pass_through_endpoints_next_to
|
|||
request.query_params = {}
|
||||
return request
|
||||
|
||||
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in
|
||||
initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
|
||||
master_key: Final = patch("litellm.proxy.proxy_server.master_key", "sk-master") # test-quality-ok: a set master key is what makes a missing Authorization header a 401
|
||||
settings: Final = patch(
|
||||
"litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}
|
||||
) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch(
|
||||
"litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]
|
||||
) # test-quality-ok: module global holding the YAML endpoints the fix merges in
|
||||
initialize: Final = patch(
|
||||
"litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()
|
||||
) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
|
||||
master_key: Final = patch(
|
||||
"litellm.proxy.proxy_server.master_key", "sk-master"
|
||||
) # test-quality-ok: a set master key is what makes a missing Authorization header a 401
|
||||
with settings, yaml_endpoints, initialize, master_key:
|
||||
await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]})
|
||||
|
||||
|
|
@ -7479,10 +7487,18 @@ async def test_update_general_settings_db_pass_through_endpoint_overrides_yaml_e
|
|||
request.headers = {}
|
||||
request.query_params = {}
|
||||
|
||||
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in
|
||||
initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
|
||||
master_key: Final = patch("litellm.proxy.proxy_server.master_key", "sk-master") # test-quality-ok: a set master key is what makes a missing Authorization header a 401
|
||||
settings: Final = patch(
|
||||
"litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}
|
||||
) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch(
|
||||
"litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]
|
||||
) # test-quality-ok: module global holding the YAML endpoints the fix merges in
|
||||
initialize: Final = patch(
|
||||
"litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()
|
||||
) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here
|
||||
master_key: Final = patch(
|
||||
"litellm.proxy.proxy_server.master_key", "sk-master"
|
||||
) # test-quality-ok: a set master key is what makes a missing Authorization header a 401
|
||||
with settings, yaml_endpoints, initialize, master_key:
|
||||
await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]})
|
||||
|
||||
|
|
@ -8597,9 +8613,7 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments():
|
|||
async def assert_reservation_not_finalized_yet(**kwargs):
|
||||
assert budget_reservation["finalized"] is False
|
||||
incremented_counters.append(kwargs["counter_key"])
|
||||
return ps.PendingSpendIncrement(
|
||||
counter_key=kwargs["counter_key"], increment=kwargs["increment"]
|
||||
)
|
||||
return ps.PendingSpendIncrement(counter_key=kwargs["counter_key"], increment=kwargs["increment"])
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
|
|
@ -10144,9 +10158,15 @@ async def _lit6973_drive_realtime_session(
|
|||
side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)
|
||||
)
|
||||
ws: Final = websocket if websocket is not None else _lit6973_fake_realtime_ws()
|
||||
can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test
|
||||
pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state
|
||||
route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object
|
||||
can_call = patch.object(
|
||||
ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)
|
||||
) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test
|
||||
pre = patch.object(
|
||||
ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call
|
||||
) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state
|
||||
route = patch.object(
|
||||
ps, "route_request", new=AsyncMock(return_value=fake_llm_call())
|
||||
) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object
|
||||
with can_call, pre, route:
|
||||
await ps.realtime_websocket_endpoint(
|
||||
websocket=ws,
|
||||
|
|
@ -10278,13 +10298,9 @@ async def _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
|
|||
from litellm.proxy.utils import InternalUsageCache
|
||||
|
||||
dual_cache: Final = DualCache()
|
||||
await dual_cache.async_set_cache(
|
||||
key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True
|
||||
)
|
||||
await dual_cache.async_set_cache(key=_LIT6463_COUNTER_KEY, value={"slot-1": 1.0, "slot-2": 2.0}, local_only=True)
|
||||
limiter: Final = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache))
|
||||
stash: Final = RequestRateLimiterStash(
|
||||
parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]}
|
||||
)
|
||||
stash: Final = RequestRateLimiterStash(parallel_slot={"slot_id": "slot-1", "counter_keys": [_LIT6463_COUNTER_KEY]})
|
||||
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
|
||||
|
||||
stash_token: Final = _request_stash.set(stash)
|
||||
|
|
@ -10336,9 +10352,7 @@ async def test_successful_realtime_session_leaves_the_max_parallel_slot_for_the_
|
|||
limiter's integer in-memory fallback, double-decrement the counter so the key
|
||||
admits more sessions than max_parallel_requests allows. With the success stamp
|
||||
present the route leaves the slot and the stash alone."""
|
||||
dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(
|
||||
backend_logged_success=True
|
||||
)
|
||||
dual_cache, stash = await _lit6463_drive_realtime_session_holding_a_max_parallel_slot(backend_logged_success=True)
|
||||
|
||||
assert await dual_cache.async_get_cache(key=_LIT6463_COUNTER_KEY, local_only=True) == {
|
||||
"slot-1": 1.0,
|
||||
|
|
@ -10384,8 +10398,12 @@ async def test_release_or_invalidate_falls_back_to_invalidating_the_counters():
|
|||
async def _record(counter_key: str) -> None:
|
||||
invalidated.append(counter_key)
|
||||
|
||||
failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated
|
||||
sink = patch.object(ps, "_invalidate_spend_counter", new=_record) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable
|
||||
failing_release = patch.object(
|
||||
br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))
|
||||
) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated
|
||||
sink = patch.object(
|
||||
ps, "_invalidate_spend_counter", new=_record
|
||||
) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable
|
||||
with failing_release, sink:
|
||||
await br.release_or_invalidate_budget_reservation(budget_reservation=reservation)
|
||||
|
||||
|
|
@ -10401,8 +10419,12 @@ async def test_release_or_invalidate_finalizes_even_when_the_invalidate_fallback
|
|||
from litellm.proxy.spend_tracking import budget_reservation as br
|
||||
|
||||
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
|
||||
failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the fallback branch
|
||||
failing_invalidate = patch.object(br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))) # test-quality-ok: forces the fallback itself to fail
|
||||
failing_release = patch.object(
|
||||
br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))
|
||||
) # test-quality-ok: forces the fallback branch
|
||||
failing_invalidate = patch.object(
|
||||
br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))
|
||||
) # test-quality-ok: forces the fallback itself to fail
|
||||
|
||||
with failing_release, failing_invalidate:
|
||||
await br.release_or_invalidate_budget_reservation(budget_reservation=reservation)
|
||||
|
|
@ -12987,9 +13009,15 @@ async def test_moderations_response_carries_litellm_call_id_header():
|
|||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", spend=0.0)
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(proxy_server_module, "route_request", new=AsyncMock(return_value=fake_llm_call())), # test-quality-ok: fakes the provider call so the response headers assembled by the real route are observable
|
||||
patch.object(proxy_server_module, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data
|
||||
), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "route_request", new=AsyncMock(return_value=fake_llm_call())
|
||||
), # test-quality-ok: fakes the provider call so the response headers assembled by the real route are observable
|
||||
patch.object(
|
||||
proxy_server_module, "proxy_logging_obj"
|
||||
) as mock_logging, # test-quality-ok: module global, no injection point
|
||||
):
|
||||
mock_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
|
||||
mock_logging.update_request_status = AsyncMock()
|
||||
|
|
@ -13026,9 +13054,15 @@ async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplo
|
|||
verbose_proxy_logger.propagate = True
|
||||
try:
|
||||
with (
|
||||
patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key"))), # test-quality-ok: fakes the provider failure so the real route's error log is observable
|
||||
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data
|
||||
), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key"))
|
||||
), # test-quality-ok: fakes the provider failure so the real route's error log is observable
|
||||
patch.object(
|
||||
proxy_server_module, "proxy_logging_obj", new=fake_logging
|
||||
), # test-quality-ok: module global, no injection point
|
||||
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
|
||||
pytest.raises(ProxyException) as raised,
|
||||
):
|
||||
|
|
@ -13061,7 +13095,9 @@ async def test_moderations_unparseable_body_bills_the_callers_litellm_call_id():
|
|||
fake_logging.post_call_failure_hook = AsyncMock()
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "proxy_logging_obj", new=fake_logging
|
||||
), # test-quality-ok: module global, no injection point
|
||||
pytest.raises(ProxyException) as raised,
|
||||
):
|
||||
await proxy_server_module.moderations(
|
||||
|
|
@ -13089,8 +13125,12 @@ async def test_moderations_already_shaped_failure_answers_with_the_callers_litel
|
|||
fake_logging.post_call_failure_hook = AsyncMock()
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)
|
||||
), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "proxy_logging_obj", new=fake_logging
|
||||
), # test-quality-ok: module global, no injection point
|
||||
pytest.raises(ProxyException) as raised,
|
||||
):
|
||||
await proxy_server_module.moderations(
|
||||
|
|
@ -13125,8 +13165,12 @@ async def test_audio_speech_already_shaped_failure_answers_with_the_callers_lite
|
|||
fake_logging.post_call_failure_hook = AsyncMock()
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)
|
||||
), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(
|
||||
proxy_server_module, "proxy_logging_obj", new=fake_logging
|
||||
), # test-quality-ok: module global, no injection point
|
||||
pytest.raises(type(exc)) as raised,
|
||||
):
|
||||
await proxy_server_module.audio_speech(
|
||||
|
|
@ -13814,6 +13858,43 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough()
|
|||
assert ps.general_settings["enable_openai_websocket_passthrough"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"db_general_settings, expected",
|
||||
[
|
||||
({"mcp_allowed_clients": ["antigravity-cli"]}, ["antigravity-cli"]),
|
||||
({"mcp_allowed_clients": []}, []),
|
||||
({}, None),
|
||||
],
|
||||
)
|
||||
async def test_update_general_settings_propagates_mcp_allowed_clients(db_general_settings, expected):
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.general_settings", {"mcp_allowed_clients": ["claude-code"]}):
|
||||
await proxy_config._update_general_settings(db_general_settings=db_general_settings)
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
assert ps.general_settings["mcp_allowed_clients"] == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_general_settings_keeps_yaml_mcp_allowed_clients():
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
proxy_config._yaml_general_settings_keys = {"mcp_allowed_clients"}
|
||||
|
||||
with patch("litellm.proxy.proxy_server.general_settings", {"mcp_allowed_clients": ["claude-code"]}):
|
||||
await proxy_config._update_general_settings(db_general_settings={"mcp_allowed_clients": ["codex-mcp-client"]})
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
assert ps.general_settings["mcp_allowed_clients"] == ["claude-code"]
|
||||
|
||||
|
||||
async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch):
|
||||
from tests.large_text import text
|
||||
from tests.test_litellm.litellm_core_utils.event_loop_lag import (
|
||||
|
|
@ -13855,14 +13936,18 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp
|
|||
{
|
||||
"model_name": "self-hosted",
|
||||
"litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"},
|
||||
"model_info": {"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}},
|
||||
"model_info": {
|
||||
"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}
|
||||
},
|
||||
}
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
response, took, lags = await timed_with_loop_lags(
|
||||
lambda: proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me off the loop"))
|
||||
lambda: proxy_server_module.token_counter(
|
||||
TokenCountRequest(model="self-hosted", prompt="count me off the loop")
|
||||
)
|
||||
)
|
||||
|
||||
assert response.tokenizer_type == "huggingface_tokenizer"
|
||||
|
|
|
|||
|
|
@ -93,7 +93,7 @@ describe("MCPNetworkSettings", () => {
|
|||
await waitFor(() =>
|
||||
expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges", ["10.0.0.0/8"]),
|
||||
);
|
||||
expect(deleteConfigFieldSetting).not.toHaveBeenCalled();
|
||||
expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges");
|
||||
});
|
||||
|
||||
it("clears the setting instead of saving an empty list", async () => {
|
||||
|
|
@ -103,4 +103,68 @@ describe("MCPNetworkSettings", () => {
|
|||
await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges"));
|
||||
expect(updateConfigFieldSetting).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("renders the stored allowed client names once settings load", async () => {
|
||||
vi.mocked(getGeneralSettingsCall).mockResolvedValue([
|
||||
{ field_name: "mcp_allowed_clients", field_value: ["antigravity-cli", "codex-mcp-client"] },
|
||||
]);
|
||||
|
||||
renderSettings();
|
||||
|
||||
expect(await screen.findByText("antigravity-cli")).toBeInTheDocument();
|
||||
expect(screen.getByText("codex-mcp-client")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("adds typed client names on Enter and saves them under mcp_allowed_clients", async () => {
|
||||
renderSettings();
|
||||
const input = await screen.findByRole("textbox", { name: "Allowed client names" });
|
||||
|
||||
await userEvent.type(input, "antigravity-cli, codex-mcp-client{Enter}");
|
||||
|
||||
expect(screen.getByText("antigravity-cli")).toBeInTheDocument();
|
||||
expect(screen.getByText("codex-mcp-client")).toBeInTheDocument();
|
||||
expect(input).toHaveValue("");
|
||||
|
||||
await userEvent.click(screen.getByRole("button", { name: /Save/ }));
|
||||
|
||||
await waitFor(() =>
|
||||
expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [
|
||||
"antigravity-cli",
|
||||
"codex-mcp-client",
|
||||
]),
|
||||
);
|
||||
expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients");
|
||||
});
|
||||
|
||||
it("removes a client name and clears the setting when the list becomes empty", async () => {
|
||||
vi.mocked(getGeneralSettingsCall).mockResolvedValue([
|
||||
{ field_name: "mcp_allowed_clients", field_value: ["claude-code"] },
|
||||
]);
|
||||
|
||||
renderSettings();
|
||||
await userEvent.click(await screen.findByRole("button", { name: "Remove claude-code" }));
|
||||
|
||||
expect(screen.queryByText("claude-code")).not.toBeInTheDocument();
|
||||
|
||||
await userEvent.click(screen.getByRole("button", { name: /Save/ }));
|
||||
|
||||
await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients"));
|
||||
expect(updateConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients", expect.anything());
|
||||
});
|
||||
|
||||
it("keeps the private ranges and the allowed clients as independent settings on save", async () => {
|
||||
vi.mocked(getGeneralSettingsCall).mockResolvedValue([
|
||||
{ field_name: "mcp_internal_ip_ranges", field_value: ["10.0.0.0/8"] },
|
||||
{ field_name: "mcp_allowed_clients", field_value: ["antigravity-cli"] },
|
||||
]);
|
||||
|
||||
renderSettings();
|
||||
await userEvent.click(await screen.findByRole("button", { name: /Save/ }));
|
||||
|
||||
await waitFor(() =>
|
||||
expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", ["antigravity-cli"]),
|
||||
);
|
||||
expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges", ["10.0.0.0/8"]);
|
||||
expect(deleteConfigFieldSetting).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -30,8 +30,10 @@ const MCPNetworkSettings: React.FC<MCPNetworkSettingsProps> = ({ accessToken })
|
|||
const [loading, setLoading] = useState(true);
|
||||
const [saving, setSaving] = useState(false);
|
||||
const [privateRanges, setPrivateRanges] = useState<string[]>([]);
|
||||
const [allowedClients, setAllowedClients] = useState<string[]>([]);
|
||||
const [currentIp, setCurrentIp] = useState<string | null>(null);
|
||||
const [rangeDraft, setRangeDraft] = useState("");
|
||||
const [clientDraft, setClientDraft] = useState("");
|
||||
|
||||
useEffect(() => {
|
||||
loadSettings();
|
||||
|
|
@ -47,6 +49,9 @@ const MCPNetworkSettings: React.FC<MCPNetworkSettingsProps> = ({ accessToken })
|
|||
if (field.field_name === "mcp_internal_ip_ranges" && field.field_value) {
|
||||
setPrivateRanges(field.field_value);
|
||||
}
|
||||
if (field.field_name === "mcp_allowed_clients" && field.field_value) {
|
||||
setAllowedClients(field.field_value);
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Failed to load MCP network settings:", error);
|
||||
|
|
@ -72,6 +77,11 @@ const MCPNetworkSettings: React.FC<MCPNetworkSettingsProps> = ({ accessToken })
|
|||
} else {
|
||||
await deleteConfigFieldSetting(accessToken, "mcp_internal_ip_ranges");
|
||||
}
|
||||
if (allowedClients.length > 0) {
|
||||
await updateConfigFieldSetting(accessToken, "mcp_allowed_clients", allowedClients);
|
||||
} else {
|
||||
await deleteConfigFieldSetting(accessToken, "mcp_allowed_clients");
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Failed to save MCP network settings:", error);
|
||||
} finally {
|
||||
|
|
@ -86,17 +96,28 @@ const MCPNetworkSettings: React.FC<MCPNetworkSettingsProps> = ({ accessToken })
|
|||
};
|
||||
|
||||
// Commas separate entries, matching the old tokenised input.
|
||||
const commitDraft = () => {
|
||||
const added = rangeDraft
|
||||
const splitDraft = (draft: string, existing: string[]) =>
|
||||
draft
|
||||
.split(",")
|
||||
.map((r) => r.trim())
|
||||
.filter((r) => r !== "" && !privateRanges.includes(r));
|
||||
.filter((r) => r !== "" && !existing.includes(r));
|
||||
|
||||
const commitDraft = () => {
|
||||
const added = splitDraft(rangeDraft, privateRanges);
|
||||
if (added.length > 0) {
|
||||
setPrivateRanges([...privateRanges, ...added]);
|
||||
}
|
||||
setRangeDraft("");
|
||||
};
|
||||
|
||||
const commitClientDraft = () => {
|
||||
const added = splitDraft(clientDraft, allowedClients);
|
||||
if (added.length > 0) {
|
||||
setAllowedClients([...allowedClients, ...added]);
|
||||
}
|
||||
setClientDraft("");
|
||||
};
|
||||
|
||||
if (loading) {
|
||||
return (
|
||||
<div className="flex justify-center py-12">
|
||||
|
|
@ -178,6 +199,56 @@ const MCPNetworkSettings: React.FC<MCPNetworkSettingsProps> = ({ accessToken })
|
|||
</p>
|
||||
</Card>
|
||||
|
||||
<div>
|
||||
<p className="text-lg font-semibold">Allowed Client Applications</p>
|
||||
<p className="mt-1 text-sm text-muted-foreground">
|
||||
Only the MCP client applications listed here can connect to the gateway. Names are matched exactly against
|
||||
the clientInfo.name each client sends in its MCP initialize request (for example claude-code or
|
||||
codex-mcp-client). Leave empty to allow every client. Clients choose the name they send, so treat this as a
|
||||
policy control rather than a security boundary.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<Card className="p-6">
|
||||
<div className="mb-2 flex items-center">
|
||||
<p className="text-sm font-medium">Allowed Client Names</p>
|
||||
</div>
|
||||
{allowedClients.length > 0 && (
|
||||
<div className="mb-2 flex flex-wrap gap-1.5">
|
||||
{allowedClients.map((client) => (
|
||||
<Badge key={client} variant="secondary" className="font-mono">
|
||||
{client}
|
||||
<button
|
||||
type="button"
|
||||
aria-label={`Remove ${client}`}
|
||||
onClick={() => setAllowedClients(allowedClients.filter((c) => c !== client))}
|
||||
className="ml-1 cursor-pointer"
|
||||
>
|
||||
<X className="size-3" />
|
||||
</button>
|
||||
</Badge>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
<Input
|
||||
aria-label="Allowed client names"
|
||||
value={clientDraft}
|
||||
placeholder="Leave empty to allow every client, e.g. claude-code, codex-mcp-client"
|
||||
onChange={(e) => setClientDraft(e.target.value)}
|
||||
onBlur={commitClientDraft}
|
||||
onKeyDown={(e) => {
|
||||
if (e.key === "Enter" || e.key === ",") {
|
||||
e.preventDefault();
|
||||
commitClientDraft();
|
||||
}
|
||||
}}
|
||||
/>
|
||||
<p className="mt-2 text-xs text-muted-foreground">
|
||||
Enter the clientInfo.name values to admit. Any other client, or one that does not identify itself, gets a
|
||||
403 on its MCP initialize request.
|
||||
</p>
|
||||
</Card>
|
||||
|
||||
<div className="flex justify-end">
|
||||
<Button onClick={handleSave} disabled={saving}>
|
||||
<Save />
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue