scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

fix(langfuse): preserve session trace provenance

Browse source
This commit is contained in:
Yucheng He 2026-09-08 22:57:15 -07:00
parent 5069d20736
commit a9c72887d6
2 changed files with 63 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/integrations/langfuse/langfuse.py
View file

@ -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

58
tests/test_litellm/integrations/test_langfuse.py
View file

@ -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}