feat(agents): durable tmux-backed terminal sessions + fit/chrome fixes
- 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.
This commit is contained in:
@@ -1,49 +1,75 @@
|
||||
"""Live agent/terminal session registry.
|
||||
|
||||
A session is a running PTY (an agent CLI or a shell) whose lifetime is decoupled
|
||||
from any WebSocket connection: closing the browser tab detaches, it does not
|
||||
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. Each session keeps a bounded scrollback buffer so a
|
||||
re-attaching client can repaint the terminal.
|
||||
explicitly terminated.
|
||||
|
||||
This is deliberately in-memory and single-node: sessions do not survive a
|
||||
castle-api restart. That's the right scope for a personal dashboard — for true
|
||||
cross-restart persistence you'd launch the agent under tmux, which breaks the
|
||||
"castle just runs a command" agnosticism.
|
||||
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 dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Protocol
|
||||
|
||||
from castle_core.config import CASTLE_HOME
|
||||
|
||||
from castle_api.pty_session import PtySession
|
||||
|
||||
# Cap the per-session replay buffer. Raw terminal bytes; enough to repaint a
|
||||
# full-screen TUI and some history without unbounded growth.
|
||||
_SCROLLBACK_BYTES = 256 * 1024
|
||||
# Reap sessions that have exited and sat idle longer than this (seconds).
|
||||
_EXITED_TTL = 3600
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SCROLLBACK_BYTES = 256 * 1024 # memory backend replay cap
|
||||
_EXITED_TTL = 3600 # memory backend: reap exited sessions after this
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentSession:
|
||||
"""One live PTY plus its scrollback and attached subscribers."""
|
||||
class Attachment(Protocol):
|
||||
"""One WebSocket's live view of a session."""
|
||||
|
||||
id: str
|
||||
agent: str # agent name, or "terminal" for the plain shell
|
||||
command: str
|
||||
cwd: str
|
||||
created_at: float
|
||||
pty: PtySession | None = None
|
||||
scrollback: bytearray = field(default_factory=bytearray)
|
||||
subscribers: set[asyncio.Queue[bytes | None]] = field(default_factory=set)
|
||||
exited: bool = False
|
||||
exit_code: int | None = None
|
||||
exited_at: float | None = None
|
||||
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)
|
||||
@@ -53,7 +79,6 @@ class AgentSession:
|
||||
try:
|
||||
q.put_nowait(data)
|
||||
except asyncio.QueueFull:
|
||||
# Slow client: drop it; the WS side will notice on its next send.
|
||||
self.subscribers.discard(q)
|
||||
|
||||
def _on_exit(self) -> None:
|
||||
@@ -61,15 +86,7 @@ class AgentSession:
|
||||
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) # EOF sentinel
|
||||
|
||||
def subscribe(self) -> asyncio.Queue[bytes | None]:
|
||||
q: asyncio.Queue[bytes | None] = asyncio.Queue(maxsize=256)
|
||||
self.subscribers.add(q)
|
||||
return q
|
||||
|
||||
def unsubscribe(self, q: asyncio.Queue[bytes | None]) -> None:
|
||||
self.subscribers.discard(q)
|
||||
q.put_nowait(None)
|
||||
|
||||
@property
|
||||
def running(self) -> bool:
|
||||
@@ -91,58 +108,78 @@ class AgentSession:
|
||||
}
|
||||
|
||||
|
||||
class SessionManager:
|
||||
"""In-memory registry of live agent sessions."""
|
||||
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, AgentSession] = {}
|
||||
self._sessions: dict[str, _MemSession] = {}
|
||||
|
||||
def _reap_exited(self) -> None:
|
||||
def _reap(self) -> None:
|
||||
now = time.time()
|
||||
stale = [
|
||||
sid
|
||||
for sid, s in self._sessions.items()
|
||||
if s.exited and s.exited_at and (now - s.exited_at) > _EXITED_TTL
|
||||
]
|
||||
for sid in stale:
|
||||
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,
|
||||
command: str,
|
||||
args: list[str] | None = None,
|
||||
cwd: str | Path | None = None,
|
||||
env: dict[str, str] | None = None,
|
||||
cols: int = 80,
|
||||
rows: int = 24,
|
||||
) -> AgentSession:
|
||||
self._reap_exited()
|
||||
session = AgentSession(
|
||||
id=uuid.uuid4().hex[:12],
|
||||
agent=agent,
|
||||
command=command,
|
||||
cwd=str(cwd) if cwd else "",
|
||||
created_at=time.time(),
|
||||
)
|
||||
session.pty = await PtySession.start(
|
||||
command,
|
||||
args=args,
|
||||
cwd=cwd,
|
||||
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=session._on_output,
|
||||
on_exit=session._on_exit,
|
||||
on_output=s._on_output,
|
||||
on_exit=s._on_exit,
|
||||
)
|
||||
self._sessions[session.id] = session
|
||||
return session
|
||||
self._sessions[sid] = s
|
||||
return sid
|
||||
|
||||
def get(self, session_id: str) -> AgentSession | None:
|
||||
return self._sessions.get(session_id)
|
||||
async def get_info(self, sid: str) -> dict | None:
|
||||
s = self._sessions.get(sid)
|
||||
return s.info() if s else None
|
||||
|
||||
def list(self) -> list[dict]:
|
||||
self._reap_exited()
|
||||
async def list(self) -> list[dict]:
|
||||
self._reap()
|
||||
return [
|
||||
s.info()
|
||||
for s in sorted(
|
||||
@@ -150,13 +187,25 @@ class SessionManager:
|
||||
)
|
||||
]
|
||||
|
||||
async def close(self, session_id: str) -> bool:
|
||||
session = self._sessions.pop(session_id, None)
|
||||
if session is None:
|
||||
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 session.pty is not None:
|
||||
await session.pty.close()
|
||||
for q in list(session.subscribers):
|
||||
if s.pty:
|
||||
await s.pty.close()
|
||||
for q in list(s.subscribers):
|
||||
q.put_nowait(None)
|
||||
return True
|
||||
|
||||
@@ -165,5 +214,235 @@ class SessionManager:
|
||||
await self.close(sid)
|
||||
|
||||
|
||||
# Module-level singleton (mirrors mesh_state / stream subscribers).
|
||||
manager = SessionManager()
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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()
|
||||
|
||||
Reference in New Issue
Block a user