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 commit a3196907e4)

* 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 commit 071cb49d32)

* 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 commit 08639fcf42)

* 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 commit 7b4fd47c6e)

---------

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:
yuneng-jiang 2026-09-26 15:03:25 -07:00 • committed by GitHub
parent 6bfc987349
commit 6076c0a058
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 1909 additions and 106 deletions

View file

@ -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:

View file

@ -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,

View file

@ -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

View file

@ -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,
)

View file

@ -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(

View file

@ -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}'."
),
)

View file

@ -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"}

View file

@ -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

View file

@ -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}"
)

View file

@ -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": {

View 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)

View file

@ -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 = {

View file

@ -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

View file

@ -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.