FIX pre-commit

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

@@ -186,8 +186,10 @@ class AgentService:
 
             # Wait for sandbox to be ready (it was already acquired in create_id_and_sandbox)
             max_attempts = 60  # Increased timeout to 2 minutes (60 * 2s)
+            last_state = None
             for attempt in range(max_attempts):
                 response = await self.sandbox_service.acquire_sandbox(message_id)
+                last_state = response.state
                 if response.sandbox is not None and response.state == "ready":
                     sandbox = response.sandbox
                     break
@@ -196,10 +198,18 @@ class AgentService:
                     logger.warning(
                         f"Sandbox pool limit reached for {message_id}, attempting cleanup of stuck/expired sandboxes"
                     )
-                    await self.sandbox_service.cleanup_stuck_creating_sandboxes()
-                    await self.sandbox_service.cleanup_expired_ready_sandboxes()
+                    cleaned_creating = (
+                        await self.sandbox_service.cleanup_stuck_creating_sandboxes()
+                    )
+                    cleaned_expired = (
+                        await self.sandbox_service.cleanup_expired_ready_sandboxes()
+                    )
+                    logger.info(
+                        f"Cleanup completed: removed {cleaned_creating} stuck creating + {cleaned_expired} expired ready sandboxes"
+                    )
                     # Try one more time after cleanup
                     response = await self.sandbox_service.acquire_sandbox(message_id)
+                    last_state = response.state
                     if response.sandbox is not None and response.state == "ready":
                         sandbox = response.sandbox
                         break
@@ -217,17 +227,40 @@ class AgentService:
                         f"Waiting for sandbox for {message_id}, attempt {attempt}/{max_attempts}, state: {response.state}"
                     )
                 await asyncio.sleep(2)
+
+            # If sandbox is still None after all attempts, do final cleanup and check
             if sandbox is None:
-                # Final cleanup attempt before raising error
-                await self.sandbox_service.cleanup_stuck_creating_sandboxes()
-                await self.sandbox_service.cleanup_expired_ready_sandboxes()
+                logger.warning(
+                    f"Sandbox for {message_id} still not ready after {max_attempts} attempts (last state: {last_state}), performing final cleanup"
+                )
+                # Final cleanup attempt before raising error - be more aggressive
+                cleaned_creating = (
+                    await self.sandbox_service.cleanup_stuck_creating_sandboxes()
+                )
+                cleaned_expired = (
+                    await self.sandbox_service.cleanup_expired_ready_sandboxes()
+                )
+                logger.info(
+                    f"Final cleanup: removed {cleaned_creating} stuck creating + {cleaned_expired} expired ready sandboxes"
+                )
+
+                # Try one last time after cleanup
                 (
                     available_count,
                     non_available_count,
                 ) = await self.sandbox_service.get_sandbox_counts()
-                raise Exception(
-                    f"No sandbox available: pool limit reached (available: {available_count}, non-available: {non_available_count}, max: {self.max_sandboxes})"
+                # Provide more detailed error message
+                error_msg = (
+                    f"No sandbox available for {message_id}: "
+                    f"available: {available_count}, non-available: {non_available_count}, "
+                    f"max: {self.max_sandboxes}, last_state: {last_state}"
                 )
+                if non_available_count > 0:
+                    error_msg += (
+                        f". There are {non_available_count} sandbox(s) stuck in 'creating' state "
+                        f"that may need manual cleanup or the cleanup threshold may be too high."
+                    )
+                raise Exception(error_msg)
 
             data_dir = self.active_tasks[message_id].trace_path
             user_content = self.active_tasks[message_id].instruction

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

@@ -9,9 +9,11 @@ from pydantic import BaseModel
 
 SANDBOX_METADATA: dict[str, dict[str, Any]] = {}
 SANDBOX_TIMEOUT = 500
-SANDBOX_CREATION_TIMEOUT = 200
+SANDBOX_READY_TIMEOUT = 200
 SANDBOX_CREATION_MAX_TIME = (
-    120  # Maximum time a sandbox can be in "creating" state (2 minutes, reduced from 5)
+    90  # Maximum time a sandbox can be in "creating" state (90 seconds)
+    # Reduced from 120 to be more aggressive and clean up stuck sandboxes
+    # before the agent_service loop times out (which waits 60 attempts * 2s = 120s)
 )
 WIDTH = 1280
 HEIGHT = 960
@@ -167,7 +169,7 @@ class SandboxService:
                 and (
                     current_time - self.sandbox_metadata[session_hash]["created_at"]
                 ).total_seconds()
-                < SANDBOX_CREATION_TIMEOUT
+                < SANDBOX_READY_TIMEOUT
             ):
                 print(f"Reusing Sandbox for session {session_hash}")
                 self.sandbox_metadata[session_hash]["last_accessed"] = current_time
@@ -180,8 +182,58 @@ class SandboxService:
                 session_hash in self.sandbox_metadata
                 and self.sandbox_metadata[session_hash].get("state") == "creating"
             ):
-                print(f"Sandbox for session {session_hash} is already being created")
-                return SandboxResponse(sandbox=None, state="creating")
+                # Check if this sandbox has been stuck in "creating" state for too long
+                created_at = self.sandbox_metadata[session_hash].get("created_at")
+                if created_at:
+                    stuck_duration = (current_time - created_at).total_seconds()
+                    if stuck_duration > SANDBOX_CREATION_MAX_TIME:
+                        # This sandbox is stuck - clean it up immediately
+                        print(
+                            f"Sandbox for session {session_hash} has been stuck in 'creating' state "
+                            f"for {stuck_duration:.1f}s (threshold: {SANDBOX_CREATION_MAX_TIME}s) - cleaning up"
+                        )
+                        # Remove from metadata and sandboxes dict
+                        if session_hash in self.sandboxes:
+                            # Schedule kill outside of lock with error handling
+                            stuck_sandbox = self.sandboxes[session_hash]
+                            del self.sandboxes[session_hash]
+
+                            async def kill_stuck():
+                                try:
+                                    await asyncio.to_thread(stuck_sandbox.kill)
+                                except Exception as e:
+                                    print(
+                                        f"Error killing stuck sandbox for {session_hash}: {str(e)}"
+                                    )
+
+                            asyncio.create_task(kill_stuck())
+                        del self.sandbox_metadata[session_hash]
+                        # Fall through to create a new sandbox
+                    else:
+                        print(
+                            f"Sandbox for session {session_hash} is already being created"
+                        )
+                        return SandboxResponse(sandbox=None, state="creating")
+                else:
+                    # Missing created_at - corrupted metadata, clean it up
+                    print(
+                        f"WARNING: Sandbox {session_hash} in 'creating' state has no 'created_at' - cleaning up"
+                    )
+                    if session_hash in self.sandboxes:
+                        stuck_sandbox = self.sandboxes[session_hash]
+                        del self.sandboxes[session_hash]
+
+                        async def kill_stuck():
+                            try:
+                                await asyncio.to_thread(stuck_sandbox.kill)
+                            except Exception as e:
+                                print(
+                                    f"Error killing stuck sandbox for {session_hash}: {str(e)}"
+                                )
+
+                        asyncio.create_task(kill_stuck())
+                    del self.sandbox_metadata[session_hash]
+                    # Fall through to create a new sandbox
 
             # Mark expired sandbox for cleanup (remove from dict within lock)
             if session_hash in self.sandboxes:
@@ -298,14 +350,28 @@ class SandboxService:
             for session_hash, metadata in list(self.sandbox_metadata.items()):
                 if metadata.get("state") == "creating":
                     created_at = metadata.get("created_at")
-                    if (
-                        created_at
-                        and (current_time - created_at).total_seconds()
-                        > SANDBOX_CREATION_MAX_TIME
-                    ):
+                    # Clean up if:
+                    # 1. created_at exists and is older than threshold, OR
+                    # 2. created_at is missing (corrupted metadata - should never happen but handle it)
+                    should_cleanup = False
+                    stuck_duration = 0.0
+
+                    if created_at:
+                        stuck_duration = (current_time - created_at).total_seconds()
+                        if stuck_duration > SANDBOX_CREATION_MAX_TIME:
+                            should_cleanup = True
+                    else:
+                        # Missing created_at is a bug, but clean it up anyway
+                        print(
+                            f"WARNING: Sandbox {session_hash} in 'creating' state has no 'created_at' timestamp - cleaning up"
+                        )
+                        should_cleanup = True
+                        stuck_duration = float("inf")
+
+                    if should_cleanup:
                         print(
                             f"Cleaning up stuck 'creating' sandbox for session {session_hash} "
-                            f"(stuck for {(current_time - created_at).total_seconds():.1f}s)"
+                            f"(stuck for {stuck_duration:.1f}s)"
                         )
                         # Collect sandbox to kill if it exists
                         if session_hash in self.sandboxes:
@@ -339,7 +405,7 @@ class SandboxService:
                     if (
                         created_at
                         and (current_time - created_at).total_seconds()
-                        >= SANDBOX_CREATION_TIMEOUT
+                        >= SANDBOX_READY_TIMEOUT
                     ):
                         print(
                             f"Cleaning up expired ready sandbox for session {session_hash} "