fix(langfuse): iterate lease renewal instead of recursing, monkeypatch update_trace_keys flag in tests

The recursive lease fallback tripped tests/code_coverage_tests/recursive_detector.py; the renewal
candidates are now walked with itertools.chain. The six update_trace_keys tests set the litellm
global through pytest monkeypatch so the TQ008 budget stays within its ceiling

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/integrations/langfuse/langfuse_sdk.py

@@ -10,6 +10,7 @@ from contextvars import ContextVar
 from datetime import datetime
 from hashlib import sha256
 from importlib.metadata import version
+from itertools import chain
 from types import MappingProxyType
 from typing import Final
 from weakref import WeakKeyDictionary, WeakSet
@@ -350,17 +351,22 @@ def lease_langfuse_client(client: Langfuse, renew: Callable[[], Langfuse]) -> Ge
     it registers as a holder and the reference count degrades the queued teardown to a flush, or a
     fresh bundle once the old one is gone.
     """
-    state: Final = _lifecycle_state(client)
-    if not state.open_lease(client):
-        with lease_langfuse_client(renew(), renew) as leased:
-            yield leased
-        return
+    leased, state = _open_lease(client, renew)
     try:
-        yield client
+        yield leased
     finally:
         _run_teardowns(state, state.release_lease())
 
 
+def _open_lease(client: Langfuse, renew: Callable[[], Langfuse]) -> tuple[Langfuse, _LangfuseLifecycleState]:
+    """The first of ``client`` then ``renew()``'s clients that is not already retired, with its lease taken."""
+    return next(
+        (candidate, state)
+        for candidate in chain((client,), iter(renew, None))
+        if (state := _lifecycle_state(candidate)).open_lease(candidate)
+    )
+
+
 def _run_teardowns(state: _LangfuseLifecycleState, clients: tuple[Langfuse, ...]) -> None:
     """Tear down ``clients``, then whatever eviction queued meanwhile, and hand the flag back.
 

tests/test_litellm/integrations/test_langfuse.py

@@ -1848,19 +1848,19 @@ def test_mask_input_from_the_request_body_is_unchanged(mask_input, expect_redact
 
 
 @pytest.mark.parametrize("flag", [True, "true"])
-def test_update_trace_keys_header_applies_every_key_when_enabled(flag):
+def test_update_trace_keys_header_applies_every_key_when_enabled(flag, monkeypatch):
     rig = _steering_logger()
 
-    with patch.object(litellm, "langfuse_enable_update_trace_keys", flag):
-        trace_params, _, span = _emit(
-            rig,
-            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",
-            },
-        )
+    monkeypatch.setattr(litellm, "langfuse_enable_update_trace_keys", flag)
+    trace_params, _, span = _emit(
+        rig,
+        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"
@@ -1893,31 +1893,31 @@ def test_update_trace_keys_is_off_by_default():
     assert "sk-canary" not in json.dumps(dict(span.attributes or {}), default=repr)
 
 
-def test_update_trace_keys_input_and_output_are_gated_too():
+def test_update_trace_keys_input_and_output_are_gated_too(monkeypatch):
     rig = _steering_logger()
 
     off, _, _ = _emit(rig, metadata={"existing_trace_id": "trace-1", "update_trace_keys": ["input", "output"]})
-    with patch.object(litellm, "langfuse_enable_update_trace_keys", True):
-        on, _, _ = _emit(rig, metadata={"existing_trace_id": "trace-1", "update_trace_keys": ["input", "output"]})
+    monkeypatch.setattr(litellm, "langfuse_enable_update_trace_keys", True)
+    on, _, _ = _emit(rig, 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_input_output_reach_the_trace_even_under_a_parent():
+def test_update_trace_keys_input_output_reach_the_trace_even_under_a_parent(monkeypatch):
     """With a real parent the generation is not the trace root, so trace-level
     I/O must be stamped explicitly; v2 updated the trace object directly."""
     rig = _steering_logger()
 
-    with patch.object(litellm, "langfuse_enable_update_trace_keys", True):
-        _, _, span = _emit(
-            rig,
-            metadata={
-                "existing_trace_id": "trace-1",
-                "parent_observation_id": "b" * 16,
-                "update_trace_keys": ["input", "output"],
-            },
-        )
+    monkeypatch.setattr(litellm, "langfuse_enable_update_trace_keys", True)
+    _, _, span = _emit(
+        rig,
+        metadata={
+            "existing_trace_id": "trace-1",
+            "parent_observation_id": "b" * 16,
+            "update_trace_keys": ["input", "output"],
+        },
+    )
 
     assert "the-input" in str(span.attributes["langfuse.trace.input"])
     assert "the-output" in str(span.attributes["langfuse.trace.output"])
@@ -1945,54 +1945,54 @@ def test_a_fresh_trace_still_claims_root_so_its_generation_names_it():
     assert span.attributes["langfuse.trace.name"] == "first-call"
 
 
-def test_trace_io_is_not_stamped_when_update_trace_keys_does_not_ask():
+def test_trace_io_is_not_stamped_when_update_trace_keys_does_not_ask(monkeypatch):
     rig = _steering_logger()
 
-    with patch.object(litellm, "langfuse_enable_update_trace_keys", True):
-        _, _, span = _emit(
-            rig,
-            metadata={
-                "existing_trace_id": "trace-1",
-                "parent_observation_id": "b" * 16,
-                "update_trace_keys": ["trace_release"],
-            },
-        )
+    monkeypatch.setattr(litellm, "langfuse_enable_update_trace_keys", True)
+    _, _, span = _emit(
+        rig,
+        metadata={
+            "existing_trace_id": "trace-1",
+            "parent_observation_id": "b" * 16,
+            "update_trace_keys": ["trace_release"],
+        },
+    )
 
     assert "langfuse.trace.input" not in (span.attributes or {})
     assert "langfuse.trace.output" not in (span.attributes or {})
 
 
-def test_update_trace_keys_from_the_request_body_list_applies_when_enabled():
+def test_update_trace_keys_from_the_request_body_list_applies_when_enabled(monkeypatch):
     rig = _steering_logger()
 
-    with patch.object(litellm, "langfuse_enable_update_trace_keys", True):
-        trace_params, _, span = _emit(
-            rig,
-            metadata={
-                "existing_trace_id": "trace-1",
-                "update_trace_keys": ["trace_release"],
-                "trace_release": "v1.2.3",
-            },
-        )
+    monkeypatch.setattr(litellm, "langfuse_enable_update_trace_keys", True)
+    trace_params, _, span = _emit(
+        rig,
+        metadata={
+            "existing_trace_id": "trace-1",
+            "update_trace_keys": ["trace_release"],
+            "trace_release": "v1.2.3",
+        },
+    )
 
     assert trace_params["release"] == "v1.2.3"
     assert span.attributes["langfuse.release"] == "v1.2.3"
 
 
-def test_update_trace_keys_trace_metadata_reaches_the_trace_not_just_the_generation():
+def test_update_trace_keys_trace_metadata_reaches_the_trace_not_just_the_generation(monkeypatch):
     """v2 updated the trace object's metadata; v4 has to propagate it as a trace attribute."""
     rig = _steering_logger()
 
-    with patch.object(litellm, "langfuse_enable_update_trace_keys", True):
-        _, _, span = _emit(
-            rig,
-            metadata={
-                "existing_trace_id": "trace-1",
-                "parent_observation_id": "b" * 16,
-                "update_trace_keys": ["trace_metadata"],
-                "trace_metadata": {"step": 2, "note": "x" * 300},
-            },
-        )
+    monkeypatch.setattr(litellm, "langfuse_enable_update_trace_keys", True)
+    _, _, span = _emit(
+        rig,
+        metadata={
+            "existing_trace_id": "trace-1",
+            "parent_observation_id": "b" * 16,
+            "update_trace_keys": ["trace_metadata"],
+            "trace_metadata": {"step": 2, "note": "x" * 300},
+        },
+    )
 
     assert span.attributes["langfuse.trace.metadata.step"] == "2"
     assert span.attributes["langfuse.trace.metadata.note"] == "x" * 200