diff --git a/plugins/observability/langfuse/__init__.py b/plugins/observability/langfuse/__init__.py index 31904d47e3052..f8fe00f7ac47b 100644 --- a/plugins/observability/langfuse/__init__.py +++ b/plugins/observability/langfuse/__init__.py @@ -19,6 +19,9 @@ HERMES_LANGFUSE_SAMPLE_RATE - sampling rate 0.0–1.0 (default: 1.0) HERMES_LANGFUSE_MAX_CHARS - max chars per field (default: 12000) HERMES_LANGFUSE_DEBUG - set to "true" for verbose logging + HERMES_LANGFUSE_END_USER_ATTRIBUTION + - set to "true" to attribute the trace userId to + the messaging end-user (sender_id); default off """ from __future__ import annotations @@ -50,6 +53,7 @@ class TraceState: pending_tools_by_name: Dict[str, list] = field(default_factory=dict) turn_tool_calls: list[dict[str, Any]] = field(default_factory=list) last_updated_at: float = field(default_factory=time.time) + session_id: str = "" _STATE_LOCK = threading.Lock() @@ -63,6 +67,10 @@ class TraceState: # is far above any realistic concurrent-live-turn working set; it exists only # to bound the leak from non-finalizing turns, not to limit concurrency. _MAX_TRACE_STATE = 256 +# session_id -> messaging end-user id (sender_id). Captured from the turn-scoped +# pre_llm_call hook so the root trace (created later from pre_api_request, which +# carries no user identity) can attribute itself to that user via Langfuse user_id. +_SENDER_BY_SESSION: Dict[str, str] = {} _LANGFUSE_CLIENT = None _READ_FILE_LINE_RE = re.compile(r"^\s*(\d+)\|(.*)$") _READ_FILE_HEAD_LINES = 25 @@ -96,6 +104,14 @@ def _debug_enabled() -> bool: return _env_bool("HERMES_LANGFUSE_DEBUG") +def _end_user_attribution_enabled() -> bool: + # Opt-in (default off): attribute the trace userId to the messaging + # end-user (sender_id). Kept off by default so it never conflicts with + # other userId semantics (e.g. assigning the active profile, #26455) — when + # unset the trace carries no user, exactly like the pre-existing behavior. + return _env_bool("HERMES_LANGFUSE_END_USER_ATTRIBUTION") + + def _debug(message: str) -> None: if _debug_enabled(): logger.info("Langfuse tracing: %s", message) @@ -602,7 +618,7 @@ def _usage_and_cost(response: Any, *, provider: str, api_mode: str, model: str, def _start_root_trace(task_key: str, *, task_id: str, session_id: str, platform: str, provider: str, model: str, api_mode: str, messages: Any, client: Langfuse, - turn_id: str = "", api_request_id: str = "") -> TraceState: + turn_id: str = "", api_request_id: str = "", user_id: str = "") -> TraceState: trace_id = client.create_trace_id(seed=f"{session_id or 'sessionless'}::{task_id or task_key}") trace_input = _extract_last_user_message(messages) metadata = { @@ -615,40 +631,19 @@ def _start_root_trace(task_key: str, *, task_id: str, session_id: str, platform: "model": model, "api_mode": api_mode, } + # Also record the end-user on trace metadata as a fallback: an older SDK + # whose propagate_attributes drops the user_id kwarg (see the retry below) + # still surfaces the user here. + if user_id: + metadata["user_id"] = user_id # session_id must be passed in trace_context for Langfuse session grouping. trace_ctx: Dict[str, Any] = {"trace_id": trace_id} if session_id: trace_ctx["session_id"] = session_id - if propagate_attributes is not None: - try: - with propagate_attributes( - session_id=session_id or task_key, - trace_name="Hermes turn", - tags=["hermes", "langfuse"], - ): - root_ctx = client.start_as_current_observation( - trace_context=trace_ctx, - name="Hermes turn", - as_type="chain", - input=trace_input, - metadata=metadata, - end_on_exit=False, - ) - root_span = root_ctx.__enter__() - except Exception: - root_ctx = client.start_as_current_observation( - trace_context=trace_ctx, - name="Hermes turn", - as_type="chain", - input=trace_input, - metadata=metadata, - end_on_exit=False, - ) - root_span = root_ctx.__enter__() - else: - root_ctx = client.start_as_current_observation( + def _open_root() -> tuple[Any, Any]: + ctx = client.start_as_current_observation( trace_context=trace_ctx, name="Hermes turn", as_type="chain", @@ -656,7 +651,34 @@ def _start_root_trace(task_key: str, *, task_id: str, session_id: str, platform: metadata=metadata, end_on_exit=False, ) - root_span = root_ctx.__enter__() + return ctx, ctx.__enter__() + + if propagate_attributes is not None: + attr_kwargs: Dict[str, Any] = { + "session_id": session_id or task_key, + "trace_name": "Hermes turn", + "tags": ["hermes", "langfuse"], + } + if user_id: + attr_kwargs["user_id"] = user_id + try: + try: + attr_cm = propagate_attributes(**attr_kwargs) + except TypeError: + # Older Langfuse SDKs don't accept user_id on + # propagate_attributes. Drop only that kwarg and retry so + # session grouping / trace_name / tags are still applied — + # without this, an unsupported user_id would fall through to + # the bare-observation fallback below and silently lose all + # three. + attr_kwargs.pop("user_id", None) + attr_cm = propagate_attributes(**attr_kwargs) + with attr_cm: + root_ctx, root_span = _open_root() + except Exception: + root_ctx, root_span = _open_root() + else: + root_ctx, root_span = _open_root() try: root_span.set_trace_io(input=trace_input) @@ -664,7 +686,7 @@ def _start_root_trace(task_key: str, *, task_id: str, session_id: str, platform: pass _debug(f"started trace {trace_id} for {task_key}") - return TraceState(trace_id=trace_id, root_ctx=root_ctx, root_span=root_span) + return TraceState(trace_id=trace_id, root_ctx=root_ctx, root_span=root_span, session_id=session_id) def _start_child_observation(state: TraceState, *, client: Langfuse, name: str, as_type: str, @@ -728,6 +750,11 @@ def _evict_stale_locked() -> None: stale = sorted(_TRACE_STATE.items(), key=lambda kv: kv[1].last_updated_at)[:over] for key, state in stale: _TRACE_STATE.pop(key, None) + # Bound _SENDER_BY_SESSION by the same eviction that bounds _TRACE_STATE: + # a turn that never reaches _finish_trace (interrupted / tool-only final + # step) would otherwise leave its sender entry behind forever. + if state.session_id: + _SENDER_BY_SESSION.pop(state.session_id, None) try: state.root_span.end() except Exception as exc: # pragma: no cover - fail-open @@ -741,6 +768,8 @@ def _finish_trace(task_key: str, *, output: Any = None) -> None: with _STATE_LOCK: state = _TRACE_STATE.pop(task_key, None) + if state is not None and state.session_id: + _SENDER_BY_SESSION.pop(state.session_id, None) if state is None: return @@ -778,7 +807,27 @@ def on_pre_llm_call(*, task_id: str = "", session_id: str = "", platform: str = provider: str = "", base_url: str = "", api_mode: str = "", api_call_count: int = 0, messages: Any = None, turn_type: str = "user", conversation_history: Any = None, user_message: Any = None, - turn_id: str = "", api_request_id: str = "", **_: Any) -> None: + turn_id: str = "", api_request_id: str = "", + sender_id: str = "", **_: Any) -> None: + client = _get_langfuse() + if client is None: + return + + # End-user attribution is opt-in (default off): when disabled, drop the + # sender so every downstream path (stash, request-shaped trace creation, + # and the pre_api_request read) sees no user and behaves like upstream. + if sender_id and not _end_user_attribution_enabled(): + sender_id = "" + + # Capture the messaging end-user (emitted on the turn-scoped pre_llm_call) + # before the early return below, keyed by session_id, so the root trace + # created later from pre_api_request — which carries no user identity — can + # attribute itself to that user. Gated on an active client so it stays + # symmetric with the cleanup in _finish_trace / _evict_stale_locked. + if session_id and sender_id: + with _STATE_LOCK: + _SENDER_BY_SESSION[session_id] = sender_id + # Older Hermes branches used pre_llm_call for request-scoped tracing and # passed the actual API messages. Current Hermes also has a turn-scoped # pre_llm_call used for context injection; tracing that hook creates an @@ -787,10 +836,6 @@ def on_pre_llm_call(*, task_id: str = "", session_id: str = "", platform: str = if not isinstance(messages, list): return - client = _get_langfuse() - if client is None: - return - # messages is a list only for legacy Hermes branches that fired # pre_llm_call with API messages directly. Current Hermes fires # pre_llm_call for context injection (conversation_history/user_message, @@ -817,6 +862,7 @@ def on_pre_llm_call(*, task_id: str = "", session_id: str = "", platform: str = client=client, turn_id=turn_id, api_request_id=api_request_id, + user_id=sender_id, ) _evict_stale_locked() _TRACE_STATE[task_key] = state @@ -881,6 +927,7 @@ def on_pre_llm_request( client=client, turn_id=turn_id, api_request_id=api_request_id, + user_id=_SENDER_BY_SESSION.get(session_id, ""), ) _evict_stale_locked() _TRACE_STATE[task_key] = state diff --git a/tests/plugins/test_langfuse_plugin.py b/tests/plugins/test_langfuse_plugin.py index dd58149eba2e5..9d878969abe61 100644 --- a/tests/plugins/test_langfuse_plugin.py +++ b/tests/plugins/test_langfuse_plugin.py @@ -1021,3 +1021,248 @@ class _Resp: assert seen["resp"] is resp assert captured["usage_details"] == {"input": 7, "output": 3} + + +# --------------------------------------------------------------------------- +# End-user attribution: the messaging sender_id should be set as the Langfuse +# trace userId (via propagate_attributes + trace metadata), with the sender +# captured on the turn-scoped pre_llm_call and read back on pre_api_request. +# --------------------------------------------------------------------------- + + +class TestEndUserAttribution: + @staticmethod + def _fresh_plugin(): + mod_name = "plugins.observability.langfuse" + sys.modules.pop(mod_name, None) + return importlib.import_module(mod_name) + + @staticmethod + def _fake_client(captured): + class _Span: + def __init__(self, kind="root"): + self._kind = kind + + def set_trace_io(self, **k): + pass + + def update(self, **k): + pass + + def start_observation(self, **k): + captured.setdefault("child_observations", []).append(k) + return _Span(kind="child") + + def end(self): + captured.setdefault("ended", []).append(self._kind) + + class _Ctx: + def __enter__(self): + return _Span() + + def __exit__(self, *a): + return False + + class _Client: + def create_trace_id(self, seed=None): + return "trace-id" + + def start_as_current_observation(self, **kwargs): + captured.setdefault("observations", []).append(kwargs) + return _Ctx() + + def flush(self): + pass + + return _Client() + + def _wire(self, monkeypatch, captured, enabled=True): + from contextlib import contextmanager + + mod = self._fresh_plugin() + client = self._fake_client(captured) + + @contextmanager + def _fake_propagate(**kwargs): + captured["propagate"] = kwargs + yield + + monkeypatch.setattr(mod, "_get_langfuse", lambda: client) + monkeypatch.setattr(mod, "propagate_attributes", _fake_propagate) + # End-user attribution is opt-in; enable it for the attribution tests. + if enabled: + monkeypatch.setenv("HERMES_LANGFUSE_END_USER_ATTRIBUTION", "true") + else: + monkeypatch.delenv("HERMES_LANGFUSE_END_USER_ATTRIBUTION", raising=False) + mod._TRACE_STATE.clear() + mod._SENDER_BY_SESSION.clear() + return mod + + def test_sender_id_sets_trace_user_id(self, monkeypatch): + captured = {} + mod = self._wire(monkeypatch, captured) + + # Legacy request-shaped pre_llm_call (messages is a list) carries the + # sender and creates the root trace in one shot. + mod.on_pre_llm_call( + task_id="t1", session_id="s1", sender_id="U12345", + messages=[{"role": "user", "content": "hi"}], + ) + + assert mod._SENDER_BY_SESSION.get("s1") == "U12345" + assert captured["propagate"]["user_id"] == "U12345" + assert captured["observations"][0]["metadata"]["user_id"] == "U12345" + + def test_pre_api_request_reads_stashed_sender_and_finish_clears_it(self, monkeypatch): + captured = {} + mod = self._wire(monkeypatch, captured) + + # Turn-scoped pre_llm_call (no messages list) only stashes the sender. + mod.on_pre_llm_call(task_id="t2", session_id="s2", sender_id="U999", messages=None) + assert mod._SENDER_BY_SESSION.get("s2") == "U999" + assert "propagate" not in captured # no trace created yet + + # The real request creates the root trace, picks up the stash, and opens + # a child generation through the real _start_child_observation path. + mod.on_pre_llm_request( + task_id="t2", session_id="s2", api_call_count=1, + request_messages=[{"role": "user", "content": "hi"}], + ) + assert captured["propagate"]["user_id"] == "U999" + assert len(captured.get("child_observations", [])) == 1 + + # Finishing the trace runs the real _finish_trace/_end_observation path: + # the child generation and the root span are both ended, and the + # per-session sender mapping is cleared. + mod._finish_trace(mod._trace_key("t2", "s2")) + assert captured.get("ended") == ["child", "root"] + assert "s2" not in mod._SENDER_BY_SESSION + + def test_no_sender_leaves_user_id_unset(self, monkeypatch): + captured = {} + mod = self._wire(monkeypatch, captured) + + mod.on_pre_llm_call( + task_id="t3", session_id="s3", + messages=[{"role": "user", "content": "hi"}], + ) + + assert "s3" not in mod._SENDER_BY_SESSION + # With no end-user, user_id is omitted entirely (not passed as None), + # so an SDK that predates the user_id kwarg is never exercised. + assert "user_id" not in captured["propagate"] + assert "user_id" not in captured["observations"][0]["metadata"] + + def test_legacy_sdk_without_user_id_kwarg_still_groups_session(self, monkeypatch): + # An older Langfuse SDK whose propagate_attributes has no user_id + # parameter must not lose session grouping / trace_name / tags: the + # plugin retries without user_id instead of falling through to the + # bare-observation fallback. + from contextlib import contextmanager + + captured = {} + mod = self._wire(monkeypatch, captured) + + @contextmanager + def _legacy_propagate(*, session_id=None, trace_name=None, tags=None): + captured["propagate"] = { + "session_id": session_id, + "trace_name": trace_name, + "tags": tags, + } + yield + + monkeypatch.setattr(mod, "propagate_attributes", _legacy_propagate) + + mod.on_pre_llm_call( + task_id="t4", session_id="s4", sender_id="U777", + messages=[{"role": "user", "content": "hi"}], + ) + + # propagate_attributes was still entered with grouping intact... + assert captured["propagate"]["session_id"] == "s4" + assert captured["propagate"]["trace_name"] == "Hermes turn" + assert "user_id" not in captured["propagate"] + # ...and the user_id still landed in trace metadata as a fallback. + assert captured["observations"][0]["metadata"]["user_id"] == "U777" + + def test_eviction_clears_stashed_sender_for_unfinalized_turn(self, monkeypatch): + # A turn that never reaches _finish_trace (interrupted / tool-only final + # step) must not leak its sender entry: _evict_stale_locked drops it + # alongside the _TRACE_STATE entry, so _SENDER_BY_SESSION is bounded by + # the same cap as _TRACE_STATE instead of growing per distinct session. + captured = {} + mod = self._wire(monkeypatch, captured) + monkeypatch.setattr(mod, "_MAX_TRACE_STATE", 1) + + # Turn A: stash sender, create its (never-finalized) root trace. + mod.on_pre_llm_call(task_id="A", session_id="sA", sender_id="UA", messages=None) + mod.on_pre_llm_request( + task_id="A", session_id="sA", api_call_count=1, + request_messages=[{"role": "user", "content": "hi"}], + ) + assert mod._SENDER_BY_SESSION.get("sA") == "UA" + + # Turn B in a different session forces eviction of turn A (cap=1). + mod.on_pre_llm_call(task_id="B", session_id="sB", sender_id="UB", messages=None) + mod.on_pre_llm_request( + task_id="B", session_id="sB", api_call_count=1, + request_messages=[{"role": "user", "content": "hi"}], + ) + + # Turn A was evicted from _TRACE_STATE AND its sender entry dropped. + assert "A" not in mod._TRACE_STATE + assert "sA" not in mod._SENDER_BY_SESSION + + def test_propagate_raises_non_typeerror_still_creates_trace(self, monkeypatch): + # The outer `except Exception` must keep creating the root observation + # (without session propagation) when propagate_attributes fails for a + # reason other than the user_id kwarg — the trace is not lost. + captured = {} + mod = self._wire(monkeypatch, captured) + + def _boom(**kwargs): + raise RuntimeError("propagate unavailable") + + monkeypatch.setattr(mod, "propagate_attributes", _boom) + + mod.on_pre_llm_call( + task_id="t", session_id="s", sender_id="U1", + messages=[{"role": "user", "content": "hi"}], + ) + + assert len(captured.get("observations", [])) == 1 + assert captured["observations"][0]["metadata"]["user_id"] == "U1" + assert "propagate" not in captured # propagate failed before entering + + def test_propagate_attributes_none_still_creates_trace(self, monkeypatch): + # When the SDK lacks propagate_attributes entirely, the bare-observation + # branch still opens the root trace with the user on metadata. + captured = {} + mod = self._wire(monkeypatch, captured) + monkeypatch.setattr(mod, "propagate_attributes", None) + + mod.on_pre_llm_call( + task_id="t", session_id="s", sender_id="U1", + messages=[{"role": "user", "content": "hi"}], + ) + + assert len(captured.get("observations", [])) == 1 + assert captured["observations"][0]["metadata"]["user_id"] == "U1" + + def test_attribution_disabled_by_default(self, monkeypatch): + # Without HERMES_LANGFUSE_END_USER_ATTRIBUTION the sender is dropped: + # nothing is stashed and the trace carries no user_id, so the plugin + # behaves exactly like upstream and never conflicts with other userId + # semantics (e.g. #26455). + captured = {} + mod = self._wire(monkeypatch, captured, enabled=False) + + mod.on_pre_llm_call( + task_id="t", session_id="s", sender_id="U1", + messages=[{"role": "user", "content": "hi"}], + ) + + assert "s" not in mod._SENDER_BY_SESSION + assert "user_id" not in captured["propagate"] + assert "user_id" not in captured["observations"][0]["metadata"]