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

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

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

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"