fix: MCP authorization parameter implementation (#4052)

# What does this PR do?
Adding a user-facing `authorization ` parameter to MCP tool definitions
that allows users to explicitly configure credentials per MCP server,
addressing GitHub Issue #4034 in a secure manner.


## Test Plan
tests/integration/responses/test_mcp_authentication.py

---------

Co-authored-by: Omar Abdelwahab <omara@fb.com>
Co-authored-by: Ashwin Bharambe <ashwin.bharambe@gmail.com>
This commit is contained in:
Omar Abdelwahab 2025-11-14 08:54:42 -08:00 committed by GitHub
parent dc49ad3f89
commit eb545034ab
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
34 changed files with 5205 additions and 62 deletions

View file

@ -37,6 +37,7 @@ def test_mcp_invocation(llama_stack_client, text_model_id, mcp_server):
mcp_endpoint=dict(uri=uri),
)
# Use old header-based approach for Phase 1 (backward compatibility)
provider_data = {
"mcp_headers": {
uri: {
@ -53,7 +54,7 @@ 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,
extra_headers=auth_headers, # Use old header-based approach
)
assert len(tools_list) == 2
assert {t.name for t in tools_list} == {"greet_everyone", "get_boiling_point"}
@ -61,7 +62,7 @@ def test_mcp_invocation(llama_stack_client, text_model_id, mcp_server):
response = llama_stack_client.tool_runtime.invoke_tool(
tool_name="greet_everyone",
kwargs=dict(url="https://www.google.com"),
extra_headers=auth_headers,
extra_headers=auth_headers, # Use old header-based approach
)
content = response.content
assert len(content) == 1
@ -76,9 +77,7 @@ def test_mcp_invocation(llama_stack_client, text_model_id, mcp_server):
"server_label": test_toolgroup_id,
"require_approval": "never",
"allowed_tools": [tool.name for tool in tools_list],
"headers": {
"Authorization": f"Bearer {AUTH_TOKEN}",
},
"authorization": AUTH_TOKEN,
}
]
agent = Agent(
@ -104,7 +103,6 @@ def test_mcp_invocation(llama_stack_client, text_model_id, mcp_server):
}
],
stream=True,
extra_headers=auth_headers,
)
)
events = [chunk.event for chunk in chunks]

View file

@ -4,8 +4,7 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
"""
Integration tests for MCP tools with complex JSON Schema support.
"""Integration tests for MCP tools with complex JSON Schema support.
Tests $ref, $defs, and other JSON Schema features through MCP integration.
"""
@ -123,7 +122,14 @@ class TestMCPSchemaPreservation:
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
# Use old header-based approach for Phase 1 (backward compatibility)
provider_data = {
"mcp_headers": {
uri: {
"Authorization": f"Bearer {AUTH_TOKEN}",
},
},
}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
@ -166,7 +172,15 @@ class TestMCPSchemaPreservation:
provider_id="model-context-protocol",
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
# Use old header-based approach for Phase 1 (backward compatibility)
provider_data = {
"mcp_headers": {
uri: {
"Authorization": f"Bearer {AUTH_TOKEN}",
},
},
}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
@ -216,7 +230,14 @@ class TestMCPSchemaPreservation:
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
# Use old header-based approach for Phase 1 (backward compatibility)
provider_data = {
"mcp_headers": {
uri: {
"Authorization": f"Bearer {AUTH_TOKEN}",
},
},
}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
@ -263,7 +284,14 @@ class TestMCPToolInvocation:
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
# Use old header-based approach for Phase 1 (backward compatibility)
provider_data = {
"mcp_headers": {
uri: {
"Authorization": f"Bearer {AUTH_TOKEN}",
},
},
}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
@ -309,7 +337,14 @@ class TestMCPToolInvocation:
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
# Use old header-based approach for Phase 1 (backward compatibility)
provider_data = {
"mcp_headers": {
uri: {
"Authorization": f"Bearer {AUTH_TOKEN}",
},
},
}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
@ -365,7 +400,14 @@ class TestAgentWithMCPTools:
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
# Use old header-based approach for Phase 1 (backward compatibility)
provider_data = {
"mcp_headers": {
uri: {
"Authorization": f"Bearer {AUTH_TOKEN}",
},
},
}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
@ -381,6 +423,7 @@ class TestAgentWithMCPTools:
"server_label": test_toolgroup_id,
"require_approval": "never",
"allowed_tools": [tool.name for tool in tools_list],
"authorization": AUTH_TOKEN,
}
]
@ -389,7 +432,6 @@ class TestAgentWithMCPTools:
model=text_model_id,
instructions="You are a helpful assistant that can process orders and book flights.",
tools=tool_defs,
extra_headers=auth_headers,
)
session_id = agent.create_session("test-session-complex")
@ -411,7 +453,6 @@ class TestAgentWithMCPTools:
}
],
stream=True,
extra_headers=auth_headers,
)
)