From 8dd04357753fbe7eb764e736abf71831b656be4a Mon Sep 17 00:00:00 2001 From: Djebran Lezzoum Date: Fri, 11 Jul 2025 08:42:06 +0200 Subject: [PATCH] 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 --- src/app/endpoints/query.py | 3 +- src/app/endpoints/streaming_query.py | 5 ++- src/utils/mcp_headers.py | 42 +++++++++++++++++++ tests/unit/app/endpoints/test_query.py | 21 ++++++++-- .../app/endpoints/test_streaming_query.py | 20 +++++++-- 5 files changed, 83 insertions(+), 8 deletions(-) diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 265f55cf..47b7f36e 100644 --- a/src/app/endpoints/query.py +++ b/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] = { diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index 39ee4b93..4b3081fd 100644 --- a/src/app/endpoints/streaming_query.py +++ b/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] = { diff --git a/src/utils/mcp_headers.py b/src/utils/mcp_headers.py index 27e54b8c..0a2eb570 100644 --- a/src/utils/mcp_headers.py +++ b/src/utils/mcp_headers.py @@ -2,8 +2,13 @@ import json import logging +from urllib.parse import urlparse + from fastapi import Request +from configuration import AppConfig + + logger = logging.getLogger("app.endpoints.dependencies") @@ -46,3 +51,40 @@ 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]]: + """Process MCP headers by converting toolgroup names to URLs. + + This function takes MCP headers where keys can be either valid URLs or + toolgroup names. For valid URLs (HTTP/HTTPS), it keeps them as-is. For + toolgroup names, it looks up the corresponding MCP server URL in the + configuration and replaces the key with the URL. Unknown toolgroup names + are filtered out. + + Args: + mcp_headers: Dictionary with keys as URLs or toolgroup names + config: Application configuration containing MCP server definitions + + Returns: + Dictionary with URLs as keys and their corresponding headers as values + """ + 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 diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 4599a391..193390a7 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/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, } ) } diff --git a/tests/unit/app/endpoints/test_streaming_query.py b/tests/unit/app/endpoints/test_streaming_query.py index ced37be7..aa90dbd2 100644 --- a/tests/unit/app/endpoints/test_streaming_query.py +++ b/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