fix: add resiliency for task cancellation (#5846)

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:

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",