diff --git a/litellm/proxy/_experimental/mcp_server/v2_port_bodies.py b/litellm/proxy/_experimental/mcp_server/v2_port_bodies.py index 36257f0a82c..f1347978ca1 100644 --- a/litellm/proxy/_experimental/mcp_server/v2_port_bodies.py +++ b/litellm/proxy/_experimental/mcp_server/v2_port_bodies.py @@ -7,9 +7,11 @@ invariant; the graft composition root injects them and v2 unit tests fake them. from __future__ import annotations +import asyncio from datetime import timedelta from typing import Optional +import httpx from pydantic import BaseModel, ConfigDict, SecretStr, ValidationError from litellm.constants import MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL @@ -17,8 +19,11 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy.gateway.mcp.outbound_credentials.clock import Clock, SystemClock from litellm.proxy.gateway.mcp.outbound_credentials.token_store import StoredToken from litellm.proxy.gateway.mcp.outbound_credentials.types import ( + AssumeRole, + AwsSigV4Config, ClientCredentialsConfig, CredError, + StaticKeys, ) from litellm.proxy.gateway.mcp.result import Error, Ok, Result from litellm.types.llms.custom_http import httpxSpecialProvider @@ -107,3 +112,61 @@ class HttpxClientCredentialsFetcher: expires_at=self._clock.now() + ttl, ) ) + + +class HttpxSigV4Signer: + """AWS SigV4 signer for the `aws_sigv4` arm, backed by botocore. + + Phase-1 graft body: maps the typed credential source onto v1's proven `MCPSigV4Auth`, so the + signed headers stay byte-identical to v1 (the signer relocates into the v2 package when v1 is + retired). Credentials resolve eagerly at build time -- an unassumable role or a missing + ambient chain fails closed here, not mid-request -- and botocore refreshes temporary STS + credentials at sign time. + """ + + async def build(self, config: AwsSigV4Config) -> Result[httpx.Auth, CredError]: + try: + auth = await asyncio.to_thread(_build_sigv4_auth, config) + except Exception as e: # classified into a CredError below + return Error(_classify_sigv4_error(e)) + return Ok(auth) + + +def _build_sigv4_auth(config: AwsSigV4Config) -> httpx.Auth: + from litellm.experimental_mcp_client.client import MCPSigV4Auth + + creds = config.credentials + if isinstance(creds, StaticKeys): + return MCPSigV4Auth( + aws_access_key_id=creds.access_key_id, + aws_secret_access_key=creds.secret_access_key.get_secret_value(), + aws_session_token=( + creds.session_token.get_secret_value() if creds.session_token else None + ), + aws_region_name=config.region, + aws_service_name=config.service, + ) + if isinstance(creds, AssumeRole): + return MCPSigV4Auth( + aws_role_name=creds.role_arn, + aws_session_name=creds.session_name, + aws_region_name=config.region, + aws_service_name=config.service, + ) + return MCPSigV4Auth(aws_region_name=config.region, aws_service_name=config.service) + + +def _classify_sigv4_error(error: Exception) -> CredError: + # botocore is an optional dependency without type stubs; match its connection-error classes + # by name rather than importing it just to isinstance-check. + if type(error).__name__ in ( + "EndpointConnectionError", + "ConnectTimeoutError", + "ConnectionError", + ): + return CredError.of_upstream_unavailable( + f"STS unreachable while resolving aws_sigv4 credentials: {error}" + ) + return CredError.of_misconfigured( + f"aws_sigv4 credentials could not be resolved: {error}" + ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py index 84448298d4c..f3d1dd11412 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_port_bodies.py @@ -10,8 +10,14 @@ from pydantic import SecretStr from litellm.proxy._experimental.mcp_server import v2_port_bodies from litellm.proxy._experimental.mcp_server.v2_port_bodies import ( HttpxClientCredentialsFetcher, + HttpxSigV4Signer, + _classify_sigv4_error, +) +from litellm.proxy.gateway.mcp.outbound_credentials.types import ( + AwsSigV4Config, + ClientCredentialsConfig, + StaticKeys, ) -from litellm.proxy.gateway.mcp.outbound_credentials.types import ClientCredentialsConfig from litellm.proxy.gateway.mcp.result import Error, Ok pytestmark = pytest.mark.asyncio @@ -101,3 +107,58 @@ async def test_fetch_missing_access_token_is_misconfigured(monkeypatch): result = await HttpxClientCredentialsFetcher().fetch(_cfg()) assert isinstance(result, Error) assert result.error.tag == "misconfigured" + + +def _sigv4_cfg(session_token=None, region="us-east-1", service="bedrock-agentcore"): + return AwsSigV4Config( + region=region, + service=service, + credentials=StaticKeys( + access_key_id="AKIATEST", + secret_access_key=SecretStr("secret"), + session_token=SecretStr(session_token) if session_token else None, + ), + ) + + +async def test_sigv4_static_keys_signs_request(): + result = await HttpxSigV4Signer().build(_sigv4_cfg()) + assert isinstance(result, Ok) + req = httpx.Request( + "POST", "https://svc.us-east-1.amazonaws.com/mcp", content=b"{}" + ) + signed = next(result.ok.auth_flow(req)) + assert signed.headers["Authorization"].startswith( + "AWS4-HMAC-SHA256 Credential=AKIATEST/" + ) + assert "X-Amz-Date" in signed.headers + assert "X-Amz-Security-Token" not in signed.headers + + +async def test_sigv4_session_token_adds_security_token(): + result = await HttpxSigV4Signer().build(_sigv4_cfg(session_token="tok")) + assert isinstance(result, Ok) + req = httpx.Request("GET", "https://svc.us-east-1.amazonaws.com/mcp") + signed = next(result.ok.auth_flow(req)) + assert "X-Amz-Security-Token" in signed.headers + + +async def test_sigv4_region_and_service_in_credential_scope(): + result = await HttpxSigV4Signer().build(_sigv4_cfg(region="eu-west-1")) + assert isinstance(result, Ok) + req = httpx.Request("GET", "https://svc.eu-west-1.amazonaws.com/mcp") + signed = next(result.ok.auth_flow(req)) + assert ( + "/eu-west-1/bedrock-agentcore/aws4_request" in signed.headers["Authorization"] + ) + + +async def test_classify_sigv4_connection_error_is_upstream_unavailable(): + from botocore.exceptions import EndpointConnectionError + + err = EndpointConnectionError(endpoint_url="https://sts.amazonaws.com") + assert _classify_sigv4_error(err).tag == "upstream_unavailable" + + +async def test_classify_sigv4_other_error_is_misconfigured(): + assert _classify_sigv4_error(ValueError("no creds")).tag == "misconfigured"