fix(memory): restore Memory and AsyncMemory context-manager protocol (#7354)

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.")

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