mirror of
https://github.com/invariantlabs-ai/invariant-gateway.git
synced 2026-07-06 02:47:50 +02:00
Add mcp.py
This commit is contained in:
Hemang
+5
-7
@@ -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
|
||||
|
||||
@@ -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)
|
||||
):
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
+2
-1
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user