add mcp-headers support (#1716)

Add mcp headers support to allow authentication to mcp servers and their target AAP services
issue: https://issues.redhat.com/browse/AAP-49395

Signed-off-by: Djebran Lezzoum <[email protected]>

Author: ldjebran

ansible_ai_connect/ai/api/model_pipelines/http/configuration.py

@@ -31,6 +31,11 @@
 # ENABLE_HEALTHCHECK_XXX
 
 
+class MCPServersSerializer(serializers.Serializer):
+    name = serializers.CharField(required=True)
+    type = serializers.CharField(required=True)
+
+
 @dataclass
 class HttpConfiguration(BaseConfig):
 
@@ -42,13 +47,16 @@ def __init__(
         enable_health_check: Optional[bool],
         verify_ssl: bool,
         stream: bool = False,
+        mcp_servers: Optional[list[dict[str, str]]] = None,
     ):
         super().__init__(inference_url, model_id, timeout, enable_health_check)
         self.verify_ssl = verify_ssl
         self.stream = stream
+        self.mcp_servers = mcp_servers or []
 
     verify_ssl: bool
     stream: bool
+    mcp_servers: Optional[list[dict[str, str]]] = None
 
 
 @Register(api_type="http")
@@ -64,6 +72,7 @@ def __init__(self, **kwargs):
                 enable_health_check=kwargs["enable_health_check"],
                 verify_ssl=kwargs["verify_ssl"],
                 stream=kwargs["stream"],
+                mcp_servers=kwargs["mcp_servers"],
             ),
         )
 
@@ -72,3 +81,6 @@ def __init__(self, **kwargs):
 class HttpConfigurationSerializer(BaseConfigSerializer):
     verify_ssl = serializers.BooleanField(required=False, default=True)
     stream = serializers.BooleanField(required=False, default=False)
+    mcp_servers = serializers.ListSerializer(
+        child=MCPServersSerializer(), allow_empty=True, required=False, default=None
+    )

ansible_ai_connect/ai/api/model_pipelines/http/pipelines.py

@@ -19,6 +19,7 @@
 
 import aiohttp
 import requests
+from django.conf import settings
 from django.http import StreamingHttpResponse
 from health_check.exceptions import ServiceUnavailable
 
@@ -176,7 +177,7 @@ def invoke(self, params: ChatBotParameters) -> ChatBotResponse:
         conversation_id = params.conversation_id
         provider = params.provider
         model_id = params.model_id
