session.py

"""Session persistence — save, load, and query agent conversation sessions.

Each session is stored as a single JSON file under a sessions directory.
A lightweight index file provides fast metadata queries without loading
every session file.  The index is a cache; ``rebuild_index()`` recreates
it from the individual session files.
"""

import json
import secrets
import threading
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from config import get_logger

logger = get_logger(__name__)


class SessionManager:
    """Manages persistent conversation sessions.

    Each session is stored as ``session_<id>.json``.  An ``index.json``
    file caches metadata for fast listing.

    Usage::

        mgr = SessionManager(Path(".sessions"))
        sid = mgr.init()
        mgr.save(messages)
        mgr.close()
    """

    def __init__(self, sessions_dir: Path) -> None:
        self.sessions_dir = sessions_dir
        self._current_session_id: Optional[str] = None
        self._lock = threading.Lock()

    # ── Public API ─────────────────────────────────────────────────────

    def init(self) -> str:
        """Create a new session and set it as the current session.

        Idempotent: if a session is already active, returns its ID.
        Call ``close()`` before ``init()`` when starting fresh.
        """
        with self._lock:
            if self._current_session_id is not None:
                return self._current_session_id

            session_id = _generate_session_id()
            now = datetime.now(timezone.utc).isoformat()

            session_data = {
                "session_id": session_id,
                "created_at": now,
                "updated_at": now,
                "message_count": 0,
                "title": "New session",
                "workdir": str(Path.cwd()),
                "is_active": True,
                "messages": [],
            }

            self.sessions_dir.mkdir(parents=True, exist_ok=True)
            self._write_session_file(session_id, session_data)
            self._append_index_entry(session_data)

            self._current_session_id = session_id
            logger.info(f"[session] created {session_id}")
            return session_id

    def save(self, messages: list) -> None:
        """Save the current session with the latest messages and metadata.

        Thread-safe.  Safe to call after every turn.  No-op if no session
        is active.
        """
        session_id = self._current_session_id
        if session_id is None:
            logger.warning("[session] save_session called with no active session")
            return

        with self._lock:
            now = datetime.now(timezone.utc).isoformat()
            title = _extract_title(messages)

            session_data = {
                "session_id": session_id,
                "created_at": self._read_created_at(session_id) or now,
                "updated_at": now,
                "message_count": len(messages),
                "title": title,
                "workdir": str(Path.cwd()),
                "is_active": True,
                "messages": messages,
            }

            self._write_session_file(session_id, session_data)
            self._upsert_index_entry(session_data)
            logger.debug(f"[session] saved {session_id} ({len(messages)} messages)")

    def load(self, session_id: str) -> Optional[list]:
        """Load the full message list for a session.

        Returns ``None`` if the session file doesn't exist.
        """
        path = self._session_path(session_id)
        with self._lock:
            if not path.exists():
                logger.warning(f"[session] session not found: {session_id}")
                return None
            try:
                data = json.loads(path.read_text())
                return data.get("messages")
            except (json.JSONDecodeError, OSError) as exc:
                logger.error(f"[session] failed to load {session_id}: {exc}")
                return None

    def close(self) -> None:
        """Mark the current session as inactive and save final state.

        No-op if no session is active.
        """
        session_id = self._current_session_id
        if session_id is None:
            return

        with self._lock:
            path = self._session_path(session_id)
            if path.exists():
                try:
                    data = json.loads(path.read_text())
                except (json.JSONDecodeError, OSError):
                    data = {}
                data["is_active"] = False
                data["updated_at"] = datetime.now(timezone.utc).isoformat()
                self._write_session_file(session_id, data)
                self._upsert_index_entry(data)

            self._current_session_id = None
            logger.info(f"[session] closed {session_id}")

    def list_sessions(self) -> list[dict]:
        """Return metadata for all sessions, newest-first.

        Fast path: reads the index file.  Falls back to scanning the
        sessions directory if the index is missing.
        """
        with self._lock:
            index = self._read_index()
            if index is not None:
                return _sort_and_dedup_index(index)
            return self._scan_directory()

    def resume(self, session_id: str) -> None:
        """Switch the current session to an existing session by ID.

        Subsequent ``save()`` calls will write to the resumed session.
        """
        with self._lock:
            path = self._session_path(session_id)
            if not path.exists():
                logger.warning(
                    f"[session] cannot resume non-existent session: {session_id}"
                )
                return
            self._current_session_id = session_id
            logger.info(f"[session] resumed {session_id}")

    def get_current_id(self) -> Optional[str]:
        """Return the current session ID, or ``None``."""
        return self._current_session_id

    def rebuild_index(self) -> None:
        """Scan the sessions directory and rebuild the index file from scratch."""
        with self._lock:
            entries = self._scan_directory()
            self._write_index(entries)
            logger.info(f"[session] rebuilt index with {len(entries)} entries")

    # ── Internal helpers ────────────────────────────────────────────────

    def _session_path(self, session_id: str) -> Path:
        return self.sessions_dir / f"session_{session_id}.json"

    def _index_path(self) -> Path:
        return self.sessions_dir / "index.json"

    def _write_session_file(self, session_id: str, data: dict) -> None:
        path = self._session_path(session_id)
        path.write_text(json.dumps(data, indent=2, default=str))

    def _read_created_at(self, session_id: str) -> Optional[str]:
        path = self._session_path(session_id)
        if path.exists():
            try:
                return json.loads(path.read_text()).get("created_at")
            except (json.JSONDecodeError, OSError):
                return None
        return None

    def _read_index(self) -> Optional[list]:
        path = self._index_path()
        if not path.exists():
            return None
        try:
            data = json.loads(path.read_text())
            return data.get("sessions", [])
        except (json.JSONDecodeError, OSError):
            return None

    def _write_index(self, entries: list) -> None:
        path = self._index_path()
        path.write_text(json.dumps({"sessions": entries}, indent=2))

    def _append_index_entry(self, entry: dict) -> None:
        entry_for_index = {k: v for k, v in entry.items() if k != "messages"}
        entries = self._read_index() or []
        entries.append(entry_for_index)
        self._write_index(entries)

    def _upsert_index_entry(self, entry: dict) -> None:
        entry_for_index = {k: v for k, v in entry.items() if k != "messages"}
        entries = self._read_index() or []
        entries = [e for e in entries if e.get("session_id") != entry["session_id"]]
        entries.append(entry_for_index)
        self._write_index(entries)

    def _scan_directory(self) -> list:
        """Scan sessions directory for session files and extract metadata."""
        if not self.sessions_dir.exists():
            return []
        entries = []
        for f in sorted(self.sessions_dir.glob("session_*.json")):
            try:
                data = json.loads(f.read_text())
                entry = {k: v for k, v in data.items() if k != "messages"}
                entries.append(entry)
            except (json.JSONDecodeError, OSError) as exc:
                logger.warning(f"[session] failed to read {f.name}: {exc}")
                continue
        entries.sort(key=lambda x: x.get("created_at", ""), reverse=True)
        return entries


