trycua--cua/libs/python/computer-server/computer_server/main.py

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": "<base64-encoded bytes>"}``
    """
    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": "<base64>"}``
    - ``data: {"type": "exit", "code": <int>}``
    """
    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": "<base64>"}``
    - ``{"type": "resize",  "cols": N, "rows": N}``
    - ``{"type": "disconnect"}``

    Server → Client messages (JSON):
    - ``{"type": "output", "data": "<base64>"}``
    - ``{"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)