fix: inject trace context headers per-request instead of at client construction (#870)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ctiliescuuipath

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",

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,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"]))
 
         request.headers.update(headers)
 

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

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)
 

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

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,
     *,