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

A node served from the node cache had its writes applied but never
saved: the loop only streamed them. Plain channels didn't notice, since
the next checkpoint stores their value, but a `DeltaChannel` is rebuilt
from the saved writes, so its saved state lost the cached node's write
for good, under every durability. With `a -> b -> c` and `b` cached, a
second thread read back `log` as `a, c` next to `plain` as `a, b, c`.

The loop now saves a cache hit with `put_writes`, like a node that ran,
in all six places it matches cached writes. A new `cached` flag keeps
the hit streamed as cached and stops `put_writes` from writing it back
to the cache, which would also restart its TTL. A hit now costs the
saver write the node makes when it runs, and a cached `@task` call's
result is saved the same way.

The sync loop also stops caching a node that failed or was interrupted,
as the async loop already did: a cached error made the next run with the
same input skip the node without raising.

Found in review of langchain-ai/langgraphjs#2969. JS twin:
langchain-ai/langgraphjs#2980

## Tests

- `test_delta_channel_node_cache.py`: `b` runs on the first thread and
is a cache hit on the second, under every durability, with `invoke` and
`ainvoke`. Every checkpoint of the second thread reads the same in the
DeltaChannel and a plain channel, and the cache is written once. All 6
cases fail on main; without the `cached` check in `put_writes`, they
fail on the cache count.
- `test_a_node_that_fails_is_not_cached`, with `invoke` and `ainvoke`: a
cached node that raises raises again on the next thread. The sync case
fails on main.
- The two tests that spy on `PregelLoop.put_writes` pass the new keyword
through.
- `make format`, `make lint_package`, `make lint_tests`, and the
`langgraph`, `prebuilt` and `cli` suites pass (without Redis).

libs/langgraph/langgraph/pregel/_loop.py

@@ -161,6 +161,12 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol:
     return StreamProtocol(__call__, {mode for s in streams for mode in s.modes})
 
 
+def _cacheable(writes: WritesT) -> bool:
+    # A cache hit skips the node, so a cached error would skip it without
+    # raising: only a node that finished goes to the cache.
+    return not any(c in (INTERRUPT, ERROR) for c, _ in writes)
+
+
 class PregelLoop:
     config: RunnableConfig
     store: BaseStore | None
@@ -443,7 +449,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
@@ -536,7 +544,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:
@@ -1710,7 +1718,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(
@@ -1747,16 +1755,18 @@ 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:
+        if task is None or task.cache_key is None or not _cacheable(writes):
             return
         self.submit(
             self.cache.set,
@@ -1964,7 +1974,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(
@@ -2001,19 +2011,18 @@ 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:
-            return
-        if writes[0][0] in (INTERRUPT, ERROR):
-            # only cache successful tasks
+        if task is None or task.cache_key is None or not _cacheable(writes):
             return
         self.submit(
             self.cache.aset,

libs/langgraph/langgraph/pregel/_retry.py

@@ -7,7 +7,7 @@ import sys
 import threading
 import time
 import weakref
-from collections.abc import Awaitable, Callable, Sequence
+from collections.abc import Callable, Sequence
 from contextlib import suppress
 from dataclasses import dataclass, replace
 from datetime import datetime, timedelta, timezone
@@ -686,8 +686,6 @@ async def arun_with_retry(
     task: PregelExecutableTask,
     retry_policy: Sequence[RetryPolicy] | None,
     stream: bool = False,
-    match_cached_writes: Callable[[], Awaitable[Sequence[PregelExecutableTask]]]
-    | None = None,
     configurable: dict[str, Any] | None = None,
 ) -> None:
     """Run a task asynchronously with retries."""
@@ -711,11 +709,6 @@ async def arun_with_retry(
                 )
             },
         )
-    if match_cached_writes is not None and task.cache_key is not None:
-        for t in await match_cached_writes():
-            if t is task:
-                # if the task is already cached, return
-                return
     while True:
         runtime = config.get(CONF, {}).get(CONFIG_KEY_RUNTIME)
         if isinstance(runtime, Runtime):

libs/langgraph/langgraph/pregel/main.py

@@ -3054,7 +3054,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,
@@ -3525,7 +3525,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,122 @@
+"""A node served from the node cache saves its writes like a node that ran,
+without writing them back to the cache, and a node that fails isn't cached."""
+
+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", [])
+
+
+def _cached_node_that_fails(runs: list[str]) -> Any:
+    def fail(state: _State) -> dict:
+        runs.append("b")
+        raise ValueError("b failed")
+
+    builder = StateGraph(_State)
+    builder.add_node("b", fail, cache_policy=CachePolicy())
+    builder.add_edge(START, "b")
+    return builder.compile(checkpointer=InMemorySaver(), cache=InMemoryCache())
+
+
+def test_a_node_that_fails_is_not_cached() -> None:
+    runs: list[str] = []
+    graph = _cached_node_that_fails(runs)
+
+    for thread in ("1", "2"):
+        with pytest.raises(ValueError, match="b failed"):
+            graph.invoke(INPUT, {"configurable": {"thread_id": thread}})
+
+    assert runs == ["b", "b"]
+
+
+async def test_a_node_that_fails_is_not_cached_async() -> None:
+    runs: list[str] = []
+    graph = _cached_node_that_fails(runs)
+
+    for thread in ("1", "2"):
+        with pytest.raises(ValueError, match="b failed"):
+            await graph.ainvoke(INPUT, {"configurable": {"thread_id": thread}})
+
+    assert runs == ["b", "b"]

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)