From 84266bf924ebafde269625cece6772a887719dbe Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 20 Jun 2026 18:49:41 -0700 Subject: [PATCH 01/50] feat(auth): resolve caller identity once into a Principal at the auth seam (#30887) Introduce a single, typed caller identity that is resolved once at the auth boundary and read by reference downstream, instead of being re-derived from a 50-field key object or rebuilt from request metadata. What this adds (litellm/proxy/auth/resolvers/), organized by responsibility: - Principal: a small, frozen, identity-only value type (user / organization / teams / project / end-user / roles / scopes / network), with its sub-models and the role mapping. No budget or policy state; those stay on the key object. - DbIdentityStore: the auth flow's resolver, owning both halves of resolving a caller. resolve_key does the one combined_view lookup (cache, then DB via the shared lower-level helpers, then write-back) and returns the key object, which still flows for budget / rate-limit / policy unchanged. principal_from_key projects the identity slice of that key object into a Principal, issuing no lookup. user_api_key_auth resolves every key through the store rather than calling get_key_object directly; auth_checks.get_key_object stays as the legacy entrypoint for its other callers until they migrate. - network: the X-Forwarded-For / trusted-proxy CIDR primitives live here in one place. trusted_proxy_utils now imports them rather than keeping a second copy. At the seam, user_api_key_auth projects one per-request Principal off the resolved key object and stamps the request network context onto it once (X-Forwarded-For is trusted only when trusted_proxy_ranges is configured). It is attached to request.state.principal for the downstream consumers later phases add. The projection is additive and defensive: a failure never rejects an already-authenticated request, and a missing principal must be treated as deny by any future reader. The Principal is always identifiable (credential_ref and a stable subject off the token), never anonymous. This is additive and changes no behavior today; it is the identity foundation the spend-attribution and authorization phases build on. --- litellm/proxy/auth/auth_method.py | 13 ++ litellm/proxy/auth/network.py | 101 +++++++++ litellm/proxy/auth/resolvers/__init__.py | 33 +++ litellm/proxy/auth/resolvers/exceptions.py | 50 +++++ litellm/proxy/auth/resolvers/models.py | 82 +++++++ litellm/proxy/auth/resolvers/store.py | 206 ++++++++++++++++++ litellm/proxy/auth/roles.py | 35 +++ litellm/proxy/auth/trusted_proxy_utils.py | 72 ++---- litellm/proxy/auth/user_api_key_auth.py | 117 +++++++--- scripts/ruff_strict_gate.py | 94 +++----- .../proxy/auth/test_user_api_key_auth.py | 25 ++- .../proxy_unit_tests/test_jwt_key_mapping.py | 31 ++- .../proxy/auth/test_auth_checks.py | 48 +++- .../auth/test_custom_auth_end_user_budget.py | 4 +- .../proxy/auth/test_handle_jwt.py | 27 ++- .../proxy/auth/test_model_checks.py | 4 +- tests/test_litellm/proxy/auth/test_network.py | 104 +++++++++ .../proxy/auth/test_onboarding.py | 1 - .../proxy/auth/test_resolvers_exceptions.py | 35 +++ .../proxy/auth/test_resolvers_models.py | 95 ++++++++ .../proxy/auth/test_resolvers_seam.py | 96 ++++++++ .../proxy/auth/test_resolvers_store.py | 70 ++++++ .../proxy/auth/test_user_api_key_auth.py | 10 +- tests/test_litellm/test_ruff_strict_gate.py | 10 - 24 files changed, 1149 insertions(+), 214 deletions(-) create mode 100644 litellm/proxy/auth/auth_method.py create mode 100644 litellm/proxy/auth/network.py create mode 100644 litellm/proxy/auth/resolvers/__init__.py create mode 100644 litellm/proxy/auth/resolvers/exceptions.py create mode 100644 litellm/proxy/auth/resolvers/models.py create mode 100644 litellm/proxy/auth/resolvers/store.py create mode 100644 litellm/proxy/auth/roles.py create mode 100644 tests/test_litellm/proxy/auth/test_network.py create mode 100644 tests/test_litellm/proxy/auth/test_resolvers_exceptions.py create mode 100644 tests/test_litellm/proxy/auth/test_resolvers_models.py create mode 100644 tests/test_litellm/proxy/auth/test_resolvers_seam.py create mode 100644 tests/test_litellm/proxy/auth/test_resolvers_store.py diff --git a/litellm/proxy/auth/auth_method.py b/litellm/proxy/auth/auth_method.py new file mode 100644 index 00000000000..a604eb563e6 --- /dev/null +++ b/litellm/proxy/auth/auth_method.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from enum import Enum + + +class AuthMethod(str, Enum): + API_KEY = "api_key" + HTTP_BASIC = "http_basic" + BEARER_JWT = "bearer_jwt" + OAUTH2_INTROSPECTION = "oauth2_introspection" + OIDC = "oidc" + SAML = "saml" + MUTUAL_TLS = "mutual_tls" diff --git a/litellm/proxy/auth/network.py b/litellm/proxy/auth/network.py new file mode 100644 index 00000000000..4eb6f1dcec2 --- /dev/null +++ b/litellm/proxy/auth/network.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import ipaddress +from typing import Any, Union + +from fastapi import Request +from pydantic import BaseModel, Field + +from litellm._logging import verbose_proxy_logger + +TrustedProxyNetwork = Union[ipaddress.IPv4Network, ipaddress.IPv6Network] + + +class NetworkContext(BaseModel): + client_ip: str | None = None + host: str | None = None + via_trusted_proxy: bool = False + + +class TrustedProxyConfig(BaseModel): + use_forwarded_for: bool = False + trusted_proxy_cidrs: list[str] = Field(default_factory=list) + + +def normalize_cidr_ranges( + configured_ranges: Any, *, setting_name: str = "trusted_proxy_cidrs" +) -> list[str]: + if not configured_ranges: + return [] + if isinstance(configured_ranges, str): + return [r.strip() for r in configured_ranges.split(",") if r.strip()] + if isinstance(configured_ranges, (list, tuple, set)): + return [str(r).strip() for r in configured_ranges if str(r).strip()] + verbose_proxy_logger.warning( + "Invalid %s value: expected a list of CIDR ranges, got %s", + setting_name, + type(configured_ranges).__name__, + ) + return [] + + +def parse_trusted_proxy_ranges( + configured_ranges: Any, *, setting_name: str = "trusted_proxy_cidrs" +) -> list[TrustedProxyNetwork]: + networks: list[TrustedProxyNetwork] = [] + for cidr in normalize_cidr_ranges(configured_ranges, setting_name=setting_name): + try: + networks.append(ipaddress.ip_network(cidr, strict=False)) + except ValueError: + verbose_proxy_logger.warning( + "Invalid CIDR in %s: %s, skipping", setting_name, cidr + ) + return networks + + +def ip_in_networks(client_ip: str | None, networks: list[TrustedProxyNetwork]) -> bool: + if not client_ip or not networks: + return False + try: + addr = ipaddress.ip_address(client_ip.strip()) + except ValueError: + return False + return any(addr in network for network in networks) + + +def _is_valid_ip(value: str) -> bool: + try: + ipaddress.ip_address(value) + return True + except ValueError: + return False + + +def resolve_client_ip( + request: Request, config: TrustedProxyConfig +) -> tuple[str | None, bool]: + """Resolve the real client IP, trusting X-Forwarded-For only when the direct + peer is itself a configured trusted proxy. Walks the header right-to-left and + returns the first hop that is not a trusted proxy, so a forged left-most entry + cannot spoof the client.""" + peer = request.client.host if request.client else None + networks = parse_trusted_proxy_ranges(config.trusted_proxy_cidrs) + if not config.use_forwarded_for or not ip_in_networks(peer, networks): + return peer, False + forwarded = request.headers.get("x-forwarded-for", "") + hops = [h.strip() for h in forwarded.split(",") if h.strip()] + for hop in reversed(hops): + if _is_valid_ip(hop) and not ip_in_networks(hop, networks): + return hop, True + return peer, True + + +def resolve_network_context( + request: Request, config: TrustedProxyConfig +) -> NetworkContext: + ip, via_proxy = resolve_client_ip(request, config) + return NetworkContext( + client_ip=ip, + host=request.headers.get("host"), + via_trusted_proxy=via_proxy, + ) diff --git a/litellm/proxy/auth/resolvers/__init__.py b/litellm/proxy/auth/resolvers/__init__.py new file mode 100644 index 00000000000..d6bfb335c09 --- /dev/null +++ b/litellm/proxy/auth/resolvers/__init__.py @@ -0,0 +1,33 @@ +from litellm.proxy.auth.resolvers.exceptions import ( + IdentityResolutionError, + KeyNotFoundError, + KeyNotInCacheError, + NoDatabaseConnectionError, + PrincipalMissingSourceKeyError, +) +from litellm.proxy.auth.resolvers.models import ( + CredentialRef, + EndUserIdentity, + OrganizationIdentity, + Principal, + PrincipalType, + ProjectIdentity, + TeamIdentity, + UserIdentity, +) + +__all__ = [ + "CredentialRef", + "EndUserIdentity", + "IdentityResolutionError", + "KeyNotFoundError", + "KeyNotInCacheError", + "NoDatabaseConnectionError", + "OrganizationIdentity", + "Principal", + "PrincipalMissingSourceKeyError", + "PrincipalType", + "ProjectIdentity", + "TeamIdentity", + "UserIdentity", +] diff --git a/litellm/proxy/auth/resolvers/exceptions.py b/litellm/proxy/auth/resolvers/exceptions.py new file mode 100644 index 00000000000..dd953e66659 --- /dev/null +++ b/litellm/proxy/auth/resolvers/exceptions.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from fastapi import status + +from litellm.proxy._types import ProxyErrorTypes, ProxyException + + +class IdentityResolutionError(Exception): + """Base for every failure raised while resolving a caller's identity.""" + + +class NoDatabaseConnectionError(IdentityResolutionError): + def __init__(self) -> None: + super().__init__( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + + +class KeyNotInCacheError(IdentityResolutionError): + def __init__(self, hashed_token: str) -> None: + super().__init__( + f"Key doesn't exist in cache + check_cache_only=True. key={hashed_token}." + ) + + +class KeyNotFoundError(IdentityResolutionError, ProxyException): + """The token matched nothing in the cache or the verification token table. + + Also a ``ProxyException`` so the auth flow keeps mapping a missing key to the + OpenAI 401 contract unchanged while callers migrate onto the typed hierarchy. + """ + + def __init__(self, hashed_token: str) -> None: + ProxyException.__init__( + self, + message="Authentication Error, Invalid proxy server token passed. key={}, not found in db. Create key via `/key/generate` call.".format( + hashed_token + ), + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, + ) + + +class PrincipalMissingSourceKeyError(IdentityResolutionError): + def __init__(self) -> None: + super().__init__( + "Principal carries no source key; it was not produced by " + "IdentityStore.resolve" + ) diff --git a/litellm/proxy/auth/resolvers/models.py b/litellm/proxy/auth/resolvers/models.py new file mode 100644 index 00000000000..97e66b8fe67 --- /dev/null +++ b/litellm/proxy/auth/resolvers/models.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from enum import Enum + +from pydantic import BaseModel, ConfigDict, Field + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.network import NetworkContext +from litellm.proxy.auth.roles import Role, TeamRole + + +class PrincipalType(str, Enum): + HUMAN = "human" + SERVICE_ACCOUNT = "service_account" + + +class UserIdentity(BaseModel): + id: str + external_id: str | None = None + user_name: str | None = None + email: str | None = None + display_name: str | None = None + + +class OrganizationIdentity(BaseModel): + id: str + name: str | None = None + + +class TeamIdentity(BaseModel): + id: str + name: str | None = None + role: TeamRole = TeamRole.MEMBER + + +class ProjectIdentity(BaseModel): + id: str + name: str | None = None + + +class EndUserIdentity(BaseModel): + id: str + + +class CredentialRef(BaseModel): + key_id: str | None = None + token_id: str | None = None + + +class Principal(BaseModel): + """Normalized caller identity, resolved once per request at the auth seam. + + Frozen and constructed fresh per request, never cached or shared. The identity + fields carry no policy, budget, or rate-limit state. ``source_key`` is a + transitional carrier for the resolved key object so ``key_from_principal`` can + hand it to the request flow that still consumes ``UserAPIKeyAuth``; it is + excluded from serialization and repr and goes away once those consumers read + identity off the Principal directly. + """ + + model_config = ConfigDict(frozen=True) + + principal_type: PrincipalType + subject: str + issuer: str | None = None + audience: list[str] = Field(default_factory=list) + + user: UserIdentity | None = None + organization: OrganizationIdentity | None = None + teams: list[TeamIdentity] = Field(default_factory=list) + project: ProjectIdentity | None = None + end_user: EndUserIdentity | None = None + + roles: list[Role] = Field(default_factory=list) + scopes: list[str] = Field(default_factory=list) + + auth_method: AuthMethod + credential_ref: CredentialRef = Field(default_factory=CredentialRef) + network: NetworkContext = Field(default_factory=NetworkContext) + + source_key: UserAPIKeyAuth | None = Field(default=None, exclude=True, repr=False) diff --git a/litellm/proxy/auth/resolvers/store.py b/litellm/proxy/auth/resolvers/store.py new file mode 100644 index 00000000000..c43d4c84ca4 --- /dev/null +++ b/litellm/proxy/auth/resolvers/store.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Sequence + +from pydantic import BaseModel + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import ( + _cache_key_object, + _copy_user_api_key_auth_for_cache, + _fetch_key_object_from_db_with_reconnect, + get_object_permission, +) +from litellm.proxy.auth.resolvers.exceptions import ( + KeyNotFoundError, + KeyNotInCacheError, + NoDatabaseConnectionError, + PrincipalMissingSourceKeyError, +) +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.network import NetworkContext +from litellm.proxy.auth.resolvers.models import ( + CredentialRef, + EndUserIdentity, + OrganizationIdentity, + Principal, + PrincipalType, + ProjectIdentity, + TeamIdentity, + UserIdentity, +) +from litellm.proxy.auth.roles import TeamRole, map_role, team_role + +if TYPE_CHECKING: + from litellm.caching.caching import DualCache + from litellm.integrations.opentelemetry import Span + from litellm.proxy.utils import PrismaClient, ProxyLogging + + +class IdentityStore: + """The auth flow's resolver: one combined_view lookup, projected into a Principal. + + ``resolve`` does the lookup (cache, then DB via the shared lower-level helpers, + then write-back) and returns the per-caller Principal. The Principal carries the + source key object so ``key_from_principal`` can hand it back to the parts of the + request flow that still consume ``UserAPIKeyAuth`` (budget, rate limits, policy); + that carrier is a stopgap until those consumers read identity off the Principal. + The Prisma client, key cache, the request's tracing span / logging sink, and + whether this store may only read the cache are injected so the composition root + can build the store once the proxy DB is connected; the span and logging sink + are infra the DB call is instrumented with and ``check_cache_only`` is a store + mode, none of them inputs to resolving identity. ``auth_checks.get_key_object`` + stays as the legacy entrypoint for its other callers until they migrate onto + this store. + """ + + def __init__( + self, + prisma_client: PrismaClient | None, + cache: DualCache, + *, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + check_cache_only: bool = False, + ) -> None: + self._prisma = prisma_client + self._cache = cache + self._parent_otel_span = parent_otel_span + self._proxy_logging_obj = proxy_logging_obj + self._check_cache_only = check_cache_only + + async def resolve( + self, + hashed_token: str, + *, + auth_method: AuthMethod = AuthMethod.API_KEY, + network: NetworkContext | None = None, + ) -> Principal: + key = await self._resolve_key(hashed_token) + return self._principal_from_key( + key, + auth_method=auth_method, + network=network, + subject_fallback=key.token, + credential_ref=CredentialRef(token_id=key.token), + ) + + @staticmethod + def key_from_principal(principal: Principal) -> UserAPIKeyAuth: + """Hand back the resolved key object carried on the Principal. + + Stopgap for the request flow that still consumes ``UserAPIKeyAuth`` for + budget, rate-limit, and policy state. Only Principals produced by + ``resolve`` carry a source key. + """ + if principal.source_key is None: + raise PrincipalMissingSourceKeyError() + return principal.source_key + + async def _resolve_key(self, hashed_token: str) -> UserAPIKeyAuth: + if self._prisma is None: + raise NoDatabaseConnectionError() + + cached = await self._cache.async_get_cache( + key=hashed_token, model_type=UserAPIKeyAuth + ) + if cached is not None: + return _copy_user_api_key_auth_for_cache(user_api_key_obj=cached) + + if self._check_cache_only: + raise KeyNotInCacheError(hashed_token) + + from_db: BaseModel | None = await _fetch_key_object_from_db_with_reconnect( + hashed_token=hashed_token, + prisma_client=self._prisma, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + if from_db is None: + raise KeyNotFoundError(hashed_token) + + key = UserAPIKeyAuth(**from_db.model_dump(exclude_none=True)) + + if key.object_permission_id and not key.object_permission: + try: + key.object_permission = await get_object_permission( + object_permission_id=key.object_permission_id, + prisma_client=self._prisma, + user_api_key_cache=self._cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + except Exception as e: + verbose_proxy_logger.debug( + f"Failed to load object_permission for key with object_permission_id={key.object_permission_id}: {e}" + ) + + await _cache_key_object( + hashed_token=hashed_token, + user_api_key_obj=key, + user_api_key_cache=self._cache, + proxy_logging_obj=self._proxy_logging_obj, + ) + return key + + @staticmethod + def _principal_from_key( + key: UserAPIKeyAuth, + *, + auth_method: AuthMethod, + issuer: str | None = None, + subject_fallback: str | None = None, + scopes: Sequence[str] = (), + credential_ref: CredentialRef | None = None, + network: NetworkContext | None = None, + ) -> Principal: + """Project the identity slice off an already-resolved key object and carry + the key on the Principal so ``key_from_principal`` can recover it. + + Pure: issues no lookup. Both ``resolve`` and the auth seam call this so + identity is projected once off whichever key object they already hold. + """ + teams: list[TeamIdentity] = [] + if key.team_id is not None: + role = ( + team_role(key.team_member.role) if key.team_member else TeamRole.MEMBER + ) + teams.append(TeamIdentity(id=key.team_id, name=key.team_alias, role=role)) + organization = ( + OrganizationIdentity(id=key.org_id, name=key.organization_alias) + if key.org_id is not None + else None + ) + user = ( + UserIdentity(id=key.user_id, email=key.user_email) + if key.user_id is not None + else None + ) + project = ( + ProjectIdentity(id=key.project_id, name=key.project_alias) + if key.project_id is not None + else None + ) + end_user = ( + EndUserIdentity(id=key.end_user_id) if key.end_user_id is not None else None + ) + mapped = map_role(key.user_role) + return Principal( + principal_type=( + PrincipalType.HUMAN if key.user_id else PrincipalType.SERVICE_ACCOUNT + ), + subject=key.user_id or key.key_alias or subject_fallback or "", + issuer=issuer, + user=user, + organization=organization, + teams=teams, + project=project, + end_user=end_user, + roles=[mapped] if mapped else [], + scopes=list(scopes), + auth_method=auth_method, + credential_ref=credential_ref or CredentialRef(), + network=network or NetworkContext(), + source_key=key, + ) diff --git a/litellm/proxy/auth/roles.py b/litellm/proxy/auth/roles.py new file mode 100644 index 00000000000..efe56a8b6b2 --- /dev/null +++ b/litellm/proxy/auth/roles.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from enum import Enum + + +class Role(str, Enum): + PLATFORM_ADMIN = "platform_admin" + PLATFORM_VIEWER = "platform_viewer" + ORG_ADMIN = "org_admin" + ORG_VIEWER = "org_viewer" + TEAM_ADMIN = "team_admin" + TEAM_MEMBER = "team_member" + + +class TeamRole(str, Enum): + ADMIN = "admin" + MEMBER = "member" + + +_ROLE_MAP: dict[str, Role] = { + "proxy_admin": Role.PLATFORM_ADMIN, + "proxy_admin_viewer": Role.PLATFORM_VIEWER, + "org_admin": Role.ORG_ADMIN, +} + + +def map_role(value: str | None) -> Role | None: + """Map a LiteLLM ``user_role`` string to a platform Role.""" + if value is None: + return None + return _ROLE_MAP.get(value) + + +def team_role(role: str | None) -> TeamRole: + return TeamRole.ADMIN if role == "admin" else TeamRole.MEMBER diff --git a/litellm/proxy/auth/trusted_proxy_utils.py b/litellm/proxy/auth/trusted_proxy_utils.py index df7b3080f28..35bb79e7efe 100644 --- a/litellm/proxy/auth/trusted_proxy_utils.py +++ b/litellm/proxy/auth/trusted_proxy_utils.py @@ -1,12 +1,15 @@ -import ipaddress -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, Optional from fastapi import Request from litellm._logging import verbose_proxy_logger +from litellm.proxy.auth.network import ( + ip_in_networks, + normalize_cidr_ranges, + parse_trusted_proxy_ranges, +) TRUSTED_PROXY_RANGES_KEY = "trusted_proxy_ranges" -TrustedProxyNetwork = Union[ipaddress.IPv4Network, ipaddress.IPv6Network] def _get_proxy_general_settings() -> Dict[str, Any]: @@ -18,43 +21,20 @@ def _get_proxy_general_settings() -> Dict[str, Any]: return {} -def _normalize_cidr_ranges(configured_ranges: Any, *, setting_name: str) -> List[str]: - if not configured_ranges: - return [] - if isinstance(configured_ranges, str): - return [ - raw_range.strip() - for raw_range in configured_ranges.split(",") - if raw_range.strip() - ] - if isinstance(configured_ranges, (list, tuple, set)): - return [ - str(raw_range).strip() - for raw_range in configured_ranges - if str(raw_range).strip() - ] - verbose_proxy_logger.warning( - "Invalid %s value: expected a list of CIDR ranges, got %s", - setting_name, - type(configured_ranges).__name__, +def get_trusted_proxy_cidrs( + general_settings: dict[str, Any] | None = None, +) -> list[str]: + """Operator-configured trusted reverse-proxy CIDRs, normalized to strings. + + Empty when none are configured, in which case X-Forwarded-For must not be + trusted and only the direct peer is authoritative. + """ + if general_settings is None: + general_settings = _get_proxy_general_settings() + return normalize_cidr_ranges( + general_settings.get(TRUSTED_PROXY_RANGES_KEY), + setting_name=TRUSTED_PROXY_RANGES_KEY, ) - return [] - - -def parse_trusted_proxy_ranges( - configured_ranges: Any, - *, - setting_name: str = TRUSTED_PROXY_RANGES_KEY, -) -> List[TrustedProxyNetwork]: - networks: List[TrustedProxyNetwork] = [] - for cidr in _normalize_cidr_ranges(configured_ranges, setting_name=setting_name): - try: - networks.append(ipaddress.ip_network(cidr, strict=False)) - except ValueError: - verbose_proxy_logger.warning( - "Invalid CIDR in %s: %s, skipping", setting_name, cidr - ) - return networks def _get_direct_client_ip(request: Request) -> Optional[str]: @@ -65,18 +45,6 @@ def _get_direct_client_ip(request: Request) -> Optional[str]: return None -def _is_ip_in_networks( - client_ip: Optional[str], networks: List[TrustedProxyNetwork] -) -> bool: - if not client_ip or not networks: - return False - try: - addr = ipaddress.ip_address(client_ip.strip()) - except ValueError: - return False - return any(addr in network for network in networks) - - def require_trusted_proxy_request( *, request: Request, @@ -105,7 +73,7 @@ def require_trusted_proxy_request( ) direct_client_ip = _get_direct_client_ip(request) - if not _is_ip_in_networks(direct_client_ip, trusted_networks): + if not ip_in_networks(direct_client_ip, trusted_networks): verbose_proxy_logger.warning( "%s rejected identity headers from untrusted direct client IP %r", feature_name, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 00d98a04a78..4a2df18b93b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -42,7 +42,6 @@ from litellm.proxy.auth.auth_checks import ( common_checks, get_end_user_object, get_jwt_key_mapping_object, - get_key_object, get_project_object, get_team_object, get_user_object, @@ -63,7 +62,12 @@ from litellm.proxy.auth.auth_utils import ( from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context +from litellm.proxy.auth.resolvers import CredentialRef, Principal +from litellm.proxy.auth.resolvers.store import IdentityStore from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, @@ -758,12 +762,13 @@ async def _auto_register_jwt_mapping( claim_value, ) - auto_registered_key = await get_key_object( - hashed_token=token_hash, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, + auto_registered_key = IdentityStore.key_from_principal( + await IdentityStore( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ).resolve(hashed_token=token_hash) ) if auto_registered_key is not None: auto_registered_key.org_id = org_id @@ -865,12 +870,13 @@ async def _resolve_jwt_to_virtual_key( ) return None elif cached_mapping is not None: - return await get_key_object( - hashed_token=cached_mapping, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, + return IdentityStore.key_from_principal( + await IdentityStore( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ).resolve(hashed_token=cached_mapping) ) # Resolve the mapping from DB, or treat prisma_client=None as a definitive @@ -889,12 +895,13 @@ async def _resolve_jwt_to_virtual_key( value=token_hash, ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl, ) - return await get_key_object( - hashed_token=token_hash, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, + return IdentityStore.key_from_principal( + await IdentityStore( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ).resolve(hashed_token=token_hash) ) # No mapping found (DB miss or no DB) — apply no-match policy. @@ -1493,13 +1500,14 @@ async def _user_api_key_auth_builder( ## Check CACHE try: with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"): - valid_token = await get_key_object( - hashed_token=hash_token(api_key), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_cache_only=True, + valid_token = IdentityStore.key_from_principal( + await IdentityStore( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_cache_only=True, + ).resolve(hashed_token=hash_token(api_key)) ) except Exception: verbose_logger.debug("api key not found in cache.") @@ -1679,12 +1687,13 @@ async def _user_api_key_auth_builder( try: with tracer.trace("litellm.proxy.auth.get_key_object_from_db"): - valid_token = await get_key_object( - hashed_token=api_key, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, + valid_token = IdentityStore.key_from_principal( + await IdentityStore( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ).resolve(hashed_token=api_key) ) except ProxyException as e: if e.code == 401 or e.code == "401": @@ -2501,6 +2510,34 @@ def _should_skip_budget_checks( return False +def _resolve_request_principal( + request: Request, valid_token: UserAPIKeyAuth +) -> Principal: + """Project the resolved identity into one per-request Principal, off the key + object the builder already fetched, and stamp the request network context + onto it once. X-Forwarded-For is only trusted when the operator configured + ``trusted_proxy_ranges``; otherwise the direct peer is authoritative. + + credential_ref and a stable subject fallback are always set off the token so + the Principal can never be anonymous, even for a keyless service-account key + with no user or alias.""" + cidrs = get_trusted_proxy_cidrs() + network = resolve_network_context( + request, + TrustedProxyConfig(use_forwarded_for=bool(cidrs), trusted_proxy_cidrs=cidrs), + ) + auth_method = ( + AuthMethod.BEARER_JWT if valid_token.jwt_claims else AuthMethod.API_KEY + ) + return IdentityStore._principal_from_key( + valid_token, + auth_method=auth_method, + network=network, + subject_fallback=valid_token.token, + credential_ref=CredentialRef(token_id=valid_token.token), + ) + + @tracer.wrap() async def user_api_key_auth( request: Request, @@ -2615,6 +2652,22 @@ async def user_api_key_auth( model=request_data.get("model") if isinstance(request_data, dict) else None, ) user_api_key_auth_obj.request_route = normalize_request_route(route) + + # Resolve caller identity once, here at the seam, into a single per-request + # Principal projected off the key object the builder already fetched (no + # second lookup). Downstream consumers read identity off this instead of + # re-resolving it. Additive and defensive: a projection failure must never + # reject an already-authenticated request, so it is left unset on failure; + # any future consumer must treat a missing principal as deny, not allow. + try: + request.state.principal = _resolve_request_principal( + request, user_api_key_auth_obj + ) + except Exception as e: + verbose_proxy_logger.warning( + "Principal projection at auth seam failed (non-fatal): %s", e + ) + return user_api_key_auth_obj diff --git a/scripts/ruff_strict_gate.py b/scripts/ruff_strict_gate.py index 9c406b8482b..5951a1215ed 100644 --- a/scripts/ruff_strict_gate.py +++ b/scripts/ruff_strict_gate.py @@ -5,13 +5,9 @@ Each rule has a hard ceiling (baseline + slack) in ruff-strict-budget.json. The gate counts each rule across the whole tree and fails when a rule is both over its ceiling and higher than the base it merges into, so a change is blamed for the violations it adds, never for drift that already exists in the base. - -The base is the merge-base of the current branch with --base; this matches CI, -which checks out the PR head sha and runs the gate against the PR's base sha. """ import argparse -import contextlib import json import re import shutil @@ -19,7 +15,6 @@ import subprocess import sys import tempfile from collections import Counter -from collections.abc import Iterator from pathlib import Path from typing import NamedTuple @@ -45,12 +40,6 @@ class Breach(NamedTuple): added: int -class GateInputs(NamedTuple): - head: list[Violation] - base: dict[str, int] - changed: dict[str, set[int]] - - def _run(cmd: list, cwd: Path = REPO_ROOT) -> str: proc = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True) if proc.returncode not in (0, 1): @@ -67,14 +56,14 @@ def _ruff_json(cwd: Path, config: Path) -> list: return json.loads(raw or "[]") -def collect_violations(root: Path, config: Path) -> list: +def head_violations() -> list: out = [] - for item in _ruff_json(root, config): + for item in _ruff_json(REPO_ROOT, STRICT_CONFIG): name = Path(item["filename"]) rel = ( - (name if name.is_absolute() else root / name) + (name if name.is_absolute() else REPO_ROOT / name) .resolve() - .relative_to(root) + .relative_to(REPO_ROOT) .as_posix() ) out.append(Violation(rel, item["location"]["row"], item["code"])) @@ -85,29 +74,27 @@ def count_by_rule(violations: list) -> dict: return dict(Counter(v.code for v in violations)) -@contextlib.contextmanager -def _temp_worktree(ref: str) -> Iterator[Path]: - parent = Path(tempfile.mkdtemp(prefix="ruff_wt_")) +def base_counts(ref: str) -> dict: + parent = Path(tempfile.mkdtemp(prefix="ruff_base_")) worktree = parent / "wt" try: _run(["git", "worktree", "add", "--detach", str(worktree), ref]) - yield worktree + shutil.copy(STRICT_CONFIG, worktree / "ruff-strict.toml") + items = _ruff_json(worktree, worktree / "ruff-strict.toml") + return dict(Counter(item["code"] for item in items)) finally: - subprocess.run( - ["git", "worktree", "remove", "--force", str(worktree)], - cwd=REPO_ROOT, - capture_output=True, - text=True, - ) + _run(["git", "worktree", "remove", "--force", str(worktree)]) shutil.rmtree(parent, ignore_errors=True) -def base_counts(ref: str) -> dict: - with _temp_worktree(ref) as worktree: - shutil.copy(STRICT_CONFIG, worktree / "ruff-strict.toml") - return count_by_rule( - collect_violations(worktree, worktree / "ruff-strict.toml") - ) +def evaluate(head: dict, base: dict, budget: dict) -> list: + breaches = [] + for rule, spec in budget.items(): + cap = spec["baseline"] + spec["slack"] + total = head.get(rule, 0) + if total > cap and total > base.get(rule, 0): + breaches.append(Breach(rule, total, cap, total - base.get(rule, 0))) + return sorted(breaches) def parse_changed_lines(diff_text: str) -> dict: @@ -123,31 +110,24 @@ def parse_changed_lines(diff_text: str) -> dict: return changed -def evaluate(head: dict, base: dict, budget: dict) -> list: - breaches = [] - for rule, spec in budget.items(): - cap = spec["baseline"] + spec["slack"] - total = head.get(rule, 0) - if total > cap and total > base.get(rule, 0): - breaches.append(Breach(rule, total, cap, total - base.get(rule, 0))) - return sorted(breaches) - - def introduced(violations: list, changed: dict) -> list: return [v for v in violations if v.line in changed.get(v.file, set())] -def gather(base: str) -> GateInputs: +def cmd_check(base: str) -> None: + budget = json.loads(BUDGET_PATH.read_text()) + head = head_violations() base_point = _run(["git", "merge-base", base, "HEAD"]).strip() or base - diff = _run(["git", "diff", base_point, "--unified=0", "--no-color", "--", TARGET]) - return GateInputs( - collect_violations(REPO_ROOT, STRICT_CONFIG), - base_counts(base_point), - parse_changed_lines(diff), + breaches = evaluate(count_by_rule(head), base_counts(base_point), budget) + if not breaches: + print(f"OK: every strict rule is within its codebase ceiling (base {base})") + return + new = introduced( + head, + parse_changed_lines( + _run(["git", "diff", base_point, "--unified=0", "--no-color", "--", TARGET]) + ), ) - - -def report(breaches: list, new: list, base: str) -> None: print(f"FAIL: strict-rule totals exceed their ceiling (base {base}):") for breach in breaches: print( @@ -158,24 +138,12 @@ def report(breaches: list, new: list, base: str) -> None: print( "Reduce the new violations or remove an equal number elsewhere; the ceiling is baseline + slack in ruff-strict-budget.json." ) - summary = "; ".join(f"{b.rule} {b.total}/{b.cap} (+{b.added})" for b in breaches) - print(f"BREACHED RULES: {summary}") - - -def cmd_check(base: str) -> None: - budget = json.loads(BUDGET_PATH.read_text()) - inputs = gather(base) - breaches = evaluate(count_by_rule(inputs.head), inputs.base, budget) - if not breaches: - print(f"OK: every strict rule is within its codebase ceiling (base {base})") - return - report(breaches, introduced(inputs.head, inputs.changed), base) raise SystemExit(1) def cmd_update() -> None: budget = json.loads(BUDGET_PATH.read_text()) - head = count_by_rule(collect_violations(REPO_ROOT, STRICT_CONFIG)) + head = count_by_rule(head_violations()) for rule in budget: budget[rule]["baseline"] = head.get(rule, 0) BUDGET_PATH.write_text(json.dumps(budget, indent=2, sort_keys=True) + "\n") diff --git a/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py b/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py index a45df5df008..fd674afd85a 100644 --- a/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py +++ b/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py @@ -4,6 +4,8 @@ import pytest from fastapi import Request from litellm_enterprise.proxy.auth.user_api_key_auth import enterprise_custom_auth +from litellm.proxy._types import UserAPIKeyAuth + @pytest.mark.asyncio async def test_enterprise_custom_auth_none_user_auth(): @@ -49,16 +51,19 @@ async def test_enterprise_custom_auth_returns_string(): mock_user_auth = AsyncMock(return_value="sk-test-key") request = MagicMock(spec=Request) - with patch( - "litellm.proxy.auth.user_api_key_auth.enterprise_custom_auth", mock_user_auth - ), patch("litellm.proxy.proxy_server.master_key", "sk-1234"), patch( - "litellm.proxy.proxy_server.prisma_client", MagicMock() + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.enterprise_custom_auth", + mock_user_auth, + ), + patch("litellm.proxy.proxy_server.master_key", "sk-1234"), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), ): # Verify the key is correctly handled in _user_api_key_auth_builder with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object" + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key" ) as mock_get_key_object: - mock_get_key_object.return_value = MagicMock( + mock_get_key_object.return_value = UserAPIKeyAuth( token="sk-test-key", user_role="internal_user", team_id=None, @@ -82,9 +87,7 @@ async def test_enterprise_custom_auth_returns_string(): except Exception as e: print("error:", e) - # Verify get_key_object was called with the correct key + # Verify the key lookup was called with the correct hashed key mock_get_key_object.assert_called_once() - # The key should be hashed before being passed to get_key_object - assert mock_get_key_object.call_args[1]["hashed_token"] == hash_token( - "sk-test-key" - ) + # The key should be hashed before being passed to the resolver + assert mock_get_key_object.call_args[0][0] == hash_token("sk-test-key") diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 61c24183964..65d075b7f99 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -60,7 +60,8 @@ async def test_jwt_to_virtual_key_mapping_resolution(): # Use patch to mock get_key_object in the module where it's used with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ) as mock_get_key: mock_get_key.return_value = mock_key_obj @@ -105,7 +106,8 @@ async def test_jwt_to_virtual_key_mapping_no_mapping(): # Mock get_key_object just in case with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ): user_api_key_cache = DualCache() @@ -481,7 +483,8 @@ async def test_reject_behavior_raises_403_on_no_mapping(): user_api_key_cache = DualCache() with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ): with pytest.raises(HTTPException) as exc_info: await _resolve_jwt_to_virtual_key( @@ -519,7 +522,8 @@ async def test_reject_behavior_caches_sentinel_after_db_miss(): user_api_key_cache = DualCache() with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ): # First call — DB miss, should raise 403 and write sentinel with pytest.raises(HTTPException) as exc_info: @@ -578,7 +582,8 @@ async def test_reject_behavior_raises_403_on_cached_no_mapping(): await user_api_key_cache.async_set_cache(cache_key, "__NO_MAPPING__") with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ): with pytest.raises(HTTPException) as exc_info: await _resolve_jwt_to_virtual_key( @@ -671,7 +676,7 @@ async def test_auto_register_creates_key_and_mapping_when_helper_invoked(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key, patch( @@ -803,7 +808,7 @@ async def test_auto_register_race_condition_unique_conflict(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key, patch( @@ -833,13 +838,7 @@ async def test_auto_register_race_condition_unique_conflict(): # Cache should hold the winner's token, not the loser's cached = await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:user-42") assert cached == "winner_token_hash" - mock_get_key.assert_called_once_with( - hashed_token="winner_token_hash", - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) + mock_get_key.assert_called_once_with("winner_token_hash") # ────────────────────────────────────────────── @@ -1093,7 +1092,7 @@ async def test_auto_register_race_conflict_tolerates_delete_failure(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key, patch( @@ -1249,7 +1248,7 @@ async def test_auto_register_helper_stamps_validated_identity_context(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key, patch( diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5ec5d12784f..52634cc25fe 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2365,7 +2365,9 @@ async def test_virtual_key_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) proxy_logging_obj.budget_alerts = AsyncMock() - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:key:test-hashed-token": return 1.5 return fallback_spend @@ -2397,7 +2399,9 @@ async def test_virtual_key_budget_check_fallback_no_counter(): proxy_logging_obj.budget_alerts = AsyncMock() # get_current_spend returns fallback_spend when no counter exists - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): return fallback_spend with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): @@ -2424,7 +2428,9 @@ async def test_team_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) proxy_logging_obj.budget_alerts = AsyncMock() - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team:test-team": return 1.5 return fallback_spend @@ -2449,7 +2455,9 @@ async def test_end_user_budget_check_reads_from_spend_counter(): litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:end_user:customer-1": return 1.5 return fallback_spend @@ -2475,7 +2483,9 @@ async def test_tag_budget_check_reads_from_spend_counter(): litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:tag:paid-tag": return 1.5 return fallback_spend @@ -2523,7 +2533,9 @@ async def test_team_member_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 1.5 return fallback_spend @@ -2756,7 +2768,9 @@ async def test_team_member_budget_check_falls_back_to_team_default_budget_id(): return_value=fake_budget_row ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 70.0 return fallback_spend @@ -2853,7 +2867,9 @@ async def test_team_member_budget_check_per_member_override_wins_over_team_defau mocked_spend = 70.0 - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return mocked_spend return fallback_spend @@ -2943,7 +2959,9 @@ async def test_team_member_budget_check_null_clone_falls_back_to_team_default(): return_value=fake_default_row ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 500.0 return fallback_spend @@ -3010,7 +3028,9 @@ async def test_team_member_budget_check_null_clone_with_null_default_skips_enfor return_value=fake_default_row ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 1000.0 return fallback_spend @@ -3077,7 +3097,9 @@ async def test_team_member_budget_check_zero_team_default_treated_as_no_cap(): return_value=fake_default_row ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 0.0 return fallback_spend @@ -3135,7 +3157,9 @@ async def test_team_member_budget_check_zero_per_member_row_still_blocks(): prisma_client = MagicMock() prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 0.0 return fallback_spend diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index e49f025df2e..3203878a1e0 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -106,7 +106,9 @@ async def test_custom_auth_enforces_end_user_budget_when_common_checks_skipped() litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:end_user:customer-1": return 5.0 return fallback_spend diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 63510086f95..b547ec877e2 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -1206,7 +1206,13 @@ async def test_auth_builder_returns_team_membership_object(): JWTAuthManager, "get_objects", new_callable=AsyncMock, - return_value=(user_object, None, None, mock_team_membership, user_object.user_id), + return_value=( + user_object, + None, + None, + mock_team_membership, + user_object.user_id, + ), ) as mock_get_objects, patch.object( JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock @@ -3509,9 +3515,7 @@ def test_canonical_user_id_no_change_when_ids_match(): user_object = LiteLLM_UserTable(user_id=same, user_email=same) assert ( - JWTAuthManager._canonical_user_id_from_db( - user_id=same, user_object=user_object - ) + JWTAuthManager._canonical_user_id_from_db(user_id=same, user_object=user_object) == same ) @@ -3802,12 +3806,15 @@ async def test_get_objects_team_membership_uses_rebound_user_id(): user_id_jwt_field="email", user_id_upsert=True ) - with patch( - "litellm.proxy.auth.handle_jwt.get_user_object", - side_effect=fake_get_user_object, - ), patch( - "litellm.proxy.auth.handle_jwt.get_team_membership", - side_effect=fake_get_team_membership, + with ( + patch( + "litellm.proxy.auth.handle_jwt.get_user_object", + side_effect=fake_get_user_object, + ), + patch( + "litellm.proxy.auth.handle_jwt.get_team_membership", + side_effect=fake_get_team_membership, + ), ): ( user_object, diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 8d686900ea6..02b1f698132 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -436,7 +436,9 @@ def test_wildcard_custom_prefix_keeps_org_segment_for_non_provider_first_segment result = get_known_models_from_wildcard( wildcard_model="my_hf/*", - litellm_params=LiteLLM_Params(model="huggingface/*", custom_llm_provider="huggingface"), + litellm_params=LiteLLM_Params( + model="huggingface/*", custom_llm_provider="huggingface" + ), ) assert result == ["my_hf/meta-llama/Llama-3-8B"] diff --git a/tests/test_litellm/proxy/auth/test_network.py b/tests/test_litellm/proxy/auth/test_network.py new file mode 100644 index 00000000000..b67723305e4 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_network.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +from fastapi import Request + +from litellm.proxy.auth.network import ( + TrustedProxyConfig, + resolve_client_ip, + resolve_network_context, +) + +TRUSTED = TrustedProxyConfig(use_forwarded_for=True, trusted_proxy_cidrs=["10.0.0.0/8"]) + + +def make_request( + *, + headers: Optional[Dict[str, str]] = None, + client: Optional[Tuple[str, int]] = ("203.0.113.7", 5555), +) -> Request: + raw_headers: List[Tuple[bytes, bytes]] = [ + (key.lower().encode(), value.encode()) for key, value in (headers or {}).items() + ] + scope: Dict[str, Any] = { + "type": "http", + "http_version": "1.1", + "method": "GET", + "path": "/", + "raw_path": b"/", + "query_string": b"", + "headers": raw_headers, + "client": client, + "server": ("testserver", 80), + "scheme": "http", + } + return Request(scope) + + +def test_xff_ignored_when_forwarding_disabled(): + config = TrustedProxyConfig( + use_forwarded_for=False, trusted_proxy_cidrs=["10.0.0.0/8"] + ) + request = make_request( + headers={"x-forwarded-for": "203.0.113.9"}, client=("10.0.0.1", 1) + ) + ip, via_proxy = resolve_client_ip(request, config) + assert ip == "10.0.0.1" + assert via_proxy is False + + +def test_xff_honored_from_trusted_peer(): + request = make_request( + headers={"x-forwarded-for": "203.0.113.9, 10.0.0.5"}, client=("10.0.0.1", 1) + ) + ip, via_proxy = resolve_client_ip(request, TRUSTED) + assert ip == "203.0.113.9" + assert via_proxy is True + + +def test_spoofed_xff_from_untrusted_peer_is_ignored(): + request = make_request( + headers={"x-forwarded-for": "203.0.113.9"}, client=("8.8.8.8", 1) + ) + ip, via_proxy = resolve_client_ip(request, TRUSTED) + assert ip == "8.8.8.8" + assert via_proxy is False + + +def test_right_to_left_parse_skips_chained_trusted_proxies(): + request = make_request( + headers={"x-forwarded-for": "198.51.100.4, 10.1.1.1, 10.0.0.9"}, + client=("10.0.0.1", 1), + ) + ip, via_proxy = resolve_client_ip(request, TRUSTED) + assert ip == "198.51.100.4" + assert via_proxy is True + + +def test_all_trusted_hops_fall_back_to_peer(): + request = make_request( + headers={"x-forwarded-for": "10.1.1.1, 10.0.0.9"}, client=("10.0.0.1", 1) + ) + ip, via_proxy = resolve_client_ip(request, TRUSTED) + assert ip == "10.0.0.1" + assert via_proxy is True + + +def test_invalid_xff_token_is_skipped(): + request = make_request( + headers={"x-forwarded-for": "not-an-ip, 203.0.113.50"}, client=("10.0.0.1", 1) + ) + ip, _ = resolve_client_ip(request, TRUSTED) + assert ip == "203.0.113.50" + + +def test_network_context_captures_host_and_proxy_flag(): + request = make_request( + headers={"x-forwarded-for": "203.0.113.9", "host": "proxy.litellm.ai"}, + client=("10.0.0.1", 1), + ) + ctx = resolve_network_context(request, TRUSTED) + assert ctx.client_ip == "203.0.113.9" + assert ctx.host == "proxy.litellm.ai" + assert ctx.via_trusted_proxy is True diff --git a/tests/test_litellm/proxy/auth/test_onboarding.py b/tests/test_litellm/proxy/auth/test_onboarding.py index c81f4cb7d66..d55a5472af1 100644 --- a/tests/test_litellm/proxy/auth/test_onboarding.py +++ b/tests/test_litellm/proxy/auth/test_onboarding.py @@ -18,7 +18,6 @@ from fastapi import HTTPException import litellm from litellm.proxy._types import InvitationClaim - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/auth/test_resolvers_exceptions.py b/tests/test_litellm/proxy/auth/test_resolvers_exceptions.py new file mode 100644 index 00000000000..f88b7e2be14 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_exceptions.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.auth.resolvers.exceptions import ( + IdentityResolutionError, + KeyNotFoundError, + KeyNotInCacheError, + NoDatabaseConnectionError, + PrincipalMissingSourceKeyError, +) + + +def test_all_resolution_errors_share_one_base(): + errors = [ + NoDatabaseConnectionError(), + KeyNotInCacheError("hashed"), + KeyNotFoundError("hashed"), + PrincipalMissingSourceKeyError(), + ] + assert all(isinstance(e, IdentityResolutionError) for e in errors) + + +def test_key_not_found_preserves_the_public_401_contract(): + # The auth seam catches ProxyException and rewrites the 401 message, so a + # missing key must keep mapping to that exact contract. + error = KeyNotFoundError("hashed-token") + + assert isinstance(error, ProxyException) + assert error.code == "401" + assert error.type == ProxyErrorTypes.token_not_found_in_db.value + assert error.param == "key" + + +def test_key_not_in_cache_names_the_token(): + assert "hashed-token" in str(KeyNotInCacheError("hashed-token")) diff --git a/tests/test_litellm/proxy/auth/test_resolvers_models.py b/tests/test_litellm/proxy/auth/test_resolvers_models.py new file mode 100644 index 00000000000..0fbe10b6f28 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_models.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.resolvers.models import ( + EndUserIdentity, + Principal, + PrincipalType, + ProjectIdentity, + TeamIdentity, + UserIdentity, +) +from litellm.proxy.auth.roles import Role, TeamRole + + +def _principal() -> Principal: + return Principal( + principal_type=PrincipalType.HUMAN, + subject="u1", + auth_method=AuthMethod.OIDC, + ) + + +def test_principal_is_frozen(): + principal = _principal() + with pytest.raises(ValidationError): + principal.subject = "mutated" + + +def test_principal_defaults_are_independent_instances(): + a = _principal() + b = _principal() + assert a.teams == [] and a.scopes == [] and a.audience == [] + assert a.teams is not b.teams + + +def test_principal_requires_identity_core_fields(): + with pytest.raises(ValidationError): + Principal(subject="u1") # missing principal_type + auth_method + + +def test_principal_roles_are_validated_against_role_enum(): + principal = Principal( + principal_type=PrincipalType.HUMAN, + subject="u1", + auth_method=AuthMethod.OIDC, + roles=["org_admin"], + ) + assert principal.roles == [Role.ORG_ADMIN] + assert isinstance(principal.roles[0], Role) + + with pytest.raises(ValidationError): + Principal( + principal_type=PrincipalType.HUMAN, + subject="u1", + auth_method=AuthMethod.OIDC, + roles=["not_a_real_role"], + ) + + +def test_principal_default_network_and_collections(): + principal = Principal( + principal_type=PrincipalType.SERVICE_ACCOUNT, + subject="svc", + auth_method=AuthMethod.MUTUAL_TLS, + ) + assert principal.teams == [] + assert principal.scopes == [] + assert principal.project is None + assert principal.end_user is None + assert principal.network.client_ip is None + assert principal.network.via_trusted_proxy is False + + +def test_team_identity_defaults_to_member_role(): + team = TeamIdentity(id="g1") + assert team.role == TeamRole.MEMBER + + +def test_user_identity_optional_fields_default_none(): + user = UserIdentity(id="u1") + assert user.email is None + assert user.external_id is None + + +def test_project_identity_name_is_optional(): + assert ProjectIdentity(id="p1").name is None + assert ProjectIdentity(id="p1", name="Acme").name == "Acme" + + +def test_end_user_identity_requires_id(): + with pytest.raises(ValidationError): + EndUserIdentity() diff --git a/tests/test_litellm/proxy/auth/test_resolvers_seam.py b/tests/test_litellm/proxy/auth/test_resolvers_seam.py new file mode 100644 index 00000000000..cb42d4d3d2d --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_seam.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +from fastapi import Request + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.resolvers.models import PrincipalType +from litellm.proxy.auth.roles import Role +from litellm.proxy.auth.user_api_key_auth import _resolve_request_principal + + +def _request( + *, + headers: Optional[Dict[str, str]] = None, + client: Optional[Tuple[str, int]] = ("203.0.113.7", 5555), +) -> Request: + raw: List[Tuple[bytes, bytes]] = [ + (k.lower().encode(), v.encode()) for k, v in (headers or {}).items() + ] + scope: Dict[str, Any] = { + "type": "http", + "http_version": "1.1", + "method": "POST", + "path": "/v1/chat/completions", + "raw_path": b"/v1/chat/completions", + "query_string": b"", + "headers": raw, + "client": client, + "server": ("testserver", 80), + "scheme": "http", + } + return Request(scope) + + +def test_seam_projects_full_identity_from_key_object(): + token = UserAPIKeyAuth( + token="hashed-token", + user_id="u-1", + user_role="org_admin", + team_id="t-1", + team_alias="Eng", + org_id="o-1", + organization_alias="Acme", + end_user_id="cust-9", + ) + + principal = _resolve_request_principal(_request(), token) + + assert principal.principal_type == PrincipalType.HUMAN + assert principal.auth_method == AuthMethod.API_KEY + assert principal.user is not None and principal.user.id == "u-1" + assert principal.roles == [Role.ORG_ADMIN] + assert [t.id for t in principal.teams] == ["t-1"] + assert principal.teams[0].name == "Eng" + assert principal.organization is not None and principal.organization.id == "o-1" + assert principal.organization.name == "Acme" + assert principal.end_user is not None and principal.end_user.id == "cust-9" + # the key is always identifiable via credential_ref, even when other ids exist + assert principal.credential_ref.token_id == "hashed-token" + + +def test_seam_principal_is_never_anonymous_for_keyless_service_account(): + # no user_id and no key_alias -> subject must still identify the key + token = UserAPIKeyAuth(token="hashed-token") + + principal = _resolve_request_principal(_request(), token) + + assert principal.principal_type == PrincipalType.SERVICE_ACCOUNT + assert principal.user is None + assert principal.subject == "hashed-token" + assert principal.credential_ref.token_id == "hashed-token" + + +def test_seam_stamps_direct_peer_when_no_trusted_proxy_configured(): + token = UserAPIKeyAuth(token="hashed-token", user_id="u-1") + + # No trusted_proxy_ranges configured -> XFF is not trusted, direct peer wins. + principal = _resolve_request_principal( + _request(headers={"x-forwarded-for": "10.9.9.9"}, client=("203.0.113.7", 5555)), + token, + ) + + assert principal.network.client_ip == "203.0.113.7" + assert principal.network.via_trusted_proxy is False + + +def test_seam_detects_jwt_auth_method(): + token = UserAPIKeyAuth( + token="hashed-token", user_id="u-2", jwt_claims={"sub": "u-2"} + ) + + principal = _resolve_request_principal(_request(), token) + + assert principal.auth_method == AuthMethod.BEARER_JWT diff --git a/tests/test_litellm/proxy/auth/test_resolvers_store.py b/tests/test_litellm/proxy/auth/test_resolvers_store.py new file mode 100644 index 00000000000..5e644a084a5 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_store.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from typing import Dict, Optional + +import pytest + +from litellm.proxy._types import UserAPIKeyAuth, hash_token +from litellm.proxy.auth.resolvers.exceptions import ( + NoDatabaseConnectionError, + PrincipalMissingSourceKeyError, +) +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.resolvers.models import Principal, PrincipalType +from litellm.proxy.auth.resolvers.store import IdentityStore + + +class _FakeCache: + """Stands in for the DualCache get_key_object reads. It returns a cache hit + before the DB is touched, so seeding it exercises resolve without a database + (a non-None prisma client is still required; it is never reached on a hit).""" + + def __init__(self, entries: Optional[Dict[str, object]] = None) -> None: + self._entries = entries or {} + + async def async_get_cache(self, key, *args, **kwargs): + return self._entries.get(key) + + async def async_set_cache(self, *args, **kwargs): + return None + + +async def test_resolve_returns_a_principal_projected_from_the_looked_up_key(): + raw = "sk-live-abc" + key = UserAPIKeyAuth(token=hash_token(raw), user_id="u-1", team_id="t-1") + store = IdentityStore(object(), _FakeCache({hash_token(raw): key})) + + principal = await store.resolve(hashed_token=hash_token(raw)) + + assert isinstance(principal, Principal) + assert principal.principal_type == PrincipalType.HUMAN + assert principal.user is not None and principal.user.id == "u-1" + assert [t.id for t in principal.teams] == ["t-1"] + + +async def test_resolve_carries_the_key_for_key_from_principal(): + raw = "sk-live-abc" + key = UserAPIKeyAuth(token=hash_token(raw), user_id="u-1", team_id="t-1") + store = IdentityStore(object(), _FakeCache({hash_token(raw): key})) + + principal = await store.resolve(hashed_token=hash_token(raw)) + recovered = IdentityStore.key_from_principal(principal) + + assert recovered.user_id == "u-1" + assert recovered.team_id == "t-1" + + +def test_key_from_principal_raises_when_no_source_key_is_carried(): + bare = Principal( + principal_type=PrincipalType.SERVICE_ACCOUNT, + subject="svc", + auth_method=AuthMethod.API_KEY, + ) + with pytest.raises(PrincipalMissingSourceKeyError): + IdentityStore.key_from_principal(bare) + + +async def test_resolve_raises_without_a_db_connection(): + store = IdentityStore(None, _FakeCache()) + with pytest.raises(NoDatabaseConnectionError): + await store.resolve(hashed_token="missing") diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 80f12d4459f..27c1b04fbd9 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -1106,7 +1106,7 @@ async def test_proxy_admin_expired_key_from_cache(): # Mock get_key_object to return expired token from cache with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key_object, patch( @@ -1261,7 +1261,7 @@ async def test_scim_deactivated_user_key_is_rejected(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, return_value=valid_token, ), @@ -2484,7 +2484,7 @@ async def test_user_api_key_auth_builder_no_blocking_calls(): stack.enter_context(p) stack.enter_context( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, return_value=valid_token, ) @@ -2598,7 +2598,7 @@ async def test_team_metadata_refreshed_from_team_object_during_auth(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, return_value=valid_token, ), @@ -3652,7 +3652,7 @@ async def _run_builder_with_key_lookup(get_key_object_mock): request._url = URL(url="/chat/completions") with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", get_key_object_mock, ), patch( diff --git a/tests/test_litellm/test_ruff_strict_gate.py b/tests/test_litellm/test_ruff_strict_gate.py index 96852e3a8a5..22255f0555e 100644 --- a/tests/test_litellm/test_ruff_strict_gate.py +++ b/tests/test_litellm/test_ruff_strict_gate.py @@ -82,13 +82,3 @@ def test_introduced_keeps_only_violations_on_changed_lines(): @pytest.mark.parametrize("hunk", ["@@ -1 +1 @@", "@@ -1,0 +1,2 @@"]) def test_parse_changed_lines_handles_single_and_ranged_hunks(hunk): assert gate.parse_changed_lines(f"+++ b/litellm/a.py\n{hunk}\n")["litellm/a.py"] - - -def test_report_emits_breached_rules_as_final_line(capsys): - # CI surfaces only the tail of the log, so the breached-rule summary (rule, - # total/cap, added) must be the last line or it gets truncated away. - breaches = sorted([gate.Breach("UP045", 530, 529, 1), gate.Breach("ANN401", 12, 10, 2)]) - new = [gate.Violation("litellm/types/llms/bedrock.py", 16, "UP045")] - gate.report(breaches, new, "origin/litellm_internal_staging") - last = capsys.readouterr().out.strip().splitlines()[-1] - assert last == "BREACHED RULES: ANN401 12/10 (+2); UP045 530/529 (+1)" From 9f97111edd736cf81e532f345663885457b916a9 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 20 Jun 2026 19:49:07 -0700 Subject: [PATCH 02/50] feat(fireworks_ai): sync chat completions endpoint with full API surface (#30885) * feat(fireworks_ai): sync chat completions endpoint with full API surface Add 23 missing request parameters to get_supported_openai_params(): seed, top_logprobs, min_p, typical_p, repetition_penalty, mirostat_target, mirostat_lr, logit_bias, echo, echo_last, ignore_eos, prompt_cache_key, prompt_cache_isolation_key, raw_output, perf_metrics_in_response, return_token_ids, safe_tokenization, service_tier, metadata, speculation, prediction, stream_options, sampling_mask. Also add reasoning_history gated on supports_reasoning. Fix prompt_truncate_length to prompt_truncate_len to match the actual API parameter name. The old name was never in DEFAULT_CHAT_COMPLETION_PARAM_VALUES, so it always went to extra_body and was rejected by Fireworks; it never actually worked. Normalize reasoning_effort boolean values to strings: True becomes "medium", False becomes "none". The Fireworks OpenAPI schema documents these as accepted types, but the server rejects non-string values with HTTP 400 in practice. Integers pass through as-is since the server is expected to validate them. Auto-inject stream_options.include_usage=true when stream=true and the user has not explicitly set stream_options. Without this, Fireworks returns null usage in all streaming chunks, which is inconsistent with the non-streaming behavior where usage is always present. If the user explicitly sets include_usage=false, it is preserved. Capture Fireworks-specific response fields in transform_response(): perf_metrics, prompt_token_ids, raw_output, and token_ids are now extracted from the response and stored in response._hidden_params (fireworks_perf_metrics, fireworks_prompt_token_ids, fireworks_raw_outputs, fireworks_token_ids) so they are accessible to logging, the proxy, and downstream consumers when the corresponding request parameters are enabled. Remove deprecated document inlining logic. Document inlining was deprecated on 2025-06-30 (https://docs.fireworks.ai/updates/changelog#-document-inlining-deprecation). This removes _add_transform_inline_image_block(), the file-to-image_url migration in _transform_messages_helper(), and the disable_add_transform_inline_image_block lookup. Current models that support image input do so natively as VLMs. cache_control, provider_specific_fields, and thinking_blocks stripping is retained. Update get_provider_info() to look up supports_vision and supports_pdf_input from the model cost map instead of hardcoding both to True (which was based on the now-deprecated document inlining). supports_prompt_caching remains True. API docs: https://docs.fireworks.ai/api-reference/post-chatcompletions Reasoning guide: https://docs.fireworks.ai/guides/reasoning Prompt caching: https://docs.fireworks.ai/guides/prompt-caching * fix fireworks chat api surface gaps * Scope Fireworks thinking param to reasoning models * style: fix black formatting * fix(test): update minimax-m3 expected_vision to True * test: cover non-dict content branch in transform_messages_helper * fix(fireworks_ai): remove metadata from supported params to prevent internal metadata disclosure * test(fireworks_ai): replace stale document-inlining capability test The CircleCI-only litellm_utils_tests suite still asserted the old behavior where document inlining made every Fireworks model report supports_pdf_input and supports_vision as True. That premise was removed in this change, so the test now reflects cost-map-driven capabilities: unmapped models no longer advertise vision/PDF support while mapped VLMs like minimax-m3 still do. * test(fireworks_ai): add end-to-end regression for native OpenAI params The existing coverage for the newly supported OpenAI-native params asserted list membership in get_supported_openai_params or called map_openai_params with a hand-built dict, both of which bypass the get_optional_params gate (DEFAULT_CHAT_COMPLETION_PARAM_VALUES). That gate is what previously raised UnsupportedParamsError for seed, top_logprobs, logit_bias, prompt_cache_key, service_tier and prediction when drop_params=False. Assert the full path so a revert of the supported-params additions fails the test instead of passing a shallow membership check. * test(fireworks_ai): fix test isolation in vision/inlining tests Use monkeypatch in test_fireworks_ai_vision_capability_from_cost_map so the LITELLM_LOCAL_MODEL_COST_MAP env var and litellm.model_cost are restored after the test instead of leaking global state into the rest of the process. Switch the document-inlining integration tests off deepseek-v3p1, whose supports_vision is null in the cost map, onto minimax-m3 which is explicitly supports_vision:true. The pass-through assertions no longer depend on a model incidentally not being marked non-vision. * fix(fireworks_ai): gate image rejection on exact vision capability The image_url rejection read supports_vision via _get_model_cost_capability, which falls back to hyphen-boundary substring matching when no exact cost-map entry exists. A custom or fine-tuned model id that merely contains a known non-vision model's short name (e.g. an id ending in -glm-5p2) inherited that entry's supports_vision:false and hard-failed valid image_url blocks on a vision-capable deployment. Split the exact candidate-key lookup into _get_model_cost_capability_exact and use it for the hard rejection so a fuzzy match can never block images; the substring fallback stays a soft signal for capability reporting. Also rewrites the fallback as a comprehension + max instead of an accumulating loop. * feat(fireworks_ai): surface response fields on streaming responses The Fireworks-specific response fields (perf_metrics, prompt_token_ids, per-choice raw_output and token_ids) were only captured into _hidden_params in transform_response, which runs for non-streaming completions; streaming chat went through the default OpenAI chunk handler and dropped them. Add a FireworksAIChatCompletionStreamingHandler that the provider now returns from get_model_response_iterator. It reuses one extraction helper with transform_response and attaches the fields to each streamed chunk's provider_specific_fields, which is the channel litellm preserves when it rebuilds streamed chunks (per-chunk _hidden_params is not carried through). Per-choice token_ids/raw_output ride the content chunks; response-level perf_metrics/prompt_token_ids ride the final usage chunk. Covered by an end-to-end streaming test through litellm.completion(stream=True). --------- Co-authored-by: Ahmad Shahzad Co-authored-by: Graham Neubig <398875+neubig@users.noreply.github.com> --- .../llms/fireworks_ai/chat/transformation.py | 276 +++++--- ...odel_prices_and_context_window_backup.json | 4 +- model_prices_and_context_window.json | 4 +- tests/litellm_utils_tests/test_utils.py | 17 +- .../test_fireworks_ai_translation.py | 92 ++- .../test_fireworks_ai_chat_transformation.py | 600 ++++++++++++++++-- tests/test_litellm/test_utils.py | 2 +- 7 files changed, 809 insertions(+), 186 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 341c2fc7350..7e4395959b9 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -1,5 +1,15 @@ import json -from typing import Any, List, Literal, Optional, Tuple, Union, cast +from typing import ( + Any, + AsyncIterator, + Iterator, + List, + Literal, + Optional, + Tuple, + Union, + cast, +) import httpx @@ -15,7 +25,6 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, - ChatCompletionImageObject, ChatCompletionToolParam, OpenAIChatCompletionToolParam, ) @@ -25,6 +34,7 @@ from litellm.types.utils import ( Function, Message, ModelResponse, + ModelResponseStream, ProviderSpecificModelInfo, ) from litellm.utils import ( @@ -34,10 +44,34 @@ from litellm.utils import ( supports_tool_choice, ) -from ...openai.chat.gpt_transformation import OpenAIGPTConfig +from ...openai.chat.gpt_transformation import ( + OpenAIChatCompletionStreamingHandler, + OpenAIGPTConfig, +) from ..common_utils import FireworksAIException +def _extract_fireworks_hidden_params(payload: dict) -> dict: + """ + Collect Fireworks-specific response fields (perf_metrics, prompt_token_ids, + per-choice raw_output and token_ids) from a non-streaming completion payload + or a single streaming chunk, so the same data lands in ``_hidden_params`` on + both response paths. + """ + choices = [c for c in (payload.get("choices") or []) if isinstance(c, dict)] + top_level = { + f"fireworks_{field}": payload[field] + for field in ("perf_metrics", "prompt_token_ids") + if field in payload + } + per_choice = { + f"fireworks_{dest}": [c[field] for c in choices if field in c] + for field, dest in (("raw_output", "raw_outputs"), ("token_ids", "token_ids")) + if any(field in c for c in choices) + } + return {**top_level, **per_choice} + + class FireworksAIConfig(OpenAIGPTConfig): """ Reference: https://docs.fireworks.ai/api-reference/post-chatcompletions @@ -60,8 +94,7 @@ class FireworksAIConfig(OpenAIGPTConfig): logprobs: Optional[int] = None reasoning_effort: Optional[str] = None - # Non OpenAI parameters - Fireworks AI only params - prompt_truncate_length: Optional[int] = None + prompt_truncate_len: Optional[int] = None context_length_exceeded_behavior: Optional[Literal["error", "truncate"]] = None def __init__( @@ -80,7 +113,7 @@ class FireworksAIConfig(OpenAIGPTConfig): user: Optional[str] = None, logprobs: Optional[int] = None, reasoning_effort: Optional[str] = None, - prompt_truncate_length: Optional[int] = None, + prompt_truncate_len: Optional[int] = None, context_length_exceeded_behavior: Optional[Literal["error", "truncate"]] = None, ) -> None: locals_ = locals().copy() @@ -108,8 +141,30 @@ class FireworksAIConfig(OpenAIGPTConfig): "response_format", "user", "logprobs", - "prompt_truncate_length", + "prompt_truncate_len", "context_length_exceeded_behavior", + "seed", + "top_logprobs", + "min_p", + "typical_p", + "repetition_penalty", + "mirostat_target", + "mirostat_lr", + "logit_bias", + "echo", + "echo_last", + "ignore_eos", + "prompt_cache_key", + "prompt_cache_isolation_key", + "raw_output", + "perf_metrics_in_response", + "return_token_ids", + "safe_tokenization", + "service_tier", + "speculation", + "prediction", + "stream_options", + "sampling_mask", ] # Only add tools for models that support function calling @@ -133,9 +188,11 @@ class FireworksAIConfig(OpenAIGPTConfig): if supports_tool_choice(model=model, custom_llm_provider="fireworks_ai"): supported_params.append("tool_choice") - # Only add reasoning_effort for models that support it + # Only add reasoning params for models that support it if supports_reasoning(model=model, custom_llm_provider="fireworks_ai"): supported_params.append("reasoning_effort") + supported_params.append("reasoning_history") + supported_params.append("thinking") return supported_params @@ -151,6 +208,18 @@ class FireworksAIConfig(OpenAIGPTConfig): param == "tools" and value is not None for param, value in non_default_params.items() ) + if ( + non_default_params.get("thinking") is not None + and non_default_params.get("reasoning_effort") is not None + ): + raise litellm.BadRequestError( + message=( + "Fireworks AI chat completions does not support specifying both " + "`thinking` and `reasoning_effort` in the same request." + ), + model=model, + llm_provider="fireworks_ai", + ) for param, value in non_default_params.items(): if param == "tool_choice": @@ -174,40 +243,19 @@ class FireworksAIConfig(OpenAIGPTConfig): optional_params["response_format"] = value elif param == "max_completion_tokens": optional_params["max_tokens"] = value + elif param == "reasoning_effort": + if value is True: + optional_params["reasoning_effort"] = "medium" + elif value is False: + optional_params["reasoning_effort"] = "none" + else: + optional_params["reasoning_effort"] = value elif param in supported_openai_params: if value is not None: optional_params[param] = value return optional_params - def _add_transform_inline_image_block( - self, - content: ChatCompletionImageObject, - model: str, - disable_add_transform_inline_image_block: Optional[bool], - ) -> ChatCompletionImageObject: - """ - Add transform_inline to the image_url (allows non-vision models to parse documents/images/etc.) - - ignore if model is a vision model - - ignore if user has disabled this feature - """ - if ( - "vision" in model or disable_add_transform_inline_image_block - ): # allow user to toggle this feature. - return content - if isinstance(content["image_url"], str): - # Skip base64 data URLs — appending #transform=inline corrupts the - # base64 payload and causes an "Incorrect padding" decode error on - # the Fireworks side. Data URLs are already inlined by definition. - # Lower-case before checking: URI schemes are case-insensitive (RFC 3986). - if not content["image_url"].lower().startswith("data:"): - content["image_url"] = f"{content['image_url']}#transform=inline" - elif isinstance(content["image_url"], dict): - url = content["image_url"]["url"] - if not url.lower().startswith("data:"): - content["image_url"]["url"] = f"{url}#transform=inline" - return content - def _transform_tools( self, tools: List[OpenAIChatCompletionToolParam] ) -> List[OpenAIChatCompletionToolParam]: @@ -225,36 +273,46 @@ class FireworksAIConfig(OpenAIGPTConfig): self, messages: List[AllMessageValues], model: str, litellm_params: dict ) -> List[AllMessageValues]: """ - Add 'transform=inline' to the url of the image_url + Strip fields not permitted by FireworksAI from messages. """ from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, - migrate_file_to_image_url, ) - disable_add_transform_inline_image_block = cast( - Optional[bool], - litellm_params.get("disable_add_transform_inline_image_block") - or litellm.disable_add_transform_inline_image_block, + supports_vision_value = self._get_model_cost_capability_exact( + model=model, capability="supports_vision" ) - ## For any 'file' message type with pdf content, move to 'image_url' message type - for message in messages: - if message["role"] == "user": - _message_content = message.get("content") - if _message_content is not None and isinstance(_message_content, list): - for idx, content in enumerate(_message_content): - if content["type"] == "file": - _message_content[idx] = migrate_file_to_image_url(content) for message in messages: if message["role"] == "user": _message_content = message.get("content") if _message_content is not None and isinstance(_message_content, list): for content in _message_content: - if content["type"] == "image_url": - content = self._add_transform_inline_image_block( - content=content, + if not isinstance(content, dict): + continue + if content.get("type") == "file": + raise litellm.BadRequestError( + message=( + "Fireworks AI chat completions does not support " + "file content blocks. For PDFs, convert pages to " + "images and send image_url blocks to a Fireworks " + "vision model, or extract text before calling a " + "text-only model." + ), model=model, - disable_add_transform_inline_image_block=disable_add_transform_inline_image_block, + llm_provider="fireworks_ai", + ) + if ( + content.get("type") == "image_url" + and supports_vision_value is False + ): + raise litellm.BadRequestError( + message=( + f"Fireworks AI model {model} does not support " + "image inputs. Use a Fireworks vision model or " + "remove image_url content blocks." + ), + model=model, + llm_provider="fireworks_ai", ) filter_value_from_dict(cast(dict, message), "cache_control") # Remove fields not permitted by FireworksAI (additionalProperties: false @@ -317,43 +375,55 @@ class FireworksAIConfig(OpenAIGPTConfig): return True return ("-" + key_short + "-") in short_name - def _get_model_cost_capability(self, model: str, capability: str) -> Optional[bool]: + @staticmethod + def _short_model_name(model: str) -> str: short_name = model if short_name.startswith("fireworks_ai/"): short_name = short_name[len("fireworks_ai/") :] if short_name.startswith("accounts/fireworks/models/"): short_name = short_name[len("accounts/fireworks/models/") :] + return short_name - candidate_keys = [ + def _get_model_cost_capability_exact( + self, model: str, capability: str + ) -> Optional[bool]: + short_name = self._short_model_name(model) + candidate_keys = ( model, f"fireworks_ai/{short_name}", f"fireworks_ai/accounts/fireworks/models/{short_name}", - ] - + ) for candidate_key in candidate_keys: model_info = litellm.model_cost.get(candidate_key) if model_info is not None and model_info.get(capability) is not None: return cast(Optional[bool], model_info.get(capability)) + return None - # Fallback: preserve historical substring matching for model name - # variants (e.g. fine-tuned or regionally-suffixed versions of a - # known model). Pick the *longest* matching entry so a more specific - # known model (e.g. "qwen3-8b-instruct") wins over a less specific - # one (e.g. "qwen3-8b") when the query model is more specific still. - # Use hyphen-aligned matching to avoid false positives where a short - # known model name is an unrelated substring of a longer one. - best_match_short: Optional[str] = None - best_match_value: Optional[bool] = None - for key_short, model_info in self._get_fireworks_index(): - if model_info.get(capability) is None: - continue - if not self._matches_on_hyphen_boundary(short_name, key_short): - continue - if best_match_short is None or len(key_short) > len(best_match_short): - best_match_short = key_short - best_match_value = cast(Optional[bool], model_info.get(capability)) + def _get_model_cost_capability(self, model: str, capability: str) -> Optional[bool]: + exact = self._get_model_cost_capability_exact( + model=model, capability=capability + ) + if exact is not None: + return exact - return best_match_value + # Fallback: substring matching for model name variants (e.g. fine-tuned + # or regionally-suffixed versions of a known model). Pick the *longest* + # matching entry so a more specific known model (e.g. "qwen3-8b-instruct") + # wins over a less specific one (e.g. "qwen3-8b"). Hyphen-aligned matching + # avoids false positives where a short known name is an unrelated + # substring of a longer one. This stays a soft signal: capability-gated + # hard rejections use the exact lookup so a fuzzy match never blocks a + # custom deployment. + short_name = self._short_model_name(model) + matches = [ + (key_short, cast(Optional[bool], model_info.get(capability))) + for key_short, model_info in self._get_fireworks_index() + if model_info.get(capability) is not None + and self._matches_on_hyphen_boundary(short_name, key_short) + ] + if not matches: + return None + return max(matches, key=lambda match: len(match[0]))[1] def get_provider_info(self, model: str) -> ProviderSpecificModelInfo: supports_function_calling_value = self._get_model_cost_capability( @@ -362,12 +432,16 @@ class FireworksAIConfig(OpenAIGPTConfig): supports_reasoning_value = self._get_model_cost_capability( model=model, capability="supports_reasoning" ) + supports_vision_value = self._get_model_cost_capability( + model=model, capability="supports_vision" + ) + supports_pdf_input_value = self._get_model_cost_capability( + model=model, capability="supports_pdf_input" + ) provider_specific_model_info: ProviderSpecificModelInfo = { "supports_function_calling": True, "supports_prompt_caching": True, # https://docs.fireworks.ai/guides/prompt-caching - "supports_pdf_input": True, # via document inlining - "supports_vision": True, # via document inlining } if supports_function_calling_value is not None: @@ -381,6 +455,14 @@ class FireworksAIConfig(OpenAIGPTConfig): supports_reasoning_value ) + if supports_vision_value is not None: + provider_specific_model_info["supports_vision"] = supports_vision_value + + if supports_pdf_input_value is not None: + provider_specific_model_info["supports_pdf_input"] = ( + supports_pdf_input_value + ) + return provider_specific_model_info def transform_request( @@ -402,6 +484,15 @@ class FireworksAIConfig(OpenAIGPTConfig): if "tools" in optional_params and optional_params["tools"] is not None: tools = self._transform_tools(tools=optional_params["tools"]) optional_params["tools"] = tools + if optional_params.get("stream"): + stream_options = optional_params.get("stream_options") + if stream_options is None: + optional_params["stream_options"] = {"include_usage": True} + elif stream_options.get("include_usage") is not False: + optional_params["stream_options"] = { + **stream_options, + "include_usage": True, + } return super().transform_request( model=model, messages=messages, @@ -494,10 +585,25 @@ class FireworksAIConfig(OpenAIGPTConfig): ) ) - response._hidden_params = {"additional_headers": additional_headers} + response._hidden_params = { + "additional_headers": additional_headers, + **_extract_fireworks_hidden_params(completion_response), + } return response + def get_model_response_iterator( + self, + streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], + sync_stream: bool, + json_mode: Optional[bool] = False, + ) -> Any: + return FireworksAIChatCompletionStreamingHandler( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) + def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: @@ -554,3 +660,15 @@ class FireworksAIConfig(OpenAIGPTConfig): or get_secret_str("FIREWORKSAI_API_KEY") or get_secret_str("FIREWORKS_AI_TOKEN") ) + + +class FireworksAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): + def chunk_parser(self, chunk: dict) -> ModelResponseStream: + parsed = super().chunk_parser(chunk) + fireworks_fields = _extract_fireworks_hidden_params(chunk) + if fireworks_fields: + parsed.provider_specific_fields = { + **(getattr(parsed, "provider_specific_fields", None) or {}), + **fireworks_fields, + } + return parsed diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7a5f8b9e1e3..5c962cf8440 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -15011,7 +15011,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, "fireworks_ai/accounts/fireworks/models/mixtral-8x22b-instruct-hf": { "input_cost_per_token": 1.2e-06, @@ -15314,7 +15314,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, "fireworks_ai/qwen3p7-plus": { "cache_read_input_token_cost": 8e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 47b7190185e..3c50dde9277 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -15019,7 +15019,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, "fireworks_ai/accounts/fireworks/models/mixtral-8x22b-instruct-hf": { "input_cost_per_token": 1.2e-06, @@ -15322,7 +15322,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, "fireworks_ai/qwen3p7-plus": { "cache_read_input_token_cost": 8e-08, diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index d64633413a0..697c3837602 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -1588,18 +1588,21 @@ def test_token_counter_with_image_url_with_detail_high(): assert _tokens == DEFAULT_IMAGE_TOKEN_COUNT + 7 -def test_fireworks_ai_document_inlining(): +def test_fireworks_ai_vision_capability_from_cost_map(monkeypatch): """ - With document inlining, all fireworks ai models are now: - - supports_pdf - - supports_vision + Fireworks deprecated document inlining on 2025-06-30, so vision/PDF support is + no longer hardcoded to True for every Fireworks model. Capabilities are read + from the model cost map: unmapped models no longer advertise vision or PDF + support, while mapped VLMs still do. """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) from litellm.utils import supports_pdf_input, supports_vision - litellm._turn_on_debug() + assert supports_vision("fireworks_ai/llama-3.1-8b-instruct") is False + assert supports_pdf_input("fireworks_ai/llama-3.1-8b-instruct") is False - assert supports_pdf_input("fireworks_ai/llama-3.1-8b-instruct") is True - assert supports_vision("fireworks_ai/llama-3.1-8b-instruct") is True + assert supports_vision("fireworks_ai/minimax-m3") is True def test_logprobs_type(): diff --git a/tests/llm_translation/test_fireworks_ai_translation.py b/tests/llm_translation/test_fireworks_ai_translation.py index 4e5bef16b8c..204f4d9e31b 100644 --- a/tests/llm_translation/test_fireworks_ai_translation.py +++ b/tests/llm_translation/test_fireworks_ai_translation.py @@ -9,7 +9,6 @@ sys.path.insert( import litellm from litellm import transcription from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig -from base_llm_unit_tests import BaseLLMChatTest from base_audio_transcription_unit_tests import BaseLLMAudioTranscriptionTest fireworks = FireworksAIConfig() @@ -146,11 +145,8 @@ class TestFireworksAIAudioTranscription(BaseLLMAudioTranscriptionTest): ) def test_document_inlining_example(disable_add_transform_inline_image_block): """ - Document inlining appends ``#transform=inline`` to image/PDF URLs in the - outgoing request unless explicitly disabled. Assert the transform on the - serialized payload rather than making a live Fireworks call — the live - call only proved the model responded and broke whenever Fireworks rotated - its serverless model catalog. + Fireworks document inlining has been removed from the platform. LiteLLM + must not append ``#transform=inline`` regardless of the legacy disable flag. """ from unittest.mock import patch @@ -163,7 +159,7 @@ def test_document_inlining_example(disable_add_transform_inline_image_block): with patch.object(client, "post") as mock_post: try: completion( - model="fireworks_ai/accounts/fireworks/models/deepseek-v3p1", + model="fireworks_ai/accounts/fireworks/models/minimax-m3", messages=[ { "role": "user", @@ -182,89 +178,80 @@ def test_document_inlining_example(disable_add_transform_inline_image_block): disable_add_transform_inline_image_block=disable_add_transform_inline_image_block, client=client, ) - except Exception as e: - print(e) + except Exception: + pass mock_post.assert_called_once() json_data = json.loads(mock_post.call_args.kwargs["data"]) sent_url = json_data["messages"][0]["content"][0]["image_url"]["url"] - if disable_add_transform_inline_image_block is True: - assert sent_url == pdf_url - assert "#transform=inline" not in sent_url - else: - assert sent_url == pdf_url + "#transform=inline" + assert sent_url == pdf_url + assert "#transform=inline" not in sent_url @pytest.mark.parametrize( - "content, model, expected_url", + "content, expected_url", [ ( {"image_url": "http://example.com/image.png"}, - "gpt-4", - "http://example.com/image.png#transform=inline", + "http://example.com/image.png", ), ( {"image_url": {"url": "http://example.com/image.png"}}, - "gpt-4", - {"url": "http://example.com/image.png#transform=inline"}, + {"url": "http://example.com/image.png"}, ), - ( - {"image_url": "http://example.com/image.png"}, - "vision-gpt", - "http://example.com/image.png", - ), - # data: URLs must never have #transform=inline appended — doing so - # corrupts the base64 payload (fixes #23583). - # URI schemes are case-insensitive (RFC 3986) so check all variants. ( {"image_url": "data:image/png;base64,iVBORw0KGgo="}, - "gpt-4", "data:image/png;base64,iVBORw0KGgo=", ), ( {"image_url": {"url": "data:image/jpeg;base64,/9j/4AAQ=="}}, - "gpt-4", {"url": "data:image/jpeg;base64,/9j/4AAQ=="}, ), ( {"image_url": "Data:image/png;base64,iVBORw0KGgo="}, - "gpt-4", "Data:image/png;base64,iVBORw0KGgo=", ), ], ) -def test_transform_inline(content, model, expected_url): +def test_transform_inline_no_longer_added(content, expected_url): + image_block = {"type": "image_url", **content} + messages = [{"role": "user", "content": [image_block]}] - result = litellm.FireworksAIConfig()._add_transform_inline_image_block( - content=content, model=model, disable_add_transform_inline_image_block=False + result = litellm.FireworksAIConfig()._transform_messages_helper( + messages=messages, + model="accounts/fireworks/models/minimax-m3", + litellm_params={}, ) + result_image_block = result[0]["content"][0] if isinstance(expected_url, str): - assert result["image_url"] == expected_url + assert result_image_block["image_url"] == expected_url else: - assert result["image_url"]["url"] == expected_url["url"] + assert result_image_block["image_url"]["url"] == expected_url["url"] @pytest.mark.parametrize( - "model, is_disabled, expected_url", - [ - ("gpt-4", True, "http://example.com/image.png"), - ("vision-gpt", False, "http://example.com/image.png"), - ("gpt-4", False, "http://example.com/image.png#transform=inline"), - ], + "is_disabled", + [True, False], ) -def test_global_disable_flag(model, is_disabled, expected_url): - content = {"image_url": "http://example.com/image.png"} - result = litellm.FireworksAIConfig()._add_transform_inline_image_block( - content=content, - model=model, - disable_add_transform_inline_image_block=is_disabled, +def test_global_disable_flag_no_longer_adds_transform_inline(is_disabled): + url = "http://example.com/image.png" + litellm.disable_add_transform_inline_image_block = is_disabled + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": url}], + } + ] + result = litellm.FireworksAIConfig()._transform_messages_helper( + messages=messages, + model="accounts/fireworks/models/minimax-m3", + litellm_params={}, ) - assert result["image_url"] == expected_url + assert result[0]["content"][0]["image_url"] == url litellm.disable_add_transform_inline_image_block = False # Reset for other tests def test_global_disable_flag_with_transform_messages_helper(monkeypatch): - from openai import OpenAI from unittest.mock import patch from litellm import completion from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -279,7 +266,7 @@ def test_global_disable_flag_with_transform_messages_helper(monkeypatch): ) as mock_post: try: completion( - model="fireworks_ai/accounts/fireworks/models/deepseek-v3p1", + model="fireworks_ai/accounts/fireworks/models/minimax-m3", messages=[ { "role": "user", @@ -296,11 +283,10 @@ def test_global_disable_flag_with_transform_messages_helper(monkeypatch): ], client=client, ) - except Exception as e: - print(e) + except Exception: + pass mock_post.assert_called_once() - print(mock_post.call_args.kwargs) json_data = json.loads(mock_post.call_args.kwargs["data"]) assert ( "#transform=inline" diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 683ec158f44..03e763a4161 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1,9 +1,8 @@ import json import os import sys -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch -import httpx import pytest import litellm @@ -12,10 +11,14 @@ sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path -from litellm import get_model_info, supports_reasoning +from litellm import get_model_info, supports_reasoning, supports_vision from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig -from litellm.types.llms.openai import ChatCompletionToolCallFunctionChunk -from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message +from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Function, + Message, + ModelResponse, +) @pytest.fixture(autouse=True) @@ -98,12 +101,12 @@ def test_supports_reasoning_effort(): for model in supported_models: assert ( - supports_reasoning(model=model, custom_llm_provider="fireworks_ai") == True + supports_reasoning(model=model, custom_llm_provider="fireworks_ai") is True ), f"{model} should support reasoning_effort" for model in unsupported_models: assert ( - supports_reasoning(model=model, custom_llm_provider="fireworks_ai") == False + supports_reasoning(model=model, custom_llm_provider="fireworks_ai") is False ), f"{model} should not support reasoning_effort" @@ -115,11 +118,13 @@ def test_get_supported_openai_params_reasoning_effort(): "fireworks_ai/accounts/fireworks/models/glm-5p1" ) assert "reasoning_effort" in supported_params + assert "thinking" in supported_params unsupported_params = config.get_supported_openai_params( "fireworks_ai/accounts/fireworks/models/llama-v3-70b-instruct" ) assert "reasoning_effort" not in unsupported_params + assert "thinking" not in unsupported_params def test_get_supported_openai_params_parallel_tool_calls(): @@ -181,41 +186,6 @@ def test_get_provider_info_omits_false_supports_reasoning(monkeypatch): assert "supports_reasoning" not in info -def test_add_transform_inline_image_block_skips_data_urls(): - """ - data: URLs must not have #transform=inline appended — doing so corrupts the - base64 payload and raises binascii.Error: Incorrect padding on the Fireworks side. - Regression test for https://github.com/BerriAI/litellm/issues/23583 - """ - config = FireworksAIConfig() - data_url = "data:image/jpeg;base64,/9j/4AAQSkZJRgAB" - - # str branch - str_content = {"type": "image_url", "image_url": data_url} - result = config._add_transform_inline_image_block( - str_content, model="gpt-4", disable_add_transform_inline_image_block=False - ) - assert result["image_url"] == data_url, "data URL must not be modified (str branch)" - - # dict branch - dict_content = {"type": "image_url", "image_url": {"url": data_url}} - result = config._add_transform_inline_image_block( - dict_content, model="gpt-4", disable_add_transform_inline_image_block=False - ) - assert ( - result["image_url"]["url"] == data_url - ), "data URL must not be modified (dict branch)" - - # regular https URL should still get the suffix - https_content = {"type": "image_url", "image_url": "https://example.com/image.jpg"} - result = config._add_transform_inline_image_block( - https_content, model="gpt-4", disable_add_transform_inline_image_block=False - ) - assert result["image_url"].endswith( - "#transform=inline" - ), "https URL should get #transform=inline" - - @pytest.mark.parametrize( "api_base, expected_url_prefix", [ @@ -582,3 +552,549 @@ def test_transform_request_routes_short_form_model_to_models_path(): headers={}, ) assert result["model"] == "accounts/fireworks/models/glm-5p2" + + +def _make_fireworks_raw_response(body: dict) -> MagicMock: + mock = MagicMock() + mock.status_code = 200 + mock.json.return_value = body + mock.text = json.dumps(body) + mock.headers = {} + return mock + + +_BASE_CHAT_COMPLETION_RESPONSE: dict = { + "id": "resp-test", + "object": "chat.completion", + "created": 1234567890, + "model": "glm-5p1", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, +} + + +def _run_transform_response(response_body: dict) -> ModelResponse: + config = FireworksAIConfig() + raw_response = _make_fireworks_raw_response(response_body) + logging_obj = MagicMock() + return config.transform_response( + model="accounts/fireworks/models/glm-5p1", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "Hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key="test-key", + ) + + +_REASONING_MODEL = "fireworks_ai/accounts/fireworks/models/glm-5p1" +_NON_REASONING_MODEL = "fireworks_ai/accounts/fireworks/models/llama-v3-70b-instruct" + + +def test_get_supported_openai_params_includes_all_fireworks_params(): + config = FireworksAIConfig() + params = config.get_supported_openai_params(_REASONING_MODEL) + + required = [ + "seed", + "top_logprobs", + "min_p", + "typical_p", + "repetition_penalty", + "mirostat_target", + "mirostat_lr", + "logit_bias", + "echo", + "echo_last", + "ignore_eos", + "prompt_cache_key", + "prompt_cache_isolation_key", + "raw_output", + "perf_metrics_in_response", + "return_token_ids", + "safe_tokenization", + "service_tier", + "speculation", + "prediction", + "stream_options", + "sampling_mask", + "thinking", + "reasoning_history", + ] + missing = [p for p in required if p not in params] + assert missing == [], f"Missing params: {missing}" + + +def test_native_openai_params_flow_end_to_end_with_drop_params_false(): + """ + The OpenAI-native params Fireworks supports (seed, top_logprobs, logit_bias, + prompt_cache_key, service_tier, prediction) previously hit + ``UnsupportedParamsError`` with ``drop_params=False`` because they were absent + from ``get_supported_openai_params``. Listing them must let them survive the + ``get_optional_params`` gate and reach the request, not just appear in the + supported list. Asserting via ``get_optional_params`` (the real gate) rather + than ``map_openai_params`` catches a revert of the supported-params additions, + which a list-membership check would not. + """ + native = { + "seed": 42, + "top_logprobs": 3, + "logit_bias": {"1": 1}, + "prompt_cache_key": "cache-key", + "service_tier": "auto", + "prediction": {"type": "content", "content": "x"}, + } + optional_params = litellm.get_optional_params( + model="accounts/fireworks/models/llama-v3-70b-instruct", + custom_llm_provider="fireworks_ai", + drop_params=False, + **native, + ) + for key, value in native.items(): + assert optional_params.get(key) == value + + +def test_prompt_truncate_len_correct_name(): + config = FireworksAIConfig() + params = config.get_supported_openai_params(_REASONING_MODEL) + assert "prompt_truncate_len" in params + assert "prompt_truncate_length" not in params + + result = config.map_openai_params( + {"prompt_truncate_len": 4096}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert result == {"prompt_truncate_len": 4096} + + +def test_stream_options_include_usage_auto_injected(): + config = FireworksAIConfig() + result = config.transform_request( + model="accounts/fireworks/models/glm-5p1", + messages=[{"role": "user", "content": "Hi"}], + optional_params={"stream": True}, + litellm_params={}, + headers={}, + ) + assert result["stream_options"] == {"include_usage": True} + + +def test_stream_options_not_injected_when_not_streaming(): + config = FireworksAIConfig() + result = config.transform_request( + model="accounts/fireworks/models/glm-5p1", + messages=[{"role": "user", "content": "Hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "stream_options" not in result + + +def test_stream_options_preserves_user_override(): + config = FireworksAIConfig() + result = config.transform_request( + model="accounts/fireworks/models/glm-5p1", + messages=[{"role": "user", "content": "Hi"}], + optional_params={"stream": True, "stream_options": {"include_usage": False}}, + litellm_params={}, + headers={}, + ) + assert result["stream_options"]["include_usage"] is False + + +def test_reasoning_history_in_supported_params(): + config = FireworksAIConfig() + reasoning_params = config.get_supported_openai_params(_REASONING_MODEL) + assert "reasoning_history" in reasoning_params + + non_reasoning_params = config.get_supported_openai_params(_NON_REASONING_MODEL) + assert "reasoning_history" not in non_reasoning_params + + +def test_thinking_param_passthrough(): + config = FireworksAIConfig() + thinking = {"type": "disabled"} + result = config.map_openai_params( + {"thinking": thinking}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert result == {"thinking": thinking} + + +def test_thinking_and_reasoning_effort_conflict_rejected(): + config = FireworksAIConfig() + with pytest.raises( + litellm.BadRequestError, + match="does not support specifying both `thinking` and `reasoning_effort`", + ): + config.map_openai_params( + { + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "reasoning_effort": "medium", + }, + {}, + _REASONING_MODEL, + drop_params=False, + ) + + +def test_minimax_m3_supports_vision_from_model_map(): + config = FireworksAIConfig() + + for model in [ + "fireworks_ai/accounts/fireworks/models/minimax-m3", + "fireworks_ai/minimax-m3", + ]: + assert supports_vision(model=model, custom_llm_provider="fireworks_ai") is True + assert config.get_provider_info(model)["supports_vision"] is True + + +def test_transform_messages_helper_rejects_file_blocks(): + config = FireworksAIConfig() + messages = [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_data": "data:application/pdf;base64,JVBERi0xLjQKJSVFT0YK", + "filename": "tiny.pdf", + }, + }, + {"type": "text", "text": "Describe this"}, + ], + } + ] + + with pytest.raises( + litellm.BadRequestError, + match="Fireworks AI chat completions does not support file content blocks", + ): + config._transform_messages_helper( + messages, model="accounts/fireworks/models/kimi-k2p6", litellm_params={} + ) + + +def test_transform_messages_helper_rejects_non_vision_image_inputs(): + config = FireworksAIConfig() + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAE=" + }, + }, + ], + } + ] + + with pytest.raises(litellm.BadRequestError, match="does not support image inputs"): + config._transform_messages_helper( + messages, model="accounts/fireworks/models/glm-5p2", litellm_params={} + ) + + +def test_transform_messages_helper_allows_vision_image_inputs(): + config = FireworksAIConfig() + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAE=" + }, + }, + ], + } + ] + + out = config._transform_messages_helper( + messages, model="accounts/fireworks/models/minimax-m3", litellm_params={} + ) + assert out == messages + + +def test_image_inputs_not_rejected_for_fuzzy_non_vision_match(): + """ + A custom/fine-tuned model id that hyphen-matches a known non-vision model + (glm-5p2 has supports_vision=False) must not inherit that False via the + substring fallback and hard-reject valid image_url blocks. The capability + gate for rejection uses an exact cost-map match; the fuzzy fallback stays a + soft signal only, so an unmapped vision-capable deployment is not blocked. + """ + config = FireworksAIConfig() + custom_model = "accounts/myorg/models/custom-glm-5p2" + + assert config._get_model_cost_capability(custom_model, "supports_vision") is False + assert ( + config._get_model_cost_capability_exact(custom_model, "supports_vision") is None + ) + + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAE=" + }, + }, + ], + } + ] + out = config._transform_messages_helper( + messages, model=custom_model, litellm_params={} + ) + assert out == messages + + +def test_transform_messages_helper_skips_non_dict_content(): + config = FireworksAIConfig() + messages = [ + { + "role": "user", + "content": ["just a string", {"type": "text", "text": "hello"}], + } + ] + + out = config._transform_messages_helper( + messages, model="accounts/fireworks/models/glm-5p2", litellm_params={} + ) + assert out == messages + + +def test_transform_messages_helper_no_transform_inline(): + config = FireworksAIConfig() + url = "https://example.com/image.jpg" + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": url}], + } + ] + out = config._transform_messages_helper( + messages, model="accounts/fireworks/models/minimax-m3", litellm_params={} + ) + block = out[0]["content"][0] + assert block["image_url"] == url + assert "#transform=inline" not in block["image_url"] + + +def test_get_provider_info_vision_from_model_cost(monkeypatch): + config = FireworksAIConfig() + + vision_model = "fireworks_ai/test-vision-from-cost" + monkeypatch.setitem( + litellm.model_cost, + vision_model, + {"supports_vision": True, "supports_pdf_input": True}, + ) + info = config.get_provider_info(vision_model) + assert info["supports_vision"] is True + assert info["supports_pdf_input"] is True + + no_vision_model = "fireworks_ai/test-no-vision-from-cost" + monkeypatch.setitem(litellm.model_cost, no_vision_model, {}) + info_no_vision = config.get_provider_info(no_vision_model) + assert info_no_vision.get("supports_vision") is not True + assert "supports_pdf_input" not in info_no_vision + + +def test_reasoning_effort_boolean_true_to_medium(): + config = FireworksAIConfig() + result = config.map_openai_params( + {"reasoning_effort": True}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert result["reasoning_effort"] == "medium" + + +def test_reasoning_effort_boolean_false_to_none(): + config = FireworksAIConfig() + result = config.map_openai_params( + {"reasoning_effort": False}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert result["reasoning_effort"] == "none" + + +def test_reasoning_effort_string_passthrough(): + config = FireworksAIConfig() + result = config.map_openai_params( + {"reasoning_effort": "high"}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert result["reasoning_effort"] == "high" + + +def test_reasoning_effort_integer_passthrough(): + config = FireworksAIConfig() + result = config.map_openai_params( + {"reasoning_effort": 1000}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert result["reasoning_effort"] == 1000 + assert isinstance(result["reasoning_effort"], int) + + +def test_transform_response_captures_perf_metrics(): + body = { + **_BASE_CHAT_COMPLETION_RESPONSE, + "perf_metrics": {"prompt-tokens": 10}, + } + result = _run_transform_response(body) + assert result._hidden_params["fireworks_perf_metrics"] == {"prompt-tokens": 10} + + +def test_transform_response_captures_prompt_token_ids(): + body = { + **_BASE_CHAT_COMPLETION_RESPONSE, + "prompt_token_ids": [1, 2, 3], + } + result = _run_transform_response(body) + assert result._hidden_params["fireworks_prompt_token_ids"] == [1, 2, 3] + + +def test_transform_response_captures_raw_output(): + raw_output = { + "prompt_fragments": [], + "prompt_token_ids": [], + "completion": "test", + } + body = { + **_BASE_CHAT_COMPLETION_RESPONSE, + "choices": [ + { + **_BASE_CHAT_COMPLETION_RESPONSE["choices"][0], + "raw_output": raw_output, + } + ], + } + result = _run_transform_response(body) + assert result._hidden_params["fireworks_raw_outputs"] == [raw_output] + + +def test_transform_response_captures_token_ids(): + body = { + **_BASE_CHAT_COMPLETION_RESPONSE, + "choices": [ + { + **_BASE_CHAT_COMPLETION_RESPONSE["choices"][0], + "token_ids": [4, 5, 6], + } + ], + } + result = _run_transform_response(body) + assert result._hidden_params["fireworks_token_ids"] == [[4, 5, 6]] + + +def test_streaming_surfaces_fireworks_response_fields(): + """ + The Fireworks-specific response fields captured into _hidden_params for + non-streaming calls must also reach streamed responses. They ride the + streamed chunks' provider_specific_fields (litellm rebuilds each streamed + chunk, so per-chunk _hidden_params does not survive): per-choice + token_ids/raw_output on the content chunk, response-level + perf_metrics/prompt_token_ids on the final usage chunk. Driving the real + litellm.completion(stream=True) path also covers the get_model_response_iterator + wiring; dropping the Fireworks iterator would leave these fields unset. + """ + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + model = "accounts/fireworks/models/llama-v3p1-8b-instruct" + raw_output = {"completion": "Hi"} + sse_lines = [ + "data: " + + json.dumps( + { + "id": "stream-1", + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "Hi"}, + "token_ids": [123], + "raw_output": raw_output, + } + ], + } + ), + "data: " + + json.dumps( + { + "id": "stream-1", + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 5, + "completion_tokens": 1, + "total_tokens": 6, + }, + "perf_metrics": {"prompt-tokens": 5}, + "prompt_token_ids": [1, 2, 3], + } + ), + "data: [DONE]", + ] + + raw_response = MagicMock() + raw_response.status_code = 200 + raw_response.headers = {} + raw_response.iter_lines = lambda: iter(sse_lines) + + client = HTTPHandler() + with patch.object(client, "post", return_value=raw_response): + stream = litellm.completion( + model=f"fireworks_ai/{model}", + messages=[{"role": "user", "content": "hi"}], + stream=True, + api_key="fw-test-key", + client=client, + ) + surfaced: dict = {} + for chunk in stream: + fields = getattr(chunk, "provider_specific_fields", None) or {} + surfaced.update( + {k: v for k, v in fields.items() if k.startswith("fireworks_")} + ) + + assert surfaced["fireworks_token_ids"] == [[123]] + assert surfaced["fireworks_raw_outputs"] == [raw_output] + assert surfaced["fireworks_perf_metrics"] == {"prompt-tokens": 5} + assert surfaced["fireworks_prompt_token_ids"] == [1, 2, 3] diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 44e0b55ee3b..d94a86d8e55 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4292,7 +4292,7 @@ _FIREWORKS_MODELS = [ 6e-08, 512000, 512000, - False, + True, True, ), ( From accbd7e5878c0f5d3e64b0614ac90516c5579473 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 20 Jun 2026 20:37:22 -0700 Subject: [PATCH 03/50] feat: litellm plugin architecture v2 (#30688) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: plugin architecture — toggle between AI Gateway and external plugins Adds a generic plugin system so any external service can register with litellm and appear as a mode in the UI alongside the AI Gateway. Backend (litellm/proxy/plugin_routes.py — new): - GET /api/plugins: returns registered plugins from config; returns plugin_key only to authenticated requests - ANY /plugin-proxy/{name}/{path}: reverse proxies API calls to plugin Config: general_settings: plugins: - name: my-plugin display_name: My Plugin url: https://my-plugin.example.com plugin_key: sk-... # plugin auth key, passed to iframe UI: - PluginModeContext.tsx: fetches /api/plugins, persists mode to localStorage - leftnav.tsx: mode switcher dropdown at top of sidebar; plugin mode shows plugin-specific nav items - layout.tsx: renders iframe to plugin URL in plugin mode; passes plugin_key as ?token= for auto sign-in Plugin contract: expose GET /api/plugin-manifest returning { name, display_name, nav_items[], capabilities[] }. No litellm changes needed to add new plugins — config only. Reference implementation: LiteLLM-Labs/litellm-agent-control-plane * feat: add Plugins tab to Admin Settings UI Allows admins to add/edit/delete plugin registrations directly in the litellm UI under Admin Settings > Plugins, instead of editing config.yaml. Uses existing /config/field/update API to persist to general_settings.plugins. Each plugin entry has: name (identifier), display_name, url, plugin_key. * fix(ci): black, prettier, eslint, async-client violations - Black: format plugin_routes.py and proxy_server.py - Prettier: format PluginModeContext.tsx and PluginSettings.tsx - ESLint: replace raw fetch() with createApiClient in PluginModeContext - ESLint: use lazy useState initializer to read localStorage instead of calling setModeState inside useEffect (react-hooks/set-state-in-effect) - code-quality: replace httpx.AsyncClient per-request with get_async_httpx_client() shared client (avoids +500ms overhead) * fix(ci): schema.d.ts regen, Black proxy_server.py, ApiClientConfig fix - Regenerate schema.d.ts for new /api/plugins routes - Re-run Black 26.3.1 on proxy_server.py (matches CI version) - Fix PluginModeContext: createApiClient requires getBaseUrl field * fix: security hardening + CI fixes Security (Greptile 1/5 → addressing all 3 findings): - plugin_routes.py: add Depends(user_api_key_auth) to both /api/plugins and /plugin-proxy/{name}/{path} — was an unauthenticated open relay - plugin_routes.py: /api/plugins now returns plugin_key only to callers with a valid litellm token (enforced by user_api_key_auth), not just any header presence - layout.tsx: replace ?token= URL param with postMessage(targetOrigin) — token no longer exposed in browser history / logs / Referer headers CI: - backend/routes/allowlist.py: add /api/plugins and /plugin-proxy/ to fix test_gateway_plus_backend_covers_full_app - schema.d.ts: regenerated with enterprise routes included - Black + Prettier formatting * fix: regenerate schema.d.ts with enterprise routes included Install litellm-enterprise workspace member before gen:api so audit and other enterprise routes appear in the generated types, matching what CI produces with uv sync --extra proxy. * fix: exclude plugin routes from OpenAPI schema, restore upstream schema.d.ts Both /api/plugins and /plugin-proxy/ are internal infrastructure routes, not part of the public litellm API surface. Marking include_in_schema=False prevents Python-version-dependent schema diffs from breaking the schema sync check across different environments. * fix: schema.d.ts - passing schema base + exact plugin route types from openapi-typescript Use the CI-correct schema from a recently passing branch as base, then inject plugin route entries (paths + operations) generated by openapi-typescript from the plugin routes' OpenAPI spec. This avoids Python-version-dependent formatting differences that made local gen:api produce incorrect output. * fix: schema.d.ts - insert plugin ops at correct route registration position Plugin operations belong after delete_memory_v1_memory__key__delete (memory_router is included immediately before plugin_router in proxy_server.py), not after list_organization which is alphabetically but not registration-order. * fix: schema.d.ts - correct op positions from hunk analysis list_plugins_api_plugins_get goes after event_logging_batch op (hunk 1: line 33583). plugin_proxy ops go after create_policy_policies_post (hunk 2: line 44634). Previous location after delete_memory_v1_memory__key__delete was wrong. * fix: schema.d.ts - proxy ops go before create_policy (after otel_spans) * fix(security): restrict plugin_key to proxy_admin role only Veria finding: plugin_key was returned to any authenticated caller. Now only proxy_admin users receive plugin credentials in /api/plugins response — regular internal users see plugin name/url but not the key. * fix: update schema.d.ts docstring for list_plugins * fix: clear plugin registry on config reload (Greptile medium) register_plugins_from_config now replaces the registry instead of merging, so plugins removed from config are unreachable immediately without requiring a process restart. * fix(security): encrypted token exchange for plugin iframe — no raw litellm credential exposure The dashboard was sending the user's litellm bearer token to the plugin iframe via postMessage, allowing a compromised plugin to act as that user. Fix: - GET /api/plugins/auth-token: proxy encrypts caller token with Fernet keyed from LITELLM_SALT_KEY, returns ciphertext only - UI postMessages the ciphertext (not raw token) to the iframe - Plugin decrypts server-side with same LITELLM_SALT_KEY via POST /api/plugin-auth - Raw litellm credential never leaves the proxy in plaintext Additional hardening already in place: - /plugin-proxy/* strips Authorization header, injects plugin_key instead - plugin_key only returned to proxy_admin role via /api/plugins - Plugin registry cleared (not merged) on config reload Adds docs/plugin_architecture.md with plugin integration guide. * fix(code-quality): use get_async_httpx_client in plugin_proxy * fix: add /api/plugins/auth-token to schema.d.ts * fix: use apiClient for auth-token fetch, copy correct layout.tsx and PluginModeContext - Replace raw fetch() with createApiClient (fixes no-restricted-syntax ESLint rule) - Copy correct layout.tsx with encrypted token + postMessage approach - Copy correct PluginModeContext.tsx with accessToken prop injection - Update schema.d.ts with auth-token path and operation entries * fix: add plugin_auth_token operation to schema.d.ts * fix(security): strip cookie/set-cookie + fix compressed response headers Veria High: cookie header was forwarded to plugin backends allowing capture of litellm JWT session cookies. Strip cookie on requests. Strip set-cookie from responses so plugins cannot overwrite litellm session cookies. Greptile P1: httpx decompresses responses but resp.headers still contained Content-Encoding/Transfer-Encoding/Content-Length from the wire. Forwarding these caused double-decompression and length errors. Now filtered via _RESPONSE_STRIP before returning to the browser. * fix: update plugin_key help text — no more ?token= reference * fix(security): disable follow_redirects to prevent SSRF follow_redirects=True allowed a plugin backend to return a 3xx to an internal URL, causing the proxy to fetch that internal service and relay the response. Disabled: clients handle their own redirects. * fix: forward user identity headers to plugin to address confused deputy Plugins receive X-LiteLLM-User-Id and X-LiteLLM-User-Role so they can enforce their own per-user access control before acting on requests that arrive with the shared plugin_key credential. * fix(security): restrict /plugin-proxy/* to proxy_admin role Closes the confused deputy gap: regular users could invoke any plugin endpoint using the shared plugin_key as a bearer credential. Now only proxy_admin callers can use the plugin proxy route. Plugin UIs communicate with the plugin service directly via the iframe (using the encrypted token exchange); this proxy route is for administrative/server-to-server access only. * fix: update schema.d.ts for admin-only proxy route docstring * fix(bug): use PassThroughEndpoint instead of None for get_async_httpx_client get_async_httpx_client(llm_provider=None) raises TypeError — the function concatenates the provider string and None is not a str. Use httpxSpecialProvider.PassThroughEndpoint, the enum value used by other internal proxy pass-through routes. * fix(security): add 30s TTL to encrypted plugin auth tokens Veria medium: encrypted tokens had no expiry, allowing indefinite replay. Fernet embeds a timestamp; decrypt_token now passes ttl=30 so tokens older than 30 seconds are rejected even with a valid HMAC. Plugin's /api/plugin-auth must call litellm within 30s of the iframe receiving the postMessage — normal browser behavior, tight enough to close the replay window. * feat(ui): topnav plugin switcher, embed plugins at their root Builds on the plugin architecture already on this branch (encrypted-token postMessage handshake, /api/plugins, PluginSettings) and removes the parts of the embed that assumed a specific plugin's shape. The mode switcher moves out of the sidebar into the topnav and lists AI Gateway plus each registered plugin by its display_name. Selecting a plugin hides litellm's sidebar entirely and renders the plugin full-bleed at its root url; the plugin draws its own navigation inside the iframe. This drops the hardcoded "Agent Control Plane" label and the hardcoded Sessions/Agents/Routines/... nav groups (agentControlPlaneMenuGroups / acpPagePaths) that only matched the agent platform and 404'd for a plugin that serves only / (e.g. the chat UI). The encrypted-token postMessage flow is unchanged. Note: embedding at root means a plugin must route internally from /; plugins that previously relied on the /sessions entrypoint should redirect from their root. * fix(security): audience-scoped identity claim replaces litellm token Veria: shared LITELLM_SALT_KEY with plugins + encrypting user bearer token created delegation/impersonation risk. Architecture change: - /api/plugins/auth-token now issues a plugin-scoped identity CLAIM {user_id, user_role, plugin, exp} encrypted with HMAC(LITELLM_SALT_KEY, plugin_name) - Each plugin holds only its own HMAC-derived key; cannot forge claims for other plugins or recover LITELLM_SALT_KEY - Claim contains NO litellm bearer token — compromised plugin learns caller identity only, cannot act as that user against the proxy - 30s TTL enforced in both Fernet header and explicit exp field - LAP /api/plugin-auth verifies claim, returns its own master key to browser (LAP key never exposed without valid claim) * fix(plugins): allow registering plugins from the admin UI Adding a plugin in the UI POSTs general_settings.plugins to /config/field/update, which rejected it with "Invalid field=plugins passed in." because `plugins` was not a field on ConfigGeneralSettings. Add a typed PluginConfig model and a `plugins` field so the update validates and persists. The in-memory plugin registry only refreshed at startup, so a plugin added via the UI did not appear in /api/plugins (the view switcher) until a restart. Refresh the registry from the new general_settings whenever the plugins field is updated. While here, type the registry as dict[str, PluginConfig] instead of raw dicts so list_plugins and plugin_proxy access typed attributes. Fix the Plugin Key field copy: it is optional and only used to authenticate litellm's server-side reverse proxy to a plugin's own backend (/plugin-proxy//*). It is not involved in iframe auth, which forwards the user's litellm token. Plugins that use the forwarded token leave it blank. * fix: regenerate schema.d.ts with PluginConfig type and updated auth-token endpoint * fix: use CI-compatible schema base for plugin entries * fix(plugins): load DB-persisted plugins on startup Plugins added through the admin UI are saved to DB general_settings, but the registry only initialised from the YAML config at boot, so UI-added plugins disappeared from the view switcher after a restart (the Plugins table still listed them since it reads the DB directly). Refresh the registry from the DB general_settings when it is merged in at startup. * fix: add PluginConfig schema, plugins field, fix list_plugins return type * fix: correct PluginConfig and plugins field positions in schema * fix: correct plugins field position in schema (after pass_through_endpoints) * fix: update PluginConfig.plugin_key description to match _types.py source * fix: move plugins field after pass_through_request_timeout (correct alphabetical position) * fix: redact plugin_key in config/field/info response Veria medium: proxy_admin_viewer could read plugin_key via GET /config/field/info?field_name=plugins. Now plugin_key is replaced with *** in the response regardless of caller role. The credential is only usable server-side. * fix(security): correct plugin docs salt-key guidance, drop iframe clipboard-read Address the two open Veria findings on the plugin architecture. The plugin docs told external services to decrypt the iframe auth payload with the proxy's LITELLM_SALT_KEY directly. That is both insecure and wrong: the running code derives a per-plugin key as HMAC-SHA256(LITELLM_SALT_KEY, plugin_name) and ships only a short-lived identity claim with no litellm bearer token. Sharing the master salt would let a compromised plugin decrypt any litellm secret recovered from a dump or backup. Rewrite the doc to match the implementation: the proxy computes the per-plugin key once and provisions it as a dedicated secret, the plugin validates the claim's audience and 30s TTL, and LITELLM_SALT_KEY never leaves the proxy. Also refresh the now-stale module and UI comments that still described the old shared-key token flow. Drop clipboard-read from the plugin iframe's allow attribute so an untrusted plugin can no longer read the user's clipboard; clipboard-write is retained. * fix(ci): modernize PluginConfig typing, refresh budget baselines via merge * fix(plugins): close iframe auth race and empty-plugins mode fallback Address the two open Greptile behavioral findings. The iframe auth handshake only posted the encrypted claim on the iframe's `load` event. When the auth-token fetch resolved after the iframe had already loaded, that listener never fired again and the plugin never received the claim. Send the claim immediately as well as on subsequent loads so both orderings are covered. The plugin mode fallback guarded on a non-empty plugins list, so removing all plugins left a user stranded on a stale mode with a blank iframe instead of returning to the AI Gateway. Track a loaded flag and fall back to ai-gateway once plugins have loaded whenever the stored mode is no longer registered, including the empty-list case. Add a PluginModeContext regression test covering the empty-list fallback and the still-registered path. * chore: re-trigger CI (GH Actions missed the prior head; re-run flaky live-API suites) * fix(plugins): scope iframe auth claim to the active plugin The iframe auth-token fetch omitted plugin_name, so the proxy always issued a claim encrypted under the default plugin's per-plugin key. For any other active plugin the iframe received a claim it could not decrypt and sign-in silently broke, and because the cached claim was posted to whichever plugin was mounted, a compromised iframe could replay the default plugin's claim. The active plugin's name was also missing from the fetch effect's dependencies, so switching plugins never refreshed the claim. Request the claim with the active plugin's name, re-fetch when the active plugin changes, and only deliver a claim while it still matches the mounted plugin so one plugin's claim is never replayed to another. * fix(plugins): never overwrite a stored plugin_key with its redaction placeholder /config/field/info redacts every plugin_key to "***", so an admin editing a plugin in the settings UI posted that placeholder straight back and the update handler persisted "***" as the real credential, permanently destroying the key. Preserve the stored credential on update: a blank or redacted plugin_key now sources the existing key from the saved config, only a real value replaces it, and a placeholder with no stored key is dropped rather than written. The edit modal also starts the key field blank so an untouched save keeps the current key, with the field labelled accordingly. * fix(security): sandbox proxied plugin responses on the dashboard origin The /plugin-proxy reverse proxy returned the plugin's body and content-type on the litellm dashboard origin, so a compromised plugin could serve an HTML/JS document that a proxy_admin navigates to and have it execute with the admin's session against same-origin management APIs. Force every proxied response inert: set Content-Security-Policy: sandbox (opaque origin, scripts disabled) and X-Content-Type-Options: nosniff, applied after the plugin's own headers so they cannot be overridden. The header construction moves to a pure helper with a unit test covering the sandbox enforcement and the existing wire/cookie header stripping. * fix(plugins): recover to ai-gateway when the plugins fetch fails The loaded flag was only set on a successful /api/plugins response, so when the fetch failed a user with a plugin mode stored in localStorage stayed on the blank plugin placeholder with no switcher to escape. Mark loaded in a finally so the stored mode still falls back to ai-gateway on failure, and add a regression test for the failed-fetch path. * fix(security): never return plugin_key from /api/plugins The plugin list endpoint returned the plaintext plugin_key to proxy_admin callers, and the dashboard fetches /api/plugins on every load into React state, so the credential was exposed to DevTools, memory snapshots, and any same-origin script. The browser never uses the key; the proxy injects it server-side from the registry and admin key management runs through the redacted /config/field/info path. Drop plugin_key from the response for every caller and update the regression test to assert it is never returned. * chore(ui): regenerate schema.d.ts for updated list_plugins docstring * fix(security): strip every litellm auth header before forwarding to plugins The plugin reverse proxy only removed Authorization and x-api-key, but user_api_key_auth also authenticates a caller via API-Key, x-goog-api-key, Ocp-Apim-Subscription-Key, x-litellm-api-key, and any configured custom key header. A malicious plugin could lure a proxy_admin into calling /plugin-proxy/... with the litellm key in one of those headers; the request authenticated locally and then forwarded the same key to the plugin, letting it impersonate the admin. Add a canonical SpecialHeaders.litellm_credential_header_names() that the auth header enum is the single source for, and strip that whole set plus the live general_settings.litellm_key_header_name from every forwarded request. New auth headers added to SpecialHeaders are now stripped automatically. Regression tests cover each credential header, the custom configured header, and the canonical list's contents. --- backend/routes/allowlist.py | 3 + docs/plugin_architecture.md | 141 ++++++ litellm/proxy/_types.py | 39 ++ litellm/proxy/plugin_routes.py | 344 +++++++++++++ litellm/proxy/proxy_server.py | 73 ++- .../test_litellm/proxy/test_plugin_routes.py | 232 +++++++++ tests/test_litellm/proxy/test_proxy_server.py | 36 ++ .../agentControlPlaneView.test.tsx | 83 ++++ .../src/app/(dashboard)/layout.tsx | 105 +++- .../src/components/AdminPanel.tsx | 6 + .../components/Navbar/ViewSwitcher.test.tsx | 76 +++ .../src/components/Navbar/ViewSwitcher.tsx | 46 ++ .../PluginSettings/PluginSettings.tsx | 173 +++++++ .../src/components/navbar.tsx | 7 + .../src/contexts/PluginModeContext.test.tsx | 64 +++ .../src/contexts/PluginModeContext.tsx | 95 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 464 ++++++++++++++++++ 17 files changed, 1977 insertions(+), 10 deletions(-) create mode 100644 docs/plugin_architecture.md create mode 100644 litellm/proxy/plugin_routes.py create mode 100644 tests/test_litellm/proxy/test_plugin_routes.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agentControlPlaneView.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/PluginSettings/PluginSettings.tsx create mode 100644 ui/litellm-dashboard/src/contexts/PluginModeContext.test.tsx create mode 100644 ui/litellm-dashboard/src/contexts/PluginModeContext.tsx diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index d1a576aeb33..2f65f99c292 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -120,6 +120,9 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = ( "/robots.txt", # Health (k8s probes) "/health", + # Plugin system + "/api/plugins", + "/plugin-proxy/", ) BACKEND_EXACT_PATHS: frozenset[str] = frozenset( diff --git a/docs/plugin_architecture.md b/docs/plugin_architecture.md new file mode 100644 index 00000000000..8801761531d --- /dev/null +++ b/docs/plugin_architecture.md @@ -0,0 +1,141 @@ +# LiteLLM Plugin Architecture + +Plugins let external services appear as selectable modes in the litellm UI sidebar alongside the AI Gateway. + +--- + +## Quick start + +### 1. Configure the plugin + +Add a `plugins` block to your litellm `config.yaml`: + +```yaml +general_settings: + master_key: sk-... + plugins: + - name: my-plugin # unique identifier (no spaces) + display_name: My Plugin # shown in the UI dropdown + url: "https://my-plugin.example.com" + plugin_key: "sk-..." # plugin's own auth credential +``` + +`plugin_key` is injected as `Authorization: Bearer ` on every +request proxied through `/plugin-proxy/my-plugin/*`. The caller's litellm +credential is stripped before forwarding so the plugin never receives a live +litellm API key. + +### 2. Implement two endpoints on your service + +| Endpoint | Method | Purpose | +|---|---|---| +| `GET /api/plugin-manifest` | public | Returns plugin metadata for the UI | +| `POST /api/plugin-auth` | public | Decrypts the identity claim for seamless sign-in | + +#### `GET /api/plugin-manifest` + +```json +{ + "name": "my-plugin", + "display_name": "My Plugin", + "version": "1.0.0", + "nav_items": [ + { "key": "home", "label": "Home", "icon": "HomeOutlined", "path": "/" }, + { "key": "reports", "label": "Reports", "icon": "BarChartOutlined", "path": "/reports" } + ], + "capabilities": ["reports", "data"] +} +``` + +#### `POST /api/plugin-auth` + +Receives `{ "session_claim": "" }`. + +The proxy never shares `LITELLM_SALT_KEY` with your plugin. Each plugin is +provisioned with its own dedicated key, derived as +`HMAC-SHA256(LITELLM_SALT_KEY, plugin_name)`. Compute it once on the proxy +host and hand the result to your plugin as a secret (e.g. `PLUGIN_AUTH_KEY`): + +```bash +python -c 'import base64,hmac,hashlib,os; \ +print(base64.urlsafe_b64encode(hmac.new(os.environ["LITELLM_SALT_KEY"].encode(), b"my-plugin", hashlib.sha256).digest()).decode())' +``` + +A compromised plugin holding only this scoped key cannot recover +`LITELLM_SALT_KEY` or decrypt any other litellm secret. + +Decrypt and validate the claim with that key: + +```python +import json, os, time +from cryptography.fernet import Fernet + +_CLAIM_TTL_SECONDS = 30 + +def plugin_auth(session_claim: str) -> dict: + cipher = Fernet(os.environ["PLUGIN_AUTH_KEY"].encode()) + claim = json.loads(cipher.decrypt(session_claim.encode(), ttl=_CLAIM_TTL_SECONDS)) + if claim.get("plugin") != "my-plugin": + raise ValueError("claim audience mismatch") + if int(claim.get("exp", 0)) < int(time.time()): + raise ValueError("claim expired") + return claim +``` + +The claim is `{ "plugin", "user_id", "user_role", "exp" }`; it carries no +litellm bearer token. Establish the plugin's own session from `user_id` / +`user_role` and authenticate API calls back to litellm through the +`/plugin-proxy/my-plugin/*` reverse proxy, which injects `plugin_key` for you. + +--- + +## How iframe auth works + +``` +litellm UI + ├─ GET /api/plugins/auth-token -> { session_claim } + └─ postMessage({ type:"litellm-auth", session_claim }, pluginOrigin) + │ + ▼ +Plugin iframe browser + └─ POST /api/plugin-auth { session_claim } + │ + ▼ +Plugin server + ├─ decrypt(session_claim, PLUGIN_AUTH_KEY) -> { user_id, user_role, exp } + └─ establish plugin session -> stored in sessionStorage +``` + +No litellm bearer token ever leaves the proxy; the claim only conveys the +caller's identity and expires after 30 seconds. A postMessage intercept +yields ciphertext that is useless without the plugin's scoped key. + +--- + +## Proxy routes + +- `GET /api/plugins` — list registered plugins (`name`, `display_name`, `url`). `plugin_key` is **never** returned; it stays server-side. Requires an authenticated caller. +- `GET /api/plugins/auth-token?plugin_name=` — short-lived encrypted identity claim for the named plugin. Requires `LITELLM_SALT_KEY` to be set (503 otherwise) and the plugin to be registered (404 otherwise). +- `ANY /plugin-proxy/{name}/{path}` — authenticated reverse proxy to the plugin backend. Restricted to `proxy_admin`. + +--- + +## Reverse proxy behaviour + +When an admin (or server-to-server caller) hits `/plugin-proxy//`, the proxy authenticates the caller locally, then rewrites the request before forwarding it to the plugin's `url`: + +- **Every litellm credential header is stripped** — `Authorization`, `x-api-key`, `API-Key`, `x-goog-api-key`, `Ocp-Apim-Subscription-Key`, `x-litellm-api-key`, any configured `litellm_key_header_name`, plus `Cookie`. The plugin can never be handed the caller's live litellm key. +- **`plugin_key` is injected** as `Authorization: Bearer ` — the only credential the plugin receives. +- **Caller identity is forwarded** as `x-litellm-user-id` and `x-litellm-user-role` so the plugin can run its own authorization. These are informational, not credentials. +- **Responses are sandboxed** — `Content-Security-Policy: sandbox` and `X-Content-Type-Options: nosniff` are set so plugin-controlled bytes served from the litellm origin cannot execute against the dashboard. + +--- + +## Security checklist + +- [ ] `LITELLM_SALT_KEY` is set on the proxy and never shared with the plugin +- [ ] The plugin holds only its derived `HMAC(LITELLM_SALT_KEY, plugin_name)` key, provisioned as a dedicated secret +- [ ] `plugin_key` is a dedicated credential scoped to the plugin (not your litellm master key) +- [ ] Plugin's `POST /api/plugin-auth` enforces the claim's `plugin` audience and `exp` (30s TTL) +- [ ] Plugin treats `x-litellm-user-id` / `x-litellm-user-role` as identity hints, not as proof of authentication +- [ ] Plugin service URL uses HTTPS in production diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8e2ec423cde..e856e5e3cdb 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2142,6 +2142,20 @@ class UserHeaderMapping(LiteLLMPydanticObjectBase): UserMCPManagementMode = Literal["restricted", "view_all"] +class PluginConfig(LiteLLMPydanticObjectBase): + """A single external service registered as an embeddable UI plugin.""" + + name: str = Field(description="unique plugin identifier (kebab-case)") + display_name: str | None = Field( + None, description="human-readable label shown in the UI view switcher" + ) + url: str = Field(description="base URL of the plugin service") + plugin_key: str | None = Field( + None, + description="plugin's own credential, injected as Bearer auth only on /plugin-proxy//* reverse-proxy calls", + ) + + class ConfigGeneralSettings(LiteLLMPydanticObjectBase): """ Documents all the fields supported by `general_settings` in config.yaml @@ -2150,6 +2164,9 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): completion_model: Optional[str] = Field( None, description="proxy level default model for all chat completion calls" ) + plugins: list[PluginConfig] | None = Field( + None, description="external services registered as embeddable UI plugins" + ) key_management_system: Optional[KeyManagementSystem] = Field( None, description="key manager to load keys from / decrypt keys with" ) @@ -3808,6 +3825,28 @@ class SpecialHeaders(enum.Enum): mcp_servers = "x-mcp-servers" mcp_access_groups = "x-mcp-access-groups" + @classmethod + def litellm_credential_header_names(cls) -> "frozenset[str]": + """Lowercased header names user_api_key_auth accepts as a litellm key. + + Every header here authenticates the caller, so any code that forwards a + request onward (e.g. the plugin reverse proxy) must strip all of them to + avoid leaking the caller's litellm credential downstream. The static + custom-key header (general_settings.litellm_key_header_name) is runtime + config and must be added on top of this set by the caller. + """ + return frozenset( + header.value.lower() + for header in ( + cls.openai_authorization, + cls.azure_authorization, + cls.anthropic_authorization, + cls.google_ai_studio_authorization, + cls.azure_apim_authorization, + cls.custom_litellm_api_key, + ) + ) + class LitellmDataForBackendLLMCall(TypedDict, total=False): headers: dict diff --git a/litellm/proxy/plugin_routes.py b/litellm/proxy/plugin_routes.py new file mode 100644 index 00000000000..6a94f78fbe7 --- /dev/null +++ b/litellm/proxy/plugin_routes.py @@ -0,0 +1,344 @@ +""" +Plugin proxy routes for litellm. + +Enables external services to register as plugins and be proxied through +the litellm proxy server. + +Config (in litellm config.yaml general_settings): + plugins: + - name: my-plugin + url: "http://localhost:3210" + display_name: "My Plugin" + plugin_key: "sk-..." # optional: plugin's own auth key + +Plugin iframe auth: + The UI calls GET /api/plugins/auth-token to receive a short-lived identity + claim ({user_id, user_role, plugin, exp}) encrypted with a per-plugin key + derived as HMAC-SHA256(LITELLM_SALT_KEY, plugin_name). The claim carries no + litellm bearer token, so a compromised plugin learns only the caller's + identity, never their credential. LITELLM_SALT_KEY itself is never shared + with plugins — each plugin holds only its own derived key. +""" + +import base64 +import hashlib +import hmac as _hmac +import json +import os +import time +from collections.abc import Mapping + +from cryptography.fernet import Fernet, InvalidToken +from fastapi import APIRouter, Depends, HTTPException, Request, Response + +from litellm.proxy._types import PluginConfig, SpecialHeaders, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.llms.custom_http import httpxSpecialProvider + +router = APIRouter() + +# Hop-by-hop headers (RFC 7230) and the litellm session cookie — never forwarded +# to a plugin backend. Credential headers are added on top per-request from the +# canonical SpecialHeaders set so the plugin only ever authenticates via its own +# injected plugin_key. +_HOP_BY_HOP_STRIP = frozenset( + { + "host", + "connection", + "transfer-encoding", + "te", + "trailers", + "upgrade", + "cookie", + } +) + + +def _configured_key_header_names() -> frozenset[str]: + """The lowercased general_settings.litellm_key_header_name, if configured. + + Read live from the proxy module (not import-time) so a custom key header set + via config is honoured without a restart. Returns empty when unset. + """ + try: + from litellm.proxy import proxy_server + except Exception: + return frozenset() + general_settings = getattr(proxy_server, "general_settings", None) + if not isinstance(general_settings, dict): + return frozenset() + name: object = general_settings.get("litellm_key_header_name") + return frozenset({name.lower()}) if isinstance(name, str) and name else frozenset() + + +def _request_strip_headers() -> frozenset[str]: + """Headers to drop before forwarding a request to a plugin backend. + + Every header user_api_key_auth accepts as a litellm credential is stripped — + Authorization, x-api-key, API-Key, x-goog-api-key, Ocp-Apim-Subscription-Key, + x-litellm-api-key, and any configured custom key header — so a plugin can + never be handed the caller's live litellm key (confused-deputy escalation). + """ + return ( + _HOP_BY_HOP_STRIP + | SpecialHeaders.litellm_credential_header_names() + | _configured_key_header_names() + ) + + +# Headers to strip from plugin RESPONSES before returning to the browser. +# httpx already decompresses and de-chunks the body, so forwarding the wire +# encoding headers causes clients to attempt double-decompression (garbage) or +# incorrect length checks. set-cookie is removed so plugins cannot overwrite +# litellm session cookies. +_RESPONSE_STRIP = { + "content-encoding", + "transfer-encoding", + "content-length", + "set-cookie", +} + + +def _safe_response_headers(raw: "Mapping[str, str]") -> dict[str, str]: + """Strip wire-encoding/cookie headers and force proxied responses inert. + + Plugin-controlled bytes are served from the litellm dashboard origin, so a + compromised plugin could return an HTML/JS document that executes with the + admin's session against same-origin management APIs. A sandbox CSP forces + the response into an opaque origin with scripts disabled, and nosniff stops + content-type confusion from re-enabling execution. Both are set last so a + plugin cannot override them with its own headers. + """ + return { + **{k: v for k, v in raw.items() if k.lower() not in _RESPONSE_STRIP}, + "content-security-policy": "sandbox", + "x-content-type-options": "nosniff", + } + + +# In-memory plugin registry — populated from general_settings at startup +_plugin_registry: dict[str, PluginConfig] = {} + + +# --------------------------------------------------------------------------- +# Key derivation — audience-scoped per plugin so compromising one plugin +# cannot be used to forge claims for another. LITELLM_SALT_KEY is NEVER +# shared with plugins; each plugin only receives a key derived from +# HMAC(LITELLM_SALT_KEY, plugin_name) which reveals nothing about the master. +# --------------------------------------------------------------------------- +def _plugin_fernet(plugin_name: str) -> Fernet: + """Return a Fernet cipher whose key is scoped to a specific plugin. + + Key material: HMAC-SHA256(LITELLM_SALT_KEY, plugin_name). + A plugin possessing its own key cannot derive the master salt or + forge claims intended for a different plugin. + """ + salt = os.getenv("LITELLM_SALT_KEY", "").encode() + derived = _hmac.new(salt, plugin_name.encode(), hashlib.sha256).digest() + return Fernet(base64.urlsafe_b64encode(derived)) + + +_CLAIM_TTL_SECONDS = 30 # identity claims expire after 30 s + + +def issue_plugin_session_claim( + plugin_name: str, user_id: str | None, user_role: str | None +) -> str: + """Issue a short-lived, audience-scoped identity claim for the plugin. + + The claim contains {user_id, user_role, plugin, exp}. Crucially it + contains NO litellm bearer token — the plugin can only derive the + caller's identity, not act as them against the proxy. + """ + claim = { + "plugin": plugin_name, + "user_id": user_id or "", + "user_role": user_role or "", + "exp": int(time.time()) + _CLAIM_TTL_SECONDS, + } + return _plugin_fernet(plugin_name).encrypt(json.dumps(claim).encode()).decode() + + +def verify_plugin_session_claim(plugin_name: str, ciphertext: str) -> dict: + """Verify and decode a plugin session claim. + + Raises ValueError if the HMAC is invalid, the audience is wrong, or + the claim is expired. Returns the decoded claim dict on success. + """ + try: + raw = _plugin_fernet(plugin_name).decrypt( + ciphertext.encode(), ttl=_CLAIM_TTL_SECONDS + ) + claim = json.loads(raw) + except (InvalidToken, Exception) as exc: + raise ValueError("Invalid, tampered, or expired plugin session claim") from exc + + if claim.get("plugin") != plugin_name: + raise ValueError("Plugin claim audience mismatch") + if int(claim.get("exp", 0)) < int(time.time()): + raise ValueError("Plugin session claim expired") + return claim + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- +def register_plugins_from_config(general_settings: dict[str, object]) -> None: + """Replace the plugin registry from general_settings. + + Replaces (not merges) so plugins removed from config are immediately + unreachable without requiring a process restart. + """ + raw = general_settings.get("plugins") + entries: list[object] = raw if isinstance(raw, list) else [] + new_registry = { + p.name: p for p in (PluginConfig.model_validate(entry) for entry in entries) + } + _plugin_registry.clear() + _plugin_registry.update(new_registry) + + +# --------------------------------------------------------------------------- +# Routes +# --------------------------------------------------------------------------- +@router.get("/api/plugins", tags=["plugins"]) +async def list_plugins( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> list[dict[str, str]]: + """Return registered plugins for authenticated UI callers. + + plugin_key is never returned — the browser never needs it (the proxy injects + it server-side from the registry), and exposing it here would leak the + credential into React state and DevTools. Admin key management goes through + the redacted /config/field/info path instead. + """ + return [ + { + "name": plugin.name, + "display_name": plugin.display_name or plugin.name, + "url": plugin.url, + } + for plugin in _plugin_registry.values() + ] + + +@router.get("/api/plugins/auth-token", tags=["plugins"]) +async def plugin_auth_token( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + plugin_name: str = "litellm-platform-plugin", +) -> dict: + """Issue a short-lived, audience-scoped plugin session claim. + + The claim contains {user_id, user_role, plugin, exp}. It does NOT + contain the caller's litellm bearer token — a compromised plugin can + only learn the caller's identity, not impersonate them against the proxy. + + Encrypted with a key derived from HMAC(LITELLM_SALT_KEY, plugin_name), + so each plugin holds only its own key and cannot forge claims for others. + + Requires LITELLM_SALT_KEY to be set; returns 503 otherwise. + """ + if not os.getenv("LITELLM_SALT_KEY"): + raise HTTPException( + status_code=503, + detail="LITELLM_SALT_KEY is not configured; plugin iframe auth unavailable.", + ) + if plugin_name not in _plugin_registry: + raise HTTPException( + status_code=404, detail=f"Plugin '{plugin_name}' is not registered." + ) + user_id = getattr(user_api_key_dict, "user_id", None) + user_role = getattr(user_api_key_dict, "user_role", None) + return { + "session_claim": issue_plugin_session_claim(plugin_name, user_id, user_role) + } + + +@router.api_route( + "/plugin-proxy/{plugin_name}/{path:path}", + methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], + tags=["plugins"], + include_in_schema=False, +) +async def plugin_proxy( + plugin_name: str, + path: str, + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> Response: + """Authenticated reverse-proxy to a registered plugin backend. + + Restricted to proxy_admin callers — the shared plugin_key must not be + usable as a confused-deputy credential by regular users. Plugin UIs + talk to the plugin service directly via the iframe; this route is for + administrative and server-to-server access only. + + The caller's litellm credential is stripped and replaced with the + plugin's own plugin_key so plugins never receive a live litellm API key. + """ + if getattr(user_api_key_dict, "user_role", None) != "proxy_admin": + return Response( + content="Plugin proxy access requires proxy_admin role.", + status_code=403, + ) + + plugin = _plugin_registry.get(plugin_name) + if not plugin: + return Response( + content=f"Plugin '{plugin_name}' not registered", + status_code=404, + ) + + target_url = f"{plugin.url.rstrip('/')}/{path}" + query = request.url.query + if query: + target_url = f"{target_url}?{query}" + + body = await request.body() + + # Strip caller credentials and hop-by-hop headers from forwarded request + strip = _request_strip_headers() + forward_headers = { + k: v for k, v in request.headers.items() if k.lower() not in strip + } + + # Inject plugin's own credential as upstream auth (if configured) + plugin_key = plugin.plugin_key + if plugin_key: + forward_headers["authorization"] = f"Bearer {plugin_key}" + + # Forward caller identity so the plugin can enforce its own access control. + # The plugin MUST NOT trust these as credentials — they are informational. + # The plugin_key above is the only authentication mechanism. + user_id = getattr(user_api_key_dict, "user_id", None) + user_role = getattr(user_api_key_dict, "user_role", None) + if user_id: + forward_headers["x-litellm-user-id"] = str(user_id) + if user_role: + forward_headers["x-litellm-user-role"] = str(user_role) + + handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.PassThroughEndpoint + ) + try: + req = handler.client.build_request( + method=request.method, + url=target_url, + headers=forward_headers, + content=body, + ) + # Do not follow redirects — a redirect to an internal URL would allow + # the plugin to SSRF the proxy into fetching arbitrary internal services. + resp = await handler.client.send(req, follow_redirects=False) + except Exception: + return Response( + content=f"Cannot connect to plugin '{plugin_name}' at {plugin.url}", + status_code=502, + ) + + return Response( + content=resp.content, + status_code=resp.status_code, + headers=_safe_response_headers(resp.headers), + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c138626a272..1bc23502165 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -419,6 +419,10 @@ from litellm.proxy.management_endpoints.workflow_management_endpoints import ( ) from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update from litellm.proxy.memory.memory_endpoints import router as memory_router +from litellm.proxy.plugin_routes import ( + router as plugin_router, + register_plugins_from_config, +) from litellm.proxy.middleware.in_flight_requests_middleware import ( InFlightRequestsMiddleware, ) @@ -4502,6 +4506,8 @@ class ProxyConfig: load_from_azure_key_vault(use_azure_key_vault=use_azure_key_vault) ### ALERTING ### self._load_alerting_settings(general_settings=general_settings) + ### PLUGINS ### + register_plugins_from_config(general_settings) ### CONNECT TO DATABASE ### database_url = general_settings.get("database_url", None) if database_url and database_url.startswith("os.environ/"): @@ -5663,6 +5669,10 @@ class ProxyConfig: llm_router=llm_router, ) + if _general_settings is not None and "plugins" in _general_settings: + general_settings["plugins"] = _general_settings["plugins"] + register_plugins_from_config(general_settings) + async def _reschedule_spend_log_cleanup_job(self): """ Reschedule the spend log cleanup job based on current general_settings. @@ -14931,6 +14941,41 @@ async def update_config( Keep it more precise, to prevent overwrite other values unintentially """ +_PLUGIN_KEY_REDACTED = "***" + + +def _preserve_redacted_plugin_keys(incoming: object, existing: object) -> object: + """Restore real plugin_key values the client never sees. + + /config/field/info redacts every plugin_key to ``"***"``, so an admin + editing a plugin posts that placeholder (or a blank, when the UI clears the + field) straight back. Treat a blank or redacted plugin_key as "keep the + stored credential" by sourcing it from the existing config; only a real, + non-redacted value replaces it, and a blank with no stored key drops the + field entirely instead of persisting the placeholder. + """ + if not isinstance(incoming, list): + return incoming + + stored_keys = { + p["name"]: p["plugin_key"] + for p in (existing if isinstance(existing, list) else []) + if isinstance(p, dict) and p.get("name") and p.get("plugin_key") + } + + def resolve(plugin: object) -> object: + if not isinstance(plugin, dict): + return plugin + key = plugin.get("plugin_key") + if key not in (None, "", _PLUGIN_KEY_REDACTED): + return plugin + name = plugin.get("name") + if name in stored_keys: + return {**plugin, "plugin_key": stored_keys[name]} + return {k: v for k, v in plugin.items() if k != "plugin_key"} + + return [resolve(p) for p in incoming] + @router.post( "/config/field/update", @@ -14997,7 +15042,13 @@ async def update_config_general_settings( ## update db - general_settings[data.field_name] = data.field_value + field_value = data.field_value + if data.field_name == "plugins": + field_value = _preserve_redacted_plugin_keys( + field_value, general_settings.get("plugins") + ) + + general_settings[data.field_name] = field_value response = await ConfigRepository(prisma_client).table.upsert( where={"param_name": "general_settings"}, @@ -15008,6 +15059,9 @@ async def update_config_general_settings( ) await invalidate_config_param("general_settings") + if data.field_name == "plugins": + register_plugins_from_config(general_settings) + return response @@ -15063,9 +15117,19 @@ async def get_config_general_settings( general_settings = dict(db_general_settings.param_value) if field_name in general_settings: - return ConfigFieldInfo( - field_name=field_name, field_value=general_settings[field_name] - ) + field_value = general_settings[field_name] + # Redact plugin_key from plugin configs so the shared credential + # is never returned even to admin-viewer callers. + if field_name == "plugins" and isinstance(field_value, list): + field_value = [ + ( + {k: ("***" if k == "plugin_key" else v) for k, v in p.items()} + if isinstance(p, dict) + else p + ) + for p in field_value + ] + return ConfigFieldInfo(field_name=field_name, field_value=field_value) else: raise HTTPException( status_code=400, @@ -16387,6 +16451,7 @@ app.include_router(model_access_group_management_router) app.include_router(tag_management_router) app.include_router(workflow_management_router) app.include_router(memory_router) +app.include_router(plugin_router) app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) diff --git a/tests/test_litellm/proxy/test_plugin_routes.py b/tests/test_litellm/proxy/test_plugin_routes.py new file mode 100644 index 00000000000..52999447179 --- /dev/null +++ b/tests/test_litellm/proxy/test_plugin_routes.py @@ -0,0 +1,232 @@ +"""Regression tests for UI-registered embed plugins. + +Covers three bugs: +1. `general_settings.plugins` was not a field on ConfigGeneralSettings, so the + admin UI's POST /config/field/update with field_name="plugins" was rejected + with "Invalid field=plugins passed in." +2. The in-memory plugin registry only refreshed at startup, so a plugin added + via the UI did not appear in /api/plugins until a restart. +3. Plugins persisted to DB general_settings were not loaded on startup (the + registry only initialised from the YAML config), so UI-added plugins vanished + after a restart. +""" + +import asyncio +from unittest.mock import MagicMock + +from litellm.proxy._types import ( + ConfigGeneralSettings, + LitellmUserRoles, + PluginConfig, + UserAPIKeyAuth, +) +from litellm.proxy.plugin_routes import list_plugins, register_plugins_from_config + + +def _admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + +def _non_admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER) + + +def test_plugins_is_a_valid_general_setting() -> None: + """The config-update endpoint gates on this exact membership check.""" + assert "plugins" in ConfigGeneralSettings.model_fields + + +def test_config_general_settings_parses_plugin_list() -> None: + """A list of plugin dicts (what the UI sends) coerces into PluginConfig.""" + settings = ConfigGeneralSettings.model_validate( + { + "plugins": [ + { + "name": "chat-ui", + "display_name": "Chat UI", + "url": "http://localhost:3300", + }, + { + "name": "agent-builder", + "url": "http://127.0.0.1:4010", + "plugin_key": "sk-secret", + }, + ] + } + ) + plugins = settings.plugins + assert plugins is not None + assert [p.name for p in plugins] == ["chat-ui", "agent-builder"] + assert isinstance(plugins[0], PluginConfig) + assert plugins[1].display_name is None + assert plugins[1].plugin_key == "sk-secret" + + +def test_registered_plugins_appear_in_list_without_restart() -> None: + """register_plugins_from_config makes UI-added plugins visible immediately, + and replaces (not merges) so removed plugins disappear.""" + register_plugins_from_config( + { + "plugins": [ + { + "name": "chat-ui", + "display_name": "Chat UI", + "url": "http://localhost:3300", + } + ] + } + ) + names = [p["name"] for p in asyncio.run(list_plugins(user_api_key_dict=_admin()))] + assert names == ["chat-ui"] + + register_plugins_from_config( + { + "plugins": [ + { + "name": "chat-ui", + "display_name": "Chat UI", + "url": "http://localhost:3300", + }, + { + "name": "agent-builder", + "display_name": "Agent Builder", + "url": "http://127.0.0.1:4010", + }, + ] + } + ) + names = sorted( + p["name"] for p in asyncio.run(list_plugins(user_api_key_dict=_admin())) + ) + assert names == ["agent-builder", "chat-ui"] + + # Removing a plugin from config drops it from the live list. + register_plugins_from_config({}) + assert asyncio.run(list_plugins(user_api_key_dict=_admin())) == [] + + +def test_plugin_key_is_never_returned_to_the_browser() -> None: + """plugin_key is a credential the UI never needs; /api/plugins must omit it + for every caller, admin included, so it never lands in browser state.""" + register_plugins_from_config( + { + "plugins": [ + { + "name": "p", + "display_name": "P", + "url": "http://localhost:9", + "plugin_key": "sk-secret", + } + ] + } + ) + + admin_entry = asyncio.run(list_plugins(user_api_key_dict=_admin()))[0] + user_entry = asyncio.run(list_plugins(user_api_key_dict=_non_admin()))[0] + + assert "plugin_key" not in admin_entry + assert "plugin_key" not in user_entry + assert admin_entry["url"] == "http://localhost:9" + + register_plugins_from_config({}) + + +def test_db_persisted_plugins_load_on_startup() -> None: + """Plugins saved to DB general_settings must register when the DB config is + merged at startup, not just when present in the YAML file.""" + from litellm.proxy.proxy_server import ProxyConfig + + register_plugins_from_config({}) # start empty (as if YAML had no plugins) + + ProxyConfig()._add_general_settings_from_db_config( + config_data={ + "general_settings": { + "plugins": [ + { + "name": "db-plugin", + "display_name": "DB Plugin", + "url": "http://localhost:5000", + } + ] + } + }, + general_settings={}, + proxy_logging_obj=MagicMock(), + ) + + names = [p["name"] for p in asyncio.run(list_plugins(user_api_key_dict=_admin()))] + assert names == ["db-plugin"] + + register_plugins_from_config({}) + + +def test_safe_response_headers_sandbox_and_strips_wire_headers() -> None: + """Proxied plugin responses must be inert and shed wire/cookie headers.""" + from litellm.proxy.plugin_routes import _safe_response_headers + + out = _safe_response_headers( + { + "content-type": "text/html", + "content-encoding": "gzip", + "content-length": "123", + "set-cookie": "session=abc", + "content-security-policy": "default-src *", + } + ) + + assert out["content-security-policy"] == "sandbox" + assert out["x-content-type-options"] == "nosniff" + assert out["content-type"] == "text/html" + for stripped in ("content-encoding", "content-length", "set-cookie"): + assert stripped not in out + + +def test_litellm_credential_header_names_covers_every_auth_header() -> None: + """The canonical strip set must list every header user_api_key_auth accepts + as a litellm key, so a new auth header can't silently start leaking.""" + from litellm.proxy._types import SpecialHeaders + + assert SpecialHeaders.litellm_credential_header_names() == { + "authorization", + "api-key", + "x-api-key", + "x-goog-api-key", + "ocp-apim-subscription-key", + "x-litellm-api-key", + } + + +def test_every_litellm_auth_header_is_stripped_before_forwarding() -> None: + """A plugin must never receive any header that authenticates against litellm, + only the hop-by-hop set and benign headers are forwarded.""" + from litellm.proxy.plugin_routes import _request_strip_headers + + strip = _request_strip_headers() + incoming = { + "Authorization": "Bearer sk-litellm", + "API-Key": "sk-litellm", + "X-Api-Key": "sk-litellm", + "X-Goog-Api-Key": "sk-litellm", + "Ocp-Apim-Subscription-Key": "sk-litellm", + "X-Litellm-Api-Key": "sk-litellm", + "Cookie": "litellm_session=abc", + "Accept": "application/json", + "X-Trace-Id": "t-1", + } + forwarded = {k: v for k, v in incoming.items() if k.lower() not in strip} + + assert forwarded == {"Accept": "application/json", "X-Trace-Id": "t-1"} + + +def test_configured_custom_key_header_is_stripped() -> None: + """A custom general_settings.litellm_key_header_name must also be stripped, + read live so config changes are honoured without a restart.""" + from litellm.proxy import proxy_server + from litellm.proxy.plugin_routes import _request_strip_headers + + original = getattr(proxy_server, "general_settings", None) + proxy_server.general_settings = {"litellm_key_header_name": "X-My-Tenant-Key"} + try: + assert "x-my-tenant-key" in _request_strip_headers() + finally: + proxy_server.general_settings = original diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 6017b9555e9..8b10539b188 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8326,3 +8326,39 @@ def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): assert fields["cancel_on_disconnect"]["field_type"] == "Boolean" finally: app.dependency_overrides.clear() + + +def test_preserve_redacted_plugin_keys_keeps_stored_credential(): + """A redacted or blank plugin_key on update must not overwrite the real key.""" + from litellm.proxy.proxy_server import _preserve_redacted_plugin_keys + + existing = [{"name": "p1", "url": "https://p1", "plugin_key": "sk-real-1"}] + + redacted = _preserve_redacted_plugin_keys( + [{"name": "p1", "url": "https://p1-new", "plugin_key": "***"}], existing + ) + assert redacted == [ + {"name": "p1", "url": "https://p1-new", "plugin_key": "sk-real-1"} + ] + + blanked = _preserve_redacted_plugin_keys( + [{"name": "p1", "url": "https://p1", "plugin_key": ""}], existing + ) + assert blanked[0]["plugin_key"] == "sk-real-1" + + +def test_preserve_redacted_plugin_keys_sets_new_and_drops_orphan_placeholder(): + """A real new key replaces; a placeholder with no stored key is dropped, never persisted.""" + from litellm.proxy.proxy_server import _preserve_redacted_plugin_keys + + existing = [{"name": "p1", "url": "https://p1", "plugin_key": "sk-real-1"}] + + rotated = _preserve_redacted_plugin_keys( + [{"name": "p1", "url": "https://p1", "plugin_key": "sk-new"}], existing + ) + assert rotated[0]["plugin_key"] == "sk-new" + + new_plugin = _preserve_redacted_plugin_keys( + [{"name": "p2", "url": "https://p2", "plugin_key": "***"}], existing + ) + assert "plugin_key" not in new_plugin[0] diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agentControlPlaneView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agentControlPlaneView.test.tsx new file mode 100644 index 00000000000..26de1eed7b8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/agentControlPlaneView.test.tsx @@ -0,0 +1,83 @@ +import { describe, it, expect, vi } from "vitest"; +import { render, waitFor } from "@testing-library/react"; +import { AgentControlPlaneView } from "./layout"; + +const { getMock } = vi.hoisted(() => ({ getMock: vi.fn(() => Promise.resolve({ session_claim: "claim" })) })); + +const pluginModeValue = { + mode: "litellm-platform-plugin" as string, + setMode: vi.fn(), + plugins: [{ name: "litellm-platform-plugin", display_name: "Chat UI", url: "http://localhost:3300" }], + activePlugin: { name: "litellm-platform-plugin", display_name: "Chat UI", url: "http://localhost:3300" } as { + name: string; + display_name: string; + url: string; + } | null, +}; + +vi.mock("@/contexts/PluginModeContext", () => ({ usePluginMode: () => pluginModeValue })); +vi.mock("@/contexts/AuthContext", () => ({ useAuth: () => ({ accessToken: "sk-test-token" }) })); + +vi.mock("@/lib/http/client", () => ({ + createApiClient: () => ({ get: getMock }), +})); +vi.mock("@/components/networking", () => ({ getProxyBaseUrl: () => "" })); + +describe("AgentControlPlaneView iframe", () => { + it("embeds the plugin at its ROOT url, never a hardcoded subpath like /sessions", () => { + const { container } = render(); + const iframe = container.querySelector("iframe"); + + expect(iframe).not.toBeNull(); + const src = iframe!.getAttribute("src")!; + expect(src).toBe("http://localhost:3300/"); + expect(src).not.toContain("/sessions"); + // title comes from the plugin's display_name, not a hardcoded label + expect(iframe!.getAttribute("title")).toBe("Chat UI"); + }); + + it("does not double the slash when the plugin url has a trailing slash", () => { + pluginModeValue.activePlugin = { + name: "litellm-platform-plugin", + display_name: "Chat UI", + url: "http://localhost:3300/", + }; + const { container } = render(); + expect(container.querySelector("iframe")!.getAttribute("src")).toBe("http://localhost:3300/"); + pluginModeValue.activePlugin = { + name: "litellm-platform-plugin", + display_name: "Chat UI", + url: "http://localhost:3300", + }; + }); + + it("does not leak the raw token in the iframe src (token goes via encrypted postMessage)", () => { + const { container } = render(); + expect(container.querySelector("iframe")!.getAttribute("src")).not.toContain("token"); + }); + + it("does not delegate clipboard-read to the untrusted plugin iframe", () => { + const { container } = render(); + const allow = container.querySelector("iframe")!.getAttribute("allow") ?? ""; + + expect(allow).not.toContain("clipboard-read"); + expect(allow).toContain("clipboard-write"); + }); + + it("requests the auth-token claim scoped to the active plugin, not a hardcoded default", async () => { + getMock.mockClear(); + pluginModeValue.activePlugin = { name: "reports-plugin", display_name: "Reports", url: "http://localhost:3300" }; + render(); + + await waitFor(() => expect(getMock).toHaveBeenCalled()); + const [path, opts] = getMock.mock.calls[0]; + expect(path).toBe("/api/plugins/auth-token"); + expect(opts.query).toEqual({ plugin_name: "reports-plugin" }); + + pluginModeValue.activePlugin = { + name: "litellm-platform-plugin", + display_name: "Chat UI", + url: "http://localhost:3300", + }; + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index b32bed44a87..a5e83436888 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -1,6 +1,6 @@ "use client"; -import React, { Suspense, useState } from "react"; +import React, { Suspense, useState, useRef, useEffect } from "react"; import Navbar from "@/components/navbar"; import LoadingScreen from "@/components/common_components/LoadingScreen"; import { ThemeProvider } from "@/contexts/ThemeContext"; @@ -9,6 +9,88 @@ import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; import { useRouter, useSearchParams, usePathname } from "next/navigation"; import { DebugWarningBanner } from "@/components/DebugWarningBanner"; import { MIGRATED_PAGES, migratedHref, legacyPageHref, legacyKeyForPathname } from "@/utils/migratedPages"; +import { PluginModeProvider, usePluginMode } from "@/contexts/PluginModeContext"; +import { createApiClient } from "@/lib/http/client"; +import { getProxyBaseUrl } from "@/components/networking"; + +const pluginApiClient = createApiClient({ getBaseUrl: () => getProxyBaseUrl() ?? "" }); + +// Wrapper so PluginModeProvider receives the live accessToken from auth context, +// which means plugin data refreshes on login/logout without stale cookie reads. +function PluginModeProviderWithAuth({ children }: { children: React.ReactNode }) { + const { accessToken } = useAuth(); + return {children}; +} + +export function AgentControlPlaneView() { + const { activePlugin } = usePluginMode(); + const activePluginName = activePlugin?.name; + const agentPlatformUrl = activePlugin?.url ?? ""; + const { accessToken } = useAuth(); + const iframeRef = useRef(null); + const [auth, setAuth] = useState<{ plugin: string; claim: string } | null>(null); + + // Fetch a short-lived identity claim scoped to the *active* plugin. The claim + // is encrypted under that plugin's own per-plugin key, so it must be requested + // per plugin and re-fetched when the user switches plugins. + useEffect(() => { + if (!accessToken || !activePluginName) return; + let cancelled = false; + pluginApiClient + .get("/api/plugins/auth-token", { accessToken, query: { plugin_name: activePluginName } }) + .then((data: { session_claim?: string }) => { + if (!cancelled && data?.session_claim) setAuth({ plugin: activePluginName, claim: data.session_claim }); + }) + .catch(() => {}); + return () => { + cancelled = true; + }; + }, [accessToken, activePluginName]); + + // Deliver the claim to the iframe via postMessage, but only while it was issued + // for the plugin currently mounted — never replay one plugin's claim to another. + // targetOrigin is the configured plugin URL — no other origin receives it. + useEffect(() => { + const iframe = iframeRef.current; + if (!iframe || !auth || auth.plugin !== activePluginName || !agentPlatformUrl) return; + const send = () => { + iframe.contentWindow?.postMessage({ type: "litellm-auth", session_claim: auth.claim }, agentPlatformUrl); + }; + // Cover both orderings: the iframe may have already fired `load` before the + // claim arrived (send now), or it may load/reload later (send on the event). + send(); + iframe.addEventListener("load", send); + return () => iframe.removeEventListener("load", send); + }, [auth, activePluginName, agentPlatformUrl]); + + if (!agentPlatformUrl) { + return ( +
+
+

Plugin

+

Configure the plugin URL in settings

+
+
+ ); + } + + // Embed the plugin at its root; the plugin renders its own full UI (incl. nav) inside. + return ( +