From a9c72887d6ff5f806f2d9f1c015ee177577ff06c Mon Sep 17 00:00:00 2001
From: Yucheng He <yucheng@berri.ai>
Date: Tue, 8 Sep 2026 22:57:15 -0700
Subject: [PATCH] fix(langfuse): preserve session trace provenance

---
 litellm/integrations/langfuse/langfuse.py     | 10 ++--
 .../integrations/test_langfuse.py             | 58 ++++++++++++++++++-
 2 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py
index daed1eab118..b75369965de 100644
--- a/litellm/integrations/langfuse/langfuse.py
+++ b/litellm/integrations/langfuse/langfuse.py
@@ -70,7 +70,7 @@ def _widened_items(mapping: Mapping[str, object]) -> Iterable[tuple[object, obje
 
 
 def _is_session_header_trace(trace_id: object, session_id: object, proxy_server_request: object) -> bool:
-    if not isinstance(session_id, str) or trace_id != session_id:
+    if not isinstance(trace_id, str) or not isinstance(session_id, str):
         return False
     request: Final = _object_mapping(proxy_server_request)
     raw_headers: Final = _object_mapping(request.get("headers")) if request is not None else None
@@ -83,14 +83,16 @@ def _is_session_header_trace(trace_id: object, session_id: object, proxy_server_
         return False
     if headers.get("langfuse_trace_id") is not None:
         return False
-    if headers.get("x-litellm-session-id") == session_id:
+    if trace_id != session_id and headers.get("langfuse_session_id") != session_id:
+        return False
+    if headers.get("x-litellm-session-id") == trace_id:
         return True
-    if re.fullmatch(r"[a-zA-Z0-9_\-]{8,}", session_id) is None:
+    if re.fullmatch(r"[a-zA-Z0-9_\-]{8,}", trace_id) is None:
         return False
     user_agent: Final = headers.get("user-agent")
     codex: Final = isinstance(user_agent, str) and re.match(r"^codex[-_ /]", user_agent, re.IGNORECASE) is not None
     return any(
-        value == session_id
+        value == trace_id
         and (
             key == "x-session-id"
             or re.fullmatch(r"x-.+-session-id", key) is not None
diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py
index d01bba18227..87e76499b84 100644
--- a/tests/test_litellm/integrations/test_langfuse.py
+++ b/tests/test_litellm/integrations/test_langfuse.py
@@ -1374,6 +1374,62 @@ def _emit(logger: LangFuseLogger, *, metadata=None, headers=None):
             {},
             "existing-trace",
         ),
+        (
+            {"x-litellm-session-id": "session-7125", "langfuse_session_id": "custom-session"},
+            {},
+            "call",
+        ),
+        (
+            {"x-litellm-session-id": "short", "langfuse_session_id": "custom-session"},
+            {},
+            "call",
+        ),
+        (
+            {"X-Claude-Code-Session-Id": "session-7125", "langfuse_session_id": "custom-session"},
+            {},
+            "call",
+        ),
+        (
+            {"x-session-id": "session-7125", "langfuse_session_id": "custom-session"},
+            {},
+            "call",
+        ),
+        (
+            {
+                "session-id": "session-7125",
+                "user-agent": "codex_cli_rs/1.0",
+                "langfuse_session_id": "custom-session",
+            },
+            {},
+            "call",
+        ),
+        (
+            {
+                "x-litellm-session-id": "session-7125",
+                "langfuse_session_id": "custom-session",
+                "x-litellm-trace-id": "explicit-trace",
+            },
+            {},
+            "explicit-trace",
+        ),
+        (
+            {
+                "x-litellm-session-id": "session-7125",
+                "langfuse_session_id": "custom-session",
+                "langfuse_trace_id": "explicit-trace",
+            },
+            {},
+            "explicit-trace",
+        ),
+        (
+            {
+                "x-litellm-session-id": "session-7125",
+                "langfuse_session_id": "custom-session",
+                "langfuse_existing_trace_id": "existing-trace",
+            },
+            {},
+            "existing-trace",
+        ),
         ({}, {"trace_id": "session-7125", "session_id": "session-7125"}, "session-7125"),
         ({}, {"trace_id": "explicit-trace", "session_id": "session-7125"}, "explicit-trace"),
         (
@@ -1443,7 +1499,7 @@ def test_session_header_trace_provenance(headers, metadata, expected_id, level):
         assert trace_params["id"] == (call_id if expected_id == "call" else expected_id)
         assert result["trace_id"] == trace_params["id"]
         if expected_id != "existing-trace":
-            assert trace_params["session_id"] == original_metadata.get("session_id")
+            assert trace_params["session_id"] == headers.get("langfuse_session_id", original_metadata.get("session_id"))
         steering = {key[len("langfuse_") :]: value for key, value in headers.items() if key.startswith("langfuse_")}
         assert data["metadata"] == {**original_metadata, **steering}