quick_start_session.py

"""
Quick Start Script for Session Isolation
=========================================

This script demonstrates how the session isolation works and provides
test utilities.

Run this file to verify the session manager is working correctly.
"""

from user_session_manager import session_manager, SessionKeys
from session_helpers import get_current_username
import gradio as gr


def demo_isolated_chat():
    """
    Minimal demo showing per-user chat isolation.
    """
    
    def chat_fn(message, request: gr.Request):
        """Chat function with per-user isolation."""
        # Get username
        username = get_current_username(request)
        
        # Get user's chat history
        history = session_manager.get_user_data(
            username, 
            SessionKeys.SIMPLE_CHAT_HISTORY, 
            default=[]
        )
        
        # Add user message and bot response
        bot_response = f"Hello {username}! You said: {message}"
        history.append([message, bot_response])
        
        # Save back to user's session
        session_manager.set_user_data(
            username,
            SessionKeys.SIMPLE_CHAT_HISTORY,
            history
        )
        
        return history, ""
    
    def clear_fn(request: gr.Request):
        """Clear chat for current user only."""
        username = get_current_username(request)
        session_manager.set_user_data(
            username,
            SessionKeys.SIMPLE_CHAT_HISTORY,
            []
        )
        return [], ""
    
    # Build simple UI
    with gr.Blocks() as demo:
        gr.Markdown("# 🔒 Per-User Session Isolation Demo")
        gr.Markdown("Each authenticated user sees only their own messages!")
        
        chatbot = gr.Chatbot(label="Your Personal Chat")
        msg = gr.Textbox(label="Message", placeholder="Type your message...")
        clear = gr.Button("Clear My Chat")
        
        # Note: No gr.State needed! Everything is per-user
        msg.submit(chat_fn, inputs=[msg], outputs=[chatbot, msg])
        clear.click(clear_fn, inputs=[], outputs=[chatbot, msg])
    
    return demo


def show_session_stats():
    """Display session statistics."""
    print("\n" + "="*50)
    print("SESSION MANAGER STATISTICS")
    print("="*50)
    
    active_users = session_manager.get_active_users()
    print(f"\n📊 Active Users: {len(active_users)}")
    
    for username in active_users:
        stats = session_manager.get_user_stats(username)
        print(f"\n👤 User: {username}")
        print(f"   Keys stored: {stats['keys']}")
        print(f"   Data size: {stats['data_size']} chars")
        
        # Show chat history if present
        chat_history = session_manager.get_user_data(
            username, 
            SessionKeys.SIMPLE_CHAT_HISTORY, 
            default=[]
        )
        if chat_history:
            print(f"   Chat messages: {len(chat_history)}")
    
    print("\n" + "="*50 + "\n")


if __name__ == "__main__":
    print("🚀 Testing Session Isolation System")
    print("-" * 50)
    
    # Simulate multiple users
    print("\n1️⃣ Simulating user 'alice'...")
    session_manager.set_user_data("alice", SessionKeys.SIMPLE_CHAT_HISTORY, [
        ["Hi", "Hello Alice!"],
        ["How are you?", "I'm great!"]
    ])
    
    print("2️⃣ Simulating user 'bob'...")
    session_manager.set_user_data("bob", SessionKeys.SIMPLE_CHAT_HISTORY, [
        ["Hey", "Hey Bob!"],
        ["What's up?", "Not much!"]
    ])
    
    print("3️⃣ Simulating user 'carol'...")
    session_manager.set_user_data("carol", SessionKeys.SIMPLE_CHAT_HISTORY, [
        ["Hello", "Hi Carol!"]
    ])
    
    # Show stats
    show_session_stats()
    
    # Verify isolation
    print("🔍 Verifying Isolation:")
    alice_history = session_manager.get_user_data("alice", SessionKeys.SIMPLE_CHAT_HISTORY)
    bob_history = session_manager.get_user_data("bob", SessionKeys.SIMPLE_CHAT_HISTORY)
    
    print(f"   Alice has {len(alice_history)} messages")
    print(f"   Bob has {len(bob_history)} messages")
    print(f"   Alice's data != Bob's data: {alice_history != bob_history}")
    
    if alice_history != bob_history:
        print("\n✅ SESSION ISOLATION WORKING CORRECTLY!")
    else:
        print("\n❌ SESSION ISOLATION FAILED!")
    
    print("\n" + "="*50)
    print("To launch the demo app with authentication:")
    print("  python quick_start_session.py --demo")
    print("="*50 + "\n")
    
    # Uncomment to launch demo:
    # demo_isolated_chat().launch(auth=[("alice", "pass1"), ("bob", "pass2")])