365 lines
16 KiB
Python
365 lines
16 KiB
Python
"""Keep-alive PTY sessions for dashboard terminals.
|
|
|
|
A PTY process outlives the WebSocket that created it: a single drain task always reads the PTY into
|
|
a bounded RingBuffer and forwards to the attached socket when present. Reconnecting with the same
|
|
opaque token replays the buffer and resumes live.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Callable, Dict, Optional, Tuple
|
|
|
|
WS_CLOSE_PROCESS_EXITED = 4410
|
|
WS_CLOSE_SUPERSEDED = 4409
|
|
TUI_FORCE_REDRAW = b"\x0c"
|
|
|
|
|
|
class RingBuffer:
|
|
"""Keeps only the most recent ``capacity`` bytes appended to it."""
|
|
|
|
def __init__(self, capacity: int) -> None:
|
|
self._cap = capacity
|
|
self._buf = bytearray()
|
|
self.truncated = False
|
|
|
|
def append(self, data: bytes) -> None:
|
|
self._buf.extend(data)
|
|
overflow = len(self._buf) - self._cap
|
|
if overflow > 0:
|
|
del self._buf[:overflow]
|
|
self.truncated = True
|
|
|
|
def snapshot(self) -> bytes:
|
|
return bytes(self._buf)
|
|
|
|
|
|
async def _close_ws(ws, code: int) -> None:
|
|
try:
|
|
if ws is not None:
|
|
await ws.close(code=code)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _process_ancestors(pid: int) -> set[int]:
|
|
"""Pids of ``pid``'s live ancestors; empty when the process is gone or unreadable."""
|
|
import psutil # type: ignore
|
|
|
|
try:
|
|
return {parent.pid for parent in psutil.Process(pid).parents()}
|
|
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess, ValueError):
|
|
return set()
|
|
|
|
|
|
def _key_segments(key: str) -> tuple[str, str]:
|
|
"""``(profile, resume)`` a keep-alive key was registered under.
|
|
|
|
Registry keys read ``token\\0profile\\0resume``; a key without those
|
|
segments — a chat that was never resumed from — yields ``("", "")``.
|
|
Session ids are only unique within a profile's store, so a resume target
|
|
is identified by both.
|
|
"""
|
|
parts = key.split("\0")
|
|
return (parts[1], parts[2]) if len(parts) >= 3 else ("", "")
|
|
|
|
|
|
class PtySession:
|
|
def __init__(self, key: str, bridge, *, buffer_cap: int, read_timeout: float, active_session_file: Optional[Path] = None) -> None:
|
|
self.key = key
|
|
self.bridge = bridge
|
|
self.active_session_file = active_session_file
|
|
self.active_session_cleanup: Optional[Callable[[], None]] = None
|
|
self.buffer = RingBuffer(buffer_cap)
|
|
self.alive = True
|
|
self.attached = False
|
|
self.last_detached_at: Optional[float] = None
|
|
self._read_timeout = read_timeout
|
|
self._ws = None
|
|
self._attach_generation = 0
|
|
self._drain_task: Optional[asyncio.Task] = None
|
|
self._write_lock = asyncio.Lock()
|
|
|
|
async def start(self) -> None:
|
|
self._drain_task = asyncio.create_task(self._drain())
|
|
|
|
async def _drain(self) -> None:
|
|
loop = asyncio.get_running_loop()
|
|
while True:
|
|
chunk = await loop.run_in_executor(None, self.bridge.read, self._read_timeout)
|
|
if chunk is None: # EOF — the agent process exited
|
|
self.alive = False
|
|
await _close_ws(self._ws, WS_CLOSE_PROCESS_EXITED)
|
|
return
|
|
if not chunk: # idle tick
|
|
await asyncio.sleep(0)
|
|
continue
|
|
self.buffer.append(chunk)
|
|
ws = self._ws
|
|
try:
|
|
if ws is not None:
|
|
await ws.send_bytes(chunk)
|
|
except Exception:
|
|
# The viewer is gone; nothing else observes this failure (the handler's finally
|
|
# only runs once ws.receive() sees the disconnect). detach() is a no-op when a
|
|
# replacement socket attached during the send, so the new viewer keeps its session.
|
|
self.detach(ws)
|
|
|
|
async def write(self, ws, data: bytes) -> bool:
|
|
"""Serialize input and discard bytes from a superseded socket."""
|
|
async with self._write_lock:
|
|
if self._ws is not ws:
|
|
return True
|
|
generation = self._attach_generation
|
|
delivered = await self.bridge.write(data)
|
|
# A replacement socket can attach while the bridge write is
|
|
# suspended on backpressure. A late failure from the superseded
|
|
# socket must not poison the replacement's shared PTY session.
|
|
if (
|
|
not delivered
|
|
and self._ws is ws
|
|
and self._attach_generation == generation
|
|
):
|
|
self.alive = False
|
|
return delivered
|
|
|
|
async def attach(self, ws, *, force_redraw: bool = False) -> bool:
|
|
"""Attach a browser terminal and replay buffered PTY output.
|
|
|
|
The TUI renders differentially on an alternate screen, so a bounded ANSI tail is not a
|
|
self-contained frame; ``force_redraw`` asks the live TUI for one full redraw after replay.
|
|
"""
|
|
if self._ws is not ws:
|
|
await _close_ws(self._ws, WS_CLOSE_SUPERSEDED)
|
|
self._ws = ws
|
|
self._attach_generation += 1
|
|
self.attached = True
|
|
self.last_detached_at = None
|
|
if snap := self.buffer.snapshot():
|
|
try:
|
|
await ws.send_bytes(snap)
|
|
except Exception:
|
|
# Client dropped mid-replay; the caller never reaches its writer loop, so undo the
|
|
# attach here or reap_idle() can never reclaim this PTY (#110849).
|
|
self.detach(ws)
|
|
return False
|
|
if force_redraw:
|
|
return await self.write(ws, TUI_FORCE_REDRAW)
|
|
return True
|
|
|
|
def detach(self, ws) -> None:
|
|
# Only the currently-attached socket may mark the session detached: a superseded socket's
|
|
# handler also calls detach on its way out (after the new tab attached), and flipping
|
|
# ``attached`` then would make a session with a live viewer look idle and reapable.
|
|
if self._ws is not ws:
|
|
return
|
|
self._ws = None
|
|
self.attached = False
|
|
self.last_detached_at = time.monotonic()
|
|
|
|
def hosts_pid(self, pid: Optional[int]) -> bool:
|
|
"""Whether ``pid`` runs inside this PTY: the leader itself or one of its descendants.
|
|
|
|
The dashboard's PTY leader is ``node entry.js``; the process that takes the
|
|
session lease is its ``tui_gateway`` child, so the lease pid is never the
|
|
bridge's own pid and only ancestry identifies the terminal. A bridge that
|
|
cannot report a pid (already reaped, or a stub) is simply not a match.
|
|
"""
|
|
if pid is None:
|
|
return False
|
|
try:
|
|
leader = int(self.bridge.pid)
|
|
except (AttributeError, TypeError, ValueError, OSError):
|
|
return False
|
|
return leader == pid or leader in _process_ancestors(pid)
|
|
|
|
async def close(self) -> None:
|
|
self.alive = False
|
|
if self._drain_task is not None:
|
|
self._drain_task.cancel()
|
|
try:
|
|
await self._drain_task
|
|
except (asyncio.CancelledError, Exception):
|
|
pass
|
|
try:
|
|
# bridge.close() joins the child — blocking; keep it off the event loop.
|
|
# See #53227.
|
|
await asyncio.to_thread(self.bridge.close)
|
|
except Exception: # health: allow BLE001 S110 -- teardown of an already-dead PTY must not mask the caller's error path
|
|
pass
|
|
try:
|
|
if self.active_session_file is not None:
|
|
self.active_session_file.unlink(missing_ok=True)
|
|
except OSError:
|
|
pass
|
|
if self.active_session_cleanup is not None:
|
|
try:
|
|
self.active_session_cleanup()
|
|
except Exception: # health: allow BLE001 S110 -- cleanup callback must not mask the close path; the PTY is dead either way
|
|
pass
|
|
self.active_session_cleanup = None
|
|
|
|
|
|
class RegistryFull(Exception):
|
|
"""Every keep-alive slot holds a PTY that some tab is still attached to."""
|
|
|
|
def __init__(self, message: str = "Too many chat terminals are open in other tabs; close one and try again.") -> None:
|
|
super().__init__(message)
|
|
|
|
|
|
async def run_reaper(registry: PtySessionRegistry, *, interval: float = 60.0) -> None:
|
|
"""Periodically reap idle/dead keep-alive sessions. Cancelled on shutdown."""
|
|
while True:
|
|
await asyncio.sleep(interval)
|
|
try:
|
|
await registry.reap_idle()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
class PtySessionRegistry:
|
|
def __init__(self, *, ttl: float, max_sessions: int, buffer_cap: int, read_timeout: float) -> None:
|
|
self._ttl = ttl
|
|
self._max = max_sessions
|
|
self._buffer_cap = buffer_cap
|
|
self._read_timeout = read_timeout
|
|
self._sessions: dict[str, PtySession] = {}
|
|
# The get-or-spawn decision spans awaits (reap_idle, the spawn thread,
|
|
# session.start), so two connections racing one attach token both saw
|
|
# "no session" and forked a PTY each: the token then mapped to whichever
|
|
# registered last while the other tab's live session fell out of the
|
|
# registry — never reaped, and a reattach landed on the wrong terminal
|
|
# (#115304). Serialize the decision so a token maps to one PTY.
|
|
# ponytail: one registry-wide lock, not per key — argv resolution is
|
|
# already serialized globally for the same reason, and a spawn only
|
|
# delays NEW chats. Per-key locks if spawn throughput ever matters.
|
|
self._attach_lock = asyncio.Lock()
|
|
# Sessions popped from the registry but still closing in the background; close_all()
|
|
# awaits them too, and holding the tasks keeps them from being garbage-collected.
|
|
self._background_closes: set[asyncio.Task] = set()
|
|
|
|
async def attach_or_spawn(self, key: str, *, spawn: Callable[[], object], active_session_file: Optional[Path] = None) -> tuple[PtySession, bool]:
|
|
await self.reap_idle()
|
|
async with self._attach_lock:
|
|
existing = self._sessions.get(key)
|
|
if existing is not None and existing.alive:
|
|
return existing, False
|
|
if existing is not None: # dead remnant
|
|
# Close in the background: ending a dead leader's helpers can take the helper
|
|
# grace, and this lock serializes every new chat.
|
|
self._sessions.pop(key, None)
|
|
self._close_in_background(existing)
|
|
if len(self._sessions) >= self._max:
|
|
self._reap_one_idle_or_raise()
|
|
# PTY spawn does blocking fork/exec work — keep it off the event loop.
|
|
# See #53227.
|
|
bridge = await asyncio.to_thread(spawn)
|
|
session = PtySession(
|
|
key,
|
|
bridge,
|
|
buffer_cap=self._buffer_cap,
|
|
read_timeout=self._read_timeout,
|
|
active_session_file=active_session_file,
|
|
)
|
|
await session.start()
|
|
self._sessions[key] = session
|
|
return session, True
|
|
|
|
async def close_other_sessions(self, prefix: str, *, keep_key: str) -> None:
|
|
"""Close sessions belonging to the same logical client except ``keep_key``.
|
|
|
|
Dashboard profile changes keep the browser's attach token but change the
|
|
canonical session key. The previous profile's detached PTY must not
|
|
remain alive long enough to hold the TUI session lease and reject a
|
|
later return to that chat.
|
|
"""
|
|
async with self._attach_lock:
|
|
keys = [
|
|
key for key in self._sessions
|
|
if key != keep_key and (key == prefix or key.startswith(prefix + "\0"))
|
|
]
|
|
for key in keys:
|
|
session = self._sessions.pop(key, None)
|
|
if session is not None:
|
|
# A sibling tab sharing the attach token may still be viewing this
|
|
# PTY: supersede it explicitly (4409) instead of leaving it silent
|
|
# until its next keystroke fails with 1013.
|
|
await _close_ws(session._ws, WS_CLOSE_SUPERSEDED)
|
|
await session.close()
|
|
|
|
async def close_orphaned_sessions(
|
|
self, resume: Optional[str], *, keep_key: str, holder_pid: Optional[int] = None,
|
|
) -> None:
|
|
"""Close a keep-alive PTY stranded under a superseded attach token.
|
|
|
|
Rotating the token — the dashboard's *New chat* — moves the tab to a
|
|
fresh PTY while the previous one keeps the TUI's single-writer lease on
|
|
its session. That terminal is out of the user's reach, so the next
|
|
return to the same chat is refused as a session held elsewhere. A PTY
|
|
with a live viewer is never a candidate: somebody is still using it.
|
|
|
|
A chat that was never resumed from carries no resume target in its key,
|
|
so ``holder_pid`` — the process holding the lease for the session being
|
|
resumed — identifies that terminal instead.
|
|
"""
|
|
if not resume:
|
|
return
|
|
# The requested chat is (profile, session): the same session id in
|
|
# another profile's store is a different chat whose terminal stays.
|
|
target = (_key_segments(keep_key)[0], resume)
|
|
async with self._attach_lock:
|
|
doomed = [
|
|
key for key, session in self._sessions.items()
|
|
if key != keep_key and not session.attached
|
|
and (_key_segments(key) == target or session.hosts_pid(holder_pid))
|
|
]
|
|
sessions = [self._sessions.pop(key) for key in doomed]
|
|
# Close outside the registry lock — a close can wait out its helpers' SIGHUP grace and
|
|
# this lock serializes every new chat — but still before the caller spawns: the child's
|
|
# session lease is only released once its close finishes.
|
|
for session in sessions:
|
|
await session.close()
|
|
|
|
def detach(self, key: str, ws) -> None:
|
|
s = self._sessions.get(key)
|
|
if s is not None:
|
|
s.detach(ws)
|
|
|
|
async def reap_idle(self, now: Optional[float] = None) -> None:
|
|
now = time.monotonic() if now is None else now
|
|
doomed = [
|
|
key for key, s in self._sessions.items()
|
|
if not s.alive
|
|
or (not s.attached and s.last_detached_at is not None and (now - s.last_detached_at) > self._ttl)
|
|
# EOF never arrives if a helper still holds the PTY slave after the child died (#76759);
|
|
# ask the process itself (a WNOHANG waitpid).
|
|
or not s.bridge.is_alive()
|
|
]
|
|
for key in doomed:
|
|
# Reaps overlap (attach_or_spawn and the background reaper) and close()
|
|
# awaits, so a concurrent reap can have popped this key already — skip
|
|
# it instead of raising KeyError into the websocket handler.
|
|
session = self._sessions.pop(key, None)
|
|
if session is not None:
|
|
await session.close()
|
|
|
|
def _reap_one_idle_or_raise(self) -> None:
|
|
idle = [s for s in self._sessions.values() if not s.attached and s.last_detached_at is not None]
|
|
if not idle:
|
|
raise RegistryFull()
|
|
oldest = min(idle, key=lambda s: s.last_detached_at or 0.0)
|
|
self._sessions.pop(oldest.key, None)
|
|
self._close_in_background(oldest)
|
|
|
|
def _close_in_background(self, session: PtySession) -> None:
|
|
task = asyncio.create_task(session.close())
|
|
self._background_closes.add(task)
|
|
task.add_done_callback(self._background_closes.discard)
|
|
|
|
async def close_all(self) -> None:
|
|
# Close concurrently: each close() may wait out its helpers' SIGHUP grace, and shutdown runs
|
|
# under the backend's SIGTERM -> SIGKILL budget (dashboard_procs._POSIX_TERM_GRACE_SECONDS).
|
|
sessions = [self._sessions.pop(key) for key in list(self._sessions)]
|
|
await asyncio.gather(*(s.close() for s in sessions), *self._background_closes, return_exceptions=True)
|