src/kvstore.py

"""
KVStore - persistent key-value cache (disk-backed JSON, TTL + LRU eviction).
Wraps cross-turn conversation context so it survives process restarts.
"""

import json
import threading
import time
from pathlib import Path
from typing import Any, Dict, List, Optional


class KVStore:
    def __init__(self, path: str, max_entries: int = 1000, default_ttl: float = 3600.0):
        self.path = Path(path)
        self.max_entries = max_entries
        self.default_ttl = default_ttl
        self._lock = threading.Lock()
        self._data: Dict[str, Any] = {}
        self._meta: Dict[str, Dict] = {}  # key -> {exp, last, size}
        if self.path.exists():
            self._load()

    def _load(self):
        try:
            raw = json.loads(self.path.read_text(encoding="utf-8"))
            self._data = raw.get("data", {})
            self._meta = raw.get("meta", {})
        except (json.JSONDecodeError, OSError):
            self._data, self._meta = {}, {}

    def _persist(self):
        self.path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self.path.with_suffix(".tmp")
        tmp.write_text(
            json.dumps({"data": self._data, "meta": self._meta}),
            encoding="utf-8",
        )
        tmp.replace(self.path)

    def _evict(self):
        now = time.time()
        expired = [k for k, m in self._meta.items() if m.get("exp", 0) < now]
        for k in expired:
            self._data.pop(k, None)
            self._meta.pop(k, None)
        while len(self._data) > self.max_entries:
            oldest = min(self._meta, key=lambda k: self._meta[k].get("last", 0))
            self._data.pop(oldest, None)
            self._meta.pop(oldest, None)

    def get(self, key: str, default: Any = None) -> Any:
        with self._lock:
            m = self._meta.get(key)
            if m is None:
                return default
            if m.get("exp", 0) < time.time():
                self._data.pop(key, None)
                self._meta.pop(key, None)
                return default
            m["last"] = time.time()
            return self._data.get(key, default)

    def set(self, key: str, value: Any, ttl: Optional[float] = None):
        with self._lock:
            exp = time.time() + (ttl if ttl is not None else self.default_ttl)
            self._data[key] = value
            self._meta[key] = {
                "exp": exp,
                "last": time.time(),
                "size": len(json.dumps(value, default=str)),
            }
            self._evict()
            self._persist()

    def delete(self, key: str):
        with self._lock:
            self._data.pop(key, None)
            self._meta.pop(key, None)
            self._persist()

    def keys(self) -> List[str]:
        with self._lock:
            return list(self._data.keys())

    def flush(self):
        with self._lock:
            self._persist()