460 lines
15 KiB
Python
460 lines
15 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Turnstone MCP Server — Improved Multi-Tool Server
|
||
|
||
Provides 4 tools:
|
||
1. execute_python_code — sandboxed Python execution with timeouts & limits
|
||
2. playwright_navigate — browser navigation with timeout & validation
|
||
3. playwright_get_page_text — rendered page text extraction with timeout
|
||
4. extract_to_markdown — format conversion with size limits & validation
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import re
|
||
import json
|
||
import logging
|
||
import asyncio
|
||
import subprocess
|
||
import tempfile
|
||
import secrets
|
||
import time
|
||
from typing import Dict, Optional
|
||
from urllib.parse import urlparse
|
||
|
||
from fastmcp import FastMCP
|
||
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
||
from starlette.middleware import Middleware as StarletteMiddleware
|
||
from starlette.middleware.cors import CORSMiddleware
|
||
from starlette.responses import FileResponse, JSONResponse
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Logging (JSON Format)
|
||
# ---------------------------------------------------------------------------
|
||
class JSONLogFormatter(logging.Formatter):
|
||
"""Custom JSON formatter for professional structured logging."""
|
||
def format(self, record):
|
||
log_record = {
|
||
"timestamp": self.formatTime(record, self.datefmt),
|
||
"level": record.levelname,
|
||
"logger": record.name,
|
||
"message": record.getMessage(),
|
||
"file": f"{record.filename}:{record.lineno}"
|
||
}
|
||
if record.exc_info:
|
||
log_record["exception"] = self.formatException(record.exc_info)
|
||
return json.dumps(log_record)
|
||
|
||
json_handler = logging.StreamHandler(sys.stdout)
|
||
json_handler.setFormatter(JSONLogFormatter())
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
handlers=[json_handler],
|
||
force=True # Forces Uvicorn/FastAPI to also use this JSON handler
|
||
)
|
||
logger = logging.getLogger("turnstone-mcp")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Configuration (env-overridable)
|
||
# ---------------------------------------------------------------------------
|
||
PYTHON_EXEC_TIMEOUT = int(os.environ.get("PYTHON_EXEC_TIMEOUT", "30")) # seconds
|
||
PYTHON_MAX_OUTPUT = int(os.environ.get("PYTHON_MAX_OUTPUT", "65536")) # bytes
|
||
PYTHON_MAX_MEMORY_MB = int(os.environ.get("PYTHON_MAX_MEMORY_MB", "256"))
|
||
|
||
PLAYWRIGHT_NAV_TIMEOUT = int(os.environ.get("PLAYWRIGHT_NAV_TIMEOUT", "30000")) # ms
|
||
PLAYWRIGHT_TEXT_TIMEOUT = int(os.environ.get("PLAYWRIGHT_TEXT_TIMEOUT", "15000")) # ms
|
||
|
||
EXTRACT_MAX_FILE_SIZE = int(os.environ.get("EXTRACT_MAX_FILE_SIZE", "10485760")) # 10 MB
|
||
|
||
# Defaults to /workspace to match Docker volumes
|
||
WORKSPACE_MOUNT = os.path.abspath(os.environ.get("WORKSPACE_MOUNT", "/workspace"))
|
||
|
||
logger.info(f"Workspace Root configured to: {WORKSPACE_MOUNT}")
|
||
|
||
PUBLIC_BASE_URL = os.environ.get("PUBLIC_BASE_URL", "http://localhost:8000")
|
||
DOWNLOAD_DEFAULT_TTL = 10
|
||
DOWNLOAD_MIN_TTL = 1
|
||
DOWNLOAD_MAX_TTL = 60
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# FastMCP app
|
||
# ---------------------------------------------------------------------------
|
||
mcp = FastMCP("turnstone-mcp")
|
||
|
||
TOKEN_STORE: Dict[str, dict] = {}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Playwright lifecycle (stateful browser session)
|
||
# ---------------------------------------------------------------------------
|
||
class PlaywrightBrowserManager:
|
||
"""Encapsulates Playwright state to prevent garbage collection teardowns."""
|
||
|
||
def __init__(self):
|
||
self.context_manager = None
|
||
self.playwright = None
|
||
self.browser = None
|
||
self.context = None
|
||
self.page = None
|
||
|
||
async def get_page(self):
|
||
"""Lazily initialize and return a headless Chromium page."""
|
||
if self.page is None:
|
||
from playwright.async_api import async_playwright
|
||
|
||
logger.info("Starting headless Playwright session…")
|
||
|
||
self.context_manager = async_playwright()
|
||
self.playwright = await self.context_manager.start()
|
||
|
||
self.browser = await self.playwright.chromium.launch(headless=True)
|
||
self.context = await self.browser.new_context(
|
||
viewport={"width": 1280, "height": 720},
|
||
user_agent=(
|
||
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
|
||
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
||
"Chrome/124.0.0.0 Safari/537.36"
|
||
),
|
||
)
|
||
self.page = await self.context.new_page()
|
||
|
||
return self.page
|
||
|
||
# Initialize singleton manager
|
||
browser_manager = PlaywrightBrowserManager()
|
||
|
||
def _resolve_safe_path(relative_path: str) -> str:
|
||
"""Resolve a relative path against WORKSPACE_MOUNT, blocking traversal."""
|
||
target = os.path.abspath(os.path.join(WORKSPACE_MOUNT, relative_path))
|
||
if not target.startswith(WORKSPACE_MOUNT):
|
||
raise PermissionError("Access denied: path escapes WORKSPACE_MOUNT")
|
||
return target
|
||
|
||
|
||
def _validate_download_path(raw_path: str) -> str:
|
||
canonical_path = os.path.realpath(raw_path)
|
||
if not canonical_path.startswith(WORKSPACE_MOUNT + "/"):
|
||
raise PermissionError("Access denied: path escapes WORKSPACE_MOUNT")
|
||
if not os.path.exists(canonical_path):
|
||
raise FileNotFoundError(f"File not found: {raw_path}")
|
||
if os.path.isdir(canonical_path):
|
||
raise IsADirectoryError(
|
||
"Target path is a directory. Please compress it to an archive first."
|
||
)
|
||
return canonical_path
|
||
|
||
|
||
def _purge_expired_tokens():
|
||
now = time.time()
|
||
expired_keys = [k for k, v in TOKEN_STORE.items() if now > v["expires_at"]]
|
||
for k in expired_keys:
|
||
TOKEN_STORE.pop(k, None)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Health endpoint
|
||
# ---------------------------------------------------------------------------
|
||
@mcp.custom_route("/health", methods=["GET"])
|
||
async def health_check(request):
|
||
return JSONResponse({"status": "healthy", "service": "turnstone-mcp"})
|
||
|
||
|
||
@mcp.custom_route("/download", methods=["GET"])
|
||
async def download_file(request):
|
||
_purge_expired_tokens()
|
||
|
||
token = request.query_params.get("token")
|
||
if not token:
|
||
return JSONResponse({"detail": "Missing token parameter."}, status_code=400)
|
||
|
||
record = TOKEN_STORE.get(token)
|
||
if not record:
|
||
return JSONResponse(
|
||
{"detail": "Invalid or previously used download link."}, status_code=403
|
||
)
|
||
|
||
del TOKEN_STORE[token]
|
||
|
||
if time.time() > record["expires_at"]:
|
||
return JSONResponse({"detail": "Download link has expired."}, status_code=410)
|
||
|
||
file_path = record["file_path"]
|
||
if not os.path.exists(file_path):
|
||
return JSONResponse(
|
||
{"detail": "File no longer exists on host."}, status_code=404
|
||
)
|
||
|
||
return FileResponse(
|
||
path=file_path,
|
||
filename=os.path.basename(file_path),
|
||
media_type="application/octet-stream",
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Semantic-logging middleware
|
||
# ---------------------------------------------------------------------------
|
||
class ToolObserverMiddleware(Middleware):
|
||
"""Log incoming tool invocations and responses."""
|
||
|
||
async def on_call_tool(self, context: MiddlewareContext, call_next):
|
||
tool_name = getattr(context.message, "name", "unknown")
|
||
arguments = getattr(context.message, "arguments", {})
|
||
|
||
logger.info(f"LLM REQUEST | Tool: '{tool_name}' | Args: {arguments}")
|
||
try:
|
||
result = await call_next(context)
|
||
content = getattr(result, "content", result)
|
||
preview = str(content)[:250].replace('\n', ' ') # Clean up newlines for JSON reading
|
||
logger.info(f"SERVER RESPONSE | Tool: '{tool_name}' completed | Preview: {preview}")
|
||
return result
|
||
except Exception as e:
|
||
logger.error(f"SERVER ERROR | Tool '{tool_name}' crashed: {e}")
|
||
raise
|
||
|
||
mcp.add_middleware(ToolObserverMiddleware())
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CORS
|
||
# ---------------------------------------------------------------------------
|
||
middleware_config = [
|
||
StarletteMiddleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_methods=["GET", "POST", "DELETE", "OPTIONS"],
|
||
allow_headers=[
|
||
"mcp-protocol-version",
|
||
"mcp-session-id",
|
||
"Authorization",
|
||
"Content-Type",
|
||
],
|
||
expose_headers=["mcp-session-id"],
|
||
)
|
||
]
|
||
|
||
# ===================================================================
|
||
# Tool 1 — execute_python_code
|
||
# ===================================================================
|
||
|
||
@mcp.tool()
|
||
def execute_python_code(code: str) -> str:
|
||
"""Execute Python code in a sandboxed environment.
|
||
|
||
Parameters
|
||
----------
|
||
code : str
|
||
The Python source to execute.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
stdout output, or error message if execution failed.
|
||
"""
|
||
if not code or not code.strip():
|
||
return "Error: code cannot be empty"
|
||
|
||
if len(code) > 50_000:
|
||
return "Error: code exceeds maximum length (50,000 chars)"
|
||
|
||
memory_limit_code = f"""
|
||
try:
|
||
import resource
|
||
desired = {PYTHON_MAX_MEMORY_MB} * 1024 * 1024
|
||
resource.setrlimit(resource.RLIMIT_AS, (desired, desired))
|
||
except Exception:
|
||
pass
|
||
"""
|
||
full_code = memory_limit_code.strip() + "\n\n" + code
|
||
|
||
try:
|
||
with tempfile.NamedTemporaryFile(
|
||
mode="w", suffix=".py", delete=False
|
||
) as tmp:
|
||
tmp.write(full_code)
|
||
tmp_path = tmp.name
|
||
|
||
proc = subprocess.run(
|
||
[sys.executable, tmp_path],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=PYTHON_EXEC_TIMEOUT,
|
||
)
|
||
except subprocess.TimeoutExpired:
|
||
return f"Error: execution timed out after {PYTHON_EXEC_TIMEOUT}s"
|
||
except Exception as e:
|
||
return f"Error: {e}"
|
||
finally:
|
||
try:
|
||
os.unlink(tmp_path)
|
||
except OSError:
|
||
pass
|
||
|
||
stdout = proc.stdout[:PYTHON_MAX_OUTPUT]
|
||
stderr = proc.stderr[:PYTHON_MAX_OUTPUT]
|
||
|
||
if proc.returncode != 0:
|
||
return f"Exit code: {proc.returncode}\nStderr: {stderr}\nStdout: {stdout}"
|
||
if stdout:
|
||
return stdout
|
||
return "Code executed successfully with no stdout output."
|
||
|
||
# ===================================================================
|
||
# Tool 2 — playwright_navigate
|
||
# ===================================================================
|
||
|
||
@mcp.tool()
|
||
async def playwright_navigate(url: str) -> str:
|
||
"""Navigate the persistent browser to a URL.
|
||
|
||
Parameters
|
||
----------
|
||
url : str
|
||
Target URL. Will be prefixed with https:// if missing scheme.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
Navigation result including page title, or error message.
|
||
"""
|
||
try:
|
||
parsed = urlparse(url)
|
||
if not parsed.scheme or not parsed.netloc:
|
||
if re.match(r"^[a-zA-Z0-9]", url):
|
||
url = "https://" + url
|
||
else:
|
||
return f"Error: invalid URL '{url}'"
|
||
|
||
page = await browser_manager.get_page()
|
||
response = await page.goto(
|
||
url, wait_until="domcontentloaded", timeout=PLAYWRIGHT_NAV_TIMEOUT
|
||
)
|
||
|
||
title = await page.title()
|
||
status = response.status if response else "unknown"
|
||
return f"Successfully loaded {url} (status {status})\nPage Title: {title}"
|
||
except Exception as e:
|
||
return f"Navigation error: {e}"
|
||
|
||
# ===================================================================
|
||
# Tool 3 — playwright_get_page_text
|
||
# ===================================================================
|
||
|
||
@mcp.tool()
|
||
async def playwright_get_page_text() -> str:
|
||
"""Retrieve the rendered page text from the current browser session.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
The innerText of document.body, or an error message.
|
||
"""
|
||
try:
|
||
page = browser_manager.page
|
||
if not page:
|
||
return "Error: no active browser session. Run playwright_navigate first."
|
||
|
||
text_content = await asyncio.wait_for(
|
||
page.evaluate("() => document.body.innerText"),
|
||
timeout=PLAYWRIGHT_TEXT_TIMEOUT / 1000.0,
|
||
)
|
||
return f"=== CURRENT BROWSER TEXT SCREEN ===\n\n{text_content.strip()}"
|
||
except asyncio.TimeoutError:
|
||
return "Error: Page text extraction timed out."
|
||
except Exception as e:
|
||
return f"Extraction error: {e}"
|
||
|
||
# ===================================================================
|
||
# Tool 4 — extract_to_markdown
|
||
# ===================================================================
|
||
|
||
SUPPORTED_EXTENSIONS = {
|
||
".pdf", ".docx", ".pptx", ".xlsx", ".html", ".htm",
|
||
".md", ".txt", ".csv", ".json", ".xml", ".epub",
|
||
".odt", ".rtf", ".jpg", ".jpeg", ".png", ".gif",
|
||
".webp", ".bmp",
|
||
}
|
||
|
||
@mcp.tool()
|
||
def extract_to_markdown(relative_path: str) -> str:
|
||
"""Convert a file to Markdown using markitdown.
|
||
|
||
Parameters
|
||
----------
|
||
relative_path : str
|
||
Path relative to WORKSPACE_MOUNT.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
Markdown content, or an error message.
|
||
"""
|
||
try:
|
||
target = _resolve_safe_path(relative_path)
|
||
except PermissionError as e:
|
||
return f"Error: {e}"
|
||
|
||
if not os.path.isfile(target):
|
||
return f"Error: file not found at {relative_path}"
|
||
|
||
file_size = os.path.getsize(target)
|
||
if file_size > EXTRACT_MAX_FILE_SIZE:
|
||
return f"Error: file size ({file_size} bytes) exceeds limit ({EXTRACT_MAX_FILE_SIZE} bytes)"
|
||
|
||
ext = os.path.splitext(target)[1].lower()
|
||
if ext and ext not in SUPPORTED_EXTENSIONS:
|
||
return f"Error: unsupported file type '{ext}'. Supported: {', '.join(sorted(SUPPORTED_EXTENSIONS))}"
|
||
|
||
try:
|
||
from markitdown import MarkItDown
|
||
mid = MarkItDown()
|
||
result = mid.convert(target)
|
||
return result.text_content
|
||
except ImportError:
|
||
return "Error: markitdown library is not installed."
|
||
except Exception as e:
|
||
return f"Extraction error: {e}"
|
||
|
||
# ===================================================================
|
||
# Tool 5 — generate_download_link
|
||
# ===================================================================
|
||
|
||
@mcp.tool()
|
||
def generate_download_link(path: str, ttl_minutes: int = DOWNLOAD_DEFAULT_TTL) -> str:
|
||
"""Generates a secure, single-use, temporary HTTPS download link for a file
|
||
stored within the /workspace directory.
|
||
|
||
Parameters
|
||
----------
|
||
path : str
|
||
The absolute path to the file inside the container (must reside within /workspace).
|
||
ttl_minutes : int
|
||
Validity period of the link in minutes. Default is 10. Clamped to 1–60.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
A clickable download URL, or an error message if validation fails.
|
||
"""
|
||
try:
|
||
abs_path = _validate_download_path(path)
|
||
except (PermissionError, FileNotFoundError, IsADirectoryError) as e:
|
||
return f"Error: {e}"
|
||
|
||
_purge_expired_tokens()
|
||
|
||
effective_ttl = max(DOWNLOAD_MIN_TTL, min(ttl_minutes, DOWNLOAD_MAX_TTL))
|
||
token = secrets.token_urlsafe(32)
|
||
expires_at = time.time() + (effective_ttl * 60)
|
||
|
||
TOKEN_STORE[token] = {
|
||
"file_path": abs_path,
|
||
"expires_at": expires_at,
|
||
}
|
||
|
||
return f"{PUBLIC_BASE_URL}/download?token={token}"
|
||
|
||
|
||
# ===================================================================
|
||
# Execution
|
||
# ===================================================================
|
||
if __name__ == "__main__":
|
||
mcp.run(transport="http", host="0.0.0.0", port=8000, middleware=middleware_config)
|