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 <omara@fb.com>
Co-authored-by: Ashwin Bharambe <ashwin.bharambe@gmail.com>
(cherry picked from commit fe91d331ef)

# 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
This commit is contained in:
Omar Abdelwahab 2025-11-17 12:16:35 -08:00 committed by Mergify
parent 49a290e53e
commit 9d54a854d9
5 changed files with 178 additions and 28 deletions

View file

@ -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 <your_access_token>",
},
},
}
),
},
tools=[
{
"type": "mcp",
"server_url": "https://mcp.deepwiki.com/sse",
"server_label": "mcp::deepwiki",
"authorization": "<your_access_token>", # OAuth token (without "Bearer " prefix)
}
],
)
agent.create_turn(...)
```

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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 = [
{