diff --git a/docs/static/deprecated-llama-stack-spec.yaml b/docs/static/deprecated-llama-stack-spec.yaml index ec3880a6b..51f6e7ecb 100644 --- a/docs/static/deprecated-llama-stack-spec.yaml +++ b/docs/static/deprecated-llama-stack-spec.yaml @@ -7711,6 +7711,41 @@ components: title: ResponseGuardrailSpec description: >- Specification for a guardrail to apply during response generation. + MCPAuthentication: + type: object + properties: + type: + type: string + enum: + - bearer + - basic + - api_key + description: >- + Authentication type ("bearer", "basic", or "api_key") + token: + type: string + description: Bearer token for bearer authentication + username: + type: string + description: Username for basic authentication + password: + type: string + description: Password for basic authentication + api_key: + type: string + description: API key for api_key authentication + header_name: + type: string + default: X-API-Key + description: >- + Custom header name for API key (default: "X-API-Key") + additionalProperties: false + required: + - type + - header_name + title: MCPAuthentication + description: >- + Authentication configuration for MCP servers. OpenAIResponseInputTool: oneOf: - $ref: '#/components/schemas/OpenAIResponseInputToolWebSearch' @@ -7750,6 +7785,10 @@ components: - type: object description: >- (Optional) HTTP headers to include when connecting to the server + authentication: + $ref: '#/components/schemas/MCPAuthentication' + description: >- + (Optional) Authentication configuration for the MCP server require_approval: oneOf: - type: string diff --git a/docs/static/llama-stack-spec.yaml b/docs/static/llama-stack-spec.yaml index e35287952..dc9178af4 100644 --- a/docs/static/llama-stack-spec.yaml +++ b/docs/static/llama-stack-spec.yaml @@ -6443,6 +6443,41 @@ components: title: ResponseGuardrailSpec description: >- Specification for a guardrail to apply during response generation. + MCPAuthentication: + type: object + properties: + type: + type: string + enum: + - bearer + - basic + - api_key + description: >- + Authentication type ("bearer", "basic", or "api_key") + token: + type: string + description: Bearer token for bearer authentication + username: + type: string + description: Username for basic authentication + password: + type: string + description: Password for basic authentication + api_key: + type: string + description: API key for api_key authentication + header_name: + type: string + default: X-API-Key + description: >- + Custom header name for API key (default: "X-API-Key") + additionalProperties: false + required: + - type + - header_name + title: MCPAuthentication + description: >- + Authentication configuration for MCP servers. OpenAIResponseInputTool: oneOf: - $ref: '#/components/schemas/OpenAIResponseInputToolWebSearch' @@ -6482,6 +6517,10 @@ components: - type: object description: >- (Optional) HTTP headers to include when connecting to the server + authentication: + $ref: '#/components/schemas/MCPAuthentication' + description: >- + (Optional) Authentication configuration for the MCP server require_approval: oneOf: - type: string diff --git a/docs/static/stainless-llama-stack-spec.yaml b/docs/static/stainless-llama-stack-spec.yaml index a1085c9eb..27fe184e6 100644 --- a/docs/static/stainless-llama-stack-spec.yaml +++ b/docs/static/stainless-llama-stack-spec.yaml @@ -7656,6 +7656,41 @@ components: title: ResponseGuardrailSpec description: >- Specification for a guardrail to apply during response generation. + MCPAuthentication: + type: object + properties: + type: + type: string + enum: + - bearer + - basic + - api_key + description: >- + Authentication type ("bearer", "basic", or "api_key") + token: + type: string + description: Bearer token for bearer authentication + username: + type: string + description: Username for basic authentication + password: + type: string + description: Password for basic authentication + api_key: + type: string + description: API key for api_key authentication + header_name: + type: string + default: X-API-Key + description: >- + Custom header name for API key (default: "X-API-Key") + additionalProperties: false + required: + - type + - header_name + title: MCPAuthentication + description: >- + Authentication configuration for MCP servers. OpenAIResponseInputTool: oneOf: - $ref: '#/components/schemas/OpenAIResponseInputToolWebSearch' @@ -7695,6 +7730,10 @@ components: - type: object description: >- (Optional) HTTP headers to include when connecting to the server + authentication: + $ref: '#/components/schemas/MCPAuthentication' + description: >- + (Optional) Authentication configuration for the MCP server require_approval: oneOf: - type: string diff --git a/src/llama_stack/apis/agents/openai_responses.py b/src/llama_stack/apis/agents/openai_responses.py index 69e2b2012..b67b1d589 100644 --- a/src/llama_stack/apis/agents/openai_responses.py +++ b/src/llama_stack/apis/agents/openai_responses.py @@ -479,6 +479,26 @@ class AllowedToolsFilter(BaseModel): tool_names: list[str] | None = None +@json_schema_type +class MCPAuthentication(BaseModel): + """Authentication configuration for MCP servers. + + :param type: Authentication type ("bearer", "basic", or "api_key") + :param token: Bearer token for bearer authentication + :param username: Username for basic authentication + :param password: Password for basic authentication + :param api_key: API key for api_key authentication + :param header_name: Custom header name for API key (default: "X-API-Key") + """ + + type: Literal["bearer", "basic", "api_key"] + token: str | None = None + username: str | None = None + password: str | None = None + api_key: str | None = None + header_name: str = "X-API-Key" + + @json_schema_type class OpenAIResponseInputToolMCP(BaseModel): """Model Context Protocol (MCP) tool configuration for OpenAI response inputs. @@ -487,6 +507,7 @@ class OpenAIResponseInputToolMCP(BaseModel): :param server_label: Label to identify this MCP server :param server_url: URL endpoint of the MCP server :param headers: (Optional) HTTP headers to include when connecting to the server + :param authentication: (Optional) Authentication configuration for the MCP server :param require_approval: Approval requirement for tool calls ("always", "never", or filter) :param allowed_tools: (Optional) Restriction on which tools can be used from this server """ @@ -495,6 +516,7 @@ class OpenAIResponseInputToolMCP(BaseModel): server_label: str server_url: str headers: dict[str, Any] | None = None + authentication: MCPAuthentication | None = None require_approval: Literal["always"] | Literal["never"] | ApprovalFilter = "never" allowed_tools: list[str] | AllowedToolsFilter | None = None 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 ef5603420..ca4c28752 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 @@ -11,6 +11,7 @@ from typing import Any from llama_stack.apis.agents.openai_responses import ( AllowedToolsFilter, ApprovalFilter, + MCPAuthentication, MCPListToolsTool, OpenAIResponseContentPartOutputText, OpenAIResponseContentPartReasoningText, @@ -80,6 +81,34 @@ from .utils import ( logger = get_logger(name=__name__, category="agents::meta_reference") +def _convert_authentication_to_headers(auth: MCPAuthentication) -> dict[str, str]: + """Convert MCPAuthentication config to HTTP headers. + + Args: + auth: Authentication configuration + + Returns: + Dictionary of HTTP headers for authentication + """ + headers = {} + + if auth.type == "bearer": + if auth.token: + headers["Authorization"] = f"Bearer {auth.token}" + elif auth.type == "basic": + if auth.username and auth.password: + import base64 + + credentials = f"{auth.username}:{auth.password}" + encoded = base64.b64encode(credentials.encode()).decode() + headers["Authorization"] = f"Basic {encoded}" + elif auth.type == "api_key": + if auth.api_key: + headers[auth.header_name] = auth.api_key + + return headers + + def convert_tooldef_to_chat_tool(tool_def): """Convert a ToolDef to OpenAI ChatCompletionToolParam format. @@ -1079,10 +1108,20 @@ class StreamingResponseOrchestrator: "server_url": mcp_tool.server_url, "mcp_list_tools_id": list_id, } + # Prepare headers with authentication from tool config + headers = dict(mcp_tool.headers or {}) + if mcp_tool.authentication: + auth_headers = _convert_authentication_to_headers(mcp_tool.authentication) + # Don't override existing headers (case-insensitive check) + existing_keys_lower = {k.lower() for k in headers.keys()} + for key, value in auth_headers.items(): + if key.lower() not in existing_keys_lower: + headers[key] = value + async with tracing.span("list_mcp_tools", attributes): tool_defs = await list_mcp_tools( endpoint=mcp_tool.server_url, - headers=mcp_tool.headers or {}, + headers=headers, ) # 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 09a161d50..10e3a1ec8 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 @@ -10,6 +10,7 @@ from collections.abc import AsyncIterator from typing import Any from llama_stack.apis.agents.openai_responses import ( + MCPAuthentication, OpenAIResponseInputToolFileSearch, OpenAIResponseInputToolMCP, OpenAIResponseObjectStreamResponseFileSearchCallCompleted, @@ -47,6 +48,34 @@ from .types import ChatCompletionContext, ToolExecutionResult logger = get_logger(name=__name__, category="agents::meta_reference") +def _convert_authentication_to_headers(auth: MCPAuthentication) -> dict[str, str]: + """Convert MCPAuthentication config to HTTP headers. + + Args: + auth: Authentication configuration + + Returns: + Dictionary of HTTP headers for authentication + """ + headers = {} + + if auth.type == "bearer": + if auth.token: + headers["Authorization"] = f"Bearer {auth.token}" + elif auth.type == "basic": + if auth.username and auth.password: + import base64 + + credentials = f"{auth.username}:{auth.password}" + encoded = base64.b64encode(credentials.encode()).decode() + headers["Authorization"] = f"Basic {encoded}" + elif auth.type == "api_key": + if auth.api_key: + headers[auth.header_name] = auth.api_key + + return headers + + class ToolExecutor: def __init__( self, @@ -299,10 +328,20 @@ class ToolExecutor: "server_url": mcp_tool.server_url, "tool_name": function_name, } + # Prepare headers with authentication from tool config + headers = dict(mcp_tool.headers or {}) + if mcp_tool.authentication: + auth_headers = _convert_authentication_to_headers(mcp_tool.authentication) + # Don't override existing headers (case-insensitive check) + existing_keys_lower = {k.lower() for k in headers.keys()} + for key, value in auth_headers.items(): + if key.lower() not in existing_keys_lower: + headers[key] = value + async with tracing.span("invoke_mcp_tool", attributes): result = await invoke_mcp_tool( endpoint=mcp_tool.server_url, - headers=mcp_tool.headers or {}, + headers=headers, tool_name=function_name, kwargs=tool_kwargs, ) diff --git a/tests/integration/responses/test_mcp_authentication.py b/tests/integration/responses/test_mcp_authentication.py new file mode 100644 index 000000000..c6df3f1e9 --- /dev/null +++ b/tests/integration/responses/test_mcp_authentication.py @@ -0,0 +1,156 @@ +# 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. + +import pytest + +from llama_stack import LlamaStackAsLibraryClient +from tests.common.mcp import make_mcp_server + +from .helpers import setup_mcp_tools + + +def test_mcp_authentication_bearer(compat_client, text_model_id): + """Test that bearer authentication is correctly applied to MCP requests.""" + if not isinstance(compat_client, LlamaStackAsLibraryClient): + pytest.skip("in-process MCP server is only supported in library client") + + test_token = "test-bearer-token-789" + with make_mcp_server(required_auth_token=test_token) as mcp_server_info: + tools = setup_mcp_tools( + [ + { + "type": "mcp", + "server_label": "auth-mcp", + "server_url": "", + "authentication": { + "type": "bearer", + "token": test_token, + }, + } + ], + mcp_server_info, + ) + + # Create response - authentication should be applied + response = compat_client.responses.create( + model=text_model_id, + input="What is the boiling point of myawesomeliquid?", + tools=tools, + stream=False, + ) + + # Verify list_tools succeeded (requires auth) + assert len(response.output) >= 3 + assert response.output[0].type == "mcp_list_tools" + assert len(response.output[0].tools) == 2 + + # Verify tool invocation succeeded (requires auth) + assert response.output[1].type == "mcp_call" + assert response.output[1].error is None + + +def test_mcp_authentication_api_key(compat_client, text_model_id): + """Test that API key authentication is correctly applied to MCP requests.""" + if not isinstance(compat_client, LlamaStackAsLibraryClient): + pytest.skip("in-process MCP server is only supported in library client") + + test_api_key = "test-api-key-456" + with make_mcp_server(required_auth_token=test_api_key, auth_header="X-API-Key") as mcp_server_info: + tools = setup_mcp_tools( + [ + { + "type": "mcp", + "server_label": "apikey-mcp", + "server_url": "", + "authentication": { + "type": "api_key", + "api_key": test_api_key, + "header_name": "X-API-Key", + }, + } + ], + mcp_server_info, + ) + + # Create response - authentication should be applied + response = compat_client.responses.create( + model=text_model_id, + input="What is the boiling point of myawesomeliquid?", + tools=tools, + stream=False, + ) + + # Verify operations succeeded + assert len(response.output) >= 3 + assert response.output[0].type == "mcp_list_tools" + assert response.output[1].type == "mcp_call" + assert response.output[1].error is None + + +def test_mcp_authentication_fallback_to_headers(compat_client, text_model_id): + """Test that authentication parameter doesn't override existing headers.""" + if not isinstance(compat_client, LlamaStackAsLibraryClient): + pytest.skip("in-process MCP server is only supported in library client") + + # Headers should take precedence - this test uses headers auth + test_token = "headers-token-123" + with make_mcp_server(required_auth_token=test_token) as mcp_server_info: + tools = setup_mcp_tools( + [ + { + "type": "mcp", + "server_label": "headers-mcp", + "server_url": "", + "headers": {"Authorization": f"Bearer {test_token}"}, + "authentication": { + "type": "bearer", + "token": "should-not-override", + }, + } + ], + mcp_server_info, + ) + + # Create response - headers should take precedence + response = compat_client.responses.create( + model=text_model_id, + input="What is the boiling point of myawesomeliquid?", + tools=tools, + stream=False, + ) + + # Verify operations succeeded with headers auth + assert len(response.output) >= 3 + assert response.output[0].type == "mcp_list_tools" + assert response.output[1].type == "mcp_call" + assert response.output[1].error is None + + +def test_mcp_authentication_backward_compatibility(compat_client, text_model_id): + """Test that MCP tools work without authentication (backward compatibility).""" + if not isinstance(compat_client, LlamaStackAsLibraryClient): + pytest.skip("in-process MCP server is only supported in library client") + + # No authentication required + with make_mcp_server(required_auth_token=None) as mcp_server_info: + tools = setup_mcp_tools( + [{"type": "mcp", "server_label": "noauth-mcp", "server_url": ""}], + mcp_server_info, + ) + + # Create response without authentication + response = compat_client.responses.create( + model=text_model_id, + input="What is the boiling point of myawesomeliquid?", + tools=tools, + stream=False, + ) + + # Verify operations succeeded without auth + assert len(response.output) >= 3 + assert response.output[0].type == "mcp_list_tools" + assert response.output[1].type == "mcp_call" + assert response.output[1].error is None