358 lines
14 KiB
Python
358 lines
14 KiB
Python
"""Session store backends — InMemory and Redis."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
from typing import Any, Protocol, runtime_checkable
|
|
|
|
# redis 可选依赖；未安装时回退为 Exception 以保留原 catch-all 语义
|
|
try:
|
|
from redis.exceptions import RedisError as _RedisError
|
|
except ImportError:
|
|
_RedisError = Exception
|
|
|
|
from agentkit.session.models import Message, Session, SessionStatus
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@runtime_checkable
|
|
class SessionStore(Protocol):
|
|
"""Protocol for session persistence backends."""
|
|
|
|
async def save_session(self, session: Session) -> None: ...
|
|
async def get_session(self, session_id: str) -> Session | None: ...
|
|
async def update_session_status(self, session_id: str, status: SessionStatus) -> Session | None: ...
|
|
async def delete_session(self, session_id: str) -> bool: ...
|
|
async def list_sessions(self, agent_name: str | None = None, limit: int = 100) -> list[Session]: ...
|
|
|
|
async def append_message(self, message: Message) -> None: ...
|
|
async def get_messages(self, session_id: str, limit: int | None = None, offset: int = 0) -> list[Message]: ...
|
|
async def count_messages(self, session_id: str) -> int: ...
|
|
|
|
async def health_check(self) -> bool: ...
|
|
|
|
|
|
class InMemorySessionStore:
|
|
"""In-memory session store for development and testing."""
|
|
|
|
def __init__(self, max_sessions: int = 10000, max_messages_per_session: int = 50000):
|
|
self._sessions: dict[str, Session] = {}
|
|
self._messages: dict[str, list[Message]] = {}
|
|
self._max_sessions = max_sessions
|
|
self._max_messages_per_session = max_messages_per_session
|
|
|
|
async def save_session(self, session: Session) -> None:
|
|
if len(self._sessions) >= self._max_sessions and session.session_id not in self._sessions:
|
|
# Evict oldest closed session
|
|
closed = [s for s in self._sessions.values() if s.status == SessionStatus.CLOSED]
|
|
if closed:
|
|
oldest = min(closed, key=lambda s: s.updated_at)
|
|
del self._sessions[oldest.session_id]
|
|
self._messages.pop(oldest.session_id, None)
|
|
else:
|
|
raise RuntimeError("SessionStore is full and no closed sessions to evict")
|
|
self._sessions[session.session_id] = session
|
|
if session.session_id not in self._messages:
|
|
self._messages[session.session_id] = []
|
|
|
|
async def get_session(self, session_id: str) -> Session | None:
|
|
return self._sessions.get(session_id)
|
|
|
|
async def update_session_status(self, session_id: str, status: SessionStatus) -> Session | None:
|
|
session = self._sessions.get(session_id)
|
|
if session is None:
|
|
return None
|
|
session.status = status
|
|
session.updated_at = datetime.now(timezone.utc)
|
|
return session
|
|
|
|
async def delete_session(self, session_id: str) -> bool:
|
|
if session_id in self._sessions:
|
|
del self._sessions[session_id]
|
|
self._messages.pop(session_id, None)
|
|
return True
|
|
return False
|
|
|
|
async def list_sessions(self, agent_name: str | None = None, limit: int = 100) -> list[Session]:
|
|
sessions = list(self._sessions.values())
|
|
if agent_name:
|
|
sessions = [s for s in sessions if s.agent_name == agent_name]
|
|
sessions.sort(key=lambda s: s.updated_at, reverse=True)
|
|
return sessions[:limit]
|
|
|
|
async def append_message(self, message: Message) -> None:
|
|
msgs = self._messages.setdefault(message.session_id, [])
|
|
if len(msgs) >= self._max_messages_per_session:
|
|
# Remove oldest messages to stay within limit
|
|
excess = len(msgs) - self._max_messages_per_session + 1
|
|
del msgs[:excess]
|
|
msgs.append(message)
|
|
|
|
async def get_messages(self, session_id: str, limit: int | None = None, offset: int = 0) -> list[Message]:
|
|
msgs = self._messages.get(session_id, [])
|
|
sliced = msgs[offset:]
|
|
if limit is not None:
|
|
sliced = sliced[:limit]
|
|
return sliced
|
|
|
|
async def count_messages(self, session_id: str) -> int:
|
|
return len(self._messages.get(session_id, []))
|
|
|
|
async def health_check(self) -> bool:
|
|
return True
|
|
|
|
|
|
class RedisSessionStore:
|
|
"""Redis-backed session store for production use.
|
|
|
|
Key patterns:
|
|
- ``agentkit:session:{session_id}`` — session metadata (JSON + TTL)
|
|
- ``agentkit:session:{session_id}:messages`` — message list (Redis list)
|
|
"""
|
|
|
|
KEY_PREFIX = "agentkit:session:"
|
|
MSG_SUFFIX = ":messages"
|
|
|
|
def __init__(self, redis_url: str = "redis://localhost:6379/0", ttl_seconds: int = 86400):
|
|
self._redis_url = redis_url
|
|
self._ttl_seconds = ttl_seconds
|
|
self._redis: Any = None
|
|
|
|
async def _get_redis(self):
|
|
if self._redis is None:
|
|
import redis.asyncio as aioredis
|
|
|
|
self._redis = aioredis.from_url(self._redis_url, decode_responses=True)
|
|
return self._redis
|
|
|
|
def _session_key(self, session_id: str) -> str:
|
|
return f"{self.KEY_PREFIX}{session_id}"
|
|
|
|
def _messages_key(self, session_id: str) -> str:
|
|
return f"{self.KEY_PREFIX}{session_id}{self.MSG_SUFFIX}"
|
|
|
|
async def save_session(self, session: Session) -> None:
|
|
redis = await self._get_redis()
|
|
key = self._session_key(session.session_id)
|
|
await redis.set(key, json.dumps(session.to_dict()), ex=self._ttl_seconds)
|
|
|
|
async def get_session(self, session_id: str) -> Session | None:
|
|
redis = await self._get_redis()
|
|
raw = await redis.get(self._session_key(session_id))
|
|
if raw is None:
|
|
return None
|
|
return Session.from_dict(json.loads(raw))
|
|
|
|
async def update_session_status(self, session_id: str, status: SessionStatus) -> Session | None:
|
|
redis = await self._get_redis()
|
|
key = self._session_key(session_id)
|
|
raw = await redis.get(key)
|
|
if raw is None:
|
|
return None
|
|
session = Session.from_dict(json.loads(raw))
|
|
session.status = status
|
|
session.updated_at = datetime.now(timezone.utc)
|
|
await redis.set(key, json.dumps(session.to_dict()), ex=self._ttl_seconds)
|
|
return session
|
|
|
|
async def delete_session(self, session_id: str) -> bool:
|
|
redis = await self._get_redis()
|
|
keys = [self._session_key(session_id), self._messages_key(session_id)]
|
|
deleted = await redis.delete(*keys)
|
|
return deleted > 0
|
|
|
|
async def list_sessions(self, agent_name: str | None = None, limit: int = 100) -> list[Session]:
|
|
redis = await self._get_redis()
|
|
sessions: list[Session] = []
|
|
cursor = 0
|
|
while True:
|
|
cursor, keys = await redis.scan(cursor, match=f"{self.KEY_PREFIX}*", count=200)
|
|
# Filter out message list keys
|
|
session_keys = [k for k in keys if not k.endswith(self.MSG_SUFFIX)]
|
|
if session_keys:
|
|
values = await redis.mget(session_keys)
|
|
for raw in values:
|
|
if raw is None:
|
|
continue
|
|
session = Session.from_dict(json.loads(raw))
|
|
if agent_name is None or session.agent_name == agent_name:
|
|
sessions.append(session)
|
|
if cursor == 0:
|
|
break
|
|
sessions.sort(key=lambda s: s.updated_at, reverse=True)
|
|
return sessions[:limit]
|
|
|
|
async def append_message(self, message: Message) -> None:
|
|
redis = await self._get_redis()
|
|
key = self._messages_key(message.session_id)
|
|
await redis.rpush(key, json.dumps(message.to_dict()))
|
|
# Set TTL on message list to match session TTL
|
|
await redis.expire(key, self._ttl_seconds)
|
|
|
|
async def get_messages(self, session_id: str, limit: int | None = None, offset: int = 0) -> list[Message]:
|
|
redis = await self._get_redis()
|
|
key = self._messages_key(session_id)
|
|
# Use LRANGE for offset-based pagination
|
|
# Redis list indices: 0-based, -1 = last element
|
|
start = offset
|
|
if limit is not None:
|
|
end = offset + limit - 1
|
|
else:
|
|
end = -1
|
|
raw_list = await redis.lrange(key, start, end)
|
|
return [Message.from_dict(json.loads(raw)) for raw in raw_list]
|
|
|
|
async def count_messages(self, session_id: str) -> int:
|
|
redis = await self._get_redis()
|
|
return await redis.llen(self._messages_key(session_id))
|
|
|
|
async def health_check(self) -> bool:
|
|
try:
|
|
redis = await self._get_redis()
|
|
return await redis.ping()
|
|
except (ImportError, OSError, _RedisError, ValueError, KeyError, RuntimeError):
|
|
return False
|
|
|
|
|
|
# Needed for from_dict deserialization
|
|
from datetime import datetime, timezone # noqa: E402
|
|
|
|
|
|
class FileSessionStore:
|
|
"""File-based session store — persists sessions to ~/.agentkit/sessions/.
|
|
|
|
Each session is stored as a JSON file containing both session metadata
|
|
and messages. Suitable for single-user GUI mode without Redis.
|
|
"""
|
|
|
|
def __init__(self, data_dir: str | None = None):
|
|
if data_dir is None:
|
|
data_dir = os.path.expanduser("~/.agentkit/sessions")
|
|
self._data_dir = data_dir
|
|
os.makedirs(self._data_dir, exist_ok=True)
|
|
|
|
def _session_path(self, session_id: str) -> str:
|
|
return os.path.join(self._data_dir, f"{session_id}.json")
|
|
|
|
def _read_session_file(self, session_id: str) -> dict | None:
|
|
path = self._session_path(session_id)
|
|
if not os.path.exists(path):
|
|
return None
|
|
with open(path, encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
def _write_session_file(self, session_id: str, data: dict) -> None:
|
|
path = self._session_path(session_id)
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
|
|
async def save_session(self, session: Session) -> None:
|
|
data = self._read_session_file(session.session_id) or {"messages": []}
|
|
data["session"] = session.to_dict()
|
|
data["session"]["updated_at"] = datetime.now(timezone.utc).isoformat()
|
|
self._write_session_file(session.session_id, data)
|
|
|
|
async def get_session(self, session_id: str) -> Session | None:
|
|
data = self._read_session_file(session_id)
|
|
if data is None:
|
|
return None
|
|
return Session.from_dict(data["session"])
|
|
|
|
async def update_session_status(self, session_id: str, status: SessionStatus) -> Session | None:
|
|
data = self._read_session_file(session_id)
|
|
if data is None:
|
|
return None
|
|
data["session"]["status"] = status.value
|
|
data["session"]["updated_at"] = datetime.now(timezone.utc).isoformat()
|
|
self._write_session_file(session_id, data)
|
|
return Session.from_dict(data["session"])
|
|
|
|
async def delete_session(self, session_id: str) -> bool:
|
|
path = self._session_path(session_id)
|
|
if os.path.exists(path):
|
|
os.remove(path)
|
|
return True
|
|
return False
|
|
|
|
async def list_sessions(self, agent_name: str | None = None, limit: int = 100) -> list[Session]:
|
|
sessions: list[Session] = []
|
|
for fname in os.listdir(self._data_dir):
|
|
if not fname.endswith(".json"):
|
|
continue
|
|
path = os.path.join(self._data_dir, fname)
|
|
try:
|
|
with open(path, encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
session = Session.from_dict(data["session"])
|
|
if agent_name is None or session.agent_name == agent_name:
|
|
sessions.append(session)
|
|
except (ValueError, KeyError, TypeError, OSError):
|
|
continue
|
|
sessions.sort(key=lambda s: s.updated_at, reverse=True)
|
|
return sessions[:limit]
|
|
|
|
async def append_message(self, message: Message) -> None:
|
|
data = self._read_session_file(message.session_id)
|
|
if data is None:
|
|
data = {"session": {"session_id": message.session_id}, "messages": []}
|
|
data.setdefault("messages", []).append(message.to_dict())
|
|
# Update session timestamp
|
|
if "session" in data:
|
|
data["session"]["updated_at"] = datetime.now(timezone.utc).isoformat()
|
|
self._write_session_file(message.session_id, data)
|
|
|
|
async def get_messages(self, session_id: str, limit: int | None = None, offset: int = 0) -> list[Message]:
|
|
data = self._read_session_file(session_id)
|
|
if data is None:
|
|
return []
|
|
msgs = data.get("messages", [])[offset:]
|
|
if limit is not None:
|
|
msgs = msgs[:limit]
|
|
return [Message.from_dict(m) for m in msgs]
|
|
|
|
async def count_messages(self, session_id: str) -> int:
|
|
data = self._read_session_file(session_id)
|
|
if data is None:
|
|
return 0
|
|
return len(data.get("messages", []))
|
|
|
|
async def health_check(self) -> bool:
|
|
return os.path.isdir(self._data_dir)
|
|
|
|
|
|
def create_session_store(
|
|
backend: str = "memory",
|
|
redis_url: str = "redis://localhost:6379/0",
|
|
ttl_seconds: int = 86400,
|
|
data_dir: str | None = None,
|
|
) -> InMemorySessionStore | RedisSessionStore | FileSessionStore:
|
|
"""Factory: create a SessionStore backed by memory, file, or Redis.
|
|
|
|
- ``memory``: In-memory (lost on restart)
|
|
- ``file``: JSON files in ``~/.agentkit/sessions/`` (persistent, no deps)
|
|
- ``redis``: Redis-backed (production, requires Redis)
|
|
|
|
Falls back to InMemorySessionStore if Redis is unavailable.
|
|
"""
|
|
if backend == "file":
|
|
store = FileSessionStore(data_dir=data_dir)
|
|
logger.info(f"SessionStore backend: file ({store._data_dir})")
|
|
return store
|
|
|
|
if backend == "redis":
|
|
try:
|
|
import redis.asyncio as aioredis # noqa: F401
|
|
|
|
store = RedisSessionStore(redis_url=redis_url, ttl_seconds=ttl_seconds)
|
|
logger.info(f"SessionStore backend: redis")
|
|
return store
|
|
except (ImportError, OSError, _RedisError, ValueError, KeyError, RuntimeError) as exc:
|
|
logger.warning(f"Failed to initialise RedisSessionStore ({exc}), falling back to InMemorySessionStore")
|
|
|
|
store = InMemorySessionStore()
|
|
logger.info("SessionStore backend: memory")
|
|
return store
|