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..98464d3a127 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1216,7 +1216,7 @@ class GenerateKeyRequest(KeyRequestBase): organization_id: str | None = None project_id: str | None = None - @field_validator("team_id", "organization_id", mode="before") + @field_validator("team_id", "organization_id", "project_id", mode="before") @classmethod def treat_cleared_id_as_unset(cls, v: object) -> object: if v == "": @@ -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/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 68704f476b5..7954a4693cc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -17505,6 +17505,18 @@ def test_generate_key_request_blank_team_id_is_personal(): assert GenerateKeyRequest(team_id="team-1").team_id == "team-1" +def test_generate_key_request_blank_organization_and_project_id_are_unset(): + from litellm.proxy._types import RegenerateKeyRequest + + cleared = GenerateKeyRequest(organization_id="", project_id="") + assert cleared.organization_id is None + assert cleared.project_id is None + assert "organization_id" not in cleared.model_dump(exclude_none=True) + assert RegenerateKeyRequest(organization_id="").organization_id is None + assert GenerateKeyRequest(organization_id="org-1", project_id="proj-1").organization_id == "org-1" + assert GenerateKeyRequest(organization_id="org-1", project_id="proj-1").project_id == "proj-1" + + def test_key_request_blank_organization_id_is_unset(): from litellm.proxy._types import RegenerateKeyRequest, UpdateKeyRequest 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/components/common_components/OrganizationDropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx index e0b2b897b36..524ec6a715e 100644 --- a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx @@ -58,6 +58,17 @@ describe("OrganizationDropdown", () => { expect(onChange.mock.calls[0][0]).toBe("org-1"); }); + it("emits null, never the empty string, when the selection is cleared", async () => { + const onChange = vi.fn(); + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("button", { name: "Clear" })); + + expect(onChange).toHaveBeenCalledTimes(1); + expect(onChange).toHaveBeenCalledWith(null); + }); + it("should filter options by organization id", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx index 8da35ecd02e..663028b2d92 100644 --- a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx @@ -5,7 +5,7 @@ import { Organization } from "../networking"; interface OrganizationDropdownProps { organizations?: Organization[] | null; value?: string; - onChange?: (value: string) => void; + onChange?: (value: string | null) => void; disabled?: boolean; loading?: boolean; style?: React.CSSProperties; @@ -32,7 +32,7 @@ const OrganizationDropdown: React.FC = ({ sublabel: org.organization_id, }))} value={value} - onValueChange={(organizationId) => onChange?.(organizationId)} + onValueChange={(organizationId) => onChange?.(organizationId || null)} placeholder={placeholder} emptyText={loading ? "Loading organizations…" : "No organizations found"} disabled={disabled} diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 5749541dcea..ee3a88acba0 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -587,9 +587,9 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp } }; - const changeOrganization = (write: FieldWrite) => (orgId: string) => { - write(orgId || undefined); - setSelectedOrganizationId(orgId || null); + const changeOrganization = (write: FieldWrite) => (orgId: string | null) => { + write(orgId ?? undefined); + setSelectedOrganizationId(orgId); // Clear team and project when org changes setSelectedCreateKeyTeam(null); setSelectedProjectId(null); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 3e772fd0e9b..1330be2788b 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -303,9 +303,9 @@ export function KeyEditView({ } }; - const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | undefined) => { - setField(orgId || null); - setSelectedOrganizationId(orgId || null); + const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => { + setField(orgId); + setSelectedOrganizationId(orgId); form.setValue("team_id", undefined); }; 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.