mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
fix(mcp): add AWS SigV4 auth for Bedrock AgentCore MCP servers (#22782)
* fix(mcp): add AWS SigV4 auth for Bedrock AgentCore MCP servers Add aws_sigv4 auth type to MCP client via httpx.Auth subclass that signs each request with SigV4 using botocore. Enables mcp_servers config to connect to AgentCore-hosted MCP servers. * docs(mcp): add AWS SigV4 auth documentation for Bedrock AgentCore Add dedicated docs page for configuring MCP servers with AWS SigV4 authentication, update MCP overview with aws_sigv4 auth type and config example, and link from Bedrock AgentCore provider docs. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix(mcp): address Greptile review — requires_request_body, full header signing, health check - Add requires_request_body = True to MCPSigV4Auth so httpx buffers the request body before calling auth_flow (prevents empty body hash for streaming requests) - Pass all request headers to AWSRequest for canonical SigV4 signing instead of only Content-Type - Exclude aws_sigv4 from health check skip logic since it has its own credential fields (not authentication_token) - Fix docs: mark aws_access_key_id/aws_secret_access_key as optional (falls back to boto3 credential chain) - Add test for requires_request_body flag Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>
This commit is contained in:
michelligabriele
GitHub
parent
9543d785b5
commit
ffc89e4ef6
9 changed files with 601 additions and 4 deletions
|
|
@ -217,6 +217,7 @@ mcp_servers:
|
|||
| `bearer_token` | `Authorization: Bearer <auth_value>` |
|
||||
| `basic` | `Authorization: Basic <auth_value>` |
|
||||
| `authorization` | `Authorization: <auth_value>` |
|
||||
| `aws_sigv4` | Per-request AWS SigV4 signature ([details](./mcp_aws_sigv4.md)) |
|
||||
|
||||
- **Extra Headers**: Optional list of additional header names that should be forwarded from client to the MCP server
|
||||
- **Static Headers**: Optional map of header key/value pairs to include every request to the MCP server.
|
||||
|
|
@ -257,6 +258,16 @@ mcp_servers:
|
|||
auth_type: "authorization"
|
||||
auth_value: "Token example123" # headers={"Authorization": "Token example123"}
|
||||
|
||||
# AWS SigV4 for Bedrock AgentCore MCP servers
|
||||
agentcore_mcp:
|
||||
url: "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/<url-encoded-ARN>/invocations"
|
||||
transport: "http"
|
||||
auth_type: "aws_sigv4"
|
||||
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
|
||||
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
|
||||
aws_region_name: us-east-1
|
||||
aws_service_name: bedrock-agentcore
|
||||
|
||||
# Example with extra headers forwarding
|
||||
github_mcp:
|
||||
url: "https://api.githubcopilot.com/mcp"
|
||||
|
|
|
|||
144
docs/my-website/docs/mcp_aws_sigv4.md
Normal file
144
docs/my-website/docs/mcp_aws_sigv4.md
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
# MCP - AWS SigV4 Auth
|
||||
|
||||
Use AWS SigV4 authentication to connect LiteLLM to MCP servers hosted on [AWS Bedrock AgentCore](https://docs.aws.amazon.com/bedrock/latest/userguide/agentcore.html).
|
||||
|
||||
## Why SigV4?
|
||||
|
||||
AWS services authenticate requests using [Signature Version 4](https://docs.aws.amazon.com/general/latest/gr/signature-version-4.html) — a per-request signing protocol that includes the request body in the cryptographic signature. This is fundamentally different from static-header auth types (`api_key`, `bearer_token`, etc.) which send the same header on every request.
|
||||
|
||||
LiteLLM's `aws_sigv4` auth type handles this automatically: every outgoing MCP request is signed with your AWS credentials before it's sent.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### 1. Set AWS credentials
|
||||
|
||||
```bash
|
||||
export AWS_ACCESS_KEY_ID="AKIA..."
|
||||
export AWS_SECRET_ACCESS_KEY="..."
|
||||
export AWS_REGION_NAME="us-east-1"
|
||||
```
|
||||
|
||||
### 2. Add your AgentCore MCP server to config.yaml
|
||||
|
||||
```yaml title="config.yaml" showLineNumbers
|
||||
model_list:
|
||||
- model_name: gpt-4o
|
||||
litellm_params:
|
||||
model: openai/gpt-4o
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
mcp_servers:
|
||||
my_agentcore_mcp:
|
||||
url: "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/<url-encoded-ARN>/invocations"
|
||||
transport: "http"
|
||||
auth_type: "aws_sigv4"
|
||||
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
|
||||
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
|
||||
aws_region_name: "us-east-1"
|
||||
aws_service_name: "bedrock-agentcore"
|
||||
```
|
||||
|
||||
:::info URL encoding
|
||||
|
||||
The AgentCore runtime ARN must be URL-encoded in the `url` field. For example:
|
||||
|
||||
```
|
||||
arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/my-mcp-server
|
||||
```
|
||||
|
||||
becomes:
|
||||
|
||||
```
|
||||
arn%3Aaws%3Abedrock-agentcore%3Aus-east-1%3A123456789012%3Aruntime%2Fmy-mcp-server
|
||||
```
|
||||
|
||||
:::
|
||||
|
||||
### 3. Start the proxy
|
||||
|
||||
```bash
|
||||
litellm --config config.yaml
|
||||
```
|
||||
|
||||
### 4. Use the MCP tools
|
||||
|
||||
Once started, your AgentCore MCP tools are available through LiteLLM like any other MCP server:
|
||||
|
||||
```bash title="List available tools"
|
||||
curl http://localhost:4000/mcp-rest/tools/list \
|
||||
-H "Authorization: Bearer sk-1234"
|
||||
```
|
||||
|
||||
```bash title="Call a tool"
|
||||
curl http://localhost:4000/mcp-rest/tools/call \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-d '{
|
||||
"name": "my_agentcore_mcp_your_tool_name",
|
||||
"arguments": {"key": "value"}
|
||||
}'
|
||||
```
|
||||
|
||||
## Config Reference
|
||||
|
||||
| Field | Required | Description |
|
||||
|-------|----------|-------------|
|
||||
| `url` | Yes | AgentCore MCP server URL (with URL-encoded ARN) |
|
||||
| `transport` | Yes | Must be `"http"` |
|
||||
| `auth_type` | Yes | Must be `"aws_sigv4"` |
|
||||
| `aws_access_key_id` | No | AWS access key. Supports `os.environ/VAR_NAME`. Falls back to boto3 credential chain if omitted |
|
||||
| `aws_secret_access_key` | No | AWS secret key. Supports `os.environ/VAR_NAME`. Falls back to boto3 credential chain if omitted |
|
||||
| `aws_region_name` | Yes | AWS region (e.g., `us-east-1`) |
|
||||
| `aws_service_name` | No | AWS service name for signing. Defaults to `bedrock-agentcore` |
|
||||
| `aws_session_token` | No | AWS session token for temporary credentials. Supports `os.environ/VAR_NAME` |
|
||||
|
||||
## How It Works
|
||||
|
||||
LiteLLM uses an `httpx.Auth` subclass (`MCPSigV4Auth`) that hooks into the HTTP request lifecycle:
|
||||
|
||||
1. For every outgoing MCP request, the auth handler computes a SHA-256 hash of the request body
|
||||
2. It creates a SigV4 signature using your AWS credentials, the request URL, headers, and body hash
|
||||
3. The signed `Authorization` and `x-amz-date` headers are added to the request
|
||||
4. AWS validates the signature and processes the MCP request
|
||||
|
||||
This happens transparently — no manual token management required.
|
||||
|
||||
## Using Temporary Credentials (STS)
|
||||
|
||||
If you use AWS STS temporary credentials (e.g., from IAM roles or SSO), include the session token:
|
||||
|
||||
```yaml title="config.yaml with STS credentials" showLineNumbers
|
||||
mcp_servers:
|
||||
my_agentcore_mcp:
|
||||
url: "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/<url-encoded-ARN>/invocations"
|
||||
transport: "http"
|
||||
auth_type: "aws_sigv4"
|
||||
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
|
||||
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
|
||||
aws_session_token: os.environ/AWS_SESSION_TOKEN
|
||||
aws_region_name: "us-east-1"
|
||||
aws_service_name: "bedrock-agentcore"
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### 403 Forbidden from AWS
|
||||
|
||||
- Verify your AWS credentials are valid and not expired
|
||||
- Check that `aws_region_name` matches the region in your AgentCore URL
|
||||
- Ensure `aws_service_name` is set to `bedrock-agentcore`
|
||||
- If using STS credentials, confirm `aws_session_token` is set and not expired
|
||||
|
||||
### Health check errors on startup
|
||||
|
||||
SigV4-authenticated MCP servers skip the standard health check on proxy startup. This is expected — the proxy will still sign requests correctly when tools are invoked.
|
||||
|
||||
### "botocore not found" error
|
||||
|
||||
Install the `botocore` package:
|
||||
|
||||
```bash
|
||||
pip install botocore
|
||||
```
|
||||
|
||||
`botocore` is used for SigV4 credential handling and is required when using `aws_sigv4` auth.
|
||||
|
|
@ -13,7 +13,7 @@ Call Bedrock AgentCore in the OpenAI Request/Response format.
|
|||
|
||||
:::info
|
||||
|
||||
This documentation is for **AgentCore Agents** (agent runtimes). If you want to use AgentCore MCP servers, add them as you would any other MCP server. See the [MCP documentation](https://docs.litellm.ai/docs/mcp) for details.
|
||||
This documentation is for **AgentCore Agents** (agent runtimes). If you want to use AgentCore MCP servers with LiteLLM, see the [MCP AWS SigV4 Auth](https://docs.litellm.ai/docs/mcp_aws_sigv4) guide for setup instructions.
|
||||
|
||||
:::
|
||||
|
||||
|
|
|
|||
|
|
@ -614,6 +614,7 @@ const sidebars = {
|
|||
"mcp_usage",
|
||||
"mcp_openapi",
|
||||
"mcp_oauth",
|
||||
"mcp_aws_sigv4",
|
||||
"mcp_public_internet",
|
||||
"mcp_semantic_filter",
|
||||
"mcp_control",
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
|
|||
|
||||
import asyncio
|
||||
import base64
|
||||
from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, TypeVar, Union
|
||||
from typing import Any, Awaitable, Callable, Dict, Generator, List, Optional, Tuple, TypeVar, Union
|
||||
|
||||
import httpx
|
||||
from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
|
||||
|
|
@ -50,6 +50,86 @@ def to_basic_auth(auth_value: str) -> str:
|
|||
TSessionResult = TypeVar("TSessionResult")
|
||||
|
||||
|
||||
class MCPSigV4Auth(httpx.Auth):
|
||||
"""
|
||||
httpx Auth class that signs each request with AWS SigV4.
|
||||
|
||||
This is used for MCP servers that require AWS SigV4 authentication,
|
||||
such as AWS Bedrock AgentCore MCP servers. httpx calls auth_flow()
|
||||
for every outgoing request, enabling per-request signature computation.
|
||||
"""
|
||||
|
||||
requires_request_body = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
aws_access_key_id: Optional[str] = None,
|
||||
aws_secret_access_key: Optional[str] = None,
|
||||
aws_session_token: Optional[str] = None,
|
||||
aws_region_name: Optional[str] = None,
|
||||
aws_service_name: Optional[str] = None,
|
||||
):
|
||||
try:
|
||||
from botocore.credentials import Credentials
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Missing botocore to use AWS SigV4 authentication. "
|
||||
"Run 'pip install boto3'."
|
||||
)
|
||||
|
||||
self.service_name = aws_service_name or "bedrock-agentcore"
|
||||
self.region_name = aws_region_name or "us-east-1"
|
||||
|
||||
# Note: os.environ/ prefixed values are already resolved by
|
||||
# ProxyConfig._check_for_os_environ_vars() at config load time.
|
||||
# Values arrive here as plain strings.
|
||||
if aws_access_key_id and aws_secret_access_key:
|
||||
self.credentials = Credentials(
|
||||
access_key=aws_access_key_id,
|
||||
secret_key=aws_secret_access_key,
|
||||
token=aws_session_token,
|
||||
)
|
||||
else:
|
||||
# Fall back to default boto3 credential chain
|
||||
import botocore.session
|
||||
|
||||
session = botocore.session.get_session()
|
||||
self.credentials = session.get_credentials()
|
||||
if self.credentials is None:
|
||||
raise ValueError(
|
||||
"No AWS credentials found. Provide aws_access_key_id and "
|
||||
"aws_secret_access_key, or configure default credentials "
|
||||
"(env vars, ~/.aws/credentials, instance profile)."
|
||||
)
|
||||
|
||||
def auth_flow(
|
||||
self, request: httpx.Request
|
||||
) -> Generator[httpx.Request, httpx.Response, None]:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
|
||||
# Build AWSRequest from the httpx Request.
|
||||
# Pass all request headers so the canonical SigV4 signature covers them.
|
||||
aws_request = AWSRequest(
|
||||
method=request.method,
|
||||
url=str(request.url),
|
||||
data=request.content,
|
||||
headers=dict(request.headers),
|
||||
)
|
||||
|
||||
# Sign the request — SigV4Auth.add_auth() adds Authorization,
|
||||
# X-Amz-Date, and X-Amz-Security-Token (if session token present).
|
||||
# Host header is derived automatically from the URL.
|
||||
sigv4 = SigV4Auth(self.credentials, self.service_name, self.region_name)
|
||||
sigv4.add_auth(aws_request)
|
||||
|
||||
# Copy SigV4 headers back to the httpx request
|
||||
for header_name, header_value in aws_request.headers.items():
|
||||
request.headers[header_name] = header_value
|
||||
|
||||
yield request
|
||||
|
||||
|
||||
class MCPClient:
|
||||
"""
|
||||
MCP Client supporting:
|
||||
|
|
@ -68,6 +148,7 @@ class MCPClient:
|
|||
stdio_config: Optional[MCPStdioConfig] = None,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
ssl_verify: Optional[VerifyTypes] = None,
|
||||
aws_auth: Optional[httpx.Auth] = None,
|
||||
):
|
||||
self.server_url: str = server_url
|
||||
self.transport_type: MCPTransport = transport_type
|
||||
|
|
@ -77,6 +158,7 @@ class MCPClient:
|
|||
self.stdio_config: Optional[MCPStdioConfig] = stdio_config
|
||||
self.extra_headers: Optional[Dict[str, str]] = extra_headers
|
||||
self.ssl_verify: Optional[VerifyTypes] = ssl_verify
|
||||
self._aws_auth: Optional[httpx.Auth] = aws_auth
|
||||
# handle the basic auth value if provided
|
||||
if auth_value:
|
||||
self.update_auth_value(auth_value)
|
||||
|
|
@ -216,6 +298,9 @@ class MCPClient:
|
|||
headers["Authorization"] = f"token {self._mcp_auth_value}"
|
||||
elif isinstance(self._mcp_auth_value, dict):
|
||||
headers.update(self._mcp_auth_value)
|
||||
# Note: aws_sigv4 auth is not handled here — SigV4 requires per-request
|
||||
# signing (including the body hash), so it uses httpx.Auth flow instead
|
||||
# of static headers. See MCPSigV4Auth and _create_httpx_client_factory().
|
||||
|
||||
# update the headers with the extra headers
|
||||
if self.extra_headers:
|
||||
|
|
@ -248,10 +333,16 @@ class MCPClient:
|
|||
f"MCP client using SSL configuration: {type(ssl_config).__name__}"
|
||||
)
|
||||
|
||||
# Use SigV4 auth if configured and no explicit auth provided.
|
||||
# The MCP SDK's sse_client and streamable_http_client call this
|
||||
# factory without passing auth=, so self._aws_auth is used.
|
||||
# For non-SigV4 clients, self._aws_auth is None — no behavior change.
|
||||
effective_auth = auth if auth is not None else self._aws_auth
|
||||
|
||||
return httpx.AsyncClient(
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
auth=auth,
|
||||
auth=effective_auth,
|
||||
verify=ssl_config,
|
||||
follow_redirects=True,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -38,7 +38,7 @@ from litellm.constants import (
|
|||
MCP_TOOL_LISTING_TIMEOUT,
|
||||
)
|
||||
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
|
||||
from litellm.experimental_mcp_client.client import MCPClient
|
||||
from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
||||
MCPRequestHandler,
|
||||
|
|
@ -340,6 +340,12 @@ class MCPServerManager:
|
|||
available_on_public_internet=bool(
|
||||
server_config.get("available_on_public_internet", True)
|
||||
),
|
||||
# AWS SigV4 fields
|
||||
aws_access_key_id=server_config.get("aws_access_key_id", None),
|
||||
aws_secret_access_key=server_config.get("aws_secret_access_key", None),
|
||||
aws_session_token=server_config.get("aws_session_token", None),
|
||||
aws_region_name=server_config.get("aws_region_name", None),
|
||||
aws_service_name=server_config.get("aws_service_name", None),
|
||||
)
|
||||
self.config_mcp_servers[server_id] = new_server
|
||||
|
||||
|
|
@ -591,6 +597,10 @@ class MCPServerManager:
|
|||
else:
|
||||
client_secret_value = encrypted_client_secret
|
||||
|
||||
# TODO: Add AWS SigV4 credential decryption here when DB-stored
|
||||
# SigV4 MCP servers are supported. Requires corresponding changes
|
||||
# to encrypt_credentials() in db.py and MCPCredentials TypedDict.
|
||||
|
||||
scopes: Optional[List[str]] = None
|
||||
if credentials_dict:
|
||||
scopes_value = credentials_dict.get("scopes")
|
||||
|
|
@ -977,6 +987,18 @@ class MCPServerManager:
|
|||
else:
|
||||
# For HTTP/SSE transports
|
||||
server_url = server.url or ""
|
||||
|
||||
# Create SigV4 auth if configured
|
||||
aws_auth = None
|
||||
if server.auth_type == MCPAuth.aws_sigv4:
|
||||
aws_auth = MCPSigV4Auth(
|
||||
aws_access_key_id=server.aws_access_key_id,
|
||||
aws_secret_access_key=server.aws_secret_access_key,
|
||||
aws_session_token=server.aws_session_token,
|
||||
aws_region_name=server.aws_region_name,
|
||||
aws_service_name=server.aws_service_name,
|
||||
)
|
||||
|
||||
return MCPClient(
|
||||
server_url=server_url,
|
||||
transport_type=transport,
|
||||
|
|
@ -984,6 +1006,7 @@ class MCPServerManager:
|
|||
auth_value=auth_value,
|
||||
timeout=MCP_CLIENT_TIMEOUT,
|
||||
extra_headers=extra_headers,
|
||||
aws_auth=aws_auth,
|
||||
)
|
||||
|
||||
async def _get_tools_from_server(
|
||||
|
|
@ -2510,9 +2533,11 @@ class MCPServerManager:
|
|||
if server.requires_per_user_auth:
|
||||
should_skip_health_check = True
|
||||
# Skip if auth_type is not none and authentication_token is missing
|
||||
# (except aws_sigv4 which uses its own credential fields)
|
||||
elif (
|
||||
server.auth_type
|
||||
and server.auth_type != MCPAuth.none
|
||||
and server.auth_type != MCPAuth.aws_sigv4
|
||||
and not server.authentication_token
|
||||
):
|
||||
should_skip_health_check = True
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ class MCPAuth(str, enum.Enum):
|
|||
basic = "basic"
|
||||
authorization = "authorization"
|
||||
oauth2 = "oauth2"
|
||||
aws_sigv4 = "aws_sigv4"
|
||||
token = "token"
|
||||
|
||||
|
||||
|
|
@ -51,6 +52,7 @@ MCPAuthType = Optional[
|
|||
MCPAuth.basic,
|
||||
MCPAuth.authorization,
|
||||
MCPAuth.oauth2,
|
||||
MCPAuth.aws_sigv4,
|
||||
MCPAuth.token,
|
||||
]
|
||||
]
|
||||
|
|
|
|||
|
|
@ -48,6 +48,12 @@ class MCPServer(BaseModel):
|
|||
authorization_url: Optional[str] = None
|
||||
token_url: Optional[str] = None
|
||||
registration_url: Optional[str] = None
|
||||
# AWS SigV4 fields
|
||||
aws_access_key_id: Optional[str] = None
|
||||
aws_secret_access_key: Optional[str] = None
|
||||
aws_session_token: Optional[str] = None
|
||||
aws_region_name: Optional[str] = None
|
||||
aws_service_name: Optional[str] = None # defaults to "bedrock-agentcore"
|
||||
# Stdio-specific fields
|
||||
command: Optional[str] = None
|
||||
args: Optional[List[str]] = None
|
||||
|
|
|
|||
|
|
@ -0,0 +1,317 @@
|
|||
"""
|
||||
Tests for AWS SigV4 authentication in MCP client.
|
||||
|
||||
Tests the MCPSigV4Auth httpx.Auth subclass that enables per-request
|
||||
SigV4 signing for Bedrock AgentCore MCP servers.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.experimental_mcp_client.client import MCPSigV4Auth, MCPClient
|
||||
from litellm.types.mcp import MCPAuth, MCPTransport
|
||||
|
||||
|
||||
class TestMCPSigV4Auth:
|
||||
"""Unit tests for the MCPSigV4Auth class."""
|
||||
|
||||
def test_init_with_explicit_credentials(self):
|
||||
"""MCPSigV4Auth initializes with explicit AWS credentials."""
|
||||
auth = MCPSigV4Auth(
|
||||
aws_access_key_id="AKIAIOSFODNN7EXAMPLE",
|
||||
aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
aws_session_token="FwoGZXIvYXdzEBYaDH...",
|
||||
aws_region_name="us-east-1",
|
||||
aws_service_name="bedrock-agentcore",
|
||||
)
|
||||
assert auth.credentials is not None
|
||||
assert auth.credentials.access_key == "AKIAIOSFODNN7EXAMPLE"
|
||||
assert auth.credentials.secret_key == "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
|
||||
assert auth.credentials.token == "FwoGZXIvYXdzEBYaDH..."
|
||||
assert auth.region_name == "us-east-1"
|
||||
assert auth.service_name == "bedrock-agentcore"
|
||||
|
||||
def test_requires_request_body_flag(self):
|
||||
"""MCPSigV4Auth sets requires_request_body so httpx buffers the body before signing."""
|
||||
assert MCPSigV4Auth.requires_request_body is True
|
||||
|
||||
def test_init_defaults(self):
|
||||
"""MCPSigV4Auth uses correct defaults for region and service."""
|
||||
auth = MCPSigV4Auth(
|
||||
aws_access_key_id="AKIAIOSFODNN7EXAMPLE",
|
||||
aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
)
|
||||
assert auth.region_name == "us-east-1"
|
||||
assert auth.service_name == "bedrock-agentcore"
|
||||
|
||||
def test_init_with_resolved_env_values(self):
|
||||
"""MCPSigV4Auth works with pre-resolved values (os.environ/ is resolved at config load time)."""
|
||||
# Values arrive already resolved by ProxyConfig._check_for_os_environ_vars(),
|
||||
# so MCPSigV4Auth receives plain strings, not os.environ/ prefixed values.
|
||||
auth = MCPSigV4Auth(
|
||||
aws_access_key_id="RESOLVED_KEY_FROM_ENV",
|
||||
aws_secret_access_key="RESOLVED_SECRET_FROM_ENV",
|
||||
aws_region_name="us-west-2",
|
||||
)
|
||||
assert auth.credentials.access_key == "RESOLVED_KEY_FROM_ENV"
|
||||
assert auth.credentials.secret_key == "RESOLVED_SECRET_FROM_ENV"
|
||||
assert auth.region_name == "us-west-2"
|
||||
|
||||
def test_init_falls_back_to_boto_session(self):
|
||||
"""MCPSigV4Auth falls back to boto3 credential chain when no explicit creds."""
|
||||
mock_creds = MagicMock()
|
||||
mock_creds.access_key = "SESSION_KEY"
|
||||
mock_creds.secret_key = "SESSION_SECRET"
|
||||
|
||||
mock_session = MagicMock()
|
||||
mock_session.get_credentials.return_value = mock_creds
|
||||
|
||||
with patch("botocore.session.get_session", return_value=mock_session):
|
||||
auth = MCPSigV4Auth(
|
||||
aws_region_name="eu-west-1",
|
||||
aws_service_name="custom-service",
|
||||
)
|
||||
assert auth.credentials == mock_creds
|
||||
assert auth.region_name == "eu-west-1"
|
||||
assert auth.service_name == "custom-service"
|
||||
|
||||
def test_init_raises_when_no_credentials(self):
|
||||
"""MCPSigV4Auth raises ValueError when no credentials are available."""
|
||||
mock_session = MagicMock()
|
||||
mock_session.get_credentials.return_value = None
|
||||
|
||||
with patch("botocore.session.get_session", return_value=mock_session):
|
||||
with pytest.raises(ValueError, match="No AWS credentials found"):
|
||||
MCPSigV4Auth()
|
||||
|
||||
def test_auth_flow_signs_request(self):
|
||||
"""MCPSigV4Auth.auth_flow adds SigV4 headers to the request."""
|
||||
auth = MCPSigV4Auth(
|
||||
aws_access_key_id="AKIAIOSFODNN7EXAMPLE",
|
||||
aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
aws_region_name="us-east-1",
|
||||
aws_service_name="bedrock-agentcore",
|
||||
)
|
||||
|
||||
request = httpx.Request(
|
||||
method="POST",
|
||||
url="https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/test/invocations",
|
||||
headers={"Content-Type": "application/json"},
|
||||
content=b'{"jsonrpc":"2.0","method":"tools/list","id":1}',
|
||||
)
|
||||
|
||||
# Execute auth_flow generator
|
||||
flow = auth.auth_flow(request)
|
||||
signed_request = next(flow)
|
||||
|
||||
# Verify SigV4 headers were added
|
||||
assert "Authorization" in signed_request.headers
|
||||
assert "AWS4-HMAC-SHA256" in signed_request.headers["Authorization"]
|
||||
assert "x-amz-date" in signed_request.headers
|
||||
assert "bedrock-agentcore" in signed_request.headers["Authorization"]
|
||||
|
||||
def test_auth_flow_different_bodies_produce_different_signatures(self):
|
||||
"""Each request gets a unique signature based on its body."""
|
||||
auth = MCPSigV4Auth(
|
||||
aws_access_key_id="AKIAIOSFODNN7EXAMPLE",
|
||||
aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
aws_region_name="us-east-1",
|
||||
)
|
||||
|
||||
request1 = httpx.Request(
|
||||
method="POST",
|
||||
url="https://example.com/mcp",
|
||||
headers={"Content-Type": "application/json"},
|
||||
content=b'{"jsonrpc":"2.0","method":"tools/list","id":1}',
|
||||
)
|
||||
request2 = httpx.Request(
|
||||
method="POST",
|
||||
url="https://example.com/mcp",
|
||||
headers={"Content-Type": "application/json"},
|
||||
content=b'{"jsonrpc":"2.0","method":"tools/call","id":2,"params":{"name":"search"}}',
|
||||
)
|
||||
|
||||
signed1 = next(auth.auth_flow(request1))
|
||||
signed2 = next(auth.auth_flow(request2))
|
||||
|
||||
# Signatures must differ because body content differs
|
||||
assert signed1.headers["Authorization"] != signed2.headers["Authorization"]
|
||||
|
||||
def test_auth_flow_includes_security_token(self):
|
||||
"""SigV4 signing includes X-Amz-Security-Token when session token is present."""
|
||||
auth = MCPSigV4Auth(
|
||||
aws_access_key_id="AKIAIOSFODNN7EXAMPLE",
|
||||
aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
aws_session_token="SESSION_TOKEN_EXAMPLE",
|
||||
aws_region_name="us-east-1",
|
||||
)
|
||||
|
||||
request = httpx.Request(
|
||||
method="POST",
|
||||
url="https://example.com/mcp",
|
||||
headers={"Content-Type": "application/json"},
|
||||
content=b'{"jsonrpc":"2.0","method":"initialize","id":0}',
|
||||
)
|
||||
|
||||
signed_request = next(auth.auth_flow(request))
|
||||
assert "x-amz-security-token" in signed_request.headers
|
||||
|
||||
|
||||
class TestMCPClientSigV4Integration:
|
||||
"""Tests for MCPClient with SigV4 auth wired through."""
|
||||
|
||||
def test_mcp_client_stores_aws_auth(self):
|
||||
"""MCPClient stores the aws_auth parameter."""
|
||||
mock_auth = MagicMock(spec=httpx.Auth)
|
||||
client = MCPClient(
|
||||
server_url="https://example.com/mcp",
|
||||
transport_type=MCPTransport.http,
|
||||
auth_type=MCPAuth.aws_sigv4,
|
||||
aws_auth=mock_auth,
|
||||
)
|
||||
assert client._aws_auth is mock_auth
|
||||
|
||||
def test_mcp_client_factory_uses_aws_auth(self):
|
||||
"""The httpx client factory uses aws_auth when no explicit auth is passed."""
|
||||
mock_auth = MCPSigV4Auth(
|
||||
aws_access_key_id="AKIAIOSFODNN7EXAMPLE",
|
||||
aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
)
|
||||
client = MCPClient(
|
||||
server_url="https://example.com/mcp",
|
||||
transport_type=MCPTransport.http,
|
||||
aws_auth=mock_auth,
|
||||
)
|
||||
|
||||
factory = client._create_httpx_client_factory()
|
||||
httpx_client = factory(
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=httpx.Timeout(30.0),
|
||||
)
|
||||
|
||||
# Verify the auth object was actually wired into the httpx client
|
||||
assert httpx_client._auth is mock_auth
|
||||
|
||||
def test_mcp_client_factory_explicit_auth_takes_precedence(self):
|
||||
"""When explicit auth= is passed to the factory, it takes precedence over aws_auth."""
|
||||
aws_auth = MCPSigV4Auth(
|
||||
aws_access_key_id="AKIAIOSFODNN7EXAMPLE",
|
||||
aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
)
|
||||
explicit_auth = MagicMock(spec=httpx.Auth)
|
||||
|
||||
client = MCPClient(
|
||||
server_url="https://example.com/mcp",
|
||||
transport_type=MCPTransport.http,
|
||||
aws_auth=aws_auth,
|
||||
)
|
||||
|
||||
factory = client._create_httpx_client_factory()
|
||||
httpx_client = factory(
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=httpx.Timeout(30.0),
|
||||
auth=explicit_auth,
|
||||
)
|
||||
|
||||
# Explicit auth should win over aws_auth
|
||||
assert httpx_client._auth is explicit_auth
|
||||
|
||||
def test_mcp_client_factory_no_aws_auth(self):
|
||||
"""The httpx client factory works normally when no aws_auth is set."""
|
||||
client = MCPClient(
|
||||
server_url="https://example.com/mcp",
|
||||
transport_type=MCPTransport.http,
|
||||
)
|
||||
|
||||
factory = client._create_httpx_client_factory()
|
||||
httpx_client = factory(
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=httpx.Timeout(30.0),
|
||||
)
|
||||
# No auth should be set when aws_auth is not configured
|
||||
assert httpx_client._auth is None
|
||||
|
||||
|
||||
class TestMCPServerManagerSigV4:
|
||||
"""Tests for MCPServerManager config loading with SigV4."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_config_with_aws_sigv4(self):
|
||||
"""Config loading correctly parses aws_sigv4 auth type and AWS fields."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
|
||||
config = {
|
||||
"agentcore_tools": {
|
||||
"url": "https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/test/invocations",
|
||||
"transport": "http",
|
||||
"auth_type": "aws_sigv4",
|
||||
"aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
|
||||
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_service_name": "bedrock-agentcore",
|
||||
}
|
||||
}
|
||||
|
||||
manager = MCPServerManager()
|
||||
await manager.load_servers_from_config(config)
|
||||
|
||||
server = next(iter(manager.config_mcp_servers.values()))
|
||||
assert server.auth_type == MCPAuth.aws_sigv4
|
||||
assert server.aws_access_key_id == "AKIAIOSFODNN7EXAMPLE"
|
||||
assert server.aws_secret_access_key == "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
|
||||
assert server.aws_region_name == "us-east-1"
|
||||
assert server.aws_service_name == "bedrock-agentcore"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_mcp_client_with_sigv4(self):
|
||||
"""_create_mcp_client creates client with SigV4 auth when auth_type is aws_sigv4."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
server = MCPServer(
|
||||
server_id="test-sigv4",
|
||||
name="test_sigv4_server",
|
||||
server_name="test_sigv4",
|
||||
url="https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/test/invocations",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.aws_sigv4,
|
||||
aws_access_key_id="AKIAIOSFODNN7EXAMPLE",
|
||||
aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
aws_region_name="us-east-1",
|
||||
)
|
||||
|
||||
manager = MCPServerManager()
|
||||
client = await manager._create_mcp_client(server=server)
|
||||
|
||||
assert client.auth_type == MCPAuth.aws_sigv4
|
||||
assert client._aws_auth is not None
|
||||
assert isinstance(client._aws_auth, MCPSigV4Auth)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_mcp_client_without_sigv4(self):
|
||||
"""_create_mcp_client does not create SigV4 auth for non-SigV4 servers."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
server = MCPServer(
|
||||
server_id="test-bearer",
|
||||
name="test_bearer_server",
|
||||
server_name="test_bearer",
|
||||
url="https://example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.bearer_token,
|
||||
authentication_token="test-token",
|
||||
)
|
||||
|
||||
manager = MCPServerManager()
|
||||
client = await manager._create_mcp_client(server=server)
|
||||
|
||||
assert client._aws_auth is None
|
||||
Loading…
Add table
Reference in a new issue