langprotect-mcp-gateway 1.0.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

langprotect_mcp_gateway/gateway.py

@@ -1,26 +1,6 @@
 #!/usr/bin/env python3
 """
-LangProtect MCP Gateway - Universal MCP Logging & Security
-
-This gateway intercepts ALL MCP tool calls from any client (Claude, VS Code, etc.)
-and logs them to LangProtect backend for security scanning and auditing.
-
-Architecture:
-1. Reads user's MCP config to discover other MCP servers
-2. Implements MCP protocol (initialize, list_tools, call_tool)
-3. Logs every tool call to LangProtect backend
-4. Forwards requests to actual MCP servers
-5. Returns responses to AI client
-
-Usage:
-    python3 langprotect_gateway.py
-
-Configuration (via environment variables):
-    LANGPROTECT_URL      - Backend API URL (default: http://localhost:8000)
-    LANGPROTECT_EMAIL    - User email for authentication
-    LANGPROTECT_PASSWORD - User password
-    MCP_CONFIG_PATH      - Path to mcp.json (auto-detected if not set)
-    DEBUG                - Enable debug logging (true/false)
+LangProtect MCP Gateway - Security Gateway for MCP Servers
 """
 
 import sys
@@ -28,86 +8,82 @@ import json
 import os
 import subprocess
 import requests
-import threading
-import time
+import argparse
 from datetime import datetime, timedelta
-from pathlib import Path
 from typing import Dict, List, Any, Optional
 import logging
 
-# Configure logging
-logging.basicConfig(
-    level=logging.DEBUG if os.getenv('DEBUG', 'false').lower() == 'true' else logging.INFO,
-    format='[%(asctime)s] %(levelname)s: %(message)s',
-    handlers=[logging.StreamHandler(sys.stderr)]
-)
+log_level = os.environ.get("LOGLEVEL", "DEBUG" if os.getenv('DEBUG', 'false').lower() == 'true' else "INFO").upper()
+logging.basicConfig(level=getattr(logging, log_level), format='[%(asctime)s] %(levelname)s: %(message)s', handlers=[logging.StreamHandler(sys.stderr)])
 logger = logging.getLogger('langprotect-gateway')
 
 
 class MCPServer:
-    """Represents a single MCP server that can be called"""
-    
     def __init__(self, name: str, config: Dict[str, Any]):
         self.name = name
         self.command = config.get('command')
         self.args = config.get('args', [])
         self.env = config.get('env', {})
-        self.tools = []  # Will be populated by discovery
-        self.process = None
-        logger.info(f"Initialized MCP server: {name}")
+        self.tools: List[Dict] = []
+        self.process: Optional[subprocess.Popen] = None
+        self._request_id = 0
+        logger.info(f"Configured MCP server: {name} ({self.command} {' '.join(self.args)})")
     
-    def call(self, method: str, params: Dict) -> Dict:
-        """Call this MCP server with a method and params"""
+    def _next_id(self) -> int:
+        self._request_id += 1
+        return self._request_id
+    
+    def start(self) -> bool:
         try:
-            # Start process for this call
             env = {**os.environ, **self.env}
-            
-            self.process = subprocess.Popen(
-                [self.command] + self.args,
-                stdin=subprocess.PIPE,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                text=True,
-                env=env
-            )
-            
-            # Send request
-            request = {
-                "jsonrpc": "2.0",
-                "id": 1,
-                "method": method,
-                "params": params
-            }
-            
-            logger.debug(f"Calling {self.name}.{method}: {json.dumps(request)}")
-            
+            self.process = subprocess.Popen([self.command] + self.args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env, bufsize=1)
+            init_request = {"jsonrpc": "2.0", "id": self._next_id(), "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "langprotect-gateway", "version": "1.0.0"}}}
+            self.process.stdin.write(json.dumps(init_request) + "\n")
+            self.process.stdin.flush()
+            response_line = self.process.stdout.readline()
+            if not response_line:
+                logger.error(f"Failed to initialize {self.name}")
+                return False
+            response = json.loads(response_line)
+            if "error" in response:
+                logger.error(f"Initialize error for {self.name}: {response['error']}")
+                return False
+            self.process.stdin.write(json.dumps({"jsonrpc": "2.0", "method": "notifications/initialized"}) + "\n")
+            self.process.stdin.flush()
+            logger.info(f"Started MCP server: {self.name}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to start {self.name}: {e}")
+            return False
+    
+    def stop(self):
+        if self.process:
+            try:
+                self.process.terminate()
+                self.process.wait(timeout=5)
+            except:
+                self.process.kill()
+            self.process = None
+    
+    def call(self, method: str, params: Dict) -> Dict:
+        if not self.process or self.process.poll() is not None:
+            raise Exception(f"Server {self.name} is not running")
+        request = {"jsonrpc": "2.0", "id": self._next_id(), "method": method, "params": params}
+        logger.debug(f"-> {self.name}.{method}")
+        try:
             self.process.stdin.write(json.dumps(request) + "\n")
             self.process.stdin.flush()
-            
-            # Read response
             response_line = self.process.stdout.readline()
             if not response_line:
-                error_output = self.process.stderr.read()
-                raise Exception(f"No response from {self.name}: {error_output}")
-            
+                raise Exception(f"No response from {self.name}")
             response = json.loads(response_line)
-            
-            # Cleanup
-            self.process.terminate()
-            self.process.wait(timeout=1)
-            
-            logger.debug(f"Response from {self.name}: {json.dumps(response)[:200]}")
-            
+            logger.debug(f"<- {self.name}: {str(response)[:200]}")
             return response
-            
         except Exception as e:
-            logger.error(f"Error calling {self.name}: {e}")
-            if self.process:
-                self.process.terminate()
+            logger.error(f"Error calling {self.name}.{method}: {e}")
             raise
     
     def discover_tools(self) -> List[Dict]:
-        """Discover what tools this MCP server provides"""
         try:
             response = self.call("tools/list", {})
             if "result" in response:
@@ -120,407 +96,286 @@ class MCPServer:
             return []
 
 
-class LangProtectGateway:
-    """Main gateway that intercepts and logs all MCP traffic"""
-    
-    def __init__(self):
-        self.langprotect_url = os.getenv('LANGPROTECT_URL', 'http://localhost:8000')
-        self.email = os.getenv('LANGPROTECT_EMAIL')
-        self.password = os.getenv('LANGPROTECT_PASSWORD')
-        self.jwt_token = None
-        self.token_expiry = None
-        self.mcp_servers = {}  # name -> MCPServer
-        self.tool_to_server = {}  # tool_name -> server_name
-        
-        # Debug: Log what we received
-        logger.debug(f"Environment variables received:")
-        logger.debug(f"  LANGPROTECT_URL: {self.langprotect_url}")
-        logger.debug(f"  LANGPROTECT_EMAIL: {self.email}")
-        logger.debug(f"  LANGPROTECT_PASSWORD: {'***' + self.password[-4:] if self.password else 'NOT SET'}")
-        
-        # Validate credentials
-        if not self.email or not self.password:
-            logger.error("LANGPROTECT_EMAIL and LANGPROTECT_PASSWORD must be set")
-            sys.exit(1)
-        
-        # Authenticate
-        self.login()
-        
-        # Discover MCP servers
-        self.discover_mcp_servers()
-        
-        # Start token refresh thread
-        self.start_token_refresh()
-        
-        logger.info("LangProtect Gateway initialized successfully")
+class LangProtectAuth:
+    def __init__(self, url: str, email: str, password: str):
+        self.url = url
+        self.email = email
+        self.password = password
+        self.jwt_token: Optional[str] = None
+        self.token_expiry: Optional[datetime] = None
     
-    def login(self):
-        """Authenticate with LangProtect backend and get JWT token"""
+    def login(self) -> bool:
         try:
-            logger.info(f"Logging in to {self.langprotect_url}...")
-            
-            response = requests.post(
-                f"{self.langprotect_url}/v1/group-users/signin",
-                json={
-                    'email': self.email,
-                    'password': self.password
-                },
-                timeout=10
-            )
-            
+            logger.info(f"Authenticating with {self.url}...")
+            response = requests.post(f"{self.url}/v1/group-users/signin", json={'email': self.email, 'password': self.password}, timeout=10)
             if response.status_code == 200:
                 data = response.json()
                 self.jwt_token = data.get('access_token')
                 self.token_expiry = datetime.now() + timedelta(days=6)
-                logger.info("✅ Authentication successful")
+                logger.info("Authentication successful")
                 return True
             else:
-                logger.error(f"❌ Login failed: {response.status_code} - {response.text}")
-                sys.exit(1)
-                
+                logger.error(f"Login failed: {response.status_code}")
+                return False
         except Exception as e:
-            logger.error(f"❌ Login error: {e}")
-            sys.exit(1)
+            logger.error(f"Login error: {e}")
+            return False
     
-    def ensure_token(self):
-        """Ensure we have a valid JWT token"""
+    def ensure_token(self) -> bool:
         if not self.jwt_token or (self.token_expiry and datetime.now() > self.token_expiry):
-            logger.info("Token expired, re-authenticating...")
             return self.login()
         return True
     
-    def start_token_refresh(self):
-        """Start background thread to refresh JWT token"""
-        def refresh_loop():
-            while True:
-                time.sleep(86400)  # Check daily
-                if self.token_expiry and datetime.now() > self.token_expiry:
-                    logger.info("Auto-refreshing token...")
-                    self.login()
+    def scan(self, tool_name: str, arguments: Dict, server_name: str) -> Dict:
+        self.ensure_token()
+        try:
+            payload = {'prompt': json.dumps({'tool': tool_name, 'arguments': arguments, 'server': server_name}), 'client_ip': '127.0.0.1', 'user_agent': f'LangProtect-MCP-Gateway/1.0 (server={server_name})', 'source': 'mcp-gateway'}
+            response = requests.post(f"{self.url}/v1/group-logs/scan", json=payload, headers={'Authorization': f'Bearer {self.jwt_token}', 'Content-Type': 'application/json'}, timeout=5)
+            if response.status_code != 200:
+                logger.warning(f"Backend returned {response.status_code}, allowing request (fail-open)")
+                return {'status': 'allowed', 'error': f'Backend error: {response.status_code}'}
+            result = response.json()
+            # Handle scan service timeout - fail open
+            if result.get('detections', {}).get('error') == 'Scan service timeout':
+                logger.warning("Scan service timeout, allowing request (fail-open)")
+                return {'status': 'allowed', 'id': result.get('id'), 'error': 'Scan timeout'}
+            return result
+        except requests.exceptions.Timeout:
+            logger.warning("Backend scan timeout, allowing request (fail-open)")
+            return {'status': 'allowed', 'error': 'Request timeout'}
+        except Exception as e:
+            logger.error(f"Scan error: {e}")
+            return {'status': 'allowed', 'error': str(e)}
+
+
+class LangProtectGateway:
+    def __init__(self, mcp_json_path: Optional[str] = None):
+        self.mcp_json_path = mcp_json_path
         
-        refresh_thread = threading.Thread(target=refresh_loop, daemon=True)
-        refresh_thread.start()
-        logger.debug("Token refresh thread started")
-    
-    def discover_mcp_servers(self):
-        """Find and load MCP servers from user's config"""
-        config_paths = [
-            os.getenv('MCP_CONFIG_PATH'),
-            os.path.expanduser('~/.cursor/mcp.json'),
-            os.path.expanduser('~/.config/Claude/claude_desktop_config.json'),
-            os.path.expanduser('~/Library/Application Support/Claude/claude_desktop_config.json')
-        ]
+        # Load credentials from env vars first, then potentially from config
+        self.langprotect_url = os.getenv('LANGPROTECT_URL', 'http://localhost:8000')
+        self.email = os.getenv('LANGPROTECT_EMAIL')
+        self.password = os.getenv('LANGPROTECT_PASSWORD')
         
-        for path in config_paths:
-            if not path:
-                continue
-            
-            path = os.path.expanduser(path)
-            if os.path.exists(path):
-                logger.info(f"Found MCP config: {path}")
-                try:
-                    with open(path, 'r') as f:
-                        config = json.load(f)
-                    
-                    servers = config.get('mcpServers', {})
-                    
-                    # Load all servers except ourselves
-                    for name, cfg in servers.items():
-                        if name == 'langprotect-gateway' or name == 'langprotect':
-                            continue
-                        
-                        try:
-                            server = MCPServer(name, cfg)
-                            self.mcp_servers[name] = server
-                            
-                            # Discover tools from this server
-                            tools = server.discover_tools()
-                            for tool in tools:
-                                tool_name = tool.get('name')
-                                if tool_name:
-                                    self.tool_to_server[tool_name] = name
-                                    logger.debug(f"  Tool: {tool_name} -> {name}")
-                        
-                        except Exception as e:
-                            logger.warning(f"Failed to load MCP server {name}: {e}")
-                    
-                    logger.info(f"✅ Loaded {len(self.mcp_servers)} MCP servers")
-                    logger.info(f"✅ Discovered {len(self.tool_to_server)} tools")
-                    return
-                    
-                except Exception as e:
-                    logger.error(f"Error reading config {path}: {e}")
+        # Try to load credentials from mcp.json env section (like Lasso)
+        if mcp_json_path and (not self.email or not self.password):
+            self._load_env_from_config(mcp_json_path)
         
-        logger.warning("⚠️  No MCP config found - gateway will not proxy any servers")
+        self.auth: Optional[LangProtectAuth] = None
+        self.mcp_servers: Dict[str, MCPServer] = {}
+        self.tool_to_server: Dict[str, str] = {}
+        self.all_tools: List[Dict] = []
+        logger.debug(f"LANGPROTECT_URL: {self.langprotect_url}")
+        logger.debug(f"LANGPROTECT_EMAIL: {self.email}")
     
-    def log_to_backend(self, method: str, params: Dict, server_name: str = "unknown") -> Dict:
-        """Log MCP request to LangProtect backend for scanning"""
-        self.ensure_token()
-        
+    def _load_env_from_config(self, path: str):
+        """Load credentials from mcp.json env section (Lasso-style)"""
         try:
-            # Build content for scanning - convert MCP request to scannable format
-            if method == 'tools/call':
-                tool_name = params.get('name', 'unknown_tool')
-                arguments = params.get('arguments', {})
-                scan_content = json.dumps({
-                    'tool': tool_name,
-                    'arguments': arguments,
-                    'server_url': server_name,
-                    'method': method
-                })
-            else:
-                scan_content = json.dumps({
-                    'method': method,
-                    'params': params,
-                    'server_url': server_name
-                })
+            expanded_path = os.path.expanduser(path)
+            with open(expanded_path, 'r') as f:
+                config = json.load(f)
             
-            # Use standard group scan endpoint (works with existing scanners + policies)
-            payload = {
-                'prompt': scan_content,  # MCP request as prompt
-                'client_ip': '127.0.0.1',
-                'user_agent': f'LangProtect-MCP-Gateway/1.0 (server={server_name})',
-                'source': 'mcp-gateway'
-            }
+            # Look for env vars in the gateway's config section
+            mcp_servers = config.get('mcpServers', {})
+            for gateway_name in ['langprotect-gateway', 'langprotect', 'mcp-gateway']:
+                gateway_config = mcp_servers.get(gateway_name, {})
+                env_section = gateway_config.get('env', {})
+                if env_section:
+                    if not self.langprotect_url or self.langprotect_url == 'http://localhost:8000':
+                        self.langprotect_url = env_section.get('LANGPROTECT_URL', self.langprotect_url)
+                    if not self.email:
+                        self.email = env_section.get('LANGPROTECT_EMAIL')
+                    if not self.password:
+                        self.password = env_section.get('LANGPROTECT_PASSWORD')
+                    logger.info(f"Loaded credentials from config env section")
+                    break
+        except Exception as e:
+            logger.debug(f"Could not load env from config: {e}")
+    
+    def initialize(self) -> bool:
+        if self.email and self.password:
+            self.auth = LangProtectAuth(self.langprotect_url, self.email, self.password)
+            if not self.auth.login():
+                logger.error("Failed to authenticate with LangProtect backend")
+                return False
+        else:
+            logger.warning("No LangProtect credentials - running in pass-through mode")
+        if not self.load_servers():
+            return False
+        if not self.start_servers():
+            return False
+        logger.info("=" * 50)
+        logger.info("LangProtect Gateway initialized")
+        logger.info(f"Backend: {self.langprotect_url}")
+        logger.info(f"Servers: {len(self.mcp_servers)}")
+        logger.info(f"Tools: {len(self.all_tools)}")
+        logger.info("=" * 50)
+        return True
+    
+    def load_servers(self) -> bool:
+        # Mode 1: Single server via environment variables (for wrapper scripts)
+        mcp_command = os.getenv('MCP_SERVER_COMMAND')
+        mcp_args = os.getenv('MCP_SERVER_ARGS')
+        if mcp_command:
+            logger.info(f"Single server mode: {mcp_command}")
+            args_list = [arg.strip() for arg in mcp_args.split(',')] if mcp_args else []
+            server_name = os.getenv('MCP_SERVER_NAME', 'proxied-server')
+            self.mcp_servers[server_name] = MCPServer(server_name, {'command': mcp_command, 'args': args_list, 'env': {}})
+            return True
+        
+        # Mode 2: Config file (mcp.json)
+        if self.mcp_json_path:
+            return self.load_from_mcp_json(self.mcp_json_path)
+        
+        logger.warning("No MCP servers configured")
+        return False
+    
+    def load_from_mcp_json(self, path: str) -> bool:
+        try:
+            expanded_path = os.path.expanduser(path)
+            with open(expanded_path, 'r') as f:
+                config = json.load(f)
             
-            logger.debug(f"Logging to backend: {method} on {server_name}")
+            # Try multiple config structures:
+            # 1. Lasso-style: mcpServers.langprotect-gateway.servers (nested)
+            # 2. VS Code style: servers (direct)
+            # 3. Claude Desktop style: mcpServers (direct)
             
-            response = requests.post(
-                f"{self.langprotect_url}/v1/group-logs/scan",
-                json=payload,
-                headers={
-                    'Authorization': f'Bearer {self.jwt_token}',
-                    'Content-Type': 'application/json'
-                },
-                timeout=30
-            )
+            servers = {}
             
-            if response.status_code != 200:
-                logger.error(f"Backend scan failed: {response.status_code} - {response.text}")
-                # Fail-open: allow request if backend error
-                return {'status': 'allowed', 'error': f'Backend error: {response.status_code}'}
+            # Check for Lasso-style nested config
+            mcp_servers = config.get('mcpServers', {})
+            for gateway_name in ['langprotect-gateway', 'langprotect', 'mcp-gateway']:
+                gateway_config = mcp_servers.get(gateway_name, {})
+                if 'servers' in gateway_config:
+                    servers = gateway_config['servers']
+                    logger.info(f"Found nested servers config under mcpServers.{gateway_name}.servers")
+                    break
             
-            result = response.json()
-            status = result.get('status', '').lower()
+            # Fallback to direct config
+            if not servers:
+                servers = config.get('servers', config.get('mcpServers', {}))
             
-            logger.debug(f"Backend response: status={status}, log_id={result.get('id')}")
+            if not servers:
+                logger.error("No servers found in config file")
+                return False
             
-            return result
+            for name, cfg in servers.items():
+                # Skip gateway self-references
+                if name in ['langprotect-gateway', 'langprotect', 'mcp-gateway']:
+                    continue
+                self.mcp_servers[name] = MCPServer(name, cfg)
             
+            logger.info(f"Loaded {len(self.mcp_servers)} servers from config")
+            return len(self.mcp_servers) > 0
         except Exception as e:
-            logger.error(f"Error logging to backend: {e}")
-            # Fail-open for now (TODO: make configurable)
-            return {'status': 'allowed', 'error': str(e)}
+            logger.error(f"Error loading {path}: {e}")
+            return False
+    
+    def start_servers(self) -> bool:
+        started = 0
+        for name, server in list(self.mcp_servers.items()):
+            if server.start():
+                tools = server.discover_tools()
+                for tool in tools:
+                    tool_name = tool.get('name')
+                    if tool_name:
+                        self.tool_to_server[tool_name] = name
+                        tool_copy = tool.copy()
+                        tool_copy['description'] = f"[{name}] {tool_copy.get('description', '')}"
+                        self.all_tools.append(tool_copy)
+                started += 1
+            else:
+                del self.mcp_servers[name]
+        return started > 0
     
-    def handle_mcp_request(self, request: Dict) -> Dict:
-        """Main MCP protocol handler"""
+    def shutdown(self):
+        for server in self.mcp_servers.values():
+            server.stop()
+    
+    def handle_request(self, request: Dict) -> Optional[Dict]:
         method = request.get('method')
         request_id = request.get('id')
-        
-        logger.info(f"📥 MCP Request: {method} (id={request_id})")
-        
+        params = request.get('params', {})
+        logger.info(f"Request: {method} (id={request_id})")
         try:
             if method == 'initialize':
-                return self.handle_initialize(request)
+                return {'jsonrpc': '2.0', 'id': request_id, 'result': {'protocolVersion': '2024-11-05', 'capabilities': {'tools': {}}, 'serverInfo': {'name': 'langprotect-gateway', 'version': '1.0.0'}}}
+            elif method == 'notifications/initialized':
+                return None
             elif method == 'tools/list':
-                return self.handle_list_tools(request)
+                logger.info(f"Returning {len(self.all_tools)} tools")
+                return {'jsonrpc': '2.0', 'id': request_id, 'result': {'tools': self.all_tools}}
             elif method == 'tools/call':
-                return self.handle_call_tool(request)
+                return self._handle_call_tool(request_id, params)
             else:
-                logger.warning(f"Unknown method: {method}")
-                return {
-                    'jsonrpc': '2.0',
-                    'id': request_id,
-                    'error': {
-                        'code': -32601,
-                        'message': f'Method not found: {method}'
-                    }
-                }
-        
+                return {'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32601, 'message': f'Method not found: {method}'}}
         except Exception as e:
-            logger.error(f"Error handling {method}: {e}", exc_info=True)
-            return {
-                'jsonrpc': '2.0',
-                'id': request_id,
-                'error': {
-                    'code': -32603,
-                    'message': f'Internal error: {str(e)}'
-                }
-            }
-    
-    def handle_initialize(self, request: Dict) -> Dict:
-        """Handle MCP initialize handshake"""
-        logger.info("Handling initialize request")
-        
-        return {
-            'jsonrpc': '2.0',
-            'id': request.get('id'),
-            'result': {
-                'protocolVersion': '2024-11-05',
-                'capabilities': {
-                    'tools': {}
-                },
-                'serverInfo': {
-                    'name': 'langprotect-gateway',
-                    'version': '1.0.0'
-                }
-            }
-        }
-    
-    def handle_list_tools(self, request: Dict) -> Dict:
-        """Return all tools from all MCP servers"""
-        logger.info(f"Listing tools from {len(self.mcp_servers)} servers")
-        
-        all_tools = []
-        
-        for server_name, server in self.mcp_servers.items():
-            # Use cached tools from discovery
-            for tool in server.tools:
-                # Add server name to tool description for clarity
-                tool_copy = tool.copy()
-                original_desc = tool_copy.get('description', '')
-                tool_copy['description'] = f"[{server_name}] {original_desc}"
-                all_tools.append(tool_copy)
-        
-        logger.info(f"✅ Returning {len(all_tools)} tools")
-        
-        return {
-            'jsonrpc': '2.0',
-            'id': request.get('id'),
-            'result': {
-                'tools': all_tools
-            }
-        }
+            logger.error(f"Error handling {method}: {e}")
+            return {'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32603, 'message': str(e)}}
     
-    def handle_call_tool(self, request: Dict) -> Dict:
-        """Intercept tool call, log it, check policy, forward to server"""
-        params = request.get('params', {})
-        tool_name = params.get('name')
+    def _handle_call_tool(self, request_id, params: Dict) -> Dict:
+        tool_name = params.get('name', '')
         arguments = params.get('arguments', {})
-        
-        logger.info(f"🔧 Tool call: {tool_name}")
-        logger.debug(f"   Arguments: {json.dumps(arguments)[:200]}")
-        
-        # Find which server owns this tool
         server_name = self.tool_to_server.get(tool_name)
         if not server_name:
-            logger.error(f"Unknown tool: {tool_name}")
-            return {
-                'jsonrpc': '2.0',
-                'id': request.get('id'),
-                'error': {
-                    'code': -32602,
-                    'message': f'Unknown tool: {tool_name}'
-                }
-            }
-        
-        # 1. Log to LangProtect backend (security scan)
-        log_result = self.log_to_backend('tools/call', params, server_name)
-        
-        # 2. Check if blocked by policy
-        status = log_result.get('status', '').lower()
-        if status == 'blocked':
-            reason = 'Policy violation'
-            detections = log_result.get('detections', {})
-            if 'MCPActionControl' in detections:
-                reason = detections['MCPActionControl'].get('reason', reason)
-            
-            logger.warning(f"🛡️  BLOCKED: {tool_name} - {reason}")
-            
-            return {
-                'jsonrpc': '2.0',
-                'id': request.get('id'),
-                'error': {
-                    'code': -32000,
-                    'message': f'🛡️ LangProtect: {reason}',
-                    'data': {
-                        'log_id': log_result.get('id'),
-                        'risk_score': log_result.get('risk_score'),
-                        'detections': detections
-                    }
-                }
-            }
-        
-        # 3. Forward to actual MCP server
-        logger.info(f"✅ ALLOWED: Forwarding to {server_name}")
-        
+            return {'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32602, 'message': f'Unknown tool: {tool_name}'}}
+        server = self.mcp_servers.get(server_name)
+        if not server:
+            return {'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32602, 'message': f'Server not found: {server_name}'}}
+        logger.info(f"Tool call: {server_name}.{tool_name}")
+        if self.auth:
+            scan_result = self.auth.scan(tool_name, arguments, server_name)
+            status = scan_result.get('status', '').lower()
+            if status == 'blocked':
+                reason = scan_result.get('detections', {}).get('MCPActionControl', {}).get('reason', 'Policy violation')
+                logger.warning(f"BLOCKED: {tool_name} - {reason}")
+                return {'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32000, 'message': f'LangProtect: {reason}'}}
+            logger.info(f"ALLOWED (log_id={scan_result.get('id')})")
         try:
-            server = self.mcp_servers[server_name]
-            response = server.call('tools/call', params)
-            
-            # Log successful response
+            response = server.call('tools/call', {'name': tool_name, 'arguments': arguments})
             if 'result' in response:
-                logger.info(f"✅ Tool {tool_name} completed successfully")
+                return {'jsonrpc': '2.0', 'id': request_id, 'result': response['result']}
             elif 'error' in response:
-                logger.warning(f"⚠️  Tool {tool_name} returned error: {response['error'].get('message')}")
-            
-            # Return response with original request ID
-            response['id'] = request.get('id')
+                return {'jsonrpc': '2.0', 'id': request_id, 'error': response['error']}
             return response
-            
         except Exception as e:
-            logger.error(f"Error executing {tool_name} on {server_name}: {e}")
-            return {
-                'jsonrpc': '2.0',
-                'id': request.get('id'),
-                'error': {
-                    'code': -32603,
-                    'message': f'Error executing tool: {str(e)}'
-                }
-            }
+            return {'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32603, 'message': f'Error executing tool: {e}'}}
     
     def run(self):
-        """Main loop: read MCP requests from stdin, write responses to stdout"""
-        logger.info("🚀 LangProtect Gateway started")
-        logger.info(f"📡 Connected to: {self.langprotect_url}")
-        logger.info(f"👤 User: {self.email}")
-        logger.info(f"🔧 Proxying {len(self.mcp_servers)} MCP servers")
-        logger.info("=" * 60)
-        
         try:
             for line in sys.stdin:
                 line = line.strip()
                 if not line:
                     continue
-                
                 try:
                     request = json.loads(line)
-                    response = self.handle_mcp_request(request)
-                    
-                    # Write response to stdout (MCP protocol)
-                    print(json.dumps(response), flush=True)
-                    
+                    response = self.handle_request(request)
+                    if response:
+                        print(json.dumps(response), flush=True)
                 except json.JSONDecodeError as e:
-                    logger.error(f"Invalid JSON: {e}")
-                    error_response = {
-                        'jsonrpc': '2.0',
-                        'error': {
-                            'code': -32700,
-                            'message': f'Parse error: {str(e)}'
-                        }
-                    }
-                    print(json.dumps(error_response), flush=True)
-        
+                    print(json.dumps({'jsonrpc': '2.0', 'error': {'code': -32700, 'message': f'Parse error: {e}'}}), flush=True)
         except KeyboardInterrupt:
-            logger.info("\n👋 Gateway shutting down...")
-        except Exception as e:
-            logger.error(f"Fatal error: {e}", exc_info=True)
-            sys.exit(1)
+            pass
+        finally:
+            self.shutdown()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='LangProtect MCP Gateway')
+    parser.add_argument('--mcp-json-path', type=str, help='Path to mcp.json configuration file')
+    parser.add_argument('--debug', action='store_true', help='Enable debug logging')
+    return parser.parse_args()
 
 
 def main():
-    """Entry point"""
-    # Check for setup mode
-    if len(sys.argv) > 1 and sys.argv[1] == 'setup':
-        logger.info("⚠️  Setup mode not yet implemented")
-        logger.info("Please set environment variables:")
-        logger.info("  LANGPROTECT_URL=https://api.langprotect.com")
-        logger.info("  LANGPROTECT_EMAIL=your@email.com")
-        logger.info("  LANGPROTECT_PASSWORD=yourpassword")
+    args = parse_args()
+    if args.debug:
+        os.environ['DEBUG'] = 'true'
+        logging.getLogger('langprotect-gateway').setLevel(logging.DEBUG)
+    gateway = LangProtectGateway(mcp_json_path=args.mcp_json_path)
+    if not gateway.initialize():
         sys.exit(1)
-    
-    # Start gateway
-    gateway = LangProtectGateway()
     gateway.run()