251 lines
9.1 KiB
Python
251 lines
9.1 KiB
Python
"""Cascade State Store — Persistent cascade detection state with Redis INCR backend.
|
|
|
|
Provides CascadeStateStore Protocol with InMemoryCascadeStateStore and
|
|
RedisCascadeStateStore backends. Enables CascadeDetector state to survive
|
|
restarts and work across multiple instances.
|
|
|
|
Key schema (Redis):
|
|
agentkit:cascade:interactions:{session_id} → INCR counter with TTL
|
|
agentkit:cascade:depths:{session_id} → SET counter with TTL
|
|
"""
|
|
|
|
import logging
|
|
import time
|
|
from typing import Any, Protocol, runtime_checkable
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@runtime_checkable
|
|
class CascadeStateStore(Protocol):
|
|
"""Persistent cascade detection state interface."""
|
|
|
|
def increment_interaction(self, session_id: str) -> int:
|
|
"""Atomically increment interaction count. Returns new count."""
|
|
...
|
|
|
|
def get_interaction(self, session_id: str) -> int:
|
|
"""Get current interaction count."""
|
|
...
|
|
|
|
def set_depth(self, session_id: str, depth: int) -> None:
|
|
"""Set loop depth for a session."""
|
|
...
|
|
|
|
def get_depth(self, session_id: str) -> int:
|
|
"""Get current loop depth."""
|
|
...
|
|
|
|
def reset(self, session_id: str) -> None:
|
|
"""Reset all counters for a session."""
|
|
...
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# InMemoryCascadeStateStore
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class InMemoryCascadeStateStore:
|
|
"""In-memory cascade state store (default, process-local).
|
|
|
|
Supports optional session TTL to prevent unbounded memory growth.
|
|
Expired entries are lazily cleaned up on access.
|
|
"""
|
|
|
|
DEFAULT_SESSION_TTL = 86400 # 24 hours
|
|
|
|
def __init__(self, session_ttl: int = 86400):
|
|
self._session_ttl = session_ttl
|
|
self._interaction_counts: dict[str, int] = {}
|
|
self._loop_depths: dict[str, int] = {}
|
|
self._timestamps: dict[str, float] = {}
|
|
|
|
def _is_expired(self, session_id: str) -> bool:
|
|
ts = self._timestamps.get(session_id)
|
|
if ts is None:
|
|
return False
|
|
return (time.monotonic() - ts) > self._session_ttl
|
|
|
|
def _cleanup_expired(self) -> None:
|
|
"""Lazy cleanup: remove expired sessions."""
|
|
expired = [sid for sid in self._timestamps if self._is_expired(sid)]
|
|
for sid in expired:
|
|
self._interaction_counts.pop(sid, None)
|
|
self._loop_depths.pop(sid, None)
|
|
self._timestamps.pop(sid, None)
|
|
|
|
def _touch(self, session_id: str) -> None:
|
|
self._timestamps[session_id] = time.monotonic()
|
|
|
|
def increment_interaction(self, session_id: str) -> int:
|
|
self._cleanup_expired()
|
|
self._interaction_counts[session_id] = self._interaction_counts.get(session_id, 0) + 1
|
|
self._touch(session_id)
|
|
return self._interaction_counts[session_id]
|
|
|
|
def get_interaction(self, session_id: str) -> int:
|
|
if self._is_expired(session_id):
|
|
self.reset(session_id)
|
|
return 0
|
|
return self._interaction_counts.get(session_id, 0)
|
|
|
|
def set_depth(self, session_id: str, depth: int) -> None:
|
|
self._touch(session_id)
|
|
self._loop_depths[session_id] = depth
|
|
|
|
def get_depth(self, session_id: str) -> int:
|
|
if self._is_expired(session_id):
|
|
self.reset(session_id)
|
|
return 0
|
|
return self._loop_depths.get(session_id, 0)
|
|
|
|
def reset(self, session_id: str) -> None:
|
|
self._interaction_counts.pop(session_id, None)
|
|
self._loop_depths.pop(session_id, None)
|
|
self._timestamps.pop(session_id, None)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# RedisCascadeStateStore
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class RedisCascadeStateStore:
|
|
"""Redis-backed cascade state store using INCR for atomic increments.
|
|
|
|
Key schema:
|
|
agentkit:cascade:interactions:{session_id} → INCR counter with TTL
|
|
agentkit:cascade:depths:{session_id} → SET counter with TTL
|
|
"""
|
|
|
|
INTER_PREFIX = "agentkit:cascade:interactions:"
|
|
DEPTH_PREFIX = "agentkit:cascade:depths:"
|
|
SESSION_TTL = 86400 # 24 hours — sessions rarely last longer
|
|
|
|
def __init__(self, redis_url: str = "redis://localhost:6379", session_ttl: int = 86400):
|
|
self._redis_url = redis_url
|
|
self._session_ttl = session_ttl
|
|
self._sync_redis: Any = None
|
|
self._fallback: InMemoryCascadeStateStore | None = None
|
|
self._degraded = False
|
|
|
|
def _get_sync_redis(self):
|
|
"""Get or create a persistent sync Redis client (connection pool backed)."""
|
|
if self._sync_redis is None:
|
|
import redis as sync_redis
|
|
self._sync_redis = sync_redis.from_url(
|
|
self._redis_url, decode_responses=True
|
|
)
|
|
return self._sync_redis
|
|
|
|
def _degrade_to_fallback(self) -> None:
|
|
if not self._degraded:
|
|
self._degraded = True
|
|
if self._fallback is None:
|
|
self._fallback = InMemoryCascadeStateStore(session_ttl=self._session_ttl)
|
|
logger.warning("Redis cascade store unreachable, degraded to in-memory")
|
|
|
|
def increment_interaction(self, session_id: str) -> int:
|
|
if self._degraded and self._fallback is not None:
|
|
return self._fallback.increment_interaction(session_id)
|
|
try:
|
|
r = self._get_sync_redis()
|
|
key = f"{self.INTER_PREFIX}{session_id}"
|
|
pipe = r.pipeline()
|
|
pipe.incr(key)
|
|
pipe.expire(key, self._session_ttl)
|
|
results = pipe.execute()
|
|
return results[0]
|
|
except Exception as e:
|
|
logger.warning(f"Redis cascade increment failed: {e}")
|
|
self._degrade_to_fallback()
|
|
if self._fallback is not None:
|
|
return self._fallback.increment_interaction(session_id)
|
|
return 0
|
|
|
|
def get_interaction(self, session_id: str) -> int:
|
|
if self._degraded and self._fallback is not None:
|
|
return self._fallback.get_interaction(session_id)
|
|
try:
|
|
r = self._get_sync_redis()
|
|
val = r.get(f"{self.INTER_PREFIX}{session_id}")
|
|
return int(val) if val is not None else 0
|
|
except Exception as e:
|
|
logger.warning(f"Redis cascade get failed: {e}")
|
|
if self._fallback is not None:
|
|
return self._fallback.get_interaction(session_id)
|
|
return 0
|
|
|
|
def set_depth(self, session_id: str, depth: int) -> None:
|
|
if self._degraded and self._fallback is not None:
|
|
self._fallback.set_depth(session_id, depth)
|
|
return
|
|
try:
|
|
r = self._get_sync_redis()
|
|
key = f"{self.DEPTH_PREFIX}{session_id}"
|
|
pipe = r.pipeline()
|
|
pipe.set(key, depth)
|
|
pipe.expire(key, self._session_ttl)
|
|
pipe.execute()
|
|
except Exception as e:
|
|
logger.warning(f"Redis cascade set_depth failed: {e}")
|
|
self._degrade_to_fallback()
|
|
if self._fallback is not None:
|
|
self._fallback.set_depth(session_id, depth)
|
|
|
|
def get_depth(self, session_id: str) -> int:
|
|
if self._degraded and self._fallback is not None:
|
|
return self._fallback.get_depth(session_id)
|
|
try:
|
|
r = self._get_sync_redis()
|
|
val = r.get(f"{self.DEPTH_PREFIX}{session_id}")
|
|
return int(val) if val is not None else 0
|
|
except Exception as e:
|
|
logger.warning(f"Redis cascade get_depth failed: {e}")
|
|
if self._fallback is not None:
|
|
return self._fallback.get_depth(session_id)
|
|
return 0
|
|
|
|
def reset(self, session_id: str) -> None:
|
|
if self._degraded and self._fallback is not None:
|
|
self._fallback.reset(session_id)
|
|
return
|
|
try:
|
|
r = self._get_sync_redis()
|
|
pipe = r.pipeline()
|
|
pipe.delete(f"{self.INTER_PREFIX}{session_id}")
|
|
pipe.delete(f"{self.DEPTH_PREFIX}{session_id}")
|
|
pipe.execute()
|
|
except Exception as e:
|
|
logger.warning(f"Redis cascade reset failed: {e}")
|
|
self._degrade_to_fallback()
|
|
if self._fallback is not None:
|
|
self._fallback.reset(session_id)
|
|
|
|
def close(self) -> None:
|
|
"""Close the Redis connection pool."""
|
|
if self._sync_redis is not None:
|
|
self._sync_redis.close()
|
|
self._sync_redis = None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Factory
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def create_cascade_state_store(
|
|
backend: str = "auto",
|
|
redis_url: str = "redis://localhost:6379",
|
|
session_ttl: int = 86400,
|
|
) -> CascadeStateStore:
|
|
"""Create a cascade state store backend."""
|
|
if backend in ("auto", "redis"):
|
|
try:
|
|
import redis # noqa: F401
|
|
return RedisCascadeStateStore(redis_url=redis_url, session_ttl=session_ttl)
|
|
except ImportError:
|
|
logger.warning("redis package not available, falling back to in-memory cascade store")
|
|
return InMemoryCascadeStateStore(session_ttl=session_ttl)
|