From e430a53ddad7c6e6441f460de4a836d795564fd9 Mon Sep 17 00:00:00 2001
From: Aditya Chhabra <aditya.chhabra@cleverthis.com>
Date: Mon, 17 Nov 2025 19:27:12 +0530
Subject: [PATCH] fix: fix deadlock in AgentWithMemory with reentrant lock
 implementation

---
 src/cleveragents/agents/base.py | 85 ++++++++++++++++++++++++++++++---
 1 file changed, 78 insertions(+), 7 deletions(-)

diff --git a/src/cleveragents/agents/base.py b/src/cleveragents/agents/base.py
index b4d4b03b2..c4651d6f0 100644
--- a/src/cleveragents/agents/base.py
+++ b/src/cleveragents/agents/base.py
@@ -235,6 +235,9 @@ class AgentWithMemory(Agent):
         super().__init__(name, config, template_renderer)
         self.memory: dict[str, Any] = {}
         self._memory_lock_instance: Optional[asyncio.Lock] = None
+        # Reentrant lock state: track which task owns the lock and reentrancy depth
+        self._lock_owner: Optional[asyncio.Task[Any]] = None
+        self._lock_count: int = 0
 
     @property
     def _memory_lock(self) -> asyncio.Lock:
@@ -261,10 +264,66 @@ class AgentWithMemory(Agent):
                     self._memory_lock_instance = asyncio.Lock()
         return self._memory_lock_instance
 
+    async def _acquire_reentrant_lock(self) -> bool:
+        """
+        Acquire memory lock with reentrant behavior.
+
+        If the current task already owns the lock, increments the reentrancy counter
+        instead of deadlocking. This allows process_message() to safely call
+        get_memory() or update_memory() while already holding the lock.
+
+        Returns:
+            True if the lock was actually acquired, False if already owned by current task.
+        """
+        current_task = asyncio.current_task()
+
+        # Check if current task already owns the lock (reentrant case)
+        if self._lock_owner == current_task:
+            self._lock_count += 1
+            return False  # Lock already owned, no need to release later
+
+        # Actually acquire the lock for the first time
+        await self._memory_lock.acquire()
+        self._lock_owner = current_task
+        self._lock_count = 1
+        return True  # Lock was acquired, must be released later
+
+    def _release_reentrant_lock(self, acquired: bool) -> None:
+        """
+        Release memory lock with reentrant behavior.
+
+        Only releases the underlying lock when reentrancy count reaches zero.
+        This ensures proper cleanup regardless of how many times the lock was
+        reentered.
+
+        Args:
+            acquired: True if this call actually acquired the lock, False if reentrant.
+        """
+        if not acquired:
+            # This was a reentrant call, just decrement the counter
+            self._lock_count -= 1
+            return
+
+        # This was the original acquirer, decrement and check if we can release
+        self._lock_count -= 1
+        if self._lock_count == 0:
+            # No more reentrant calls, safe to release
+            self._lock_owner = None
+            self._memory_lock.release()
+
     async def _process_wrapper(self, message_data: tuple[str, dict[str, Any]]) -> str:
-        """Wrapper that manages memory access."""
-        async with self._memory_lock:
+        """
+        Wrapper that manages memory access with reentrant locking.
+
+        Acquires the memory lock before processing to ensure all memory operations
+        within a single message are atomic. Uses reentrant locking so that
+        process_message() can safely call get_memory()/update_memory() without deadlock.
+        """
+        acquired = await self._acquire_reentrant_lock()
+        try:
             return await super()._process_wrapper(message_data)
+        finally:
+            self._release_reentrant_lock(acquired)
 
     def save_memory(self) -> dict[str, Any]:
         """
@@ -291,18 +350,25 @@ class AgentWithMemory(Agent):
 
     async def update_memory(self, key: str, value: Any) -> None:
         """
-        Update a memory value asynchronously.
+        Update a memory value asynchronously with reentrant locking.
+
+        Safe to call from within process_message() even when the lock is already held.
 
         Args:
             key: The memory key to update.
             value: The new value.
         """
-        async with self._memory_lock:
+        acquired = await self._acquire_reentrant_lock()
+        try:
             self.memory[key] = value
+        finally:
+            self._release_reentrant_lock(acquired)
 
     async def get_memory(self, key: str, default: Any = None) -> Any:
         """
-        Get a memory value asynchronously.
+        Get a memory value asynchronously with reentrant locking.
+
+        Safe to call from within process_message() even when the lock is already held.
 
         Args:
             key: The memory key to retrieve.
@@ -311,13 +377,18 @@ class AgentWithMemory(Agent):
         Returns:
             The memory value or default.
         """
-        async with self._memory_lock:
+        acquired = await self._acquire_reentrant_lock()
+        try:
             return self.memory.get(key, default)
+        finally:
+            self._release_reentrant_lock(acquired)
 
     def dispose(self) -> None:
         """Clean up the agent's resources including memory lock."""
-        # Clear the lock instance to allow proper cleanup
+        # Clear the lock instance and reentrant state to allow proper cleanup
         self._memory_lock_instance = None
+        self._lock_owner = None
+        self._lock_count = 0
 
         # Call parent dispose
         super().dispose()
-- 
2.52.0