Merge main into codex/migrate-poetry-to-uv

Author: hassiebp

langfuse/_client/propagation.py

@@ -7,7 +7,9 @@
 
 from typing import Any, Dict, Generator, List, Literal, Optional, TypedDict, Union, cast
 
-from opentelemetry import baggage
+from opentelemetry import (
+    baggage,
+)
 from opentelemetry import (
     baggage as otel_baggage_api,
 )
@@ -17,6 +19,7 @@
 from opentelemetry import (
     trace as otel_trace_api,
 )
+from opentelemetry.context import _RUNTIME_CONTEXT
 from opentelemetry.util._decorator import (
     _AgnosticContextManager,
     _agnosticcontextmanager,
@@ -72,6 +75,22 @@ class PropagatedExperimentAttributes(TypedDict):
     experiment_item_root_observation_id: str
 
 
+def _detach_context_token_safely(token: Any) -> None:
+    """Detach a context token without emitting noisy async teardown errors.
+
+    OpenTelemetry tokens are backed by ``contextvars`` and must be detached in the
+    same execution context where they were attached. Async frameworks can legitimately
+    end spans or unwind context managers in a different task/context, in which case
+    detach raises and the public OpenTelemetry helper logs an error. At that point the
+    observation is already completed, so the mismatch is safe to ignore.
+    """
+
+    try:
+        _RUNTIME_CONTEXT.detach(token)
+    except Exception:
+        pass
+
+
 def propagate_attributes(
     *,
     user_id: Optional[str] = None,
@@ -272,7 +291,7 @@ def _propagate_attributes(
         yield
 
     finally:
-        otel_context_api.detach(token)
+        _detach_context_token_safely(token)
 
 
 def _get_propagated_attributes_from_context(

langfuse/_client/utils.py

@@ -5,6 +5,7 @@
 """
 
 import asyncio
+import contextvars
 import json
 import threading
 from hashlib import sha256
@@ -69,13 +70,14 @@ class _RunAsyncThread(threading.Thread):
 
     def __init__(self, coro: Coroutine[Any, Any, Any]) -> None:
         self.coro = coro
+        self.context = contextvars.copy_context()
         self.result: Any = None
         self.exception: Exception | None = None
         super().__init__()
 
     def run(self) -> None:
         try:
-            self.result = asyncio.run(self.coro)
+            self.result = self.context.run(asyncio.run, self.coro)
         except Exception as e:
             self.exception = e
 

langfuse/langchain/CallbackHandler.py

@@ -15,13 +15,13 @@
 
 import pydantic
 from opentelemetry import context, trace
-from opentelemetry.context import _RUNTIME_CONTEXT
 from opentelemetry.util._decorator import _AgnosticContextManager
 
 from langfuse import propagate_attributes
 from langfuse._client.attributes import LangfuseOtelSpanAttributes
 from langfuse._client.client import Langfuse
 from langfuse._client.get_client import get_client
+from langfuse._client.propagation import _detach_context_token_safely
 from langfuse._client.span import (
     LangfuseAgent,
     LangfuseChain,
@@ -458,18 +458,7 @@ def _detach_observation(
         token = self._context_tokens.pop(run_id, None)
 
         if token:
-            try:
-                # Directly detach from runtime context to avoid error logging
-                _RUNTIME_CONTEXT.detach(token)
-            except Exception:
-                # Context detach can fail in async scenarios - this is expected and safe to ignore
-                # The span itself was properly ended and tracing data is correctly captured
-                #
-                # Examples:
-                # 1. Token created in one async task/thread, detached in another
-                # 2. Context already detached by framework or other handlers
-                # 3. Runtime context state mismatch in concurrent execution
-                pass
+            _detach_context_token_safely(token)
 
         return cast(
             Union[
@@ -564,11 +553,8 @@ def on_chain_end(
                     input=kwargs.get("inputs"),
                 )
 
-                if (
-                    parent_run_id is None
-                    and self._propagation_context_manager is not None
-                ):
-                    self._propagation_context_manager.__exit__(None, None, None)
+                if parent_run_id is None:
+                    self._exit_propagation_context()
 
                 span.end()
 
@@ -579,6 +565,7 @@ def on_chain_end(
 
         finally:
             if parent_run_id is None:
+                self._exit_propagation_context()
                 self._reset()
 
     def on_chain_error(
@@ -608,10 +595,19 @@ def on_chain_error(
                     status_message=str(error) if level else None,
                     input=kwargs.get("inputs"),
                     cost_details={"total": 0},
-                ).end()
+                )
+
+                if parent_run_id is None:
+                    self._exit_propagation_context()
+
+                observation.end()
 
         except Exception as e:
             langfuse_logger.exception(e)
+        finally:
+            if parent_run_id is None:
+                self._exit_propagation_context()
+                self._reset()
 
     def on_chat_model_start(
         self,
@@ -1026,6 +1022,15 @@ def on_llm_error(
     def _reset(self) -> None:
         self._child_to_parent_run_id_map = {}
 
+    def _exit_propagation_context(self) -> None:
+        manager = self._propagation_context_manager
+
+        if manager is None:
+            return
+
+        self._propagation_context_manager = None
+        manager.__exit__(None, None, None)
+
     def __join_tags_and_metadata(
         self,
         tags: Optional[List[str]] = None,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "langfuse"
-version = "4.0.1"
+version = "4.0.3"
 description = "A client library for accessing langfuse"
 readme = "README.md"
 authors = [{ name = "langfuse", email = "developers@langfuse.com" }]

tests/test_propagate_attributes.py

@@ -2295,6 +2295,36 @@ def test_tags_attribute_key_format(self, langfuse_client, memory_exporter):
 class TestPropagateAttributesExperiment(TestPropagateAttributesBase):
     """Tests for experiment attribute propagation."""
 
+    @pytest.mark.asyncio
+    async def test_experiment_propagates_user_id_in_async_context(
+        self, langfuse_client, memory_exporter
+    ):
+        """Verify run_experiment keeps propagated attributes when called from async code."""
+        import asyncio
+
+        local_data = [{"input": "test input", "expected_output": "expected output"}]
+
+        async def async_task(*, item, **kwargs):
+            await asyncio.sleep(0.01)
+            return f"processed: {item['input']}"
+
+        with propagate_attributes(user_id="async-experiment-user"):
+            langfuse_client.run_experiment(
+                name="Async Experiment",
+                data=local_data,
+                task=async_task,
+            )
+
+        langfuse_client.flush()
+        time.sleep(0.1)
+
+        root_span = self.get_span_by_name(memory_exporter, "experiment-item-run")
+        self.verify_span_attribute(
+            root_span,
+            LangfuseOtelSpanAttributes.TRACE_USER_ID,
+            "async-experiment-user",
+        )
+
     def test_experiment_attributes_propagate_without_dataset(
         self, langfuse_client, memory_exporter
     ):

tests/test_utils.py

@@ -1,6 +1,7 @@
 """Test suite for utility functions in langfuse._client.utils module."""
 
 import asyncio
+import contextvars
 import threading
 from unittest import mock
 
@@ -81,6 +82,22 @@ async def check_thread_isolation():
         result = run_async_safely(check_thread_isolation())
         assert result == "isolated"
 
+    @pytest.mark.asyncio
+    async def test_run_async_context_preserves_contextvars(self):
+        """Test that threaded execution preserves the caller's contextvars."""
+        request_id = contextvars.ContextVar("request_id")
+        token = request_id.set("req-123")
+
+        async def read_contextvar():
+            await asyncio.sleep(0.001)
+            return request_id.get()
+
+        try:
+            result = run_async_safely(read_contextvar())
+            assert result == "req-123"
+        finally:
+            request_id.reset(token)
+
     def test_multiple_calls_sync_context(self):
         """Test multiple sequential calls in sync context."""