mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
fix(jwt,otel): backport session conversation id and JWT team header selection to rc/1.103.0 (#43343)
* feat(otel): emit gen_ai.conversation.id from the caller's session id on v2 LLM spans (#42486) * feat(otel): emit gen_ai.conversation.id from the caller's session id on v2 LLM spans Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): keep the caller's header session under missing_session_id: generate and read replayed payload session ids Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): drop only the proxy-minted session id so a caller id on the other metadata key survives Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): keep a replayed session id hidden when it only echoes the payload trace id Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): keep a replayed session id even when the payload trace id fell back to it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): stop reading the replayed payload's session id, the generated marker does not survive replay Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): audit gen_ai.conversation.id on otel v2 spans through a real proxy, sink and postgres Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): keep otel conversation rigs alive for the whole session so shuffled shards do not reboot the proxy per test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(otel): stop the audit rig proxies from probing sibling test peers for model info Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(otel): record accepted OTLP batches in the sink instead of mutating the collector Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(otel): guard the accepted batch deque so snapshots cannot race sink appends Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mrinal <mrinal@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yucheng <yucheng@berri.ai> (cherry picked from commita3196907e4) * fix(jwt): accept a team alias in x-litellm-team-id (#42445) * 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> (cherry picked from commit071cb49d32) * fix(jwt): say x-litellm-team-id matched no team id or alias in the 403 (#42495) * fix(jwt): say x-litellm-team-id matched no team id or alias in the 403 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(jwt): tell the caller when x-litellm-team-id names an alias shared by several teams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(jwt): deny a shared x-litellm-team-id alias exactly like an unknown value A distinct 403 for an alias several teams share was raised before the allowed-teams check, so any JWT could probe which aliases exist. The alias lookup now treats the duplicate as a miss, and both denials say the value does not resolve to a team id or a unique team alias, which is true for unknown, unauthorized and duplicate values alike 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> (cherry picked from commit08639fcf42) * fix(jwt): let x-litellm-team-id select DB membership teams when the token also carries a team claim (#43206) * fix(jwt): let x-litellm-team-id select DB membership teams when the token also carries a team claim Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(jwt): describe header team selection under fallback_to_db_teams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit7b4fd47c6e) --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mrinal <mrinal@berri.ai> Co-authored-by: yucheng <yucheng@berri.ai> Co-authored-by: ryan <ryan@berri.ai> Co-authored-by: yassin <yassin@berri.ai>
This commit is contained in:
parent
6bfc987349
commit
6076c0a058
14 changed files with 1909 additions and 106 deletions
|
|
@ -561,6 +561,7 @@ class OpenTelemetryV2(CustomLogger):
|
|||
time_to_first_chunk_seconds=call.time_to_first_chunk_seconds,
|
||||
request_route=request_root_http_route(),
|
||||
trace=call.trace,
|
||||
session_id=call.session_id,
|
||||
)
|
||||
end_time_ns: Final = to_ns(end_time)
|
||||
if carrier is not None and carrier.span is not None:
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ class GenAIMapper:
|
|||
GenAI.OPERATION_NAME: lambda d: d.operation.value,
|
||||
GenAI.PROVIDER_NAME: lambda d: d.provider or None,
|
||||
GenAI.OUTPUT_TYPE: lambda d: d.output_type.value if d.output_type else None,
|
||||
GenAI.CONVERSATION_ID: lambda d: d.session_id,
|
||||
GenAI.REQUEST_MODEL: lambda d: d.request_model or None,
|
||||
GenAI.REQUEST_TEMPERATURE: lambda d: d.request_params.temperature,
|
||||
GenAI.REQUEST_TOP_P: lambda d: d.request_params.top_p,
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ from dataclasses import dataclass, field
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
|
||||
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, SESSION_ID_GENERATED_METADATA_KEY
|
||||
from litellm.integrations.otel.model.semconv import resolve_operation
|
||||
from litellm.integrations.otel.model.trace_controls import TraceControls, caller_trace_controls
|
||||
from litellm.integrations.otel.model.utils import as_str, as_str_mapping, to_seconds
|
||||
|
|
@ -226,6 +226,7 @@ class LLMCallEvent:
|
|||
provisional_span_name: str
|
||||
time_to_first_chunk_seconds: float | None
|
||||
trace: TraceControls
|
||||
session_id: str | None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, kwargs: Mapping[str, Any]) -> LLMCallEvent:
|
||||
|
|
@ -233,6 +234,7 @@ class LLMCallEvent:
|
|||
payload: Final = cast("StandardLoggingPayload", raw_payload) if raw_payload else None
|
||||
operation: Final = resolve_operation(as_str(kwargs.get("call_type")))
|
||||
model: Final = as_str(kwargs.get("model")) or ""
|
||||
trace: Final = caller_trace_controls(kwargs)
|
||||
return cls(
|
||||
call_id=_call_id(payload, kwargs),
|
||||
payload=payload,
|
||||
|
|
@ -242,10 +244,40 @@ class LLMCallEvent:
|
|||
upstream_started=kwargs.get("api_call_start_time") is not None,
|
||||
provisional_span_name=f"{operation.value} {model}".strip(),
|
||||
time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs),
|
||||
trace=caller_trace_controls(kwargs),
|
||||
trace=trace,
|
||||
session_id=caller_session_id(kwargs, trace),
|
||||
)
|
||||
|
||||
|
||||
def caller_session_id(kwargs: Mapping[str, object], trace: TraceControls) -> str | None:
|
||||
"""The conversation id the caller sent (``litellm_session_id``, else the
|
||||
``session_id`` trace control); ``None`` when the request carried none.
|
||||
|
||||
``get_litellm_params`` back-fills ``litellm_session_id`` from ``metadata.trace_id``
|
||||
(which the proxy stamps with the OTel trace id) and ``missing_session_id: generate``
|
||||
mints one into the body; neither is a caller conversation, so both are ignored,
|
||||
while a ``langfuse_session_id`` header still counts under the generate policy.
|
||||
``StandardLoggingPayload.session_id`` is never read: the payload drops the
|
||||
generated marker, so a replayed minted id would pass for a caller's."""
|
||||
params: Final[Mapping[str, object]] = as_str_mapping(kwargs.get("litellm_params")) or MappingProxyType({})
|
||||
bodies: Final = tuple(
|
||||
metadata
|
||||
for key in ("metadata", "litellm_metadata")
|
||||
if (metadata := as_str_mapping(params.get(key))) is not None
|
||||
)
|
||||
from_body: Final = tuple(session for body in bodies if (session := as_str(body.get("session_id"))))
|
||||
minted: Final = frozenset(
|
||||
session
|
||||
for body in bodies
|
||||
if body.get(SESSION_ID_GENERATED_METADATA_KEY) and (session := as_str(body.get("session_id")))
|
||||
)
|
||||
if minted:
|
||||
return next((session for session in (trace.session_id, *from_body) if session and session not in minted), None)
|
||||
explicit: Final = as_str(params.get("litellm_session_id"))
|
||||
echoes_trace_id: Final = explicit is not None and any(as_str(body.get("trace_id")) == explicit for body in bodies)
|
||||
return (None if echoes_trace_id else explicit) or trace.session_id or None
|
||||
|
||||
|
||||
def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None:
|
||||
"""Seconds from the upstream request being issued (``api_call_start_time``)
|
||||
to the first streamed chunk (``completion_start_time``); ``None`` for
|
||||
|
|
|
|||
|
|
@ -407,6 +407,7 @@ class LLMCallSpanData:
|
|||
call_type: str | None = None
|
||||
request_route: str | None = None
|
||||
trace: TraceControls = field(default_factory=TraceControls)
|
||||
session_id: str | None = None
|
||||
embedding_output: EmbeddingOutput | None = None
|
||||
|
||||
@classmethod
|
||||
|
|
@ -417,6 +418,7 @@ class LLMCallSpanData:
|
|||
time_to_first_chunk_seconds: float | None = None,
|
||||
request_route: str | None = None,
|
||||
trace: TraceControls | None = None,
|
||||
session_id: str | None = None,
|
||||
) -> LLMCallSpanData:
|
||||
params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {})
|
||||
# The single parse of the request's metadata — the request-vs-provider
|
||||
|
|
@ -463,6 +465,7 @@ class LLMCallSpanData:
|
|||
call_type=call_type or None,
|
||||
request_route=request_route or context.identity.request_route,
|
||||
trace=trace or TraceControls(),
|
||||
session_id=session_id or None,
|
||||
embedding_output=embedding_output if capture_content else None,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -5236,11 +5236,12 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
|
|||
default=False,
|
||||
description=(
|
||||
"When True, users whose JWT contains no team claims are authenticated "
|
||||
"using their database team memberships instead of receiving HTTP 403. "
|
||||
"Usage is attributed to the user's first resolvable DB team, or to the "
|
||||
"team specified via the x-litellm-team-id request header (validated "
|
||||
"against DB membership). Requires user_id_upsert=True so that user "
|
||||
"records exist before the fallback runs."
|
||||
"using their database team memberships instead of receiving HTTP 403, "
|
||||
"with usage attributed to the user's first resolvable DB team. Whether or "
|
||||
"not the JWT carries team claims, the x-litellm-team-id request header may "
|
||||
"select any team the user is a member of in the database (validated against "
|
||||
"DB membership); without the header the JWT team stays the default. Requires "
|
||||
"user_id_upsert=True so that user records exist before the fallback runs."
|
||||
),
|
||||
)
|
||||
issuers: list[JWTIssuerConfig] | None = Field(
|
||||
|
|
|
|||
|
|
@ -70,6 +70,7 @@ from litellm.types.agents import AgentResponse
|
|||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
|
||||
from .auth_checks import (
|
||||
TeamNotFoundError,
|
||||
_allowed_routes_check,
|
||||
allowed_routes_check,
|
||||
get_actual_routes,
|
||||
|
|
@ -147,6 +148,12 @@ class _JWTProvisioning:
|
|||
team_id_upsert: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class HeaderTeam:
|
||||
header_value: str
|
||||
team_id: str
|
||||
|
||||
|
||||
class AgentLookup(Protocol):
|
||||
"""The registered-agent lookups a JWT agent claim is matched against."""
|
||||
|
||||
|
|
@ -1866,48 +1873,108 @@ class JWTAuthManager:
|
|||
return True
|
||||
|
||||
@staticmethod
|
||||
def get_team_id_from_header(
|
||||
request_headers: Mapping[str, str] | None,
|
||||
allowed_team_ids: set[str],
|
||||
fallback_to_db_teams: bool = False,
|
||||
) -> str | None:
|
||||
"""
|
||||
Extract team_id from x-litellm-team-id header if present.
|
||||
Validates that the team is in the user's allowed teams from JWT.
|
||||
|
||||
Args:
|
||||
request_headers: Dictionary of request headers
|
||||
allowed_team_ids: Set of team IDs the user is allowed to access (from JWT)
|
||||
fallback_to_db_teams: When True and the JWT carries no team claims
|
||||
(allowed_team_ids is empty), the header value is returned
|
||||
provisionally and validated against DB memberships later in
|
||||
auth_builder instead of being rejected here.
|
||||
|
||||
Returns:
|
||||
The team_id from header if valid, None otherwise
|
||||
|
||||
Raises:
|
||||
HTTPException: If team_id is provided but not in allowed_team_ids
|
||||
"""
|
||||
def _team_header_value(request_headers: Mapping[str, str] | None) -> str | None:
|
||||
if not request_headers:
|
||||
return None
|
||||
|
||||
# Normalize headers to lowercase for case-insensitive lookup
|
||||
normalized_headers: Final = {k.lower(): v for k, v in request_headers.items()}
|
||||
header_team_id: Final = normalized_headers.get("x-litellm-team-id")
|
||||
return normalized_headers.get("x-litellm-team-id")
|
||||
|
||||
if not header_team_id:
|
||||
@staticmethod
|
||||
def _raise_header_team_not_allowed(header_value: str, allowed_team_ids: set[str]) -> NoReturn:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=(
|
||||
f"x-litellm-team-id '{header_value}' does not resolve to a team id or a unique team alias in your "
|
||||
f"JWT's allowed teams. Allowed team ids: {sorted(allowed_team_ids)}"
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _team_id_by_alias(
|
||||
team_alias: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> str | None:
|
||||
if prisma_client is None:
|
||||
return None
|
||||
try:
|
||||
team: Final = await get_team_object_by_alias(
|
||||
team_alias=team_alias,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except HTTPException as exc:
|
||||
if exc.status_code >= 500:
|
||||
raise
|
||||
return None
|
||||
return team.team_id
|
||||
|
||||
@staticmethod
|
||||
async def resolve_team_from_header(
|
||||
request_headers: Mapping[str, str] | None,
|
||||
allowed_team_ids: set[str],
|
||||
fallback_to_db_teams: bool,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> HeaderTeam | None:
|
||||
"""
|
||||
The team named by x-litellm-team-id, which may carry a team id or a team
|
||||
alias. A value that is already an allowed team id never costs a lookup;
|
||||
under the DB fallback any other value is accepted provisionally, by id
|
||||
or alias, for the membership check auth_builder runs later. Under the
|
||||
DB fallback only a team row that is provably absent falls through to
|
||||
the alias lookup; a read that failed for any other reason keeps the
|
||||
membership denial the id path already gives.
|
||||
|
||||
Raises:
|
||||
HTTPException: 403 when neither the value nor the team it aliases is
|
||||
an allowed team, or the DB fallback's membership denial when the
|
||||
value names no team at all; an alias several teams share resolves
|
||||
to no team and is denied like an unknown value; a 5xx from the
|
||||
alias lookup itself is re-raised rather than reported as a denial
|
||||
"""
|
||||
header_value: Final = JWTAuthManager._team_header_value(request_headers)
|
||||
if not header_value:
|
||||
return None
|
||||
|
||||
defer_to_db_membership: Final = fallback_to_db_teams and not allowed_team_ids
|
||||
if not defer_to_db_membership and header_team_id not in allowed_team_ids:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"Team '{header_team_id}' from x-litellm-team-id header is not in your JWT's allowed teams. Allowed teams: {list(allowed_team_ids)}",
|
||||
)
|
||||
if header_value in allowed_team_ids:
|
||||
verbose_proxy_logger.debug("Using team_id from x-litellm-team-id header: %s", header_value)
|
||||
return HeaderTeam(header_value=header_value, team_id=header_value)
|
||||
|
||||
verbose_proxy_logger.debug("Using team_id from x-litellm-team-id header: %s", header_team_id)
|
||||
return header_team_id
|
||||
if fallback_to_db_teams:
|
||||
try:
|
||||
await get_team_object(
|
||||
team_id=header_value,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=False,
|
||||
)
|
||||
except TeamNotFoundError:
|
||||
aliased_team_id: Final = await JWTAuthManager._team_id_by_alias(
|
||||
header_value, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj
|
||||
)
|
||||
if aliased_team_id is None:
|
||||
JWTAuthManager._raise_header_team_membership_denial(header_value)
|
||||
return HeaderTeam(header_value=header_value, team_id=aliased_team_id)
|
||||
except HTTPException:
|
||||
JWTAuthManager._raise_header_team_membership_denial(header_value)
|
||||
return HeaderTeam(header_value=header_value, team_id=header_value)
|
||||
|
||||
team_id_by_alias: Final = await JWTAuthManager._team_id_by_alias(
|
||||
header_value, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj
|
||||
)
|
||||
if team_id_by_alias is None or team_id_by_alias not in allowed_team_ids:
|
||||
JWTAuthManager._raise_header_team_not_allowed(header_value, allowed_team_ids)
|
||||
verbose_proxy_logger.debug("Using team_id %s for x-litellm-team-id alias: %s", team_id_by_alias, header_value)
|
||||
return HeaderTeam(header_value=header_value, team_id=team_id_by_alias)
|
||||
|
||||
@staticmethod
|
||||
async def map_user_to_teams(
|
||||
|
|
@ -2256,31 +2323,37 @@ class JWTAuthManager:
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _raise_header_team_membership_denial(team_id: str) -> NoReturn:
|
||||
def _raise_header_team_membership_denial(header_value: str) -> NoReturn:
|
||||
"""
|
||||
The single denial shape for a provisional x-litellm-team-id header,
|
||||
raised identically for nonexistent teams and for teams the user is not
|
||||
a member of, so the response does not reveal whether a team id exists.
|
||||
a member of, and naming only the value the caller sent, so the response
|
||||
reveals neither whether a team exists nor which id an alias maps to.
|
||||
"""
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=(f"Team '{team_id}' (from x-litellm-team-id header) is not in your team memberships."),
|
||||
detail=(
|
||||
f"x-litellm-team-id '{header_value}' does not resolve to a team id or a unique team alias among your "
|
||||
"team memberships."
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _validate_header_team_in_db_membership(
|
||||
team_id: str,
|
||||
user_object: LiteLLM_UserTable | None,
|
||||
header_value: str,
|
||||
) -> None:
|
||||
"""
|
||||
A provisional team_id from the x-litellm-team-id header (accepted without
|
||||
JWT-team validation when the JWT carries no team claims) must exist in the
|
||||
user's DB team memberships before it becomes request context.
|
||||
A provisional team_id from the x-litellm-team-id header (accepted under
|
||||
fallback_to_db_teams because it is outside the JWT's teams) must exist in
|
||||
the user's DB team memberships before it becomes request context. The denial
|
||||
names `header_value`, the id or alias the caller sent, not `team_id`.
|
||||
"""
|
||||
user_team_ids: Final = user_object.teams if user_object else []
|
||||
if team_id in user_team_ids:
|
||||
return
|
||||
JWTAuthManager._raise_header_team_membership_denial(team_id)
|
||||
JWTAuthManager._raise_header_team_membership_denial(header_value)
|
||||
|
||||
@staticmethod
|
||||
async def auth_builder(
|
||||
|
|
@ -2491,18 +2564,30 @@ class JWTAuthManager:
|
|||
if specific_team_id and not db_team_fallback:
|
||||
all_team_ids.add(specific_team_id)
|
||||
|
||||
header_team_id: Final = JWTAuthManager.get_team_id_from_header(
|
||||
header_db_fallback: Final = handler.litellm_jwtauth.fallback_to_db_teams and team_id is None
|
||||
|
||||
header_team: Final = await JWTAuthManager.resolve_team_from_header(
|
||||
request_headers=request_headers,
|
||||
allowed_team_ids=all_team_ids,
|
||||
fallback_to_db_teams=db_team_fallback,
|
||||
fallback_to_db_teams=header_db_fallback,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if header_team_id:
|
||||
team_id = header_team_id
|
||||
# A provisional header team (accepted only because the JWT carries no
|
||||
# team claims) is validated against DB membership further down; never
|
||||
# upsert it here or an attacker-supplied x-litellm-team-id would create
|
||||
# an orphaned team row before that check runs. A genuine membership team
|
||||
# already exists, so suppressing the upsert in that case costs nothing.
|
||||
provisional_header_team: Final = (
|
||||
header_team
|
||||
if header_team is not None and header_db_fallback and header_team.team_id not in all_team_ids
|
||||
else None
|
||||
)
|
||||
if header_team:
|
||||
team_id = header_team.team_id
|
||||
# A provisional header team (accepted because it is outside the
|
||||
# JWT's teams under fallback_to_db_teams) is validated against DB
|
||||
# membership further down; never upsert it here or an
|
||||
# attacker-supplied x-litellm-team-id would create an orphaned team
|
||||
# row before that check runs. A genuine membership team already
|
||||
# exists, so suppressing the upsert in that case costs nothing.
|
||||
try:
|
||||
team_object = await get_team_object(
|
||||
team_id=team_id,
|
||||
|
|
@ -2510,12 +2595,12 @@ class JWTAuthManager:
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=(team_id_upsert and not db_team_fallback),
|
||||
team_id_upsert=(team_id_upsert and provisional_header_team is None),
|
||||
)
|
||||
except HTTPException:
|
||||
if not db_team_fallback:
|
||||
if provisional_header_team is None:
|
||||
raise
|
||||
JWTAuthManager._raise_header_team_membership_denial(team_id)
|
||||
JWTAuthManager._raise_header_team_membership_denial(header_team.header_value)
|
||||
elif not team_id and not db_team_fallback:
|
||||
## SPECIFIC TEAM ID
|
||||
(
|
||||
|
|
@ -2654,10 +2739,11 @@ class JWTAuthManager:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id_upsert=team_id_upsert,
|
||||
)
|
||||
elif db_team_fallback and team_id == header_team_id:
|
||||
elif provisional_header_team is not None and team_id == provisional_header_team.team_id:
|
||||
JWTAuthManager._validate_header_team_in_db_membership(
|
||||
team_id=team_id,
|
||||
user_object=user_object,
|
||||
header_value=provisional_header_team.header_value,
|
||||
)
|
||||
if not JWTAuthManager._is_team_route_allowed(
|
||||
route=route,
|
||||
|
|
@ -2667,7 +2753,8 @@ class JWTAuthManager:
|
|||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=(
|
||||
f"Team '{team_id}' (from x-litellm-team-id header) is not allowed to access route '{route}'."
|
||||
f"Team '{provisional_header_team.header_value}' (from x-litellm-team-id header) "
|
||||
f"is not allowed to access route '{route}'."
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -15,6 +15,8 @@
|
|||
- {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:1158-1166 _decode_jwt_with_public_key", rationale: "A genuine token whose signature bytes were altered fails verification with 401"}
|
||||
- {id: other.auth.jwt.unknown_team_denied, module: other, tier: P0, area: auth, assertions: [unknown_team_denied], source: "handle_jwt.py:1549-1632 find_team_with_model_access", rationale: "A verified JWT whose groups claim resolves to no existing team is denied with 403 naming the unresolved team, never silently admitted without a team. The proxy words it as a model-access denial, the same body an existing team without model access gets"}
|
||||
- {id: other.auth.jwt.virtual_key_unaffected, module: other, tier: P0, area: auth, assertions: [virtual_key_unaffected], source: "handle_jwt.py:213 is_jwt / user_api_key_auth.py:1332-1333", rationale: "enable_jwt_auth only routes three-segment bearer tokens into the JWT branch, so sk- virtual keys keep working on the same proxy"}
|
||||
- {id: other.auth.jwt.team_header_alias_binds_team, module: other, tier: P0, area: auth, assertions: [team_header_alias_binds_team], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", fail_before_fix: proven, rationale: "x-litellm-team-id carrying the team alias binds and attributes the same team as the team id, so a managed client can pin a stable alias instead of a uuid"}
|
||||
- {id: other.auth.jwt.team_header_non_member_alias_denied, module: other, tier: P0, area: auth, assertions: [team_header_non_member_alias_denied], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", rationale: "x-litellm-team-id naming the alias of a team the JWT does not grant is denied 403 with the same body as an unknown value, so the response does not reveal whether that team exists"}
|
||||
- {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"}
|
||||
- {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"}
|
||||
- {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"}
|
||||
|
|
|
|||
|
|
@ -15,15 +15,25 @@ from __future__ import annotations
|
|||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from e2e_http import NoBody, ProbeResult, Result
|
||||
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)
|
||||
|
|
@ -66,6 +76,18 @@ class OtherClient:
|
|||
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
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from __future__ import annotations
|
|||
|
||||
import base64
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
|
@ -52,6 +53,31 @@ def identity(client: OtherClient, resources: ResourceManager) -> Identity:
|
|||
return provisioned
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BoundTeam:
|
||||
identity: Identity
|
||||
team_id: str
|
||||
team_alias: str
|
||||
|
||||
|
||||
def _team(client: OtherClient, resources: ResourceManager, *, marker: str, team_id: str) -> str:
|
||||
"""A litellm team whose alias differs from its id, so a header naming one
|
||||
cannot accidentally match the other."""
|
||||
team_alias: Final = f"e2e-jwt-alias-{marker}"
|
||||
created: Final = client.proxy.create_team(TeamNewBody(team_alias=team_alias, team_id=team_id))
|
||||
resources.defer(lambda: client.proxy.delete_team(created))
|
||||
return team_alias
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bound_team(client: OtherClient, resources: ResourceManager) -> BoundTeam:
|
||||
"""An identity whose single group is a real team, plus that team's alias."""
|
||||
marker: Final = unique_marker()
|
||||
provisioned: Final = _provision(client, resources, marker=marker)
|
||||
team_alias: Final = _team(client, resources, marker=marker, team_id=provisioned.group)
|
||||
return BoundTeam(identity=provisioned, team_id=provisioned.group, team_alias=team_alias)
|
||||
|
||||
|
||||
def _ping() -> ChatBody:
|
||||
return ChatBody(
|
||||
model=CHEAP_OPENAI_MODEL,
|
||||
|
|
@ -155,3 +181,58 @@ class TestJwtAuth:
|
|||
def test_plain_virtual_key_still_works_with_jwt_auth_enabled(self, client: OtherClient, scoped_key: str) -> None:
|
||||
response: Final = unwrap(client.proxy.chat(scoped_key, _ping()))
|
||||
assert response.choices, f"an sk- key must keep working on a proxy with enable_jwt_auth, got {response}"
|
||||
|
||||
|
||||
def _team_of_request(client: OtherClient, token: str, team: str) -> str | None:
|
||||
response: Final = unwrap(client.chat_as_team(token, team, _ping()))
|
||||
assert response.id is not None and response.choices, (
|
||||
f"chat with x-litellm-team-id={team!r} returned no completion: {response}"
|
||||
)
|
||||
rows: Final = client.proxy.poll_logs_for_request_id(response.id)
|
||||
assert rows, f"no spend log row for request {response.id} within the poll deadline"
|
||||
return rows[0].team_id
|
||||
|
||||
|
||||
def _denial(client: OtherClient, token: str, team: str) -> str:
|
||||
result: Final = client.chat_as_team(token, team, _ping())
|
||||
assert isinstance(result, UnknownApiError) and result.status_code == 403, (
|
||||
f"x-litellm-team-id={team!r} names no team the caller is in, so it must be rejected with 403, got {result}"
|
||||
)
|
||||
assert team in result.body, f"the 403 must name the header value it rejected ({team!r}), got {result.body[:300]}"
|
||||
return result.body
|
||||
|
||||
|
||||
class TestJwtTeamHeader:
|
||||
@pytest.mark.covers("other.auth.jwt.team_header_alias_binds_team")
|
||||
def test_team_header_with_the_team_alias_binds_the_same_team_as_the_team_id(
|
||||
self, client: OtherClient, bound_team: BoundTeam
|
||||
) -> None:
|
||||
token: Final = client.idp.access_token(bound_team.identity)
|
||||
assert bound_team.team_alias != bound_team.team_id
|
||||
|
||||
by_id: Final = _team_of_request(client, token, bound_team.team_id)
|
||||
assert by_id == bound_team.team_id, (
|
||||
f"precondition: x-litellm-team-id with the team id must bind {bound_team.team_id!r}, got {by_id!r}"
|
||||
)
|
||||
|
||||
by_alias: Final = _team_of_request(client, token, bound_team.team_alias)
|
||||
assert by_alias == bound_team.team_id, (
|
||||
f"x-litellm-team-id={bound_team.team_alias!r} must bind the same team as its id "
|
||||
f"{bound_team.team_id!r}, got {by_alias!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("other.auth.jwt.team_header_non_member_alias_denied")
|
||||
def test_team_header_with_the_alias_of_a_team_the_caller_is_not_in_is_rejected_like_an_unknown_value(
|
||||
self, client: OtherClient, resources: ResourceManager, bound_team: BoundTeam
|
||||
) -> None:
|
||||
token: Final = client.idp.access_token(bound_team.identity)
|
||||
other_marker: Final = unique_marker()
|
||||
other_alias: Final = _team(client, resources, marker=other_marker, team_id=f"e2e-jwt-other-{other_marker}")
|
||||
unknown: Final = f"e2e-jwt-unknown-{unique_marker()}"
|
||||
|
||||
for_other_alias: Final = _denial(client, token, other_alias)
|
||||
for_unknown: Final = _denial(client, token, unknown)
|
||||
assert for_other_alias.replace(other_alias, "<value>") == for_unknown.replace(unknown, "<value>"), (
|
||||
"a non-member alias and an unknown value must get the same denial body, so the response does not "
|
||||
f"reveal whether the team exists; got {for_other_alias[:300]!r} vs {for_unknown[:300]!r}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1311,6 +1311,75 @@
|
|||
],
|
||||
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[us.anthropic.claude-opus-5-v1:0-stream_full_usage]": [
|
||||
"quota_management.spend_tracking.scripted_wire.logs_cost"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_chat_completion_sdk_body_litellm_session_id_lands_as_conversation_id": [
|
||||
"other.observability.otel.conversation_id_from_body_session_id_chat_sdk"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_chat_stream_async_sdk_x_litellm_session_id_header_lands_as_conversation_id": [
|
||||
"other.observability.otel.conversation_id_from_header_chat_stream_async_sdk"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_messages_sdk_x_litellm_session_id_header_lands_as_conversation_id": [
|
||||
"other.observability.otel.conversation_id_from_header_messages_sdk"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_messages_stream_async_sdk_langfuse_session_id_header_lands_as_conversation_id": [
|
||||
"other.observability.otel.conversation_id_from_langfuse_header_messages_stream_async_sdk"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_responses_sdk_x_litellm_session_id_header_lands_as_conversation_id": [
|
||||
"other.observability.otel.conversation_id_from_header_responses_sdk"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_responses_stream_raw_metadata_session_id_lands_as_conversation_id": [
|
||||
"other.observability.otel.conversation_id_from_metadata_responses_stream_raw"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_chat_raw_metadata_session_id_lands_as_conversation_id": [
|
||||
"other.observability.otel.conversation_id_from_metadata_chat_raw"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_integer_and_list_litellm_session_id_match_the_spend_row_or_are_dropped_together": [
|
||||
"other.observability.otel.conversation_id_non_string_session_ids_match_spend_row"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_empty_string_litellm_session_id_leaves_the_span_without_a_conversation_id": [
|
||||
"other.observability.otel.conversation_id_empty_string_session_id_is_omitted"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_five_kilobyte_session_header_round_trips_to_the_span_and_the_spend_row": [
|
||||
"other.observability.otel.conversation_id_five_kilobyte_header_round_trips"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_duplicate_session_header_lands_once_and_unchanged": [
|
||||
"other.observability.otel.conversation_id_duplicate_header_lands_once"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_unauthenticated_request_with_session_header_is_rejected_and_leaves_no_span": [
|
||||
"other.observability.otel.conversation_id_unauthenticated_request_leaves_no_span"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_sink_rejecting_with_403_drops_those_spans_and_later_spans_still_land": [
|
||||
"other.observability.otel.conversation_id_survives_sink_rejection"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_request_without_any_session_input_has_no_conversation_id": [
|
||||
"other.observability.otel.conversation_id_absent_without_caller_session"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_generate_policy_minted_session_id_reaches_the_spend_row_but_not_the_span": [
|
||||
"other.observability.otel.conversation_id_ignores_generated_session_id"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_generate_policy_keeps_the_langfuse_session_header_as_conversation_id": [
|
||||
"other.observability.otel.conversation_id_langfuse_header_wins_over_generated"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_header_body_and_metadata_session_ids_resolve_to_the_same_id_as_the_spend_row": [
|
||||
"other.observability.otel.conversation_id_header_precedence_matches_spend_row"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_three_identical_requests_produce_one_span_each_with_the_same_conversation_id": [
|
||||
"other.observability.otel.conversation_id_repeated_requests_log_once_each"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_metadata_trace_id_alone_fills_the_spend_row_but_not_the_span": [
|
||||
"other.observability.otel.conversation_id_ignores_trace_id_backfill"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_sink_outage_during_a_mixed_burst_lands_every_response_exactly_once_after_recovery": [
|
||||
"other.observability.otel.conversation_id_sink_outage_recovers_exactly_once"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_slow_sink_during_a_burst_lands_every_response_exactly_once": [
|
||||
"other.observability.otel.conversation_id_slow_sink_no_duplicates"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_killing_one_of_two_workers_mid_burst_keeps_serving_and_never_duplicates_a_span": [
|
||||
"other.observability.otel.conversation_id_survives_worker_kill"
|
||||
],
|
||||
"tests/integration/observability/test_otel_conversation_id.py::test_terminating_the_proxy_right_after_a_burst_flushes_every_span_before_exit": [
|
||||
"other.observability.otel.conversation_id_flushes_on_shutdown"
|
||||
]
|
||||
},
|
||||
"browser": {
|
||||
|
|
|
|||
830
tests/integration/observability/test_otel_conversation_id.py
Normal file
830
tests/integration/observability/test_otel_conversation_id.py
Normal file
|
|
@ -0,0 +1,830 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import threading
|
||||
import uuid
|
||||
from collections import deque
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
MARKER: Final = re.compile(rb"otelconv-[0-9a-f]{32}")
|
||||
CONVERSATION: Final = "gen_ai.conversation.id"
|
||||
|
||||
|
||||
def _marker() -> str:
|
||||
return "otelconv-" + uuid.uuid4().hex
|
||||
|
||||
|
||||
def _chat_reply(identity: str, stream: bool) -> Reply:
|
||||
if not stream:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "conversation ok"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"}
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(
|
||||
b"data: "
|
||||
+ json.dumps(
|
||||
{**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "conversation"}}]}
|
||||
).encode()
|
||||
+ b"\n\n",
|
||||
b"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
**chunk,
|
||||
"choices": [{"index": 0, "delta": {"content": " ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
).encode()
|
||||
+ b"\n\n",
|
||||
b"data: [DONE]\n\n",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _responses_reply(identity: str, stream: bool) -> Reply:
|
||||
response: Final = {
|
||||
"id": identity,
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [
|
||||
{
|
||||
"id": "msg_" + identity,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": "conversation ok", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
if not stream:
|
||||
return Reply(body=json.dumps(response).encode())
|
||||
events: Final = (
|
||||
{
|
||||
"type": "response.created",
|
||||
"sequence_number": 0,
|
||||
"response": {**response, "status": "in_progress", "output": []},
|
||||
},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": "msg_" + identity,
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "conversation ok",
|
||||
},
|
||||
{"type": "response.completed", "sequence_number": 2, "response": response},
|
||||
)
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events),
|
||||
)
|
||||
|
||||
|
||||
def _decoded_responses_id(identity: str) -> str:
|
||||
try:
|
||||
return base64.b64decode(identity.removeprefix("resp_").encode()).decode()
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
return identity
|
||||
|
||||
|
||||
def _canonical_id(identity: str) -> str:
|
||||
return _decoded_responses_id(identity).rpartition("response_id:")[2]
|
||||
|
||||
|
||||
def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
json.loads(line[6:]) for line in text.splitlines() if line.startswith("data: ") and line != "data: [DONE]"
|
||||
)
|
||||
|
||||
|
||||
def _upstream(request: Request) -> Reply:
|
||||
found: Final = MARKER.search(request.body)
|
||||
if found is None:
|
||||
return Reply(status=404, body=b'{"error":"no marker"}')
|
||||
marker: Final = found.group(0).decode()
|
||||
stream: Final = json.loads(request.body).get("stream") is True
|
||||
if request.target.endswith("/responses"):
|
||||
return _responses_reply(f"resp_{marker}", stream)
|
||||
return _chat_reply(f"chatcmpl-{marker}", stream)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Collector:
|
||||
wire: Wire
|
||||
outage: threading.Event
|
||||
rejection: threading.Event
|
||||
slow: threading.Event
|
||||
accepted: Sequence[Request]
|
||||
guard: threading.Lock
|
||||
|
||||
def attributes(self) -> tuple[dict[str, dict[str, JsonValue]], ...]:
|
||||
with self.guard:
|
||||
batches: Final = tuple(self.accepted)
|
||||
return tuple(
|
||||
{attribute["key"]: attribute["value"] for attribute in span.get("attributes", ())}
|
||||
for batch in batches
|
||||
for resource in json.loads(batch.body)["resourceSpans"]
|
||||
for scope in resource["scopeSpans"]
|
||||
for span in scope["spans"]
|
||||
)
|
||||
|
||||
def spans(self, response_id: str) -> tuple[dict[str, dict[str, JsonValue]], ...]:
|
||||
return tuple(
|
||||
attributes
|
||||
for attributes in self.attributes()
|
||||
if isinstance(logged := attributes.get("gen_ai.response.id", {}).get("stringValue"), str)
|
||||
and _canonical_id(logged) == _canonical_id(response_id)
|
||||
)
|
||||
|
||||
def conversation_ids(self, response_id: str) -> tuple[str | None, ...]:
|
||||
return tuple(
|
||||
attributes[CONVERSATION]["stringValue"] if CONVERSATION in attributes else None
|
||||
for attributes in self.spans(response_id)
|
||||
)
|
||||
|
||||
def single_span(self, response_id: str) -> str | None:
|
||||
return eventually(lambda: self.conversation_ids(response_id), lambda values: len(values) == 1, seconds=30)[0]
|
||||
|
||||
def logged_id(self, response_id: str) -> str:
|
||||
spans: Final = eventually(lambda: self.spans(response_id), lambda values: len(values) == 1, seconds=30)
|
||||
return str(spans[0]["gen_ai.response.id"]["stringValue"])
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def collector() -> Iterator[Collector]:
|
||||
outage: Final = threading.Event()
|
||||
rejection: Final = threading.Event()
|
||||
slow: Final = threading.Event()
|
||||
accepted: Final[deque[Request]] = deque() # mutable-ok: sink thread appends each accepted batch
|
||||
guard: Final = threading.Lock()
|
||||
|
||||
def sink(request: Request) -> Reply:
|
||||
if slow.is_set():
|
||||
threading.Event().wait(1.5)
|
||||
if outage.is_set():
|
||||
return Reply(status=503, body=b'{"error":"sink down"}')
|
||||
if rejection.is_set():
|
||||
return Reply(status=403, body=b'{"error":"forbidden"}')
|
||||
with guard:
|
||||
accepted.append(request)
|
||||
return Reply()
|
||||
|
||||
with wire_server(sink) as wire:
|
||||
yield Collector(wire, outage, rejection, slow, accepted, guard)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def provider() -> Iterator[Wire]:
|
||||
with wire_server(_upstream) as wire:
|
||||
yield wire
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: Gateway
|
||||
process: OwnedProxy
|
||||
model: str
|
||||
upstream: Wire
|
||||
sink: Collector
|
||||
|
||||
def openai_client(self) -> openai.OpenAI:
|
||||
return openai.OpenAI(base_url=str(self.proxy.client.base_url) + "/v1", api_key=self.proxy.key, max_retries=0)
|
||||
|
||||
def async_openai_client(self) -> openai.AsyncOpenAI:
|
||||
return openai.AsyncOpenAI(
|
||||
base_url=str(self.proxy.client.base_url) + "/v1", api_key=self.proxy.key, max_retries=0
|
||||
)
|
||||
|
||||
def anthropic_client(self) -> anthropic.Anthropic:
|
||||
return anthropic.Anthropic(base_url=str(self.proxy.client.base_url), api_key=self.proxy.key, max_retries=0)
|
||||
|
||||
def async_anthropic_client(self) -> anthropic.AsyncAnthropic:
|
||||
return anthropic.AsyncAnthropic(base_url=str(self.proxy.client.base_url), api_key=self.proxy.key, max_retries=0)
|
||||
|
||||
def chat(
|
||||
self,
|
||||
marker: str,
|
||||
*,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
key: str | None = None,
|
||||
**extra: JsonValue,
|
||||
) -> httpx.Response:
|
||||
return self.proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
"cache": {"no-cache": True},
|
||||
**extra,
|
||||
},
|
||||
headers=headers,
|
||||
key=key,
|
||||
)
|
||||
|
||||
def upstream_bodies(self, marker: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(json.loads(request.body) for request in self.upstream.drain() if marker.encode() in request.body)
|
||||
|
||||
def spend_session(self, response_id: str) -> str | None:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows('SELECT session_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response_id,)),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
value: Final = rows[0]["session_id"]
|
||||
assert value is None or isinstance(value, str), rows
|
||||
return value
|
||||
|
||||
def spend_request_ids(self, session: str) -> tuple[str, ...]:
|
||||
rows: Final = read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE session_id=%s', (session,))
|
||||
return tuple(str(row["request_id"]) for row in rows)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RigFactory:
|
||||
provider: Wire
|
||||
sink: Collector
|
||||
directory: Path
|
||||
settings: Mapping[str, JsonValue]
|
||||
workers: int
|
||||
|
||||
def start(self) -> Iterator[Rig]:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["litellm_settings"].update({"callbacks": ["otel"]})
|
||||
config["general_settings"].update({"disable_model_info_refresh": True, **self.settings})
|
||||
config["callback_settings"] = {
|
||||
"otel": {"exporter": "http/json", "endpoint": self.sink.wire.url, "mapper_names": ["genai"]},
|
||||
}
|
||||
path: Final = self.directory / f"otel-{uuid.uuid4().hex}.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
overrides: Final = {"LITELLM_OTEL_V2": "1", "OTEL_BSP_SCHEDULE_DELAY": "300"}
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
owned_proxy_process(gateway, self.directory, overrides, config=path, workers=self.workers) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(api_base=self.provider.url + "/v1")
|
||||
yield Rig(owned.gateway, owned, model, self.provider, self.sink)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def rig(provider: Wire, collector: Collector, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
yield from RigFactory(provider, collector, tmp_path_factory.mktemp("otel"), {}, 2).start()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def generating_rig(provider: Wire, collector: Collector, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
factory: Final = RigFactory(
|
||||
provider, collector, tmp_path_factory.mktemp("otel-generate"), {"missing_session_id": "generate"}, 2
|
||||
)
|
||||
yield from factory.start()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def two_worker_rig(provider: Wire, collector: Collector, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
yield from RigFactory(provider, collector, tmp_path_factory.mktemp("otel-workers"), {}, 2).start()
|
||||
|
||||
|
||||
def _assert_upstream_clean(rig: Rig, marker: str, session: str) -> None:
|
||||
bodies: Final = rig.upstream_bodies(marker)
|
||||
assert len(bodies) == 1, bodies
|
||||
assert session not in json.dumps(bodies[0]), bodies[0]
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_from_body_session_id_chat_sdk")
|
||||
def test_chat_completion_sdk_body_litellm_session_id_lands_as_conversation_id(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = "conv-" + uuid.uuid4().hex
|
||||
completion: Final = rig.openai_client().chat.completions.create(
|
||||
model=rig.model,
|
||||
messages=[{"role": "user", "content": marker}],
|
||||
extra_body={"litellm_session_id": session, "cache": {"no-cache": True}},
|
||||
)
|
||||
assert completion.id == f"chatcmpl-{marker}", completion
|
||||
assert completion.choices[0].message.content == "conversation ok", completion
|
||||
assert rig.sink.single_span(completion.id) == session
|
||||
assert rig.spend_session(completion.id) == session
|
||||
_assert_upstream_clean(rig, marker, session)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_from_header_chat_stream_async_sdk")
|
||||
def test_chat_stream_async_sdk_x_litellm_session_id_header_lands_as_conversation_id(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = "conv-" + uuid.uuid4().hex
|
||||
|
||||
async def consume() -> tuple[str, str]:
|
||||
stream: Final = await rig.async_openai_client().chat.completions.create(
|
||||
model=rig.model,
|
||||
messages=[{"role": "user", "content": marker}],
|
||||
stream=True,
|
||||
extra_headers={"x-litellm-session-id": session},
|
||||
extra_body={"cache": {"no-cache": True}},
|
||||
)
|
||||
chunks: Final = [chunk async for chunk in stream]
|
||||
return chunks[0].id, "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
|
||||
|
||||
identity, text = asyncio.run(consume())
|
||||
assert identity == f"chatcmpl-{marker}", identity
|
||||
assert text == "conversation ok", text
|
||||
assert rig.sink.single_span(identity) == session
|
||||
assert rig.spend_session(identity) == session
|
||||
_assert_upstream_clean(rig, marker, session)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_from_header_messages_sdk")
|
||||
def test_messages_sdk_x_litellm_session_id_header_lands_as_conversation_id(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = "conv-" + uuid.uuid4().hex
|
||||
message: Final = rig.anthropic_client().messages.create(
|
||||
model=rig.model,
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": marker}],
|
||||
extra_headers={"x-litellm-session-id": session},
|
||||
)
|
||||
assert message.content[0].type == "text" and message.content[0].text == "conversation ok", message
|
||||
assert rig.sink.single_span(message.id) == session
|
||||
assert rig.spend_session(message.id) == session
|
||||
_assert_upstream_clean(rig, marker, session)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_from_langfuse_header_messages_stream_async_sdk")
|
||||
def test_messages_stream_async_sdk_langfuse_session_id_header_lands_as_conversation_id(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = "conv-" + uuid.uuid4().hex
|
||||
|
||||
async def consume() -> tuple[str, str]:
|
||||
async with rig.async_anthropic_client().messages.stream(
|
||||
model=rig.model,
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": marker}],
|
||||
extra_headers={"langfuse_session_id": session},
|
||||
) as stream:
|
||||
text: Final = "".join([piece async for piece in stream.text_stream])
|
||||
return (await stream.get_final_message()).id, text
|
||||
|
||||
identity, text = asyncio.run(consume())
|
||||
assert text == "conversation ok", text
|
||||
assert rig.sink.single_span(identity) == session
|
||||
assert rig.spend_session(identity), identity
|
||||
_assert_upstream_clean(rig, marker, session)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_from_header_responses_sdk")
|
||||
def test_responses_sdk_x_litellm_session_id_header_lands_as_conversation_id(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = "conv-" + uuid.uuid4().hex
|
||||
response: Final = rig.openai_client().responses.create(
|
||||
model=rig.model, input=marker, extra_headers={"x-litellm-session-id": session}
|
||||
)
|
||||
assert response.output[0].id == f"msg_resp_{marker}", response
|
||||
assert response.output_text == "conversation ok", response
|
||||
assert rig.sink.single_span(response.id) == session
|
||||
assert rig.spend_session(response.id) == session
|
||||
_assert_upstream_clean(rig, marker, session)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_from_metadata_responses_stream_raw")
|
||||
def test_responses_stream_raw_metadata_session_id_lands_as_conversation_id(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = "conv-" + uuid.uuid4().hex
|
||||
with rig.proxy.client.stream(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
json={"model": rig.model, "input": marker, "stream": True, "metadata": {"session_id": session}},
|
||||
headers={"Authorization": f"Bearer {rig.proxy.key}"},
|
||||
) as response:
|
||||
body: Final = response.read().decode()
|
||||
assert response.status_code == 200, body
|
||||
events: Final = _sse_events(body)
|
||||
completed: Final = tuple(event for event in events if event["type"] == "response.completed")
|
||||
assert len(completed) == 1, events
|
||||
assert completed[0]["response"]["output"][0]["id"] == f"msg_resp_{marker}", completed
|
||||
assert str(completed[0]["response"]["id"]).startswith("resp_"), completed
|
||||
assert rig.upstream_bodies(marker) == (
|
||||
{"model": "gpt-4o-mini", "input": marker, "metadata": {"session_id": session}, "stream": True},
|
||||
)
|
||||
assert rig.sink.single_span(f"resp_{marker}") == session
|
||||
assert rig.spend_session(rig.sink.logged_id(f"resp_{marker}")) == session
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_from_metadata_chat_raw")
|
||||
def test_chat_raw_metadata_session_id_lands_as_conversation_id(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = "conv-" + uuid.uuid4().hex
|
||||
response: Final = rig.chat(marker, metadata={"session_id": session})
|
||||
assert response.status_code == 200, response.text
|
||||
identity: Final = response.json()["id"]
|
||||
assert identity == f"chatcmpl-{marker}", response.text
|
||||
assert rig.sink.single_span(identity) == session
|
||||
assert rig.spend_session(identity) == session
|
||||
_assert_upstream_clean(rig, marker, session)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_non_string_session_ids_match_spend_row")
|
||||
def test_integer_and_list_litellm_session_id_match_the_spend_row_or_are_dropped_together(rig: Rig) -> None:
|
||||
for odd in (123, ["a", "b"]):
|
||||
marker: Final = _marker()
|
||||
response: Final = rig.chat(marker, litellm_session_id=odd)
|
||||
assert response.status_code == 200, response.text
|
||||
identity: Final = response.json()["id"]
|
||||
assert rig.sink.single_span(identity) == rig.spend_session(identity), (odd, rig.sink.conversation_ids(identity))
|
||||
assert len(rig.upstream_bodies(marker)) == 1
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_empty_string_session_id_is_omitted")
|
||||
def test_empty_string_litellm_session_id_leaves_the_span_without_a_conversation_id(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
response: Final = rig.chat(marker, litellm_session_id="")
|
||||
assert response.status_code == 200, response.text
|
||||
identity: Final = response.json()["id"]
|
||||
assert rig.sink.single_span(identity) is None
|
||||
assert rig.spend_session(identity), response.text
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_five_kilobyte_header_round_trips")
|
||||
def test_five_kilobyte_session_header_round_trips_to_the_span_and_the_spend_row(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = ("s" * 5000) + uuid.uuid4().hex
|
||||
response: Final = rig.chat(marker, headers={"x-litellm-session-id": session})
|
||||
assert response.status_code == 200, response.text
|
||||
identity: Final = response.json()["id"]
|
||||
assert rig.sink.single_span(identity) == session
|
||||
assert rig.spend_session(identity) == session
|
||||
_assert_upstream_clean(rig, marker, session)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_duplicate_header_lands_once")
|
||||
def test_duplicate_session_header_lands_once_and_unchanged(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = "conv-" + uuid.uuid4().hex
|
||||
response: Final = rig.proxy.client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"model": rig.model, "messages": [{"role": "user", "content": marker}], "cache": {"no-cache": True}},
|
||||
headers=[
|
||||
("Authorization", f"Bearer {rig.proxy.key}"),
|
||||
("x-litellm-session-id", session),
|
||||
("x-litellm-session-id", session),
|
||||
],
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
identity: Final = response.json()["id"]
|
||||
assert rig.sink.single_span(identity) == session
|
||||
assert rig.spend_session(identity) == session
|
||||
_assert_upstream_clean(rig, marker, session)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_unauthenticated_request_leaves_no_span")
|
||||
def test_unauthenticated_request_with_session_header_is_rejected_and_leaves_no_span(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
response: Final = rig.chat(marker, headers={"x-litellm-session-id": "conv-" + uuid.uuid4().hex}, key="sk-wrong")
|
||||
assert response.status_code == 401, response.text
|
||||
assert rig.upstream_bodies(marker) == ()
|
||||
control: Final = rig.chat(marker)
|
||||
assert control.status_code == 200, control.text
|
||||
assert rig.sink.single_span(control.json()["id"]) is None
|
||||
assert rig.spend_session(control.json()["id"]), control.text
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_survives_sink_rejection")
|
||||
def test_sink_rejecting_with_403_drops_those_spans_and_later_spans_still_land(rig: Rig) -> None:
|
||||
rig.sink.rejection.set()
|
||||
try:
|
||||
rejected: Final = rig.chat(_marker(), headers={"x-litellm-session-id": "conv-rejected"})
|
||||
assert rejected.status_code == 200, rejected.text
|
||||
eventually(lambda: any(request.body for request in rig.sink.wire.drain()), lambda seen: seen, seconds=30)
|
||||
finally:
|
||||
rig.sink.rejection.clear()
|
||||
marker: Final = _marker()
|
||||
session: Final = "conv-" + uuid.uuid4().hex
|
||||
response: Final = rig.chat(marker, headers={"x-litellm-session-id": session})
|
||||
assert response.status_code == 200, response.text
|
||||
assert rig.sink.single_span(response.json()["id"]) == session
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_absent_without_caller_session")
|
||||
def test_request_without_any_session_input_has_no_conversation_id(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
response: Final = rig.chat(marker)
|
||||
assert response.status_code == 200, response.text
|
||||
identity: Final = response.json()["id"]
|
||||
assert rig.sink.single_span(identity) is None
|
||||
assert rig.spend_session(identity), response.text
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_ignores_generated_session_id")
|
||||
def test_generate_policy_minted_session_id_reaches_the_spend_row_but_not_the_span(generating_rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
response: Final = generating_rig.chat(marker)
|
||||
assert response.status_code == 200, response.text
|
||||
identity: Final = response.json()["id"]
|
||||
minted: Final = generating_rig.spend_session(identity)
|
||||
assert minted, response.text
|
||||
assert generating_rig.sink.single_span(identity) is None
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_langfuse_header_wins_over_generated")
|
||||
def test_generate_policy_keeps_the_langfuse_session_header_as_conversation_id(generating_rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = "conv-" + uuid.uuid4().hex
|
||||
response: Final = generating_rig.chat(marker, headers={"langfuse_session_id": session})
|
||||
assert response.status_code == 200, response.text
|
||||
assert generating_rig.sink.single_span(response.json()["id"]) == session
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_header_precedence_matches_spend_row")
|
||||
def test_header_body_and_metadata_session_ids_resolve_to_the_same_id_as_the_spend_row(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
header: Final = "conv-header-" + uuid.uuid4().hex
|
||||
response: Final = rig.chat(
|
||||
marker,
|
||||
headers={"x-litellm-session-id": header},
|
||||
litellm_session_id="conv-body-" + uuid.uuid4().hex,
|
||||
metadata={"session_id": "conv-meta-" + uuid.uuid4().hex},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
identity: Final = response.json()["id"]
|
||||
assert rig.sink.single_span(identity) == header
|
||||
assert rig.spend_session(identity) == header
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_repeated_requests_log_once_each")
|
||||
def test_three_identical_requests_produce_one_span_each_with_the_same_conversation_id(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
session: Final = "conv-" + uuid.uuid4().hex
|
||||
responses: Final = tuple(rig.chat(marker, headers={"x-litellm-session-id": session}) for _ in range(3))
|
||||
assert all(response.status_code == 200 for response in responses), [response.text for response in responses]
|
||||
identity: Final = f"chatcmpl-{marker}"
|
||||
spans: Final = eventually(lambda: rig.sink.conversation_ids(identity), lambda values: len(values) == 3, seconds=30)
|
||||
assert spans == (session, session, session), spans
|
||||
assert len(rig.upstream_bodies(marker)) == 3
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_ignores_trace_id_backfill")
|
||||
def test_metadata_trace_id_alone_fills_the_spend_row_but_not_the_span(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
trace: Final = "trace-" + uuid.uuid4().hex
|
||||
response: Final = rig.chat(marker, metadata={"trace_id": trace})
|
||||
assert response.status_code == 200, response.text
|
||||
identity: Final = response.json()["id"]
|
||||
assert rig.sink.single_span(identity) is None
|
||||
assert rig.spend_session(identity) == trace
|
||||
|
||||
|
||||
def _chat_id(response: httpx.Response) -> str:
|
||||
if not response.headers.get("content-type", "").startswith("text/event-stream"):
|
||||
return response.json()["id"]
|
||||
identities: Final = frozenset(str(event["id"]) for event in _sse_events(response.text))
|
||||
assert len(identities) == 1, response.text
|
||||
return next(iter(identities))
|
||||
|
||||
|
||||
def _responses_id(response: httpx.Response) -> str:
|
||||
if not response.headers.get("content-type", "").startswith("text/event-stream"):
|
||||
return response.json()["id"]
|
||||
completed: Final = tuple(
|
||||
event["response"]["id"] for event in _sse_events(response.text) if event.get("type") == "response.completed"
|
||||
)
|
||||
assert len(completed) == 1, response.text
|
||||
return str(completed[0])
|
||||
|
||||
|
||||
def _message_id(response: httpx.Response) -> str:
|
||||
if not response.headers.get("content-type", "").startswith("text/event-stream"):
|
||||
return response.json()["id"]
|
||||
starts: Final = tuple(
|
||||
event["message"]["id"] for event in _sse_events(response.text) if event.get("type") == "message_start"
|
||||
)
|
||||
assert len(starts) == 1, response.text
|
||||
return starts[0]
|
||||
|
||||
|
||||
def _burst(rig: Rig, count: int, session_for: Mapping[int, str]) -> tuple[tuple[int, str, str | None], ...]:
|
||||
markers: Final = tuple(_marker() for _ in range(count))
|
||||
|
||||
def one(index: int) -> tuple[int, str, str | None]:
|
||||
marker: Final = markers[index]
|
||||
headers: Final = {"Authorization": f"Bearer {rig.proxy.key}", "x-litellm-session-id": session_for[index]}
|
||||
route: Final = index % 3
|
||||
try:
|
||||
if route == 0:
|
||||
response: Final = rig.proxy.client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"model": rig.model,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
"stream": index % 2 == 0,
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
response.read()
|
||||
if response.status_code != 200:
|
||||
return index, marker, response.text
|
||||
return index, _chat_id(response), None
|
||||
if route == 1:
|
||||
response = rig.proxy.client.post(
|
||||
"/v1/responses",
|
||||
json={"model": rig.model, "input": marker, "stream": index % 2 == 0},
|
||||
headers=headers,
|
||||
)
|
||||
response.read()
|
||||
if response.status_code != 200:
|
||||
return index, marker, response.text
|
||||
return index, _responses_id(response), None
|
||||
response = rig.proxy.client.post(
|
||||
"/v1/messages",
|
||||
json={
|
||||
"model": rig.model,
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
"stream": index % 2 == 0,
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
response.read()
|
||||
if response.status_code != 200:
|
||||
return index, marker, response.text
|
||||
return index, _message_id(response), None
|
||||
except httpx.HTTPError as error:
|
||||
return index, marker, repr(error)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=10) as pool:
|
||||
return tuple(pool.map(one, range(count)))
|
||||
|
||||
|
||||
def _is_encrypted_responses_id(identity: str) -> bool:
|
||||
return identity.startswith("resp_") and _decoded_responses_id(identity) == identity
|
||||
|
||||
|
||||
def _landed(rig: Rig, expected: Mapping[str, str]) -> dict[str, tuple[str, ...]]:
|
||||
spans: Final = rig.sink.attributes()
|
||||
return {
|
||||
session: tuple(
|
||||
_canonical_id(str(attributes["gen_ai.response.id"]["stringValue"]))
|
||||
for attributes in spans
|
||||
if attributes.get(CONVERSATION, {}).get("stringValue") == session and "gen_ai.response.id" in attributes
|
||||
)
|
||||
for session in expected.values()
|
||||
}
|
||||
|
||||
|
||||
def _assert_exactly_once(rig: Rig, expected: Mapping[str, str], landed: Mapping[str, tuple[str, ...]]) -> None:
|
||||
spend: Final = eventually(
|
||||
lambda: {
|
||||
session: tuple(_canonical_id(identity) for identity in rig.spend_request_ids(session))
|
||||
for session in expected.values()
|
||||
},
|
||||
lambda rows: all(len(values) >= 1 for values in rows.values()),
|
||||
seconds=70,
|
||||
)
|
||||
assert landed == spend, (landed, spend)
|
||||
assert all(len(values) == 1 for values in landed.values()), landed
|
||||
caller_visible: Final = {
|
||||
session: (_canonical_id(identity),)
|
||||
for identity, session in expected.items()
|
||||
if not _is_encrypted_responses_id(identity)
|
||||
}
|
||||
assert {session: landed[session] for session in caller_visible} == caller_visible, landed
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_sink_outage_recovers_exactly_once")
|
||||
def test_sink_outage_during_a_mixed_burst_lands_every_response_exactly_once_after_recovery(rig: Rig) -> None:
|
||||
sessions: Final = {index: f"conv-{index}-{uuid.uuid4().hex}" for index in range(30)}
|
||||
rig.sink.outage.set()
|
||||
try:
|
||||
health_down: Final = rig.proxy.request("GET", "/health/services", params={"service": "otel"})
|
||||
results: Final = _burst(rig, 30, sessions)
|
||||
assert all(error is None for _, _, error in results), [error for _, _, error in results if error]
|
||||
eventually(lambda: any(True for _ in rig.sink.wire.drain()), lambda seen: seen, seconds=30)
|
||||
finally:
|
||||
rig.sink.outage.clear()
|
||||
assert health_down.status_code == 200, health_down.text
|
||||
expected: Final = {identity: sessions[index] for index, identity, _ in results}
|
||||
landed: Final = eventually(
|
||||
lambda: _landed(rig, expected), lambda seen: all(len(values) >= 1 for values in seen.values()), seconds=80
|
||||
)
|
||||
_assert_exactly_once(rig, expected, landed)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_slow_sink_no_duplicates")
|
||||
def test_slow_sink_during_a_burst_lands_every_response_exactly_once(rig: Rig) -> None:
|
||||
sessions: Final = {index: f"conv-{index}-{uuid.uuid4().hex}" for index in range(20)}
|
||||
rig.sink.slow.set()
|
||||
try:
|
||||
results: Final = _burst(rig, 20, sessions)
|
||||
assert all(error is None for _, _, error in results), [error for _, _, error in results if error]
|
||||
expected: Final = {identity: sessions[index] for index, identity, _ in results}
|
||||
landed: Final = eventually(
|
||||
lambda: _landed(rig, expected), lambda seen: all(len(values) >= 1 for values in seen.values()), seconds=80
|
||||
)
|
||||
finally:
|
||||
rig.sink.slow.clear()
|
||||
_assert_exactly_once(rig, expected, landed)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_survives_worker_kill")
|
||||
def test_killing_one_of_two_workers_mid_burst_keeps_serving_and_never_duplicates_a_span(two_worker_rig: Rig) -> None:
|
||||
rig: Final = two_worker_rig
|
||||
root: Final = psutil.Process(rig.process.process.pid)
|
||||
workers: Final = eventually(
|
||||
lambda: tuple(child for child in root.children() if "resource_tracker" not in " ".join(child.cmdline())),
|
||||
lambda found: len(found) == 2,
|
||||
seconds=30,
|
||||
)
|
||||
sessions: Final = {index: f"conv-{index}-{uuid.uuid4().hex}" for index in range(24)}
|
||||
markers: Final = tuple(_marker() for _ in range(24))
|
||||
|
||||
def one(index: int) -> tuple[str, str | None]:
|
||||
if index == 8:
|
||||
os.kill(workers[0].pid, signal.SIGKILL)
|
||||
try:
|
||||
response: Final = rig.chat(markers[index], headers={"x-litellm-session-id": sessions[index]})
|
||||
return f"chatcmpl-{markers[index]}", None if response.status_code == 200 else response.text
|
||||
except httpx.HTTPError as error:
|
||||
return f"chatcmpl-{markers[index]}", repr(error)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=6) as pool:
|
||||
results: Final = tuple(pool.map(one, range(24)))
|
||||
assert rig.process.process.poll() is None, "Proxy root exited after a worker was killed"
|
||||
after: Final = rig.chat(_marker(), headers={"x-litellm-session-id": "conv-after-kill"})
|
||||
assert after.status_code == 200, after.text
|
||||
assert rig.sink.single_span(after.json()["id"]) == "conv-after-kill"
|
||||
failures: Final = tuple(error for _, error in results if error)
|
||||
assert all(error.startswith(("ReadError(", "RemoteProtocolError(", "ConnectError(")) for error in failures), (
|
||||
failures
|
||||
)
|
||||
assert len(failures) <= 6, failures
|
||||
served: Final = {identity: sessions[index] for index, (identity, error) in enumerate(results) if error is None}
|
||||
assert len(served) >= 18, results
|
||||
settled: Final = {
|
||||
identity: sessions[index] for index, (identity, error) in enumerate(results) if index > 14 and not error
|
||||
}
|
||||
landed: Final = eventually(
|
||||
lambda: _landed(rig, settled), lambda seen: all(len(values) >= 1 for values in seen.values()), seconds=60
|
||||
)
|
||||
_assert_exactly_once(rig, settled, landed)
|
||||
assert all(len(values) <= 1 for values in _landed(rig, served).values()), _landed(rig, served)
|
||||
lost: Final = _landed(rig, {identity: sessions[index] for index, (identity, error) in enumerate(results) if error})
|
||||
assert all(values == () for values in lost.values()), lost
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.otel.conversation_id_flushes_on_shutdown")
|
||||
def test_terminating_the_proxy_right_after_a_burst_flushes_every_span_before_exit(
|
||||
provider: Wire, collector: Collector, tmp_path_factory: pytest.TempPathFactory
|
||||
) -> None:
|
||||
factory: Final = RigFactory(provider, collector, tmp_path_factory.mktemp("otel-shutdown"), {}, 2)
|
||||
started: Final = factory.start()
|
||||
rig: Final = next(started)
|
||||
sessions: Final = {index: f"conv-{index}-{uuid.uuid4().hex}" for index in range(10)}
|
||||
markers: Final = tuple(_marker() for _ in range(10))
|
||||
responses: Final = tuple(
|
||||
rig.chat(markers[index], headers={"x-litellm-session-id": sessions[index]}) for index in range(10)
|
||||
)
|
||||
assert all(response.status_code == 200 for response in responses), [response.text for response in responses]
|
||||
expected: Final = {f"chatcmpl-{markers[index]}": sessions[index] for index in range(10)}
|
||||
drained: Final = eventually(
|
||||
lambda: _landed(rig, expected), lambda seen: all(len(values) >= 1 for values in seen.values()), seconds=60
|
||||
)
|
||||
assert drained == {session: (identity,) for identity, session in expected.items()}, drained
|
||||
rig.process.process.terminate()
|
||||
assert rig.process.process.wait(timeout=40) in (0, -signal.SIGTERM)
|
||||
assert _landed(rig, expected) == drained
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
next(started)
|
||||
|
|
@ -23,6 +23,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E4
|
|||
from opentelemetry.trace import SpanKind # noqa: E402
|
||||
from opentelemetry.trace.status import StatusCode # noqa: E402
|
||||
|
||||
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY # noqa: E402
|
||||
from litellm.integrations.otel import ( # noqa: E402
|
||||
GenAI,
|
||||
LiteLLM,
|
||||
|
|
@ -175,6 +176,64 @@ def test_async_log_success_event_emits_llm_call_span():
|
|||
assert span.status.status_code is StatusCode.UNSET
|
||||
|
||||
|
||||
def test_llm_call_span_carries_the_callers_conversation_id():
|
||||
logger, exporter = _logger()
|
||||
kwargs = {**_kwargs(), "litellm_params": {"litellm_session_id": "conv-42", "metadata": {}}}
|
||||
_emit_llm(logger, kwargs)
|
||||
(span,) = exporter.get_finished_spans()
|
||||
assert span.attributes[GenAI.CONVERSATION_ID] == "conv-42"
|
||||
|
||||
|
||||
def test_llm_call_span_without_a_caller_session_has_no_conversation_id():
|
||||
"""The proxy stamps ``metadata.trace_id`` with the OTel trace id and
|
||||
``get_litellm_params`` back-fills ``litellm_session_id`` from it."""
|
||||
logger, exporter = _logger()
|
||||
otel_trace_id = "6ca5745ef6780d958f62925747f7a5ee"
|
||||
kwargs = {
|
||||
**_kwargs(payload=_payload(trace_id=otel_trace_id)),
|
||||
"litellm_trace_id": otel_trace_id,
|
||||
"litellm_params": {
|
||||
"litellm_session_id": otel_trace_id,
|
||||
"litellm_trace_id": otel_trace_id,
|
||||
"metadata": {"trace_id": otel_trace_id},
|
||||
},
|
||||
}
|
||||
_emit_llm(logger, kwargs)
|
||||
(span,) = exporter.get_finished_spans()
|
||||
assert GenAI.CONVERSATION_ID not in span.attributes
|
||||
|
||||
|
||||
def test_llm_call_span_keeps_the_header_session_when_the_proxy_generated_a_body_one():
|
||||
"""``missing_session_id: generate`` mints a body session and marks it, but the
|
||||
caller's ``langfuse_session_id`` header is still their conversation."""
|
||||
logger, exporter = _logger()
|
||||
kwargs = {
|
||||
**_kwargs(),
|
||||
"litellm_params": {
|
||||
"litellm_session_id": "minted-by-proxy",
|
||||
"metadata": {"session_id": "minted-by-proxy", SESSION_ID_GENERATED_METADATA_KEY: True},
|
||||
"proxy_server_request": {"headers": {"langfuse_session_id": "conv-header"}},
|
||||
},
|
||||
}
|
||||
_emit_llm(logger, kwargs)
|
||||
(span,) = exporter.get_finished_spans()
|
||||
assert span.attributes[GenAI.CONVERSATION_ID] == "conv-header"
|
||||
|
||||
|
||||
def test_replayed_llm_call_span_does_not_take_the_payloads_session_id():
|
||||
"""``/callback_logs`` replays a finished payload whose ``litellm_params`` hold
|
||||
only key metadata; a session minted under ``missing_session_id: generate``
|
||||
lands there without its marker, so ``payload.session_id`` is never trusted."""
|
||||
logger, exporter = _logger()
|
||||
kwargs = {
|
||||
**_kwargs(payload=_payload(session_id="minted-then-replayed", trace_id="minted-then-replayed")),
|
||||
"litellm_params": {"metadata": {"user_api_key_hash": "hsh"}},
|
||||
}
|
||||
_emit_llm(logger, kwargs)
|
||||
(span,) = exporter.get_finished_spans()
|
||||
assert GenAI.CONVERSATION_ID not in span.attributes
|
||||
|
||||
|
||||
def test_streaming_span_carries_time_to_first_chunk():
|
||||
logger, exporter = _logger()
|
||||
kwargs = {
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from typing import Final
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
|
||||
from litellm.integrations.otel import (
|
||||
BAGGAGE_PROMOTED_KEYS,
|
||||
DB,
|
||||
|
|
@ -978,6 +979,148 @@ def test_llm_span_data_carries_the_caller_trace_controls():
|
|||
assert LLMCallSpanData.from_standard_logging_payload(_sample_payload()).trace == TraceControls()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("litellm_params", "expected"),
|
||||
[
|
||||
({"litellm_session_id": "conv-body"}, "conv-body"),
|
||||
({"metadata": {"session_id": "conv-meta"}}, "conv-meta"),
|
||||
({"litellm_metadata": {"session_id": "conv-anthropic"}}, "conv-anthropic"),
|
||||
({"proxy_server_request": {"headers": {"langfuse_session_id": "conv-header"}}}, "conv-header"),
|
||||
({"litellm_session_id": "conv-body", "metadata": {"session_id": "conv-meta"}}, "conv-body"),
|
||||
({"litellm_session_id": "", "metadata": {"session_id": ""}}, None),
|
||||
({"litellm_trace_id": "trace-only", "metadata": {"trace_id": "trace-only"}}, None),
|
||||
(
|
||||
{
|
||||
"litellm_session_id": "0" * 32,
|
||||
"litellm_trace_id": "0" * 32,
|
||||
"metadata": {"trace_id": "0" * 32},
|
||||
},
|
||||
None,
|
||||
),
|
||||
(
|
||||
{
|
||||
"litellm_session_id": "0" * 32,
|
||||
"litellm_trace_id": "0" * 32,
|
||||
"metadata": {"trace_id": "0" * 32},
|
||||
"proxy_server_request": {"headers": {"langfuse_session_id": "conv-header"}},
|
||||
},
|
||||
"conv-header",
|
||||
),
|
||||
(
|
||||
{
|
||||
"litellm_session_id": "minted-by-proxy",
|
||||
"metadata": {"session_id": "minted-by-proxy", SESSION_ID_GENERATED_METADATA_KEY: True},
|
||||
},
|
||||
None,
|
||||
),
|
||||
(
|
||||
{
|
||||
"litellm_session_id": "minted-by-proxy",
|
||||
"metadata": {"session_id": "minted-by-proxy", SESSION_ID_GENERATED_METADATA_KEY: True},
|
||||
"proxy_server_request": {"headers": {"langfuse_session_id": "conv-header"}},
|
||||
},
|
||||
"conv-header",
|
||||
),
|
||||
(
|
||||
{
|
||||
"litellm_session_id": "minted-by-proxy",
|
||||
"metadata": {"session_id": "conv-other-key"},
|
||||
"litellm_metadata": {"session_id": "minted-by-proxy", SESSION_ID_GENERATED_METADATA_KEY: True},
|
||||
},
|
||||
"conv-other-key",
|
||||
),
|
||||
(
|
||||
{
|
||||
"litellm_session_id": "minted-by-proxy",
|
||||
"metadata": {"session_id": "minted-by-proxy", SESSION_ID_GENERATED_METADATA_KEY: True},
|
||||
"litellm_metadata": {"session_id": "conv-other-key"},
|
||||
},
|
||||
"conv-other-key",
|
||||
),
|
||||
(
|
||||
{
|
||||
"litellm_session_id": "conv-x-header",
|
||||
"litellm_trace_id": "conv-x-header",
|
||||
"metadata": {"trace_id": "conv-x-header", "session_id": "conv-x-header"},
|
||||
},
|
||||
"conv-x-header",
|
||||
),
|
||||
({}, None),
|
||||
],
|
||||
ids=[
|
||||
"litellm_session_id",
|
||||
"metadata",
|
||||
"anthropic-metadata",
|
||||
"langfuse-header",
|
||||
"litellm_session_id-beats-metadata",
|
||||
"blank-values",
|
||||
"trace-id-is-not-a-session",
|
||||
"backfilled-from-otel-trace-id-is-not-a-conversation",
|
||||
"backfilled-trace-id-does-not-shadow-the-header",
|
||||
"proxy-generated-is-not-a-conversation",
|
||||
"proxy-generated-does-not-shadow-the-header",
|
||||
"proxy-generated-on-litellm_metadata-does-not-shadow-metadata",
|
||||
"proxy-generated-on-metadata-does-not-shadow-litellm_metadata",
|
||||
"x-litellm-session-id-header-sets-trace-and-session",
|
||||
"empty",
|
||||
],
|
||||
)
|
||||
def test_llm_call_event_resolves_the_callers_conversation_id(litellm_params, expected):
|
||||
kwargs: Final = {"litellm_params": litellm_params, "litellm_trace_id": "per-request-uuid"}
|
||||
assert LLMCallEvent.from_dict(kwargs).session_id == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("litellm_params", "payload", "expected"),
|
||||
[
|
||||
(
|
||||
{"metadata": {"user_api_key_hash": "hsh"}},
|
||||
{"session_id": "minted-then-replayed", "trace_id": "minted-then-replayed"},
|
||||
None,
|
||||
),
|
||||
(
|
||||
{"metadata": {"user_api_key_hash": "hsh"}},
|
||||
{"session_id": "conv-replayed", "trace_id": "0af7651916cd43dd8448eb211c80319c"},
|
||||
None,
|
||||
),
|
||||
({"litellm_session_id": "conv-live"}, {"session_id": "conv-replayed"}, "conv-live"),
|
||||
(
|
||||
{
|
||||
"litellm_session_id": "minted-by-proxy",
|
||||
"metadata": {"session_id": "minted-by-proxy", SESSION_ID_GENERATED_METADATA_KEY: True},
|
||||
},
|
||||
{"session_id": "minted-by-proxy"},
|
||||
None,
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"replayed-minted-session-stays-hidden",
|
||||
"replayed-payload-is-not-a-source",
|
||||
"live-params-win",
|
||||
"generated-stays-hidden",
|
||||
],
|
||||
)
|
||||
def test_llm_call_event_never_reads_the_replayed_payloads_session_id(litellm_params, payload, expected):
|
||||
"""``/callback_logs`` rebuilds ``litellm_params`` with key metadata only, so a
|
||||
``StandardLoggingPayload`` minted under ``missing_session_id: generate`` arrives
|
||||
without its generated marker and is indistinguishable from a caller's session;
|
||||
the payload is therefore never a source for the conversation id."""
|
||||
kwargs: Final = {
|
||||
"litellm_params": litellm_params,
|
||||
"standard_logging_object": _sample_payload(**payload),
|
||||
}
|
||||
assert LLMCallEvent.from_dict(kwargs).session_id == expected
|
||||
|
||||
|
||||
def test_llm_span_stamps_gen_ai_conversation_id_only_when_the_caller_sent_one():
|
||||
with_session: Final = LLMCallSpanData.from_standard_logging_payload(_sample_payload(), session_id="conv-1")
|
||||
assert GenAIMapper().map(with_session)[GenAI.CONVERSATION_ID] == "conv-1"
|
||||
|
||||
without: Final = LLMCallSpanData.from_standard_logging_payload(_sample_payload(trace_id="per-request-uuid"))
|
||||
assert without.session_id is None
|
||||
assert GenAI.CONVERSATION_ID not in GenAIMapper().map(without)
|
||||
|
||||
|
||||
def test_llm_span_carries_proxy_request_route():
|
||||
"""The LLM span records the proxy route the request arrived on, so it can be
|
||||
filtered by endpoint (``/v1/responses`` vs ``/v1/chat/completions``) without
|
||||
|
|
|
|||
|
|
@ -28,10 +28,12 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
|
||||
from litellm.proxy.auth.auth_checks import TeamNotFoundError
|
||||
from litellm.proxy.auth.handle_jwt import (
|
||||
JWKS_FETCH_ATTEMPTS,
|
||||
STALE_CACHE_KEY_PREFIX,
|
||||
STALE_WRITTEN_AT_CACHE_KEY_PREFIX,
|
||||
HeaderTeam,
|
||||
JWKSUnreachableError,
|
||||
JWTAuthManager,
|
||||
JWTHandler,
|
||||
|
|
@ -1893,29 +1895,41 @@ async def test_auth_builder_oidc_enabled_falls_back_to_jwt_auth_for_jwt_tokens()
|
|||
assert result["user_object"] == user_object
|
||||
|
||||
|
||||
def test_get_team_id_from_header():
|
||||
"""Test get_team_id_from_header returns team when valid, None when missing, raises on invalid."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
# Valid team in allowed list
|
||||
result = JWTAuthManager.get_team_id_from_header(
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_returns_allowed_id_none_without_header_and_403_on_invalid():
|
||||
"""Without a DB, x-litellm-team-id resolves to the team when it names an allowed
|
||||
team id, to None when the header is absent, and to a 403 for any other value."""
|
||||
allowed = await JWTAuthManager.resolve_team_from_header(
|
||||
request_headers={"x-litellm-team-id": "team-1"},
|
||||
allowed_team_ids={"team-1", "team-2"},
|
||||
fallback_to_db_teams=False,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert result == "team-1"
|
||||
assert allowed == HeaderTeam(header_value="team-1", team_id="team-1")
|
||||
|
||||
# No header returns None
|
||||
result = JWTAuthManager.get_team_id_from_header(
|
||||
absent = await JWTAuthManager.resolve_team_from_header(
|
||||
request_headers={"authorization": "Bearer token"},
|
||||
allowed_team_ids={"team-1"},
|
||||
fallback_to_db_teams=False,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert result is None
|
||||
assert absent is None
|
||||
|
||||
# Invalid team raises 403
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
JWTAuthManager.get_team_id_from_header(
|
||||
await JWTAuthManager.resolve_team_from_header(
|
||||
request_headers={"x-litellm-team-id": "invalid-team"},
|
||||
allowed_team_ids={"team-1", "team-2"},
|
||||
fallback_to_db_teams=False,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
|
|
@ -5205,32 +5219,29 @@ def test_build_decode_kwargs_warns_for_unscoped_global_fallback_in_mixed_deploym
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_get_team_id_from_header_defers_to_db_membership_only_without_jwt_claims():
|
||||
"""With fallback_to_db_teams=True, an x-litellm-team-id header is accepted
|
||||
provisionally only when the JWT carries no team claims (allowed set empty).
|
||||
When the JWT does carry team claims, the header must still be validated
|
||||
against them, and the flag-off behavior must keep rejecting unknown teams."""
|
||||
deferred = JWTAuthManager.get_team_id_from_header(
|
||||
request_headers={"x-litellm-team-id": "team-from-db"},
|
||||
allowed_team_ids=set(),
|
||||
fallback_to_db_teams=True,
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_accepts_db_teams_provisionally_under_fallback_even_with_jwt_claims():
|
||||
"""With fallback_to_db_teams=True, an x-litellm-team-id header naming an existing
|
||||
team is accepted provisionally whether or not the JWT carries team claims; the
|
||||
union of JWT teams and DB memberships is enforced by auth_builder's later
|
||||
membership check. Unknown values still 403, and the flag-off behavior keeps
|
||||
rejecting teams outside the JWT's allowed set."""
|
||||
known_ids = frozenset({"team-from-db"})
|
||||
|
||||
deferred, _, _ = await _resolve_header("team-from-db", set(), True, _teams_by_id(known_ids), _team_alias_lookup_404)
|
||||
assert deferred == HeaderTeam(header_value="team-from-db", team_id="team-from-db")
|
||||
|
||||
deferred_with_claims, _, _ = await _resolve_header(
|
||||
"team-from-db", {"team-1"}, True, _teams_by_id(known_ids), _team_alias_lookup_404
|
||||
)
|
||||
assert deferred == "team-from-db"
|
||||
assert deferred_with_claims == HeaderTeam(header_value="team-from-db", team_id="team-from-db")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
JWTAuthManager.get_team_id_from_header(
|
||||
request_headers={"x-litellm-team-id": "team-x"},
|
||||
allowed_team_ids={"team-1", "team-2"},
|
||||
fallback_to_db_teams=True,
|
||||
)
|
||||
await _resolve_header("team-x", {"team-1", "team-2"}, True, _teams_by_id(known_ids), _team_alias_lookup_404)
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
JWTAuthManager.get_team_id_from_header(
|
||||
request_headers={"x-litellm-team-id": "team-from-db"},
|
||||
allowed_team_ids=set(),
|
||||
fallback_to_db_teams=False,
|
||||
)
|
||||
await _resolve_header("team-from-db", set(), False, _teams_by_id(known_ids), _team_alias_lookup_404)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -5684,6 +5695,7 @@ def test_validate_header_team_in_db_membership_does_not_leak_team_ids():
|
|||
JWTAuthManager._validate_header_team_in_db_membership(
|
||||
team_id="outsider_team",
|
||||
user_object=user_object,
|
||||
header_value="outsider_team",
|
||||
)
|
||||
|
||||
detail = exc_info.value.detail
|
||||
|
|
@ -5693,6 +5705,43 @@ def test_validate_header_team_in_db_membership_does_not_leak_team_ids():
|
|||
assert "outsider_team" in detail
|
||||
|
||||
|
||||
async def _team_lookup_404(team_id, **kwargs):
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call.",
|
||||
)
|
||||
|
||||
|
||||
async def _team_alias_lookup_404(team_alias, **kwargs):
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Team with alias '{team_alias}' doesn't exist in db. Create team via `/team/new` call."},
|
||||
)
|
||||
|
||||
|
||||
def _teams_by_alias(aliases: Mapping[str, str]):
|
||||
"""An alias lookup over `aliases` (alias -> team_id) that 404s like the real one otherwise."""
|
||||
|
||||
async def lookup(team_alias, **kwargs):
|
||||
if team_alias not in aliases:
|
||||
return await _team_alias_lookup_404(team_alias)
|
||||
return LiteLLM_TeamTable(team_id=aliases[team_alias], team_alias=team_alias)
|
||||
|
||||
return lookup
|
||||
|
||||
|
||||
def _teams_by_id(team_ids: frozenset[str]):
|
||||
"""A team lookup that knows exactly `team_ids` and, like the real one, reports
|
||||
any other id as provably absent."""
|
||||
|
||||
async def lookup(team_id, **kwargs):
|
||||
if team_id not in team_ids:
|
||||
raise TeamNotFoundError(team_id=team_id)
|
||||
return LiteLLM_TeamTable(team_id=team_id)
|
||||
|
||||
return lookup
|
||||
|
||||
|
||||
async def _run_auth_builder_with_header_team(
|
||||
jwt_auth_config: LiteLLM_JWTAuth,
|
||||
token: dict,
|
||||
|
|
@ -5700,6 +5749,9 @@ async def _run_auth_builder_with_header_team(
|
|||
user_object: LiteLLM_UserTable,
|
||||
fake_get_team,
|
||||
allowed_team_ids: set,
|
||||
fake_get_team_by_alias=_team_alias_lookup_404,
|
||||
route: str = "/chat/completions",
|
||||
send_header: bool = True,
|
||||
):
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.litellm_jwtauth = jwt_auth_config
|
||||
|
|
@ -5744,28 +5796,26 @@ async def _run_auth_builder_with_header_team(
|
|||
new_callable=AsyncMock,
|
||||
side_effect=fake_get_team,
|
||||
),
|
||||
patch( # test-quality-ok: handle_jwt calls the team lookups as module globals with no injection seam
|
||||
"litellm.proxy.auth.handle_jwt.get_team_object_by_alias",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=fake_get_team_by_alias,
|
||||
),
|
||||
):
|
||||
return await JWTAuthManager.auth_builder(
|
||||
api_key="test_jwt_token",
|
||||
jwt_handler=jwt_handler,
|
||||
request_data={"model": "gpt-4"},
|
||||
general_settings={"enforce_rbac": False},
|
||||
route="/chat/completions",
|
||||
prisma_client=None,
|
||||
route=route,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=None,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
request_headers={"x-litellm-team-id": header_team_id},
|
||||
request_headers={"x-litellm-team-id": header_team_id} if send_header else {},
|
||||
)
|
||||
|
||||
|
||||
async def _team_lookup_404(team_id, **kwargs):
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_header_team_not_found_matches_non_membership_denial() -> (
|
||||
None
|
||||
|
|
@ -6754,6 +6804,428 @@ async def test_auth_builder_header_team_enforces_team_allowed_routes_under_db_fa
|
|||
assert result["team_id"] == header_team
|
||||
|
||||
|
||||
async def _resolve_header(
|
||||
header_value: str,
|
||||
allowed_team_ids: set[str],
|
||||
fallback_to_db_teams: bool,
|
||||
fake_get_team,
|
||||
fake_get_team_by_alias,
|
||||
) -> tuple[HeaderTeam | None, AsyncMock, AsyncMock]:
|
||||
with (
|
||||
patch( # test-quality-ok: handle_jwt calls the team lookups as module globals with no injection seam
|
||||
"litellm.proxy.auth.handle_jwt.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=fake_get_team,
|
||||
) as by_id,
|
||||
patch( # test-quality-ok: handle_jwt calls the team lookups as module globals with no injection seam
|
||||
"litellm.proxy.auth.handle_jwt.get_team_object_by_alias",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=fake_get_team_by_alias,
|
||||
) as by_alias,
|
||||
):
|
||||
resolved = await JWTAuthManager.resolve_team_from_header(
|
||||
request_headers={"X-LiteLLM-Team-Id": header_value},
|
||||
allowed_team_ids=allowed_team_ids,
|
||||
fallback_to_db_teams=fallback_to_db_teams,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
return resolved, by_id, by_alias
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_accepts_the_alias_of_an_allowed_team():
|
||||
"""x-litellm-team-id may carry the team alias instead of the team id (LIT-7181).
|
||||
The alias resolves to its team id before the allowed-teams check, so a
|
||||
caller whose JWT grants team_a gets team_a whether it sends the id or the
|
||||
alias, and the id path never pays for an alias lookup."""
|
||||
aliases = {"alias_a": "team_a", "alias_b": "team_b"}
|
||||
|
||||
by_alias_value, lookups_by_id, lookups_by_alias = await _resolve_header(
|
||||
"alias_a", {"team_a"}, False, _team_lookup_404, _teams_by_alias(aliases)
|
||||
)
|
||||
assert by_alias_value == HeaderTeam(header_value="alias_a", team_id="team_a")
|
||||
lookups_by_alias.assert_awaited_once()
|
||||
assert lookups_by_alias.await_args.kwargs["team_alias"] == "alias_a"
|
||||
lookups_by_id.assert_not_awaited()
|
||||
|
||||
by_id_value, lookups_by_id, lookups_by_alias = await _resolve_header(
|
||||
"team_a", {"team_a"}, False, _team_lookup_404, _teams_by_alias(aliases)
|
||||
)
|
||||
assert by_id_value == HeaderTeam(header_value="team_a", team_id="team_a")
|
||||
lookups_by_alias.assert_not_awaited()
|
||||
lookups_by_id.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_denies_aliases_of_teams_the_jwt_does_not_grant():
|
||||
"""An alias that exists but names a team outside the JWT's allowed teams is
|
||||
refused with the same 403 as an unknown value, and the detail names only
|
||||
what the caller sent, so the response reveals neither that the alias exists
|
||||
nor which team id it maps to. Because the id and the alias lookups both ran
|
||||
before the denial, the detail says the value resolved to neither form and
|
||||
lists the team ids the JWT does allow."""
|
||||
aliases = {"alias_a": "team_a", "alias_b": "team_b"}
|
||||
|
||||
with pytest.raises(HTTPException) as other_team:
|
||||
await _resolve_header("alias_b", {"team_a"}, False, _team_lookup_404, _teams_by_alias(aliases))
|
||||
with pytest.raises(HTTPException) as unknown:
|
||||
await _resolve_header("no_such", {"team_a"}, False, _team_lookup_404, _teams_by_alias(aliases))
|
||||
|
||||
assert other_team.value.status_code == 403
|
||||
assert unknown.value.status_code == 403
|
||||
assert "team_b" not in other_team.value.detail
|
||||
assert other_team.value.detail.replace("alias_b", "<value>") == unknown.value.detail.replace("no_such", "<value>")
|
||||
assert unknown.value.detail == (
|
||||
"x-litellm-team-id 'no_such' does not resolve to a team id or a unique team alias in your JWT's allowed "
|
||||
"teams. Allowed team ids: ['team_a']"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_under_db_fallback_tries_the_id_before_the_alias():
|
||||
"""Under fallback_to_db_teams a claimless JWT's header is provisional: a value
|
||||
that is an existing team id resolves to itself without an alias lookup, a
|
||||
value that is only an alias resolves to that team's id, and a value that is
|
||||
neither gets the membership denial the id path already uses."""
|
||||
known_ids = frozenset({"team_a"})
|
||||
aliases = {"alias_a": "team_a"}
|
||||
|
||||
as_id, _, lookups_by_alias = await _resolve_header(
|
||||
"team_a", set(), True, _teams_by_id(known_ids), _teams_by_alias(aliases)
|
||||
)
|
||||
assert as_id == HeaderTeam(header_value="team_a", team_id="team_a")
|
||||
lookups_by_alias.assert_not_awaited()
|
||||
|
||||
as_alias, _, _ = await _resolve_header("alias_a", set(), True, _teams_by_id(known_ids), _teams_by_alias(aliases))
|
||||
assert as_alias == HeaderTeam(header_value="alias_a", team_id="team_a")
|
||||
|
||||
with pytest.raises(HTTPException) as neither:
|
||||
await _resolve_header("ghost", set(), True, _teams_by_id(known_ids), _teams_by_alias(aliases))
|
||||
assert neither.value.status_code == 403
|
||||
assert neither.value.detail == (
|
||||
"x-litellm-team-id 'ghost' does not resolve to a team id or a unique team alias among your team memberships."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_under_db_fallback_never_aliases_a_team_id_it_could_not_read():
|
||||
"""Only a team row the database provably lacks falls through to the alias
|
||||
lookup. When the id read fails for any other reason (the generic 404 the
|
||||
team lookup uses for an unreadable database) the value keeps the id path's
|
||||
membership denial, so an outage can never turn a team id into the team
|
||||
that happens to carry it as an alias."""
|
||||
aliases = {"team_a": "team_b"}
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, side_effect=_team_lookup_404), # test-quality-ok: handle_jwt calls the team lookups as module globals with no injection seam
|
||||
patch( # test-quality-ok: handle_jwt calls the team lookups as module globals with no injection seam
|
||||
"litellm.proxy.auth.handle_jwt.get_team_object_by_alias",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=_teams_by_alias(aliases),
|
||||
) as lookups_by_alias,
|
||||
pytest.raises(HTTPException) as unreadable,
|
||||
):
|
||||
await JWTAuthManager.resolve_team_from_header(
|
||||
request_headers={"x-litellm-team-id": "team_a"},
|
||||
allowed_team_ids=set(),
|
||||
fallback_to_db_teams=True,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert unreadable.value.status_code == 403
|
||||
assert unreadable.value.detail == (
|
||||
"x-litellm-team-id 'team_a' does not resolve to a team id or a unique team alias among your team memberships."
|
||||
)
|
||||
lookups_by_alias.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_from_header_denies_a_duplicate_alias_like_an_unknown_value_but_surfaces_lookup_errors():
|
||||
"""An alias two teams share resolves to no single team, so it is denied with
|
||||
the very 403 an unknown value gets, under claims and under the DB fallback
|
||||
alike: the caller cannot be checked against either team, so telling it the
|
||||
alias is shared would let any JWT probe which aliases exist. The detail says
|
||||
the value resolved to no team id or unique alias, which is true for both.
|
||||
A lookup failure (5xx) is not disguised as a denial and propagates as is."""
|
||||
|
||||
async def duplicate_alias(team_alias, **kwargs):
|
||||
raise HTTPException(status_code=400, detail={"error": f"Multiple teams found with alias '{team_alias}'."})
|
||||
|
||||
async def db_down(team_alias, **kwargs):
|
||||
raise HTTPException(status_code=500, detail={"error": f"Error looking up team by alias '{team_alias}'"})
|
||||
|
||||
with pytest.raises(HTTPException) as duplicate:
|
||||
await _resolve_header("shared_alias", {"team_a"}, False, _team_lookup_404, duplicate_alias)
|
||||
with pytest.raises(HTTPException) as unknown:
|
||||
await _resolve_header("shared_alias", {"team_a"}, False, _team_lookup_404, _team_alias_lookup_404)
|
||||
assert duplicate.value.status_code == 403
|
||||
assert duplicate.value.detail == unknown.value.detail
|
||||
assert duplicate.value.detail == (
|
||||
"x-litellm-team-id 'shared_alias' does not resolve to a team id or a unique team alias in your JWT's "
|
||||
"allowed teams. Allowed team ids: ['team_a']"
|
||||
)
|
||||
|
||||
known_ids = frozenset({"team_a"})
|
||||
with pytest.raises(HTTPException) as duplicate_under_fallback:
|
||||
await _resolve_header("shared_alias", set(), True, _teams_by_id(known_ids), duplicate_alias)
|
||||
with pytest.raises(HTTPException) as unknown_under_fallback:
|
||||
await _resolve_header("shared_alias", set(), True, _teams_by_id(known_ids), _team_alias_lookup_404)
|
||||
assert duplicate_under_fallback.value.status_code == 403
|
||||
assert duplicate_under_fallback.value.detail == unknown_under_fallback.value.detail
|
||||
assert duplicate_under_fallback.value.detail == (
|
||||
"x-litellm-team-id 'shared_alias' does not resolve to a team id or a unique team alias among your team "
|
||||
"memberships."
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as failure:
|
||||
await _resolve_header("alias_a", {"team_a"}, False, _team_lookup_404, db_down)
|
||||
assert failure.value.status_code == 500
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_header_alias_binds_the_aliased_team_under_claims_and_db_fallback():
|
||||
"""End to end through auth_builder, x-litellm-team-id carrying a team alias
|
||||
binds the request to the aliased team (result team_id is the canonical id)
|
||||
both when the JWT grants that team by claim and when a claimless JWT relies
|
||||
on fallback_to_db_teams and DB membership."""
|
||||
aliases = {"alias_member": "team_member"}
|
||||
user_object = LiteLLM_UserTable(
|
||||
user_id="u_alias",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
teams=["team_member"],
|
||||
)
|
||||
|
||||
claims_config = LiteLLM_JWTAuth(team_ids_jwt_field="team_ids")
|
||||
by_claim = await _run_auth_builder_with_header_team(
|
||||
claims_config,
|
||||
{"sub": "u_alias", "scope": "", "team_ids": ["team_member"]},
|
||||
"alias_member",
|
||||
user_object,
|
||||
_teams_by_id(frozenset({"team_member"})),
|
||||
{"team_member"},
|
||||
_teams_by_alias(aliases),
|
||||
)
|
||||
assert by_claim["team_id"] == "team_member"
|
||||
assert by_claim["team_object"].team_id == "team_member"
|
||||
|
||||
fallback_config = LiteLLM_JWTAuth(fallback_to_db_teams=True)
|
||||
by_membership = await _run_auth_builder_with_header_team(
|
||||
fallback_config,
|
||||
{"sub": "u_alias", "scope": ""},
|
||||
"alias_member",
|
||||
user_object,
|
||||
_teams_by_id(frozenset({"team_member"})),
|
||||
set(),
|
||||
_teams_by_alias(aliases),
|
||||
)
|
||||
assert by_membership["team_id"] == "team_member"
|
||||
assert by_membership["team_object"].team_id == "team_member"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_header_alias_of_a_non_member_team_is_denied_like_an_unknown_value_under_db_fallback():
|
||||
"""Under fallback_to_db_teams, an alias naming a team the user is not a member
|
||||
of is denied with the exact same 403 as an unknown value, naming the alias
|
||||
the caller sent rather than the team id it resolved to."""
|
||||
aliases = {"alias_other": "team_other"}
|
||||
known_ids = frozenset({"team_member", "team_other"})
|
||||
user_object = LiteLLM_UserTable(
|
||||
user_id="u_alias_outsider",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
teams=["team_member"],
|
||||
)
|
||||
config = LiteLLM_JWTAuth(fallback_to_db_teams=True)
|
||||
token = {"sub": "u_alias_outsider", "scope": ""}
|
||||
|
||||
with pytest.raises(HTTPException) as outsider_alias:
|
||||
await _run_auth_builder_with_header_team(
|
||||
config, token, "alias_other", user_object, _teams_by_id(known_ids), set(), _teams_by_alias(aliases)
|
||||
)
|
||||
with pytest.raises(HTTPException) as unknown:
|
||||
await _run_auth_builder_with_header_team(
|
||||
config, token, "alias_ghost", user_object, _teams_by_id(known_ids), set(), _teams_by_alias(aliases)
|
||||
)
|
||||
|
||||
assert outsider_alias.value.status_code == 403
|
||||
assert unknown.value.status_code == 403
|
||||
assert "team_other" not in outsider_alias.value.detail
|
||||
assert outsider_alias.value.detail.replace("alias_other", "<value>") == unknown.value.detail.replace(
|
||||
"alias_ghost", "<value>"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_header_alias_under_db_fallback_keeps_the_team_allowed_routes_gate():
|
||||
"""Under fallback_to_db_teams, a member team selected by alias is still held
|
||||
to team_allowed_routes, and the denial names the alias the caller sent."""
|
||||
aliases = {"alias_member": "team_member"}
|
||||
user_object = LiteLLM_UserTable(
|
||||
user_id="u_alias_routes",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
teams=["team_member"],
|
||||
)
|
||||
config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_allowed_routes=["openai_routes"])
|
||||
token = {"sub": "u_alias_routes", "scope": ""}
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _run_auth_builder_with_header_team(
|
||||
config,
|
||||
token,
|
||||
"alias_member",
|
||||
user_object,
|
||||
_teams_by_id(frozenset({"team_member"})),
|
||||
set(),
|
||||
_teams_by_alias(aliases),
|
||||
route="/key/info",
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == (
|
||||
"Team 'alias_member' (from x-litellm-team-id header) is not allowed to access route '/key/info'."
|
||||
)
|
||||
|
||||
allowed = await _run_auth_builder_with_header_team(
|
||||
config,
|
||||
token,
|
||||
"alias_member",
|
||||
user_object,
|
||||
_teams_by_id(frozenset({"team_member"})),
|
||||
set(),
|
||||
_teams_by_alias(aliases),
|
||||
route="/chat/completions",
|
||||
)
|
||||
assert allowed["team_id"] == "team_member"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_header_selects_db_membership_team_when_jwt_also_carries_a_team_claim() -> None:
|
||||
"""Under fallback_to_db_teams, x-litellm-team-id may name a DB-membership
|
||||
team the JWT does not claim (LIT-8656): the allowed set is the JWT teams
|
||||
union the user's DB memberships, not the JWT teams alone. The flag-off
|
||||
path keeps rejecting the same header against the JWT's allowed teams."""
|
||||
user_object = LiteLLM_UserTable(
|
||||
user_id="u_mixed",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
teams=["team_member"],
|
||||
)
|
||||
config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_jwt_field="appid")
|
||||
token = {"sub": "u_mixed", "scope": "", "appid": "team_claimed"}
|
||||
fake_get_team = _teams_by_id(frozenset({"team_claimed", "team_member"}))
|
||||
|
||||
by_membership = await _run_auth_builder_with_header_team(
|
||||
config, token, "team_member", user_object, fake_get_team, {"team_claimed"}
|
||||
)
|
||||
assert by_membership["team_id"] == "team_member"
|
||||
assert by_membership["team_object"].team_id == "team_member"
|
||||
|
||||
by_claim = await _run_auth_builder_with_header_team(
|
||||
config, token, "team_claimed", user_object, fake_get_team, {"team_claimed"}
|
||||
)
|
||||
assert by_claim["team_id"] == "team_claimed"
|
||||
|
||||
flag_off = LiteLLM_JWTAuth(fallback_to_db_teams=False, team_id_jwt_field="appid")
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _run_auth_builder_with_header_team(
|
||||
flag_off, token, "team_member", user_object, fake_get_team, {"team_claimed"}
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "JWT's allowed teams" in exc_info.value.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_header_non_member_team_is_denied_when_jwt_also_carries_a_team_claim() -> None:
|
||||
"""A header naming a team the user does not belong to stays a membership
|
||||
denial even when the JWT carries a team claim, and an existing but
|
||||
non-member team produces the exact same 403 shape as a nonexistent one so
|
||||
the response is no oracle for which team ids exist."""
|
||||
user_object = LiteLLM_UserTable(
|
||||
user_id="u_mixed",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
teams=["team_member"],
|
||||
)
|
||||
config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_jwt_field="appid")
|
||||
token = {"sub": "u_mixed", "scope": "", "appid": "team_claimed"}
|
||||
fake_get_team = _teams_by_id(frozenset({"team_claimed", "team_member", "team_other"}))
|
||||
|
||||
with pytest.raises(HTTPException) as outsider_exc:
|
||||
await _run_auth_builder_with_header_team(
|
||||
config, token, "team_other", user_object, fake_get_team, {"team_claimed"}
|
||||
)
|
||||
with pytest.raises(HTTPException) as missing_exc:
|
||||
await _run_auth_builder_with_header_team(
|
||||
config, token, "team_ghost", user_object, fake_get_team, {"team_claimed"}
|
||||
)
|
||||
|
||||
assert outsider_exc.value.status_code == 403
|
||||
assert missing_exc.value.status_code == 403
|
||||
assert outsider_exc.value.detail == (
|
||||
"x-litellm-team-id 'team_other' does not resolve to a team id or a unique team alias among your "
|
||||
"team memberships."
|
||||
)
|
||||
assert missing_exc.value.detail.replace("team_ghost", "<team>") == outsider_exc.value.detail.replace(
|
||||
"team_other", "<team>"
|
||||
)
|
||||
assert "exist" not in missing_exc.value.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_no_header_keeps_the_jwt_team_when_fallback_to_db_teams_is_on() -> None:
|
||||
"""With no x-litellm-team-id header, fallback_to_db_teams must not disturb
|
||||
the claim path: the JWT's own team claim still binds the request."""
|
||||
user_object = LiteLLM_UserTable(
|
||||
user_id="u_mixed",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
teams=["team_member"],
|
||||
)
|
||||
config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_jwt_field="appid")
|
||||
token = {"sub": "u_mixed", "scope": "", "appid": "team_claimed"}
|
||||
|
||||
result = await _run_auth_builder_with_header_team(
|
||||
config,
|
||||
token,
|
||||
"team_member",
|
||||
user_object,
|
||||
_teams_by_id(frozenset({"team_claimed", "team_member"})),
|
||||
{"team_claimed"},
|
||||
send_header=False,
|
||||
)
|
||||
assert result["team_id"] == "team_claimed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_team_id_default_does_not_widen_the_header_allowed_set() -> None:
|
||||
"""team_id_default fills in a team for claimless tokens but must not widen
|
||||
the header's allowed set: a header naming the default team is still held
|
||||
to DB membership under fallback_to_db_teams."""
|
||||
user_object = LiteLLM_UserTable(
|
||||
user_id="u_default",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
teams=["team_member"],
|
||||
)
|
||||
config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_id_default="team_default")
|
||||
token = {"sub": "u_default", "scope": ""}
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _run_auth_builder_with_header_team(
|
||||
config,
|
||||
token,
|
||||
"team_default",
|
||||
user_object,
|
||||
_teams_by_id(frozenset({"team_default", "team_member"})),
|
||||
set(),
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == (
|
||||
"x-litellm-team-id 'team_default' does not resolve to a team id or a unique team alias among your "
|
||||
"team memberships."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_flag():
|
||||
"""Reading the singular team claim during sync is scoped to fallback_to_db_teams.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue