lorax-arg 0.1__py3-none-any.whl

lorax/session_manager.py

@@ -0,0 +1,206 @@
+
+import os
+import json
+from uuid import uuid4
+from datetime import datetime, timezone
+from typing import Optional, Dict, Tuple
+
+from lorax.redis_utils import create_redis_client
+from fastapi import Request, Response
+
+from lorax.constants import (
+    SESSION_COOKIE,
+    COOKIE_MAX_AGE,
+    MAX_SOCKETS_PER_SESSION,
+    ENFORCE_CONNECTION_LIMITS,
+)
+
+def _is_https_request(request: Request) -> bool:
+    """
+    Determine whether the *original* request was HTTPS.
+
+    In production, Lorax typically runs behind an HTTPS load balancer / proxy.
+    In that setup, the app may see an internal hop as http://, but the proxy
+    sets X-Forwarded-Proto=https (or Forwarded: proto=https).
+    """
+    xf_proto = (request.headers.get("x-forwarded-proto") or "").split(",")[0].strip().lower()
+    if xf_proto:
+        return xf_proto == "https"
+
+    forwarded = request.headers.get("forwarded") or ""
+    # Minimal parse: look for "proto=https" token anywhere.
+    if "proto=https" in forwarded.lower():
+        return True
+
+    return request.url.scheme == "https"
+
+
+class Session:
+    """Per-user session with socket connection tracking."""
+
+    def __init__(self, sid, file_path=None, socket_connections=None, last_activity=None):
+        self.sid = sid
+        self.file_path = file_path
+        self.created_at = datetime.now(timezone.utc).isoformat()
+        self.last_activity = last_activity or self.created_at
+        # socket_connections: {socket_sid: connected_at_iso_string}
+        self.socket_connections: Dict[str, str] = socket_connections or {}
+
+    def to_dict(self): 
+        return {
+            "sid": self.sid,
+            "file_path": self.file_path,
+            "created_at": self.created_at,
+            "last_activity": self.last_activity,
+            "socket_connections": self.socket_connections,
+        }
+
+    @staticmethod
+    def from_dict(data):
+        session = Session(
+            sid=data["sid"],
+            file_path=data.get("file_path"),
+            socket_connections=data.get("socket_connections", {}),
+            last_activity=data.get("last_activity"),
+        )
+        session.created_at = data.get("created_at", session.created_at)
+        return session
+
+    def update_activity(self):
+        """Update last activity timestamp."""
+        self.last_activity = datetime.now(timezone.utc).isoformat()
+
+    def add_socket(self, socket_sid: str) -> Optional[str]:
+        """
+        Add a socket connection.
+
+        Returns:
+            socket_sid of oldest connection to replace, or None if under limit
+        """
+        self.update_activity()
+
+        # Check if we need to replace an existing connection
+        if ENFORCE_CONNECTION_LIMITS and len(self.socket_connections) >= MAX_SOCKETS_PER_SESSION:
+            # Find oldest connection by connected_at timestamp
+            oldest_socket = min(
+                self.socket_connections.items(),
+                key=lambda x: x[1]  # Sort by timestamp
+            )
+            oldest_socket_sid = oldest_socket[0]
+            # Remove oldest
+            del self.socket_connections[oldest_socket_sid]
+            # Add new
+            self.socket_connections[socket_sid] = datetime.now(timezone.utc).isoformat()
+            return oldest_socket_sid
+
+        # Under limit, just add
+        self.socket_connections[socket_sid] = datetime.now(timezone.utc).isoformat()
+        return None
+
+    def remove_socket(self, socket_sid: str):
+        """Remove a socket connection."""
+        self.socket_connections.pop(socket_sid, None)
+        self.update_activity()
+
+    def get_socket_count(self) -> int:
+        """Get current socket connection count."""
+        return len(self.socket_connections)
+
+    def is_at_connection_limit(self) -> bool:
+        """Check if session is at connection limit."""
+        if not ENFORCE_CONNECTION_LIMITS:
+            return False
+        return len(self.socket_connections) >= MAX_SOCKETS_PER_SESSION
+
+class SessionManager:
+    def __init__(
+        self,
+        redis_url: Optional[str] = None,
+        *,
+        redis_client=None,
+        redis_cluster: bool = False,
+    ):
+        self.redis_url = redis_url
+        self.redis_client = redis_client
+        self.redis_cluster = redis_cluster
+        self.memory_sessions: Dict[str, Session] = {}
+
+        if self.redis_client is None and self.redis_url:
+            self.redis_client = create_redis_client(
+                self.redis_url,
+                decode_responses=True,
+                cluster=self.redis_cluster,
+            )
+
+        if self.redis_client:
+            print(f"✅ SessionManager using Redis at {self.redis_url}")
+        else:
+            print("⚠️ SessionManager running in in-memory mode")
+
+    async def get_session(self, sid: str) -> Optional[Session]:
+        """Retrieve a session by SID."""
+        if not sid:
+            return None
+            
+        if self.redis_client:
+            data = await self.redis_client.get(f"sessions:{sid}")
+            if data:
+                return Session.from_dict(json.loads(data))
+            return None
+        else:
+            return self.memory_sessions.get(sid)
+
+    async def create_session(self, sid: str = None) -> Session:
+        """Create a new session. If SID provided, verify uniqueness/overwrite if needed."""
+        if not sid:
+            sid = str(uuid4())
+        
+        session = Session(sid)
+        await self.save_session(session)
+        return session
+
+    async def save_session(self, session: Session):
+        """Persist session state."""
+        if self.redis_client:
+            await self.redis_client.setex(
+                f"sessions:{session.sid}", 
+                COOKIE_MAX_AGE, 
+                json.dumps(session.to_dict())
+            )
+        else:
+            self.memory_sessions[session.sid] = session
+
+    async def get_or_create_session(self, request: Request, response: Response):
+        """Helper to handle cookie extraction and setting."""
+        sid = request.cookies.get(SESSION_COOKIE)
+        # Cross-site usage (e.g. lorax.ucsc.edu -> api.lorax.in) requires SameSite=None,
+        # and browsers require SameSite=None cookies to be Secure.
+        # Detect original HTTPS behind proxies via X-Forwarded-Proto / Forwarded.
+        secure = _is_https_request(request)
+        
+        session = None
+        if sid:
+            session = await self.get_session(sid)
+        
+        if not session:
+            # Create new if missing or expired
+            session = await self.create_session()
+            # SameSite=None is required for cross-site usage (e.g. lorax.ucsc.edu -> api.lorax.in),
+            # but browsers require SameSite=None cookies to be Secure (HTTPS).
+            # For local HTTP dev, fall back to Lax so the cookie is accepted.
+            samesite = "none" if secure else "lax"
+            response.set_cookie(
+                key=SESSION_COOKIE, 
+                value=session.sid, 
+                httponly=True, 
+                samesite=samesite,
+                max_age=COOKIE_MAX_AGE, 
+                secure=secure
+            )
+        
+        return session.sid, session
+
+    async def health_check(self):
+        if self.redis_client:
+            return await self.redis_client.ping()
+        return True

lorax/sockets/__init__.py

@@ -0,0 +1,55 @@
+"""
+Lorax Socket Event Handlers.
+
+This package provides modularized socket event handling:
+- connection: connect, disconnect, ping events
+- file_ops: load_file, details, query events
+- tree_layout: process_postorder_layout, cache_trees events
+- metadata: fetch_metadata_*, search_metadata events
+- mutations: query_mutations_window, search_mutations events
+- node_search: search_nodes, get_highlight_positions events
+- lineage: ancestors, descendants, mrca, subtree events
+- debug: cache_stats events
+"""
+
+from lorax.sockets.utils import is_csv_session_file
+from lorax.sockets.decorators import (
+    require_session,
+    with_session,
+    with_file_loaded,
+    csv_not_supported,
+    socket_error_handler,
+)
+
+# Will be populated by individual modules after split
+from lorax.sockets.connection import register_connection_events
+from lorax.sockets.file_ops import register_file_events
+from lorax.sockets.tree_layout import register_tree_layout_events
+from lorax.sockets.metadata import register_metadata_events
+from lorax.sockets.mutations import register_mutations_events
+from lorax.sockets.node_search import register_node_search_events
+from lorax.sockets.lineage import register_lineage_events
+from lorax.sockets.debug import register_debug_events
+
+
+def register_socket_events(sio):
+    """Register all socket event handlers."""
+    register_connection_events(sio)
+    register_file_events(sio)
+    register_tree_layout_events(sio)
+    register_metadata_events(sio)
+    register_mutations_events(sio)
+    register_node_search_events(sio)
+    register_lineage_events(sio)
+    register_debug_events(sio)
+
+
+__all__ = [
+    "register_socket_events",
+    "is_csv_session_file",
+    "require_session",
+    "with_session",
+    "with_file_loaded",
+    "csv_not_supported",
+    "socket_error_handler",
+]

