From 2fc39cde18db6da24f2ee44fb9efb93f540f5476 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 14 Aug 2026 16:05:28 -0700 Subject: [PATCH] fix(langfuse): gate update_trace_keys behind an operator setting (#36862) update_trace_keys lets a caller name which request metadata entries get copied onto an existing trace, and the name is unrestricted. Sending update_trace_keys: ["user_api_key_auth"] with existing_trace_id serializes the resolved auth object, including the team callback credentials it carries, onto the trace through Langfuse.trace(**trace_params). TraceBody is Extra.allow, so an unexpected key ships rather than being dropped. Any holder of a team key can do this and read the result in the destination the team already logs to, so the feature is now inert unless an operator turns it on with langfuse_enable_update_trace_keys. --- litellm/__init__.py | 1 + litellm/integrations/langfuse/langfuse.py | 5 +- .../integrations/test_langfuse.py | 62 +++++++++++++++---- 3 files changed, 55 insertions(+), 13 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 056dd532f5f..8961de940a0 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -172,6 +172,7 @@ callbacks: List[ callback_settings: Dict[str, Dict[str, Any]] = {} initialized_langfuse_clients: int = 0 langfuse_default_tags: Optional[List[str]] = None +langfuse_enable_update_trace_keys: bool = False langsmith_batch_size: Optional[int] = None prometheus_initialize_budget_metrics: Optional[bool] = False prometheus_latency_buckets: Optional[List[float]] = None diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 8720f561e14..6d31f22b422 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -568,7 +568,10 @@ 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 - update_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ())) + 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 () + ) debug: Final = clean_metadata.pop("debug_langfuse", None) mask_input: Final = _as_steering_flag(clean_metadata.pop("mask_input", False)) mask_output: Final = _as_steering_flag(clean_metadata.pop("mask_output", False)) diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index de04a65c310..6e57a36c5b6 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -1,4 +1,5 @@ import datetime +import json import os import sys import types @@ -1332,35 +1333,72 @@ def test_mask_input_from_the_request_body_is_unchanged(mask_input, expect_redact assert (trace_params["input"] == _LANGFUSE_REDACTED) is expect_redacted -def test_update_trace_keys_header_applies_every_key(): +@pytest.mark.parametrize("flag", [True, "true"]) +def test_update_trace_keys_header_applies_every_key_when_enabled(flag): logger = _steering_logger() - trace_params, _ = _emit( - logger, - headers={ - "langfuse_existing_trace_id": "trace-1", - "langfuse_update_trace_keys": "trace_release, trace_tail", - "langfuse_trace_release": "v1.2.3", - "langfuse_trace_tail": "last", - }, - ) + with patch.object(litellm, "langfuse_enable_update_trace_keys", flag): + trace_params, _ = _emit( + logger, + headers={ + "langfuse_existing_trace_id": "trace-1", + "langfuse_update_trace_keys": "trace_release, trace_tail", + "langfuse_trace_release": "v1.2.3", + "langfuse_trace_tail": "last", + }, + ) assert trace_params["release"] == "v1.2.3" assert trace_params["tail"] == "last" -def test_update_trace_keys_from_the_request_body_list_is_unchanged(): +def test_update_trace_keys_is_off_by_default(): + """ + The caller picks the key name, so while the feature is on they can name + user_api_key_auth and have the resolved auth object, including team callback + credentials, serialized onto the trace. It stays inert until an operator opts in. + """ logger = _steering_logger() trace_params, _ = _emit( logger, metadata={ "existing_trace_id": "trace-1", - "update_trace_keys": ["trace_release"], + "update_trace_keys": ["user_api_key_auth", "trace_release"], + "user_api_key_auth": {"team_metadata": {"logging": [{"callback_vars": {"secret": "sk-canary"}}]}}, "trace_release": "v1.2.3", }, ) + assert "user_api_key_auth" not in trace_params + assert "release" not in trace_params + assert "sk-canary" not in json.dumps(trace_params, default=repr) + + +def test_update_trace_keys_input_and_output_are_gated_too(): + logger = _steering_logger() + + off, _ = _emit(logger, metadata={"existing_trace_id": "trace-1", "update_trace_keys": ["input", "output"]}) + with patch.object(litellm, "langfuse_enable_update_trace_keys", True): + on, _ = _emit(logger, metadata={"existing_trace_id": "trace-1", "update_trace_keys": ["input", "output"]}) + + assert "input" not in off and "output" not in off + assert "input" in on and "output" in on + + +def test_update_trace_keys_from_the_request_body_list_applies_when_enabled(): + logger = _steering_logger() + + with patch.object(litellm, "langfuse_enable_update_trace_keys", True): + trace_params, _ = _emit( + logger, + metadata={ + "existing_trace_id": "trace-1", + "update_trace_keys": ["trace_release"], + "trace_release": "v1.2.3", + }, + ) + assert trace_params["release"] == "v1.2.3"