From 7b9e0be45d9892feddc0204b7c010a1dce476681 Mon Sep 17 00:00:00 2001 From: Pukaphol Thienpreecha Date: Sun, 26 Apr 2026 16:55:28 -0700 Subject: [PATCH] Handle MCP HTTP request failures --- agentic_security/mcp/main.py | 74 ++++++++++++++++++------- tests/unit/test_mcp_main.py | 102 +++++++++++++++++++++++++++++++++++ 2 files changed, 156 insertions(+), 20 deletions(-) create mode 100644 tests/unit/test_mcp_main.py diff --git a/agentic_security/mcp/main.py b/agentic_security/mcp/main.py index 7e19e2c..f863a59 100644 --- a/agentic_security/mcp/main.py +++ b/agentic_security/mcp/main.py @@ -1,6 +1,10 @@ +from typing import Any + import httpx from mcp.server.fastmcp import FastMCP +from agentic_security.logutils import logger + # Initialize MCP server mcp = FastMCP( name="Agentic Security MCP Server", @@ -11,6 +15,51 @@ AGENTIC_SECURITY = "http://0.0.0.0:8718" +def _api_error(error_type: str, message: str, **details: Any) -> dict[str, Any]: + error = {"type": error_type, "message": message} + error.update(details) + return {"error": error} + + +def _response_body(response: httpx.Response) -> Any: + try: + return response.json() + except ValueError: + return response.text + + +async def _request_api( + method: str, + path: str, + *, + json: dict[str, Any] | None = None, +) -> dict[str, Any] | list[Any]: + url = f"{AGENTIC_SECURITY}{path}" + try: + async with httpx.AsyncClient() as client: + request_kwargs = {} + if json is not None: + request_kwargs["json"] = json + response = await client.request(method, url, **request_kwargs) + response.raise_for_status() + return response.json() + except httpx.HTTPStatusError as exc: + response = exc.response + logger.error("MCP backend returned an error: %s", exc) + return _api_error( + "http_status", + str(exc), + status_code=response.status_code, + response=_response_body(response), + ) + except httpx.RequestError as exc: + logger.error("MCP backend request failed: %s", exc) + return _api_error("request", str(exc)) + except ValueError as exc: + logger.error("MCP backend returned invalid JSON: %s", exc) + return _api_error("invalid_json", str(exc)) + + @mcp.tool() async def verify_llm(spec: str) -> dict: """ @@ -22,10 +71,7 @@ async def verify_llm(spec: str) -> dict: Args: spect(str): The specification of the LLM model to verify. """ - url = f"{AGENTIC_SECURITY}/verify" - async with httpx.AsyncClient() as client: - response = await client.post(url, json={"spec": spec}) - return response.json() + return await _request_api("POST", "/verify", json={"spec": spec}) @mcp.tool() @@ -47,7 +93,6 @@ async def start_scan( enableMultiStepAttack (bool, optional): Whether to enable multi-step attack """ - url = f"{AGENTIC_SECURITY}/scan" payload = { "llmSpec": llmSpec, "maxBudget": maxBudget, @@ -57,9 +102,7 @@ async def start_scan( "probe_datasets": [], "secrets": {}, } - async with httpx.AsyncClient() as client: - response = await client.post(url, json=payload) - return response.json() + return await _request_api("POST", "/scan", json=payload) @mcp.tool() @@ -69,10 +112,7 @@ async def stop_scan() -> dict: Returns: dict: The confirmation from the FastAPI server that the scan has been stopped. """ - url = f"{AGENTIC_SECURITY}/stop" - async with httpx.AsyncClient() as client: - response = await client.post(url) - return response.json() + return await _request_api("POST", "/stop") @mcp.tool() @@ -83,10 +123,7 @@ async def get_data_config() -> list: Returns: list: The response from the FastAPI server, confirming the scan has been stopped. """ - url = f"{AGENTIC_SECURITY}/v1/data-config" - async with httpx.AsyncClient() as client: - response = await client.get(url) - return response.json() + return await _request_api("GET", "/v1/data-config") @mcp.tool() @@ -97,10 +134,7 @@ async def get_spec_templates() -> list: Returns: list: The LLM specification templates from the FastAPI server. """ - url = f"{AGENTIC_SECURITY}/v1/llm-specs" - async with httpx.AsyncClient() as client: - response = await client.get(url) - return response.json() + return await _request_api("GET", "/v1/llm-specs") # Run the MCP server diff --git a/tests/unit/test_mcp_main.py b/tests/unit/test_mcp_main.py new file mode 100644 index 0000000..cf164df --- /dev/null +++ b/tests/unit/test_mcp_main.py @@ -0,0 +1,102 @@ +import httpx +import pytest + +from agentic_security.mcp import main as mcp_main + + +class MockResponse: + def __init__(self, status_code=200, payload=None, text=""): + self.status_code = status_code + self.payload = payload + self.text = text + self.request = httpx.Request("GET", "http://testserver") + + def json(self): + if isinstance(self.payload, Exception): + raise self.payload + return self.payload + + def raise_for_status(self): + if self.status_code >= 400: + raise httpx.HTTPStatusError( + f"{self.status_code} error", + request=self.request, + response=self, + ) + + +class MockAsyncClient: + def __init__(self, response=None, exc=None): + self.response = response + self.exc = exc + self.calls = [] + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, traceback): + return False + + async def request(self, method, url, json=None): + self.calls.append((method, url, json)) + if self.exc: + raise self.exc + return self.response + + +def use_client(monkeypatch, client): + monkeypatch.setattr(mcp_main.httpx, "AsyncClient", lambda: client) + return client + + +@pytest.mark.asyncio +async def test_verify_llm_returns_backend_json(monkeypatch): + client = use_client( + monkeypatch, MockAsyncClient(MockResponse(payload={"ok": True})) + ) + + result = await mcp_main.verify_llm("openai:gpt-4") + + assert result == {"ok": True} + assert client.calls == [ + ( + "POST", + "http://0.0.0.0:8718/verify", + {"spec": "openai:gpt-4"}, + ) + ] + + +@pytest.mark.asyncio +async def test_start_scan_returns_http_status_error(monkeypatch): + response = MockResponse(status_code=503, payload={"detail": "backend unavailable"}) + use_client(monkeypatch, MockAsyncClient(response)) + + result = await mcp_main.start_scan("openai:gpt-4", 10) + + assert result["error"]["type"] == "http_status" + assert result["error"]["status_code"] == 503 + assert result["error"]["response"] == {"detail": "backend unavailable"} + + +@pytest.mark.asyncio +async def test_get_data_config_returns_request_error(monkeypatch): + request = httpx.Request("GET", "http://0.0.0.0:8718/v1/data-config") + error = httpx.ConnectError("connection refused", request=request) + use_client(monkeypatch, MockAsyncClient(exc=error)) + + result = await mcp_main.get_data_config() + + assert result["error"]["type"] == "request" + assert "connection refused" in result["error"]["message"] + + +@pytest.mark.asyncio +async def test_get_spec_templates_returns_invalid_json_error(monkeypatch): + response = MockResponse(payload=ValueError("bad json")) + use_client(monkeypatch, MockAsyncClient(response)) + + result = await mcp_main.get_spec_templates() + + assert result["error"]["type"] == "invalid_json" + assert "bad json" in result["error"]["message"]