lorax/sockets/connection.py

@@ -0,0 +1,99 @@
+"""
+Connection event handlers for Lorax Socket.IO.
+
+Handles connect, disconnect, and ping events.
+"""
+
+from datetime import datetime
+from http.cookies import SimpleCookie
+
+from lorax.context import session_manager
+from lorax.constants import ERROR_SESSION_NOT_FOUND, ERROR_CONNECTION_REPLACED
+
+
+# Mapping from socket_sid to lorax_sid for disconnect handling
+_socket_to_session: dict = {}
+
+
+def register_connection_events(sio):
+    """Register connection-related socket events."""
+
+    @sio.event
+    async def connect(sid, environ, auth=None):
+        print(f"🔌 Socket.IO connected: {sid}")
+
+        cookie = environ.get("HTTP_COOKIE", "")
+        cookies = SimpleCookie()
+        cookies.load(cookie)
+
+        lorax_sid_cookie = cookies.get("lorax_sid")
+        session_id = lorax_sid_cookie.value if lorax_sid_cookie else None
+
+        if not session_id:
+            print(f"⚠️ No lorax_sid cookie found for socket {sid}")
+            await sio.emit("error", {
+                "code": ERROR_SESSION_NOT_FOUND,
+                "message": "Session not found. Please refresh the page."
+            }, to=sid)
+            return
+
+        # Validate session exists
+        session = await session_manager.get_session(session_id)
+        if not session:
+            print(f"⚠️ Session not found: {session_id}")
+            await sio.emit("error", {
+                "code": ERROR_SESSION_NOT_FOUND,
+                "message": "Session expired. Please refresh the page."
+            }, to=sid)
+            return
+
+        # Track this socket connection
+        replaced_socket_sid = session.add_socket(sid)
+        await session_manager.save_session(session)
+
+        # Store mapping for disconnect handling
+        _socket_to_session[sid] = session_id
+
+        # If we replaced an old connection, notify it
+        if replaced_socket_sid:
+            print(f"🔄 Replacing old socket {replaced_socket_sid} with new socket {sid}")
+            await sio.emit("connection-replaced", {
+                "code": ERROR_CONNECTION_REPLACED,
+                "message": "This connection was replaced by a newer tab. Please use the new tab.",
+            }, to=replaced_socket_sid)
+            # Disconnect the old socket
+            try:
+                await sio.disconnect(replaced_socket_sid)
+            except Exception as e:
+                print(f"Warning: Failed to disconnect old socket: {e}")
+
+        # Send session state
+        if session.file_path:
+            await sio.emit("session-restored", {
+                "lorax_sid": session_id,
+                "file_path": session.file_path
+            }, to=sid)
+        else:
+            await sio.emit("status", {
+                "message": "Connected",
+                "lorax_sid": session_id
+            }, to=sid)
+
+    @sio.event
+    async def disconnect(sid):
+        print(f"🔌 Socket.IO disconnected: {sid}")
+
+        # Remove socket from session tracking
+        session_id = _socket_to_session.pop(sid, None)
+        if session_id:
+            session = await session_manager.get_session(session_id)
+            if session:
+                session.remove_socket(sid)
+                await session_manager.save_session(session)
+
+    @sio.event
+    async def ping(sid, data):
+        await sio.emit("pong", {
+            "type": "pong",
+            "time": datetime.utcnow().isoformat()
+        }, to=sid)

lorax/sockets/debug.py

