Allow mcp headers to contain mcp servers names

https://issues.redhat.com/browse/LCORE-353
this PR allow the client to provide mcp headers with mcp servers names
mcp_headers = {
    "mcp-server-url-1": {"HEADER-1": "HEADER-1-VALUE-1", "HEADER-1": "HEADER-1-VALUE-2"},
    "mcp-server-name": {"HEADER-2": "HEADER-2-VALUE"}
}

the header with "mcp-server-name" will be replaced with it corresponding url using the configuration.mcp_servers, if it does exit it's getting ignored.

Signed-off-by: Djebran Lezzoum <ldjebran@gmail.com>

Author: ldjebran

src/app/endpoints/query.py

@@ -29,7 +29,7 @@
 from auth import get_auth_dependency
 from utils.common import retrieve_user_id
 from utils.endpoints import check_configuration_loaded, get_system_prompt
-from utils.mcp_headers import mcp_headers_dependency
+from utils.mcp_headers import mcp_headers_dependency, handle_mcp_headers_with_toolgroups
 from utils.suid import get_suid
 from utils.types import GraniteToolParser
 
@@ -231,6 +231,7 @@ def retrieve_response(
     # preserve compatibility when mcp_headers is not provided
     if mcp_headers is None:
         mcp_headers = {}
+    mcp_headers = handle_mcp_headers_with_toolgroups(mcp_headers, configuration)
     if not mcp_headers and token:
         for mcp_server in configuration.mcp_servers:
             mcp_headers[mcp_server.url] = {

src/app/endpoints/streaming_query.py

@@ -22,7 +22,7 @@
 from models.requests import QueryRequest
 from utils.endpoints import check_configuration_loaded, get_system_prompt
 from utils.common import retrieve_user_id
-from utils.mcp_headers import mcp_headers_dependency
+from utils.mcp_headers import mcp_headers_dependency, handle_mcp_headers_with_toolgroups
 from utils.suid import get_suid
 from utils.types import GraniteToolParser
 
@@ -290,6 +290,9 @@ async def retrieve_response(
     # preserve compatibility when mcp_headers is not provided
     if mcp_headers is None:
         mcp_headers = {}
+
+    mcp_headers = handle_mcp_headers_with_toolgroups(mcp_headers, configuration)
+
     if not mcp_headers and token:
         for mcp_server in configuration.mcp_servers:
             mcp_headers[mcp_server.url] = {

src/utils/mcp_headers.py

@@ -2,6 +2,10 @@
 
 import json
 import logging
+from urllib.parse import urlparse
+
+from configuration import AppConfig
+
 from fastapi import Request
 
 logger = logging.getLogger("app.endpoints.dependencies")
@@ -46,3 +50,25 @@ def extract_mcp_headers(request: Request) -> dict[str, dict[str, str]]:
             )
             mcp_headers = {}
     return mcp_headers
+
+
+def handle_mcp_headers_with_toolgroups(
+    mcp_headers: dict[str, dict[str, str]], config: AppConfig
+) -> dict[str, dict[str, str]]:
+    converted_mcp_headers = {}
+
+    for key, item in mcp_headers.items():
+        key_url_parsed = urlparse(key)
+        if key_url_parsed.scheme in ("http", "https") and key_url_parsed.netloc:
+            # a valid url is supplied, deliver it as is
+            converted_mcp_headers[key] = item
+        else:
+            # assume the key is a toolgroup name
+            # look for toolgroups name in mcp_servers configuration
+            # if the mcp server is not found, the mcp header gets ignored
+            for mcp_server in config.mcp_servers:
+                if mcp_server.name == key and mcp_server.url:
+                    converted_mcp_headers[mcp_server.url] = item
+                    break
+
+    return converted_mcp_headers

tests/unit/app/endpoints/test_query.py

@@ -694,8 +694,14 @@ def test_retrieve_response_with_mcp_servers_and_mcp_headers(mocker):
     model_id = "fake_model_id"
     access_token = ""
     mcp_headers = {
-        "http://localhost:3000": {"Authorization": "Bearer test_token_123"},
-        "https://git.example.com/mcp": {"Authorization": "Bearer test_token_123"},
+        "filesystem-server": {"Authorization": "Bearer test_token_123"},
+        "git-server": {"Authorization": "Bearer test_token_456"},
+        "http://another-server-mcp-server:3000": {
+            "Authorization": "Bearer test_token_789"
+        },
+        "unknown-mcp-server": {
+            "Authorization": "Bearer test_token_for_unknown-mcp-server"
+        },
     }
 
     response, conversation_id = retrieve_response(
@@ -718,11 +724,20 @@ def test_retrieve_response_with_mcp_servers_and_mcp_headers(mocker):
         None,  # conversation_id
     )
 
+    expected_mcp_headers = {
+        "http://localhost:3000": {"Authorization": "Bearer test_token_123"},
+        "https://git.example.com/mcp": {"Authorization": "Bearer test_token_456"},
+        "http://another-server-mcp-server:3000": {
+            "Authorization": "Bearer test_token_789"
+        },
+        # we do not put "unknown-mcp-server" url as it's unknown to lightspeed-stack
+    }
+
     # Check that the agent's extra_headers property was set correctly
     expected_extra_headers = {
         "X-LlamaStack-Provider-Data": json.dumps(
             {
-                "mcp_headers": mcp_headers,
+                "mcp_headers": expected_mcp_headers,
             }
         )
     }

tests/unit/app/endpoints/test_streaming_query.py

@@ -762,8 +762,14 @@ async def test_retrieve_response_with_mcp_servers_and_mcp_headers(mocker):
     model_id = "fake_model_id"
     access_token = ""
     mcp_headers = {
-        "http://localhost:3000": {"Authorization": "Bearer test_token_123"},
-        "https://git.example.com/mcp": {"Authorization": "Bearer test_token_456"},
+        "filesystem-server": {"Authorization": "Bearer test_token_123"},
+        "git-server": {"Authorization": "Bearer test_token_456"},
+        "http://another-server-mcp-server:3000": {
+            "Authorization": "Bearer test_token_789"
+        },
+        "unknown-mcp-server": {
+            "Authorization": "Bearer test_token_for_unknown-mcp-server"
+        },
     }
 
     response, conversation_id = await retrieve_response(
@@ -786,9 +792,17 @@ async def test_retrieve_response_with_mcp_servers_and_mcp_headers(mocker):
         None,  # conversation_id
     )
 
+    expected_mcp_headers = {
+        "http://localhost:3000": {"Authorization": "Bearer test_token_123"},
+        "https://git.example.com/mcp": {"Authorization": "Bearer test_token_456"},
+        "http://another-server-mcp-server:3000": {
+            "Authorization": "Bearer test_token_789"
+        },
+        # we do not put "unknown-mcp-server" url as it's unknown to lightspeed-stack
+    }
     # Check that the agent's extra_headers property was set correctly
     expected_extra_headers = {
-        "X-LlamaStack-Provider-Data": json.dumps({"mcp_headers": mcp_headers})
+        "X-LlamaStack-Provider-Data": json.dumps({"mcp_headers": expected_mcp_headers})
     }
     assert mock_agent.extra_headers == expected_extra_headers