tomasz-syn-grzegorza
73e4fc005e
Server UI gains Update on stack cards, ComfyUI Models tab with workflow scan and downloads, and centralized comfyui_config. Model catalog and download scripts move from stacks/comfyui to server-ui so ComfyUI stays a minimal Docker wrapper for easier image updates.
364 lines
11 KiB
Python
364 lines
11 KiB
Python
"""PTY shell sessions for browser CLI (WebSocket bridge)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import fcntl
|
|
import json
|
|
import logging
|
|
import os
|
|
import pty
|
|
import shlex
|
|
import signal
|
|
import struct
|
|
import termios
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from starlette.websockets import WebSocket, WebSocketDisconnect
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
TIOCSWINSZ = termios.TIOCSWINSZ
|
|
OUTPUT_BUFFER_MAX = 256 * 1024
|
|
|
|
|
|
class CliSessionLimitError(Exception):
|
|
pass
|
|
|
|
|
|
class CliSessionNotFoundError(Exception):
|
|
pass
|
|
|
|
|
|
def _set_winsize(fd: int, rows: int, cols: int) -> None:
|
|
winsize = struct.pack("HHHH", rows, cols, 0, 0)
|
|
fcntl.ioctl(fd, TIOCSWINSZ, winsize)
|
|
|
|
|
|
def _parse_shell(shell: str) -> list[str]:
|
|
parts = shlex.split(shell.strip())
|
|
return parts if parts else ["/bin/bash"]
|
|
|
|
|
|
async def _kill_process(process: asyncio.subprocess.Process) -> None:
|
|
if process.returncode is not None:
|
|
return
|
|
try:
|
|
os.killpg(process.pid, signal.SIGTERM)
|
|
except ProcessLookupError:
|
|
return
|
|
try:
|
|
await asyncio.wait_for(process.wait(), timeout=2.0)
|
|
except asyncio.TimeoutError:
|
|
try:
|
|
os.killpg(process.pid, signal.SIGKILL)
|
|
except ProcessLookupError:
|
|
pass
|
|
await process.wait()
|
|
|
|
|
|
@dataclass
|
|
class CliSession:
|
|
id: str
|
|
master_fd: int
|
|
process: asyncio.subprocess.Process
|
|
cols: int = 80
|
|
rows: int = 24
|
|
ws: WebSocket | None = None
|
|
output_buffer: bytearray = field(default_factory=bytearray)
|
|
running: asyncio.Event = field(default_factory=asyncio.Event)
|
|
read_task: asyncio.Task[None] | None = None
|
|
detach_timer: asyncio.Task[None] | None = None
|
|
closed: bool = False
|
|
|
|
def append_output(self, data: bytes) -> None:
|
|
self.output_buffer.extend(data)
|
|
overflow = len(self.output_buffer) - OUTPUT_BUFFER_MAX
|
|
if overflow > 0:
|
|
del self.output_buffer[:overflow]
|
|
|
|
async def send_to_ws(self, data: bytes) -> bool:
|
|
if not self.ws:
|
|
return False
|
|
try:
|
|
await self.ws.send_bytes(data)
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
async def send_json(self, obj: dict[str, Any]) -> bool:
|
|
if not self.ws:
|
|
return False
|
|
try:
|
|
await self.ws.send_text(json.dumps(obj))
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
class CliSessionManager:
|
|
def __init__(self, *, session_ttl: int = 300, max_sessions: int = 5) -> None:
|
|
self.session_ttl = session_ttl
|
|
self.max_sessions = max_sessions
|
|
self._sessions: dict[str, CliSession] = {}
|
|
self._lock = asyncio.Lock()
|
|
|
|
def sessions_available(self) -> bool:
|
|
return len(self._sessions) < self.max_sessions
|
|
|
|
def get(self, session_id: str) -> CliSession | None:
|
|
return self._sessions.get(session_id)
|
|
|
|
async def create_session(
|
|
self,
|
|
*,
|
|
shell: str,
|
|
cwd: str,
|
|
env: dict[str, str],
|
|
) -> CliSession:
|
|
async with self._lock:
|
|
if len(self._sessions) >= self.max_sessions:
|
|
raise CliSessionLimitError(
|
|
f"Limit sesji CLI ({self.max_sessions}) — spróbuj później"
|
|
)
|
|
|
|
master_fd, slave_fd = pty.openpty()
|
|
shell_cmd = _parse_shell(shell)
|
|
proc_env = {**env, "TERM": "xterm-256color"}
|
|
|
|
process = await asyncio.create_subprocess_exec(
|
|
*shell_cmd,
|
|
stdin=slave_fd,
|
|
stdout=slave_fd,
|
|
stderr=slave_fd,
|
|
cwd=cwd,
|
|
env=proc_env,
|
|
preexec_fn=os.setsid,
|
|
)
|
|
os.close(slave_fd)
|
|
|
|
session = CliSession(
|
|
id=str(uuid.uuid4()),
|
|
master_fd=master_fd,
|
|
process=process,
|
|
)
|
|
session.running.set()
|
|
|
|
async with self._lock:
|
|
self._sessions[session.id] = session
|
|
|
|
session.read_task = asyncio.create_task(self._read_pty_loop(session))
|
|
log.info("CLI session created: %s", session.id)
|
|
return session
|
|
|
|
async def _read_pty_loop(self, session: CliSession) -> None:
|
|
loop = asyncio.get_running_loop()
|
|
while session.running.is_set() and not session.closed:
|
|
try:
|
|
data = await loop.run_in_executor(None, os.read, session.master_fd, 4096)
|
|
except OSError:
|
|
break
|
|
if not data:
|
|
break
|
|
session.append_output(data)
|
|
await session.send_to_ws(data)
|
|
if session.process.returncode is not None:
|
|
break
|
|
|
|
if not session.closed:
|
|
log.info("CLI session ended (shell exit): %s", session.id)
|
|
await self.close_session(session.id)
|
|
|
|
async def attach_session(self, session_id: str, websocket: WebSocket) -> CliSession:
|
|
session = self.get(session_id)
|
|
if session is None or session.closed:
|
|
raise CliSessionNotFoundError(session_id)
|
|
|
|
await self._cancel_detach_timer(session)
|
|
|
|
if session.ws is not None and session.ws is not websocket:
|
|
try:
|
|
await session.ws.close()
|
|
except Exception:
|
|
pass
|
|
|
|
session.ws = websocket
|
|
|
|
if session.output_buffer:
|
|
await session.send_to_ws(bytes(session.output_buffer))
|
|
|
|
await session.send_json({"type": "replay", "complete": True})
|
|
_set_winsize(session.master_fd, session.rows, session.cols)
|
|
log.info("CLI session reattached: %s", session.id)
|
|
return session
|
|
|
|
async def detach_session(self, session_id: str) -> None:
|
|
session = self.get(session_id)
|
|
if session is None or session.closed:
|
|
return
|
|
session.ws = None
|
|
await self._schedule_detach_timer(session)
|
|
log.info("CLI session detached (TTL %ds): %s", self.session_ttl, session.id)
|
|
|
|
async def close_session(self, session_id: str) -> None:
|
|
async with self._lock:
|
|
session = self._sessions.pop(session_id, None)
|
|
if session is None or session.closed:
|
|
return
|
|
|
|
session.closed = True
|
|
session.running.clear()
|
|
await self._cancel_detach_timer(session)
|
|
|
|
if session.read_task is not None:
|
|
session.read_task.cancel()
|
|
try:
|
|
await session.read_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
await _kill_process(session.process)
|
|
try:
|
|
os.close(session.master_fd)
|
|
except OSError:
|
|
pass
|
|
|
|
if session.ws is not None:
|
|
try:
|
|
await session.ws.close()
|
|
except Exception:
|
|
pass
|
|
|
|
log.info("CLI session closed: %s", session_id)
|
|
|
|
async def _cancel_detach_timer(self, session: CliSession) -> None:
|
|
if session.detach_timer is not None:
|
|
session.detach_timer.cancel()
|
|
try:
|
|
await session.detach_timer
|
|
except asyncio.CancelledError:
|
|
pass
|
|
session.detach_timer = None
|
|
|
|
async def _schedule_detach_timer(self, session: CliSession) -> None:
|
|
await self._cancel_detach_timer(session)
|
|
|
|
async def _expire() -> None:
|
|
try:
|
|
await asyncio.sleep(self.session_ttl)
|
|
except asyncio.CancelledError:
|
|
return
|
|
log.info("CLI session expired after detach TTL: %s", session.id)
|
|
await self.close_session(session.id)
|
|
|
|
session.detach_timer = asyncio.create_task(_expire())
|
|
|
|
async def handle_resize(self, session: CliSession, rows: int, cols: int) -> None:
|
|
session.rows = rows
|
|
session.cols = cols
|
|
_set_winsize(session.master_fd, rows, cols)
|
|
|
|
async def write_input(self, session: CliSession, payload: bytes) -> None:
|
|
try:
|
|
os.write(session.master_fd, payload)
|
|
except OSError:
|
|
await self.close_session(session.id)
|
|
|
|
|
|
_manager: CliSessionManager | None = None
|
|
|
|
|
|
def init_session_manager(*, session_ttl: int = 300, max_sessions: int = 5) -> CliSessionManager:
|
|
global _manager
|
|
_manager = CliSessionManager(session_ttl=session_ttl, max_sessions=max_sessions)
|
|
return _manager
|
|
|
|
|
|
def get_session_manager() -> CliSessionManager:
|
|
if _manager is None:
|
|
raise RuntimeError("CLI session manager not initialized")
|
|
return _manager
|
|
|
|
|
|
def sessions_available(max_sessions: int) -> bool:
|
|
if _manager is not None:
|
|
return _manager.sessions_available()
|
|
return True
|
|
|
|
|
|
async def _handle_ws_input(session: CliSession, message: dict[str, Any], mgr: CliSessionManager) -> bool:
|
|
"""Process one WS message. Returns False to end the connection loop."""
|
|
msg_type = message.get("type")
|
|
if msg_type == "websocket.disconnect":
|
|
return False
|
|
|
|
if message.get("bytes"):
|
|
await mgr.write_input(session, message["bytes"])
|
|
return True
|
|
|
|
if message.get("text"):
|
|
text = message["text"]
|
|
if text.startswith("{"):
|
|
try:
|
|
obj: dict[str, Any] = json.loads(text)
|
|
cmd = obj.get("type")
|
|
if cmd == "resize":
|
|
rows = int(obj.get("rows", 24))
|
|
cols = int(obj.get("cols", 80))
|
|
await mgr.handle_resize(session, rows, cols)
|
|
return True
|
|
if cmd == "close":
|
|
await mgr.close_session(session.id)
|
|
return False
|
|
except (json.JSONDecodeError, TypeError, ValueError):
|
|
pass
|
|
await mgr.write_input(session, text.encode("utf-8"))
|
|
|
|
if session.closed or session.process.returncode is not None:
|
|
return False
|
|
return True
|
|
|
|
|
|
async def run_pty_session(
|
|
websocket: WebSocket,
|
|
*,
|
|
shell: str,
|
|
cwd: str,
|
|
env: dict[str, str],
|
|
max_sessions: int,
|
|
session_id: str | None = None,
|
|
session_ttl: int = 300,
|
|
) -> None:
|
|
if _manager is None:
|
|
init_session_manager(session_ttl=session_ttl, max_sessions=max_sessions)
|
|
mgr = get_session_manager()
|
|
|
|
session: CliSession
|
|
|
|
if session_id:
|
|
try:
|
|
session = await mgr.attach_session(session_id, websocket)
|
|
except CliSessionNotFoundError:
|
|
await websocket.close(code=1008, reason="Session expired")
|
|
return
|
|
else:
|
|
session = await mgr.create_session(shell=shell, cwd=cwd, env=env)
|
|
session.ws = websocket
|
|
await websocket.send_text(json.dumps({"type": "session", "id": session.id}))
|
|
_set_winsize(session.master_fd, session.rows, session.cols)
|
|
|
|
try:
|
|
while not session.closed:
|
|
message = await websocket.receive()
|
|
if not await _handle_ws_input(session, message, mgr):
|
|
break
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
if not session.closed:
|
|
if session.ws is websocket:
|
|
session.ws = None
|
|
await mgr.detach_session(session.id)
|