letsbe-sysadmin/app/main.py

"""Main entry point for the LetsBe SysAdmin Agent."""

import asyncio
import signal
import sys
from typing import Optional

from app import __version__
from app.agent import Agent
from app.clients.orchestrator_client import OrchestratorClient
from app.config import get_settings
from app.task_manager import TaskManager
from app.utils.logger import configure_logging, get_logger


def print_banner() -> None:
    """Print startup banner."""
    settings = get_settings()
    banner = f"""
+==============================================================+
|            LetsBe SysAdmin Agent v{__version__:<24}|
+==============================================================+
|  Hostname:     {settings.hostname:<45}|
|  Orchestrator: {settings.orchestrator_url:<45}|
|  Log Level:    {settings.log_level:<45}|
+==============================================================+
"""
    print(banner)


async def main() -> int:
    """Main async entry point.

    Returns:
        Exit code (0 for success, non-zero for failure)
    """
    settings = get_settings()

    # Configure logging
    configure_logging(settings.log_level, settings.log_json)
    logger = get_logger("main")

    print_banner()

    logger.info(
        "agent_starting",
        version=__version__,
        hostname=settings.hostname,
        orchestrator_url=settings.orchestrator_url,
    )

    # Create components
    client = OrchestratorClient(settings)
    agent = Agent(client, settings)
    task_manager = TaskManager(client, settings)

    # Shutdown handler
    shutdown_event = asyncio.Event()

    def handle_signal(sig: int) -> None:
        """Handle shutdown signals."""
        sig_name = signal.Signals(sig).name
        logger.info("signal_received", signal=sig_name)
        shutdown_event.set()

    # Register signal handlers (Unix)
    if sys.platform != "win32":
        loop = asyncio.get_running_loop()
        for sig in (signal.SIGTERM, signal.SIGINT):
            loop.add_signal_handler(sig, lambda s=sig: handle_signal(s))
    else:
        # Windows: Use default CTRL+C handling
        pass

    try:
        # Register with orchestrator
        if not await agent.register():
            logger.error("registration_failed_exit")
            return 1

        # Start background tasks
        heartbeat_task = asyncio.create_task(
            agent.heartbeat_loop(),
            name="heartbeat",
        )
        poll_task = asyncio.create_task(
            task_manager.poll_loop(),
            name="poll",
        )

        logger.info("agent_running")

        # Wait for shutdown signal
        await shutdown_event.wait()

        logger.info("shutdown_initiated")

        # Graceful shutdown
        await task_manager.shutdown()
        await agent.shutdown()

        # Cancel background tasks
        heartbeat_task.cancel()
        poll_task.cancel()

        # Wait for tasks to finish
        await asyncio.gather(
            heartbeat_task,
            poll_task,
            return_exceptions=True,
        )

        logger.info("agent_stopped")
        return 0

    except Exception as e:
        logger.error("agent_fatal_error", error=str(e))
        await client.close()
        return 1


def run() -> None:
    """Entry point for CLI."""
    try:
        exit_code = asyncio.run(main())
        sys.exit(exit_code)
    except KeyboardInterrupt:
        print("\nAgent interrupted by user")
        sys.exit(130)


if __name__ == "__main__":
    run()