1"""Small MCP Streamable HTTP client used by the scanner.
2
3Only the lifecycle and tool operations needed by the Actor are implemented:
4``initialize``, ``notifications/initialized``, ``tools/list``, and
5``tools/call``. Keeping this client local also lets the scanner report useful
6errors from servers that are not fully SDK-compatible.
7"""
8from __future__ import annotations
9
10import json
11import uuid
12from dataclasses import dataclass, field
13from typing import Any
14
15import httpx
16
17PREFERRED_PROTOCOL_VERSION = "2025-11-25"
18SUPPORTED_PROTOCOL_VERSIONS = frozenset(
19 {
20 PREFERRED_PROTOCOL_VERSION,
21 "2025-06-18",
22 "2025-03-26",
23 }
24)
25
26
27class MCPClientError(Exception):
28 """Base error raised for invalid MCP transport or response behavior."""
29
30
31class MCPProtocolError(MCPClientError):
32 """A top-level JSON-RPC error returned by the MCP server."""
33
34 def __init__(self, code: int | None, message: str):
35 self.code = code
36 self.message = message
37 super().__init__(f"JSON-RPC error {code}: {message}")
38
39
40@dataclass
41class MCPTool:
42 name: str
43 description: str = ""
44 input_schema: dict[str, Any] = field(default_factory=dict)
45
46
47@dataclass
48class MCPProbeResult:
49 reachable: bool
50 protocol_version: str | None = None
51 server_name: str | None = None
52 tools: list[MCPTool] = field(default_factory=list)
53 raw_initialize_response: dict[str, Any] | None = None
54 transport_notes: list[str] = field(default_factory=list)
55 error: str | None = None
56
57
58def is_tool_call_success(http_status: int, parsed: dict[str, Any] | None) -> bool:
59 """Return true only for a successful HTTP and MCP tool result.
60
61 MCP distinguishes top-level protocol errors from tool execution errors in
62 ``result.isError``. Both must be considered failures.
63 """
64
65 if not 200 <= http_status < 300 or not isinstance(parsed, dict):
66 return False
67 if parsed.get("error") is not None:
68 return False
69 result = parsed.get("result")
70 return isinstance(result, dict) and result.get("isError") is not True
71
72
73class MCPClient:
74 def __init__(
75 self,
76 base_url: str,
77 auth_header: str | None = None,
78 timeout: float = 15.0,
79 transport: httpx.AsyncBaseTransport | None = None,
80 ):
81 self.base_url = base_url.rstrip("/")
82 self.timeout = timeout
83 self.transport = transport
84 self.protocol_version: str | None = None
85 self.session_id: str | None = None
86 self.initialized = False
87 self.headers = {
88 "Content-Type": "application/json",
89 "Accept": "application/json, text/event-stream",
90 }
91 if auth_header:
92 self.headers["Authorization"] = auth_header
93
94 @staticmethod
95 def _rpc_body(method: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
96 return {
97 "jsonrpc": "2.0",
98 "id": str(uuid.uuid4()),
99 "method": method,
100 "params": params or {},
101 }
102
103 @staticmethod
104 def _notification_body(
105 method: str, params: dict[str, Any] | None = None
106 ) -> dict[str, Any]:
107 body: dict[str, Any] = {"jsonrpc": "2.0", "method": method}
108 if params:
109 body["params"] = params
110 return body
111
112 def _request_headers(
113 self,
114 *,
115 subsequent: bool,
116 extra_headers: dict[str, str] | None = None,
117 omit_authorization: bool = False,
118 ) -> dict[str, str]:
119 headers = dict(self.headers)
120 if omit_authorization:
121 headers = {k: v for k, v in headers.items() if k.lower() != "authorization"}
122 if subsequent:
123 if self.protocol_version:
124 headers["MCP-Protocol-Version"] = self.protocol_version
125 if self.session_id:
126 headers["Mcp-Session-Id"] = self.session_id
127 if extra_headers:
128 headers.update(extra_headers)
129 return headers
130
131 async def _post(
132 self,
133 client: httpx.AsyncClient,
134 body: dict[str, Any],
135 *,
136 subsequent: bool = True,
137 headers: dict[str, str] | None = None,
138 omit_authorization: bool = False,
139 ) -> httpx.Response:
140 return await client.post(
141 self.base_url,
142 json=body,
143 headers=self._request_headers(
144 subsequent=subsequent,
145 extra_headers=headers,
146 omit_authorization=omit_authorization,
147 ),
148 timeout=self.timeout,
149 )
150
151 @staticmethod
152 def _select_message(payload: Any, request_id: str | int | None) -> dict[str, Any]:
153 messages = payload if isinstance(payload, list) else [payload]
154 candidates = [message for message in messages if isinstance(message, dict)]
155 if request_id is None:
156 if not candidates:
157 raise MCPClientError("Response did not contain a JSON-RPC object")
158 return candidates[0]
159 for message in candidates:
160 if message.get("id") == request_id:
161 return message
162 raise MCPClientError(f"Response did not contain JSON-RPC id {request_id!r}")
163
164 @staticmethod
165 def _parse_json_or_sse(
166 resp: httpx.Response, request_id: str | int | None = None
167 ) -> dict[str, Any]:
168 """Parse JSON or SSE and select the response matching ``request_id``."""
169
170 content_type = resp.headers.get("content-type", "").lower()
171 if "text/event-stream" not in content_type:
172 try:
173 return MCPClient._select_message(resp.json(), request_id)
174 except (json.JSONDecodeError, UnicodeDecodeError, ValueError) as exc:
175 raise MCPClientError(f"Invalid JSON response: {exc}") from exc
176
177 data_lines: list[str] = []
178 events: list[str] = []
179
180 def finish_event() -> None:
181 if data_lines:
182 events.append("\n".join(data_lines))
183 data_lines.clear()
184
185 for line in resp.text.splitlines():
186 if not line:
187 finish_event()
188 elif line.startswith("data:"):
189 data_lines.append(line[5:].lstrip())
190 finish_event()
191
192 if not events:
193 raise MCPClientError("SSE response contained no data events")
194
195 parse_errors: list[str] = []
196 for event_data in events:
197 if event_data == "[DONE]":
198 continue
199 try:
200 payload = json.loads(event_data)
201 except json.JSONDecodeError as exc:
202 parse_errors.append(str(exc))
203 continue
204 try:
205 return MCPClient._select_message(payload, request_id)
206 except MCPClientError:
207 continue
208
209 if request_id is not None:
210 raise MCPClientError(
211 f"SSE response contained no JSON-RPC response matching id {request_id!r}"
212 )
213 detail = f": {parse_errors[0]}" if parse_errors else ""
214 raise MCPClientError(f"SSE response contained no JSON-RPC object{detail}")
215
216 @staticmethod
217 def _result_or_raise(message: dict[str, Any], operation: str) -> dict[str, Any]:
218 error = message.get("error")
219 if error is not None:
220 if isinstance(error, dict):
221 code = error.get("code")
222 raw_message = error.get("message", "Unknown protocol error")
223 else:
224 code = None
225 raw_message = "Malformed JSON-RPC error"
226 raise MCPProtocolError(code if isinstance(code, int) else None, str(raw_message))
227 result = message.get("result")
228 if not isinstance(result, dict):
229 raise MCPClientError(f"{operation} response did not contain an object result")
230 return result
231
232 async def initialize(self, client: httpx.AsyncClient) -> dict[str, Any]:
233 """Establish the MCP session and send ``notifications/initialized``."""
234
235 init_body = self._rpc_body(
236 "initialize",
237 {
238 "protocolVersion": PREFERRED_PROTOCOL_VERSION,
239 "capabilities": {},
240 "clientInfo": {"name": "mcp-security-scanner", "version": "0.2"},
241 },
242 )
243 init_resp = await self._post(client, init_body, subsequent=False)
244 if not 200 <= init_resp.status_code < 300:
245 raise MCPClientError(f"initialize returned HTTP {init_resp.status_code}")
246
247 init_json = self._parse_json_or_sse(init_resp, init_body["id"])
248 init_result = self._result_or_raise(init_json, "initialize")
249 negotiated_version = init_result.get("protocolVersion")
250 if not isinstance(negotiated_version, str) or not negotiated_version:
251 raise MCPClientError("initialize result omitted protocolVersion")
252 if negotiated_version not in SUPPORTED_PROTOCOL_VERSIONS:
253 raise MCPClientError(
254 "initialize negotiated an unsupported MCP protocol version"
255 )
256
257 self.protocol_version = negotiated_version
258 self.session_id = init_resp.headers.get("Mcp-Session-Id") or None
259
260 initialized_resp = await self._post(
261 client,
262 self._notification_body("notifications/initialized"),
263 subsequent=True,
264 )
265 if not 200 <= initialized_resp.status_code < 300:
266 raise MCPClientError(
267 f"notifications/initialized returned HTTP {initialized_resp.status_code}"
268 )
269 self.initialized = True
270 return init_json
271
272 async def list_tools(self, client: httpx.AsyncClient) -> list[MCPTool]:
273 """List all tools, following opaque ``nextCursor`` values."""
274
275 tools: list[MCPTool] = []
276 cursor: str | None = None
277 seen_cursors: set[str] = set()
278
279 while True:
280 params = {"cursor": cursor} if cursor is not None else {}
281 list_body = self._rpc_body("tools/list", params)
282 list_resp = await self._post(client, list_body, subsequent=True)
283 if not 200 <= list_resp.status_code < 300:
284 raise MCPClientError(f"tools/list returned HTTP {list_resp.status_code}")
285
286 list_json = self._parse_json_or_sse(list_resp, list_body["id"])
287 page = self._result_or_raise(list_json, "tools/list")
288 raw_tools = page.get("tools", [])
289 if not isinstance(raw_tools, list):
290 raise MCPClientError("tools/list result contained a non-array tools value")
291
292 for raw_tool in raw_tools:
293 if not isinstance(raw_tool, dict):
294 continue
295 schema = raw_tool.get("inputSchema")
296 tools.append(
297 MCPTool(
298 name=str(raw_tool.get("name", "<unnamed>")),
299 description=str(raw_tool.get("description", "") or ""),
300 input_schema=schema if isinstance(schema, dict) else {},
301 )
302 )
303
304 next_cursor = page.get("nextCursor")
305 if not isinstance(next_cursor, str) or not next_cursor:
306 break
307 if next_cursor in seen_cursors:
308 raise MCPClientError("tools/list returned a repeated nextCursor")
309 seen_cursors.add(next_cursor)
310 cursor = next_cursor
311
312 return tools
313
314 async def probe(self) -> MCPProbeResult:
315 result = MCPProbeResult(reachable=False)
316 try:
317 async with httpx.AsyncClient(
318 follow_redirects=True, transport=self.transport
319 ) as client:
320 try:
321 init_json = await self.initialize(client)
322 result.reachable = True
323 result.raw_initialize_response = init_json
324 init_result = init_json["result"]
325 server_info = init_result.get("serverInfo", {})
326 if isinstance(server_info, dict):
327 server_name = server_info.get("name")
328 result.server_name = str(server_name) if server_name else None
329 result.protocol_version = self.protocol_version
330 except MCPClientError as exc:
331 result.reachable = True
332 result.error = str(exc)
333 return result
334
335 try:
336 result.tools = await self.list_tools(client)
337 except MCPClientError as exc:
338 result.error = str(exc)
339 return result
340
341 except httpx.RequestError as exc:
342 result.error = f"Connection failed: {type(exc).__name__}"
343 return result
344 except Exception as exc:
345 result.error = f"Unexpected error: {exc}"
346 return result
347
348 async def call_tool(
349 self,
350 client: httpx.AsyncClient,
351 tool_name: str,
352 arguments: dict[str, Any],
353 headers: dict[str, str] | None = None,
354 *,
355 omit_authorization: bool = False,
356 ) -> tuple[int, dict[str, Any] | None, str | None]:
357 """Return ``(HTTP status, parsed JSON-RPC response, short raw text)``."""
358
359 body = self._rpc_body("tools/call", {"name": tool_name, "arguments": arguments})
360 try:
361 resp = await self._post(
362 client,
363 body,
364 subsequent=True,
365 headers=headers,
366 omit_authorization=omit_authorization,
367 )
368 try:
369 parsed = self._parse_json_or_sse(resp, body["id"])
370 except MCPClientError:
371 parsed = None
372 return resp.status_code, parsed, resp.text[:500]
373 except httpx.RequestError as exc:
374 return -1, None, str(exc)
