Capture debug screenshot on failure

backend/core/trae_bot.py

@@ -10,6 +10,8 @@ import logging
 import threading
 import json
 import os
+import base64
+import traceback
 
 logger = logging.getLogger("TraeBot")
 
@@ -24,6 +26,11 @@ class TraeBot:
         self.logs = []
         self.status = "idle"  # idle, running, waiting_for_scan, success, error
         self.current_stage = "Initializing"
+        self.last_error = None
+        self.last_traceback = None
+        self.last_url = None
+        self.last_html = None
+        self.last_screenshot_png_b64 = None
 
     def log(self, message: str, level="info"):
         entry = {"timestamp": time.time(), "message": message, "level": level}
@@ -40,6 +47,27 @@ class TraeBot:
     def update_stage(self, stage: str):
         self.current_stage = stage
 
+    def _capture_debug(self, page, exc: Exception):
+        self.last_error = f"{type(exc).__name__}: {exc}"
+        self.last_traceback = traceback.format_exc()
+        try:
+            self.last_url = getattr(page, "url", None)
+        except Exception:
+            self.last_url = None
+        try:
+            html = page.content()
+            if isinstance(html, str) and len(html) > 100000:
+                html = html[:100000]
+            self.last_html = html
+        except Exception:
+            self.last_html = None
+        try:
+            png = page.screenshot(full_page=True, type="png")
+            if isinstance(png, (bytes, bytearray)):
+                self.last_screenshot_png_b64 = base64.b64encode(png).decode("ascii")
+        except Exception:
+            self.last_screenshot_png_b64 = None
+
     def generate_password(self, length: int = 12) -> str:
         characters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()"
         return "".join(random.choice(characters) for _ in range(length))
@@ -48,6 +76,7 @@ class TraeBot:
         self.status = "running"
         self.update_stage("Starting Browser")
         start_time = time.time()
+        page = None
         
         try:
             self.log("Starting Browser and navigate to Trae.ai")
@@ -58,7 +87,7 @@ class TraeBot:
                 window=(1280, 720),
                 humanize=False
             ) as browser:
-                page: Page = browser.new_page()
+                page = browser.new_page()
 
                 def check_timeout():
                     if self.status != "waiting_for_scan" and time.time() - start_time > 60:
@@ -254,6 +283,8 @@ class TraeBot:
                     page.goto("https://www.trae.ai/account-settings")
                     
                 except Exception as e:
+                    if page:
+                        self._capture_debug(page, e)
                     self.log(f"Process failed inside browser: {e}", "error")
                     self.status = "error"
                     raise e
@@ -268,6 +299,8 @@ class TraeBot:
                 self.log("Registration process completed successfully")
 
         except Exception as e:
+            if page and not self.last_error:
+                self._capture_debug(page, e)
             self.log(f"Process failed: {e}", "error")
             self.status = "error"
             self.update_stage(f"Failed: {str(e)}")

backend/main.py

@@ -48,6 +48,10 @@ class StatusResponse(BaseModel):
     queue_position: int | None = None
     active_slots: int
     max_slots: int
+    last_error: str | None = None
+    last_traceback: str | None = None
+    last_url: str | None = None
+    has_screenshot: bool = False
 
 def run_bot(session_id: str):
     with sessions_lock:
@@ -162,8 +166,51 @@ async def get_status(session_id: str):
             "queue_position": queue_position,
             "active_slots": active_slots,
             "max_slots": max_slots,
+            "last_error": getattr(bot, "last_error", None),
+            "last_traceback": getattr(bot, "last_traceback", None),
+            "last_url": getattr(bot, "last_url", None),
+            "has_screenshot": bool(getattr(bot, "last_screenshot_png_b64", None)),
         }
 
+@app.get("/api/debug/{session_id}")
+async def get_debug(session_id: str):
+    with sessions_lock:
+        if session_id not in sessions:
+            raise HTTPException(status_code=404, detail="Session not found")
+        bot = sessions[session_id]["bot"]
+        return {
+            "session_id": session_id,
+            "status": bot.status,
+            "current_stage": bot.current_stage,
+            "last_error": getattr(bot, "last_error", None),
+            "last_traceback": getattr(bot, "last_traceback", None),
+            "last_url": getattr(bot, "last_url", None),
+            "last_html": getattr(bot, "last_html", None),
+            "has_screenshot": bool(getattr(bot, "last_screenshot_png_b64", None)),
+        }
+
+@app.get("/api/screenshot/{session_id}")
+async def get_screenshot(session_id: str):
+    with sessions_lock:
+        if session_id not in sessions:
+            raise HTTPException(status_code=404, detail="Session not found")
+        bot = sessions[session_id]["bot"]
+        b64 = getattr(bot, "last_screenshot_png_b64", None)
+        if not b64:
+            raise HTTPException(status_code=404, detail="Screenshot not available")
+        try:
+            import base64
+
+            png = base64.b64decode(b64.encode("ascii"))
+        except Exception:
+            raise HTTPException(status_code=500, detail="Screenshot decode failed")
+
+    return Response(
+        content=png,
+        media_type="image/png",
+        headers={"Content-Disposition": f'inline; filename="debug_{session_id}.png"'},
+    )
+
 @app.get("/api/download/{session_id}")
 async def download_token(session_id: str):
     if session_id not in sessions:

backend/tests/test_debug_capture.py

@@ -0,0 +1,45 @@
+import base64
+import unittest
+
+from core.trae_bot import TraeBot
+
+
+class DummyPage:
+    def __init__(self):
+        self.url = "https://example.com/x"
+
+    def content(self):
+        return "<html><body>hi</body></html>"
+
+    def screenshot(self, full_page=True, type="png"):
+        return b"\x89PNG\r\n\x1a\nfake"
+
+
+class DummyPageNoScreenshot(DummyPage):
+    def screenshot(self, full_page=True, type="png"):
+        raise RuntimeError("no screenshot")
+
+
+class TestDebugCapture(unittest.TestCase):
+    def test_capture_sets_fields(self):
+        bot = TraeBot("sid")
+        try:
+            raise ValueError("boom")
+        except Exception as e:
+            bot._capture_debug(DummyPage(), e)
+        self.assertIn("ValueError", bot.last_error)
+        self.assertIsInstance(bot.last_traceback, str)
+        self.assertEqual(bot.last_url, "https://example.com/x")
+        self.assertIn("hi", bot.last_html)
+        png = base64.b64decode(bot.last_screenshot_png_b64.encode("ascii"))
+        self.assertTrue(png.startswith(b"\x89PNG"))
+
+    def test_capture_survives_screenshot_failure(self):
+        bot = TraeBot("sid")
+        try:
+            raise RuntimeError("boom")
+        except Exception as e:
+            bot._capture_debug(DummyPageNoScreenshot(), e)
+        self.assertIn("RuntimeError", bot.last_error)
+        self.assertIsNone(bot.last_screenshot_png_b64)
+