utils/session_manager.py

"""
Session state management utilities
"""

import streamlit as st
from typing import Any, Dict, Optional

class SessionManager:
    """Manage session state with better error handling"""
    
    @staticmethod
    def init_session_state():
        """Initialize session state with default values"""
        defaults = {
            "results": [],
            "jobs": [],
            "job_idx": 0,
            "running": False,
            "api_key_val": "",
            "last_error": None,
            "retry_count": 0,
            "backend_loaded": False,
            "backend": None
        }
        
        for key, value in defaults.items():
            if key not in st.session_state:
                st.session_state[key] = value
    
    @staticmethod
    def safe_update(updates: Dict[str, Any]):
        """Safely update session state"""
        try:
            st.session_state.update(updates)
        except Exception as e:
            st.error(f"Session update failed: {e}")
    
    @staticmethod
    def reset_job_state():
        """Reset job-related state"""
        SessionManager.safe_update({
            "results": [],
            "jobs": [],
            "job_idx": 0,
            "running": False,
            "last_error": None,
            "retry_count": 0
        })
    
    @staticmethod
    def get_safe(key: str, default: Any = None) -> Any:
        """Safely get session state value"""
        return st.session_state.get(key, default)