Coverage for mcpgateway/translate.py: 80%

158 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-22 14:46 +0100

1# -*- coding: utf-8 -*- 

2""" mcpgateway.translate - bridges local JSON-RPC/stdio servers to HTTP/SSE 

3 

4Copyright 2025 

5SPDX-License-Identifier: Apache-2.0 

6Authors: Mihai Criveti 

7 

8Only the stdio→SSE direction is implemented for now. 

9 

10Usage 

11----- 

12# 1. expose an MCP server that talks JSON-RPC on stdio at :9000/sse 

13python -m mcpgateway.translate --stdio "uvenv run mcp-server-git" --port 9000 

14 

15# 2. from another shell / browser subscribe to the SSE stream 

16curl -N http://localhost:9000/sse # receive the stream 

17 

18# 3. send a test echo request 

19curl -X POST http://localhost:9000/message \ 

20 -H 'Content-Type: application/json' \ 

21 -d '{"jsonrpc":"2.0","id":1,"method":"echo","params":{"value":"hi"}}' 

22 

23# 4. proper MCP handshake and tool listing 

24curl -X POST http://localhost:9000/message \ 

25 -H 'Content-Type: application/json' \ 

26 -d '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"demo","version":"0.0.1"}}}' 

27 

28curl -X POST http://localhost:9000/message \ 

29 -H 'Content-Type: application/json' \ 

30 -d '{"jsonrpc":"2.0","id":2,"method":"tools/list"}' 

31 

32The SSE stream now emits JSON-RPC responses as `event: message` frames and sends 

33regular `event: keepalive` frames (default every 30s) so that proxies and 

34clients never time out. Each client receives a unique *session-id* that is 

35appended as a query parameter to the back-channel `/message` URL. 

36""" 

37 

38from __future__ import annotations 

39 

40import argparse 

41import asyncio 

42import json 

43import logging 

44import shlex 

45import signal 

46import sys 

47import uuid 

48from contextlib import suppress 

49from typing import AsyncIterator, List, Optional, Sequence 

50 

51import uvicorn 

52from fastapi import FastAPI, Request, Response, status 

53from fastapi.responses import PlainTextResponse 

54from sse_starlette.sse import EventSourceResponse 

55 

56LOGGER = logging.getLogger("mcpgateway.translate") 

57KEEP_ALIVE_INTERVAL = 30 # seconds ── matches the reference implementation 

58__all__ = ["main"] # for console-script entry-point 

59 

60 

61# ---------------------------------------------------------------------------# 

62# Helpers - trivial in-process Pub/Sub # 

63# ---------------------------------------------------------------------------# 

64class _PubSub: 

65 """Very small fan-out helper - one async Queue per subscriber.""" 

66 

67 def __init__(self) -> None: 

68 self._subscribers: List[asyncio.Queue[str]] = [] 

69 

70 async def publish(self, data: str) -> None: 

71 dead: List[asyncio.Queue[str]] = [] 

72 for q in self._subscribers: 

73 try: 

74 q.put_nowait(data) 

75 except asyncio.QueueFull: 

76 dead.append(q) 

77 for q in dead: 

78 with suppress(ValueError): 

79 self._subscribers.remove(q) 

80 

81 def subscribe(self) -> "asyncio.Queue[str]": 

82 q: asyncio.Queue[str] = asyncio.Queue(maxsize=1024) 

83 self._subscribers.append(q) 

84 return q 

85 

86 def unsubscribe(self, q: "asyncio.Queue[str]") -> None: 

87 with suppress(ValueError): 

88 self._subscribers.remove(q) 

89 

90 

91# ---------------------------------------------------------------------------# 

92# StdIO endpoint (child process ↔ async queues) # 

93# ---------------------------------------------------------------------------# 

94class StdIOEndpoint: 

95 """Wrap a child process whose stdin/stdout speak line-delimited JSON-RPC.""" 

96 

97 def __init__(self, cmd: str, pubsub: _PubSub) -> None: 

98 self._cmd = cmd 

99 self._pubsub = pubsub 

100 self._proc: Optional[asyncio.subprocess.Process] = None 

101 self._stdin: Optional[asyncio.StreamWriter] = None 

102 self._pump_task: Optional[asyncio.Task[None]] = None 

103 

104 async def start(self) -> None: 

105 LOGGER.info("Starting stdio subprocess: %s", self._cmd) 

