From cdaa7ba00389c758445fff4015481e8e38f461ef Mon Sep 17 00:00:00 2001
From: William FH <13333726+hinthornw@users.noreply.github.com>
Date: Thu, 24 Jul 2025 18:00:44 -0700
Subject: [PATCH] chore: typing for headers in remote graph (#5653)

---
 libs/langgraph/langgraph/pregel/remote.py | 32 ++++++--
 libs/langgraph/tests/test_remote_graph.py | 91 ++++++++++++++++++++---
 2 files changed, 104 insertions(+), 19 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py
index fbdc1615b..611c56d28 100644
--- a/libs/langgraph/langgraph/pregel/remote.py
+++ b/libs/langgraph/langgraph/pregel/remote.py
@@ -633,6 +633,7 @@ class RemoteGraph(PregelProtocol):
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
         subgraphs: bool = False,
+        headers: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> Iterator[dict[str, Any] | Any]:
         """Create a run and stream the results.
@@ -648,6 +649,7 @@ class RemoteGraph(PregelProtocol):
             interrupt_before: Interrupt the graph before these nodes.
             interrupt_after: Interrupt the graph after these nodes.
             subgraphs: Stream from subgraphs.
+            headers: Additional headers to pass to the request.
             **kwargs: Additional params to pass to client.runs.stream.
 
         Yields:
@@ -676,7 +678,9 @@ class RemoteGraph(PregelProtocol):
             interrupt_after=interrupt_after,
             stream_subgraphs=subgraphs or stream is not None,
             if_not_exists="create",
-            headers=self._merge_tracing_headers(kwargs.pop("headers", None) or {}),
+            headers=_merge_tracing_headers(headers)
+            if self.distributed_tracing
+            else headers,
             **kwargs,
         ):
             # split mode and ns
@@ -736,6 +740,7 @@ class RemoteGraph(PregelProtocol):
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
         subgraphs: bool = False,
+        headers: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> AsyncIterator[dict[str, Any] | Any]:
         """Create a run and stream the results.
@@ -751,6 +756,7 @@ class RemoteGraph(PregelProtocol):
             interrupt_before: Interrupt the graph before these nodes.
             interrupt_after: Interrupt the graph after these nodes.
             subgraphs: Stream from subgraphs.
+            headers: Additional headers to pass to the request.
             **kwargs: Additional params to pass to client.runs.stream.
 
         Yields:
@@ -779,7 +785,9 @@ class RemoteGraph(PregelProtocol):
             interrupt_after=interrupt_after,
             stream_subgraphs=subgraphs or stream is not None,
             if_not_exists="create",
-            headers=self._merge_tracing_headers(kwargs.pop("headers", None) or {}),
+            headers=_merge_tracing_headers(headers)
+            if self.distributed_tracing
+            else headers,
             **kwargs,
         ):
             # split mode and ns
@@ -853,6 +861,7 @@ class RemoteGraph(PregelProtocol):
         *,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
+        headers: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> dict[str, Any] | Any:
         """Create a run, wait until it finishes and return the final state.
@@ -862,6 +871,7 @@ class RemoteGraph(PregelProtocol):
             config: A `RunnableConfig` for graph invocation.
             interrupt_before: Interrupt the graph before these nodes.
             interrupt_after: Interrupt the graph after these nodes.
+            headers: Additional headers to pass to the request.
             **kwargs: Additional params to pass to RemoteGraph.stream.
 
         Returns:
@@ -872,6 +882,7 @@ class RemoteGraph(PregelProtocol):
             config=config,
             interrupt_before=interrupt_before,
             interrupt_after=interrupt_after,
+            headers=headers,
             stream_mode="values",
             **kwargs,
         ):
@@ -888,6 +899,7 @@ class RemoteGraph(PregelProtocol):
         *,
         interrupt_before: All | Sequence[str] | None = None,
         interrupt_after: All | Sequence[str] | None = None,
+        headers: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> dict[str, Any] | Any:
         """Create a run, wait until it finishes and return the final state.
@@ -897,6 +909,7 @@ class RemoteGraph(PregelProtocol):
             config: A `RunnableConfig` for graph invocation.
             interrupt_before: Interrupt the graph before these nodes.
             interrupt_after: Interrupt the graph after these nodes.
+            headers: Additional headers to pass to the request.
             **kwargs: Additional params to pass to RemoteGraph.astream.
 
         Returns:
@@ -907,6 +920,7 @@ class RemoteGraph(PregelProtocol):
             config=config,
             interrupt_before=interrupt_before,
             interrupt_after=interrupt_after,
+            headers=headers,
             stream_mode="values",
             **kwargs,
         ):
@@ -916,12 +930,16 @@ class RemoteGraph(PregelProtocol):
         except UnboundLocalError:
             return None
 
-    def _merge_tracing_headers(self, headers: dict[str, str]) -> dict[str, str]:
-        if rt := ls.get_current_run_tree():
-            tracing_headers = rt.to_headers()
-            baggage = tracing_headers.pop("baggage")
+
+def _merge_tracing_headers(headers: dict[str, str] | None) -> dict[str, str] | None:
+    if rt := ls.get_current_run_tree():
+        tracing_headers = rt.to_headers()
+        baggage = tracing_headers.pop("baggage")
+        if headers:
             if "baggage" in headers:
                 baggage = headers["baggage"] + "," + baggage
             tracing_headers["baggage"] = baggage
             headers.update(tracing_headers)
