Add get_state() and update_state() methods to get and update checkpoint in between runs

Makefile

@@ -18,7 +18,7 @@ test:
 	poetry run pytest
 
 test_watch:
-	poetry run ptw --snapshot-update --now . -- -vv -x --ff tests
+	poetry run ptw tests
 
 ######################
 # LINTING AND FORMATTING

langgraph/graph/graph.py

@@ -11,9 +11,10 @@ from langchain_core.runnables.base import (
 )
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.runnables.graph import Graph as RunnableGraph
+from langgraph.channels.ephemeral_value import EphemeralValue
 
 from langgraph.checkpoint import BaseCheckpointSaver
-from langgraph.pregel import Channel, Pregel
+from langgraph.pregel import Channel, Pregel, StateSnapshot
 
 logger = logging.getLogger(__name__)
 
@@ -190,6 +191,11 @@ class Graph:
             key: (Channel.subscribe_to(f"{key}:inbox") | node | Channel.write_to(key))
             for key, node in self.nodes.items()
         }
+        node_outboxes = {
+            # we clear outbox channels after each step
+            key: EphemeralValue(Any)
+            for key in self.nodes
+        }
 
         for key in self.nodes:
             outgoing = outgoing_edges[key]
@@ -216,6 +222,7 @@ class Graph:
         return CompiledGraph(
             graph=self,
             nodes=nodes,
+            channels={**node_outboxes},
             input=f"{self.entry_point}:inbox" if self.entry_point else START,
             output=END,
             hidden=[f"{node}:inbox" for node in self.nodes],
@@ -272,3 +279,19 @@ class CompiledGraph(Pregel):
             graph.add_edge(graph.nodes[START], graph.nodes[self.graph.entry_point])
 
         return graph
+
+    def get_state(self, config: RunnableConfig) -> StateSnapshot:
+        snapshot = super().get_state(config)
+
+        return StateSnapshot(
+            values={k: v for k, v in snapshot.values.items() if k in self.graph.nodes},
+            next=snapshot.next,
+        )
+
+    async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
+        snapshot = await super().aget_state(config)
+
+        return StateSnapshot(
+            values={k: v for k, v in snapshot.values.items() if k in self.graph.nodes},
+            next=snapshot.next,
+        )

langgraph/graph/state.py

@@ -1,10 +1,11 @@
 from collections import defaultdict
 from functools import partial
 from inspect import signature
-from typing import Any, Optional, Sequence, Type
+from typing import Any, Optional, Sequence, Type, Union
 
 from langchain_core.runnables import RunnableLambda
 from langchain_core.runnables.base import RunnableLike
+from langchain_core.runnables.config import RunnableConfig
 
 from langgraph.channels.any_value import AnyValue
 from langgraph.channels.base import BaseChannel, InvalidUpdateError
@@ -13,7 +14,7 @@ from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.channels.last_value import LastValue
 from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.graph.graph import END, START, CompiledGraph, Graph
-from langgraph.pregel import Channel
+from langgraph.pregel import Channel, StateSnapshot
 from langgraph.pregel.read import ChannelInvoke
 from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
 
@@ -134,7 +135,7 @@ class StateGraph(Graph):
         else:
             raise ValueError("No entry point set")
 
-        return CompiledGraph(
+        return CompiledStateGraph(
             graph=self,
             nodes=nodes,
             channels={
@@ -204,3 +205,43 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
             ):
                 return BinaryOperatorAggregate(typ, meta[0])
     return None
+
+
+class CompiledStateGraph(CompiledGraph):
+    graph: StateGraph
+
+    def get_state(self, config: RunnableConfig) -> StateSnapshot:
+        snapshot = super(CompiledGraph, self).get_state(config)
+
+        return StateSnapshot(
+            values=snapshot.values.get("__root__")
+            if "__root__" in self.graph.channels
+            else {k: v for k, v in snapshot.values.items() if k in self.graph.channels},
+            next=snapshot.next,
+        )
+
+    async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
+        snapshot = await super(CompiledGraph, self).aget_state(config)
+
+        return StateSnapshot(
+            values=snapshot.values.get("__root__")
+            if "__root__" in self.graph.channels
+            else {k: v for k, v in snapshot.values.items() if k in self.graph.channels},
+            next=snapshot.next,
+        )
+
+    def update_state(
+        self, config: RunnableConfig, values: Union[Any, dict[str, Any]]
+    ) -> None:
+        return super(CompiledGraph, self).update_state(
+            config,
+            {"__root__": values} if "__root__" in self.graph.channels else values,
+        )
+
+    async def aupdate_state(
+        self, config: RunnableConfig, values: Union[Any, dict[str, Any]]
+    ) -> None:
+        return await super(CompiledGraph, self).aupdate_state(
+            config,
+            {"__root__": values} if "__root__" in self.graph.channels else values,
+        )

langgraph/pregel/__init__.py

@@ -11,6 +11,7 @@ from typing import (
     Callable,
     Iterator,
     Mapping,
+    NamedTuple,
     Optional,
     Sequence,
     Type,
@@ -40,6 +41,7 @@ from langchain_core.runnables.utils import (
     get_unique_config_specs,
 )
 from langchain_core.tracers.log_stream import LogStreamCallbackHandler
+from langgraph.channels.any_value import AnyValue
 
 from langgraph.channels.base import (
     AsyncChannelsManager,
@@ -49,6 +51,7 @@ from langgraph.channels.base import (
     InvalidUpdateError,
     create_checkpoint,
 )
+from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.channels.last_value import LastValue
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
@@ -158,6 +161,13 @@ class Channel:
         )
 
 
+class StateSnapshot(NamedTuple):
+    values: dict[str, Any]
+    """Current values of channels"""
+    next: tuple[str]
+    """Nodes to execute in the next step, if any"""
+
+
 class Pregel(
     RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]
 ):
@@ -247,6 +257,72 @@ class Pregel(
                 **{k: (self.channels[k].ValueType, None) for k in self.output},
             )
 
+    def get_state(self, config: RunnableConfig) -> StateSnapshot:
+        if not self.checkpointer:
+            raise ValueError("No checkpointer set")
+
+        checkpoint = self.checkpointer.get(config)
+        checkpoint = checkpoint or empty_checkpoint()
+        with ChannelsManager(self.channels, checkpoint) as channels:
+            next_tasks = _prepare_next_tasks(
+                checkpoint, self.nodes, channels, update_seen=False
+            )
+            return StateSnapshot(
+                {
+                    k: _read_channel(channels, k)
+                    for k in channels
+                    if k not in [k.value for k in ReservedChannels]
+                },
+                tuple(name for _, _, name in next_tasks),
+            )
+
+    async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
+        if not self.checkpointer:
+            raise ValueError("No checkpointer set")
+
+        checkpoint = await self.checkpointer.aget(config)
+        checkpoint = checkpoint or empty_checkpoint()
+        async with AsyncChannelsManager(self.channels, checkpoint) as channels:
+            next_tasks = _prepare_next_tasks(
+                checkpoint, self.nodes, channels, update_seen=False
+            )
+            return StateSnapshot(
+                {
+                    k: _read_channel(channels, k)
+                    for k in channels
+                    if k not in [k.value for k in ReservedChannels]
+                },
+                tuple(name for _, _, name in next_tasks),
+            )
+
+    def update_state(self, config: RunnableConfig, values: dict[str, Any]) -> None:
+        if not self.checkpointer:
+            raise ValueError("No checkpointer set")
+
+        checkpoint = self.checkpointer.get(config)
+        checkpoint = checkpoint or empty_checkpoint()
+        with ChannelsManager(self.channels, checkpoint) as channels:
+            for k, v in values.items():
+                channels[k].update([v])
+                checkpoint["channel_versions"][k] += 1
+            self.checkpointer.put(config, create_checkpoint(checkpoint, channels))
+
+    async def aupdate_state(
+        self, config: RunnableConfig, values: dict[str, Any]
+    ) -> None:
+        if not self.checkpointer:
+            raise ValueError("No checkpointer set")
+
+        checkpoint = await self.checkpointer.aget(config)
+        checkpoint = checkpoint or empty_checkpoint()
+        async with AsyncChannelsManager(self.channels, checkpoint) as channels:
+            for k, v in values.items():
+                channels[k].update([v])
+                checkpoint["channel_versions"][k] += 1
+            await self.checkpointer.aput(
+                config, create_checkpoint(checkpoint, channels)
+            )
+
     def _transform(
         self,
         input: Iterator[Union[dict[str, Any], Any]],
@@ -768,7 +844,7 @@ def _apply_writes_from_view(
         if value == _read_channel(channels, chan):
             continue
 
-        assert isinstance(channels[chan], LastValue), (
+        assert isinstance(channels[chan], (LastValue, EphemeralValue, AnyValue)), (
             f"Can't modify channel {chan} of type "
             f"{channels[chan].__class__.__name__}"
         )
@@ -780,6 +856,7 @@ def _prepare_next_tasks(
     checkpoint: Checkpoint,
     processes: Mapping[str, Union[ChannelInvoke, ChannelBatch]],
     channels: Mapping[str, BaseChannel],
+    update_seen: bool = True,
 ) -> list[tuple[Runnable, Any, str]]:
     tasks: list[tuple[Runnable, Any, str]] = []
     # Check if any processes should be run in next step
@@ -814,12 +891,13 @@ def _prepare_next_tasks(
                     val = val[None]
 
                 # update seen versions
-                seen.update(
-                    {
-                        chan: checkpoint["channel_versions"][chan]
-                        for chan in proc.triggers
-                    }
-                )
+                if update_seen:
+                    seen.update(
+                        {
+                            chan: checkpoint["channel_versions"][chan]
+                            for chan in proc.triggers
+                        }
+                    )
 
                 # skip if condition is not met
                 if proc.when is None or proc.when(val):
@@ -836,7 +914,8 @@ def _prepare_next_tasks(
                     val = [{proc.key: v} for v in val]
 
                 tasks.append((proc, val, name))
-                seen[proc.channel] = checkpoint["channel_versions"][proc.channel]
+                if update_seen:
+                    seen[proc.channel] = checkpoint["channel_versions"][proc.channel]
 
     return tasks
 

pyproject.toml

@@ -55,6 +55,12 @@ exclude = ["notebooks", "examples", "example_data"]
 [tool.coverage.run]
 omit = ["tests/*"]
 
+[tool.pytest-watcher]
+now = true
+delay = 0.1
+runner_args = ["-x", "--ff", "-vv", "--snapshot-update"]
+patterns = ["*.py"]
+
 [build-system]
 requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"

tests/test_pregel.py

@@ -26,7 +26,7 @@ from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
 )
 from langgraph.prebuilt.tool_executor import ToolExecutor
-from langgraph.pregel import Channel, GraphRecursionError, Pregel
+from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
 from langgraph.pregel.reserved import ReservedChannels
 
 
@@ -282,6 +282,22 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
     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
+
+    # inbox == 21
+    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}}, {"inbox": 25})
+    assert app.invoke(None, {"configurable": {"thread_id": 2}}) == 26
+
+    # no pending tasks
+    snapshot = app.get_state({"configurable": {"thread_id": 2}})
+    assert snapshot.next == ()
+
 
 def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
@@ -761,6 +777,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
         ),
     }
 
+    # deepcopy because the nodes mutate the data
     assert [deepcopy(c) for c in app.stream({"input": "what is weather in sf"})] == [
         {
             "agent": {
@@ -882,6 +899,151 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
         },
     ]
 
+    # test state get/update methods
+
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaver(), interrupt_after=["agent"]
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert [
+        c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
+    ] == [
+        {
+            "agent": {
+                "input": "what is weather in sf",
+                "agent_outcome": AgentAction(
+                    tool="search_api", tool_input="query", log="tool:search_api:query"
+                ),
+            }
+        }
+    ]
+
+    assert app_w_interrupt.get_state(config) == StateSnapshot(
+        values={
+            "agent": {
+                "input": "what is weather in sf",
+                "agent_outcome": AgentAction(
+                    tool="search_api", tool_input="query", log="tool:search_api:query"
+                ),
+            },
+            "tools": None,
+        },
+        next=("agent:edges",),
+    )
+
+    app_w_interrupt.update_state(
+        config,
+        {
+            "agent": {
+                "agent_outcome": AgentAction(
+                    tool="search_api",
+                    tool_input="query",
+                    log="tool:search_api:a different query",
+                ),
+                "input": "what is weather in sf",
+            },
+        },
+    )
+
+    assert app_w_interrupt.get_state(config) == StateSnapshot(
+        values={
+            "agent": {
+                "agent_outcome": AgentAction(
+                    tool="search_api",
+                    tool_input="query",
+                    log="tool:search_api:a different query",
+                ),
+                "input": "what is weather in sf",
+            },
+            "tools": None,
+        },
+        next=("agent:edges",),
+    )
+
+    assert [c for c in app_w_interrupt.stream(None, config)] == [
+        {
+            "tools": {
+                "input": "what is weather in sf",
+                "intermediate_steps": [
+                    (
+                        AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:a different query",
+                        ),
+                        "result for query",
+                    )
+                ],
+            }
+        },
+        {
+            "agent": {
+                "input": "what is weather in sf",
+                "intermediate_steps": [
+                    (
+                        AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:a different query",
+                        ),
+                        "result for query",
+                    )
+                ],
+                "agent_outcome": AgentAction(
+                    tool="search_api",
+                    tool_input="another",
+                    log="tool:search_api:another",
+                ),
+            }
+        },
+    ]
+
+    app_w_interrupt.update_state(
+        config,
+        {
+            "agent": {
+                "input": "what is weather in sf",
+                "intermediate_steps": [
+                    (
+                        AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:a different query",
+                        ),
+                        "result for query",
+                    )
+                ],
+                "agent_outcome": AgentFinish(
+                    return_values={"answer": "a really nice answer"},
+                    log="finish:a really nice answer",
+                ),
+            }
+        },
+    )
+
+    assert [c for c in app_w_interrupt.stream(None, config)] == [
+        {
+            "__end__": {
+                "input": "what is weather in sf",
+                "agent_outcome": AgentFinish(
+                    return_values={"answer": "a really nice answer"},
+                    log="finish:a really nice answer",
+                ),
+                "intermediate_steps": [
+                    (
+                        AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:a different query",
+                        ),
+                        "result for query",
+                    )
+                ],
+            }
+        }
+    ]
+
 
 def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
     from langchain.llms.fake import FakeStreamingListLLM