106 self._proc = await asyncio.create_subprocess_exec( 

107 *shlex.split(self._cmd), 

108 stdin=asyncio.subprocess.PIPE, 

109 stdout=asyncio.subprocess.PIPE, 

110 stderr=sys.stderr, # passthrough for visibility 

111 ) 

112 assert self._proc.stdin and self._proc.stdout 

113 self._stdin = self._proc.stdin 

114 self._pump_task = asyncio.create_task(self._pump_stdout()) 

115 

116 async def stop(self) -> None: 

117 if self._proc is None: 117 ↛ 118line 117 didn't jump to line 118 because the condition on line 117 was never true

118 return 

119 LOGGER.info("Stopping subprocess (pid=%s)", self._proc.pid) 

120 self._proc.terminate() 

121 with suppress(asyncio.TimeoutError): 

122 await asyncio.wait_for(self._proc.wait(), timeout=5) 

123 if self._pump_task: 123 ↛ exitline 123 didn't return from function 'stop' because the condition on line 123 was always true

124 self._pump_task.cancel() 

125 

126 async def send(self, raw: str) -> None: 

127 if not self._stdin: 

128 raise RuntimeError("stdio endpoint not started") 

129 LOGGER.debug("→ stdio: %s", raw.strip()) 

130 self._stdin.write(raw.encode()) 

131 await self._stdin.drain() 

132 

133 async def _pump_stdout(self) -> None: 

134 assert self._proc and self._proc.stdout 

135 reader = self._proc.stdout 

136 try: 

137 while True: 

138 line = await reader.readline() 

139 if not line: # EOF 139 ↛ 140line 139 didn't jump to line 140 because the condition on line 139 was never true

140 break 

141 text = line.decode(errors="replace") 

142 LOGGER.debug("← stdio: %s", text.strip()) 

143 await self._pubsub.publish(text) 

144 except asyncio.CancelledError: 

145 raise 

146 except Exception: # pragma: no cover --best-effort logging 

147 LOGGER.exception("stdout pump crashed - terminating bridge") 

148 raise 

149 

150 

151# ---------------------------------------------------------------------------# 

152# FastAPI app exposing /sse & /message # 

153# ---------------------------------------------------------------------------# 

154 

155 

156def _build_fastapi( 

157 pubsub: _PubSub, 

158 stdio: StdIOEndpoint, 

159 keep_alive: int = KEEP_ALIVE_INTERVAL, 

160 sse_path: str = "/sse", 

161 message_path: str = "/message", 

162) -> FastAPI: 

163 app = FastAPI() 

164 

165 # ----- GET /sse ---------------------------------------------------------# 

166 @app.get(sse_path) 

167 async def get_sse(request: Request) -> EventSourceResponse: # noqa: D401 

168 """Fan-out *stdout* of the subprocess to any number of SSE clients.""" 

169 queue = pubsub.subscribe() 

170 session_id = uuid.uuid4().hex 

171 

172 async def event_gen() -> AsyncIterator[dict]: 

173 # 1️⃣ Mandatory "endpoint" bootstrap required by the MCP spec 

174 endpoint_url = f"{str(request.base_url).rstrip('/')}{message_path}?session_id={session_id}" 

175 yield { 

176 "event": "endpoint", 

177 "data": endpoint_url, 

178 "retry": int(keep_alive * 1000), 

179 } 

180 

181 # 2️⃣ Immediate keepalive so clients know the stream is alive 

182 yield {"event": "keepalive", "data": "{}", "retry": keep_alive * 1000} 

183 

184 try: 

185 while True: 

186 if await request.is_disconnected(): 

187 break 

188 

189 try: 

190 msg = await asyncio.wait_for(queue.get(), keep_alive) 

191 yield {"event": "message", "data": msg.rstrip()} 

192 except asyncio.TimeoutError: 

193 yield { 

194 "event": "keepalive", 

195 "data": "{}", 

196 "retry": keep_alive * 1000, 

197 } 

198 finally: 

199 pubsub.unsubscribe(queue) 

200 

201 return EventSourceResponse( 

202 event_gen(), 

203 headers={ 

204 "Cache-Control": "no-cache", 

205 "Connection": "keep-alive", 

206 "X-Accel-Buffering": "no", # disable proxy buffering 

207 }, 

208 ) 

209 

210 # ----- POST /message ----------------------------------------------------# 

211 @app.post(message_path, status_code=status.HTTP_202_ACCEPTED) 

212 async def post_message(raw: Request, session_id: str | None = None) -> Response: # noqa: D401 

213 """Forward the raw JSON body to the stdio process without modification.""" 

214 payload = await raw.body() 

215 try: 

216 json.loads(payload) # validate 

217 except Exception as exc: # noqa: BLE001 

218 return PlainTextResponse( 

219 f"Invalid JSON payload: {exc}", 

220 status_code=status.HTTP_400_BAD_REQUEST, 

221 ) 

222 await stdio.send(payload.decode().rstrip() + "\n") 

223 return PlainTextResponse("forwarded", status_code=status.HTTP_202_ACCEPTED) 

224 

225 # ----- Liveness ---------------------------------------------------------# 

226 @app.get("/healthz") 

227 async def health() -> Response: # noqa: D401 

228 return PlainTextResponse("ok") 

229 

230 return app 

231 

232 

233# ---------------------------------------------------------------------------# 

234# CLI & orchestration # 

235# ---------------------------------------------------------------------------# 

236 

237 

238def _parse_args(argv: Sequence[str]) -> argparse.Namespace: 

239 p = argparse.ArgumentParser( 

240 prog="mcpgateway.translate", 

241 description="Bridges stdio JSON-RPC to SSE.", 

242 ) 

243 src = p.add_mutually_exclusive_group(required=True) 

244 src.add_argument("--stdio", help='Command to run, e.g. "uv run mcp-server-git"') 

245 src.add_argument("--sse", help="[NOT IMPLEMENTED]") 

246 src.add_argument("--streamableHttp", help="[NOT IMPLEMENTED]") 

247 

248 p.add_argument("--port", type=int, default=8000, help="HTTP port to bind") 

249 p.add_argument( 

250 "--logLevel", 

251 default="info", 

252 choices=["debug", "info", "warning", "error", "critical"], 

253 help="Log level", 

254 ) 

255 

256 args = p.parse_args(argv) 

257 if args.sse or args.streamableHttp: 

258 raise NotImplementedError("Only --stdio → SSE is available in this build.") 

259 return args 

260 

261 

262async def _run_stdio_to_sse(cmd: str, port: int, log_level: str = "info") -> None: 

263 pubsub = _PubSub() 

264 stdio = StdIOEndpoint(cmd, pubsub) 

265 await stdio.start() 

266 

267 app = _build_fastapi(pubsub, stdio) 

268 config = uvicorn.Config( 

269 app, 

270 host="0.0.0.0", 

271 port=port, 

272 log_level=log_level, 

273 lifespan="off", 

274 ) 

275 server = uvicorn.Server(config) 

276 

277 shutting_down = asyncio.Event() # 🔄 make shutdown idempotent 

278 

279 async def _shutdown() -> None: 

280 if shutting_down.is_set(): 280 ↛ 281line 280 didn't jump to line 281 because the condition on line 280 was never true

281 return 

282 shutting_down.set() 

283 LOGGER.info("Shutting down …") 

284 await stdio.stop() 

285 await server.shutdown() 

286 

287 loop = asyncio.get_running_loop() 

288 for sig in (signal.SIGINT, signal.SIGTERM): 

289 with suppress(NotImplementedError): # Windows lacks add_signal_handler 

290 loop.add_signal_handler(sig, lambda _s=sig: asyncio.create_task(_shutdown())) 

291 

292 LOGGER.info("Bridge ready → http://127.0.0.1:%s/sse", port) 

293 await server.serve() 

294 await _shutdown() # final cleanup 

295 

296 

297def main(argv: Optional[Sequence[str]] | None = None) -> None: # entry-point 

298 args = _parse_args(argv or sys.argv[1:]) 

299 logging.basicConfig( 

300 level=getattr(logging, args.logLevel.upper(), logging.INFO), 

301 format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", 

302 ) 

303 try: 

304 asyncio.run(_run_stdio_to_sse(args.stdio, args.port, args.logLevel)) 

305 except KeyboardInterrupt: 

306 print("") # restore shell prompt 

307 sys.exit(0) 

308 except NotImplementedError as exc: 

309 print(exc, file=sys.stderr) 

310 sys.exit(1) 

311 

312 

313if __name__ == "__main__": # python -m mcpgateway.translate … 

314 main()