diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 5ec709dee..f27a2e740 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/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. diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index ed129838c..322a6f2d3 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/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).