260 lines
8.2 KiB
Python
260 lines
8.2 KiB
Python
"""Pool of live Agent SDK sessions across every Claude agent (§3.2).
|
|
|
|
One :class:`Session` is one ``ClaudeSDKClient`` (one claude subprocess).
|
|
The pool owns the two decisions the adapters used to make on their own:
|
|
when a session is closed for idleness (TTL by conversation kind) and
|
|
which one goes when memory runs out (measured RSS of the subprocess tree
|
|
against the cgroup limit, ``max_live`` where there is no limit). Eviction
|
|
only ever picks ``idle && !running_turn && !pending_question`` sessions
|
|
that are neither pinned (the master) nor ``dirty`` (mirror gap not yet
|
|
repaired); forks and jobs go first.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
import os
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any, Protocol
|
|
|
|
import psutil
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator, Iterator, Mapping
|
|
|
|
__all__ = ["DEFAULT_TTL", "Session", "SessionClient", "SessionPool", "cgroup_limit"]
|
|
|
|
_log = logging.getLogger("beaver_gateway.core.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
|
|
|
|
|
|
class SessionClient(Protocol):
|
|
async def connect(self) -> None: ...
|
|
async def query(self, prompt: str) -> None: ...
|
|
def receive_response(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 (``core/policy``); dies with the process."""
|
|
|
|
@property
|
|
def busy(self) -> bool:
|
|
return self.lock.locked() or self.running_turn is not None
|
|
|
|
@property
|
|
def evictable(self) -> bool:
|
|
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
|
|
|
|
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
|
|
|
|
@staticmethod
|
|
def rss_of(session: Session) -> int | None:
|
|
pid = session.pid
|
|
if pid is None:
|
|
return None
|
|
try:
|
|
process = psutil.Process(pid)
|
|
return process.memory_info().rss + sum(
|
|
c.memory_info().rss for c in process.children(recursive=True)
|
|
)
|
|
except psutil.Error:
|
|
return None
|
|
|
|
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]:
|
|
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)
|
|
try:
|
|
await session.client.disconnect()
|
|
except Exception: # noqa: BLE001
|
|
_log.exception("disconnect failed for session %s", session.session_id)
|
|
|
|
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]]:
|
|
now = time.monotonic()
|
|
return [
|
|
{
|
|
"key": s.key,
|
|
"agent": s.agent,
|
|
"kind": s.kind,
|
|
"session_id": s.session_id,
|
|
"pid": s.pid,
|
|
"rss": self.rss_of(s),
|
|
"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_question": s.pending_question,
|
|
"pinned": s.pinned,
|
|
"dirty": s.dirty,
|
|
}
|
|
for s in self._sessions.values()
|
|
]
|
|
|
|
|
|
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
|