langgraph: update get_state to handle nested subgraph state

libs/langgraph/langgraph/checkpoint/base.py

@@ -231,6 +231,11 @@ class BaseCheckpointSaver(ABC):
         """
         raise NotImplementedError
 
+    def list_subgraph_checkpoints(
+        self, config: RunnableConfig
+    ) -> Iterator[CheckpointTuple]:
+        raise NotImplementedError
+
     def put(
         self,
         config: RunnableConfig,

libs/langgraph/langgraph/checkpoint/memory.py

@@ -171,6 +171,39 @@ class MemorySaver(BaseCheckpointSaver):
                     else None,
                 )
 
+    def list_subgraph_checkpoints(
+        self, config: RunnableConfig
+    ) -> Iterator[CheckpointTuple]:
+        thread_id_prefix = config["configurable"]["thread_id"]
+        matching_thread_ids = [
+            key for key in self.storage.keys() if key.startswith(thread_id_prefix)
+        ]
+        for thread_id in matching_thread_ids:
+            ts = config["configurable"].get("thread_ts")
+            if not ts:
+                if checkpoints := self.storage[thread_id]:
+                    ts = max(checkpoints.keys())
+
+            if saved := self.storage[thread_id].get(ts):
+                checkpoint, metadata, parent_ts = saved
+                writes = self.writes[(thread_id, ts)]
+                yield CheckpointTuple(
+                    config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
+                    checkpoint=self.serde.loads(checkpoint),
+                    metadata=self.serde.loads(metadata),
+                    pending_writes=[
+                        (id, c, self.serde.loads(v)) for id, c, v in writes
+                    ],
+                    parent_config={
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "thread_ts": parent_ts,
+                        }
+                    }
+                    if parent_ts
+                    else None,
+                )
+
     def put(
         self,
         config: RunnableConfig,

libs/langgraph/langgraph/checkpoint/sqlite.py

@@ -357,6 +357,49 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
                     ),
                 )
 
+    def list_subgraph_checkpoints(
+        self, config: RunnableConfig
+    ) -> Iterator[CheckpointTuple]:
+        with self.cursor(transaction=False) as cur:
+            if config["configurable"].get("thread_ts"):
+                cur.execute(
+                    "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id LIKE ? || '%' AND thread_ts = ?",
+                    (
+                        str(config["configurable"]["thread_id"]),
+                        str(config["configurable"]["thread_ts"]),
+                    ),
+                )
+            else:
+                cur.execute(
+                    """SELECT checkpoints.thread_id, checkpoints.thread_ts, checkpoints.parent_ts, checkpoints.checkpoint, checkpoints.metadata
+                    FROM checkpoints
+                    INNER JOIN (
+                        SELECT thread_id, MAX(thread_ts) as thread_ts
+                        FROM checkpoints
+                        WHERE thread_id LIKE ? || '%'
+                        GROUP BY thread_id
+                    ) latest_checkpoints
+                    ON checkpoints.thread_id = latest_checkpoints.thread_id AND checkpoints.thread_ts = latest_checkpoints.thread_ts
+                    ORDER BY checkpoints.thread_id, checkpoints.thread_ts DESC""",
+                    (str(config["configurable"]["thread_id"]),),
+                )
+            for thread_id, thread_ts, parent_ts, value, metadata in cur:
+                yield CheckpointTuple(
+                    {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
+                    self.serde.loads(value),
+                    self.serde.loads(metadata) if metadata is not None else {},
+                    (
+                        {
+                            "configurable": {
+                                "thread_id": thread_id,
+                                "thread_ts": parent_ts,
+                            }
+                        }
+                        if parent_ts
+                        else None
+                    ),
+                )
+
     def put(
         self,
         config: RunnableConfig,

libs/langgraph/langgraph/constants.py

@@ -18,6 +18,8 @@ RESERVED = {
 }
 TAG_HIDDEN = "langsmith:hidden"
 
+THREAD_ID_SEPARATOR = "__"
+
 START = "__start__"
 END = "__end__"
 

libs/langgraph/langgraph/pregel/__init__.py

@@ -59,6 +59,7 @@ from langgraph.channels.manager import (
 )
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
+    CheckpointTuple,
     copy_checkpoint,
     empty_checkpoint,
 )
@@ -68,6 +69,7 @@ from langgraph.constants import (
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_SEND,
     INTERRUPT,
+    THREAD_ID_SEPARATOR,
 )
 from langgraph.errors import GraphRecursionError, InvalidUpdateError
 from langgraph.managed.base import (
@@ -350,12 +352,9 @@ class Pregel(
             if is_managed_value(v)
         }
 
-    def get_state(self, config: RunnableConfig) -> StateSnapshot:
-        """Get the current state of the graph."""
-        if not self.checkpointer:
-            raise ValueError("No checkpointer set")
-
-        saved = self.checkpointer.get_tuple(config)
+    def _prepare_state_snapshot(
+        self, saved: CheckpointTuple, config: RunnableConfig
+    ) -> StateSnapshot:
         checkpoint = saved.checkpoint if saved else empty_checkpoint()
         config = saved.config if saved else config
         with ChannelsManager(
@@ -373,14 +372,58 @@ class Pregel(
                 for_execution=False,
             )
             return StateSnapshot(
-                read_channels(channels, self.stream_channels_asis),
-                tuple(name for name, _ in next_tasks),
-                saved.config if saved else config,
-                saved.metadata if saved else None,
-                saved.checkpoint["ts"] if saved else None,
-                saved.parent_config if saved else None,
+                values=read_channels(channels, self.stream_channels_asis),
+                next=tuple(name for name, _ in next_tasks),
+                config=saved.config if saved else config,
+                metadata=saved.metadata if saved else None,
+                created_at=saved.checkpoint["ts"] if saved else None,
+                parent_config=saved.parent_config if saved else None,
             )
 
+    @staticmethod
+    def _assemble_state_snapshot_hierarchy(
+        root_thread_id: str, subgraph_state_snapshots: dict[str, StateSnapshot]
+    ) -> StateSnapshot:
+        thread_ids_to_visit = sorted(
+            subgraph_state_snapshots.keys(),
+            key=lambda x: len(x.split(THREAD_ID_SEPARATOR)),
+        )
+        while thread_ids_to_visit:
+            thread_id = thread_ids_to_visit.pop()
+            state_snapshot = subgraph_state_snapshots[thread_id]
+            *path, subgraph_node = thread_id.split(THREAD_ID_SEPARATOR)
+            parent_thread_id = THREAD_ID_SEPARATOR.join(path)
+            if parent_thread_id and THREAD_ID_SEPARATOR in parent_thread_id:
+                parent_subgraph_snapshots = (
+                    subgraph_state_snapshots[parent_thread_id].subgraph_state_snapshots
+                    or {}
+                )
+                parent_subgraph_snapshots[subgraph_node] = state_snapshot
+                subgraph_state_snapshots[parent_thread_id] = subgraph_state_snapshots[
+                    parent_thread_id
+                ]._replace(subgraph_state_snapshots=parent_subgraph_snapshots)
+
+        state_snapshot = subgraph_state_snapshots.pop(root_thread_id)
+        return state_snapshot
+
+    def get_state(self, config: RunnableConfig) -> StateSnapshot:
+        """Get the current state of the graph."""
+        if not self.checkpointer:
+            raise ValueError("No checkpointer set")
+
+        subgraph_state_snapshots: dict[str, StateSnapshot] = {
+            checkpoint.config["configurable"][
+                "thread_id"
+            ]: self._prepare_state_snapshot(checkpoint, config)
+            for checkpoint in self.checkpointer.list_subgraph_checkpoints(config)
+        }
+
+        thread_id = config["configurable"]["thread_id"]
+        state_snapshot = self._assemble_state_snapshot_hierarchy(
+            thread_id, subgraph_state_snapshots
+        )
+        return state_snapshot
+
     async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
         """Get the current state of the graph."""
         if not self.checkpointer:

libs/langgraph/langgraph/pregel/algo.py

@@ -36,6 +36,7 @@ from langgraph.constants import (
     RESERVED,
     TAG_HIDDEN,
     TASKS,
+    THREAD_ID_SEPARATOR,
     Send,
 )
 from langgraph.errors import EmptyChannelError, InvalidUpdateError
@@ -345,7 +346,9 @@ def prepare_next_tasks(
                     if parent_thread_id := config.get("configurable", {}).get(
                         "thread_id"
                     ):
-                        thread_id: Optional[str] = f"{parent_thread_id}-{name}"
+                        thread_id: Optional[
+                            str
+                        ] = f"{parent_thread_id}{THREAD_ID_SEPARATOR}{name}"
                     else:
                         thread_id = None
                     writes = deque()

libs/langgraph/langgraph/pregel/types.py

@@ -85,6 +85,8 @@ class StateSnapshot(NamedTuple):
     """Timestamp of snapshot creation"""
     parent_config: Optional[RunnableConfig] = None
     """Config used to fetch the parent snapshot, if any"""
+    subgraph_state_snapshots: Optional[dict[str, "StateSnapshot"]] = None
+    """State snapshots of subgraphs represented as a mapping from thread ID suffix to snapshot."""
 
 
 All = Literal["*"]

libs/langgraph/tests/test_pregel.py

@@ -568,53 +568,53 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
     )
 
     # start execution, stop at inbox
-    assert app.invoke(2, {"configurable": {"thread_id": 1}}) is None
+    assert app.invoke(2, {"configurable": {"thread_id": "1"}}) is None
 
     # inbox == 3
-    checkpoint = memory.get({"configurable": {"thread_id": 1}})
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
     assert checkpoint is not None
     assert checkpoint["channel_values"]["inbox"] == 3
 
     # resume execution, finish
-    assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 4
+    assert app.invoke(None, {"configurable": {"thread_id": "1"}}) == 4
 
     # start execution again, stop at inbox
-    assert app.invoke(20, {"configurable": {"thread_id": 1}}) is None
+    assert app.invoke(20, {"configurable": {"thread_id": "1"}}) is None
 
     # inbox == 21
-    checkpoint = memory.get({"configurable": {"thread_id": 1}})
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
     assert checkpoint is not None
     assert checkpoint["channel_values"]["inbox"] == 21
 
     # send a new value in, interrupting the previous execution
-    assert app.invoke(3, {"configurable": {"thread_id": 1}}) is None
-    assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 5
+    assert app.invoke(3, {"configurable": {"thread_id": "1"}}) is None
+    assert app.invoke(None, {"configurable": {"thread_id": "1"}}) == 5
 
     # start execution again, stopping at inbox
-    assert app.invoke(20, {"configurable": {"thread_id": 2}}) is None
+    assert app.invoke(20, {"configurable": {"thread_id": "2"}}) is None
 
     # inbox == 21
-    snapshot = app.get_state({"configurable": {"thread_id": 2}})
+    snapshot = app.get_state({"configurable": {"thread_id": "2"}})
     assert snapshot.values["inbox"] == 21
     assert snapshot.next == ("two",)
 
     # update the state, resume
-    app.update_state({"configurable": {"thread_id": 2}}, 25, as_node="one")
-    assert app.invoke(None, {"configurable": {"thread_id": 2}}) == 26
+    app.update_state({"configurable": {"thread_id": "2"}}, 25, as_node="one")
+    assert app.invoke(None, {"configurable": {"thread_id": "2"}}) == 26
 
     # no pending tasks
-    snapshot = app.get_state({"configurable": {"thread_id": 2}})
+    snapshot = app.get_state({"configurable": {"thread_id": "2"}})
     assert snapshot.next == ()
 
     # list history
-    thread1 = {"configurable": {"thread_id": 1}}
+    thread1 = {"configurable": {"thread_id": "1"}}
     assert [c for c in app.get_state_history(thread1)] == [
         StateSnapshot(
             values={"inbox": 4, "output": 5, "input": 3},
             next=(),
             config={
                 "configurable": {
-                    "thread_id": 1,
+                    "thread_id": "1",
                     "thread_ts": AnyStr(),
                 }
             },
@@ -627,7 +627,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
             next=("two",),
             config={
                 "configurable": {
-                    "thread_id": 1,
+                    "thread_id": "1",
                     "thread_ts": AnyStr(),
                 }
             },
@@ -640,7 +640,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
             next=("one",),
             config={
                 "configurable": {
-                    "thread_id": 1,
+                    "thread_id": "1",
                     "thread_ts": AnyStr(),
                 }
             },
@@ -653,7 +653,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
             next=("two",),
             config={
                 "configurable": {
-                    "thread_id": 1,
+                    "thread_id": "1",
                     "thread_ts": AnyStr(),
                 }
             },
@@ -666,7 +666,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
             next=("one",),
             config={
                 "configurable": {
-                    "thread_id": 1,
+                    "thread_id": "1",
                     "thread_ts": AnyStr(),
                 }
             },
@@ -679,7 +679,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
             next=(),
             config={
                 "configurable": {
-                    "thread_id": 1,
+                    "thread_id": "1",
                     "thread_ts": AnyStr(),
                 }
             },
@@ -692,7 +692,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
             next=("two",),
             config={
                 "configurable": {
-                    "thread_id": 1,
+                    "thread_id": "1",
                     "thread_ts": AnyStr(),
                 }
             },
@@ -705,7 +705,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
             next=("one",),
             config={
                 "configurable": {
-                    "thread_id": 1,
+                    "thread_id": "1",
                     "thread_ts": AnyStr(),
                 }
             },
@@ -1080,7 +1080,7 @@ def test_pending_writes_resume(checkpointer: BaseCheckpointSaver) -> None:
         builder.add_edge(START, "two")
         graph = builder.compile(checkpointer=checkpointer)
 
-        thread1: RunnableConfig = {"configurable": {"thread_id": 1}}
+        thread1: RunnableConfig = {"configurable": {"thread_id": "1"}}
         with pytest.raises(ConnectionError, match="I'm not good"):
             graph.invoke({"value": 1}, thread1)