From 7d53fb29c282d3a1dcd2fd125d7b7c274a375a25 Mon Sep 17 00:00:00 2001
From: ykd007 <ykd007@users.noreply.github.com>
Date: Fri, 15 May 2026 09:45:39 +0530
Subject: [PATCH] Tighten type hints in mcp main and client modules

Use dict[str, Any] and list[dict[str, Any]] instead of plain dict/list
for the MCP tool return types, and add a proper return type annotation
to client.run(). Also imports typing.Any where needed so IDEs and mypy
can reason about the shapes without losing flexibility.
---
 agentic_security/mcp/client.py |  3 ++-
 agentic_security/mcp/main.py   | 14 ++++++++------
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/agentic_security/mcp/client.py b/agentic_security/mcp/client.py
index e8a29cc9..5b2a9a52 100644
--- a/agentic_security/mcp/client.py
+++ b/agentic_security/mcp/client.py
@@ -1,4 +1,5 @@
 import asyncio
+from typing import Any
 
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
@@ -13,7 +14,7 @@
 )
 
 
-async def run() -> None:
+async def run() -> tuple[Any, Any, Any] | None:
     try:
         logger.info(
             "Starting stdio client session with server parameters: %s", server_params
diff --git a/agentic_security/mcp/main.py b/agentic_security/mcp/main.py
index 7e19e2c6..4a374cd4 100644
--- a/agentic_security/mcp/main.py
+++ b/agentic_security/mcp/main.py
@@ -1,3 +1,5 @@
+from typing import Any
+
 import httpx
 from mcp.server.fastmcp import FastMCP
 
@@ -12,7 +14,7 @@
 
 
 @mcp.tool()
-async def verify_llm(spec: str) -> dict:
+async def verify_llm(spec: str) -> dict[str, Any]:
     """
     Verify an LLM model specification using the FastAPI server
 
@@ -34,7 +36,7 @@ async def start_scan(
     maxBudget: int,
     optimize: bool = False,
     enableMultiStepAttack: bool = False,
-) -> dict:
+) -> dict[str, Any]:
     """
     Start an LLM security scan via the FastAPI server.
     Returns:
@@ -48,7 +50,7 @@ async def start_scan(
 
     """
     url = f"{AGENTIC_SECURITY}/scan"
-    payload = {
+    payload: dict[str, Any] = {
         "llmSpec": llmSpec,
         "maxBudget": maxBudget,
         "datasets": [],
@@ -63,7 +65,7 @@ async def start_scan(
 
 
 @mcp.tool()
-async def stop_scan() -> dict:
+async def stop_scan() -> dict[str, Any]:
     """Stop an ongoing scan via the FastAPI server.
 
     Returns:
@@ -76,7 +78,7 @@ async def stop_scan() -> dict:
 
 
 @mcp.tool()
-async def get_data_config() -> list:
+async def get_data_config() -> list[dict[str, Any]]:
     """
     Retrieve data configuration from the FastAPI server.
 
@@ -90,7 +92,7 @@ async def get_data_config() -> list:
 
 
 @mcp.tool()
-async def get_spec_templates() -> list:
+async def get_spec_templates() -> list[dict[str, Any]]:
     """
     Retrieve data configuration from the FastAPI server.