Spaces:
Paused
Paused
File size: 4,068 Bytes
ea3710d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import sys
import io
import traceback
import base64
import os
import tempfile
import signal
from contextlib import redirect_stdout, redirect_stderr
from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationCapabilities
from mcp.server.stdio import stdio_server
import mcp.types as types
server = Server("python-runner")
def run_code(code: str) -> dict:
"""Execute Python code and return stdout, stderr, and images."""
stdout_capture = io.StringIO()
stderr_capture = io.StringIO()
images = []
tmpdir = tempfile.mkdtemp()
original_dir = os.getcwd()
os.chdir(tmpdir)
# Prepare a globals dict with matplotlib import
exec_globals = {"__builtins__": __builtins__}
try:
import matplotlib
matplotlib.use('Agg') # non-interactive backend
exec_globals["plt"] = __import__("matplotlib.pyplot")
except ImportError:
pass
def save_fig():
"""Helper to save current matplotlib figure."""
try:
import matplotlib.pyplot as plt
fig = plt.gcf()
if fig.get_axes():
buf = io.BytesIO()
fig.savefig(buf, format='png', bbox_inches='tight')
buf.seek(0)
img_b64 = base64.b64encode(buf.read()).decode()
images.append(img_b64)
plt.close(fig)
except Exception:
pass
# Timeout handler
def timeout_handler(signum, frame):
raise TimeoutError("Code execution timed out (10s)")
signal.signal(signal.SIGALRM, timeout_handler)
signal.alarm(10)
try:
with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
exec(code, exec_globals)
save_fig()
output = stdout_capture.getvalue()
errors = stderr_capture.getvalue()
signal.alarm(0)
except Exception as e:
save_fig()
output = stdout_capture.getvalue()
errors = stderr_capture.getvalue() + "\n" + traceback.format_exc()
signal.alarm(0)
finally:
os.chdir(original_dir)
# Clean up temp dir
for f in os.listdir(tmpdir):
os.unlink(os.path.join(tmpdir, f))
os.rmdir(tmpdir)
return {
"stdout": output,
"stderr": errors,
"images": images
}
@server.list_tools()
async def list_tools() -> list[types.Tool]:
return [
types.Tool(
name="execute_code",
description="Execute Python code and return output and any matplotlib figures as base64 PNG images.",
inputSchema={
"type": "object",
"properties": {
"code": {"type": "string", "description": "Python code to run"}
},
"required": ["code"]
}
)
]
@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[types.TextContent | types.ImageContent]:
if name == "execute_code":
result = run_code(arguments["code"])
contents = []
if result["stdout"]:
contents.append(types.TextContent(type="text", text=f"STDOUT:\n{result['stdout']}"))
if result["stderr"]:
contents.append(types.TextContent(type="text", text=f"STDERR:\n{result['stderr']}"))
for img in result["images"]:
contents.append(types.ImageContent(type="image", data=img, mimeType="image/png"))
if not contents:
contents.append(types.TextContent(type="text", text="(no output)"))
return contents
raise ValueError(f"Unknown tool: {name}")
async def main():
async with stdio_server() as (read_stream, write_stream):
await server.run(
read_stream,
write_stream,
InitializationCapabilities(
sampling=None,
experimental=None,
roots=None
),
notification_options=NotificationOptions(),
)
if __name__ == "__main__":
import asyncio
asyncio.run(main()) |