FIX race condition

cua2-core/src/cua2_core/models/models.py

@@ -1,6 +1,6 @@
+import asyncio
 import json
 import os
-import threading
 from datetime import datetime
 from typing import Annotated, Literal, Optional
 from uuid import uuid4
@@ -269,51 +269,58 @@ class ActiveTask(BaseModel):
     timestamp: datetime = datetime.now()
     steps: list[AgentStep] = []
     traceMetadata: AgentTraceMetadata = AgentTraceMetadata()
-    _file_lock: threading.Lock = PrivateAttr(default_factory=threading.Lock)
+    _file_lock: asyncio.Lock | None = PrivateAttr(default=None)
+
+    def _get_lock(self) -> asyncio.Lock:
+        """Get or create the async lock (lazy initialization)"""
+        if self._file_lock is None:
+            self._file_lock = asyncio.Lock()
+        return self._file_lock
 
     @property
     def trace_path(self):
         """Trace path"""
         return f"data/trace-{self.message_id}-{self.model_id.replace('/', '-')}"
 
+    def _write_to_file_sync(self):
+        """Synchronous file write helper (used in async context via to_thread)"""
+        self.traceMetadata.traceId = self.message_id
+        os.makedirs(self.trace_path, exist_ok=True)
+        with open(f"{self.trace_path}/tasks.json", "w") as f:
+            json.dump(
+                self.model_dump(
+                    mode="json",
+                    exclude={"_file_lock", "_lock_initialized"},
+                    context={"actions_as_json": True, "image_as_path": True},
+                ),
+                f,
+                indent=2,
+            )
+
     @model_validator(mode="after")
     def store_model(self):
-        """Validate model ID"""
-        with self._file_lock:
-            self.traceMetadata.traceId = self.message_id
-            os.makedirs(self.trace_path, exist_ok=True)
-            with open(f"{self.trace_path}/tasks.json", "w") as f:
-                json.dump(
-                    self.model_dump(
-                        mode="json",
-                        exclude={"_file_locks"},
-                        context={"actions_as_json": True, "image_as_path": True},
-                    ),
-                    f,
-                    indent=2,
-                )
+        """Validate model ID - creates directory, but file write is deferred to async method"""
+        self.traceMetadata.traceId = self.message_id
+        os.makedirs(self.trace_path, exist_ok=True)
         return self
 
-    def update_step(self, step: AgentStep):
+    async def save_to_file(self):
+        """Async method to save task data to file"""
+        async with self._get_lock():
+            await asyncio.to_thread(self._write_to_file_sync)
+
+    async def update_step(self, step: AgentStep):
         """Update step"""
-        with self._file_lock:
+        async with self._get_lock():
             if int(step.stepId) <= len(self.steps):
                 self.steps[int(step.stepId) - 1] = step
             else:
                 self.steps.append(step)
                 self.traceMetadata.numberOfSteps = len(self.steps)
