From b5cd3cbccbb9c11b655bb6b34c7a7437f3dc52ce Mon Sep 17 00:00:00 2001 From: iliescucristian Date: Thu, 21 May 2026 16:07:59 +0300 Subject: [PATCH 1/3] fix: inject trace context headers per-request instead of at client construction Move W3C trace context header injection from `build_uipath_headers()` (called once at client init) to per-request hooks in each transport, so headers reflect the active OTEL span at call time. - Legacy OpenAI/Bedrock/Vertex: inject in transport `handle_request` - New clients: inject via `_TraceContextHeadersCallback` + dynamic headers ContextVar - Remove trace context from shared `build_uipath_headers()` Co-Authored-By: Claude Opus 4.6 (1M context) --- src/uipath_langchain/chat/_legacy/bedrock.py | 4 + .../chat/_legacy/http_client/headers.py | 3 - src/uipath_langchain/chat/_legacy/openai.py | 11 +++ src/uipath_langchain/chat/_legacy/vertex.py | 11 +++ .../chat/chat_model_factory.py | 51 ++++++++++++- tests/chat/test_trace_context_callback.py | 74 +++++++++++++++++++ 6 files changed, 150 insertions(+), 4 deletions(-) create mode 100644 tests/chat/test_trace_context_callback.py diff --git a/src/uipath_langchain/chat/_legacy/bedrock.py b/src/uipath_langchain/chat/_legacy/bedrock.py index eb8ce9889..3a3fdb107 100644 --- a/src/uipath_langchain/chat/_legacy/bedrock.py +++ b/src/uipath_langchain/chat/_legacy/bedrock.py @@ -7,6 +7,7 @@ from langchain_core.messages import BaseMessage from langchain_core.outputs import ChatGenerationChunk, ChatResult from tenacity import AsyncRetrying, Retrying +from uipath.platform.chat.llm_trace_context import build_trace_context_headers from uipath.platform.common import ( EndpointManager, get_ca_bundle_path, @@ -181,6 +182,9 @@ def _modify_request(self, request, **kwargs): ) headers["X-UiPath-LlmGateway-ApiFlavor"] = self.api_flavor headers["X-UiPath-Streaming-Enabled"] = streaming + headers.update( + build_trace_context_headers(extra_baggage=["source=agents"]) + ) request.headers.update(headers) diff --git a/src/uipath_langchain/chat/_legacy/http_client/headers.py b/src/uipath_langchain/chat/_legacy/http_client/headers.py index 64a110729..bd642b8c7 100644 --- a/src/uipath_langchain/chat/_legacy/http_client/headers.py +++ b/src/uipath_langchain/chat/_legacy/http_client/headers.py @@ -3,7 +3,6 @@ import os from urllib.parse import quote -from uipath.platform.chat.llm_trace_context import build_trace_context_headers from uipath.platform.common._config import UiPathConfig from uipath.platform.common.constants import ( ENV_FOLDER_KEY, @@ -65,6 +64,4 @@ def build_uipath_headers( if organization_id := os.getenv(ENV_ORGANIZATION_ID): headers[HEADER_INTERNAL_ACCOUNT_ID] = organization_id - headers.update(build_trace_context_headers(extra_baggage=["source=agents"])) - return headers diff --git a/src/uipath_langchain/chat/_legacy/openai.py b/src/uipath_langchain/chat/_legacy/openai.py index 8a6da57b5..93cf89b56 100644 --- a/src/uipath_langchain/chat/_legacy/openai.py +++ b/src/uipath_langchain/chat/_legacy/openai.py @@ -5,6 +5,7 @@ import httpx from langchain_openai import AzureChatOpenAI from pydantic import PrivateAttr +from uipath.platform.chat.llm_trace_context import build_trace_context_headers from uipath.platform.common import ( EndpointManager, get_httpx_client_kwargs, @@ -53,6 +54,14 @@ def _inject_license_ref_id(request: httpx.Request) -> None: request.headers["X-UiPath-License-RefId"] = license_ref_id +def _inject_trace_context_headers(request: httpx.Request) -> None: + """Inject trace context headers per-request from the active OTEL span.""" + for key, value in build_trace_context_headers( + extra_baggage=["source=agents"] + ).items(): + request.headers[key] = value + + class UiPathURLRewriteTransport(httpx.AsyncHTTPTransport): def __init__(self, verify: bool = True, **kwargs): super().__init__(verify=verify, **kwargs) @@ -62,6 +71,7 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response: if new_url: request.url = new_url _inject_license_ref_id(request) + _inject_trace_context_headers(request) return await super().handle_async_request(request) @@ -75,6 +85,7 @@ def handle_request(self, request: httpx.Request) -> httpx.Response: if new_url: request.url = new_url _inject_license_ref_id(request) + _inject_trace_context_headers(request) return super().handle_request(request) diff --git a/src/uipath_langchain/chat/_legacy/vertex.py b/src/uipath_langchain/chat/_legacy/vertex.py index a8a577aff..bef142b4b 100644 --- a/src/uipath_langchain/chat/_legacy/vertex.py +++ b/src/uipath_langchain/chat/_legacy/vertex.py @@ -13,6 +13,7 @@ from tenacity import AsyncRetrying, Retrying from uipath._utils import resource_override from uipath._utils._ssl_context import get_httpx_client_kwargs +from uipath.platform.chat.llm_trace_context import build_trace_context_headers from uipath.platform.common import EndpointManager from .http_client import build_uipath_headers, resolve_gateway_url @@ -96,6 +97,11 @@ def handle_request(self, request: httpx.Request) -> httpx.Response: request.headers["host"] = new_url.host request.url = new_url + for key, value in build_trace_context_headers( + extra_baggage=["source=agents"] + ).items(): + request.headers[key] = value + response = super().handle_request(request) if self.header_capture: self.header_capture.set(dict(response.headers)) @@ -129,6 +135,11 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response: request.headers["host"] = new_url.host request.url = new_url + for key, value in build_trace_context_headers( + extra_baggage=["source=agents"] + ).items(): + request.headers[key] = value + response = await super().handle_async_request(request) if self.header_capture: self.header_capture.set(dict(response.headers)) diff --git a/src/uipath_langchain/chat/chat_model_factory.py b/src/uipath_langchain/chat/chat_model_factory.py index 30db14709..c22e16f76 100644 --- a/src/uipath_langchain/chat/chat_model_factory.py +++ b/src/uipath_langchain/chat/chat_model_factory.py @@ -12,8 +12,13 @@ from typing import Any, Final -from langchain_core.callbacks import Callbacks +from langchain_core.callbacks import BaseCallbackHandler, Callbacks from langchain_core.language_models import BaseChatModel +from uipath.llm_client.utils.headers import ( + get_dynamic_request_headers, + set_dynamic_request_headers, +) +from uipath.platform.chat.llm_trace_context import build_trace_context_headers from uipath_langchain_client.base_client import UiPathBaseChatModel from uipath_langchain_client.factory import get_chat_model as get_chat_model_factory from uipath_langchain_client.settings import ( @@ -23,6 +28,33 @@ VendorType, ) + +class _TraceContextHeadersCallback(BaseCallbackHandler): + """Inject W3C-style trace context headers into each LLM gateway request. + + Merges into the existing dynamic-headers ContextVar so that headers + set by earlier callbacks (e.g. ``LicenseRefIdHeadersCallback``) are + preserved instead of overwritten. + """ + + run_inline: bool = True + + def _merge_headers(self) -> None: + existing = get_dynamic_request_headers() + existing.update(build_trace_context_headers(extra_baggage=["source=agents"])) + set_dynamic_request_headers(existing) + + def on_chat_model_start( + self, serialized: dict[str, Any], messages: list[list[Any]], **kwargs: Any + ) -> None: + self._merge_headers() + + def on_llm_start( + self, serialized: dict[str, Any], prompts: list[str], **kwargs: Any + ) -> None: + self._merge_headers() + + _UNSET: Final[Any] = object() DEFAULT_TIMEOUT_SECONDS: Final[float] = 895.0 DEFAULT_MAX_TOKENS: Final[int] = 1000 @@ -84,6 +116,12 @@ def get_chat_model( Returns: A configured ``BaseChatModel`` instance. """ + # Always inject trace context headers per-request via a dynamic-headers + # callback. For the new path the UiPathHttpxClient reads the ContextVar + # set by the callback; for the legacy path the callback is a no-op but + # keeps the wiring consistent. + callbacks = _ensure_trace_context_callback(callbacks) + if not use_new_llm_clients: return _legacy_chat_model( model, @@ -120,6 +158,17 @@ def get_chat_model( ) +def _ensure_trace_context_callback(callbacks: Callbacks) -> list[BaseCallbackHandler]: + """Append a ``_TraceContextHeadersCallback`` if one is not already present.""" + if callbacks is _UNSET or callbacks is None: + cb_list: list[BaseCallbackHandler] = [] + else: + cb_list = list(callbacks) # type: ignore[arg-type] + if not any(isinstance(cb, _TraceContextHeadersCallback) for cb in cb_list): + cb_list.append(_TraceContextHeadersCallback()) + return cb_list + + def _legacy_chat_model( model: str, *, diff --git a/tests/chat/test_trace_context_callback.py b/tests/chat/test_trace_context_callback.py new file mode 100644 index 000000000..8ceff9e7e --- /dev/null +++ b/tests/chat/test_trace_context_callback.py @@ -0,0 +1,74 @@ +"""Tests for _TraceContextHeadersCallback in chat_model_factory.""" + +from unittest.mock import MagicMock, patch + +import pytest +from uipath.core.feature_flags import FeatureFlags + +from uipath_langchain.chat.chat_model_factory import ( + _TraceContextHeadersCallback, + _ensure_trace_context_callback, +) + +_UNSET = object() + + +class TestTraceContextHeadersCallback: + """The callback produces the expected headers from an active OTEL span.""" + + def setup_method(self) -> None: + FeatureFlags.configure_flags({"EnableTraceContextHeaders": True}) + + def teardown_method(self) -> None: + FeatureFlags.reset_flags() + + def test_returns_headers_with_active_span(self) -> None: + mock_span = MagicMock() + ctx = MagicMock() + ctx.trace_id = 0xAABBCCDDEEFF00112233445566778899 + ctx.span_id = 0x1122334455667788 + mock_span.get_span_context.return_value = ctx + + cb = _TraceContextHeadersCallback() + with patch( + "uipath.platform.chat.llm_trace_context.UiPathSpanUtils" + ".get_external_current_span", + return_value=mock_span, + ): + headers = cb.get_headers() + + assert "x-uipath-traceparent-id" in headers + assert headers["x-uipath-traceparent-id"].startswith("00-") + assert "x-uipath-tracebaggage" in headers + assert "source=agents" in headers["x-uipath-tracebaggage"] + + def test_returns_empty_when_flag_disabled(self) -> None: + FeatureFlags.configure_flags({"EnableTraceContextHeaders": False}) + cb = _TraceContextHeadersCallback() + assert cb.get_headers() == {} + + +class TestEnsureTraceContextCallback: + """_ensure_trace_context_callback always appends the callback.""" + + def test_adds_callback_to_empty_list(self) -> None: + result = _ensure_trace_context_callback(None) + assert any(isinstance(cb, _TraceContextHeadersCallback) for cb in result) + + def test_adds_callback_when_unset(self) -> None: + from uipath_langchain.chat.chat_model_factory import _UNSET + + result = _ensure_trace_context_callback(_UNSET) + assert any(isinstance(cb, _TraceContextHeadersCallback) for cb in result) + + def test_does_not_duplicate(self) -> None: + existing = [_TraceContextHeadersCallback()] + result = _ensure_trace_context_callback(existing) + count = sum(1 for cb in result if isinstance(cb, _TraceContextHeadersCallback)) + assert count == 1 + + def test_preserves_existing_callbacks(self) -> None: + sentinel = MagicMock() + result = _ensure_trace_context_callback([sentinel]) + assert sentinel in result + assert any(isinstance(cb, _TraceContextHeadersCallback) for cb in result) From bac327c7db5eddebd6a75c1fc5123f44aa8ded09 Mon Sep 17 00:00:00 2001 From: iliescucristian Date: Fri, 22 May 2026 10:10:18 +0300 Subject: [PATCH 2/3] fix: update trace context callback tests to use _merge_headers and ContextVar Tests were calling a non-existent get_headers() method. Updated to call _merge_headers() and verify results via the ContextVar, matching the actual callback implementation. Also fixed mypy variance error. Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/chat/test_trace_context_callback.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/chat/test_trace_context_callback.py b/tests/chat/test_trace_context_callback.py index 8ceff9e7e..41422ad4f 100644 --- a/tests/chat/test_trace_context_callback.py +++ b/tests/chat/test_trace_context_callback.py @@ -2,25 +2,29 @@ from unittest.mock import MagicMock, patch -import pytest +from langchain_core.callbacks import BaseCallbackHandler from uipath.core.feature_flags import FeatureFlags +from uipath.llm_client.utils.headers import ( + get_dynamic_request_headers, + set_dynamic_request_headers, +) from uipath_langchain.chat.chat_model_factory import ( _TraceContextHeadersCallback, _ensure_trace_context_callback, ) -_UNSET = object() - class TestTraceContextHeadersCallback: """The callback produces the expected headers from an active OTEL span.""" def setup_method(self) -> None: FeatureFlags.configure_flags({"EnableTraceContextHeaders": True}) + set_dynamic_request_headers({}) def teardown_method(self) -> None: FeatureFlags.reset_flags() + set_dynamic_request_headers({}) def test_returns_headers_with_active_span(self) -> None: mock_span = MagicMock() @@ -31,12 +35,15 @@ def test_returns_headers_with_active_span(self) -> None: cb = _TraceContextHeadersCallback() with patch( - "uipath.platform.chat.llm_trace_context.UiPathSpanUtils" + "uipath.core.tracing.span_utils.UiPathSpanUtils" ".get_external_current_span", return_value=mock_span, + ), patch.dict( + "os.environ", {"UIPATH_TRACE_ID": "aabbccddeeff00112233445566778899"} ): - headers = cb.get_headers() + cb._merge_headers() + headers = get_dynamic_request_headers() assert "x-uipath-traceparent-id" in headers assert headers["x-uipath-traceparent-id"].startswith("00-") assert "x-uipath-tracebaggage" in headers @@ -45,7 +52,8 @@ def test_returns_headers_with_active_span(self) -> None: def test_returns_empty_when_flag_disabled(self) -> None: FeatureFlags.configure_flags({"EnableTraceContextHeaders": False}) cb = _TraceContextHeadersCallback() - assert cb.get_headers() == {} + cb._merge_headers() + assert get_dynamic_request_headers() == {} class TestEnsureTraceContextCallback: @@ -62,7 +70,7 @@ def test_adds_callback_when_unset(self) -> None: assert any(isinstance(cb, _TraceContextHeadersCallback) for cb in result) def test_does_not_duplicate(self) -> None: - existing = [_TraceContextHeadersCallback()] + existing: list[BaseCallbackHandler] = [_TraceContextHeadersCallback()] result = _ensure_trace_context_callback(existing) count = sum(1 for cb in result if isinstance(cb, _TraceContextHeadersCallback)) assert count == 1 From e00e9685de938703ac9243aa815ef52f49bf485d Mon Sep 17 00:00:00 2001 From: iliescucristian Date: Mon, 25 May 2026 13:55:33 +0300 Subject: [PATCH 3/3] fix: add tests --- pyproject.toml | 4 +- src/uipath_langchain/chat/_legacy/bedrock.py | 4 +- tests/chat/test_trace_context_callback.py | 260 ++++++++++++++++++- uv.lock | 10 +- 4 files changed, 260 insertions(+), 18 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2126afc09..e5fa00c20 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,13 @@ [project] name = "uipath-langchain" -version = "0.11.6" +version = "0.11.7" description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" dependencies = [ "uipath>=2.10.70, <2.11.0", "uipath-core>=0.5.15, <0.6.0", - "uipath-platform>=0.1.45, <0.2.0", + "uipath-platform>=0.1.58, <0.2.0", "uipath-runtime>=0.10.0, <0.11.0", "langgraph>=1.1.8, <2.0.0", "langchain-core>=1.2.11, <2.0.0", diff --git a/src/uipath_langchain/chat/_legacy/bedrock.py b/src/uipath_langchain/chat/_legacy/bedrock.py index 3a3fdb107..dc99e3757 100644 --- a/src/uipath_langchain/chat/_legacy/bedrock.py +++ b/src/uipath_langchain/chat/_legacy/bedrock.py @@ -182,9 +182,7 @@ def _modify_request(self, request, **kwargs): ) headers["X-UiPath-LlmGateway-ApiFlavor"] = self.api_flavor headers["X-UiPath-Streaming-Enabled"] = streaming - headers.update( - build_trace_context_headers(extra_baggage=["source=agents"]) - ) + headers.update(build_trace_context_headers(extra_baggage=["source=agents"])) request.headers.update(headers) diff --git a/tests/chat/test_trace_context_callback.py b/tests/chat/test_trace_context_callback.py index 41422ad4f..5b82e0f2e 100644 --- a/tests/chat/test_trace_context_callback.py +++ b/tests/chat/test_trace_context_callback.py @@ -1,7 +1,10 @@ -"""Tests for _TraceContextHeadersCallback in chat_model_factory.""" +"""Tests for _TraceContextHeadersCallback and per-request trace context injection.""" +import os from unittest.mock import MagicMock, patch +import httpx +import pytest from langchain_core.callbacks import BaseCallbackHandler from uipath.core.feature_flags import FeatureFlags from uipath.llm_client.utils.headers import ( @@ -9,11 +12,43 @@ set_dynamic_request_headers, ) +from uipath_langchain.chat._legacy.http_client import build_uipath_headers +from uipath_langchain.chat._legacy.openai import _inject_trace_context_headers from uipath_langchain.chat.chat_model_factory import ( - _TraceContextHeadersCallback, _ensure_trace_context_callback, + _TraceContextHeadersCallback, ) +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_MOCK_TRACE_HEADERS = { + "x-uipath-traceparent-id": "00-aabbccddeeff00112233445566778899-1122334455667788-01", + "x-uipath-tracebaggage": "source=agents", +} + + +def _patch_build_trace_context_headers(*extra_targets: str): + """Return a combined context-manager that patches build_trace_context_headers + at both the canonical location and any re-import locations so tests don't + need a live OTEL span.""" + from contextlib import ExitStack, contextmanager + + targets = [ + "uipath.platform.chat.llm_trace_context.build_trace_context_headers", + *extra_targets, + ] + + @contextmanager + def _cm(): + with ExitStack() as stack: + for t in targets: + stack.enter_context(patch(t, return_value=dict(_MOCK_TRACE_HEADERS))) + yield + + return _cm() + class TestTraceContextHeadersCallback: """The callback produces the expected headers from an active OTEL span.""" @@ -34,12 +69,15 @@ def test_returns_headers_with_active_span(self) -> None: mock_span.get_span_context.return_value = ctx cb = _TraceContextHeadersCallback() - with patch( - "uipath.core.tracing.span_utils.UiPathSpanUtils" - ".get_external_current_span", - return_value=mock_span, - ), patch.dict( - "os.environ", {"UIPATH_TRACE_ID": "aabbccddeeff00112233445566778899"} + with ( + patch( + "uipath.core.tracing.span_utils.UiPathSpanUtils" + ".get_external_current_span", + return_value=mock_span, + ), + patch.dict( + "os.environ", {"UIPATH_TRACE_ID": "aabbccddeeff00112233445566778899"} + ), ): cb._merge_headers() @@ -55,6 +93,24 @@ def test_returns_empty_when_flag_disabled(self) -> None: cb._merge_headers() assert get_dynamic_request_headers() == {} + def test_on_chat_model_start_delegates_to_merge_headers(self) -> None: + cb = _TraceContextHeadersCallback() + with _patch_build_trace_context_headers( + "uipath_langchain.chat.chat_model_factory.build_trace_context_headers", + ): + cb.on_chat_model_start(serialized={}, messages=[[]]) + headers = get_dynamic_request_headers() + assert "x-uipath-traceparent-id" in headers + + def test_on_llm_start_delegates_to_merge_headers(self) -> None: + cb = _TraceContextHeadersCallback() + with _patch_build_trace_context_headers( + "uipath_langchain.chat.chat_model_factory.build_trace_context_headers", + ): + cb.on_llm_start(serialized={}, prompts=["hello"]) + headers = get_dynamic_request_headers() + assert "x-uipath-traceparent-id" in headers + class TestEnsureTraceContextCallback: """_ensure_trace_context_callback always appends the callback.""" @@ -80,3 +136,191 @@ def test_preserves_existing_callbacks(self) -> None: result = _ensure_trace_context_callback([sentinel]) assert sentinel in result assert any(isinstance(cb, _TraceContextHeadersCallback) for cb in result) + + +# --------------------------------------------------------------------------- +# Legacy OpenAI transport — per-request trace context injection +# --------------------------------------------------------------------------- + + +class TestOpenAITransportTraceContextHeaders: + """_inject_trace_context_headers stamps headers on every httpx.Request.""" + + def test_inject_trace_context_headers_adds_headers(self) -> None: + request = httpx.Request("POST", "https://example.com/completions") + with _patch_build_trace_context_headers( + "uipath_langchain.chat._legacy.openai.build_trace_context_headers", + ): + _inject_trace_context_headers(request) + + for key, value in _MOCK_TRACE_HEADERS.items(): + assert request.headers[key] == value + + def test_sync_transport_calls_inject_trace_context(self) -> None: + from uipath_langchain.chat._legacy.openai import UiPathSyncURLRewriteTransport + + transport = UiPathSyncURLRewriteTransport() + request = httpx.Request("POST", "https://example.com/completions") + + with ( + _patch_build_trace_context_headers( + "uipath_langchain.chat._legacy.openai.build_trace_context_headers", + ), + patch.object( + httpx.HTTPTransport, + "handle_request", + return_value=httpx.Response(200), + ), + ): + transport.handle_request(request) + + for key, value in _MOCK_TRACE_HEADERS.items(): + assert request.headers[key] == value + + async def test_async_transport_calls_inject_trace_context(self) -> None: + from uipath_langchain.chat._legacy.openai import UiPathURLRewriteTransport + + transport = UiPathURLRewriteTransport() + request = httpx.Request("POST", "https://example.com/completions") + + with ( + _patch_build_trace_context_headers( + "uipath_langchain.chat._legacy.openai.build_trace_context_headers", + ), + patch.object( + httpx.AsyncHTTPTransport, + "handle_async_request", + return_value=httpx.Response(200), + ), + ): + await transport.handle_async_request(request) + + for key, value in _MOCK_TRACE_HEADERS.items(): + assert request.headers[key] == value + + def test_inject_trace_context_headers_noop_when_empty(self) -> None: + """When build_trace_context_headers returns {}, no headers are added.""" + request = httpx.Request("POST", "https://example.com/completions") + with patch( + "uipath_langchain.chat._legacy.openai.build_trace_context_headers", + return_value={}, + ): + _inject_trace_context_headers(request) + + assert "x-uipath-traceparent-id" not in request.headers + + +# --------------------------------------------------------------------------- +# Legacy Vertex transport — per-request trace context injection +# --------------------------------------------------------------------------- + + +class TestVertexTransportTraceContextHeaders: + """_UrlRewriteTransport and _AsyncUrlRewriteTransport inject trace headers.""" + + def test_sync_transport_injects_trace_headers(self) -> None: + pytest.importorskip("google.genai") + from uipath_langchain.chat._legacy.vertex import _UrlRewriteTransport + + transport = _UrlRewriteTransport( + gateway_url="https://gateway.example.com/completions" + ) + request = httpx.Request( + "POST", + "https://generativelanguage.googleapis.com/v1beta/models/gemini:generateContent", + ) + + with ( + _patch_build_trace_context_headers( + "uipath_langchain.chat._legacy.vertex.build_trace_context_headers", + ), + patch.object( + httpx.HTTPTransport, "handle_request", return_value=httpx.Response(200) + ), + ): + transport.handle_request(request) + + for key, value in _MOCK_TRACE_HEADERS.items(): + assert request.headers[key] == value + + async def test_async_transport_injects_trace_headers(self) -> None: + pytest.importorskip("google.genai") + from uipath_langchain.chat._legacy.vertex import _AsyncUrlRewriteTransport + + transport = _AsyncUrlRewriteTransport( + gateway_url="https://gateway.example.com/completions" + ) + request = httpx.Request( + "POST", + "https://generativelanguage.googleapis.com/v1beta/models/gemini:generateContent", + ) + + with ( + _patch_build_trace_context_headers( + "uipath_langchain.chat._legacy.vertex.build_trace_context_headers", + ), + patch.object( + httpx.AsyncHTTPTransport, + "handle_async_request", + return_value=httpx.Response(200), + ), + ): + await transport.handle_async_request(request) + + for key, value in _MOCK_TRACE_HEADERS.items(): + assert request.headers[key] == value + + +# --------------------------------------------------------------------------- +# Legacy Bedrock transport — per-request trace context injection +# --------------------------------------------------------------------------- + + +class TestBedrockTransportTraceContextHeaders: + """_modify_request in AwsBedrockCompletionsPassthroughClient injects trace headers.""" + + def test_modify_request_injects_trace_headers(self) -> None: + pytest.importorskip("botocore") + from uipath_langchain.chat._legacy.bedrock import ( + AwsBedrockCompletionsPassthroughClient, + ) + + passthrough = AwsBedrockCompletionsPassthroughClient( + model="anthropic.claude-haiku-4-5-20251001", + token="test-token", + api_flavor="converse", + ) + + class _Req: + url = "https://bedrock.example/foo/converse" + headers: dict[str, str] = {} + + request = _Req() + with ( + _patch_build_trace_context_headers( + "uipath_langchain.chat._legacy.bedrock.build_trace_context_headers", + ), + patch.object( + passthrough, "_resolve_url", return_value=("https://gateway/x", False) + ), + ): + passthrough._modify_request(request) + + for key, value in _MOCK_TRACE_HEADERS.items(): + assert request.headers[key] == value + + +# --------------------------------------------------------------------------- +# build_uipath_headers no longer includes trace context headers +# --------------------------------------------------------------------------- + + +class TestBuildUiPathHeadersNoTraceContext: + """After the refactor, build_uipath_headers must NOT contain trace headers.""" + + def test_no_trace_context_headers_in_build_uipath_headers(self) -> None: + with patch.dict(os.environ, {}, clear=True): + headers = build_uipath_headers() + + assert "x-uipath-traceparent-id" not in headers + assert "x-uipath-tracebaggage" not in headers diff --git a/uv.lock b/uv.lock index e99046256..66ae9f583 100644 --- a/uv.lock +++ b/uv.lock @@ -4388,7 +4388,7 @@ wheels = [ [[package]] name = "uipath-langchain" -version = "0.11.6" +version = "0.11.7" source = { editable = "." } dependencies = [ { name = "a2a-sdk" }, @@ -4472,7 +4472,7 @@ requires-dist = [ { name = "uipath-langchain-client", extras = ["google"], marker = "extra == 'vertex'", specifier = ">=1.12.2,<1.13.0" }, { name = "uipath-langchain-client", extras = ["openai"], specifier = ">=1.12.2,<1.13.0" }, { name = "uipath-langchain-client", extras = ["vertexai"], marker = "extra == 'vertex'", specifier = ">=1.12.2,<1.13.0" }, - { name = "uipath-platform", specifier = ">=0.1.45,<0.2.0" }, + { name = "uipath-platform", specifier = ">=0.1.58,<0.2.0" }, { name = "uipath-runtime", specifier = ">=0.10.0,<0.11.0" }, ] provides-extras = ["anthropic", "vertex", "bedrock", "fireworks", "all"] @@ -4556,7 +4556,7 @@ wheels = [ [[package]] name = "uipath-platform" -version = "0.1.54" +version = "0.1.58" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "httpx" }, @@ -4566,9 +4566,9 @@ dependencies = [ { name = "truststore" }, { name = "uipath-core" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/07/3f/d642609355d60618a71626e67475508aacd5c29ea295e0d23423619c1658/uipath_platform-0.1.54.tar.gz", hash = "sha256:5c49c4303a8a40dcc7df4e0ab14c8d64c307ce3555e5ab71545b3de7b59ab2fd", size = 339248, upload-time = "2026-05-19T16:57:37.753Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0b/74/ef0e57fb583c4b4c485da2957a5a5bcadc5d525009c79ab97057838a7236/uipath_platform-0.1.58.tar.gz", hash = "sha256:5eb76dd94e639965c7c64175b208d12f81bcedb349a268593670afda0d5c8c0a", size = 365298, upload-time = "2026-05-25T10:51:07.277Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1f/03/2e32d0efabd3dded465ff27042f2af39cfd7ecab72b89cd9f1fddc04eb3e/uipath_platform-0.1.54-py3-none-any.whl", hash = "sha256:64e4c9f8848b4e0124747cb9049875f4ad534d01ac7672460d3ed1ccc305556e", size = 221582, upload-time = "2026-05-19T16:57:35.171Z" }, + { url = "https://files.pythonhosted.org/packages/49/24/71ef531865ae345e491d60d8105ef4505f82ce57ae5fc88bdeb1c03ed493/uipath_platform-0.1.58-py3-none-any.whl", hash = "sha256:192e9cd1853d183e4ca12f8c651a9a6d9e065b84c430d6e25a93cf63fee55fd9", size = 242630, upload-time = "2026-05-25T10:51:05.38Z" }, ] [[package]]