error handling for httpx requests in MCP

agentic_security/http_spec.py

@@ -18,16 +18,34 @@ class Modality(Enum):
 
 def encode_image_base64_by_url(url: str = "https://github.com/fluidicon.png") -> str:
     """Encode image data to base64 from a URL"""
-    response = httpx.get(url)
-    encoded_content = base64.b64encode(response.content).decode("utf-8")
-    return "data:image/jpeg;base64," + encoded_content
+    try:
+        response = httpx.get(url)
+        encoded_content = base64.b64encode(response.content).decode("utf-8")
+        return "data:image/jpeg;base64," + encoded_content
+    except httpx.TimeoutException as e:
+        raise ValueError(f"Image fetch timed out: {str(e)}")
+    except httpx.HTTPStatusError as e:
+        raise ValueError(f"HTTP error {e.response.status_code}: {str(e)}")
+    except httpx.RequestError as e:
+        raise ValueError(f"Image fetch failed: {str(e)}")
+    except Exception as e:
+        raise ValueError(f"Unexpected error fetching image: {str(e)}")
 
 
 def encode_audio_base64_by_url(url: str) -> str:
     """Encode audio data to base64 from a URL"""
-    response = httpx.get(url)
-    encoded_content = base64.b64encode(response.content).decode("utf-8")
-    return "data:audio/mpeg;base64," + encoded_content
+    try:
+        response = httpx.get(url)
+        encoded_content = base64.b64encode(response.content).decode("utf-8")
+        return "data:audio/mpeg;base64," + encoded_content
+    except httpx.TimeoutException as e:
+        raise ValueError(f"Audio fetch timed out: {str(e)}")
+    except httpx.HTTPStatusError as e:
+        raise ValueError(f"HTTP error {e.response.status_code}: {str(e)}")
+    except httpx.RequestError as e:
+        raise ValueError(f"Audio fetch failed: {str(e)}")
+    except Exception as e:
+        raise ValueError(f"Unexpected error fetching audio: {str(e)}")
 
 
 class InvalidHTTPSpecError(Exception):
@@ -58,16 +76,25 @@ def timeout(self):
 
     async def _probe_with_files(self, files):
         transport = httpx.AsyncHTTPTransport(retries=settings_var("network.retry", 3))
-        async with httpx.AsyncClient(transport=transport) as client:
-            response = await client.request(
-                method=self.method,
-                url=self.url,
-                headers=self.headers,
-                files=files,
-                timeout=self.timeout(),
-            )
-
-        return response
+        try:
+            async with httpx.AsyncClient(transport=transport) as client:
+                response = await client.request(
+                    method=self.method,
+                    url=self.url,
+                    headers=self.headers,
+                    files=files,
+                    timeout=self.timeout(),
+                )
+                response.raise_for_status()
+                return response
+        except httpx.TimeoutException as e:
+            raise ValueError(f"Request timed out: {str(e)}")
+        except httpx.HTTPStatusError as e:
+            raise ValueError(f"HTTP error {e.response.status_code}: {str(e)}")
+        except httpx.RequestError as e:
+            raise ValueError(f"Request failed: {str(e)}")
+        except Exception as e:
+            raise ValueError(f"Unexpected error: {str(e)}")
 
     def validate(self, prompt, encoded_image, encoded_audio, files) -> None:
         if self.has_files and not files:
@@ -102,16 +129,26 @@ async def probe(
         content = content.replace("<<BASE64_AUDIO>>", encoded_audio)
 
         transport = httpx.AsyncHTTPTransport(retries=settings_var("network.retry", 3))
-        async with httpx.AsyncClient(transport=transport) as client:
-            response = await client.request(
-                method=self.method,
-                url=self.url,
-                headers=self.headers,
-                content=content,
-                timeout=self.timeout(),
-            )
-
-        return response
+        try:
+            async with httpx.AsyncClient(transport=transport) as client:
+                response = await client.request(
+                    method=self.method,
+                    url=self.url,
+                    headers=self.headers,
+                    content=content,
+                    timeout=self.timeout(),
+                )
+                response.raise_for_status()
+                return response
+
+        except httpx.TimeoutException as e:
+            raise ValueError(f"Request timed out: {str(e)}")
+        except httpx.HTTPStatusError as e:
+            raise ValueError(f"HTTP error {e.response.status_code}: {str(e)}")
+        except httpx.RequestError as e:
+            raise ValueError(f"Request failed: {str(e)}")
+        except Exception as e:
+            raise ValueError(f"Unexpected error: {str(e)}")
 
     async def verify(self) -> httpx.Response:
         match self:

agentic_security/mcp/main.py

@@ -16,9 +16,19 @@
 async def verify_llm(spec: str) -> dict:
     """Verify an LLM model specification using the FastAPI server."""
     url = f"{AGENTIC_SECURITY}/verify"
-    async with httpx.AsyncClient() as client:
-        response = await client.post(url, json={"spec": spec})
-        return response.json()
+    try:
+        async with httpx.AsyncClient(timeout=10.0) as client:
+            response = await client.post(url, json={"spec": spec})
+            response.raise_for_status()
+            return response.json()
+    except httpx.TimeoutException as e:
+        return {"error": f"Request timed out: {str(e)}"}
+    except httpx.HTTPStatusError as e:
+        return {"error": f"HTTP error {e.response.status_code}: {str(e)}"}
+    except httpx.RequestError as e:
+        return {"error": f"Request failed: {str(e)}"}
+    except Exception as e:
+        return {"error": f"Unexpected error: {str(e)}"}
 
 
 @mcp.tool()
@@ -39,36 +49,77 @@ async def start_scan(
         "probe_datasets": [],
         "secrets": {},
     }
-    async with httpx.AsyncClient() as client:
-        response = await client.post(url, json=payload)
-        return response.json()
+    try:
+        async with httpx.AsyncClient(timeout=10.0) as client:
+            response = await client.post(url, json=payload)
+            response.raise_for_status()
+            return response.json()
+    except httpx.TimeoutException as e:
+        return {"error": f"Request timed out: {str(e)}"}
+    except httpx.HTTPStatusError as e:
+        return {"error": f"HTTP error {e.response.status_code}: {str(e)}"}
+    except httpx.RequestError as e:
+        return {"error": f"Request failed: {str(e)}"}
+    except Exception as e:
+        return {"error": f"Unexpected error: {str(e)}"}
 
 
 @mcp.tool()
 async def stop_scan() -> dict:
     """Stop an ongoing scan via the FastAPI server."""
     url = f"{AGENTIC_SECURITY}/stop"
-    async with httpx.AsyncClient() as client:
-        response = await client.post(url)
-        return response.json()
+    try:
+        async with httpx.AsyncClient(timeout=10.0) as client:
+            response = await client.post(url)
+            response.raise_for_status()
+            return response.json()
+    except httpx.TimeoutException as e:
+        return {"error": f"Request timed out: {str(e)}"}
+    except httpx.HTTPStatusError as e:
+        return {"error": f"HTTP error {e.response.status_code}: {str(e)}"}
+    except httpx.RequestError as e:
+        return {"error": f"Request failed: {str(e)}"}
+    except Exception as e:
+        return {"error": f"Unexpected error: {str(e)}"}
 
 
 @mcp.tool()
 async def get_data_config() -> list:
     """Retrieve data configuration from the FastAPI server."""
     url = f"{AGENTIC_SECURITY}/v1/data-config"
-    async with httpx.AsyncClient() as client:
-        response = await client.get(url)
-        return response.json()
+    try:
+        async with httpx.AsyncClient(timeout=10.0) as client:
+            response = await client.get(url)
+            response.raise_for_status()
+            return response.json()
+    except httpx.TimeoutException as e:
+        return {"error": f"Request timed out: {str(e)}"}
+    except httpx.HTTPStatusError as e:
+        return {"error": f"HTTP error {e.response.status_code}: {str(e)}"}
+    except httpx.RequestError as e:
+        return {"error": f"Request failed: {str(e)}"}
+    except Exception as e:
+        return {"error": f"Unexpected error: {str(e)}"}
 
 
 @mcp.tool()
 async def get_spec_templates() -> list:
     """Retrieve data configuration 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()
+    try:
+        async with httpx.AsyncClient(timeout=10.0) as client:
+            response = await client.get(url)
+            response.raise_for_status()
+            return response.json()
+
+    except httpx.TimeoutException as e:
+        return {"error": f"Request timed out: {str(e)}"}
+    except httpx.HTTPStatusError as e:
+        return {"error": f"HTTP error {e.response.status_code}: {str(e)}"}
+    except httpx.RequestError as e:
+        return {"error": f"Request failed: {str(e)}"}
+    except Exception as e:
+        return {"error": f"Unexpected error: {str(e)}"}
 
 
 # Run the MCP server