diff --git a/langfuse/langchain/CallbackHandler.py b/langfuse/langchain/CallbackHandler.py index 65ce1b2d8..0db58cd6e 100644 --- a/langfuse/langchain/CallbackHandler.py +++ b/langfuse/langchain/CallbackHandler.py @@ -852,6 +852,8 @@ def on_chain_error( except Exception as e: langfuse_logger.exception(e) finally: + self._deregister_langfuse_prompt(run_id) + if parent_run_id is None: self._exit_propagation_context(run_id) self._reset(run_id) diff --git a/tests/unit/test_langchain.py b/tests/unit/test_langchain.py index dbbb3d93b..6b6143fcc 100644 --- a/tests/unit/test_langchain.py +++ b/tests/unit/test_langchain.py @@ -401,6 +401,82 @@ def test_streamed_llm_error_cleans_only_its_completion_start_time_memo( ) +def test_chain_error_deregisters_prompt_registered_for_the_run( + langfuse_memory_client, +): + handler = CallbackHandler() + + root_run_id = uuid4() + prompt_chain_run_id = uuid4() + + handler.on_chain_start( + {"name": "RootChain"}, + {"input": "hi"}, + run_id=root_run_id, + ) + handler.on_chain_start( + {"name": "PromptChain"}, + {"input": "hi"}, + run_id=prompt_chain_run_id, + parent_run_id=root_run_id, + metadata={"langfuse_prompt": "sentinel-prompt"}, + ) + + assert handler._prompt_to_parent_run_map.get(root_run_id) == "sentinel-prompt" + + handler.on_chain_error( + RuntimeError("synthetic failure"), + run_id=root_run_id, + ) + + assert root_run_id not in handler._prompt_to_parent_run_map + + +def test_nested_chain_error_deregisters_propagated_prompt( + langfuse_memory_client, +): + handler = CallbackHandler() + + root_run_id = uuid4() + prompt_chain_run_id = uuid4() + nested_chain_run_id = uuid4() + + handler.on_chain_start( + {"name": "RootChain"}, + {"input": "hi"}, + run_id=root_run_id, + ) + handler.on_chain_start( + {"name": "PromptChain"}, + {"input": "hi"}, + run_id=prompt_chain_run_id, + parent_run_id=root_run_id, + metadata={"langfuse_prompt": "sentinel-prompt"}, + ) + handler.on_chain_start( + {"name": "NestedChain"}, + {"input": "hi"}, + run_id=nested_chain_run_id, + parent_run_id=root_run_id, + ) + + assert ( + handler._prompt_to_parent_run_map.get(nested_chain_run_id) == "sentinel-prompt" + ) + + handler.on_chain_error( + RuntimeError("synthetic failure"), + run_id=nested_chain_run_id, + parent_run_id=root_run_id, + ) + + assert nested_chain_run_id not in handler._prompt_to_parent_run_map + + # The parent-level registration survives so a following LLM sibling + # can still consume it, matching the success-path lifecycle. + assert handler._prompt_to_parent_run_map.get(root_run_id) == "sentinel-prompt" + + def test_control_flow_resume_uses_thread_keyed_explicit_resume_context( memory_exporter, langfuse_memory_client, monkeypatch ):