diff --git a/packaging/litellm-core/pyproject.toml b/packaging/litellm-core/pyproject.toml index 27681d9116c..1b802f7dc1b 100644 --- a/packaging/litellm-core/pyproject.toml +++ b/packaging/litellm-core/pyproject.toml @@ -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", ] diff --git a/tests/unit/proxy/admin_mcp/conftest.py b/tests/unit/proxy/admin_mcp/conftest.py index 4a0d9889e08..43f565f75d6 100644 --- a/tests/unit/proxy/admin_mcp/conftest.py +++ b/tests/unit/proxy/admin_mcp/conftest.py @@ -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) diff --git a/tests/unit/proxy/admin_mcp/test_gateway.py b/tests/unit/proxy/admin_mcp/test_gateway.py index cb2bd4f81cc..21349fe8ee6 100644 --- a/tests/unit/proxy/admin_mcp/test_gateway.py +++ b/tests/unit/proxy/admin_mcp/test_gateway.py @@ -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") diff --git a/tests/unit/proxy/admin_mcp/test_progressive.py b/tests/unit/proxy/admin_mcp/test_progressive.py index 334e9883c89..521be9cd051 100644 --- a/tests/unit/proxy/admin_mcp/test_progressive.py +++ b/tests/unit/proxy/admin_mcp/test_progressive.py @@ -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) diff --git a/tests/unit/proxy/admin_mcp/test_results.py b/tests/unit/proxy/admin_mcp/test_results.py index 35d6e788049..97550226b7f 100644 --- a/tests/unit/proxy/admin_mcp/test_results.py +++ b/tests/unit/proxy/admin_mcp/test_results.py @@ -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 diff --git a/tests/unit/proxy/test_component_allowlists.py b/tests/unit/proxy/test_component_allowlists.py index eb058d01770..d2460f93b23 100644 --- a/tests/unit/proxy/test_component_allowlists.py +++ b/tests/unit/proxy/test_component_allowlists.py @@ -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)