"""SQLite-backed command/agent mode storage.""" from __future__ import annotations import os import json import sqlite3 import threading import time from pathlib import Path from typing import Optional, Union DEFAULT_DB_PATH = Path(__file__).resolve().parents[1] / "data" / "agent.db" _VALID_MODES = frozenset(("command", "agent", "chat")) class ModeStore: """Persist conversation mode with private and group scopes kept separate. A private conversation is keyed by ``(channel, uid)``. A group conversation is keyed by ``(channel, group_id)`` and therefore shared by its members. Every operation uses a short-lived SQLite connection, making one instance safe to share between the event loop and worker threads. """ def __init__(self, db_path: Optional[Union[str, os.PathLike]] = None) -> None: self.db_path = Path(db_path) if db_path is not None else DEFAULT_DB_PATH self.db_path.parent.mkdir(parents=True, exist_ok=True) self._lock = threading.RLock() with self._connect() as conn: conn.execute( """ CREATE TABLE IF NOT EXISTS agent_modes ( channel TEXT NOT NULL, scope TEXT NOT NULL, subject_id TEXT NOT NULL, mode TEXT NOT NULL CHECK(mode IN ('command', 'agent', 'chat')), updated_at REAL NOT NULL, PRIMARY KEY (channel, scope, subject_id) ) """ ) def _connect(self) -> sqlite3.Connection: conn = sqlite3.connect(str(self.db_path), timeout=10) conn.execute("PRAGMA busy_timeout = 10000") return conn @staticmethod def _key(channel: str, uid: int, group_id: Optional[int]) -> tuple: if not isinstance(channel, str) or not channel.strip(): raise ValueError("channel must be a non-empty string") if group_id is not None: return channel.strip(), "group", str(group_id) return channel.strip(), "private", str(uid) def _configured_default(self, channel: str, uid: int, group_id: Optional[int]) -> str: """Return the settings.json default without overriding saved preferences. Example: {"chat_mode": {"default": "command", "admin": "chat", "group": "command"}}. Administrator chat is intended for private conversations; groups stay in command unless explicitly configured. """ if group_id is not None: scope = "group" else: try: import identity scope = "admin" if identity.is_admin(uid) else "default" except Exception: scope = "default" try: settings_path = Path(__file__).resolve().parents[1] / "settings.json" with settings_path.open("r", encoding="utf-8") as fh: config = json.load(fh).get("chat_mode", {}) mode = config.get(scope, config.get("default", "command")) return mode if mode in ("chat", "agent") else "command" except (OSError, ValueError, TypeError, json.JSONDecodeError): return "command" def get_mode(self, channel: str, uid: int, group_id: Optional[int] = None) -> str: """Return saved mode, then the config default; ultimate fallback is command.""" key = self._key(channel, uid, group_id) with self._lock, self._connect() as conn: row = conn.execute( "SELECT mode FROM agent_modes WHERE channel=? AND scope=? AND subject_id=?", key, ).fetchone() if row and row[0] in _VALID_MODES: return row[0] return self._configured_default(channel, uid, group_id) def set_mode( self, channel: str, uid: int, mode: str, group_id: Optional[int] = None, ) -> str: """Set and return the mode for a private user or an entire group.""" if mode not in _VALID_MODES: raise ValueError("mode must be 'command' or 'chat'") key = self._key(channel, uid, group_id) with self._lock, self._connect() as conn: conn.execute( """ INSERT INTO agent_modes(channel, scope, subject_id, mode, updated_at) VALUES (?, ?, ?, ?, ?) ON CONFLICT(channel, scope, subject_id) DO UPDATE SET mode=excluded.mode, updated_at=excluded.updated_at """, (*key, mode, time.time()), ) return mode _default_store: Optional[ModeStore] = None _default_lock = threading.Lock() def _default() -> ModeStore: global _default_store if _default_store is None: with _default_lock: if _default_store is None: _default_store = ModeStore() return _default_store def get_mode(channel: str, uid: int, group_id: Optional[int] = None) -> str: """Module-level convenience wrapper using ``bot-utility/data/agent.db``.""" return _default().get_mode(channel, uid, group_id) def set_mode( channel: str, uid: int, mode: str, group_id: Optional[int] = None, ) -> str: """Module-level convenience wrapper using ``bot-utility/data/agent.db``.""" return _default().set_mode(channel, uid, mode, group_id)