Paul Payne
bd644ea905
- Run agent sessions under an isolated tmux server (its own systemd scope, with a keepalive session) so they survive a castle-api restart and stay rediscoverable; each WebSocket is a tmux attach client. Falls back to the in-memory pty backend when tmux is absent or CASTLE_API_AGENT_BACKEND=memory. - Refactor the session layer behind one Attachment interface (create/list/attach/ close), keeping resume-by-id, agent-native history, and continue modes. - Fix the terminal not refilling when the dock is maximized (explicit refit signal + instant size transition so xterm measures the final box). - On narrow screens, go full-bleed and strip chrome (header/border/padding) when maximized.
449 lines
15 KiB
Python
449 lines
15 KiB
Python
"""Live agent/terminal session registry.
|
|
|
|
A session is a running terminal (an agent CLI or a shell) whose lifetime is
|
|
decoupled from any WebSocket connection — closing a tab detaches, it does not
|
|
kill the process. Sessions can be listed, re-attached (resumed) by id, and
|
|
explicitly terminated.
|
|
|
|
Two backends behind one interface, chosen at startup:
|
|
|
|
- **tmux** (default when `tmux` is on PATH): each session is a tmux session, so
|
|
it **survives a castle-api restart** and is rediscoverable. The tmux server is
|
|
started inside its own systemd scope so it isn't in castle-api's cgroup (which
|
|
systemd would kill on restart); tmux's own systemd integration puts each pane
|
|
in its own scope too. Each WebSocket is a `tmux attach` client, so multi-client
|
|
and repaint-on-attach come for free.
|
|
- **memory** (fallback when tmux is absent, or `CASTLE_API_AGENT_BACKEND=memory`):
|
|
castle-api owns the pty directly and fans output out to subscribers with a
|
|
scrollback buffer. Simpler, but sessions die on restart.
|
|
|
|
The WebSocket handler is backend-agnostic: it calls `manager.create/get_info/
|
|
list/close`, then `manager.attach(id)` and drives the returned `Attachment`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import shlex
|
|
import shutil
|
|
import time
|
|
import uuid
|
|
from collections.abc import AsyncIterator
|
|
from typing import Protocol
|
|
|
|
from castle_core.config import CASTLE_HOME
|
|
|
|
from castle_api.pty_session import PtySession
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_SCROLLBACK_BYTES = 256 * 1024 # memory backend replay cap
|
|
_EXITED_TTL = 3600 # memory backend: reap exited sessions after this
|
|
|
|
|
|
class Attachment(Protocol):
|
|
"""One WebSocket's live view of a session."""
|
|
|
|
def reader(self) -> AsyncIterator[bytes]: ...
|
|
def write(self, data: bytes) -> None: ...
|
|
def resize(self, cols: int, rows: int) -> None: ...
|
|
async def aclose(self) -> None: ...
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# In-memory backend
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _MemSession:
|
|
def __init__(self, sid: str, agent: str, command: str, cwd: str) -> None:
|
|
self.id = sid
|
|
self.agent = agent
|
|
self.command = command
|
|
self.cwd = cwd
|
|
self.created_at = time.time()
|
|
self.pty: PtySession | None = None
|
|
self.scrollback = bytearray()
|
|
self.subscribers: set[asyncio.Queue[bytes | None]] = set()
|
|
self.exited = False
|
|
self.exit_code: int | None = None
|
|
self.exited_at: float | None = None
|
|
|
|
def _on_output(self, data: bytes) -> None:
|
|
self.scrollback.extend(data)
|
|
if len(self.scrollback) > _SCROLLBACK_BYTES:
|
|
del self.scrollback[: len(self.scrollback) - _SCROLLBACK_BYTES]
|
|
for q in list(self.subscribers):
|
|
try:
|
|
q.put_nowait(data)
|
|
except asyncio.QueueFull:
|
|
self.subscribers.discard(q)
|
|
|
|
def _on_exit(self) -> None:
|
|
self.exited = True
|
|
self.exit_code = self.pty.proc.returncode if self.pty else None
|
|
self.exited_at = time.time()
|
|
for q in list(self.subscribers):
|
|
q.put_nowait(None)
|
|
|
|
@property
|
|
def running(self) -> bool:
|
|
return not self.exited and self.pty is not None and self.pty.running
|
|
|
|
def info(self) -> dict:
|
|
return {
|
|
"id": self.id,
|
|
"agent": self.agent,
|
|
"command": self.command,
|
|
"cwd": self.cwd,
|
|
"created_at": self.created_at,
|
|
"running": self.running,
|
|
"exited": self.exited,
|
|
"exit_code": self.exit_code,
|
|
"cols": self.pty.cols if self.pty else None,
|
|
"rows": self.pty.rows if self.pty else None,
|
|
"clients": len(self.subscribers),
|
|
}
|
|
|
|
|
|
class _MemAttachment:
|
|
def __init__(
|
|
self, session: _MemSession, q: asyncio.Queue[bytes | None], snap: bytes
|
|
):
|
|
self._s = session
|
|
self._q = q
|
|
self._snap = snap
|
|
|
|
async def reader(self) -> AsyncIterator[bytes]:
|
|
if self._snap:
|
|
yield self._snap
|
|
while (chunk := await self._q.get()) is not None:
|
|
yield chunk
|
|
|
|
def write(self, data: bytes) -> None:
|
|
if self._s.pty:
|
|
self._s.pty.write(data)
|
|
|
|
def resize(self, cols: int, rows: int) -> None:
|
|
if self._s.pty:
|
|
self._s.pty.resize(cols, rows)
|
|
|
|
async def aclose(self) -> None:
|
|
self._s.subscribers.discard(self._q)
|
|
|
|
|
|
class MemoryManager:
|
|
backend = "memory"
|
|
|
|
def __init__(self) -> None:
|
|
self._sessions: dict[str, _MemSession] = {}
|
|
|
|
def _reap(self) -> None:
|
|
now = time.time()
|
|
for sid in [
|
|
s.id
|
|
for s in self._sessions.values()
|
|
if s.exited and s.exited_at and now - s.exited_at > _EXITED_TTL
|
|
]:
|
|
self._sessions.pop(sid, None)
|
|
|
|
async def create(
|
|
self,
|
|
agent: str,
|
|
argv: list[str],
|
|
cwd: str,
|
|
env: dict[str, str],
|
|
cols: int,
|
|
rows: int,
|
|
) -> str:
|
|
self._reap()
|
|
sid = uuid.uuid4().hex[:12]
|
|
s = _MemSession(sid, agent, shlex.join(argv), cwd)
|
|
s.pty = await PtySession.start(
|
|
argv[0],
|
|
args=argv[1:],
|
|
cwd=cwd or None,
|
|
env=env,
|
|
cols=cols,
|
|
rows=rows,
|
|
on_output=s._on_output,
|
|
on_exit=s._on_exit,
|
|
)
|
|
self._sessions[sid] = s
|
|
return sid
|
|
|
|
async def get_info(self, sid: str) -> dict | None:
|
|
s = self._sessions.get(sid)
|
|
return s.info() if s else None
|
|
|
|
async def list(self) -> list[dict]:
|
|
self._reap()
|
|
return [
|
|
s.info()
|
|
for s in sorted(
|
|
self._sessions.values(), key=lambda s: s.created_at, reverse=True
|
|
)
|
|
]
|
|
|
|
async def attach(self, sid: str, cols: int, rows: int) -> Attachment | None:
|
|
s = self._sessions.get(sid)
|
|
if s is None or not s.running or s.pty is None:
|
|
return None
|
|
# Subscribe + snapshot atomically (no await between → the output callback,
|
|
# in this same loop, can't interleave), so replay has no gap/dup.
|
|
q: asyncio.Queue[bytes | None] = asyncio.Queue(maxsize=256)
|
|
s.subscribers.add(q)
|
|
snap = bytes(s.scrollback)
|
|
s.pty.resize(cols, rows)
|
|
return _MemAttachment(s, q, snap)
|
|
|
|
async def close(self, sid: str) -> bool:
|
|
s = self._sessions.pop(sid, None)
|
|
if s is None:
|
|
return False
|
|
if s.pty:
|
|
await s.pty.close()
|
|
for q in list(s.subscribers):
|
|
q.put_nowait(None)
|
|
return True
|
|
|
|
async def close_all(self) -> None:
|
|
for sid in list(self._sessions):
|
|
await self.close(sid)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# tmux backend
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_TMUX_SOCK = CASTLE_HOME / "run" / "agents.sock"
|
|
_PREFIX = "castle-"
|
|
_KEEPALIVE = "__keepalive__" # hidden session that keeps the isolated server alive
|
|
|
|
|
|
class _TmuxAttachment:
|
|
def __init__(self, pty: PtySession, q: asyncio.Queue[bytes | None]):
|
|
self._pty = pty
|
|
self._q = q
|
|
|
|
@classmethod
|
|
async def start(
|
|
cls, tmux_bin: str, sock: str, name: str, cols: int, rows: int
|
|
) -> _TmuxAttachment:
|
|
q: asyncio.Queue[bytes | None] = asyncio.Queue(maxsize=256)
|
|
pty = await PtySession.start(
|
|
tmux_bin,
|
|
args=["-S", sock, "attach-session", "-t", name],
|
|
cols=cols,
|
|
rows=rows,
|
|
on_output=lambda d: q.put_nowait(d),
|
|
on_exit=lambda: q.put_nowait(None),
|
|
)
|
|
return cls(pty, q)
|
|
|
|
async def reader(self) -> AsyncIterator[bytes]:
|
|
while (chunk := await self._q.get()) is not None:
|
|
yield chunk
|
|
|
|
def write(self, data: bytes) -> None:
|
|
self._pty.write(data)
|
|
|
|
def resize(self, cols: int, rows: int) -> None:
|
|
self._pty.resize(cols, rows)
|
|
|
|
async def aclose(self) -> None:
|
|
await self._pty.close() # kills the attach client → detaches; session lives
|
|
|
|
|
|
class TmuxManager:
|
|
backend = "tmux"
|
|
|
|
def __init__(self, tmux_bin: str) -> None:
|
|
self._bin = tmux_bin
|
|
self._sock = str(_TMUX_SOCK)
|
|
self._ready = False
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def _raw(self, *args: str) -> tuple[int, str, str]:
|
|
proc = await asyncio.create_subprocess_exec(
|
|
self._bin,
|
|
"-S",
|
|
self._sock,
|
|
*args,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
out, err = await proc.communicate()
|
|
return (
|
|
proc.returncode or 0,
|
|
out.decode(errors="replace"),
|
|
err.decode(errors="replace"),
|
|
)
|
|
|
|
async def _server_running(self) -> bool:
|
|
_, out, err = await self._raw("list-sessions")
|
|
text = (out + err).lower()
|
|
return not any(
|
|
m in text for m in ("no server", "error connecting", "no such file")
|
|
)
|
|
|
|
async def _ensure(self) -> None:
|
|
if self._ready:
|
|
return
|
|
async with self._lock:
|
|
if self._ready:
|
|
return
|
|
_TMUX_SOCK.parent.mkdir(parents=True, exist_ok=True)
|
|
if not await self._server_running():
|
|
# Start the server inside its OWN systemd scope so it is NOT in
|
|
# castle-api's cgroup (which a restart kills). It must be born with
|
|
# a keepalive session: an empty server races its own empty-exit
|
|
# check against `exit-empty off` and dies, after which a plain
|
|
# `new-session` would restart it inside castle-api's cgroup. The
|
|
# keepalive (a detached `sleep infinity`, name-filtered from
|
|
# listings) guarantees the isolated server persists.
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"systemd-run",
|
|
"--user",
|
|
"--scope",
|
|
"--quiet",
|
|
"--collect",
|
|
self._bin,
|
|
"-S",
|
|
self._sock,
|
|
"new-session",
|
|
"-d",
|
|
"-s",
|
|
_KEEPALIVE,
|
|
"sleep infinity",
|
|
stdout=asyncio.subprocess.DEVNULL,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
_, err = await proc.communicate()
|
|
if proc.returncode != 0 or not await self._server_running():
|
|
logger.warning(
|
|
"agent tmux server not isolated (systemd-run: %s); sessions "
|
|
"will still work but won't survive a castle-api restart",
|
|
err.decode(errors="replace").strip() or "unknown error",
|
|
)
|
|
# Clean, terminal-like defaults (idempotent).
|
|
for opt in (
|
|
("set", "-g", "status", "off"), # no tmux status bar
|
|
("set", "-g", "window-size", "latest"), # newest client's size wins
|
|
("set", "-g", "escape-time", "0"),
|
|
("set", "-g", "destroy-unattached", "off"),
|
|
("set", "-s", "exit-empty", "off"),
|
|
):
|
|
await self._raw(*opt)
|
|
self._ready = True
|
|
|
|
async def create(
|
|
self,
|
|
agent: str,
|
|
argv: list[str],
|
|
cwd: str,
|
|
env: dict[str, str],
|
|
cols: int,
|
|
rows: int,
|
|
) -> str:
|
|
await self._ensure()
|
|
sid = uuid.uuid4().hex[:12]
|
|
name = f"{_PREFIX}{sid}"
|
|
cmd = shlex.join(argv) # tmux runs this as the pane's shell-command
|
|
args = [
|
|
"new-session",
|
|
"-d",
|
|
"-s",
|
|
name,
|
|
"-x",
|
|
str(cols),
|
|
"-y",
|
|
str(rows),
|
|
]
|
|
if cwd:
|
|
args += ["-c", cwd]
|
|
for k, v in env.items():
|
|
args += ["-e", f"{k}={v}"]
|
|
args += [cmd]
|
|
rc, _, err = await self._raw(*args)
|
|
if rc != 0:
|
|
raise RuntimeError(f"tmux new-session failed: {err.strip()}")
|
|
await self._raw("set-option", "-t", name, "@agent", agent)
|
|
return sid
|
|
|
|
async def _list_raw(self) -> list[dict]:
|
|
fmt = "#{session_name}\t#{session_created}\t#{@agent}\t#{session_attached}\t#{window_width}\t#{window_height}"
|
|
rc, out, _ = await self._raw("list-sessions", "-F", fmt)
|
|
if rc != 0:
|
|
return []
|
|
rows: list[dict] = []
|
|
for line in out.splitlines():
|
|
parts = line.split("\t")
|
|
if len(parts) < 6 or not parts[0].startswith(_PREFIX):
|
|
continue
|
|
name, created, agent, attached, w, h = parts[:6]
|
|
rows.append(
|
|
{
|
|
"id": name[len(_PREFIX) :],
|
|
"agent": agent or name[len(_PREFIX) :],
|
|
"command": "",
|
|
"cwd": "",
|
|
"created_at": float(created) if created.isdigit() else None,
|
|
"running": True,
|
|
"exited": False,
|
|
"exit_code": None,
|
|
"cols": int(w) if w.isdigit() else None,
|
|
"rows": int(h) if h.isdigit() else None,
|
|
"clients": int(attached) if attached.isdigit() else 0,
|
|
}
|
|
)
|
|
return rows
|
|
|
|
async def get_info(self, sid: str) -> dict | None:
|
|
for row in await self._list_raw():
|
|
if row["id"] == sid:
|
|
return row
|
|
return None
|
|
|
|
async def list(self) -> list[dict]:
|
|
rows = await self._list_raw()
|
|
rows.sort(key=lambda r: r["created_at"] or 0, reverse=True)
|
|
return rows
|
|
|
|
async def attach(self, sid: str, cols: int, rows: int) -> Attachment | None:
|
|
if await self.get_info(sid) is None:
|
|
return None
|
|
return await _TmuxAttachment.start(
|
|
self._bin, self._sock, f"{_PREFIX}{sid}", cols, rows
|
|
)
|
|
|
|
async def close(self, sid: str) -> bool:
|
|
rc, _, _ = await self._raw("kill-session", "-t", f"{_PREFIX}{sid}")
|
|
return rc == 0
|
|
|
|
async def close_all(self) -> None:
|
|
# No-op: tmux sessions are meant to SURVIVE a castle-api shutdown.
|
|
return None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Backend selection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_manager() -> MemoryManager | TmuxManager:
|
|
choice = os.environ.get("CASTLE_API_AGENT_BACKEND", "auto").lower()
|
|
tmux_bin = shutil.which("tmux")
|
|
if choice == "memory":
|
|
return MemoryManager()
|
|
if choice == "tmux" and not tmux_bin:
|
|
raise RuntimeError("CASTLE_API_AGENT_BACKEND=tmux but tmux is not on PATH")
|
|
if tmux_bin and choice in ("auto", "tmux"):
|
|
return TmuxManager(tmux_bin)
|
|
return MemoryManager()
|
|
|
|
|
|
manager: MemoryManager | TmuxManager = _make_manager()
|