diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index e581a8bc5..2728a06ad 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -64,14 +64,21 @@ class AsyncBatchedBaseStore(BaseStore): super().__init__() self._loop = asyncio.get_running_loop() self._aqueue: asyncio.Queue[tuple[asyncio.Future, Op]] = asyncio.Queue() - self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self))) + self._task: asyncio.Task | None = None + self._ensure_task() def __del__(self) -> None: try: - self._task.cancel() + if self._task: + self._task.cancel() except RuntimeError: pass + def _ensure_task(self) -> None: + """Ensure the background processing loop is running.""" + if self._task is None or self._task.done(): + self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self))) + async def aget( self, namespace: tuple[str, ...], @@ -79,7 +86,7 @@ class AsyncBatchedBaseStore(BaseStore): *, refresh_ttl: bool | None = None, ) -> Item | None: - assert not self._task.done() + self._ensure_task() fut = self._loop.create_future() self._aqueue.put_nowait( ( @@ -104,7 +111,7 @@ class AsyncBatchedBaseStore(BaseStore): offset: int = 0, refresh_ttl: bool | None = None, ) -> list[SearchItem]: - assert not self._task.done() + self._ensure_task() fut = self._loop.create_future() self._aqueue.put_nowait( ( @@ -130,7 +137,7 @@ class AsyncBatchedBaseStore(BaseStore): *, ttl: float | None | NotProvided = NOT_PROVIDED, ) -> None: - assert not self._task.done() + self._ensure_task() _validate_namespace(namespace) fut = self._loop.create_future() self._aqueue.put_nowait( @@ -148,7 +155,7 @@ class AsyncBatchedBaseStore(BaseStore): namespace: tuple[str, ...], key: str, ) -> None: - assert not self._task.done() + self._ensure_task() fut = self._loop.create_future() self._aqueue.put_nowait((fut, PutOp(namespace, key, None))) return await fut @@ -162,7 +169,7 @@ class AsyncBatchedBaseStore(BaseStore): limit: int = 100, offset: int = 0, ) -> list[tuple[str, ...]]: - assert not self._task.done() + self._ensure_task() fut = self._loop.create_future() match_conditions = [] if prefix: diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py index 75774ef09..589eec0e4 100644 --- a/libs/checkpoint/tests/test_store.py +++ b/libs/checkpoint/tests/test_store.py @@ -34,6 +34,42 @@ class MockAsyncBatchedStore(AsyncBatchedBaseStore): return self._store.batch(ops) +async def test_async_batch_store_resilience() -> None: + """Test that AsyncBatchedBaseStore recovers gracefully from task cancellation.""" + doc = {"foo": "bar"} + async_store = MockAsyncBatchedStore() + + await async_store.aput(("foo", "langgraph", "foo"), "bar", doc) + + # Store the original task reference + original_task = async_store._task + assert original_task is not None + assert not original_task.done() + + # Cancel the background task + original_task.cancel() + await asyncio.sleep(0.01) + assert original_task.cancelled() + + # Perform a new operation - this should trigger _ensure_task() to create a new task + result = await async_store.asearch(("foo", "langgraph", "foo")) + assert len(result) > 0 + assert result[0].value == doc + + # Verify a new task was created + new_task = async_store._task + assert new_task is not None + assert new_task is not original_task + assert not new_task.done() + + # Test that operations continue to work with the new task + doc2 = {"baz": "qux"} + await async_store.aput(("test", "namespace"), "key", doc2) + result2 = await async_store.aget(("test", "namespace"), "key") + assert result2 is not None + assert result2.value == doc2 + + def test_get_text_at_path() -> None: nested_data = { "name": "test",