-        return headers
+        else:
+            headers = tracing_headers
+    return headers
diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py
index 615a5a833..737a9f750 100644
--- a/libs/langgraph/tests/test_remote_graph.py
+++ b/libs/langgraph/tests/test_remote_graph.py
@@ -3,6 +3,7 @@ import sys
 from typing import Annotated, Union
 from unittest.mock import AsyncMock, MagicMock
 
+import langsmith as ls
 import pytest
 from langchain_core.messages import AnyMessage, BaseMessage
 from langchain_core.runnables import RunnableConfig
@@ -899,21 +900,20 @@ async def test_langgraph_cloud_integration():
     }
 
     # test invoke
-    response = app.invoke(
+    app.invoke(
         input,
         config={"configurable": {"thread_id": "39a6104a-34e7-4f83-929c-d9eb163003c9"}},
         interrupt_before=["agent"],
     )
-    print("response:", response["messages"][-1].content)
 
     # test stream
-    async for chunk in app.astream(
+    async for _ in app.astream(
         input,
         config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
         subgraphs=True,
         stream_mode=["debug", "messages"],
     ):
-        print("chunk:", chunk)
+        pass
 
     # test stream events
     async for chunk in remote_pregel.astream_events(
@@ -923,17 +923,16 @@ async def test_langgraph_cloud_integration():
         subgraphs=True,
         stream_mode=[],
     ):
-        print("chunk:", chunk)
+        pass
 
     # test get state
-    state_snapshot = await remote_pregel.aget_state(
+    await remote_pregel.aget_state(
         config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
         subgraphs=True,
     )
-    print("state snapshot:", state_snapshot)
 
     # test update state
-    response = await remote_pregel.aupdate_state(
+    await remote_pregel.aupdate_state(
         config={"configurable": {"thread_id": "6645e002-ed50-4022-92a3-d0d186fdf812"}},
         values={
             "messages": [
@@ -944,18 +943,16 @@ async def test_langgraph_cloud_integration():
             ]
         },
     )
-    print("response:", response)
 
     # test get history
     async for state in remote_pregel.aget_state_history(
         config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
     ):
-        print("state snapshot:", state)
+        pass
 
     # test get graph
     remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca"  # must be UUID
-    graph = await remote_pregel.aget_graph(xray=True)
-    print("graph:", graph)
+    await remote_pregel.aget_graph(xray=True)
 
 
 def test_sanitize_config():
@@ -1181,3 +1178,73 @@ async def test_remote_graph_stream_messages_tuple(
     assert coerced_events == coerced_inmem_events
     # TODO: Fix the namespace matching in the next api release.
     # assert namespaces == inmem_namespaces
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("distributed_tracing", [False, True])
+@pytest.mark.parametrize("stream", [False, True])
+async def test_include_headers(distributed_tracing: bool, stream: bool):
+    mock_async_client = MagicMock()
+    async_iter = MagicMock()
+    return_value = [
+        StreamPart(event="values", data={"chunk": "data1"}),
+    ]
+    async_iter.__aiter__.return_value = return_value
+    astream_mock = mock_async_client.runs.stream
+    astream_mock.return_value = async_iter
+
+    mock_sync_client = MagicMock()
+    sync_iter = MagicMock()
+    sync_iter.__iter__.return_value = return_value
+    stream_mock = mock_sync_client.runs.stream
+    stream_mock.return_value = async_iter
+
+    remote_pregel = RemoteGraph(
+        "test_graph_id",
+        client=mock_async_client,
+        sync_client=mock_sync_client,
+        distributed_tracing=distributed_tracing,
+    )
+
+    config = {"configurable": {"thread_id": "thread_1"}}
+    with ls.tracing_context(enabled=True, client=MagicMock()):
+        with ls.trace("foo"):
+            if stream:
+                async for _ in remote_pregel.astream(
+                    {"input": {"messages": [{"type": "human", "content": "hello"}]}},
+                    config,
+                    headers={"foo": "bar"},
+                ):
+                    pass
+
+            else:
+                await remote_pregel.ainvoke(
+                    {"input": {"messages": [{"type": "human", "content": "hello"}]}},
+                    config,
+                    headers={"foo": "bar"},
+                )
+    expected = {"foo": "bar"}
+    if distributed_tracing:
+        expected["langsmith-trace"] = AnyStr()
+        expected["baggage"] = AnyStr()
+
+    assert astream_mock.call_args.kwargs["headers"] == expected
+    stream_mock.assert_not_called()
+
+    with ls.tracing_context(enabled=True, client=MagicMock()):
+        with ls.trace("foo"):
+            if stream:
+                for _ in remote_pregel.stream(
+                    {"input": {"messages": [{"type": "human", "content": "hello"}]}},
+                    config,
+                    headers={"foo": "bar"},
+                ):
+                    pass
+
+            else:
+                remote_pregel.invoke(
+                    {"input": {"messages": [{"type": "human", "content": "hello"}]}},
+                    config,
+                    headers={"foo": "bar"},
+                )
+    assert stream_mock.call_args.kwargs["headers"] == expected