fix(langchain): propagate root trace attributes to child runs under ainvoke

langfuse/langchain/CallbackHandler.py

@@ -1,9 +1,11 @@
 from collections import OrderedDict
+from contextlib import contextmanager
 from contextvars import Token
 from dataclasses import dataclass, field
 from typing import (
     Any,
     Dict,
+    Iterator,
     List,
     Literal,
     Optional,
@@ -23,7 +25,11 @@
 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.propagation import (
+    _detach_context_token_safely,
+    _get_propagated_context_key,
+    propagated_keys,
+)
 from langfuse._client.span import (
     LangfuseAgent,
     LangfuseChain,
@@ -90,6 +96,7 @@ class _RootRunState:
     run_ids: Set[UUID] = field(default_factory=set)
     resume_key: Optional[str] = None
     propagation_context_manager: Optional[_AgnosticContextManager] = None
+    propagation_context: Optional[context.Context] = None
 
 
 class _PendingResumeTraceContextStore:
@@ -580,6 +587,9 @@ def on_chain_start(
 
                 propagation_context_manager.__enter__()
 
+                if root_run_state is not None:
+                    root_run_state.propagation_context = context.get_current()
+
             obs = self._get_parent_observation(parent_run_id)
             if isinstance(obs, Langfuse):
                 resume_key, trace_context = self._take_root_trace_context(
@@ -597,16 +607,17 @@ def on_chain_start(
                     ),
                 )
             else:
-                span = obs.start_observation(
-                    name=span_name,
-                    as_type=observation_type,
-                    metadata=span_metadata,
-                    input=inputs,
-                    level=cast(
-                        Literal["DEBUG", "DEFAULT", "WARNING", "ERROR"] | None,
-                        span_level,
-                    ),
-                )
+                with self._use_root_propagated_attributes(run_id):
+                    span = obs.start_observation(
+                        name=span_name,
+                        as_type=observation_type,
+                        metadata=span_metadata,
+                        input=inputs,
+                        level=cast(
+                            Literal["DEBUG", "DEFAULT", "WARNING", "ERROR"] | None,
+                            span_level,
+                        ),
+                    )
 
             self._attach_observation(run_id, span)
 
@@ -974,13 +985,16 @@ def on_tool_start(
                     level="DEBUG" if tags and LANGSMITH_TAG_HIDDEN in tags else None,
                 )
             else:
-                span = parent_observation.start_observation(
-                    name=self.get_langchain_run_name(serialized, **kwargs),
-                    as_type=observation_type,
-                    input=tool_input,
-                    metadata=meta,
-                    level="DEBUG" if tags and LANGSMITH_TAG_HIDDEN in tags else None,
-                )
+                with self._use_root_propagated_attributes(run_id):
+                    span = parent_observation.start_observation(
+                        name=self.get_langchain_run_name(serialized, **kwargs),
+                        as_type=observation_type,
+                        input=tool_input,
+                        metadata=meta,
+                        level="DEBUG"
+                        if tags and LANGSMITH_TAG_HIDDEN in tags
+                        else None,
+                    )
 
             self._attach_observation(run_id, span)
 
@@ -1029,16 +1043,17 @@ def on_retriever_start(
                     ),
                 )
             else:
-                span = parent_observation.start_observation(
-                    name=span_name,
-                    as_type=observation_type,
-                    metadata=span_metadata,
-                    input=query,
-                    level=cast(
-                        Optional[Literal["DEBUG", "DEFAULT", "WARNING", "ERROR"]],
-                        span_level,
-                    ),
-                )
+                with self._use_root_propagated_attributes(run_id):
+                    span = parent_observation.start_observation(
+                        name=span_name,
+                        as_type=observation_type,
+                        metadata=span_metadata,
+                        input=query,
+                        level=cast(
+                            Optional[Literal["DEBUG", "DEFAULT", "WARNING", "ERROR"]],
+                            span_level,
+                        ),
+                    )
 
             self._attach_observation(run_id, span)
 
@@ -1219,9 +1234,10 @@ def __on_llm_action(
                     **content,
                 )  # type: ignore
             else:
-                generation = parent_observation.start_observation(
-                    as_type="generation", **content
-                )  # type: ignore
+                with self._use_root_propagated_attributes(run_id):
+                    generation = parent_observation.start_observation(
+                        as_type="generation", **content
+                    )  # type: ignore
             self._attach_observation(run_id, generation)
 
             self.last_trace_id = self._runs[run_id].trace_id
@@ -1389,6 +1405,59 @@ def _reset(self, root_run_id: UUID) -> None:
         for run_id in root_run_state.run_ids:
             self._run_states.pop(run_id, None)
 
+    @contextmanager
+    def _use_root_propagated_attributes(self, run_id: UUID) -> Iterator[None]:
+        # Async LangChain runs sync handlers in a copied context per callback, so
+        # attributes the root run propagated are missing when children start.
+        root_run_state = self._get_root_run_state(run_id)
+        root_context = (
+            root_run_state.propagation_context if root_run_state is not None else None
+        )
+
+        if root_context is None:
+            yield
+            return
+
+        current_context = context.get_current()
+        merged_context = current_context
+
+        for key in propagated_keys:
+            context_key = _get_propagated_context_key(key)
+            root_value = context.get_value(context_key, root_context)
+
+            if root_value is None:
+                continue
+
+            current_value = context.get_value(context_key, current_context)
+            merged_value: Any
+
+            if isinstance(root_value, dict) and isinstance(current_value, dict):
+                merged_value = root_value | current_value
+            elif isinstance(root_value, list) and isinstance(current_value, list):
+                merged_value = root_value + [
+                    item for item in current_value if item not in root_value
+                ]
+            elif current_value is None:
+                merged_value = root_value
+            else:
+                continue
+
+            if merged_value != current_value:
+                merged_context = context.set_value(
+                    context_key, merged_value, merged_context
+                )
+
+        if merged_context is current_context:
+            yield
+            return
+
+        token = context.attach(merged_context)
+
+        try:
+            yield
+        finally:
+            _detach_context_token_safely(token)
+
     def _exit_propagation_context(self, run_id: UUID) -> None:
         root_run_state = self._get_root_run_state(run_id)
 

tests/unit/test_langchain.py

@@ -247,6 +247,96 @@ def test_root_chain_metadata_propagates_trace_name(
     assert len(find_spans("ChatOpenAI")) == 1
 
 
+@pytest.mark.asyncio
+async def test_root_chain_metadata_propagates_session_id_to_children_on_ainvoke(
+    langfuse_memory_client, get_span
+):
+    response = ChatResult(
+        generations=[
+            ChatGeneration(
+                message=AIMessage(content="knock knock"),
+                text="knock knock",
+            )
+        ],
+        llm_output={"model_name": "gpt-4o-mini"},
+    )
+
+    with patch.object(ChatOpenAI, "_agenerate", return_value=response):
+        handler = CallbackHandler()
+        prompt = ChatPromptTemplate.from_template("tell me a joke about {topic}")
+        chain = prompt | ChatOpenAI(api_key="test", temperature=0) | StrOutputParser()
+
+        result = await chain.ainvoke(
+            {"topic": "otters"},
+            config={
+                "callbacks": [handler],
+                "metadata": {
+                    "langfuse_session_id": "async-session",
+                    "langfuse_user_id": "async-user",
+                },
+            },
+        )
+
+    assert result == "knock knock"
+
+    langfuse_memory_client.flush()
+
+    for span_name in (
+        "RunnableSequence",
+        "ChatPromptTemplate",
+        "ChatOpenAI",
+        "StrOutputParser",
+    ):
+        span = get_span(span_name)
+        assert (
+            span.attributes[LangfuseOtelSpanAttributes.TRACE_SESSION_ID]
+            == "async-session"
+        ), span_name
+        assert (
+            span.attributes[LangfuseOtelSpanAttributes.TRACE_USER_ID] == "async-user"
+        ), span_name
+
+
+def test_child_run_keeps_attributes_propagated_in_its_own_context(
+    langfuse_memory_client, get_span
+):
+    handler = CallbackHandler()
+    root_run_id = uuid4()
+    child_run_id = uuid4()
+
+    handler.on_chain_start(
+        {"id": ["RunnableSequence"]},
+        {"topic": "otters"},
+        run_id=root_run_id,
+        metadata={"langfuse_session_id": "root-session"},
+    )
+
+    def start_child_with_own_session():
+        with callback_handler_module.propagate_attributes(session_id="child-session"):
+            handler.on_chain_start(
+                {"id": ["ChildChain"]},
+                {"topic": "otters"},
+                run_id=child_run_id,
+                parent_run_id=root_run_id,
+            )
+
+    copy_context().run(start_child_with_own_session)
+    copy_context().run(
+        handler.on_chain_end,
+        {"output": "done"},
+        run_id=child_run_id,
+        parent_run_id=root_run_id,
+    )
+    handler.on_chain_end({"output": "done"}, run_id=root_run_id)
+
+    langfuse_memory_client.flush()
+
+    assert (
+        get_span("ChildChain").attributes[LangfuseOtelSpanAttributes.TRACE_SESSION_ID]
+        == "child-session"
+    )
+
+
 def test_root_chain_exports_when_end_runs_in_copied_context(
     langfuse_memory_client, get_span
 ):