From 0d8dfa7bbaeb6fad807eaa3ec0147e2e3297948d Mon Sep 17 00:00:00 2001
From: Darren Clark <darrenjoelclark@icloud.com>
Date: Mon, 14 Jul 2025 16:45:25 -0400
Subject: [PATCH] fix(checkpoint): fix AsyncBatchedBaseStore getting stuck

This commit fixes #5503

Gist of it is:

- `asyncio.exception.InvalidStateError` were being raised when the
  future was cancelled
- this exception bubbled up and killed the background task
- `AsyncBatchedBaseStore` stopped doing queries because the background
  task wasn't running anymore

This commit adds some "if future is not done" checks to guard against
this.
---
 libs/checkpoint/langgraph/store/base/batch.py |  8 +++-
 libs/checkpoint/tests/test_store.py           | 37 +++++++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py
index 7d9daad4a..e581a8bc5 100644
--- a/libs/checkpoint/langgraph/store/base/batch.py
+++ b/libs/checkpoint/langgraph/store/base/batch.py
@@ -343,10 +343,14 @@ async def _run(
 
                     # set the results of each operation
                     for fut, result in zip(futs, results):
-                        fut.set_result(result)
+                        # guard against future being done (e.g. cancelled)
+                        if not fut.done():
+                            fut.set_result(result)
                 except Exception as e:
                     for fut in futs:
-                        fut.set_exception(e)
+                        # guard against future being done (e.g. cancelled)
+                        if not fut.done():
+                            fut.set_exception(e)
             finally:
                 # remove strong ref to store
                 del s
diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py
index 8bdabf80d..75774ef09 100644
--- a/libs/checkpoint/tests/test_store.py
+++ b/libs/checkpoint/tests/test_store.py
@@ -155,6 +155,43 @@ async def test_async_batch_store(mocker: MockerFixture) -> None:
     ]
 
 
+async def test_async_batch_store_handles_cancellation() -> None:
+    class MockStore(AsyncBatchedBaseStore):
+        def batch(self, ops: Iterable[Op]) -> list[Result]:
+            raise NotImplementedError
+
+        async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+            assert all(isinstance(op, GetOp) for op in ops)
+            return [
+                Item(
+                    value={},
+                    key=getattr(op, "key", ""),
+                    namespace=getattr(op, "namespace", ()),
+                    created_at=datetime(2024, 9, 24, 17, 29, 10, 128397),
+                    updated_at=datetime(2024, 9, 24, 17, 29, 10, 128397),
+                )
+                for op in ops
+            ]
+
+    store = MockStore()
+
+    # Simulate cancellation
+    task = asyncio.create_task(store.aget(namespace=("a",), key="b"))
+    await asyncio.sleep(0)
+    task.cancel()
+    await asyncio.sleep(0)
+
+    # Cancelling individual queries against the store should not break the store
+    result = await store.aget(namespace=("c",), key="d")
+    assert result == Item(
+        value={},
+        key="d",
+        namespace=("c",),
+        created_at=datetime(2024, 9, 24, 17, 29, 10, 128397),
+        updated_at=datetime(2024, 9, 24, 17, 29, 10, 128397),
+    )
+
+
 def test_list_namespaces_basic() -> None:
     store = InMemoryStore()