From cac041c9447a798fe611ef4f433142ed90c63cbc Mon Sep 17 00:00:00 2001
From: Harshit28j <Harshitjain0562@gmail.com>
Date: Mon, 2 Mar 2026 20:55:31 +0530
Subject: [PATCH] Align Langfuse trace_id fallback with DB session_id for
 failed requests

When standard_logging_object is None (failure case), Langfuse was falling
back to litellm_call_id while the DB used litellm_trace_id as session_id.
This caused the Session ID in LiteLLM logs to not match the trace in
Langfuse. Now Langfuse checks litellm_trace_id first, matching the DB.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 litellm/integrations/langfuse/langfuse.py     |  5 +-
 .../integrations/test_langfuse.py             | 73 +++++++++++++++++++
 2 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py
index 7bf97665fd2..afad3f50940 100644
--- a/litellm/integrations/langfuse/langfuse.py
+++ b/litellm/integrations/langfuse/langfuse.py
@@ -610,9 +610,10 @@ class LangFuseLogger:
                 trace_id = cast(
                     Optional[str], standard_logging_object.get("trace_id")
                 )
-            # Fallback to litellm_call_id if no trace_id found
+            # Fallback: use litellm_trace_id from kwargs (matches DB session_id),
+            # then litellm_call_id as last resort
             if trace_id is None:
-                trace_id = litellm_call_id
+                trace_id = kwargs.get("litellm_trace_id") or litellm_call_id
             existing_trace_id = clean_metadata.pop("existing_trace_id", None)
             # If existing_trace_id is provided, use it as the trace_id to return
             # This allows continuing an existing trace while still returning the correct trace_id
diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py
index 084fd7d0480..b4028709218 100644
--- a/tests/test_litellm/integrations/test_langfuse.py
+++ b/tests/test_litellm/integrations/test_langfuse.py
@@ -467,6 +467,79 @@ class TestLangfuseUsageDetails(unittest.TestCase):
 
         assert self.last_trace_kwargs.get("id") == "call-id-xyz"
 
+    def test_log_langfuse_v2_uses_litellm_trace_id_fallback_over_call_id(self):
+        """
+        When standard_logging_object has no trace_id, but kwargs contains
+        litellm_trace_id (the same ID the DB stores as Session ID), Langfuse
+        should use litellm_trace_id — NOT litellm_call_id. This ensures the
+        trace_id in Langfuse matches the Session ID shown in LiteLLM logs.
+        """
+        payload = self._build_standard_logging_payload()  # no trace_id
+        kwargs = self._build_langfuse_kwargs(payload)
+        kwargs["litellm_trace_id"] = "trace-id-from-kwargs"
+        self.last_trace_kwargs = {}
+
+        with patch(
+            "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
+            side_effect=lambda generation_params, **kwargs: generation_params,
+            create=True,
+        ):
+            self.logger._log_langfuse_v2(
+                user_id="user-1",
+                metadata={},
+                litellm_params={"metadata": {}},
+                output=None,
+                start_time=datetime.datetime.utcnow(),
+                end_time=datetime.datetime.utcnow(),
+                kwargs=kwargs,
+                optional_params={},
+                input=None,
+                response_obj=None,
+                level="ERROR",
+                litellm_call_id="call-id-xyz",
+            )
+
+        # litellm_trace_id should be preferred over litellm_call_id
+        assert self.last_trace_kwargs.get("id") == "trace-id-from-kwargs"
+
+    def test_log_langfuse_v2_uses_litellm_trace_id_when_standard_logging_object_none(self):
+        """
+        When standard_logging_object is None (failure case where
+        get_standard_logging_object_payload threw), litellm_trace_id from kwargs
+        should be used as the Langfuse trace_id. This matches the DB Session ID.
+        """
+        kwargs = {
+            "standard_logging_object": None,
+            "model": "gpt-4",
+            "call_type": "completion",
+            "cache_hit": False,
+            "messages": [],
+            "litellm_trace_id": "trace-id-failure",
+        }
+        self.last_trace_kwargs = {}
+
+        with patch(
+            "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params",
+            side_effect=lambda generation_params, **kwargs: generation_params,
+            create=True,
+        ):
+            self.logger._log_langfuse_v2(
+                user_id="user-1",
+                metadata={},
+                litellm_params={"metadata": {}},
+                output=None,
+                start_time=datetime.datetime.utcnow(),
+                end_time=datetime.datetime.utcnow(),
+                kwargs=kwargs,
+                optional_params={},
+                input=None,
+                response_obj=None,
+                level="ERROR",
+                litellm_call_id="call-id-different",
+            )
+
+        # Must use litellm_trace_id, not litellm_call_id
+        assert self.last_trace_kwargs.get("id") == "trace-id-failure"
 
     def test_log_langfuse_v2_session_id_passed_as_trace_session_id(self):
         """