"""SQLite-backed external-conversation to Agent-session mapping.""" from __future__ import annotations import json import sqlite3 import threading from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Any, Dict, Optional, Protocol, Union class SessionStoreError(RuntimeError): """Raised when persisted session state is invalid or unavailable.""" @dataclass(frozen=True) class SessionRecord: conversation_id: str session_id: str external_subject_id: Optional[str] metadata: Dict[str, Any] created_at: str updated_at: str def to_dict(self) -> Dict[str, Any]: return { "conversation_id": self.conversation_id, "session_id": self.session_id, "external_subject_id": self.external_subject_id, "metadata": self.metadata, "created_at": self.created_at, "updated_at": self.updated_at, } class SessionStore(Protocol): """Storage contract that can also be implemented by a shared database adapter.""" def get(self, conversation_id: str) -> Optional[SessionRecord]: ... def put( self, *, conversation_id: str, session_id: str, external_subject_id: Optional[str], metadata: Optional[Dict[str, Any]] = None, ) -> SessionRecord: ... def delete(self, conversation_id: str) -> bool: ... class SQLiteSessionStore: """Small thread-safe SQLite store suitable for a single application instance.""" def __init__(self, path: Union[str, Path]) -> None: self.path = Path(path) self._lock = threading.RLock() try: if str(self.path) != ":memory:": self.path.parent.mkdir(parents=True, exist_ok=True) self._connection = sqlite3.connect( str(self.path), timeout=5.0, check_same_thread=False, ) self._connection.row_factory = sqlite3.Row self._connection.execute("PRAGMA busy_timeout = 5000") self._initialize() except SessionStoreError: connection = getattr(self, "_connection", None) if connection is not None: connection.close() raise except (OSError, sqlite3.Error) as exc: connection = getattr(self, "_connection", None) if connection is not None: connection.close() raise SessionStoreError(f"cannot initialize session database: {exc}") from exc def __enter__(self) -> "SQLiteSessionStore": return self def __exit__(self, exc_type: object, exc: object, traceback: object) -> None: self.close() def close(self) -> None: try: with self._lock: self._connection.close() except sqlite3.Error as exc: raise SessionStoreError(f"cannot close session database: {exc}") from exc def _initialize(self) -> None: try: with self._lock, self._connection: self._connection.execute( """ CREATE TABLE IF NOT EXISTS agent_sessions ( conversation_id TEXT PRIMARY KEY NOT NULL, session_id TEXT NOT NULL, external_subject_id TEXT, metadata_json TEXT NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """ ) except sqlite3.Error as exc: raise SessionStoreError(f"cannot create session database schema: {exc}") from exc def get(self, conversation_id: str) -> Optional[SessionRecord]: conversation_id = self._required("conversation_id", conversation_id) try: with self._lock: row = self._connection.execute( "SELECT * FROM agent_sessions WHERE conversation_id = ?", (conversation_id,), ).fetchone() except sqlite3.Error as exc: raise SessionStoreError(f"cannot read session mapping: {exc}") from exc return self._record_from_row(row) if row is not None else None def put( self, *, conversation_id: str, session_id: str, external_subject_id: Optional[str], metadata: Optional[Dict[str, Any]] = None, ) -> SessionRecord: conversation_id = self._required("conversation_id", conversation_id) session_id = self._required("session_id", session_id) metadata_value = dict(metadata or {}) try: metadata_json = json.dumps( metadata_value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), ) except (TypeError, ValueError) as exc: raise SessionStoreError("session metadata must be JSON serializable") from exc now = datetime.now(timezone.utc).isoformat() try: with self._lock, self._connection: self._connection.execute( """ INSERT INTO agent_sessions ( conversation_id, session_id, external_subject_id, metadata_json, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT(conversation_id) DO UPDATE SET session_id = excluded.session_id, external_subject_id = excluded.external_subject_id, metadata_json = excluded.metadata_json, updated_at = excluded.updated_at """, ( conversation_id, session_id, external_subject_id, metadata_json, now, now, ), ) row = self._connection.execute( "SELECT * FROM agent_sessions WHERE conversation_id = ?", (conversation_id,), ).fetchone() except sqlite3.Error as exc: raise SessionStoreError(f"cannot persist session mapping: {exc}") from exc if row is None: raise SessionStoreError("session mapping was not persisted") return self._record_from_row(row) def delete(self, conversation_id: str) -> bool: conversation_id = self._required("conversation_id", conversation_id) try: with self._lock, self._connection: cursor = self._connection.execute( "DELETE FROM agent_sessions WHERE conversation_id = ?", (conversation_id,), ) except sqlite3.Error as exc: raise SessionStoreError(f"cannot delete session mapping: {exc}") from exc return cursor.rowcount > 0 @staticmethod def _required(name: str, value: str) -> str: normalized = str(value).strip() if not normalized: raise ValueError(f"{name} must not be empty") return normalized @staticmethod def _record_from_row(row: sqlite3.Row) -> SessionRecord: try: metadata = json.loads(row["metadata_json"]) except (TypeError, ValueError) as exc: raise SessionStoreError("persisted session metadata is invalid JSON") from exc if not isinstance(metadata, dict): raise SessionStoreError("persisted session metadata must be an object") return SessionRecord( conversation_id=row["conversation_id"], session_id=row["session_id"], external_subject_id=row["external_subject_id"], metadata=metadata, created_at=row["created_at"], updated_at=row["updated_at"], )