Update app.py

app.py

@@ -1,65 +1,41 @@
 """
-THE Z AI — Computer Mode Server (fixed)
-======================================
-
-Fixes included:
-- Robust Linux/Xvfb screenshot capture with scrot fallback
-- Explicit DISPLAY/XAUTHORITY handling
-- Safer Firefox launch helper (optional)
-- Clearer logging and error reporting
-- Better WebSocket resilience
-
-Compatible with the existing front-end WebSocket protocol.
+THE Z AI — Computer Mode Server
+================================
+سيرفر التحكم بالكمبيوتر عن بعد للذكاء الاصطناعي
+يوفر:
+  - WebSocket للتحكم اللحظي (الماوس، لوحة المفاتيح، الترمينال)
+  - تصوير سكرين شوت وإرساله للعميل
+  - بث شاشة مباشر (base64 frames)
+  - تنفيذ أوامر bash وقراءة الناتج
 """
 
-from __future__ import annotations
-
 import asyncio
 import base64
 import io
 import json
-import logging
 import os
-import shutil
 import subprocess
-import tempfile
+import threading
 import time
+import traceback
 from pathlib import Path
-from typing import Any, Dict, Optional
 
 import pyautogui
 import pyperclip
+from PIL import Image, ImageGrab
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
-from PIL import Image, ImageGrab
 import uvicorn
 
-# ---------------------------------------------------------------------
-# Environment / logging
-# ---------------------------------------------------------------------
-
-logging.basicConfig(
-    level=os.environ.get("LOG_LEVEL", "INFO").upper(),
-    format="[%(asctime)s] %(levelname)s %(name)s: %(message)s",
-)
-logger = logging.getLogger("z-computer-mode")
-
+# ─── ENV ────────────────────────────────────────────
 DISPLAY = os.environ.get("DISPLAY", ":1")
-XAUTHORITY = os.environ.get("XAUTHORITY", "/home/zai/.Xauthority")
 os.environ["DISPLAY"] = DISPLAY
-os.environ["XAUTHORITY"] = XAUTHORITY
-os.environ.setdefault("PYTHONUNBUFFERED", "1")
-
-# PyAutoGUI can be too aggressive in headless environments.
 pyautogui.FAILSAFE = False
-pyautogui.PAUSE = 0.03
-
-# ---------------------------------------------------------------------
-# App setup
-# ---------------------------------------------------------------------
+pyautogui.PAUSE = 0.05
 
 app = FastAPI(title="Z-Computer-Mode API")
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -68,141 +44,54 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# ─── Active Connections ──────────────────────────────
 active_connections: list[WebSocket] = []
 stream_active = False
-stream_quality = int(os.environ.get("STREAM_QUALITY", "60"))
-stream_fps = int(os.environ.get("STREAM_FPS", "3"))
-stream_scale = float(os.environ.get("STREAM_SCALE", "0.5"))
-
+stream_quality = 60   # JPEG quality
+stream_fps    = 3     # frames per second (low to save bandwidth)
+stream_scale  = 0.5   # downscale factor
 
-# ---------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------
 
-def _env_with_display(extra: Optional[Dict[str, str]] = None) -> Dict[str, str]:
-    env = dict(os.environ)
-    env["DISPLAY"] = DISPLAY
-    env["XAUTHORITY"] = XAUTHORITY
-    if extra:
-        env.update(extra)
-    return env
+# ─── Utility ─────────────────────────────────────────
 
-
-def _ensure_xauthority() -> None:
+def capture_screen(scale=stream_scale, quality=stream_quality) -> str:
+    """Capture screen → base64 JPEG string"""
     try:
