app.py

"""SoniCoder v2.1 — Enhanced Entry Point.

Gradio Server (Gradio 6.x) serves index.html from project root
with Gradio's queuing, concurrency, and API infrastructure.

Enhanced version with:
- Better startup logging
- Graceful shutdown handling
- Health check endpoint
- Session management initialization
- Model optimization on startup
- **Hugging Face Hub cloud storage integration**
- Auto-detection for HF Spaces environment

Blog: https://huggingface.co/blog/introducing-gradio-server
"""

from __future__ import annotations

import logging
import os
import signal
import sys
import time

# Ensure workspace + config dirs exist
from code.config import WORKSPACE_ROOT, CONFIG_DIR, APP_VERSION

WORKSPACE_ROOT.mkdir(parents=True, exist_ok=True)
CONFIG_DIR.mkdir(parents=True, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)


def setup_signal_handlers(app):
    """Setup graceful shutdown handlers."""
    def shutdown_handler(signum, frame):
        logger.info("Received shutdown signal, cleaning up...")
        
        # Sync to Hub before shutdown if enabled
        try:
            from code.hub_storage import is_hub_available, get_hub_storage
            if is_hub_available():
                hub = get_hub_storage()
                hub.sync_to_hub(commit_message="Auto-sync before shutdown")
                logger.info("Workspace synced to Hub before shutdown")
        except Exception as e:
            logger.warning("Failed to sync to Hub: %s", e)
        
        try:
            # Save sessions
            from code.server.session_manager import get_session_manager
            manager = get_session_manager()
            saved = manager.save_all()
            if saved > 0:
                logger.info("Saved %d sessions before shutdown", saved)
        except Exception as e:
            logger.warning("Failed to save sessions: %s", e)
        
        try:
            # Clear model memory
            from code.model.loader import _unload_model
            _unload_model()
            logger.info("Model unloaded")
        except Exception as e:
            logger.warning("Failed to unload model: %e", e)
        
        logger.info("Shutdown complete")
        sys.exit(0)
    
    signal.signal(signal.SIGTERM, shutdown_handler)
    signal.signal(signal.SIGINT, shutdown_handler)


def initialize_session_manager():
    """Initialize session manager with persistence."""
    try:
        from code.server.session_manager import get_session_manager
        
        storage_dir = CONFIG_DIR / "sessions"
        manager = get_session_manager(storage_dir=str(storage_dir))
        
        # Cleanup old sessions (older than 30 days)
        removed = manager.cleanup_old_sessions(max_age_days=30)
        if removed > 0:
            logger.info("Cleaned up %d old sessions", removed)
            
        logger.info("Session manager initialized with storage at %s", storage_dir)
        return manager
    except Exception as exc:
        logger.warning("Session manager initialization skipped: %s", exc)
        return None


def initialize_hub_storage():
    """
    Initialize Hugging Face Hub cloud storage.
    
    Auto-detects environment:
    - On HF Spaces: Uses SPACE_ID, HF_TOKEN env vars
    - Local: Requires explicit configuration or HF_TOKEN
    
    Returns:
        Initialized HubStorage instance or None if not configured.
    """
    try:
        from code.hub_storage import init_hub_storage, is_hub_available
        
        # Get configuration from environment
        token = os.environ.get("HF_TOKEN", "") or os.environ.get("HF_AUTH_TOKEN", "")
        space_id = os.environ.get("SPACE_ID", "")
        
        # Determine mode based on environment
        if space_id:
            # Running on HF Spaces
            logger.info("Detected HF Spaces environment: %s", space_id)
            
            # Extract repo_id from space_id (format: "user/space-name")
            if space_id and token:
                # Use a dataset repo for workspace persistence
                # This allows the AI's code to persist across Space restarts
                user_name = space_id.split("/")[0] if "/" in space_id else "unknown"
                dataset_repo_id = os.environ.get(
                    "HUB_DATASET_REPO",
                    f"{user_name}/sonicoder-workspace"
                )
                
                mode = os.environ.get("HUB_STORAGE_MODE", "hybrid")
                
                storage = init_hub_storage(
                    token=token,
                    repo_id=space_id,
                    dataset_repo_id=dataset_repo_id,
                    mode=mode,
                )
                
                stats = storage.get_storage_stats()
                logger.info(
                    "Hub storage initialized: mode=%s, hub=%s, dataset=%s",
                    stats["mode"],
                    stats["hub_enabled"],
                    stats["dataset_enabled"]
                )
                
                return storage
            else:
                logger.warning("HF Spaces detected but no token provided, using local storage only")
                return None
        elif token:
            # Local mode with token
            repo_id = os.environ.get("HUB_REPO_ID", "")
            dataset_repo_id = os.environ.get("HUB_DATASET_REPO", "")
            
            if repo_id or dataset_repo_id:
                storage = init_hub_storage(
                    token=token,
                    repo_id=repo_id,
                    dataset_repo_id=dataset_repo_id,
                    mode="hybrid",
                )
                
                logger.info("Hub storage initialized in local mode")
                return storage
        
        # No configuration available
        logger.info("Hub storage not configured (set HF_TOKEN to enable)")
        return None
        
    except ImportError:
        logger.debug("huggingface_hub not installed, skipping Hub storage")
        return None
    except Exception as exc:
        logger.warning("Hub storage initialization failed: %s", exc)
        return None


def main():
    start_time = time.time()
    logger.info("=" * 50)
    logger.info("SoniCoder v%s starting...", APP_VERSION)
    logger.info("=" * 50)

    # Detect and log environment
    space_id = os.environ.get("SPACE_ID", None)
    if space_id:
        logger.info("Running on Hugging Face Spaces: %s", space_id)
    else:
        logger.info("Running in local mode")

    # Initialize session manager
    session_mgr = initialize_session_manager()

    # Initialize Hub storage (cloud persistence)
    hub_storage = initialize_hub_storage()

    # Start loading default model in background
    from code.model.loader import start_background_load, add_progress_callback

    def on_progress(progress):
        if progress.progress_percent % 25 < 5:  # Log at ~25%, 50%, 75%, 100%
            logger.info(
                "Model loading: %.0f%% - %s",
                progress.progress_percent,
                progress.message,
            )

    add_progress_callback(on_progress)
    
    load_thread = start_background_load()
    logger.info("Background model loading started")

    # Set up MCP servers from config (best-effort)
    try:
        from code.config import load_settings
        from code.mcp import get_mcp_manager

        settings = load_settings()
        if settings.mcp_servers:
            mcp = get_mcp_manager()
            mcp.configure(settings.mcp_servers)
            results = mcp.connect_all()
            for name, error in results.items():
                if error:
                    logger.warning("MCP server '%s' failed: %s", name, error)
                else:
                    logger.info("MCP server '%s' connected", name)
    except Exception as exc:
        logger.warning("MCP setup skipped: %s", exc)

    # Create the server (gradio.Server — custom HTML + Gradio backend)
    from code.server import create_app

    app = create_app()

    # Setup signal handlers for graceful shutdown
    setup_signal_handlers(app)

    # Log startup time
    startup_time = time.time() - start_time
    logger.info("Startup completed in %.2f seconds", startup_time)
    
    # Log storage info
    if hub_storage:
        logger.info("Cloud storage: ENABLED (files will persist on HF Hub)")
    else:
        logger.info("Cloud storage: DISABLED (local files only)")
    
    logger.info("Launching...")

    # Launch the server
    app.launch(
        show_error=True,
        server_name="0.0.0.0",
        server_port=7860,
    )


if __name__ == "__main__":
    main()