-        system_prompt = params.system_prompt
+        system_prompt = params.system_prompt or settings.CHATBOT_DEFAULT_SYSTEM_PROMPT
 
         data = {
             "query": query,
@@ -188,6 +189,10 @@ def invoke(self, params: ChatBotParameters) -> ChatBotResponse:
         if system_prompt:
             data["system_prompt"] = str(system_prompt)
 
+        headers = self.headers or {}
+        if params.mcp_headers:
+            headers["MCP-HEADERS"] = json.dumps(params.mcp_headers)
+
         response = requests.post(
             self.config.inference_url + "/v1/query",
             headers=self.headers,
@@ -262,11 +267,14 @@ async def async_invoke(self, params: StreamingChatBotParameters) -> AsyncGenerat
                 "Accept": "application/json,text/event-stream",
             }
 
+            if params.mcp_headers:
+                headers["MCP-HEADERS"] = json.dumps(params.mcp_headers)
+
             query = params.query
             conversation_id = params.conversation_id
             provider = params.provider
             model_id = params.model_id
-            system_prompt = params.system_prompt
+            system_prompt = params.system_prompt or settings.CHATBOT_DEFAULT_SYSTEM_PROMPT
             media_type = params.media_type
 
             data = {

ansible_ai_connect/ai/api/model_pipelines/pipelines.py

@@ -16,7 +16,7 @@
 from abc import ABCMeta, abstractmethod
 from typing import Any, Dict, Generic, Optional
 
-from attrs import define
+from attrs import define, field
 from django.conf import settings
 from django.http import StreamingHttpResponse
 from rest_framework import serializers
@@ -231,6 +231,7 @@ class ChatBotParameters:
     model_id: str
     conversation_id: Optional[str]
     system_prompt: str
+    mcp_headers: Optional[dict[str, dict[str, str]]] = field(kw_only=True, default=None)
 
     @classmethod
     def init(
@@ -240,13 +241,15 @@ def init(
         model_id: Optional[str] = None,
         conversation_id: Optional[str] = None,
         system_prompt: Optional[str] = None,
+        mcp_headers: Optional[dict[str, dict[str, str]]] = None,
     ):
         return cls(
             query=query,
             provider=provider,
             model_id=model_id,
             conversation_id=conversation_id,
             system_prompt=system_prompt,
+            mcp_headers=mcp_headers,
         )
 
 
@@ -266,6 +269,7 @@ def init(
         conversation_id: Optional[str] = None,
         system_prompt: Optional[str] = None,
         media_type: Optional[str] = None,
+        mcp_headers: Optional[dict[str, dict[str, str]]] = None,
     ):
         return cls(
             query=query,
@@ -274,6 +278,7 @@ def init(
             conversation_id=conversation_id,
             system_prompt=system_prompt,
             media_type=media_type,
+            mcp_headers=mcp_headers,
         )
 
 

ansible_ai_connect/ai/api/model_pipelines/tests/__init__.py

@@ -81,6 +81,7 @@ def mock_pipeline_config(pipeline_provider: t_model_mesh_api_type, **kwargs):
                 enable_health_check=extract("enable_health_check", False, **kwargs),
                 verify_ssl=extract("verify_ssl", False, **kwargs),
                 stream=extract("stream", False, **kwargs),
+                mcp_servers=extract("mcp_servers", [], **kwargs),
             )
         case "llamacpp":
             return LlamaCppConfiguration(

ansible_ai_connect/ai/api/versions/v1/tests/test_jwt_authentication.py

@@ -5,16 +5,20 @@
 from datetime import datetime, timedelta
 from http import HTTPStatus
 from unittest import mock
+from unittest.mock import Mock, patch
 
 import jwt
 import requests
 from ansible_base.resource_registry.models.service_identifier import service_id
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
+from django.apps import apps
 from django.test import override_settings
 from rest_framework.test import APITransactionTestCase
 
+from ansible_ai_connect.ai.api.model_pipelines.http.pipelines import HttpChatBotPipeline
+from ansible_ai_connect.ai.api.model_pipelines.tests import mock_pipeline_config
 from ansible_ai_connect.ai.api.versions.v1.test_base import API_VERSION
 from ansible_ai_connect.test_utils import APIVersionTestCaseBase
 from ansible_ai_connect.users.models import User
@@ -40,8 +44,35 @@
 ISSUER = "ansible-issuer"
 AUDIENCE = "ansible-services"
 
+ANSIBLE_AI_MODEL_MESH_CONFIG = {
+    "ModelPipelineChatBot": {
+        "provider": "http",
+        "config": {
+            "inference_url": "http://localhost:8080",
+            "model_id": "granite-3.3-8b-instruct",
+            "enable_health_check": True,
+            "mcp_servers": [
+                {"name": "mcp::aap-controller", "type": "controller"},
+                {"name": "mcp::aap-gateway", "type": "gateway"},
+                {"name": "mcp::aap-lightspeed", "type": "lightspeed"},
+            ],
+        },
+    },
+}
+
 
 @override_settings(ANSIBLE_BASE_JWT_KEY=test_encryption_public_key)
+@patch.object(
+    apps.get_app_config("ai"),
+    "get_model_pipeline",
+    Mock(
+        return_value=HttpChatBotPipeline(
+            mock_pipeline_config(
+                "http", **ANSIBLE_AI_MODEL_MESH_CONFIG["ModelPipelineChatBot"]["config"]
+            ),
+        )
+    ),
+)
 class TestJWTAuthentication(APIVersionTestCaseBase, APITransactionTestCase):
     api_version = API_VERSION
 
@@ -113,5 +144,23 @@ def json(self):
         response = self.jwt_client.post(
             self.api_version_reverse("chat"), {"query": "Hello"}, format="json"
         )
+
+        mock_requests_post.assert_called_once()
+        _, kwargs = mock_requests_post.call_args
+
+        headers = kwargs.get("headers", None)
+        self.assertIsNotNone(headers)
+
+        mcp_headers_string = headers.get("MCP-HEADERS", None)
+        self.assertIsNotNone(headers)
+
+        expected_mcp_headers = {
+            "mcp::aap-controller": {"X-DAB-JW-TOKEN": self.encrypted_token},
+            "mcp::aap-lightspeed": {"X-DAB-JW-TOKEN": self.encrypted_token},
+        }
+        mcp_headers = json.loads(mcp_headers_string)
+
+        self.assertDictEqual(mcp_headers, expected_mcp_headers)
+
         self.assertEqual(response.status_code, HTTPStatus.OK)
         self.assertDictEqual(response.data, expected_response)

ansible_ai_connect/ai/api/views.py

@@ -29,6 +29,7 @@
 from rest_framework import permissions, serializers
 from rest_framework import status as rest_framework_status
 from rest_framework.generics import GenericAPIView
+from rest_framework.request import Request
 from rest_framework.response import Response
 from rest_framework.views import APIView
 
@@ -108,6 +109,7 @@
 from ..feature_flags import FeatureFlags
 from .data.data_model import ContentMatchPayloadData, ContentMatchResponseDto
 from .model_pipelines.exceptions import ModelTimeoutError
+from .model_pipelines.http.configuration import HttpConfiguration
 from .permissions import (
     BlockUserWithoutSeat,
     BlockUserWithoutSeatAndWCAReadyOrg,
@@ -265,6 +267,27 @@ def finalize_response(self, request, response, *args, **kwargs):
         send_schema1_event(self.event)
         return response
 
+    @staticmethod
+    def get_mcp_headers(request: Request, config: HttpConfiguration) -> dict:
+        mcp_headers = {}
+        jwt_header_name = "X-DAB-JW-TOKEN"
+        token = request.headers.get(jwt_header_name, None)
+        user = request.user
+        if token and user.is_authenticated and user.aap_user:
+            for mcp_server in config.mcp_servers:
+                if mcp_server["type"] in ["controller", "eda", "hub", "lightspeed"]:
+                    mcp_headers[mcp_server["name"]] = {jwt_header_name: token}
+                # This functionality seems experimental for gateway and does not allow the user to
+                # access wide range of api endpoints, we need to find a solution for gateway,
+                # but for the moment comment this code.
+                # elif mcp_server["type"] == "gateway":
+                #     from ansible_base.resource_registry.resource_server import get_service_token
+                #     user_ansible_id = str(user.ansible_id_for_filter)
+                #     token = get_service_token(user_id=user_ansible_id, expiration=3600)
+                #     mcp_headers[mcp_server["name"]] = {"X-ANSIBLE-SERVICE-AUTH": token}
+
+        return mcp_headers
+
 
 class Completions(APIView):
     """
@@ -1106,13 +1129,16 @@ def post(self, request) -> Response:
         self.event.conversation_id = conversation_id
         self.event.modelName = self.req_model_id or self.llm.config.model_id
 
+        mcp_headers = self.get_mcp_headers(request, self.llm.config)
+
         data = self.llm.invoke(
             ChatBotParameters.init(
                 query=req_query,
                 system_prompt=req_system_prompt,
                 model_id=self.req_model_id or self.llm.config.model_id,
                 provider=req_provider,
                 conversation_id=conversation_id,
+                mcp_headers=mcp_headers,
             )
         )
 
@@ -1122,10 +1148,10 @@ def post(self, request) -> Response:
             raise ChatbotInvalidResponseException()
 
         # Finalise Segment Event with response details
-        self.event.chat_truncated = bool(data["truncated"])
+        self.event.chat_truncated = bool(data.get("truncated", False))
         self.event.chat_referenced_documents = [
             ChatBotResponseDocsReferences(docs_url=rd["docs_url"], title=rd["title"])
-            for rd in data["referenced_documents"]
+            for rd in data.get("referenced_documents", [])
         ]
         self.event.chat_response = anonymize_struct(data["response"])
         self.event.chat_response = (
@@ -1203,6 +1229,8 @@ def post(self, request) -> StreamingHttpResponse:
         self.event.conversation_id = conversation_id
         self.event.modelName = self.req_model_id or self.llm.config.model_id
 
+        mcp_headers = self.get_mcp_headers(request, self.llm.config)
+
         return self.llm.invoke(
             StreamingChatBotParameters.init(
                 query=req_query,
@@ -1211,5 +1239,6 @@ def post(self, request) -> StreamingHttpResponse:
                 provider=req_provider,
                 conversation_id=conversation_id,
                 media_type=media_type,
+                mcp_headers=mcp_headers,
             )
         )

ansible_ai_connect/main/settings/base.py

@@ -595,7 +595,7 @@ def is_ssl_enabled(value: str) -> bool:
 # ------------------------------------------
 CHATBOT_DEFAULT_PROVIDER = os.getenv("CHATBOT_DEFAULT_PROVIDER")
 CHATBOT_DEBUG_UI = os.getenv("CHATBOT_DEBUG_UI", "False").lower() == "true"
-LIGHTSPEED_TOOL_GROUP_TOKEN = os.getenv("LIGHTSPEED_TOOL_GROUP_TOKEN")
+CHATBOT_DEFAULT_SYSTEM_PROMPT = os.getenv("CHATBOT_DEFAULT_SYSTEM_PROMPT")
 # ==========================================
 
 # ==========================================

tools/docker-compose/compose.yaml

@@ -76,6 +76,7 @@ services:
       - CHATBOT_URL=${CHATBOT_URL}
       - CHATBOT_DEFAULT_PROVIDER=${CHATBOT_DEFAULT_PROVIDER}
       - CHATBOT_DEFAULT_MODEL=${CHATBOT_DEFAULT_MODEL}
+      - CHATBOT_DEFAULT_SYSTEM_PROMPT=${CHATBOT_DEFAULT_SYSTEM_PROMPT}
       - ANSIBLE_AI_MODEL_MESH_CONFIG=${ANSIBLE_AI_MODEL_MESH_CONFIG}
       - ANSIBLE_AI_ENABLE_ROLE_GEN_ENDPOINT=${ANSIBLE_AI_ENABLE_ROLE_GEN_ENDPOINT}
       - AAP_API_URL=${AAP_API_URL}