-        Path(XAUTHORITY).parent.mkdir(parents=True, exist_ok=True)
-        Path(XAUTHORITY).touch(exist_ok=True)
-    except Exception as exc:
-        logger.warning("Could not prepare XAUTHORITY file: %s", exc)
-
-
-def _image_to_b64(img: Image.Image, quality: int = 70) -> str:
-    buf = io.BytesIO()
-    img.convert("RGB").save(buf, format="JPEG", quality=quality, optimize=True)
-    return base64.b64encode(buf.getvalue()).decode("ascii")
-
-
-def _capture_with_scrot() -> Optional[Image.Image]:
-    scrot = shutil.which("scrot")
-    if not scrot:
-        return None
-
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-        tmp_path = tmp.name
-
-    try:
-        # -z hides the cursor; remove it if you want cursor visible.
-        proc = subprocess.run(
-            [scrot, "-z", tmp_path],
-            capture_output=True,
-            text=True,
-            timeout=15,
-            env=_env_with_display(),
-        )
-        if proc.returncode != 0:
-            logger.debug("scrot failed: %s %s", proc.stdout, proc.stderr)
-            return None
-        if not Path(tmp_path).exists() or Path(tmp_path).stat().st_size == 0:
-            return None
-        return Image.open(tmp_path)
-    except Exception as exc:
-        logger.debug("scrot capture failed: %s", exc)
-        return None
-    finally:
-        try:
-            Path(tmp_path).unlink(missing_ok=True)
-        except Exception:
-            pass
-
-
-def _capture_with_pil() -> Optional[Image.Image]:
-    # ImageGrab on Linux often needs an explicit xdisplay.
-    try:
-        try:
-            return ImageGrab.grab(xdisplay=DISPLAY)
-        except TypeError:
-            # Older Pillow versions may not accept xdisplay.
-            return ImageGrab.grab()
-    except Exception as exc:
-        logger.debug("ImageGrab failed: %s", exc)
-        return None
-
-
-def capture_screen(scale: float = 0.5, quality: int = 60) -> str:
-    """Capture the Xvfb display and return a base64 JPEG string.
-
-    Strategy:
-    1) scrot (most reliable on Linux/Xvfb)
-    2) PIL.ImageGrab as fallback
-    3) return empty string on failure
-    """
-    try:
-        _ensure_xauthority()
-
-        img = _capture_with_scrot()
-        if img is None:
-            img = _capture_with_pil()
-        if img is None:
-            return ""
-
-        # Normalize and optionally resize.
-        if scale and 0 < scale < 1.0:
-            w = max(1, int(img.width * scale))
-            h = max(1, int(img.height * scale))
+        img = ImageGrab.grab()
+        if scale < 1.0:
+            w = int(img.width * scale)
+            h = int(img.height * scale)
             img = img.resize((w, h), Image.LANCZOS)
-
-        return _image_to_b64(img, quality=quality)
-    except Exception as exc:
-        logger.exception("capture_screen failed: %s", exc)
+        buf = io.BytesIO()
+        img.convert("RGB").save(buf, format="JPEG", quality=quality, optimize=True)
+        return base64.b64encode(buf.getvalue()).decode()
+    except Exception as e:
         return ""
 
 
-def run_command(cmd: str, timeout: int = 30) -> Dict[str, Any]:
-    """Run a shell command and return stdout/stderr/returncode."""
+def run_command(cmd: str, timeout: int = 30) -> dict:
+    """Run a shell command and return stdout/stderr"""
     try:
         result = subprocess.run(
-            cmd,
-            shell=True,
-            capture_output=True,
-            text=True,
-            timeout=timeout,
-            env=_env_with_display(),
+            cmd, shell=True, capture_output=True, text=True,
+            timeout=timeout, env={**os.environ, "DISPLAY": DISPLAY}
         )
         return {
-            "stdout": result.stdout[-8000:] if result.stdout else "",
-            "stderr": result.stderr[-4000:] if result.stderr else "",
+            "stdout": result.stdout[-8000:] if len(result.stdout) > 8000 else result.stdout,
+            "stderr": result.stderr[-4000:] if len(result.stderr) > 4000 else result.stderr,
             "returncode": result.returncode,
         }
     except subprocess.TimeoutExpired:
         return {"stdout": "", "stderr": f"Command timed out after {timeout}s", "returncode": -1}
