fix(langfuse): give each call in a session its own trace instead of upserting one trace per session

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/integrations/langfuse/langfuse.py

@@ -63,6 +63,20 @@ def _object_mapping(value: object) -> Mapping[str, object] | None:
     return value if isinstance(value, dict) else None
 
 
+def _resolve_langfuse_trace_id(
+    metadata_trace_id: object,
+    standard_trace_id: object,
+    litellm_trace_id: object,
+    litellm_call_id: str | None,
+    session_id: object,
+) -> str | None:
+    """Resolve the Langfuse trace id; a session id never doubles as the trace id."""
+    resolved: Final = metadata_trace_id or standard_trace_id or litellm_trace_id or litellm_call_id
+    if session_id is not None and resolved == session_id:
+        return litellm_call_id
+    return str(resolved) if resolved is not None else None
+
+
 class _UsageObject(Protocol):
     """Token-count surface the Langfuse logger reads off a response usage payload."""
 
@@ -596,19 +610,26 @@ class LangFuseLogger:
 
             session_id: Final = clean_metadata.pop("session_id", None)
             trace_name = cast(str | None, clean_metadata.pop("trace_name", None))
-            trace_id = clean_metadata.pop("trace_id", None)
-            # Use standard_logging_object.trace_id if available (when trace_id from metadata is None)
-            # This allows standard trace_id to be used when provided in standard_logging_object
-            if trace_id is None and standard_logging_object is not None:
-                trace_id = cast(str | None, standard_logging_object.get("trace_id"))
-            # Fallback to litellm_call_id if no trace_id found
-            if trace_id is None:
-                trace_id = kwargs.get("litellm_trace_id") or litellm_call_id
+            metadata_trace_id: Final = clean_metadata.pop("trace_id", None)
+            standard_trace_id: Final = (
+                cast(str | None, standard_logging_object.get("trace_id"))
+                if standard_logging_object is not None
+                else None
+            )
             existing_trace_id: Final = 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
-            if existing_trace_id is not None:
-                trace_id = existing_trace_id
+            trace_id: Final = (
+                existing_trace_id
+                if existing_trace_id is not None
+                else _resolve_langfuse_trace_id(
+                    metadata_trace_id=metadata_trace_id,
+                    standard_trace_id=standard_trace_id,
+                    litellm_trace_id=kwargs.get("litellm_trace_id"),
+                    litellm_call_id=litellm_call_id,
+                    session_id=session_id,
+                )
+            )
             requested_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ()))
             update_trace_keys: Final = (
                 requested_trace_keys if _as_steering_flag(litellm.langfuse_enable_update_trace_keys) else ()

tests/test_litellm/integrations/test_langfuse.py

@@ -902,6 +902,69 @@ class TestLangfuseUsageDetails(unittest.TestCase):
         # session_id must still be set for session grouping
         assert self.last_trace_kwargs.get("session_id") == "session-999"
 
+    def _drive_session_call(self, litellm_call_id, metadata_trace_id="sess-abc"):
+        """One _log_langfuse_v2 call shaped like a proxied request that carries a session id."""
+        payload = self._build_standard_logging_payload(trace_id="sess-abc")
+        kwargs = self._build_langfuse_kwargs(payload)
+        kwargs["litellm_trace_id"] = "sess-abc"
+        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={"session_id": "sess-abc", "trace_id": metadata_trace_id},
+                litellm_params={"metadata": {"session_id": "sess-abc", "trace_id": metadata_trace_id}},
+                output=None,
+                start_time=datetime.datetime.utcnow(),
+                end_time=datetime.datetime.utcnow(),
+                kwargs=kwargs,
+                optional_params={},
+                input=None,
+                response_obj=None,
+                level="INFO",
+                litellm_call_id=litellm_call_id,
+            )
+        return self.last_trace_kwargs
+
+    def test_log_langfuse_v2_session_header_trace_id_falls_back_to_call_id(self):
+        """
+        When the proxy stamps the session id into every trace_id slot
+        (metadata, standard_logging_object, litellm_trace_id), the Langfuse
+        trace id must be the per-call litellm_call_id so each call gets its
+        own trace inside the shared session.
+        """
+        trace_kwargs = self._drive_session_call(litellm_call_id="call-1")
+
+        assert trace_kwargs.get("id") == "call-1"
+        assert trace_kwargs.get("session_id") == "sess-abc"
+
+    def test_log_langfuse_v2_session_calls_produce_distinct_trace_ids(self):
+        """
+        Two calls in one session must land as two Langfuse traces. With the
+        session id used as the trace id, Langfuse upserts both calls onto a
+        single trace, which is the regression this guards.
+        """
+        first = self._drive_session_call(litellm_call_id="call-1")
+        second = self._drive_session_call(litellm_call_id="call-2")
+
+        assert first.get("id") == "call-1"
+        assert second.get("id") == "call-2"
+        assert first.get("id") != second.get("id")
+
+    def test_log_langfuse_v2_explicit_distinct_trace_id_kept_with_session_id(self):
+        """
+        An explicit metadata.trace_id that differs from the session id is a
+        real trace id and must be kept, with the session id still applied.
+        """
+        trace_kwargs = self._drive_session_call(litellm_call_id="call-1", metadata_trace_id="trace-xyz")
+
+        assert trace_kwargs.get("id") == "trace-xyz"
+        assert trace_kwargs.get("session_id") == "sess-abc"
+
 
 def test_failure_handler_langfuse_kwargs_excludes_original_response():
     """