diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 20cdf9662e1..bab0d7ec092 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 ad513968b45..ede8ac99467 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, object]) -> 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 2f337c59148..ea4ded90480 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/tests/integration/contracts.json b/tests/integration/contracts.json index a8d9cf1df8b..605f78b71be 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -2020,6 +2020,75 @@ "tests/integration/observability/test_callback_delivery.py::test_response_survives_raising_success_deployment_hook[videos]": [ "other.observability.callbacks.raising_success_deployment_hook_keeps_response" ], + "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" + ], "tests/integration/management/test_user_updates_wedged_coordination_redis.py::test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged": [ "mgmt.user.update.budget_change_returns_promptly_with_wedged_coordination_redis", "mgmt.user.bulk_update.budget_change_returns_promptly_with_wedged_coordination_redis", 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 95df3709ab8..7e93d3d67a7 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 @@ -11,6 +11,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, @@ -1345,6 +1346,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