From 66c6dcb8f92ee479ee618a96641f66826396fb6b Mon Sep 17 00:00:00 2001 From: Hemang Date: Fri, 11 Apr 2025 06:09:44 +0200 Subject: [PATCH] Add mcp.py --- gateway/__main__.py | 12 +- gateway/common/request_context.py | 2 +- gateway/mcp/__init__.py | 0 gateway/mcp/mcp.py | 302 ++++++++++++++++++++++++++++++ gateway/mcp/mcp_context.py | 54 ++++++ pyproject.toml | 3 +- 6 files changed, 364 insertions(+), 9 deletions(-) create mode 100644 gateway/mcp/__init__.py create mode 100644 gateway/mcp/mcp.py create mode 100644 gateway/mcp/mcp_context.py diff --git a/gateway/__main__.py b/gateway/__main__.py index b351c45..b8c3961 100644 --- a/gateway/__main__.py +++ b/gateway/__main__.py @@ -2,6 +2,8 @@ import sys +from gateway.mcp import mcp + def main(): """Entry point for the Invariant Gateway.""" @@ -12,8 +14,7 @@ def main(): } def _help(): - """Prints the help message.""" - print("\nSupported Commands by invariant-gateway:\n") + """_prints the help message.""" for verb, description in actions.items(): print(f"{verb}: {description}") @@ -23,14 +24,11 @@ def main(): verb = sys.argv[1] if verb == "mcp": - # Use sys.argv[2:] to pass arguments to the MCP gateway - return 0 + return mcp.execute(sys.argv[2:]) if verb == "llm": - # Use sys.argv[2:] to pass arguments to the LLM gateway - print("LLM gateway via the invariant-gateway command is not implemented yet.") return 1 if verb == "help": _help() return 0 - print(f"Unknown action: {verb}") + print(f"[gateway/__main__.py] Unknown action: {verb}") return 1 diff --git a/gateway/common/request_context.py b/gateway/common/request_context.py index 945c8b1..c91a01a 100644 --- a/gateway/common/request_context.py +++ b/gateway/common/request_context.py @@ -82,7 +82,7 @@ class RequestContext: # if additionally provided, extract separate API key to use with guardrailing service guardrail_service_authorization = None - if ( + if request and ( guardrail_authorization := extract_guardrail_service_authorization_from_headers(request) ): diff --git a/gateway/mcp/__init__.py b/gateway/mcp/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/gateway/mcp/mcp.py b/gateway/mcp/mcp.py new file mode 100644 index 0000000..015d8fd --- /dev/null +++ b/gateway/mcp/mcp.py @@ -0,0 +1,302 @@ +"""Gateway for MCP (Model Context Protocol) integration with Invariant.""" + +import argparse +import asyncio +import sys +import subprocess +import json +import os +import threading +import signal + +from builtins import print as builtins_print +from contextlib import redirect_stdout + +from gateway.common.request_context import RequestContext +from gateway.integrations.guardrails import check_guardrails +from gateway.integrations.explorer import ( + fetch_guardrails_from_explorer, +) +from gateway.mcp.mcp_context import McpContext + + +def custom_print(ctx, *args, **kwargs): + """Custom print function to redirect output to log_out.""" + builtins_print(*args, **kwargs, file=ctx.log_out, flush=True) + + +def append_and_push_trace(ctx, message): + """ + Append a message to the trace if it exists or create a new one + and push it to the Invariant Explorer. + """ + try: + if ctx.trace_id is None: + ctx.trace.append(message) + response = ctx.client.create_request_and_push_trace( + messages=[ctx.trace], + dataset=ctx.explorer_dataset, + metadata=[{"source": "mcp", "tools": ctx.tools}], + ) + ctx.trace_id = response.id[0] + ctx.last_trace_length = len(ctx.trace) + else: + ctx.trace.append(message) + ctx.client.create_request_and_append_messages( + trace_id=ctx.trace_id, messages=ctx.trace[ctx.last_trace_length :] + ) + ctx.last_trace_length = len(ctx.trace) + except Exception as e: + custom_print(ctx, "Error pushing trace:", e) + + +def fetch_guardrails(ctx, dataset): + """Fetch guardrails from the Invariant Explorer.""" + # Use async fetch_guardrails_from_explorer in a thread + return asyncio.run( + fetch_guardrails_from_explorer( + dataset, "Bearer " + os.getenv("INVARIANT_API_KEY") + ) + ) + + +def check_blocking_guardrails(ctx, message, request): + """Check against blocking guardrails.""" + try: + guardrails = fetch_guardrails(ctx, ctx.explorer_dataset) + + custom_print(ctx, "Here are the guardrails: ", guardrails) + + context = RequestContext.create( + request_json=request, + dataset_name=ctx.explorer_dataset, + invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), + guardrails=guardrails, + ) + + if guardrails.blocking_guardrails: + with redirect_stdout(ctx.log_out): + return asyncio.run( + check_guardrails( + messages=ctx.trace + [message], + guardrails=guardrails.blocking_guardrails, + context=context, + ) + ) + else: + return {} + except Exception as e: + custom_print(ctx, "Error checking blocking guardrails:", e) + + +def hook_tool_call(ctx, request): + """ + Hook function to intercept tool calls. + Modify this function to change behavior for tool calls. + Returns the potentially modified request. + """ + tool_call = { + "id": f"call_{request.get('id')}", + "type": "function", + "function": { + "name": request["params"]["name"], + "arguments": request["params"]["arguments"], + }, + } + + message = {"role": "assistant", "content": "", "tool_calls": [tool_call]} + + # Check for blocking guardrails + result = check_blocking_guardrails(ctx, message, request) + append_and_push_trace(ctx, message) + return request + + +def hook_tool_result(ctx, result): + """ + Hook function to intercept tool results. + Modify this function to change behavior for tool results. + Returns the potentially modified result. + """ + method = ctx.id_to_method_mapping.get(result.get("id")) + call_id = f"call_{result.get('id')}" + + if method is None: + return result + elif method == "tools/call": + message = { + "role": "tool", + "content": {"result": result.get("result").get("content")}, + "error": result.get("result").get("error"), + "tool_call_id": call_id, + } + + # Check for blocking guardrails + guardrailing_result = check_blocking_guardrails(ctx, message, result) + + if guardrailing_result and guardrailing_result.get("errors", []): + result["result"]["content"] = [ + { + "type": "text", + "text": "[Invariant] Your MCP tool call was blocked for security reasons. Do not attempt to circumvent this block, rather explain to the user based on the following output what went wrong: \n" + + json.dumps(guardrailing_result["errors"]), + } + ] + + append_and_push_trace(ctx, message) + return result + elif method == "tools/list": + ctx.tools = result.get("result").get("tools") + return result + else: + return result + + +def forward_stdout(process, ctx, buffer_size=1): + """Read from the process stdout, parse JSON chunks, and forward to sys.stdout""" + buffer = b"" + + while True: + chunk = process.stdout.read(buffer_size) + if not chunk: + break + buffer += chunk + + try: + # Try parsing full JSON object from buffer + text = buffer.decode("utf-8") + obj = json.loads(text) + + obj = hook_tool_result(ctx, obj) + # clear the buffer + buffer = b"" + + # Forward the original JSON to stdout + json_output = json.dumps(obj).encode("utf-8") + b"\n" + sys.stdout.buffer.write(json_output) + sys.stdout.buffer.flush() + except (json.JSONDecodeError, UnicodeDecodeError): + # Wait for more data + continue + + +def forward_stderr(process, ctx, buffer_size=1): + """Read from the process stderr and write to sys.stderr""" + for line in iter(lambda: process.stderr.read(buffer_size), b""): + ctx.log_out.buffer.write(line) + ctx.log_out.buffer.flush() + + +def execute(args=None): + """Main function to execute the MCP gateway.""" + if "INVARIANT_API_KEY" not in os.environ: + print("[ERROR] INVARIANT_API_KEY environment variable is not set.") + sys.exit(1) + + # Split args at the "--exec" boundary + if args and "--exec" in args: + exec_index = args.index("--exec") + pre_exec_args = args[:exec_index] + post_exec_args = args[exec_index + 1 :] + else: + pre_exec_args = args or [] + post_exec_args = [] + + if not post_exec_args: + print("[ERROR] No command provided after --exec.") + sys.exit(1) + + # Parse pre-exec args using argparse + parser = argparse.ArgumentParser(description="MCP Gateway") + parser.add_argument("--directory", help="Working directory") + parser.add_argument("--verbose", action="store_true", help="Enable verbose output") + config = parser.parse_args(pre_exec_args) + # Initialize the singleton context using config + ctx = McpContext() + + # Can now use post_exec_args as your cmd + cmd = post_exec_args + + process = subprocess.Popen( + cmd, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + bufsize=0, # No buffering + ) + + # Start threads to forward stdout and stderr + stdout_thread = threading.Thread( + target=forward_stdout, args=(process, ctx), daemon=True + ) + stderr_thread = threading.Thread( + target=forward_stderr, args=(process, ctx), daemon=True + ) + stdout_thread.start() + stderr_thread.start() + + # Handle forwarding stdin and intercept tool calls + try: + current_chunk = b"" + + while True: + data = sys.stdin.buffer.read(1) + current_chunk += data + + if not data: + break + + # Try to decode and parse as JSON to check for tool calls + try: + text = current_chunk.decode("utf-8") + obj = json.loads(text) + # clear the current chunk + current_chunk = b"" + + if obj.get("method") is not None: + ctx.id_to_method_mapping[obj.get("id")] = obj.get("method") + + # Check if this is a tool call request + if obj.get("method") == "tools/call": + # Intercept and potentially modify the request + obj = hook_tool_call(ctx, obj) + # Convert back to bytes + data = json.dumps(obj).encode("utf-8") + + # Forward to the process + process.stdin.write(data + b"\n") + process.stdin.flush() + continue + else: + process.stdin.write(json.dumps(obj).encode("utf-8") + b"\n") + process.stdin.flush() + continue + except Exception: + # Not a complete or valid JSON, just pass through + pass + + except BrokenPipeError: + pass + except KeyboardInterrupt: + process.terminate() + + # Wait for process to terminate + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + +# Handle signals to ensure clean shutdown +def signal_handler(sig, frame): + """Handle signals for graceful shutdown.""" + ctx = McpContext() + custom_print(ctx, f"Received signal {sig}, shutting down...") + sys.exit(0) + + +if __name__ == "__main__": + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + execute(sys.argv) diff --git a/gateway/mcp/mcp_context.py b/gateway/mcp/mcp_context.py new file mode 100644 index 0000000..43eb849 --- /dev/null +++ b/gateway/mcp/mcp_context.py @@ -0,0 +1,54 @@ +"""Context manager for MCP (Model Context Protocol) gateway.""" + +import atexit +import os +import sys +from invariant_sdk.client import Client + + +class McpContext: + """Singleton class to manage MCP context and state.""" + + _instance = None + + def __new__(cls): + """Control instance creation to ensure only one instance exists.""" + if cls._instance is None: + cls._instance = super(McpContext, cls).__new__(cls) + cls._instance._initialized = False + return cls._instance + + def __init__(self): + """Initialize the singleton instance with default values (only once).""" + # Define _initialized attribute explicitly at the beginning to avoid warnings + # This is redundant but prevents warnings about accessing before definition + if not hasattr(self, "_initialized"): + self._initialized = False + + if self._initialized: + return + + def setup_logging(self): + """Set up logging to a file in the user's home directory. + + Uses proper resource management to ensure the file is closed on program exit. + """ + os.makedirs( + os.path.join(os.path.expanduser("~"), ".invariant"), exist_ok=True + ) + log_path = os.path.join(os.path.expanduser("~"), ".invariant", "mcp.log") + self.log_out = open(log_path, "a", buffering=1, encoding="utf-8") + atexit.register(self.log_out.close) + sys.stderr = self.log_out + + self.client = Client() + self.explorer_dataset = "mcp-capture" + self.trace = [] + self.tools = [] + self.trace_id = None + self.last_trace_length = 0 + self.guardrails = None + self.id_to_method_mapping = {} + setup_logging(self) + # Mark as initialized + self._initialized = True diff --git a/pyproject.toml b/pyproject.toml index 61edb73..dba21dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,13 +8,14 @@ dependencies = [ "fastapi==0.115.7", "httpx==0.28.1", "invariant-ai>=0.2.1", - "invariant-sdk>=0.0.10", + "invariant-sdk>=0.0.11", "starlette-compress==1.4.0", "uvicorn==0.34.0" ] [tool.setuptools.packages.find] where = ["."] +exclude = ["tests", "tests.*"] [project.scripts] invariant-gateway = "gateway.__main__:main"