# ── Module-level shims (delegate to a default manager) ───────────────────

_default_mgr: SessionManager | None = None


def _get_default() -> SessionManager:
    global _default_mgr
    if _default_mgr is None:
        from config import SESSIONS_DIR
        _default_mgr = SessionManager(SESSIONS_DIR)
    return _default_mgr


def init_session() -> str:
    return _get_default().init()


def save_session(messages: list) -> None:
    _get_default().save(messages)


def load_session(session_id: str) -> Optional[list]:
    return _get_default().load(session_id)


def close_session() -> None:
    _get_default().close()


def list_sessions() -> list[dict]:
    return _get_default().list_sessions()


def get_current_session_id() -> Optional[str]:
    return _get_default().get_current_id()


def resume_session(session_id: str) -> None:
    _get_default().resume(session_id)


# ── Helpers (module-level, stateless) ────────────────────────────────────


def _generate_session_id() -> str:
    """Generate a unique session ID: ``YYYYMMDD_HHMMSS_XXXX``."""
    now = datetime.now()
    date_part = now.strftime("%Y%m%d_%H%M%S")
    rand_part = secrets.token_hex(2)
    return f"{date_part}_{rand_part}"


def _extract_title(messages: list) -> str:
    """Extract title from first non-system, non-tool user message."""
    for m in messages:
        if m.get("role") == "user" and isinstance(m.get("content"), str):
            text = m["content"].strip()
            if text:
                return text[:80] + ("..." if len(text) > 80 else "")
    return "New session"


def _sort_and_dedup_index(entries: list) -> list:
    """Sort newest-first.  Remove duplicate session_ids (keep last occurrence)."""
    seen = {}
    for e in entries:
        sid = e.get("session_id")
        if sid:
            seen[sid] = e
    deduped = list(seen.values())
    deduped.sort(key=lambda x: x.get("created_at", ""), reverse=True)
    return deduped