diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index e342548a7..1097757e9 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -389,7 +389,7 @@ async def test_delta_channel_chain_reconstruction(saver_name: str) -> None: from typing_extensions import TypedDict class State(TypedDict): - messages: Annotated[list, DeltaChannel(add_messages)] + messages: Annotated[list, DeltaChannel(list, add_messages)] def respond(state: State) -> dict: n = len(state["messages"]) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 00b910dca..cfeacda15 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -25,7 +25,6 @@ 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, @@ -69,6 +68,7 @@ from langgraph.callbacks import ( GraphResumeEvent, ) from langgraph.channels.base import BaseChannel +from langgraph.channels.delta import DeltaChannel from langgraph.channels.untracked_value import UntrackedValue from langgraph.constants import TAG_HIDDEN from langgraph.errors import ( @@ -199,7 +199,6 @@ 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", @@ -423,7 +422,10 @@ class PregelLoop: writes_to_save, task_id, ) - self._pending_write_futs.append(fut) + if hasattr(self, "_delta_write_futs") and any( + isinstance(self.specs.get(c), DeltaChannel) for c, _ in writes_to_save + ): + self._delta_write_futs.append(fut) # output writes if hasattr(self, "tasks"): self.output_writes(task_id, writes) @@ -934,17 +936,6 @@ 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 @@ -1289,7 +1280,6 @@ 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, @@ -1390,6 +1380,11 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): metadata: CheckpointMetadata, new_versions: ChannelVersions, ) -> RunnableConfig: + # Drain DeltaChannel write futures before committing the checkpoint so + # DELTA_SENTINEL blobs are never saved ahead of their backing writes. + if self._delta_write_futs: + futs, self._delta_write_futs = self._delta_write_futs, [] + await asyncio.gather(*futs) try: if prev is not None: await prev @@ -1495,7 +1490,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): if saved.pending_writes is not None else [] ) - self._pending_write_futs = [] + self._delta_write_futs: list[asyncio.Future[Any]] = [] 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 5ccb347ab..f898ac4a1 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9586,14 +9586,9 @@ async def test_delta_channel_update_by_id_end_to_end() -> None: 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 +async def test_delta_channel_async_write_ordering() -> None: + """In async mode, DeltaChannel write futures are awaited before the checkpoint + is committed, so aput_writes always precedes aput for sentinel checkpoints.""" from typing import Annotated from langchain_core.messages import AIMessage, HumanMessage @@ -9611,63 +9606,52 @@ async def test_delta_channel_write_flushed_before_put() -> None: 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 + original_aput_writes = InMemorySaver.aput_writes + original_aput = InMemorySaver.aput - 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") + async def tracked_aput_writes(self, config, writes, task_id, task_path=""): + result = await original_aput_writes(self, config, writes, task_id, task_path) + order.append("aput_writes") return result - def tracked_put(self, config, checkpoint, metadata, new_versions): - # Check if this checkpoint has any DELTA_SENTINEL blobs + async def tracked_aput(self, config, checkpoint, metadata, new_versions): has_sentinel = any( - v is DELTA_SENTINEL for v in checkpoint.get("channel_values", {}).values() + 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) + order.append("aput_sentinel" if has_sentinel else "aput_other") + return await original_aput(self, config, checkpoint, metadata, new_versions) - InMemorySaver.put_writes = tracked_put_writes - InMemorySaver.put = tracked_put + InMemorySaver.aput_writes = tracked_aput_writes + InMemorySaver.aput = tracked_aput 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"}} + graph = builder.compile(checkpointer=InMemorySaver()) + config = {"configurable": {"thread_id": "async-ordering-test"}} for i in range(3): - graph.invoke( + await graph.ainvoke( {"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config ) - # For every sentinel put, all preceding put_writes must already be in order + # Every aput_sentinel must be preceded by at least one aput_writes for i, event in enumerate(order): - if event == "put_sentinel": - # All put_writes before this index must appear before this sentinel + if event == "aput_sentinel": preceding = order[:i] - assert "put_writes" in preceding, ( - f"put_sentinel at index {i} had no preceding put_writes: {order}" + assert "aput_writes" in preceding, ( + f"aput_sentinel at {i} had no preceding aput_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" + j for j, e in enumerate(order[:i]) if e == "aput_writes" ) assert last_write_idx < i, ( - f"put_writes at {last_write_idx} not before put_sentinel at {i}" + f"aput_writes at {last_write_idx} should precede aput_sentinel at {i}: {order}" ) finally: - InMemorySaver.put_writes = original_put_writes - InMemorySaver.put = original_put + InMemorySaver.aput_writes = original_aput_writes + InMemorySaver.aput = original_aput - # Final state must still be correct - state = graph.get_state(config) + state = await graph.aget_state(config) assert len(state.values["messages"]) == 6 # 3 human + 3 AI diff --git a/new_pr_desc.md b/new_pr_desc.md index 23a0ea627..1be12ac31 100644 --- a/new_pr_desc.md +++ b/new_pr_desc.md @@ -56,10 +56,11 @@ checkpoint-loaded clones. ### Serializer support (`libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py`) -- `dumps_typed`: `DELTA_SENTINEL → ("delta", b"")` (zero bytes, no payload). -- `loads_typed`: `"delta"` tag → `DELTA_SENTINEL`. -- `_DeltaSnapshot` → msgpack ext code 7 (value packed as nested msgpack). -- `loads_typed` ext hook: ext 7 → `_DeltaSnapshot(unpacked_value)`. +Both delta types serialize through msgpack, keeping them in the same codec path with no special string type tags: + +- `DELTA_SENTINEL` → msgpack ext code 8 (`EXT_DELTA_SENTINEL`, zero data bytes). +- `_DeltaSnapshot` → msgpack ext code 7 (`EXT_DELTA_SNAPSHOT`, value packed as nested msgpack). +- Ext hooks decode both back to their singleton / NamedTuple counterparts. ### Ancestor-walk API on `BaseCheckpointSaver` (`libs/checkpoint/langgraph/checkpoint/base/__init__.py`) @@ -131,12 +132,16 @@ is 3–100× faster in the realistic depth range. - `SyncPregelLoop.__enter__` passes `saver` + `config` to `channels_from_checkpoint`. - `AsyncPregelLoop.__aenter__` calls `achannels_from_checkpoint`. -- **Write-ordering safety**: `_pending_write_futs` tracks in-flight - `put_writes` futures. Before committing a checkpoint that contains any - `DELTA_SENTINEL` blob, the loop flushes all pending write futures - synchronously. This ensures checkpoint_writes are durable before the - sentinel blob is stored — a sentinel backed by missing writes would cause - silent data loss on replay. +- **Async write-ordering safety**: `AsyncPregelLoop` maintains + `_delta_write_futs`, a list of in-flight `aput_writes` futures for + DeltaChannel channels. In `accept_writes`, any write to a `DeltaChannel` + appends its future to this list. In `_checkpointer_put_after_previous`, + the list is drained via `await asyncio.gather()` before `aput()` is + called. This ensures checkpoint_writes are durable before the sentinel + blob is stored — a sentinel backed by missing writes would cause silent + data loss on replay. The sync loop does not need this guard: all + background tasks complete before `invoke()` returns via + `BackgroundExecutor.__exit__`. ### `binop.py` refactoring @@ -157,3 +162,6 @@ is 3–100× faster in the realistic depth range. ancestor walk collecting writes for all sentinel channels at once would reduce roundtrips proportionally to the number of DeltaChannel fields in a state schema. +- **Store delta epoch ids in writes table**: this would allow for more targeted +reads of the writes table when reconstructing delta channels (especialy valuable +for postgres).