feat: propagate reference hierarchy context from LangGraph runtime

pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "uipath-langchain"
-version = "0.14.1"
+version = "0.14.2"
 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.13.3, <2.14.0",
+    "uipath>=2.13.5, <2.14.0",
     "uipath-core>=0.5.29, <0.6.0",
-    "uipath-platform>=0.2.2, <0.3.0",
+    "uipath-platform>=0.2.4, <0.3.0",
     "uipath-runtime>=0.12.1, <0.13.0",
     "uipath-llm-client>=1.16.2, <1.17.0",
     "langgraph>=1.1.8, <2.0.0",

src/uipath_langchain/runtime/runtime.py

@@ -1,3 +1,4 @@
+import contextvars
 import logging
 import os
 from collections.abc import Iterator
@@ -42,6 +43,39 @@
 
 from ._serialize import serialize_output
 
+# Guarded import: ReferenceContext was added to uipath.tracing in a later release.
+# Older installed packages lack it; the noop shim keeps the runtime loadable while
+# silently disabling span propagation until the package is updated.
+_noop_ref_ctx: contextvars.ContextVar[None] = contextvars.ContextVar(
+    "_uipath_langchain_noop_ref_ctx", default=None
+)
+
+
+class _NoopReferenceContextAccessor:
+    @staticmethod
+    def get() -> None:
+        return None
+
+    @staticmethod
+    def set(value: Any) -> "contextvars.Token[None]":
+        return _noop_ref_ctx.set(None)
+
+    @staticmethod
+    def reset(token: Any) -> None:
+        _noop_ref_ctx.reset(token)
+
+
+ReferenceContext: Any = None
+ReferenceContextAccessor: Any = _NoopReferenceContextAccessor
+
+try:
+    from uipath.tracing import (
+        ReferenceContext,
+        ReferenceContextAccessor,
+    )
+except ImportError:
+    pass
+
 logger = logging.getLogger(__name__)
 
 
