From 338a634762d7dac2801a169dae7dbe9c7098aa3b Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Mon, 2 Mar 2026 21:31:35 +0530 Subject: [PATCH] Fix root cause: DB spend log session_id didn't match Langfuse trace_id MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The proxy has two separate failure paths: 1. async_failure_handler → Langfuse callback (uses model_call_details with standard_logging_object containing the correct trace_id) 2. post_call_failure_hook → _ProxyDBLogger → spend log (uses request_data which did NOT have standard_logging_object, so session_id fell to random uuid4()) These two paths used different data dicts, so the DB session_id was a random UUID unrelated to the Langfuse trace_id. Users could not search by the Session ID from LiteLLM logs in Langfuse for failed requests. Fix: In _ProxyDBLogger.async_post_call_failure_hook, propagate standard_logging_object and litellm_trace_id from the litellm_logging_obj (already present in request_data) before writing the spend log. Co-Authored-By: Claude Opus 4.6 --- .../proxy/hooks/proxy_track_cost_callback.py | 16 +++++ .../hooks/test_proxy_track_cost_callback.py | 62 +++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0734756d8ed..9a806fa4f87 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -110,6 +110,22 @@ 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. + _litellm_logging_obj = request_data.get("litellm_logging_obj") + if _litellm_logging_obj is not None: + if "standard_logging_object" not in request_data: + request_data["standard_logging_object"] = getattr( + _litellm_logging_obj, "model_call_details", {} + ).get("standard_logging_object") + if request_data.get("litellm_trace_id") is None: + request_data["litellm_trace_id"] = getattr( + _litellm_logging_obj, "litellm_trace_id", None + ) + await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key_dict.api_key, response_cost=0.0, diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index c46b8df5efc..d269a9531fd 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -169,6 +169,68 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj(): + """ + When an LLM call fails, the proxy calls post_call_failure_hook with + request_data that doesn't contain standard_logging_object. But the + litellm_logging_obj (set by function_setup) is in request_data and + holds the standard_logging_object with the correct trace_id. + + The failure hook should propagate this so the DB spend log's session_id + matches the Langfuse trace_id. + """ + logger = _ProxyDBLogger() + + user_api_key_dict = UserAPIKeyAuth( + api_key="test_api_key", + user_id="test_user_id", + team_id="test_team_id", + ) + + # Simulate a litellm_logging_obj with model_call_details containing + # the standard_logging_object (as set by _failure_handler_helper_fn) + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_trace_id = "trace-id-from-logging-obj" + mock_logging_obj.model_call_details = { + "standard_logging_object": { + "trace_id": "trace-id-from-logging-obj", + "error_str": "InternalServerError", + } + } + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + "litellm_logging_obj": mock_logging_obj, + # Note: no "standard_logging_object" and no "litellm_trace_id" + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("Provider error"), + user_api_key_dict=user_api_key_dict, + ) + + mock_update_database.assert_called_once() + call_kwargs = mock_update_database.call_args[1]["kwargs"] + + # standard_logging_object should have been propagated from logging obj + assert call_kwargs.get("standard_logging_object") is not None + assert ( + call_kwargs["standard_logging_object"]["trace_id"] + == "trace-id-from-logging-obj" + ) + # litellm_trace_id should also be propagated as a fallback + assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj" + + @pytest.mark.asyncio async def test_enrich_failure_metadata_with_team_alias(): """