mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
567 lines
27 KiB
Python
567 lines
27 KiB
Python
import asyncio
|
|
import json
|
|
import sys
|
|
from typing import Final
|
|
|
|
import httpx2
|
|
import pytest
|
|
from fastapi import FastAPI, HTTPException, Request
|
|
from fastapi.routing import APIRoute
|
|
from pydantic import BaseModel
|
|
from starlette.testclient import TestClient
|
|
|
|
from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS
|
|
from litellm.proxy import proxy_server
|
|
from litellm.proxy._types import LiteLLM_UserTable, ProxyException, SpecialHeaders
|
|
from litellm.proxy.admin_mcp import admin_mcp_lifespan
|
|
from litellm.proxy.auth.user_api_key_auth import get_api_key, get_api_key_from_custom_header
|
|
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
|
from litellm.proxy.management_endpoints import internal_user_endpoints
|
|
from litellm.proxy.middleware.admission_control_middleware import (
|
|
AdmissionControlMiddleware,
|
|
AdmissionControlSettings,
|
|
AdmissionControlState,
|
|
AdmissionControlStats,
|
|
)
|
|
from litellm.proxy.middleware.per_request_root_path_middleware import PerRequestRootPathMiddleware
|
|
|
|
|
|
class KeyRequest(BaseModel):
|
|
key_alias: str
|
|
|
|
|
|
def _native_credential(request: Request) -> str:
|
|
key, _ = get_api_key(
|
|
custom_litellm_key_header=request.headers.get("x-litellm-api-key"),
|
|
api_key=request.headers.get("authorization", ""),
|
|
azure_api_key_header=request.headers.get("api-key"),
|
|
anthropic_api_key_header=request.headers.get("x-api-key"),
|
|
google_ai_studio_api_key_header=request.headers.get("x-goog-api-key"),
|
|
azure_apim_header=request.headers.get("ocp-apim-subscription-key"),
|
|
pass_through_endpoints=None, route=request.url.path, request=request,
|
|
)
|
|
configured: Final = proxy_server.general_settings.get("litellm_key_header_name")
|
|
return get_api_key_from_custom_header(request, configured) if configured is not None else key
|
|
|
|
|
|
@pytest.fixture
|
|
def management_app(monkeypatch: pytest.MonkeyPatch) -> FastAPI:
|
|
monkeypatch.setattr(proxy_server, "premium_user", True)
|
|
monkeypatch.setattr(proxy_server, "general_settings", {})
|
|
monkeypatch.setenv("LITELLM_ENABLE_ADMIN_MCP", "true")
|
|
monkeypatch.setenv("LITELLM_ADMIN_TOOLS", "list_keys")
|
|
for name in (
|
|
"PROXY_BASE_URL", "LITELLM_MCP_PUBLIC_URL", "LITELLM_ADMIN_READ_ONLY",
|
|
"LITELLM_ADMIN_RESPONSE_VIEW", "LITELLM_ADMIN_SCHEMA_MODE",
|
|
):
|
|
monkeypatch.delenv(name, raising=False)
|
|
app: Final = FastAPI(lifespan=admin_mcp_lifespan)
|
|
|
|
@app.get("/user/info")
|
|
async def user_info(request: Request) -> dict[str, object]:
|
|
await asyncio.sleep(0)
|
|
credential: Final = _native_credential(request)
|
|
if credential == "team-key":
|
|
raise HTTPException(status_code=404, detail="User None not found")
|
|
if credential not in {"admin-a", "admin-b", "admin-b-limited", "member", "viewer"}:
|
|
raise HTTPException(status_code=401)
|
|
user_id: Final = "admin-b" if credential == "admin-b-limited" else credential
|
|
return {
|
|
"user_id": user_id,
|
|
"user_info": {
|
|
"user_id": user_id,
|
|
"user_role": {"member": "internal_user", "viewer": "proxy_admin_viewer"}.get(user_id, "proxy_admin"),
|
|
},
|
|
}
|
|
|
|
@app.get("/key/list", operation_id="list_keys_key_list_get")
|
|
async def list_keys(request: Request, large: bool = False) -> dict[str, object]:
|
|
credential: Final = _native_credential(request)
|
|
if credential == "admin-b-limited":
|
|
raise HTTPException(403, "This key cannot list keys")
|
|
if large:
|
|
return {"keys": [{"key_alias": ("a" if credential == "admin-a" else "b") * 20000}]}
|
|
return {
|
|
"keys": ["owned-by-a" if credential == "admin-a" else "owned-by-b"],
|
|
"client": request.client.host if request.client else None,
|
|
"scheme": request.url.scheme,
|
|
"forwarded_for": request.headers.get("x-forwarded-for"),
|
|
"cookie": request.headers.get("cookie"),
|
|
"caller_is_changed_by": request.headers.get("litellm-changed-by")
|
|
== credential,
|
|
"policy_team": request.headers.get("x-litellm-team-id"),
|
|
"alternate_credentials_absent": not any(
|
|
name in request.headers
|
|
for name in SpecialHeaders.litellm_credential_header_names() - {"authorization"}
|
|
),
|
|
}
|
|
|
|
app.state.created_aliases = []
|
|
|
|
@app.post("/key/generate", operation_id="generate_key_fn_key_generate_post")
|
|
async def create_key(payload: KeyRequest) -> dict[str, str]:
|
|
app.state.created_aliases.append(payload.key_alias)
|
|
if payload.key_alias == "fail-after-write":
|
|
raise HTTPException(status_code=500)
|
|
return {"key": "sk-new-key", "key_alias": payload.key_alias}
|
|
|
|
@app.post("/mcp")
|
|
async def existing_mcp() -> dict[str, str]:
|
|
return {"server": "existing"}
|
|
|
|
@app.post("/{server_name}/mcp")
|
|
async def existing_namespace(server_name: str) -> dict[str, str]:
|
|
return {"server": server_name}
|
|
|
|
return app
|
|
|
|
|
|
@pytest.mark.parametrize("enabled", [None, "false", "0", "off", "No"])
|
|
def test_disabled_preserves_existing_admin_namespace(monkeypatch: pytest.MonkeyPatch, enabled: str | None) -> None:
|
|
monkeypatch.setattr(proxy_server, "premium_user", False)
|
|
monkeypatch.setitem(sys.modules, "litellm_admin_mcp.config", None)
|
|
if enabled is None:
|
|
monkeypatch.delenv("LITELLM_ENABLE_ADMIN_MCP", raising=False)
|
|
else:
|
|
monkeypatch.setenv("LITELLM_ENABLE_ADMIN_MCP", enabled)
|
|
app: Final = FastAPI(lifespan=admin_mcp_lifespan)
|
|
|
|
@app.post("/{server_name}/mcp")
|
|
async def namespace(server_name: str) -> dict[str, str]:
|
|
return {"server": server_name}
|
|
|
|
with TestClient(app) as client:
|
|
response: Final = client.post("/admin/mcp")
|
|
assert response.status_code == 200
|
|
assert response.json() == {"server": "admin"}
|
|
|
|
|
|
def test_enabled_without_connector_explains_installation(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(proxy_server, "premium_user", True)
|
|
monkeypatch.setenv("LITELLM_ENABLE_ADMIN_MCP", "true")
|
|
monkeypatch.setitem(sys.modules, "litellm_admin_mcp.config", None)
|
|
with pytest.raises(RuntimeError, match="admin-mcp dependency group"):
|
|
with TestClient(FastAPI(lifespan=admin_mcp_lifespan)):
|
|
pytest.fail("Enabling the connector without its dependency must fail startup")
|
|
|
|
|
|
@pytest.mark.parametrize("enabled", ["true", "Yes", "on"])
|
|
def test_unlicensed_opt_in_fails_before_loading_connector(
|
|
management_app: FastAPI, monkeypatch: pytest.MonkeyPatch, enabled: str
|
|
) -> None:
|
|
monkeypatch.setattr(proxy_server, "premium_user", False)
|
|
monkeypatch.setenv("LITELLM_ENABLE_ADMIN_MCP", enabled)
|
|
monkeypatch.setitem(sys.modules, "litellm_admin_mcp.config", None)
|
|
with pytest.raises(HTTPException) as exc:
|
|
with TestClient(management_app):
|
|
pytest.fail("An unlicensed deployment must not serve the hosted admin connector")
|
|
assert exc.value.status_code == 403
|
|
assert "LITELLM_LICENSE" in str(exc.value.detail)
|
|
assert all(route.name != "admin_mcp" for route in management_app.routes)
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
def test_losing_enterprise_status_blocks_admin_tool_calls(
|
|
management_app: FastAPI, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_ADMIN_TOOLS", "create_key")
|
|
headers: Final = {"Authorization": "Bearer admin-a", "Accept": "application/json, text/event-stream"}
|
|
payload: Final = {
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "tools/call",
|
|
"params": {"name": "create_key", "arguments": {"body": {"key_alias": "licensed-write"}}},
|
|
}
|
|
with TestClient(management_app, base_url="http://localhost:4000") as client:
|
|
licensed: Final = client.post("/admin/mcp", headers=headers, json=payload)
|
|
assert licensed.status_code == 200, licensed.text
|
|
assert licensed.json()["result"]["isError"] is False
|
|
monkeypatch.setattr(proxy_server, "premium_user", False)
|
|
denied: Final = client.post("/admin/mcp", headers=headers, json=payload)
|
|
assert denied.status_code == 403, denied.text
|
|
assert "LITELLM_LICENSE" in denied.text
|
|
assert management_app.state.created_aliases == ["licensed-write"]
|
|
|
|
|
|
def test_invalid_flag_fails_startup(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LITELLM_ENABLE_ADMIN_MCP", "treu")
|
|
with pytest.raises(ValueError, match="LITELLM_ENABLE_ADMIN_MCP"):
|
|
with TestClient(FastAPI(lifespan=admin_mcp_lifespan)):
|
|
pytest.fail("Invalid flag must fail startup")
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
@pytest.mark.parametrize(
|
|
"authorization,status",
|
|
[(None, 401), ("Bearer invalid", 401), ("Bearer member", 403), ("Bearer viewer", 403), ("Bearer team-key", 403)],
|
|
)
|
|
def test_admin_endpoint_rejects_unauthorized_callers(
|
|
management_app: FastAPI, authorization: str | None, status: int
|
|
) -> None:
|
|
headers: Final = {
|
|
**({"Authorization": authorization} if authorization else {}),
|
|
"x-litellm-api-key": "Bearer admin-a",
|
|
}
|
|
with TestClient(management_app, base_url="http://localhost:4000") as client:
|
|
response: Final = client.post("/admin/mcp", headers=headers, json={})
|
|
assert response.status_code == status
|
|
assert response.headers["cache-control"] == "no-store"
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
def test_mount_keeps_existing_mcp_and_manages_restart(management_app: FastAPI) -> None:
|
|
for _ in range(2):
|
|
with TestClient(management_app, base_url="http://localhost:4000") as client:
|
|
assert client.post("/mcp").json() == {"server": "existing"}
|
|
assert client.post("/tools/mcp").json() == {"server": "tools"}
|
|
response: Final = client.post(
|
|
"/admin/mcp",
|
|
headers={"Authorization": "Bearer admin-a", "Accept": "application/json, text/event-stream"},
|
|
json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert {tool["name"] for tool in response.json()["result"]["tools"]} == {
|
|
"list_keys",
|
|
"describe_admin_tool",
|
|
"read_admin_result",
|
|
}
|
|
assert all(route.name != "admin_mcp" for route in management_app.routes)
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
@pytest.mark.parametrize("root_path", ["", "/gateway"])
|
|
@pytest.mark.parametrize("prefix_mode", ["ingress", "scalar", "multiple"])
|
|
async def test_concurrent_calls_preserve_identity_network_context_and_strip_cookies(
|
|
management_app: FastAPI, monkeypatch: pytest.MonkeyPatch, root_path: str, prefix_mode: str
|
|
) -> None:
|
|
monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com" + root_path)
|
|
monkeypatch.setenv("LITELLM_BASE_URL", "https://must-not-call.example.com")
|
|
monkeypatch.setenv("LITELLM_API_KEY", "must-not-use-shared-credential")
|
|
if prefix_mode == "scalar":
|
|
management_app.root_path = root_path
|
|
elif prefix_mode == "multiple":
|
|
management_app.add_middleware(PerRequestRootPathMiddleware, root_paths=("/other", root_path))
|
|
|
|
async def call(credential: str) -> dict[str, object]:
|
|
async with httpx2.AsyncClient(
|
|
transport=httpx2.ASGITransport(
|
|
app=management_app,
|
|
root_path=root_path if prefix_mode == "ingress" else "",
|
|
client=("198.51.100.7", 4567),
|
|
),
|
|
base_url="https://gateway.example.com",
|
|
) as client:
|
|
response: Final = await client.post(
|
|
root_path + "/admin/mcp",
|
|
headers={
|
|
**{
|
|
name: "Bearer member"
|
|
for name in SpecialHeaders.litellm_credential_header_names() - {"authorization"}
|
|
},
|
|
"Authorization": "Bearer " + credential,
|
|
"Accept": "application/json, text/event-stream",
|
|
"X-Forwarded-For": "203.0.113.8",
|
|
"Cookie": "session=must-not-forward",
|
|
"litellm-changed-by": "forged-actor",
|
|
"x-litellm-team-id": "policy-team",
|
|
},
|
|
json={"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "list_keys"}},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
return response.json()
|
|
|
|
async with management_app.router.lifespan_context(management_app):
|
|
results: Final = await asyncio.gather(call("admin-a"), call("admin-b"))
|
|
|
|
for credential, result in zip(("admin-a", "admin-b"), results):
|
|
assert json.loads(result["result"]["content"][0]["text"]) == {
|
|
"keys": ["owned-by-a" if credential == "admin-a" else "owned-by-b"],
|
|
"client": "198.51.100.7",
|
|
"scheme": "https",
|
|
"forwarded_for": "203.0.113.8",
|
|
"cookie": "",
|
|
"caller_is_changed_by": True,
|
|
"policy_team": "policy-team",
|
|
"alternate_credentials_absent": True,
|
|
}
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
@pytest.mark.parametrize("header_name", ["X-Custom-Key", "Authorization", "x-litellm-api-key"])
|
|
@pytest.mark.parametrize("caller_value", [None, "Bearer member"])
|
|
def test_configured_key_header_uses_mcp_bearer_after_settings_change(
|
|
management_app: FastAPI, monkeypatch: pytest.MonkeyPatch, header_name: str, caller_value: str | None
|
|
) -> None:
|
|
with TestClient(management_app, base_url="http://localhost:4000") as client:
|
|
monkeypatch.setitem(proxy_server.general_settings, "litellm_key_header_name", header_name)
|
|
response: Final = client.post(
|
|
"/admin/mcp",
|
|
headers={
|
|
**({header_name: caller_value} if caller_value else {}),
|
|
"Authorization": "Bearer admin-a",
|
|
"Accept": "application/json, text/event-stream",
|
|
},
|
|
json={"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "list_keys"}},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert json.loads(response.json()["result"]["content"][0]["text"])["keys"] == ["owned-by-a"]
|
|
|
|
|
|
@pytest.mark.parametrize("header_name", [
|
|
"", "bad name", "x-ключ", 7, "Cookie", "litellm-changed-by",
|
|
"Content-Type", "Host", "X-Forwarded-For", "x-litellm-team-id",
|
|
*STANDARD_CUSTOMER_ID_HEADERS,
|
|
])
|
|
def test_configured_key_header_rejects_invalid_or_reserved_names(
|
|
management_app: FastAPI, monkeypatch: pytest.MonkeyPatch, header_name: object
|
|
) -> None:
|
|
monkeypatch.setitem(proxy_server.general_settings, "litellm_key_header_name", header_name)
|
|
with pytest.raises(ValueError, match="litellm_key_header_name"):
|
|
with TestClient(management_app):
|
|
pytest.fail("An ambiguous or malformed credential header must fail startup")
|
|
|
|
|
|
@pytest.mark.parametrize("policy", [
|
|
{"user_header_name": "X-Custom-Key"},
|
|
{"user_header_mappings": {"header_name": "X-Custom-Key", "litellm_user_role": "customer"}},
|
|
{"user_header_mappings": [{"header_name": "X-Custom-Key", "litellm_user_role": "internal_user"}]},
|
|
{"enable_oauth2_proxy_auth": True, "oauth2_config_mappings": {"user_id": "X-Custom-Key"}},
|
|
{"mcp_client_id_header": "X-Custom-Key"},
|
|
])
|
|
def test_configured_key_header_cannot_replace_configured_identity(
|
|
management_app: FastAPI, monkeypatch: pytest.MonkeyPatch, policy: dict[str, object]
|
|
) -> None:
|
|
monkeypatch.setattr(proxy_server, "general_settings", {**policy, "litellm_key_header_name": "x-custom-key"})
|
|
with pytest.raises(ValueError, match="litellm_key_header_name"):
|
|
with TestClient(management_app):
|
|
pytest.fail("A credential header must not replace a configured identity header")
|
|
|
|
|
|
@pytest.mark.parametrize("policy", [
|
|
{"user_header_name": "x-api-key"},
|
|
{"user_header_mappings": {"header_name": "Cookie", "litellm_user_role": "customer"}},
|
|
{"enable_oauth2_proxy_auth": True, "oauth2_config_mappings": {"user_id": "x-litellm-api-key"}},
|
|
])
|
|
def test_policy_headers_cannot_share_overwritten_slots_without_a_custom_key(
|
|
management_app: FastAPI, monkeypatch: pytest.MonkeyPatch, policy: dict[str, object]
|
|
) -> None:
|
|
monkeypatch.setattr(proxy_server, "general_settings", policy)
|
|
with pytest.raises(ValueError, match="configured identity headers"):
|
|
with TestClient(management_app):
|
|
pytest.fail("Native credentials must not overwrite configured policy headers")
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
def test_compact_results_bind_to_bearer_despite_conflicting_alternate_keys(management_app: FastAPI) -> None:
|
|
headers: Final = {
|
|
"Authorization": "Bearer admin-a", "x-litellm-api-key": "Bearer admin-b",
|
|
"Accept": "application/json, text/event-stream",
|
|
}
|
|
with TestClient(management_app, base_url="http://localhost:4000") as client:
|
|
saved: Final = client.post(
|
|
"/admin/mcp", headers=headers,
|
|
json={"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {
|
|
"name": "list_keys", "arguments": {"query": {"large": True}, "response": {"view": "compact"}},
|
|
}},
|
|
)
|
|
assert saved.status_code == 200, saved.text
|
|
result_id: Final = json.loads(saved.json()["result"]["content"][0]["text"])["result_id"]
|
|
read_payload: Final = {"jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": {
|
|
"name": "read_admin_result", "arguments": {"result_id": result_id, "view": "full"},
|
|
}}
|
|
replayed: Final = client.post(
|
|
"/admin/mcp", headers={**headers, "x-litellm-api-key": "Bearer admin-b-limited"}, json=read_payload,
|
|
)
|
|
other_bearer: Final = client.post(
|
|
"/admin/mcp", headers={**headers, "Authorization": "Bearer admin-b"}, json=read_payload,
|
|
)
|
|
limited: Final = client.post(
|
|
"/admin/mcp",
|
|
headers={"Authorization": "Bearer admin-b-limited", "Accept": headers["Accept"]},
|
|
json={"jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": {"name": "list_keys"}},
|
|
)
|
|
assert replayed.status_code == 200, replayed.text
|
|
assert json.loads(replayed.json()["result"]["content"][0]["text"]) == {"keys": [{"key_alias": "a" * 20000}]}
|
|
assert other_bearer.status_code == 200, other_bearer.text
|
|
assert other_bearer.json()["result"]["isError"] is True
|
|
assert limited.status_code == 200, limited.text
|
|
assert limited.json()["result"]["isError"] is True
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
@pytest.mark.parametrize("public_url", ["gateway.example.com", "https:/broken", "http://gateway.example.com"])
|
|
def test_invalid_public_origin_fails_startup(
|
|
management_app: FastAPI, monkeypatch: pytest.MonkeyPatch, public_url: str
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_MCP_PUBLIC_URL", public_url)
|
|
with pytest.raises(ValueError, match="HTTPS gateway origin"):
|
|
with TestClient(management_app):
|
|
pytest.fail("An invalid trusted public origin must fail startup")
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
def test_public_host_and_origin_are_checked(management_app: FastAPI, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com")
|
|
headers: Final = {"Authorization": "Bearer admin-a", "Accept": "application/json, text/event-stream"}
|
|
with TestClient(management_app, base_url="https://gateway.example.com") as client:
|
|
hostile_host: Final = client.post("/admin/mcp", headers={**headers, "Host": "hostile.example.com"}, json={})
|
|
hostile_origin: Final = client.post(
|
|
"/admin/mcp", headers={**headers, "Origin": "https://hostile.example.com"}, json={}
|
|
)
|
|
assert hostile_host.status_code == 421
|
|
assert hostile_origin.status_code == 403
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
@pytest.mark.parametrize("read_only,alias", [(False, "created"), (False, "fail-after-write"), (True, "denied")])
|
|
def test_writes_respect_read_only_and_are_never_retried(
|
|
management_app: FastAPI, monkeypatch: pytest.MonkeyPatch, read_only: bool, alias: str
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_ADMIN_TOOLS", "list_keys,create_key")
|
|
monkeypatch.setenv("LITELLM_ADMIN_READ_ONLY", str(read_only).lower())
|
|
with TestClient(management_app, base_url="http://localhost:4000") as client:
|
|
response: Final = client.post(
|
|
"/admin/mcp",
|
|
headers={"Authorization": "Bearer admin-a", "Accept": "application/json, text/event-stream"},
|
|
json={
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "tools/call",
|
|
"params": {"name": "create_key", "arguments": {"body": {"key_alias": alias}}},
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert management_app.state.created_aliases == ([] if read_only else [alias])
|
|
assert response.json()["result"]["isError"] == (read_only or alias == "fail-after-write")
|
|
if alias == "created":
|
|
assert json.loads(response.json()["result"]["content"][0]["text"]) == {"key": "sk-new-key", "key_alias": alias}
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
def test_full_results_do_not_require_worker_affinity(management_app: FastAPI) -> None:
|
|
with TestClient(management_app, base_url="http://localhost:4000") as client:
|
|
response: Final = client.post(
|
|
"/admin/mcp",
|
|
headers={"Authorization": "Bearer admin-a", "Accept": "application/json, text/event-stream"},
|
|
json={
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "tools/call",
|
|
"params": {"name": "list_keys", "arguments": {"query": {"large": True}}},
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert json.loads(response.json()["result"]["content"][0]["text"]) == {"keys": [{"key_alias": "a" * 20000}]}
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
def test_nested_management_calls_share_one_admission_slot(management_app: FastAPI) -> None:
|
|
def no_metrics() -> None:
|
|
return None
|
|
|
|
def single_slot() -> AdmissionControlSettings:
|
|
return AdmissionControlSettings(1, 0, 1.0)
|
|
|
|
state: Final = AdmissionControlState(no_metrics)
|
|
management_app.add_middleware(
|
|
AdmissionControlMiddleware,
|
|
get_settings=single_slot,
|
|
state=state,
|
|
)
|
|
with TestClient(management_app, base_url="http://localhost:4000") as client:
|
|
response: Final = client.post(
|
|
"/admin/mcp",
|
|
headers={"Authorization": "Bearer admin-a", "Accept": "application/json, text/event-stream"},
|
|
json={"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "list_keys"}},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["result"]["isError"] is False
|
|
assert json.loads(response.json()["result"]["content"][0]["text"])["keys"] == ["owned-by-a"]
|
|
assert state.get_stats() == AdmissionControlStats(0, 0, 0)
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
def test_oversized_mcp_body_is_rejected_before_a_write(
|
|
management_app: FastAPI, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_ADMIN_TOOLS", "create_key")
|
|
with TestClient(management_app, base_url="http://localhost:4000") as client:
|
|
response: Final = client.post(
|
|
"/admin/mcp",
|
|
headers={"Authorization": "Bearer admin-a", "Accept": "application/json, text/event-stream"},
|
|
json={
|
|
"jsonrpc": "2.0", "id": 1, "method": "tools/call",
|
|
"params": {"name": "create_key", "arguments": {"body": {"key_alias": "x" * 300_000}}},
|
|
},
|
|
)
|
|
assert response.status_code == 413, response.text
|
|
assert management_app.state.created_aliases == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
|
|
@pytest.mark.parametrize("peer,mapped_user,bearer,status", [
|
|
("10.0.0.8", "admin-a", "member", 200),
|
|
("10.0.0.8", "member", "admin-a", 403),
|
|
("203.0.113.8", "admin-a", "admin-a", 401),
|
|
])
|
|
async def test_oauth2_proxy_auth_preserves_native_identity_and_peer_trust(
|
|
management_app: FastAPI, monkeypatch: pytest.MonkeyPatch,
|
|
peer: str, mapped_user: str, bearer: str, status: int,
|
|
) -> None:
|
|
users: Final = {
|
|
user_id: LiteLLM_UserTable(user_id=user_id, user_role=role, teams=[])
|
|
for user_id, role in (("admin-a", "proxy_admin"), ("member", "internal_user"))
|
|
}
|
|
|
|
class EmptyTable:
|
|
async def find_many(self, **_: object) -> list[object]:
|
|
return []
|
|
|
|
class DatabaseTables:
|
|
litellm_teamtable: Final = EmptyTable()
|
|
litellm_teammembership: Final = EmptyTable()
|
|
|
|
class UserInfoDatabase:
|
|
db: Final = DatabaseTables()
|
|
|
|
async def get_data(
|
|
self, *, user_id: str | None = None, table_name: str | None = None, **_: object
|
|
) -> LiteLLM_UserTable | list[object] | None:
|
|
return users.get(user_id) if user_id is not None and table_name is None else []
|
|
|
|
cache: Final = UserApiKeyCache()
|
|
for user_id, user in users.items():
|
|
cache.set_cache(user_id, user)
|
|
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
|
|
monkeypatch.setattr(proxy_server, "prisma_client", UserInfoDatabase())
|
|
monkeypatch.setattr(proxy_server, "general_settings", {
|
|
"enable_oauth2_proxy_auth": True,
|
|
"trusted_proxy_ranges": ["10.0.0.0/24"],
|
|
"oauth2_config_mappings": {"user_id": "X-Authenticated-User"},
|
|
})
|
|
management_app.router.routes[:] = [
|
|
route for route in management_app.router.routes
|
|
if not (isinstance(route, APIRoute) and route.path == "/user/info")
|
|
]
|
|
management_app.add_api_route("/user/info", internal_user_endpoints.user_info, methods=["GET"])
|
|
management_app.add_exception_handler(ProxyException, proxy_server.openai_exception_handler)
|
|
|
|
async with management_app.router.lifespan_context(management_app):
|
|
async with httpx2.AsyncClient(
|
|
transport=httpx2.ASGITransport(app=management_app, client=(peer, 12345)),
|
|
base_url="http://localhost:4000",
|
|
) as client:
|
|
response: Final = await client.post(
|
|
"/admin/mcp",
|
|
headers={
|
|
"Authorization": "Bearer " + bearer,
|
|
"X-Authenticated-User": mapped_user,
|
|
"X-Forwarded-For": "10.0.0.8",
|
|
"Accept": "application/json, text/event-stream",
|
|
},
|
|
json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
|
|
)
|
|
assert response.status_code == status, response.text
|
|
if status == 200:
|
|
assert "list_keys" in {tool["name"] for tool in response.json()["result"]["tools"]}
|