From 9c582abcd887bb7db885a350dfdb75e7d0290dc7 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 28 Oct 2026 09:54:40 +0000 Subject: [PATCH] fix(langchain): propagate root trace attributes to child runs under ainvoke Co-authored-by: hassieb.pakzad --- langfuse/langchain/CallbackHandler.py | 131 ++++++++++++++++++++------ tests/unit/test_langchain.py | 90 ++++++++++++++++++ 2 files changed, 190 insertions(+), 31 deletions(-) diff --git a/langfuse/langchain/CallbackHandler.py b/langfuse/langchain/CallbackHandler.py index 0db58cd6e..83ba97bc0 100644 --- a/langfuse/langchain/CallbackHandler.py +++ b/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) diff --git a/tests/unit/test_langchain.py b/tests/unit/test_langchain.py index 6b6143fcc..eab8d34ff 100644 --- a/tests/unit/test_langchain.py +++ b/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 ):