fix(langgraph): save a cached node's writes like a node that ran

A cache hit's writes were applied but never saved, so a DeltaChannel,
rebuilt from saved writes, lost them for good under every durability.
The loop now saves a hit with put_writes like a node that ran, in all
six places it matches cached writes. A cached flag keeps it streamed as
cached and stops put_writes from writing it back to the cache.

libs/langgraph/langgraph/pregel/_loop.py

@@ -439,7 +439,9 @@ class PregelLoop:
             return None
         return self._graph_lifecycle_events.popleft()
 
-    def put_writes(self, task_id: str, writes: WritesT) -> None:
+    def put_writes(
+        self, task_id: str, writes: WritesT, *, cached: bool = False
+    ) -> None:
         """Put writes for a task, to be read by the next tick."""
         if not writes:
             return
@@ -532,7 +534,7 @@ class PregelLoop:
                 self._error_handler_write_futs.append(fut)
         # output writes
         if hasattr(self, "tasks"):
-            self.output_writes(task_id, writes)
+            self.output_writes(task_id, writes, cached=cached)
 
     def _put_pending_writes(self) -> None:
         if self.checkpointer_put_writes is None:
@@ -1684,7 +1686,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
     ) -> PregelExecutableTask | None:
         if pushed := super().accept_push(task, write_idx, call):
             for task in self.match_cached_writes():
-                self.output_writes(task.id, task.writes, cached=True)
+                self.put_writes(task.id, task.writes, cached=True)
         return pushed
 
     def schedule_error_handler(
@@ -1721,13 +1723,15 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
         if self._reapplies_pending_writes:
             self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task})
         for task in self.match_cached_writes():
-            self.output_writes(task.id, task.writes, cached=True)
+            self.put_writes(task.id, task.writes, cached=True)
         return handler_task
 
-    def put_writes(self, task_id: str, writes: WritesT) -> None:
+    def put_writes(
+        self, task_id: str, writes: WritesT, *, cached: bool = False
+    ) -> None:
         """Put writes for a task, to be read by the next tick."""
-        super().put_writes(task_id, writes)
-        if not writes or self.cache is None or not hasattr(self, "tasks"):
+        super().put_writes(task_id, writes, cached=cached)
+        if cached or not writes or self.cache is None or not hasattr(self, "tasks"):
             return
         task = self.tasks.get(task_id)
         if task is None or task.cache_key is None:
@@ -1937,7 +1941,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
     ) -> PregelExecutableTask | None:
         if pushed := super().accept_push(task, write_idx, call):
             for task in await self.amatch_cached_writes():
-                self.output_writes(task.id, task.writes, cached=True)
+                self.put_writes(task.id, task.writes, cached=True)
         return pushed
 
     async def aschedule_error_handler(
@@ -1974,13 +1978,15 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
         if self._reapplies_pending_writes:
             self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task})
         for task in await self.amatch_cached_writes():
-            self.output_writes(task.id, task.writes, cached=True)
+            self.put_writes(task.id, task.writes, cached=True)
         return handler_task
 
-    def put_writes(self, task_id: str, writes: WritesT) -> None:
+    def put_writes(
+        self, task_id: str, writes: WritesT, *, cached: bool = False
+    ) -> None:
         """Put writes for a task, to be read by the next tick."""
-        super().put_writes(task_id, writes)
-        if not writes or self.cache is None or not hasattr(self, "tasks"):
+        super().put_writes(task_id, writes, cached=cached)
+        if cached or not writes or self.cache is None or not hasattr(self, "tasks"):
             return
         task = self.tasks.get(task_id)
         if task is None or task.cache_key is None:

libs/langgraph/langgraph/pregel/main.py

