mirror of
https://github.com/trycua/computer.git
synced 2026-01-01 11:00:31 -06:00
124 lines
4.6 KiB
Python
124 lines
4.6 KiB
Python
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
|
from typing import List, Dict, Any
|
|
import uvicorn
|
|
import logging
|
|
import asyncio
|
|
import json
|
|
import traceback
|
|
from contextlib import redirect_stdout, redirect_stderr
|
|
from io import StringIO
|
|
from .handlers.factory import HandlerFactory
|
|
|
|
# Set up logging with more detail
|
|
logging.basicConfig(
|
|
level=logging.DEBUG, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Configure WebSocket with larger message size
|
|
WEBSOCKET_MAX_SIZE = 1024 * 1024 * 10 # 10MB limit
|
|
|
|
# Configure application with WebSocket settings
|
|
app = FastAPI(
|
|
title="Computer API",
|
|
description="API for the Computer project",
|
|
version="0.1.0",
|
|
websocket_max_size=WEBSOCKET_MAX_SIZE,
|
|
)
|
|
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
self.active_connections: List[WebSocket] = []
|
|
# Create OS-specific handlers
|
|
self.accessibility_handler, self.automation_handler = HandlerFactory.create_handlers()
|
|
|
|
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()
|
|
|
|
|
|
@app.websocket("/ws", name="websocket_endpoint")
|
|
async def websocket_endpoint(websocket: WebSocket):
|
|
# WebSocket message size is configured at the app or endpoint level, not on the instance
|
|
await manager.connect(websocket)
|
|
|
|
# Map commands to appropriate handler methods
|
|
handlers = {
|
|
# Accessibility commands
|
|
"get_accessibility_tree": manager.accessibility_handler.get_accessibility_tree,
|
|
"find_element": manager.accessibility_handler.find_element,
|
|
# Automation commands
|
|
"screenshot": manager.automation_handler.screenshot,
|
|
"left_click": manager.automation_handler.left_click,
|
|
"right_click": manager.automation_handler.right_click,
|
|
"double_click": manager.automation_handler.double_click,
|
|
"scroll_down": manager.automation_handler.scroll_down,
|
|
"scroll_up": manager.automation_handler.scroll_up,
|
|
"move_cursor": manager.automation_handler.move_cursor,
|
|
"type_text": manager.automation_handler.type_text,
|
|
"press_key": manager.automation_handler.press_key,
|
|
"drag_to": manager.automation_handler.drag_to,
|
|
"hotkey": manager.automation_handler.hotkey,
|
|
"get_cursor_position": manager.automation_handler.get_cursor_position,
|
|
"get_screen_size": manager.automation_handler.get_screen_size,
|
|
"copy_to_clipboard": manager.automation_handler.copy_to_clipboard,
|
|
"set_clipboard": manager.automation_handler.set_clipboard,
|
|
"run_command": manager.automation_handler.run_command,
|
|
}
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
data = await websocket.receive_json()
|
|
command = data.get("command")
|
|
params = data.get("params", {})
|
|
|
|
if command not in handlers:
|
|
await websocket.send_json(
|
|
{"success": False, "error": f"Unknown command: {command}"}
|
|
)
|
|
continue
|
|
|
|
try:
|
|
result = await handlers[command](**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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|