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:
@@ -4,13 +4,14 @@ import { Terminal } from "@xterm/xterm"
|
||||
import { FitAddon } from "@xterm/addon-fit"
|
||||
import "@xterm/xterm/css/xterm.css"
|
||||
import { apiClient } from "@/services/api/client"
|
||||
import { cn } from "@/lib/utils"
|
||||
|
||||
type Status = "connecting" | "connected" | "exited" | "closed" | "error"
|
||||
|
||||
interface AgentTerminalProps {
|
||||
// Agent name to launch, or "terminal" for the plain login shell.
|
||||
agent: string
|
||||
// When set, resume this existing session instead of creating a new one.
|
||||
// When set, resume this existing live session (replaying its scrollback).
|
||||
resumeId?: string
|
||||
// Fired once the backend confirms the session id (new or resumed).
|
||||
onSession?: (id: string) => void
|
||||
@@ -20,6 +21,13 @@ interface AgentTerminalProps {
|
||||
mode?: "continue"
|
||||
// Resume one of the agent's OWN past sessions by its id (agent-native).
|
||||
resumeSessionId?: string
|
||||
// Bump this whenever the terminal's container is resized by an ancestor (e.g.
|
||||
// the dock expanding). Triggers a refit that ResizeObserver alone can miss
|
||||
// when the size change is driven by a class swap rather than layout reflow.
|
||||
fitSignal?: number
|
||||
// Strip chrome (header, border, padding) on narrow screens — used when the
|
||||
// dock is maximized so the terminal goes edge-to-edge on a phone.
|
||||
compact?: boolean
|
||||
}
|
||||
|
||||
export function AgentTerminal({
|
||||
@@ -29,12 +37,32 @@ export function AgentTerminal({
|
||||
fill,
|
||||
mode,
|
||||
resumeSessionId,
|
||||
fitSignal,
|
||||
compact,
|
||||
}: AgentTerminalProps) {
|
||||
const containerRef = useRef<HTMLDivElement>(null)
|
||||
const termRef = useRef<Terminal | null>(null)
|
||||
const fitRef = useRef<FitAddon | null>(null)
|
||||
const wsRef = useRef<WebSocket | null>(null)
|
||||
const [status, setStatus] = useState<Status>("connecting")
|
||||
const [detail, setDetail] = useState<string>("")
|
||||
const [maximized, setMaximized] = useState(false)
|
||||
|
||||
// Fit to the container, then tell the backend pty the new size.
|
||||
const refit = () => {
|
||||
const term = termRef.current
|
||||
const ws = wsRef.current
|
||||
if (!term) return
|
||||
try {
|
||||
fitRef.current?.fit()
|
||||
} catch {
|
||||
/* container not laid out yet */
|
||||
}
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
ws.send(JSON.stringify({ type: "resize", cols: term.cols, rows: term.rows }))
|
||||
}
|
||||
}
|
||||
|
||||
useEffect(() => {
|
||||
const term = new Terminal({
|
||||
convertEol: false,
|
||||
@@ -46,6 +74,8 @@ export function AgentTerminal({
|
||||
const fit = new FitAddon()
|
||||
term.loadAddon(fit)
|
||||
term.open(containerRef.current!)
|
||||
termRef.current = term
|
||||
fitRef.current = fit
|
||||
try {
|
||||
fit.fit()
|
||||
} catch {
|
||||
@@ -61,17 +91,12 @@ export function AgentTerminal({
|
||||
: `/agents/${agent}/session`
|
||||
const ws = new WebSocket(apiClient.wsUrl(path))
|
||||
ws.binaryType = "arraybuffer"
|
||||
wsRef.current = ws
|
||||
const enc = new TextEncoder()
|
||||
|
||||
const sendResize = () => {
|
||||
if (ws.readyState === WebSocket.OPEN) {
|
||||
ws.send(JSON.stringify({ type: "resize", cols: term.cols, rows: term.rows }))
|
||||
}
|
||||
}
|
||||
|
||||
ws.onopen = () => {
|
||||
setStatus("connected")
|
||||
sendResize()
|
||||
refit()
|
||||
// On resume, full-screen TUIs (opencode, claude) don't know a new client
|
||||
// attached and won't repaint. A size change forces SIGWINCH → a full
|
||||
// redraw of the current frame, so the prior session becomes visible.
|
||||
@@ -81,7 +106,7 @@ export function AgentTerminal({
|
||||
ws.send(
|
||||
JSON.stringify({ type: "resize", cols: Math.max(1, term.cols - 1), rows: term.rows }),
|
||||
)
|
||||
setTimeout(sendResize, 60)
|
||||
setTimeout(refit, 60)
|
||||
}, 80)
|
||||
}
|
||||
term.focus()
|
||||
@@ -116,14 +141,7 @@ export function AgentTerminal({
|
||||
if (ws.readyState === WebSocket.OPEN) ws.send(enc.encode(d))
|
||||
})
|
||||
|
||||
const ro = new ResizeObserver(() => {
|
||||
try {
|
||||
fit.fit()
|
||||
} catch {
|
||||
/* ignore */
|
||||
}
|
||||
sendResize()
|
||||
})
|
||||
const ro = new ResizeObserver(() => refit())
|
||||
if (containerRef.current) ro.observe(containerRef.current)
|
||||
|
||||
return () => {
|
||||
@@ -131,10 +149,27 @@ export function AgentTerminal({
|
||||
dataDisp.dispose()
|
||||
ws.close()
|
||||
term.dispose()
|
||||
termRef.current = null
|
||||
fitRef.current = null
|
||||
wsRef.current = null
|
||||
}
|
||||
// Remounting (new agent / resume target / mode) tears this down and restarts.
|
||||
}, [agent, resumeId, mode, resumeSessionId, onSession])
|
||||
|
||||
// An ancestor changed our size (dock expanded/restored/reopened). ResizeObserver
|
||||
// usually catches it, but a class-swap + CSS can settle a frame late, so refit
|
||||
// across a few frames until the new size stabilizes.
|
||||
useEffect(() => {
|
||||
let raf = 0
|
||||
const start = performance.now()
|
||||
const tick = () => {
|
||||
refit()
|
||||
if (performance.now() - start < 400) raf = requestAnimationFrame(tick)
|
||||
}
|
||||
raf = requestAnimationFrame(tick)
|
||||
return () => cancelAnimationFrame(raf)
|
||||
}, [fitSignal])
|
||||
|
||||
const statusColor =
|
||||
status === "connected"
|
||||
? "text-green-400"
|
||||
@@ -149,15 +184,21 @@ export function AgentTerminal({
|
||||
|
||||
return (
|
||||
<div
|
||||
className={
|
||||
className={cn(
|
||||
maximized
|
||||
? "fixed inset-0 z-50 bg-[var(--background)] flex flex-col"
|
||||
: fill
|
||||
? "h-full flex flex-col bg-black/40 border border-[var(--border)] rounded-lg overflow-hidden"
|
||||
: "bg-black/40 border border-[var(--border)] rounded-lg overflow-hidden"
|
||||
}
|
||||
: "bg-black/40 border border-[var(--border)] rounded-lg overflow-hidden",
|
||||
compact && "max-sm:border-0 max-sm:rounded-none",
|
||||
)}
|
||||
>
|
||||
<div className="flex items-center justify-between px-3 py-1.5 border-b border-[var(--border)] text-xs text-[var(--muted)]">
|
||||
<div
|
||||
className={cn(
|
||||
"flex items-center justify-between px-3 py-1.5 border-b border-[var(--border)] text-xs text-[var(--muted)]",
|
||||
compact && "max-sm:hidden",
|
||||
)}
|
||||
>
|
||||
<span className="flex items-center gap-2">
|
||||
<span>terminal: {agent}</span>
|
||||
<span className={statusColor}>
|
||||
@@ -180,7 +221,10 @@ export function AgentTerminal({
|
||||
</div>
|
||||
<div
|
||||
ref={containerRef}
|
||||
className={maximized || fill ? "flex-1 min-h-0 p-2" : "h-[560px] p-2"}
|
||||
className={cn(
|
||||
maximized || fill ? "flex-1 min-h-0 p-2" : "h-[560px] p-2",
|
||||
compact && "max-sm:p-0",
|
||||
)}
|
||||
/>
|
||||
</div>
|
||||
)
|
||||
|
||||
@@ -99,12 +99,15 @@ export function AssistantDock() {
|
||||
terminal's WebSocket survives minimize. */}
|
||||
<div
|
||||
className={cn(
|
||||
"fixed z-40 flex flex-col",
|
||||
"fixed z-40 flex flex-col bg-[var(--background)] shadow-2xl",
|
||||
expanded
|
||||
? "inset-4"
|
||||
: "bottom-4 right-4 w-[760px] max-w-[calc(100vw-2rem)] h-[600px] max-h-[calc(100vh-2rem)]",
|
||||
"bg-[var(--background)] border border-[var(--border)] rounded-xl shadow-2xl",
|
||||
"transition-all duration-200",
|
||||
? // Maximized: full-bleed on narrow screens (no margin/border/radius),
|
||||
// ins-4 with chrome on sm+.
|
||||
"inset-0 rounded-none border-0 sm:inset-4 sm:rounded-xl sm:border sm:border-[var(--border)]"
|
||||
: "bottom-4 right-4 w-[760px] max-w-[calc(100vw-2rem)] h-[600px] max-h-[calc(100vh-2rem)] rounded-xl border border-[var(--border)]",
|
||||
// Animate only open/close (opacity + translate); size changes (expand)
|
||||
// are instant so the terminal's fit measures the final box immediately.
|
||||
"transition-[opacity,transform] duration-200",
|
||||
open
|
||||
? "opacity-100 translate-y-0"
|
||||
: "opacity-0 translate-y-[calc(100%+2rem)] pointer-events-none",
|
||||
@@ -114,7 +117,14 @@ export function AssistantDock() {
|
||||
{/* Header */}
|
||||
<div className="flex items-center gap-2 px-3 py-2 border-b border-[var(--border)]">
|
||||
<Bot size={16} className="text-[var(--primary)] shrink-0" />
|
||||
<div className="flex flex-wrap items-center gap-1.5 min-w-0">
|
||||
<div
|
||||
className={cn(
|
||||
"flex flex-wrap items-center gap-1.5 min-w-0",
|
||||
// Maximized on a narrow screen: drop the launcher pills for space
|
||||
// (shrink to switch agents). Visible again at sm+.
|
||||
expanded && "max-sm:hidden",
|
||||
)}
|
||||
>
|
||||
<button
|
||||
onClick={() => launch(SHELL)}
|
||||
className={cn(
|
||||
@@ -190,7 +200,7 @@ export function AssistantDock() {
|
||||
</div>
|
||||
|
||||
{/* Body */}
|
||||
<div className="relative flex-1 min-h-0 p-2">
|
||||
<div className={cn("relative flex-1 min-h-0 p-2", expanded && "max-sm:p-0")}>
|
||||
{active ? (
|
||||
<AgentTerminal
|
||||
key={active.key}
|
||||
@@ -200,6 +210,8 @@ export function AssistantDock() {
|
||||
resumeSessionId={active.resumeSessionId}
|
||||
onSession={onSession}
|
||||
fill
|
||||
compact={expanded}
|
||||
fitSignal={(open ? 1 : 0) + (expanded ? 2 : 0)}
|
||||
/>
|
||||
) : (
|
||||
<div className="h-full flex items-center justify-center text-sm text-[var(--muted)]">
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
"""Agent terminal UX — list agents, run/resume them in a pty over WebSocket.
|
||||
"""Agent terminal UX — list agents, run/resume them in a terminal over WebSocket.
|
||||
|
||||
Castle is assistant-agnostic: it launches a configured command (or a shell) in a
|
||||
real tty and streams raw bytes to an xterm.js terminal in the browser. Sessions
|
||||
outlive their WebSocket connection, so they can be listed, resumed, and killed.
|
||||
outlive their WebSocket connection (and, with the tmux backend, a castle-api
|
||||
restart), so they can be listed, resumed, and killed. The session backend lives
|
||||
in `agent_sessions.py`; this module is the HTTP/WebSocket surface.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -24,7 +26,7 @@ from castle_api.agent_registry import (
|
||||
resolve_agent,
|
||||
resume_argv,
|
||||
)
|
||||
from castle_api.agent_sessions import AgentSession, manager
|
||||
from castle_api.agent_sessions import manager
|
||||
from castle_api.config import get_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -75,9 +77,9 @@ def get_agents() -> list[dict]:
|
||||
|
||||
|
||||
@router.get("/sessions")
|
||||
def get_sessions() -> list[dict]:
|
||||
"""List live terminal sessions (running or recently exited)."""
|
||||
return manager.list()
|
||||
async def get_sessions() -> list[dict]:
|
||||
"""List live terminal sessions (running, backend-managed)."""
|
||||
return await manager.list()
|
||||
|
||||
|
||||
@router.get("/history")
|
||||
@@ -97,7 +99,7 @@ async def get_history() -> list[dict]:
|
||||
|
||||
@router.delete("/sessions/{session_id}")
|
||||
async def delete_session(session_id: str) -> dict:
|
||||
"""Kill a session's process group and drop it from the registry."""
|
||||
"""Kill a session and drop it from the backend."""
|
||||
ok = await manager.close(session_id)
|
||||
if not ok:
|
||||
raise HTTPException(
|
||||
@@ -113,121 +115,118 @@ def _login_shell() -> str:
|
||||
return os.environ.get("SHELL") or "/bin/bash"
|
||||
|
||||
|
||||
async def _new_session(
|
||||
async def _build_launch(
|
||||
name: str, cont: bool = False, resume_session: str | None = None
|
||||
) -> AgentSession:
|
||||
"""Create a session for a named agent, or the reserved login shell.
|
||||
) -> tuple[str, list[str], str, dict[str, str]]:
|
||||
"""Return (agent_label, argv, cwd, env) for a new session.
|
||||
|
||||
Everything runs under the user's *interactive login* shell (`-lic`) so it
|
||||
inherits exactly the environment a real terminal would — the systemd service
|
||||
env is stripped down and would otherwise miss vars set in shell rc files
|
||||
(API keys, etc.). Agents are `exec`'d so the pty's foreground process is the
|
||||
agent itself (clean signals + exit).
|
||||
|
||||
``cont`` appends the agent's ``resume_args`` (open its own picker / continue).
|
||||
``resume_session`` builds the agent's resume-by-id argv for a specific past
|
||||
session. Either way castle only passes declared flags through — it never
|
||||
reads the agent's session storage.
|
||||
env is stripped down and would otherwise miss vars from shell rc files (API
|
||||
keys, etc.). Agents are `exec`'d so the terminal's foreground process is the
|
||||
agent itself. `cont` appends the agent's `resume_args`; `resume_session`
|
||||
builds its resume-by-id argv — castle only passes declared flags through.
|
||||
"""
|
||||
shell = _login_shell()
|
||||
if name == TERMINAL_AGENT:
|
||||
return await manager.create(
|
||||
TERMINAL_AGENT, shell, args=["-l"], cwd=default_cwd()
|
||||
)
|
||||
return TERMINAL_AGENT, [shell, "-l"], default_cwd(), {}
|
||||
agent = resolve_agent(name)
|
||||
if agent is None:
|
||||
raise LookupError(f"unknown agent: {name}")
|
||||
if not agent.available:
|
||||
raise LookupError(f"agent not installed: {name}")
|
||||
if resume_session:
|
||||
argv = resume_argv(agent, resume_session)
|
||||
if argv is None:
|
||||
inner = resume_argv(agent, resume_session)
|
||||
if inner is None:
|
||||
raise LookupError(f"agent cannot resume by id: {name}")
|
||||
else:
|
||||
argv = [agent.resolved_command or agent.command, *agent.args]
|
||||
inner = [agent.resolved_command or agent.command, *agent.args]
|
||||
if cont and agent.resume_args:
|
||||
argv += agent.resume_args
|
||||
launch = shlex.join(argv)
|
||||
return await manager.create(
|
||||
name,
|
||||
shell,
|
||||
args=["-lic", f"exec {launch}"],
|
||||
cwd=agent.cwd,
|
||||
env=agent.env, # explicit per-agent overrides; the login shell sets PATH
|
||||
)
|
||||
inner += agent.resume_args
|
||||
argv = [shell, "-lic", f"exec {shlex.join(inner)}"]
|
||||
return name, argv, agent.cwd, agent.env
|
||||
|
||||
|
||||
@router.websocket("/{name}/session")
|
||||
async def agent_session(ws: WebSocket, name: str) -> None:
|
||||
"""Run (or resume) an agent in a pty; stream it to the browser terminal.
|
||||
"""Run (or resume) an agent in a terminal; stream it to the browser.
|
||||
|
||||
Query param `session=<id>` resumes an existing live session (replaying its
|
||||
scrollback); otherwise a new session is created for `name` (or the reserved
|
||||
`terminal` shell). Disconnecting detaches — it does NOT kill the session.
|
||||
Query params (mutually exclusive, all optional):
|
||||
- `session=<id>` resume an existing live session.
|
||||
- `resume_session=<id>` start the agent resumed to one of *its own* past
|
||||
sessions (agent-native, by id).
|
||||
- `mode=continue` start the agent with its `resume_args`.
|
||||
Disconnecting detaches — it does NOT kill the session.
|
||||
"""
|
||||
if not _origin_allowed(ws.headers.get("origin")):
|
||||
await ws.close(code=1008) # policy violation
|
||||
return
|
||||
|
||||
resume_id = ws.query_params.get("session")
|
||||
session_id: str
|
||||
is_resume = False
|
||||
agent_label = name
|
||||
|
||||
# Resolve the target session (before accept only for hard rejects).
|
||||
session: AgentSession | None = None
|
||||
if resume_id:
|
||||
session = manager.get(resume_id)
|
||||
if session is None or not session.running:
|
||||
await ws.accept()
|
||||
await ws.send_json({"type": "error", "error": "session not resumable"})
|
||||
await ws.close()
|
||||
return
|
||||
else:
|
||||
cont = ws.query_params.get("mode") == "continue"
|
||||
resume_session = ws.query_params.get("resume_session")
|
||||
try:
|
||||
session = await _new_session(name, cont=cont, resume_session=resume_session)
|
||||
except LookupError as e:
|
||||
await ws.accept()
|
||||
await ws.send_json({"type": "error", "error": str(e)})
|
||||
await ws.close()
|
||||
return
|
||||
except Exception:
|
||||
logger.exception("failed to launch agent %s", name)
|
||||
await ws.accept()
|
||||
await ws.send_json({"type": "error", "error": "failed to launch"})
|
||||
await ws.close()
|
||||
return
|
||||
try:
|
||||
if resume_id:
|
||||
info = await manager.get_info(resume_id)
|
||||
if not info or not info.get("running"):
|
||||
await ws.accept()
|
||||
await ws.send_json({"type": "error", "error": "session not resumable"})
|
||||
await ws.close()
|
||||
return
|
||||
session_id = resume_id
|
||||
is_resume = True
|
||||
agent_label = info.get("agent") or name
|
||||
else:
|
||||
cont = ws.query_params.get("mode") == "continue"
|
||||
resume_session = ws.query_params.get("resume_session")
|
||||
agent_label, argv, cwd, env = await _build_launch(
|
||||
name, cont, resume_session
|
||||
)
|
||||
session_id = await manager.create(
|
||||
agent_label, argv, cwd, env, cols=80, rows=24
|
||||
)
|
||||
except LookupError as e:
|
||||
await ws.accept()
|
||||
await ws.send_json({"type": "error", "error": str(e)})
|
||||
await ws.close()
|
||||
return
|
||||
except Exception:
|
||||
logger.exception("failed to launch agent %s", name)
|
||||
await ws.accept()
|
||||
await ws.send_json({"type": "error", "error": "failed to launch"})
|
||||
await ws.close()
|
||||
return
|
||||
|
||||
assert session is not None # narrowed by the branches above
|
||||
await ws.accept()
|
||||
|
||||
# Subscribe + snapshot scrollback atomically (no await between → the output
|
||||
# callback, which runs in this same event loop, cannot interleave), so a
|
||||
# resumed client gets the buffer once with no duplication or gap.
|
||||
q = session.subscribe()
|
||||
snapshot = bytes(session.scrollback)
|
||||
att = await manager.attach(session_id, cols=80, rows=24)
|
||||
if att is None:
|
||||
await ws.send_json({"type": "error", "error": "session not attachable"})
|
||||
await ws.close()
|
||||
return
|
||||
|
||||
await ws.send_json(
|
||||
{
|
||||
"type": "session",
|
||||
"id": session.id,
|
||||
"agent": session.agent,
|
||||
"resumed": bool(resume_id),
|
||||
"id": session_id,
|
||||
"agent": agent_label,
|
||||
"resumed": is_resume,
|
||||
}
|
||||
)
|
||||
if snapshot:
|
||||
await ws.send_bytes(snapshot)
|
||||
|
||||
async def pump() -> None:
|
||||
try:
|
||||
while True:
|
||||
chunk = await q.get()
|
||||
if chunk is None: # child exited
|
||||
if ws.application_state == WebSocketState.CONNECTED:
|
||||
await ws.send_json({"type": "exit", "code": session.exit_code})
|
||||
break
|
||||
async for chunk in att.reader():
|
||||
if ws.application_state != WebSocketState.CONNECTED:
|
||||
break
|
||||
await ws.send_bytes(chunk)
|
||||
# reader ended without cancellation → the session ended
|
||||
if ws.application_state == WebSocketState.CONNECTED:
|
||||
info = await manager.get_info(session_id)
|
||||
code = info.get("exit_code") if info else None
|
||||
await ws.send_json({"type": "exit", "code": code})
|
||||
await ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -239,22 +238,20 @@ async def agent_session(ws: WebSocket, name: str) -> None:
|
||||
break
|
||||
data = msg.get("bytes")
|
||||
text = msg.get("text")
|
||||
if data is not None and session.pty is not None:
|
||||
session.pty.write(data)
|
||||
elif text is not None and session.pty is not None:
|
||||
if data is not None:
|
||||
att.write(data)
|
||||
elif text is not None:
|
||||
try:
|
||||
ctrl = json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if ctrl.get("type") == "resize":
|
||||
session.pty.resize(int(ctrl["cols"]), int(ctrl["rows"]))
|
||||
att.resize(int(ctrl["cols"]), int(ctrl["rows"]))
|
||||
elif ctrl.get("type") == "input":
|
||||
session.pty.write(str(ctrl.get("data", "")).encode())
|
||||
att.write(str(ctrl.get("data", "")).encode())
|
||||
except WebSocketDisconnect:
|
||||
pass
|
||||
finally:
|
||||
pump_task.cancel()
|
||||
session.unsubscribe(q)
|
||||
# Intentionally do NOT close the session here — it lives on for resume.
|
||||
# Sessions are ended via DELETE /agents/sessions/{id}, when the child
|
||||
# exits, or by idle reaping.
|
||||
await att.aclose()
|
||||
# Intentionally do NOT close the session — it lives on for resume.
|
||||
|
||||
Reference in New Issue
Block a user