diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 6b673967427..9f14b9ec819 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -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: diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index 33457f5de16..1a9b897ca28 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -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, diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index dd4247ad3d0..3600e7ea663 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -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 diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index c23b3291365..fa05cd3f361 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -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, ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d24e451cfc1..3e6e297f6b6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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( diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 803093ff93a..2f328aef3fd 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -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}'." ), ) diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 9292e5f07db..1706b2f8f25 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -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"} diff --git a/tests/e2e/other/other_client.py b/tests/e2e/other/other_client.py index 4313bbe4068..057a735e5e5 100644 --- a/tests/e2e/other/other_client.py +++ b/tests/e2e/other/other_client.py @@ -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 diff --git a/tests/e2e/other/test_jwt_auth_e2e.py b/tests/e2e/other/test_jwt_auth_e2e.py index 2bed40f4d69..de06c376f76 100644 --- a/tests/e2e/other/test_jwt_auth_e2e.py +++ b/tests/e2e/other/test_jwt_auth_e2e.py @@ -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, "") == for_unknown.replace(unknown, ""), ( + "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}" + ) diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 91bf7a1ca18..db3f07a3acd 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -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": { diff --git a/tests/integration/observability/test_otel_conversation_id.py b/tests/integration/observability/test_otel_conversation_id.py new file mode 100644 index 00000000000..78ff40927e5 --- /dev/null +++ b/tests/integration/observability/test_otel_conversation_id.py @@ -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) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 9b5abae60cc..287f15a7183 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -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 = { diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index 17de3cf1e8a..6de9eab2d0c 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -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 diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 15defb196af..8f8dcd7852d 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -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", "") == unknown.value.detail.replace("no_such", "") + 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", "") == unknown.value.detail.replace( + "alias_ghost", "" + ) + + +@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", "") == outsider_exc.value.detail.replace( + "team_other", "" + ) + 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.