Coverage for mcpgateway/translate.py: 80%
158 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-22 14:46 +0100
« 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
4Copyright 2025
5SPDX-License-Identifier: Apache-2.0
6Authors: Mihai Criveti
8Only the stdio→SSE direction is implemented for now.
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
15# 2. from another shell / browser subscribe to the SSE stream
16curl -N http://localhost:9000/sse # receive the stream
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"}}'
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"}}}'
28curl -X POST http://localhost:9000/message \
29 -H 'Content-Type: application/json' \
30 -d '{"jsonrpc":"2.0","id":2,"method":"tools/list"}'
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"""
38from __future__ import annotations
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
51import uvicorn
52from fastapi import FastAPI, Request, Response, status
53from fastapi.responses import PlainTextResponse
54from sse_starlette.sse import EventSourceResponse
56LOGGER = logging.getLogger("mcpgateway.translate")
57KEEP_ALIVE_INTERVAL = 30 # seconds ── matches the reference implementation
58__all__ = ["main"] # for console-script entry-point
61# ---------------------------------------------------------------------------#
62# Helpers - trivial in-process Pub/Sub #
63# ---------------------------------------------------------------------------#
64class _PubSub:
65 """Very small fan-out helper - one async Queue per subscriber."""
67 def __init__(self) -> None:
68 self._subscribers: List[asyncio.Queue[str]] = []
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)
81 def subscribe(self) -> "asyncio.Queue[str]":
82 q: asyncio.Queue[str] = asyncio.Queue(maxsize=1024)
83 self._subscribers.append(q)
84 return q
86 def unsubscribe(self, q: "asyncio.Queue[str]") -> None:
87 with suppress(ValueError):
88 self._subscribers.remove(q)
91# ---------------------------------------------------------------------------#
92# StdIO endpoint (child process ↔ async queues) #
93# ---------------------------------------------------------------------------#
94class StdIOEndpoint:
95 """Wrap a child process whose stdin/stdout speak line-delimited JSON-RPC."""
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
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())
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()
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()
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
151# ---------------------------------------------------------------------------#
152# FastAPI app exposing /sse & /message #
153# ---------------------------------------------------------------------------#
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()
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
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 }
181 # 2️⃣ Immediate keepalive so clients know the stream is alive
182 yield {"event": "keepalive", "data": "{}", "retry": keep_alive * 1000}
184 try:
185 while True:
186 if await request.is_disconnected():
187 break
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)
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 )
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)
225 # ----- Liveness ---------------------------------------------------------#
226 @app.get("/healthz")
227 async def health() -> Response: # noqa: D401
228 return PlainTextResponse("ok")
230 return app
233# ---------------------------------------------------------------------------#
234# CLI & orchestration #
235# ---------------------------------------------------------------------------#
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]")
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 )
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
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()
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)
277 shutting_down = asyncio.Event() # 🔄 make shutdown idempotent
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()
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()))
292 LOGGER.info("Bridge ready → http://127.0.0.1:%s/sse", port)
293 await server.serve()
294 await _shutdown() # final cleanup
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)
313if __name__ == "__main__": # python -m mcpgateway.translate …
314 main()