mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-03 09:53:45 +00:00
MCP authentication parameter implementation
This commit is contained in:
parent
da57b51fb6
commit
d0a8878337
7 changed files with 375 additions and 2 deletions
39
docs/static/deprecated-llama-stack-spec.yaml
vendored
39
docs/static/deprecated-llama-stack-spec.yaml
vendored
|
|
@ -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
|
||||
|
|
|
|||
39
docs/static/llama-stack-spec.yaml
vendored
39
docs/static/llama-stack-spec.yaml
vendored
|
|
@ -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
|
||||
|
|
|
|||
39
docs/static/stainless-llama-stack-spec.yaml
vendored
39
docs/static/stainless-llama-stack-spec.yaml
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
156
tests/integration/responses/test_mcp_authentication.py
Normal file
156
tests/integration/responses/test_mcp_authentication.py
Normal file
|
|
@ -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": "<FILLED_BY_TEST_RUNNER>",
|
||||
"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": "<FILLED_BY_TEST_RUNNER>",
|
||||
"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": "<FILLED_BY_TEST_RUNNER>",
|
||||
"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": "<FILLED_BY_TEST_RUNNER>"}],
|
||||
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
|
||||
Loading…
Add table
Add a link
Reference in a new issue