Update app.py

app.py

@@ -30,52 +30,109 @@ SCREEN_W, SCREEN_H = 800, 600
 SNIPPETS_FILE = "/tmp/axon_snippets.json"
 
 # ═══════════════════════════════════════
-# Virtual Display
+# Virtual Display (Optimized for pygame/tkinter/turtle)
 # ═══════════════════════════════════════
 class VirtualDisplay:
     def __init__(self):
         self.xvfb_proc = None
         self.display = DISPLAY_NUM
+        self._last_capture = None
         self._start_xvfb()
 
     def _start_xvfb(self):
         try:
-            subprocess.run(["pkill", "-f", f"Xvfb {self.display}"], capture_output=True, timeout=5)
-            time.sleep(0.2)
-            self.xvfb_proc = subprocess.Popen(
-                ["Xvfb", self.display, "-screen", "0", f"{SCREEN_W}x{SCREEN_H}x24",
-                 "-ac", "-nolisten", "tcp"],
-                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
-            time.sleep(0.5)
+            subprocess.run(["pkill", "-f", f"Xvfb {self.display}"],
+                           capture_output=True, timeout=5)
+            time.sleep(0.3)
+            # Full-featured Xvfb: GLX for OpenGL, RENDER for anti-aliasing,
+            # 24-bit color, no access control
+            self.xvfb_proc = subprocess.Popen([
+                "Xvfb", self.display,
+                "-screen", "0", f"{SCREEN_W}x{SCREEN_H}x24",
+                "-ac",                    # No access control
+                "-nolisten", "tcp",       # Security
+                "+extension", "GLX",      # OpenGL support (pygame)
+                "+extension", "RENDER",   # Anti-aliased rendering
+            ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+            time.sleep(0.8)
             if self.xvfb_proc.poll() is None:
-                print(f"[Xvfb] Running on {self.display}")
                 os.environ["DISPLAY"] = self.display
+                print(f"[Xvfb] Running on {self.display} ({SCREEN_W}x{SCREEN_H})")
             else:
-                self.xvfb_proc = None
+                # Fallback: simpler Xvfb without extensions
+                self.xvfb_proc = subprocess.Popen([
+                    "Xvfb", self.display,
+                    "-screen", "0", f"{SCREEN_W}x{SCREEN_H}x24", "-ac",
+                ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+                time.sleep(0.5)
+                if self.xvfb_proc.poll() is None:
+                    os.environ["DISPLAY"] = self.display
+                    print(f"[Xvfb] Running (fallback mode)")
+                else:
+                    self.xvfb_proc = None
+                    print("[Xvfb] Failed to start")
         except Exception as e:
             print(f"[Xvfb] {e}"); self.xvfb_proc = None
 
     def capture(self):
+        """Capture the virtual display. Returns filepath or None."""
         if not self.is_running: return None
-        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False); tmp.close()
-        for cmd in [
-            ["scrot", "--display", self.display, tmp.name],
-            f"import -window root -display {self.display} {tmp.name}",
-        ]:
-            try:
-                kw = {"shell": True} if isinstance(cmd, str) else {}
-                r = subprocess.run(cmd, capture_output=True, timeout=3,
-                                   env={**os.environ, "DISPLAY": self.display}, **kw)
-                if r.returncode == 0 and os.path.exists(tmp.name) and os.path.getsize(tmp.name) > 100:
-                    return tmp.name
+
+        # Clean up previous capture to avoid /tmp filling up
+        if self._last_capture:
+            try: os.unlink(self._last_capture)
             except: pass
+
+        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+        tmp.close()
+        cap_env = {**os.environ, "DISPLAY": self.display}
+
+        # Method 1: xwd + convert (most reliable for Xvfb)
+        try:
+            xwd = subprocess.run(
+                ["xwd", "-root", "-display", self.display, "-silent"],
+                capture_output=True, timeout=3, env=cap_env)
+            if xwd.returncode == 0 and xwd.stdout:
+                conv = subprocess.run(
+                    ["convert", "xwd:-", tmp.name],
+                    input=xwd.stdout, capture_output=True, timeout=3)
+                if conv.returncode == 0 and os.path.getsize(tmp.name) > 100:
+                    self._last_capture = tmp.name
+                    return tmp.name
+        except: pass
+
+        # Method 2: scrot
+        try:
+            r = subprocess.run(
+                ["scrot", tmp.name],
+                capture_output=True, timeout=3, env=cap_env)
+            if r.returncode == 0 and os.path.exists(tmp.name) and os.path.getsize(tmp.name) > 100:
+                self._last_capture = tmp.name
+                return tmp.name
+        except: pass
+
+        # Method 3: import (ImageMagick)
+        try:
+            r = subprocess.run(
+                f"import -window root -display {self.display} {tmp.name}",
+                shell=True, capture_output=True, timeout=3, env=cap_env)
+            if r.returncode == 0 and os.path.exists(tmp.name) and os.path.getsize(tmp.name) > 100:
+                self._last_capture = tmp.name
+                return tmp.name
+        except: pass
+
         try: os.unlink(tmp.name)
         except: pass
         return None
 
     @property
-    def is_running(self): return self.xvfb_proc is not None and self.xvfb_proc.poll() is None
+    def is_running(self):
+        return self.xvfb_proc is not None and self.xvfb_proc.poll() is None
+
     def cleanup(self):
+        if self._last_capture:
+            try: os.unlink(self._last_capture)
+            except: pass
         if self.xvfb_proc:
             self.xvfb_proc.terminate()
             try: self.xvfb_proc.wait(timeout=3)
@@ -84,8 +141,10 @@ class VirtualDisplay:
 vdisplay = VirtualDisplay()
 
 # ═══════════════════════════════════════
-# GUI Process Manager
+# GUI Process Manager (pygame/tkinter/turtle)
 # ═══════════════════════════════════════
+GUI_LOG = "/tmp/_axon_gui.log"
+
 class GUIProcessManager:
     def __init__(self):
         self.process = None
@@ -94,28 +153,63 @@ class GUIProcessManager:
         self.stop()
         tmp = "/tmp/_axon_gui_run.py"
         with open(tmp, "w") as f: f.write(code)
+
+        # Full environment for GUI frameworks
         env = os.environ.copy()
-        env.update({"DISPLAY": DISPLAY_NUM, "SDL_VIDEODRIVER": "x11"})
+        env.update({
+            "DISPLAY": DISPLAY_NUM,
+            # SDL / pygame
+            "SDL_VIDEODRIVER": "x11",
+            "SDL_AUDIODRIVER": "dummy",        # No audio device in headless
+            "PYGAME_HIDE_SUPPORT_PROMPT": "1",  # Suppress pygame welcome
+            # Tkinter / general X11
+            "XDG_RUNTIME_DIR": "/tmp",
+            "MESA_GL_VERSION_OVERRIDE": "3.3",  # OpenGL compat
+            "LIBGL_ALWAYS_SOFTWARE": "1",        # Software rendering (no GPU)
+            # Python
+            "PYTHONUNBUFFERED": "1",
+            "PYTHONPATH": "/tmp/axon_workspace:" + env.get("PYTHONPATH", ""),
+        })
+
         try:
-            self.process = subprocess.Popen([sys.executable, tmp], stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE, env=env, preexec_fn=os.setsid)
-            time.sleep(0.5)
+            # Log file instead of PIPE — PIPE deadlocks pygame's output
+            log_f = open(GUI_LOG, "w")
+            self.process = subprocess.Popen(
+                [sys.executable, "-u", tmp],  # -u = unbuffered
+                stdout=log_f, stderr=log_f,
+                env=env, preexec_fn=os.setsid, cwd="/tmp/axon_workspace")
+
+            time.sleep(1.0)  # Give pygame time to create the window
+
             if self.process.poll() is not None:
-                out, err = self.process.communicate(timeout=2)
-                msg = (out.decode(errors="replace") + "\n" + err.decode(errors="replace")).strip()
-                return f"[Exited immediately]\n{msg}" if msg else "[Exited]"
-            return f"[GUI launched — PID {self.process.pid}] Switch to DISPLAY tab to view."
-        except Exception as e: return f"[Error] {e}"
+                log_f.close()
+                try:
+                    with open(GUI_LOG) as f: msg = f.read().strip()
+                except: msg = ""
+                return f"[Exited immediately]\n{msg}" if msg else "[Exited — no output]"
+
+            return f"[GUI launched — PID {self.process.pid}] Auto-refresh enabled."
+        except Exception as e:
+            return f"[Error] {e}"
+
+    def get_log(self):
+        try:
+            with open(GUI_LOG) as f: return f.read()[-2000:]
+        except: return ""
 
     def stop(self):
         if self.process and self.process.poll() is None:
             pid = self.process.pid
-            try: os.killpg(os.getpgid(pid), signal.SIGTERM); self.process.wait(timeout=3)
+            try:
+                os.killpg(os.getpgid(pid), signal.SIGTERM)
+                self.process.wait(timeout=3)
             except:
                 try: os.killpg(os.getpgid(pid), signal.SIGKILL)
                 except: pass
-            self.process = None; return f"[Stopped PID {pid}]"
-        self.process = None; return "[No process]"
+            self.process = None
+            return f"[Stopped PID {pid}]"
+        self.process = None
+        return "[No process]"
 
     @property
     def is_running(self): return self.process is not None and self.process.poll() is None
@@ -479,8 +573,12 @@ def ai_gen(system_prompt, code, max_tokens=300):
 # ═══════════════════════════════════════
 # GUI code detection
 # ═══════════════════════════════════════
-GUI_HINTS = ["pygame", "tkinter", "turtle", "pyglet", "arcade", "kivy",
-             "display.set_mode", "mainloop()"]
+GUI_HINTS = [
+    "pygame", "tkinter", "turtle", "pyglet", "arcade", "kivy",
+    "display.set_mode", "mainloop()", "Tk()", "Canvas(",
+    "pygame.init", "pygame.display", "screen.fill",
+    "from turtle import", "import turtle",
+]
 def is_gui_code(c): return any(h in c for h in GUI_HINTS)
 
 # ═══════════════════════════════════════
@@ -689,10 +787,13 @@ with gr.Blocks(title="Axon Pro") as demo:
                 display_image = gr.Image(label="", type="filepath", interactive=False, height=420)
                 with gr.Row():
                     capture_btn = gr.Button("📸 CAPTURE", size="sm", elem_classes="tb")
-                    auto_capture = gr.Checkbox(label="Auto (2s)", value=False)
+                    auto_capture = gr.Checkbox(label="Auto Refresh", value=False)
+                    refresh_speed = gr.Dropdown(
+                        choices=["0.5s", "1s", "2s", "5s"], value="1s",
+                        label="", container=False, scale=1, interactive=True)
                 gui_status = gr.Markdown(
                     f"<small>Xvfb: {'● ON' if vdisplay.is_running else '○ OFF'} "
-                    f"| {SCREEN_W}x{SCREEN_H}</small>")
+                    f"| {SCREEN_W}x{SCREEN_H} | SDL: x11 + dummy audio</small>")
 
     status_bar = gr.Markdown(
         f"**AXON PRO v4.0** │ Python {sys.version.split()[0]} │ CPU │ "
@@ -702,7 +803,7 @@ with gr.Blocks(title="Axon Pro") as demo:
     # State
     diff_original = gr.State("")
     diff_modified = gr.State("")
-    auto_timer = gr.Timer(2, active=False)
+    auto_timer = gr.Timer(1, active=False)
 
     # ═══════════════════════════════════════
     # HANDLERS
@@ -745,9 +846,12 @@ with gr.Blocks(title="Axon Pro") as demo:
         fs.save_file(code)
 
         if is_gui_code(code):
-            # GUI app — launch as background process, capture display
+            # GUI app — launch, take first screenshot, auto-enable refresh
             msg = gui_mgr.launch(code)
-            return msg, vdisplay.capture(), f"<small>Xvfb: ● ON | {gui_mgr.get_status()}</small>"
+            time.sleep(0.5)  # Extra time for first frame
+            ss = vdisplay.capture()
+            status = f"<small>Xvfb: ● ON | {gui_mgr.get_status()}</small>"
+            return msg, ss, status, gr.update(value=True)  # Enable auto-refresh
 
         # Normal script — run with subprocess, output goes to OUTPUT tab
         tmp = "/tmp/_axon_run.py"
@@ -768,11 +872,12 @@ with gr.Blocks(title="Axon Pro") as demo:
             output = "[Timed out after 30s]"
         except Exception as e:
             output = f"[Error] {e}"
-        return output, gr.update(), gr.update()
+        return output, gr.update(), gr.update(), gr.update()
 
     def on_stop():
-        msg = gui_mgr.stop(); terminal._append(msg)
-        return terminal.get_log(), f"<small>Xvfb: {'● ON' if vdisplay.is_running else '○ OFF'}</small>"
+        msg = gui_mgr.stop()
+        status = f"<small>Xvfb: {'● ON' if vdisplay.is_running else '○ OFF'} | ○ idle</small>"
+        return msg, status, gr.update(value=False)  # Disable auto-refresh
 
     # --- Terminal ---
     def on_term_cmd(cmd):
@@ -851,12 +956,19 @@ with gr.Blocks(title="Axon Pro") as demo:
         return vdisplay.capture()
 
     def on_auto_tick():
-        if gui_mgr.process and gui_mgr.process.poll() is None:
-            return vdisplay.capture()
-        return gr.update()
+        ss = vdisplay.capture()
+        # Auto-disable if GUI process died
+        if not gui_mgr.is_running:
+            return ss if ss else gr.update()
+        return ss if ss else gr.update()
 
     def on_auto_toggle(checked):
-        return gr.Timer(2, active=checked)
+        return gr.Timer(1, active=checked)
+
+    def on_speed_change(speed, auto_on):
+        speed_map = {"0.5s": 0.5, "1s": 1, "2s": 2, "5s": 5}
+        interval = speed_map.get(speed, 1)
+        return gr.Timer(interval, active=auto_on)
 
     # ═══════════════════════════════════════
     # WIRING
@@ -870,9 +982,9 @@ with gr.Blocks(title="Axon Pro") as demo:
     del_btn.click(on_delete, None, [file_list, editor, editor, structure_view])
 
     # Run
-    run_btn.click(on_run, editor, [run_output, display_image, gui_status]
+    run_btn.click(on_run, editor, [run_output, display_image, gui_status, auto_capture]
     ).then(lambda: gr.Tabs(selected="output-tab"), None, bottom_tabs)
-    stop_btn.click(on_stop, None, [term_out, gui_status])
+    stop_btn.click(on_stop, None, [run_output, gui_status, auto_capture])
 
     # Terminal
     term_in.submit(on_term_cmd, term_in, [term_out, term_in])
@@ -914,6 +1026,7 @@ with gr.Blocks(title="Axon Pro") as demo:
     # Display
     capture_btn.click(on_capture, None, display_image)
     auto_capture.change(on_auto_toggle, auto_capture, auto_timer)
+    refresh_speed.change(on_speed_change, [refresh_speed, auto_capture], auto_timer)
     auto_timer.tick(on_auto_tick, None, display_image)