"""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 stdio as mcp_stdio from gateway.mcp.log import mcp_log 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 == "help": print_help() return 0 if "INVARIANT_API_KEY" not in os.environ: print("[ERROR] INVARIANT_API_KEY environment variable is not set.") mcp_log("[ERROR] INVARIANT_API_KEY environment variable is not set.") sys.exit(1) if verb == "mcp": return asyncio.run(mcp_stdio.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 print(f"[gateway/__main__.py] Unknown action: {verb}") return 1 if __name__ == "__main__": sys.exit(main())