mute logs (#13)

* NEW agent

* Handle stop agent feature

* CHG agent max step

* FIX model

* ADD files

* ADD error status (#14)

* CHG error status

* CHG error status

* CHG error status

cua2-core/src/cua2_core/app.py

@@ -21,11 +21,13 @@ async def lifespan(app: FastAPI):
     if not os.getenv("HF_TOKEN"):
         raise ValueError("HF_TOKEN is not set")
 
+    num_workers = int(os.getenv("NUM_WORKERS", "1"))
+
     websocket_manager = WebSocketManager()
 
     sandbox_service = SandboxService()
 
-    agent_service = AgentService(websocket_manager, sandbox_service)
+    agent_service = AgentService(websocket_manager, sandbox_service, num_workers)
 
     # Store services in app state for access in routes
     app.state.websocket_manager = websocket_manager

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

@@ -126,6 +126,10 @@ class AgentTraceMetadata(BaseModel):
     numberOfSteps: int = 0
     maxSteps: int = 0
     completed: bool = False
+    final_state: (
+        Literal["success", "stopped", "max_steps_reached", "error", "sandbox_timeout"]
+        | None
+    ) = None
 
 
 class AgentTrace(BaseModel):
@@ -157,6 +161,7 @@ class AgentStartEvent(BaseModel):
 
     type: Literal["agent_start"] = "agent_start"
     agentTrace: AgentTrace
+    status: Literal["max_sandboxes_reached", "success"] = "success"
 
 
 class AgentProgressEvent(BaseModel):
@@ -172,7 +177,9 @@ class AgentCompleteEvent(BaseModel):
 
     type: Literal["agent_complete"] = "agent_complete"
     traceMetadata: AgentTraceMetadata
-    final_state: Literal["success", "stopped", "max_steps_reached", "error"]
+    final_state: Literal[
+        "success", "stopped", "max_steps_reached", "error", "sandbox_timeout"
+    ]
 
 
 class AgentErrorEvent(BaseModel):
@@ -292,6 +299,10 @@ class ActiveTask(BaseModel):
         step_duration: float | None = None,
         step_numberOfSteps: int | None = None,
         completed: bool | None = None,
+        final_state: Literal[
+            "success", "stopped", "max_steps_reached", "error", "sandbox_timeout"
+        ]
+        | None = None,
     ):
         """Update trace metadata"""
         with self._file_lock:
@@ -305,6 +316,8 @@ class ActiveTask(BaseModel):
                 self.traceMetadata.numberOfSteps += step_numberOfSteps
             if completed is not None:
                 self.traceMetadata.completed = completed
+            if final_state is not None:
+                self.traceMetadata.final_state = final_state
 
 
 #################### API Routes Models ########################

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

@@ -19,7 +19,7 @@ from cua2_core.services.agent_utils.function_parser import parse_function_call
 from cua2_core.services.agent_utils.get_model import get_model
 from cua2_core.services.sandbox_service import SandboxService
 from cua2_core.websocket.websocket_manager import WebSocketException, WebSocketManager
-from e2b_desktop import Sandbox
+from e2b_desktop import Sandbox, TimeoutException
 from fastapi import WebSocket
 from PIL import Image
 from smolagents import ActionStep, AgentImage, AgentMaxStepsError, TaskStep
@@ -37,33 +37,50 @@ class AgentService:
     """Service for handling agent tasks and processing"""
 
     def __init__(
-        self, websocket_manager: WebSocketManager, sandbox_service: SandboxService
+        self,
+        websocket_manager: WebSocketManager,
+        sandbox_service: SandboxService,
+        num_workers: int,
     ):
         self.active_tasks: dict[str, ActiveTask] = {}
         self.websocket_manager: WebSocketManager = websocket_manager
         self.task_websockets: dict[str, WebSocket] = {}
         self.sandbox_service: SandboxService = sandbox_service
-        self.last_screenshot: dict[str, AgentImage] = {}
+        self.last_screenshot: dict[str, AgentImage | None] = {}
+        self._lock = asyncio.Lock()
+        self.max_sandboxes = int(600 / num_workers)
 
-    async def process_user_task(self, trace: AgentTrace, websocket: WebSocket) -> str:
+    async def process_user_task(
+        self, trace: AgentTrace, websocket: WebSocket
+    ) -> str | None:
         """Process a user task and return the trace ID"""
 
         trace_id = trace.id
         trace.steps = []
         trace.traceMetadata = AgentTraceMetadata(traceId=trace_id)
 
-        # Store the task
-        self.active_tasks[trace_id] = ActiveTask(
-            message_id=trace_id,
-            instruction=trace.instruction,
-            model_id=trace.modelId,
-            timestamp=trace.timestamp,
-            steps=trace.steps,
-            traceMetadata=trace.traceMetadata,
-        )
+        async with self._lock:
+            active_task = ActiveTask(
+                message_id=trace_id,
+                instruction=trace.instruction,
+                model_id=trace.modelId,
+                timestamp=trace.timestamp,
+                steps=trace.steps,
+                traceMetadata=trace.traceMetadata,
+            )
 
-        # Store the websocket for this task
-        self.task_websockets[trace_id] = websocket
+            if len(self.active_tasks) >= self.max_sandboxes:
+                await self.websocket_manager.send_agent_start(
+                    active_task=active_task,
+                    status="max_sandboxes_reached",
+                    websocket=websocket,
+                )
+                return trace_id
+
+            # Store the task and websocket for this task
+            self.active_tasks[trace_id] = active_task
+            self.task_websockets[trace_id] = websocket
+            self.last_screenshot[trace_id] = None
 
         asyncio.create_task(self._agent_processing(trace_id))
 
@@ -87,7 +104,9 @@ class AgentService:
             websocket = self.task_websockets.get(message_id)
 
             await self.websocket_manager.send_agent_start(
-                active_task=self.active_tasks[message_id], websocket=websocket
+                active_task=self.active_tasks[message_id],
+                websocket=websocket,
+                status="success",
             )
 
             model = get_model(self.active_tasks[message_id].model_id)
@@ -145,6 +164,9 @@ class AgentService:
         except WebSocketException:
             websocket_exception = True
 
+        except TimeoutException:
+            final_state = "sandbox_timeout"
+
         except (Exception, KeyboardInterrupt):
             import traceback
 
@@ -170,17 +192,23 @@ class AgentService:
 
             novnc_active = False
 
-            # Clean up
+            self.active_tasks[message_id].update_trace_metadata(
+                final_state=final_state,
+            )
+
             if message_id in self.active_tasks:
                 self.active_tasks[message_id].store_model()
-                del self.active_tasks[message_id]
 
-            # Clean up websocket reference
-            if message_id in self.task_websockets:
-                del self.task_websockets[message_id]
+            # Clean up
+            async with self._lock:
+                if message_id in self.active_tasks:
+                    del self.active_tasks[message_id]
+
+                if message_id in self.task_websockets:
+                    del self.task_websockets[message_id]
 
-            if message_id in self.last_screenshot:
-                del self.last_screenshot[message_id]
+                if message_id in self.last_screenshot:
+                    del self.last_screenshot[message_id]
 
             # Release sandbox back to the pool
             if sandbox:
@@ -211,7 +239,7 @@ class AgentService:
             time.sleep(3)
 
             image = self.last_screenshot[message_id]
-            # agent.last_marked_screenshot = AgentImage(screenshot_path)
+            assert image is not None
 
             for previous_memory_step in (
                 agent.memory.steps
@@ -262,8 +290,6 @@ class AgentService:
             else:
                 image_base64 = None
 
-            logger.info(memory_step)
-
             step = AgentStep(
                 traceId=message_id,
                 stepId=str(memory_step.step_number),

cua2-core/src/cua2_core/websocket/websocket_manager.py

@@ -62,7 +62,12 @@ class WebSocketManager:
                 self.disconnect(websocket)
             raise WebSocketException()
 
-    async def send_agent_start(self, active_task: ActiveTask, websocket: WebSocket):
+    async def send_agent_start(
+        self,
+        active_task: ActiveTask,
+        websocket: WebSocket,
+        status: Literal["max_sandboxes_reached", "success"],
+    ):
         """Send agent start event"""
         event = AgentStartEvent(
             agentTrace=AgentTrace(
@@ -74,6 +79,7 @@ class WebSocketManager:
                 traceMetadata=active_task.traceMetadata,
                 isRunning=True,
             ),
+            status=status,
         )
         await self.send_message(event, websocket)
 
@@ -94,7 +100,9 @@ class WebSocketManager:
         self,
         metadata: AgentTraceMetadata,
         websocket: WebSocket,
-        final_state: Literal["success", "stopped", "max_steps_reached", "error"],
+        final_state: Literal[
+            "success", "stopped", "max_steps_reached", "error", "timeout"
+        ],
     ):
         """Send agent complete event"""
         event = AgentCompleteEvent(traceMetadata=metadata, final_state=final_state)

cua2-front/src/types/agent.ts

@@ -35,6 +35,7 @@ export interface AgentTraceMetadata {
   numberOfSteps: number;
   maxSteps: number;
   completed: boolean;
+  final_state: 'success' | 'stopped' | 'max_steps_reached' | 'error' | 'sandbox_timeout' | null;
 }
 
 export interface FinalStep {
@@ -48,6 +49,7 @@ export interface FinalStep {
 interface AgentStartEvent {
   type: 'agent_start';
   agentTrace: AgentTrace;
+  status: 'max_sandboxes_reached' | 'success';
 }
 
 interface AgentProgressEvent {
@@ -59,7 +61,7 @@ interface AgentProgressEvent {
 interface AgentCompleteEvent {
   type: 'agent_complete';
   traceMetadata: AgentTraceMetadata;
-  final_state: 'success' | 'stopped' | 'max_steps_reached' | 'error';
+  final_state: 'success' | 'stopped' | 'max_steps_reached' | 'error' | 'sandbox_timeout';
 }
 
 interface AgentErrorEvent {

entrypoint.sh

@@ -19,9 +19,10 @@ echo "nginx started successfully"
 cd $HOME/app/cua2-core
 
 # Set default number of workers if not specified
-WORKERS=${WORKERS:-1}
+NUM_WORKERS=${NUM_WORKERS:-1}
 
 echo "Starting backend with $WORKERS worker(s)..."
 
 # Use uv to run the application
-exec uv run uvicorn cua2_core.main:app --host 0.0.0.0 --port 8000 --workers $WORKERS --log-level info
+echo "uv run uvicorn cua2_core.main:app --host 0.0.0.0 --port 8000 --workers $NUM_WORKERS --log-level error > /dev/null"
+exec uv run uvicorn cua2_core.main:app --host 0.0.0.0 --port 8000 --workers $NUM_WORKERS --log-level error > /dev/null