From 51ff15fb655d5631fdf7af34f2b592462a68c08b Mon Sep 17 00:00:00 2001 From: Tin Date: Mon, 26 Oct 2026 11:06:06 -0700 Subject: [PATCH] refactor: simplify built-in admin MCP and close dependency gaps --- Dockerfile | 2 - backend/Dockerfile | 2 - docker/Dockerfile.database | 2 - docker/Dockerfile.non_root | 2 - gateway/Dockerfile | 2 - litellm/proxy/admin_mcp/__init__.py | 4 +- litellm/proxy/admin_mcp/_types.py | 6 +- litellm/proxy/admin_mcp/config.py | 3 +- litellm/proxy/admin_mcp/gateway.py | 60 +++--- litellm/proxy/admin_mcp/results.py | 102 +++++---- litellm/proxy/admin_mcp/schema.py | 197 ++++++++++++------ litellm/proxy/admin_mcp/server.py | 28 +-- osv-scanner.toml | 5 - pyproject.toml | 12 +- .../code_coverage_tests/recursive_detector.py | 1 + tests/unit/proxy/admin_mcp/conftest.py | 18 +- tests/unit/proxy/admin_mcp/test_gateway.py | 57 ++++- .../unit/proxy/admin_mcp/test_progressive.py | 39 ++-- tests/unit/proxy/admin_mcp/test_results.py | 23 +- tests/unit/proxy/admin_mcp/test_schema.py | 172 +++++++++++++++ tests/unit/proxy/test_admin_mcp.py | 43 ++-- tests/unit/proxy/test_proxy_server.py | 4 +- uv.lock | 51 +++-- vscode-extension/package-lock.json | 18 +- vscode-extension/package.json | 5 + 25 files changed, 589 insertions(+), 269 deletions(-) create mode 100644 tests/unit/proxy/admin_mcp/test_schema.py diff --git a/Dockerfile b/Dockerfile index c46f1461984..f1b23f094e8 100644 --- a/Dockerfile +++ b/Dockerfile @@ -81,7 +81,6 @@ COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ - --group admin-mcp \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ @@ -104,7 +103,6 @@ RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui. RUN uv sync --frozen --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ - --group admin-mcp \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ diff --git a/backend/Dockerfile b/backend/Dockerfile index 46d65c432bb..f9e401a8c3c 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -71,7 +71,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ - --group admin-mcp \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ @@ -84,7 +83,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv sync --frozen --no-install-project --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ - --group admin-mcp \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 69d9c26b2a8..89efbbcd507 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -79,7 +79,6 @@ COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ - --group admin-mcp \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ @@ -102,7 +101,6 @@ RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui. RUN uv sync --frozen --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ - --group admin-mcp \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 54bf8c807aa..43ddffb3107 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -83,7 +83,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ - --group admin-mcp \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ @@ -111,7 +110,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ uv sync --frozen --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ - --group admin-mcp \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 390a5f5666f..8e3883b6ad1 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -87,7 +87,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ - --group admin-mcp \ --extra extra_proxy \ --extra semantic-router \ --extra bedrock-realtime \ @@ -100,7 +99,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv sync --frozen --no-install-project --no-default-groups --no-editable \ --extra proxy \ --extra proxy-runtime \ - --group admin-mcp \ --extra extra_proxy \ --extra semantic-router \ --extra bedrock-realtime \ diff --git a/litellm/proxy/admin_mcp/__init__.py b/litellm/proxy/admin_mcp/__init__.py index 11d85e48cde..35f0a50a2b2 100644 --- a/litellm/proxy/admin_mcp/__init__.py +++ b/litellm/proxy/admin_mcp/__init__.py @@ -144,8 +144,8 @@ async def admin_mcp_lifespan(app: FastAPI) -> AsyncGenerator[None, None]: ) except ImportError as exc: raise RuntimeError( - "Admin MCP requires Python 3.12+ and the admin-mcp dependency group. " - "Use a LiteLLM image that bundles it, or run uv sync --extra proxy --group admin-mcp." + "Admin MCP requires Python 3.12+ and the proxy dependencies. " + "Use a LiteLLM image that bundles them, or run uv sync --extra proxy." ) from exc configured_url: Final = os.environ.get("LITELLM_MCP_PUBLIC_URL") or os.environ.get("PROXY_BASE_URL", "") diff --git a/litellm/proxy/admin_mcp/_types.py b/litellm/proxy/admin_mcp/_types.py index 3c97b787073..e3f14bf72d9 100644 --- a/litellm/proxy/admin_mcp/_types.py +++ b/litellm/proxy/admin_mcp/_types.py @@ -11,8 +11,10 @@ STRING: Final = TypeAdapter(str) INTEGER: Final = TypeAdapter(int) -def object_value(value: object) -> JsonObject: - return JSON_OBJECT.validate_python(value, strict=True) +def object_value(value: JsonValue | JsonObject) -> JsonObject: + if isinstance(value, Mapping): + return value + raise ValueError("Expected a JSON object") def string_value(value: object) -> str: diff --git a/litellm/proxy/admin_mcp/config.py b/litellm/proxy/admin_mcp/config.py index 94bc6afa67c..6bcb705fbe3 100644 --- a/litellm/proxy/admin_mcp/config.py +++ b/litellm/proxy/admin_mcp/config.py @@ -1,6 +1,6 @@ import ipaddress import os -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Final from urllib.parse import urlsplit @@ -49,7 +49,6 @@ def env_bool(name: str) -> bool: @dataclass(frozen=True, slots=True) class Config: base_url: str - api_key: str = field(default="", repr=False) read_only: bool = False allowed_tools: frozenset[str] = frozenset() public_url: str = "" diff --git a/litellm/proxy/admin_mcp/gateway.py b/litellm/proxy/admin_mcp/gateway.py index 9b5ec8510a1..52719b6874d 100644 --- a/litellm/proxy/admin_mcp/gateway.py +++ b/litellm/proxy/admin_mcp/gateway.py @@ -13,7 +13,7 @@ import jsonschema from mcp import types from pydantic import JsonValue -from ._types import INTEGER, JSON_VALUE, JsonObject, object_value, string_value +from ._types import INTEGER, JSON_OBJECT, JSON_VALUE, JsonObject, object_value, string_value from .catalog import BY_NAME, OPERATIONS, Operation from .config import Config from .results import ( @@ -38,6 +38,7 @@ DESCRIBE_SCHEMA: Final[JsonObject] = { "response": dict(RESPONSE_SCHEMA), }, } +MAX_REDACTION_DEPTH: Final = 64 HELPERS: Final = { "describe_admin_tool": ( DESCRIBE_SCHEMA, @@ -78,10 +79,12 @@ def _model_params(item: JsonValue) -> JsonObject: return {} -def _safe_field(key: str, item: JsonValue, credential: str) -> JsonValue: +def _safe_field(key: str, item: JsonValue, credential: str, depth: int) -> JsonValue: if key.lower() == "litellm_params": parsed: Final = _model_params(item) - return safe_result({k: v for k, v in parsed.items() if k in {"model", "litellm_credential_name"}}, credential) + return safe_result( + {k: v for k, v in parsed.items() if k in {"model", "litellm_credential_name"}}, credential, depth=depth + ) if key.lower() in { "api_key", "token", @@ -95,22 +98,28 @@ def _safe_field(key: str, item: JsonValue, credential: str) -> JsonValue: "credential_values", } and not (key.lower() in {"api_key", "token"} and isinstance(item, str) and re.fullmatch(r"[a-f0-9]{64}", item)): return "[redacted]" - return safe_result(item, credential) + return safe_result(item, credential, depth=depth) -def safe_result(value: JsonValue, credential: str) -> JsonValue: +def safe_result(value: JsonValue, credential: str, *, depth: int = 0) -> JsonValue: + if depth > MAX_REDACTION_DEPTH: + raise AdminError("Response nesting exceeds the safety limit. Inspect gateway state before retrying a change.") if isinstance(value, dict): - return {key: _safe_field(key, item, credential) for key, item in value.items()} + return {key: _safe_field(key, item, credential, depth + 1) for key, item in value.items()} if isinstance(value, list): - return [safe_result(item, credential) for item in value] + return [safe_result(item, credential, depth=depth + 1) for item in value] if not isinstance(value, str): return value try: decoded: Final = JSON_VALUE.validate_json(value) except ValueError: + if value.lstrip().startswith(("{", "[")): + raise AdminError( + "Response contains unreadable JSON. Inspect gateway state before retrying a change." + ) from None return value.replace(credential, "[credential redacted]") if isinstance(decoded, (dict, list)): - return json.dumps(safe_result(decoded, credential)) + return json.dumps(safe_result(decoded, credential, depth=depth + 1)) return value.replace(credential, "[credential redacted]") @@ -136,18 +145,15 @@ def _response_view(arguments: JsonObject, default: str) -> str: class Gateway: - def __init__(self, config: Config, client: httpx2.AsyncClient | None = None) -> None: + def __init__(self, config: Config, client: httpx2.AsyncClient, *, results: ResultStore | None = None) -> None: self.config = config self.client = client - self._owns_client = client is None self._tools: Mapping[str, types.Tool] | None = None self._loaded_at = 0.0 self._schema_lock = asyncio.Lock() - self.results = ResultStore() + self.results = results if results is not None else ResultStore() async def __aenter__(self) -> "Gateway": - if self.client is None: - self.client = httpx2.AsyncClient(timeout=httpx2.Timeout(30, read=120), follow_redirects=False) return self async def __aexit__( @@ -157,9 +163,6 @@ class Gateway: traceback: TracebackType | None, ) -> None: self.results.clear() - if self._owns_client and self.client is not None: - await self.client.aclose() - self.client = None async def _request( self, @@ -173,7 +176,6 @@ class Gateway: params: tuple[tuple[str, str], ...] = (), content: bytes | None = None, ) -> JsonValue: - assert self.client is not None write: Final = method != "GET" if read_only is None else not read_only uncertain: Final = " Check gateway state before retrying this change." if write else "" chunks: Final = BytesIO() @@ -203,25 +205,19 @@ class Gateway: async def authorize(self, credential: str) -> str: try: - data: Final = await self._request("GET", "/user/info", credential) + data: Final = await self._request("GET", "/v2/user/info", credential) except AdminError as exc: - if exc.status != 404: + if exc.status not in {400, 404}: raise raise AdminError("Only current LiteLLM proxy admins may use these tools.", 403) from None - user: Final = data.get("user_info") if isinstance(data, dict) else None if ( not isinstance(data, dict) - or not isinstance(user, dict) - or user.get("user_role") != "proxy_admin" - or not isinstance(user.get("user_id"), str) - or not user["user_id"] - or data.get("user_id") != user["user_id"] - or user.get("blocked") - or user.get("deleted") - or user.get("is_active") is False + or data.get("user_role") != "proxy_admin" + or not isinstance(data.get("user_id"), str) + or not data["user_id"] ): raise AdminError("Only current LiteLLM proxy admins may use these tools.", 403) - return string_value(user["user_id"]) + return string_value(data["user_id"]) async def tools(self, credential: str) -> Mapping[str, types.Tool]: await self.authorize(credential) @@ -230,7 +226,7 @@ class Gateway: **{ name: tool.model_copy( update={ - "input_schema": dict(discovery_schema(object_value(tool.input_schema), name)), + "input_schema": dict(discovery_schema(JSON_OBJECT.validate_python(tool.input_schema), name)), "description": (tool.description or "") + " Call describe_admin_tool for the complete argument schema before use.", } @@ -312,7 +308,9 @@ class Gateway: if name not in available and name not in HELPERS: raise AdminError("This tool is not enabled for this connector.", 403) try: - schema: Final = HELPERS[name][0] if name in HELPERS else object_value(available[name].input_schema) + schema: Final = ( + HELPERS[name][0] if name in HELPERS else JSON_OBJECT.validate_python(available[name].input_schema) + ) validation_error: Final = next(jsonschema.Draft202012Validator(schema).iter_errors(arguments), None) if validation_error is not None: raise validation_error diff --git a/litellm/proxy/admin_mcp/results.py b/litellm/proxy/admin_mcp/results.py index 67f42253937..c0158c99564 100644 --- a/litellm/proxy/admin_mcp/results.py +++ b/litellm/proxy/admin_mcp/results.py @@ -4,10 +4,11 @@ import json import re import secrets import time -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterable, Iterator, Mapping from dataclasses import dataclass from functools import reduce -from typing import Final +from itertools import islice +from typing import Final, TypeAlias from pydantic import JsonValue @@ -17,6 +18,7 @@ PAGE_BYTES: Final = 16_000 ITEM_BYTES: Final = PAGE_BYTES - 1000 PREVIEW_BYTES: Final = 6_000 MAX_RESULT_BYTES: Final = 16_000_000 +CredentialNode: TypeAlias = tuple[tuple[str | int, ...], JsonValue] SECRET_RESULTS: Final = frozenset( {"create_key", "create_service_account_key", "regenerate_key", "regenerate_key_by_id"} ) @@ -34,14 +36,12 @@ def pointer(parent: str, key: str | int) -> str: return parent + "/" + str(key).replace("~", "~0").replace("/", "~1") -def _select(value: JsonValue, parts: tuple[str, ...]) -> JsonValue: - if not parts: - return value - part: Final = parts[0].replace("~1", "/").replace("~0", "~") +def _select(value: JsonValue, segment: str) -> JsonValue: + part: Final = segment.replace("~1", "/").replace("~0", "~") if isinstance(value, list) and re.fullmatch(r"0|[1-9][0-9]*", part): - return _select(value[int(part)], parts[1:]) + return value[int(part)] if isinstance(value, dict): - return _select(value[part], parts[1:]) + return value[part] raise ValueError("Invalid JSON Pointer segment") @@ -51,7 +51,7 @@ def select(value: JsonValue, path: str) -> JsonValue: if not path.startswith("/") or re.search(r"~(?![01])", path): raise ResultError("Use a JSON Pointer path returned by the result index.") try: - return _select(value, tuple(path[1:].split("/"))) + return reduce(_select, path[1:].split("/"), value) except (KeyError, IndexError, ValueError): raise ResultError("The result does not contain that path; browse its parent first.") from None @@ -101,23 +101,41 @@ def _credential_field(key: str, child: JsonValue, source: str, path: tuple[str | ) -def contains_credential(value: JsonValue, *, source: str = "", path: tuple[str | int, ...] = ()) -> bool: +def _children(value: JsonValue, path: tuple[str | int, ...]) -> Iterator[CredentialNode]: if isinstance(value, dict): - return any( - _credential_field(key, child, source, path) or contains_credential(child, source=source, path=(*path, key)) - for key, child in value.items() - ) - if isinstance(value, list): - return any(contains_credential(child, source=source, path=(*path, i)) for i, child in enumerate(value)) - if not isinstance(value, str): - return False - if re.search(r"\bsk-[A-Za-z0-9_+-]+", value): - return True - try: - nested: Final = JSON_VALUE.validate_json(value) - except ValueError: - return False - return contains_credential(nested, source=source, path=path) if isinstance(nested, (dict, list)) else False + for key, child in value.items(): + yield (*path, key), child + elif isinstance(value, list): + for index, child in enumerate(value): + yield (*path, index), child + + +def contains_credential(value: JsonValue, *, source: str = "") -> bool: + pending: Final[list[Iterator[CredentialNode]]] = [iter((((), value),))] # mutable-ok: depth-first iterator stack + while pending: + node = next(pending[-1], None) + if node is None: + pending.pop() + continue + path, child = node + if isinstance(child, dict) and any(_credential_field(key, item, source, path) for key, item in child.items()): + return True + if isinstance(child, (dict, list)): + pending.append(_children(child, path)) + continue + if not isinstance(child, str): + continue + if re.search(r"\bsk-[A-Za-z0-9_+-]+", child): + return True + try: + nested = JSON_VALUE.validate_json(child) + except ValueError: + if child.lstrip().startswith(("{", "[")): + return True + continue + if isinstance(nested, (dict, list)): + pending.append(iter(((path, nested),))) + return False @dataclass(frozen=True, slots=True) @@ -234,22 +252,20 @@ def _entry(key: str | int, child: JsonValue, path: str, token: str) -> JsonObjec def _entries( - children: tuple[tuple[str | int, JsonValue], ...], + children: Iterable[tuple[str | int, JsonValue]], path: str, token: str, - remaining: int = PAGE_BYTES, - first: bool = True, -) -> tuple[JsonObject, ...]: - if not children: - return () - key, child = children[0] - if not first and len(encode(child)) > ITEM_BYTES and remaining < PREVIEW_BYTES + 1000: - return () - entry: Final = _entry(key, child, path, token) - cost: Final = len(encode(entry)) - if not first and cost > remaining: - return () - return (entry, *_entries(children[1:], path, token, remaining - cost, False)) +) -> Iterator[JsonObject]: + remaining: int = PAGE_BYTES # rebind-ok: byte budget is consumed as the page generator advances + for index, (key, child) in enumerate(children): + if index and len(encode(child)) > ITEM_BYTES and remaining < PREVIEW_BYTES + 1000: + return + entry = _entry(key, child, path, token) + cost = len(encode(entry)) + if index and cost > remaining: + return + remaining -= cost + yield entry def page( @@ -309,8 +325,12 @@ def _page_contents( "complete": offset == 0 and end == len(value), } assert isinstance(value, (dict, list)) - children: Final = tuple(value.items()) if isinstance(value, dict) else tuple(enumerate(value)) - entries: Final = _entries(children[offset : offset + limit], path, token) + children: Final = ( + islice(value.items(), offset, offset + limit) + if isinstance(value, dict) + else ((index, value[index]) for index in range(offset, min(offset + limit, len(value)))) + ) + entries: Final = tuple(_entries(children, path, token)) last: Final = offset + len(entries) return last, { "entries": [dict(entry) for entry in entries], diff --git a/litellm/proxy/admin_mcp/schema.py b/litellm/proxy/admin_mcp/schema.py index 631f6cb3c5a..a305748514e 100644 --- a/litellm/proxy/admin_mcp/schema.py +++ b/litellm/proxy/admin_mcp/schema.py @@ -1,38 +1,59 @@ +import re from collections.abc import Iterable, Mapping from dataclasses import dataclass +from functools import reduce from typing import Final from pydantic import JsonValue from ._types import JSON_ARRAY, JsonObject, object_value, string_value +_SINGLE_SCHEMAS: Final = frozenset( + { + "additionalProperties", + "unevaluatedProperties", + "items", + "contains", + "propertyNames", + "not", + "if", + "then", + "else", + "unevaluatedItems", + "additionalItems", + "contentSchema", + } +) +_SCHEMA_MAPS: Final = frozenset({"properties", "patternProperties", "$defs", "definitions", "dependentSchemas"}) +_SCHEMA_ARRAYS: Final = frozenset({"allOf", "anyOf", "oneOf", "prefixItems"}) + class SchemaError(ValueError): pass -def _reference_target(spec: JsonObject, reference: str) -> JsonObject: - def descend(value: JsonObject, parts: tuple[str, ...]) -> JsonObject: - if not parts: - return value - return descend(object_value(value[parts[0].replace("~1", "/").replace("~0", "~")]), parts[1:]) +def _pointer_step(value: JsonObject, part: str) -> JsonObject: + return object_value(value[part.replace("~1", "/").replace("~0", "~")]) + +def _reference_target(spec: JsonObject, reference: str) -> JsonObject: try: - return descend(spec, tuple(reference[2:].split("/"))) + return reduce(_pointer_step, reference[2:].split("/"), spec) except (KeyError, ValueError, TypeError): raise SchemaError("Missing or invalid OpenAPI reference") from None -def resolve(spec: JsonObject, value: JsonObject, seen: frozenset[str] = frozenset()) -> JsonObject: - if "$ref" not in value: - return value - reference: Final = value["$ref"] - if not isinstance(reference, str) or not reference.startswith("#/components/") or reference in seen: - raise SchemaError("Unsupported OpenAPI reference") - target: Final = _reference_target(spec, reference) - return resolve( - spec, {**target, **{key: child for key, child in value.items() if key != "$ref"}}, seen | {reference} - ) +def resolve(spec: JsonObject, value: JsonObject) -> JsonObject: + current: JsonObject = value # rebind-ok: iterative reference resolution avoids Python stack growth + seen: Final[set[str]] = set() # mutable-ok: local visited set bounds reference-chain traversal + while "$ref" in current: + reference = current["$ref"] + if not isinstance(reference, str) or not reference.startswith("#/components/") or reference in seen: + raise SchemaError("Unsupported OpenAPI reference") + seen.add(reference) + target = _reference_target(spec, reference) + current = {**target, **{key: child for key, child in current.items() if key != "$ref"}} + return current def _parameters(spec: JsonObject, path_item: JsonObject, operation: JsonObject) -> Mapping[tuple[str, str], JsonObject]: @@ -71,70 +92,114 @@ def _body_schema(spec: JsonObject, operation: JsonObject) -> tuple[JsonValue, bo return object_value(media).get("schema", {"type": "object"}), bool(body.get("required")) +def _schema_children(value: JsonValue) -> Iterable[JsonValue]: + if not isinstance(value, dict): + return + for key, child in value.items(): + if key in _SINGLE_SCHEMAS: + yield child + elif key in _SCHEMA_MAPS and isinstance(child, dict): + yield from child.values() + elif key in _SCHEMA_ARRAYS and isinstance(child, list): + yield from child + + +def _rewritten_keyword(key: str, value: JsonValue, rewritten: Mapping[int, JsonValue]) -> JsonValue: + if key in _SINGLE_SCHEMAS: + return rewritten[id(value)] + if key in _SCHEMA_MAPS and isinstance(value, dict): + return {name: rewritten[id(child)] for name, child in value.items()} + if key in _SCHEMA_ARRAYS and isinstance(value, list): + return [rewritten[id(child)] for child in value] + return value + + +def _schema_reference(value: JsonObject) -> str | None: + if any(key in value for key in ("$dynamicRef", "$recursiveRef", "$id")): + raise SchemaError("Schema reference scopes and dynamic references are unsupported") + if "$ref" not in value: + return None + reference: Final = value["$ref"] + if not isinstance(reference, str) or not reference.startswith("#/components/schemas/"): + raise SchemaError("Only local schema references are supported") + pointer: Final = reference.removeprefix("#/components/schemas/") + if "/" in pointer or re.search(r"~(?![01])", pointer): + raise SchemaError("Only component schema references are supported") + return pointer + + @dataclass(frozen=True, slots=True) class Rewritten: value: JsonValue - references: frozenset[str] = frozenset() + references: frozenset[str] def _rewrite(value: JsonValue) -> Rewritten: - if isinstance(value, list): - children: Final = tuple(_rewrite(child) for child in value) - return Rewritten([child.value for child in children], frozenset[str]().union(*(c.references for c in children))) - if not isinstance(value, dict): - return Rewritten(value) - if "$ref" in value and ( - not isinstance(value["$ref"], str) or not value["$ref"].startswith("#/components/schemas/") - ): - raise SchemaError("Only local schema references are supported") - reference: Final = string_value(value["$ref"]).removeprefix("#/components/schemas/") if "$ref" in value else None - children_by_key: Final = {key: _rewrite(child) for key, child in value.items() if key not in {"$ref", "nullable"}} - result: Final = { - **{key: child.value for key, child in children_by_key.items()}, - **({"$ref": "#/$defs/" + reference} if reference is not None else {}), - } - references: Final = frozenset[str]().union(*(child.references for child in children_by_key.values())) - return Rewritten( - {"anyOf": [result, {"type": "null"}]} if value.get("nullable") else result, - references | ({reference} if reference is not None else frozenset[str]()), - ) + pending: Final[list[tuple[JsonValue, bool]]] = [(value, False)] # mutable-ok: iterative postorder + rewritten: Final[dict[int, JsonValue]] = {} # mutable-ok: local results let each schema node be assembled once + references: Final[set[str]] = set() # mutable-ok: local reference collection avoids copying a set per node + while pending: + node, ready = pending.pop() + if id(node) in rewritten: + continue + if not isinstance(node, dict): + rewritten[id(node)] = node + continue + if not ready: + pending.append((node, True)) + pending.extend((child, False) for child in _schema_children(node)) + continue + reference = _schema_reference(node) + if reference is not None: + references.add(reference) + result = { + **{ + key: _rewritten_keyword(key, child, rewritten) + for key, child in node.items() + if key not in {"$ref", "nullable"} + }, + **({"$ref": "#/$defs/" + reference} if reference is not None else {}), + } + rewritten[id(node)] = {"anyOf": [result, {"type": "null"}]} if node.get("nullable") else result + return Rewritten(rewritten[id(value)], frozenset(references)) -def _definitions(spec: JsonObject, references: frozenset[str], known: JsonObject) -> JsonObject: - missing: Final = references - known.keys() - if not missing: - return known - rewritten: Final = { - pointer: _rewrite(dict(_reference_target(spec, "#/components/schemas/" + pointer))) for pointer in missing - } - definitions: Final = { - **known, - **{pointer.replace("~1", "/").replace("~0", "~"): item.value for pointer, item in rewritten.items()}, - } - next_refs: Final = frozenset[str]().union(*(item.references for item in rewritten.values())) - encoded_known: Final = {name.replace("~", "~0").replace("/", "~1"): value for name, value in definitions.items()} - unresolved: Final = next_refs - encoded_known.keys() - return _definitions(spec, unresolved, definitions) if unresolved else definitions +def _definitions(spec: JsonObject, references: frozenset[str]) -> JsonObject: + pending: Final[list[str]] = list(references) # mutable-ok: worklist resolves each referenced definition once + definitions: Final[dict[str, JsonValue]] = {} # mutable-ok: visited definitions terminate cyclic schema graphs + while pending: + pointer = pending.pop() + name = pointer.replace("~1", "/").replace("~0", "~") + if name in definitions: + continue + item = _rewrite(dict(_reference_target(spec, "#/components/schemas/" + pointer))) + definitions[name] = item.value + pending.extend(item.references) + return definitions def _references(value: JsonValue) -> Iterable[str]: - if isinstance(value, list): - for child in value: - yield from _references(child) - elif isinstance(value, dict): - reference: Final = value.get("$ref") + pending: Final[list[JsonValue]] = [value] # mutable-ok: iterative schema walk + while pending: + node = pending.pop() + if not isinstance(node, dict): + continue + reference = node.get("$ref") if isinstance(reference, str) and reference.startswith("#/$defs/"): yield reference.removeprefix("#/$defs/").replace("~1", "/").replace("~0", "~") - for child in value.values(): - yield from _references(child) + pending.extend(_schema_children(node)) -def _reachable(value: JsonValue, definitions: JsonObject, seen: frozenset[str] = frozenset()) -> frozenset[str]: - names: Final = frozenset(_references(value)) - seen - if not names: - return seen - children: Final = [definitions[name] for name in names] - return _reachable(children, definitions, seen | names) +def _reachable(value: JsonValue, definitions: JsonObject) -> frozenset[str]: + pending: Final[list[str]] = list(_references(value)) # mutable-ok: worklist visits the reachable definition graph + seen: Final[set[str]] = set() # mutable-ok: local visited set terminates cycles without recursive calls + while pending: + name = pending.pop() + if name in seen: + continue + seen.add(name) + pending.extend(_references(definitions[name])) + return frozenset(seen) def tool_schema(spec: JsonObject, path_item: JsonObject, operation: JsonObject) -> JsonObject: @@ -155,7 +220,7 @@ def tool_schema(spec: JsonObject, path_item: JsonObject, operation: JsonObject) **({"required": required_groups} if required_groups else {}), } ) - definitions: Final = _definitions(spec, root.references, {}) + definitions: Final = _definitions(spec, root.references) result: Final = object_value(root.value) properties: Final = object_value(result["properties"]) body_value: Final = properties.get("body", {}) diff --git a/litellm/proxy/admin_mcp/server.py b/litellm/proxy/admin_mcp/server.py index 82bc79bdf7d..b658bf60c0c 100644 --- a/litellm/proxy/admin_mcp/server.py +++ b/litellm/proxy/admin_mcp/server.py @@ -1,7 +1,7 @@ """Standard MCP transports; no agent, model inference or Slack dependency.""" import json -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator, Callable, Mapping from contextlib import asynccontextmanager from typing import Final @@ -17,7 +17,7 @@ from starlette.types import ASGIApp, Message, Receive, Scope, Send from litellm._version import version as __version__ -from ._types import object_value +from ._types import JSON_OBJECT from .gateway import AdminError, Gateway, valid_credential @@ -28,15 +28,7 @@ def bearer(headers: Mapping[str, str]) -> str: return valid_credential(value) -def create_server(gateway: Gateway, *, http: bool = False) -> Server[None]: - def credential(context: ServerRequestContext[None]) -> str: - if http: - request: Final[object] = context.request - if not isinstance(request, Request): - raise AdminError("Caller authentication required.", 401) - return bearer(request.headers) - return valid_credential(gateway.config.api_key) - +def create_server(gateway: Gateway, credential: Callable[[ServerRequestContext[None]], str]) -> Server[None]: async def list_tools( context: ServerRequestContext[None], params: types.PaginatedRequestParams | None ) -> types.ListToolsResult: @@ -50,7 +42,8 @@ def create_server(gateway: Gateway, *, http: bool = False) -> Server[None]: context: ServerRequestContext[None], params: types.CallToolRequestParams ) -> types.CallToolResult: try: - result: Final = await gateway.call(params.name, object_value(params.arguments or {}), credential(context)) + arguments: Final = JSON_OBJECT.validate_python(params.arguments or {}, strict=True) + result: Final = await gateway.call(params.name, arguments, credential(context)) return types.CallToolResult( content=[ types.TextContent(type="text", text=json.dumps(result, ensure_ascii=False, separators=(",", ":"))) @@ -113,10 +106,11 @@ class RequireAdmin: def create_http_app(gateway: Gateway) -> Starlette: - if gateway.config.api_key: - raise ValueError( - "Do not set LITELLM_API_KEY on an HTTP server; each client must send its own bearer credential." - ) + def credential(context: ServerRequestContext[None]) -> str: + request: Final[object] = context.request + if not isinstance(request, Request): + raise AdminError("Caller authentication required.", 401) + return bearer(request.headers) async def health(request: Request) -> JSONResponse: return JSONResponse({"status": "ok", "version": __version__}) @@ -127,7 +121,7 @@ def create_http_app(gateway: Gateway) -> Starlette: allowed_hosts=["localhost:*", "127.0.0.1:*", "[::1]:*", *([public.split("://", 1)[1]] if public else [])], allowed_origins=["http://localhost:*", "http://127.0.0.1:*", "http://[::1]:*", *([public] if public else [])], ) - app: Final = create_server(gateway, http=True).streamable_http_app( + app: Final = create_server(gateway, credential).streamable_http_app( stateless_http=True, json_response=True, max_request_body_size=256_000, diff --git a/osv-scanner.toml b/osv-scanner.toml index 33e2f614edc..bb802137f6e 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -3,11 +3,6 @@ id = "GHSA-w8v5-vhqr-4h9v" ignoreUntil = 2026-11-01 reason = "diskcache has no fixed release published; remove this entry once one exists" -[[IgnoredVulns]] -id = "GHSA-v4x8-3wr5-3wq8" -ignoreUntil = 2026-11-20 -reason = "the only fixed filelock is 4.2.1, a new major that requires Python >= 3.11 while litellm still supports 3.10, and the advisory is in ReadWriteLock, which litellm never uses (only FileLock); bump filelock and remove this entry once 3.10 is dropped or a 3.x fix ships" - [[IgnoredVulns]] id = "GHSA-vfj7-8cjw-p6xm" ignoreUntil = 2026-11-03 diff --git a/pyproject.toml b/pyproject.toml index 3fa50128a6c..50238423def 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,9 @@ dependencies = [ # When changing a floor, verify it installs + imports on every supported # Python with: `uv pip install --resolution=lowest-direct .` "fastuuid>=0.14.0,<1.0", - "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-dotenv>=1.0.0,<2.0", @@ -34,7 +36,7 @@ dependencies = [ "pydantic>=2.11.0,<3.0.0; python_version < '3.14'", "pydantic>=2.12.0,<3.0.0; python_version >= '3.14'", "pydantic-settings>=2.14.1,<3.0", - "jsonschema>=4.0.0,<5.0", + "jsonschema>=4.22.0,<5.0", "boto3>=1.43.117,<2.0", "certifi>=2026.10.19", "typing-extensions>=4.13.0,<5.0", @@ -246,13 +248,7 @@ e2e-dev = [ "psutil==7.2.2", "mcp>=2.3.0,<3", ] -admin-mcp = [ - "mcp>=2.3.0,<3", - "httpx2>=2.10.0,<3", - "jsonschema>=4.22.0,<5", -] proxy-dev = [ - { include-group = "admin-mcp" }, "prisma==0.11.0", "hypercorn==0.17.3", "prometheus-client==0.20.0", diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 7610d602fff..6378e3daeb7 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -2,6 +2,7 @@ import ast import os IGNORE_FUNCTIONS = [ + "safe_result", # Admin MCP redaction fails closed at MAX_REDACTION_DEPTH, including JSON embedded in strings. "_json_cost", # bounded at depth 32 and consumed under byte/node limits. "_format_type", "remove_additional_properties", diff --git a/tests/unit/proxy/admin_mcp/conftest.py b/tests/unit/proxy/admin_mcp/conftest.py index 08965f9ca0e..4a0d9889e08 100644 --- a/tests/unit/proxy/admin_mcp/conftest.py +++ b/tests/unit/proxy/admin_mcp/conftest.py @@ -100,15 +100,15 @@ class StubGateway: def handle(self, request): self.requests.append(request) token = request.headers.get("authorization", "").removeprefix("Bearer ") - if request.url.path == "/user/info" and token == "team-key": - return httpx2.Response(404, json={"error": {"message": "User None not found"}}) - if request.url.path == "/user/info": + 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 = ( "proxy_admin" if token in {"alice-admin", "bob-admin"} and token not in self.revoked else "internal_user" ) - return httpx2.Response(200, json={"user_id": token, "user_info": {"user_id": token, "user_role": role}}) + return httpx2.Response(200, json={"user_id": token, "user_role": role}) if request.url.path == "/openapi.json": return httpx2.Response(200, json=self.spec) assert token in {"alice-admin", "bob-admin"} @@ -158,8 +158,8 @@ def stub(): @asynccontextmanager -async def connected(gateway): - server = create_server(gateway) +async def connected(gateway, api_key="alice-admin"): + server = 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,10 +170,10 @@ async def connected(gateway): @asynccontextmanager -async def session_for(stub, **overrides): - config = Config("https://gateway.example.com", **{"api_key": "alice-admin", **overrides}) +async def session_for(stub, api_key="alice-admin", **overrides): + config = Config("https://gateway.example.com", **overrides) async with httpx2.AsyncClient(transport=httpx2.MockTransport(stub.handle)) as client: - async with connected(Gateway(config, client)) as session: + async with connected(Gateway(config, client), api_key) as session: yield session diff --git a/tests/unit/proxy/admin_mcp/test_gateway.py b/tests/unit/proxy/admin_mcp/test_gateway.py index e0dfd95af8e..cb2bd4f81cc 100644 --- a/tests/unit/proxy/admin_mcp/test_gateway.py +++ b/tests/unit/proxy/admin_mcp/test_gateway.py @@ -8,7 +8,7 @@ import pytest from tests.unit.proxy.admin_mcp.conftest import SPEC, 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 Gateway, HELPERS +from litellm.proxy.admin_mcp.gateway import AdminError, Gateway, HELPERS, MAX_REDACTION_DEPTH, safe_result from litellm.proxy.admin_mcp.schema import tool_schema from litellm.proxy.admin_mcp.server import create_http_app @@ -61,7 +61,7 @@ async def test_unverified_user_cannot_discover_or_call(stub): with pytest.raises(MCPError): await session.list_tools() assert (await session.call_tool("create_key", {"body": {"key_alias": "forbidden"}})).is_error - assert all(r.url.path == "/user/info" for r in stub.requests) + assert all(r.url.path == "/v2/user/info" for r in stub.requests) @pytest.mark.parametrize( @@ -129,9 +129,25 @@ def test_gateway_must_be_a_fixed_trusted_origin(url): Config(url) -def test_http_server_refuses_shared_process_credential(): - with pytest.raises(ValueError, match="Do not set"): - create_http_app(Gateway(Config("https://gateway.example.com", api_key="shared"))) +async def test_http_server_requires_each_request_credential(stub, monkeypatch): + 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)) + async with app.router.lifespan_context(app): + async with httpx2.AsyncClient( + transport=httpx2.ASGITransport(app=app), base_url="http://localhost:4000" + ) as caller: + payload = {"jsonrpc": "2.0", "id": 1, "method": "tools/list"} + headers = {"Accept": "application/json, text/event-stream"} + assert (await caller.post("/mcp", json=payload, headers=headers)).status_code == 401 + for credential in ("alice-admin", "bob-admin"): + response = await caller.post( + "/mcp", json=payload, headers={**headers, "Authorization": "Bearer " + credential} + ) + assert response.status_code == 200 + assert not response.json().get("error") + assert (await caller.post("/mcp", json=payload, headers=headers)).status_code == 401 + 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(): @@ -194,3 +210,34 @@ async def test_shared_http_client_never_replays_another_callers_gateway_cookie(s for key in ("alice-admin", "bob-admin"): await gateway.call("create_key", {"body": {"key_alias": key}}, key) assert len(stub.writes) == 2 + + +@pytest.mark.parametrize("container", ["object", "array", "encoded"]) +def test_response_redaction_fails_closed_for_excessive_nesting(container): + value = {"echo": "personal-admin-credential", "password": "provider-secret"} + for _ in range(MAX_REDACTION_DEPTH): + value = [value] if container == "array" else {"metadata": value} + payload = json.dumps(value) if container == "encoded" else value + with pytest.raises(AdminError, match="Inspect gateway state before retrying"): + safe_result(payload, "personal-admin-credential") + + +def test_response_redaction_handles_encoded_secrets_and_preserves_other_values(): + credential = "personal-admin-credential" + value = {"metadata": json.dumps({"nested": [{"echo": credential, "password": "provider-secret", "count": 7}]})} + assert safe_result(value, credential) == { + "metadata": json.dumps({"nested": [{"echo": "[credential redacted]", "password": "[redacted]", "count": 7}]}) + } + + +@pytest.mark.parametrize( + "payload", + [ + "[" * 250 + '{"password":"provider-secret"}' + "]" * 250, + ' {"password":"provider-secret"', + '[{"password":"provider-secret"}', + ], +) +def test_response_redaction_rejects_unparseable_json_containers(payload): + 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 fc32e28c48d..334e9883c89 100644 --- a/tests/unit/proxy/admin_mcp/test_progressive.py +++ b/tests/unit/proxy/admin_mcp/test_progressive.py @@ -203,7 +203,7 @@ async def test_large_new_key_response_is_delivered_inline_and_never_cached(stub) return response async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: - gateway = Gateway(Config("https://gateway.example.com", api_key="alice-admin", response_view="compact"), client) + gateway = Gateway(Config("https://gateway.example.com", response_view="compact"), client) async with connected(gateway) as session: result = data(await session.call_tool("create_key", {"body": {"key_alias": "new"}})) assert result["key"].startswith("sk-created") @@ -227,17 +227,19 @@ async def test_result_reads_require_same_credential_principal_and_current_role(s async def test_revocation_during_result_read_suppresses_snapshot(stub): stub.keys = [{"data": "private" * 4000}] - async with httpx2.AsyncClient(transport=httpx2.MockTransport(stub.handle)) as client: - gateway = Gateway(Config("https://gateway.example.com", response_view="compact"), client) - result = await gateway.call("list_keys", {}, "alice-admin") - original_get = gateway.results.get + revoke = False - def get(*args): - item = original_get(*args) + def clock(): + if revoke: stub.revoked.add("alice-admin") - return item + return 0.0 - gateway.results.get = get + async with httpx2.AsyncClient(transport=httpx2.MockTransport(stub.handle)) as client: + gateway = Gateway( + Config("https://gateway.example.com", response_view="compact"), client, results=ResultStore(clock=clock) + ) + result = await gateway.call("list_keys", {}, "alice-admin") + revoke = True with pytest.raises(AdminError, match="proxy admins"): await gateway.call("read_admin_result", result["full_read"], "alice-admin") @@ -256,7 +258,7 @@ async def test_warm_call_retains_pre_post_checks_without_discovery_auth(stub): await session.list_tools() stub.requests.clear() await session.call_tool("list_keys", {}) - assert [r.url.path for r in stub.requests] == ["/user/info", "/key/list", "/user/info"] + 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): @@ -268,10 +270,10 @@ async def test_identity_change_during_cold_discovery_prevents_write(stub): result = original(request) if request.url.path == "/openapi.json": changed = True - elif request.url.path == "/user/info" and changed: + elif request.url.path == "/v2/user/info" and changed: return httpx2.Response( 200, - json={"user_id": "new-identity", "user_info": {"user_id": "new-identity", "user_role": "proxy_admin"}}, + json={"user_id": "new-identity", "user_role": "proxy_admin"}, ) return result @@ -291,8 +293,9 @@ async def test_capacity_falls_back_to_full_write_response_without_retry(stub): return httpx2.Response(200, json=payload) if request.url.path.endswith("/reset_spend") else result async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: - gateway = Gateway(Config("https://gateway.example.com", response_view="compact"), client) - gateway.results = ResultStore(max_bytes=1) + gateway = Gateway( + Config("https://gateway.example.com", response_view="compact"), client, results=ResultStore(max_bytes=1) + ) assert await gateway.call("reset_key_spend", {"path": {"key": "fixture"}}, "alice-admin") == payload assert len(stub.writes) == 1 @@ -382,16 +385,16 @@ async def test_every_reviewed_operation_retains_mapping_advanced_inputs_and_resu def handle(request): if request.url.path == "/openapi.json": return httpx2.Response(200, json=spec) - if request.url.path == "/user/info" and "litellm-changed-by" not in request.headers: + if request.url.path == "/v2/user/info" and "litellm-changed-by" not in request.headers: return httpx2.Response( - 200, json={"user_id": "admin", "user_info": {"user_id": "admin", "user_role": "proxy_admin"}} + 200, json={"user_id": "admin", "user_role": "proxy_admin"} ) calls.append(request) return httpx2.Response(200, json=result) async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: - gateway = Gateway(Config("https://gateway.example.com", api_key="admin", schema_mode=mode), client) - async with connected(gateway) as session: + gateway = Gateway(Config("https://gateway.example.com", schema_mode=mode), client) + async with connected(gateway, "admin") as session: listed = {t.name: t for t in (await session.list_tools()).tools} assert set(listed) - HELPERS.keys() == {op.name for op in OPERATIONS} for op in OPERATIONS: diff --git a/tests/unit/proxy/admin_mcp/test_results.py b/tests/unit/proxy/admin_mcp/test_results.py index 7811dea8461..35d6e788049 100644 --- a/tests/unit/proxy/admin_mcp/test_results.py +++ b/tests/unit/proxy/admin_mcp/test_results.py @@ -2,7 +2,7 @@ import json import pytest -from litellm.proxy.admin_mcp.results import ResultError, ResultStore, encode, page, select +from litellm.proxy.admin_mcp.results import ResultError, ResultStore, contains_credential, encode, page, select def rebuild(item, token, path=""): @@ -162,3 +162,24 @@ def test_empty_values_and_pagination_edges(): assert rebuild(item, token) == value with pytest.raises(ResultError): page(item, token, offset=1) + + +def test_deep_result_paths_and_credential_scan_do_not_use_python_call_stack(): + value = {"password": "private-token"} + for _ in range(1200): + value = {"nested": value} + assert select(value, "/nested" * 1200 + "/password") == "private-token" + assert contains_credential(value) + + +@pytest.mark.parametrize( + "payload", + [ + "[" * 250 + '{"password":"provider-secret"}' + "]" * 250, + ' {"password":"provider-secret"', + '[{"password":"provider-secret"}', + ], +) +def test_unparseable_json_containers_are_never_retained(payload): + assert contains_credential({"metadata": payload}) + assert ResultStore().put("get_user", {"metadata": payload}, "alice", "user-1") is None diff --git a/tests/unit/proxy/admin_mcp/test_schema.py b/tests/unit/proxy/admin_mcp/test_schema.py new file mode 100644 index 00000000000..0a4abec1e43 --- /dev/null +++ b/tests/unit/proxy/admin_mcp/test_schema.py @@ -0,0 +1,172 @@ +from typing import Final + +import jsonschema +import pytest +from pydantic import JsonValue + +from litellm.proxy.admin_mcp._types import JsonObject, object_value +from litellm.proxy.admin_mcp.schema import SchemaError, resolve, tool_schema + + +def _operation(schema: JsonValue) -> JsonObject: + return {"requestBody": {"required": True, "content": {"application/json": {"schema": schema}}}} + + +def test_schema_keywords_in_property_names_and_literal_values_are_preserved() -> None: + literal: Final[JsonValue] = { + "$ref": "#/components/schemas/NotAReference", + "$dynamicRef": "https://untrusted.example/schema.json", + "$recursiveRef": "#", + "$id": "https://untrusted.example/", + "nullable": True, + } + body: Final[JsonValue] = { + "type": "object", + "properties": { + "$ref": {"type": "string"}, + "nullable": {"type": "boolean"}, + "metadata": {"const": literal, "default": literal, "examples": [literal], "enum": [literal]}, + }, + "required": ["$ref", "nullable", "metadata"], + "additionalProperties": False, + } + schema: Final = tool_schema({}, {}, _operation(body)) + assert schema == { + "type": "object", + "properties": {"body": body}, + "additionalProperties": False, + "required": ["body"], + } + validator: Final = jsonschema.Draft202012Validator(schema) + validator.validate({"body": {"$ref": "literal", "nullable": True, "metadata": literal}}) + with pytest.raises(jsonschema.ValidationError): + validator.validate({"body": {"$ref": "literal", "nullable": "wrong", "metadata": literal}}) + + +@pytest.mark.parametrize( + "body,accepted,rejected", + [ + ({"type": "object", "additionalProperties": {"$ref": "#/components/schemas/Value"}}, {"a": 3}, {"a": 2}), + ( + {"type": "object", "patternProperties": {"^item": {"$ref": "#/components/schemas/Value"}}}, + {"item_a": 3}, + {"item_a": 2}, + ), + ({"type": "array", "items": {"$ref": "#/components/schemas/Value"}}, [3], [2]), + ({"type": "array", "prefixItems": [{"$ref": "#/components/schemas/Value"}]}, [3], [2]), + ({"anyOf": [{"$ref": "#/components/schemas/Value"}, {"type": "null"}]}, 3, 2), + ({"if": {"type": "integer"}, "then": {"$ref": "#/components/schemas/Value"}, "else": False}, 3, 2), + ], +) +def test_references_inside_schema_keywords_are_retained( + body: JsonValue, accepted: JsonValue, rejected: JsonValue +) -> None: + spec: Final[JsonObject] = {"components": {"schemas": {"Value": {"type": "integer", "minimum": 3}}}} + schema: Final = tool_schema(spec, {}, _operation(body)) + assert schema["$defs"] == {"Value": {"type": "integer", "minimum": 3}} + validator: Final = jsonschema.Draft202012Validator(schema) + validator.validate({"body": accepted}) + with pytest.raises(jsonschema.ValidationError): + validator.validate({"body": rejected}) + + +def test_mutually_recursive_references_with_escaped_names_terminate_and_validate() -> None: + spec: Final[JsonObject] = { + "components": { + "schemas": { + "A/B~C": {"type": "object", "properties": {"next": {"$ref": "#/components/schemas/Other"}}}, + "Other": {"anyOf": [{"$ref": "#/components/schemas/A~1B~0C"}, {"type": "null"}]}, + "Unused": {"type": "string"}, + } + } + } + schema: Final = tool_schema(spec, {}, _operation({"$ref": "#/components/schemas/A~1B~0C"})) + assert set(object_value(schema["$defs"])) == {"A/B~C", "Other"} + validator: Final = jsonschema.Draft202012Validator(schema) + validator.validate({"body": {"next": {"next": None}}}) + with pytest.raises(jsonschema.ValidationError): + validator.validate({"body": {"next": "wrong"}}) + + +def test_long_reference_graph_does_not_depend_on_python_recursion_limit() -> None: + count: Final = 1200 + components: Final[JsonValue] = { + **{f"Node{i}": {"$ref": f"#/components/schemas/Node{i + 1}"} for i in range(count)}, + f"Node{count}": {"type": "integer"}, + } + spec: Final[JsonObject] = {"components": {"schemas": components}} + schema: Final = tool_schema(spec, {}, _operation({"$ref": "#/components/schemas/Node0"})) + definitions: Final = object_value(schema["$defs"]) + assert len(definitions) == count + assert definitions[f"Node{count}"] == {"type": "integer"} + assert object_value(schema["properties"])["body"] == {"$ref": "#/$defs/Node1"} + + +def test_long_parameter_reference_chain_is_resolved_without_recursion() -> None: + count: Final = 1200 + parameter: Final[JsonValue] = {"name": "limit", "in": "query", "schema": {"type": "integer"}} + parameters: Final[JsonValue] = { + **{f"P{i}": {"$ref": f"#/components/parameters/P{i + 1}"} for i in range(count)}, + f"P{count}": parameter, + } + spec: Final[JsonObject] = {"components": {"parameters": parameters}} + assert resolve(spec, {"$ref": "#/components/parameters/P0"}) == parameter + + +def test_cyclic_parameter_reference_chain_fails_closed() -> None: + spec: Final[JsonObject] = { + "components": { + "parameters": {"A": {"$ref": "#/components/parameters/B"}, "B": {"$ref": "#/components/parameters/A"}} + } + } + with pytest.raises(SchemaError, match="Unsupported OpenAPI reference"): + resolve(spec, {"$ref": "#/components/parameters/A"}) + + +@pytest.mark.parametrize( + "reference", + [ + "https://untrusted.example/schema.json", + "#/components/schemas/Missing", + "#/components/schemas/Parent/properties/child", + "#/components/schemas/Bad~2Escape", + ], +) +def test_unsupported_or_missing_schema_reference_fails_closed(reference: str) -> None: + with pytest.raises(SchemaError): + tool_schema({}, {}, _operation({"$ref": reference})) + + +def test_escaped_component_names_cannot_hide_a_linked_definition() -> None: + spec: Final[JsonObject] = { + "components": { + "schemas": { + "A~1B": {"type": "object", "properties": {"next": {"$ref": "#/components/schemas/A~1B"}}}, + "A/B": {"type": "integer"}, + } + } + } + schema: Final = tool_schema(spec, {}, _operation({"$ref": "#/components/schemas/A~01B"})) + assert schema["$defs"] == {"A/B": {"type": "integer"}} + validator: Final = jsonschema.Draft202012Validator(schema) + validator.validate({"body": {"next": 3}}) + with pytest.raises(jsonschema.ValidationError): + validator.validate({"body": {"next": "wrong"}}) + + +@pytest.mark.parametrize("keyword", ["$dynamicRef", "$recursiveRef", "$id"]) +def test_reference_scopes_and_dynamic_references_are_rejected(keyword: str) -> None: + with pytest.raises(SchemaError, match="unsupported"): + tool_schema({}, {}, _operation({"type": "string", keyword: "https://untrusted.example/schema.json"})) + + +@pytest.mark.parametrize( + "body", [{"type": "integer", "nullable": True}, {"$ref": "#/components/schemas/Number", "nullable": True}] +) +def test_nullable_schemas_accept_null_without_losing_the_original_constraint(body: JsonValue) -> None: + spec: Final[JsonObject] = {"components": {"schemas": {"Number": {"type": "integer"}}}} + validator: Final = jsonschema.Draft202012Validator(tool_schema(spec, {}, _operation(body))) + validator.validate({"body": None}) + validator.validate({"body": 3}) + with pytest.raises(jsonschema.ValidationError): + validator.validate({"body": "wrong"}) diff --git a/tests/unit/proxy/test_admin_mcp.py b/tests/unit/proxy/test_admin_mcp.py index 43a81d5caa6..94ce6328dcd 100644 --- a/tests/unit/proxy/test_admin_mcp.py +++ b/tests/unit/proxy/test_admin_mcp.py @@ -1,6 +1,7 @@ import asyncio import json import sys +from collections.abc import Mapping from typing import Final import httpx2 @@ -57,7 +58,7 @@ def management_app(monkeypatch: pytest.MonkeyPatch) -> FastAPI: monkeypatch.delenv(name, raising=False) app: Final = FastAPI(lifespan=admin_mcp_lifespan) - @app.get("/user/info") + @app.get("/v2/user/info") async def user_info(request: Request) -> dict[str, object]: await asyncio.sleep(0) credential: Final = _native_credential(request) @@ -68,10 +69,7 @@ def management_app(monkeypatch: pytest.MonkeyPatch) -> FastAPI: 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"), - }, + "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") @@ -140,7 +138,7 @@ def test_enabled_without_mcp_dependencies_explains_installation(monkeypatch: pyt monkeypatch.setattr(proxy_server, "premium_user", True) monkeypatch.setenv("LITELLM_ENABLE_ADMIN_MCP", "true") monkeypatch.setitem(sys.modules, "httpx2", None) - with pytest.raises(RuntimeError, match="admin-mcp dependency group"): + with pytest.raises(RuntimeError, match="proxy dependencies"): with TestClient(FastAPI(lifespan=admin_mcp_lifespan)): pytest.fail("Enabling the connector without its dependency must fail startup") @@ -500,36 +498,35 @@ def test_oversized_mcp_body_is_rejected_before_a_write( @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), +@pytest.mark.parametrize("peer,mapped_user,bearer,database_role,status", [ + ("10.0.0.8", "admin-a", "member", "proxy_admin", 200), + ("10.0.0.8", "member", "admin-a", "internal_user", 403), + ("203.0.113.8", "admin-a", "admin-a", "proxy_admin", 401), + ("10.0.0.8", "admin-a", "admin-a", "internal_user", 403), + ("10.0.0.8", "admin-a", "admin-a", None, 403), ]) 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, + peer: str, mapped_user: str, bearer: str, database_role: str | None, 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 UserTable: + async def find_unique(self, *, where: Mapping[str, object], **_: object) -> LiteLLM_UserTable | None: + user_id: Final = where.get("user_id") + if user_id != mapped_user or database_role is None: + return None + return LiteLLM_UserTable(user_id=mapped_user, user_role=database_role, teams=[]) class DatabaseTables: - litellm_teamtable: Final = EmptyTable() - litellm_teammembership: Final = EmptyTable() + litellm_usertable: Final = UserTable() 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) @@ -542,9 +539,9 @@ async def test_oauth2_proxy_auth_preserves_native_identity_and_peer_trust( }) management_app.router.routes[:] = [ route for route in management_app.router.routes - if not (isinstance(route, APIRoute) and route.path == "/user/info") + if not (isinstance(route, APIRoute) and route.path == "/v2/user/info") ] - management_app.add_api_route("/user/info", internal_user_endpoints.user_info, methods=["GET"]) + management_app.add_api_route("/v2/user/info", internal_user_endpoints.user_info_v2, methods=["GET"]) management_app.add_exception_handler(ProxyException, proxy_server.openai_exception_handler) async with management_app.router.lifespan_context(management_app): diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 04781c23e28..75d6ec38201 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -142,9 +142,9 @@ async def test_proxy_shutdown_drains_active_admin_tool_before_closing_connector( started: Final = asyncio.Event() release: Final = asyncio.Event() - @admin_mcp_proxy.get("/user/info") + @admin_mcp_proxy.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"} @admin_mcp_proxy.get("/team/list", operation_id="list_team_team_list_get") async def list_teams() -> dict[str, object]: diff --git a/uv.lock b/uv.lock index f55e01426d7..436e67a838b 100644 --- a/uv.lock +++ b/uv.lock @@ -2126,11 +2126,29 @@ wheels = [ [[package]] name = "filelock" -version = "3.32.7" +version = "3.20.3" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0f/59/e19834834cb01a32febfbb0f8a23a9088088f5d45991824ff2bc3b5e8acb/filelock-3.32.7.tar.gz", hash = "sha256:37b8a3d9811b0f9aef7e5ec5c71bb320de52df51e6ca9bcd6f5ad81187660da7", size = 225154, upload-time = "2026-09-16T00:24:20.907Z" } +resolution-markers = [ + "python_full_version < '3.11'", +] +sdist = { url = "https://files.pythonhosted.org/packages/1d/65/ce7f1b70157833bf3cb851b556a37d4547ceafc158aa9b34b36782f23696/filelock-3.20.3.tar.gz", hash = "sha256:18c57ee915c7ec61cff0ecf7f0f869936c7c30191bb0cf406f1341778d0834e1", size = 19485, upload-time = "2026-01-09T17:55:05.421Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/15/df/31098c5aeb4d966b553641472bd55fcf5fdfac953549894b8a765ba44e91/filelock-3.32.7-py3-none-any.whl", hash = "sha256:65ff0d0190ea42038b32bda4b77834fb05be2cad4c5b9b01aa4dfb3614536e52", size = 100157, upload-time = "2026-09-16T00:24:19.543Z" }, + { url = "https://files.pythonhosted.org/packages/b5/36/7fb70f04bf00bc646cd5bb45aa9eddb15e19437a28b8fb2b4a5249fac770/filelock-3.20.3-py3-none-any.whl", hash = "sha256:4b0dda527ee31078689fc205ec4f1c1bf7d56cf88b6dc9426c4f230e46c2dce1", size = 16701, upload-time = "2026-01-09T17:55:04.334Z" }, +] + +[[package]] +name = "filelock" +version = "4.2.4" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14'", + "python_full_version == '3.13.*'", + "python_full_version == '3.12.*'", + "python_full_version == '3.11.*'", +] +sdist = { url = "https://files.pythonhosted.org/packages/d8/28/3893324bcadf8f89b1616637468d766af3f625008c0735bf704fdfa86617/filelock-4.2.4.tar.gz", hash = "sha256:c6d8663535e1650ee430507b4ffbac7b464dd64b7647ea0112a9f406499d163d", size = 652444, upload-time = "2026-10-19T01:57:51.001Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/1b/5e74289a1878e666bb4b65fe5805a04cafcac2aa5f519ef1ead67ef91c56/filelock-4.2.4-py3-none-any.whl", hash = "sha256:301241438c57936ef93f4243a94f2559022268bfcf9b375a9a72eb0f8db8358e", size = 162475, upload-time = "2026-10-19T01:57:49.419Z" }, ] [[package]] @@ -3514,7 +3532,8 @@ version = "1.32.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, - { name = "filelock" }, + { name = "filelock", version = "3.20.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "filelock", version = "4.2.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "fsspec" }, { name = "hf-xet", marker = "platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, { name = "httpx" }, @@ -4650,7 +4669,8 @@ dependencies = [ { name = "certifi" }, { name = "click" }, { name = "fastuuid" }, - { name = "filelock" }, + { name = "filelock", version = "3.20.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "filelock", version = "4.2.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "httpx", extra = ["http2"] }, { name = "huggingface-hub" }, { name = "importlib-metadata" }, @@ -4788,11 +4808,6 @@ utils = [ ] [package.dev-dependencies] -admin-mcp = [ - { name = "httpx2" }, - { name = "jsonschema" }, - { name = "mcp" }, -] benchmarks = [ { name = "a2a-sdk" }, { name = "mcp" }, @@ -4888,10 +4903,7 @@ proxy-dev = [ { name = "a2a-sdk" }, { name = "azure-identity" }, { name = "backports-zstd", marker = "python_full_version < '3.14'" }, - { name = "httpx2" }, { name = "hypercorn" }, - { name = "jsonschema" }, - { name = "mcp" }, { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp" }, { name = "opentelemetry-instrumentation-fastapi" }, @@ -4932,7 +4944,8 @@ requires-dist = [ { name = "fastapi", marker = "extra == 'proxy'", specifier = ">=0.136.3,<1.0" }, { name = "fastapi-sso", marker = "extra == 'proxy'", specifier = ">=0.19.0,<1.0" }, { name = "fastuuid", specifier = ">=0.14.0,<1.0" }, - { name = "filelock", specifier = ">=3.16.1,<4.0" }, + { name = "filelock", marker = "python_full_version < '3.11'", specifier = ">=3.20.3,<3.20.4" }, + { name = "filelock", marker = "python_full_version >= '3.11'", specifier = ">=4.2.1,<5.0" }, { name = "google-cloud-aiplatform", marker = "extra == 'google'", specifier = ">=1.133.0,<2.0" }, { name = "google-cloud-aiplatform", marker = "extra == 'proxy-runtime'", specifier = ">=1.133.0,<2.0" }, { name = "google-cloud-iam", marker = "extra == 'extra-proxy'", specifier = ">=2.19.1,<3.0" }, @@ -4953,7 +4966,7 @@ requires-dist = [ { name = "inquirerpy", marker = "extra == 'cli'", specifier = ">=0.3.4,<1.0" }, { name = "inquirerpy", marker = "extra == 'proxy'", specifier = ">=0.3.4,<1.0" }, { name = "jinja2", specifier = ">=3.1.6,<4.0" }, - { name = "jsonschema", specifier = ">=4.0.0,<5.0" }, + { name = "jsonschema", specifier = ">=4.22.0,<5.0" }, { name = "keyring", marker = "extra == 'cli'", specifier = ">=25.6.0,<26.0" }, { name = "langfuse", marker = "extra == 'proxy-runtime'", specifier = ">=4.7,<5.0" }, { name = "litellm-enterprise", marker = "extra == 'proxy'", editable = "enterprise" }, @@ -5018,11 +5031,6 @@ requires-dist = [ provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "saml", "semantic-router", "mlflow", "grpc", "stt-vertex-chirp", "stt-nvidia-riva", "google", "bedrock-realtime", "proxy-runtime"] [package.metadata.requires-dev] -admin-mcp = [ - { name = "httpx2", specifier = ">=2.10.0,<3" }, - { name = "jsonschema", specifier = ">=4.22.0,<5" }, - { name = "mcp", specifier = ">=2.3.0,<3" }, -] benchmarks = [ { name = "a2a-sdk", specifier = "==1.1.0" }, { name = "mcp", specifier = ">=2.3.0,<3" }, @@ -5117,10 +5125,7 @@ proxy-dev = [ { name = "a2a-sdk", specifier = "==1.1.0" }, { name = "azure-identity", specifier = "==1.25.2" }, { name = "backports-zstd", marker = "python_full_version < '3.14'", specifier = "==1.8.0" }, - { name = "httpx2", specifier = ">=2.10.0,<3" }, { name = "hypercorn", specifier = "==0.17.3" }, - { name = "jsonschema", specifier = ">=4.22.0,<5" }, - { name = "mcp", specifier = ">=2.3.0,<3" }, { name = "opentelemetry-api", specifier = "==1.45.1" }, { name = "opentelemetry-exporter-otlp", specifier = "==1.45.1" }, { name = "opentelemetry-instrumentation-fastapi", specifier = "==0.66b1" }, diff --git a/vscode-extension/package-lock.json b/vscode-extension/package-lock.json index 81e9425ef4e..536b48502fd 100644 --- a/vscode-extension/package-lock.json +++ b/vscode-extension/package-lock.json @@ -189,19 +189,29 @@ } }, "node_modules/@azure/msal-node": { - "version": "6.0.1", - "resolved": "https://registry.npmjs.org/@azure/msal-node/-/msal-node-6.0.1.tgz", - "integrity": "sha512-ixSO1Y/kCVRthRs+hSx/5qkwaunX1/RAePhlMN0wIpIQ4WEZ6AREGGnGd1AP0qspHVsAwzWQKGubpjh50JyJIQ==", + "version": "7.1.0", + "resolved": "https://registry.npmjs.org/@azure/msal-node/-/msal-node-7.1.0.tgz", + "integrity": "sha512-s1O6hSY/C/dmqykTmmeHY6GhSkDuY8N1gxyhuLyaCwbY7N+UBvU2Wu4flas7+NxZiAPcXNPVpgt3n/HP/5yL+w==", "dev": true, "license": "MIT", "dependencies": { - "@azure/msal-common": "16.14.1", + "@azure/msal-common": "16.14.2", "jsonwebtoken": "^9.0.0" }, "engines": { "node": ">=20" } }, + "node_modules/@azure/msal-node/node_modules/@azure/msal-common": { + "version": "16.14.2", + "resolved": "https://registry.npmjs.org/@azure/msal-common/-/msal-common-16.14.2.tgz", + "integrity": "sha512-HOBR0Gfo39TMrkxNG8k6cfjRkrRjXZc3JQs3ebMWQc9S4cihuknMYra3sgmzbZB752oQBqBN9+WF7c826p2vIQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.8.0" + } + }, "node_modules/@esbuild/aix-ppc64": { "version": "0.28.2", "resolved": "https://registry.npmjs.org/@esbuild/aix-ppc64/-/aix-ppc64-0.28.2.tgz", diff --git a/vscode-extension/package.json b/vscode-extension/package.json index a3cf2dc7a5f..ccc26c741cd 100644 --- a/vscode-extension/package.json +++ b/vscode-extension/package.json @@ -83,5 +83,10 @@ "esbuild": "^0.28.2", "typescript": "^5.9.3", "vitest": "^4.1.12" + }, + "overrides": { + "@azure/identity": { + "@azure/msal-node": "7.1.0" + } } }