diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index d7157e239..cddc0a8ca 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/libs/langgraph/langgraph/pregel/_algo.py @@ -210,7 +210,7 @@ def local_read( return values -def increment(current: int | None, channel: None) -> int: +def increment(current: int | None, channel: None = None) -> int: """Default channel versioning function, increments the current int version.""" return current + 1 if current is not None else 1 diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 0ebb14042..00b910dca 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -25,6 +25,7 @@ from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager from langchain_core.runnables import RunnableConfig from langgraph.cache.base import BaseCache from langgraph.checkpoint.base import ( + DELTA_SENTINEL, WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, @@ -198,6 +199,7 @@ class PregelLoop: checkpoint_pending_writes: list[PendingWrite] checkpoint_previous_versions: dict[str, str | float | int] prev_checkpoint_config: RunnableConfig | None + _pending_write_futs: list[concurrent.futures.Future] status: Literal[ "input", @@ -407,7 +409,7 @@ class PregelLoop: task = self.tasks.get(task_id) else: task = None - self.submit( + fut = self.submit( self.checkpointer_put_writes, config, writes_to_save, @@ -415,12 +417,13 @@ class PregelLoop: task_path_str(task.path) if task else "", ) else: - self.submit( + fut = self.submit( self.checkpointer_put_writes, config, writes_to_save, task_id, ) + self._pending_write_futs.append(fut) # output writes if hasattr(self, "tasks"): self.output_writes(task_id, writes) @@ -931,6 +934,17 @@ class PregelLoop: ) self.checkpoint_previous_versions = channel_versions + # If the checkpoint has any DELTA_SENTINEL blobs, the sentinel is + # only meaningful if checkpoint_writes are durable first. Flush + # pending write futures synchronously before committing the blob so + # we never end up with a sentinel blob backed by missing writes. + if self._pending_write_futs and any( + v is DELTA_SENTINEL for v in self.checkpoint["channel_values"].values() + ): + for fut in self._pending_write_futs: + fut.result() + self._pending_write_futs.clear() + # save it, without blocking # if there's a previous checkpoint save in progress, wait for it # ensuring checkpointers receive checkpoints in order @@ -1275,6 +1289,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): if saved.pending_writes is not None else [] ) + self._pending_write_futs = [] self.submit = self.stack.enter_context(BackgroundExecutor(self.config)) self.channels, self.managed = channels_from_checkpoint( self.specs, @@ -1480,6 +1495,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): if saved.pending_writes is not None else [] ) + self._pending_write_futs = [] self.submit = await self.stack.enter_async_context( AsyncBackgroundExecutor(self.config) ) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0f1e4e196..52057abe5 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9584,3 +9584,90 @@ async def test_delta_channel_update_by_id_end_to_end() -> None: assert "h1" in ids # h1 persists (updated, not duplicated) assert "h2" in ids assert ids.count("h1") == 1, "h1 must not be duplicated" + + +async def test_delta_channel_write_flushed_before_put() -> None: + """checkpoint_writes are flushed synchronously before put when DELTA_SENTINEL + is present, ensuring writes are durable before the sentinel blob is committed. + + We verify this by intercepting put_writes and put calls and confirming + put_writes always completes before put is called for sentinel checkpoints. + """ + import threading + from typing import Annotated + + from langchain_core.messages import AIMessage, HumanMessage + from langgraph.checkpoint.base import DELTA_SENTINEL + from langgraph.checkpoint.memory import InMemorySaver + + from langgraph.graph import START, StateGraph + from langgraph.graph.message import add_messages + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(add_messages)] + + def respond(state: State) -> dict: + i = len(state["messages"]) + return {"messages": [AIMessage(content=f"r{i}", id=f"ai{i}")]} + + order: list[str] = [] + lock = threading.Lock() + original_put_writes = InMemorySaver.put_writes + original_put = InMemorySaver.put + + def tracked_put_writes(self, config, writes, task_id, task_path=""): + result = original_put_writes(self, config, writes, task_id, task_path) + with lock: + order.append("put_writes") + return result + + def tracked_put(self, config, checkpoint, metadata, new_versions): + # Check if this checkpoint has any DELTA_SENTINEL blobs + has_sentinel = any( + v is DELTA_SENTINEL for v in checkpoint.get("channel_values", {}).values() + ) + if has_sentinel: + with lock: + order.append("put_sentinel") + else: + with lock: + order.append("put_snapshot") + return original_put(self, config, checkpoint, metadata, new_versions) + + InMemorySaver.put_writes = tracked_put_writes + InMemorySaver.put = tracked_put + try: + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_edge(START, "respond") + saver = InMemorySaver() + graph = builder.compile(checkpointer=saver) + config = {"configurable": {"thread_id": "flush-test"}} + + for i in range(3): + graph.invoke( + {"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config + ) + + # For every sentinel put, all preceding put_writes must already be in order + for i, event in enumerate(order): + if event == "put_sentinel": + # All put_writes before this index must appear before this sentinel + preceding = order[:i] + assert "put_writes" in preceding, ( + f"put_sentinel at index {i} had no preceding put_writes: {order}" + ) + # And the most recent put_writes must come before this sentinel + last_write_idx = max( + j for j, e in enumerate(order[:i]) if e == "put_writes" + ) + assert last_write_idx < i, ( + f"put_writes at {last_write_idx} not before put_sentinel at {i}" + ) + finally: + InMemorySaver.put_writes = original_put_writes + InMemorySaver.put = original_put + + # Final state must still be correct + state = graph.get_state(config) + assert len(state.values["messages"]) == 6 # 3 human + 3 AI