From 9d54a854d9e36f07e9155f936b2544e68831255d Mon Sep 17 00:00:00 2001 From: Omar Abdelwahab Date: Mon, 17 Nov 2025 12:16:35 -0800 Subject: [PATCH] fix: Remove authorization from provider data (#4161) # What does this PR do? - Remove backward compatibility for authorization in mcp_headers - Enforce authorization must use dedicated parameter - Add validation error if Authorization found in provider_data headers - Update test_mcp.py to use authorization parameter - Update test_mcp_json_schema.py to use authorization parameter - Update test_tools_with_schemas.py to use authorization parameter - Update documentation to show the change in the authorization approach Breaking Change: - Authorization can no longer be passed via mcp_headers in provider_data - Users must use the dedicated 'authorization' parameter instead - Clear error message guides users to the new approach" ## Test Plan CI --------- Co-authored-by: Omar Abdelwahab Co-authored-by: Ashwin Bharambe (cherry picked from commit fe91d331efb98193ebed6208a6f9ecaa52cbceda) # Conflicts: # src/llama_stack/providers/remote/tool_runtime/model_context_protocol/model_context_protocol.py # tests/integration/inference/test_tools_with_schemas.py # tests/integration/tool_runtime/test_mcp.py # tests/integration/tool_runtime/test_mcp_json_schema.py --- docs/docs/building_applications/tools.mdx | 22 ++-- .../model_context_protocol.py | 115 ++++++++++++++++++ .../inference/test_tools_with_schemas.py | 8 +- tests/integration/tool_runtime/test_mcp.py | 15 ++- .../tool_runtime/test_mcp_json_schema.py | 46 +++++-- 5 files changed, 178 insertions(+), 28 deletions(-) create mode 100644 src/llama_stack/providers/remote/tool_runtime/model_context_protocol/model_context_protocol.py diff --git a/docs/docs/building_applications/tools.mdx b/docs/docs/building_applications/tools.mdx index 3b78ec57b..f7b913fef 100644 --- a/docs/docs/building_applications/tools.mdx +++ b/docs/docs/building_applications/tools.mdx @@ -104,23 +104,19 @@ client.toolgroups.register( ) ``` -Note that most of the more useful MCP servers need you to authenticate with them. Many of them use OAuth2.0 for authentication. You can provide authorization headers to send to the MCP server using the "Provider Data" abstraction provided by Llama Stack. When making an agent call, +Note that most of the more useful MCP servers need you to authenticate with them. Many of them use OAuth2.0 for authentication. You can provide the authorization token when creating the Agent: ```python agent = Agent( ..., - tools=["mcp::deepwiki"], - extra_headers={ - "X-LlamaStack-Provider-Data": json.dumps( - { - "mcp_headers": { - "http://mcp.deepwiki.com/sse": { - "Authorization": "Bearer ", - }, - }, - } - ), - }, + tools=[ + { + "type": "mcp", + "server_url": "https://mcp.deepwiki.com/sse", + "server_label": "mcp::deepwiki", + "authorization": "", # OAuth token (without "Bearer " prefix) + } + ], ) agent.create_turn(...) ``` diff --git a/src/llama_stack/providers/remote/tool_runtime/model_context_protocol/model_context_protocol.py b/src/llama_stack/providers/remote/tool_runtime/model_context_protocol/model_context_protocol.py new file mode 100644 index 000000000..97b044dbf --- /dev/null +++ b/src/llama_stack/providers/remote/tool_runtime/model_context_protocol/model_context_protocol.py @@ -0,0 +1,115 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +from typing import Any +from urllib.parse import urlparse + +from llama_stack.core.request_headers import NeedsRequestProviderData +from llama_stack.log import get_logger +from llama_stack.providers.utils.tools.mcp import invoke_mcp_tool, list_mcp_tools +from llama_stack_api import ( + URL, + Api, + ListToolDefsResponse, + ToolGroup, + ToolGroupsProtocolPrivate, + ToolInvocationResult, + ToolRuntime, +) + +from .config import MCPProviderConfig + +logger = get_logger(__name__, category="tools") + + +class ModelContextProtocolToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, NeedsRequestProviderData): + def __init__(self, config: MCPProviderConfig, _deps: dict[Api, Any]): + self.config = config + + async def initialize(self): + pass + + async def register_toolgroup(self, toolgroup: ToolGroup) -> None: + pass + + async def unregister_toolgroup(self, toolgroup_id: str) -> None: + return + + async def list_runtime_tools( + self, + tool_group_id: str | None = None, + mcp_endpoint: URL | None = None, + authorization: str | None = None, + ) -> ListToolDefsResponse: + # this endpoint should be retrieved by getting the tool group right? + if mcp_endpoint is None: + raise ValueError("mcp_endpoint is required") + + # Get other headers from provider data (but NOT authorization) + provider_headers = await self.get_headers_from_request(mcp_endpoint.uri) + + return await list_mcp_tools(endpoint=mcp_endpoint.uri, headers=provider_headers, authorization=authorization) + + async def invoke_tool( + self, tool_name: str, kwargs: dict[str, Any], authorization: str | None = None + ) -> ToolInvocationResult: + tool = await self.tool_store.get_tool(tool_name) + if tool.metadata is None or tool.metadata.get("endpoint") is None: + raise ValueError(f"Tool {tool_name} does not have metadata") + endpoint = tool.metadata.get("endpoint") + if urlparse(endpoint).scheme not in ("http", "https"): + raise ValueError(f"Endpoint {endpoint} is not a valid HTTP(S) URL") + + # Get other headers from provider data (but NOT authorization) + provider_headers = await self.get_headers_from_request(endpoint) + + return await invoke_mcp_tool( + endpoint=endpoint, + tool_name=tool_name, + kwargs=kwargs, + headers=provider_headers, + authorization=authorization, + ) + + async def get_headers_from_request(self, mcp_endpoint_uri: str) -> dict[str, str]: + """ + Extract headers from request provider data, excluding authorization. + + Authorization must be provided via the dedicated authorization parameter. + If Authorization is found in mcp_headers, raise an error to guide users to the correct approach. + + Args: + mcp_endpoint_uri: The MCP endpoint URI to match against provider data + + Returns: + dict[str, str]: Headers dictionary (without Authorization) + + Raises: + ValueError: If Authorization header is found in mcp_headers + """ + + def canonicalize_uri(uri: str) -> str: + return f"{urlparse(uri).netloc or ''}/{urlparse(uri).path or ''}" + + headers = {} + + provider_data = self.get_request_provider_data() + if provider_data and hasattr(provider_data, "mcp_headers") and provider_data.mcp_headers: + for uri, values in provider_data.mcp_headers.items(): + if canonicalize_uri(uri) != canonicalize_uri(mcp_endpoint_uri): + continue + + # Reject Authorization in mcp_headers - must use authorization parameter + for key in values.keys(): + if key.lower() == "authorization": + raise ValueError( + "Authorization cannot be provided via mcp_headers in provider_data. " + "Please use the dedicated 'authorization' parameter instead. " + "Example: tool_runtime.invoke_tool(..., authorization='your-token')" + ) + headers[key] = values[key] + + return headers diff --git a/tests/integration/inference/test_tools_with_schemas.py b/tests/integration/inference/test_tools_with_schemas.py index b144a5196..96a3a6a01 100644 --- a/tests/integration/inference/test_tools_with_schemas.py +++ b/tests/integration/inference/test_tools_with_schemas.py @@ -9,8 +9,6 @@ Integration tests for inference/chat completion with JSON Schema-based tools. Tests that tools pass through correctly to various LLM providers. """ -import json - import pytest from llama_stack import LlamaStackAsLibraryClient @@ -193,15 +191,19 @@ class TestMCPToolsInChatCompletion: mcp_endpoint=dict(uri=uri), ) +<<<<<<< HEAD provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}} auth_headers = { "X-LlamaStack-Provider-Data": json.dumps(provider_data), } +======= + # Use the dedicated authorization parameter +>>>>>>> fe91d331 (fix: Remove authorization from provider data (#4161)) # Get the tools from MCP tools_response = llama_stack_client.tool_runtime.list_tools( tool_group_id=test_toolgroup_id, - extra_headers=auth_headers, + authorization=AUTH_TOKEN, ) # Convert to OpenAI format for inference diff --git a/tests/integration/tool_runtime/test_mcp.py b/tests/integration/tool_runtime/test_mcp.py index 59f558d2c..64f5546ac 100644 --- a/tests/integration/tool_runtime/test_mcp.py +++ b/tests/integration/tool_runtime/test_mcp.py @@ -4,8 +4,6 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -import json - import pytest from llama_stack_client.lib.agents.agent import Agent from llama_stack_client.lib.agents.turn_events import StepCompleted, StepProgress, ToolCallIssuedDelta @@ -42,6 +40,7 @@ def test_mcp_invocation(llama_stack_client, text_model_id, mcp_server): mcp_endpoint=dict(uri=uri), ) +<<<<<<< HEAD provider_data = { "mcp_headers": { uri: { @@ -59,14 +58,26 @@ def test_mcp_invocation(llama_stack_client, text_model_id, mcp_server): tools_list = llama_stack_client.tools.list( toolgroup_id=test_toolgroup_id, extra_headers=auth_headers, +======= + # Use the dedicated authorization parameter (no more provider_data headers) + # This tests direct tool_runtime.invoke_tool API calls + tools_list = llama_stack_client.tool_runtime.list_tools( + tool_group_id=test_toolgroup_id, + authorization=AUTH_TOKEN, # Use dedicated authorization parameter +>>>>>>> fe91d331 (fix: Remove authorization from provider data (#4161)) ) assert len(tools_list) == 2 assert {t.name for t in tools_list} == {"greet_everyone", "get_boiling_point"} + # Invoke tool with authorization parameter response = llama_stack_client.tool_runtime.invoke_tool( tool_name="greet_everyone", kwargs=dict(url="https://www.google.com"), +<<<<<<< HEAD extra_headers=auth_headers, +======= + authorization=AUTH_TOKEN, # Use dedicated authorization parameter +>>>>>>> fe91d331 (fix: Remove authorization from provider data (#4161)) ) content = response.content assert len(content) == 1 diff --git a/tests/integration/tool_runtime/test_mcp_json_schema.py b/tests/integration/tool_runtime/test_mcp_json_schema.py index 240ec403a..744b63722 100644 --- a/tests/integration/tool_runtime/test_mcp_json_schema.py +++ b/tests/integration/tool_runtime/test_mcp_json_schema.py @@ -9,8 +9,6 @@ Integration tests for MCP tools with complex JSON Schema support. Tests $ref, $defs, and other JSON Schema features through MCP integration. """ -import json - import pytest from llama_stack import LlamaStackAsLibraryClient @@ -123,15 +121,19 @@ class TestMCPSchemaPreservation: mcp_endpoint=dict(uri=uri), ) +<<<<<<< HEAD provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}} auth_headers = { "X-LlamaStack-Provider-Data": json.dumps(provider_data), } +======= + # Use the dedicated authorization parameter +>>>>>>> fe91d331 (fix: Remove authorization from provider data (#4161)) # List runtime tools response = llama_stack_client.tool_runtime.list_tools( tool_group_id=test_toolgroup_id, - extra_headers=auth_headers, + authorization=AUTH_TOKEN, ) tools = response @@ -166,15 +168,20 @@ class TestMCPSchemaPreservation: provider_id="model-context-protocol", mcp_endpoint=dict(uri=uri), ) +<<<<<<< HEAD provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}} auth_headers = { "X-LlamaStack-Provider-Data": json.dumps(provider_data), } +======= + + # Use the dedicated authorization parameter +>>>>>>> fe91d331 (fix: Remove authorization from provider data (#4161)) # List tools response = llama_stack_client.tool_runtime.list_tools( tool_group_id=test_toolgroup_id, - extra_headers=auth_headers, + authorization=AUTH_TOKEN, ) # Find book_flight tool (which should have $ref/$defs) @@ -216,14 +223,18 @@ class TestMCPSchemaPreservation: mcp_endpoint=dict(uri=uri), ) +<<<<<<< HEAD provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}} auth_headers = { "X-LlamaStack-Provider-Data": json.dumps(provider_data), } +======= + # Use the dedicated authorization parameter +>>>>>>> fe91d331 (fix: Remove authorization from provider data (#4161)) response = llama_stack_client.tool_runtime.list_tools( tool_group_id=test_toolgroup_id, - extra_headers=auth_headers, + authorization=AUTH_TOKEN, ) # Find get_weather tool @@ -263,15 +274,19 @@ class TestMCPToolInvocation: mcp_endpoint=dict(uri=uri), ) +<<<<<<< HEAD provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}} auth_headers = { "X-LlamaStack-Provider-Data": json.dumps(provider_data), } # List tools to populate the tool index +======= + # Use the dedicated authorization parameter +>>>>>>> fe91d331 (fix: Remove authorization from provider data (#4161)) llama_stack_client.tool_runtime.list_tools( tool_group_id=test_toolgroup_id, - extra_headers=auth_headers, + authorization=AUTH_TOKEN, ) # Invoke tool with complex nested data @@ -283,7 +298,7 @@ class TestMCPToolInvocation: "shipping": {"address": {"street": "123 Main St", "city": "San Francisco", "zipcode": "94102"}}, } }, - extra_headers=auth_headers, + authorization=AUTH_TOKEN, ) # Should succeed without schema validation errors @@ -309,22 +324,26 @@ class TestMCPToolInvocation: mcp_endpoint=dict(uri=uri), ) +<<<<<<< HEAD provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}} auth_headers = { "X-LlamaStack-Provider-Data": json.dumps(provider_data), } # List tools to populate the tool index +======= + # Use the dedicated authorization parameter +>>>>>>> fe91d331 (fix: Remove authorization from provider data (#4161)) llama_stack_client.tool_runtime.list_tools( tool_group_id=test_toolgroup_id, - extra_headers=auth_headers, + authorization=AUTH_TOKEN, ) # Test with email format result_email = llama_stack_client.tool_runtime.invoke_tool( tool_name="flexible_contact", kwargs={"contact_info": "user@example.com"}, - extra_headers=auth_headers, + authorization=AUTH_TOKEN, ) assert result_email.error_message is None @@ -333,7 +352,7 @@ class TestMCPToolInvocation: result_phone = llama_stack_client.tool_runtime.invoke_tool( tool_name="flexible_contact", kwargs={"contact_info": "+15551234567"}, - extra_headers=auth_headers, + authorization=AUTH_TOKEN, ) assert result_phone.error_message is None @@ -365,6 +384,7 @@ class TestAgentWithMCPTools: mcp_endpoint=dict(uri=uri), ) +<<<<<<< HEAD provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}} auth_headers = { "X-LlamaStack-Provider-Data": json.dumps(provider_data), @@ -373,6 +393,12 @@ class TestAgentWithMCPTools: tools_list = llama_stack_client.tools.list( toolgroup_id=test_toolgroup_id, extra_headers=auth_headers, +======= + # Use the dedicated authorization parameter + tools_list = llama_stack_client.tool_runtime.list_tools( + tool_group_id=test_toolgroup_id, + authorization=AUTH_TOKEN, +>>>>>>> fe91d331 (fix: Remove authorization from provider data (#4161)) ) tool_defs = [ {