Add stack Update, ComfyUI model manager, and slim ComfyUI stack.
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.
This commit is contained in:
tomasz-syn-grzegorza
+286
-97
@@ -12,6 +12,8 @@ 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
|
||||
@@ -19,15 +21,17 @@ from starlette.websockets import WebSocket, WebSocketDisconnect
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
TIOCSWINSZ = termios.TIOCSWINSZ
|
||||
|
||||
_active_sessions = 0
|
||||
_session_lock = asyncio.Lock()
|
||||
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)
|
||||
@@ -38,25 +42,6 @@ def _parse_shell(shell: str) -> list[str]:
|
||||
return parts if parts else ["/bin/bash"]
|
||||
|
||||
|
||||
async def _read_pty_loop(
|
||||
master_fd: int,
|
||||
websocket: WebSocket,
|
||||
running: asyncio.Event,
|
||||
) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
while running.is_set():
|
||||
try:
|
||||
data = await loop.run_in_executor(None, os.read, master_fd, 4096)
|
||||
except OSError:
|
||||
break
|
||||
if not data:
|
||||
break
|
||||
try:
|
||||
await websocket.send_bytes(data)
|
||||
except Exception:
|
||||
break
|
||||
|
||||
|
||||
async def _kill_process(process: asyncio.subprocess.Process) -> None:
|
||||
if process.returncode is not None:
|
||||
return
|
||||
@@ -74,34 +59,71 @@ async def _kill_process(process: asyncio.subprocess.Process) -> None:
|
||||
await process.wait()
|
||||
|
||||
|
||||
def sessions_available(max_sessions: int) -> bool:
|
||||
return _active_sessions < max_sessions
|
||||
|
||||
|
||||
async def run_pty_session(
|
||||
websocket: WebSocket,
|
||||
*,
|
||||
shell: str,
|
||||
cwd: str,
|
||||
env: dict[str, str],
|
||||
max_sessions: int,
|
||||
) -> None:
|
||||
global _active_sessions
|
||||
|
||||
async with _session_lock:
|
||||
if _active_sessions >= max_sessions:
|
||||
raise CliSessionLimitError(
|
||||
f"Limit sesji CLI ({max_sessions}) — spróbuj później"
|
||||
)
|
||||
_active_sessions += 1
|
||||
|
||||
master_fd: int | None = None
|
||||
process: asyncio.subprocess.Process | None = None
|
||||
@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
|
||||
running = asyncio.Event()
|
||||
running.set()
|
||||
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"
|
||||
)
|
||||
|
||||
try:
|
||||
master_fd, slave_fd = pty.openpty()
|
||||
shell_cmd = _parse_shell(shell)
|
||||
proc_env = {**env, "TERM": "xterm-256color"}
|
||||
@@ -116,59 +138,226 @@ async def run_pty_session(
|
||||
preexec_fn=os.setsid,
|
||||
)
|
||||
os.close(slave_fd)
|
||||
slave_fd = -1
|
||||
|
||||
read_task = asyncio.create_task(_read_pty_loop(master_fd, websocket, running))
|
||||
session = CliSession(
|
||||
id=str(uuid.uuid4()),
|
||||
master_fd=master_fd,
|
||||
process=process,
|
||||
)
|
||||
session.running.set()
|
||||
|
||||
while True:
|
||||
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()
|
||||
msg_type = message.get("type")
|
||||
if msg_type == "websocket.disconnect":
|
||||
if not await _handle_ws_input(session, message, mgr):
|
||||
break
|
||||
|
||||
payload: bytes | None = None
|
||||
if message.get("bytes"):
|
||||
payload = message["bytes"]
|
||||
elif message.get("text"):
|
||||
text = message["text"]
|
||||
if text.startswith("{"):
|
||||
try:
|
||||
obj: dict[str, Any] = json.loads(text)
|
||||
if obj.get("type") == "resize":
|
||||
rows = int(obj.get("rows", 24))
|
||||
cols = int(obj.get("cols", 80))
|
||||
if master_fd is not None:
|
||||
_set_winsize(master_fd, rows, cols)
|
||||
continue
|
||||
except (json.JSONDecodeError, TypeError, ValueError):
|
||||
pass
|
||||
payload = text.encode("utf-8")
|
||||
|
||||
if payload and master_fd is not None:
|
||||
try:
|
||||
os.write(master_fd, payload)
|
||||
except OSError:
|
||||
break
|
||||
|
||||
if process.returncode is not None:
|
||||
break
|
||||
|
||||
except WebSocketDisconnect:
|
||||
pass
|
||||
finally:
|
||||
running.clear()
|
||||
if read_task is not None:
|
||||
read_task.cancel()
|
||||
try:
|
||||
await read_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
if process is not None:
|
||||
await _kill_process(process)
|
||||
if master_fd is not None:
|
||||
try:
|
||||
os.close(master_fd)
|
||||
except OSError:
|
||||
pass
|
||||
async with _session_lock:
|
||||
_active_sessions = max(0, _active_sessions - 1)
|
||||
if not session.closed:
|
||||
if session.ws is websocket:
|
||||
session.ws = None
|
||||
await mgr.detach_session(session.id)
|
||||
|
||||
Reference in New Issue
Block a user