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 82

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

Browse source 
* fix(langfuse): give session-header calls their own trace id

A client that sends only a session header (x-litellm-session-id, a vendor
x-<name>-session-id such as Claude Code's X-Claude-Code-Session-Id, a bare
x-session-id, or the Codex session/thread/conversation family) has that value
stamped into both trace_id and session_id by the proxy. Langfuse upserts a
trace by id, so every turn of a session collapsed into one growing trace and
the Sessions view showed "Total traces: 1"

Detect that aliasing in the Langfuse callback from the request headers the
callback already receives, and use litellm_call_id as the trace id for those
calls. session_id still carries the header value, so the turns stay grouped
under one session. An explicit x-litellm-trace-id, langfuse_trace_id, or
langfuse_existing_trace_id keeps its trace id, including when the caller sets
it to the same value as the session id

Co-authored-by: jesus <jesus@berri.ai>

* test(langfuse): cover direct-SDK callers without proxy request headers

* fix(langfuse): preserve session trace provenance

---------

Co-authored-by: jesus <jesus@berri.ai>
This commit is contained in:
yucheng-berri 2026-09-09 15:56:38 -07:00 • committed by GitHub
parent eb45a088d3
commit fc161faa96
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 306 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -2,6 +2,7 @@
# On success, logs events to Langfuse
import inspect
import os
import re
import traceback
from collections.abc import Callable, Iterable, Mapping
from datetime import datetime
@ -63,6 +64,44 @@ def _object_mapping(value: object) -> Mapping[str, object] | None:
return value if isinstance(value, dict) else None
def _widened_items(mapping: Mapping[str, object]) -> Iterable[tuple[object, object]]:
"""Header pairs with the key type widened back to what a caller-supplied dict can actually hold."""
return mapping.items()
def _is_session_header_trace(trace_id: object, session_id: object, proxy_server_request: object) -> bool:
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
if raw_headers is None:
return False
headers: Final = MappingProxyType(
{key.lower(): value for key, value in _widened_items(raw_headers) if isinstance(key, str)}
)
if headers.get("x-litellm-trace-id"):
return False
if headers.get("langfuse_trace_id") is not None:
return False
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,}", 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 == trace_id
and (
key == "x-session-id"
or re.fullmatch(r"x-.+-session-id", key) is not None
or (codex and key in ("session-id", "session_id", "thread-id", "conversation_id"))
)
for key, value in headers.items()
)
class _UsageObject(Protocol):
"""Token-count surface the Langfuse logger reads off a response usage payload."""
@ -609,6 +648,18 @@ class LangFuseLogger:
# 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
resolved_trace_id: Final = (
litellm_call_id or trace_id
if existing_trace_id is None
and _is_session_header_trace(trace_id, session_id, litellm_params.get("proxy_server_request"))
else trace_id
)
if resolved_trace_id != trace_id:
verbose_logger.debug(
"Langfuse: trace_id %s came from a session header; using call id %s so each call gets its own trace",
trace_id,
resolved_trace_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 ()
@ -663,7 +714,7 @@ class LangFuseLogger:
trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm"
else: # don't overwrite an existing trace
trace_params = {
"id": trace_id,
"id": resolved_trace_id,
"name": trace_name,
"session_id": session_id,
"input": masked_input if not mask_input else "redacted-by-litellm",
@ -845,13 +896,13 @@ class LangFuseLogger:
# Verify langfuse accepted our trace_id; if it differs, log a warning but still return our intended value
# to match expected test behavior
if hasattr(generation_client, "trace_id") and generation_client.trace_id:
if generation_client.trace_id != trace_id:
if generation_client.trace_id != resolved_trace_id:
verbose_logger.warning(
"Langfuse trace_id mismatch: set %s, but langfuse returned %s. Using our intended trace_id for consistency.",
trace_id,
resolved_trace_id,
generation_client.trace_id,
)
return trace_id, generation_id
return resolved_trace_id, generation_id
except Exception:
verbose_logger.error("Langfuse Layer Error - %s", traceback.format_exc())
return None, None

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

@ -1341,6 +1341,257 @@ def _emit(logger: LangFuseLogger, *, metadata=None, headers=None):
)
@pytest.mark.parametrize("level", ["DEFAULT", "ERROR"])
@pytest.mark.parametrize(
"headers,metadata,expected_id",
[
({"x-litellm-session-id": "session-7125"}, {}, "call"),
({"X-Claude-Code-Session-Id": "session-7125"}, {}, "call"),
({"x-session-id": "session-7125"}, {}, "call"),
({"session-id": "session-7125", "user-agent": "codex_cli_rs/1.0"}, {}, "call"),
({"thread-id": "session-7125", "user-agent": "codex-tui"}, {}, "call"),
({"session_id": "session-7125", "user-agent": "Codex 1.0"}, {}, "call"),
({"conversation_id": "session-7125", "user-agent": "codex_vscode/1.0"}, {}, "call"),
({"x-litellm-session-id": "short"}, {}, "call"),
({"x-litellm-trace-id": "session-7125"}, {}, "session-7125"),
(
{"X-LiteLLM-Trace-Id": "session-7125", "x-litellm-session-id": "session-7125"},
{},
"session-7125",
),
(
{"x-litellm-session-id": "session-7125", "langfuse_trace_id": "session-7125"},
{},
"session-7125",
),
(
{"x-litellm-session-id": "session-7125", "langfuse_trace_id": "explicit-trace"},
{},
"explicit-trace",
),
(
{"x-litellm-session-id": "session-7125", "langfuse_existing_trace_id": "existing-trace"},
{},
"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"),
(
{"x-vendor-session-id": "short"},
{"trace_id": "short", "session_id": "short"},
"short",
),
(
{"x-session-id": "invalid value"},
{"trace_id": "invalid value", "session_id": "invalid value"},
"invalid value",
),
(
{"session-id": "session-7125", "user-agent": "codexfoo/1.0"},
{"trace_id": "session-7125", "session_id": "session-7125"},
"session-7125",
),
(
{"x-vendor-session-id": "short"},
{"trace_id": "session-7125", "session_id": "session-7125"},
"session-7125",
),
({}, {}, "call"),
],
)
def test_session_header_trace_provenance(headers, metadata, expected_id, level):
from starlette.datastructures import Headers
from litellm.proxy.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
clean_headers,
redact_credential_headers,
)
logger: Final = _steering_logger()
for turn in range(2):
call_id = f"call-{turn}"
request_headers = Headers(headers)
data = LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
headers=request_headers, data={"metadata": dict(metadata)}, _metadata_variable_name="metadata"
)
original_metadata = dict(data["metadata"])
now = datetime.datetime.now()
result = logger.log_event_on_langfuse(
kwargs={
"call_type": "completion",
"litellm_call_id": call_id,
"litellm_trace_id": data.get("litellm_trace_id"),
"litellm_params": {
"metadata": data["metadata"],
"proxy_server_request": {"headers": redact_credential_headers(clean_headers(request_headers))},
},
"messages": [{"role": "user", "content": f"turn {turn}"}],
"optional_params": {},
},
response_obj=(
None
if level == "ERROR"
else litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}])
),
start_time=now,
end_time=now,
level=level,
status_message="provider error" if level == "ERROR" else None,
)
trace_params = logger.Langfuse.trace.call_args.kwargs
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"] == 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}
def test_session_header_trace_without_call_id_keeps_session_alias():
logger: Final = _steering_logger()
now: Final = datetime.datetime.now()
result: Final = logger.log_event_on_langfuse(
kwargs={
"call_type": "completion",
"litellm_call_id": "",
"litellm_params": {
"metadata": {"trace_id": "session-7125", "session_id": "session-7125"},
"proxy_server_request": {"headers": {"x-litellm-session-id": "session-7125"}},
},
"messages": [{"role": "user", "content": "no call id"}],
"optional_params": {},
},
response_obj=litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]),
start_time=now,
end_time=now,
)
assert logger.Langfuse.trace.call_args.kwargs["id"] == "session-7125"
assert result["trace_id"] == "session-7125"
def test_every_proxy_session_header_shape_is_classified_as_a_session_alias():
"""The classifier must cover every header shape the proxy turns into a chain id."""
from litellm.integrations.langfuse.langfuse import _is_session_header_trace
from litellm.proxy.litellm_pre_call_utils import (
_CODEX_SESSION_ID_HEADERS,
get_chain_id_from_headers,
)
session: Final = "session-7125-abcdef"
session_shapes: Final = (
{"x-litellm-session-id": session},
{"X-Claude-Code-Session-Id": session},
{"x-session-id": session},
*({header: session, "user-agent": "codex_cli_rs/1.0"} for header in _CODEX_SESSION_ID_HEADERS),
)
for headers in session_shapes:
assert get_chain_id_from_headers(dict(headers)) == session, headers
assert _is_session_header_trace(session, session, {"headers": headers}) is True, headers
explicit_trace: Final = {"x-litellm-trace-id": session, "x-litellm-session-id": session}
assert get_chain_id_from_headers(dict(explicit_trace)) == session
assert _is_session_header_trace(session, session, {"headers": explicit_trace}) is False
@pytest.mark.parametrize(
"proxy_server_request",
[None, {}, {"headers": None}],
ids=["no-proxy-request", "no-headers-key", "null-headers"],
)
def test_sdk_caller_without_request_headers_keeps_its_trace(proxy_server_request):
"""A direct SDK caller has no request headers, so a session-shaped trace id stays the caller's."""
logger: Final = _steering_logger()
now: Final = datetime.datetime.now()
result: Final = logger.log_event_on_langfuse(
kwargs={
"call_type": "completion",
"litellm_call_id": "call-0",
"litellm_params": {
"metadata": {"trace_id": "session-7125", "session_id": "session-7125"},
"proxy_server_request": proxy_server_request,
},
"messages": [{"role": "user", "content": "sdk turn"}],
"optional_params": {},
},
response_obj=litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "OK"}}]),
start_time=now,
end_time=now,
)
assert logger.Langfuse.trace.call_args.kwargs["id"] == "session-7125"
assert result["trace_id"] == "session-7125"
def test_session_header_classifier_survives_non_string_header_keys():
"""A non-string header key must not cost the caller its whole trace."""
from litellm.integrations.langfuse.langfuse import _is_session_header_trace
session: Final = "session-7125-abcdef"
headers: Final = {7: "numeric key", "x-litellm-session-id": session}
assert _is_session_header_trace(session, session, {"headers": headers}) is True
assert _is_session_header_trace(session, session, {"headers": {7: "numeric key"}}) is False
def test_mask_input_header_false_keeps_the_prompt():
logger = _steering_logger()