@@ -1073,6 +1235,116 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
         },
     ]
 
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaver(), interrupt_after=["agent"]
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert [
+        c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
+    ] == [
+        {
+            "agent": {
+                "agent_outcome": AgentAction(
+                    tool="search_api", tool_input="query", log="tool:search_api:query"
+                ),
+            }
+        }
+    ]
+
+    assert app_w_interrupt.get_state(config) == StateSnapshot(
+        values={
+            "input": "what is weather in sf",
+            "agent_outcome": AgentAction(
+                tool="search_api", tool_input="query", log="tool:search_api:query"
+            ),
+            "intermediate_steps": [],
+        },
+        next=("agent:edges",),
+    )
+
+    app_w_interrupt.update_state(
+        config,
+        {
+            "agent_outcome": AgentAction(
+                tool="search_api",
+                tool_input="query",
+                log="tool:search_api:a different query",
+            )
+        },
+    )
+
+    assert app_w_interrupt.get_state(config) == StateSnapshot(
+        values={
+            "input": "what is weather in sf",
+            "agent_outcome": AgentAction(
+                tool="search_api",
+                tool_input="query",
+                log="tool:search_api:a different query",
+            ),
+            "intermediate_steps": [],
+        },
+        next=("agent:edges",),
+    )
+
+    assert [c for c in app_w_interrupt.stream(None, config)] == [
+        {
+            "tools": {
+                "intermediate_steps": [
+                    (
+                        AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:a different query",
+                        ),
+                        "result for query",
+                    )
+                ],
+            }
+        },
+        {
+            "agent": {
+                "agent_outcome": AgentAction(
+                    tool="search_api",
+                    tool_input="another",
+                    log="tool:search_api:another",
+                ),
+            }
+        },
+    ]
+
+    app_w_interrupt.update_state(
+        config,
+        {
+            "agent_outcome": AgentFinish(
+                return_values={"answer": "a really nice answer"},
+                log="finish:a really nice answer",
+            )
+        },
+    )
+
+    assert [c for c in app_w_interrupt.stream(None, config)] == [
+        {
+            "__end__": {
+                "input": "what is weather in sf",
+                "agent_outcome": AgentFinish(
+                    return_values={"answer": "a really nice answer"},
+                    log="finish:a really nice answer",
+                ),
+                "intermediate_steps": [
+                    (
+                        AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:a different query",
+                        ),
+                        "result for query",
+                    )
+                ],
+            }
+        }
+    ]
+
 
 def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
     def left(data: str) -> str:
