fix: preserve JSON-safe checkpoint serialization for msgpack async saves (#182)

* fix: normalize msgpack checkpoints for RedisJSON

* fix: preserve msgpack checkpoints with non-string keys

Author: vishal-bala

langgraph/checkpoint/redis/base.py

@@ -276,17 +276,14 @@ def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]:
         if type_ == "json":
             checkpoint_data = cast(dict, orjson.loads(data))
         else:
-            # For msgpack or other types, deserialize with loads_typed
             checkpoint_data = cast(dict, self.serde.loads_typed((type_, data)))
-
-            # When using msgpack, bytes are preserved - but Redis JSON.SET can't handle them
-            # Encode bytes in channel_values with type marker for JSON storage
-            if "channel_values" in checkpoint_data:
-                for key, value in checkpoint_data["channel_values"].items():
-                    if isinstance(value, bytes):
-                        checkpoint_data["channel_values"][key] = {
-                            "__bytes__": self._encode_blob(value)
-                        }
+            if type_ == "msgpack":
+                # Msgpack fallback can rehydrate LangChain messages as live Python
+                # objects. Normalize the checkpoint back through the JSON serializer
+                # so RedisJSON only sees JSON-safe constructor dictionaries.
+                checkpoint_data = cast(
+                    dict, self._msgpack_to_redis_json(checkpoint_data)
+                )
 
         # Ensure channel_versions are always strings to fix issue #40
         if "channel_versions" in checkpoint_data:
@@ -296,6 +293,30 @@ def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]:
 
         return {"type": type_, **checkpoint_data, "pending_sends": []}
 
+    def _msgpack_to_redis_json(self, value: Any) -> dict[str, Any]:
+        """Convert a msgpack-deserialized checkpoint into Redis JSON-safe data."""
+        binary_safe = self._replace_binary_markers(value)
+        serializer = cast(JsonPlusRedisSerializer, self.serde)
+        processed = serializer._preprocess_interrupts(binary_safe)
+        json_bytes = orjson.dumps(
+            processed,
+            default=serializer._default_handler,
+            option=orjson.OPT_NON_STR_KEYS,
+        )
+        return cast(dict, orjson.loads(json_bytes))
+
+    def _replace_binary_markers(self, value: Any) -> Any:
+        """Recursively replace binary values with JSON-safe markers."""
+        if isinstance(value, bytes):
+            return {"__bytes__": self._encode_blob(value)}
+        if isinstance(value, dict):
+            return {k: self._replace_binary_markers(v) for k, v in value.items()}
+        if isinstance(value, list):
+            return [self._replace_binary_markers(item) for item in value]
+        if isinstance(value, tuple):
+            return tuple(self._replace_binary_markers(item) for item in value)
+        return value
+
     def _deserialize_channel_values(
         self, channel_values: dict[str, Any]
     ) -> dict[str, Any]:

tests/test_checkpoint_serialization.py

@@ -561,6 +561,91 @@ def test_checkpoint_with_messages(redis_url: str) -> None:
         assert loaded_messages[1].content == "Let me check that for you."
 
 
