user_session_manager.py

"""
User Session Manager for ID Agents
===================================

Manages isolated user sessions so each authenticated user has their own
chat histories, agents, and application state. This prevents users from
seeing or interfering with each other's work.

Usage:
    from user_session_manager import session_manager
    
    # In a Gradio function with request parameter:
    def my_function(user_input, request: gr.Request):
        username = request.username
        
        # Get user-specific data
        user_data = session_manager.get_user_data(username, "my_key", default_value=[])
        
        # Update user-specific data
        session_manager.set_user_data(username, "my_key", new_value)
"""

import threading
from typing import Dict, Any, Optional
from collections import defaultdict
import json
import logging

logger = logging.getLogger(__name__)


class UserSessionManager:
    """
    Thread-safe manager for user-specific session data.
    Each authenticated user gets their own isolated storage space.
    """
    
    def __init__(self):
        self._sessions: Dict[str, Dict[str, Any]] = defaultdict(dict)
        self._lock = threading.Lock()
        logger.info("UserSessionManager initialized")
    
    def get_user_data(self, username: str, key: str, default: Any = None) -> Any:
        """
        Get user-specific data by key.
        
        Args:
            username: The authenticated username
            key: The data key (e.g., 'chat_history', 'agents', etc.)
            default: Default value if key doesn't exist
            
        Returns:
            The stored value or default
        """
        with self._lock:
            if username not in self._sessions:
                self._sessions[username] = {}
            return self._sessions[username].get(key, default)
    
    def set_user_data(self, username: str, key: str, value: Any) -> None:
        """
        Set user-specific data.
        
        Args:
            username: The authenticated username
            key: The data key
            value: The value to store
        """
        with self._lock:
            if username not in self._sessions:
                self._sessions[username] = {}
            self._sessions[username][key] = value
            logger.debug(f"Set {key} for user {username}")
    
    def update_user_data(self, username: str, key: str, update_func: callable) -> Any:
        """
        Atomically update user-specific data using a function.
        
        Args:
            username: The authenticated username
            key: The data key
            update_func: Function that takes current value and returns new value
            
        Returns:
            The updated value
        """
        with self._lock:
            if username not in self._sessions:
                self._sessions[username] = {}
            
            current = self._sessions[username].get(key)
            updated = update_func(current)
            self._sessions[username][key] = updated
            return updated
    
    def clear_user_data(self, username: str, key: Optional[str] = None) -> None:
        """
        Clear user-specific data.
        
        Args:
            username: The authenticated username
            key: Specific key to clear, or None to clear all user data
        """
        with self._lock:
            if username not in self._sessions:
                return
            
            if key is None:
                # Clear all data for this user
                self._sessions[username].clear()
                logger.info(f"Cleared all data for user {username}")
            else:
                # Clear specific key
                if key in self._sessions[username]:
                    del self._sessions[username][key]
                    logger.debug(f"Cleared {key} for user {username}")
    
    def get_all_user_keys(self, username: str) -> list:
        """Get all keys stored for a user."""
        with self._lock:
            if username not in self._sessions:
                return []
            return list(self._sessions[username].keys())
    
    def user_exists(self, username: str) -> bool:
        """Check if user has any session data."""
        with self._lock:
            return username in self._sessions and bool(self._sessions[username])
    
    def get_active_users(self) -> list:
        """Get list of users with active sessions."""
        with self._lock:
            return [u for u, data in self._sessions.items() if data]
    
    def get_user_stats(self, username: str) -> Dict[str, Any]:
        """Get statistics about user's session."""
        with self._lock:
            if username not in self._sessions:
                return {"exists": False}
            
            data = self._sessions[username]
            return {
                "exists": True,
                "keys": list(data.keys()),
                "data_size": len(str(data))
            }


# Global singleton instance
session_manager = UserSessionManager()


def get_username_from_request(request: Any) -> str:
    """
    Extract username from Gradio request object.
    
    Args:
        request: Gradio gr.Request object
        
    Returns:
        Username string, or "anonymous" if not authenticated
    """
    if request is None:
        return "anonymous"
    
    # Gradio stores username in request.username after basic auth
    username = getattr(request, "username", None)
    
    if username is None or username == "":
        return "anonymous"
    
    return str(username)


# Session data keys (constants for consistency)
class SessionKeys:
    """Constants for session data keys."""
    SIMPLE_CHAT_HISTORY = "simple_chat_history"
    BUILDER_CHAT_HISTORIES = "builder_chat_histories"
    DEPLOYED_CHAT_HISTORIES = "deployed_chat_histories"
    ACTIVE_CHILDREN = "active_children"
    PATIENT_DATA = "patient_data"
    AGENT_OUTPUT = "agent_output"
    PREFILL_FLAG = "prefill_flag"


if __name__ == "__main__":
    # Simple test
    print("Testing UserSessionManager...")
    
    # Test basic operations
    session_manager.set_user_data("user1", "chat_history", [["Hi", "Hello"]])
    session_manager.set_user_data("user2", "chat_history", [["Hey", "Hi there"]])
    
    print("User1 chat:", session_manager.get_user_data("user1", "chat_history"))
    print("User2 chat:", session_manager.get_user_data("user2", "chat_history"))
    
    print("Active users:", session_manager.get_active_users())
    print("User1 stats:", session_manager.get_user_stats("user1"))
    
    print("✅ UserSessionManager test complete")