Merge pull request #1143 from tisnik/lcore-1173

LCORE-1173: named type for MCP headers complicated structure

Author: tisnik

src/app/endpoints/a2a.py

@@ -44,7 +44,7 @@
 from configuration import configuration
 from models.config import Action
 from models.requests import QueryRequest
-from utils.mcp_headers import mcp_headers_dependency
+from utils.mcp_headers import mcp_headers_dependency, McpHeaders
 from utils.responses import (
     extract_text_from_response_output_item,
     prepare_responses_params,
@@ -183,17 +183,15 @@ class A2AAgentExecutor(AgentExecutor):
     routing queries to the LLM backend using the Responses API.
     """
 
-    def __init__(
-        self, auth_token: str, mcp_headers: Optional[dict[str, dict[str, str]]] = None
-    ):
+    def __init__(self, auth_token: str, mcp_headers: Optional[McpHeaders] = None):
         """Initialize the A2A agent executor.
 
         Args:
             auth_token: Authentication token for the request
             mcp_headers: MCP headers for context propagation
         """
         self.auth_token: str = auth_token
-        self.mcp_headers: dict[str, dict[str, str]] = mcp_headers or {}
+        self.mcp_headers: McpHeaders = mcp_headers or {}
 
     async def execute(
         self,
@@ -648,9 +646,7 @@ async def get_agent_card(  # pylint: disable=unused-argument
         raise
 
 
-async def _create_a2a_app(
-    auth_token: str, mcp_headers: dict[str, dict[str, str]]
-) -> Any:
+async def _create_a2a_app(auth_token: str, mcp_headers: McpHeaders) -> Any:
     """Create an A2A Starlette application instance with auth context.
 
     Args:
@@ -681,7 +677,7 @@ async def _create_a2a_app(
 async def handle_a2a_jsonrpc(  # pylint: disable=too-many-locals,too-many-statements
     request: Request,
     auth: Annotated[AuthTuple, Depends(auth_dependency)],
-    mcp_headers: dict[str, dict[str, str]] = Depends(mcp_headers_dependency),
+    mcp_headers: McpHeaders = Depends(mcp_headers_dependency),
 ) -> Response | StreamingResponse:
     """
     Handle A2A JSON-RPC requests following the A2A protocol specification.

src/app/endpoints/query.py

@@ -40,7 +40,7 @@
     check_configuration_loaded,
     validate_and_retrieve_conversation,
 )
-from utils.mcp_headers import mcp_headers_dependency
+from utils.mcp_headers import mcp_headers_dependency, McpHeaders
 from utils.query import (
     consume_query_tokens,
     handle_known_apistatus_errors,
@@ -93,7 +93,7 @@ async def query_endpoint_handler(
     request: Request,
     query_request: QueryRequest,
     auth: Annotated[AuthTuple, Depends(get_auth_dependency())],
-    mcp_headers: dict[str, dict[str, str]] = Depends(mcp_headers_dependency),
+    mcp_headers: McpHeaders = Depends(mcp_headers_dependency),
 ) -> QueryResponse:
     """
     Handle request to the /query endpoint using Responses API.

src/app/endpoints/streaming_query.py

@@ -57,7 +57,7 @@
     check_configuration_loaded,
     validate_and_retrieve_conversation,
 )
-from utils.mcp_headers import mcp_headers_dependency
+from utils.mcp_headers import mcp_headers_dependency, McpHeaders
 from utils.query import (
     consume_query_tokens,
     extract_provider_and_model_from_model_id,
@@ -118,7 +118,7 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
     request: Request,
     query_request: QueryRequest,
     auth: Annotated[AuthTuple, Depends(get_auth_dependency())],
-    mcp_headers: dict[str, dict[str, str]] = Depends(mcp_headers_dependency),
+    mcp_headers: McpHeaders = Depends(mcp_headers_dependency),
 ) -> StreamingResponse:
     """
     Handle request to the /streaming_query endpoint using Responses API.

src/utils/mcp_headers.py

@@ -10,8 +10,10 @@
 
 logger = logging.getLogger("app.endpoints.dependencies")
 
+type McpHeaders = dict[str, dict[str, str]]
 
-async def mcp_headers_dependency(request: Request) -> dict[str, dict[str, str]]:
+
+async def mcp_headers_dependency(request: Request) -> McpHeaders:
     """Get the MCP headers dependency to passed to mcp servers.
 
     mcp headers is a json dictionary or mcp url paths and their respective headers
@@ -25,7 +27,7 @@ async def mcp_headers_dependency(request: Request) -> dict[str, dict[str, str]]:
     return extract_mcp_headers(request)
 
 
-def extract_mcp_headers(request: Request) -> dict[str, dict[str, str]]:
+def extract_mcp_headers(request: Request) -> McpHeaders:
     """Extract mcp headers from MCP-HEADERS header.
 
     If the header is missing, contains invalid JSON, or the decoded
@@ -56,8 +58,8 @@ def extract_mcp_headers(request: Request) -> dict[str, dict[str, str]]:
 
 
 def handle_mcp_headers_with_toolgroups(
-    mcp_headers: dict[str, dict[str, str]], config: AppConfig
-) -> dict[str, dict[str, str]]:
+    mcp_headers: McpHeaders, config: AppConfig
+) -> McpHeaders:
     """Process MCP headers by converting toolgroup names to URLs.
 
     This function takes MCP headers where keys can be either valid URLs or

src/utils/responses.py

@@ -37,6 +37,7 @@
     prepare_input,
     select_model_and_provider_id,
 )
+from utils.mcp_headers import McpHeaders
 from utils.suid import to_llama_stack_conversation_id
 from utils.token_counter import TokenCounter
 from utils.types import (
@@ -141,7 +142,7 @@ async def prepare_tools(
     query_request: QueryRequest,
     token: str,
     config: AppConfig,
-    mcp_headers: Optional[dict[str, dict[str, str]]] = None,
+    mcp_headers: Optional[McpHeaders] = None,
 ) -> Optional[list[dict[str, Any]]]:
     """Prepare tools for Responses API including RAG and MCP tools.
 
@@ -202,7 +203,7 @@ async def prepare_responses_params(  # pylint: disable=too-many-arguments,too-ma
     query_request: QueryRequest,
     user_conversation: Optional[UserConversation],
     token: str,
-    mcp_headers: Optional[dict[str, dict[str, str]]] = None,
+    mcp_headers: Optional[McpHeaders] = None,
     stream: bool = False,
     store: bool = True,
 ) -> ResponsesApiParams:
@@ -315,7 +316,7 @@ def get_rag_tools(vector_store_ids: list[str]) -> Optional[list[dict[str, Any]]]
 def get_mcp_tools(
     mcp_servers: list[ModelContextProtocolServer],
     token: str | None = None,
-    mcp_headers: dict[str, dict[str, str]] | None = None,
+    mcp_headers: Optional[McpHeaders] = None,
 ) -> list[dict[str, Any]]:
     """Convert MCP servers to tools format for Responses API.