butler/gateway/session_registry.py
Rama main · réplica de lectura del repositorio original.
"""Per-session AgentLoop lifecycle registry for Gateway handlers."""
from __future__ import annotations
import logging
import time
import threading
from collections.abc import Callable, Iterable, Mapping
from typing import Any
logger = logging.getLogger(__name__)
_evict_notify_hook: Callable[[str], None] | None = None
def set_evict_notify_hook(hook: Callable[[str], None] | None) -> None:
global _evict_notify_hook
_evict_notify_hook = hook
class _TrackedSessionDict(dict[str, Any]):
def __init__(self, registry: "GatewaySessionRegistry", initial: dict[str, Any] | None = None) -> None:
self._registry = registry
super().__init__()
if initial:
self.update(initial)
def __setitem__(self, key: str, value: Any) -> None:
normalized = str(key or "default")
with self._registry._lock:
super().__setitem__(normalized, value)
self._registry._last_active_at[normalized] = self._registry._now()
def update( # type: ignore[override]
self,
other: Mapping[str, Any] | Iterable[tuple[str, Any]] = (),
/,
**kwargs: Any,
) -> None:
items = dict(other, **kwargs)
for key, value in items.items():
self[key] = value
class GatewaySessionRegistry:
"""Track Gateway AgentLoop instances, health, and lifecycle finalization."""
def __init__(
self,
loop_factory: Callable[[str], Any],
*,
finalize: Callable[[Any], Any] | None = None,
on_session_removed: Callable[[str], None] | None = None,
max_sessions: int = 128,
idle_ttl_seconds: float = 3600,
now: Callable[[], float] | None = None,
) -> None:
self._loop_factory = loop_factory
self._finalize = finalize
self._on_session_removed = on_session_removed
self.max_sessions = max(1, int(max_sessions or 1))
self.idle_ttl_seconds = max(0.0, float(idle_ttl_seconds or 0))
self._now = now or time.monotonic
self._lock = threading.RLock()
self.sessions: dict[str, Any] = _TrackedSessionDict(self)
self.health_by_session: dict[str, dict[str, Any]] = {}
self._last_active_at: dict[str, float] = {}
self._session_locks: dict[str, threading.RLock] = {}
self._active_sessions: set[str] = set()
self._pending_session_entries = 0
self._resetting_all = False
self._reset_condition = threading.Condition(self._lock)
def get_or_create(self, session_key: str) -> Any:
key = str(session_key or "default")
with self._lock:
self._wait_for_reset_all_locked(key)
if key not in self.sessions:
self.sessions[key] = self._loop_factory(key)
self.touch(key)
loop = self.sessions[key]
self.enforce_lru()
self._publish_session_gauges()
return loop
def session_lock(self, session_key: str) -> threading.RLock:
key = str(session_key or "default")
with self._lock:
self._wait_for_reset_all_locked()
lock = self._session_locks.get(key)
if lock is None:
lock = threading.RLock()
self._session_locks[key] = lock
return lock
def enter_session(self, session_key: str) -> threading.RLock:
"""Enter a session turn and mark it active before work begins."""
key = str(session_key or "default")
while True:
with self._lock:
self._wait_for_reset_all_locked()
lock = self._session_locks.get(key)
if lock is None:
lock = threading.RLock()
self._session_locks[key] = lock
self._pending_session_entries += 1
lock.acquire()
with self._lock:
self._pending_session_entries -= 1
self._reset_condition.notify_all()
if not self._resetting_all:
self._active_sessions.add(key)
return lock
lock.release()
def exit_session(self, session_key: str, lock: threading.RLock) -> None:
key = str(session_key or "default")
try:
with self._lock:
self._active_sessions.discard(key)
self._reset_condition.notify_all()
finally:
lock.release()
def mark_active(self, session_key: str) -> None:
with self._lock:
self._wait_for_reset_all_locked()
self._active_sessions.add(str(session_key or "default"))
self._publish_session_gauges()
def mark_inactive(self, session_key: str) -> None:
with self._lock:
self._active_sessions.discard(str(session_key or "default"))
self._reset_condition.notify_all()
self._publish_session_gauges()
def _publish_session_gauges(self) -> None:
from butler.gateway.session_registry_ops import publish_session_gauges_safe
with self._lock:
publish_session_gauges_safe(
session_count=len(self.sessions),
active_turns=len(self._active_sessions),
)
def is_session_active(self, session_key: str) -> bool:
"""True while a gateway turn holds the session lock (AgentLoop running)."""
key = str(session_key or "default")
with self._lock:
return key in self._active_sessions
def touch(self, session_key: str) -> None:
with self._lock:
self._last_active_at[str(session_key or "default")] = self._now()
def last_active_at(self, session_key: str) -> float:
with self._lock:
return self._last_active_at.get(str(session_key or "default"), 0.0)
def set_health(self, session_key: str, health: dict[str, Any]) -> None:
with self._lock:
self.health_by_session[str(session_key or "default")] = dict(health)
def get_health(self, session_key: str) -> dict[str, Any]:
with self._lock:
return dict(self.health_by_session.get(str(session_key or "default"), {}))
def reset_sessions_for_chat(self, *, platform: str, chat_id: str) -> list[str]:
"""Drop all cached loops for ``platform:chat_id:*`` (e.g. after /切换 project)."""
plat = str(platform or "unknown").strip() or "unknown"
cid = str(chat_id or "default").strip() or "default"
prefix = f"{plat}:{cid}:"
with self._lock:
keys = [k for k in list(self.sessions.keys()) if k.startswith(prefix)]
cleared: list[str] = []
for key in keys:
self.reset(key)
cleared.append(key)
if cleared:
logger.info(
"Reset %d gateway session(s) for chat %s after project switch",
len(cleared),
prefix,
)
return cleared
def reset(self, session_key: str, *, skip_finalize: bool = False) -> None:
key = str(session_key or "default")
with self.session_lock(key):
with self._lock:
loop = self.sessions.pop(key, None)
self.health_by_session.pop(key, None)
self._last_active_at.pop(key, None)
self._active_sessions.discard(key)
self._session_locks.pop(key, None)
self._notify_session_removed(key)
if not skip_finalize:
self._finalize_loop(loop)
def reset_all(self, *, wait_timeout: float = 120.0) -> None:
finalized: set[int] = set()
with self._lock:
self._wait_for_reset_all_locked()
self._resetting_all = True
deadline = self._now() + max(1.0, float(wait_timeout))
try:
with self._lock:
while (self._active_sessions or self._pending_session_entries) and self._now() < deadline:
self._reset_condition.wait(timeout=1.0)
if self._active_sessions or self._pending_session_entries:
logger.error(
"reset_all timed out after %.0fs (active=%s pending=%d); forcing clear",
wait_timeout,
sorted(self._active_sessions),
self._pending_session_entries,
)
self._active_sessions.clear()
self._pending_session_entries = 0
items = list(self.sessions.items())
self.sessions.clear()
self.health_by_session.clear()
self._last_active_at.clear()
self._active_sessions.clear()
self._session_locks.clear()
finally:
with self._lock:
self._resetting_all = False
self._reset_condition.notify_all()
for key, _loop in items:
self._notify_session_removed(key)
for _key, loop in items:
if id(loop) in finalized:
continue
finalized.add(id(loop))
self._finalize_loop(loop)
def evict_idle(self) -> list[str]:
with self._lock:
if self.idle_ttl_seconds <= 0 or self._resetting_all:
return []
cutoff = self._now() - self.idle_ttl_seconds
expired = [
key for key in self.sessions
if key not in self._active_sessions
and self._last_active_at.get(key, 0.0) < cutoff
]
evicted: list[str] = []
for key in expired:
if self._reset_if_still_idle(key, cutoff):
evicted.append(key)
if evicted and _evict_notify_hook is not None:
from butler.gateway.session_registry_ops import run_evict_notify_hook_safe
for key in evicted:
run_evict_notify_hook_safe(_evict_notify_hook, key)
return evicted
def enforce_lru(self) -> list[str]:
loops_to_finalize: list[Any] = []
with self._lock:
if self._resetting_all:
return []
evicted: list[str] = []
while len(self.sessions) > self.max_sessions:
candidates = [key for key in self.sessions if key not in self._active_sessions]
if not candidates:
break
oldest = min(candidates, key=lambda key: self._last_active_at.get(key, 0.0))
evicted.append(oldest)
loop = self.sessions.pop(oldest, None)
self.health_by_session.pop(oldest, None)
self._last_active_at.pop(oldest, None)
self._session_locks.pop(oldest, None)
self._notify_session_removed(oldest)
loops_to_finalize.append(loop)
for loop in loops_to_finalize:
self._finalize_loop(loop)
return evicted
def _finalize_loop(self, loop: Any) -> None:
if loop is None or self._finalize is None:
return
self._finalize(loop)
def _wait_for_reset_all_locked(self, session_key: str | None = None) -> None:
while self._resetting_all and (
session_key is None or session_key not in self._active_sessions
):
self._reset_condition.wait()
def _reset_if_still_idle(self, session_key: str, cutoff: float) -> bool:
key = str(session_key or "default")
loop = None
with self.session_lock(key):
with self._lock:
if (
self._resetting_all
or key in self._active_sessions
or self._last_active_at.get(key, 0.0) >= cutoff
):
return False
loop = self.sessions.pop(key, None)
self.health_by_session.pop(key, None)
self._last_active_at.pop(key, None)
self._session_locks.pop(key, None)
if loop is not None:
self._notify_session_removed(key)
self._finalize_loop(loop)
return loop is not None
def _notify_session_removed(self, session_key: str) -> None:
from butler.gateway.session_registry_ops import notify_session_removed_safe
notify_session_removed_safe(self._on_session_removed, session_key)