mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
feat(mcp): add botocore SigV4 signer body for the v2 aws_sigv4 arm
HttpxSigV4Signer is the real SignerFactory body for the aws_sigv4 arm: it maps the typed credential source (StaticKeys / AssumeRole / Ambient) onto v1's MCPSigV4Auth, so the signed headers stay byte-identical to v1 by construction (the signer relocates into the v2 package when v1 is retired). Credentials resolve eagerly off the event loop (asyncio.to_thread) so an unassumable role or a missing ambient chain fails closed at build time rather than mid-request; STS connection errors map to upstream_unavailable, everything else to misconfigured. This is the seam-agnostic body only. Wiring it into the bridge and grafting it at _create_mcp_client's aws_auth seam (aws_sigv4 signs per request, so it cannot ride the header seam the other modes use) is a follow-up. Tests cover signing a real request (AWS4-HMAC-SHA256 + X-Amz-Date), the session-token path (X-Amz-Security-Token), region/service in the credential scope, and the error classification.
This commit is contained in:
parent
8bf3594c4a
commit
5dd99ca52f
2 changed files with 125 additions and 1 deletions
|
|
@ -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}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue