NEW workflow

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

@@ -127,10 +127,32 @@ class AgentService:
         trace.steps = []
         trace.traceMetadata = AgentTraceMetadata(traceId=trace_id)
 
+        trace_id_to_release = None
         async with self._lock:
             if self.task_websockets[trace_id] != websocket:
-                raise WebSocketException("WebSocket mismatch")
+                # Release sandbox before raising exception to prevent leak
+                # Do this outside the lock to avoid deadlock
+                trace_id_to_release = trace_id
+                # Remove from task_websockets since we're rejecting this
+                if trace_id in self.task_websockets:
+                    del self.task_websockets[trace_id]
+
+        # Release sandbox outside of lock if there was a mismatch
+        if trace_id_to_release:
+            try:
+                await self.sandbox_service.release_sandbox(trace_id_to_release)
+                logger.info(
+                    f"Released sandbox for {trace_id_to_release} due to WebSocket mismatch"
+                )
+            except Exception as e:
+                logger.error(
+                    f"Error releasing sandbox for {trace_id_to_release}: {e}",
+                    exc_info=True,
+                )
+            raise WebSocketException("WebSocket mismatch")
 
+        # Continue with normal processing if no mismatch
+        async with self._lock:
             active_task = ActiveTask(
                 message_id=trace_id,
                 instruction=trace.instruction,
@@ -184,83 +206,67 @@ class AgentService:
 
             model = get_model(self.active_tasks[message_id].model_id)
 
-            # 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
+            # Wait for sandbox to be ready
+            max_attempts = 60  # 2 minutes timeout (60 * 2s)
+            sandbox = None
             for attempt in range(max_attempts):
                 response = await self.sandbox_service.acquire_sandbox(message_id)
-                last_state = response.state
+
+                # Check for creation errors
+                if response.error:
+                    logger.error(
+                        f"Sandbox creation failed for {message_id}: {response.error}",
+                        exc_info=False,
+                    )
+                    # Continue retrying - might succeed on next attempt
+                    await asyncio.sleep(2)
+                    continue
+
                 if response.sandbox is not None and response.state == "ready":
                     sandbox = response.sandbox
                     break
-                elif response.state == "max_sandboxes_reached":
-                    # Trigger cleanup of stuck and expired sandboxes before giving up
+
+                if response.state == "max_sandboxes_reached":
+                    # Service handles cleanup automatically, but log the state
+                    (
+                        available_count,
+                        pending_count,
+                    ) = await self.sandbox_service.get_sandbox_counts()
                     logger.warning(
-                        f"Sandbox pool limit reached for {message_id}, attempting cleanup of stuck/expired sandboxes"
-                    )
-                    cleaned_creating = (
-                        await self.sandbox_service.cleanup_stuck_creating_sandboxes()
+                        f"Sandbox pool at capacity for {message_id}: "
+                        f"{available_count} ready, {pending_count} pending, max: {self.max_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
-                    elif response.state == "max_sandboxes_reached":
-                        (
-                            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})"
-                        )
+                    # Wait a bit and retry - cleanup may free up space
+                    await asyncio.sleep(2)
+                    continue
+
                 # Log progress every 10 attempts
                 if attempt > 0 and attempt % 10 == 0:
                     logger.info(
                         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
+            # Check if we got a sandbox
             if sandbox is None:
-                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,
+                    pending_count,
                 ) = await self.sandbox_service.get_sandbox_counts()
-                # 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."
+                # Check for any final errors
+                final_response = await self.sandbox_service.acquire_sandbox(message_id)
+                error_info = ""
+                if final_response.error:
+                    error_info = f" Last creation error: {final_response.error}"
+                    logger.error(
+                        f"Sandbox creation failed for {message_id} after {max_attempts} attempts: {final_response.error}",
+                        exc_info=False,
                     )
-                raise Exception(error_msg)
+                raise Exception(
+                    f"No sandbox available for {message_id} after {max_attempts} attempts: "
+                    f"{available_count} ready, {pending_count} pending, max: {self.max_sandboxes}.{error_info}"
+                )
 
             data_dir = self.active_tasks[message_id].trace_path
             user_content = self.active_tasks[message_id].instruction
@@ -317,9 +323,14 @@ class AgentService:
                 f"Error processing task: {traceback.format_exc()}", exc_info=True
             )
             final_state = "error"
-            await self.websocket_manager.send_agent_error(
-                error="Error processing task", websocket=websocket
-            )
+            if (
+                not websocket_exception
+                and websocket
+                and websocket.client_state == WebSocketState.CONNECTED
+            ):
+                await self.websocket_manager.send_agent_error(
+                    error="Error processing task", websocket=websocket
+                )
 
         finally:
             # Send completion event

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

@@ -2,19 +2,13 @@ import asyncio
 import os
 import time
 from datetime import datetime
-from typing import Any, Literal
+from typing import Literal
 
 from e2b_desktop import Sandbox
 from pydantic import BaseModel
 
-SANDBOX_METADATA: dict[str, dict[str, Any]] = {}
 SANDBOX_TIMEOUT = 500
-SANDBOX_READY_TIMEOUT = 200
-SANDBOX_CREATION_MAX_TIME = (
-    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)
-)
+SANDBOX_READY_TIMEOUT = 200  # Seconds before a sandbox expires
 WIDTH = 1280
 HEIGHT = 960
 
@@ -24,116 +18,242 @@ class SandboxResponse(BaseModel):
 
     sandbox: Sandbox | None
     state: Literal["creating", "ready", "max_sandboxes_reached"]
+    error: str | None = None  # Error message if creation failed
+
+
+class SandboxEntry:
+    """Simple container for sandbox and its metadata"""
+
+    def __init__(self, sandbox: Sandbox):
+        self.sandbox = sandbox
+        self.created_at = datetime.now()
+        self.last_accessed = datetime.now()
+
+    def is_expired(self) -> bool:
+        """Check if sandbox has expired"""
+        age = (datetime.now() - self.created_at).total_seconds()
+        return age >= SANDBOX_READY_TIMEOUT
+
+    def update_access(self):
+        """Update last access time"""
+        self.last_accessed = datetime.now()
 
 
 class SandboxService:
+    """
+    Simplified sandbox service for production use.
+
+    Key simplifications:
+    - Non-blocking sandbox creation (background tasks)
+    - Simple expiration-based cleanup
+    - Minimal state tracking (pending vs ready)
+    - Straightforward locking
+    """
+
     def __init__(self, max_sandboxes: int = 50):
         if not os.getenv("E2B_API_KEY"):
             raise ValueError("E2B_API_KEY is not set")
         self.max_sandboxes = max_sandboxes
-        self.sandboxes: dict[str, Sandbox] = {}
-        self.sandbox_metadata: dict[str, dict[str, Any]] = {}
-        self.sandbox_lock = asyncio.Lock()
+        self.sandboxes: dict[str, SandboxEntry] = {}  # Ready sandboxes
+        self.pending: set[str] = set()  # Session hashes currently being created
+        self.creation_errors: dict[
+            str, str
+        ] = {}  # Track creation errors by session_hash
+        self.lock = asyncio.Lock()
         self._cleanup_task: asyncio.Task | None = None
 
-    async def _create_sandbox_background(
-        self, session_hash: str, expired_sandbox: Sandbox | None
-    ):
-        """Background task to create and setup a sandbox."""
-        # Kill expired sandbox first
-        if expired_sandbox:
-            try:
-                print(f"Closing expired sandbox for session {session_hash}")
-                await asyncio.to_thread(expired_sandbox.kill)
-            except Exception as e:
-                print(f"Error closing expired sandbox: {str(e)}")
-
-        def create_and_setup_sandbox():
-            desktop = Sandbox.create(
-                api_key=os.getenv("E2B_API_KEY"),
-                resolution=(WIDTH, HEIGHT),
-                dpi=96,
-                timeout=SANDBOX_TIMEOUT,
-                template="k0wmnzir0zuzye6dndlw",
-            )
-            desktop.stream.start(require_auth=True)
-            setup_cmd = """sudo mkdir -p /usr/lib/firefox-esr/distribution && echo '{"policies":{"OverrideFirstRunPage":"","OverridePostUpdatePage":"","DisableProfileImport":true,"DontCheckDefaultBrowser":true}}' | sudo tee /usr/lib/firefox-esr/distribution/policies.json > /dev/null"""
-            desktop.commands.run(setup_cmd)
-            time.sleep(3)
-            return desktop
+    def _create_and_setup_sandbox(self) -> Sandbox:
+        """Create and setup a sandbox (synchronous operation)"""
+        desktop = Sandbox.create(
+            api_key=os.getenv("E2B_API_KEY"),
+            resolution=(WIDTH, HEIGHT),
+            dpi=96,
+            timeout=SANDBOX_TIMEOUT,
+            template="k0wmnzir0zuzye6dndlw",
+        )
+        desktop.stream.start(require_auth=True)
+        setup_cmd = """sudo mkdir -p /usr/lib/firefox-esr/distribution && echo '{"policies":{"OverrideFirstRunPage":"","OverridePostUpdatePage":"","DisableProfileImport":true,"DontCheckDefaultBrowser":true}}' | sudo tee /usr/lib/firefox-esr/distribution/policies.json > /dev/null"""
+        desktop.commands.run(setup_cmd)
+        time.sleep(3)
+        return desktop
+
+    async def acquire_sandbox(self, session_hash: str) -> SandboxResponse:
+        """
+        Acquire a sandbox for a session.
+        Returns immediately - either with ready sandbox, or "creating" if one is being created.
+        """
+        async with self.lock:
+            # Check if we have a valid sandbox for this session
+            if session_hash in self.sandboxes:
+                entry = self.sandboxes[session_hash]
+                if not entry.is_expired():
+                    entry.update_access()
+                    print(f"Reusing sandbox for session {session_hash}")
+                    return SandboxResponse(sandbox=entry.sandbox, state="ready")
+                else:
+                    # Expired - remove it
+                    print(f"Removing expired sandbox for session {session_hash}")
+                    old_entry = self.sandboxes.pop(session_hash)
+                    # Schedule cleanup outside lock
+                    asyncio.create_task(
+                        self._kill_sandbox_safe(old_entry.sandbox, session_hash)
+                    )
+
+            # Check if already being created
+            if session_hash in self.pending:
+                print(f"Sandbox for session {session_hash} is already being created")
+                # Check if there was a creation error
+                if session_hash in self.creation_errors:
+                    error_msg = self.creation_errors.pop(session_hash)
+                    return SandboxResponse(
+                        sandbox=None, state="creating", error=error_msg
+                    )
+                return SandboxResponse(sandbox=None, state="creating")
 
+            # Check if there's a previous creation error (shouldn't happen, but handle it)
+            if session_hash in self.creation_errors:
+                error_msg = self.creation_errors.pop(session_hash)
+                return SandboxResponse(sandbox=None, state="creating", error=error_msg)
+
+            # Check capacity (count both ready and pending)
+            total_count = len(self.sandboxes) + len(self.pending)
+            if total_count >= self.max_sandboxes:
+                print(
+                    f"Sandbox pool at capacity: {len(self.sandboxes)} ready + {len(self.pending)} pending = {total_count}/{self.max_sandboxes}"
+                )
+                # Try to clean up expired sandboxes first
+                await self._cleanup_expired_internal()
+                # Recheck capacity after cleanup
+                total_count = len(self.sandboxes) + len(self.pending)
+                if total_count >= self.max_sandboxes:
+                    return SandboxResponse(sandbox=None, state="max_sandboxes_reached")
+
+            # Mark as pending and start creation in background
+            self.pending.add(session_hash)
+            print(f"Starting creation of sandbox for session {session_hash}")
+
+        # Start creation in background (non-blocking)
+        asyncio.create_task(self._create_sandbox_background(session_hash))
+        return SandboxResponse(sandbox=None, state="creating")
+
+    async def _create_sandbox_background(self, session_hash: str):
+        """Background task to create a sandbox"""
         desktop = None
         try:
-            desktop = await asyncio.to_thread(create_and_setup_sandbox)
-            print(f"Sandbox ID for session {session_hash} is {desktop.sandbox_id}.")
-
-            # Update sandbox state under lock
-            async with self.sandbox_lock:
-                # Double-check metadata still exists and is in "creating" state
-                # (it might have been released while we were creating)
-                if (
-                    session_hash in self.sandbox_metadata
-                    and self.sandbox_metadata[session_hash].get("state") == "creating"
-                ):
-                    self.sandboxes[session_hash] = desktop
-                    self.sandbox_metadata[session_hash]["state"] = "ready"
-                    print(f"Sandbox {session_hash} is now ready")
-                else:
-                    # Sandbox was released while creating, kill it immediately
+            desktop = await asyncio.to_thread(self._create_and_setup_sandbox)
+            print(
+                f"Sandbox created for session {session_hash}, ID: {desktop.sandbox_id}"
+            )
+
+            async with self.lock:
+                # Check if session was released while creating (removed from pending)
+                was_released = session_hash not in self.pending
+                self.pending.discard(session_hash)
+
+                if was_released:
+                    # Session was released - kill the sandbox
                     print(
-                        f"Sandbox {session_hash} was released during creation, killing it"
+                        f"Session {session_hash} was released during creation, killing sandbox"
                     )
-                    try:
-                        await asyncio.to_thread(desktop.kill)
-                    except Exception as kill_error:
-                        print(f"Error killing orphaned sandbox: {str(kill_error)}")
+                    asyncio.create_task(self._kill_sandbox_safe(desktop, session_hash))
+                    return
 
-        except Exception as e:
-            print(f"Error creating sandbox for session {session_hash}: {str(e)}")
-            # Clean up metadata on failure - CRITICAL to prevent leaks
-            async with self.sandbox_lock:
-                if session_hash in self.sandbox_metadata:
-                    state = self.sandbox_metadata[session_hash].get("state")
+                # Check capacity before adding
+                if len(self.sandboxes) >= self.max_sandboxes:
                     print(
-                        f"Cleaning up failed sandbox creation for {session_hash} (state was: {state})"
+                        f"Pool at capacity, killing newly created sandbox for {session_hash}"
                     )
-                    del self.sandbox_metadata[session_hash]
-                # Also remove from sandboxes dict if it somehow got added
-                if session_hash in self.sandboxes:
-                    del self.sandboxes[session_hash]
-            # Kill the sandbox if it was partially created
-            if desktop is not None:
-                try:
-                    await asyncio.to_thread(desktop.kill)
-                    print(f"Killed partially created sandbox for {session_hash}")
-                except Exception as kill_error:
-                    print(f"Error killing partially created sandbox: {str(kill_error)}")
+                    asyncio.create_task(self._kill_sandbox_safe(desktop, session_hash))
+                    return
+
+                # Add to pool
+                self.sandboxes[session_hash] = SandboxEntry(desktop)
+                print(f"Sandbox {session_hash} is now ready")
+
+        except Exception as e:
+            error_msg = str(e)
+            import traceback
+
+            error_details = traceback.format_exc()
+            print(f"Error creating sandbox for session {session_hash}: {error_msg}")
+            print(f"Full traceback: {error_details}")
+
+            async with self.lock:
+                self.pending.discard(session_hash)
+                # Store error so agent service can retrieve it
+                self.creation_errors[session_hash] = error_msg
+            if desktop:
+                asyncio.create_task(self._kill_sandbox_safe(desktop, session_hash))
+
+    async def release_sandbox(self, session_hash: str):
+        """Release a sandbox for a session"""
+        sandbox = None
+        async with self.lock:
+            if session_hash in self.sandboxes:
+                entry = self.sandboxes.pop(session_hash)
+                sandbox = entry.sandbox
+            # Also remove from pending if it's there (creation will be cancelled)
+            self.pending.discard(session_hash)
+            # Clean up any stored errors
+            self.creation_errors.pop(session_hash, None)
+
+        if sandbox:
+            await self._kill_sandbox_safe(sandbox, session_hash)
+            print(f"Released sandbox for session {session_hash}")
+
+    async def _kill_sandbox_safe(self, sandbox: Sandbox, session_hash: str):
+        """Safely kill a sandbox with error handling"""
+        try:
+            await asyncio.to_thread(sandbox.kill)
+        except Exception as e:
+            print(f"Error killing sandbox for session {session_hash}: {str(e)}")
+
+    async def _cleanup_expired_internal(self) -> int:
+        """Internal cleanup of expired sandboxes (must be called with lock held)"""
+        expired = []
+        for session_hash, entry in list(self.sandboxes.items()):
+            if entry.is_expired():
+                expired.append((session_hash, entry.sandbox))
+                del self.sandboxes[session_hash]
+
+        # Kill expired sandboxes outside lock
+        for session_hash, sandbox in expired:
+            await self._kill_sandbox_safe(sandbox, session_hash)
+            print(f"Cleaned up expired sandbox for session {session_hash}")
+
+        return len(expired)
+
+    async def cleanup_expired_ready_sandboxes(self) -> int:
+        """Clean up expired ready sandboxes"""
+        async with self.lock:
+            return await self._cleanup_expired_internal()
+
+    async def get_sandbox_counts(self) -> tuple[int, int]:
+        """
+        Get the count of available (ready) and non-available (pending) sandboxes.
+        Returns: (available_count, non_available_count)
+        """
+        async with self.lock:
+            available = len(self.sandboxes)
+            non_available = len(self.pending)
+            return (available, non_available)
 
     async def _periodic_cleanup(self):
-        """Background task to periodically clean up stuck creating sandboxes and expired ready sandboxes"""
+        """Periodic cleanup task"""
         while True:
             try:
-                await asyncio.sleep(30)  # Run every 30 seconds (more aggressive)
-                cleaned_creating = await self.cleanup_stuck_creating_sandboxes()
-                cleaned_expired = await self.cleanup_expired_ready_sandboxes()
-                if cleaned_creating > 0 or cleaned_expired > 0:
-                    print(
-                        f"Periodic cleanup: removed {cleaned_creating} stuck creating + "
-                        f"{cleaned_expired} expired ready = {cleaned_creating + cleaned_expired} total"
-                    )
-                # Log sandbox pool state periodically for debugging
-                async with self.sandbox_lock:
+                await asyncio.sleep(60)  # Run every minute
+                async with self.lock:
+                    cleaned = await self._cleanup_expired_internal()
+                    if cleaned > 0:
+                        print(f"Periodic cleanup: removed {cleaned} expired sandboxes")
+                    # Log pool state
                     ready_count = len(self.sandboxes)
-                    creating_count = sum(
-                        1
-                        for meta in self.sandbox_metadata.values()
-                        if meta.get("state") == "creating"
-                    )
-                    total_count = ready_count + creating_count
-                    if total_count > 0:
+                    pending_count = len(self.pending)
+                    total = ready_count + pending_count
+                    if total > 0:
                         print(
-                            f"Sandbox pool state: {ready_count} ready, {creating_count} creating, "
-                            f"{total_count}/{self.max_sandboxes} total"
+                            f"Sandbox pool: {ready_count} ready, {pending_count} pending, {total}/{self.max_sandboxes} total"
                         )
             except asyncio.CancelledError:
                 break
@@ -146,7 +266,6 @@ class SandboxService:
             try:
                 self._cleanup_task = asyncio.create_task(self._periodic_cleanup())
             except RuntimeError as e:
-                # If called outside event loop, log but don't crash
                 print(f"Warning: Cannot start periodic cleanup (no event loop): {e}")
 
     def stop_periodic_cleanup(self):
@@ -154,314 +273,14 @@ class SandboxService:
         if self._cleanup_task and not self._cleanup_task.done():
             self._cleanup_task.cancel()
 
