""" utils/session_manager.py — Persistent session save/load Stores active symbols and their "I own this stock" state to a JSON file. Supports per-user sessions stored in utils/sessions/{username}.json. On HF Spaces each user's file is additionally synced to a private HF dataset. """ import json import logging import re import shutil import threading from pathlib import Path from typing import Dict, List, Tuple from datetime import datetime, timezone from utils.config import UTILS_DIR, IS_HF_SPACE, HF_TOKEN, HF_USER logger = logging.getLogger(__name__) _HF_DATASET_REPO = f"{HF_USER}/StocksAnalysisDashboard_dataset" SESSIONS_DIR = UTILS_DIR / "sessions" SESSIONS_DIR.mkdir(exist_ok=True) _USERNAME_RE = re.compile(r'^[\w\-]{1,32}$') DEFAULT_USER = "Default User" _DEFAULT_FILE_STEM = "_default_" # ─── Migrate legacy anonymous session (.json → _default_.json) ─────────────── _legacy = SESSIONS_DIR / ".json" if _legacy.exists(): _legacy.rename(SESSIONS_DIR / f"{_DEFAULT_FILE_STEM}.json") # ─── Ensure Default User file always exists (first run / fresh HF Space) ───── _default_file = SESSIONS_DIR / f"{_DEFAULT_FILE_STEM}.json" if not _default_file.exists(): try: with open(_default_file, "w") as _f: json.dump({ "version": "1.0", "saved_at": None, "symbols": [], "owned": {}, "watchlist": [], "refresh_interval": "Off", "snapshots": {}, }, _f, indent=2) except Exception: pass # ─── Helpers ────────────────────────────────────────────────────────────────── def _user_file(username: str) -> Path: if not username or username == DEFAULT_USER: return SESSIONS_DIR / f"{_DEFAULT_FILE_STEM}.json" return SESSIONS_DIR / f"{username}.json" def _hf_user_path(username: str) -> str: return f"sessions/{username}.json" # ─── HF Hub sync (per-user) ─────────────────────────────────────────────────── def _hf_pull_all_users() -> None: """On HF Space startup: pull ALL session files from the dataset so list_users() works. HF Space wipes the local filesystem on every redeploy, so without this every user except Default User disappears after a deploy.""" if not IS_HF_SPACE or not HF_TOKEN: return try: from huggingface_hub import HfApi, hf_hub_download api = HfApi(token=HF_TOKEN) try: all_files = list(api.list_repo_files( repo_id=_HF_DATASET_REPO, repo_type="dataset", token=HF_TOKEN )) except Exception as e: logger.warning(f"HF pull-all-users: could not list repo files: {e}") return session_files = [f for f in all_files if f.startswith("sessions/") and f.endswith(".json")] for hf_path in session_files: stem = Path(hf_path).stem local = SESSIONS_DIR / f"{stem}.json" try: cached = hf_hub_download( repo_id=_HF_DATASET_REPO, filename=hf_path, repo_type="dataset", token=HF_TOKEN, force_download=True, ) shutil.copy2(cached, local) logger.info(f"Startup pull: '{stem}' restored from HF dataset") except Exception as e: logger.warning(f"Startup pull: could not restore '{stem}': {e}") except Exception as e: logger.warning(f"HF pull-all-users failed: {e}") def _hf_pull_user(username: str) -> bool: if not IS_HF_SPACE or not HF_TOKEN: return False import concurrent.futures def _do_pull(): from huggingface_hub import hf_hub_download cached = hf_hub_download( repo_id=_HF_DATASET_REPO, filename=_hf_user_path(username), repo_type="dataset", token=HF_TOKEN, force_download=True, ) shutil.copy2(cached, _user_file(username)) try: with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex: ex.submit(_do_pull).result(timeout=5) logger.info(f"Session pulled for '{username}' from HF Hub") return True except concurrent.futures.TimeoutError: logger.warning(f"HF pull timed out for '{username}', using local file") return False except Exception as e: logger.warning(f"Could not pull session for '{username}' from HF Hub: {e}") return False def _hf_push_user(username: str) -> None: if not IS_HF_SPACE or not HF_TOKEN: return ufile = _user_file(username) hf_path = _hf_user_path(username) def _push(): try: from huggingface_hub import HfApi api = HfApi(token=HF_TOKEN) api.create_repo(repo_id=_HF_DATASET_REPO, repo_type="dataset", private=True, exist_ok=True) api.upload_file(path_or_fileobj=str(ufile), path_in_repo=hf_path, repo_id=_HF_DATASET_REPO, repo_type="dataset") logger.info(f"Session pushed for '{username}' to HF Hub") except Exception as e: logger.warning(f"Could not push session for '{username}' to HF Hub: {e}") threading.Thread(target=_push, daemon=True).start() # ─── Public API ─────────────────────────────────────────────────────────────── def list_users() -> List[str]: """Return sorted list of existing usernames (DEFAULT_USER listed first).""" names = [] for p in SESSIONS_DIR.glob("*.json"): names.append(DEFAULT_USER if p.stem == _DEFAULT_FILE_STEM else p.stem) others = sorted(n for n in names if n != DEFAULT_USER) return ([DEFAULT_USER] if DEFAULT_USER in names else []) + others def create_user(username: str) -> Tuple[bool, str]: """Create a new empty session for username. Returns (ok, error_message).""" username = username.strip() if not username: return False, "Username cannot be empty." if username == DEFAULT_USER: return False, f"'{DEFAULT_USER}' is a reserved name." if not _USERNAME_RE.match(username): return False, "Use only letters, numbers, underscores, hyphens (max 32 chars)." ufile = _user_file(username) if ufile.exists(): return False, f"User '{username}' already exists." try: with open(ufile, "w") as f: json.dump(_empty_session(), f, indent=2) logger.info(f"Created user '{username}'") return True, "" except Exception as e: return False, str(e) def delete_user(username: str) -> Tuple[bool, str]: """Delete a user's session file. Returns (ok, error_message).""" if username == DEFAULT_USER: return False, f"'{DEFAULT_USER}' cannot be deleted." ufile = _user_file(username) if not ufile.exists(): return False, f"User '{username}' not found." try: ufile.unlink() logger.info(f"Deleted user '{username}'") return True, "" except Exception as e: return False, str(e) def rename_user(old_name: str, new_name: str) -> Tuple[bool, str]: """Rename a user. Returns (ok, error_message).""" if old_name == DEFAULT_USER: return False, f"'{DEFAULT_USER}' cannot be renamed." new_name = new_name.strip() if not new_name: return False, "New username cannot be empty." if new_name == DEFAULT_USER: return False, f"'{DEFAULT_USER}' is a reserved name." if not _USERNAME_RE.match(new_name): return False, "Use only letters, numbers, underscores, hyphens (max 32 chars)." old_file = _user_file(old_name) new_file = _user_file(new_name) if not old_file.exists(): return False, f"User '{old_name}' not found." if new_file.exists(): return False, f"User '{new_name}' already exists." try: old_file.rename(new_file) logger.info(f"Renamed user '{old_name}' → '{new_name}'") return True, "" except Exception as e: return False, str(e) def load_session(username: str) -> Dict: """Load saved session for the given username.""" if IS_HF_SPACE: _hf_pull_user(username) ufile = _user_file(username) if ufile.exists(): try: with open(ufile, "r") as f: data = json.load(f) logger.info(f"Session loaded for '{username}': {len(data.get('symbols', []))} symbols") return data except Exception as e: logger.warning(f"Could not load session for '{username}': {e}") return _empty_session() def save_session(symbols: List[str], owned: Dict[str, bool], watchlist: List[str] = None, refresh_interval: str = "Off", snapshots: Dict[str, dict] = None, username: str = "") -> Tuple[bool, str]: """Persist dashboard state for the given username.""" # Default User is read-only — never persist changes if username == DEFAULT_USER: return True, "" data = { "version": "1.0", "saved_at": datetime.now(timezone.utc).isoformat(), "symbols": symbols, "owned": owned, "watchlist": watchlist or [], "refresh_interval": refresh_interval, "snapshots": {}, } ufile = _user_file(username) try: with open(ufile, "w") as f: json.dump(data, f, indent=2) logger.info(f"Session saved for '{username}': {symbols}") _hf_push_user(username) return True, "" except Exception as e: logger.error(f"Could not save session for '{username}': {e}") return False, str(e) def _empty_session() -> Dict: return { "version": "1.0", "saved_at": None, "symbols": [], "owned": {}, "watchlist": [], "refresh_interval": "Off", "snapshots": {}, } # ─── HF Space startup: restore all user sessions ───────────────────────────── # HF Space wipes the local filesystem on every redeploy. Session files saved # by users in prior runs survive only in the private HF dataset. Pulling them # all here means list_users() finds them immediately on the first page load. _hf_pull_all_users()