test: cover async access rejection during active deepcopy

Co-authored-by: Mervin Praison <[email protected]>

src/praisonai-agents/praisonaiagents/agent/async_safety.py

@@ -110,14 +110,32 @@ class AsyncSafeState:
     def __init__(self, initial_value: Any = None):
         self.value = initial_value
         self._lock = DualLock()
+        self._activity_lock = threading.Lock()
+        self._async_holders = 0
+        self._copying = False
 
     def __deepcopy__(self, memo):
-        """Copy the protected value while replacing all lock state."""
+        """Copy protected state, rejecting overlap with asynchronous access."""
         result = type(self).__new__(type(self))
         memo[id(self)] = result
-        with self.lock():
-            result.value = copy.deepcopy(self.value, memo)
+        with self._activity_lock:
+            if self._async_holders:
+                raise RuntimeError(
+                    "Cannot deepcopy AsyncSafeState during asynchronous access"
+                )
+            if self._copying:
+                raise RuntimeError("AsyncSafeState deepcopy is already in progress")
+            self._copying = True
+        try:
+            with self.lock():
+                result.value = copy.deepcopy(self.value, memo)
+        finally:
+            with self._activity_lock:
+                self._copying = False
         result._lock = DualLock()
+        result._activity_lock = threading.Lock()
+        result._async_holders = 0
+        result._copying = False
         return result
         
     @contextmanager 
@@ -130,7 +148,25 @@ class AsyncSafeState:
     async def async_lock(self):
         """Acquire lock in async context."""
         async with self._lock.async_lock():
-            yield self.value
+            self._begin_async_access()
+            try:
+                yield self.value
+            finally:
+                self._end_async_access()
+
+    def _begin_async_access(self) -> None:
+        """Register async access unless a synchronous copy is in progress."""
+        with self._activity_lock:
+            if self._copying:
+                raise RuntimeError(
+                    "Cannot access AsyncSafeState during synchronous deepcopy"
+                )
+            self._async_holders += 1
+
+    def _end_async_access(self) -> None:
+        """Unregister one active asynchronous accessor."""
+        with self._activity_lock:
+            self._async_holders -= 1
             
     def __enter__(self):
         """Support for synchronous context manager protocol (backward compatibility)."""
@@ -146,12 +182,20 @@ class AsyncSafeState:
         """Support for asynchronous context manager protocol."""
         async_lock = self._lock._get_async_lock()
         await async_lock.acquire()
+        try:
+            self._begin_async_access()
+        except Exception:
+            async_lock.release()
+            raise
         return self.value
         
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         """Support for asynchronous context manager protocol."""
         async_lock = self._lock._get_async_lock()
-        async_lock.release()
+        try:
+            self._end_async_access()
+        finally:
+            async_lock.release()
         return None
             
     def get(self) -> Any:

src/praisonai-agents/tests/unit/agent/test_agent_clone.py

@@ -11,6 +11,8 @@ import asyncio
 import copy
 import threading
 
+import pytest
+
 
 class TestAgentDeepCopy:
     """Test Agent.__deepcopy__ and clone_for_channel()."""
@@ -104,3 +106,92 @@ class TestAgentDeepCopy:
         assert cloned._approvals_lock is not agent._approvals_lock
         assert agent._approvals_lock.locked() is True
         assert cloned._approvals_lock.locked() is False
+
+    @pytest.mark.parametrize(
+        "use_async_lock_helper",
+        [True, False],
+        ids=["async-lock-helper", "async-context-manager"],
+    )
+    def test_deepcopy_rejects_active_async_state_access(
+        self,
+        use_async_lock_helper,
+    ):
+        from praisonaiagents.agent.async_safety import AsyncSafeState
+
+        async def _exercise():
+            state = AsyncSafeState(["stable"])
+            entered = asyncio.Event()
+            release = asyncio.Event()
+
+            async def _hold_state():
+                context = state.async_lock() if use_async_lock_helper else state
+                async with context:
+                    entered.set()
+                    await release.wait()
+
+            holder = asyncio.create_task(_hold_state())
+            await entered.wait()
+            try:
+                with pytest.raises(
+                    RuntimeError,
+                    match="during asynchronous access",
+                ):
+                    copy.deepcopy(state)
+            finally:
+                release.set()
+                await holder
+
+            cloned = copy.deepcopy(state)
+            assert cloned.value == ["stable"]
+
+        asyncio.run(_exercise())
+
+    @pytest.mark.parametrize(
+        "use_async_lock_helper",
+        [True, False],
+        ids=["async-lock-helper", "async-context-manager"],
+    )
+    def test_async_access_rejected_during_active_deepcopy(
+        self,
+        use_async_lock_helper,
+    ):
+        from praisonaiagents.agent.async_safety import AsyncSafeState
+
+        class _BlockingValue:
+            def __init__(self, entered, release):
+                self.entered = entered
+                self.release = release
+
+            def __deepcopy__(self, memo):
+                self.entered.set()
+                self.release.wait()
+                return _BlockingValue(self.entered, self.release)
+
+        async def _exercise():
+            entered = threading.Event()
+            release = threading.Event()
+            state = AsyncSafeState(_BlockingValue(entered, release))
+
+            def _copy_state():
+                copy.deepcopy(state)
+
+            copier = threading.Thread(target=_copy_state)
+            copier.start()
+            try:
+                await asyncio.to_thread(entered.wait)
+
+                context = state.async_lock() if use_async_lock_helper else state
+                with pytest.raises(
+                    RuntimeError,
+                    match="during synchronous deepcopy",
+                ):
+                    async with context:
+                        pass
+            finally:
+                release.set()
+                copier.join()
+
+            assert state._async_holders == 0
+            assert state._copying is False
+
+        asyncio.run(_exercise())