fix: header callbacks merge instead of overwriting

Author: andreiancuta-uipath

packages/uipath_langchain_client/CHANGELOG.md

@@ -2,6 +2,11 @@
 
 All notable changes to `uipath_langchain_client` will be documented in this file.
 
+## [1.11.1] - 2026-05-13
+
+### Fixed
+- `UiPathDynamicHeadersCallback` now merges `get_headers()` into the dynamic-headers ContextVar instead of replacing it wholesale. This prevents two stacked callbacks from overwriting each other.
+
 ## [1.11.0] - 2026-05-08
 
 ### Changed

packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py

@@ -1,3 +1,3 @@
 __title__ = "UiPath LangChain Client"
 __description__ = "A Python client for interacting with UiPath's LLM services via LangChain."
-__version__ = "1.11.0"
+__version__ = "1.11.1"

packages/uipath_langchain_client/src/uipath_langchain_client/callbacks.py

@@ -5,7 +5,10 @@
 
 from langchain_core.callbacks import BaseCallbackHandler
 
-from uipath.llm_client.utils.headers import set_dynamic_request_headers
+from uipath.llm_client.utils.headers import (
+    get_dynamic_request_headers,
+    set_dynamic_request_headers,
+)
 
 
 class UiPathDynamicHeadersCallback(BaseCallbackHandler, ABC):
@@ -37,21 +40,26 @@ def get_headers(self) -> dict[str, str]:
         """Return headers to inject into the next LLM gateway request."""
         ...
 
+    def _merge_headers(self) -> None:
+        merged = get_dynamic_request_headers()
+        merged.update(self.get_headers())
+        set_dynamic_request_headers(merged)
+
     def on_chat_model_start(
         self,
         serialized: dict[str, Any],
         messages: list[list[Any]],
         **kwargs: Any,
     ) -> None:
-        set_dynamic_request_headers(self.get_headers())
+        self._merge_headers()
 
     def on_llm_start(
         self,
         serialized: dict[str, Any],
         prompts: list[str],
         **kwargs: Any,
     ) -> None:
-        set_dynamic_request_headers(self.get_headers())
+        self._merge_headers()
 
     def on_llm_end(self, response: Any, **kwargs: Any) -> None:
         set_dynamic_request_headers({})

tests/langchain/features/test_dynamic_headers.py

@@ -235,10 +235,42 @@ def test_on_chat_model_start_sets_headers(self, tracer):
             cb.on_chat_model_start({}, [[]])
         assert "x-trace-id" in get_dynamic_request_headers()
 
-    def test_on_chat_model_start_no_span_sets_empty(self):
-        """When there is no active span, on_chat_model_start clears the ContextVar."""
-        set_dynamic_request_headers({"x-stale": "value"})
+    def test_on_chat_model_start_no_span_injects_nothing(self):
+        """When there is no active span, on_chat_model_start adds no headers of its own."""
         OtelHeadersCallback().on_chat_model_start({}, [[]])
+        assert "x-trace-id" not in get_dynamic_request_headers()
+        assert "x-span-id" not in get_dynamic_request_headers()
+
+    def test_on_chat_model_start_preserves_other_callbacks_headers(self):
+        """Merge semantics: a callback returning {} leaves existing headers intact."""
+        set_dynamic_request_headers({"x-other-callback": "value"})
+        OtelHeadersCallback().on_chat_model_start({}, [[]])
+        assert get_dynamic_request_headers() == {"x-other-callback": "value"}
+
+    def test_two_callbacks_compose_without_clobbering(self, tracer):
+        """Two callbacks in a row produce the union of their headers."""
+
+        class StaticHeadersCallback(UiPathDynamicHeadersCallback):
+            def __init__(self, headers: dict[str, str]):
+                super().__init__()
+                self._headers = headers
+
+            def get_headers(self) -> dict[str, str]:
+                return self._headers
+
+        first = StaticHeadersCallback({"x-previous-header": "abc"})
+        second = OtelHeadersCallback()
+        with active_span(tracer, "llm-call"):
+            first.on_chat_model_start({}, [[]])
+            second.on_chat_model_start({}, [[]])
+            headers = get_dynamic_request_headers()
+        assert headers.get("x-previous-header") == "abc"
+        assert "x-trace-id" in headers
+
+    def test_on_llm_end_clears_all_headers(self):
+        """on_llm_end resets the ContextVar wholesale for the next call."""
+        set_dynamic_request_headers({"x-header": "1", "x-other-header": "2"})
+        OtelHeadersCallback().on_llm_end(None)
         assert get_dynamic_request_headers() == {}