+def test_issue_181_msgpack_checkpoint_keeps_messages_json_safe() -> None:
+    """Ensure msgpack fallback does not leak live message objects into Redis JSON."""
+    from langchain_core.messages import AIMessage, HumanMessage
+
+    from langgraph.checkpoint.redis.base import BaseRedisSaver
+
+    class DummySaver(BaseRedisSaver):
+        def create_indexes(self) -> None:
+            pass
+
+        def configure_client(self, **kwargs) -> None:
+            self._redis = None
+
+    saver = DummySaver(redis_client=object())
+    messages = [
+        HumanMessage(content="What is the weather in SF?"),
+        AIMessage(content="Let me check that for you."),
+    ]
+
+    checkpoint = create_checkpoint(
+        checkpoint=empty_checkpoint(),
+        channels={"messages": messages, "binary": b"abc"},
+        step=1,
+    )
+    checkpoint["channel_values"]["messages"] = messages
+    checkpoint["channel_values"]["binary"] = b"abc"
+
+    dumped = saver._dump_checkpoint(checkpoint)
+
+    # The mixed bytes payload forces msgpack fallback, but the Redis document
+    # must still be fully JSON-safe.
+    assert dumped["type"] == "msgpack"
+    assert isinstance(dumped["channel_values"]["messages"][0], dict)
+    assert dumped["channel_values"]["messages"][0]["id"][-1] == "HumanMessage"
+    assert dumped["channel_values"]["binary"] == {"__bytes__": "YWJj"}
+
+    restored = saver._load_checkpoint(
+        dumped,
+        saver._deserialize_channel_values(dumped["channel_values"]),
+        [],
+    )
+
+    restored_messages = restored["channel_values"]["messages"]
+    assert isinstance(restored_messages[0], HumanMessage)
+    assert isinstance(restored_messages[1], AIMessage)
+    assert restored["channel_values"]["binary"] == b"abc"
+
+
+def test_msgpack_checkpoint_with_non_string_keys_remains_json_safe() -> None:
+    """Ensure msgpack checkpoints with non-string keys still normalize for Redis JSON."""
+    from langchain_core.messages import HumanMessage
+
+    from langgraph.checkpoint.redis.base import BaseRedisSaver
+
+    class DummySaver(BaseRedisSaver):
+        def create_indexes(self) -> None:
+            pass
+
+        def configure_client(self, **kwargs) -> None:
+            self._redis = None
+
+    saver = DummySaver(redis_client=object())
+    messages = [HumanMessage(content="hello")]
+
+    checkpoint = create_checkpoint(
+        checkpoint=empty_checkpoint(),
+        channels={
+            "messages": messages,
+            "binary": b"abc",
+            "mapping": {1: "one", 2: "two"},
+        },
+        step=1,
+    )
+    checkpoint["channel_values"]["messages"] = messages
+    checkpoint["channel_values"]["binary"] = b"abc"
+    checkpoint["channel_values"]["mapping"] = {1: "one", 2: "two"}
+
+    dumped = saver._dump_checkpoint(checkpoint)
+
+    assert dumped["type"] == "msgpack"
+    assert dumped["channel_values"]["mapping"] == {"1": "one", "2": "two"}
+    assert isinstance(dumped["channel_values"]["messages"][0], dict)
+    assert dumped["channel_values"]["binary"] == {"__bytes__": "YWJj"}
+
+
 def test_subgraph_state_history_pending_sends(redis_url: str) -> None:
     """Test that get_state_history with subgraphs properly handles pending_sends.
 

tests/test_issue_87_async_deserialization.py

@@ -115,6 +115,38 @@ async def test_async_deserializes_langchain_messages(redis_url: str):
         assert loaded_messages[3].tool_call_id == "call-1"
 
 
+@pytest.mark.asyncio
+async def test_issue_181_async_checkpoint_with_messages_and_bytes(redis_url: str):
+    """Ensure mixed binary payloads do not break async message checkpoint writes."""
+    async with AsyncRedisSaver.from_conn_string(redis_url) as saver:
+        thread_id = str(uuid4())
+        messages = [
+            HumanMessage(content="What's the weather like?", id="human-1"),
+            AIMessage(content="I'll help you check the weather.", id="ai-1"),
+        ]
+
+        checkpoint = create_checkpoint(
+            checkpoint=empty_checkpoint(),
+            channels={"messages": messages, "binary": b"abc"},
+            step=1,
+        )
+        checkpoint["channel_values"]["messages"] = messages
+        checkpoint["channel_values"]["binary"] = b"abc"
+
+        config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
+
+        saved_config = await saver.aput(
+            config, checkpoint, {"source": "test", "step": 1, "writes": {}}, {}
+        )
+        loaded_tuple = await saver.aget_tuple(saved_config)
+
+        assert loaded_tuple is not None
+        loaded_messages = loaded_tuple.checkpoint["channel_values"]["messages"]
+        assert isinstance(loaded_messages[0], HumanMessage)
+        assert isinstance(loaded_messages[1], AIMessage)
+        assert loaded_tuple.checkpoint["channel_values"]["binary"] == b"abc"
+
+
 @pytest.mark.asyncio
 async def test_async_handles_serialized_langchain_format(redis_url: str):
     """Test that async handles the serialized LangChain format that causes MESSAGE_COERCION_FAILURE.