Merge pull request #20 from lokumai/feat/auth

README.md

@@ -13,37 +13,44 @@ short_description: Defect Solver - MCP Server
 This project implements a Model Context Protocol (MCP) server for defect solver bug localization using FastMCP. It provides multiple bug localization tools that forward bug report data to remote endpoints for prediction and analysis.
 
 ## Features
+
 - Exposes three MCP tools for bug localization:
-   - `multi_module_bug_localization`: Localizes bugs across the entire microservice architecture
-   - `single_module_bug_localization`: Localizes bugs within a specific module
-   - `search_space_routing`: Identifies likely microservices (search spaces) for a reported bug
+  - `multi_module_bug_localization`: Localizes bugs across the entire microservice architecture
+  - `single_module_bug_localization`: Localizes bugs within a specific module
+  - `search_space_routing`: Identifies likely microservices (search spaces) for a reported bug
 - Forwards issue data (key, summary, description, and optionally module) to configurable remote endpoints
 - Returns predicted source code locations and detailed workflow prompts
 - Uses async HTTP requests and robust error handling
 - Includes workflow and prompt templates for end-to-end bug localization and resolution
 
 ## Requirements
+
 - Python >= 3.13
 - See `pyproject.toml` for dependencies:
-   - fastmcp
-   - httpx
-   - loguru
-   - pydantic
-   - dotenv
-   - uv
+  - fastmcp
+  - httpx
+  - loguru
+  - pydantic
+  - dotenv
+  - uv
 
-### Note: The tools implemented in this server require a DEFECT SOLVER API KEY to send requests. To get your key, please contact [Lokum AI](https://github.com/lokumai).
+### Note: The tools implemented in this server require a DEFECT SOLVER API KEY to send requests. To get your key, please contact [Lokum AI](https://github.com/lokumai)
 
 ## Setup
+
 1. Clone the repository:
+
     ```sh
     git clone <your-repo-url>
     cd defect_solver_mcp_server
     ```
+
 2. Install dependencies using uv:
+
     ```sh
     uv sync
     ```
+
 3. Configure environment variables:
     - Copy `.env.example` to `.env` and set the following as needed:
        - `DS_API_BASE_URL`
@@ -54,32 +61,41 @@ This project implements a Model Context Protocol (MCP) server for defect solver
        - `TIMEOUT` (optional)
 
 ## Usage
+
 Run the MCP server:
+
 ```sh
 python main.py
 ```
+
 The server will start and expose the bug localization tools via MCP.
 
 ## Available Tools
+
 - **multi_module_bug_localization**: Use when the responsible microservice is unknown. Returns likely microservices and localized files.
 - **single_module_bug_localization**: Use when the affected module is known. Returns localized files within the specified module.
 - **search_space_routing**: Use to identify the most likely microservices (search spaces) for a reported bug.
 
 ## Prompts & Workflow
+
 The server provides prompt templates for:
+
 - Tool selection and workflow planning
 - Bug report augmentation
 - Bug localization and result explanation
 - Automated bug fixing and full resolution workflow
 
-### NOTE: There is a `chatmode.md` file personalized for DNext development that provides detailed instructions and rules on how to use the MCP server with chat-based interactions in IDEs like VSCode, Cursor, Windsurf and etc. Please refer to the [dnext-chatmode.md](https://github.com/pia-team/defect_solver_api/blob/main/.github/chatmodes/dnext-defect-solver.chatmode.md) for more information. 
+### NOTE: There is a `chatmode.md` file personalized for DNext development that provides detailed instructions and rules on how to use the MCP server with chat-based interactions in IDEs like VSCode, Cursor, Windsurf and etc. Please refer to the [dnext-chatmode.md](https://github.com/pia-team/defect_solver_api/blob/main/.github/chatmodes/dnext-defect-solver.chatmode.md) for more information
 
 ## Configuration
+
 - API base URL and endpoints are configured in `utils/config.py` and via environment variables:
-   - `DS_API_BASE_URL`, `DS_API_MULTIMODULE_ENDPOINT`, `DS_API_SINGLEMODULE_ENDPOINT`, `DS_API_SEARCHSPACE_ENDPOINT`
+  - `DS_API_BASE_URL`, `DS_API_MULTIMODULE_ENDPOINT`, `DS_API_SINGLEMODULE_ENDPOINT`, `DS_API_SEARCHSPACE_ENDPOINT`
 - Timeout and HuggingFace access token are also configurable.
