mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-03 09:53:45 +00:00
feat(tool-runtime): Add authorization parameter with backward compatibility
Implement Phase 1 of MCP auth migration: - Add authorization parameter to list_runtime_tools() and invoke_tool() - Maintain backward compatibility with X-LlamaStack-Provider-Data header - Tests use old header-based auth to avoid client SDK dependency - New parameter takes precedence when both methods provided Phase 2 will migrate tests to new parameter after Stainless SDK release. Related: PR #4052
This commit is contained in:
Omar Abdelwahab
parent
fa2b361f46
commit
8783255bc3
4 changed files with 161 additions and 40 deletions
|
|
@ -4,6 +4,8 @@
|
|||
# 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
|
||||
|
|
@ -35,13 +37,24 @@ def test_mcp_invocation(llama_stack_client, text_model_id, mcp_server):
|
|||
mcp_endpoint=dict(uri=uri),
|
||||
)
|
||||
|
||||
# Authorization now passed as request body parameter (not provider-data)
|
||||
# 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),
|
||||
}
|
||||
|
||||
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,
|
||||
authorization=AUTH_TOKEN, # Pass authorization as parameter
|
||||
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"}
|
||||
|
|
@ -49,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"),
|
||||
authorization=AUTH_TOKEN, # Pass authorization as parameter
|
||||
extra_headers=auth_headers, # Use old header-based approach
|
||||
)
|
||||
content = response.content
|
||||
assert len(content) == 1
|
||||
|
|
|
|||
|
|
@ -4,11 +4,12 @@
|
|||
# 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.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from llama_stack.core.library_client import LlamaStackAsLibraryClient
|
||||
|
|
@ -121,10 +122,22 @@ class TestMCPSchemaPreservation:
|
|||
mcp_endpoint=dict(uri=uri),
|
||||
)
|
||||
|
||||
# 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),
|
||||
}
|
||||
|
||||
# List runtime tools
|
||||
response = llama_stack_client.tool_runtime.list_tools(
|
||||
tool_group_id=test_toolgroup_id,
|
||||
authorization=AUTH_TOKEN,
|
||||
extra_headers=auth_headers,
|
||||
)
|
||||
|
||||
tools = response
|
||||
|
|
@ -160,10 +173,22 @@ class TestMCPSchemaPreservation:
|
|||
mcp_endpoint=dict(uri=uri),
|
||||
)
|
||||
|
||||
# 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),
|
||||
}
|
||||
|
||||
# List tools
|
||||
response = llama_stack_client.tool_runtime.list_tools(
|
||||
tool_group_id=test_toolgroup_id,
|
||||
authorization=AUTH_TOKEN,
|
||||
extra_headers=auth_headers,
|
||||
)
|
||||
|
||||
# Find book_flight tool (which should have $ref/$defs)
|
||||
|
|
@ -205,9 +230,21 @@ class TestMCPSchemaPreservation:
|
|||
mcp_endpoint=dict(uri=uri),
|
||||
)
|
||||
|
||||
# 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),
|
||||
}
|
||||
|
||||
response = llama_stack_client.tool_runtime.list_tools(
|
||||
tool_group_id=test_toolgroup_id,
|
||||
authorization=AUTH_TOKEN,
|
||||
extra_headers=auth_headers,
|
||||
)
|
||||
|
||||
# Find get_weather tool
|
||||
|
|
@ -247,10 +284,22 @@ class TestMCPToolInvocation:
|
|||
mcp_endpoint=dict(uri=uri),
|
||||
)
|
||||
|
||||
# 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),
|
||||
}
|
||||
|
||||
# List tools to populate the tool index
|
||||
llama_stack_client.tool_runtime.list_tools(
|
||||
tool_group_id=test_toolgroup_id,
|
||||
authorization=AUTH_TOKEN,
|
||||
extra_headers=auth_headers,
|
||||
)
|
||||
|
||||
# Invoke tool with complex nested data
|
||||
|
|
@ -262,7 +311,7 @@ class TestMCPToolInvocation:
|
|||
"shipping": {"address": {"street": "123 Main St", "city": "San Francisco", "zipcode": "94102"}},
|
||||
}
|
||||
},
|
||||
authorization=AUTH_TOKEN,
|
||||
extra_headers=auth_headers,
|
||||
)
|
||||
|
||||
# Should succeed without schema validation errors
|
||||
|
|
@ -288,17 +337,29 @@ class TestMCPToolInvocation:
|
|||
mcp_endpoint=dict(uri=uri),
|
||||
)
|
||||
|
||||
# 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),
|
||||
}
|
||||
|
||||
# List tools to populate the tool index
|
||||
llama_stack_client.tool_runtime.list_tools(
|
||||
tool_group_id=test_toolgroup_id,
|
||||
authorization=AUTH_TOKEN,
|
||||
extra_headers=auth_headers,
|
||||
)
|
||||
|
||||
# Test with email format
|
||||
result_email = llama_stack_client.tool_runtime.invoke_tool(
|
||||
tool_name="flexible_contact",
|
||||
kwargs={"contact_info": "user@example.com"},
|
||||
authorization=AUTH_TOKEN,
|
||||
extra_headers=auth_headers,
|
||||
)
|
||||
|
||||
assert result_email.error_message is None
|
||||
|
|
@ -307,7 +368,7 @@ class TestMCPToolInvocation:
|
|||
result_phone = llama_stack_client.tool_runtime.invoke_tool(
|
||||
tool_name="flexible_contact",
|
||||
kwargs={"contact_info": "+15551234567"},
|
||||
authorization=AUTH_TOKEN,
|
||||
extra_headers=auth_headers,
|
||||
)
|
||||
|
||||
assert result_phone.error_message is None
|
||||
|
|
@ -339,9 +400,21 @@ class TestAgentWithMCPTools:
|
|||
mcp_endpoint=dict(uri=uri),
|
||||
)
|
||||
|
||||
# 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),
|
||||
}
|
||||
|
||||
tools_list = llama_stack_client.tools.list(
|
||||
toolgroup_id=test_toolgroup_id,
|
||||
authorization=AUTH_TOKEN,
|
||||
extra_headers=auth_headers,
|
||||
)
|
||||
tool_defs = [
|
||||
{
|
||||
|
|
@ -350,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,
|
||||
}
|
||||
]
|
||||
|
||||
|
|
@ -358,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,
|
||||
authorization=AUTH_TOKEN,
|
||||
)
|
||||
|
||||
session_id = agent.create_session("test-session-complex")
|
||||
|
|
@ -380,7 +453,6 @@ class TestAgentWithMCPTools:
|
|||
}
|
||||
],
|
||||
stream=True,
|
||||
authorization=AUTH_TOKEN,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue