fix(langchain): deregister langfuse prompts when a chain run errors

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)

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
 ):