WORKSPACE_DIR env var was never set in docker-compose.yml, so the default always applied anyway. The container-internal mount point is always /workspace, so there's no reason to make it configurable.
422 lines
14 KiB
Python
422 lines
14 KiB
Python
"""
|
|
MCP HTTP Server — persistent PTY sessions + background jobs
|
|
Implements MCP Streamable HTTP transport (JSON-RPC 2.0)
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import signal
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from typing import Any
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
|
log = logging.getLogger("mcp")
|
|
|
|
WORKSPACE = "/workspace"
|
|
|
|
import pyte
|
|
import ptyprocess
|
|
import uvicorn
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import JSONResponse
|
|
|
|
app = FastAPI(title="MCP Bash Server", version="1.0.0")
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
# ─── 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)},
|
|
cwd=WORKSPACE,
|
|
)
|
|
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 = "".join(self.screen.buffer[y][x].data for x in range(self.cols))
|
|
lines.append(row.rstrip())
|
|
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",
|
|
cwd=WORKSPACE,
|
|
)
|
|
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] = {}
|
|
|
|
# ─── Tool definitions ─────────────────────────────────────────────────────────
|
|
|
|
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": {}},
|
|
},
|
|
]
|
|
|
|
# ─── Tool dispatch ────────────────────────────────────────────────────────────
|
|
|
|
async def _dispatch(name: str, args: dict) -> dict:
|
|
if name == "exec":
|
|
cmd = args.get("cmd", "")
|
|
timeout = int(args.get("timeout", 30))
|
|
log.info("exec: %s", cmd)
|
|
try:
|
|
r = subprocess.run(
|
|
cmd, shell=True, capture_output=True, text=True,
|
|
timeout=timeout, executable="/bin/bash",
|
|
cwd=WORKSPACE,
|
|
)
|
|
return {"stdout": r.stdout, "stderr": r.stderr, "exit_code": r.returncode}
|
|
except subprocess.TimeoutExpired:
|
|
return {"error": f"Command timed out after {timeout}s"}
|
|
|
|
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)
|
|
return {"session_id": sess.id, "cols": cols, "rows": rows}
|
|
|
|
elif name == "pty_send":
|
|
sess = _get_pty(args["session_id"])
|
|
text = _decode_escapes(args.get("text", ""))
|
|
log.info("pty_send [%s]: %r", sess.id[:8], text)
|
|
sess.send(text)
|
|
return {"sent": len(text)}
|
|
|
|
elif name == "pty_read":
|
|
sess = _get_pty(args["session_id"])
|
|
wait_ms = int(args.get("wait_ms", 200))
|
|
if wait_ms > 0:
|
|
await asyncio.sleep(wait_ms / 1000)
|
|
return sess.snapshot()
|
|
|
|
elif name == "pty_kill":
|
|
sess = _get_pty(args["session_id"])
|
|
sess.kill()
|
|
del _pty_sessions[sess.id]
|
|
return {"killed": sess.id}
|
|
|
|
elif name == "pty_list":
|
|
return {
|
|
"sessions": [
|
|
{"session_id": sid, "alive": s.proc.isalive(), "created_at": s.created_at}
|
|
for sid, s in _pty_sessions.items()
|
|
]
|
|
}
|
|
|
|
elif name == "job_run":
|
|
job = BackgroundJob(args.get("cmd", ""))
|
|
log.info("job_run [%s]: %s", job.id[:8], job.cmd)
|
|
_jobs[job.id] = job
|
|
return {"job_id": job.id, "status": job.status}
|
|
|
|
elif name == "job_output":
|
|
job = _get_job(args["job_id"])
|
|
offset = int(args.get("offset", 0))
|
|
lines, total = job.output(offset)
|
|
return {"job_id": job.id, "offset": offset, "lines": lines, "total_lines": total}
|
|
|
|
elif name == "job_status":
|
|
job = _get_job(args["job_id"])
|
|
_, total = job.output(0)
|
|
return {
|
|
"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,
|
|
}
|
|
|
|
elif name == "job_kill":
|
|
job = _get_job(args["job_id"])
|
|
job.kill()
|
|
return {"killed": job.id}
|
|
|
|
elif name == "job_list":
|
|
return {
|
|
"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 ValueError(f"Unknown tool: {name}")
|
|
|
|
|
|
# ─── MCP JSON-RPC 2.0 endpoint ────────────────────────────────────────────────
|
|
|
|
async def _handle_rpc(body: dict) -> dict | None:
|
|
method = body.get("method", "")
|
|
params = body.get("params") or {}
|
|
req_id = body.get("id")
|
|
|
|
try:
|
|
if method == "initialize":
|
|
result = {
|
|
"protocolVersion": "2024-11-05",
|
|
"capabilities": {"tools": {}},
|
|
"serverInfo": {"name": "mcp-bash", "version": "1.0.0"},
|
|
}
|
|
elif method in ("notifications/initialized", "notifications/cancelled"):
|
|
return None # notifications need no response
|
|
elif method == "ping":
|
|
result = {}
|
|
elif method == "tools/list":
|
|
result = {"tools": TOOLS}
|
|
elif method == "tools/call":
|
|
name = params.get("name", "")
|
|
args = params.get("arguments") or {}
|
|
data = await _dispatch(name, args)
|
|
result = {"content": [{"type": "text", "text": json.dumps(data)}]}
|
|
else:
|
|
return {"jsonrpc": "2.0", "id": req_id,
|
|
"error": {"code": -32601, "message": f"Method not found: {method}"}}
|
|
except Exception as exc:
|
|
return {"jsonrpc": "2.0", "id": req_id,
|
|
"error": {"code": -32603, "message": str(exc)}}
|
|
|
|
if req_id is None:
|
|
return None
|
|
return {"jsonrpc": "2.0", "id": req_id, "result": result}
|
|
|
|
|
|
@app.post("/")
|
|
@app.post("/mcp")
|
|
async def mcp_endpoint(request: Request):
|
|
body = await request.json()
|
|
response = await _handle_rpc(body)
|
|
if response is None:
|
|
return JSONResponse({}, status_code=202)
|
|
return JSONResponse(response)
|
|
|
|
|
|
@app.get("/")
|
|
def health():
|
|
return {"service": "mcp-bash", "version": "1.0.0"}
|
|
|
|
|
|
# ─── 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:
|
|
return m.group(0).encode("utf-8").decode("unicode_escape")
|
|
return _ESCAPE_RE.sub(_replace, text)
|
|
|
|
|
|
def _get_pty(session_id: str) -> PtySession:
|
|
sess = _pty_sessions.get(session_id)
|
|
if not sess:
|
|
raise ValueError(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 ValueError(f"Job not found: {job_id}")
|
|
return job
|
|
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|