-    except Exception as exc:
-        return {"stdout": "", "stderr": str(exc), "returncode": -1}
+    except Exception as e:
+        return {"stdout": "", "stderr": str(e), "returncode": -1}
 
 
-async def broadcast(msg: Dict[str, Any]) -> None:
-    txt = json.dumps(msg, ensure_ascii=False)
-    disconnected: list[WebSocket] = []
-    for ws in list(active_connections):
+async def broadcast(msg: dict):
+    """Send JSON message to all connected WebSocket clients"""
+    txt = json.dumps(msg)
+    disconnected = []
+    for ws in active_connections:
         try:
             await ws.send_text(txt)
         except Exception:
@@ -212,263 +101,233 @@ async def broadcast(msg: Dict[str, Any]) -> None:
             active_connections.remove(ws)
 
 
-async def screen_stream_loop() -> None:
+# ─── Screen Streaming Loop ───────────────────────────
+
+async def screen_stream_loop():
     global stream_active
-    interval = 1.0 / max(1, stream_fps)
+    interval = 1.0 / stream_fps
     while stream_active and active_connections:
         try:
-            frame = capture_screen(scale=stream_scale, quality=stream_quality)
+            frame = capture_screen()
             if frame:
                 await broadcast({"type": "frame", "data": frame})
-        except Exception as exc:
-            logger.debug("screen_stream_loop error: %s", exc)
+        except Exception:
+            pass
         await asyncio.sleep(interval)
     stream_active = False
 
 
-def _open_app_process(app_cmd: str) -> None:
-    if not app_cmd:
-        return
-    subprocess.Popen(
-        app_cmd,
-        shell=True,
-        stdout=subprocess.DEVNULL,
-        stderr=subprocess.DEVNULL,
-        env=_env_with_display(),
-        start_new_session=True,
-    )
-
-
-# Optional helper if you want the server to launch Firefox itself.
-# Not required by the protocol, but useful when start.sh is misordered.
-def launch_firefox_if_missing() -> None:
-    if not os.environ.get("AUTO_LAUNCH_FIREFOX", "0").lower() in {"1", "true", "yes"}:
-        return
+# ─── Action Handler ──────────────────────────────────
 
-    try:
-        # Heuristic only; do not fail if Firefox is already running.
-        subprocess.Popen(
-            "DISPLAY='{}' firefox --no-sandbox --new-instance about:blank".format(DISPLAY),
-            shell=True,
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.DEVNULL,
-            env=_env_with_display(),
-            start_new_session=True,
-        )
-        logger.info("Firefox launch requested by AUTO_LAUNCH_FIREFOX")
-    except Exception as exc:
-        logger.warning("Firefox auto-launch failed: %s", exc)
-
-
-# ---------------------------------------------------------------------
-# Action handler
-# ---------------------------------------------------------------------
-
-async def handle_action(ws: WebSocket, msg: Dict[str, Any]) -> None:
-    global stream_active, stream_fps, stream_quality, stream_scale
+async def handle_action(ws: WebSocket, msg: dict):
     action = msg.get("action", "")
-    data = msg.get("data", {}) or {}
+    data   = msg.get("data", {})
 
+    # ── SCREENSHOT ──────────────────────────────────
     if action == "screenshot":
         frame = capture_screen(scale=0.8, quality=75)
-        await ws.send_text(json.dumps({"type": "screenshot", "data": frame, "ts": int(time.time() * 1000)}, ensure_ascii=False))
+        await ws.send_text(json.dumps({
+            "type": "screenshot",
+            "data": frame,
+            "ts": int(time.time() * 1000)
+        }))
 
+    # ── TERMINAL COMMAND ─────────────────────────────
     elif action == "terminal":
-        cmd = str(data.get("cmd", ""))
+        cmd = data.get("cmd", "")
         if cmd:
-            result = run_command(cmd, timeout=int(data.get("timeout", 30)))
+            result = run_command(cmd, timeout=data.get("timeout", 30))
             await ws.send_text(json.dumps({
                 "type": "terminal_result",
                 "cmd": cmd,
                 "stdout": result["stdout"],
                 "stderr": result["stderr"],
                 "returncode": result["returncode"],
-            }, ensure_ascii=False))
+            }))
+            # Auto-screenshot after terminal command
             await asyncio.sleep(0.5)
             frame = capture_screen(scale=0.7, quality=70)
             if frame:
-                await ws.send_text(json.dumps({"type": "screenshot", "data": frame, "ts": int(time.time() * 1000), "auto": True}, ensure_ascii=False))
-
+                await ws.send_text(json.dumps({
+                    "type": "screenshot",
+                    "data": frame,
+                    "ts": int(time.time() * 1000),
+                    "auto": True
+                }))
+
+    # ── MOUSE MOVE ───────────────────────────────────
     elif action == "mouse_move":
         x, y = int(data.get("x", 0)), int(data.get("y", 0))
         pyautogui.moveTo(x, y, duration=0.1)
-        await ws.send_text(json.dumps({"type": "ack", "action": "mouse_move"}, ensure_ascii=False))
+        await ws.send_text(json.dumps({"type": "ack", "action": "mouse_move"}))
 
+    # ── MOUSE CLICK ──────────────────────────────────
     elif action == "mouse_click":
-        x = int(data.get("x", 0))
-        y = int(data.get("y", 0))
-        button = str(data.get("button", "left"))
-        double = bool(data.get("double", False))
+        x      = int(data.get("x", 0))
+        y      = int(data.get("y", 0))
+        button = data.get("button", "left")
+        double = data.get("double", False)
         pyautogui.moveTo(x, y, duration=0.08)
         if double:
             pyautogui.doubleClick(x, y, button=button)
         else:
             pyautogui.click(x, y, button=button)
-        await ws.send_text(json.dumps({"type": "ack", "action": "mouse_click"}, ensure_ascii=False))
+        await ws.send_text(json.dumps({"type": "ack", "action": "mouse_click"}))
 
+    # ── MOUSE DRAG ───────────────────────────────────
     elif action == "mouse_drag":
         x1, y1 = int(data.get("x1", 0)), int(data.get("y1", 0))
         x2, y2 = int(data.get("x2", 0)), int(data.get("y2", 0))
         pyautogui.moveTo(x1, y1)
         pyautogui.dragTo(x2, y2, duration=0.3, button="left")
-        await ws.send_text(json.dumps({"type": "ack", "action": "mouse_drag"}, ensure_ascii=False))
+        await ws.send_text(json.dumps({"type": "ack", "action": "mouse_drag"}))
 
+    # ── KEYBOARD TYPE ────────────────────────────────
     elif action == "keyboard_type":
-        text = str(data.get("text", ""))
+        text = data.get("text", "")
         pyautogui.typewrite(text, interval=0.03)
-        await ws.send_text(json.dumps({"type": "ack", "action": "keyboard_type"}, ensure_ascii=False))
+        await ws.send_text(json.dumps({"type": "ack", "action": "keyboard_type"}))
 
+    # ── KEYBOARD HOTKEY ──────────────────────────────
     elif action == "keyboard_hotkey":
         keys = data.get("keys", [])
         if keys:
             pyautogui.hotkey(*keys)
-        await ws.send_text(json.dumps({"type": "ack", "action": "keyboard_hotkey"}, ensure_ascii=False))
+        await ws.send_text(json.dumps({"type": "ack", "action": "keyboard_hotkey"}))
 
+    # ── KEYBOARD PRESS ───────────────────────────────
     elif action == "keyboard_press":
-        key = str(data.get("key", ""))
+        key = data.get("key", "")
         if key:
             pyautogui.press(key)
-        await ws.send_text(json.dumps({"type": "ack", "action": "keyboard_press"}, ensure_ascii=False))
+        await ws.send_text(json.dumps({"type": "ack", "action": "keyboard_press"}))
 
+    # ── CLIPBOARD WRITE ──────────────────────────────
     elif action == "clipboard_write":
-        text = str(data.get("text", ""))
+        text = data.get("text", "")
         pyperclip.copy(text)
-        await ws.send_text(json.dumps({"type": "ack", "action": "clipboard_write"}, ensure_ascii=False))
+        await ws.send_text(json.dumps({"type": "ack", "action": "clipboard_write"}))
 
+    # ── CLIPBOARD READ ───────────────────────────────
     elif action == "clipboard_read":
         text = pyperclip.paste()
-        await ws.send_text(json.dumps({"type": "clipboard_content", "text": text}, ensure_ascii=False))
+        await ws.send_text(json.dumps({"type": "clipboard_content", "text": text}))
 
+    # ── SCROLL ───────────────────────────────────────
     elif action == "scroll":
-        x = int(data.get("x", 0))
-        y = int(data.get("y", 0))
+        x      = int(data.get("x", 0))
+        y      = int(data.get("y", 0))
         clicks = int(data.get("clicks", 3))
         pyautogui.scroll(clicks, x=x, y=y)
-        await ws.send_text(json.dumps({"type": "ack", "action": "scroll"}, ensure_ascii=False))
+        await ws.send_text(json.dumps({"type": "ack", "action": "scroll"}))
 
+    # ── OPEN APP ─────────────────────────────────────
     elif action == "open_app":
-        app_cmd = str(data.get("cmd", ""))
+        app_cmd = data.get("cmd", "")
         if app_cmd:
-            _open_app_process(app_cmd)
-        await ws.send_text(json.dumps({"type": "ack", "action": "open_app"}, ensure_ascii=False))
+            subprocess.Popen(
+                app_cmd, shell=True,
+                env={**os.environ, "DISPLAY": DISPLAY},
+                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
+            )
+        await ws.send_text(json.dumps({"type": "ack", "action": "open_app"}))
+        # Auto-screenshot after app opens (wait for it to load)
         await asyncio.sleep(3)
         frame = capture_screen(scale=0.7, quality=70)
         if frame:
-            await ws.send_text(json.dumps({"type": "screenshot", "data": frame, "ts": int(time.time() * 1000), "auto": True}, ensure_ascii=False))
+            await ws.send_text(json.dumps({
+                "type": "screenshot",
+                "data": frame,
+                "ts": int(time.time() * 1000),
+                "auto": True
+            }))
 
+    # ── START STREAM ─────────────────────────────────
     elif action == "start_stream":
-        stream_fps = int(data.get("fps", 3))
+        global stream_active, stream_fps, stream_quality, stream_scale
+        stream_fps     = int(data.get("fps", 3))
         stream_quality = int(data.get("quality", 60))
-        stream_scale = float(data.get("scale", 0.5))
+        stream_scale   = float(data.get("scale", 0.5))
         if not stream_active:
             stream_active = True
             asyncio.create_task(screen_stream_loop())
-        await ws.send_text(json.dumps({"type": "ack", "action": "start_stream"}, ensure_ascii=False))
+        await ws.send_text(json.dumps({"type": "ack", "action": "start_stream"}))
 
+    # ── STOP STREAM ──────────────────────────────────
     elif action == "stop_stream":
         stream_active = False
-        await ws.send_text(json.dumps({"type": "ack", "action": "stop_stream"}, ensure_ascii=False))
+        await ws.send_text(json.dumps({"type": "ack", "action": "stop_stream"}))
 
+    # ── GET SCREEN INFO ──────────────────────────────
     elif action == "screen_info":
         try:
             w, h = pyautogui.size()
             mx, my = pyautogui.position()
