diff --git a/src/llama_stack/apis/agents/openai_responses.py b/src/llama_stack/apis/agents/openai_responses.py index d562411ec..09ece328e 100644 --- a/src/llama_stack/apis/agents/openai_responses.py +++ b/src/llama_stack/apis/agents/openai_responses.py @@ -496,8 +496,6 @@ class OpenAIResponseInputToolMCP(BaseModel): server_label: str server_url: str headers: dict[str, Any] | None = None - # OAuth access token for MCP server authentication - # Provide just the token (e.g., "my-secret-token"), the "Bearer " prefix will be added automatically authorization: str | None = None require_approval: Literal["always"] | Literal["never"] | ApprovalFilter = "never" diff --git a/src/llama_stack/providers/inline/agents/meta_reference/responses/streaming.py b/src/llama_stack/providers/inline/agents/meta_reference/responses/streaming.py index 029ba7b89..ea98d19cd 100644 --- a/src/llama_stack/providers/inline/agents/meta_reference/responses/streaming.py +++ b/src/llama_stack/providers/inline/agents/meta_reference/responses/streaming.py @@ -1079,23 +1079,12 @@ class StreamingResponseOrchestrator: "server_url": mcp_tool.server_url, "mcp_list_tools_id": list_id, } - # Prepare headers with authorization from tool config - headers = dict(mcp_tool.headers or {}) - if mcp_tool.authorization: - # Check if Authorization header already exists (case-insensitive check) - existing_keys_lower = {k.lower() for k in headers.keys()} - if "authorization" in existing_keys_lower: - raise ValueError( - "Cannot specify Authorization in both 'headers' and 'authorization' fields. " - "Please use only the 'authorization' field." - ) - # OAuth access token - add "Bearer " prefix - headers["Authorization"] = f"Bearer {mcp_tool.authorization}" - + # List MCP tools with authorization from tool config async with tracing.span("list_mcp_tools", attributes): tool_defs = await list_mcp_tools( endpoint=mcp_tool.server_url, - headers=headers, + headers=mcp_tool.headers, + authorization=mcp_tool.authorization, ) # Create the MCP list tools message diff --git a/src/llama_stack/providers/inline/agents/meta_reference/responses/tool_executor.py b/src/llama_stack/providers/inline/agents/meta_reference/responses/tool_executor.py index bea3e720c..47ca500a9 100644 --- a/src/llama_stack/providers/inline/agents/meta_reference/responses/tool_executor.py +++ b/src/llama_stack/providers/inline/agents/meta_reference/responses/tool_executor.py @@ -299,25 +299,14 @@ class ToolExecutor: "server_url": mcp_tool.server_url, "tool_name": function_name, } - # Prepare headers with authorization from tool config - headers = dict(mcp_tool.headers or {}) - if mcp_tool.authorization: - # Check if Authorization header already exists (case-insensitive check) - existing_keys_lower = {k.lower() for k in headers.keys()} - if "authorization" in existing_keys_lower: - raise ValueError( - "Cannot specify Authorization in both 'headers' and 'authorization' fields. " - "Please use only the 'authorization' field." - ) - # OAuth access token - add "Bearer " prefix - headers["Authorization"] = f"Bearer {mcp_tool.authorization}" - + # Invoke MCP tool with authorization from tool config async with tracing.span("invoke_mcp_tool", attributes): result = await invoke_mcp_tool( endpoint=mcp_tool.server_url, - headers=headers, tool_name=function_name, kwargs=tool_kwargs, + headers=mcp_tool.headers, + authorization=mcp_tool.authorization, ) elif function_name == "knowledge_search": response_file_search_tool = ( diff --git a/src/llama_stack/providers/utils/tools/mcp.py b/src/llama_stack/providers/utils/tools/mcp.py index a271cb959..309e3bf5d 100644 --- a/src/llama_stack/providers/utils/tools/mcp.py +++ b/src/llama_stack/providers/utils/tools/mcp.py @@ -27,6 +27,36 @@ from llama_stack.providers.utils.tools.ttl_dict import TTLDict logger = get_logger(__name__, category="tools") + +def prepare_mcp_headers(base_headers: dict[str, str] | None, authorization: str | None) -> dict[str, str]: + """Prepare headers for MCP requests with authorization handling. + + Args: + base_headers: Base headers to use (e.g., from mcp_tool.headers) + authorization: OAuth access token (just the token, not "Bearer ") + + Returns: + Final headers dict with Authorization header if authorization is provided + + Raises: + ValueError: If both base_headers contains Authorization and authorization parameter is provided + """ + headers = dict(base_headers or {}) + + if authorization: + # Check if Authorization header already exists (case-insensitive check) + existing_keys_lower = {k.lower() for k in headers.keys()} + if "authorization" in existing_keys_lower: + raise ValueError( + "Cannot specify Authorization in both 'headers' and 'authorization' fields. " + "Please use only the 'authorization' field." + ) + # OAuth access token - add "Bearer " prefix + headers["Authorization"] = f"Bearer {authorization}" + + return headers + + protocol_cache = TTLDict(ttl_seconds=3600) @@ -109,9 +139,29 @@ async def client_wrapper(endpoint: str, headers: dict[str, str]) -> AsyncGenerat raise -async def list_mcp_tools(endpoint: str, headers: dict[str, str]) -> ListToolDefsResponse: +async def list_mcp_tools( + endpoint: str, + headers: dict[str, str] | None = None, + authorization: str | None = None, +) -> ListToolDefsResponse: + """List tools available from an MCP server. + + Args: + endpoint: MCP server endpoint URL + headers: Optional base headers to include + authorization: Optional OAuth access token (just the token, not "Bearer ") + + Returns: + List of tool definitions from the MCP server + + Raises: + ValueError: If both headers contains Authorization and authorization parameter is provided + """ + # Prepare headers with authorization handling + final_headers = prepare_mcp_headers(headers, authorization) + tools = [] - async with client_wrapper(endpoint, headers) as session: + async with client_wrapper(endpoint, final_headers) as session: tools_result = await session.list_tools() for tool in tools_result.tools: tools.append( @@ -129,9 +179,31 @@ async def list_mcp_tools(endpoint: str, headers: dict[str, str]) -> ListToolDefs async def invoke_mcp_tool( - endpoint: str, headers: dict[str, str], tool_name: str, kwargs: dict[str, Any] + endpoint: str, + tool_name: str, + kwargs: dict[str, Any], + headers: dict[str, str] | None = None, + authorization: str | None = None, ) -> ToolInvocationResult: - async with client_wrapper(endpoint, headers) as session: + """Invoke an MCP tool with the given arguments. + + Args: + endpoint: MCP server endpoint URL + tool_name: Name of the tool to invoke + kwargs: Tool invocation arguments + headers: Optional base headers to include + authorization: Optional OAuth access token (just the token, not "Bearer ") + + Returns: + Tool invocation result with content and error information + + Raises: + ValueError: If both headers contains Authorization and authorization parameter is provided + """ + # Prepare headers with authorization handling + final_headers = prepare_mcp_headers(headers, authorization) + + async with client_wrapper(endpoint, final_headers) as session: result = await session.call_tool(tool_name, kwargs) content: list[InterleavedContentItem] = []