session_sync_daemon.py

#!/usr/bin/env python3
"""
Session sync daemon - Periodically backs up sessions to HF Dataset.

Runs as a background process managed by process-compose.
Syncs session files every 5 minutes to ensure persistence.

Implements graceful shutdown with signal handling:
- SIGTERM: Docker/process-compose shutdown (triggers final backup)
- SIGINT: Ctrl+C / manual interrupt (triggers final backup)
"""

import os
import sys
import time
import signal
import logging
from pathlib import Path

# Add hf_space directory to path
sys.path.insert(0, str(Path(__file__).parent))

from session_sync import backup_sessions_to_dataset, SessionSyncError

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='[%(asctime)s] %(levelname)s: %(message)s'
)
logger = logging.getLogger(__name__)

# Configuration
SYNC_INTERVAL_SECONDS = 120  # 2 minutes (increased from 5 for faster persistence)

# Global flag for graceful shutdown
shutdown_requested = False


def signal_handler(signum, frame):
    """
    Handle shutdown signals (SIGTERM, SIGINT).

    Triggers immediate backup before exit to preserve session state.

    Signal sources:
    - SIGTERM: Docker stop, process-compose shutdown, HF Space restart
    - SIGINT: Ctrl+C, manual interrupt
    """
    global shutdown_requested

    signal_name = signal.Signals(signum).name
    logger.info("=" * 60)
    logger.info(f"Received {signal_name} signal - initiating graceful shutdown")
    logger.info("=" * 60)

    shutdown_requested = True

    # Perform final backup if HF_TOKEN is available
    if os.environ.get("HF_TOKEN"):
        try:
            logger.info("Triggering final session backup...")
            backup_sessions_to_dataset()
            logger.info("✓ Final backup complete")
        except SessionSyncError as e:
            logger.error(f"🚨 SECURITY ERROR during shutdown: {e}")
            # Don't exit(1) here - we're already shutting down
        except Exception as e:
            logger.warning(f"Final backup failed: {e}")
    else:
        logger.info("HF_TOKEN not set - skipping final backup")

    logger.info("=" * 60)
    logger.info("Session Sync Daemon - Stopped")
    logger.info("=" * 60)
    sys.exit(0)


def main():
    """Run periodic session sync."""

    # Register signal handlers BEFORE starting loop
    # SIGTERM: Docker/process-compose shutdown
    # SIGINT: Ctrl+C / manual interrupt
    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)

    logger.info("=" * 60)
    logger.info("Session Sync Daemon - Started")
    logger.info("=" * 60)
    logger.info("Signal handlers registered (SIGTERM, SIGINT)")

    # Check if HF_TOKEN is set
    if not os.environ.get("HF_TOKEN"):
        logger.warning("HF_TOKEN not set - daemon will not sync sessions")
        logger.info("Set HF_TOKEN in Space settings to enable persistence")
        logger.info("Daemon will sleep indefinitely (required by process-compose)")
        # Keep daemon running but do nothing (still responsive to signals)
        while not shutdown_requested:
            time.sleep(3600)
        return

    logger.info(f"Syncing sessions every {SYNC_INTERVAL_SECONDS} seconds")

    while not shutdown_requested:
        try:
            # Sleep in small intervals to be responsive to shutdown signals
            # Split SYNC_INTERVAL into 1-second chunks for faster signal response
            for _ in range(SYNC_INTERVAL_SECONDS):
                if shutdown_requested:
                    break
                time.sleep(1)

            if shutdown_requested:
                break

            logger.info("Starting session backup...")
            backup_sessions_to_dataset()
            logger.info("✓ Session backup complete")

        except SessionSyncError as e:
            logger.error(f"🚨 SECURITY ERROR: {e}")
            logger.error("Daemon stopping to prevent data leak")
            sys.exit(1)

        except Exception as e:
            logger.warning(f"Session sync failed: {e}")
            logger.info("Will retry on next interval...")

    # Graceful exit (signal handler already performed final backup)
    logger.info("Main loop exited gracefully")


if __name__ == "__main__":
    main()