chore: import upstream snapshot with attribution

libs/python/agent/cua_agent/playground/server.py

@@ -0,0 +1,303 @@
+"""Playground server implementation for Cua agents."""
+
+import asyncio
+import logging
+import os
+import platform
+import socket
+import traceback
+import webbrowser
+from typing import Any, Dict, List, Optional, Union
+from urllib.parse import quote
+
+import uvicorn
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
+logger = logging.getLogger(__name__)
+
+
+class PlaygroundServer:
+    """Playground server for running Cua agents via HTTP API."""
+
+    def __init__(self, agent_instance=None):
+        """
+        Initialize the playground server.
+
+        Args:
+            agent_instance: Optional pre-configured agent instance to use
+        """
+        self.agent_instance = agent_instance
+        self.app = FastAPI(
+            title="Cua Playground Server",
+            description="Playground server for Cua agents",
+            version="0.1.0",
+        )
+        self._setup_middleware()
+        self._setup_routes()
+        self.server = None
+        self.port = None
+
+    def _setup_middleware(self):
+        """Setup CORS middleware."""
+        self.app.add_middleware(
+            CORSMiddleware,
+            allow_origins=["*"],
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        )
+
+    def _setup_routes(self):
+        """Setup API routes."""
+
+        @self.app.get("/status")
+        async def status():
+            """Health check endpoint."""
+            sys = platform.system().lower()
+            if "darwin" in sys or sys in ("macos", "mac"):
+                os_type = "macos"
+            elif "windows" in sys:
+                os_type = "windows"
+            else:
+                os_type = "linux"
+
+            return {
+                "status": "ok",
+                "os_type": os_type,
+                "features": ["agent", "playground"],
+            }
+
+        @self.app.post("/responses")
+        async def responses_endpoint(request: Request):
+            """
+            Run ComputerAgent for up to 2 turns.
+
+            Body JSON:
+            {
+              "model": "...",                 # required
+              "input": "... or messages[]",   # required
+              "agent_kwargs": { ... },         # optional, passed directly to ComputerAgent
+              "env": { ... }                   # optional env overrides for agent
+            }
+            """
+            # Import here to avoid circular imports
+            try:
+                from cua_agent import ComputerAgent
+            except ImportError:
+                raise HTTPException(status_code=501, detail="ComputerAgent not available")
+
+            # Parse request body
+            try:
+                body = await request.json()
+            except Exception as e:
+                raise HTTPException(status_code=400, detail=f"Invalid JSON body: {str(e)}")
+
+            model = body.get("model")
+            input_data = body.get("input")
+            if not model or input_data is None:
+                raise HTTPException(status_code=400, detail="'model' and 'input' are required")
+
+            agent_kwargs: Dict[str, Any] = body.get("agent_kwargs") or {}
+            env_overrides: Dict[str, str] = body.get("env") or {}
+
+            # Simple env override context
+            class _EnvOverride:
+                def __init__(self, overrides: Dict[str, str]):
+                    self.overrides = overrides
+                    self._original: Dict[str, Optional[str]] = {}
+
+                def __enter__(self):
+                    for k, v in (self.overrides or {}).items():
+                        self._original[k] = os.environ.get(k)
+                        os.environ[k] = str(v)
+
+                def __exit__(self, exc_type, exc, tb):
+                    for k, old in self._original.items():
+                        if old is None:
+                            os.environ.pop(k, None)
+                        else:
+                            os.environ[k] = old
+
+            # Convert input to messages
+            def _to_messages(data: Union[str, List[Dict[str, Any]]]) -> List[Dict[str, Any]]:
+                if isinstance(data, str):
+                    return [{"role": "user", "content": data}]
+                if isinstance(data, list):
+                    return data
+                return []
+
+            messages = _to_messages(input_data)
+
+            error = None
+
+            with _EnvOverride(env_overrides):
+                # Use pre-configured agent only if model matches, otherwise create new one
+                # This ensures the model picker in the UI is respected
+                if self.agent_instance and self.agent_instance.model == model:
+                    agent = self.agent_instance
+                else:
+                    # Model changed or no pre-configured agent - create new agent with requested model
+                    agent = ComputerAgent(model=model, **agent_kwargs)  # type: ignore[arg-type]
+
+                total_output: List[Any] = []
+                total_usage: Dict[str, Any] = {}
+
+                pending_computer_call_ids = set()
+                try:
+                    async for result in agent.run(messages):
+                        total_output += result["output"]
+                        # Try to collect usage if present
+                        if (
+                            isinstance(result, dict)
+                            and "usage" in result
+                            and isinstance(result["usage"], dict)
+                        ):
+                            # Merge usage counters
+                            for k, v in result["usage"].items():
+                                if isinstance(v, (int, float)):
+                                    total_usage[k] = total_usage.get(k, 0) + v
+                                else:
+                                    total_usage[k] = v
+                        for msg in result.get("output", []):
+                            if msg.get("type") == "computer_call":
+                                pending_computer_call_ids.add(msg["call_id"])
+                            elif msg.get("type") == "computer_call_output":
+                                pending_computer_call_ids.discard(msg["call_id"])
+                            elif msg.get("type") == "function_call":
+                                pending_computer_call_ids.add(msg["call_id"])
+                            elif msg.get("type") == "function_call_output":
+                                pending_computer_call_ids.discard(msg["call_id"])
+                        # exit if no pending computer calls
+                        if not pending_computer_call_ids:
+                            break
+                except Exception as e:
+                    logger.error(f"Error running agent: {str(e)}")
+                    logger.error(traceback.format_exc())
+                    error = str(e)
+
+            # Build response payload
+            payload = {
+                "model": model,
+                "error": error,
+                "output": total_output,
+                "usage": total_usage,
+                "status": "completed" if not error else "failed",
+            }
+
+            # CORS: allow any origin
+            headers = {
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+            }
+
+            return JSONResponse(content=payload, headers=headers)
+
+    def _find_available_port(self, start_port: int = 8000, max_attempts: int = 100) -> int:
+        """Find an available port starting from start_port."""
+        for port in range(start_port, start_port + max_attempts):
+            try:
+                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                    s.bind(("127.0.0.1", port))
+                    return port
+            except OSError:
+                continue
+        raise RuntimeError(
+            f"Could not find an available port in range {start_port}-{start_port + max_attempts}"
+        )
+
+    async def start_async(self, port: Optional[int] = None, open_browser: bool = False):
+        """
+        Start the playground server asynchronously.
+
+        Args:
+            port: Port to run the server on. If None, finds an available port.
+            open_browser: Whether to open the browser automatically.
+        """
+        if port is None:
+            port = self._find_available_port()
+
+        self.port = port
+        host = f"http://localhost:{port}"
+
+        logger.info(f"Starting playground server on {host}")
+
+        if open_browser:
+            # Construct the playground URL
+            encoded_host = quote(host, safe="")
+            encoded_model = quote(self.agent_instance.model, safe="")
+            encoded_vnc_url = quote("http://localhost:8006/?autoconnect=true", safe="")
+
+            # Build URL with custom_model if agent instance is configured
+            playground_url = (
+                # f"http://cua.ai/dashboard/playground"
+                f"http://localhost:3000/dashboard/playground"
+                f"?host={encoded_host}"
+                f"&port={port}"
+                f"&id=localhost"
+                f"&name=localhost"
+                f"&custom_model={encoded_model}"
+                f"&custom_vnc_url={encoded_vnc_url}"
+                f"&vnc_password=null"
+                f"&resize=scale"
+                f"&fullscreen=true"
+            )
+
+            logger.info(f"Opening browser at: {playground_url}")
+            webbrowser.open(playground_url)
+
+        config = uvicorn.Config(
+            self.app,
+            host="0.0.0.0",
+            port=port,
+            log_level="info",
+        )
+        self.server = uvicorn.Server(config)
+        await self.server.serve()
+
+    def start(self, port: Optional[int] = None, open_browser: bool = False):
+        """
+        Start the playground server (blocking).
+
+        Args:
+            port: Port to run the server on. If None, finds an available port.
+            open_browser: Whether to open the browser automatically.
+        """
+        # Check if there's already a running event loop
+        try:
+            loop = asyncio.get_running_loop()
+            # If we're in an async context, schedule as a task
+            import threading
+
+            # Run the server in a separate thread to avoid blocking
+            server_thread = threading.Thread(
+                target=self._run_in_new_loop,
+                args=(port, open_browser),
+                daemon=True,
+            )
+            server_thread.start()
+
+            # Give the server a moment to start and open browser
+            import time
+
+            time.sleep(1)
+
+        except RuntimeError:
+            # No running loop, can use asyncio.run() safely
+            asyncio.run(self.start_async(port=port, open_browser=open_browser))
+
+    def _run_in_new_loop(self, port: Optional[int] = None, open_browser: bool = False):
+        """Helper to run server in a new event loop (for threading)."""
+        new_loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(new_loop)
+        try:
+            new_loop.run_until_complete(self.start_async(port=port, open_browser=open_browser))
+        finally:
+            new_loop.close()
+
+    async def stop(self):
+        """Stop the playground server."""
+        if self.server:
+            logger.info("Stopping playground server")
+            await self.server.shutdown()