Add "node finished" callback

libs/langgraph/langgraph/constants.py

@@ -65,6 +65,8 @@ CONFIG_KEY_CHECKPOINT_ID = sys.intern("checkpoint_id")
 # holds the current checkpoint_id, if any
 CONFIG_KEY_CHECKPOINT_NS = sys.intern("checkpoint_ns")
 # holds the current checkpoint_ns, "" for root graph
+CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished")
+# callback to be called when a node is finished
 
 # --- Other constants ---
 PUSH = sys.intern("__pregel_push")

libs/langgraph/langgraph/pregel/__init__.py

@@ -56,6 +56,7 @@ from langgraph.constants import (
     CONF,
     CONFIG_KEY_CHECKPOINT_NS,
     CONFIG_KEY_CHECKPOINTER,
+    CONFIG_KEY_NODE_FINISHED,
     CONFIG_KEY_READ,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_SEND,
@@ -1288,6 +1289,7 @@ class Pregel(PregelProtocol):
                 runner = PregelRunner(
                     submit=loop.submit,
                     put_writes=loop.put_writes,
+                    node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED),
                 )
                 # enable subgraph streaming
                 if subgraphs:
@@ -1509,6 +1511,7 @@ class Pregel(PregelProtocol):
                     submit=loop.submit,
                     put_writes=loop.put_writes,
                     use_astream=do_stream is not None,
+                    node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED),
                 )
                 # enable subgraph streaming
                 if subgraphs:

libs/langgraph/langgraph/pregel/runner.py

@@ -14,7 +14,7 @@ from typing import (
     cast,
 )
 
-from langgraph.constants import ERROR, INTERRUPT, NO_WRITES
+from langgraph.constants import ERROR, INTERRUPT, NO_WRITES, TAG_HIDDEN
 from langgraph.errors import GraphDelegate, GraphInterrupt
 from langgraph.pregel.executor import Submit
 from langgraph.pregel.retry import arun_with_retry, run_with_retry
@@ -32,10 +32,12 @@ class PregelRunner:
         submit: Submit,
         put_writes: Callable[[str, Sequence[tuple[str, Any]]], None],
         use_astream: bool = False,
+        node_finished: Optional[Callable[[str], None]] = None,
     ) -> None:
         self.submit = submit
         self.put_writes = put_writes
         self.use_astream = use_astream
+        self.node_finished = node_finished
 
     def tick(
         self,
@@ -209,6 +211,10 @@ class PregelRunner:
                 # save error to checkpointer
                 self.put_writes(task.id, [(ERROR, exception)])
         else:
+            if self.node_finished and (
+                task.config is None or TAG_HIDDEN not in task.config.get("tags", [])
+            ):
+                self.node_finished(task.name)
             if not task.writes:
                 # add no writes marker
                 task.writes.append((NO_WRITES, None))

libs/langgraph/tests/test_pregel.py

@@ -54,7 +54,7 @@ from langgraph.checkpoint.base import (
     CheckpointTuple,
 )
 from langgraph.checkpoint.memory import MemorySaver
-from langgraph.constants import ERROR, PULL, PUSH
+from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH
 from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
 from langgraph.graph import END, Graph
 from langgraph.graph.graph import START
@@ -8917,15 +8917,27 @@ def test_doubly_nested_graph_interrupts(
     }
 
     # test stream updates w/ nested interrupt
-    config = {"configurable": {"thread_id": "2"}}
+    nodes: list[str] = []
+    config = {
+        "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append}
+    }
     assert [*app.stream({"my_key": "my value"}, config)] == [
         {"parent_1": {"my_key": "hi my value"}},
         {"__interrupt__": ()},
     ]
+    assert nodes == ["parent_1", "grandchild_1"]
     assert [*app.stream(None, config)] == [
         {"child": {"my_key": "hi my value here and there"}},
         {"parent_2": {"my_key": "hi my value here and there and back again"}},
     ]
+    assert nodes == [
+        "parent_1",
+        "grandchild_1",
+        "grandchild_2",
+        "child_1",
+        "child",
+        "parent_2",
+    ]
 
     # test stream values w/ nested interrupt
     config = {"configurable": {"thread_id": "3"}}

libs/langgraph/tests/test_pregel_async.py

@@ -50,7 +50,7 @@ from langgraph.checkpoint.base import (
     CheckpointTuple,
 )
 from langgraph.checkpoint.memory import MemorySaver
-from langgraph.constants import ERROR, PULL, PUSH
+from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH
 from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.graph import START
@@ -7605,15 +7605,27 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None:
         }
 
         # test stream updates w/ nested interrupt
-        config = {"configurable": {"thread_id": "2"}}
+        nodes: list[str] = []
+        config = {
+            "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append}
+        }
         assert [c async for c in app.astream({"my_key": "my value"}, config)] == [
             {"parent_1": {"my_key": "hi my value"}},
             {"__interrupt__": ()},
         ]
+        assert nodes == ["parent_1", "grandchild_1"]
         assert [c async for c in app.astream(None, config)] == [
             {"child": {"my_key": "hi my value here and there"}},
             {"parent_2": {"my_key": "hi my value here and there and back again"}},
         ]
+        assert nodes == [
+            "parent_1",
+            "grandchild_1",
+            "grandchild_2",
+            "child_1",
+            "child",
+            "parent_2",
+        ]
 
         # test stream values w/ nested interrupt
         config = {"configurable": {"thread_id": "3"}}