Fix for update_state mistakenly resetting barrier channels

- This introduces a new optional method BaseChannel.consume, which gets called when a channel triggers a node
- This is an alternative place to clean up channel state, in addition to the existing pattern of clearing state in the update() call for the next step
- Channels should implement consume() when they want to clean up state exactly if and only if the channel triggered a node
- Channels should clean up state in update() when they instead need to guarantee their value is available for a single step, irrespective of whether it was actually read

langgraph/channels/base.py

@@ -65,7 +65,8 @@ class BaseChannel(Generic[Value, Update, C], ABC):
     def update(self, values: Sequence[Update]) -> None:
         """Update the channel's value with the given sequence of updates.
         The order of the updates in the sequence is arbitrary.
-
+        This method is called by Pregel for all channels at the end of each step.
+        If there are no updates, it is called with an empty sequence.
         Raises InvalidUpdateError if the sequence of updates is invalid."""
 
     @abstractmethod
@@ -74,6 +75,13 @@ class BaseChannel(Generic[Value, Update, C], ABC):
 
         Raises EmptyChannelError if the channel is empty (never updated yet)."""
 
+    def consume(self) -> None:
+        """Mark the current value of the channel as consumed. By default, no-op.
+        This is called by Pregel before the start of the next step, for all
+        channels that triggered a node.
+        """
+        pass
+
 
 @contextmanager
 def ChannelsManager(

langgraph/channels/dynamic_barrier_value.py

@@ -60,11 +60,6 @@ class DynamicBarrierValue(
             pass
 
     def update(self, values: Sequence[Union[Value, WaitForNames]]) -> None:
-        # switch to "priming" state after reading
-        if self.seen == self.names:
-            self.seen = set()
-            self.names = None
-
         if wait_for_names := [v for v in values if isinstance(v, WaitForNames)]:
             if len(wait_for_names) > 1:
                 raise InvalidUpdateError(
@@ -83,3 +78,8 @@ class DynamicBarrierValue(
         if self.seen != self.names:
             raise EmptyChannelError()
         return None
+
+    def consume(self) -> None:
+        if self.seen == self.names:
+            self.seen = set()
+            self.names = None

langgraph/channels/named_barrier_value.py

@@ -42,8 +42,6 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
             pass
 
     def update(self, values: Sequence[Value]) -> None:
-        if self.seen == self.names:
-            self.seen = set()
         for value in values:
             if value in self.names:
                 self.seen.add(value)
@@ -54,3 +52,7 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
         if self.seen != self.names:
             raise EmptyChannelError()
         return None
+
+    def consume(self) -> None:
+        if self.seen == self.names:
+            self.seen = set()

langgraph/pregel/__init__.py

@@ -1685,6 +1685,8 @@ def _prepare_next_tasks(
             tasks.append(PregelTaskDescription(packet.node, packet.arg))
     if for_execution:
         checkpoint["pending_sends"].clear()
+    # Collect channels to consume
+    channels_to_consume = set()
     # Check if any processes should be run in next step
     # If so, prepare the values to be passed to them
     for name, proc in processes.items():
@@ -1698,6 +1700,7 @@ def _prepare_next_tasks(
             )
             and checkpoint["channel_versions"][chan] > seen[chan]
         ]:
+            channels_to_consume.update(triggers)
             try:
                 val = next(_proc_input(step, name, proc, managed, channels))
             except StopIteration:
@@ -1754,6 +1757,10 @@ def _prepare_next_tasks(
                     )
             else:
                 tasks.append(PregelTaskDescription(name, val))
+    # Consume all channels that were read
+    if for_execution:
+        for chan in channels_to_consume:
+            channels[chan].consume()
     return checkpoint, tasks
 
 
@@ -1763,7 +1770,6 @@ def _proc_input(
     proc: PregelNode,
     managed: ManagedValueMapping,
     channels: Mapping[str, BaseChannel],
-    catch: bool = False,
 ) -> Iterator[Any]:
     # If all trigger channels subscribed by this process are not empty
     # then invoke the process with the values of all non-empty channels
@@ -1771,7 +1777,9 @@ def _proc_input(
         try:
             val: dict = {
                 k: read_channel(
-                    channels, chan, catch=catch or chan not in proc.triggers
+                    channels,
+                    chan,
+                    catch=chan not in proc.triggers,
                 )
                 for k, chan in proc.channels.items()
                 if isinstance(chan, str)

tests/test_pregel.py

@@ -6224,6 +6224,52 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
             parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config,
         )
 
+    with SqliteSaver.from_conn_string(":memory:") as saver:
+        tool_two = tool_two_graph.compile(
+            checkpointer=saver, interrupt_before=["finish"]
+        )
+
+        thread1 = {"configurable": {"thread_id": "1"}}
+
+        # stop when about to enter node
+        assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == {
+            "my_key": "value prepared slow",
+            "market": "DE",
+        }
+        assert tool_two.get_state(thread1) == StateSnapshot(
+            values={
+                "my_key": "value prepared slow",
+                "market": "DE",
+            },
+            next=("finish",),
+            config=tool_two.checkpointer.get_tuple(thread1).config,
+            created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
+            metadata={
+                "source": "loop",
+                "step": 2,
+                "writes": {"tool_two_slow": {"my_key": " slow"}},
+            },
+            parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config,
+        )
+
+        # update state
+        tool_two.update_state(thread1, {"my_key": "er"})
+        assert tool_two.get_state(thread1) == StateSnapshot(
+            values={
+                "my_key": "value prepared slower",
+                "market": "DE",
+            },
+            next=("finish",),
+            config=tool_two.checkpointer.get_tuple(thread1).config,
+            created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
+            metadata={
+                "source": "update",
+                "step": 3,
+                "writes": {"tool_two_slow": {"my_key": "er"}},
+            },
+            parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config,
+        )
+
     with SqliteSaver.from_conn_string(":memory:") as saver:
         tool_two = tool_two_graph.compile(
             checkpointer=saver, interrupt_after=["prepare"]
@@ -6440,6 +6486,41 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) ->
         {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
     ]
 
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaverAssertImmutable(),
+        interrupt_before=["qa"],
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert [
+        c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
+    ] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+        {"retriever_two": {"docs": ["doc3", "doc4"]}},
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+    ]
+
+    app_w_interrupt.update_state(config, {"docs": ["doc5"]})
+    assert app_w_interrupt.get_state(config) == StateSnapshot(
+        values={
+            "query": "analyzed: query: what is weather in sf",
+            "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"],
+        },
+        next=("qa",),
+        config=app_w_interrupt.checkpointer.get_tuple(config).config,
+        created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
+        metadata={
+            "source": "update",
+            "step": 4,
+            "writes": {"retriever_one": {"docs": ["doc5"]}},
+        },
+    )
+
+    assert [c for c in app_w_interrupt.stream(None, config)] == [
+        {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}},
+    ]
+
 
 def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
     snapshot: SnapshotAssertion,