init
This commit is contained in:
Sidney Marvin Fricke
+394
@@ -0,0 +1,394 @@
|
||||
"""
|
||||
MCP HTTP Server — persistent PTY sessions + background jobs
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
import pyte
|
||||
import ptyprocess
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
app = FastAPI(title="MCP Bash Server", version="1.0.0")
|
||||
|
||||
# ─── Stores ──────────────────────────────────────────────────────────────────
|
||||
|
||||
class PtySession:
|
||||
def __init__(self, cols: int = 220, rows: int = 50):
|
||||
self.id = str(uuid.uuid4())
|
||||
self.cols = cols
|
||||
self.rows = rows
|
||||
self.screen = pyte.Screen(cols, rows)
|
||||
self.stream = pyte.ByteStream(self.screen)
|
||||
self._lock = threading.Lock()
|
||||
self.created_at = time.time()
|
||||
|
||||
self.proc = ptyprocess.PtyProcess.spawn(
|
||||
["/bin/bash", "--login"],
|
||||
dimensions=(rows, cols),
|
||||
env={**os.environ, "TERM": "xterm-256color", "COLUMNS": str(cols), "LINES": str(rows)},
|
||||
)
|
||||
self._reader = threading.Thread(target=self._read_loop, daemon=True)
|
||||
self._reader.start()
|
||||
|
||||
def _read_loop(self):
|
||||
while self.proc.isalive():
|
||||
try:
|
||||
data = self.proc.read(4096)
|
||||
except EOFError:
|
||||
break
|
||||
with self._lock:
|
||||
self.stream.feed(data)
|
||||
|
||||
def send(self, text: str):
|
||||
self.proc.write(text.encode())
|
||||
|
||||
def snapshot(self) -> dict[str, Any]:
|
||||
with self._lock:
|
||||
lines = []
|
||||
for y in range(self.rows):
|
||||
row = ""
|
||||
for x in range(self.cols):
|
||||
char = self.screen.buffer[y][x]
|
||||
row += char.data
|
||||
lines.append(row.rstrip())
|
||||
# trim trailing blank lines
|
||||
while lines and lines[-1] == "":
|
||||
lines.pop()
|
||||
cursor = {"x": self.screen.cursor.x, "y": self.screen.cursor.y}
|
||||
return {"lines": lines, "cursor": cursor, "alive": self.proc.isalive()}
|
||||
|
||||
def kill(self):
|
||||
try:
|
||||
self.proc.terminate(force=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class BackgroundJob:
|
||||
def __init__(self, cmd: str):
|
||||
self.id = str(uuid.uuid4())
|
||||
self.cmd = cmd
|
||||
self.output_lines: list[str] = []
|
||||
self.exit_code: int | None = None
|
||||
self.started_at = time.time()
|
||||
self.finished_at: float | None = None
|
||||
self._lock = threading.Lock()
|
||||
|
||||
self._proc = subprocess.Popen(
|
||||
cmd,
|
||||
shell=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
executable="/bin/bash",
|
||||
)
|
||||
self._reader = threading.Thread(target=self._read_loop, daemon=True)
|
||||
self._reader.start()
|
||||
|
||||
def _read_loop(self):
|
||||
for line in self._proc.stdout:
|
||||
with self._lock:
|
||||
self.output_lines.append(line.rstrip("\n"))
|
||||
self._proc.wait()
|
||||
with self._lock:
|
||||
self.exit_code = self._proc.returncode
|
||||
self.finished_at = time.time()
|
||||
|
||||
@property
|
||||
def status(self) -> str:
|
||||
if self.exit_code is None:
|
||||
return "running"
|
||||
return "done" if self.exit_code == 0 else "failed"
|
||||
|
||||
def output(self, offset: int = 0) -> tuple[list[str], int]:
|
||||
with self._lock:
|
||||
return self.output_lines[offset:], len(self.output_lines)
|
||||
|
||||
def kill(self):
|
||||
try:
|
||||
self._proc.send_signal(signal.SIGTERM)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
_pty_sessions: dict[str, PtySession] = {}
|
||||
_jobs: dict[str, BackgroundJob] = {}
|
||||
|
||||
# ─── MCP Tool Discovery ───────────────────────────────────────────────────────
|
||||
|
||||
TOOLS = [
|
||||
{
|
||||
"name": "exec",
|
||||
"description": "Run a shell command and return stdout+stderr synchronously (max 30 s).",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"cmd": {"type": "string", "description": "Shell command to execute"},
|
||||
"timeout": {"type": "integer", "description": "Timeout in seconds (default 30)", "default": 30},
|
||||
},
|
||||
"required": ["cmd"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "pty_create",
|
||||
"description": "Open a persistent PTY (interactive terminal) session. Returns session_id.",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"cols": {"type": "integer", "default": 220},
|
||||
"rows": {"type": "integer", "default": 50},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "pty_send",
|
||||
"description": "Send text/keystrokes to a PTY session. Use \\n for Enter, \\x03 for Ctrl-C, \\x04 for Ctrl-D.",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"session_id": {"type": "string"},
|
||||
"text": {"type": "string", "description": "Raw text to send (escape sequences allowed)"},
|
||||
},
|
||||
"required": ["session_id", "text"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "pty_read",
|
||||
"description": "Read the current screen state of a PTY session (rendered terminal grid).",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"session_id": {"type": "string"},
|
||||
"wait_ms": {"type": "integer", "default": 200, "description": "Wait N ms before reading (let output settle)"},
|
||||
},
|
||||
"required": ["session_id"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "pty_kill",
|
||||
"description": "Terminate and remove a PTY session.",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {"session_id": {"type": "string"}},
|
||||
"required": ["session_id"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "pty_list",
|
||||
"description": "List all open PTY sessions.",
|
||||
"inputSchema": {"type": "object", "properties": {}},
|
||||
},
|
||||
{
|
||||
"name": "job_run",
|
||||
"description": "Start a long-running background job. Returns job_id immediately.",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"cmd": {"type": "string", "description": "Shell command to run in background"},
|
||||
},
|
||||
"required": ["cmd"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "job_output",
|
||||
"description": "Retrieve output lines of a background job (optionally from an offset).",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"job_id": {"type": "string"},
|
||||
"offset": {"type": "integer", "default": 0, "description": "Start reading from this line index"},
|
||||
},
|
||||
"required": ["job_id"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "job_status",
|
||||
"description": "Get status (running/done/failed), exit code, and line count of a background job.",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {"job_id": {"type": "string"}},
|
||||
"required": ["job_id"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "job_kill",
|
||||
"description": "Send SIGTERM to a running background job.",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {"job_id": {"type": "string"}},
|
||||
"required": ["job_id"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "job_list",
|
||||
"description": "List all known background jobs with their current status.",
|
||||
"inputSchema": {"type": "object", "properties": {}},
|
||||
},
|
||||
]
|
||||
|
||||
# ─── MCP Endpoints ────────────────────────────────────────────────────────────
|
||||
|
||||
@app.get("/")
|
||||
def root():
|
||||
return {"service": "mcp-bash", "version": "1.0.0"}
|
||||
|
||||
|
||||
@app.get("/mcp/tools")
|
||||
def list_tools():
|
||||
return {"tools": TOOLS}
|
||||
|
||||
|
||||
class ToolCallRequest(BaseModel):
|
||||
name: str
|
||||
arguments: dict[str, Any] = {}
|
||||
|
||||
|
||||
@app.post("/mcp/call")
|
||||
async def call_tool(req: ToolCallRequest):
|
||||
name = req.name
|
||||
args = req.arguments
|
||||
|
||||
# ── exec ─────────────────────────────────────────────────────────────────
|
||||
if name == "exec":
|
||||
cmd = args.get("cmd", "")
|
||||
timeout = int(args.get("timeout", 30))
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd, shell=True, capture_output=True, text=True,
|
||||
timeout=timeout, executable="/bin/bash",
|
||||
)
|
||||
return _ok({"stdout": result.stdout, "stderr": result.stderr, "exit_code": result.returncode})
|
||||
except subprocess.TimeoutExpired:
|
||||
return _err(f"Command timed out after {timeout}s")
|
||||
|
||||
# ── pty_create ────────────────────────────────────────────────────────────
|
||||
elif name == "pty_create":
|
||||
cols = int(args.get("cols", 220))
|
||||
rows = int(args.get("rows", 50))
|
||||
sess = PtySession(cols=cols, rows=rows)
|
||||
_pty_sessions[sess.id] = sess
|
||||
await asyncio.sleep(0.3) # let bash prompt appear
|
||||
return _ok({"session_id": sess.id, "cols": cols, "rows": rows})
|
||||
|
||||
# ── pty_send ──────────────────────────────────────────────────────────────
|
||||
elif name == "pty_send":
|
||||
sess = _get_pty(args.get("session_id", ""))
|
||||
text = _decode_escapes(args.get("text", ""))
|
||||
sess.send(text)
|
||||
return _ok({"sent": len(text)})
|
||||
|
||||
# ── pty_read ──────────────────────────────────────────────────────────────
|
||||
elif name == "pty_read":
|
||||
sess = _get_pty(args.get("session_id", ""))
|
||||
wait_ms = int(args.get("wait_ms", 200))
|
||||
if wait_ms > 0:
|
||||
await asyncio.sleep(wait_ms / 1000)
|
||||
snap = sess.snapshot()
|
||||
return _ok(snap)
|
||||
|
||||
# ── pty_kill ──────────────────────────────────────────────────────────────
|
||||
elif name == "pty_kill":
|
||||
sess = _get_pty(args.get("session_id", ""))
|
||||
sess.kill()
|
||||
del _pty_sessions[sess.id]
|
||||
return _ok({"killed": sess.id})
|
||||
|
||||
# ── pty_list ──────────────────────────────────────────────────────────────
|
||||
elif name == "pty_list":
|
||||
return _ok({
|
||||
"sessions": [
|
||||
{"session_id": sid, "alive": s.proc.isalive(), "created_at": s.created_at}
|
||||
for sid, s in _pty_sessions.items()
|
||||
]
|
||||
})
|
||||
|
||||
# ── job_run ───────────────────────────────────────────────────────────────
|
||||
elif name == "job_run":
|
||||
job = BackgroundJob(args.get("cmd", ""))
|
||||
_jobs[job.id] = job
|
||||
return _ok({"job_id": job.id, "status": job.status})
|
||||
|
||||
# ── job_output ────────────────────────────────────────────────────────────
|
||||
elif name == "job_output":
|
||||
job = _get_job(args.get("job_id", ""))
|
||||
offset = int(args.get("offset", 0))
|
||||
lines, total = job.output(offset)
|
||||
return _ok({"job_id": job.id, "offset": offset, "lines": lines, "total_lines": total})
|
||||
|
||||
# ── job_status ────────────────────────────────────────────────────────────
|
||||
elif name == "job_status":
|
||||
job = _get_job(args.get("job_id", ""))
|
||||
_, total = job.output(0)
|
||||
return _ok({
|
||||
"job_id": job.id, "status": job.status,
|
||||
"exit_code": job.exit_code,
|
||||
"total_lines": total,
|
||||
"started_at": job.started_at, "finished_at": job.finished_at,
|
||||
})
|
||||
|
||||
# ── job_kill ──────────────────────────────────────────────────────────────
|
||||
elif name == "job_kill":
|
||||
job = _get_job(args.get("job_id", ""))
|
||||
job.kill()
|
||||
return _ok({"killed": job.id})
|
||||
|
||||
# ── job_list ──────────────────────────────────────────────────────────────
|
||||
elif name == "job_list":
|
||||
return _ok({
|
||||
"jobs": [
|
||||
{"job_id": jid, "cmd": j.cmd, "status": j.status,
|
||||
"exit_code": j.exit_code, "total_lines": j.output(0)[1]}
|
||||
for jid, j in _jobs.items()
|
||||
]
|
||||
})
|
||||
|
||||
raise HTTPException(status_code=404, detail=f"Unknown tool: {name}")
|
||||
|
||||
|
||||
# ─── Helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
_ESCAPE_RE = re.compile(r"\\(n|r|t|\\|x[0-9a-fA-F]{2}|u[0-9a-fA-F]{4})")
|
||||
|
||||
def _decode_escapes(text: str) -> str:
|
||||
"""Decode only ASCII escape sequences — leaves UTF-8 text intact."""
|
||||
def _replace(m: re.Match) -> str:
|
||||
s = m.group(0)
|
||||
return s.encode("utf-8").decode("unicode_escape")
|
||||
return _ESCAPE_RE.sub(_replace, text)
|
||||
|
||||
|
||||
def _ok(data: dict) -> JSONResponse:
|
||||
return JSONResponse({"ok": True, **data})
|
||||
|
||||
|
||||
def _err(msg: str, status: int = 400) -> JSONResponse:
|
||||
return JSONResponse({"ok": False, "error": msg}, status_code=status)
|
||||
|
||||
|
||||
def _get_pty(session_id: str) -> PtySession:
|
||||
sess = _pty_sessions.get(session_id)
|
||||
if not sess:
|
||||
raise HTTPException(status_code=404, detail=f"PTY session not found: {session_id}")
|
||||
return sess
|
||||
|
||||
|
||||
def _get_job(job_id: str) -> BackgroundJob:
|
||||
job = _jobs.get(job_id)
|
||||
if not job:
|
||||
raise HTTPException(status_code=404, detail=f"Job not found: {job_id}")
|
||||
return job
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
Reference in New Issue
Block a user