-            with open(f"{self.trace_path}/tasks.json", "w") as f:
-                json.dump(
-                    self.model_dump(
-                        mode="json",
-                        exclude={"_file_locks"},
-                        context={"actions_as_json": True, "image_as_path": True},
-                    ),
-                    f,
-                    indent=2,
-                )
-
-    def update_trace_metadata(
+            # Use to_thread for file I/O to avoid blocking
+            await asyncio.to_thread(self._write_to_file_sync)
+
+    async def update_trace_metadata(
         self,
         step_input_tokens_used: int | None = None,
         step_output_tokens_used: int | None = None,
@@ -327,7 +334,7 @@ class ActiveTask(BaseModel):
         user_evaluation: Literal["success", "failed", "not_evaluated"] | None = None,
     ):
         """Update trace metadata"""
-        with self._file_lock:
+        async with self._get_lock():
             if step_input_tokens_used is not None:
                 self.traceMetadata.inputTokensUsed += step_input_tokens_used
             if step_output_tokens_used is not None:

cua2-core/src/cua2_core/routes/routes.py

@@ -74,7 +74,7 @@ async def update_trace_step(
 ):
     """Update a specific step in a trace (e.g., update step evaluation)"""
     try:
-        agent_service.update_trace_step(
+        await agent_service.update_trace_step(
             trace_id=trace_id,
             step_id=step_id,
             step_evaluation=request.step_evaluation,
@@ -99,7 +99,7 @@ async def update_trace_evaluation(
 ):
     """Update the user evaluation for a trace (overall task feedback)"""
     try:
-        agent_service.update_trace_evaluation(
+        await agent_service.update_trace_evaluation(
             trace_id=trace_id,
             user_evaluation=request.user_evaluation,
         )

cua2-core/src/cua2_core/services/agent_service.py

@@ -104,9 +104,13 @@ class AgentService:
         """
         Update the archival service with current active task IDs.
         Should be called whenever tasks are added or removed.
+        Note: This should be called while holding self._lock to ensure consistent snapshot.
+        The archival service update itself is fast and non-blocking.
         """
         if self.archival_service.is_alive():
-            self.archival_service.update_active_tasks(set(self.active_tasks.keys()))
+            # Create a snapshot of active task IDs (should be called with lock held)
+            active_task_ids = set(self.active_tasks.keys())
+            self.archival_service.update_active_tasks(active_task_ids)
 
     async def create_id_and_sandbox(self, websocket: WebSocket) -> str:
         """Create a new ID and sandbox"""
@@ -174,8 +178,8 @@ class AgentService:
             self.active_tasks[trace_id] = active_task
             self.last_screenshot[trace_id] = None
 
-        # Update archival service with new active task
-        self._update_archival_active_tasks()
+            # Update archival service with new active task (while holding lock)
+            self._update_archival_active_tasks()
 
         asyncio.create_task(self._agent_processing(trace_id))
 
@@ -351,13 +355,13 @@ class AgentService:
 
             novnc_active = False
 
-            self.active_tasks[message_id].update_trace_metadata(
+            await self.active_tasks[message_id].update_trace_metadata(
                 final_state=final_state,
                 completed=True,
             )
 
             if message_id in self.active_tasks:
-                self.active_tasks[message_id].store_model()
+                await self.active_tasks[message_id].save_to_file()
 
             # Clean up
             async with self._lock:
@@ -370,8 +374,8 @@ class AgentService:
                 if message_id in self.last_screenshot:
                     del self.last_screenshot[message_id]
 
-            # Update archival service after task removal
-            self._update_archival_active_tasks()
+                # Update archival service after task removal (while holding lock)
+                self._update_archival_active_tasks()
 
             # Always release sandbox back to the pool, even if it's still in "creating" state
             # This handles cases where acquire_sandbox was called but sandbox never became ready
@@ -469,14 +473,23 @@ class AgentService:
                         step_evaluation="neutral",
                     )
 
-                    self.active_tasks[message_id].update_trace_metadata(
-                        step_input_tokens_used=memory_step.token_usage.input_tokens,
-                        step_output_tokens_used=memory_step.token_usage.output_tokens,
-                        step_duration=memory_step.timing.duration,
-                        step_numberOfSteps=1,
+                    # Schedule async operations in the event loop (callback runs in worker thread)
+                    future1 = asyncio.run_coroutine_threadsafe(
+                        self.active_tasks[message_id].update_trace_metadata(
+                            step_input_tokens_used=memory_step.token_usage.input_tokens,
+                            step_output_tokens_used=memory_step.token_usage.output_tokens,
+                            step_duration=memory_step.timing.duration,
+                            step_numberOfSteps=1,
+                        ),
+                        loop,
                     )
-
-                    self.active_tasks[message_id].update_step(step)
+                    future2 = asyncio.run_coroutine_threadsafe(
+                        self.active_tasks[message_id].update_step(step),
+                        loop,
+                    )
+                    # Wait for both to complete
+                    future1.result()
+                    future2.result()
 
                     websocket = self.task_websockets.get(message_id)
                     if websocket and websocket.client_state == WebSocketState.CONNECTED:
@@ -529,7 +542,7 @@ class AgentService:
             # Re-raise to ensure error is logged
             raise
 
-    def update_trace_step(
+    async def update_trace_step(
         self,
         trace_id: str,
         step_id: str,
@@ -559,7 +572,8 @@ class AgentService:
                 step_index = int(step_id) - 1
                 if 0 <= step_index < len(active_task.steps):
                     active_task.steps[step_index].step_evaluation = step_evaluation
-                    active_task.update_step(active_task.steps[step_index])
+                    await active_task.update_step(active_task.steps[step_index])
+                    return active_task.steps[step_index]
                 else:
                     raise ValueError(f"Step {step_id} not found in trace")
             except (ValueError, TypeError) as e:
@@ -602,7 +616,7 @@ class AgentService:
             except (ValueError, KeyError, TypeError) as e:
                 raise ValueError(f"Error processing step update: {e}")
 
-    def update_trace_evaluation(
+    async def update_trace_evaluation(
         self,
         trace_id: str,
         user_evaluation: Literal["success", "failed", "not_evaluated"],
@@ -622,7 +636,7 @@ class AgentService:
 
         if active_task:
             # Task is still active
-            active_task.update_trace_metadata(user_evaluation=user_evaluation)
+            await active_task.update_trace_metadata(user_evaluation=user_evaluation)
         else:
             # Task is not active, try to load from file
             data_dir = "data"
@@ -657,7 +671,7 @@ class AgentService:
     async def stop_task(self, trace_id: str):
         """Stop a task"""
         if trace_id in self.active_tasks:
-            self.active_tasks[trace_id].update_trace_metadata(
+            await self.active_tasks[trace_id].update_trace_metadata(
                 completed=True,
             )
 
@@ -687,7 +701,7 @@ class AgentService:
             try:
                 # Mark task as completed to stop the agent (if task exists)
                 if message_id in self.active_tasks:
-                    self.active_tasks[message_id].update_trace_metadata(
+                    await self.active_tasks[message_id].update_trace_metadata(
                         completed=True,
                     )
                     logger.info(

cua2-core/src/cua2_core/services/sandbox_service.py

@@ -158,10 +158,14 @@ class SandboxService:
                     asyncio.create_task(self._kill_sandbox_safe(desktop, session_hash))
                     return
 
-                # Check capacity before adding
-                if len(self.sandboxes) >= self.max_sandboxes:
+                # Check total capacity before adding (sandboxes + other pending creations)
+                # Note: We already removed this session_hash from pending, so we check
+                # if adding it to sandboxes would exceed capacity
+                total_count = len(self.sandboxes) + len(self.pending)
+                if total_count >= self.max_sandboxes:
                     print(
-                        f"Pool at capacity, killing newly created sandbox for {session_hash}"
+                        f"Pool at capacity ({total_count}/{self.max_sandboxes}), "
+                        f"killing newly created sandbox for {session_hash}"
                     )
                     asyncio.create_task(self._kill_sandbox_safe(desktop, session_hash))
                     return

cua2-core/tests/test_routes.py

@@ -1,4 +1,4 @@
-from unittest.mock import Mock
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 from cua2_core.models.models import AvailableModelsResponse, UpdateStepResponse
@@ -15,7 +15,9 @@ def mock_agent_service():
     """Fixture to create a mocked AgentService"""
     service = Mock(spec=AgentService)
     service.active_tasks = {}
-    service.update_trace_step = Mock()
+    # update_trace_step is now async, so use AsyncMock
+    service.update_trace_step = AsyncMock(return_value=None)
+    service.update_trace_evaluation = AsyncMock(return_value=None)
     return service
 
 
@@ -112,8 +114,8 @@ class TestUpdateTraceStep:
         step_id = "1"
         request_data = {"step_evaluation": "like"}
 
-        # Mock the service method to succeed
-        mock_agent_service.update_trace_step.return_value = None
+        # Mock the service method to succeed (already set up as AsyncMock in fixture)
+        pass
 
         response = client.patch(
             f"/traces/{trace_id}/steps/{step_id}", json=request_data
@@ -136,8 +138,6 @@ class TestUpdateTraceStep:
         step_id = "2"
         request_data = {"step_evaluation": "dislike"}
 
-        mock_agent_service.update_trace_step.return_value = None
-
         response = client.patch(
             f"/traces/{trace_id}/steps/{step_id}", json=request_data
         )
@@ -154,8 +154,6 @@ class TestUpdateTraceStep:
         step_id = "3"
         request_data = {"step_evaluation": "neutral"}
 
-        mock_agent_service.update_trace_step.return_value = None
-
         response = client.patch(
             f"/traces/{trace_id}/steps/{step_id}", json=request_data
         )
@@ -186,8 +184,8 @@ class TestUpdateTraceStep:
         request_data = {"step_evaluation": "like"}
 
         # Mock the service to raise ValueError
-        mock_agent_service.update_trace_step.side_effect = ValueError(
-            "Invalid step_id format"
+        mock_agent_service.update_trace_step = AsyncMock(
+            side_effect=ValueError("Invalid step_id format")
         )
 
         response = client.patch(
@@ -204,8 +202,8 @@ class TestUpdateTraceStep:
         request_data = {"step_evaluation": "like"}
 
         # Mock the service to raise FileNotFoundError
-        mock_agent_service.update_trace_step.side_effect = FileNotFoundError(
-            "Trace not found"
+        mock_agent_service.update_trace_step = AsyncMock(
+            side_effect=FileNotFoundError("Trace not found")
         )
 
         response = client.patch(
@@ -222,8 +220,8 @@ class TestUpdateTraceStep:
         request_data = {"step_evaluation": "like"}
 
         # Mock the service to raise ValueError for step not found
-        mock_agent_service.update_trace_step.side_effect = ValueError(
-            "Step 999 not found in trace"
+        mock_agent_service.update_trace_step = AsyncMock(
+            side_effect=ValueError("Step 999 not found in trace")
         )
 
         response = client.patch(
@@ -251,8 +249,6 @@ class TestUpdateTraceStep:
         step_id = "1"
         request_data = {"step_evaluation": "like"}
 
-        mock_agent_service.update_trace_step.return_value = None
-
         response = client.patch(
             f"/traces/{trace_id}/steps/{step_id}", json=request_data
         )
@@ -269,8 +265,6 @@ class TestUpdateTraceStep:
         step_id = "1"
         request_data = {"step_evaluation": "like"}
 
-        mock_agent_service.update_trace_step.return_value = None
-
         response = client.patch(
             f"/traces/{trace_id}/steps/{step_id}", json=request_data
         )
@@ -294,8 +288,7 @@ class TestRoutesIntegration:
 
     def test_update_step_endpoint_available(self, client, mock_agent_service):
         """Test that update step endpoint is available"""
-        mock_agent_service.update_trace_step.return_value = None
-
+        # Mock is already set up as AsyncMock in fixture
         response = client.patch(
             "/traces/test/steps/1", json={"step_evaluation": "like"}
         )