diff --git a/litellm/constants.py b/litellm/constants.py index 271e42ee661..eb03d034c7e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1629,6 +1629,7 @@ INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted" +CALLER_SESSION_ID_METADATA_KEY: Final = "litellm_caller_session_id" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 6bd31a42e9e..675bf0830f9 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -27,6 +27,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm._uuid import uuid from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -187,6 +188,30 @@ def _caller_trace_field(data: Mapping[str, object], metadata_variable_name: str, return requester_value if isinstance(requester_value, str) and requester_value else None +def _caller_body_session_id(data: Mapping[str, object], metadata_variable_name: str) -> str | None: + top_level: Final = data.get("litellm_session_id") + if isinstance(top_level, str) and top_level: + return top_level + return _caller_trace_field(data, metadata_variable_name, "session_id") + + +def _caller_session_beside_trace_header( + normalized_headers: Mapping[str, str], data: Mapping[str, object], metadata_variable_name: str +) -> str | None: + session_headers: Final = {key: value for key, value in normalized_headers.items() if key != "x-litellm-trace-id"} + return get_chain_id_from_headers(session_headers) or _caller_body_session_id(data, metadata_variable_name) + + +def _drop_client_caller_session(request_body: Mapping[str, object]) -> None: + for metadata_key in ("metadata", "litellm_metadata"): + if isinstance(metadata := request_body.get(metadata_key), dict): + metadata.pop(CALLER_SESSION_ID_METADATA_KEY, None) + + +def _requester_metadata_snapshot(metadata: Mapping[str, object]) -> Mapping[str, object]: + return copy.deepcopy({key: value for key, value in metadata.items() if key != CALLER_SESSION_ID_METADATA_KEY}) + + def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None: """Only proxy-validated keys are stamped, proven by the unforgeable via_virtual_key marker AND a known non-secret shape: the sha256 hex digest @@ -1623,6 +1648,8 @@ class LiteLLMProxyRequestSetup: from litellm.proxy._types import LitellmMetadataFromRequestHeaders metadata_from_headers: Final = LitellmMetadataFromRequestHeaders() + request_body: Final = cast(Mapping[str, object], data) # cast-ok: request body is a str-keyed JSON object + _drop_client_caller_session(request_body) spend_logs_metadata: Final = LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(headers) if spend_logs_metadata is not None: metadata_from_headers["spend_logs_metadata"] = spend_logs_metadata @@ -1631,6 +1658,9 @@ class LiteLLMProxyRequestSetup: # Finally update the requests metadata with the `metadata_from_headers` ######################################################################################### + normalized_headers: Final = MappingProxyType( + {k.lower(): v for k, v in headers.items() if isinstance(k, str) and isinstance(v, str)} + ) agent_id_from_header: Final = headers.get("x-litellm-agent-id") # Explicit litellm headers take precedence; fall back to any x-*-session-id header. chain_id: Final = get_chain_id_from_headers(dict(headers)) @@ -1639,6 +1669,13 @@ class LiteLLMProxyRequestSetup: metadata_from_headers["agent_id"] = agent_id_from_header verbose_proxy_logger.debug("Extracted agent_id from header: %s", agent_id_from_header) + spend_log_session_id: Final = ( + _caller_session_beside_trace_header(normalized_headers, request_body, _metadata_variable_name) + if litellm.request_correlation_in_logs + and chain_id is not None + and chain_id == normalized_headers.get("x-litellm-trace-id") + else None + ) if chain_id: metadata_from_headers["trace_id"] = chain_id metadata_from_headers["session_id"] = chain_id @@ -1663,7 +1700,6 @@ class LiteLLMProxyRequestSetup: # - but lets a caller's existing traceparent/baggage headers (from # real OTel instrumentation) correlate with litellm's own logs instead # of generating an unrelated trace_id. - normalized_headers: Final = MappingProxyType({k.lower(): v for k, v in headers.items() if isinstance(k, str)}) if "litellm_trace_id" not in data: traceparent: Final = normalized_headers.get("traceparent") if isinstance(traceparent, str): @@ -1689,11 +1725,7 @@ class LiteLLMProxyRequestSetup: baggage: Final = normalized_headers.get("baggage") if isinstance(baggage, str): session_id_from_baggage: Final = _session_id_from_baggage(baggage) - caller_session_id: Final = _caller_trace_field( - cast(Mapping[str, object], data), # cast-ok: request body is a str-keyed JSON object - _metadata_variable_name, - "session_id", - ) + caller_session_id: Final = _caller_trace_field(request_body, _metadata_variable_name, "session_id") if session_id_from_baggage and ( caller_session_id is None or caller_session_id == session_id_from_baggage ): @@ -1702,7 +1734,11 @@ class LiteLLMProxyRequestSetup: verbose_proxy_logger.debug("Extracted session_id from W3C baggage header") if isinstance(data[_metadata_variable_name], dict): - data[_metadata_variable_name].update(metadata_from_headers) + data[_metadata_variable_name].update( + {**metadata_from_headers, CALLER_SESSION_ID_METADATA_KEY: spend_log_session_id} + if spend_log_session_id + else metadata_from_headers + ) return data @staticmethod @@ -2417,7 +2453,7 @@ async def add_litellm_data_to_request( # them — from leaking into requester_metadata where guardrails and audit # paths may read from it. if "metadata" in data and isinstance(data["metadata"], dict): - data[_metadata_variable_name]["requester_metadata"] = copy.deepcopy(data["metadata"]) + data[_metadata_variable_name]["requester_metadata"] = _requester_metadata_snapshot(data["metadata"]) if _metadata_variable_name == "litellm_metadata": data[_metadata_variable_name].update( _promoted_trace_control_fields( diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 49644f1aa12..fe55092591e 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -15,6 +15,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, CLI_SESSION_KEY_PREFIX, EMPTY_MAPPING, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -33,6 +34,8 @@ from litellm.constants import ( from litellm.constants import ( MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB, ) +from litellm.integrations.otel.model.metadata import caller_session_id +from litellm.integrations.otel.model.trace_controls import TraceControls from litellm.litellm_core_utils.classifier_logging import ( CLASSIFIER_AUDIT_FIELDS, classifier_audit_fields, @@ -930,6 +933,16 @@ def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) - return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX) +def _caller_session_for_spend_log(kwargs: Mapping[str, object], metadata: Mapping[str, object] | None) -> str | None: + beside_trace_header: Final = metadata.get(CALLER_SESSION_ID_METADATA_KEY) if metadata else None + caller_session: Final = ( + beside_trace_header + if isinstance(beside_trace_header, str) and beside_trace_header + else caller_session_id(kwargs, TraceControls()) + ) + return caller_session if caller_session and "\x00" not in caller_session else None + + def _get_session_id_for_spend_log( kwargs: Mapping[str, object], metadata: Mapping[str, object] | None, @@ -937,9 +950,12 @@ def _get_session_id_for_spend_log( omit_when_missing: bool, batch_trace_session_id: str | None = None, ) -> str | None: - """Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may - be a copied trace id. Batch call types carry a deterministic session derived from the batch id, which outranks - the per-request trace ids because those differ between the create call and the cost poller's row.""" + """Resolve the session id for the spend log row: `omit` honors only metadata.session_id, + batch sessions outrank everything, then the conversation id the caller sent (also one an x-litellm-trace-id header + displaced from the request's session), then trace ids. + + A session minted by `missing_session_id: generate` or copied from the request's trace id is not a conversation, so + it yields to the trace id, which the Responses bridge points at the previous turn's session to keep a chain whole.""" if omit_when_missing: session_id: Final = metadata.get("session_id") if metadata else None return str(session_id) if session_id else None @@ -948,11 +964,14 @@ def _get_session_id_for_spend_log( if batch_trace_session_id is not None: return batch_trace_session_id - if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None: - return str(standard_logging_payload.get("trace_id")) - if kwargs.get("litellm_trace_id") is not None: - return str(kwargs.get("litellm_trace_id")) - return str(uuid.uuid4()) + if litellm.request_correlation_in_logs and (caller_session := _caller_session_for_spend_log(kwargs, metadata)): + return caller_session + match standard_logging_payload: + case {"trace_id": trace_id} if trace_id is not None: + return str(trace_id) + case _: + litellm_trace_id: Final = kwargs.get("litellm_trace_id") + return str(litellm_trace_id) if litellm_trace_id is not None else str(uuid.uuid4()) def _get_request_duration_ms(start_time: datetime, end_time: datetime) -> int | None: diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 732c32ff2cb..e78ec5e5d3e 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -281,7 +281,13 @@ class ResponsesSessionHandler: SQL query - SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id + SELECT session_id, api_key, "user", team_id FROM spend_logs WHERE request_id = previous_response_id, then + SELECT * FROM spend_logs with that session_id and the same owner + + Only rows with the referenced response's owner are returned: the same user and team, and for a key with + neither, the same key. The session id can be caller supplied, so matching on it alone would mix owners. A + user's or team's regenerated or second key keeps the history; a key with neither does not, and turns written + by another owner (an admin continuing the chain) are not replayed. A just-finished turn gets a short second chance: the worker that served it may still be writing its spend log when the follow-up arrives, and an empty result @@ -303,13 +309,24 @@ class ResponsesSessionHandler: query: Final = """ WITH matching_session AS ( - SELECT session_id + SELECT session_id, api_key, "user", team_id FROM "LiteLLM_SpendLogs" WHERE request_id = $1 ) SELECT * - FROM "LiteLLM_SpendLogs" - WHERE session_id IN (SELECT session_id FROM matching_session) + FROM "LiteLLM_SpendLogs" AS logs + WHERE EXISTS ( + SELECT 1 + FROM matching_session + WHERE matching_session.session_id = logs.session_id + AND matching_session."user" IS NOT DISTINCT FROM logs."user" + AND matching_session.team_id IS NOT DISTINCT FROM logs.team_id + AND ( + COALESCE(matching_session."user", '') <> '' + OR COALESCE(matching_session.team_id, '') <> '' + OR matching_session.api_key = logs.api_key + ) + ) ORDER BY "endTime" ASC; """ diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 2e7fb6da194..587cdfcee25 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -39,7 +39,11 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.caching import InMemoryCache -from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER +from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, + REDACTED_BY_LITELLM, + REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER, +) from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.get_supported_openai_params import ( @@ -563,6 +567,14 @@ class ChatCompletionSession(TypedDict, total=False): ########### End of Initialize Classes used for Responses API ########### +def _chain_session_metadata(chain_session: str, request_kwargs: Mapping[str, object]) -> Mapping[str, object]: + key: Final = "litellm_metadata" if request_kwargs.get("litellm_metadata") else "metadata" + metadata: Final = request_kwargs.get(key) + if not isinstance(metadata, dict): + return MappingProxyType({}) + return {key: {**_STR_KEY_DICT_ADAPTER.validate_python(metadata), CALLER_SESSION_ID_METADATA_KEY: chain_session}} + + class LiteLLMCompletionResponsesConfig: @staticmethod def get_supported_openai_params(model: str) -> list: @@ -1078,8 +1090,6 @@ class LiteLLMCompletionResponsesConfig: else: # Both are empty - this likely means function_call_output had empty/invalid call_id # Provide a helpful error message - import litellm - raise litellm.BadRequestError( message=( f"Unable to create messages for completion request. " @@ -1094,8 +1104,14 @@ class LiteLLMCompletionResponsesConfig: ) litellm_completion_request["messages"] = combined_messages - litellm_completion_request["litellm_trace_id"] = chat_completion_session.get("litellm_session_id") - return litellm_completion_request + chain_session: Final = chat_completion_session.get("litellm_session_id") + litellm_completion_request["litellm_trace_id"] = chain_session + chain_session_metadata: Final[Mapping[str, object]] = ( + _chain_session_metadata(chain_session, request_kwargs) + if chain_session and litellm.request_correlation_in_logs + else MappingProxyType({}) + ) + return {**litellm_completion_request, **chain_session_metadata} @staticmethod def _transform_response_input_param_to_chat_completion_message( diff --git a/tests/integration/spend/test_spend_log_caller_session.py b/tests/integration/spend/test_spend_log_caller_session.py new file mode 100644 index 00000000000..213583d21a1 --- /dev/null +++ b/tests/integration/spend/test_spend_log_caller_session.py @@ -0,0 +1,783 @@ +import itertools +import json +import re +import threading +import uuid +from urllib.parse import urlsplit +from collections import deque +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from dataclasses import dataclass, replace +from datetime import datetime, timedelta, timezone +from hashlib import sha256 +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import openai +import psutil +import pytest +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + list_value, + object_value, + string_value, +) +from integration._support.budget import readiness_seconds +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * readiness_seconds(2) + 2 * graceful_stop_seconds() + 240) + +MARKER: Final = re.compile(rb"callersess-[0-9a-f]{32}") +CALLER_SESSION_KEY: Final = "litellm_caller_session_id" +PATHS: Final = MappingProxyType( + {"chat": "/v1/chat/completions", "messages": "/v1/messages", "responses": "/v1/responses"} +) +ENDPOINT_STREAMS: Final = tuple(itertools.product(PATHS, (False, True))) +Headers = Mapping[str, str] | Sequence[tuple[str, str]] + + +def _marker() -> str: + return "callersess-" + uuid.uuid4().hex + + +def _token(label: str) -> str: + return f"{label}-{uuid.uuid4().hex}" + + +def _frame(payload: Mapping[str, JsonValue]) -> bytes: + return b"data: " + json.dumps(payload).encode() + b"\n\n" + + +def _chat_reply(text: str, stream: bool) -> Reply: + identity: Final = "chatcmpl-" + uuid.uuid4().hex + usage: Final[dict[str, JsonValue]] = {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9} + 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": text}, "finish_reason": "stop"} + ], + "usage": usage, + } + ).encode() + ) + chunk: Final[dict[str, JsonValue]] = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + } + return Reply( + content_type="text/event-stream", + chunks=( + _frame({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}}]}), + _frame({**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": usage}), + b"data: [DONE]\n\n", + ), + ) + + +def _upstream(request: Request) -> Reply: + markers: Final = tuple(found.group(0).decode() for found in MARKER.finditer(request.body)) + if request.method != "POST" or not request.target.endswith("/chat/completions") or not markers: + return Reply(status=404, body=b'{"error":"unexpected request"}') + return _chat_reply(f"answer {markers[-1]}", JSON_OBJECT.validate_json(request.body).get("stream") is True) + + +def _user_markers(body: Mapping[str, JsonValue]) -> tuple[str, ...]: + user_contents: Final = tuple( + json.dumps(object_value(message)["content"]) + for message in list_value(body.get("messages", [])) + if object_value(message)["role"] == "user" + ) + return tuple(found.group(0).decode() for content in user_contents if (found := MARKER.search(content.encode()))) + + +def _spans(batch: Request) -> Iterator[dict[str, JsonValue]]: + for resource in list_value(JSON_OBJECT.validate_json(batch.body)["resourceSpans"]): + for scope in list_value(object_value(resource)["scopeSpans"]): + yield from (object_value(span) for span in list_value(object_value(scope)["spans"])) + + +def _attributes(span: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + string_value(object_value(attribute)["key"]): object_value(attribute)["value"] + for attribute in list_value(span.get("attributes", [])) + } + + +def _batch_attributes(batches: Sequence[Request]) -> Iterator[dict[str, JsonValue]]: + return (_attributes(span) for span in itertools.chain.from_iterable(_spans(batch) for batch in batches)) + + +@dataclass(frozen=True, slots=True) +class Collector: + wire: Wire + accepted: deque[Request] + guard: threading.Lock + + def span_attributes(self, response_id: str) -> tuple[dict[str, JsonValue], ...]: + with self.guard: + batches: Final = tuple(self.accepted) + attributes: Final = tuple(_batch_attributes(batches)) + return tuple( + found + for found in attributes + if object_value(found.get("gen_ai.response.id", {})).get("stringValue") == response_id + ) + + def conversation_ids(self, conversations: frozenset[str]) -> tuple[str, ...]: + with self.guard: + batches: Final = tuple(self.accepted) + responded: Final = ( + object_value(found.get("gen_ai.conversation.id", {})).get("stringValue") + for found in _batch_attributes(batches) + if "gen_ai.response.id" in found + ) + return tuple( + conversation + for conversation in responded + if isinstance(conversation, str) and conversation in conversations + ) + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + scenario: Scenario + model: str + upstream: Wire + process: OwnedProxy + + def body(self, endpoint: str, marker: str, extra: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + match endpoint: + case "chat": + return {"model": self.model, "messages": [{"role": "user", "content": marker}], **extra} + case "messages": + return { + "model": self.model, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + **extra, + } + case _: + return {"model": self.model, "input": marker, **extra} + + def post( + self, + endpoint: str, + marker: str, + fields: Mapping[str, JsonValue] = MappingProxyType({}), + *, + stream: bool = False, + headers: Headers = (), + key: str | None = None, + ) -> httpx.Response: + payload: Final = self.body(endpoint, marker, fields) | ({"stream": True} if stream else {}) + request_headers: Final = ( + ("Authorization", f"Bearer {self.proxy.key if key is None else key}"), + *(headers.items() if isinstance(headers, Mapping) else headers), + ) + with self.proxy.client.stream("POST", PATHS[endpoint], json=payload, headers=request_headers) as response: + response.read() + return response + + def send( + self, + endpoint: str, + marker: str, + fields: Mapping[str, JsonValue] = MappingProxyType({}), + *, + stream: bool = False, + headers: Headers = (), + key: str | None = None, + ) -> httpx.Response: + response: Final = self.post(endpoint, marker, fields, stream=stream, headers=headers, key=key) + assert response.status_code == 200, response.text + return response + + def upstream_bodies(self, marker: str) -> tuple[dict[str, JsonValue], ...]: + bodies: Final = tuple(JSON_OBJECT.validate_json(request.body) for request in self.upstream.drain()) + return tuple(body for body in bodies if _user_markers(body)[-1:] == (marker,)) + + +@dataclass(frozen=True, slots=True) +class Rigs: + flag_off: Rig + flag_on: Rig + sink: Collector + + +def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT session_id, request_id, metadata::text AS metadata_text FROM "LiteLLM_SpendLogs" ' + "WHERE litellm_call_id=%s", + (call_id,), + ) + + +def _spend_row(response: httpx.Response) -> dict[str, JsonValue]: + call_id: Final = response.headers["x-litellm-call-id"] + rows: Final = eventually(lambda: _spend_rows(call_id), lambda values: len(values) >= 1, seconds=70) + assert len(rows) == 1, rows + return rows[0] + + +def _session(response: httpx.Response) -> JsonValue: + return _spend_row(response)["session_id"] + + +def _model(scenario: Scenario, upstream: Wire) -> str: + return scenario.model( + model="hosted_vllm/gpt-4o-mini", + api_base=f"{upstream.url}/v1", + use_chat_completions_api=True, + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + + +def _rig_config( + directory: Path, name: str, litellm_settings: Mapping[str, JsonValue], extra: Mapping[str, JsonValue] +) -> Path: + shared: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "general_settings": {**object_value(shared["general_settings"]), "store_prompts_in_spend_logs": True}, + "litellm_settings": {**object_value(shared["litellm_settings"]), **litellm_settings}, + **extra, + } + path: Final = directory / f"{name}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def rigs(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rigs]: + accepted: Final[deque[Request]] = deque() # mutable-ok: sink thread appends each accepted OTLP batch + guard: Final = threading.Lock() + + def sink(request: Request) -> Reply: + with guard: + accepted.append(request) + return Reply() + + directory: Final = tmp_path_factory.mktemp("caller-session") + with ExitStack() as stack: + upstream: Final = stack.enter_context(wire_server(_upstream)) + collector: Final = Collector(stack.enter_context(wire_server(sink)), accepted, guard) + shared: Final = stack.enter_context(gateway_from_environment()) + flag_off: Final = stack.enter_context( + owned_proxy_process(shared, directory, {}, config=_rig_config(directory, "flag-off", {}, {}), workers=2) + ) + flag_on: Final = stack.enter_context( + owned_proxy_process( + shared, + directory, + {"LITELLM_OTEL_V2": "1", "OTEL_BSP_SCHEDULE_DELAY": "300"}, + config=_rig_config( + directory, + "flag-on", + {"request_correlation_in_logs": True, "callbacks": ["otel"]}, + { + "callback_settings": { + "otel": {"exporter": "http/json", "endpoint": collector.wire.url, "mapper_names": ["genai"]} + } + }, + ), + workers=2, + ) + ) + off_scenario: Final = stack.enter_context(flag_off.gateway.scenario()) + on_scenario: Final = stack.enter_context(flag_on.gateway.scenario()) + yield Rigs( + Rig(flag_off.gateway, off_scenario, _model(off_scenario, upstream), upstream, flag_off), + Rig(flag_on.gateway, on_scenario, _model(on_scenario, upstream), upstream, flag_on), + collector, + ) + + +@pytest.mark.parametrize(("endpoint", "stream"), ENDPOINT_STREAMS) +def test_spend_log_session_is_the_body_session_beside_a_trace_header(rigs: Rigs, endpoint: str, stream: bool) -> None: + marker: Final = _marker() + session: Final = _token("convo") + response: Final = rigs.flag_on.send( + endpoint, + marker, + stream=stream, + headers={"x-litellm-trace-id": _token("trace")}, + fields={"litellm_session_id": session}, + ) + + row: Final = _spend_row(response) + upstream: Final = rigs.flag_on.upstream_bodies(marker) + + assert row["session_id"] == session, row + assert CALLER_SESSION_KEY not in string_value(row["metadata_text"]), row + assert len(upstream) == 1, upstream + assert CALLER_SESSION_KEY not in json.dumps(upstream[0]), upstream + + +def test_sdk_clients_record_the_body_session_beside_a_trace_header(rigs: Rigs) -> None: + rig: Final = rigs.flag_on + base_url: Final = str(rig.proxy.client.base_url) + traces: Final = tuple(_token("trace") for _ in range(3)) + sessions: Final = tuple(_token("convo") for _ in range(3)) + markers: Final = tuple(_marker() for _ in range(3)) + chat_client: Final = openai.OpenAI(base_url=base_url + "/v1", api_key=rig.proxy.key, max_retries=0) + messages_client: Final = anthropic.Anthropic(base_url=base_url, api_key=rig.proxy.key, max_retries=0) + + with chat_client.chat.completions.with_streaming_response.create( + model=rig.model, + messages=[{"role": "user", "content": markers[0]}], + stream=True, + extra_headers={"x-litellm-trace-id": traces[0]}, + extra_body={"litellm_session_id": sessions[0]}, + ) as chat_stream: + chat_call: Final = chat_stream.headers["x-litellm-call-id"] + chat_text: Final = "".join( + chunk.choices[0].delta.content or "" for chunk in chat_stream.parse() if chunk.choices + ) + message: Final = messages_client.messages.with_raw_response.create( + model=rig.model, + max_tokens=16, + messages=[{"role": "user", "content": markers[1]}], + extra_headers={"x-litellm-trace-id": traces[1]}, + extra_body={"litellm_session_id": sessions[1]}, + ) + with chat_client.responses.with_streaming_response.create( + model=rig.model, + input=markers[2], + stream=True, + extra_headers={"x-litellm-trace-id": traces[2]}, + extra_body={"litellm_session_id": sessions[2]}, + ) as responses_stream: + responses_call: Final = responses_stream.headers["x-litellm-call-id"] + completed: Final = tuple(event for event in responses_stream.parse() if event.type == "response.completed") + + first_block: Final = message.parse().content[0] + assert chat_text == f"answer {markers[0]}", chat_text + assert isinstance(first_block, anthropic.types.TextBlock) and first_block.text == f"answer {markers[1]}", message + assert len(completed) == 1, completed + calls: Final = (chat_call, message.headers["x-litellm-call-id"], responses_call) + landed: Final = tuple( + eventually(lambda call=call: _spend_rows(call), lambda values: len(values) == 1, seconds=70)[0]["session_id"] + for call in calls + ) + assert landed == sessions, (landed, calls) + + +def test_conversation_span_keeps_the_trace_header_and_no_internal_key(rigs: Rigs) -> None: + trace: Final = _token("trace") + response: Final = rigs.flag_on.send( + "chat", _marker(), headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": _token("convo")} + ) + response_id: Final = string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + spans: Final = eventually(lambda: rigs.sink.span_attributes(response_id), lambda found: len(found) == 1, seconds=30) + + assert spans[0]["gen_ai.conversation.id"] == {"stringValue": trace}, spans + assert CALLER_SESSION_KEY not in json.dumps(spans[0]), spans + + +def test_a_continued_chains_conversation_spans_keep_each_turns_trace_header(rigs: Rigs) -> None: + traces: Final = (_token("trace"), _token("trace")) + first: Final = _turn(rigs.flag_on, _marker(), headers={"x-litellm-trace-id": traces[0]}) + _turn( + rigs.flag_on, + _marker(), + first, + headers={"x-litellm-trace-id": traces[1]}, + fields={"litellm_session_id": _token("convo")}, + ) + + conversations: Final = eventually( + lambda: rigs.sink.conversation_ids(frozenset(traces)), lambda found: len(found) == 2, seconds=30 + ) + + assert sorted(conversations) == sorted(traces), conversations + + +@pytest.mark.parametrize("endpoint", list(PATHS)) +def test_flag_off_keeps_the_trace_id_beside_a_body_session(rigs: Rigs, endpoint: str) -> None: + trace: Final = _token("trace") + + response: Final = rigs.flag_off.send( + endpoint, _marker(), headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": _token("convo")} + ) + + assert _session(response) == trace + + +@pytest.mark.parametrize("endpoint", list(PATHS)) +def test_response_cache_hit_records_the_body_session_beside_a_trace_header(rigs: Rigs, endpoint: str) -> None: + marker: Final = _marker() + sessions: Final = (_token("convo"), _token("convo")) + responses: Final = tuple( + rigs.flag_on.send( + endpoint, marker, headers={"x-litellm-trace-id": _token("trace")}, fields={"litellm_session_id": session} + ) + for session in sessions + ) + + assert tuple(_session(response) for response in responses) == sessions + assert len(rigs.flag_on.upstream_bodies(marker)) == 1 + + +def test_metadata_session_beside_a_new_trace_per_turn_groups_both_turns(rigs: Rigs) -> None: + session: Final = _token("convo") + responses: Final = tuple( + rigs.flag_on.send( + "chat", + _marker(), + headers={"x-litellm-trace-id": _token("trace")}, + fields={"metadata": {"session_id": session}}, + ) + for _ in range(2) + ) + + assert tuple(_session(response) for response in responses) == (session, session) + now: Final = datetime.now(timezone.utc) + page: Final = rigs.flag_on.proxy.get( + "/spend/logs/ui", + params={ + "session_id": session, + "start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + "end_date": (now + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + }, + ) + listed: Final = sorted(string_value(object_value(row)["request_id"]) for row in list_value(page["data"])) + assert listed == sorted( + string_value(JSON_OBJECT.validate_json(response.content)["id"]) for response in responses + ), page + + +def _shape(shape: str, trace: str, session: str, header: str) -> tuple[Headers, dict[str, JsonValue], str]: + traced: Final = {"x-litellm-trace-id": trace} + match shape: + case "session-header-wins": + return {"x-litellm-session-id": header}, {"litellm_session_id": session}, header + case "trace-header-only": + return traced, {}, trace + case "body-session-only": + return {}, {"litellm_session_id": session}, session + case "int-session": + return traced, {"litellm_session_id": 7}, trace + case "list-session": + return traced, {"litellm_session_id": [session]}, trace + case "empty-session": + return traced, {"litellm_session_id": ""}, trace + case "null-session": + return traced, {"litellm_session_id": None}, trace + case "nul-session": + return traced, {"litellm_session_id": session + "\x00x"}, trace + case "two-trace-headers": + return ( + (("x-litellm-trace-id", trace), ("x-litellm-trace-id", header)), + {"litellm_session_id": session}, + session, + ) + case _: + return traced, {"litellm_session_id": session + "s" * 5120}, session + "s" * 5120 + + +@pytest.mark.parametrize( + "shape", + [ + "session-header-wins", + "trace-header-only", + "body-session-only", + "int-session", + "list-session", + "empty-session", + "null-session", + "nul-session", + "two-trace-headers", + "5kb-session", + ], +) +def test_spend_log_session_for_each_session_shape(rigs: Rigs, shape: str) -> None: + headers, fields, expected = _shape(shape, _token("trace"), _token("convo"), _token("header")) + neighbor_trace: Final = _token("neighbor") + + response: Final = rigs.flag_on.send("chat", _marker(), fields, headers=headers) + neighbor: Final = rigs.flag_on.send("chat", _marker(), headers={"x-litellm-trace-id": neighbor_trace}) + + assert _session(response) == expected + assert _session(neighbor) == neighbor_trace + + +@pytest.mark.parametrize("endpoint", ["chat", "messages"]) +def test_client_sent_internal_session_key_is_ignored(rigs: Rigs, endpoint: str) -> None: + forged: Final = _token("forged") + trace: Final = _token("trace") + forged_metadata: Final[dict[str, JsonValue]] = { + "metadata" if endpoint == "chat" else "litellm_metadata": {CALLER_SESSION_KEY: forged} + } + + untraced: Final = rigs.flag_on.send(endpoint, _marker(), forged_metadata) + traced: Final = rigs.flag_on.send(endpoint, _marker(), forged_metadata, headers={"x-litellm-trace-id": trace}) + + assert _session(untraced) not in (forged, None) + assert _session(traced) == trace + + +def test_identical_requests_each_log_once_with_the_body_session(rigs: Rigs) -> None: + session: Final = _token("convo") + trace: Final = _token("trace") + marker: Final = _marker() + + responses: Final = tuple( + rigs.flag_on.send("chat", marker, headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": session}) + for _ in range(3) + ) + + assert len({response.headers["x-litellm-call-id"] for response in responses}) == 3 + assert tuple(_session(response) for response in responses) == (session, session, session) + + +def _pass_through_route(rig: Rig, target: str) -> str: + path: Final = f"/caller-session-{uuid.uuid4().hex}" + created: Final = rig.proxy.post("/config/pass_through_endpoint", {"path": path, "target": target, "auth": True}) + endpoint_id: Final = string_value(object_value(list_value(created["endpoints"])[0])["id"]) + rig.scenario.cleanups.callback( + lambda: rig.proxy.request("DELETE", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id}) + ) + return path + + +@pytest.mark.parametrize("field", ["litellm_session_id", "metadata"]) +def test_pass_through_spend_log_session_is_the_body_session_beside_a_trace_header(rigs: Rigs, field: str) -> None: + marker: Final = _marker() + session: Final = _token("convo") + rig: Final = rigs.flag_on + path: Final = _pass_through_route(rig, f"{rig.upstream.url}/v1/chat/completions") + key: Final = rig.scenario.key(allowed_passthrough_routes=[path]) + fields: Final[dict[str, JsonValue]] = ( + {"litellm_session_id": session} if field == "litellm_session_id" else {"metadata": {"session_id": session}} + ) + + response: Final = rig.proxy.request( + "POST", + path, + rig.body("chat", marker, fields), + key=key, + headers={"x-litellm-trace-id": _token("trace")}, + ) + + rows: Final = eventually( + lambda: read_rows( + 'SELECT session_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),) + ), + lambda found: len(found) >= 1, + seconds=70, + ) + assert response.status_code == 200, response.text + assert f"answer {marker}" in response.text + assert len(rig.upstream_bodies(marker)) == 1 + assert rows == [{"session_id": session}] + + +def _turn( + rig: Rig, + marker: str, + previous: str | None = None, + *, + key: str | None = None, + headers: Headers = (), + fields: Mapping[str, JsonValue] = MappingProxyType({}), +) -> str: + chained: Final[dict[str, JsonValue]] = {"previous_response_id": previous} if previous else {} + response: Final = rig.send("responses", marker, {**chained, **fields}, headers=headers, key=key) + _spend_row(response) + return string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + +def _turn_fields(shape: str, session: str) -> tuple[dict[str, JsonValue], dict[str, JsonValue]]: + match shape: + case "session-every-turn": + return {"litellm_session_id": session}, {"litellm_session_id": session} + case "session-from-turn-2": + return {}, {"litellm_session_id": session} + case _: + return {"litellm_session_id": _token("convo")}, {"litellm_session_id": _token("convo")} + + +@pytest.mark.parametrize( + ("correlated", "shape", "replays_turn_1"), + [ + (True, "session-every-turn", True), + (True, "session-from-turn-2", True), + (True, "new-session-per-turn", True), + (False, "session-from-turn-2", False), + ], +) +def test_bridged_chain_history(rigs: Rigs, correlated: bool, shape: str, replays_turn_1: bool) -> None: + rig: Final = rigs.flag_on if correlated else rigs.flag_off + first_fields, second_fields = _turn_fields(shape, _token("convo")) + markers: Final = (_marker(), _marker(), _marker()) + first: Final = _turn(rig, markers[0], headers={"x-litellm-trace-id": _token("trace")}, fields=first_fields) + second: Final = _turn(rig, markers[1], first, headers={"x-litellm-trace-id": _token("trace")}, fields=second_fields) + + _turn(rig, markers[2], second) + + replayed: Final = rig.upstream_bodies(markers[2]) + assert len(replayed) == 1, replayed + assert _user_markers(replayed[0]) == (markers if replays_turn_1 else markers[1:]) + + +def test_bridged_chains_sharing_a_metadata_session_replay_each_others_turns(rigs: Rigs) -> None: + shared: Final[dict[str, JsonValue]] = {"metadata": {"session_id": _token("convo")}} + markers: Final = (_marker(), _marker(), _marker()) + first: Final = _turn(rigs.flag_on, markers[0], fields=shared) + _turn(rigs.flag_on, markers[1], fields=shared) + + _turn(rigs.flag_on, markers[2], first) + + replayed: Final = rigs.flag_on.upstream_bodies(markers[2]) + assert len(replayed) == 1, replayed + assert _user_markers(replayed[0]) == markers + + +@dataclass(frozen=True, slots=True) +class Owners: + first: str + second: str + + +def _owners(scenario: Scenario, model: str, shape: str) -> Owners: + user: Final = _token("user") + match shape: + case "other-user": + return Owners( + scenario.key(user_id=user, models=[model]), scenario.key(user_id=_token("user"), models=[model]) + ) + case "same-user": + return Owners(scenario.key(user_id=user, models=[model]), scenario.key(user_id=user, models=[model])) + case "same-team": + team: Final = scenario.team() + return Owners(scenario.key(team_id=team, models=[model]), scenario.key(team_id=team, models=[model])) + case _: + return Owners(scenario.key(models=[model]), scenario.key(models=[model])) + + +@pytest.mark.parametrize( + ("shape", "replays_first_owner"), + [("other-user", False), ("same-user", True), ("same-team", True), ("ownerless-keys", False)], +) +def test_bridged_history_follows_the_previous_responses_owner( + rigs: Rigs, shape: str, replays_first_owner: bool +) -> None: + rig: Final = rigs.flag_off + owners: Final = _owners(rig.scenario, rig.model, shape) + session: Final = _token("convo") + markers: Final = (_marker(), _marker(), _marker()) + _turn(rig, markers[0], key=owners.first, fields={"litellm_session_id": session}) + second: Final = _turn(rig, markers[1], key=owners.second, fields={"litellm_session_id": session}) + + _turn(rig, markers[2], second, key=owners.second) + + replayed: Final = rig.upstream_bodies(markers[2]) + assert len(replayed) == 1, replayed + assert _user_markers(replayed[0]) == (markers if replays_first_owner else markers[1:]) + + +BURST: Final = ENDPOINT_STREAMS * 4 + + +def test_concurrent_burst_across_endpoints_logs_each_session_once(rigs: Rigs) -> None: + sessions: Final = tuple(_token("convo") for _ in BURST) + + def one(index: int) -> httpx.Response: + endpoint, stream = BURST[index] + return rigs.flag_on.send( + endpoint, + _marker(), + stream=stream, + headers={"x-litellm-trace-id": _token("trace")}, + fields={"litellm_session_id": sessions[index]}, + ) + + with ThreadPoolExecutor(max_workers=12) as pool: + responses: Final = tuple(pool.map(one, range(len(BURST)))) + + assert tuple(_session(response) for response in responses) == sessions + + +def _open_upstream_connections(pid: int, upstream: Wire) -> int: + port: Final = urlsplit(upstream.url).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def _post_or_none(rig: Rig, session: str) -> httpx.Response | None: + try: + return rig.post( + "chat", _marker(), headers={"x-litellm-trace-id": _token("trace")}, fields={"litellm_session_id": session} + ) + except httpx.HTTPError: + return None + + +def test_killing_one_worker_mid_burst_keeps_logging_the_body_session(rigs: Rigs) -> None: + release: Final = threading.Event() + arrived: Final[SimpleQueue[bytes]] = SimpleQueue() + + def held(request: Request) -> Reply: + arrived.put(request.body) + assert release.wait(timeout=60), "the burst was never released" + return _upstream(request) + + process: Final = rigs.flag_on.process + sessions: Final = tuple(_token("convo") for _ in range(20)) + after_session: Final = _token("convo") + with wire_server(held) as wire: + rig: Final = replace(rigs.flag_on, model=_model(rigs.flag_on.scenario, wire), upstream=wire) + workers: Final = eventually(process.worker_pids, lambda found: len(found) == 2, seconds=30) + with ThreadPoolExecutor(max_workers=len(sessions)) as pool: + pending: Final = tuple(pool.submit(_post_or_none, rig, session) for session in sessions) + eventually(arrived.qsize, lambda size: size == len(sessions), seconds=60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire) for pid in workers}) + victim, survivor = sorted(workers, key=held_by.__getitem__) + with process.replacing_worker(victim) as kill: + kill.freeze() + kill() + release.set() + results: Final = tuple(future.result() for future in pending) + after: Final = rig.send( + "chat", + _marker(), + headers={"x-litellm-trace-id": _token("trace")}, + fields={"litellm_session_id": after_session}, + ) + + served: Final = tuple( + (response, session) + for response, session in zip(results, sessions, strict=True) + if response is not None and response.status_code == 200 + ) + assert sum(held_by.values()) == len(sessions), held_by + assert len(served) == held_by[survivor] >= len(sessions) // 2, (held_by, len(served)) + assert _session(after) == after_session + assert tuple(_session(response) for response, _ in served) == tuple(session for _, session in served) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 7291f49af99..de2676bafaa 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -14,8 +14,8 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm import litellm.constants as litellm_constants import litellm.proxy.spend_tracking.spend_tracking_utils as spend_tracking_utils -from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL -from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY +from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL +from litellm.constants import CALLER_SESSION_ID_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -243,6 +243,82 @@ def test_batch_session_outranks_the_per_request_trace_id(): assert session_id == "batch-uid-1" +def _chat_session_id( + monkeypatch: pytest.MonkeyPatch, + litellm_params: Mapping[str, object], + correlation_in_logs: bool = True, + batch_trace_session_id: str | None = None, + metadata: Mapping[str, object] = MappingProxyType({}), +) -> str | None: + monkeypatch.setattr(litellm, "request_correlation_in_logs", correlation_in_logs) + return _get_session_id_for_spend_log( + kwargs={"litellm_trace_id": "trace-abc", "litellm_params": litellm_params}, + metadata=metadata, + standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING, + omit_when_missing=False, + batch_trace_session_id=batch_trace_session_id, + ) + + +_CALLER_SESSION_PARAMS: Final = MappingProxyType({"litellm_session_id": "sess-1", "metadata": {"trace_id": "otel-1"}}) + + +def test_caller_session_wins_over_the_per_request_trace_id(monkeypatch: pytest.MonkeyPatch): + assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS) == "sess-1" + + +def test_batch_session_outranks_a_caller_session(monkeypatch: pytest.MonkeyPatch): + assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS, batch_trace_session_id="batch-uid-1") == "batch-uid-1" + + +def test_caller_session_keeps_the_legacy_trace_id_with_correlation_logging_off(monkeypatch: pytest.MonkeyPatch): + assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS, correlation_in_logs=False) == "trace-abc" + + +@pytest.mark.parametrize(("correlation_in_logs", "expected_session_id"), [(True, "convo-1"), (False, "trace-abc")]) +def test_session_displaced_by_a_trace_header_wins_only_with_correlation_logging( + monkeypatch: pytest.MonkeyPatch, correlation_in_logs: bool, expected_session_id: str +): + header_chain: Final = {"litellm_session_id": "trace-abc", "metadata": {"trace_id": "trace-abc", "session_id": "trace-abc"}} + session_id: Final = _chat_session_id( + monkeypatch, + header_chain, + correlation_in_logs=correlation_in_logs, + metadata={CALLER_SESSION_ID_METADATA_KEY: "convo-1"}, + ) + assert session_id == expected_session_id + + +@pytest.mark.parametrize( + ("litellm_params", "metadata"), + [ + pytest.param({"litellm_session_id": "convo\x00x"}, {}, id="body-session"), + pytest.param({"litellm_session_id": "trace-abc"}, {CALLER_SESSION_ID_METADATA_KEY: "convo\x00x"}, id="beside-trace-header"), + ], +) +def test_caller_session_postgres_cannot_store_keeps_the_trace_id( + monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object], metadata: Mapping[str, object] +): + assert _chat_session_id(monkeypatch, litellm_params, metadata=metadata) == "trace-abc" + + +@pytest.mark.parametrize( + "litellm_params", + [ + pytest.param({"metadata": {}}, id="no-session"), + pytest.param( + {"litellm_session_id": "minted-1", "metadata": {"session_id": "minted-1", SESSION_ID_GENERATED_METADATA_KEY: True}}, + id="minted-by-generate-policy", + ), + pytest.param({"litellm_session_id": "job-7", "metadata": {"trace_id": "job-7"}}, id="copied-from-trace-id"), + ], +) +def test_session_that_is_not_a_caller_conversation_keeps_the_trace_id( + monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object] +): + assert _chat_session_id(monkeypatch, litellm_params) == "trace-abc" + + def test_omit_policy_still_suppresses_batch_sessions(): session_id: Final = _get_session_id_for_spend_log( kwargs={}, diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 96c11be2650..93202d93566 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -17,6 +17,7 @@ from starlette.datastructures import Headers import litellm from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, SERVER_STREAMING_CLASSIFICATION_KEY, SERVER_STREAMING_CLASSIFICATION_MARKER, @@ -3308,6 +3309,93 @@ def test_add_litellm_metadata_from_request_headers_x_litellm_trace_id_sets_chain assert data["litellm_trace_id"] == "foo" +@pytest.mark.parametrize(("correlation_in_logs", "expected_spend_log_session"), [(True, "convo-1"), (False, None)]) +@pytest.mark.parametrize("session_in_metadata", [False, True], ids=["top-level", "metadata"]) +def test_x_litellm_trace_id_keeps_the_body_session_for_spend_logs_only( + monkeypatch: pytest.MonkeyPatch, + correlation_in_logs: bool, + expected_spend_log_session: str | None, + session_in_metadata: bool, +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", correlation_in_logs) + metadata: Final[dict[str, str]] = {"session_id": "convo-1"} if session_in_metadata else {} + data: Final[dict[str, object]] = ( + {"metadata": metadata} if session_in_metadata else {"metadata": metadata, "litellm_session_id": "convo-1"} + ) + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={"x-litellm-trace-id": "req-trace-1"}, data=data, _metadata_variable_name="metadata" + ) + assert metadata.get(CALLER_SESSION_ID_METADATA_KEY) == expected_spend_log_session + assert data["litellm_session_id"] == "req-trace-1" + assert metadata["session_id"] == "req-trace-1" + assert data["litellm_trace_id"] == "req-trace-1" + + +@pytest.mark.parametrize( + "session_header", ["x-litellm-session-id", "x-claude-code-session-id", "x-session-id"] +) +@pytest.mark.parametrize("with_trace_header", [False, True], ids=["session-header-only", "beside-trace-header"]) +def test_session_headers_still_beat_the_body_session_with_correlation_logging( + monkeypatch: pytest.MonkeyPatch, session_header: str, with_trace_header: bool +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + metadata: Final[dict[str, str]] = {} + data: Final[dict[str, object]] = {"metadata": metadata, "litellm_session_id": "convo-1"} + trace_header: Final = {"x-litellm-trace-id": "req-trace-1"} if with_trace_header else {} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={**trace_header, session_header: "hdr-session-1"}, data=data, _metadata_variable_name="metadata" + ) + spend_log_session: Final = metadata.get(CALLER_SESSION_ID_METADATA_KEY) or data["litellm_session_id"] + assert spend_log_session == "hdr-session-1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("forged_in", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize( + "headers", + [{"x-litellm-session-id": "hdr-session-1"}, {}], + ids=["session-header", "no-headers"], +) +async def test_client_metadata_cannot_set_the_spend_log_caller_session( + monkeypatch: pytest.MonkeyPatch, forged_in: str, headers: dict[str, str] +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + updated: Final = await add_litellm_data_to_request( + data={ + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "hi"}], + forged_in: {CALLER_SESSION_ID_METADATA_KEY: "forged-session"}, + }, + request=_make_request_mock("/v1/chat/completions", headers), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + assert CALLER_SESSION_ID_METADATA_KEY not in updated["metadata"] + + +@pytest.mark.asyncio +async def test_spend_log_caller_session_stays_out_of_requester_metadata(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + updated: Final = await add_litellm_data_to_request( + data={ + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "hi"}], + "litellm_session_id": "convo-1", + "metadata": {"team_label": "kept"}, + }, + request=_make_request_mock("/v1/chat/completions", {"x-litellm-trace-id": "req-trace-1"}), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + assert updated["metadata"][CALLER_SESSION_ID_METADATA_KEY] == "convo-1" + assert updated["metadata"]["requester_metadata"]["team_label"] == "kept" + assert CALLER_SESSION_ID_METADATA_KEY not in updated["metadata"]["requester_metadata"] + + def test_add_litellm_metadata_from_request_headers_x_litellm_session_id_sets_chain_id(): """x-litellm-session-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining.""" headers = {"x-litellm-session-id": "bar"} diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 19bbb6481fb..cc5d5c62739 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -17,6 +17,7 @@ from openai.types.responses.response_function_web_search import ( from pydantic import TypeAdapter import litellm +from litellm.constants import CALLER_SESSION_ID_METADATA_KEY from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -6499,3 +6500,69 @@ def test_transform_request_reads_service_tier_from_responses_api_request() -> No ) assert result["service_tier"] == "priority" + + +def _continued_chain_request(request_kwargs: Mapping[str, object]) -> Mapping[str, object]: + async def session_lookup(previous_response_id: str) -> ChatCompletionSession: + assert previous_response_id == "resp_turn_1" + return ChatCompletionSession( + messages=[{"role": "user", "content": "turn 1"}], litellm_session_id="chain-session" + ) + + return asyncio.run( + LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( + previous_response_id="resp_turn_1", + litellm_completion_request={ + "model": _CLAUDE_ON_BEDROCK, + "custom_llm_provider": "bedrock", + "messages": [{"role": "user", "content": "turn 2"}], + "metadata": request_kwargs.get("metadata"), + }, + instructions=None, + session_lookup=session_lookup, + request_kwargs=request_kwargs, + ) + ) + + +@pytest.mark.parametrize( + ("request_kwargs", "logged_metadata_key"), + [ + ( + { + "metadata": {"tag": "client"}, + "litellm_metadata": {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"}, + }, + "litellm_metadata", + ), + ({"metadata": {"user_api_key": "hashed"}}, "metadata"), + ], +) +def test_a_continued_chain_logs_its_spend_under_the_chain_session_with_request_correlation( + monkeypatch: pytest.MonkeyPatch, request_kwargs: Mapping[str, object], logged_metadata_key: str +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + sent_kwargs: Final = deepcopy(request_kwargs) + + request: Final = _continued_chain_request(request_kwargs) + + assert request["litellm_trace_id"] == "chain-session" + assert "litellm_session_id" not in request + assert request[logged_metadata_key] == { + **TypeAdapter(dict[str, object]).validate_python(request_kwargs[logged_metadata_key]), + CALLER_SESSION_ID_METADATA_KEY: "chain-session", + } + assert request_kwargs == sent_kwargs + + +def test_a_continued_chain_leaves_the_turns_metadata_alone_without_request_correlation( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + caller_metadata: Final = {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"} + + request: Final = _continued_chain_request({"metadata": caller_metadata, "litellm_metadata": caller_metadata}) + + assert request["litellm_trace_id"] == "chain-session" + assert "litellm_metadata" not in request + assert request["metadata"] == {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"} diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 002a595a5e7..5db563be4b0 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1,4 +1,5 @@ import json +import sqlite3 from typing import Final from unittest.mock import AsyncMock, patch @@ -787,3 +788,80 @@ async def test_message_history_replays_real_key_named_tool_payloads() -> None: tool_message: Final = result["messages"][2] assert json.loads(assistant_message["tool_calls"][0]["function"]["arguments"]) == function_arguments assert json.loads(tool_message["content"]) == function_output + + +class _SqliteBackedPrismaDB: + def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None: + self._connection = sqlite3.connect(":memory:") + self._connection.row_factory = sqlite3.Row + self._connection.execute( + 'CREATE TABLE "LiteLLM_SpendLogs" ' + '(request_id TEXT, api_key TEXT, session_id TEXT, "endTime" TEXT, "user" TEXT, team_id TEXT)' + ) + self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows) + + async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]: + return [dict(row) for row in self._connection.execute(query.replace("$1", "?"), args)] + + +@pytest.mark.asyncio +async def test_session_lookup_does_not_return_another_keys_rows_for_a_shared_session_id(): + """A caller-supplied litellm_session_id can collide with another key's session; history must stay per key.""" + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("victim-1", "victim-key", "shared-session", "2026-01-01T00:00:01", "", ""), + ("victim-2", "victim-key", "shared-session", "2026-01-01T00:00:02", "", ""), + ("attacker-1", "attacker-key", "shared-session", "2026-01-01T00:00:03", "", ""), + ("attacker-2", "attacker-key", "other-session", "2026-01-01T00:00:04", "", ""), + ] + ) + fake_prisma_client = _FakePrismaClient(results=[]) + fake_prisma_client.db = sqlite_db + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id("attacker-1") + victim_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id("victim-1") + + assert [row["request_id"] for row in spend_logs] == ["attacker-1"] + assert [row["request_id"] for row in victim_logs] == ["victim-1", "victim-2"] + + +async def _history_request_ids(sqlite_db: _SqliteBackedPrismaDB, response_id: str) -> list[str]: + fake_prisma_client: Final = _FakePrismaClient(results=[]) + fake_prisma_client.db = sqlite_db + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + spend_logs: Final = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(response_id) + return [row["request_id"] for row in spend_logs] + + +@pytest.mark.asyncio +async def test_session_lookup_keeps_jwt_users_and_teams_apart_when_they_share_an_empty_api_key(): + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("alice-1", "", "shared-session", "2026-01-01T00:00:01", "alice", "team-a"), + ("bob-1", "", "shared-session", "2026-01-01T00:00:02", "bob", "team-a"), + ("service-a-1", "", "shared-session", "2026-01-01T00:00:03", "", "team-a"), + ("service-b-1", "", "shared-session", "2026-01-01T00:00:04", "", "team-b"), + ("alice-2", "", "shared-session", "2026-01-01T00:00:05", "alice", "team-a"), + ] + ) + + assert await _history_request_ids(sqlite_db, "alice-1") == ["alice-1", "alice-2"] + assert await _history_request_ids(sqlite_db, "bob-1") == ["bob-1"] + assert await _history_request_ids(sqlite_db, "service-a-1") == ["service-a-1"] + assert await _history_request_ids(sqlite_db, "service-b-1") == ["service-b-1"] + + +@pytest.mark.asyncio +async def test_session_lookup_follows_the_owner_across_keys_and_keeps_ownerless_keys_apart(): + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("alice-old-key", "hashed-old", "shared-session", "2026-01-01T00:00:01", "alice", "team-a"), + ("ownerless-1", "hashed-other", "shared-session", "2026-01-01T00:00:02", "", ""), + ("alice-new-key", "hashed-new", "shared-session", "2026-01-01T00:00:03", "alice", "team-a"), + ("ownerless-2", "hashed-another", "shared-session", "2026-01-01T00:00:04", "", ""), + ] + ) + + assert await _history_request_ids(sqlite_db, "alice-new-key") == ["alice-old-key", "alice-new-key"] + assert await _history_request_ids(sqlite_db, "ownerless-2") == ["ownerless-2"]