Merge pull request #367 from langchain-ai/nc/30apr/interrupt-before-all

Implement interrupt_before=* and interrupt_after=*

langgraph/graph/graph.py

@@ -28,6 +28,7 @@ from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.constants import TAG_HIDDEN
 from langgraph.pregel import Channel, Pregel
 from langgraph.pregel.read import PregelNode
+from langgraph.pregel.types import All
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.utils import DrawableGraph, RunnableCallable, coerce_to_runnable
 
@@ -295,8 +296,8 @@ class Graph:
     def compile(
         self,
         checkpointer: Optional[BaseCheckpointSaver] = None,
-        interrupt_before: Optional[Sequence[str]] = None,
-        interrupt_after: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: bool = False,
     ) -> "CompiledGraph":
         # assign default values
@@ -304,11 +305,16 @@ class Graph:
         interrupt_after = interrupt_after or []
 
         # validate the graph
-        self.validate(interrupt=interrupt_before + interrupt_after)
+        self.validate(
+            interrupt=(interrupt_before if interrupt_before != "*" else [])
+            + interrupt_after
+            if interrupt_after != "*"
+            else []
+        )
 
         # create empty compiled graph
         compiled = CompiledGraph(
-            graph=self,
+            builder=self,
             nodes={},
             channels={START: EphemeralValue(Any), END: EphemeralValue(Any)},
             input_channels=START,
@@ -338,7 +344,7 @@ class Graph:
 
 
 class CompiledGraph(Pregel):
-    graph: Graph
+    builder: Graph
 
     def attach_node(self, key: str, node: Runnable) -> None:
         self.channels[key] = EphemeralValue(Any)
@@ -400,7 +406,7 @@ class CompiledGraph(Pregel):
             END: graph.add_node(self.get_output_schema(config), END)
         }
 