@@ -75,12 +109,36 @@ def __init__(
         self.chat.client_side_tools = self._get_client_side_tools()
         self._middleware_node_names: set[str] = self._detect_middleware_nodes()
 
+    def _push_reference_context(self) -> "contextvars.Token[Any]":
+        """Append this runtime's own entry to the ambient ReferenceContext.
+
+        Reads any parent context already in the accessor (e.g. set by an
+        upstream middleware or the agents-python runtime), then appends a
+        ``langgraph`` entry for this runtime.  Returns the ContextVar token
+        so the caller can reset in a ``finally`` block.
+
+        Returns a no-op token when the installed uipath package predates
+        reference-context support.
+        """
+        if ReferenceContext is None:
+            return ReferenceContextAccessor.set(None)
+        agent_id = os.environ.get("UIPATH_AGENT_ID")
+        agent_version = os.environ.get("UIPATH_PROCESS_VERSION") or None
+        parent_ctx = ReferenceContextAccessor.get() or ReferenceContext.Empty
+        ref_ctx = (
+            parent_ctx.add("langgraph", agent_id, agent_version)
+            if agent_id
+            else parent_ctx
+        )
+        return ReferenceContextAccessor.set(ref_ctx)
+
     async def execute(
         self,
         input: dict[str, Any] | None = None,
         options: UiPathExecuteOptions | None = None,
     ) -> UiPathRuntimeResult:
         """Execute the graph with the provided input and configuration."""
+        ref_ctx_token = self._push_reference_context()
         try:
             graph_input = await self._get_graph_input(input, options)
             graph_config = self._get_graph_config()
@@ -99,6 +157,8 @@ async def execute(
 
         except Exception as e:
             raise self.create_runtime_error(e) from e
+        finally:
+            ReferenceContextAccessor.reset(ref_ctx_token)
 
     async def stream(
         self,
@@ -133,6 +193,7 @@ async def stream(
         Raises:
             LangGraphRuntimeError: If execution fails
         """
+        ref_ctx_token = self._push_reference_context()
         try:
             graph_input = await self._get_graph_input(input, options)
             graph_config = self._get_graph_config()
@@ -230,6 +291,8 @@ async def stream(
 
         except Exception as e:
             raise self.create_runtime_error(e) from e
+        finally:
+            ReferenceContextAccessor.reset(ref_ctx_token)
 
     async def get_schema(self) -> UiPathRuntimeSchema:
         """Get schema for this LangGraph runtime."""

tests/runtime/test_reference_context_wiring.py

@@ -0,0 +1,242 @@
+"""Tests for ReferenceContext wiring in UiPathLangGraphRuntime."""
+
+from typing import Any, TypedDict
+
+import pytest
+from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+from langgraph.graph import END, START, StateGraph
+
+from uipath_langchain.runtime.errors import LangGraphRuntimeError
+from uipath_langchain.runtime.runtime import UiPathLangGraphRuntime
+
+ReferenceContext: Any = None
+ReferenceContextAccessor: Any = None
+_reference_context_available = False
+
+try:
+    from uipath.tracing import (
+        ReferenceContext,
+        ReferenceContextAccessor,
+    )
+
+    _reference_context_available = True
+except ImportError:
+    pass
+
+pytestmark = pytest.mark.skipif(
+    not _reference_context_available,
+    reason="installed uipath does not export ReferenceContext",
+)
+
+
+# ---------------------------------------------------------------------------
+# Minimal graph fixture
+# ---------------------------------------------------------------------------
+
+
+class _State(TypedDict):
+    value: str
+
+
+def _build_graph() -> Any:
+    graph = StateGraph(_State)
+    graph.add_node("step", lambda s: {"value": s.get("value", "") + "_done"})
+    graph.add_edge(START, "step")
+    graph.add_edge("step", END)
+    return graph
+
+
+def _clear_accessor() -> None:
+    ReferenceContextAccessor.set(None)
+
+
+# ---------------------------------------------------------------------------
+# _push_reference_context — unit tests (no graph needed)
+# ---------------------------------------------------------------------------
+
+
+class TestPushReferenceContext:
+    def setup_method(self) -> None:
+        _clear_accessor()
+
+    def teardown_method(self) -> None:
+        _clear_accessor()
+
+    def test_sets_langgraph_entry_when_agent_id_present(
+        self, monkeypatch: pytest.MonkeyPatch, tmp_path
+    ) -> None:
+        monkeypatch.setenv("UIPATH_AGENT_ID", "550e8400-e29b-41d4-a716-446655440020")
+        monkeypatch.delenv("UIPATH_PROCESS_VERSION", raising=False)
+
+        graph = _build_graph().compile()
+        runtime = UiPathLangGraphRuntime(graph=graph, runtime_id="t")
+
+        token = runtime._push_reference_context()
+        try:
+            ctx = ReferenceContextAccessor.get()
+            assert ctx is not None
+            assert len(ctx) == 1
+            assert ctx.entries[0].service_type == "langgraph"
+            assert ctx.entries[0].reference_id == "550e8400-e29b-41d4-a716-446655440020"
+            assert ctx.entries[0].version is None
+        finally:
+            ReferenceContextAccessor.reset(token)
+
+    def test_includes_version_when_env_set(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setenv("UIPATH_AGENT_ID", "550e8400-e29b-41d4-a716-446655440020")
+        monkeypatch.setenv("UIPATH_PROCESS_VERSION", "3.1.0")
+
+        graph = _build_graph().compile()
+        runtime = UiPathLangGraphRuntime(graph=graph, runtime_id="t")
+
+        token = runtime._push_reference_context()
+        try:
+            ctx = ReferenceContextAccessor.get()
+            assert ctx is not None
+            assert ctx.entries[0].version == "3.1.0"
+        finally:
+            ReferenceContextAccessor.reset(token)
+
+    def test_no_entry_when_agent_id_absent(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.delenv("UIPATH_AGENT_ID", raising=False)
+        monkeypatch.delenv("UIPATH_PROCESS_VERSION", raising=False)
+
+        graph = _build_graph().compile()
+        runtime = UiPathLangGraphRuntime(graph=graph, runtime_id="t")
+
+        token = runtime._push_reference_context()
+        try:
+            ctx = ReferenceContextAccessor.get()
+            assert ctx is not None
+            assert len(ctx) == 0
+        finally:
+            ReferenceContextAccessor.reset(token)
+
+    def test_stacks_on_top_of_parent_context(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setenv("UIPATH_AGENT_ID", "550e8400-e29b-41d4-a716-446655440020")
+        monkeypatch.delenv("UIPATH_PROCESS_VERSION", raising=False)
+
+        parent = ReferenceContext.Empty.add(
+            "agent", "550e8400-e29b-41d4-a716-446655440001", "1.0"
+        )
+        parent_token = ReferenceContextAccessor.set(parent)
+
+        graph = _build_graph().compile()
+        runtime = UiPathLangGraphRuntime(graph=graph, runtime_id="t")
+
+        token = runtime._push_reference_context()
+        try:
+            ctx = ReferenceContextAccessor.get()
+            assert ctx is not None
+            assert len(ctx) == 2
+            assert ctx.entries[0].service_type == "agent"
+            assert ctx.entries[1].service_type == "langgraph"
+        finally:
+            ReferenceContextAccessor.reset(token)
+            ReferenceContextAccessor.reset(parent_token)
+
+
+# ---------------------------------------------------------------------------
+# execute() — context cleared after run
+# ---------------------------------------------------------------------------
+
+
+async def test_context_cleared_after_execute(
+    monkeypatch: pytest.MonkeyPatch, tmp_path
+) -> None:
+    _clear_accessor()
+    monkeypatch.setenv("UIPATH_AGENT_ID", "550e8400-e29b-41d4-a716-446655440020")
+    monkeypatch.delenv("UIPATH_PROCESS_VERSION", raising=False)
+
+    async with AsyncSqliteSaver.from_conn_string(str(tmp_path / "mem.db")) as memory:
+        await memory.setup()
+        graph = _build_graph().compile(checkpointer=memory)
+        runtime = UiPathLangGraphRuntime(graph=graph, runtime_id="exec-run")
+        await runtime.execute(input={"value": "hello"})
+
+    assert ReferenceContextAccessor.get() is None
+
+
+async def test_context_cleared_after_execute_on_error(
+    monkeypatch: pytest.MonkeyPatch, tmp_path
+) -> None:
+    _clear_accessor()
+    monkeypatch.setenv("UIPATH_AGENT_ID", "550e8400-e29b-41d4-a716-446655440020")
+
+    class _S(TypedDict):
+        v: str
+
+    def _boom(s: _S) -> _S:
+        raise ValueError("explode")
+
+    g = StateGraph(_S)
+    g.add_node("boom", _boom)  # type: ignore[arg-type]
+    g.add_edge(START, "boom")
+    g.add_edge("boom", END)
+
+    async with AsyncSqliteSaver.from_conn_string(str(tmp_path / "mem.db")) as memory:
+        await memory.setup()
+        compiled = g.compile(checkpointer=memory)
+        runtime = UiPathLangGraphRuntime(graph=compiled, runtime_id="err-run")
+        with pytest.raises(LangGraphRuntimeError):
+            await runtime.execute(input={"v": "x"})
+
+    assert ReferenceContextAccessor.get() is None
+
+
+# ---------------------------------------------------------------------------
+# stream() — context cleared after run
+# ---------------------------------------------------------------------------
+
+
+async def test_context_cleared_after_stream(
+    monkeypatch: pytest.MonkeyPatch, tmp_path
+) -> None:
+    _clear_accessor()
+    monkeypatch.setenv("UIPATH_AGENT_ID", "550e8400-e29b-41d4-a716-446655440020")
+    monkeypatch.delenv("UIPATH_PROCESS_VERSION", raising=False)
+
+    async with AsyncSqliteSaver.from_conn_string(str(tmp_path / "mem.db")) as memory:
+        await memory.setup()
+        graph = _build_graph().compile(checkpointer=memory)
+        runtime = UiPathLangGraphRuntime(graph=graph, runtime_id="stream-run")
+        async for _ in runtime.stream(input={"value": "hi"}):
+            pass
+
+    assert ReferenceContextAccessor.get() is None
+
+
+async def test_context_cleared_after_stream_on_error(
+    monkeypatch: pytest.MonkeyPatch, tmp_path
+) -> None:
+    _clear_accessor()
+    monkeypatch.setenv("UIPATH_AGENT_ID", "550e8400-e29b-41d4-a716-446655440020")
+
+    class _S(TypedDict):
+        v: str
+
+    def _boom(s: _S) -> _S:
+        raise ValueError("explode")
+
+    g = StateGraph(_S)
+    g.add_node("boom", _boom)  # type: ignore[arg-type]
+    g.add_edge(START, "boom")
+    g.add_edge("boom", END)
+
+    async with AsyncSqliteSaver.from_conn_string(
+        str(tmp_path / "stream-err.db")
+    ) as memory:
+        await memory.setup()
+        compiled = g.compile(checkpointer=memory)
+        runtime = UiPathLangGraphRuntime(graph=compiled, runtime_id="stream-err-run")
+        with pytest.raises(LangGraphRuntimeError):
+            async for _ in runtime.stream(input={"v": "x"}):
+                pass
+
+    assert ReferenceContextAccessor.get() is None