fix(checkpoint): InMemorySaver context managers should return self in… (#6529)

h/t to @lexi-k for openening. merged here to check/fix CI

Co-authored-by: lexi-k <69981673+lexi-k@users.noreply.github.com>

libs/checkpoint/langgraph/checkpoint/memory/__init__.py

@@ -97,7 +97,8 @@ class InMemorySaver(
             self.stack.enter_context(self.blobs)  # type: ignore[arg-type]
 
     def __enter__(self) -> InMemorySaver:
-        return self.stack.__enter__()
+        self.stack.__enter__()
+        return self
 
     def __exit__(
         self,
@@ -108,7 +109,8 @@ class InMemorySaver(
         return self.stack.__exit__(exc_type, exc_value, traceback)
 
     async def __aenter__(self) -> InMemorySaver:
-        return self.stack.__enter__()
+        self.stack.__enter__()
+        return self
 
     async def __aexit__(
         self,

libs/checkpoint/tests/test_memory.py

@@ -188,7 +188,14 @@ class TestMemorySaver:
         assert len(search_results_4) == 0
 
 
-def test_memory_saver() -> None:
+async def test_memory_saver() -> None:
     from langgraph.checkpoint.memory import InMemorySaver
 
-    assert isinstance(InMemorySaver(), InMemorySaver)
+    memory_saver = InMemorySaver()
+    assert isinstance(memory_saver, InMemorySaver)
+
+    async with memory_saver as async_memory_saver:
+        assert async_memory_saver is memory_saver
+
+    with memory_saver as sync_memory_saver:
+        assert sync_memory_saver is memory_saver