-        for key, node in self.graph.nodes.items():
+        for key, node in self.builder.nodes.items():
             if xray:
                 subgraph = (
                     node.get_graph(
@@ -424,11 +430,11 @@ class CompiledGraph(Pregel):
                 n = graph.add_node(node, key)
                 start_nodes[key] = n
                 end_nodes[key] = n
-        for start, end in sorted(self.graph._all_edges):
+        for start, end in sorted(self.builder._all_edges):
             graph.add_edge(start_nodes[start], end_nodes[end])
-        for start, branches in self.graph.branches.items():
+        for start, branches in self.builder.branches.items():
             default_ends = {
-                **{k: k for k in self.graph.nodes if k != start},
+                **{k: k for k in self.builder.nodes if k != start},
                 END: END,
             }
             for _, branch in branches.items():

langgraph/graph/state.py

@@ -16,6 +16,7 @@ from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.constants import TAG_HIDDEN
 from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph
 from langgraph.pregel.read import ChannelRead, PregelNode
+from langgraph.pregel.types import All
 from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
 from langgraph.utils import RunnableCallable
 
@@ -100,8 +101,8 @@ class StateGraph(Graph):
     def compile(
         self,
         checkpointer: Optional[BaseCheckpointSaver] = None,
-        interrupt_before: Optional[Sequence[str]] = None,
-        interrupt_after: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: bool = False,
     ) -> CompiledGraph:
         """Compiles the state graph into a `CompiledGraph` object.
@@ -120,14 +121,19 @@ class StateGraph(Graph):
         interrupt_after = interrupt_after or []
 
         # validate the graph
-        self.validate(interrupt=interrupt_before + interrupt_after)
+        self.validate(
+            interrupt=(interrupt_before if interrupt_before != "*" else [])
+            + interrupt_after
+            if interrupt_after != "*"
+            else []
+        )
 
         # prepare output channels
         state_keys = list(self.channels)
         output_channels = state_keys[0] if state_keys == ["__root__"] else state_keys
 
         compiled = CompiledStateGraph(
-            graph=self,
+            builder=self,
             nodes={},
             channels={**self.channels, START: EphemeralValue(self.schema)},
             input_channels=START,
@@ -159,7 +165,7 @@ class StateGraph(Graph):
 
 
 class CompiledStateGraph(CompiledGraph):
-    graph: StateGraph
+    builder: StateGraph
 
     def attach_node(self, key: str, node: Optional[Runnable]) -> None:
         def _get_state_key(input: dict, config: RunnableConfig, *, key: str) -> Any:
@@ -170,7 +176,7 @@ class CompiledStateGraph(CompiledGraph):
             else:
                 return input.get(key, SKIP_WRITE)
 
-        state_keys = list(self.graph.channels)
+        state_keys = list(self.builder.channels)
         # state updaters
         state_write_entries = [
             (
@@ -210,7 +216,7 @@ class CompiledStateGraph(CompiledGraph):
                 mapper=(
                     None
                     if state_keys == ["__root__"]
-                    else partial(_coerce_state, self.graph.schema)
+                    else partial(_coerce_state, self.builder.schema)
                 ),
                 writers=[
                     # publish to this channel and state keys
@@ -265,13 +271,13 @@ class CompiledStateGraph(CompiledGraph):
                 return ChannelWrite(writes, tags=[TAG_HIDDEN])
 
         # attach branch publisher
-        self.nodes[start] |= branch.run(branch_writer, _get_state_reader(self.graph))
+        self.nodes[start] |= branch.run(branch_writer, _get_state_reader(self.builder))
 
         # attach branch subscribers
         ends = (
             branch.ends.values()
             if branch.ends
-            else [node for node in self.graph.nodes if node != branch.then]
+            else [node for node in self.builder.nodes if node != branch.then]
         )
         for end in ends:
             if end != END:

langgraph/pregel/__init__.py

@@ -65,6 +65,7 @@ from langgraph.constants import (
     CONFIG_KEY_READ,
     CONFIG_KEY_SEND,
     INTERRUPT,
+    TAG_HIDDEN,
 )
 from langgraph.pregel.debug import (
     print_step_checkpoint,
@@ -81,6 +82,7 @@ from langgraph.pregel.io import (
 from langgraph.pregel.log import logger
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.types import (
+    All,
     PregelExecutableTask,
     PregelTaskDescription,
     StateSnapshot,
@@ -194,9 +196,9 @@ class Pregel(
     stream_channels: Optional[Union[str, Sequence[str]]] = None
     """Channels to stream, defaults to all channels not in reserved channels"""
 
-    interrupt_after_nodes: Sequence[str] = Field(default_factory=list)
+    interrupt_after_nodes: Union[All, Sequence[str]] = Field(default_factory=list)
 
-    interrupt_before_nodes: Sequence[str] = Field(default_factory=list)
+    interrupt_before_nodes: Union[All, Sequence[str]] = Field(default_factory=list)
 
     input_channels: Union[str, Sequence[str]]
 
@@ -518,8 +520,8 @@ class Pregel(
         stream_mode: Optional[StreamMode] = None,
         input_keys: Optional[Union[str, Sequence[str]]] = None,
         output_keys: Optional[Union[str, Sequence[str]]] = None,
-        interrupt_before: Optional[Sequence[str]] = None,
-        interrupt_after: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: Optional[bool] = None,
     ) -> tuple[
         bool,
@@ -565,8 +567,8 @@ class Pregel(
         stream_mode: Optional[StreamMode] = None,
         output_keys: Optional[Union[str, Sequence[str]]] = None,
         input_keys: Optional[Union[str, Sequence[str]]] = None,
-        interrupt_before: Optional[Sequence[str]] = None,
-        interrupt_after: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: Optional[bool] = None,
     ) -> Iterator[Union[dict[str, Any], Any]]:
         """Stream graph steps for a single input."""
@@ -775,8 +777,8 @@ class Pregel(
         stream_mode: Optional[StreamMode] = None,
         output_keys: Optional[Union[str, Sequence[str]]] = None,
         input_keys: Optional[Union[str, Sequence[str]]] = None,
-        interrupt_before: Optional[Sequence[str]] = None,
-        interrupt_after: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: Optional[bool] = None,
     ) -> AsyncIterator[Union[dict[str, Any], Any]]:
         config = ensure_config(config)
@@ -1002,8 +1004,8 @@ class Pregel(
         stream_mode: StreamMode = "values",
         output_keys: Optional[Union[str, Sequence[str]]] = None,
         input_keys: Optional[Union[str, Sequence[str]]] = None,
-        interrupt_before_nodes: Optional[Sequence[str]] = None,
-        interrupt_after_nodes: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: Optional[bool] = None,
         **kwargs: Any,
     ) -> Union[dict[str, Any], Any]:
@@ -1015,8 +1017,8 @@ class Pregel(
             stream_mode: Optional[str]. The stream mode for the graph run. Default is "values".
             output_keys: Optional. The output keys to retrieve from the graph run.
             input_keys: Optional. The input keys to provide for the graph run.
-            interrupt_before_nodes: Optional. The nodes to interrupt the graph run before.
-            interrupt_after_nodes: Optional. The nodes to interrupt the graph run after.
+            interrupt_before: Optional. The nodes to interrupt the graph run before.
+            interrupt_after: Optional. The nodes to interrupt the graph run after.
             debug: Optional. Enable debug mode for the graph run.
             **kwargs: Additional keyword arguments to pass to the graph run.
 
@@ -1035,8 +1037,8 @@ class Pregel(
             stream_mode=stream_mode,
             output_keys=output_keys,
             input_keys=input_keys,
-            interrupt_before=interrupt_before_nodes,
-            interrupt_after=interrupt_after_nodes,
+            interrupt_before=interrupt_before,
+            interrupt_after=interrupt_after,
             debug=debug,
             **kwargs,
         ):
@@ -1057,8 +1059,8 @@ class Pregel(
         stream_mode: StreamMode = "values",
         output_keys: Optional[Union[str, Sequence[str]]] = None,
         input_keys: Optional[Union[str, Sequence[str]]] = None,
-        interrupt_before_nodes: Optional[Sequence[str]] = None,
-        interrupt_after_nodes: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Union[All, Sequence[str]]] = None,
+        interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: Optional[bool] = None,
         **kwargs: Any,
     ) -> Union[dict[str, Any], Any]:
@@ -1070,8 +1072,8 @@ class Pregel(
             stream_mode: Optional. The stream mode for the computation. Default is "values".
             output_keys: Optional. The output keys to include in the result. Default is None.
             input_keys: Optional. The input keys to include in the result. Default is None.
-            interrupt_before_nodes: Optional. The nodes to interrupt before. Default is None.
-            interrupt_after_nodes: Optional. The nodes to interrupt after. Default is None.
+            interrupt_before: Optional. The nodes to interrupt before. Default is None.
+            interrupt_after: Optional. The nodes to interrupt after. Default is None.
             debug: Optional. Whether to enable debug mode. Default is None.
             **kwargs: Additional keyword arguments.
 
@@ -1091,8 +1093,8 @@ class Pregel(
             stream_mode=stream_mode,
             output_keys=output_keys,
             input_keys=input_keys,
-            interrupt_before=interrupt_before_nodes,
-            interrupt_after=interrupt_after_nodes,
+            interrupt_before=interrupt_before,
+            interrupt_after=interrupt_after,
             debug=debug,
             **kwargs,
         ):
@@ -1132,7 +1134,7 @@ def _panic_or_proceed(
 
 def _should_interrupt(
     checkpoint: Checkpoint,
-    interrupt_nodes: Sequence[str],
+    interrupt_nodes: Union[All, Sequence[str]],
     snapshot_channels: Sequence[str],
     tasks: list[PregelExecutableTask],
 ) -> bool:
@@ -1145,7 +1147,15 @@ def _should_interrupt(
             for chan in snapshot_channels
         )
         # and any channel written to is in interrupt_nodes list
-        and any(node for node, _, _, _, _ in tasks if node in interrupt_nodes)
+        and any(
+            node
+            for node, _, _, _, config in tasks
+            if (
+                (not config or TAG_HIDDEN not in config.get("tags"))
+                if interrupt_nodes == "*"
+                else node in interrupt_nodes
+            )
+        )
     )
 
 

langgraph/pregel/types.py

@@ -1,5 +1,5 @@
 from collections import deque
-from typing import Any, NamedTuple, Optional, Union
+from typing import Any, Literal, NamedTuple, Optional, Union
 
 from langchain_core.runnables import Runnable, RunnableConfig
 
@@ -26,3 +26,6 @@ class StateSnapshot(NamedTuple):
     """Config used to fetch this snapshot"""
     parent_config: Optional[RunnableConfig] = None
     """Config used to fetch the parent snapshot, if any"""
+
+
+All = Literal["*"]

langgraph/pregel/validate.py

@@ -3,6 +3,7 @@ from typing import Mapping, Optional, Sequence, Union
 from langgraph.channels.base import BaseChannel
 from langgraph.constants import INTERRUPT
 from langgraph.pregel.read import PregelNode
+from langgraph.pregel.types import All
 
 
 def validate_graph(
@@ -11,8 +12,8 @@ def validate_graph(
     input_channels: Union[str, Sequence[str]],
     output_channels: Union[str, Sequence[str]],
     stream_channels: Optional[Union[str, Sequence[str]]],
-    interrupt_after_nodes: Sequence[str],
-    interrupt_before_nodes: Sequence[str],
+    interrupt_after_nodes: Union[All, Sequence[str]],
+    interrupt_before_nodes: Union[All, Sequence[str]],
 ) -> None:
     subscribed_channels = set[str]()
     for name, node in nodes.items():
@@ -59,12 +60,14 @@ def validate_graph(
         if chan not in channels:
             raise ValueError(f"Output channel '{chan}' not in 'channels'")
 
-    for node in interrupt_after_nodes:
-        if node not in nodes:
-            raise ValueError(f"Node {node} not in nodes")
-    for node in interrupt_before_nodes:
-        if node not in nodes:
-            raise ValueError(f"Node {node} not in nodes")
+    if interrupt_after_nodes != "*":
+        for node in interrupt_after_nodes:
+            if node not in nodes:
+                raise ValueError(f"Node {node} not in nodes")
+    if interrupt_before_nodes != "*":
+        for node in interrupt_before_nodes:
+            if node not in nodes:
+                raise ValueError(f"Node {node} not in nodes")
 
 
 def validate_keys(

langgraph/serde/jsonplus.py

@@ -38,7 +38,13 @@ class JsonPlusSerializer(SerializerProtocol):
         if isinstance(obj, Serializable):
             return obj.to_json()
         elif isinstance(obj, (BaseModel, LcBaseModel)):
-            return self._encode_constructor_args(obj.__class__, kwargs=obj.dict())
+            # prefer non-deprecated method if available
+            if hasattr(obj, "model_dump"):
+                return self._encode_constructor_args(
+                    obj.__class__, kwargs=obj.model_dump()
+                )
+            else:
+                return self._encode_constructor_args(obj.__class__, kwargs=obj.dict())
         elif isinstance(obj, UUID):
             return self._encode_constructor_args(UUID, args=[obj.hex])
         elif isinstance(obj, (set, frozenset)):

tests/__snapshots__/test_pregel.ambr

@@ -687,13 +687,13 @@
   
   '''
 # ---
-# name: test_conditional_graph_state[end_of_run]
+# name: test_conditional_state_graph[end_of_run]
   '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}}}}'
 # ---
-# name: test_conditional_graph_state[end_of_run].1
+# name: test_conditional_state_graph[end_of_run].1
   '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
 # ---
-# name: test_conditional_graph_state[end_of_run].2
+# name: test_conditional_state_graph[end_of_run].2
   '''
   {
     "nodes": [
@@ -758,7 +758,7 @@
   }
   '''
 # ---
-# name: test_conditional_graph_state[end_of_run].3
+# name: test_conditional_state_graph[end_of_run].3
   '''
   graph TD;
   	__start__ --> agent;
@@ -768,13 +768,13 @@
   
   '''
 # ---
-# name: test_conditional_graph_state[end_of_step]
+# name: test_conditional_state_graph[end_of_step]
   '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}}}}'
 # ---
-# name: test_conditional_graph_state[end_of_step].1
+# name: test_conditional_state_graph[end_of_step].1
   '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
 # ---
-# name: test_conditional_graph_state[end_of_step].2
+# name: test_conditional_state_graph[end_of_step].2
   '''
   {
     "nodes": [
@@ -839,7 +839,7 @@
   }
   '''
 # ---
-# name: test_conditional_graph_state[end_of_step].3
+# name: test_conditional_state_graph[end_of_step].3
   '''
   graph TD;
   	__start__ --> agent;

tests/test_pregel.py

@@ -1574,7 +1574,7 @@ def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
 @pytest.mark.parametrize(
     "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
 )
-def test_conditional_graph_state(
+def test_conditional_state_graph(
     snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt
 ) -> None:
     from langchain.llms.fake import FakeStreamingListLLM
@@ -1972,6 +1972,182 @@ def test_conditional_graph_state(
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
     )
 
+    # test w interrupt before all
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
+        interrupt_before="*",
+        debug=True,
+    )
+    config = {"configurable": {"thread_id": "3"}}
+    llm.i = 0  # reset the llm
+
+    assert [
+        c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
+    ] == []
+
+    assert app_w_interrupt.get_state(config) == StateSnapshot(
+        values={
+            "input": "what is weather in sf",
+            "intermediate_steps": [],
+        },
+        next=("agent",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
+    )
+
+    assert [c for c in app_w_interrupt.stream(None, 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=("tools",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
+    )
+
+    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:query",
+                        ),
+                        "result for 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": [
+                (
+                    AgentAction(
+                        tool="search_api",
+                        tool_input="query",
+                        log="tool:search_api:query",
+                    ),
+                    "result for query",
+                )
+            ],
+        },
+        next=("agent",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
+    )
+
+    assert [c for c in app_w_interrupt.stream(None, config)] == [
+        {
+            "agent": {
+                "agent_outcome": AgentAction(
+                    tool="search_api",
+                    tool_input="another",
+                    log="tool:search_api:another",
+                ),
+            }
+        },
+    ]
+
+    # test w interrupt after all
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
+        interrupt_after="*",
+    )
+    config = {"configurable": {"thread_id": "4"}}
+    llm.i = 0  # reset the llm
+
+    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=("tools",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
+    )
+
+    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:query",
+                        ),
+                        "result for 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": [
+                (
+                    AgentAction(
+                        tool="search_api",
+                        tool_input="query",
+                        log="tool:search_api:query",
+                    ),
+                    "result for query",
+                )
+            ],
+        },
+        next=("agent",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
+    )
+
+    assert [c for c in app_w_interrupt.stream(None, config)] == [
+        {
+            "agent": {
+                "agent_outcome": AgentAction(
+                    tool="search_api",
+                    tool_input="another",
+                    log="tool:search_api:another",
+                ),
+            }
+        },
+    ]
+
 
 def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None:
     class AgentState(TypedDict, total=False):