From 010564cbb392d8107d8648b74bdbf447b78aee4a Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 5 Nov 2024 15:24:20 -0800
Subject: [PATCH] lib: Make Control object serializable

---
 .../langgraph/checkpoint/serde/jsonplus.py      | 17 ++++++++++++++++-
 .../langgraph/checkpoint/serde/types.py         | 11 +++++++++++
 libs/langgraph/langgraph/types.py               |  8 +++++++-
 3 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
index f8d280b96..10908eb87 100644
--- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
+++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
@@ -25,7 +25,7 @@ from langchain_core.load.serializable import Serializable
 from zoneinfo import ZoneInfo
 
 from langgraph.checkpoint.serde.base import SerializerProtocol
-from langgraph.checkpoint.serde.types import SendProtocol
+from langgraph.checkpoint.serde.types import ControlProtocol, SendProtocol
 from langgraph.store.base import Item
 
 LC_REVIVER = Reviver()
@@ -402,6 +402,21 @@ def _msgpack_default(obj: Any) -> Union[str, msgpack.ExtType]:
                 (obj.__class__.__module__, obj.__class__.__name__, (obj.node, obj.arg)),
             ),
         )
+    elif isinstance(obj, ControlProtocol):
+        return msgpack.ExtType(
+            EXT_CONSTRUCTOR_KW_ARGS,
+            _msgpack_enc(
+                (
+                    obj.__class__.__module__,
+                    obj.__class__.__name__,
+                    {
+                        "update_state": obj.update_state,
+                        "trigger": obj.trigger,
+                        "send": obj.send,
+                    },
+                ),
+            ),
+        )
     elif dataclasses.is_dataclass(obj):
         # doesn't use dataclasses.asdict to avoid deepcopy and recursion
         return msgpack.ExtType(
diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py
index 43a5bf878..862cbe83f 100644
--- a/libs/checkpoint/langgraph/checkpoint/serde/types.py
+++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py
@@ -4,6 +4,7 @@ from typing import (
     Protocol,
     Sequence,
     TypeVar,
+    Union,
     runtime_checkable,
 )
 
@@ -48,3 +49,13 @@ class SendProtocol(Protocol):
     def __repr__(self) -> str: ...
 
     def __eq__(self, value: object) -> bool: ...
+
+
+@runtime_checkable
+class ControlProtocol(Protocol):
+    # Mirrors langgraph.constants.Control
+    update_state: Optional[dict[str, Any]]
+    trigger: Union[str, Sequence[str]]
+    send: Union[Any, Sequence[Any]]
+
+    def __repr__(self) -> str: ...
diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py
index 1ed6e43ad..b1a40f1a1 100644
--- a/libs/langgraph/langgraph/types.py
+++ b/libs/langgraph/langgraph/types.py
@@ -227,6 +227,10 @@ N = TypeVar("N")
 
 
 class Control(Generic[N]):
+    """A control object to update the graph's state, trigger nodes, and send messages."""
+
+    __slots__ = ("update_state", "trigger", "send")
+
     def __init__(
         self,
         *,
@@ -240,7 +244,9 @@ class Control(Generic[N]):
 
     def __repr__(self) -> str:
         contents = ", ".join(
-            f"{key}={value!r}" for key, value in self.__dict__.items() if value
+            f"{key}={value!r}"
+            for key in self.__slots__
+            if (value := getattr(self, key))
         )
         return f"Control({contents})"