diff --git a/server/main.py b/server/main.py index d09a91b..638378d 100644 --- a/server/main.py +++ b/server/main.py @@ -1,8 +1,10 @@ """ MCP HTTP Server — persistent PTY sessions + background jobs +Implements MCP Streamable HTTP transport (JSON-RPC 2.0) """ import asyncio +import json import os import re import signal @@ -15,10 +17,9 @@ from typing import Any import pyte import ptyprocess import uvicorn -from fastapi import FastAPI, HTTPException +from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse -from pydantic import BaseModel app = FastAPI(title="MCP Bash Server", version="1.0.0") app.add_middleware( @@ -64,12 +65,8 @@ class PtySession: 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 + row = "".join(self.screen.buffer[y][x].data for x in range(self.cols)) 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} @@ -132,7 +129,7 @@ class BackgroundJob: _pty_sessions: dict[str, PtySession] = {} _jobs: dict[str, BackgroundJob] = {} -# ─── MCP Tool Discovery ─────────────────────────────────────────────────────── +# ─── Tool definitions ───────────────────────────────────────────────────────── TOOLS = [ { @@ -244,123 +241,143 @@ TOOLS = [ }, ] -# ─── MCP Endpoints ──────────────────────────────────────────────────────────── +# ─── Tool dispatch ──────────────────────────────────────────────────────────── -@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 ───────────────────────────────────────────────────────────────── +async def _dispatch(name: str, args: dict) -> dict: if name == "exec": cmd = args.get("cmd", "") timeout = int(args.get("timeout", 30)) try: - result = subprocess.run( + r = 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}) + return {"stdout": r.stdout, "stderr": r.stderr, "exit_code": r.returncode} except subprocess.TimeoutExpired: - return _err(f"Command timed out after {timeout}s") + return {"error": 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}) + await asyncio.sleep(0.3) + return {"session_id": sess.id, "cols": cols, "rows": rows} - # ── pty_send ────────────────────────────────────────────────────────────── elif name == "pty_send": - sess = _get_pty(args.get("session_id", "")) + sess = _get_pty(args["session_id"]) text = _decode_escapes(args.get("text", "")) sess.send(text) - return _ok({"sent": len(text)}) + return {"sent": len(text)} - # ── pty_read ────────────────────────────────────────────────────────────── elif name == "pty_read": - sess = _get_pty(args.get("session_id", "")) + sess = _get_pty(args["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) + return sess.snapshot() - # ── pty_kill ────────────────────────────────────────────────────────────── elif name == "pty_kill": - sess = _get_pty(args.get("session_id", "")) + sess = _get_pty(args["session_id"]) sess.kill() del _pty_sessions[sess.id] - return _ok({"killed": sess.id}) + return {"killed": sess.id} - # ── pty_list ────────────────────────────────────────────────────────────── elif name == "pty_list": - return _ok({ + return { "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}) + return {"job_id": job.id, "status": job.status} - # ── job_output ──────────────────────────────────────────────────────────── elif name == "job_output": - job = _get_job(args.get("job_id", "")) + job = _get_job(args["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}) + return {"job_id": job.id, "offset": offset, "lines": lines, "total_lines": total} - # ── job_status ──────────────────────────────────────────────────────────── elif name == "job_status": - job = _get_job(args.get("job_id", "")) + job = _get_job(args["job_id"]) _, total = job.output(0) - return _ok({ + return { "job_id": job.id, "status": job.status, - "exit_code": job.exit_code, - "total_lines": total, + "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 = _get_job(args["job_id"]) job.kill() - return _ok({"killed": job.id}) + return {"killed": job.id} - # ── job_list ────────────────────────────────────────────────────────────── elif name == "job_list": - return _ok({ + 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 HTTPException(status_code=404, detail=f"Unknown tool: {name}") + 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 ────────────────────────────────────────────────────────────────── @@ -370,30 +387,21 @@ _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 m.group(0).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}") + 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 HTTPException(status_code=404, detail=f"Job not found: {job_id}") + raise ValueError(f"Job not found: {job_id}") return job