@@ -1749,6 +2021,42 @@ def test_message_graph(snapshot: SnapshotAssertion) -> None:
         },
     ]
 
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaver(), interrupt_after=["agent"]
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert [
+        c
+        for c in app_w_interrupt.stream(
+            HumanMessage(content="what is weather in sf"), config
+        )
+    ] == [
+        {
+            "agent": AIMessage(
+                content="",
+                additional_kwargs={
+                    "function_call": {"name": "search_api", "arguments": '"query"'}
+                },
+            )
+        }
+    ]
+
+    assert app_w_interrupt.get_state(config) == StateSnapshot(
+        values=[
+            HumanMessage(content="what is weather in sf"),
+            AIMessage(
+                content="",
+                additional_kwargs={
+                    "function_call": {"name": "search_api", "arguments": '"query"'}
+                },
+            ),
+        ],
+        next=("agent:edges",),
+    )
+
+    # TODO use update_state once we have message ids
+
 
 def test_in_one_fan_out_out_one_graph_state() -> None:
     def sorted_add(x: list[str], y: list[str]) -> list[str]:

tests/test_pregel_async.py

@@ -31,7 +31,7 @@ from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
 )
 from langgraph.prebuilt.tool_executor import ToolExecutor
-from langgraph.pregel import Channel, GraphRecursionError, Pregel
+from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
 from langgraph.pregel.reserved import ReservedChannels
 
 
