contextvar

libs/checkpoint/langgraph/checkpoint/base/__init__.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
+import contextvars
 import copy
 import logging
-import threading
 from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
 from typing import (
     Any,
@@ -39,7 +39,13 @@ from langgraph.checkpoint.serde.types import (
 V = TypeVar("V", int, float, str)
 PendingWrite = tuple[str, str, Any]
 
-_DELTA_RECONSTRUCTION: threading.local = threading.local()
+# Task-local guard: ContextVar is copied per asyncio Task, so concurrent
+# requests on the same event-loop thread do not share this flag. A plain
+# `threading.local()` would leak across tasks and let one in-flight
+# reconstruction silently short-circuit another.
+_DELTA_RECONSTRUCTION: contextvars.ContextVar[bool] = contextvars.ContextVar(
+    "_DELTA_RECONSTRUCTION", default=False
+)
 
 
 def _overwrite_types() -> tuple[type, ...]:
@@ -538,11 +544,11 @@ class BaseCheckpointSaver(Generic[V]):
         # reconstruction which calls get_tuple() again, the inner call
         # returns tuples with DELTA_SENTINEL in channel_values (which this
         # method ignores — it only reads pending_writes).
-        if getattr(_DELTA_RECONSTRUCTION, "active", False):
+        if _DELTA_RECONSTRUCTION.get():
             return DeltaChannelWrites(writes=[])
         overwrite_types = _overwrite_types()
 
-        _DELTA_RECONSTRUCTION.active = True
+        token = _DELTA_RECONSTRUCTION.set(True)
         try:
             collected: list[Any] = []  # newest first
             target_tuple = self.get_tuple(config)
@@ -567,17 +573,17 @@ class BaseCheckpointSaver(Generic[V]):
             collected.reverse()
             return DeltaChannelWrites(writes=collected)
         finally:
-            _DELTA_RECONSTRUCTION.active = False
+            _DELTA_RECONSTRUCTION.reset(token)
 
     async def aget_channel_writes(
         self, config: RunnableConfig, channel: str
     ) -> DeltaChannelWrites:
         """Async version of `get_channel_writes`. See docstring there."""
-        if getattr(_DELTA_RECONSTRUCTION, "active", False):
+        if _DELTA_RECONSTRUCTION.get():
             return DeltaChannelWrites(writes=[])
         overwrite_types = _overwrite_types()
 
-        _DELTA_RECONSTRUCTION.active = True
+        token = _DELTA_RECONSTRUCTION.set(True)
         try:
             collected: list[Any] = []
             target_tuple = await self.aget_tuple(config)
@@ -600,7 +606,7 @@ class BaseCheckpointSaver(Generic[V]):
             collected.reverse()
             return DeltaChannelWrites(writes=collected)
         finally:
-            _DELTA_RECONSTRUCTION.active = False
+            _DELTA_RECONSTRUCTION.reset(token)
 
     def get_next_version(self, current: V | None, channel: None) -> V:
         """Generate the next version ID for a channel.

libs/checkpoint/tests/test_memory.py

@@ -512,6 +512,49 @@ class TestBaseFallbackGetChannelWrites:
         )
         assert result.seed is SEED_UNSET
 
+    async def test_async_fallback_concurrent_tasks_do_not_interfere(self) -> None:
+        """Regression: the re-entrancy guard must be task-local, not thread-local.
+
+        Two concurrent `aget_channel_writes` calls on the same event-loop
+        thread must each see their full reconstructed writes. A
+        `threading.local()` guard would let whichever task set it first
+        short-circuit the other to `writes=[]`.
+        """
+        import asyncio
+
+        saver, thread_id, ns = self._build_saver_with_chain()
+
+        # Force the two tasks to interleave across the `set(True)` boundary:
+        # each `aget_tuple` yields control, so if the guard were thread-local
+        # the second task would observe `active=True` set by the first.
+        orig_aget_tuple = saver.aget_tuple
+
+        async def slow_aget_tuple(config: RunnableConfig) -> Any:
+            await asyncio.sleep(0)
+            return await orig_aget_tuple(config)
+
+        saver.aget_tuple = slow_aget_tuple  # type: ignore[method-assign]
+
+        target_id = "00000000000000000000000000000003.0000000000000000"
+        config: RunnableConfig = {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": ns,
+                "checkpoint_id": target_id,
+            }
+        }
+
+        results = await asyncio.gather(
+            saver.aget_channel_writes(config, "messages"),
+            saver.aget_channel_writes(config, "messages"),
+        )
+
+        expected = DeltaChannelWrites(
+            writes=[{"content": "first"}, {"content": "second"}]
+        )
+        assert results[0] == expected
+        assert results[1] == expected
+
     def test_fallback_stops_at_first_overwrite(self) -> None:
         """An `Overwrite` dominates older history: scan newest→oldest stops at
         the first one (so `snapshot_every` / user Overwrites bound replay cost).