mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-03 01:48:05 +00:00
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>
This commit is contained in:
parent
0128effbf7
commit
fe91d331ef
5 changed files with 60 additions and 171 deletions
|
|
@ -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
|
```python
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
...,
|
...,
|
||||||
tools=["mcp::deepwiki"],
|
tools=[
|
||||||
extra_headers={
|
{
|
||||||
"X-LlamaStack-Provider-Data": json.dumps(
|
"type": "mcp",
|
||||||
{
|
"server_url": "https://mcp.deepwiki.com/sse",
|
||||||
"mcp_headers": {
|
"server_label": "mcp::deepwiki",
|
||||||
"http://mcp.deepwiki.com/sse": {
|
"authorization": "<your_access_token>", # OAuth token (without "Bearer " prefix)
|
||||||
"Authorization": "Bearer <your_access_token>",
|
}
|
||||||
},
|
],
|
||||||
},
|
|
||||||
}
|
|
||||||
),
|
|
||||||
},
|
|
||||||
)
|
)
|
||||||
agent.create_turn(...)
|
agent.create_turn(...)
|
||||||
```
|
```
|
||||||
|
|
|
||||||
|
|
@ -48,16 +48,10 @@ class ModelContextProtocolToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime
|
||||||
if mcp_endpoint is None:
|
if mcp_endpoint is None:
|
||||||
raise ValueError("mcp_endpoint is required")
|
raise ValueError("mcp_endpoint is required")
|
||||||
|
|
||||||
# Phase 1: Support both old header-based auth AND new authorization parameter
|
# Get other headers from provider data (but NOT authorization)
|
||||||
# Get headers and auth from provider data (old approach)
|
provider_headers = await self.get_headers_from_request(mcp_endpoint.uri)
|
||||||
provider_headers, provider_auth = await self.get_headers_from_request(mcp_endpoint.uri)
|
|
||||||
|
|
||||||
# New authorization parameter takes precedence over provider data
|
return await list_mcp_tools(endpoint=mcp_endpoint.uri, headers=provider_headers, authorization=authorization)
|
||||||
final_authorization = authorization or provider_auth
|
|
||||||
|
|
||||||
return await list_mcp_tools(
|
|
||||||
endpoint=mcp_endpoint.uri, headers=provider_headers, authorization=final_authorization
|
|
||||||
)
|
|
||||||
|
|
||||||
async def invoke_tool(
|
async def invoke_tool(
|
||||||
self, tool_name: str, kwargs: dict[str, Any], authorization: str | None = None
|
self, tool_name: str, kwargs: dict[str, Any], authorization: str | None = None
|
||||||
|
|
@ -69,39 +63,38 @@ class ModelContextProtocolToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime
|
||||||
if urlparse(endpoint).scheme not in ("http", "https"):
|
if urlparse(endpoint).scheme not in ("http", "https"):
|
||||||
raise ValueError(f"Endpoint {endpoint} is not a valid HTTP(S) URL")
|
raise ValueError(f"Endpoint {endpoint} is not a valid HTTP(S) URL")
|
||||||
|
|
||||||
# Phase 1: Support both old header-based auth AND new authorization parameter
|
# Get other headers from provider data (but NOT authorization)
|
||||||
# Get headers and auth from provider data (old approach)
|
provider_headers = await self.get_headers_from_request(endpoint)
|
||||||
provider_headers, provider_auth = await self.get_headers_from_request(endpoint)
|
|
||||||
|
|
||||||
# New authorization parameter takes precedence over provider data
|
|
||||||
final_authorization = authorization or provider_auth
|
|
||||||
|
|
||||||
return await invoke_mcp_tool(
|
return await invoke_mcp_tool(
|
||||||
endpoint=endpoint,
|
endpoint=endpoint,
|
||||||
tool_name=tool_name,
|
tool_name=tool_name,
|
||||||
kwargs=kwargs,
|
kwargs=kwargs,
|
||||||
headers=provider_headers,
|
headers=provider_headers,
|
||||||
authorization=final_authorization,
|
authorization=authorization,
|
||||||
)
|
)
|
||||||
|
|
||||||
async def get_headers_from_request(self, mcp_endpoint_uri: str) -> tuple[dict[str, str], str | None]:
|
async def get_headers_from_request(self, mcp_endpoint_uri: str) -> dict[str, str]:
|
||||||
"""
|
"""
|
||||||
Extract headers and authorization from request provider data (Phase 1 backward compatibility).
|
Extract headers from request provider data, excluding authorization.
|
||||||
|
|
||||||
Phase 1: Temporarily allows Authorization to be passed via mcp_headers for backward compatibility.
|
Authorization must be provided via the dedicated authorization parameter.
|
||||||
Phase 2: Will enforce that Authorization should use the dedicated authorization parameter instead.
|
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:
|
Returns:
|
||||||
Tuple of (headers_dict, authorization_token)
|
dict[str, str]: Headers dictionary (without Authorization)
|
||||||
- headers_dict: All headers except Authorization
|
|
||||||
- authorization_token: Token from Authorization header (with "Bearer " prefix removed), or None
|
Raises:
|
||||||
|
ValueError: If Authorization header is found in mcp_headers
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def canonicalize_uri(uri: str) -> str:
|
def canonicalize_uri(uri: str) -> str:
|
||||||
return f"{urlparse(uri).netloc or ''}/{urlparse(uri).path or ''}"
|
return f"{urlparse(uri).netloc or ''}/{urlparse(uri).path or ''}"
|
||||||
|
|
||||||
headers = {}
|
headers = {}
|
||||||
authorization = None
|
|
||||||
|
|
||||||
provider_data = self.get_request_provider_data()
|
provider_data = self.get_request_provider_data()
|
||||||
if provider_data and hasattr(provider_data, "mcp_headers") and provider_data.mcp_headers:
|
if provider_data and hasattr(provider_data, "mcp_headers") and provider_data.mcp_headers:
|
||||||
|
|
@ -109,17 +102,14 @@ class ModelContextProtocolToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime
|
||||||
if canonicalize_uri(uri) != canonicalize_uri(mcp_endpoint_uri):
|
if canonicalize_uri(uri) != canonicalize_uri(mcp_endpoint_uri):
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Phase 1: Extract Authorization from mcp_headers for backward compatibility
|
# Reject Authorization in mcp_headers - must use authorization parameter
|
||||||
# (Phase 2 will reject this and require the dedicated authorization parameter)
|
|
||||||
for key in values.keys():
|
for key in values.keys():
|
||||||
if key.lower() == "authorization":
|
if key.lower() == "authorization":
|
||||||
# Extract authorization token and strip "Bearer " prefix if present
|
raise ValueError(
|
||||||
auth_value = values[key]
|
"Authorization cannot be provided via mcp_headers in provider_data. "
|
||||||
if auth_value.startswith("Bearer "):
|
"Please use the dedicated 'authorization' parameter instead. "
|
||||||
authorization = auth_value[7:] # Remove "Bearer " prefix
|
"Example: tool_runtime.invoke_tool(..., authorization='your-token')"
|
||||||
else:
|
)
|
||||||
authorization = auth_value
|
headers[key] = values[key]
|
||||||
else:
|
|
||||||
headers[key] = values[key]
|
|
||||||
|
|
||||||
return headers, authorization
|
return headers
|
||||||
|
|
|
||||||
|
|
@ -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.
|
Tests that tools pass through correctly to various LLM providers.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import json
|
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from llama_stack.core.library_client import LlamaStackAsLibraryClient
|
from llama_stack.core.library_client import LlamaStackAsLibraryClient
|
||||||
|
|
@ -193,22 +191,11 @@ class TestMCPToolsInChatCompletion:
|
||||||
mcp_endpoint=dict(uri=uri),
|
mcp_endpoint=dict(uri=uri),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Use old header-based approach for Phase 1 (backward compatibility)
|
# Use the dedicated authorization parameter
|
||||||
provider_data = {
|
|
||||||
"mcp_headers": {
|
|
||||||
uri: {
|
|
||||||
"Authorization": f"Bearer {AUTH_TOKEN}",
|
|
||||||
},
|
|
||||||
},
|
|
||||||
}
|
|
||||||
auth_headers = {
|
|
||||||
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
|
|
||||||
}
|
|
||||||
|
|
||||||
# Get the tools from MCP
|
# Get the tools from MCP
|
||||||
tools_response = llama_stack_client.tool_runtime.list_tools(
|
tools_response = llama_stack_client.tool_runtime.list_tools(
|
||||||
tool_group_id=test_toolgroup_id,
|
tool_group_id=test_toolgroup_id,
|
||||||
extra_headers=auth_headers,
|
authorization=AUTH_TOKEN,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Convert to OpenAI format for inference
|
# Convert to OpenAI format for inference
|
||||||
|
|
|
||||||
|
|
@ -4,8 +4,6 @@
|
||||||
# This source code is licensed under the terms described in the LICENSE file in
|
# This source code is licensed under the terms described in the LICENSE file in
|
||||||
# the root directory of this source tree.
|
# the root directory of this source tree.
|
||||||
|
|
||||||
import json
|
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from llama_stack_client.lib.agents.agent import Agent
|
from llama_stack_client.lib.agents.agent import Agent
|
||||||
from llama_stack_client.lib.agents.turn_events import StepCompleted, StepProgress, ToolCallIssuedDelta
|
from llama_stack_client.lib.agents.turn_events import StepCompleted, StepProgress, ToolCallIssuedDelta
|
||||||
|
|
@ -37,32 +35,20 @@ def test_mcp_invocation(llama_stack_client, text_model_id, mcp_server):
|
||||||
mcp_endpoint=dict(uri=uri),
|
mcp_endpoint=dict(uri=uri),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Use old header-based approach for Phase 1 (backward compatibility)
|
# Use the dedicated authorization parameter (no more provider_data headers)
|
||||||
provider_data = {
|
# This tests direct tool_runtime.invoke_tool API calls
|
||||||
"mcp_headers": {
|
tools_list = llama_stack_client.tool_runtime.list_tools(
|
||||||
uri: {
|
tool_group_id=test_toolgroup_id,
|
||||||
"Authorization": f"Bearer {AUTH_TOKEN}",
|
authorization=AUTH_TOKEN, # Use dedicated authorization parameter
|
||||||
},
|
|
||||||
},
|
|
||||||
}
|
|
||||||
auth_headers = {
|
|
||||||
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
|
|
||||||
}
|
|
||||||
|
|
||||||
with pytest.raises(Exception, match="Unauthorized"):
|
|
||||||
llama_stack_client.tools.list(toolgroup_id=test_toolgroup_id)
|
|
||||||
|
|
||||||
tools_list = llama_stack_client.tools.list(
|
|
||||||
toolgroup_id=test_toolgroup_id,
|
|
||||||
extra_headers=auth_headers, # Use old header-based approach
|
|
||||||
)
|
)
|
||||||
assert len(tools_list) == 2
|
assert len(tools_list) == 2
|
||||||
assert {t.name for t in tools_list} == {"greet_everyone", "get_boiling_point"}
|
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(
|
response = llama_stack_client.tool_runtime.invoke_tool(
|
||||||
tool_name="greet_everyone",
|
tool_name="greet_everyone",
|
||||||
kwargs=dict(url="https://www.google.com"),
|
kwargs=dict(url="https://www.google.com"),
|
||||||
extra_headers=auth_headers, # Use old header-based approach
|
authorization=AUTH_TOKEN, # Use dedicated authorization parameter
|
||||||
)
|
)
|
||||||
content = response.content
|
content = response.content
|
||||||
assert len(content) == 1
|
assert len(content) == 1
|
||||||
|
|
|
||||||
|
|
@ -8,8 +8,6 @@
|
||||||
Tests $ref, $defs, and other JSON Schema features through MCP integration.
|
Tests $ref, $defs, and other JSON Schema features through MCP integration.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import json
|
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from llama_stack.core.library_client import LlamaStackAsLibraryClient
|
from llama_stack.core.library_client import LlamaStackAsLibraryClient
|
||||||
|
|
@ -122,22 +120,11 @@ class TestMCPSchemaPreservation:
|
||||||
mcp_endpoint=dict(uri=uri),
|
mcp_endpoint=dict(uri=uri),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Use old header-based approach for Phase 1 (backward compatibility)
|
# Use the dedicated authorization parameter
|
||||||
provider_data = {
|
|
||||||
"mcp_headers": {
|
|
||||||
uri: {
|
|
||||||
"Authorization": f"Bearer {AUTH_TOKEN}",
|
|
||||||
},
|
|
||||||
},
|
|
||||||
}
|
|
||||||
auth_headers = {
|
|
||||||
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
|
|
||||||
}
|
|
||||||
|
|
||||||
# List runtime tools
|
# List runtime tools
|
||||||
response = llama_stack_client.tool_runtime.list_tools(
|
response = llama_stack_client.tool_runtime.list_tools(
|
||||||
tool_group_id=test_toolgroup_id,
|
tool_group_id=test_toolgroup_id,
|
||||||
extra_headers=auth_headers,
|
authorization=AUTH_TOKEN,
|
||||||
)
|
)
|
||||||
|
|
||||||
tools = response
|
tools = response
|
||||||
|
|
@ -173,22 +160,11 @@ class TestMCPSchemaPreservation:
|
||||||
mcp_endpoint=dict(uri=uri),
|
mcp_endpoint=dict(uri=uri),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Use old header-based approach for Phase 1 (backward compatibility)
|
# Use the dedicated authorization parameter
|
||||||
provider_data = {
|
|
||||||
"mcp_headers": {
|
|
||||||
uri: {
|
|
||||||
"Authorization": f"Bearer {AUTH_TOKEN}",
|
|
||||||
},
|
|
||||||
},
|
|
||||||
}
|
|
||||||
auth_headers = {
|
|
||||||
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
|
|
||||||
}
|
|
||||||
|
|
||||||
# List tools
|
# List tools
|
||||||
response = llama_stack_client.tool_runtime.list_tools(
|
response = llama_stack_client.tool_runtime.list_tools(
|
||||||
tool_group_id=test_toolgroup_id,
|
tool_group_id=test_toolgroup_id,
|
||||||
extra_headers=auth_headers,
|
authorization=AUTH_TOKEN,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Find book_flight tool (which should have $ref/$defs)
|
# Find book_flight tool (which should have $ref/$defs)
|
||||||
|
|
@ -230,21 +206,10 @@ class TestMCPSchemaPreservation:
|
||||||
mcp_endpoint=dict(uri=uri),
|
mcp_endpoint=dict(uri=uri),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Use old header-based approach for Phase 1 (backward compatibility)
|
# Use the dedicated authorization parameter
|
||||||
provider_data = {
|
|
||||||
"mcp_headers": {
|
|
||||||
uri: {
|
|
||||||
"Authorization": f"Bearer {AUTH_TOKEN}",
|
|
||||||
},
|
|
||||||
},
|
|
||||||
}
|
|
||||||
auth_headers = {
|
|
||||||
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
|
|
||||||
}
|
|
||||||
|
|
||||||
response = llama_stack_client.tool_runtime.list_tools(
|
response = llama_stack_client.tool_runtime.list_tools(
|
||||||
tool_group_id=test_toolgroup_id,
|
tool_group_id=test_toolgroup_id,
|
||||||
extra_headers=auth_headers,
|
authorization=AUTH_TOKEN,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Find get_weather tool
|
# Find get_weather tool
|
||||||
|
|
@ -284,22 +249,10 @@ class TestMCPToolInvocation:
|
||||||
mcp_endpoint=dict(uri=uri),
|
mcp_endpoint=dict(uri=uri),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Use old header-based approach for Phase 1 (backward compatibility)
|
# Use the dedicated authorization parameter
|
||||||
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
|
|
||||||
llama_stack_client.tool_runtime.list_tools(
|
llama_stack_client.tool_runtime.list_tools(
|
||||||
tool_group_id=test_toolgroup_id,
|
tool_group_id=test_toolgroup_id,
|
||||||
extra_headers=auth_headers,
|
authorization=AUTH_TOKEN,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Invoke tool with complex nested data
|
# Invoke tool with complex nested data
|
||||||
|
|
@ -311,7 +264,7 @@ class TestMCPToolInvocation:
|
||||||
"shipping": {"address": {"street": "123 Main St", "city": "San Francisco", "zipcode": "94102"}},
|
"shipping": {"address": {"street": "123 Main St", "city": "San Francisco", "zipcode": "94102"}},
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
extra_headers=auth_headers,
|
authorization=AUTH_TOKEN,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Should succeed without schema validation errors
|
# Should succeed without schema validation errors
|
||||||
|
|
@ -337,29 +290,17 @@ class TestMCPToolInvocation:
|
||||||
mcp_endpoint=dict(uri=uri),
|
mcp_endpoint=dict(uri=uri),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Use old header-based approach for Phase 1 (backward compatibility)
|
# Use the dedicated authorization parameter
|
||||||
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
|
|
||||||
llama_stack_client.tool_runtime.list_tools(
|
llama_stack_client.tool_runtime.list_tools(
|
||||||
tool_group_id=test_toolgroup_id,
|
tool_group_id=test_toolgroup_id,
|
||||||
extra_headers=auth_headers,
|
authorization=AUTH_TOKEN,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Test with email format
|
# Test with email format
|
||||||
result_email = llama_stack_client.tool_runtime.invoke_tool(
|
result_email = llama_stack_client.tool_runtime.invoke_tool(
|
||||||
tool_name="flexible_contact",
|
tool_name="flexible_contact",
|
||||||
kwargs={"contact_info": "user@example.com"},
|
kwargs={"contact_info": "user@example.com"},
|
||||||
extra_headers=auth_headers,
|
authorization=AUTH_TOKEN,
|
||||||
)
|
)
|
||||||
|
|
||||||
assert result_email.error_message is None
|
assert result_email.error_message is None
|
||||||
|
|
@ -368,7 +309,7 @@ class TestMCPToolInvocation:
|
||||||
result_phone = llama_stack_client.tool_runtime.invoke_tool(
|
result_phone = llama_stack_client.tool_runtime.invoke_tool(
|
||||||
tool_name="flexible_contact",
|
tool_name="flexible_contact",
|
||||||
kwargs={"contact_info": "+15551234567"},
|
kwargs={"contact_info": "+15551234567"},
|
||||||
extra_headers=auth_headers,
|
authorization=AUTH_TOKEN,
|
||||||
)
|
)
|
||||||
|
|
||||||
assert result_phone.error_message is None
|
assert result_phone.error_message is None
|
||||||
|
|
@ -400,21 +341,10 @@ class TestAgentWithMCPTools:
|
||||||
mcp_endpoint=dict(uri=uri),
|
mcp_endpoint=dict(uri=uri),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Use old header-based approach for Phase 1 (backward compatibility)
|
# Use the dedicated authorization parameter
|
||||||
provider_data = {
|
tools_list = llama_stack_client.tool_runtime.list_tools(
|
||||||
"mcp_headers": {
|
tool_group_id=test_toolgroup_id,
|
||||||
uri: {
|
authorization=AUTH_TOKEN,
|
||||||
"Authorization": f"Bearer {AUTH_TOKEN}",
|
|
||||||
},
|
|
||||||
},
|
|
||||||
}
|
|
||||||
auth_headers = {
|
|
||||||
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
|
|
||||||
}
|
|
||||||
|
|
||||||
tools_list = llama_stack_client.tools.list(
|
|
||||||
toolgroup_id=test_toolgroup_id,
|
|
||||||
extra_headers=auth_headers,
|
|
||||||
)
|
)
|
||||||
tool_defs = [
|
tool_defs = [
|
||||||
{
|
{
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue