feat: enhance configuration and API key handling in MCP server

.env.example

@@ -5,5 +5,14 @@ DS_API_MULTIMODULE_ENDPOINT=/mcp_multi_module_bug_localization
 DS_API_SINGLEMODULE_ENDPOINT=/mcp_single_module_bug_localization
 DS_API_SEARCHSPACE_ENDPOINT=/mcp_search_space_routing
 
+# Hugging Face Access Token for when the API is hosted on Hugging Face Spaces
+HF_ACCESS_TOKEN=your_hf_access_token_here
+
 # Request timeout in seconds (for requests between MCP Server and Defect Solver API)
-TIMEOUT=120
+TIMEOUT=120
+
+# MCP Server Configuration
+TRANSPORT_MODE=streamable-http
+HOST=0.0.0.0
+PORT=8000
+LOG_LEVEL=INFO

README.md

@@ -1,3 +1,13 @@
+---
+title: Defect Solver API
+emoji: 🛠️
+colorFrom: orange
+colorTo: yellow
+sdk: docker
+pinned: false
+short_description: Defect Solver - MCP Server 
+---
+
 # Defect Solver MCP Server
 
 This project implements a Model Context Protocol (MCP) server for defect solver bug localization using FastMCP. It forwards bug description data (ticket) to a remote endpoint for prediction and analysis.

main.py

@@ -9,6 +9,9 @@ import uuid
 from utils.config import API_CONFIG
 from utils.model import MultiModuleRequest, SingleModuleRequest, SearchSpaceRoutingRequest
 
+from fastmcp.server.dependencies import get_http_request
+from starlette.requests import Request
+
 
 @dataclass
 class DSContext:
@@ -56,7 +59,21 @@ async def multi_module_bug_localization(request: MultiModuleRequest, ctx: Contex
     api_url = API_CONFIG["api_base_url"] + API_CONFIG["api_multimodule_endpoint"]
     hf_token = API_CONFIG.get("hf_access_token", None)
 
-    headers = {"Authorization": f"Bearer {hf_token}", "Content-Type": "application/json"}
+    # Get per-user API key from MCP Context
+    user_request: Request = get_http_request()
+    defect_solver_api_key = user_request.headers.get("DS-API-Key", None)
+    logger.info(
+        "Using API key: {}".format(defect_solver_api_key)
+        if defect_solver_api_key
+        else "No Defect Solver API key provided"
+    )
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {hf_token}" if hf_token else None,
+        "DS-API-Key": defect_solver_api_key,
+    }
+    headers = {k: v for k, v in headers.items() if v is not None}
 
     try:
         # Generate a UUID for issue_key even if provided (overwrite)
@@ -90,7 +107,7 @@ async def multi_module_bug_localization(request: MultiModuleRequest, ctx: Contex
         logger.error(f"Request error while calling multimodule bug localization endpoint: {e}")
         return {"error": "Request error", "details": str(e)}
     except Exception as e:
-        logger.exception("Unexpected error during multimodule bug localization request")
+        logger.error("Unexpected error during multimodule bug localization request")
         return {"error": "Unexpected error", "details": str(e)}
 
 
@@ -118,7 +135,21 @@ async def single_module_bug_localization(request: SingleModuleRequest, ctx: Cont
     api_url = API_CONFIG["api_base_url"] + API_CONFIG["api_singlemodule_endpoint"]
     hf_token = API_CONFIG.get("hf_access_token", None)
 
-    headers = {"Authorization": f"Bearer {hf_token}", "Content-Type": "application/json"}
+    # Get per-user API key from MCP Context
+    user_request: Request = get_http_request()
+    defect_solver_api_key = user_request.headers.get("DS-API-Key", None)
+    logger.info(
+        "Using API key: {}".format(defect_solver_api_key)
+        if defect_solver_api_key
+        else "No Defect Solver API key provided"
+    )
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {hf_token}" if hf_token else None,
+        "DS-API-Key": defect_solver_api_key,
+    }
+    headers = {k: v for k, v in headers.items() if v is not None}
 
     try:
         # Generate a UUID for issue_key even if provided (overwrite)
@@ -151,7 +182,7 @@ async def single_module_bug_localization(request: SingleModuleRequest, ctx: Cont
         logger.error(f"Request error while calling endpoint: {e}")
         return {"error": "Request error", "details": str(e)}
     except Exception as e:
-        logger.exception("Unexpected error during single-module bug localization")
+        logger.error("Unexpected error during single-module bug localization")
         return {"error": "Unexpected error", "details": str(e)}
 
 
@@ -179,7 +210,21 @@ async def search_space_routing(request: SearchSpaceRoutingRequest, ctx: Context)
     api_url = API_CONFIG["api_base_url"] + API_CONFIG["api_searchspace_endpoint"]
     hf_token = API_CONFIG.get("hf_access_token", None)
 
-    headers = {"Authorization": f"Bearer {hf_token}", "Content-Type": "application/json"}
+    # Get per-user API key from MCP Context
+    user_request: Request = get_http_request()
+    defect_solver_api_key = user_request.headers.get("DS-API-Key", None)
+    logger.info(
+        "Using API key: {}".format(defect_solver_api_key)
+        if defect_solver_api_key
+        else "No Defect Solver API key provided"
+    )
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {hf_token}" if hf_token else None,
+        "DS-API-Key": defect_solver_api_key,
+    }
+    headers = {k: v for k, v in headers.items() if v is not None}
 
     try:
         # Generate a UUID for issue_key even if provided (overwrite)
@@ -211,7 +256,7 @@ async def search_space_routing(request: SearchSpaceRoutingRequest, ctx: Context)
         logger.error(f"Request error while calling endpoint: {e}")
         return {"error": "Request error", "details": str(e)}
     except Exception as e:
-        logger.exception("Unexpected error during search space routing")
+        logger.error("Unexpected error during search space routing")
         return {"error": "Unexpected error", "details": str(e)}
 
 
@@ -340,11 +385,10 @@ Return a 2-3x more detailed/enhanced bug report with precise technical terms for
 """
 
 
-
-
-
-
-
-
 if __name__ == "__main__":
-    mcp.run(transport="streamable-http")
+    TRANSPORT_MODE = API_CONFIG["transport_mode"]
+    HOST = API_CONFIG["host"]
+    PORT = API_CONFIG["port"]
+    LOG_LEVEL = API_CONFIG["log_level"]
+
+    mcp.run(transport=TRANSPORT_MODE, host=HOST, port=PORT, log_level=LOG_LEVEL)

utils/config.py

@@ -11,4 +11,8 @@ API_CONFIG = {
     "api_searchspace_endpoint": os.getenv("DS_API_SEARCHSPACE_ENDPOINT", "/mcp_search_space_routing"),
     "hf_access_token": os.getenv("HF_ACCESS_TOKEN", None),
     "timeout": os.getenv("TIMEOUT", 120),
+    "transport_mode": os.getenv("TRANSPORT_MODE", "streamable-http"),
+    "host": os.getenv("HOST", "0.0.0.0"),
+    "port": os.getenv("PORT", 8000),
+    "log_level": os.getenv("LOG_LEVEL", "INFO"),
 }