ADD agent stop feature (#11)

* NEW agent

* Handle stop agent feature

* CHG agent max step

* FIX model

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

@@ -82,7 +82,7 @@ class AgentAction(FunctionCall):
             return f"Open: {url}"
 
         elif action_type == "launch":
-            url = args.get("url") or args.get("arg_0")
+            url = args.get("app") or args.get("arg_0")
             return f"Open: {url}"
 
         elif action_type == "final_answer":
@@ -172,6 +172,7 @@ class AgentCompleteEvent(BaseModel):
 
     type: Literal["agent_complete"] = "agent_complete"
     traceMetadata: AgentTraceMetadata
+    final_state: Literal["success", "stopped", "max_steps_reached", "error"]
 
 
 class AgentErrorEvent(BaseModel):
@@ -222,6 +223,13 @@ class UserTaskMessage(BaseModel):
     agent_trace: AgentTrace | None = None
 
 
+class StopTask(BaseModel):
+    """Stop task message"""
+
+    event_type: Literal["stop_task"]
+    traceId: str
+
+
 ##################### Agent Service ########################
 
 
@@ -256,6 +264,7 @@ class ActiveTask(BaseModel):
                     f,
                     indent=2,
                 )
+        return self
 
     def update_step(self, step: AgentStep):
         """Update step"""
@@ -276,6 +285,27 @@ class ActiveTask(BaseModel):
                     indent=2,
                 )
 
+    def update_trace_metadata(
+        self,
+        step_input_tokens_used: int | None = None,
+        step_output_tokens_used: int | None = None,
+        step_duration: float | None = None,
+        step_numberOfSteps: int | None = None,
+        completed: bool | None = None,
+    ):
+        """Update trace metadata"""
+        with self._file_lock:
+            if step_input_tokens_used is not None:
+                self.traceMetadata.inputTokensUsed += step_input_tokens_used
+            if step_output_tokens_used is not None:
+                self.traceMetadata.outputTokensUsed += step_output_tokens_used
+            if step_duration is not None:
+                self.traceMetadata.duration += step_duration
+            if step_numberOfSteps is not None:
+                self.traceMetadata.numberOfSteps += step_numberOfSteps
+            if completed is not None:
+                self.traceMetadata.completed = completed
+
 
 #################### API Routes Models ########################
 

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

@@ -59,6 +59,16 @@ async def websocket_endpoint(websocket: WebSocket):
                         else:
                             print("No trace data in message")
 
+                    elif message_data.get("type") == "stop_task":
+                        # Extract and parse the trace
+                        trace_id = message_data.get("trace_id")
+                        if trace_id:
+                            # Stop the task
+                            await agent_service.stop_task(trace_id)
+                            print(f"Stopped task: {trace_id}")
+                        else:
+                            print("No trace ID in message")
+
                 except json.JSONDecodeError as e:
                     print(f"JSON decode error: {e}")
                     from cua2_core.models.models import AgentErrorEvent

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

@@ -27,6 +27,12 @@ from smolagents import ActionStep, AgentImage, AgentMaxStepsError, TaskStep
 logger = logging.getLogger(__name__)
 
 
+class AgentStopException(Exception):
+    """Exception for agent stop"""
+
+    pass
+
+
 class AgentService:
     """Service for handling agent tasks and processing"""
 
@@ -74,6 +80,7 @@ class AgentService:
         agent = None
         novnc_active = False
         websocket_exception = False
+        final_state = "success"
 
         try:
             # Get the websocket for this task
@@ -129,9 +136,14 @@ class AgentService:
 
             self.active_tasks[message_id].traceMetadata.completed = True
 
+        except AgentStopException as e:
+            if str(e) == "Max steps reached":
+                final_state = "max_steps_reached"
+            elif str(e) == "Task not completed":
+                final_state = "stopped"
+
         except WebSocketException:
             websocket_exception = True
-            pass
 
         except (Exception, KeyboardInterrupt):
             import traceback
@@ -139,6 +151,7 @@ class AgentService:
             logger.error(
                 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
             )
@@ -149,6 +162,7 @@ class AgentService:
                 await self.websocket_manager.send_agent_complete(
                     metadata=self.active_tasks[message_id].traceMetadata,
                     websocket=websocket,
+                    final_state=final_state,
                 )
 
                 if novnc_active:
@@ -191,6 +205,9 @@ class AgentService:
         def step_callback(memory_step: ActionStep, agent: E2BVisionAgent):
             assert memory_step.step_number is not None
 
