ad4e349df7
Review findings on #126144 (independent review, all minor): - pty_session.close_other_sessions closed a sibling tab's PTY under it without closing its websocket: two tabs sharing one attach token on different profiles left tab A silent until its next keystroke returned 1013. Close the displaced viewer with the documented 4409 before closing the session. Still closes attached siblings (skipping them would reopen the #125287 lease hole via the old handler's finally: detach). - openNewSessionTile paired projectProfile() with $newChatConnectionId — the source captured with a stale new-chat pin — so a pin made on another connection produced the right profile on the wrong host. Build the route from the active source (resolveActiveSourceOwnerRoute) for the project-derived profile. - windows.test.ts: restore the no-parentSessionId case asserting the viewed-profile fallback alongside the parent-routed case.
262 lines
11 KiB
Python
262 lines
11 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 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
|
|
|
|
|
|
class PtySession:
|
|
def __init__(self, key: str, bridge, *, buffer_cap: int, read_timeout: float) -> None:
|
|
self.key = key
|
|
self.bridge = bridge
|
|
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()
|
|
|
|
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:
|
|
pass
|
|
|
|
|
|
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()
|
|
|
|
async def attach_or_spawn(self, key: str, *, spawn: Callable[[], object]) -> 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
|
|
await existing.close()
|
|
self._sessions.pop(key, None)
|
|
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)
|
|
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()
|
|
|
|
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)
|
|
]
|
|
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)
|
|
asyncio.create_task(oldest.close())
|
|
|
|
async def close_all(self) -> None:
|
|
for key in list(self._sessions):
|
|
# Same overlap window as reap_idle: an in-flight reap may have popped
|
|
# a snapshot key while we awaited an earlier close().
|
|
session = self._sessions.pop(key, None)
|
|
if session is not None:
|
|
await session.close()
|