Fix: restore original pty speed + keep sessions alive

app.py

@@ -17,11 +17,11 @@ app = FastAPI()
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 PERSISTENT_DIR = Path("/persistent")
-SESSION_DIR = PERSISTENT_DIR / "sessions"
 PKGS_FILE = PERSISTENT_DIR / "installed_packages.txt"
-SESSION_STATE = PERSISTENT_DIR / "session_state.json"
+PERSISTENT_DIR.mkdir(parents=True, exist_ok=True)
 
-SESSION_DIR.mkdir(parents=True, exist_ok=True)
+# In-memory session storage - ptys stay alive even after WebSocket disconnects
+sessions = {}
 
 @app.get("/")
 async def index():
@@ -33,20 +33,11 @@ async def get_pass():
 
 @app.get("/api/sessions")
 async def list_sessions():
-    sessions = []
-    for f in SESSION_DIR.glob("*.json"):
-        try:
-            data = json.loads(f.read_text())
-            sid = data.get("id", f.stem)
-            alive = subprocess.run(
-                ["screen", "-ls", f"term_{sid}"],
-                capture_output=True
-            ).returncode == 0
-            data["alive"] = alive
-            sessions.append(data)
-        except:
-            pass
-    return JSONResponse(sessions)
+    alive = []
+    for sid, s in sessions.items():
+        if s.get("proc") and s["proc"].returncode is None:
+            alive.append({"id": sid, "alive": True})
+    return JSONResponse(alive)
 
 @app.websocket("/ws")
 async def terminal(websocket: WebSocket):
@@ -54,6 +45,8 @@ async def terminal(websocket: WebSocket):
 
     session_id = None
     read_task = None
+    master_fd = None
+    proc = None
 
     try:
         while True:
@@ -62,113 +55,112 @@ async def terminal(websocket: WebSocket):
 
             if msg.get("type") == "attach":
                 session_id = msg.get("session_id")
-                if not session_id or not (SESSION_DIR / f"{session_id}.json").exists():
-                    session_id = str(uuid.uuid4())[:8]
-                    (SESSION_DIR / f"{session_id}.json").write_text(json.dumps({
-                        "id": session_id,
-                        "created": str(subprocess.getoutput("date -Iseconds")),
-                    }))
 
-                await websocket.send_text(json.dumps({
-                    "type": "session_info",
-                    "session_id": session_id
-                }))
-
-                alive = subprocess.run(
-                    ["screen", "-ls", f"term_{session_id}"],
-                    capture_output=True
-                ).returncode == 0
-
-                if not alive:
-                    subprocess.Popen(
-                        ["screen", "-dmS", f"term_{session_id}", "-L", "-Logfile",
-                         f"/persistent/sessions/{session_id}.log", "/bin/bash", "--login"],
-                        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
-                    )
-                    await asyncio.sleep(0.3)
-
-                    state_file = PERSISTENT_DIR / "session_state.json"
-                    if state_file.exists():
+                # Reuse existing session if alive
+                if session_id and session_id in sessions:
+                    s = sessions[session_id]
+                    if s.get("proc") and s["proc"].returncode is None:
+                        master_fd = s["master_fd"]
+                        proc = s["proc"]
+                        await websocket.send_text(json.dumps({
+                            "type": "session_info",
+                            "session_id": session_id
+                        }))
+                        # Send current terminal state
                         try:
-                            state = json.loads(state_file.read_text())
-                            cwd = state.get("cwd", "/persistent")
-                            subprocess.run(
-                                ["screen", "-S", f"term_{session_id}", "-X", "stuff",
-                                 f"cd {cwd}\nclear\n"],
-                                capture_output=True, timeout=3
-                            )
+                            output = s.get("last_output", "")
+                            if output:
+                                await websocket.send_text(json.dumps({
+                                    "type": "output",
+                                    "data": output[-4000:]
+                                }))
                         except:
                             pass
