Move dockerfiles inside gateway/ and update main CLI script to be able to run build, up, down and logs on a local gateway server instance.

.env

@@ -1,6 +0,0 @@
-# This specifies the Invariant Explorer instance where the gateway will push the traces 
-# Set this to https://preview-explorer.invariantlabs.ai if you want to push to Preview.
-# If you want to push to a local instance of explorer, then specify the app-api docker container name like: 
-# http://<app-api-docker-container-name>:8000 to push to the local explorer instance.
-INVARIANT_API_URL=https://explorer.invariantlabs.ai
-GUARDRAILS_API_URL=https://explorer.invariantlabs.ai

.github/workflows/publish-images.yml

@@ -48,7 +48,7 @@ jobs:
         uses: docker/build-push-action@v5
         with:
           context: ./gateway
-          file: Dockerfile.gateway
+          file: ./gateway/Dockerfile.gateway
           platforms: linux/amd64
           push: true
           tags: ${{ env.tags }}

MANIFEST.in

@@ -0,0 +1,2 @@
+include gateway/docker-compose.local.yml
+include gateway/Dockerfile.gateway

gateway/.env

@@ -1,4 +1,11 @@
 POSTGRES_USER=postgres
 POSTGRES_PASSWORD=postgres
 POSTGRES_DB=invariantmonitor
-POSTGRES_HOST=database
+POSTGRES_HOST=database
+
+# This specifies the Invariant Explorer instance where the gateway will push the traces 
+# Set this to https://preview-explorer.invariantlabs.ai if you want to push to Preview.
+# If you want to push to a local instance of explorer, then specify the app-api docker container name like: 
+# http://<app-api-docker-container-name>:8000 to push to the local explorer instance.
+INVARIANT_API_URL=https://explorer.invariantlabs.ai
+GUARDRAILS_API_URL=https://explorer.invariantlabs.ai

gateway/Dockerfile.gateway

@@ -2,11 +2,11 @@ FROM python:3.12
 
 WORKDIR /srv/gateway
 
-COPY pyproject.toml ./
+COPY ../pyproject.toml ./pyproject.toml
 
-COPY gateway/ ./
+COPY . .
 
-COPY README.md /srv/gateway/README.md
+COPY ../README.md ./README.md
 
 RUN pip install --no-cache-dir .
 

gateway/__main__.py

@@ -1,12 +1,20 @@
 """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."""
@@ -16,8 +24,14 @@ def signal_handler(sig, frame):
 def print_help():
     """Prints the help message."""
     actions = {
-        "mcp": "Runs the Invariant Gateway against MCP (Model Context Protocol) servers with guardrailing and push to Explorer features.",
-        "llm": "Runs the Invariant Gateway against LLM providers with guardrailing and push to Explorer features. Not implemented yet.",
+        "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.",
     }
 
@@ -25,6 +39,181 @@ def print_help():
         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)
@@ -37,9 +226,19 @@ def main():
     verb = sys.argv[1]
     if verb == "mcp":
         return asyncio.run(mcp.execute(sys.argv[2:]))
-    if verb == "llm":
-        print("[gateway/__main__.py] 'llm' action is not implemented yet.")
-        return 1
+    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

gateway/docker-compose.local.yml

@@ -2,8 +2,8 @@ services:
   invariant-gateway:
     container_name: invariant-gateway
     build:
-      context: .
-      dockerfile: Dockerfile.gateway
+      context: ..
+      dockerfile: gateway/Dockerfile.gateway
     working_dir: /srv/gateway
     env_file:
       - .env
@@ -13,7 +13,7 @@ services:
       - ${INVARIANT_API_KEY:+INVARIANT_API_KEY=${INVARIANT_API_KEY}}
     volumes:
       - type: bind
-        source: ./gateway
+        source: ../gateway
         target: /srv/gateway
       - type: bind
         source: ${GUARDRAILS_FILE_PATH:-/dev/null}

gateway/integrations/explorer.py

@@ -164,7 +164,7 @@ async def fetch_guardrails_from_explorer(
     )
 
     # Get the user details.
-    user_info_response = await client.get("/api/v1/user/identity")
+    user_info_response = await client.get("/api/v1/user/identity", timeout=5)
     if user_info_response.status_code == 401:
         raise HTTPException(
             status_code=401,
@@ -179,7 +179,7 @@ async def fetch_guardrails_from_explorer(
 
     # Get the dataset policies.
     policies_response = await client.get(
-        f"/api/v1/dataset/byuser/{username}/{dataset_name}/policy"
+        f"/api/v1/dataset/byuser/{username}/{dataset_name}/policy", timeout=5
     )
     if policies_response.status_code != 200:
         if policies_response.status_code == 404:

gateway/integrations/guardrails.py

@@ -370,6 +370,7 @@ async def check_guardrails(
                     "Authorization": context.get_guardrailing_authorization(),
                     "Accept": "application/json",
                 },
+                timeout=5,
             )
             if not result.is_success:
                 if result.status_code == 401:

run.sh

@@ -37,7 +37,7 @@ up() {
   fi
 
   # Start Docker Compose with the correct environment variable
-  docker compose -f docker-compose.local.yml up -d
+  docker compose -f gateway/docker-compose.local.yml up -d
 
   # Get the status of the container
   sleep 2
@@ -58,12 +58,12 @@ up() {
 
 build() {
   # Build local services
-  docker compose -f docker-compose.local.yml build
+  docker compose -f gateway/docker-compose.local.yml build
 }
 
 down() {
   # Bring down local services
-  docker compose -f docker-compose.local.yml down
+  docker compose -f gateway/docker-compose.local.yml down
   GATEWAY_ROOT_PATH=$(pwd) docker compose -f tests/integration/docker-compose.test.yml down
 }
 
@@ -143,6 +143,7 @@ integration_tests() {
   # Generate latest whl file for the invariant-gateway package. 
   # This is required to run the integration tests.
   pip install build
+  rm -rf dist
   python -m build
   WHEEL_FILE=$(ls dist/*.whl | head -n 1)
   echo "WHEEL_FILE: $WHEEL_FILE"
@@ -151,9 +152,13 @@ integration_tests() {
 
   docker build -t 'invariant-gateway-tests' -f ./tests/integration/Dockerfile.test ./tests
 
-  docker run \
+  docker rm -f invariant-gateway-integration-tests-container >/dev/null 2>&1 || true
+  mkdir -p /tmp/docker-logs/.invariant
+
+  docker run --name invariant-gateway-integration-tests-container \
     --mount type=bind,source=./tests/integration,target=/tests \
     --mount type=bind,source=$(realpath $WHEEL_FILE),target=/package/$(basename $WHEEL_FILE) \
+    --mount type=bind,source=/tmp/docker-logs/.invariant,target=/root/.invariant \
     --network invariant-gateway-web-test \
     -e OPENAI_API_KEY="$OPENAI_API_KEY" \
     -e ANTHROPIC_API_KEY="$ANTHROPIC_API_KEY"\
@@ -185,7 +190,7 @@ case "$1" in
     down
     ;;
   "logs")
-    docker compose -f docker-compose.local.yml logs -f
+    docker compose -f gateway/docker-compose.local.yml logs -f
     ;;
   "unit-tests")
     shift

tests/integration/docker-compose.test.yml

@@ -25,7 +25,7 @@ services:
     container_name: invariant-gateway-test
     build:
       context: ${GATEWAY_ROOT_PATH}
-      dockerfile: ${GATEWAY_ROOT_PATH}/Dockerfile.gateway
+      dockerfile: ${GATEWAY_ROOT_PATH}/gateway/Dockerfile.gateway
     depends_on:
       app-api:
         condition: service_healthy

tests/integration/resources/mcp/stdio/client/main.py

@@ -2,7 +2,7 @@
 
 import os
 import time
-
+from datetime import timedelta
 from contextlib import AsyncExitStack
 from typing import Any, Optional
 
@@ -68,7 +68,9 @@ class MCPClient:
         )
         self.stdio, self.write = stdio_transport
         self.session = await self.exit_stack.enter_async_context(
-            ClientSession(self.stdio, self.write)
+            ClientSession(
+                self.stdio, self.write, read_timeout_seconds=timedelta(minutes=0.5)
+            )
         )
 
         await self.session.initialize()
@@ -130,5 +132,6 @@ async def run(
     finally:
         # Sleep for a while to allow the server to process the background tasks
         # like pushing traces to the explorer
-        time.sleep(2)
+        if push_to_explorer:
+            time.sleep(2)
         await client.cleanup()