fix(langgraph): Remote Baggage (#5964)

Fix baggage propagation for opt-in distributed tracing when no
additional headers are provided.

Author: hinthornw

libs/langgraph/langgraph/pregel/remote.py

@@ -733,9 +733,9 @@ def stream(
             interrupt_after=interrupt_after,
             stream_subgraphs=subgraphs or stream is not None,
             if_not_exists="create",
-            headers=_merge_tracing_headers(headers)
-            if self.distributed_tracing
-            else headers,
+            headers=(
+                _merge_tracing_headers(headers) if self.distributed_tracing else headers
+            ),
             params=params,
             **kwargs,
         ):
@@ -842,9 +842,9 @@ async def astream(
             interrupt_after=interrupt_after,
             stream_subgraphs=subgraphs or stream is not None,
             if_not_exists="create",
-            headers=_merge_tracing_headers(headers)
-            if self.distributed_tracing
-            else headers,
+            headers=(
+                _merge_tracing_headers(headers) if self.distributed_tracing else headers
+            ),
             params=params,
             **kwargs,
         ):
@@ -996,11 +996,11 @@ async def ainvoke(
 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
+                tracing_headers["baggage"] = (
+                    f"{headers['baggage']},{tracing_headers['baggage']}"
+                )
             headers.update(tracing_headers)
         else:
             headers = tracing_headers

libs/langgraph/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph"
-version = "0.6.5"
+version = "0.6.6"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 requires-python = ">=3.9"

libs/langgraph/tests/test_remote_graph.py

@@ -1,6 +1,6 @@
 import re
 import sys
-from typing import Annotated, Union
+from typing import Annotated, Optional, Union
 from unittest.mock import AsyncMock, MagicMock
 
 import langsmith as ls
@@ -1183,7 +1183,10 @@ async def test_remote_graph_stream_messages_tuple(
 @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):
+@pytest.mark.parametrize("headers", [None, {"foo": "bar"}])
+async def test_include_headers(
+    distributed_tracing: bool, stream: bool, headers: Optional[dict[str, str]]
+):
     mock_async_client = MagicMock()
     async_iter = MagicMock()
     return_value = [
@@ -1213,20 +1216,22 @@ async def test_include_headers(distributed_tracing: bool, stream: bool):
                 async for _ in remote_pregel.astream(
                     {"input": {"messages": [{"type": "human", "content": "hello"}]}},
                     config,
-                    headers={"foo": "bar"},
+                    headers=headers,
                 ):
                     pass
 
             else:
                 await remote_pregel.ainvoke(
                     {"input": {"messages": [{"type": "human", "content": "hello"}]}},
                     config,
-                    headers={"foo": "bar"},
+                    headers=headers,
                 )
-    expected = {"foo": "bar"}
+    expected = headers.copy() if headers else None
     if distributed_tracing:
+        if expected is None:
+            expected = {}
         expected["langsmith-trace"] = AnyStr()
-        expected["baggage"] = AnyStr()
+        expected["baggage"] = AnyStr("langsmith-metadata=")
 
     assert astream_mock.call_args.kwargs["headers"] == expected
     stream_mock.assert_not_called()
@@ -1237,14 +1242,14 @@ async def test_include_headers(distributed_tracing: bool, stream: bool):
                 for _ in remote_pregel.stream(
                     {"input": {"messages": [{"type": "human", "content": "hello"}]}},
                     config,
-                    headers={"foo": "bar"},
+                    headers=headers,
                 ):
                     pass
 
             else:
                 remote_pregel.invoke(
                     {"input": {"messages": [{"type": "human", "content": "hello"}]}},
                     config,
-                    headers={"foo": "bar"},
+                    headers=headers,
                 )
     assert stream_mock.call_args.kwargs["headers"] == expected