|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import json |
| 4 | +import secrets |
| 5 | +import threading |
| 6 | +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer |
| 7 | +from typing import Any, Optional |
| 8 | + |
| 9 | +from teaagent.mcp_server import handle_mcp_request |
| 10 | +from teaagent.tools import ToolRegistry |
| 11 | + |
| 12 | +MCP_PATH = "/mcp" |
| 13 | +SESSION_HEADER = "Mcp-Session-Id" |
| 14 | +DEFAULT_PORT = 7330 |
| 15 | + |
| 16 | + |
| 17 | +class MCPSessionStore: |
| 18 | + """In-memory session store for the Streamable HTTP transport.""" |
| 19 | + |
| 20 | + def __init__(self) -> None: |
| 21 | + self._lock = threading.Lock() |
| 22 | + self._sessions: set[str] = set() |
| 23 | + |
| 24 | + def create(self) -> str: |
| 25 | + session_id = secrets.token_urlsafe(24) |
| 26 | + with self._lock: |
| 27 | + self._sessions.add(session_id) |
| 28 | + return session_id |
| 29 | + |
| 30 | + def has(self, session_id: str) -> bool: |
| 31 | + with self._lock: |
| 32 | + return session_id in self._sessions |
| 33 | + |
| 34 | + def remove(self, session_id: str) -> bool: |
| 35 | + with self._lock: |
| 36 | + if session_id in self._sessions: |
| 37 | + self._sessions.remove(session_id) |
| 38 | + return True |
| 39 | + return False |
| 40 | + |
| 41 | + |
| 42 | +def build_mcp_http_server( |
| 43 | + registry: ToolRegistry, |
| 44 | + *, |
| 45 | + host: str = "127.0.0.1", |
| 46 | + port: int = DEFAULT_PORT, |
| 47 | + auth_token: Optional[str] = None, |
| 48 | + allowed_origins: Optional[list[str]] = None, |
| 49 | +) -> tuple[ThreadingHTTPServer, MCPSessionStore]: |
| 50 | + sessions = MCPSessionStore() |
| 51 | + origins = frozenset(allowed_origins) if allowed_origins else None |
| 52 | + handler_cls = _make_handler(registry, sessions, auth_token, origins) |
| 53 | + server = ThreadingHTTPServer((host, port), handler_cls) |
| 54 | + return server, sessions |
| 55 | + |
| 56 | + |
| 57 | +def serve_mcp_http( |
| 58 | + registry: ToolRegistry, |
| 59 | + *, |
| 60 | + host: str = "127.0.0.1", |
| 61 | + port: int = DEFAULT_PORT, |
| 62 | + auth_token: Optional[str] = None, |
| 63 | + allowed_origins: Optional[list[str]] = None, |
| 64 | +) -> int: |
| 65 | + server, _sessions = build_mcp_http_server( |
| 66 | + registry, |
| 67 | + host=host, |
| 68 | + port=port, |
| 69 | + auth_token=auth_token, |
| 70 | + allowed_origins=allowed_origins, |
| 71 | + ) |
| 72 | + try: |
| 73 | + server.serve_forever() |
| 74 | + except KeyboardInterrupt: |
| 75 | + pass |
| 76 | + finally: |
| 77 | + server.server_close() |
| 78 | + return 0 |
| 79 | + |
| 80 | + |
| 81 | +def _make_handler( |
| 82 | + registry: ToolRegistry, |
| 83 | + sessions: MCPSessionStore, |
| 84 | + auth_token: Optional[str], |
| 85 | + allowed_origins: Optional[frozenset[str]], |
| 86 | +) -> type[BaseHTTPRequestHandler]: |
| 87 | + class MCPHandler(BaseHTTPRequestHandler): |
| 88 | + protocol_version = "HTTP/1.1" |
| 89 | + |
| 90 | + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 |
| 91 | + return |
| 92 | + |
| 93 | + def _check_auth(self) -> bool: |
| 94 | + if auth_token is None: |
| 95 | + return True |
| 96 | + header = self.headers.get("Authorization", "") |
| 97 | + if not header.startswith("Bearer "): |
| 98 | + return False |
| 99 | + return secrets.compare_digest(header[len("Bearer ") :], auth_token) |
| 100 | + |
| 101 | + def _check_origin(self) -> bool: |
| 102 | + if allowed_origins is None: |
| 103 | + return True |
| 104 | + origin = self.headers.get("Origin") |
| 105 | + if origin is None: |
| 106 | + return True |
| 107 | + return origin in allowed_origins |
| 108 | + |
| 109 | + def _send_json( |
| 110 | + self, |
| 111 | + status: int, |
| 112 | + payload: Any, |
| 113 | + *, |
| 114 | + extra_headers: Optional[dict[str, str]] = None, |
| 115 | + ) -> None: |
| 116 | + body = json.dumps(payload, ensure_ascii=False).encode("utf-8") |
| 117 | + self.send_response(status) |
| 118 | + self.send_header("Content-Type", "application/json") |
| 119 | + self.send_header("Content-Length", str(len(body))) |
| 120 | + for key, value in (extra_headers or {}).items(): |
| 121 | + self.send_header(key, value) |
| 122 | + self.end_headers() |
| 123 | + self.wfile.write(body) |
| 124 | + |
| 125 | + def _send_status(self, status: int, message: Optional[str] = None) -> None: |
| 126 | + body = (message or "").encode("utf-8") |
| 127 | + self.send_response(status) |
| 128 | + self.send_header("Content-Type", "text/plain") |
| 129 | + self.send_header("Content-Length", str(len(body))) |
| 130 | + self.end_headers() |
| 131 | + if body: |
| 132 | + self.wfile.write(body) |
| 133 | + |
| 134 | + def _gate(self) -> bool: |
| 135 | + if self.path != MCP_PATH: |
| 136 | + self._send_status(404, "not found") |
| 137 | + return False |
| 138 | + if not self._check_origin(): |
| 139 | + self._send_status(403, "forbidden origin") |
| 140 | + return False |
| 141 | + if not self._check_auth(): |
| 142 | + self._send_status(401, "unauthorized") |
| 143 | + return False |
| 144 | + return True |
| 145 | + |
| 146 | + def _read_body(self) -> tuple[Optional[Any], Optional[str]]: |
| 147 | + length = int(self.headers.get("Content-Length", "0") or "0") |
| 148 | + if length <= 0: |
| 149 | + return None, "missing JSON-RPC body" |
| 150 | + raw = self.rfile.read(length) |
| 151 | + try: |
| 152 | + return json.loads(raw.decode("utf-8")), None |
| 153 | + except (json.JSONDecodeError, UnicodeDecodeError): |
| 154 | + return None, "invalid JSON" |
| 155 | + |
| 156 | + def do_POST(self) -> None: |
| 157 | + if not self._gate(): |
| 158 | + return |
| 159 | + payload, error = self._read_body() |
| 160 | + if error is not None: |
| 161 | + self._send_status(400, error) |
| 162 | + return |
| 163 | + if not isinstance(payload, (dict, list)): |
| 164 | + self._send_status(400, "JSON-RPC payload must be object or array") |
| 165 | + return |
| 166 | + |
| 167 | + session_header = self.headers.get(SESSION_HEADER) |
| 168 | + |
| 169 | + if isinstance(payload, dict) and payload.get("method") == "initialize": |
| 170 | + response = handle_mcp_request(registry, payload) |
| 171 | + if response is None: |
| 172 | + self._send_status(400, "initialize requires an id") |
| 173 | + return |
| 174 | + session_id = sessions.create() |
| 175 | + self._send_json(200, response, extra_headers={SESSION_HEADER: session_id}) |
| 176 | + return |
| 177 | + |
| 178 | + if not session_header or not sessions.has(session_header): |
| 179 | + self._send_status(400, "missing or invalid Mcp-Session-Id") |
| 180 | + return |
| 181 | + |
| 182 | + if isinstance(payload, list): |
| 183 | + responses: list[dict[str, Any]] = [] |
| 184 | + for item in payload: |
| 185 | + if not isinstance(item, dict): |
| 186 | + continue |
| 187 | + response = handle_mcp_request(registry, item) |
| 188 | + if response is not None: |
| 189 | + responses.append(response) |
| 190 | + if not responses: |
| 191 | + self._send_status(202) |
| 192 | + return |
| 193 | + self._send_json(200, responses) |
| 194 | + return |
| 195 | + |
| 196 | + response = handle_mcp_request(registry, payload) |
| 197 | + if response is None: |
| 198 | + self._send_status(202) |
| 199 | + return |
| 200 | + self._send_json(200, response) |
| 201 | + |
| 202 | + def do_GET(self) -> None: |
| 203 | + if not self._gate(): |
| 204 | + return |
| 205 | + session_header = self.headers.get(SESSION_HEADER) |
| 206 | + if not session_header or not sessions.has(session_header): |
| 207 | + self._send_status(400, "missing or invalid Mcp-Session-Id") |
| 208 | + return |
| 209 | + self.send_response(200) |
| 210 | + self.send_header("Content-Type", "text/event-stream") |
| 211 | + self.send_header("Cache-Control", "no-cache") |
| 212 | + self.send_header("Connection", "close") |
| 213 | + self.end_headers() |
| 214 | + self.wfile.write(b": teaagent mcp stream\n\n") |
| 215 | + self.wfile.flush() |
| 216 | + |
| 217 | + def do_DELETE(self) -> None: |
| 218 | + if not self._gate(): |
| 219 | + return |
| 220 | + session_header = self.headers.get(SESSION_HEADER) |
| 221 | + if not session_header or not sessions.has(session_header): |
| 222 | + self._send_status(404, "session not found") |
| 223 | + return |
| 224 | + sessions.remove(session_header) |
| 225 | + self._send_status(204) |
| 226 | + |
| 227 | + def do_OPTIONS(self) -> None: |
| 228 | + self._send_status(204) |
| 229 | + |
| 230 | + return MCPHandler |
0 commit comments