import asyncio import base64 import hashlib import inspect import json import logging import os import platform import time import traceback from contextlib import redirect_stderr, redirect_stdout from io import StringIO from typing import Any, Dict, List, Literal, Optional, Union, cast import aiohttp import uvicorn from cua_core.telemetry import record_event from fastapi import ( FastAPI, Header, HTTPException, Request, WebSocket, WebSocketDisconnect, ) from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, StreamingResponse from .browser import get_browser_manager from .handlers.factory import OS_TYPE, HandlerFactory # Try to import MCP server for SSE integration try: from .mcp_server import create_mcp_server HAS_MCP = True except ImportError: HAS_MCP = False # Status code returned when UNAVAILABLE_WITHOUT_CONTAINER_NAME is set and CONTAINER_NAME is missing. DEFAULT_UNAVAILABLE_STATUS_CODE: int = 503 def _parse_bool_env(name: str) -> bool: return os.environ.get(name, "").lower().strip() in ("1", "true", "yes", "y", "on") def _unavailable_status_code() -> Optional[int]: """Return the HTTP status code to use when CONTAINER_NAME is required but unset. When ``UNAVAILABLE_WITHOUT_CONTAINER_NAME`` is truthy and ``CONTAINER_NAME`` is not set, the server should reject requests with the configured status code (``UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE``, default 503) rather than passing through to local dev mode. Returns ``None`` when the server should proceed normally (either because ``CONTAINER_NAME`` is set, or because the unavailable-without-container flag is not enabled). """ if os.environ.get("CONTAINER_NAME"): return None if not _parse_bool_env("UNAVAILABLE_WITHOUT_CONTAINER_NAME"): return None try: return int( os.environ.get( "UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", str(DEFAULT_UNAVAILABLE_STATUS_CODE), ) ) except ValueError: return DEFAULT_UNAVAILABLE_STATUS_CODE try: from cua_agent import ComputerAgent HAS_AGENT = True except ImportError: HAS_AGENT = False # Set up logging with more detail logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) # Configure WebSocket with larger message size WEBSOCKET_MAX_SIZE = 1024 * 1024 * 10 # 10MB limit # Create MCP app first if available (needed for lifespan) _mcp_http_app = None if HAS_MCP: try: mcp_server = create_mcp_server() # Configure FastMCP to use "/" as its internal path, then mount at /mcp # This results in endpoint at /mcp (mount) + / (internal) = /mcp _mcp_http_app = mcp_server.http_app(path="/") logger.info("MCP server created for /mcp endpoint (streamable HTTP transport)") except Exception as e: logger.warning(f"Failed to create MCP server: {e}") # Configure application with WebSocket settings and MCP lifespan app = FastAPI( title="Computer API", description="API for the Computer project", version="0.1.0", websocket_max_size=WEBSOCKET_MAX_SIZE, lifespan=_mcp_http_app.lifespan if _mcp_http_app else None, redirect_slashes=False, ) class UnavailableWithoutContainerMiddleware: """ASGI middleware that rejects all requests when CONTAINER_NAME is required but unset. Controlled by env vars (read per-request so tests and dynamic config work): - ``UNAVAILABLE_WITHOUT_CONTAINER_NAME``: if truthy and ``CONTAINER_NAME`` is unset, every HTTP and WebSocket request is rejected. - ``UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE``: HTTP status code for rejections (default 503). When disabled (either env var not set), requests pass through unchanged, preserving the original "local development mode" behavior for backwards compatibility. """ _DETAIL = "Service unavailable: CONTAINER_NAME is required but not configured" def __init__(self, app): self.app = app async def __call__(self, scope, receive, send): scope_type = scope.get("type") if scope_type in ("http", "websocket"): status_code = _unavailable_status_code() if status_code is not None: if scope_type == "http": await self._reject_http(send, status_code) else: await self._reject_websocket(receive, send, status_code) return await self.app(scope, receive, send) @classmethod async def _reject_http(cls, send, status_code): body = json.dumps({"detail": cls._DETAIL}).encode() await send( { "type": "http.response.start", "status": status_code, "headers": [ (b"content-type", b"application/json"), (b"content-length", str(len(body)).encode()), ], } ) await send({"type": "http.response.body", "body": body}) @classmethod async def _reject_websocket(cls, receive, send, status_code): # Accept first so we can send a structured JSON error before closing — this # preserves the existing error shape that clients already handle. event = await receive() if event.get("type") != "websocket.connect": return await send({"type": "websocket.accept"}) await send( { "type": "websocket.send", "text": json.dumps( { "success": False, "error": cls._DETAIL, "status_code": status_code, } ), } ) # 1008 = Policy Violation await send({"type": "websocket.close", "code": 1008}) app.add_middleware(UnavailableWithoutContainerMiddleware) # CORS configuration origins = ["*"] app.add_middleware( CORSMiddleware, allow_origins=origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) class McpBarePathRewrite: """Serve the MCP app at both /mcp and /mcp/. Starlette's Mount only matches "/mcp/..." — a bare "/mcp" resolves to inner path "" and 404s, and redirect_slashes=False on the outer app disables the 307 rescue. Most MCP clients (claude.ai included) POST to the bare path, so rewrite it before routing. A rewrite (not a redirect) keeps POST bodies and SSE streaming intact; pure ASGI (not BaseHTTPMiddleware) so responses aren't buffered. """ def __init__(self, app): self.app = app async def __call__(self, scope, receive, send): if scope["type"] == "http" and scope["path"] == "/mcp": scope = dict(scope, path="/mcp/") await self.app(scope, receive, send) class McpAcceptHeaderShim: """Make the MCP streamable-HTTP endpoint tolerant of the Accept header. The MCP Python SDK's StreamableHTTP transport rejects requests with JSON-RPC -32600 ("Not Acceptable: Client must accept ...") unless the request's Accept header advertises BOTH ``application/json`` and ``text/event-stream`` (POST path), or ``text/event-stream`` (GET SSE path). claude.ai's MCP connector does not always send both, so its handshake POST to /mcp is rejected before it ever reaches a tool. This pure-ASGI shim normalizes the inbound Accept header for /mcp requests so it always contains both media types, which lets the SDK's Accept check pass for every client regardless of what it sent. It only rewrites the REQUEST header — it does not touch ``json_response`` or any response semantics, so the server still negotiates JSON vs. SSE exactly as it would otherwise. This keeps existing MCP clients (Claude Code, SDK) unaffected while admitting stricter/looser connectors. Scoped to the /mcp prefix only; all other routes are passed through untouched. Pure ASGI (not BaseHTTPMiddleware) so streaming responses are not buffered. """ _COMBINED = b"application/json, text/event-stream" def __init__(self, app): self.app = app async def __call__(self, scope, receive, send): path = scope.get("path", "") # Match the bare /mcp endpoint and anything under /mcp/, but NOT # unrelated routes that merely share the prefix (e.g. /mcpx). if scope["type"] == "http" and (path == "/mcp" or path.startswith("/mcp/")): headers = [(k, v) for (k, v) in scope.get("headers", []) if k != b"accept"] headers.append((b"accept", self._COMBINED)) scope = dict(scope, headers=headers) await self.app(scope, receive, send) # Mount MCP server at /mcp - FastMCP's internal path is "/" so endpoint is /mcp. # # Middleware ordering note: Starlette's app.add_middleware() prepends, so the # LAST one added runs FIRST (outermost). We want, outermost -> innermost: # McpBarePathRewrite (fix the path) -> McpAcceptHeaderShim (fix Accept) -> app # so add McpAcceptHeaderShim first, then McpBarePathRewrite. if _mcp_http_app: app.mount("/mcp", _mcp_http_app) app.add_middleware(McpAcceptHeaderShim) app.add_middleware(McpBarePathRewrite) protocol_version = 1 try: from importlib.metadata import version package_version = version("cua-computer-server") except Exception: # Fallback for cases where package is not installed or importlib.metadata is not available try: import pkg_resources package_version = pkg_resources.get_distribution("cua-computer-server").version except Exception: package_version = "unknown" ( accessibility_handler, automation_handler, diorama_handler, file_handler, desktop_handler, window_handler, ) = HandlerFactory.create_handlers() # Helper function for direction-based scrolling async def _scroll_direction_handler(direction: str, clicks: int = 1) -> Dict[str, Any]: """ Scroll in a specified direction. Args: direction: One of 'up', 'down', 'left', 'right' clicks: Number of scroll clicks (default: 1) """ direction = direction.lower() if direction == "down": return await automation_handler.scroll_down(clicks) elif direction == "up": return await automation_handler.scroll_up(clicks) elif direction in ("left", "right"): # For horizontal scrolling, use scroll with x amount x_amount = -300 if direction == "left" else 300 return await automation_handler.scroll(x_amount * clicks, 0) else: raise ValueError(f"Invalid direction: {direction}. Use 'up', 'down', 'left', or 'right'.") # Command aliases for common alternative names COMMAND_ALIASES = { "type": "type_text", "key": "press_key", "tap": "left_click", "click": "left_click", "shell": "run_command", "exec": "run_command", "read_file": "read_text", "write_file": "write_text", "ls": "list_dir", "mkdir": "create_dir", "rm": "delete_file", "rmdir": "delete_dir", } handlers = { "version": lambda: {"protocol": protocol_version, "package": package_version}, # App-Use commands "diorama_cmd": diorama_handler.diorama_cmd, # Accessibility commands "get_accessibility_tree": accessibility_handler.get_accessibility_tree, "find_element": accessibility_handler.find_element, # Shell commands "run_command": automation_handler.run_command, # File system commands "file_exists": file_handler.file_exists, "directory_exists": file_handler.directory_exists, "list_dir": file_handler.list_dir, "read_text": file_handler.read_text, "write_text": file_handler.write_text, "read_bytes": file_handler.read_bytes, "write_bytes": file_handler.write_bytes, "get_file_size": file_handler.get_file_size, "delete_file": file_handler.delete_file, "create_dir": file_handler.create_dir, "delete_dir": file_handler.delete_dir, # Desktop commands "get_desktop_environment": desktop_handler.get_desktop_environment, "set_wallpaper": desktop_handler.set_wallpaper, # Window management "open": window_handler.open, "launch": window_handler.launch, "get_current_window_id": window_handler.get_current_window_id, "get_application_windows": window_handler.get_application_windows, "get_window_name": window_handler.get_window_name, "get_window_size": window_handler.get_window_size, "get_window_position": window_handler.get_window_position, "set_window_size": window_handler.set_window_size, "set_window_position": window_handler.set_window_position, "maximize_window": window_handler.maximize_window, "minimize_window": window_handler.minimize_window, "activate_window": window_handler.activate_window, "close_window": window_handler.close_window, # Mouse commands "mouse_down": automation_handler.mouse_down, "mouse_up": automation_handler.mouse_up, "left_click": automation_handler.left_click, "right_click": automation_handler.right_click, "double_click": automation_handler.double_click, "move_cursor": automation_handler.move_cursor, "drag_to": automation_handler.drag_to, "drag": automation_handler.drag, # Keyboard commands "key_down": automation_handler.key_down, "key_up": automation_handler.key_up, "type_text": automation_handler.type_text, "press_key": automation_handler.press_key, "hotkey": automation_handler.hotkey, # Scrolling actions "scroll": automation_handler.scroll, "scroll_down": automation_handler.scroll_down, "scroll_up": automation_handler.scroll_up, "scroll_direction": _scroll_direction_handler, # Screen actions "screenshot": automation_handler.screenshot, "get_cursor_position": automation_handler.get_cursor_position, "get_screen_size": automation_handler.get_screen_size, # Clipboard actions "copy_to_clipboard": automation_handler.copy_to_clipboard, "set_clipboard": automation_handler.set_clipboard, } # Android-only commands — registered only when the Android handler is active # so non-Android server instances don't fail at startup with AttributeError. if hasattr(automation_handler, "multitouch_gesture"): handlers["multitouch_gesture"] = automation_handler.multitouch_gesture class AuthenticationManager: def __init__(self): self.container_name = os.environ.get("CONTAINER_NAME") self.api_base_url = os.environ.get("CUA_BASE_URL_AUTH", "https://www.cua.ai").rstrip("/") async def auth(self, container_name: str, api_key: str) -> bool: """Authenticate container name and API key against the TryCUA API. Every call hits the API directly — results are not cached. The previous implementation cached both successes and failures with a 60 s TTL (``CUA_AUTH_TTL_SECONDS``), but that caused transient network/DNS errors immediately after VM boot to lock out legitimate clients for the full TTL window. Since the gate sits in front of every SDK request on a freshly provisioned VM, a single cold-start DNS hiccup would flip the VM into a perpetually-401 state until the TTL expired. """ # If no CONTAINER_NAME is set, always allow access (local development) if not self.container_name: logger.info( "No CONTAINER_NAME set in environment. Allowing access (local development mode)" ) return True # Layer 1: VM Identity Verification if container_name != self.container_name: logger.warning( f"VM name mismatch. Expected: {self.container_name}, Got: {container_name}" ) return False logger.info(f"Authenticating with TryCUA API for container: {container_name}") try: from cua_core.http import cua_version_headers async with aiohttp.ClientSession() as session: headers = {"Authorization": f"Bearer {api_key}", **cua_version_headers()} async with session.get( f"{self.api_base_url}/api/vm/auth?container_name={container_name}", headers=headers, ) as resp: is_valid = resp.status == 200 and bool((await resp.text()).strip()) if is_valid: logger.info(f"Authentication successful for container: {container_name}") else: logger.warning( f"Authentication failed for container: {container_name}. Status: {resp.status}" ) return is_valid except aiohttp.ClientError as e: logger.error(f"Failed to validate API key with TryCUA API: {str(e)}") return False except Exception as e: logger.error(f"Unexpected error during authentication: {str(e)}") return False class ConnectionManager: def __init__(self): self.active_connections: List[WebSocket] = [] async def connect(self, websocket: WebSocket): await websocket.accept() self.active_connections.append(websocket) def disconnect(self, websocket: WebSocket): self.active_connections.remove(websocket) manager = ConnectionManager() auth_manager = AuthenticationManager() # PTY session manager (lazy-initialised) from .pty_manager import PtyManager pty_manager = PtyManager() def _resolve_command(command: str) -> str: """Resolve command aliases to their canonical names.""" return COMMAND_ALIASES.get(command, command) def _suggest_similar_commands(command: str) -> List[str]: """Suggest similar commands for typos.""" all_commands = list(handlers.keys()) + list(COMMAND_ALIASES.keys()) suggestions = [] command_lower = command.lower() for cmd in all_commands: # Simple similarity: starts with same letter or contains the input if cmd.lower().startswith(command_lower[:2]) or command_lower in cmd.lower(): suggestions.append(cmd) return suggestions[:5] @app.get("/status") async def status(): # get computer-server features features = [] if HAS_AGENT: features.append("agent") if HAS_MCP: features.append("mcp") return {"status": "ok", "os_type": OS_TYPE, "features": features} @app.get("/commands") async def list_commands(): """List all available commands and their aliases.""" commands_info = {} for cmd_name, handler_func in handlers.items(): try: sig = inspect.signature(handler_func) params = [ { "name": p.name, "required": p.default == inspect.Parameter.empty, "default": None if p.default == inspect.Parameter.empty else p.default, } for p in sig.parameters.values() ] except (ValueError, TypeError): params = [] commands_info[cmd_name] = {"params": params} # Add alias information aliases_by_command: Dict[str, List[str]] = {} for alias, canonical in COMMAND_ALIASES.items(): if canonical not in aliases_by_command: aliases_by_command[canonical] = [] aliases_by_command[canonical].append(alias) for cmd_name in commands_info: if cmd_name in aliases_by_command: commands_info[cmd_name]["aliases"] = aliases_by_command[cmd_name] return { "commands": commands_info, "aliases": COMMAND_ALIASES, } @app.websocket("/ws", name="websocket_endpoint") async def websocket_endpoint(websocket: WebSocket): global handlers # WebSocket message size is configured at the app or endpoint level, not on the instance await manager.connect(websocket) # Check if CONTAINER_NAME is set (indicating cloud provider) server_container_name = os.environ.get("CONTAINER_NAME") # If cloud provider, perform authentication handshake if server_container_name: try: logger.info( f"Cloud provider detected. CONTAINER_NAME: {server_container_name}. Waiting for authentication..." ) # Wait for authentication message auth_data = await websocket.receive_json() # Validate auth message format if auth_data.get("command") != "authenticate": await websocket.send_json( {"success": False, "error": "First message must be authentication"} ) await websocket.close() manager.disconnect(websocket) return # Extract credentials client_api_key = auth_data.get("params", {}).get("api_key") client_container_name = auth_data.get("params", {}).get("container_name") # Validate credentials using AuthenticationManager if not client_api_key: await websocket.send_json({"success": False, "error": "API key required"}) await websocket.close() manager.disconnect(websocket) return if not client_container_name: await websocket.send_json({"success": False, "error": "Container name required"}) await websocket.close() manager.disconnect(websocket) return # Use AuthenticationManager for validation is_authenticated = await auth_manager.auth(client_container_name, client_api_key) if not is_authenticated: await websocket.send_json({"success": False, "error": "Authentication failed"}) await websocket.close() manager.disconnect(websocket) return logger.info(f"Authentication successful for VM: {client_container_name}") await websocket.send_json({"success": True, "message": "Authentication successful"}) # Emit vm_session_started event for funnel tracking api_key_hash = hashlib.sha256(client_api_key.encode()).hexdigest() record_event( "vm_session_started", { "api_key_hash": api_key_hash, "vm_id": client_container_name, "connection_type": "websocket", }, ) except Exception as e: logger.error(f"Error during authentication handshake: {str(e)}") await websocket.send_json({"success": False, "error": "Authentication failed"}) await websocket.close() manager.disconnect(websocket) return try: while True: try: data = await websocket.receive_json() command = data.get("command") params = data.get("params", {}) # Resolve command aliases command = _resolve_command(command) if command not in handlers: suggestions = _suggest_similar_commands(command) error_msg = f"Unknown command: {command}" if suggestions: error_msg += f". Did you mean: {', '.join(suggestions)}?" await websocket.send_json({"success": False, "error": error_msg}) continue try: # Filter params to only include those accepted by the handler function handler_func = handlers[command] sig = inspect.signature(handler_func) filtered_params = {k: v for k, v in params.items() if k in sig.parameters} # Handle both sync and async functions if asyncio.iscoroutinefunction(handler_func): result = await handler_func(**filtered_params) else: # Run sync functions in thread pool to avoid blocking event loop result = await asyncio.to_thread(handler_func, **filtered_params) await websocket.send_json({"success": True, **result}) except Exception as cmd_error: logger.error(f"Error executing command {command}: {str(cmd_error)}") logger.error(traceback.format_exc()) await websocket.send_json({"success": False, "error": str(cmd_error)}) except WebSocketDisconnect: raise except json.JSONDecodeError as json_err: logger.error(f"JSON decode error: {str(json_err)}") await websocket.send_json( {"success": False, "error": f"Invalid JSON: {str(json_err)}"} ) except Exception as loop_error: logger.error(f"Error in message loop: {str(loop_error)}") logger.error(traceback.format_exc()) await websocket.send_json({"success": False, "error": str(loop_error)}) except WebSocketDisconnect: logger.info("Client disconnected") manager.disconnect(websocket) except Exception as e: logger.error(f"Fatal error in websocket connection: {str(e)}") logger.error(traceback.format_exc()) try: await websocket.close() except: pass manager.disconnect(websocket) @app.post("/cmd") async def cmd_endpoint( request: Request, container_name: Optional[str] = Header(None, alias="X-Container-Name"), api_key: Optional[str] = Header(None, alias="X-API-Key"), ): """ Backup endpoint for when WebSocket connections fail. Accepts commands via HTTP POST with streaming response. Headers: - X-Container-Name: Container name for cloud authentication - X-API-Key: API key for cloud authentication Body: { "command": "command_name", "params": {...} } """ global handlers # Parse request body try: body = await request.json() command = body.get("command") params = body.get("params", {}) except Exception as e: raise HTTPException(status_code=400, detail=f"Invalid JSON body: {str(e)}") if not command: raise HTTPException(status_code=400, detail="Command is required") # Resolve command aliases command = _resolve_command(command) # Check if CONTAINER_NAME is set (indicating cloud provider) server_container_name = os.environ.get("CONTAINER_NAME") # If cloud provider, perform authentication if server_container_name: logger.info( f"Cloud provider detected. CONTAINER_NAME: {server_container_name}. Performing authentication..." ) # Validate required headers if not container_name: raise HTTPException(status_code=401, detail="Container name required") if not api_key: raise HTTPException(status_code=401, detail="API key required") # Validate with AuthenticationManager is_authenticated = await auth_manager.auth(container_name, api_key) if not is_authenticated: raise HTTPException(status_code=401, detail="Authentication failed") # Emit vm_session_started event for funnel tracking api_key_hash = hashlib.sha256(api_key.encode()).hexdigest() record_event( "vm_session_started", { "api_key_hash": api_key_hash, "vm_id": container_name, "connection_type": "http", }, ) if command not in handlers: suggestions = _suggest_similar_commands(command) error_msg = f"Unknown command: {command}" if suggestions: error_msg += f". Did you mean: {', '.join(suggestions)}?" raise HTTPException(status_code=400, detail=error_msg) async def generate_response(): """Generate streaming response for the command execution""" try: # Filter params to only include those accepted by the handler function handler_func = handlers[command] sig = inspect.signature(handler_func) filtered_params = {k: v for k, v in params.items() if k in sig.parameters} # Handle both sync and async functions if asyncio.iscoroutinefunction(handler_func): result = await handler_func(**filtered_params) else: # Run sync functions in thread pool to avoid blocking event loop result = await asyncio.to_thread(handler_func, **filtered_params) # Stream the successful result response_data = {"success": True, **result} yield f"data: {json.dumps(response_data)}\n\n" except Exception as cmd_error: logger.error(f"Error executing command {command}: {str(cmd_error)}") logger.error(traceback.format_exc()) # Stream the error result error_data = {"success": False, "error": str(cmd_error)} yield f"data: {json.dumps(error_data)}\n\n" return StreamingResponse( generate_response(), media_type="text/plain", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", }, ) async def _require_auth( container_name: Optional[str], api_key: Optional[str], ) -> None: """Raise HTTPException(401) when cloud auth is configured and credentials are invalid.""" server_container_name = os.environ.get("CONTAINER_NAME") if not server_container_name: return # local development — no auth required if not container_name: raise HTTPException(status_code=401, detail="Container name required") if not api_key: raise HTTPException(status_code=401, detail="API key required") if not await auth_manager.auth(container_name, api_key): raise HTTPException(status_code=401, detail="Authentication failed") # --------------------------------------------------------------------------- # PTY endpoints # --------------------------------------------------------------------------- @app.post("/pty") async def pty_create( request: Request, container_name: Optional[str] = Header(None, alias="X-Container-Name"), api_key: Optional[str] = Header(None, alias="X-API-Key"), ): """Spawn a new PTY session. Body (JSON, all fields optional): ``{"command": str, "cols": int, "rows": int, "cwd": str, "envs": dict}`` Returns: ``{"pid": int, "cols": int, "rows": int}`` """ await _require_auth(container_name, api_key) try: body = await request.json() except Exception: body = {} # On Android, always use adb shell regardless of any caller-supplied command, # so that PTY sessions run inside the emulator rather than on the host. if os.environ.get("IS_CUA_ANDROID") == "true": command = "adb shell" else: command = body.get("command") info = await pty_manager.create( command=command, cols=int(body.get("cols", 80)), rows=int(body.get("rows", 24)), cwd=body.get("cwd"), envs=body.get("envs"), ) return info @app.get("/pty/{pid}") async def pty_info( pid: int, container_name: Optional[str] = Header(None, alias="X-Container-Name"), api_key: Optional[str] = Header(None, alias="X-API-Key"), ): """Return metadata for PTY session *pid*. Returns: ``{"pid": int, "cols": int, "rows": int}`` """ await _require_auth(container_name, api_key) info = pty_manager.get_info(pid) if info is None: from fastapi import HTTPException raise HTTPException(status_code=404, detail=f"PTY session {pid} not found") return info @app.delete("/pty/{pid}") async def pty_kill( pid: int, container_name: Optional[str] = Header(None, alias="X-Container-Name"), api_key: Optional[str] = Header(None, alias="X-API-Key"), ): """Kill PTY session *pid*. Returns: ``{"killed": bool}`` """ await _require_auth(container_name, api_key) killed = await pty_manager.kill(pid) return {"killed": killed} @app.post("/pty/{pid}/stdin") async def pty_stdin( pid: int, request: Request, container_name: Optional[str] = Header(None, alias="X-Container-Name"), api_key: Optional[str] = Header(None, alias="X-API-Key"), ): """Write data to stdin of PTY session *pid*. Body: ``{"data": ""}`` """ await _require_auth(container_name, api_key) body = await request.json() raw = base64.b64decode(body.get("data", "")) await pty_manager.send_stdin(pid, raw) return {"ok": True} @app.post("/pty/{pid}/resize") async def pty_resize( pid: int, request: Request, container_name: Optional[str] = Header(None, alias="X-Container-Name"), api_key: Optional[str] = Header(None, alias="X-API-Key"), ): """Resize the terminal for PTY session *pid*. Body: ``{"cols": int, "rows": int}`` """ await _require_auth(container_name, api_key) body = await request.json() await pty_manager.resize(pid, int(body.get("cols", 80)), int(body.get("rows", 24))) return {"ok": True} @app.get("/pty/{pid}/stream") async def pty_stream( pid: int, container_name: Optional[str] = Header(None, alias="X-Container-Name"), api_key: Optional[str] = Header(None, alias="X-API-Key"), ): """SSE stream for PTY session *pid*. Events: - ``data: {"type": "output", "data": ""}`` - ``data: {"type": "exit", "code": }`` """ await _require_auth(container_name, api_key) q = pty_manager.subscribe(pid) async def _generate(): try: while True: msg = await q.get() if msg["type"] == "output": payload = {"type": "output", "data": base64.b64encode(msg["data"]).decode()} else: payload = {"type": "exit", "code": msg.get("code", 0)} yield f"data: {json.dumps(payload)}\n\n" if msg["type"] == "exit": break finally: pty_manager.unsubscribe(pid, q) return StreamingResponse( _generate(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) @app.websocket("/pty/{pid}/ws") async def pty_ws( pid: int, websocket: WebSocket, ): """WebSocket endpoint for interactive PTY session *pid*. Auth (when CONTAINER_NAME is set): pass ``api_key`` and ``container_name`` as query parameters, e.g. ``/pty/123/ws?api_key=…&container_name=…``. Client → Server messages (JSON): - ``{"type": "stdin", "data": ""}`` - ``{"type": "resize", "cols": N, "rows": N}`` - ``{"type": "disconnect"}`` Server → Client messages (JSON): - ``{"type": "output", "data": ""}`` - ``{"type": "exit", "code": N}`` """ container_name = websocket.query_params.get("container_name") api_key = websocket.query_params.get("api_key") try: await _require_auth(container_name, api_key) except HTTPException: await websocket.close(code=1008) # 1008 = Policy Violation return await websocket.accept() q = pty_manager.subscribe(pid) async def _send_output(): """Forward PTY output to the WebSocket client.""" try: while True: msg = await q.get() if msg["type"] == "output": payload = {"type": "output", "data": base64.b64encode(msg["data"]).decode()} else: payload = {"type": "exit", "code": msg.get("code", 0)} await websocket.send_text(json.dumps(payload)) if msg["type"] == "exit": break except Exception: pass finally: pty_manager.unsubscribe(pid, q) output_task = asyncio.create_task(_send_output()) try: while True: try: raw = await websocket.receive_text() except WebSocketDisconnect: break except Exception: break try: msg = json.loads(raw) except json.JSONDecodeError: continue msg_type = msg.get("type") if msg_type == "stdin": data = base64.b64decode(msg.get("data", "")) await pty_manager.send_stdin(pid, data) elif msg_type == "resize": await pty_manager.resize(pid, int(msg.get("cols", 80)), int(msg.get("rows", 24))) elif msg_type == "disconnect": break finally: output_task.cancel() pty_manager.unsubscribe(pid, q) try: await websocket.close() except Exception: pass @app.post("/responses") async def agent_response_endpoint( request: Request, api_key: Optional[str] = Header(None, alias="X-API-Key"), ): """ Minimal proxy to run ComputerAgent for up to 2 turns. Security: - If CONTAINER_NAME is set on the server, require X-API-Key and validate using AuthenticationManager unless CUA_ENABLE_PUBLIC_PROXY is true. Body JSON: { "model": "...", # required "input": "... or messages[]", # required "agent_kwargs": { ... }, # optional, passed directly to ComputerAgent "env": { ... } # optional env overrides for agent } """ if not HAS_AGENT: raise HTTPException(status_code=501, detail="ComputerAgent not available") # Authenticate via AuthenticationManager if running in cloud (CONTAINER_NAME set) container_name = os.environ.get("CONTAINER_NAME") if container_name: is_public = os.environ.get("CUA_ENABLE_PUBLIC_PROXY", "").lower().strip() in [ "1", "true", "yes", "y", "on", ] if not is_public: if not api_key: raise HTTPException(status_code=401, detail="Missing AGENT PROXY auth headers") ok = await auth_manager.auth(container_name, api_key) if not ok: raise HTTPException(status_code=401, detail="Unauthorized") # 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 messages = _to_messages(input_data) # Define a direct computer tool that implements the AsyncComputerHandler protocol # and delegates to our existing automation/file/accessibility handlers. from cua_agent.computers import AsyncComputerHandler # runtime-checkable Protocol class DirectComputerInterface: """Interface wrapper providing BrowserTool compatibility. Matches the same interface shape as Computer.interface so BrowserTool works identically with both Computer (cloud) and DirectComputer (local). """ def __init__(self, automation_handler, browser_manager): self._auto = automation_handler self._browser = browser_manager @property def interface(self): """Return automation handler for hotkey, move_cursor, etc.""" return self._auto async def playwright_exec(self, command: str, params: dict) -> dict: """Execute browser command via browser_manager.""" return await self._browser.execute_command(command, params) class DirectComputer(AsyncComputerHandler): def __init__(self): # use module-scope handler singletons created by HandlerFactory self._auto = automation_handler self._file = file_handler self._access = accessibility_handler # Create interface for BrowserTool compatibility self._interface = DirectComputerInterface(automation_handler, get_browser_manager()) @property def interface(self): """Return interface compatible with BrowserTool. This matches Computer.interface shape so BrowserTool works with either: - computer.interface.interface.hotkey() -> automation - computer.interface.playwright_exec() -> browser commands - direct_computer.interface.interface.hotkey() -> automation - direct_computer.interface.playwright_exec() -> browser commands """ return self._interface async def get_environment(self) -> Literal["windows", "mac", "linux", "browser"]: sys = platform.system().lower() if "darwin" in sys or sys in ("macos", "mac"): return "mac" if "windows" in sys: return "windows" return "linux" async def get_dimensions(self) -> tuple[int, int]: size = await self._auto.get_screen_size() return size["width"], size["height"] async def screenshot(self) -> str: img_b64 = await self._auto.screenshot() return img_b64["image_data"] async def click(self, x: int, y: int, button: str = "left") -> None: if button == "left": await self._auto.left_click(x, y) elif button == "right": await self._auto.right_click(x, y) else: await self._auto.left_click(x, y) async def double_click(self, x: int, y: int) -> None: await self._auto.double_click(x, y) async def scroll(self, x: int, y: int, scroll_x: int, scroll_y: int) -> None: await self._auto.move_cursor(x, y) await self._auto.scroll(scroll_x, scroll_y) async def type(self, text: str) -> None: await self._auto.type_text(text) async def wait(self, ms: int = 1000) -> None: await asyncio.sleep(ms / 1000.0) async def move(self, x: int, y: int) -> None: await self._auto.move_cursor(x, y) async def keypress(self, keys: Union[List[str], str]) -> None: from computer.interface.models import Key def normalize_key(key: str) -> str: """Normalize key name using SDK's Key.from_string().""" result = Key.from_string(key) if isinstance(result, Key): # SDK mapped it to a Key enum, use its value return result.value else: # SDK didn't map it - lowercase for pynput compatibility # (pynput uses lowercase: Key.tab, Key.space, etc.) return key.lower() if len(key) > 1 else key if isinstance(keys, str): parts = keys.replace("-", "+").split("+") if len(keys) > 1 else [keys] else: parts = keys # Normalize all key parts using SDK's Key mapping parts = [normalize_key(p) for p in parts] if len(parts) == 1: key = parts[0] # Route single printable characters through type_text so the # insertion is layout-independent. press_key uses a physical-key # path (pynput keyboard.press/release) that follows the active # input source — under a Russian or CJK layout it inserts the # layout-mapped character instead of the intended ASCII one. # type_text uses pynput keyboard.type() which bypasses the # layout and inserts the literal Unicode codepoint. # # A key is "printable" when it is exactly one character long # and unicodedata.category is not a control category (Cc/Cs). # Special keys (return, tab, escape, arrows, f1-f12, …) are # multi-character strings or map to a Key enum — those still # go through press_key unchanged. # # See: https://github.com/trycua/cua/issues/1605 import unicodedata if len(key) == 1 and unicodedata.category(key) not in ("Cc", "Cs", "Cn"): await self._auto.type_text(key) else: await self._auto.press_key(key) else: await self._auto.hotkey(parts) async def drag(self, path: List[Dict[str, int]]) -> None: if not path: return start = path[0] await self._auto.mouse_down(start["x"], start["y"]) for pt in path[1:]: await self._auto.move_cursor(pt["x"], pt["y"]) end = path[-1] await self._auto.mouse_up(end["x"], end["y"]) async def get_current_url(self) -> str: # Not available in this server context return "" async def left_mouse_down(self, x: Optional[int] = None, y: Optional[int] = None) -> None: await self._auto.mouse_down(x, y, button="left") async def left_mouse_up(self, x: Optional[int] = None, y: Optional[int] = None) -> None: await self._auto.mouse_up(x, y, button="left") # # Inline image URLs to base64 # import base64, mimetypes, requests # # Use a browser-like User-Agent to avoid 403s from some CDNs (e.g., Wikimedia) # HEADERS = { # "User-Agent": ( # "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " # "AppleWebKit/537.36 (KHTML, like Gecko) " # "Chrome/124.0.0.0 Safari/537.36" # ) # } # def _to_data_url(content_bytes: bytes, url: str, resp: requests.Response) -> str: # ctype = resp.headers.get("Content-Type") or mimetypes.guess_type(url)[0] or "application/octet-stream" # b64 = base64.b64encode(content_bytes).decode("utf-8") # return f"data:{ctype};base64,{b64}" # def inline_image_urls(messages): # # messages: List[{"role": "...","content":[...]}] # out = [] # for m in messages: # if not isinstance(m.get("content"), list): # out.append(m) # continue # new_content = [] # for part in (m.get("content") or []): # if part.get("type") == "input_image" and (url := part.get("image_url")): # resp = requests.get(url, headers=HEADERS, timeout=30) # resp.raise_for_status() # new_content.append({ # "type": "input_image", # "image_url": _to_data_url(resp.content, url, resp) # }) # else: # new_content.append(part) # out.append({**m, "content": new_content}) # return out # messages = inline_image_urls(messages) error = None with _EnvOverride(env_overrides): # Prepare tools: if caller did not pass tools, inject our DirectComputer tools = agent_kwargs.get("tools") if not tools: tools = [DirectComputer()] agent_kwargs = {**agent_kwargs, "tools": tools} # Instantiate agent with our tools 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"]) # 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) @app.post("/playwright_exec") async def playwright_exec_endpoint( request: Request, container_name: Optional[str] = Header(None, alias="X-Container-Name"), api_key: Optional[str] = Header(None, alias="X-API-Key"), ): """ Execute Playwright browser commands. Headers: - X-Container-Name: Container name for cloud authentication - X-API-Key: API key for cloud authentication Body: { "command": "visit_url|click|type|scroll|web_search", "params": {...} } """ # Parse request body try: body = await request.json() command = body.get("command") params = body.get("params", {}) except Exception as e: raise HTTPException(status_code=400, detail=f"Invalid JSON body: {str(e)}") if not command: raise HTTPException(status_code=400, detail="Command is required") # Check if CONTAINER_NAME is set (indicating cloud provider) server_container_name = os.environ.get("CONTAINER_NAME") # If cloud provider, perform authentication if server_container_name: logger.info( f"Cloud provider detected. CONTAINER_NAME: {server_container_name}. Performing authentication..." ) # Validate required headers if not container_name: raise HTTPException(status_code=401, detail="Container name required") if not api_key: raise HTTPException(status_code=401, detail="API key required") # Validate with AuthenticationManager is_authenticated = await auth_manager.auth(container_name, api_key) if not is_authenticated: raise HTTPException(status_code=401, detail="Authentication failed") # Get browser manager and execute command try: browser_manager = get_browser_manager() result = await browser_manager.execute_command(command, params) if result.get("success"): return JSONResponse(content=result) else: raise HTTPException(status_code=400, detail=result.get("error", "Command failed")) except Exception as e: logger.error(f"Error executing playwright command: {str(e)}") logger.error(traceback.format_exc()) raise HTTPException(status_code=500, detail=str(e)) if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=8000)