@@ -0,0 +1,47 @@
+"""
+Debug event handlers for Lorax Socket.IO.
+
+Handles cache statistics and debugging events.
+"""
+
+from lorax.context import tree_graph_cache
+from lorax.sockets.decorators import require_session
+
+
+def register_debug_events(sio):
+    """Register debug-related socket events."""
+
+    @sio.event
+    async def get_cache_stats(sid, data):
+        """Socket event to get TreeGraph cache statistics for debugging.
+
+        data: {
+            lorax_sid: str
+        }
+
+        Returns: {
+            mode: str,
+            session_trees: int,  # Trees cached for this session
+            stats: dict          # Additional stats
+        }
+        """
+        try:
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return {"error": "Session not found"}
+
+            # Get session-specific stats
+            cached_trees = await tree_graph_cache.get_all_for_session(lorax_sid)
+
+            # Get global stats
+            global_stats = tree_graph_cache.get_stats()
+
+            return {
+                "session_trees": len(cached_trees),
+                "cached_tree_indices": list(cached_trees.keys()),
+                "stats": global_stats
+            }
+        except Exception as e:
+            print(f"❌ Get cache stats error: {e}")
+            return {"error": str(e)}

lorax/sockets/decorators.py

@@ -0,0 +1,112 @@
+"""
+Socket event decorators for Lorax.
+
+Provides common patterns as decorators to reduce boilerplate:
+- Session validation
+- File loaded checks
+- CSV not supported handling
+- Error wrapping
+"""
+
+import functools
+from typing import Callable, Any
+
+from lorax.context import session_manager
+from lorax.constants import ERROR_SESSION_NOT_FOUND, ERROR_NO_FILE_LOADED
+from lorax.sockets.utils import is_csv_session_file
+
+
+async def require_session(lorax_sid: str, socket_sid: str, sio) -> Any:
+    """
+    Get session or emit error to client.
+
+    Returns the session if found, None otherwise.
+    """
+    session = await session_manager.get_session(lorax_sid)
+    if not session:
+        await sio.emit("error", {
+            "code": ERROR_SESSION_NOT_FOUND,
+            "message": "Session expired. Please refresh the page."
+        }, to=socket_sid)
+        return None
+    return session
+
+
+def with_session(sio):
+    """
+    Decorator that validates session before executing handler.
+
+    The decorated function must accept (sid, data) and data must contain 'lorax_sid'.
+    Passes session as third argument to the handler.
+    """
+    def decorator(func: Callable):
+        @functools.wraps(func)
+        async def wrapper(sid, data):
+            lorax_sid = data.get("lorax_sid")
+            session = await require_session(lorax_sid, sid, sio)
+            if not session:
+                return
+            return await func(sid, data, session)
+        return wrapper
+    return decorator
+
+
+def with_file_loaded(sio, error_event: str = "error"):
+    """
+    Decorator that validates file is loaded before executing handler.
+
+    Must be used after @with_session. The handler receives (sid, data, session).
+    """
+    def decorator(func: Callable):
+        @functools.wraps(func)
+        async def wrapper(sid, data, session):
+            if not session.file_path:
+                print(f"⚠️ No file loaded for session {data.get('lorax_sid')}")
+                await sio.emit(error_event, {
+                    "code": ERROR_NO_FILE_LOADED,
+                    "message": "No file loaded. Please load a file first."
+                }, to=sid)
+                return
+            return await func(sid, data, session)
+        return wrapper
+    return decorator
+
+
+def csv_not_supported(sio, result_event: str, empty_result: dict = None):
+    """
+    Decorator that returns early with message if file is CSV.
+
+    Must be used after @with_session. The handler receives (sid, data, session).
+    """
+    def decorator(func: Callable):
+        @functools.wraps(func)
+        async def wrapper(sid, data, session):
+            if is_csv_session_file(session.file_path):
+                result = empty_result if empty_result else {"error": f"{func.__name__} is not supported for CSV yet."}
+                await sio.emit(result_event, result, to=sid)
+                return
+            return await func(sid, data, session)
+        return wrapper
+    return decorator
+
+
+def socket_error_handler(sio, result_event: str = "error"):
+    """
+    Decorator that wraps handler in try/except and emits errors.
+    """
+    def decorator(func: Callable):
+        @functools.wraps(func)
+        async def wrapper(*args, **kwargs):
+            try:
+                return await func(*args, **kwargs)
+            except Exception as e:
+                print(f"❌ {func.__name__} error: {e}")
+                # For callback-style handlers (return dict), return error
+                # For emit-style handlers, emit error
+                if result_event:
+                    sid = args[0] if args else None
+                    if sid:
+                        await sio.emit(result_event, {"error": str(e)}, to=sid)
+                return {"error": str(e)}
+        return wrapper
+    return decorator