feat(core,backends,frontends,storage): conversations, inject queue, session pool, gateway mcp tools, api frontend
This commit is contained in:
@@ -0,0 +1,257 @@
|
||||
"""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
|
||||
|
||||
@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
|
||||
Reference in New Issue
Block a user