+            if memory_step.step_number > agent.max_steps:
+                raise AgentStopException("Max steps reached")
+
             time.sleep(3)
 
             image = self.last_screenshot[message_id]
@@ -215,9 +232,7 @@ class AgentService:
                 if memory_step.model_output_message
                 else None
             )
-            if model_output is None and isinstance(
-                memory_step.error, AgentMaxStepsError
-            ):
+            if isinstance(memory_step.error, AgentMaxStepsError):
                 model_output = memory_step.action_output
 
             thought = (
@@ -229,11 +244,14 @@ class AgentService:
                 )
                 else None
             )
-            action_sequence = (
-                model_output.split("```")[1]
-                if model_output and memory_step.error is None
-                else None
-            )
+
+            if model_output is not None:
+                action_sequence = model_output.split("```")[1]
+            else:
+                action_sequence = (
+                    """The task failed due to an error"""  # TODO: To Handle in front
+                )
+
             if memory_step.observations_images:
                 image = memory_step.observations_images[0]
                 buffered = BytesIO()
@@ -244,6 +262,8 @@ class AgentService:
             else:
                 image_base64 = None
 
+            logger.info(memory_step)
+
             step = AgentStep(
                 traceId=message_id,
                 stepId=str(memory_step.step_number),
@@ -260,18 +280,14 @@ class AgentService:
                 outputTokensUsed=memory_step.token_usage.output_tokens,
                 step_evaluation="neutral",
             )
-            self.active_tasks[
-                message_id
-            ].traceMetadata.inputTokensUsed += memory_step.token_usage.input_tokens
-            self.active_tasks[
-                message_id
-            ].traceMetadata.outputTokensUsed += memory_step.token_usage.output_tokens
-            self.active_tasks[message_id].traceMetadata.numberOfSteps += 1
-            self.active_tasks[
-                message_id
-            ].traceMetadata.duration += memory_step.timing.duration
-
-            # Add step to active task
+
+            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,
+            )
+
             self.active_tasks[message_id].update_step(step)
 
             websocket = self.task_websockets.get(message_id)
@@ -285,6 +301,9 @@ class AgentService:
             )
             future.result()
 
+            if self.active_tasks[message_id].traceMetadata.completed:
+                raise AgentStopException("Task not completed")
+
             step_filename = f"{message_id}-{memory_step.step_number}"
             screenshot_bytes = agent.desktop.screenshot()
             image = Image.open(BytesIO(screenshot_bytes))
@@ -367,3 +386,10 @@ class AgentService:
                     raise ValueError(f"Step {step_id} not found in trace")
             except (ValueError, KeyError, TypeError) as e:
                 raise ValueError(f"Error processing step update: {e}")
+
+    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(
+                completed=True,
+            )

cua2-core/src/cua2_core/services/agent_utils/desktop_agent.py

@@ -20,7 +20,7 @@ class E2BVisionAgent(CodeAgent):
         model: Model,
         data_dir: str,
         desktop: Sandbox,
-        max_steps: int = 200,
+        max_steps: int = 30,
         verbosity_level: LogLevel = 2,
         planning_interval: int | None = None,
         use_v1_prompt: bool = False,

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

@@ -1,6 +1,6 @@
 import asyncio
 import json
-from typing import Dict, Set
+from typing import Dict, Literal, Set
 
 from cua2_core.models.models import (
     ActiveTask,
@@ -91,10 +91,13 @@ class WebSocketManager:
         await self.send_message(event, websocket)
 
     async def send_agent_complete(
-        self, metadata: AgentTraceMetadata, websocket: WebSocket
+        self,
+        metadata: AgentTraceMetadata,
+        websocket: WebSocket,
+        final_state: Literal["success", "stopped", "max_steps_reached", "error"],
     ):
         """Send agent complete event"""
-        event = AgentCompleteEvent(traceMetadata=metadata)
+        event = AgentCompleteEvent(traceMetadata=metadata, final_state=final_state)
         await self.send_message(event, websocket)
 
     async def send_agent_error(self, error: str, websocket: WebSocket):

cua2-front/src/types/agent.ts

@@ -59,6 +59,7 @@ interface AgentProgressEvent {
 interface AgentCompleteEvent {
   type: 'agent_complete';
   traceMetadata: AgentTraceMetadata;
+  final_state: 'success' | 'stopped' | 'max_steps_reached' | 'error';
 }
 
 interface AgentErrorEvent {
@@ -95,3 +96,9 @@ export interface UserTaskMessage {
   type: 'user_task';
   trace: AgentTrace;
 }
+
+
+export interface StopTaskMessage {
+  type: 'stop_task';
+  traceId: string;
+}