invariant-gateway/gateway/__main__.py

"""Script is used to run actions using the Invariant Gateway."""

import asyncio
import os
import signal
import subprocess
import sys
import time

from typing import Optional

from gateway.mcp import mcp


LOCAL_COMPOSE_FILE = "gateway/docker-compose.local.yml"


# Handle signals to ensure clean shutdown
def signal_handler(sig, frame):
    """Handle signals for graceful shutdown."""
    sys.exit(0)


def print_help():
    """Prints the help message."""
    actions = {
        "mcp": """
        Runs the Invariant Gateway against MCP (Model Context Protocol) stdio servers with guardrailing and push to Explorer features.
        """,
        "server": """
        Runs the Invariant Gateway server locally providing guardrailing and push to Explorer features.
        Should be called with one of the following subcommands: build, up, down, logs.
        A guardrails file can be passed with the flag: --guardrails-file=/path/to/guardrails/file.
        """,
        "help": "Shows this help message.",
    }

    for verb, description in actions.items():
        print(f"{verb}: {description}")


def ensure_network_exists(network_name: str = "invariant-explorer-web") -> bool:
    """Ensure the Docker network exists."""
    try:
        # Check if network exists
        result = subprocess.run(
            ["docker", "network", "inspect", network_name],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=False,
        )

        if result.returncode != 0:
            print(f"Creating Docker network: {network_name}")
            subprocess.run(
                ["docker", "network", "create", network_name],
                check=True,
            )

        return True
    except subprocess.CalledProcessError as e:
        print(f"Error creating Docker network: {e}")
        return False


def setup_guardrails(guardrails_file_path: Optional[str] = None) -> bool:
    """Configure guardrails if specified."""
    if not guardrails_file_path:
        return True

    if not os.path.isfile(guardrails_file_path):
        print(
            f"Error: Specified guardrails file does not exist: {guardrails_file_path}"
        )
        return False

    # Convert to absolute path
    guardrails_file_path = os.path.realpath(guardrails_file_path)
    os.environ["GUARDRAILS_FILE_PATH"] = guardrails_file_path

    # Check if INVARIANT_API_KEY is set
    if not os.environ.get("INVARIANT_API_KEY"):
        print(
            "Error: A guardrails file is specified, but INVARIANT_API_KEY env var is not set. "
            "This is required to validate guardrails."
        )
        return False

    return True


def build():
    """Build Docker containers using docker-compose."""
    try:
        print(f"Building using docker-compose file: {LOCAL_COMPOSE_FILE}")
        subprocess.run(
            ["docker", "compose", "-f", str(LOCAL_COMPOSE_FILE), "build"],
            check=True,
        )
        print("Build completed successfully")
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error building containers: {e}")
        return False


def up(guardrails_file_path: Optional[str] = None):
    """Set up the local server for the Invariant Gateway."""
    # Ensure network exists
    if not ensure_network_exists():
        return 1

    # Set up guardrails
    if not setup_guardrails(guardrails_file_path=guardrails_file_path):
        return 1

    # Run the server
    try:
        # Start containers
        print(f"Starting containers using docker-compose file: {LOCAL_COMPOSE_FILE}")
        subprocess.run(
            ["docker", "compose", "-f", str(LOCAL_COMPOSE_FILE), "up", "-d"],
            check=True,
        )

        # Wait for containers to start
        time.sleep(2)

        # Check if gateway container is running
        result = subprocess.run(
            ["docker", "ps", "-qf", "name=invariant-gateway"],
            capture_output=True,
            text=True,
            check=True,
        )

        if not result.stdout.strip():
            print("The invariant-gateway container failed to start.")
            logs = subprocess.run(
                ["docker", "logs", "invariant-gateway"],
                capture_output=True,
                text=True,
                check=False,
            )
            print("Last 20 lines of logs:")
            print("\n".join(logs.stdout.strip().split("\n")[-20:]))
            return False

        print("Gateway started at http://localhost:8005/api/v1/gateway/")
        print("See http://localhost:8005/api/v1/gateway/docs for API documentation")

        if guardrails_file_path:
            print(f"Using Guardrails File: {guardrails_file_path}")

        return True
    except subprocess.CalledProcessError as e:
        print(f"Error starting containers: {e}")
        return False


def down():
    """Stop the Docker containers."""
    try:
        print(f"Stopping containers using docker-compose file: {LOCAL_COMPOSE_FILE}")
        subprocess.run(
            ["docker", "compose", "-f", str(LOCAL_COMPOSE_FILE), "down"],
            check=True,
        )
        print("Containers stopped successfully")
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error stopping containers: {e}")
        return False


def logs():
    """Show container logs."""
    try:
        subprocess.run(
            ["docker", "compose", "-f", str(LOCAL_COMPOSE_FILE), "logs", "-f"],
            check=True,
        )
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error showing logs: {e}")
        return False
    except KeyboardInterrupt:
        print("\nExiting logs view")
        return True


def run_server_command(command, args=None):
    """Run a server command."""
    if args is None:
        args = []

    if command == "build":
        return build()
    elif command == "up":
        # Parse guardrails file from args
        guardrails_file = None
        for arg in args:
            if arg.startswith("--guardrails-file="):
                guardrails_file = arg.split("=", 1)[1]

        return up(guardrails_file)
    elif command == "down":
        return down()
    elif command == "logs":
        return logs()
    else:
        print(f"Unknown server command: {command}")
        print("Available commands: build, up, down, logs")
        return False


def main():
    """Entry point for the Invariant Gateway."""
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    if len(sys.argv) < 2:
        print_help()
        sys.exit(1)

    verb = sys.argv[1]
    if verb == "mcp":
        return asyncio.run(mcp.execute(sys.argv[2:]))
    if verb == "server":
        if len(sys.argv) < 3:
            print(
                "Error: Missing command for server. Should be one of: build, up, down, logs"
            )
            print_help()
            return 1

        command = sys.argv[2]
        args = sys.argv[3:]
        if not run_server_command(command, args):
            return 1
        return 0
    if verb == "help":
        print_help()
        return 0
    print(f"[gateway/__main__.py] Unknown action: {verb}")
    return 1


if __name__ == "__main__":
    sys.exit(main())