@@ -289,6 +289,22 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N
     assert await app.ainvoke(3, {"configurable": {"thread_id": 1}}) is None
     assert await app.ainvoke(None, {"configurable": {"thread_id": 1}}) == 5
 
+    # start execution again, stopping at inbox
+    assert await app.ainvoke(20, {"configurable": {"thread_id": 2}}) is None
+
+    # inbox == 21
+    snapshot = await app.aget_state({"configurable": {"thread_id": 2}})
+    assert snapshot.values["inbox"] == 21
+    assert snapshot.next == ("two",)
+
+    # update the state, resume
+    await app.aupdate_state({"configurable": {"thread_id": 2}}, {"inbox": 25})
+    assert await app.ainvoke(None, {"configurable": {"thread_id": 2}}) == 26
+
+    # no pending tasks
+    snapshot = await app.aget_state({"configurable": {"thread_id": 2}})
+    assert snapshot.next == ()
+
 
 async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
@@ -798,6 +814,7 @@ async def test_conditional_graph() -> None:
         ),
     }
 
+    # deepcopy because the nodes mutate the data
     assert [
         deepcopy(c) async for c in app.astream({"input": "what is weather in sf"})
     ] == [
@@ -927,6 +944,154 @@ async def test_conditional_graph() -> None:
     # Check that agent (one of the nodes) has its output streamed to the logs
     assert "/logs/agent/streamed_output/-" in patch_paths
 
+    # test state get/update methods
+
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaver(), interrupt_after=["agent"]
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert [
+        c
+        async for c in app_w_interrupt.astream(
+            {"input": "what is weather in sf"}, config
+        )
+    ] == [
+        {
+            "agent": {
+                "input": "what is weather in sf",
+                "agent_outcome": AgentAction(
+                    tool="search_api", tool_input="query", log="tool:search_api:query"
+                ),
+            }
+        }
+    ]
+
+    assert await app_w_interrupt.aget_state(config) == StateSnapshot(
+        values={
+            "agent": {
+                "input": "what is weather in sf",
+                "agent_outcome": AgentAction(
+                    tool="search_api", tool_input="query", log="tool:search_api:query"
+                ),
+            },
+            "tools": None,
+        },
+        next=("agent:edges",),
+    )
+
+    await app_w_interrupt.aupdate_state(
+        config,
+        {
+            "agent": {
+                "agent_outcome": AgentAction(
+                    tool="search_api",
+                    tool_input="query",
+                    log="tool:search_api:a different query",
+                ),
+                "input": "what is weather in sf",
+            },
+        },
+    )
+
+    assert await app_w_interrupt.aget_state(config) == StateSnapshot(
+        values={
+            "agent": {
+                "agent_outcome": AgentAction(
+                    tool="search_api",
+                    tool_input="query",
+                    log="tool:search_api:a different query",
+                ),
+                "input": "what is weather in sf",
+            },
+            "tools": None,
+        },
+        next=("agent:edges",),
+    )
+
+    assert [c async for c in app_w_interrupt.astream(None, config)] == [
+        {
+            "tools": {
+                "input": "what is weather in sf",
+                "intermediate_steps": [
+                    (
+                        AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:a different query",
+                        ),
+                        "result for query",
+                    )
+                ],
+            }
+        },
+        {
+            "agent": {
+                "input": "what is weather in sf",
+                "intermediate_steps": [
+                    (
+                        AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:a different query",
+                        ),
+                        "result for query",
+                    )
+                ],
+                "agent_outcome": AgentAction(
+                    tool="search_api",
+                    tool_input="another",
+                    log="tool:search_api:another",
+                ),
+            }
+        },
+    ]
+
+    await app_w_interrupt.aupdate_state(
+        config,
+        {
+            "agent": {
+                "input": "what is weather in sf",
+                "intermediate_steps": [
+                    (
+                        AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:a different query",
+                        ),
+                        "result for query",
+                    )
+                ],
+                "agent_outcome": AgentFinish(
+                    return_values={"answer": "a really nice answer"},
+                    log="finish:a really nice answer",
+                ),
+            }
+        },
+    )
+
+    assert [c async for c in app_w_interrupt.astream(None, config)] == [
+        {
+            "__end__": {
+                "input": "what is weather in sf",
+                "agent_outcome": AgentFinish(
+                    return_values={"answer": "a really nice answer"},
+                    log="finish:a really nice answer",
+                ),
+                "intermediate_steps": [
+                    (
+                        AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:a different query",
+                        ),
+                        "result for query",
+                    )
+                ],
+            }
+        }
+    ]
+
 
 async def test_conditional_graph_state() -> None:
     from langchain.llms.fake import FakeStreamingListLLM
