diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index 57bc1d57ffd..600f69547d4 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -217,6 +217,7 @@ mcp_servers: | `bearer_token` | `Authorization: Bearer ` | | `basic` | `Authorization: Basic ` | | `authorization` | `Authorization: ` | + | `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//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" diff --git a/docs/my-website/docs/mcp_aws_sigv4.md b/docs/my-website/docs/mcp_aws_sigv4.md new file mode 100644 index 00000000000..e00cee4fd52 --- /dev/null +++ b/docs/my-website/docs/mcp_aws_sigv4.md @@ -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//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//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. diff --git a/docs/my-website/docs/providers/bedrock_agentcore.md b/docs/my-website/docs/providers/bedrock_agentcore.md index e3e352f7ab6..7802624fccd 100644 --- a/docs/my-website/docs/providers/bedrock_agentcore.md +++ b/docs/my-website/docs/providers/bedrock_agentcore.md @@ -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. ::: diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index b4a1337d54e..8b7aad29a34 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -614,6 +614,7 @@ const sidebars = { "mcp_usage", "mcp_openapi", "mcp_oauth", + "mcp_aws_sigv4", "mcp_public_internet", "mcp_semantic_filter", "mcp_control", diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index e4f241880d8..30a1ac20d0c 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -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, ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 19582ff63fb..2bdc47bf2c2 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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 diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 7b2ea820acc..33e55f9bed9 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -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, ] ] diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 2af2dcb88b5..511cfc958a2 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py new file mode 100644 index 00000000000..715bb8e8aee --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py @@ -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