317 lines
11 KiB
Python
317 lines
11 KiB
Python
"""Pool of live Agent SDK sessions across every Claude agent.
|
|
|
|
One :class:`Session` is one ``ClaudeSDKClient`` (one claude subprocess). The
|
|
pool decides when a session is closed for idleness (TTL by conversation
|
|
kind) and which one goes when memory runs out (RSS of the subprocess tree
|
|
against the cgroup limit, or ``max_live`` with no limit).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
import os
|
|
import time
|
|
from collections import deque
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any, Protocol
|
|
|
|
import psutil
|
|
|
|
from beaver_gateway.backends.processes import Processes, subtree
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import (
|
|
AsyncIterable,
|
|
AsyncIterator,
|
|
Callable,
|
|
Iterator,
|
|
Mapping,
|
|
)
|
|
|
|
__all__ = ["DEFAULT_TTL", "Session", "SessionClient", "SessionPool", "cgroup_limit"]
|
|
|
|
_log = logging.getLogger("beaver_gateway.backends.sessions")
|
|
|
|
DEFAULT_TTL: Mapping[str, float | None] = {
|
|
"master": None,
|
|
"branch": 7200.0,
|
|
"deep": 1800.0,
|
|
"job": 0.0,
|
|
"fork": 0.0,
|
|
}
|
|
"""Idle seconds before a session is closed; ``None`` = never (pinned kinds)."""
|
|
|
|
_EVICT_ORDER = {"fork": 0, "job": 0, "deep": 1, "branch": 2, "master": 3}
|
|
_RSS_HEADROOM = 0.8
|
|
STDERR_LINES = 200
|
|
|
|
|
|
class SessionClient(Protocol):
|
|
async def connect(self) -> None: ...
|
|
async def query(self, prompt: str | AsyncIterable[dict[str, Any]]) -> None: ...
|
|
def receive_messages(self) -> AsyncIterator[Any]: ...
|
|
async def interrupt(self) -> None: ...
|
|
async def disconnect(self) -> None: ...
|
|
|
|
|
|
@dataclass
|
|
class Session:
|
|
key: str
|
|
agent: str
|
|
kind: str
|
|
client: SessionClient
|
|
session_id: str | None
|
|
resumed: bool
|
|
pinned: bool = False
|
|
dirty: bool = False
|
|
running_turn: str | None = None
|
|
pending_question: bool = False
|
|
interrupt_requested: bool = False
|
|
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
|
|
last_used: float = field(default_factory=time.monotonic)
|
|
created_at: float = field(default_factory=time.monotonic)
|
|
turns: int = 0
|
|
state: dict[str, Any] = field(default_factory=dict)
|
|
"""Scratch for policy rules (``agents/policy.py``); dies with the process."""
|
|
|
|
reader: asyncio.Task[None] | None = None
|
|
"""Reads the client for the session's whole life, not only during a turn."""
|
|
inbox: asyncio.Queue[Any] = field(default_factory=asyncio.Queue)
|
|
"""Where the reader puts messages while a turn is attached."""
|
|
pending: list[Any] = field(default_factory=list)
|
|
"""A turn the CLI started on its own (a background task reported back),
|
|
buffered until a consumer attaches."""
|
|
attached: bool = False
|
|
woken: bool = False
|
|
note: str | None = None
|
|
"""Summary of the last task notification, the pending turn's prompt."""
|
|
observer: Callable[[Any], None] | None = None
|
|
"""The last turn's observer; sees subagent traffic between turns."""
|
|
stderr: deque[str] = field(default_factory=lambda: deque(maxlen=STDERR_LINES))
|
|
|
|
@property
|
|
def busy(self) -> bool:
|
|
return self.lock.locked() or self.running_turn is not None
|
|
|
|
@property
|
|
def evictable(self) -> bool:
|
|
"""Never evict a pinned (master) session or a ``dirty`` one.
|
|
|
|
``dirty`` means a mirror gap not yet repaired.
|
|
"""
|
|
return not (self.pinned or self.dirty or self.busy or self.pending_question)
|
|
|
|
@property
|
|
def pid(self) -> int | None:
|
|
transport = getattr(self.client, "_transport", None)
|
|
process = getattr(transport, "_process", None)
|
|
pid = getattr(process, "pid", None)
|
|
return pid if isinstance(pid, int) else None
|
|
|
|
|
|
class SessionPool:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
rss_limit: int | None = None,
|
|
max_live: int = 8,
|
|
ttl: Mapping[str, float | None] = DEFAULT_TTL,
|
|
reap_interval: float = 60.0,
|
|
) -> None:
|
|
self._sessions: dict[str, Session] = {}
|
|
self._rss_limit = rss_limit if rss_limit is not None else cgroup_limit()
|
|
self._max_live = max_live
|
|
self._ttl = dict(ttl)
|
|
self._reap_interval = reap_interval
|
|
self._processes = Processes()
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._sessions)
|
|
|
|
def __iter__(self) -> Iterator[Session]:
|
|
return iter(list(self._sessions.values()))
|
|
|
|
def __contains__(self, key: object) -> bool:
|
|
return key in self._sessions
|
|
|
|
@property
|
|
def rss_limit(self) -> int | None:
|
|
return self._rss_limit
|
|
|
|
def get(self, key: str) -> Session | None:
|
|
return self._sessions.get(key)
|
|
|
|
def add(self, session: Session) -> Session:
|
|
self._sessions[session.key] = session
|
|
return session
|
|
|
|
def pop(self, key: str) -> Session | None:
|
|
return self._sessions.pop(key, None)
|
|
|
|
def rekey(self, old: str, new: str) -> Session | None:
|
|
session = self._sessions.pop(old, None)
|
|
if session is None:
|
|
return None
|
|
stale = self._sessions.pop(new, None)
|
|
session.key = new
|
|
self._sessions[new] = session
|
|
return stale if stale is not session else None
|
|
|
|
def ttl_for(self, kind: str) -> float | None:
|
|
return self._ttl.get(kind, self._ttl.get("deep"))
|
|
|
|
def rss(self) -> int:
|
|
try:
|
|
children = psutil.Process(os.getpid()).children(recursive=True)
|
|
except psutil.Error:
|
|
return 0
|
|
total = 0
|
|
for child in children:
|
|
with contextlib.suppress(psutil.Error):
|
|
total += child.memory_info().rss
|
|
return total
|
|
|
|
def over_limit(self) -> bool:
|
|
if self._rss_limit is not None:
|
|
return self.rss() > self._rss_limit * _RSS_HEADROOM
|
|
return len(self._sessions) >= self._max_live
|
|
|
|
def victims(self) -> list[Session]:
|
|
"""Evictable sessions, forks/jobs first, then oldest idle."""
|
|
candidates = [s for s in self._sessions.values() if s.evictable]
|
|
candidates.sort(key=lambda s: (_EVICT_ORDER.get(s.kind, 1), s.last_used))
|
|
return candidates
|
|
|
|
async def make_room(self) -> int:
|
|
closed = 0
|
|
while self.over_limit():
|
|
victims = self.victims()
|
|
if not victims:
|
|
_log.warning(
|
|
"session pool over limit (%d live, rss=%d) but nothing evictable",
|
|
len(self._sessions),
|
|
self.rss(),
|
|
)
|
|
break
|
|
await self.close(victims[0].key)
|
|
closed += 1
|
|
return closed
|
|
|
|
async def close(self, key: str) -> None:
|
|
session = self._sessions.pop(key, None)
|
|
if session is None:
|
|
return
|
|
_log.info("closing session %s (%s, %s)", session.session_id, session.kind, key)
|
|
if session.reader is not None and session.reader is not asyncio.current_task():
|
|
session.reader.cancel()
|
|
if session.attached:
|
|
session.inbox.put_nowait(RuntimeError(f"session {key} was closed"))
|
|
try:
|
|
await session.client.disconnect()
|
|
except Exception: # noqa: BLE001
|
|
_log.exception("disconnect failed for session %s", session.session_id)
|
|
|
|
async def interrupt(self, key: str) -> bool:
|
|
session = self._sessions.get(key)
|
|
if session is None or not session.busy:
|
|
return False
|
|
session.interrupt_requested = True
|
|
await session.client.interrupt()
|
|
return True
|
|
|
|
async def close_all(self, *, agent: str | None = None) -> None:
|
|
for session in list(self._sessions.values()):
|
|
if agent is None or session.agent == agent:
|
|
await self.close(session.key)
|
|
|
|
async def reap_once(self) -> int:
|
|
now = time.monotonic()
|
|
closed = 0
|
|
for session in list(self._sessions.values()):
|
|
ttl = self.ttl_for(session.kind)
|
|
if ttl is None or not session.evictable:
|
|
continue
|
|
if now - session.last_used > ttl:
|
|
await self.close(session.key)
|
|
closed += 1
|
|
return closed
|
|
|
|
async def reap_loop(self) -> None:
|
|
while True:
|
|
await asyncio.sleep(self._reap_interval)
|
|
try:
|
|
await self.reap_once()
|
|
await self.make_room()
|
|
except Exception: # noqa: BLE001
|
|
_log.exception("session reaper failed")
|
|
|
|
def snapshot(self) -> list[dict[str, Any]]:
|
|
return self.census()["sessions"]
|
|
|
|
def census(self) -> dict[str, Any]:
|
|
now = time.monotonic()
|
|
rows = self._processes.tree_parents_first()
|
|
owned: set[int] = set()
|
|
sessions: list[dict[str, Any]] = []
|
|
for s in self._sessions.values():
|
|
tree = subtree(rows, s.pid)
|
|
owned.update(row["pid"] for row in tree)
|
|
sessions.append(
|
|
{
|
|
"key": s.key,
|
|
"agent": s.agent,
|
|
"kind": s.kind,
|
|
"session_id": s.session_id,
|
|
"pid": s.pid,
|
|
"rss": sum(row["rss"] for row in tree) if tree else None,
|
|
"cpu": _cpu(tree),
|
|
"idle_seconds": round(now - s.last_used, 1),
|
|
"age_seconds": round(now - s.created_at, 1),
|
|
"turns": s.turns,
|
|
"busy": s.busy,
|
|
"running_turn": s.running_turn,
|
|
"pending": len(s.pending),
|
|
"pending_question": s.pending_question,
|
|
"pinned": s.pinned,
|
|
"dirty": s.dirty,
|
|
"processes": tree,
|
|
}
|
|
)
|
|
return {
|
|
"rss": sum(row["rss"] for row in rows),
|
|
"rss_limit": self._rss_limit,
|
|
"cpu": _cpu(rows),
|
|
"sessions": sessions,
|
|
"others": [row for row in rows if row["pid"] not in owned],
|
|
}
|
|
|
|
async def kill(self, pid: int) -> bool:
|
|
if not self._processes.kill_own_descendant(pid):
|
|
return False
|
|
for session in list(self._sessions.values()):
|
|
if session.pid == pid:
|
|
await self.close(session.key)
|
|
return True
|
|
|
|
|
|
def _cpu(rows: list[dict[str, Any]]) -> float | None:
|
|
rates = [row["cpu"] for row in rows if row["cpu"] is not None]
|
|
return round(sum(rates), 1) if rates else None
|
|
|
|
|
|
def cgroup_limit() -> int | None:
|
|
for path in (
|
|
"/sys/fs/cgroup/memory.max",
|
|
"/sys/fs/cgroup/memory/memory.limit_in_bytes",
|
|
):
|
|
try:
|
|
raw = Path(path).read_text(encoding="ascii").strip()
|
|
except OSError:
|
|
continue
|
|
if raw.isdigit() and int(raw) < 1 << 60:
|
|
return int(raw)
|
|
return None
|