-- 
+-
+
 ## MCP Inspector
+
 The MCP Inspector is a tool to help you visualize and debug the interactions with the MCP server.
 To use the MCP Inspector, simply run the command below in your terminal:
 
@@ -92,6 +108,7 @@ The GUI will open in your default web browser, allowing you to monitor the commu
 <img src="https://raw.githubusercontent.com/modelcontextprotocol/inspector/main/mcp-inspector.png" alt="MCP Inspector">
 
 ## Project Structure
+
 ```
 main.py                # Entry point and MCP server with tool and prompt definitions
 utils/config.py        # API configuration and environment variable loading
@@ -102,9 +119,10 @@ README.md              # Project documentation
 ```
 
 ## `mcp.json` Configuration Example for Client Usage
+
 To use this MCP server in a client application that supports MCP, you add the config below to `mcp.json`.
 
-#### PRIVATE MCP SERVER: This config is required for when deploying the MCP server in a private Hugging Face Space that requires authentication. 
+#### PRIVATE MCP SERVER: This config is required for when deploying the MCP server in a private Hugging Face Space that requires authentication
 
 ```json
 {
@@ -113,7 +131,7 @@ To use this MCP server in a client application that supports MCP, you add the co
             "url": "https://dnext-ds-mcp-server.hf.space/mcp/",
             "type": "http",
             "headers": {
-                "DS-API-Key": "${input:defect-solver-api-key}",
+                "DS-API-KEY": "${input:defect-solver-api-key}",
                 "Authorization": "Bearer ${input:hf-access-token}"
             }
         }
@@ -135,7 +153,7 @@ To use this MCP server in a client application that supports MCP, you add the co
 }
 ```
 
-#### PUBLIC MCP SERVER: For public spaces, you can omit the `hf-access-token` input or `Authorization` header.
+#### PUBLIC MCP SERVER: For public spaces, you can omit the `hf-access-token` input or `Authorization` header
 
 ```json
 {
@@ -144,7 +162,7 @@ To use this MCP server in a client application that supports MCP, you add the co
             "url": "https://dnext-ds-mcp-server.hf.space/mcp/",
             "type": "http",
             "headers": {
-                "DS-API-Key": "${input:defect-solver-api-key}"
+                "DS-API-KEY": "${input:defect-solver-api-key}"
             }
         }
     },
@@ -157,4 +175,4 @@ To use this MCP server in a client application that supports MCP, you add the co
         }
     ]
 }
-```
+```

main.py

@@ -3,13 +3,13 @@ from contextlib import asynccontextmanager
 import httpx
 from fastmcp import FastMCP, Context
 from loguru import logger
-from dataclasses import dataclass
 import uuid
 import os
 import sys
 
 from utils.config import API_CONFIG
 from utils.model import MultiModuleRequest, SingleModuleRequest, SearchSpaceRoutingRequest
+from utils.auth import DSContext, authenticate
 
 from fastmcp.server.dependencies import get_http_request
 from starlette.requests import Request
@@ -19,18 +19,15 @@ from dotenv import load_dotenv
 
 load_dotenv()
 
-
-@dataclass
-class DSContext:
-    """Typed context for DS MCP server"""
-
-    client: httpx.AsyncClient
-
-
 @asynccontextmanager
 async def ds_lifespan(server: FastMCP) -> AsyncIterator[DSContext]:
     client = httpx.AsyncClient(timeout=API_CONFIG["timeout"])
-    ctx = DSContext(client=client)
+    token = await authenticate()
+
+    if token:
+        client.headers.update({"Authorization": f"Bearer {token}"})
+
+    ctx = DSContext(client=client, access_token=token)
 
     try:
         yield ctx
@@ -70,9 +67,8 @@ async def multi_module_bug_localization(request: MultiModuleRequest, ctx: Contex
 
     # Get per-user API key from MCP Context
     user_request: Request = get_http_request()
-    logger.info("Received request headers: {}".format(user_request.headers))
-    defect_solver_api_key = user_request.headers.get("DS-API-Key", None)
-    logger.info("defect_solver_api_key: {}".format(defect_solver_api_key) )
+
+    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
@@ -82,7 +78,7 @@ async def multi_module_bug_localization(request: MultiModuleRequest, ctx: Contex
     headers = {
         "Content-Type": "application/json",
         "Authorization": f"Bearer {hf_token}" if hf_token else None,
-        "DS-API-Key": defect_solver_api_key,
+        "DS-API-KEY": defect_solver_api_key,
     }
     headers = {k: v for k, v in headers.items() if v is not None}
 
@@ -153,7 +149,7 @@ async def single_module_bug_localization(request: SingleModuleRequest, ctx: Cont
 
     # 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)
+    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
@@ -163,7 +159,7 @@ async def single_module_bug_localization(request: SingleModuleRequest, ctx: Cont
     headers = {
         "Content-Type": "application/json",
         "Authorization": f"Bearer {hf_token}" if hf_token else None,
-        "DS-API-Key": defect_solver_api_key,
+        "DS-API-KEY": defect_solver_api_key,
     }
     headers = {k: v for k, v in headers.items() if v is not None}
 
@@ -233,7 +229,7 @@ async def search_space_routing(request: SearchSpaceRoutingRequest, ctx: Context)
 
     # 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)
+    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
@@ -243,7 +239,7 @@ async def search_space_routing(request: SearchSpaceRoutingRequest, ctx: Context)
     headers = {
         "Content-Type": "application/json",
         "Authorization": f"Bearer {hf_token}" if hf_token else None,
-        "DS-API-Key": defect_solver_api_key,
+        "DS-API-KEY": defect_solver_api_key,
     }
     headers = {k: v for k, v in headers.items() if v is not None}
 

utils/auth.py

@@ -0,0 +1,34 @@
+from typing import Optional
+import httpx
+from dataclasses import dataclass
+
+from utils.config import KEYCLOAK_CONFIG, TOKEN_ENDPOINT
+
+
+@dataclass
+class DSContext:
+    """Typed context for DS MCP server"""
+
+    client: httpx.AsyncClient
+
+
+
+async def authenticate() -> Optional[str]:
+    """Authenticate with Keycloak and return access token"""
+    data = {
+        "grant_type": KEYCLOAK_CONFIG["grant_type"],
+        "client_id": KEYCLOAK_CONFIG["client_id"],
+        "username": KEYCLOAK_CONFIG["username"],
+        "password": KEYCLOAK_CONFIG["password"],
+    }
+
+    if KEYCLOAK_CONFIG["client_secret"]:
+        data["client_secret"] = KEYCLOAK_CONFIG["client_secret"]
+
+    async with httpx.AsyncClient() as client:
+        response = await client.post(TOKEN_ENDPOINT, data=data)
+        if response.status_code == 200:
+            return response.json().get("access_token")
+        else:
+            print(f"Authentication failed: {response.status_code} - {response.text}")
+            return None

utils/config.py

@@ -11,4 +11,19 @@ API_CONFIG = {
     "api_searchspace_endpoint": str(os.getenv("DS_API_SEARCHSPACE_ENDPOINT", "/mcp_search_space_routing")),
     "hf_access_token": os.getenv("HF_ACCESS_TOKEN", None),
     "timeout": int(os.getenv("TIMEOUT", 120)),
-}
+}
+
+
+# Keycloak Configuration
+KEYCLOAK_CONFIG = {
+    "base_url": os.getenv("KEYCLOAK_BASE_URL", "https://training.test.orbitant.dev"),
+    "realm": os.getenv("KEYCLOAK_REALM", "orbitant-realm"),
+    "client_id": os.getenv("KEYCLOAK_CLIENT_ID", "orbitant-backend-client"),
+    "client_secret": os.getenv("KEYCLOAK_CLIENT_SECRET", " "),
+    "grant_type": os.getenv("KEYCLOAK_GRANT_TYPE", "password"),
+    "username": os.getenv("KEYCLOAK_USERNAME", "admin"),
+    "password": os.getenv("KEYCLOAK_PASSWORD", "admin"),
+}
+
+# Token Endpoint
+TOKEN_ENDPOINT = f"{KEYCLOAK_CONFIG['base_url']}/realms/{KEYCLOAK_CONFIG['realm']}/protocol/openid-connect/token"