From 47cb6fe67e7ed4cfec5dd90992660bdc2a3253ec Mon Sep 17 00:00:00 2001 From: Andrei Ancuta Date: Wed, 13 May 2026 12:11:17 +0300 Subject: [PATCH] fix: header callbacks merge instead of overwriting --- packages/uipath_langchain_client/CHANGELOG.md | 5 +++ .../uipath_langchain_client/__version__.py | 2 +- .../src/uipath_langchain_client/callbacks.py | 14 +++++-- .../features/test_dynamic_headers.py | 38 +++++++++++++++++-- 4 files changed, 52 insertions(+), 7 deletions(-) diff --git a/packages/uipath_langchain_client/CHANGELOG.md b/packages/uipath_langchain_client/CHANGELOG.md index a9e0cb3..9768bcd 100644 --- a/packages/uipath_langchain_client/CHANGELOG.md +++ b/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 diff --git a/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py b/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py index ba2c515..7d3ecfc 100644 --- a/packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py +++ b/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" diff --git a/packages/uipath_langchain_client/src/uipath_langchain_client/callbacks.py b/packages/uipath_langchain_client/src/uipath_langchain_client/callbacks.py index 431b49c..beb2b6f 100644 --- a/packages/uipath_langchain_client/src/uipath_langchain_client/callbacks.py +++ b/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,13 +40,18 @@ 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, @@ -51,7 +59,7 @@ def on_llm_start( 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({}) diff --git a/tests/langchain/features/test_dynamic_headers.py b/tests/langchain/features/test_dynamic_headers.py index 4472844..dfc6f46 100644 --- a/tests/langchain/features/test_dynamic_headers.py +++ b/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() == {}