mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-29 01:42:19 +00:00
071cb49d32
* fix(jwt): accept a team alias in x-litellm-team-id The header only matched canonical team ids, so a JWT caller selecting one of their teams by its alias got a 403 even though they belonged to it. The header value is now resolved through the existing alias lookup before the JWT allowed-team check and the DB membership fallback, while a value that is already a team id never costs an alias lookup and denials keep naming the value the caller sent Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(jwt): only alias a header team id the database provably lacks Under fallback_to_db_teams a header value whose team row read fails for any reason other than TeamNotFoundError now keeps the membership denial instead of falling through to the alias lookup, so a degraded read cannot select a different team that carries the value as an alias. Drops the HeaderTeam docstring that only restated its fields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: ryan <ryan@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
104 lines
3.9 KiB
Python
104 lines
3.9 KiB
Python
"""Client for the `other` holding-pen suite: the auth gate (master key vs an
|
|
invalid key on an admin route), JWT auth against the suite's Keycloak realm
|
|
(idp.py), and the process-lifecycle health probes (liveness, public readiness,
|
|
authenticated readiness diagnostics).
|
|
|
|
Holds the shared ProxyClient so `resources` / `scoped_key` still clean up, and
|
|
adds only the routes these behaviors need. The health probes deliberately send
|
|
no auth header (public routes), so they go through the transport with an empty
|
|
headers model rather than a bearer. JWT tests reach the identity provider
|
|
through `idp`, which provisions identities and mints tokens through Keycloak's
|
|
own endpoints, so no test ever holds a signing key.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from e2e_http import AuthHeaders, NoBody, ProbeResult, Result
|
|
from idp import Keycloak, keycloak_from_env
|
|
from models import (
|
|
ChatBody,
|
|
ChatResponse,
|
|
ReadinessDetailsResponse,
|
|
ReadinessResponse,
|
|
UserListParams,
|
|
UserListResponse,
|
|
)
|
|
from proxy_client import ProxyClient
|
|
from pydantic import Field
|
|
|
|
|
|
class TeamHeaders(AuthHeaders):
|
|
"""Bearer auth plus ``x-litellm-team-id``, the header a JWT caller sends to
|
|
pick one of the teams it belongs to."""
|
|
|
|
x_litellm_team_id: str = Field(serialization_alias="x-litellm-team-id")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class OtherClient:
|
|
proxy: ProxyClient
|
|
|
|
@property
|
|
def idp(self) -> Keycloak:
|
|
"""Resolved per use, so the suite's non-JWT tests never need the IdP env."""
|
|
return keycloak_from_env()
|
|
|
|
def liveness(self) -> ProbeResult:
|
|
"""GET /health/liveliness. Unauthenticated; the probe returns status +
|
|
raw body so the test can assert the worker reports itself alive."""
|
|
return self.proxy.transport.probe("/health/liveliness", params=NoBody())
|
|
|
|
def readiness_public(self) -> Result[ReadinessResponse]:
|
|
"""GET /health/readiness with no credential at all, proving the probe is
|
|
safe to expose to an unauthenticated load balancer."""
|
|
return self.proxy.transport.get(
|
|
"/health/readiness",
|
|
headers=NoBody(),
|
|
params=NoBody(),
|
|
response_type=ReadinessResponse,
|
|
)
|
|
|
|
def readiness_details(self, key: str) -> Result[ReadinessDetailsResponse]:
|
|
return self.proxy.transport.get(
|
|
"/health/readiness/details",
|
|
headers=self.proxy.transport.bearer(key),
|
|
params=NoBody(),
|
|
response_type=ReadinessDetailsResponse,
|
|
)
|
|
|
|
def readiness_details_unauthenticated(self) -> Result[ReadinessDetailsResponse]:
|
|
return self.proxy.transport.get(
|
|
"/health/readiness/details",
|
|
headers=NoBody(),
|
|
params=NoBody(),
|
|
response_type=ReadinessDetailsResponse,
|
|
)
|
|
|
|
def chat_as_team(self, token: str, team: str, body: ChatBody) -> Result[ChatResponse]:
|
|
"""POST /chat/completions under `token` with `x-litellm-team-id: team`."""
|
|
return self.proxy.transport.post(
|
|
"/chat/completions",
|
|
headers=TeamHeaders(
|
|
authorization=self.proxy.transport.bearer(token).authorization,
|
|
x_litellm_team_id=team,
|
|
),
|
|
json=body,
|
|
response_type=ChatResponse,
|
|
)
|
|
|
|
def list_users_as(self, key: str) -> Result[UserListResponse]:
|
|
"""GET /user/list under `key`. Admin-only, so it doubles as the master
|
|
key's authorization proof: the master key (proxy admin) reads it, a
|
|
non-matching key is rejected before it ever reaches the handler."""
|
|
return self.proxy.transport.get(
|
|
"/user/list",
|
|
headers=self.proxy.transport.bearer(key),
|
|
params=UserListParams(user_ids="e2e-test-user"),
|
|
response_type=UserListResponse,
|
|
)
|
|
|
|
|
|
def build_client(proxy: ProxyClient) -> OtherClient:
|
|
return OtherClient(proxy=proxy)
|