mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
fix: align core dependencies and complete admin MCP regression coverage
This commit is contained in:
Tin
parent
24ccdf2564
commit
f97ea56af2
6 changed files with 190 additions and 105 deletions
|
|
@ -11,7 +11,9 @@ authors = [
|
|||
]
|
||||
dependencies = [
|
||||
"fastuuid>=0.14.0,<1.0; python_version < '3.15'",
|
||||
"filelock>=3.16.1,<4.0",
|
||||
# TODO: Drop the 3.10 cap when a compatible release fixes the APIs introduced in 3.20.4 (GHSA-v4x8-3wr5-3wq8)
|
||||
"filelock>=3.20.3,<3.20.4; python_version < '3.11'",
|
||||
"filelock>=4.2.1,<5.0; python_version >= '3.11'",
|
||||
"httpx[http2]>=0.28.0,<1.0",
|
||||
"openai>=2.20.0,<4.0.0",
|
||||
"python-dateutil>=2.8.2,<3.0",
|
||||
|
|
@ -31,7 +33,7 @@ dependencies = [
|
|||
"pydantic>=2.12.0,<3.0.0; python_version == '3.14'",
|
||||
"pydantic>=2.14.0,<3.0.0; python_version >= '3.15'",
|
||||
"pydantic-settings>=2.14.1,<3.0",
|
||||
"jsonschema>=4.0.0,<5.0",
|
||||
"jsonschema>=4.22.0,<5.0",
|
||||
"certifi>=2026.10.19",
|
||||
"typing-extensions>=4.13.0,<5.0",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,18 +1,23 @@
|
|||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
import json
|
||||
from typing import Final
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import anyio
|
||||
import httpx2
|
||||
from mcp import ClientSession
|
||||
from mcp import ClientSession, types
|
||||
from mcp.shared.memory import create_client_server_memory_streams
|
||||
from pydantic import JsonValue
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.admin_mcp._types import JSON_OBJECT, JSON_VALUE, object_value
|
||||
from litellm.proxy.admin_mcp.config import Config
|
||||
from litellm.proxy.admin_mcp.gateway import Gateway
|
||||
from litellm.proxy.admin_mcp.server import create_server
|
||||
|
||||
|
||||
SPEC = {
|
||||
SPEC: Final[dict[str, JsonValue]] = {
|
||||
"openapi": "3.1.0",
|
||||
"paths": {
|
||||
"/key/generate": {
|
||||
|
|
@ -86,24 +91,24 @@ SPEC = {
|
|||
|
||||
|
||||
class StubGateway:
|
||||
def __init__(self):
|
||||
self.requests = []
|
||||
self.writes = []
|
||||
self.keys = []
|
||||
self.models = []
|
||||
def __init__(self) -> None:
|
||||
self.requests: list[httpx2.Request] = []
|
||||
self.writes: list[httpx2.Request] = []
|
||||
self.keys: list[JsonValue] = []
|
||||
self.models: list[JsonValue] = []
|
||||
self.status = 200
|
||||
self.timeout = False
|
||||
self.revoke = False
|
||||
self.revoked = set()
|
||||
self.revoked: set[str] = set()
|
||||
self.spec = SPEC
|
||||
|
||||
def handle(self, request):
|
||||
def handle(self, request: httpx2.Request) -> httpx2.Response:
|
||||
self.requests.append(request)
|
||||
token = request.headers.get("authorization", "").removeprefix("Bearer ")
|
||||
token: Final = request.headers.get("authorization", "").removeprefix("Bearer ")
|
||||
if request.url.path == "/v2/user/info" and token == "team-key":
|
||||
return httpx2.Response(400, json={"error": {"message": "A user-bound key is required"}})
|
||||
if request.url.path == "/v2/user/info":
|
||||
role = (
|
||||
role: Final = (
|
||||
"proxy_admin"
|
||||
if token in {"alice-admin", "bob-admin"} and token not in self.revoked
|
||||
else "internal_user"
|
||||
|
|
@ -124,19 +129,19 @@ class StubGateway:
|
|||
self.status, json={"secret": "upstream-secret"}, headers={"Location": "https://untrusted.example"}
|
||||
)
|
||||
if request.url.path == "/key/generate":
|
||||
body = json.loads(request.content)
|
||||
value = {**body, "key": "sk-created-for-" + token, "echo": token}
|
||||
key_body: Final = JSON_OBJECT.validate_json(request.content)
|
||||
value: Final[dict[str, JsonValue]] = {**key_body, "key": "sk-created-for-" + token, "echo": token}
|
||||
self.keys.append(value)
|
||||
return httpx2.Response(200, json=value)
|
||||
if request.url.path == "/key/list":
|
||||
return httpx2.Response(200, json={"keys": self.keys})
|
||||
if request.url.path == "/model/new":
|
||||
body = json.loads(request.content)
|
||||
model = {
|
||||
**body,
|
||||
model_body: Final = JSON_OBJECT.validate_json(request.content)
|
||||
model: Final[dict[str, JsonValue]] = {
|
||||
**model_body,
|
||||
"model_info": {"id": "model-1"},
|
||||
"litellm_params": {
|
||||
**body["litellm_params"],
|
||||
**object_value(model_body["litellm_params"]),
|
||||
"api_key": "provider-secret",
|
||||
"aws_session_token": "aws-secret",
|
||||
},
|
||||
|
|
@ -152,14 +157,21 @@ class StubGateway:
|
|||
raise AssertionError("Unexpected gateway route")
|
||||
|
||||
|
||||
class SessionOptions(TypedDict, total=False):
|
||||
read_only: ReadOnly[bool]
|
||||
allowed_tools: ReadOnly[frozenset[str]]
|
||||
response_view: ReadOnly[str]
|
||||
schema_mode: ReadOnly[str]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stub():
|
||||
def stub() -> StubGateway:
|
||||
return StubGateway()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def connected(gateway, api_key="alice-admin"):
|
||||
server = create_server(gateway, lambda context: api_key)
|
||||
async def connected(gateway: Gateway, api_key: str = "alice-admin") -> AsyncGenerator[ClientSession, None]:
|
||||
server: Final = create_server(gateway, lambda context: api_key)
|
||||
async with create_client_server_memory_streams() as (client_streams, server_streams):
|
||||
async with anyio.create_task_group() as group:
|
||||
group.start_soon(server.run, *server_streams, server.create_initialization_options())
|
||||
|
|
@ -170,13 +182,29 @@ async def connected(gateway, api_key="alice-admin"):
|
|||
|
||||
|
||||
@asynccontextmanager
|
||||
async def session_for(stub, api_key="alice-admin", **overrides):
|
||||
config = Config("https://gateway.example.com", **overrides)
|
||||
async def session_for(
|
||||
stub: StubGateway,
|
||||
api_key: str = "alice-admin",
|
||||
*,
|
||||
read_only: bool = False,
|
||||
allowed_tools: frozenset[str] = frozenset(),
|
||||
response_view: str = "compact",
|
||||
schema_mode: str = "full",
|
||||
) -> AsyncGenerator[ClientSession, None]:
|
||||
config: Final = Config(
|
||||
"https://gateway.example.com",
|
||||
read_only=read_only,
|
||||
allowed_tools=allowed_tools,
|
||||
response_view=response_view,
|
||||
schema_mode=schema_mode,
|
||||
)
|
||||
async with httpx2.AsyncClient(transport=httpx2.MockTransport(stub.handle)) as client:
|
||||
async with connected(Gateway(config, client), api_key) as session:
|
||||
yield session
|
||||
|
||||
|
||||
def data(result):
|
||||
def data(result: types.CallToolResult) -> JsonValue:
|
||||
assert not result.is_error, result.content
|
||||
return json.loads(result.content[0].text)
|
||||
content: Final = result.content[0]
|
||||
assert isinstance(content, types.TextContent), content
|
||||
return JSON_VALUE.validate_json(content.text)
|
||||
|
|
|
|||
|
|
@ -1,11 +1,14 @@
|
|||
import copy
|
||||
import json
|
||||
from typing import Final
|
||||
|
||||
import httpx2
|
||||
from mcp import types
|
||||
from mcp.shared.exceptions import MCPError
|
||||
from pydantic import JsonValue
|
||||
import pytest
|
||||
|
||||
from tests.unit.proxy.admin_mcp.conftest import SPEC, data, session_for
|
||||
from tests.unit.proxy.admin_mcp.conftest import SPEC, SessionOptions, StubGateway, connected, data, session_for
|
||||
from litellm.proxy.admin_mcp.catalog import OPERATIONS
|
||||
from litellm.proxy.admin_mcp.config import Config
|
||||
from litellm.proxy.admin_mcp.gateway import AdminError, Gateway, HELPERS, MAX_REDACTION_DEPTH, safe_result
|
||||
|
|
@ -13,7 +16,7 @@ from litellm.proxy.admin_mcp.schema import tool_schema
|
|||
from litellm.proxy.admin_mcp.server import create_http_app
|
||||
|
||||
|
||||
async def test_real_mcp_create_key_and_read_back(stub):
|
||||
async def test_real_mcp_create_key_and_read_back(stub: StubGateway) -> None:
|
||||
async with session_for(stub) as session:
|
||||
tools = (await session.list_tools()).tools
|
||||
assert {t.name for t in tools} == {
|
||||
|
|
@ -32,7 +35,7 @@ async def test_real_mcp_create_key_and_read_back(stub):
|
|||
assert all(r.headers["Authorization"] == "Bearer alice-admin" for r in stub.requests)
|
||||
|
||||
|
||||
async def test_real_mcp_add_model_verify_and_strip_provider_secrets(stub):
|
||||
async def test_real_mcp_add_model_verify_and_strip_provider_secrets(stub: StubGateway) -> None:
|
||||
body = {
|
||||
"model_name": "support",
|
||||
"litellm_params": {"model": "openai/gpt-6.1-sol", "litellm_credential_name": "production"},
|
||||
|
|
@ -49,14 +52,14 @@ async def test_real_mcp_add_model_verify_and_strip_provider_secrets(stub):
|
|||
|
||||
|
||||
@pytest.mark.parametrize("options", [{"read_only": True}, {"allowed_tools": frozenset({"list_keys"})}])
|
||||
async def test_policy_blocks_hidden_write_even_when_called_directly(stub, options):
|
||||
async def test_policy_blocks_hidden_write_even_when_called_directly(stub: StubGateway, options: SessionOptions) -> None:
|
||||
async with session_for(stub, **options) as session:
|
||||
assert "create_key" not in {t.name for t in (await session.list_tools()).tools}
|
||||
assert (await session.call_tool("create_key", {"body": {"key_alias": "forbidden"}})).is_error
|
||||
assert not stub.writes
|
||||
|
||||
|
||||
async def test_unverified_user_cannot_discover_or_call(stub):
|
||||
async def test_unverified_user_cannot_discover_or_call(stub: StubGateway) -> None:
|
||||
async with session_for(stub, api_key="viewer") as session:
|
||||
with pytest.raises(MCPError):
|
||||
await session.list_tools()
|
||||
|
|
@ -68,21 +71,25 @@ async def test_unverified_user_cannot_discover_or_call(stub):
|
|||
"arguments",
|
||||
[{}, {"body": {}}, {"body": {"key_alias": 7}}, {"body": {"key_alias": "a"}, "headers": {"Authorization": "spoof"}}],
|
||||
)
|
||||
async def test_invalid_arguments_do_not_reach_management_api(stub, arguments):
|
||||
async def test_invalid_arguments_do_not_reach_management_api(
|
||||
stub: StubGateway, arguments: dict[str, JsonValue]
|
||||
) -> None:
|
||||
async with session_for(stub) as session:
|
||||
assert (await session.call_tool("create_key", arguments)).is_error
|
||||
assert not stub.writes
|
||||
|
||||
|
||||
@pytest.mark.parametrize("identifier", ["..", "a/b", "a\\b", "a\n", "a%2Fb", ""])
|
||||
async def test_path_parameters_cannot_escape_route(stub, identifier):
|
||||
async def test_path_parameters_cannot_escape_route(stub: StubGateway, identifier: str) -> None:
|
||||
async with session_for(stub) as session:
|
||||
assert (await session.call_tool("reset_key_spend", {"path": {"key": identifier}})).is_error
|
||||
assert not stub.writes
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status,timeout", [(307, False), (403, False), (500, False), (200, True)])
|
||||
async def test_failed_write_is_not_retried_or_redirected_and_does_not_echo_body(stub, status, timeout):
|
||||
async def test_failed_write_is_not_retried_or_redirected_and_does_not_echo_body(
|
||||
stub: StubGateway, status: int, timeout: bool
|
||||
) -> None:
|
||||
stub.status, stub.timeout = status, timeout
|
||||
async with session_for(stub) as session:
|
||||
result = await session.call_tool("create_key", {"body": {"key_alias": "one"}})
|
||||
|
|
@ -91,7 +98,27 @@ async def test_failed_write_is_not_retried_or_redirected_and_does_not_echo_body(
|
|||
assert len(stub.writes) == 1
|
||||
|
||||
|
||||
async def test_revocation_during_write_suppresses_returned_key(stub):
|
||||
async def test_successful_write_with_truncated_json_is_not_retried_or_exposed(stub: StubGateway) -> None:
|
||||
def handle(request: httpx2.Request) -> httpx2.Response:
|
||||
response: Final = stub.handle(request)
|
||||
if request.url.path == "/key/generate":
|
||||
return httpx2.Response(200, content=response.content[:-1])
|
||||
return response
|
||||
|
||||
async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client:
|
||||
gateway: Final = Gateway(Config("https://gateway.example.com"), client)
|
||||
async with connected(gateway) as session:
|
||||
result: Final = await session.call_tool("create_key", {"body": {"key_alias": "one"}})
|
||||
assert result.is_error
|
||||
message: Final = result.content[0]
|
||||
assert isinstance(message, types.TextContent)
|
||||
assert message.text == (
|
||||
"Gateway returned an invalid JSON response. Check gateway state before retrying this change."
|
||||
)
|
||||
assert len(stub.writes) == len(stub.keys) == 1
|
||||
|
||||
|
||||
async def test_revocation_during_write_suppresses_returned_key(stub: StubGateway) -> None:
|
||||
stub.revoke = True
|
||||
async with session_for(stub) as session:
|
||||
result = await session.call_tool("create_key", {"body": {"key_alias": "one"}})
|
||||
|
|
@ -99,7 +126,7 @@ async def test_revocation_during_write_suppresses_returned_key(stub):
|
|||
assert len(stub.writes) == 1
|
||||
|
||||
|
||||
async def test_operation_id_drift_fails_closed(stub):
|
||||
async def test_operation_id_drift_fails_closed(stub: StubGateway) -> None:
|
||||
stub.spec = copy.deepcopy(SPEC)
|
||||
stub.spec["paths"]["/model/new"]["post"]["operationId"] = "changed"
|
||||
async with session_for(stub) as session:
|
||||
|
|
@ -108,7 +135,7 @@ async def test_operation_id_drift_fails_closed(stub):
|
|||
assert not stub.writes
|
||||
|
||||
|
||||
def test_catalog_has_unique_stable_tools_and_preserves_existing_surface():
|
||||
def test_catalog_has_unique_stable_tools_and_preserves_existing_surface() -> None:
|
||||
assert len(OPERATIONS) == 65
|
||||
assert len({o.name for o in OPERATIONS}) == len(OPERATIONS)
|
||||
assert len({(o.method, o.path) for o in OPERATIONS}) == len(OPERATIONS)
|
||||
|
|
@ -124,12 +151,12 @@ def test_catalog_has_unique_stable_tools_and_preserves_existing_surface():
|
|||
"https://example.com:bad",
|
||||
],
|
||||
)
|
||||
def test_gateway_must_be_a_fixed_trusted_origin(url):
|
||||
def test_gateway_must_be_a_fixed_trusted_origin(url: str) -> None:
|
||||
with pytest.raises(ValueError, match=r"gateway|HTTPS"):
|
||||
Config(url)
|
||||
|
||||
|
||||
async def test_http_server_requires_each_request_credential(stub, monkeypatch):
|
||||
async def test_http_server_requires_each_request_credential(stub: StubGateway, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("LITELLM_API_KEY", "alice-admin")
|
||||
async with httpx2.AsyncClient(transport=httpx2.MockTransport(stub.handle)) as client:
|
||||
app = create_http_app(Gateway(Config("https://gateway.example.com"), client))
|
||||
|
|
@ -150,7 +177,7 @@ async def test_http_server_requires_each_request_credential(stub, monkeypatch):
|
|||
assert {request.headers["authorization"] for request in stub.requests} == {"Bearer alice-admin", "Bearer bob-admin"}
|
||||
|
||||
|
||||
def test_recursive_refs_and_overridden_parameters_have_valid_schema():
|
||||
def test_recursive_refs_and_overridden_parameters_have_valid_schema() -> None:
|
||||
import jsonschema
|
||||
|
||||
spec = {
|
||||
|
|
@ -178,7 +205,7 @@ def test_recursive_refs_and_overridden_parameters_have_valid_schema():
|
|||
validator.validate({"query": {"limit": "wrong"}, "body": {}})
|
||||
|
||||
|
||||
def test_schema_names_with_json_pointer_escapes():
|
||||
def test_schema_names_with_json_pointer_escapes() -> None:
|
||||
import jsonschema
|
||||
|
||||
spec = {"components": {"schemas": {"A/B~C": {"type": "integer"}}}}
|
||||
|
|
@ -194,12 +221,12 @@ def test_schema_names_with_json_pointer_escapes():
|
|||
jsonschema.Draft202012Validator(schema).validate({"body": "wrong"})
|
||||
|
||||
|
||||
async def test_shared_http_client_never_replays_another_callers_gateway_cookie(stub):
|
||||
async def test_shared_http_client_never_replays_another_callers_gateway_cookie(stub: StubGateway) -> None:
|
||||
from litellm.proxy.admin_mcp.gateway import Gateway
|
||||
|
||||
original = stub.handle
|
||||
|
||||
def handle(request):
|
||||
def handle(request: httpx2.Request) -> httpx2.Response:
|
||||
assert not request.headers.get("cookie")
|
||||
result = original(request)
|
||||
result.headers["set-cookie"] = "gateway-session=first-caller; Path=/"
|
||||
|
|
@ -213,7 +240,7 @@ async def test_shared_http_client_never_replays_another_callers_gateway_cookie(s
|
|||
|
||||
|
||||
@pytest.mark.parametrize("container", ["object", "array", "encoded"])
|
||||
def test_response_redaction_fails_closed_for_excessive_nesting(container):
|
||||
def test_response_redaction_fails_closed_for_excessive_nesting(container: str) -> None:
|
||||
value = {"echo": "personal-admin-credential", "password": "provider-secret"}
|
||||
for _ in range(MAX_REDACTION_DEPTH):
|
||||
value = [value] if container == "array" else {"metadata": value}
|
||||
|
|
@ -222,7 +249,7 @@ def test_response_redaction_fails_closed_for_excessive_nesting(container):
|
|||
safe_result(payload, "personal-admin-credential")
|
||||
|
||||
|
||||
def test_response_redaction_handles_encoded_secrets_and_preserves_other_values():
|
||||
def test_response_redaction_handles_encoded_secrets_and_preserves_other_values() -> None:
|
||||
credential = "personal-admin-credential"
|
||||
value = {"metadata": json.dumps({"nested": [{"echo": credential, "password": "provider-secret", "count": 7}]})}
|
||||
assert safe_result(value, credential) == {
|
||||
|
|
@ -238,6 +265,6 @@ def test_response_redaction_handles_encoded_secrets_and_preserves_other_values()
|
|||
'[{"password":"provider-secret"}',
|
||||
],
|
||||
)
|
||||
def test_response_redaction_rejects_unparseable_json_containers(payload):
|
||||
def test_response_redaction_rejects_unparseable_json_containers(payload: str) -> None:
|
||||
with pytest.raises(AdminError, match="Inspect gateway state before retrying"):
|
||||
safe_result({"metadata": payload}, "personal-admin-credential")
|
||||
|
|
|
|||
|
|
@ -1,19 +1,21 @@
|
|||
from collections.abc import AsyncIterator
|
||||
import copy
|
||||
import json
|
||||
import re
|
||||
|
||||
import httpx2
|
||||
import jsonschema
|
||||
from pydantic import JsonValue
|
||||
import pytest
|
||||
|
||||
from tests.unit.proxy.admin_mcp.conftest import SPEC, connected, data, session_for
|
||||
from tests.unit.proxy.admin_mcp.conftest import SPEC, SessionOptions, StubGateway, connected, data, session_for
|
||||
from litellm.proxy.admin_mcp.catalog import OPERATIONS
|
||||
from litellm.proxy.admin_mcp.config import Config
|
||||
from litellm.proxy.admin_mcp.gateway import AdminError, Gateway, HELPERS
|
||||
from litellm.proxy.admin_mcp.results import ResultStore
|
||||
|
||||
|
||||
async def test_discovery_keeps_named_tools_annotations_and_complete_validation(stub):
|
||||
async def test_discovery_keeps_named_tools_annotations_and_complete_validation(stub: StubGateway) -> None:
|
||||
async with session_for(stub, schema_mode="discovery") as session:
|
||||
listed = {t.name: t for t in (await session.list_tools()).tools}
|
||||
assert not listed["create_key"].annotations.read_only_hint
|
||||
|
|
@ -28,7 +30,7 @@ async def test_discovery_keeps_named_tools_annotations_and_complete_validation(s
|
|||
assert result["max_budget"] == 0
|
||||
|
||||
|
||||
async def test_full_schema_exposes_body_fields_without_losing_nested_definitions(stub):
|
||||
async def test_full_schema_exposes_body_fields_without_losing_nested_definitions(stub: StubGateway) -> None:
|
||||
async with session_for(stub) as session:
|
||||
tools = {t.name: t for t in (await session.list_tools()).tools}
|
||||
key_schema = tools["create_key"].input_schema
|
||||
|
|
@ -40,7 +42,7 @@ async def test_full_schema_exposes_body_fields_without_losing_nested_definitions
|
|||
assert "Deployment" not in model_schema["$defs"]
|
||||
|
||||
|
||||
async def test_schema_itself_can_be_paged_and_retrieved_in_full(stub):
|
||||
async def test_schema_itself_can_be_paged_and_retrieved_in_full(stub: StubGateway) -> None:
|
||||
stub.spec = copy.deepcopy(SPEC)
|
||||
stub.spec["components"]["schemas"]["Key"]["properties"]["key_alias"]["description"] = (
|
||||
"Long but useful documentation. " * 2000
|
||||
|
|
@ -62,13 +64,15 @@ async def test_schema_itself_can_be_paged_and_retrieved_in_full(stub):
|
|||
({"type": "array", "items": {"type": "string"}}, ["uncommon", "argument"]),
|
||||
],
|
||||
)
|
||||
async def test_discovery_preserves_non_object_body_shapes(stub, schema, value):
|
||||
async def test_discovery_preserves_non_object_body_shapes(
|
||||
stub: StubGateway, schema: dict[str, JsonValue], value: JsonValue
|
||||
) -> None:
|
||||
stub.spec = copy.deepcopy(SPEC)
|
||||
stub.spec["components"]["schemas"]["Key"] = schema
|
||||
original = stub.handle
|
||||
seen = []
|
||||
|
||||
def handle(request):
|
||||
def handle(request: httpx2.Request) -> httpx2.Response:
|
||||
if request.url.path == "/key/generate":
|
||||
seen.append(json.loads(request.content))
|
||||
return httpx2.Response(200, json={"accepted": seen[-1]})
|
||||
|
|
@ -82,14 +86,14 @@ async def test_discovery_preserves_non_object_body_shapes(stub, schema, value):
|
|||
assert seen == [value]
|
||||
|
||||
|
||||
async def test_post_reads_are_available_in_read_only_mode_but_health_is_not(stub):
|
||||
async def test_post_reads_are_available_in_read_only_mode_but_health_is_not(stub: StubGateway) -> None:
|
||||
stub.spec = copy.deepcopy(SPEC)
|
||||
for op in OPERATIONS:
|
||||
if op.name in {"get_budgets", "get_organization_legacy", "check_key_health"}:
|
||||
stub.spec["paths"][op.path] = {op.method.lower(): {"operationId": op.operation_id}}
|
||||
original = stub.handle
|
||||
|
||||
def handle(request):
|
||||
def handle(request: httpx2.Request) -> httpx2.Response:
|
||||
if request.url.path in {"/budget/info", "/organization/info"}:
|
||||
return httpx2.Response(200, json=[])
|
||||
return original(request)
|
||||
|
|
@ -104,7 +108,7 @@ async def test_post_reads_are_available_in_read_only_mode_but_health_is_not(stub
|
|||
assert data(await session.call_tool(name, {})) == []
|
||||
|
||||
|
||||
async def test_full_mode_is_unchanged_and_compact_pages_do_not_refetch(stub):
|
||||
async def test_full_mode_is_unchanged_and_compact_pages_do_not_refetch(stub: StubGateway) -> None:
|
||||
stub.keys = [
|
||||
{"id": i, "max_budget": None, "spend": 0, "blocked": False, "extra": {"rare_setting": "z" * 8000}}
|
||||
for i in range(15)
|
||||
|
|
@ -121,7 +125,7 @@ async def test_full_mode_is_unchanged_and_compact_pages_do_not_refetch(stub):
|
|||
assert sum(r.url.path == "/key/list" for r in stub.requests) == count
|
||||
|
||||
|
||||
async def test_compact_default_keeps_small_results_and_full_override(stub):
|
||||
async def test_compact_default_keeps_small_results_and_full_override(stub: StubGateway) -> None:
|
||||
async with session_for(stub, response_view="compact") as session:
|
||||
assert data(await session.call_tool("list_keys", {})) == {"keys": []}
|
||||
stub.keys = [{"id": i, "metadata": "a" * 1000} for i in range(40)]
|
||||
|
|
@ -130,7 +134,7 @@ async def test_compact_default_keeps_small_results_and_full_override(stub):
|
|||
|
||||
|
||||
@pytest.mark.parametrize("name", ["list_models", "get_model"])
|
||||
async def test_model_catalog_identifiers_allow_lossless_paging(stub, name):
|
||||
async def test_model_catalog_identifiers_allow_lossless_paging(stub: StubGateway, name: str) -> None:
|
||||
operation = next(op for op in OPERATIONS if op.name == name)
|
||||
stub.spec = copy.deepcopy(SPEC)
|
||||
stub.spec["paths"][operation.path] = {"get": {"operationId": operation.operation_id}}
|
||||
|
|
@ -150,7 +154,7 @@ async def test_model_catalog_identifiers_allow_lossless_paging(stub, name):
|
|||
expected = {"data": models, "total_count": len(models), "current_page": 1, "total_pages": 1}
|
||||
original = stub.handle
|
||||
|
||||
def handle(request):
|
||||
def handle(request: httpx2.Request) -> httpx2.Response:
|
||||
if request.url.path == operation.path:
|
||||
stub.requests.append(request)
|
||||
return httpx2.Response(200, json=expected)
|
||||
|
|
@ -177,11 +181,11 @@ async def test_model_catalog_identifiers_allow_lossless_paging(stub, name):
|
|||
assert data(await session.call_tool(name, {"response": {"view": "full"}})) == expected
|
||||
|
||||
|
||||
async def test_full_write_result_can_be_read_without_repeating_write(stub):
|
||||
async def test_full_write_result_can_be_read_without_repeating_write(stub: StubGateway) -> None:
|
||||
original_handler = stub.handle
|
||||
payload = {"status": "ok", "previous": {"spend": 20, "custom": "x" * 22000}}
|
||||
|
||||
def handle(request):
|
||||
def handle(request: httpx2.Request) -> httpx2.Response:
|
||||
response = original_handler(request)
|
||||
return httpx2.Response(200, json=payload) if request.url.path.endswith("/reset_spend") else response
|
||||
|
||||
|
|
@ -193,10 +197,10 @@ async def test_full_write_result_can_be_read_without_repeating_write(stub):
|
|||
assert len(stub.writes) == 1
|
||||
|
||||
|
||||
async def test_large_new_key_response_is_delivered_inline_and_never_cached(stub):
|
||||
async def test_large_new_key_response_is_delivered_inline_and_never_cached(stub: StubGateway) -> None:
|
||||
original = stub.handle
|
||||
|
||||
def handle(request):
|
||||
def handle(request: httpx2.Request) -> httpx2.Response:
|
||||
response = original(request)
|
||||
if request.url.path == "/key/generate":
|
||||
return httpx2.Response(200, json={**response.json(), "metadata": "x" * 20000})
|
||||
|
|
@ -212,7 +216,7 @@ async def test_large_new_key_response_is_delivered_inline_and_never_cached(stub)
|
|||
assert len(stub.writes) == 1
|
||||
|
||||
|
||||
async def test_result_reads_require_same_credential_principal_and_current_role(stub):
|
||||
async def test_result_reads_require_same_credential_principal_and_current_role(stub: StubGateway) -> None:
|
||||
stub.keys = [{"data": "x" * 20000}]
|
||||
async with httpx2.AsyncClient(transport=httpx2.MockTransport(stub.handle)) as client:
|
||||
gateway = Gateway(Config("https://gateway.example.com", response_view="compact"), client)
|
||||
|
|
@ -225,11 +229,11 @@ async def test_result_reads_require_same_credential_principal_and_current_role(s
|
|||
await gateway.call("read_admin_result", result["full_read"], "alice-admin")
|
||||
|
||||
|
||||
async def test_revocation_during_result_read_suppresses_snapshot(stub):
|
||||
async def test_revocation_during_result_read_suppresses_snapshot(stub: StubGateway) -> None:
|
||||
stub.keys = [{"data": "private" * 4000}]
|
||||
revoke = False
|
||||
|
||||
def clock():
|
||||
def clock() -> float:
|
||||
if revoke:
|
||||
stub.revoked.add("alice-admin")
|
||||
return 0.0
|
||||
|
|
@ -244,7 +248,7 @@ async def test_revocation_during_result_read_suppresses_snapshot(stub):
|
|||
await gateway.call("read_admin_result", result["full_read"], "alice-admin")
|
||||
|
||||
|
||||
async def test_schema_helper_respects_allowlist_and_cannot_execute_hidden_operations(stub):
|
||||
async def test_schema_helper_respects_allowlist_and_cannot_execute_hidden_operations(stub: StubGateway) -> None:
|
||||
async with session_for(stub, allowed_tools=frozenset({"list_keys"}), schema_mode="discovery") as session:
|
||||
for name in ("create_key", "flush_everything", "read_admin_result"):
|
||||
assert (await session.call_tool("describe_admin_tool", {"name": name})).is_error
|
||||
|
|
@ -253,7 +257,7 @@ async def test_schema_helper_respects_allowlist_and_cannot_execute_hidden_operat
|
|||
assert not stub.writes
|
||||
|
||||
|
||||
async def test_warm_call_retains_pre_post_checks_without_discovery_auth(stub):
|
||||
async def test_warm_call_retains_pre_post_checks_without_discovery_auth(stub: StubGateway) -> None:
|
||||
async with session_for(stub) as session:
|
||||
await session.list_tools()
|
||||
stub.requests.clear()
|
||||
|
|
@ -261,11 +265,11 @@ async def test_warm_call_retains_pre_post_checks_without_discovery_auth(stub):
|
|||
assert [r.url.path for r in stub.requests] == ["/v2/user/info", "/key/list", "/v2/user/info"]
|
||||
|
||||
|
||||
async def test_identity_change_during_cold_discovery_prevents_write(stub):
|
||||
async def test_identity_change_during_cold_discovery_prevents_write(stub: StubGateway) -> None:
|
||||
original = stub.handle
|
||||
changed = False
|
||||
|
||||
def handle(request):
|
||||
def handle(request: httpx2.Request) -> httpx2.Response:
|
||||
nonlocal changed
|
||||
result = original(request)
|
||||
if request.url.path == "/openapi.json":
|
||||
|
|
@ -284,11 +288,11 @@ async def test_identity_change_during_cold_discovery_prevents_write(stub):
|
|||
assert not stub.writes
|
||||
|
||||
|
||||
async def test_capacity_falls_back_to_full_write_response_without_retry(stub):
|
||||
async def test_capacity_falls_back_to_full_write_response_without_retry(stub: StubGateway) -> None:
|
||||
original = stub.handle
|
||||
payload = {"status": "ok", "previous": "x" * 20000}
|
||||
|
||||
def handle(request):
|
||||
def handle(request: httpx2.Request) -> httpx2.Response:
|
||||
result = original(request)
|
||||
return httpx2.Response(200, json=payload) if request.url.path.endswith("/reset_spend") else result
|
||||
|
||||
|
|
@ -300,21 +304,21 @@ async def test_capacity_falls_back_to_full_write_response_without_retry(stub):
|
|||
assert len(stub.writes) == 1
|
||||
|
||||
|
||||
async def test_views_support_large_results_and_stop_at_stream_limit(stub):
|
||||
async def test_views_support_large_results_and_stop_at_stream_limit(stub: StubGateway) -> None:
|
||||
original = stub.handle
|
||||
|
||||
class Chunks(httpx2.AsyncByteStream):
|
||||
def __init__(self):
|
||||
def __init__(self) -> None:
|
||||
self.count = 0
|
||||
|
||||
async def __aiter__(self):
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
for _ in range(30):
|
||||
self.count += 1
|
||||
yield b"x" * 1_000_000
|
||||
|
||||
stream = Chunks()
|
||||
|
||||
def handle(request):
|
||||
def handle(request: httpx2.Request) -> httpx2.Response:
|
||||
if request.url.path == "/key/list":
|
||||
if request.url.params.get("page") == "2":
|
||||
return httpx2.Response(200, stream=stream)
|
||||
|
|
@ -331,7 +335,7 @@ async def test_views_support_large_results_and_stop_at_stream_limit(stub):
|
|||
assert stream.count == 17
|
||||
|
||||
|
||||
async def test_full_schemas_support_compact_large_results(stub):
|
||||
async def test_full_schemas_support_compact_large_results(stub: StubGateway) -> None:
|
||||
config = Config("https://gateway.example.com")
|
||||
assert config.response_view == "compact" and config.schema_mode == "full"
|
||||
stub.keys = [{"data": "x" * 20000}]
|
||||
|
|
@ -342,7 +346,7 @@ async def test_full_schemas_support_compact_large_results(stub):
|
|||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["full", "discovery"])
|
||||
async def test_every_reviewed_operation_retains_mapping_advanced_inputs_and_result(mode):
|
||||
async def test_every_reviewed_operation_retains_mapping_advanced_inputs_and_result(mode: str) -> None:
|
||||
spec = {
|
||||
"paths": {},
|
||||
"components": {
|
||||
|
|
@ -382,13 +386,11 @@ async def test_every_reviewed_operation_retains_mapping_advanced_inputs_and_resu
|
|||
calls = []
|
||||
result = {"budget": None, "spend": 0, "blocked": False, "metadata": {"unexpected_field": [1, "a", None]}}
|
||||
|
||||
def handle(request):
|
||||
def handle(request: httpx2.Request) -> httpx2.Response:
|
||||
if request.url.path == "/openapi.json":
|
||||
return httpx2.Response(200, json=spec)
|
||||
if request.url.path == "/v2/user/info" and "litellm-changed-by" not in request.headers:
|
||||
return httpx2.Response(
|
||||
200, json={"user_id": "admin", "user_role": "proxy_admin"}
|
||||
)
|
||||
return httpx2.Response(200, json={"user_id": "admin", "user_role": "proxy_admin"})
|
||||
calls.append(request)
|
||||
return httpx2.Response(200, json=result)
|
||||
|
||||
|
|
@ -425,6 +427,6 @@ async def test_every_reviewed_operation_retains_mapping_advanced_inputs_and_resu
|
|||
|
||||
|
||||
@pytest.mark.parametrize("config", [{"schema_mode": "unknown"}, {"response_view": "unknown"}])
|
||||
def test_unknown_modes_fail_instead_of_silently_changing_behavior(config):
|
||||
def test_unknown_modes_fail_instead_of_silently_changing_behavior(config: SessionOptions) -> None:
|
||||
with pytest.raises(ValueError, match="LITELLM_ADMIN_"):
|
||||
Config("https://gateway.example.com", **config)
|
||||
|
|
|
|||
|
|
@ -1,14 +1,25 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.admin_mcp.results import ResultError, ResultStore, contains_credential, encode, page, select
|
||||
from litellm.proxy.admin_mcp._types import INTEGER, JSON_ARRAY, object_value, string_value
|
||||
from litellm.proxy.admin_mcp.results import (
|
||||
ResultError,
|
||||
ResultStore,
|
||||
Snapshot,
|
||||
contains_credential,
|
||||
encode,
|
||||
page,
|
||||
select,
|
||||
)
|
||||
|
||||
|
||||
def rebuild(item, token, path=""):
|
||||
def rebuild(item: Snapshot, token: str, path: str = "") -> JsonValue:
|
||||
"""Use only the public index/continuations, without knowing field names."""
|
||||
first = page(item, token, path=path, limit=3, max_chars=257)
|
||||
result = (
|
||||
first: Final = object_value(page(item, token, path=path, limit=3, max_chars=257))
|
||||
result: JsonValue = (
|
||||
{}
|
||||
if first["type"] == "object"
|
||||
else []
|
||||
|
|
@ -20,22 +31,37 @@ def rebuild(item, token, path=""):
|
|||
current = first
|
||||
while True:
|
||||
if first["type"] in {"object", "array"}:
|
||||
for entry in current["entries"]:
|
||||
value = entry["value"] if entry["complete"] else rebuild(item, token, entry["read"]["path"])
|
||||
for raw_entry in JSON_ARRAY.validate_python(current["entries"]):
|
||||
entry = object_value(raw_entry)
|
||||
value = (
|
||||
entry["value"]
|
||||
if entry["complete"]
|
||||
else rebuild(item, token, string_value(object_value(entry["read"])["path"]))
|
||||
)
|
||||
if isinstance(result, dict):
|
||||
result[entry["key"]] = value
|
||||
result[string_value(entry["key"])] = value
|
||||
else:
|
||||
assert isinstance(result, list), result
|
||||
result.append(value)
|
||||
elif first["type"] == "string":
|
||||
result += current["value"]
|
||||
assert isinstance(result, str), result
|
||||
result += string_value(current["value"])
|
||||
if not current.get("next"):
|
||||
return result
|
||||
args = dict(current["next"])
|
||||
args.pop("result_id")
|
||||
current = page(item, token, **args)
|
||||
args = object_value(current["next"])
|
||||
current = object_value(
|
||||
page(
|
||||
item,
|
||||
token,
|
||||
path=string_value(args["path"]),
|
||||
offset=INTEGER.validate_python(args["offset"], strict=True),
|
||||
limit=INTEGER.validate_python(args["limit"], strict=True),
|
||||
max_chars=INTEGER.validate_python(args["max_chars"], strict=True),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_index_and_pages_reconstruct_every_value_without_prior_field_knowledge():
|
||||
def test_index_and_pages_reconstruct_every_value_without_prior_field_knowledge() -> None:
|
||||
original = {
|
||||
"user_info": {"role": "proxy_admin", "max_budget": None, "spend": 0, "blocked": False},
|
||||
"keys": [
|
||||
|
|
@ -55,7 +81,7 @@ def test_index_and_pages_reconstruct_every_value_without_prior_field_knowledge()
|
|||
assert len(encode(first)) < len(item.raw) / 4
|
||||
|
||||
|
||||
def test_record_previews_include_all_small_fields_including_uncommon_ones():
|
||||
def test_record_previews_include_all_small_fields_including_uncommon_ones() -> None:
|
||||
value = [
|
||||
{
|
||||
"id": "team-1",
|
||||
|
|
@ -75,7 +101,7 @@ def test_record_previews_include_all_small_fields_including_uncommon_ones():
|
|||
assert entry["omitted_fields"] == ["keys"]
|
||||
|
||||
|
||||
def test_page_boundary_moves_whole_records_to_next_page_instead_of_fragmenting():
|
||||
def test_page_boundary_moves_whole_records_to_next_page_instead_of_fragmenting() -> None:
|
||||
value = [{"id": n, "data": "x" * 7000} for n in range(7)]
|
||||
store = ResultStore()
|
||||
token = store.put("list_keys", value, "alice", "user-1")
|
||||
|
|
@ -93,7 +119,7 @@ def test_page_boundary_moves_whole_records_to_next_page_instead_of_fragmenting()
|
|||
assert seen == value and calls == 4
|
||||
|
||||
|
||||
def test_snapshot_is_immutable_and_bound_to_both_credential_and_principal():
|
||||
def test_snapshot_is_immutable_and_bound_to_both_credential_and_principal() -> None:
|
||||
value = {"data": [1, 2]}
|
||||
store = ResultStore()
|
||||
token = store.put("list_keys", value, "alice-key", "user-1")
|
||||
|
|
@ -104,7 +130,7 @@ def test_snapshot_is_immutable_and_bound_to_both_credential_and_principal():
|
|||
store.get(token, credential, principal)
|
||||
|
||||
|
||||
def test_expiration_capacity_and_clear_do_not_evict_unexpired_results():
|
||||
def test_expiration_capacity_and_clear_do_not_evict_unexpired_results() -> None:
|
||||
now = [100]
|
||||
store = ResultStore(ttl=10, max_bytes=100, max_per_owner=1, clock=lambda: now[0], wall_clock=lambda: 1000)
|
||||
token = store.put("get_user", {"a": 1}, "alice", "user-1")
|
||||
|
|
@ -140,7 +166,7 @@ def test_expiration_capacity_and_clear_do_not_evict_unexpired_results():
|
|||
("list_request_logs", {"data": [{"model_info": {"key": "custom-secret"}}]}),
|
||||
],
|
||||
)
|
||||
def test_credentials_are_not_retained_in_result_store(source, value):
|
||||
def test_credentials_are_not_retained_in_result_store(source: str, value: JsonValue) -> None:
|
||||
store = ResultStore()
|
||||
assert store.put(source, value, "alice", "user-1") is None
|
||||
|
||||
|
|
@ -148,12 +174,12 @@ def test_credentials_are_not_retained_in_result_store(source, value):
|
|||
@pytest.mark.parametrize(
|
||||
"path", ["not-a-pointer", "/bad~2escape", "/array/-1", "/array/01", "/array/4", "/missing", "/scalar/x"]
|
||||
)
|
||||
def test_invalid_paths_fail_explicitly(path):
|
||||
def test_invalid_paths_fail_explicitly(path: str) -> None:
|
||||
with pytest.raises(ResultError):
|
||||
select({"array": [1], "scalar": 0}, path)
|
||||
|
||||
|
||||
def test_empty_values_and_pagination_edges():
|
||||
def test_empty_values_and_pagination_edges() -> None:
|
||||
store = ResultStore()
|
||||
for value in ([], {}, "", False, 0, None):
|
||||
token = store.put("get_user", value, "alice", "user-1")
|
||||
|
|
@ -164,7 +190,7 @@ def test_empty_values_and_pagination_edges():
|
|||
page(item, token, offset=1)
|
||||
|
||||
|
||||
def test_deep_result_paths_and_credential_scan_do_not_use_python_call_stack():
|
||||
def test_deep_result_paths_and_credential_scan_do_not_use_python_call_stack() -> None:
|
||||
value = {"password": "private-token"}
|
||||
for _ in range(1200):
|
||||
value = {"nested": value}
|
||||
|
|
@ -180,6 +206,6 @@ def test_deep_result_paths_and_credential_scan_do_not_use_python_call_stack():
|
|||
'[{"password":"provider-secret"}',
|
||||
],
|
||||
)
|
||||
def test_unparseable_json_containers_are_never_retained(payload):
|
||||
def test_unparseable_json_containers_are_never_retained(payload: str) -> None:
|
||||
assert contains_credential({"metadata": payload})
|
||||
assert ResultStore().put("get_user", {"metadata": payload}, "alice", "user-1") is None
|
||||
|
|
|
|||
|
|
@ -450,9 +450,9 @@ def test_admin_mcp_survives_only_management_component_lifespans(
|
|||
|
||||
application: Final = FastAPI(lifespan=lifespan)
|
||||
|
||||
@application.get("/user/info")
|
||||
@application.get("/v2/user/info")
|
||||
async def user_info() -> dict[str, object]:
|
||||
return {"user_id": "admin", "user_info": {"user_id": "admin", "user_role": "proxy_admin"}}
|
||||
return {"user_id": "admin", "user_role": "proxy_admin"}
|
||||
|
||||
if component_lifespan is not None:
|
||||
application.router.lifespan_context = partial(component_lifespan, lifespan=application.router.lifespan_context)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue