fix: include entire state snapshot in values stream chunk alongside interrupt

Author: casparb

libs/langgraph/langgraph/pregel/_loop.py

@@ -766,7 +766,7 @@ def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
             ]
             self.checkpoint["channel_values"][TASKS] = sanitized_tasks
         # bail if no checkpointer
-        
+
         if do_checkpoint and self._checkpointer_put_after_previous is not None:
             self.prev_checkpoint_config = (
                 self.checkpoint_config
@@ -938,9 +938,11 @@ def output_writes(
                     self._emit("updates", lambda: iter(interrupts))
                 if "values" in stream_modes:
                     current_values = read_channels(self.channels, self.output_keys)
+                    # self.output_keys is a sequence, stream chunk conntains entire state and interrupts
                     if isinstance(current_values, dict):
                         current_values[INTERRUPT] = interrupts[0][INTERRUPT]
                         self._emit("values", lambda: iter([current_values]))
+                    # self.output_keys is a string, stream chunk contains only interrupts
                     else:
                         self._emit("values", lambda: iter(interrupts))
             elif writes[0][0] != ERROR:

libs/langgraph/tests/test_pregel.py

@@ -830,10 +830,9 @@ def raise_if_above_10(input: int) -> int:
     assert checkpoint["channel_values"].get("total") == 5
 
 
-def test_pending_writes_resume(
-    sync_checkpointer: BaseCheckpointSaver
-) -> None:
+def test_pending_writes_resume(sync_checkpointer: BaseCheckpointSaver) -> None:
     durability = "exit"
+
     class State(TypedDict):
         value: Annotated[int, operator.add]
 
@@ -8424,8 +8423,8 @@ def human_node_2(state: State):
     }
 
 
-def test_interrupt_stream_mode_values():
-    """Test that interrupts are surfaced when steam_mode='values'"""
+def test_interrupt_stream_mode_values(sync_checkpointer: BaseCheckpointSaver):
+    """Test that interrupts are surfaced on 'values' stream mode"""
 
     class State(TypedDict):
         robot_input: str
@@ -8443,19 +8442,36 @@ def human_input_node(state: State) -> Command:
     builder.add_node(human_input_node)
     builder.add_edge(START, "robot_input_node")
     builder.add_edge("robot_input_node", "human_input_node")
-    app = builder.compile()
+    app = builder.compile(checkpointer=sync_checkpointer)
+    config = {"configurable": {"thread_id": str(uuid.uuid4())}}
 
-    result = [*app.stream(State(), stream_mode=["updates", "values"])]
-    print("PY_DEBUG: result", result)
+    result = [*app.stream(State(), config, stream_mode=["updates", "values"])]
     assert len(result) == 4
-   
-
-    assert result = [
-        ('updates', {'robot_input_node': {'robot_input': 'beep boop i am a robot'}}), 
-        ('values', {'robot_input': 'beep boop i am a robot'}), 
-        ('updates', {'__interrupt__': (Interrupt(value='interrupt', id=AnyStr()),)}), 
-        ('values', {'robot_input': 'beep boop i am a robot', '__interrupt__': (Interrupt(value='interrupt', id=AnyStr()),)})]
-    assert result[1] == {"robot_input": "beep boop i am a robot", "__interrupt__": (Interrupt(value="interrupt", id=AnyStr()),)}
+    assert result == [
+        ("updates", {"robot_input_node": {"robot_input": "beep boop i am a robot"}}),
+        ("values", {"robot_input": "beep boop i am a robot"}),
+        ("updates", {"__interrupt__": (Interrupt(value="interrupt", id=AnyStr()),)}),
+        (
+            "values",
+            {
+                "robot_input": "beep boop i am a robot",
+                "__interrupt__": (Interrupt(value="interrupt", id=AnyStr()),),
+            },
+        ),
+    ]
+    resume_result = [
+        *app.stream(
+            Command(resume="i am a human"), config, stream_mode=["updates", "values"]
+        )
+    ]
+    assert resume_result == [
+        ("values", {"robot_input": "beep boop i am a robot"}),
+        ("updates", {"human_input_node": {"human_input": "i am a human"}}),
+        (
+            "values",
+            {"robot_input": "beep boop i am a robot", "human_input": "i am a human"},
+        ),
+    ]
 
 
 def test_supersteps_populate_task_results(

libs/langgraph/tests/test_pregel_async.py

@@ -9196,6 +9196,63 @@ async def consumer() -> None:
     assert all(t.done() for t in recorded_tasks)
 
 
+async def test_interrupt_stream_mode_values(async_checkpointer: BaseCheckpointSaver):
+    """Test that interrupts are surfaced on 'values' stream mode"""
+
+    class State(TypedDict):
+        robot_input: str
+        human_input: str
+
+    def robot_input_node(state: State) -> State:
+        return {"robot_input": "beep boop i am a robot"}
+
+    def human_input_node(state: State) -> Command:
+        human_input = interrupt("interrupt")
+        return Command(update={"human_input": human_input})
+
+    builder = StateGraph(State)
+    builder.add_node(robot_input_node)
+    builder.add_node(human_input_node)
+    builder.add_edge(START, "robot_input_node")
+    builder.add_edge("robot_input_node", "human_input_node")
+    app = builder.compile(checkpointer=async_checkpointer)
+    config = {"configurable": {"thread_id": str(uuid.uuid4())}}
+
+    result = [
+        (mode, e)
+        async for mode, e in app.astream(
+            State(), config, stream_mode=["updates", "values"]
+        )
+    ]
+    assert len(result) == 4
+    assert result == [
+        ("updates", {"robot_input_node": {"robot_input": "beep boop i am a robot"}}),
+        ("values", {"robot_input": "beep boop i am a robot"}),
+        ("updates", {"__interrupt__": (Interrupt(value="interrupt", id=AnyStr()),)}),
+        (
+            "values",
+            {
+                "robot_input": "beep boop i am a robot",
+                "__interrupt__": (Interrupt(value="interrupt", id=AnyStr()),),
+            },
+        ),
+    ]
+    resume_result = [
+        (mode, e)
+        async for mode, e in app.astream(
+            Command(resume="i am a human"), config, stream_mode=["updates", "values"]
+        )
+    ]
+    assert resume_result == [
+        ("values", {"robot_input": "beep boop i am a robot"}),
+        ("updates", {"human_input_node": {"human_input": "i am a human"}}),
+        (
+            "values",
+            {"robot_input": "beep boop i am a robot", "human_input": "i am a human"},
+        ),
+    ]
+
+
 async def test_supersteps_populate_task_results(
     async_checkpointer: BaseCheckpointSaver,
 ) -> None: