From 9772497b60d270c51d3c29175f1ca5c479daa74a Mon Sep 17 00:00:00 2001
From: krypticmouse <herumbshandilya123@gmail.com>
Date: Tue, 3 Mar 2026 11:33:17 -0800
Subject: [PATCH] Convert classes to context handlers

---
 src/openjarvis/sdk.py    | 20 ++++++++++++++++++--
 src/openjarvis/system.py |  6 ++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/src/openjarvis/sdk.py b/src/openjarvis/sdk.py
index dfb9e0fa..95d3f497 100644
--- a/src/openjarvis/sdk.py
+++ b/src/openjarvis/sdk.py
@@ -108,6 +108,12 @@ class MemoryHandle:
                 self._backend.close()
             self._backend = None
 
+    def __enter__(self) -> MemoryHandle:
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        self.close()
+
 
 class Jarvis:
     """High-level OpenJarvis SDK.
@@ -116,9 +122,13 @@ class Jarvis:
 
         from openjarvis import Jarvis
 
+        with Jarvis() as j:
+            response = j.ask("Hello, what can you do?")
+            print(response)
+
+        # Or without context manager:
         j = Jarvis()
-        response = j.ask("Hello, what can you do?")
-        print(response)
+        response = j.ask("Hello")
         j.close()
     """
 
@@ -479,5 +489,11 @@ class Jarvis:
             self._audit_logger = None
         self._engine = None
 
+    def __enter__(self) -> Jarvis:
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        self.close()
+
 
 __all__ = ["Jarvis", "JarvisSystem", "MemoryHandle", "SystemBuilder"]
diff --git a/src/openjarvis/system.py b/src/openjarvis/system.py
index 3046ec7a..aab2fc7c 100644
--- a/src/openjarvis/system.py
+++ b/src/openjarvis/system.py
@@ -274,6 +274,12 @@ class JarvisSystem:
         if self.trace_store and hasattr(self.trace_store, "close"):
             self.trace_store.close()
 
+    def __enter__(self) -> JarvisSystem:
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        self.close()
+
 
 class SystemBuilder:
     """Config-driven fluent builder for JarvisSystem."""