From d675cf68adea399b475b71b38832361c34a432c8 Mon Sep 17 00:00:00 2001
From: Mohd Quamar Tyagi <104281681+Tyagiquamar@users.noreply.github.com>
Date: Thu, 24 Sep 2026 22:04:57 +0530
Subject: [PATCH] fix(memory): restore Memory and AsyncMemory context-manager
 protocol (#7354)

---
 mem0/memory/main.py     | 14 ++++++++++++++
 tests/test_telemetry.py | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/mem0/memory/main.py b/mem0/memory/main.py
index 05a55c876..e750edbf1 100644
--- a/mem0/memory/main.py
+++ b/mem0/memory/main.py
@@ -2179,6 +2179,13 @@ class Memory(MemoryBase):
             self.db.close()
             self.db = None
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False
+
     def chat(self, query):
         raise NotImplementedError("Chat function not implemented yet.")
 
@@ -3894,5 +3901,12 @@ class AsyncMemory(MemoryBase):
             self.db.close()
             self.db = None
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False
+
     async def chat(self, query):
         raise NotImplementedError("Chat function not implemented yet.")
diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py
index fb7001a62..7c1df5c01 100644
--- a/tests/test_telemetry.py
+++ b/tests/test_telemetry.py
@@ -333,6 +333,25 @@ class TestMemoryLifecycle:
             # db attribute not set at all
             m.close()  # should not raise due to hasattr guard
 
+    def test_context_manager_returns_self_and_closes(self):
+        """with-block should yield the instance and close() it on exit."""
+        m = self._make_mock_memory()
+        db = m.db
+        with m as ctx:
+            assert ctx is m
+        db.close.assert_called_once_with()
+        assert m.db is None
+
+    def test_context_manager_closes_on_exception(self):
+        """Exceptions from the with-body should propagate after close()."""
+        m = self._make_mock_memory()
+        db = m.db
+        with pytest.raises(RuntimeError, match="boom"):
+            with m:
+                raise RuntimeError("boom")
+        db.close.assert_called_once_with()
+        assert m.db is None
+
 
 class TestAsyncMemoryLifecycle:
     """Verify AsyncMemory.close() and async context manager support."""
@@ -357,6 +376,27 @@ class TestAsyncMemoryLifecycle:
             m = AsyncMemory.__new__(AsyncMemory)
             m.close()  # should not raise
 
+    @pytest.mark.asyncio
+    async def test_async_context_manager_returns_self_and_closes(self):
+        """async with-block should yield the instance and close() it on exit."""
+        m = self._make_mock_async_memory()
+        db = m.db
+        async with m as ctx:
+            assert ctx is m
+        db.close.assert_called_once_with()
+        assert m.db is None
+
+    @pytest.mark.asyncio
+    async def test_async_context_manager_closes_on_exception(self):
+        """Exceptions from the async with-body should propagate after close()."""
+        m = self._make_mock_async_memory()
+        db = m.db
+        with pytest.raises(RuntimeError, match="boom"):
+            async with m:
+                raise RuntimeError("boom")
+        db.close.assert_called_once_with()
+        assert m.db is None
+
 
 class TestTelemetryEnvVar:
     """Verify the MEM0_TELEMETRY env var parsing logic."""