@@ -1113,6 +1278,119 @@ async def test_conditional_graph_state() -> None:
         },
     ]
 
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaver(), interrupt_after=["agent"]
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert [
+        c
+        async for c in app_w_interrupt.astream(
+            {"input": "what is weather in sf"}, config
+        )
+    ] == [
+        {
+            "agent": {
+                "agent_outcome": AgentAction(
+                    tool="search_api", tool_input="query", log="tool:search_api:query"
+                ),
+            }
+        }
+    ]
+
+    assert await app_w_interrupt.aget_state(config) == StateSnapshot(
+        values={
+            "input": "what is weather in sf",
+            "agent_outcome": AgentAction(
+                tool="search_api", tool_input="query", log="tool:search_api:query"
+            ),
+            "intermediate_steps": [],
+        },
+        next=("agent:edges",),
+    )
+
+    await app_w_interrupt.aupdate_state(
+        config,
+        {
+            "agent_outcome": AgentAction(
+                tool="search_api",
+                tool_input="query",
+                log="tool:search_api:a different query",
+            )
+        },
+    )
+
+    assert await app_w_interrupt.aget_state(config) == StateSnapshot(
+        values={
+            "input": "what is weather in sf",
+            "agent_outcome": AgentAction(
+                tool="search_api",
+                tool_input="query",
+                log="tool:search_api:a different query",
+            ),
+            "intermediate_steps": [],
+        },
+        next=("agent:edges",),
+    )
+
+    assert [c async for c in app_w_interrupt.astream(None, config)] == [
+        {
+            "tools": {
+                "intermediate_steps": [
+                    (
+                        AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:a different query",
+                        ),
+                        "result for query",
+                    )
+                ],
+            }
+        },
+        {
+            "agent": {
+                "agent_outcome": AgentAction(
+                    tool="search_api",
+                    tool_input="another",
+                    log="tool:search_api:another",
+                ),
+            }
+        },
+    ]
+
+    await app_w_interrupt.aupdate_state(
+        config,
+        {
+            "agent_outcome": AgentFinish(
+                return_values={"answer": "a really nice answer"},
+                log="finish:a really nice answer",
+            )
+        },
+    )
+
+    assert [c async for c in app_w_interrupt.astream(None, config)] == [
+        {
+            "__end__": {
+                "input": "what is weather in sf",
+                "agent_outcome": AgentFinish(
+                    return_values={"answer": "a really nice answer"},
+                    log="finish:a really nice answer",
+                ),
+                "intermediate_steps": [
+                    (
+                        AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:a different query",
+                        ),
+                        "result for query",
+                    )
+                ],
+            }
+        }
+    ]
+
 
 async def test_conditional_entrypoint_graph() -> None:
     async def left(data: str) -> str:
@@ -1772,6 +2050,42 @@ async def test_message_graph() -> None:
         },
     ]
 
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaver(), interrupt_after=["agent"]
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert [
+        c
+        async for c in app_w_interrupt.astream(
+            HumanMessage(content="what is weather in sf"), config
+        )
+    ] == [
+        {
+            "agent": AIMessage(
+                content="",
+                additional_kwargs={
+                    "function_call": {"name": "search_api", "arguments": '"query"'}
+                },
+            )
+        }
+    ]
+
+    assert await app_w_interrupt.aget_state(config) == StateSnapshot(
+        values=[
+            HumanMessage(content="what is weather in sf"),
+            AIMessage(
+                content="",
+                additional_kwargs={
+                    "function_call": {"name": "search_api", "arguments": '"query"'}
+                },
+            ),
+        ],
+        next=("agent:edges",),
+    )
+
+    # TODO use update_state once we have message ids
+
 
 async def test_in_one_fan_out_out_one_graph_state() -> None:
     def sorted_add(x: list[str], y: list[str]) -> list[str]: