feat(auth): implement Keycloak authentication and configuration

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

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"