From e046aee3d52e2308d94399b033893fe77674ee50 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:57:56 -0700 Subject: [PATCH] fix(spend_tracking): add missing_session_id: omit to leave SpendLogs.session_id null without a client session (#39458) * fix(spend_tracking): leave SpendLogs.session_id null when no client session id was established Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(lint): ratchet basedpyright budget after session_id fix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_tracking): ignore trace ids as session ids Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_tracking): gate null SpendLogs.session_id behind missing_session_id: omit Unset, generate and reject keep the legacy trace id fallback. omit records only metadata.session_id, the key Langfuse reads, so a trace id copied into litellm_session_id by get_litellm_params never becomes a session. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_tracking): stamp the omit decision on the request so a config reload cannot fabricate a session Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_tracking): keep omit covering requests the pre-call stamp never reaches Router-model provider pass-through calls allm_passthrough_route directly and skips add_litellm_data_to_request, so those requests never run the pre-call helper and carry no omit stamp. Reading only the stamp made POST /anthropic/v1/messages write a fabricated uuid into SpendLogs.session_id under missing_session_id: omit while its Langfuse trace had no session, the exact divergence the policy exists to remove. The stamp now only pins omit on, and an unstamped request falls back to the configured policy, so a config reload still cannot fabricate a session for a request that was decided pre-call. * fix(spend_tracking): make the session-omission marker proxy-owned so clients cannot forge it * fix(spend_tracking): strip the client-sent omission marker from both metadata buckets before they merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_tracking): strip the session-omission marker from both metadata buckets The pre-call policy ran before litellm_metadata is merged into metadata, so a client that planted the marker in litellm_metadata had it copied back into the route's own bucket after the strip and still got a null SpendLogs.session_id. --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- basedpyright-code-budget.json | 6 +- litellm/constants.py | 1 + litellm/proxy/_types.py | 4 +- .../proxy/hooks/proxy_track_cost_callback.py | 7 +- litellm/proxy/litellm_pre_call_utils.py | 12 +- .../pass_through_endpoints.py | 6 +- .../spend_tracking/spend_tracking_utils.py | 38 +- .../test_pass_through_endpoints.py | 35 + .../test_spend_tracking_utils.py | 477 ++++++------- .../proxy/test_litellm_pre_call_utils.py | 649 ++++++------------ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 11 files changed, 519 insertions(+), 720 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 2967fc2a505..a37cb194757 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15288 + "limit": 15287 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,10 +105,10 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38324 + "limit": 38323 }, "reportUnknownParameterType": { - "limit": 19625 + "limit": 19624 }, "reportUnknownVariableType": { "limit": 29861 diff --git a/litellm/constants.py b/litellm/constants.py index ef9329b9dfc..be13d9aac5f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1450,6 +1450,7 @@ SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affin CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" +SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 97f5f59d2dc..0aea72be1e2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2608,9 +2608,9 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.", ) - missing_session_id: Literal["generate", "reject"] | None = Field( + missing_session_id: Literal["generate", "reject", "omit"] | None = Field( None, - description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.", + description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400; 'omit' leaves SpendLogs.session_id null, matching callbacks such as Langfuse that only record a client-established metadata.session_id. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.", ) enable_public_model_hub: bool = Field( default=False, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 47aafda2337..7254b05db2e 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -168,11 +168,8 @@ class _ProxyDBLogger(CustomLogger): "custom_llm_provider" ) or request_data.get("custom_llm_provider", "") - # Propagate standard_logging_object and litellm_trace_id from the - # Logging instance so that _get_session_id_for_spend_log uses the same - # trace_id that Langfuse received (via async_failure_handler). - # Without this, the DB session_id would be a random UUID that doesn't - # match the Langfuse trace_id, making failed requests unsearchable. + # Propagate standard_logging_object and litellm_trace_id from the Logging + # instance so the failure row carries the same trace_id Langfuse received. _litellm_logging_obj: Final = request_data.get("litellm_logging_obj") if _litellm_logging_obj is not None: if not request_data.get("standard_logging_object"): diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 1d440448c2f..f752d7cfa89 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -25,6 +25,7 @@ from litellm.constants import ( PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, + SESSION_ID_OMITTED_METADATA_KEY, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -733,12 +734,18 @@ def apply_missing_session_id_policy( general_settings: Mapping[str, object] | None, request: Request, ) -> None: + for metadata_key in ("metadata", "litellm_metadata"): + if isinstance(client_metadata := data.get(metadata_key), dict): + client_metadata.pop(SESSION_ID_OMITTED_METADATA_KEY, None) + metadata: Final = data.get(_metadata_variable_name) policy: Final = general_settings.get("missing_session_id") if general_settings else None if policy is None or not _is_llm_inference_route(request): return - metadata: Final = data.get(_metadata_variable_name) if not isinstance(metadata, dict): return + if policy == "omit": + metadata[SESSION_ID_OMITTED_METADATA_KEY] = True + return if data.get("litellm_session_id") or metadata.get("session_id"): return match policy: @@ -760,7 +767,8 @@ def apply_missing_session_id_policy( ) case _: verbose_proxy_logger.warning( - "Ignoring unknown general_settings.missing_session_id=%r; expected 'generate' or 'reject'", policy + "Ignoring unknown general_settings.missing_session_id=%r; expected 'generate', 'reject' or 'omit'", + policy, ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 79d5d0a016f..323756bf204 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -40,6 +40,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( MAXIMUM_TRACEBACK_LINES_TO_LOG, + SESSION_ID_OMITTED_METADATA_KEY, WEBSOCKET_CLOSE_REASON_MAX_BYTES, ) from litellm.integrations.custom_guardrail import CustomGuardrail @@ -581,8 +582,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): ) # Set internal keys after merging client-supplied metadata so a request - # body that mirrors them cannot clobber the authenticated key or the - # real parent span. + # body that mirrors them cannot clobber the authenticated key, the real + # parent span, or the proxy's own session-id decision. + _metadata.pop(SESSION_ID_OMITTED_METADATA_KEY, None) _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span _metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 7442d71bd96..a37c3ba4405 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -15,6 +15,7 @@ from litellm.constants import ( LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, REDACTED_BY_LITELM_STRING, + SESSION_ID_OMITTED_METADATA_KEY, ) from litellm.constants import ( MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB, @@ -578,7 +579,9 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs ), session_id=_get_session_id_for_spend_log( kwargs=kwargs, + metadata=metadata, standard_logging_payload=standard_logging_payload, + omit_when_missing=_omits_session_id_when_missing(metadata), ), request_duration_ms=_get_request_duration_ms(start_time, end_time), status=_get_status_for_spend_log( @@ -602,26 +605,39 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs raise e +def _omits_session_id_when_missing(metadata: Mapping[str, object] | None) -> bool: + """The pre-call stamp pins `omit` on for the requests that carry it, so a config reload between pre-call and spend + logging cannot fabricate a session. `apply_missing_session_id_policy` drops any client-supplied copy of the key + from both metadata buckets before stamping, which the merge of `litellm_metadata` into `metadata` makes + necessary, so a caller cannot forge it. Requests that never reach the pre-call helper, router-model + passthrough among them, carry no stamp, so they fall back to the configured policy and `omit` still covers their + spend logs.""" + if metadata is not None and metadata.get(SESSION_ID_OMITTED_METADATA_KEY): + return True + + from litellm.proxy.proxy_server import general_settings + + return general_settings.get("missing_session_id") == "omit" + + def _get_session_id_for_spend_log( - kwargs: dict, + kwargs: Mapping[str, object], + metadata: Mapping[str, object] | None, standard_logging_payload: StandardLoggingPayload | None, -) -> str: - """ - Get the session id for the spend log. + omit_when_missing: bool, +) -> 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.""" + if omit_when_missing: + session_id: Final = metadata.get("session_id") if metadata else None + return str(session_id) if session_id else None - This ensures each spend log is associated with a unique session id. - - """ from litellm._uuid import uuid 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")) - - # Users can dynamically set the trace_id for each request by passing `litellm_trace_id` in kwargs if kwargs.get("litellm_trace_id") is not None: return str(kwargs.get("litellm_trace_id")) - - # Ensure we always have a session id, if none is provided return str(uuid.uuid4()) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d3f17c73499..91367507247 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -5730,3 +5730,38 @@ async def test_pass_through_request_leaves_cost_router_logger_working(): verbose_logger.removeHandler(recorder) assert not raised, f"cost router logger raised on the passthrough logging params: {raised[0].exc_info}" + + +@pytest.mark.parametrize("client_metadata_key", ["litellm_metadata", "metadata"]) +def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key: str): + """The omit marker is proxy-owned: only the pre-call policy may set it. A pass-through body that carries + it in its own metadata must not null out SpendLogs.session_id on a request the proxy never omitted.""" + from litellm.constants import SESSION_ID_OMITTED_METADATA_KEY + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_session_id_for_spend_log + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://0.0.0.0:4000/gemini/v1beta/models/gemini-2.5-flash:generateContent" + mock_request.headers = Headers({}) + mock_request.scope = {} + + kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + request=mock_request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + passthrough_logging_payload=MagicMock(), + logging_obj=MagicMock(), + _parsed_body={client_metadata_key: {SESSION_ID_OMITTED_METADATA_KEY: True}}, + litellm_call_id="lit-6694-call-id", + ) + + metadata = kwargs["litellm_params"]["metadata"] + assert SESSION_ID_OMITTED_METADATA_KEY not in metadata + assert ( + _get_session_id_for_spend_log( + kwargs={}, + metadata=metadata, + standard_logging_payload={"trace_id": "per-call-random-trace-id"}, + omit_when_missing=bool(metadata.get(SESSION_ID_OMITTED_METADATA_KEY)), + ) + == "per-call-random-trace-id" + ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 9e5917637a8..323930eee60 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -14,6 +14,7 @@ from litellm.constants import ( LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, REDACTED_BY_LITELM_STRING, + SESSION_ID_OMITTED_METADATA_KEY, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.spend_tracking.spend_tracking_utils import ( @@ -21,6 +22,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_proxy_server_request_for_spend_logs_payload, _get_request_duration_ms, _get_response_for_spend_logs_payload, + _get_session_id_for_spend_log, _get_spend_logs_metadata, _get_vector_store_request_for_spend_logs_payload, _is_master_key, @@ -33,6 +35,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_logging_payload, get_spend_logs_id, ) +from litellm.proxy._types import SpendLogsPayload from litellm.proxy.utils import hash_token from litellm.types.utils import ( StandardLoggingHiddenParams, @@ -74,6 +77,110 @@ def test_get_logging_payload_maps_openai_cached_tokens_to_cache_read_input_token assert additional_usage_values["prompt_tokens_details"]["cached_tokens"] == 123 +_TRACE_ONLY_STANDARD_LOGGING: Final = cast( + StandardLoggingPayload, + { + "trace_id": "trace-abc", + "session_id": "trace-abc", + "metadata": {}, + "model_map_information": None, + "request_tags": [], + }, +) + + +def _trace_only_session_id(omit_when_missing: bool) -> str | None: + """get_litellm_params copies metadata.trace_id into litellm_session_id, so every field echoes the trace id.""" + return _get_session_id_for_spend_log( + kwargs={"litellm_trace_id": "trace-abc", "litellm_session_id": "trace-abc"}, + metadata={"trace_id": "trace-abc"}, + standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING, + omit_when_missing=omit_when_missing, + ) + + +def test_omit_leaves_session_id_none_when_only_a_trace_id_exists(): + assert _trace_only_session_id(omit_when_missing=True) is None + + +def test_omit_leaves_session_id_none_without_any_ids(): + assert ( + _get_session_id_for_spend_log(kwargs={}, metadata=None, standard_logging_payload=None, omit_when_missing=True) + is None + ) + + +def test_omit_records_metadata_session_id(): + session_id: Final = _get_session_id_for_spend_log( + kwargs={"litellm_session_id": "chain-1"}, + metadata={"trace_id": "chain-1", "session_id": "chain-1"}, + standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING, + omit_when_missing=True, + ) + assert session_id == "chain-1" + + +def test_legacy_policy_keeps_trace_id_fallback(): + assert _trace_only_session_id(omit_when_missing=False) == "trace-abc" + generated: Final = _get_session_id_for_spend_log( + kwargs={}, metadata=None, standard_logging_payload=None, omit_when_missing=False + ) + assert len(str(generated)) == 36 + + +@pytest.mark.parametrize( + ("request_metadata", "expected"), + [ + ({"trace_id": "trace-abc"}, "trace-abc"), + ({"trace_id": "trace-abc", SESSION_ID_OMITTED_METADATA_KEY: True}, None), + ({"trace_id": "trace-abc", "session_id": "chain-1", SESSION_ID_OMITTED_METADATA_KEY: True}, "chain-1"), + ], +) +def test_get_logging_payload_reads_omit_decision_stamped_on_request( + request_metadata: dict[str, object], expected: str | None +): + """The pre-call stamp, not the live general_settings, decides the policy, so a config reload between + pre-call and spend logging cannot fabricate a session for a request accepted under `omit`.""" + with patch( # test-quality-ok: proves log time ignores proxy config; general_settings is yaml, not an HTTP boundary + "litellm.proxy.proxy_server.general_settings", {"missing_session_id": "generate"} + ): + payload: SpendLogsPayload = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "litellm_trace_id": "trace-abc", + "litellm_params": {"litellm_session_id": "trace-abc", "metadata": request_metadata}, + "standard_logging_object": _TRACE_ONLY_STANDARD_LOGGING, + }, + response_obj=litellm.ModelResponse(id="chatcmpl-test", choices=[]), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["session_id"] == expected + + +@pytest.mark.parametrize("policy", ["omit", "generate", None]) +def test_get_logging_payload_applies_omit_to_requests_that_carry_no_stamp(policy: str | None): + """Router-model passthrough calls `allm_passthrough_route` directly and never reaches the pre-call helper that + stamps the omit decision, so an unstamped request falls back to the configured policy. Without that fallback + `missing_session_id: omit` would fabricate a uuid session id on every passthrough spend log while its Langfuse + trace has none, which is the divergence the policy exists to remove.""" + with patch( # test-quality-ok: general_settings is proxy config, loaded from yaml, not an HTTP boundary + "litellm.proxy.proxy_server.general_settings", {} if policy is None else {"missing_session_id": policy} + ): + payload: SpendLogsPayload = get_logging_payload( + kwargs={ + "model": "claude-opus-4", + "litellm_trace_id": "trace-abc", + "litellm_params": {"litellm_session_id": "trace-abc", "metadata": {"trace_id": "trace-abc"}}, + "standard_logging_object": _TRACE_ONLY_STANDARD_LOGGING, + }, + response_obj=litellm.ModelResponse(id="chatcmpl-test", choices=[]), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["session_id"] == (None if policy == "omit" else "trace-abc") + + def test_get_logging_payload_preserves_anthropic_cache_read_input_tokens(): additional_usage_values = _get_additional_usage_values_for_usage( litellm.Usage( @@ -277,9 +384,7 @@ def test_sanitize_request_body_for_spend_logs_payload_long_string(): from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB # Create a string longer than MAX_STRING_LENGTH_PROMPT_IN_DB (2048) - long_string = ( - "a" * 3000 - ) # Create a string longer than MAX_STRING_LENGTH_PROMPT_IN_DB + long_string = "a" * 3000 # Create a string longer than MAX_STRING_LENGTH_PROMPT_IN_DB request_body = {"text": long_string, "normal_text": "short text"} sanitized = _sanitize_request_body_for_spend_logs_payload(request_body) @@ -329,9 +434,7 @@ def test_sanitize_request_body_for_spend_logs_payload_nested_list(): # Create a string longer than MAX_STRING_LENGTH_PROMPT_IN_DB long_string = "a" * (MAX_STRING_LENGTH_PROMPT_IN_DB + 500) - request_body = { - "items": [{"text": long_string}, {"text": "short"}, [{"text": long_string}]] - } + request_body = {"items": [{"text": long_string}, {"text": "short"}, [{"text": long_string}]]} sanitized = _sanitize_request_body_for_spend_logs_payload(request_body) # Calculate expected lengths based on actual MAX_STRING_LENGTH_PROMPT_IN_DB @@ -415,14 +518,10 @@ def test_sanitize_request_body_for_spend_logs_payload_circular_reference(): # Test that it handles circular reference without infinite recursion sanitized = _sanitize_request_body_for_spend_logs_payload(a) - assert sanitized == { - "b": {"a": {}} - } # Should return empty dict for circular reference + assert sanitized == {"b": {"a": {}}} # Should return empty dict for circular reference -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_vector_store_request_for_spend_logs_payload_store_prompts_true( mock_should_store, ): @@ -431,27 +530,16 @@ def test_get_vector_store_request_for_spend_logs_payload_store_prompts_true( # Sample vector store request metadata vector_store_request = [ - { - "vector_store_search_response": { - "data": [ - {"content": [{"text": "sensitive information", "type": "text"}]} - ] - } - } + {"vector_store_search_response": {"data": [{"content": [{"text": "sensitive information", "type": "text"}]}]}} ] # When store_prompts is True, the original data should be returned unchanged result = _get_vector_store_request_for_spend_logs_payload(vector_store_request) assert result == vector_store_request - assert ( - result[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] - == "sensitive information" - ) + assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] == "sensitive information" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_vector_store_request_for_spend_logs_payload_store_prompts_false( mock_should_store, ): @@ -460,32 +548,18 @@ def test_get_vector_store_request_for_spend_logs_payload_store_prompts_false( # Sample vector store request metadata vector_store_request = [ - { - "vector_store_search_response": { - "data": [ - {"content": [{"text": "sensitive information", "type": "text"}]} - ] - } - } + {"vector_store_search_response": {"data": [{"content": [{"text": "sensitive information", "type": "text"}]}]}} ] # When store_prompts is False, text should be redacted result = _get_vector_store_request_for_spend_logs_payload(vector_store_request) assert result is not None - assert ( - result[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] - == REDACTED_BY_LITELM_STRING - ) + assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] == REDACTED_BY_LITELM_STRING # Ensure other fields are unchanged - assert ( - result[0]["vector_store_search_response"]["data"][0]["content"][0]["type"] - == "text" - ) + assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["type"] == "text" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_vector_store_request_for_spend_logs_payload_null_input(mock_should_store): # When input is None mock_should_store.return_value = False @@ -493,9 +567,7 @@ def test_get_vector_store_request_for_spend_logs_payload_null_input(mock_should_ assert result is None -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_realtime_returns_messages(mock_should_store): """ Test that _get_messages_for_spend_logs_payload returns messages @@ -522,9 +594,7 @@ def test_get_messages_for_spend_logs_realtime_returns_messages(mock_should_store assert parsed[1]["content"] == "What is the weather today?" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_strips_null_bytes(mock_should_store): """Regression for PostgreSQL 22P05: NUL bytes must be stripped from messages.""" mock_should_store.return_value = True @@ -541,9 +611,7 @@ def test_get_messages_for_spend_logs_strips_null_bytes(mock_should_store): assert parsed[0]["content"] == "helloworld" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_realtime_empty_when_disabled(mock_should_store): """ Test that _get_messages_for_spend_logs_payload returns '{}' for realtime calls @@ -561,9 +629,7 @@ def test_get_messages_for_spend_logs_realtime_empty_when_disabled(mock_should_st assert result == "{}" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_non_realtime_returns_empty(mock_should_store): """ Test that _get_messages_for_spend_logs_payload returns '{}' for non-realtime @@ -581,9 +647,7 @@ def test_get_messages_for_spend_logs_non_realtime_returns_empty(mock_should_stor assert result == "{}" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_truncates_large_base64(mock_should_store): from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB @@ -611,9 +675,7 @@ def test_get_response_for_spend_logs_payload_truncates_large_base64(mock_should_ assert parsed["data"][0]["other_field"] == "value" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_strips_null_bytes(mock_should_store): """Regression for PostgreSQL 22P05: NUL bytes must be stripped from response.""" mock_should_store.return_value = True @@ -626,18 +688,14 @@ def test_get_response_for_spend_logs_payload_strips_null_bytes(mock_should_store assert json.loads(response_json)["content"] == "answerhere" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_truncates_large_embedding( mock_should_store, ): from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB mock_should_store.return_value = True - embedding_values = [ - round(i * 0.0001, 6) for i in range(MAX_STRING_LENGTH_PROMPT_IN_DB + 500) - ] + embedding_values = [round(i * 0.0001, 6) for i in range(MAX_STRING_LENGTH_PROMPT_IN_DB + 500)] large_embedding = json.dumps(embedding_values) payload = cast( StandardLoggingPayload, @@ -685,9 +743,7 @@ def test_truncation_includes_db_safeguard_note(): ) -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_response_truncation_logs_info_message(mock_should_store): """ Test that when response is truncated before DB storage, an info log is emitted @@ -702,18 +758,14 @@ def test_response_truncation_logs_info_message(mock_should_store): {"response": {"data": [{"content": large_text}]}}, ) - with patch( - "litellm.proxy.spend_tracking.spend_tracking_utils.verbose_proxy_logger" - ) as mock_logger: + with patch("litellm.proxy.spend_tracking.spend_tracking_utils.verbose_proxy_logger") as mock_logger: _get_response_for_spend_logs_payload(payload) mock_logger.info.assert_called_once() log_msg = mock_logger.info.call_args[0][0] assert "response was truncated" in log_msg -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_request_body_truncation_logs_info_message(mock_should_store): """ Test that when request body is truncated before DB storage, an info log is emitted. @@ -722,18 +774,10 @@ def test_request_body_truncation_logs_info_message(mock_should_store): mock_should_store.return_value = True large_prompt = "C" * (MAX_STRING_LENGTH_PROMPT_IN_DB + 500) - litellm_params = { - "proxy_server_request": { - "body": {"messages": [{"role": "user", "content": large_prompt}]} - } - } + litellm_params = {"proxy_server_request": {"body": {"messages": [{"role": "user", "content": large_prompt}]}}} - with patch( - "litellm.proxy.spend_tracking.spend_tracking_utils.verbose_proxy_logger" - ) as mock_logger: - _get_proxy_server_request_for_spend_logs_payload( - metadata={}, litellm_params=litellm_params, kwargs={} - ) + with patch("litellm.proxy.spend_tracking.spend_tracking_utils.verbose_proxy_logger") as mock_logger: + _get_proxy_server_request_for_spend_logs_payload(metadata={}, litellm_params=litellm_params, kwargs={}) mock_logger.info.assert_called_once() log_msg = mock_logger.info.call_args[0][0] assert "request body was truncated" in log_msg @@ -870,14 +914,10 @@ def test_get_logging_payload_api_key_preserved_when_standard_logging_payload_is_ ) # The api_key should be hashed (not the raw key) - assert ( - payload["api_key"] != test_api_key - ), "api_key should be hashed, not the raw key" + assert payload["api_key"] != test_api_key, "api_key should be hashed, not the raw key" # The api_key should be a valid hash (64 character hex string for SHA256) - assert ( - len(payload["api_key"]) == 64 - ), f"Expected 64 character hash, got {len(payload['api_key'])} characters" + assert len(payload["api_key"]) == 64, f"Expected 64 character hash, got {len(payload['api_key'])} characters" # Verify other fields are set correctly assert payload["model"] == "openai/gpt-4.1" @@ -1019,9 +1059,7 @@ async def test_api_key_preserved_through_failure_hook_to_database(): assert payload_api_key is not None, "🚨 CRITICAL: payload['api_key'] is None!" - assert ( - payload_api_key == hashed_key - ), f"🚨 CRITICAL: Expected api_key={hashed_key}, got {payload_api_key}" + assert payload_api_key == hashed_key, f"🚨 CRITICAL: Expected api_key={hashed_key}, got {payload_api_key}" # Verify token parameter matches assert data["token"] == hashed_key, f"Token parameter should be {hashed_key}" @@ -1066,9 +1104,7 @@ def test_get_logging_payload_includes_agent_id_from_kwargs(): end_time=end_time, ) - assert ( - payload["agent_id"] == test_agent_id - ), f"Expected agent_id '{test_agent_id}', got '{payload.get('agent_id')}'" + assert payload["agent_id"] == test_agent_id, f"Expected agent_id '{test_agent_id}', got '{payload.get('agent_id')}'" @patch("litellm.proxy.proxy_server.master_key", None) @@ -1093,9 +1129,7 @@ def test_get_logging_payload_includes_overhead_in_spend_logs_metadata(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=None, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai", @@ -1173,9 +1207,9 @@ def test_get_logging_payload_includes_overhead_in_spend_logs_metadata(): metadata = json.loads(metadata_json) # Verify overhead is stored directly in metadata - assert ( - metadata.get("litellm_overhead_time_ms") == test_overhead_ms - ), f"Expected overhead '{test_overhead_ms}', got '{metadata.get('litellm_overhead_time_ms')}'" + assert metadata.get("litellm_overhead_time_ms") == test_overhead_ms, ( + f"Expected overhead '{test_overhead_ms}', got '{metadata.get('litellm_overhead_time_ms')}'" + ) @patch("litellm.proxy.proxy_server.master_key", None) @@ -1228,9 +1262,7 @@ def test_get_logging_payload_handles_missing_overhead_gracefully(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=None, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai", @@ -1309,14 +1341,12 @@ def test_get_logging_payload_handles_missing_overhead_gracefully(): metadata = json.loads(metadata_json) # When overhead is None, litellm_overhead_time_ms should be None or not present - assert ( - metadata.get("litellm_overhead_time_ms") is None - ), "litellm_overhead_time_ms should be None when overhead is not provided" + assert metadata.get("litellm_overhead_time_ms") is None, ( + "litellm_overhead_time_ms should be None when overhead is not provided" + ) -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_spend_logs_redacts_request_and_response_when_turn_off_message_logging_enabled( mock_should_store, ): @@ -1347,9 +1377,7 @@ def test_spend_logs_redacts_request_and_response_when_turn_off_message_logging_e ) parsed_request = json.loads(request_result) - assert parsed_request["messages"] == [ - {"role": "user", "content": "redacted-by-litellm"} - ] + assert parsed_request["messages"] == [{"role": "user", "content": "redacted-by-litellm"}] assert parsed_request["model"] == "gpt-4" # Test response redaction - use dict response to verify redaction @@ -1368,9 +1396,7 @@ def test_spend_logs_redacts_request_and_response_when_turn_off_message_logging_e {"response": response_dict}, ) - response_result = _get_response_for_spend_logs_payload( - payload=payload, kwargs=kwargs - ) + response_result = _get_response_for_spend_logs_payload(payload=payload, kwargs=kwargs) # When redaction is enabled and response is a dict (not ModelResponse), # perform_redaction redacts content in-place within the choices structure @@ -1415,30 +1441,22 @@ def test_should_store_prompts_and_responses_in_spend_logs_case_insensitive_strin # When env var is True, should return True mock_get_secret_bool.return_value = True result = _should_store_prompts_and_responses_in_spend_logs() - assert ( - result is True - ), f"Expected True (from env var) for '{false_value}', got {result}" + assert result is True, f"Expected True (from env var) for '{false_value}', got {result}" # When env var is False, should return False mock_get_secret_bool.return_value = False result = _should_store_prompts_and_responses_in_spend_logs() - assert ( - result is False - ), f"Expected False (from env var) for '{false_value}', got {result}" + assert result is False, f"Expected False (from env var) for '{false_value}', got {result}" # Test when general_settings doesn't have the key at all with patch("litellm.proxy.proxy_server.general_settings", {}): mock_get_secret_bool.return_value = True result = _should_store_prompts_and_responses_in_spend_logs() - assert ( - result is True - ), "Expected True (from env var) when key missing, got False" + assert result is True, "Expected True (from env var) when key missing, got False" mock_get_secret_bool.return_value = False result = _should_store_prompts_and_responses_in_spend_logs() - assert ( - result is False - ), "Expected False (from env var) when key missing, got True" + assert result is False, "Expected False (from env var) when key missing, got True" def test_get_spend_logs_metadata_guardrail_info_fallback_from_metadata(): @@ -1831,9 +1849,7 @@ def test_get_logging_payload_includes_retry_info_in_spend_logs_metadata(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=None, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai", @@ -1897,12 +1913,10 @@ def test_get_logging_payload_includes_retry_info_in_spend_logs_metadata(): metadata = json.loads(payload["metadata"]) - assert ( - metadata.get("attempted_retries") == 2 - ), f"Expected attempted_retries=2, got {metadata.get('attempted_retries')}" - assert ( - metadata.get("max_retries") == 3 - ), f"Expected max_retries=3, got {metadata.get('max_retries')}" + assert metadata.get("attempted_retries") == 2, ( + f"Expected attempted_retries=2, got {metadata.get('attempted_retries')}" + ) + assert metadata.get("max_retries") == 3, f"Expected max_retries=3, got {metadata.get('max_retries')}" @patch("litellm.proxy.proxy_server.master_key", None) @@ -1930,9 +1944,7 @@ def test_get_logging_payload_handles_missing_retry_info_gracefully(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=None, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai", @@ -1996,20 +2008,14 @@ def test_get_logging_payload_handles_missing_retry_info_gracefully(): metadata = json.loads(payload["metadata"]) - assert ( - metadata.get("attempted_retries") is None - ), "attempted_retries should be None when not provided" - assert ( - metadata.get("max_retries") is None - ), "max_retries should be None when not provided" + assert metadata.get("attempted_retries") is None, "attempted_retries should be None when not provided" + assert metadata.get("max_retries") is None, "max_retries should be None when not provided" def test_get_request_duration_ms_normal(): """Test that request duration is correctly computed in milliseconds.""" start = datetime.datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc) - end = datetime.datetime( - 2025, 1, 1, 0, 0, 2, 500000, tzinfo=timezone.utc - ) # 2.5s later + end = datetime.datetime(2025, 1, 1, 0, 0, 2, 500000, tzinfo=timezone.utc) # 2.5s later result = _get_request_duration_ms(start, end) assert result == 2500 @@ -2039,9 +2045,7 @@ def test_get_logging_payload_includes_request_duration_ms(): "litellm_params": {"api_base": "https://api.openai.com"}, "standard_logging_object": None, } - response_obj = { - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} - } + response_obj = {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}} with ( patch("litellm.proxy.proxy_server.master_key", None), @@ -2107,16 +2111,12 @@ def test_sanitize_request_body_strips_secret_fields(): } sanitized = _sanitize_request_body_for_spend_logs_payload(request_body) - assert ( - "secret_fields" not in sanitized - ), "secret_fields must be stripped from the sanitized request body" + assert "secret_fields" not in sanitized, "secret_fields must be stripped from the sanitized request body" assert sanitized["model"] == "gpt-4" assert sanitized["messages"] == [{"role": "user", "content": "hi"}] -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_proxy_server_request_payload_excludes_secret_fields(mock_should_store): """ End-to-end test: when the proxy_server_request body contains @@ -2140,14 +2140,10 @@ def test_proxy_server_request_payload_excludes_secret_fields(mock_should_store): } } - result = _get_proxy_server_request_for_spend_logs_payload( - metadata={}, litellm_params=litellm_params, kwargs={} - ) + result = _get_proxy_server_request_for_spend_logs_payload(metadata={}, litellm_params=litellm_params, kwargs={}) parsed = json.loads(result) - assert ( - "secret_fields" not in parsed - ), "secret_fields must never appear in the spend-log proxy_server_request column" + assert "secret_fields" not in parsed, "secret_fields must never appear in the spend-log proxy_server_request column" assert parsed["model"] == "gpt-4" assert parsed["messages"] == [{"role": "user", "content": "hello"}] @@ -2176,10 +2172,7 @@ def test_redact_prompt_leaks_strips_input_value_python_repr(): def test_redact_prompt_leaks_strips_input_value_json(): - error_text = ( - '{"error":{"message":"validation failed",' - '"input":[{"role":"user","content":"top-secret-content"}]}}' - ) + error_text = '{"error":{"message":"validation failed","input":[{"role":"user","content":"top-secret-content"}]}}' redacted = _redact_prompt_leaks_in_error_string(error_text) assert "top-secret-content" not in redacted assert REDACTED_BY_LITELM_STRING in redacted @@ -2203,9 +2196,7 @@ def test_redact_prompt_leaks_empty_string(): assert _redact_prompt_leaks_in_error_string("") == "" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_redacts_when_not_storing_prompts( mock_should_store, ): @@ -2233,9 +2224,7 @@ def test_sanitize_error_information_redacts_when_not_storing_prompts( assert sanitized["llm_provider"] == "openai" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_skips_redaction_when_storing_prompts( mock_should_store, ): @@ -2246,9 +2235,7 @@ def test_sanitize_error_information_skips_redaction_when_storing_prompts( "error_class": "RateLimitError", "llm_provider": "openai", "traceback": "", - "error_message": ( - 'OpenAIException - {"error":{"input":[{"role":"user","content":"kept"}]}}' - ), + "error_message": ('OpenAIException - {"error":{"input":[{"role":"user","content":"kept"}]}}'), } sanitized = _sanitize_error_information_for_spend_logs(error_info) @@ -2259,9 +2246,7 @@ def test_sanitize_error_information_skips_redaction_when_storing_prompts( assert REDACTED_BY_LITELM_STRING not in sanitized["error_message"] -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_caps_size_regardless_of_prompt_flag( mock_should_store, ): @@ -2292,9 +2277,7 @@ def test_sanitize_error_information_none_passthrough(): assert _sanitize_error_information_for_spend_logs(None) is None -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_reproduces_lit_2992(mock_should_store): # Mirrors the reproduced row body from LIT-2992 — a RateLimitError whose # message embeds 178 pydantic validation errors, each carrying a full @@ -2335,10 +2318,7 @@ def test_redact_prompt_leaks_handles_nested_multimodal_content(): # Multi-modal payload: 'content' is itself a list. The depth-1 regex # would stop at the inner '['; the parser-based scanner must walk # through balanced nested brackets. - error_text = ( - '{"error":{"messages":[{"role":"user",' - '"content":[{"type":"text","text":"top-secret-multimodal"}]}]}}' - ) + error_text = '{"error":{"messages":[{"role":"user","content":[{"type":"text","text":"top-secret-multimodal"}]}]}}' redacted = _redact_prompt_leaks_in_error_string(error_text) assert "top-secret-multimodal" not in redacted assert REDACTED_BY_LITELM_STRING in redacted @@ -2347,9 +2327,7 @@ def test_redact_prompt_leaks_handles_nested_multimodal_content(): def test_redact_prompt_leaks_handles_bracket_in_prompt_text(): # Prompt text contains a literal '[' — the depth-1 regex would close # the outer ']' prematurely. The parser must respect string quoting. - error_text = ( - '{"error":{"input":[{"role":"user","content":"secret[123 still secret"}]}}' - ) + error_text = '{"error":{"input":[{"role":"user","content":"secret[123 still secret"}]}}' redacted = _redact_prompt_leaks_in_error_string(error_text) assert "secret[123" not in redacted assert "still secret" not in redacted @@ -2368,8 +2346,7 @@ def test_redact_prompt_leaks_handles_escaped_quote_in_prompt_text(): def test_redact_prompt_leaks_handles_nested_input_python_repr(): # Python dict-repr with nested list inside 'input' — single quotes. error_text = ( - "validation error: {'input': [{'role': 'user', " - "'content': [{'type': 'text', 'text': 'leaked-nested-text'}]}]}" + "validation error: {'input': [{'role': 'user', 'content': [{'type': 'text', 'text': 'leaked-nested-text'}]}]}" ) redacted = _redact_prompt_leaks_in_error_string(error_text) assert "leaked-nested-text" not in redacted @@ -2385,9 +2362,7 @@ def test_redact_prompt_leaks_handles_unterminated_value(): assert REDACTED_BY_LITELM_STRING in redacted -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_redacts_traceback_when_not_storing_prompts( mock_should_store, ): @@ -2419,9 +2394,7 @@ def test_sanitize_error_information_redacts_traceback_when_not_storing_prompts( assert "ValueError: invalid request" in sanitized["traceback"] -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_skips_traceback_redaction_when_storing_prompts( mock_should_store, ): @@ -2431,9 +2404,7 @@ def test_sanitize_error_information_skips_traceback_redaction_when_storing_promp "error_code": "500", "error_class": "ValueError", "llm_provider": "", - "traceback": ( - 'raise ValueError({"input":[{"role":"user","content":"tb-kept"}]})' - ), + "traceback": ('raise ValueError({"input":[{"role":"user","content":"tb-kept"}]})'), "error_message": "invalid request", } @@ -2448,20 +2419,14 @@ def test_redact_prompt_leaks_strips_prompt_key_completions_payload(): # /v1/completions echoes the user input under the top-level 'prompt' key # rather than 'messages'. Without 'prompt' coverage the body would survive # the redactor when store_prompts_in_spend_logs is False. - error_text = ( - '{"error":{"message":"validation failed",' - '"prompt":"super-secret-completion-text"}}' - ) + error_text = '{"error":{"message":"validation failed","prompt":"super-secret-completion-text"}}' redacted = _redact_prompt_leaks_in_error_string(error_text) assert "super-secret-completion-text" not in redacted assert REDACTED_BY_LITELM_STRING in redacted def test_redact_prompt_leaks_strips_prompt_key_python_repr(): - error_text = ( - "{'model': 'gpt-3.5-turbo-instruct', " - "'prompt': 'leaked-completion-prompt-body'}" - ) + error_text = "{'model': 'gpt-3.5-turbo-instruct', 'prompt': 'leaked-completion-prompt-body'}" redacted = _redact_prompt_leaks_in_error_string(error_text) assert "leaked-completion-prompt-body" not in redacted assert REDACTED_BY_LITELM_STRING in redacted @@ -2495,11 +2460,7 @@ def test_redact_prompt_leaks_strips_pydantic_input_value_list(): def test_redact_prompt_leaks_strips_pydantic_input_value_dict(): - error_text = ( - "[type=dict_type, " - "input_value={'role': 'user', 'content': 'leaked-dict-content'}, " - "input_type=dict]" - ) + error_text = "[type=dict_type, input_value={'role': 'user', 'content': 'leaked-dict-content'}, input_type=dict]" redacted = _redact_prompt_leaks_in_error_string(error_text) assert "leaked-dict-content" not in redacted assert REDACTED_BY_LITELM_STRING in redacted @@ -2541,9 +2502,7 @@ def test_redact_prompt_leaks_combined_quoted_key_and_pydantic_assignment(): assert redacted.count(REDACTED_BY_LITELM_STRING) >= 2 -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_redacts_pydantic_assignment_form( mock_should_store, ): @@ -2741,9 +2700,7 @@ def test_get_spend_logs_metadata_non_sk_raw_key_hashed(): def test_get_spend_logs_metadata_already_hashed_unchanged_with_provenance(): already_hashed = hash_token("sk-some-key") - meta = _get_spend_logs_metadata( - {"user_api_key": already_hashed, "user_api_key_hash": already_hashed} - ) + meta = _get_spend_logs_metadata({"user_api_key": already_hashed, "user_api_key_hash": already_hashed}) assert meta["user_api_key"] == already_hashed assert hash_token(already_hashed) != meta["user_api_key"] # no double-hash @@ -2758,9 +2715,7 @@ def test_get_spend_logs_metadata_already_hashed_no_provenance_is_rehashed(): def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match(): already_hashed = hash_token("sk-some-key") different_hash = hash_token("sk-other-key") - meta = _get_spend_logs_metadata( - {"user_api_key": already_hashed, "user_api_key_hash": different_hash} - ) + meta = _get_spend_logs_metadata({"user_api_key": already_hashed, "user_api_key_hash": different_hash}) assert meta["user_api_key"] == hash_token(already_hashed) @@ -2797,16 +2752,12 @@ def test_get_logging_payload_uses_recovered_combined_usage_on_failure(): "model": "anthropic/claude-haiku-4-5", "call_type": "acompletion", "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, - "combined_usage_object": Usage( - prompt_tokens=30, completion_tokens=1, total_tokens=31 - ), + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), } response_obj = Exception("MidStreamFallbackError: read timeout") now = datetime.datetime.now(timezone.utc) - payload = get_logging_payload( - kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now - ) + payload = get_logging_payload(kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now) assert payload["prompt_tokens"] == 30 assert payload["completion_tokens"] == 1 @@ -2825,9 +2776,7 @@ def test_get_logging_payload_failure_without_recovered_usage_is_zero(): response_obj = Exception("BadRequestError") now = datetime.datetime.now(timezone.utc) - payload = get_logging_payload( - kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now - ) + payload = get_logging_payload(kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now) assert payload["total_tokens"] == 0 @@ -2853,9 +2802,7 @@ def test_get_logging_payload_sets_litellm_call_id_for_correlation(): } now = datetime.datetime.now(timezone.utc) - payload = get_logging_payload( - kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now - ) + payload = get_logging_payload(kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now) metadata = json.loads(payload["metadata"]) assert payload["request_id"] == provider_response_id @@ -2882,9 +2829,7 @@ def test_get_logging_payload_litellm_call_id_falls_back_to_litellm_params(): } now = datetime.datetime.now(timezone.utc) - payload = get_logging_payload( - kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now - ) + payload = get_logging_payload(kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now) assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id @@ -2901,14 +2846,10 @@ def test_get_logging_payload_litellm_call_id_when_response_has_no_id(): "litellm_call_id": trace_call_id, "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, } - response_obj = { - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} - } + response_obj = {"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}} now = datetime.datetime.now(timezone.utc) - payload = get_logging_payload( - kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now - ) + payload = get_logging_payload(kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now) assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id assert payload["request_id"] == trace_call_id @@ -2932,9 +2873,7 @@ def test_get_logging_payload_cache_hit_keeps_raw_litellm_call_id(): } now = datetime.datetime.now(timezone.utc) - payload = get_logging_payload( - kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now - ) + payload = get_logging_payload(kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now) assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id assert "_cache_hit" in payload["request_id"] @@ -3074,9 +3013,7 @@ def test_get_logging_payload_hashes_bearer_prefixed_api_key(): assert not payload["api_key"].startswith("Bearer"), ( f"api_key column contains plaintext Bearer key: {payload['api_key']}" ) - assert not payload["api_key"].startswith("sk-"), ( - f"api_key column contains unhashed key: {payload['api_key']}" - ) + assert not payload["api_key"].startswith("sk-"), f"api_key column contains unhashed key: {payload['api_key']}" metadata_dict = json.loads(payload["metadata"]) assert not metadata_dict["user_api_key"].startswith("Bearer"), ( @@ -3747,9 +3684,7 @@ def test_get_logging_payload_includes_fallback_info_in_spend_logs_metadata(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=None, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai", @@ -3813,12 +3748,12 @@ def test_get_logging_payload_includes_fallback_info_in_spend_logs_metadata(): metadata = json.loads(payload["metadata"]) - assert ( - metadata.get("attempted_fallbacks") == 2 - ), f"Expected attempted_fallbacks=2, got {metadata.get('attempted_fallbacks')}" - assert ( - metadata.get("original_model_group") == "azure-gpt-fallback" - ), f"Expected original_model_group=azure-gpt-fallback, got {metadata.get('original_model_group')}" + assert metadata.get("attempted_fallbacks") == 2, ( + f"Expected attempted_fallbacks=2, got {metadata.get('attempted_fallbacks')}" + ) + assert metadata.get("original_model_group") == "azure-gpt-fallback", ( + f"Expected original_model_group=azure-gpt-fallback, got {metadata.get('original_model_group')}" + ) def test_get_logging_payload_handles_missing_fallback_info_gracefully(): @@ -3844,9 +3779,7 @@ def test_get_logging_payload_handles_missing_fallback_info_gracefully(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=None, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai", @@ -3910,12 +3843,10 @@ def test_get_logging_payload_handles_missing_fallback_info_gracefully(): metadata = json.loads(payload["metadata"]) - assert ( - metadata.get("attempted_fallbacks") is None - ), "attempted_fallbacks should be None when not provided" - assert ( - metadata.get("original_model_group") is None - ), "original_model_group should be None when not provided" + assert metadata.get("attempted_fallbacks") is None, "attempted_fallbacks should be None when not provided" + assert metadata.get("original_model_group") is None, "original_model_group should be None when not provided" + + @pytest.mark.parametrize("bucket", ["metadata", "litellm_metadata"]) def test_injected_cache_breakpoints_survive_into_spend_log_metadata(bucket): """The injection marker only gates savings if it reaches the spend-log row. diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 8366e5546a9..72d37650963 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -41,21 +41,18 @@ from litellm.litellm_core_utils.get_provider_specific_headers import ( from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( TRUSTED_CALLBACK_VARS_FIELD, ) -from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import CredentialItem - def test_check_if_token_is_service_account(): """ Test that only keys with `service_account_id` in metadata are considered service accounts """ # Test case 1: Service account token - service_account_token = UserAPIKeyAuth( - api_key="test-key", metadata={"service_account_id": "test-service-account"} - ) + service_account_token = UserAPIKeyAuth(api_key="test-key", metadata={"service_account_id": "test-service-account"}) assert check_if_token_is_service_account(service_account_token) == True # Test case 2: Regular user token @@ -63,9 +60,7 @@ def test_check_if_token_is_service_account(): assert check_if_token_is_service_account(regular_token) == False # Test case 3: Token with other metadata - other_metadata_token = UserAPIKeyAuth( - api_key="test-key", metadata={"user_id": "test-user"} - ) + other_metadata_token = UserAPIKeyAuth(api_key="test-key", metadata={"user_id": "test-user"}) assert check_if_token_is_service_account(other_metadata_token) == False @@ -112,15 +107,11 @@ class TestGetMetadataVariableName: def test_returns_litellm_metadata_for_bedrock_invoke(self): # GH#30629: bedrock passthrough must use litellm_metadata # to prevent key-level tags from leaking into provider body - request = self._make_request( - "/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke" - ) + request = self._make_request("/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke") assert _get_metadata_variable_name(request) == "litellm_metadata" def test_returns_litellm_metadata_for_bedrock_converse(self): - request = self._make_request( - "/bedrock/model/us.anthropic.claude-sonnet-4-6/converse" - ) + request = self._make_request("/bedrock/model/us.anthropic.claude-sonnet-4-6/converse") assert _get_metadata_variable_name(request) == "litellm_metadata" @@ -128,9 +119,7 @@ def test_get_enforced_params_for_service_account_settings(): """ Test that service account enforced params are only added to service account keys """ - service_account_token = UserAPIKeyAuth( - api_key="test-key", metadata={"service_account_id": "test-service-account"} - ) + service_account_token = UserAPIKeyAuth(api_key="test-key", metadata={"service_account_id": "test-service-account"}) general_settings_with_service_account_settings = { "service_account_settings": {"enforced_params": ["metadata.service"]}, } @@ -140,9 +129,7 @@ def test_get_enforced_params_for_service_account_settings(): ) assert result == ["metadata.service"] - regular_token = UserAPIKeyAuth( - api_key="test-key", metadata={"enforced_params": ["user"]} - ) + regular_token = UserAPIKeyAuth(api_key="test-key", metadata={"enforced_params": ["user"]}) result = _get_enforced_params( general_settings=general_settings_with_service_account_settings, user_api_key_dict=regular_token, @@ -155,9 +142,7 @@ def test_get_enforced_params_for_service_account_settings(): [ ( {"enforced_params": ["param1", "param2"]}, - UserAPIKeyAuth( - api_key="test_api_key", user_id="test_user_id", org_id="test_org_id" - ), + UserAPIKeyAuth(api_key="test_api_key", user_id="test_user_id", org_id="test_org_id"), ["param1", "param2"], ), ( @@ -183,9 +168,7 @@ def test_get_enforced_params_for_service_account_settings(): ), ], ) -def test_get_enforced_params( - general_settings, user_api_key_dict, expected_enforced_params -): +def test_get_enforced_params(general_settings, user_api_key_dict, expected_enforced_params): from litellm.proxy.litellm_pre_call_utils import _get_enforced_params enforced_params = _get_enforced_params(general_settings, user_api_key_dict) @@ -441,9 +424,7 @@ async def test_add_litellm_data_to_request_strips_admin_injection_slots(): populated = updated["metadata"] assert populated["user_api_key_metadata"] == real_admin_metadata assert populated["user_api_key_team_metadata"] == real_admin_metadata - assert "_pipeline_managed_guardrails" not in populated or populated[ - "_pipeline_managed_guardrails" - ] != ["evaded"] + assert "_pipeline_managed_guardrails" not in populated or populated["_pipeline_managed_guardrails"] != ["evaded"] other = updated.get("litellm_metadata") or {} assert other.get("user_api_key_metadata") in (None, {}, real_admin_metadata) @@ -697,9 +678,7 @@ async def test_add_litellm_data_to_request_proxy_server_request_body_is_post_str snapshot_body = updated["proxy_server_request"]["body"] assert snapshot_body is not None snapshot_metadata = snapshot_body.get("metadata") or {} - assert "user_api_key_user_id" not in snapshot_metadata or ( - snapshot_metadata["user_api_key_user_id"] != "victim" - ) + assert "user_api_key_user_id" not in snapshot_metadata or (snapshot_metadata["user_api_key_user_id"] != "victim") @pytest.mark.asyncio @@ -754,9 +733,7 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_secret_fields( ) # secret_fields must exist on the live data dict - assert ( - "secret_fields" in updated - ), "secret_fields must still be present on the live data dict" + assert "secret_fields" in updated, "secret_fields must still be present on the live data dict" assert "raw_headers" in updated["secret_fields"] # But the body snapshot must NOT contain secret_fields @@ -815,8 +792,7 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_r snapshot_body = updated["proxy_server_request"]["body"] assert "proxy_server_request" not in snapshot_body, ( - "proxy_server_request must be excluded from its own body snapshot " - "to prevent the body from self-referencing" + "proxy_server_request must be excluded from its own body snapshot to prevent the body from self-referencing" ) @@ -1344,23 +1320,18 @@ async def test_add_litellm_data_to_request_strips_client_redaction_bypass_contro assert "turn_off_message_logging" not in (updated.get("litellm_params") or {}).get("metadata", {}) assert "turn_off_message_logging" not in updated["metadata"] assert "turn_off_message_logging" not in (updated.get("litellm_metadata") or {}) + assert "litellm-disable-message-redaction" not in {header.lower() for header in updated["metadata"]["headers"]} assert "litellm-disable-message-redaction" not in { - header.lower() for header in updated["metadata"]["headers"] - } - assert "litellm-disable-message-redaction" not in { - header.lower() - for header in updated["metadata"]["requester_metadata"].get("headers", {}) + header.lower() for header in updated["metadata"]["requester_metadata"].get("headers", {}) } assert "litellm-disable-message-redaction" not in { header.lower() for header in updated["proxy_server_request"]["headers"] } assert "litellm-disable-message-redaction" not in { - header.lower() - for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] + header.lower() for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] } assert "litellm-disable-message-redaction" not in { - header.lower() - for header in (updated.get("litellm_metadata") or {}).get("headers", {}) + header.lower() for header in (updated.get("litellm_metadata") or {}).get("headers", {}) } @@ -1430,12 +1401,7 @@ async def test_add_litellm_data_to_request_admin_callback_vars_turn_off_message_ dynamic_params = initialize_standard_callback_dynamic_params(updated) assert dynamic_params.get("turn_off_message_logging") == "False" - assert ( - should_redact_message_logging( - {"standard_callback_dynamic_params": dynamic_params} - ) - is False - ) + assert should_redact_message_logging({"standard_callback_dynamic_params": dynamic_params}) is False finally: litellm.turn_off_message_logging = original_turn_off_message_logging @@ -1506,12 +1472,7 @@ async def test_add_litellm_data_to_request_admin_callback_vars_turn_off_message_ dynamic_params = initialize_standard_callback_dynamic_params(updated) assert dynamic_params.get("turn_off_message_logging") == "True" - assert ( - should_redact_message_logging( - {"standard_callback_dynamic_params": dynamic_params} - ) - is True - ) + assert should_redact_message_logging({"standard_callback_dynamic_params": dynamic_params}) is True finally: litellm.turn_off_message_logging = original_turn_off_message_logging @@ -1552,9 +1513,7 @@ async def test_add_litellm_data_to_request_allows_redaction_opt_out_with_admin_o "headers": {"litellm-disable-message-redaction": "true"}, "turn_off_message_logging": False, }, - "litellm_metadata": json.dumps( - {"headers": {"LiteLLM-Disable-Message-Redaction": "true"}} - ), + "litellm_metadata": json.dumps({"headers": {"LiteLLM-Disable-Message-Redaction": "true"}}), }, request=request_mock, user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", **auth_kwargs), @@ -1567,19 +1526,15 @@ async def test_add_litellm_data_to_request_allows_redaction_opt_out_with_admin_o assert updated["turn_off_message_logging"] is False assert updated["metadata"]["turn_off_message_logging"] is False + assert "litellm-disable-message-redaction" in {header.lower() for header in updated["metadata"]["headers"]} assert "litellm-disable-message-redaction" in { - header.lower() for header in updated["metadata"]["headers"] - } - assert "litellm-disable-message-redaction" in { - header.lower() - for header in updated["metadata"]["requester_metadata"].get("headers", {}) + header.lower() for header in updated["metadata"]["requester_metadata"].get("headers", {}) } assert "litellm-disable-message-redaction" in { header.lower() for header in updated["proxy_server_request"]["headers"] } assert "litellm-disable-message-redaction" in { - header.lower() - for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] + header.lower() for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] } assert "litellm_metadata" not in updated @@ -1870,9 +1825,7 @@ async def test_add_litellm_data_to_request_audio_transcription_multipart(): request_mock.client.host = "127.0.0.1" # Simulate multipart data (metadata as string) - metadata_dict = { - "tags": ["jobID:214590dsff09fds", "taskName:run_page_classification"] - } + metadata_dict = {"tags": ["jobID:214590dsff09fds", "taskName:run_page_classification"]} stringified_metadata = json.dumps(metadata_dict) data = { @@ -2200,23 +2153,15 @@ def test_key_dynamic_logging_settings(): # Test with langfuse logging key_with_langfuse = UserAPIKeyAuth( api_key="test-key", - metadata={ - "logging": [{"callback_name": "langfuse", "callback_type": "success"}] - }, + metadata={"logging": [{"callback_name": "langfuse", "callback_type": "success"}]}, team_metadata={}, ) - result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings( - key_with_langfuse - ) + result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_langfuse) assert result == [{"callback_name": "langfuse", "callback_type": "success"}] # Test with no logging metadata - key_without_logging = UserAPIKeyAuth( - api_key="test-key", metadata={}, team_metadata={} - ) - result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings( - key_without_logging - ) + key_without_logging = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={}) + result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_without_logging) assert result is None @@ -2228,35 +2173,23 @@ def test_team_dynamic_logging_settings(): key_with_team_arize = UserAPIKeyAuth( api_key="test-key", metadata={}, - team_metadata={ - "logging": [{"callback_name": "arize", "callback_type": "failure"}] - }, - ) - result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( - key_with_team_arize + team_metadata={"logging": [{"callback_name": "arize", "callback_type": "failure"}]}, ) + result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_with_team_arize) assert result == [{"callback_name": "arize", "callback_type": "failure"}] # Test with langfuse team logging key_with_team_langfuse = UserAPIKeyAuth( api_key="test-key", metadata={}, - team_metadata={ - "logging": [{"callback_name": "langfuse", "callback_type": "success"}] - }, - ) - result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( - key_with_team_langfuse + team_metadata={"logging": [{"callback_name": "langfuse", "callback_type": "success"}]}, ) + result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_with_team_langfuse) assert result == [{"callback_name": "langfuse", "callback_type": "success"}] # Test with no team logging metadata - key_without_team_logging = UserAPIKeyAuth( - api_key="test-key", metadata={}, team_metadata={} - ) - result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( - key_without_team_logging - ) + key_without_team_logging = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={}) + result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_without_team_logging) assert result is None @@ -2337,9 +2270,7 @@ def test_get_dynamic_logging_metadata_with_arize_team_logging(): mock_proxy_config = MagicMock() # Call the function - result = _get_dynamic_logging_metadata( - user_api_key_dict=user_api_key_dict, proxy_config=mock_proxy_config - ) + result = _get_dynamic_logging_metadata(user_api_key_dict=user_api_key_dict, proxy_config=mock_proxy_config) # Verify the result assert result is not None @@ -2355,9 +2286,7 @@ def test_add_team_callback_rejects_env_reference(): AddTeamCallback( callback_name="langfuse", callback_type="success", - callback_vars={ - "langfuse_secret_key": "os.environ/LANGFUSE_SECRET_KEY_TEMP" - }, + callback_vars={"langfuse_secret_key": "os.environ/LANGFUSE_SECRET_KEY_TEMP"}, ) assert "os.environ/" in str(exc_info.value) @@ -2388,9 +2317,7 @@ def test_get_dynamic_logging_metadata_ignores_env_reference_from_key_metadata( team_metadata={}, ) - result = _get_dynamic_logging_metadata( - user_api_key_dict=user_api_key_dict, proxy_config=MagicMock() - ) + result = _get_dynamic_logging_metadata(user_api_key_dict=user_api_key_dict, proxy_config=MagicMock()) assert result is None @@ -2401,16 +2328,12 @@ def test_get_num_retries_from_request(): """ # Test case 1: Header is present with valid integer string headers_with_retries = {"x-litellm-num-retries": "3"} - result = LiteLLMProxyRequestSetup._get_num_retries_from_request( - headers_with_retries - ) + result = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_retries) assert result == 3 # Test case 2: Header is not present headers_without_retries = {"Content-Type": "application/json"} - result = LiteLLMProxyRequestSetup._get_num_retries_from_request( - headers_without_retries - ) + result = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_without_retries) assert result is None # Test case 3: Empty headers dictionary @@ -2425,9 +2348,7 @@ def test_get_num_retries_from_request(): # Test case 5: Header present with large number headers_with_large_number = {"x-litellm-num-retries": "100"} - result = LiteLLMProxyRequestSetup._get_num_retries_from_request( - headers_with_large_number - ) + result = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_large_number) assert result == 100 # Test case 6: Multiple headers with num retries header @@ -2441,19 +2362,17 @@ def test_get_num_retries_from_request(): # Test case 7: Header present with invalid value (should raise ValueError when int() is called) headers_with_invalid = {"x-litellm-num-retries": "invalid"} - with pytest.raises(ValueError, match='invalid literal for int\\(\\) with base'): + with pytest.raises(ValueError, match="invalid literal for int\\(\\) with base"): LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_invalid) # Test case 8: Header present with float string (should raise ValueError when int() is called) headers_with_float = {"x-litellm-num-retries": "3.5"} - with pytest.raises(ValueError, match='invalid literal for int\\(\\) with base'): + with pytest.raises(ValueError, match="invalid literal for int\\(\\) with base"): LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_float) # Test case 9: Header present with negative number headers_with_negative = {"x-litellm-num-retries": "-1"} - result = LiteLLMProxyRequestSetup._get_num_retries_from_request( - headers_with_negative - ) + result = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_negative) assert result == -1 @@ -2463,15 +2382,11 @@ def test_get_keepalive_seconds_from_request(): """ # Header present with valid float string headers_with_keepalive = {"x-litellm-keepalive-seconds": "15"} - result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( - headers_with_keepalive - ) + result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request(headers_with_keepalive) assert result == 15.0 # Header not present - result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( - {"Content-Type": "application/json"} - ) + result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request({"Content-Type": "application/json"}) assert result is None # Empty headers dictionary @@ -2479,17 +2394,13 @@ def test_get_keepalive_seconds_from_request(): assert result is None # Header present with a fractional value - result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( - {"x-litellm-keepalive-seconds": "1.5"} - ) + result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request({"x-litellm-keepalive-seconds": "1.5"}) assert result == 1.5 # Header present with invalid value raises ValueError, matching the other # x-litellm-* numeric header helpers (_get_timeout_from_request, etc.) with pytest.raises(ValueError, match="could not convert string to float: 'not-a-number"): - LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( - {"x-litellm-keepalive-seconds": "not-a-number"} - ) + LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request({"x-litellm-keepalive-seconds": "not-a-number"}) def test_add_litellm_data_for_backend_llm_call_merges_keepalive_seconds_header(): @@ -2728,9 +2639,7 @@ def test_management_endpoint_metadata_drops_callback_credentials(): ), ], ) -def test_add_headers_to_llm_call_by_model_group( - data, model_group_settings, expected_headers_added -): +def test_add_headers_to_llm_call_by_model_group(data, model_group_settings, expected_headers_added): """ Test LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_group method @@ -2751,9 +2660,7 @@ def test_add_headers_to_llm_call_by_model_group( "X-Custom-Header": "custom-value", } - user_api_key_dict = UserAPIKeyAuth( - api_key="test-key", user_id="test-user", org_id="test-org" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key", user_id="test-user", org_id="test-org") # Mock the model_group_settings original_model_group_settings = getattr(litellm, "model_group_settings", None) @@ -2771,7 +2678,6 @@ def test_add_headers_to_llm_call_by_model_group( "add_headers_to_llm_call", return_value=expected_returned_headers if expected_headers_added else {}, ) as mock_add_headers: - # Make a copy of original data to verify it's not mutated unexpectedly original_data = copy.deepcopy(data) @@ -2828,7 +2734,6 @@ def test_add_headers_to_llm_call_by_model_group_empty_headers_returned(): "add_headers_to_llm_call", return_value={}, # Return empty dict ) as mock_add_headers: - result = LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_group( data=data, headers=headers, user_api_key_dict=user_api_key_dict ) @@ -2876,7 +2781,6 @@ def test_add_headers_to_llm_call_by_model_group_existing_headers_in_data(): "add_headers_to_llm_call", return_value=new_headers, ) as mock_add_headers: - result = LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_group( data=data, headers=headers, user_api_key_dict=user_api_key_dict ) @@ -2990,13 +2894,9 @@ async def test_add_litellm_metadata_from_request_headers(): general_settings = {} # Create mock select_data_generator with correct signature - def mock_select_data_generator( - response=None, user_api_key_dict=None, request_data=None - ): + def mock_select_data_generator(response=None, user_api_key_dict=None, request_data=None): async def mock_generator(): - yield "data: " + json.dumps( - {"choices": [{"delta": {"content": "Hello"}}]} - ) + "\n\n" + yield "data: " + json.dumps({"choices": [{"delta": {"content": "Hello"}}]}) + "\n\n" yield "data: [DONE]\n\n" return mock_generator() @@ -3023,21 +2923,19 @@ async def test_add_litellm_metadata_from_request_headers(): await asyncio.sleep(3) # Check if standard_logging_object was set - assert ( - test_logger.standard_logging_object is not None - ), "standard_logging_object should be populated after LLM request" + assert test_logger.standard_logging_object is not None, ( + "standard_logging_object should be populated after LLM request" + ) # Verify the logging object contains expected metadata standard_logging_obj = test_logger.standard_logging_object - print( - f"Standard logging object captured: {json.dumps(standard_logging_obj, indent=4, default=str)}" - ) + print(f"Standard logging object captured: {json.dumps(standard_logging_obj, indent=4, default=str)}") SPEND_LOGS_METADATA = standard_logging_obj["metadata"]["spend_logs_metadata"] - assert SPEND_LOGS_METADATA == dict( - json.loads(headers["x-litellm-spend-logs-metadata"]) - ), "spend_logs_metadata should be the same as the headers" + assert SPEND_LOGS_METADATA == dict(json.loads(headers["x-litellm-spend-logs-metadata"])), ( + "spend_logs_metadata should be the same as the headers" + ) finally: litellm.callbacks = original_callbacks @@ -3188,11 +3086,7 @@ def test_add_litellm_metadata_from_request_headers_generic_session_id_header(): def test_add_litellm_metadata_from_anthropic_user_id_sets_session_id(): - data = { - "metadata": { - "user_id": "user_abc123_account__session_e96634a3-fa28-4083-b354-55542e2dca01" - } - } + data = {"metadata": {"user_id": "user_abc123_account__session_e96634a3-fa28-4083-b354-55542e2dca01"}} LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( headers={}, data=data, _metadata_variable_name="metadata" ) @@ -3308,9 +3202,7 @@ def test_get_chain_id_from_headers_generic_vendor_session_id(): from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers assert ( - get_chain_id_from_headers( - {"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01"} - ) + get_chain_id_from_headers({"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01"}) == "e96634a3-fa28-4083-b354-55542e2dca01" ) # Short / non-alphanumeric values should be ignored @@ -3600,19 +3492,13 @@ def test_get_internal_user_header_from_mapping_returns_expected_header(): {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}, ] - header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping( - mappings - ) + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) assert header_name == "X-OpenWebUI-User-Id" def test_get_internal_user_header_from_mapping_none_when_absent(): - mappings = [ - {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"} - ] - header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping( - mappings - ) + mappings = [{"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}] + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) assert header_name is None single = {"header_name": "X-Only-Customer", "litellm_user_role": "customer"} @@ -3633,9 +3519,7 @@ def test_add_internal_user_from_user_mapping_sets_user_id_when_header_present(): ] } - result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( - general_settings, user_api_key_dict, headers - ) + result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(general_settings, user_api_key_dict, headers) assert result is user_api_key_dict assert user_api_key_dict.user_id == "internal-user-123" @@ -3651,9 +3535,7 @@ def test_add_internal_user_from_user_mapping_no_header_or_mapping_returns_unchan assert user_api_key_dict.user_id is None general_settings = { - "user_header_mappings": [ - {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"} - ] + "user_header_mappings": [{"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}] } result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( general_settings, user_api_key_dict, {"Other": "value"} @@ -3673,9 +3555,7 @@ def test_get_sanitized_user_information_from_key_includes_guardrails_metadata(): metadata={"guardrails": ["presidio", "aporia"], "other_field": "value"}, ) - result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) assert result["user_api_key_auth_metadata"] is not None assert "guardrails" in result["user_api_key_auth_metadata"] @@ -3704,9 +3584,7 @@ def test_user_and_team_spend_and_budget_flow_to_standard_logging_metadata(): team_max_budget=1000.0, ) - sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) assert sanitized["user_api_key_spend"] == 1.5 assert sanitized["user_api_key_max_budget"] == 10.0 @@ -3715,9 +3593,7 @@ def test_user_and_team_spend_and_budget_flow_to_standard_logging_metadata(): assert sanitized["user_api_key_team_spend"] == 250.75 assert sanitized["user_api_key_team_max_budget"] == 1000.0 - logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( - dict(sanitized) - ) + logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata(dict(sanitized)) assert logging_metadata["user_api_key_user_spend"] == 25.5 assert logging_metadata["user_api_key_user_max_budget"] == 100.0 @@ -3734,12 +3610,8 @@ def test_user_and_team_spend_and_budget_default_to_none_in_standard_logging_meta user_api_key_dict = UserAPIKeyAuth(api_key="test-key-hash") - sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) - logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( - dict(sanitized) - ) + sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) + logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata(dict(sanitized)) assert logging_metadata["user_api_key_user_spend"] is None assert logging_metadata["user_api_key_user_max_budget"] is None @@ -4071,22 +3943,16 @@ async def test_embedding_header_forwarding_with_model_group(): # Verify that only x- prefixed headers (except x-stainless) were forwarded forwarded_headers = updated_data["headers"] - assert ( - "X-Custom-Header" in forwarded_headers - ), "X-Custom-Header should be forwarded" + assert "X-Custom-Header" in forwarded_headers, "X-Custom-Header should be forwarded" assert forwarded_headers["X-Custom-Header"] == "custom-value" assert "X-Request-ID" in forwarded_headers, "X-Request-ID should be forwarded" assert forwarded_headers["X-Request-ID"] == "test-request-123" # Verify that authorization header was NOT forwarded (sensitive header) - assert ( - "Authorization" not in forwarded_headers - ), "Authorization header should not be forwarded" + assert "Authorization" not in forwarded_headers, "Authorization header should not be forwarded" # Verify that Content-Type was NOT forwarded (doesn't start with x-) - assert ( - "Content-Type" not in forwarded_headers - ), "Content-Type should not be forwarded" + assert "Content-Type" not in forwarded_headers, "Content-Type should not be forwarded" # Verify original data fields are preserved assert updated_data["model"] == "local-openai/text-embedding-3-small" @@ -4142,9 +4008,9 @@ async def test_embedding_header_forwarding_without_model_group_config(): ) # Verify that headers were NOT added since model is not in forward list - assert ( - "headers" not in updated_data or updated_data.get("headers") is None - ), "Headers should not be forwarded for models not in forward_client_headers_to_llm_api list" + assert "headers" not in updated_data or updated_data.get("headers") is None, ( + "Headers should not be forwarded for models not in forward_client_headers_to_llm_api list" + ) # Verify original data fields are preserved assert updated_data["model"] == "text-embedding-ada-002" @@ -4198,9 +4064,7 @@ async def test_add_guardrails_from_policy_engine(): attachment_registry = get_attachment_registry() attachment_registry._attachments = [ PolicyAttachment(policy="global-baseline", scope="*"), # applies to all - PolicyAttachment( - policy="healthcare", teams=["healthcare-team"] - ), # applies to healthcare team + PolicyAttachment(policy="healthcare", teams=["healthcare-team"]), # applies to healthcare team ] attachment_registry._initialized = True @@ -4269,9 +4133,9 @@ async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_po ) # Verify that 'policies' was removed from the request body - assert ( - "policies" not in data - ), "'policies' should be removed from request body to prevent forwarding to LLM provider" + assert "policies" not in data, ( + "'policies' should be removed from request body to prevent forwarding to LLM provider" + ) # Verify that other fields are preserved assert "model" in data @@ -4316,9 +4180,7 @@ async def test_api_created_global_policy_applies_to_new_key_without_restart(): "runtime-global-policy", Policy(guardrails=PolicyGuardrails(add=["runtime-guardrail"])), ) - attachment_registry.add_attachment( - PolicyAttachment(policy="runtime-global-policy", scope="*") - ) + attachment_registry.add_attachment(PolicyAttachment(policy="runtime-global-policy", scope="*")) await add_guardrails_from_policy_engine( data=data, @@ -4415,9 +4277,7 @@ async def test_bearer_token_not_in_debug_logs(): from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import ProxyConfig - secret_token = ( - "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.fakesignature" - ) + secret_token = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.fakesignature" mock_request = MagicMock(spec=Request) mock_request.headers = { @@ -4463,8 +4323,7 @@ async def test_bearer_token_not_in_debug_logs(): log_output = log_capture.getvalue() assert secret_token not in log_output, ( - f"Bearer token leaked in debug logs. " - f"Found token in log output:\n{log_output[:500]}" + f"Bearer token leaked in debug logs. Found token in log output:\n{log_output[:500]}" ) @@ -4629,9 +4488,7 @@ def test_apply_overrides_project_model_specific(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, } }, @@ -4642,9 +4499,7 @@ def test_apply_overrides_project_model_specific(setup_test_credentials): } }, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-rec-vision.openai.azure.com/" assert data["api_key"] == "key-hotel-rec-vision" assert data["api_version"] == "2024-06-01" @@ -4657,9 +4512,7 @@ def test_apply_overrides_project_default(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, } }, @@ -4670,9 +4523,7 @@ def test_apply_overrides_project_default(setup_test_credentials): } }, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-rec-app.openai.azure.com/" assert data["api_key"] == "key-hotel-rec" @@ -4684,17 +4535,13 @@ def test_apply_overrides_team_model_specific(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, } }, project_metadata={}, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-westus.openai.azure.com/" assert data["api_key"] == "key-hotel-westus" @@ -4706,17 +4553,13 @@ def test_apply_overrides_team_default(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, } }, project_metadata={}, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" @@ -4729,9 +4572,7 @@ def test_apply_overrides_no_config(setup_test_credentials): team_metadata={}, project_metadata={}, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data assert "api_key" not in data @@ -4747,17 +4588,9 @@ def test_apply_overrides_clientside_credentials_take_precedence( } user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - } - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://my-custom-endpoint.openai.azure.com/" assert data["api_key"] == "my-custom-key" @@ -4767,15 +4600,9 @@ def test_apply_overrides_missing_credential_name(setup_test_credentials): data = {"model": "gpt-4"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "gpt-4": {"azure": {"litellm_credentials": "nonexistent-credential"}} - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"gpt-4": {"azure": {"litellm_credentials": "nonexistent-credential"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data assert "api_key" not in data @@ -4785,17 +4612,9 @@ def test_apply_overrides_api_version_only_if_present(setup_test_credentials): data = {"model": "gpt-3.5"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - } - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" assert "api_version" not in data @@ -4806,15 +4625,9 @@ def test_apply_overrides_no_model_in_data(setup_test_credentials): data = {"messages": [{"role": "user", "content": "hello"}]} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "defaultconfig": {"azure": {"litellm_credentials": "some-cred"}} - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"defaultconfig": {"azure": {"litellm_credentials": "some-cred"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data @@ -4826,9 +4639,7 @@ def test_apply_overrides_none_metadata(setup_test_credentials): team_metadata=None, project_metadata=None, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data @@ -4837,15 +4648,9 @@ def test_apply_overrides_clientside_api_version_preserved(setup_test_credentials data = {"model": "gpt-4-vision", "api_version": "2025-01-01"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "gpt-4-vision": {"azure": {"litellm_credentials": "hotel-rec-vision"}} - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"gpt-4-vision": {"azure": {"litellm_credentials": "hotel-rec-vision"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) # api_base and api_key should be set from credential assert data["api_base"] == "https://hotel-rec-vision.openai.azure.com/" assert data["api_key"] == "key-hotel-rec-vision" @@ -4858,9 +4663,7 @@ def test_resolve_non_dict_model_config_ignored(): result = _resolve_credential_from_model_config("gpt-4", "not-a-dict", None) assert result is None - result = _resolve_credential_from_model_config( - "gpt-4", None, ["also", "not", "a", "dict"] - ) + result = _resolve_credential_from_model_config("gpt-4", None, ["also", "not", "a", "dict"]) assert result is None # Valid config still works alongside invalid one @@ -4878,9 +4681,7 @@ def test_resolve_pre_alias_model_name_fallback(): "gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}}, } # Post-alias name doesn't match, but pre-alias does (team scope) - result = _resolve_credential_from_model_config( - "azure/gpt-4-0613", None, team_config, pre_alias_model_name="gpt-4" - ) + result = _resolve_credential_from_model_config("azure/gpt-4-0613", None, team_config, pre_alias_model_name="gpt-4") assert result == "team-gpt4" # Same test for project scope @@ -4900,15 +4701,11 @@ def test_resolve_post_alias_name_takes_priority(): "gpt-4o-team-1": {"azure": {"litellm_credentials": "post-alias-cred"}}, } # Team scope - result = _resolve_credential_from_model_config( - "gpt-4o-team-1", None, team_config, pre_alias_model_name="gpt-4" - ) + result = _resolve_credential_from_model_config("gpt-4o-team-1", None, team_config, pre_alias_model_name="gpt-4") assert result == "post-alias-cred" # Project scope - result = _resolve_credential_from_model_config( - "gpt-4o-team-1", team_config, None, pre_alias_model_name="gpt-4" - ) + result = _resolve_credential_from_model_config("gpt-4o-team-1", team_config, None, pre_alias_model_name="gpt-4") assert result == "post-alias-cred" @@ -4940,15 +4737,9 @@ def test_apply_overrides_feature_flag_disabled_by_default(): data = {"model": "gpt-4"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-eastus"}} - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"gpt-4": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data assert "api_key" not in data @@ -5108,9 +4899,7 @@ async def test_team_guardrail_merges_with_global_policy(): policy_registry = get_policy_registry() policy_registry._policies = { "global-policy": Policy( - guardrails=PolicyGuardrails( - add=["policy-guardrail-1", "policy-guardrail-2"] - ), + guardrails=PolicyGuardrails(add=["policy-guardrail-1", "policy-guardrail-2"]), ), } policy_registry._initialized = True @@ -5131,18 +4920,10 @@ async def test_team_guardrail_merges_with_global_policy(): guardrails = data["metadata"].get("guardrails", []) - assert ( - "team-direct-guardrail" in guardrails - ), f"Team guardrail missing from merged list: {guardrails}" - assert ( - "policy-guardrail-1" in guardrails - ), f"policy-guardrail-1 missing: {guardrails}" - assert ( - "policy-guardrail-2" in guardrails - ), f"policy-guardrail-2 missing: {guardrails}" - assert len(guardrails) == len( - set(guardrails) - ), f"Duplicates in guardrails list: {guardrails}" + assert "team-direct-guardrail" in guardrails, f"Team guardrail missing from merged list: {guardrails}" + assert "policy-guardrail-1" in guardrails, f"policy-guardrail-1 missing: {guardrails}" + assert "policy-guardrail-2" in guardrails, f"policy-guardrail-2 missing: {guardrails}" + assert len(guardrails) == len(set(guardrails)), f"Duplicates in guardrails list: {guardrails}" # Verify get_guardrail_from_metadata returns the merged list even # when litellm_metadata is present (the bug: it returned [] before fix) @@ -5153,9 +4934,9 @@ async def test_team_guardrail_merges_with_global_policy(): dummy = _DummyGuardrail(guardrail_name="team-direct-guardrail") returned = dummy.get_guardrail_from_metadata(data) - assert ( - "team-direct-guardrail" in returned - ), f"get_guardrail_from_metadata shadowed by litellm_metadata; got: {returned}" + assert "team-direct-guardrail" in returned, ( + f"get_guardrail_from_metadata shadowed by litellm_metadata; got: {returned}" + ) finally: policy_registry._policies = {} @@ -5208,9 +4989,7 @@ def test_get_guardrail_from_metadata_reads_litellm_metadata_when_no_metadata(): } result = dummy.get_guardrail_from_metadata(data) - assert result == [ - "my-guardrail" - ], f"Expected guardrails from litellm_metadata fallback, got: {result}" + assert result == ["my-guardrail"], f"Expected guardrails from litellm_metadata fallback, got: {result}" def _build_request_mock_with_headers(headers: dict) -> Request: @@ -5237,9 +5016,7 @@ class TestApplyClientTagPolicyPreAuth: """ def test_merges_header_tags_into_metadata(self): - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme,env:prod"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme,env:prod"}) data = {"model": "gpt-3.5-turbo"} user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", @@ -5256,9 +5033,7 @@ class TestApplyClientTagPolicyPreAuth: assert data["metadata"]["tags"] == ["tenant:acme", "env:prod"] def test_unions_header_tags_with_existing_metadata_tags(self): - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme,env:prod"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme,env:prod"}) data = { "model": "gpt-3.5-turbo", "metadata": {"tags": ["env:prod", "team:platform"]}, @@ -5283,9 +5058,7 @@ class TestApplyClientTagPolicyPreAuth: # (inside common_checks) enforces per-tag budgets on whatever tags # it sees in request_data, including body tags. The helper only # adds header tags to metadata.tags. - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme"}) data = { "model": "gpt-3.5-turbo", "tags": ["root-tag"], @@ -5312,9 +5085,7 @@ class TestApplyClientTagPolicyPreAuth: ] def test_uses_litellm_metadata_when_present(self): - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme"}) data = { "model": "gpt-3.5-turbo", "litellm_metadata": {"foo": "bar"}, @@ -5409,9 +5180,7 @@ class TestApplyClientTagPolicyPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:paid": return 0.50 return fallback_spend @@ -5446,9 +5215,7 @@ class TestApplyClientTagPolicyPreAuth: from litellm.proxy.auth.auth_checks import _tag_max_budget_check from litellm.proxy.utils import ProxyLogging - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme"}) data = {"model": "gpt-3.5-turbo"} user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", @@ -5468,9 +5235,7 @@ class TestApplyClientTagPolicyPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:tenant:acme": return 0.50 return fallback_spend @@ -5505,9 +5270,7 @@ class TestApplyClientTagPolicyPreAuth: "/v1/messages", ], ) - async def test_header_tags_visible_to_tag_max_budget_check_on_metadata_route( - self, route - ): + async def test_header_tags_visible_to_tag_max_budget_check_on_metadata_route(self, route): """Regression: on LITELLM_METADATA_ROUTES (bedrock, /v1/messages, ...), common_checks pre-seeds ``litellm_metadata`` and writes key tags there before ``_tag_max_budget_check`` reads from the same key. The auth wrapper @@ -5521,9 +5284,7 @@ class TestApplyClientTagPolicyPreAuth: from litellm.proxy.auth.auth_checks import common_checks from litellm.proxy.utils import ProxyLogging - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme"}) data = {"model": "us.anthropic.claude-sonnet-4-6"} valid_token = UserAPIKeyAuth( token="test-token", @@ -5548,9 +5309,7 @@ class TestApplyClientTagPolicyPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:tenant:acme": return 0.50 return fallback_spend @@ -5718,9 +5477,7 @@ class TestApplyKeyTagsPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:engineering": return 0.50 return fallback_spend @@ -5771,9 +5528,7 @@ class TestApplyKeyTagsPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:engineering": return 0.05 return fallback_spend @@ -5854,9 +5609,7 @@ def test_resolve_provider_from_deployment_falls_back_to_pre_alias(): router.get_deployment_by_model_group_name.side_effect = lookup - result = _resolve_provider_from_deployment( - router, "post-alias-name", pre_alias_model_name="pre-alias-name" - ) + result = _resolve_provider_from_deployment(router, "post-alias-name", pre_alias_model_name="pre-alias-name") assert result == "bedrock" @@ -5921,17 +5674,9 @@ def test_apply_overrides_no_router_keeps_legacy_behaviour(setup_test_credentials data = {"model": "gpt-4"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - } - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict, llm_router=None + team_metadata={"model_config": {"defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict, llm_router=None) assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" @@ -5957,9 +5702,7 @@ def test_apply_overrides_provider_prefix_in_model_skips_router_lookup( ) router = MagicMock() - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict, llm_router=router - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict, llm_router=router) assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" router.get_deployment_by_model_group_name.assert_not_called() @@ -6338,9 +6081,7 @@ def test_get_sanitized_user_information_from_key_drops_callback_config(): }, ) - result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) auth_metadata = result["user_api_key_auth_metadata"] assert "logging" not in auth_metadata @@ -6380,9 +6121,7 @@ def test_team_alias_targeting_deleted_team_deployment_keeps_requested_model(monk ) with patch("litellm.proxy.proxy_server.llm_router", _MockRouter()): - _update_model_if_team_alias_exists( - data=test_data, user_api_key_dict=user_api_key_dict - ) + _update_model_if_team_alias_exists(data=test_data, user_api_key_dict=user_api_key_dict) assert test_data.get("model") == "gpt-4" @@ -6406,9 +6145,7 @@ def test_team_alias_targeting_live_team_deployment_still_rewrites(monkeypatch): ) with patch("litellm.proxy.proxy_server.llm_router", _MockRouter()): - _update_model_if_team_alias_exists( - data=test_data, user_api_key_dict=user_api_key_dict - ) + _update_model_if_team_alias_exists(data=test_data, user_api_key_dict=user_api_key_dict) assert test_data.get("model") == "model_name_team-1_live-uuid" @@ -6545,7 +6282,6 @@ async def test_add_litellm_data_to_request_keeps_every_forwarded_credential_out_ assert value not in logged - @pytest.mark.parametrize( "header, expected_redacted", [ @@ -6571,7 +6307,6 @@ def test_redact_credential_headers_classifies_each_header(header, expected_redac assert headers[header] == "secret-value" - @pytest.mark.asyncio async def test_add_litellm_data_to_request_debug_log_does_not_print_credentials(): """The request-header debug line carries values the stdout secret filter does not match.""" @@ -7191,8 +6926,7 @@ AUTHORIZATION_HEADER_CASINGS = ["authorization", "Authorization", "AUTHORIZATION LEAK_TARGET_PROVIDERS = ["bedrock", "bedrock_converse", "vertex_ai"] BEDROCK_ENDPOINT = ( - "https://bedrock-runtime.us-west-2.amazonaws.com" - "/model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke" + "https://bedrock-runtime.us-west-2.amazonaws.com/model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke" ) BEDROCK_REGION = "us-west-2" BEDROCK_REQUEST_DATA = {"messages": [{"role": "user", "content": "Say OK"}], "max_tokens": 32} @@ -7250,9 +6984,7 @@ def _signed_headers_component(signature: str, component: str) -> str: @pytest.mark.parametrize("authorization_header_name", AUTHORIZATION_HEADER_CASINGS) @pytest.mark.parametrize("custom_llm_provider", LEAK_TARGET_PROVIDERS) -def test_oauth_credential_is_never_forwarded_to_bedrock_or_vertex( - authorization_header_name, custom_llm_provider -): +def test_oauth_credential_is_never_forwarded_to_bedrock_or_vertex(authorization_header_name, custom_llm_provider): """ A client's Anthropic OAuth credential is meaningless to AWS and Google, and sending it there both breaks the request and hands a third-party cloud a credential it should @@ -7280,9 +7012,7 @@ def test_oauth_credential_entry_is_scoped_to_anthropic_alone(): if not isinstance(scoped_headers, list): scoped_headers = [scoped_headers] - credential_entries = [ - entry for entry in scoped_headers if OAUTH_TOKEN in entry["extra_headers"].values() - ] + credential_entries = [entry for entry in scoped_headers if OAUTH_TOKEN in entry["extra_headers"].values()] assert [entry["custom_llm_provider"] for entry in credential_entries] == ["anthropic"] @@ -7344,9 +7074,7 @@ def test_bedrock_get_request_headers_keeps_the_sigv4_signature(): def test_bedrock_api_key_deployment_keeps_its_own_bearer_token(): forwarded = _headers_forwarded_to(_client_headers(), "bedrock") - signed = _signed_headers_for_bedrock( - {"Content-Type": "application/json", **forwarded}, api_key=BEDROCK_API_KEY - ) + signed = _signed_headers_for_bedrock({"Content-Type": "application/json", **forwarded}, api_key=BEDROCK_API_KEY) assert _authorization_values(signed) == [f"Bearer {BEDROCK_API_KEY}"] @@ -7369,6 +7097,8 @@ def test_vertex_sends_exactly_one_authorization_header(): vertex_request_headers.update(forwarded) assert _authorization_values(vertex_request_headers) == [GOOGLE_ACCESS_TOKEN] + + @pytest.mark.asyncio async def test_newrelic_team_callback_vars_reach_trusted_field(): """A key with a newrelic team callback stamps its vars into the proxy-owned @@ -7737,13 +7467,13 @@ def _request_for(path: str) -> MagicMock: return request -def _spend_log_session_id(data: dict[str, object]) -> str: - """Resolve session_id the way LiteLLM_SpendLogs does: standard_logging_payload.trace_id.""" +def _spend_log_session_id(data: dict[str, object], metadata_key: str = "metadata") -> str | None: + """Resolve session_id the way LiteLLM_SpendLogs does, reading the omit decision stamped on the request.""" from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.proxy.spend_tracking.spend_tracking_utils import _get_session_id_for_spend_log - metadata = data["metadata"] + metadata = data[metadata_key] assert isinstance(metadata, dict) litellm_params = get_litellm_params( litellm_session_id=str(data["litellm_session_id"]) if "litellm_session_id" in data else None, @@ -7754,7 +7484,12 @@ def _spend_log_session_id(data: dict[str, object]) -> str: logging_obj=SimpleNamespace(litellm_trace_id="per-call-random-trace-id"), litellm_params=litellm_params, ) - return _get_session_id_for_spend_log(kwargs={}, standard_logging_payload={"trace_id": trace_id}) + return _get_session_id_for_spend_log( + kwargs={}, + metadata=metadata, + standard_logging_payload={"trace_id": trace_id}, + omit_when_missing=bool(metadata.get(SESSION_ID_OMITTED_METADATA_KEY)), + ) @pytest.mark.asyncio @@ -7780,9 +7515,10 @@ async def test_missing_session_id_generate_makes_spend_log_and_callback_session_ assert isinstance(callback_session_id, str) and len(callback_session_id) == 36 assert _spend_log_session_id(updated) == callback_session_id assert updated["metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True - assert get_fireworks_session_id( - {"litellm_session_id": updated["litellm_session_id"], "metadata": updated["metadata"]} - ) is None + assert ( + get_fireworks_session_id({"litellm_session_id": updated["litellm_session_id"], "metadata": updated["metadata"]}) + is None + ) @pytest.mark.asyncio @@ -7800,6 +7536,44 @@ async def test_missing_session_id_unset_keeps_legacy_divergence(): assert _spend_log_session_id(updated) == "per-call-random-trace-id" +@pytest.mark.asyncio +async def test_missing_session_id_omit_leaves_spend_log_session_id_null(): + """Under `omit` a traceparent still becomes the trace id but never a session id, so SpendLogs and + Langfuse agree on having no session.""" + request = _request_for("/v1/chat/completions") + request.headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert "session_id" not in updated["metadata"] + assert updated["litellm_trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert updated["metadata"][SESSION_ID_OMITTED_METADATA_KEY] is True + assert _spend_log_session_id(updated) is None + + +@pytest.mark.asyncio +async def test_missing_session_id_omit_keeps_client_supplied_session_id(): + request = _request_for("/v1/chat/completions") + request.headers = {"x-litellm-session-id": "client-session-1"} + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["metadata"]["session_id"] == "client-session-1" + assert _spend_log_session_id(updated) == "client-session-1" + + @pytest.mark.asyncio async def test_missing_session_id_generate_reuses_traceparent_trace_id(): """A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it.""" @@ -7895,3 +7669,38 @@ async def test_missing_session_id_unknown_value_is_ignored(): ) assert "session_id" not in updated["metadata"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path, sent_in, metadata_key, general_settings", + [ + ("/v1/chat/completions", "metadata", "metadata", {}), + ("/v1/chat/completions", "litellm_metadata", "metadata", {}), + ("/v1/chat/completions", "litellm_metadata", "metadata", {"missing_session_id": "generate"}), + ("/v1/messages", "litellm_metadata", "litellm_metadata", {}), + ("/v1/messages", "metadata", "litellm_metadata", {}), + ("/mcp/tools", "metadata", "metadata", {"missing_session_id": "omit"}), + ("/mcp/tools", "litellm_metadata", "metadata", {"missing_session_id": "omit"}), + ], +) +async def test_client_supplied_omit_marker_never_reaches_the_spend_log( + path: str, sent_in: str, metadata_key: str, general_settings: dict[str, str] +): + """The omit marker is proxy-owned: only the pre-call policy may set it. A caller that sends it in either + metadata bucket, including the one later merged into the route's bucket, must not be able to null out + SpendLogs.session_id on a request the proxy did not omit.""" + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": [], sent_in: {SESSION_ID_OMITTED_METADATA_KEY: True}}, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings=general_settings, + ) + + assert SESSION_ID_OMITTED_METADATA_KEY not in updated[metadata_key] + assert _spend_log_session_id(updated, metadata_key) == ( + updated[metadata_key]["session_id"] + if general_settings.get("missing_session_id") == "generate" + else "per-call-random-trace-id" + ) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 098b43f6433..5246aaf6904 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25779,9 +25779,9 @@ export interface components { mcp_xff_num_trusted_hops?: number | null; /** * Missing Session Id - * @description What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id. + * @description What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400; 'omit' leaves SpendLogs.session_id null, matching callbacks such as Langfuse that only record a client-established metadata.session_id. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id. */ - missing_session_id?: ("generate" | "reject") | null; + missing_session_id?: ("generate" | "reject" | "omit") | null; /** * Model List Healthy Only * @description When true, `/models`, `/v1/models/{id}` and `/model/info` hide models whose backing deployments are all unhealthy, for every caller, without needing `healthy_only=true` per request. Requires `background_health_checks: true`, and keeps deployment health state cached without turning on `enable_health_check_routing`, so routing is unaffected. With no health state nothing is hidden. Hiding is presentation-only, a hidden model can still be called.