-        except Exception:
+        except:
             w, h, mx, my = 1920, 1080, 0, 0
         await ws.send_text(json.dumps({
             "type": "screen_info",
-            "width": w,
-            "height": h,
-            "mouse_x": mx,
-            "mouse_y": my,
-        }, ensure_ascii=False))
+            "width": w, "height": h,
+            "mouse_x": mx, "mouse_y": my,
+        }))
 
     else:
-        await ws.send_text(json.dumps({"type": "error", "msg": f"Unknown action: {action}"}, ensure_ascii=False))
+        await ws.send_text(json.dumps({"type": "error", "msg": f"Unknown action: {action}"}))
 
 
-# ---------------------------------------------------------------------
-# WebSocket + REST endpoints
-# ---------------------------------------------------------------------
+# ─── WebSocket Endpoint ──────────────────────────────
 
 @app.websocket("/ws")
 async def websocket_endpoint(ws: WebSocket):
     await ws.accept()
     active_connections.append(ws)
-
+    # Send welcome
     try:
         w, h = pyautogui.size()
-    except Exception:
+    except:
         w, h = 1920, 1080
-
     await ws.send_text(json.dumps({
         "type": "connected",
         "screen_width": w,
         "screen_height": h,
-        "display": DISPLAY,
-        "msg": "Z Computer Mode — Connected",
-    }, ensure_ascii=False))
-
+        "msg": "Z Computer Mode — Connected"
+    }))
     try:
         while True:
             raw = await ws.receive_text()
-            try:
-                msg = json.loads(raw)
-            except json.JSONDecodeError:
-                await ws.send_text(json.dumps({"type": "error", "msg": "Invalid JSON"}, ensure_ascii=False))
-                continue
+            msg = json.loads(raw)
             await handle_action(ws, msg)
     except WebSocketDisconnect:
         pass
-    except Exception as exc:
-        logger.exception("WebSocket error: %s", exc)
+    except Exception as e:
+        pass
     finally:
         if ws in active_connections:
             active_connections.remove(ws)
 
 
+# ─── REST Fallback ───────────────────────────────────
+
 @app.get("/screenshot")
 async def rest_screenshot():
     frame = capture_screen(scale=0.75, quality=70)
-    return JSONResponse({"image": frame, "ts": int(time.time() * 1000), "display": DISPLAY})
+    return JSONResponse({"image": frame, "ts": int(time.time() * 1000)})
 
 
 @app.post("/terminal")
 async def rest_terminal(body: dict):
-    cmd = str(body.get("cmd", ""))
-    result = run_command(cmd, timeout=int(body.get("timeout", 30)))
+    cmd    = body.get("cmd", "")
+    result = run_command(cmd, timeout=body.get("timeout", 30))
     return JSONResponse(result)
 
 
 @app.get("/health")
 async def health():
-    try:
-        w, h = pyautogui.size()
-    except Exception:
-        w, h = 1920, 1080
-    return {
-        "status": "ok",
-        "display": DISPLAY,
-        "xauthority": XAUTHORITY,
-        "xauthority_exists": Path(XAUTHORITY).exists(),
-        "scrot": shutil.which("scrot"),
-        "firefox": shutil.which("firefox"),
-        "screen_width": w,
-        "screen_height": h,
-    }
-
-
-@app.on_event("startup")
-async def _startup() -> None:
-    _ensure_xauthority()
-    launch_firefox_if_missing()
-    logger.info("Z Computer Mode API started on DISPLAY=%s", DISPLAY)
+    return {"status": "ok", "display": DISPLAY}
 
 
-# ---------------------------------------------------------------------
-# Entry
-# ---------------------------------------------------------------------
+# ─── Entry ───────────────────────────────────────────
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))
-    uvicorn.run("app_fixed:app", host="0.0.0.0", port=port, log_level=os.environ.get("UVICORN_LOG_LEVEL", "warning"))
+    uvicorn.run("app:app", host="0.0.0.0", port=port, log_level="warning")