-
-                subprocess.run(
-                    ["screen", "-S", f"term_{session_id}", "-X", "stuff", " \n"],
-                    capture_output=True, timeout=3
+                        if read_task:
+                            read_task.cancel()
+                        read_task = asyncio.create_task(read_pty(websocket, master_fd, session_id))
+                        continue
+
+                # Create new session
+                session_id = session_id or str(uuid.uuid4())[:8]
+                master_fd, slave_fd = pty.openpty()
+
+                env = {
+                    **os.environ,
+                    "TERM": "xterm-256color",
+                    "HOME": "/persistent",
+                    "USER": "ubuntu",
+                    "SHELL": "/bin/bash",
+                }
+
+                proc = await asyncio.create_subprocess_exec(
+                    "/bin/bash", "--login",
+                    stdin=slave_fd, stdout=slave_fd, stderr=slave_fd,
+                    env=env
                 )
+                os.close(slave_fd)
+
+                sessions[session_id] = {
+                    "master_fd": master_fd,
+                    "proc": proc,
+                    "last_output": ""
+                }
+
+                await websocket.send_text(json.dumps({
+                    "type": "session_info",
+                    "session_id": session_id
+                }))
 
                 if read_task:
                     read_task.cancel()
-                read_task = asyncio.create_task(read_screen(websocket, session_id))
-                continue
-
-            if msg.get("type") == "input" and session_id:
-                data = msg["data"]
-                subprocess.run(
-                    ["screen", "-S", f"term_{session_id}", "-X", "stuff", data],
-                    capture_output=True, timeout=5
-                )
+                read_task = asyncio.create_task(read_pty(websocket, master_fd, session_id))
 
-            elif msg.get("type") == "resize" and session_id:
-                cols = msg.get("cols", 200)
-                rows = msg.get("rows", 50)
-                subprocess.run(
-                    ["screen", "-S", f"term_{session_id}", "-X", "width", "cols", str(cols)],
-                    capture_output=True, timeout=3
-                )
-                subprocess.run(
-                    ["screen", "-S", f"term_{session_id}", "-X", "width", "lines", str(rows)],
-                    capture_output=True, timeout=3
-                )
+            elif msg.get("type") == "input":
+                if master_fd is not None:
+                    try:
+                        os.write(master_fd, msg["data"].encode("utf-8"))
+                    except OSError:
+                        pass
 
-            elif msg.get("type") == "key" and session_id:
-                key = msg.get("key", "")
-                subprocess.run(
-                    ["screen", "-S", f"term_{session_id}", "-X", "stuff", key],
-                    capture_output=True, timeout=3
-                )
+            elif msg.get("type") == "resize":
+                if master_fd is not None:
+                    try:
+                        fcntl.ioctl(
+                            master_fd, termios.TIOCSWINSZ,
+                            struct.pack("HHHH", msg.get("rows", 24), msg.get("cols", 80), 0, 0)
+                        )
+                    except OSError:
+                        pass
 
     except (WebSocketDisconnect, Exception):
         pass
     finally:
         if read_task:
             read_task.cancel()
+        # DON'T kill the process - keep it alive for reconnection
+        # Just clean up the WebSocket reference
 
-async def read_screen(websocket: WebSocket, session_id: str):
-    last_content = ""
+async def read_pty(websocket: WebSocket, master_fd: int, session_id: str):
+    loop = asyncio.get_event_loop()
     try:
         while True:
-            await asyncio.sleep(0.05)
-            try:
-                result = subprocess.run(
-                    ["screen", "-S", f"term_{session_id}", "-X", "hardcopy", "-h", "/tmp/screen_out.txt"],
-                    capture_output=True, timeout=3
-                )
-                if result.returncode == 0:
-                    try:
-                        content = Path("/tmp/screen_out.txt").read_text(errors="ignore")
-                        if content != last_content:
-                            if len(content) > len(last_content):
-                                new_data = content[len(last_content):]
-                            else:
-                                new_data = content
-                            last_content = content
-                            if new_data:
-                                await websocket.send_text(json.dumps({
-                                    "type": "output",
-                                    "data": new_data
-                                }))
-                    except:
-                        pass
-            except:
-                break
-    except asyncio.CancelledError:
+            data = await loop.run_in_executor(None, lambda: os.read(master_fd, 4096))
+            if data:
+                try:
+                    text = data.decode("utf-8")
+                except UnicodeDecodeError:
+                    text = data.decode("utf-8", errors="ignore")
+
+                # Cache last output for reconnect
+                if session_id in sessions:
+                    sessions[session_id]["last_output"] = sessions[session_id].get("last_output", "") + text
+                    # Keep only last 8000 chars
+                    if len(sessions[session_id]["last_output"]) > 8000:
+                        sessions[session_id]["last_output"] = sessions[session_id]["last_output"][-4000:]
+
+                try:
+                    await websocket.send_text(json.dumps({"type": "output", "data": text}))
+                except:
+                    break
+    except (OSError, asyncio.CancelledError):
         pass