@@ -3038,7 +3038,7 @@ class Pregel(
                 # with channel updates applied only at the transition between steps.
                 while loop.tick():
                     for task in loop.match_cached_writes():
-                        loop.output_writes(task.id, task.writes, cached=True)
+                        loop.put_writes(task.id, task.writes, cached=True)
                     for _ in runner.tick(
                         [t for t in loop.tasks.values() if not t.writes],
                         timeout=self.step_timeout,
@@ -3509,7 +3509,7 @@ class Pregel(
                 try:
                     while loop.tick():
                         for task in await loop.amatch_cached_writes():
-                            loop.output_writes(task.id, task.writes, cached=True)
+                            loop.put_writes(task.id, task.writes, cached=True)
                         async for _ in runner.atick(
                             [t for t in loop.tasks.values() if not t.writes],
                             timeout=self.step_timeout,

libs/langgraph/tests/test_delta_channel_node_cache.py

@@ -0,0 +1,89 @@
+"""A node served from the node cache saves its writes like a node that ran,
+without writing them back to the cache."""
+
+import operator
+from typing import Annotated, Any
+
+import pytest
+from langgraph.cache.memory import InMemoryCache
+from langgraph.checkpoint.memory import InMemorySaver
+from typing_extensions import TypedDict
+
+from langgraph.channels.delta import DeltaChannel
+from langgraph.graph import START, StateGraph
+from langgraph.types import CachePolicy, Durability
+
+pytestmark = pytest.mark.anyio
+
+INPUT = {"log": [], "plain": []}
+
+
+def _append(current: list, writes: list) -> list:
+    return [*current, *(item for write in writes for item in write)]
+
+
+class _State(TypedDict):
+    log: Annotated[list, DeltaChannel(_append)]
+    plain: Annotated[list, operator.add]
+
+
+class _CountsSets(InMemoryCache):
+    sets = 0
+
+    def set(self, keys: Any) -> None:
+        self.sets += 1
+        super().set(keys)
+
+
+def _a_then_cached_b_then_c(runs: list[str], cache: InMemoryCache) -> Any:
+    def node(name: str) -> Any:
+        def run(state: _State) -> dict:
+            runs.append(name)
+            return {"log": [name], "plain": [name]}
+
+        return run
+
+    builder = StateGraph(_State)
+    builder.add_node("a", node("a"))
+    builder.add_node("b", node("b"), cache_policy=CachePolicy())
+    builder.add_node("c", node("c"))
+    builder.add_edge(START, "a")
+    builder.add_edge("a", "b")
+    builder.add_edge("b", "c")
+    return builder.compile(checkpointer=InMemorySaver(), cache=cache)
+
+
+def test_a_cache_hit_saves_its_writes_without_caching_them_again(
+    durability: Durability,
+) -> None:
+    runs: list[str] = []
+    cache = _CountsSets()
+    graph = _a_then_cached_b_then_c(runs, cache)
+    graph.invoke(INPUT, {"configurable": {"thread_id": "1"}}, durability=durability)
+    config = {"configurable": {"thread_id": "2"}}
+
+    graph.invoke(INPUT, config, durability=durability)
+
+    assert runs == ["a", "b", "c", "a", "c"]
+    assert cache.sets == 1, "the cache hit was written back to the cache"
+    for state in graph.get_state_history(config):
+        assert state.values.get("log", []) == state.values.get("plain", [])
+
+
+async def test_a_cache_hit_saves_its_writes_without_caching_them_again_async(
+    durability: Durability,
+) -> None:
+    runs: list[str] = []
+    cache = _CountsSets()
+    graph = _a_then_cached_b_then_c(runs, cache)
+    await graph.ainvoke(
+        INPUT, {"configurable": {"thread_id": "1"}}, durability=durability
+    )
+    config = {"configurable": {"thread_id": "2"}}
+
+    await graph.ainvoke(INPUT, config, durability=durability)
+
+    assert runs == ["a", "b", "c", "a", "c"]
+    assert cache.sets == 1, "the cache hit was written back to the cache"
+    async for state in graph.aget_state_history(config):
+        assert state.values.get("log", []) == state.values.get("plain", [])

libs/langgraph/tests/test_pregel.py

@@ -5932,9 +5932,9 @@ def test_no_redundant_put_writes_for_cached_task(
     put_writes_task_ids: list[str] = []
     orig = PregelLoop.put_writes
 
-    def spy(self, task_id, writes):
+    def spy(self, task_id, writes, **kwargs):
         put_writes_task_ids.append(task_id)
-        return orig(self, task_id, writes)
+        return orig(self, task_id, writes, **kwargs)
 
     with patch.object(PregelLoop, "put_writes", spy):
         result = workflow.invoke(Command(resume="ans"), config=config)

libs/langgraph/tests/test_pregel_async.py

@@ -8157,9 +8157,9 @@ async def test_no_redundant_put_writes_for_cached_task(
     put_writes_task_ids: list[str] = []
     orig = PregelLoop.put_writes
 
-    def spy(self, task_id, writes):
+    def spy(self, task_id, writes, **kwargs):
         put_writes_task_ids.append(task_id)
-        return orig(self, task_id, writes)
+        return orig(self, task_id, writes, **kwargs)
 
     with patch.object(PregelLoop, "put_writes", spy):
         result = await workflow.ainvoke(Command(resume="ans"), config=config)