session_helpers.py

"""
Session Helper Functions for Per-User Isolation
================================================

Helper functions that wrap the original app functions to provide
per-user session isolation using the UserSessionManager.

These wrappers intercept gr.Request to identify the user and 
manage their isolated session data.
"""

import gradio as gr
from typing import Any, Dict, List, Tuple
from user_session_manager import session_manager, get_username_from_request, SessionKeys
import logging

logger = logging.getLogger(__name__)


def get_user_simple_chat_history(request: gr.Request) -> List:
    """Get simple chat history for current user."""
    username = get_username_from_request(request)
    return session_manager.get_user_data(username, SessionKeys.SIMPLE_CHAT_HISTORY, default=[])


def set_user_simple_chat_history(request: gr.Request, history: List) -> None:
    """Set simple chat history for current user."""
    username = get_username_from_request(request)
    session_manager.set_user_data(username, SessionKeys.SIMPLE_CHAT_HISTORY, history)


def get_user_builder_chat_histories(request: gr.Request) -> Dict:
    """Get builder chat histories for current user."""
    username = get_username_from_request(request)
    return session_manager.get_user_data(username, SessionKeys.BUILDER_CHAT_HISTORIES, default={})


def set_user_builder_chat_histories(request: gr.Request, histories: Dict) -> None:
    """Set builder chat histories for current user."""
    username = get_username_from_request(request)
    session_manager.set_user_data(username, SessionKeys.BUILDER_CHAT_HISTORIES, histories)


def get_user_deployed_chat_histories(request: gr.Request) -> Dict:
    """Get deployed chat histories for current user."""
    username = get_username_from_request(request)
    return session_manager.get_user_data(username, SessionKeys.DEPLOYED_CHAT_HISTORIES, default={})


def set_user_deployed_chat_histories(request: gr.Request, histories: Dict) -> None:
    """Set deployed chat histories for current user."""
    username = get_username_from_request(request)
    session_manager.set_user_data(username, SessionKeys.DEPLOYED_CHAT_HISTORIES, histories)


def get_user_active_children(request: gr.Request) -> List:
    """Get active children for current user."""
    username = get_username_from_request(request)
    return session_manager.get_user_data(username, SessionKeys.ACTIVE_CHILDREN, default=[])


def set_user_active_children(request: gr.Request, children: List) -> None:
    """Set active children for current user."""
    username = get_username_from_request(request)
    session_manager.set_user_data(username, SessionKeys.ACTIVE_CHILDREN, children)


def get_user_patient_data(request: gr.Request) -> Dict:
    """Get patient data for current user."""
    username = get_username_from_request(request)
    return session_manager.get_user_data(username, SessionKeys.PATIENT_DATA, default={})


def set_user_patient_data(request: gr.Request, data: Dict) -> None:
    """Set patient data for current user."""
    username = get_username_from_request(request)
    session_manager.set_user_data(username, SessionKeys.PATIENT_DATA, data)


def get_user_prefill_flag(request: gr.Request) -> bool:
    """Get prefill flag for current user."""
    username = get_username_from_request(request)
    return session_manager.get_user_data(username, SessionKeys.PREFILL_FLAG, default=False)


def set_user_prefill_flag(request: gr.Request, flag: bool) -> None:
    """Set prefill flag for current user."""
    username = get_username_from_request(request)
    session_manager.set_user_data(username, SessionKeys.PREFILL_FLAG, flag)


def clear_user_session(request: gr.Request) -> None:
    """Clear all session data for current user."""
    username = get_username_from_request(request)
    session_manager.clear_user_data(username)
    logger.info(f"Cleared session for user: {username}")


def get_current_username(request: gr.Request) -> str:
    """Get the current authenticated username."""
    username = get_username_from_request(request)
    logger.debug(f"Current user: {username}")
    return username


# Wrapper functions that maintain compatibility with existing code
# while adding per-user session management

def wrap_simple_chat_response(original_func):
    """Wrap simple_chat_response to use per-user history."""
    def wrapper(user_message: str, request: gr.Request):
        # Get user-specific history
        history = get_user_simple_chat_history(request)
        
        # Call original function
        updated_history, empty_input = original_func(user_message, history)
        
        # Save updated history
        set_user_simple_chat_history(request, updated_history)
        
        return updated_history, empty_input
    
    return wrapper


def wrap_clear_simple_chat(original_func):
    """Wrap clear_simple_chat to use per-user history."""
    def wrapper(request: gr.Request):
        # Call original function with empty history
        empty_history, empty_input = original_func()
        
        # Save to user session
        set_user_simple_chat_history(request, empty_history)
        
        return empty_history, empty_input
    
    return wrapper


def wrap_chatpanel_handle(original_func):
    """Wrap chatpanel_handle to use per-user histories."""
    def wrapper(agent_name: str, user_text: str, request: gr.Request):
        # Get user-specific histories
        histories = get_user_builder_chat_histories(request)
        
        # Call original function
        updated_history, log, empty_input, updated_histories = original_func(
            agent_name, user_text, histories
        )
        
        # Save updated histories
        set_user_builder_chat_histories(request, updated_histories)
        
        return updated_history, log, empty_input, updated_histories
    
    return wrapper


def wrap_load_history(original_func):
    """Wrap load_history to use per-user histories."""
    def wrapper(agent_name: str, request: gr.Request):
        # Get user-specific histories
        histories = get_user_builder_chat_histories(request)
        
        # Call original function
        return original_func(agent_name, histories)
    
    return wrapper


def log_user_access(request: gr.Request, action: str):
    """Log user access for debugging/auditing."""
    username = get_username_from_request(request)
    logger.info(f"User '{username}' performed action: {action}")


if __name__ == "__main__":
    print("Session helpers module loaded successfully")
    print(f"Available session keys: {[k for k in dir(SessionKeys) if not k.startswith('_')]}")