-    async def acquire_sandbox(self, session_hash: str) -> SandboxResponse:
-        current_time = datetime.now()
-        should_create = False
-        expired_sandbox = None
-
-        # Quick check under lock - only check state and mark creation
-        async with self.sandbox_lock:
-            # Check if sandbox exists and is ready
-            if (
-                session_hash in self.sandboxes
-                and session_hash in self.sandbox_metadata
-                and self.sandbox_metadata[session_hash].get("state") == "ready"
-                and (
-                    current_time - self.sandbox_metadata[session_hash]["created_at"]
-                ).total_seconds()
-                < SANDBOX_READY_TIMEOUT
-            ):
-                print(f"Reusing Sandbox for session {session_hash}")
-                self.sandbox_metadata[session_hash]["last_accessed"] = current_time
-                return SandboxResponse(
-                    sandbox=self.sandboxes[session_hash], state="ready"
-                )
-
-            # Check if sandbox is already being created
-            if (
-                session_hash in self.sandbox_metadata
-                and self.sandbox_metadata[session_hash].get("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:
-                print(f"Marking expired sandbox for session {session_hash} for cleanup")
-                expired_sandbox = self.sandboxes[session_hash]
-                del self.sandboxes[session_hash]
-                if session_hash in self.sandbox_metadata:
-                    del self.sandbox_metadata[session_hash]
-
-            # Check if we have capacity
-            # Count both ready sandboxes and sandboxes in "creating" state
-            # We count BEFORE adding this one to ensure we don't exceed the limit
-            creating_count = sum(
-                1
-                for meta in self.sandbox_metadata.values()
-                if meta.get("state") == "creating"
-            )
-            ready_count = len(self.sandboxes)
-            total_count = ready_count + creating_count
-            # Check capacity BEFORE adding this session_hash to metadata
-            if total_count >= self.max_sandboxes:
-                print(
-                    f"Sandbox pool at capacity: {ready_count} ready + {creating_count} creating = "
-                    f"{total_count}/{self.max_sandboxes}"
-                )
-                # CRITICAL: If we have an expired sandbox but can't create a new one,
-                # we must still kill the expired sandbox to prevent leaks
-                if expired_sandbox:
-                    print(
-                        f"Killing expired sandbox for {session_hash} even though pool is at capacity"
-                    )
-
-                    async def kill_expired():
-                        try:
-                            await asyncio.to_thread(expired_sandbox.kill)
-                        except Exception as e:
-                            print(
-                                f"Error killing expired sandbox for {session_hash}: {str(e)}"
-                            )
-
-                    asyncio.create_task(kill_expired())
-                return SandboxResponse(sandbox=None, state="max_sandboxes_reached")
-
-            # Mark that we're creating this sandbox
-            # This happens atomically within the lock, so no race condition
-            print(f"Creating new sandbox for session {session_hash}")
-            self.sandbox_metadata[session_hash] = {
-                "state": "creating",
-                "created_at": current_time,
-                "last_accessed": current_time,
-            }
-            should_create = True
-
-        # Start sandbox creation in background without waiting
-        if should_create:
-            asyncio.create_task(
-                self._create_sandbox_background(session_hash, expired_sandbox)
-            )
-        elif expired_sandbox:
-            # If we're not creating but have an expired sandbox, kill it
-            # This shouldn't normally happen, but handle it defensively
-            print(f"Killing expired sandbox for {session_hash} (not creating new one)")
-            try:
-                await asyncio.to_thread(expired_sandbox.kill)
-            except Exception as e:
-                print(f"Error killing expired sandbox: {str(e)}")
-
-        # Check state after starting background task (it might complete very quickly)
-        async with self.sandbox_lock:
-            if session_hash in self.sandbox_metadata:
-                state = self.sandbox_metadata[session_hash].get("state")
-                if state == "creating":
-                    return SandboxResponse(sandbox=None, state="creating")
-                if state == "ready" and session_hash in self.sandboxes:
-                    return SandboxResponse(
-                        sandbox=self.sandboxes[session_hash], state="ready"
-                    )
-
-        # If metadata doesn't exist, it means creation failed immediately
-        # Return "creating" anyway as the caller will retry
-        return SandboxResponse(sandbox=None, state="creating")
-
-    async def release_sandbox(self, session_hash: str):
-        sandbox_to_kill = None
-
-        # Remove from dictionaries under lock
-        async with self.sandbox_lock:
-            if session_hash in self.sandboxes:
-                print(f"Releasing sandbox for session {session_hash}")
-                sandbox_to_kill = self.sandboxes[session_hash]
-                del self.sandboxes[session_hash]
-            # Always clean up metadata, even if sandbox is still in "creating" state
-            if session_hash in self.sandbox_metadata:
-                state = self.sandbox_metadata[session_hash].get("state")
-                if state == "creating":
-                    print(
-                        f"Cleaning up stuck 'creating' sandbox for session {session_hash}"
-                    )
-                del self.sandbox_metadata[session_hash]
-
-        # Kill sandbox outside of lock
-        if sandbox_to_kill:
-            try:
-                await asyncio.to_thread(sandbox_to_kill.kill)
-            except Exception as e:
-                print(f"Error killing sandbox for session {session_hash}: {str(e)}")
-
-    async def cleanup_stuck_creating_sandboxes(self):
-        """Clean up sandboxes that have been stuck in 'creating' state for too long"""
-        current_time = datetime.now()
-        stuck_sandboxes_to_kill = []
-
-        async with self.sandbox_lock:
-            for session_hash, metadata in list(self.sandbox_metadata.items()):
-                if metadata.get("state") == "creating":
-                    created_at = metadata.get("created_at")
-                    # 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 {stuck_duration:.1f}s)"
-                        )
-                        # Collect sandbox to kill if it exists
-                        if session_hash in self.sandboxes:
-                            stuck_sandboxes_to_kill.append(
-                                (session_hash, self.sandboxes[session_hash])
-                            )
-                            del self.sandboxes[session_hash]
-                        del self.sandbox_metadata[session_hash]
-
-        # Kill stuck sandboxes outside of lock
-        for session_hash, sandbox in stuck_sandboxes_to_kill:
-            try:
-                await asyncio.to_thread(sandbox.kill)
-                print(f"Killed stuck sandbox for session {session_hash}")
-            except Exception as e:
-                print(
-                    f"Error killing stuck sandbox for session {session_hash}: {str(e)}"
-                )
-
-        return len(stuck_sandboxes_to_kill)
-
-    async def cleanup_expired_ready_sandboxes(self):
-        """Clean up ready sandboxes that have expired (not accessed for too long)"""
-        current_time = datetime.now()
-        expired_sandboxes_to_kill = []
-
-        async with self.sandbox_lock:
-            for session_hash, metadata in list(self.sandbox_metadata.items()):
-                if metadata.get("state") == "ready" and session_hash in self.sandboxes:
-                    created_at = metadata.get("created_at")
-                    if (
-                        created_at
-                        and (current_time - created_at).total_seconds()
-                        >= SANDBOX_READY_TIMEOUT
-                    ):
-                        print(
-                            f"Cleaning up expired ready sandbox for session {session_hash} "
-                            f"(age: {(current_time - created_at).total_seconds():.1f}s)"
-                        )
-                        expired_sandboxes_to_kill.append(
-                            (session_hash, self.sandboxes[session_hash])
-                        )
-                        del self.sandboxes[session_hash]
-                        del self.sandbox_metadata[session_hash]
-
-        # Kill expired sandboxes outside of lock
-        for session_hash, sandbox in expired_sandboxes_to_kill:
-            try:
-                await asyncio.to_thread(sandbox.kill)
-                print(f"Killed expired ready sandbox for session {session_hash}")
-            except Exception as e:
-                print(
-                    f"Error killing expired ready sandbox for session {session_hash}: {str(e)}"
-                )
-
-        return len(expired_sandboxes_to_kill)
-
-    async def get_sandbox_counts(self) -> tuple[int, int]:
-        """
-        Get the count of available (ready) and non-available (creating) sandboxes.
-
-        Returns:
-            Tuple of (available_count, non_available_count)
-        """
-        async with self.sandbox_lock:
-            ready_count = len(self.sandboxes)
-            creating_count = sum(
-                1
-                for meta in self.sandbox_metadata.values()
-                if meta.get("state") == "creating"
-            )
-            return (ready_count, creating_count)
-
     async def cleanup_sandboxes(self):
-        sandboxes_to_kill = []
+        """Clean up all sandboxes"""
+        async with self.lock:
+            sandboxes_to_kill = list(self.sandboxes.values())
+            self.sandboxes.clear()
 
-        # Collect sandboxes under lock
-        async with self.sandbox_lock:
-            for session_hash in list(self.sandboxes.keys()):
-                sandboxes_to_kill.append((session_hash, self.sandboxes[session_hash]))
-                del self.sandboxes[session_hash]
-                if session_hash in self.sandbox_metadata:
-                    del self.sandbox_metadata[session_hash]
-
-        # Kill all sandboxes outside of lock
-        for session_hash, sandbox in sandboxes_to_kill:
-            try:
-                await asyncio.to_thread(sandbox.kill)
-            except Exception as e:
-                print(f"Error killing sandbox for session {session_hash}: {str(e)}")
+        for entry in sandboxes_to_kill:
+            await self._kill_sandbox_safe(entry.sandbox, "cleanup")
 
 
 if __name__ == "__main__":