diff --git a/litellm/proxy/gateway/mcp/outbound_credentials/resolver.py b/litellm/proxy/gateway/mcp/outbound_credentials/resolver.py index a88a0a72c1c..2e752fc9ff6 100644 --- a/litellm/proxy/gateway/mcp/outbound_credentials/resolver.py +++ b/litellm/proxy/gateway/mcp/outbound_credentials/resolver.py @@ -13,9 +13,9 @@ Implemented: `none`, `passthrough`, `api_key` (shared from config, or per-user / from the injected `CredentialStore`), `authorization_code` (per-user token read from the injected `TokenStore`, refreshed proactively via the `TokenRefresher`), and `client_credentials` (shared service-account token cached in the `ServiceTokenStore`, minted by the -`ClientCredentialsFetcher`), and `token_exchange` (RFC 8693 OBO: the caller's inbound token is -swapped for an upstream-audience token via the `TokenExchanger`, cached per-user). The -`aws_sigv4` mode is a typed stub that fails closed until its signer is injected. +`ClientCredentialsFetcher`), `token_exchange` (RFC 8693 OBO: the caller's inbound token is +swapped for an upstream-audience token via the `TokenExchanger`, cached per-user), and +`aws_sigv4` (per-request SigV4 signing via the `SignerFactory`). All seven arms are live. """ from __future__ import annotations @@ -31,13 +31,13 @@ from .clock import Clock from .credential_store import CredentialKey, CredentialStore from .httpx_auth import NoOpAuth, StaticHeaderAuth from .service_token_store import ServiceTokenKey, ServiceTokenStore +from .signer_factory import SignerFactory from .token_exchanger import TokenExchanger from .token_refresher import TokenRefresher from .token_store import StoredToken, TokenKey, TokenStore from .types import ( ApiKeyConfig, AuthorizationCodeConfig, - AuthSpecKind, AwsSigV4Config, Byok, ClientCredentialsConfig, @@ -67,6 +67,7 @@ class UpstreamCredentialProvider: service_token_store: ServiceTokenStore, client_credentials_fetcher: ClientCredentialsFetcher, token_exchanger: TokenExchanger, + signer_factory: SignerFactory, ) -> None: self._credential_store = credential_store self._token_store = token_store @@ -75,6 +76,7 @@ class UpstreamCredentialProvider: self._service_token_store = service_token_store self._client_credentials_fetcher = client_credentials_fetcher self._token_exchanger = token_exchanger + self._signer_factory = signer_factory async def resolve( self, subject: Subject, server: ServerSpec @@ -264,18 +266,13 @@ class UpstreamCredentialProvider: ) # best-effort; token is re-exchangeable return Ok(_bearer(minted)) - # --- arms awaiting their collaborators (typed stubs, fail closed) ---------------------- async def _aws_sigv4( self, subject: Subject, server: ServerSpec, config: AwsSigV4Config ) -> Result[httpx.Auth, CredError]: - return _todo(AuthSpecKind.aws_sigv4) + # SigV4 signs every request from the gateway's own AWS identity (never the caller's), + # so the whole arm is the signer the factory builds for this server's credential source. + return await self._signer_factory.build(config) def _bearer(token: StoredToken) -> StaticHeaderAuth: return StaticHeaderAuth(f"Bearer {token.access_token.get_secret_value()}") - - -def _todo(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]: - return Error( - CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet") - ) diff --git a/litellm/proxy/gateway/mcp/outbound_credentials/signer_factory.py b/litellm/proxy/gateway/mcp/outbound_credentials/signer_factory.py new file mode 100644 index 00000000000..5a313d58dcd --- /dev/null +++ b/litellm/proxy/gateway/mcp/outbound_credentials/signer_factory.py @@ -0,0 +1,28 @@ +"""The AWS SigV4 signer factory, isolated behind a port. + +`resolve()` delegates the whole `aws_sigv4` arm here: given the server's AWS config, return an +`httpx.Auth` that signs each outbound request with SigV4, or fail closed. The production body is +backed by botocore - it matches on the credential source (static keys / assumed role / ambient +chain), eagerly resolves the credentials so failures surface here rather than mid-request, and +lets botocore refresh temporary (STS) credentials at sign time. Faked in tests. +""" + +from __future__ import annotations + +from typing import Protocol + +import httpx + +from ..result import Result +from .types import AwsSigV4Config, CredError + + +class SignerFactory(Protocol): + """Builds the per-request SigV4 `httpx.Auth` for an AWS-hosted upstream, or fails closed. + + Returns `misconfigured` when the credentials cannot be resolved (an unassumable role, no + ambient credentials) and `upstream_unavailable` when STS cannot be reached. The gateway + signs with its own AWS identity; the caller's identity is never involved. + """ + + async def build(self, config: AwsSigV4Config) -> Result[httpx.Auth, CredError]: ... diff --git a/litellm/proxy/gateway/mcp/outbound_credentials/types.py b/litellm/proxy/gateway/mcp/outbound_credentials/types.py index 21cf6d1597f..36d0f0008ab 100644 --- a/litellm/proxy/gateway/mcp/outbound_credentials/types.py +++ b/litellm/proxy/gateway/mcp/outbound_credentials/types.py @@ -249,20 +249,48 @@ class NoneConfig(BaseModel): kind: Literal[AuthSpecKind.none] = AuthSpecKind.none +class StaticKeys(BaseModel): + """Long-lived AWS access keys configured on the server.""" + + model_config = ConfigDict(frozen=True) + source: Literal["static_keys"] = "static_keys" + access_key_id: str + secret_access_key: SecretStr + session_token: SecretStr | None = None + + +class AssumeRole(BaseModel): + """An IAM role the gateway assumes via STS for short-lived, auto-refreshed credentials.""" + + model_config = ConfigDict(frozen=True) + source: Literal["assume_role"] = "assume_role" + role_arn: str + session_name: str | None = None + external_id: str | None = None + + +class Ambient(BaseModel): + """The environment's default AWS credential chain (instance profile, IRSA, env vars).""" + + model_config = ConfigDict(frozen=True) + source: Literal["ambient"] = "ambient" + + +AwsCredentialSource = Annotated[ + StaticKeys | AssumeRole | Ambient, Field(discriminator="source") +] + + class AwsSigV4Config(BaseModel): - """AWS SigV4 per-request signing. Creds come from static keys, an assumed role, or - the ambient environment; that source is left loose here and tightened when the arm lands. - """ + """AWS SigV4 per-request signing for an AWS-hosted upstream (e.g. Bedrock AgentCore). The + gateway signs with its own AWS identity, never the caller's; `credentials` selects how that + identity is obtained, defaulting to the ambient credential chain.""" model_config = ConfigDict(frozen=True) kind: Literal[AuthSpecKind.aws_sigv4] = AuthSpecKind.aws_sigv4 region: str service: str = "bedrock-agentcore" - access_key_id: str | None = None - secret_access_key: SecretStr | None = None - session_token: SecretStr | None = None - role_arn: str | None = None - session_name: str | None = None + credentials: AwsCredentialSource = Ambient() AuthConfig = Annotated[ diff --git a/tests/mcp_tests/gateway/test_resolver.py b/tests/mcp_tests/gateway/test_resolver.py index 5ae04953ce3..590fc8d9098 100644 --- a/tests/mcp_tests/gateway/test_resolver.py +++ b/tests/mcp_tests/gateway/test_resolver.py @@ -32,6 +32,9 @@ from litellm.proxy.gateway.mcp.outbound_credentials.service_token_store import ( ServiceTokenKey, ServiceTokenStore, ) +from litellm.proxy.gateway.mcp.outbound_credentials.signer_factory import ( + SignerFactory, +) from litellm.proxy.gateway.mcp.outbound_credentials.token_exchanger import ( TokenExchanger, ) @@ -45,9 +48,12 @@ from litellm.proxy.gateway.mcp.outbound_credentials.token_store import ( TokenStore, ) from litellm.proxy.gateway.mcp.outbound_credentials.types import ( + Ambient, ApiKeyConfig, + AssumeRole, AuthorizationCodeConfig, AuthSpecKind, + AwsSigV4Config, Byok, ClientCredentialsConfig, CredError, @@ -56,6 +62,7 @@ from litellm.proxy.gateway.mcp.outbound_credentials.types import ( PerUserEnvVar, ServerSpec, SharedKey, + StaticKeys, Subject, TokenExchangeConfig, ) @@ -144,6 +151,16 @@ class FakeExchanger: return self._result +class FakeSignerFactory: + def __init__(self, result: Result[httpx.Auth, CredError]) -> None: + self._result = result + self.calls = 0 + + async def build(self, config: AwsSigV4Config) -> Result[httpx.Auth, CredError]: + self.calls += 1 + return self._result + + def _provider( *, credential_store: CredentialStore | None = None, @@ -153,6 +170,7 @@ def _provider( service_token_store: ServiceTokenStore | None = None, fetcher: ClientCredentialsFetcher | None = None, token_exchanger: TokenExchanger | None = None, + signer_factory: SignerFactory | None = None, ) -> UpstreamCredentialProvider: return UpstreamCredentialProvider( credential_store=credential_store or InMemoryCredentialStore(), @@ -165,6 +183,8 @@ def _provider( or FakeFetcher(Error(CredError.of_upstream_unavailable("unused"))), token_exchanger=token_exchanger or FakeExchanger(Error(CredError.of_upstream_unavailable("unused"))), + signer_factory=signer_factory + or FakeSignerFactory(Error(CredError.of_upstream_unavailable("unused"))), ) @@ -202,6 +222,10 @@ def _tx_cfg() -> TokenExchangeConfig: return TokenExchangeConfig() +def _aws_cfg() -> AwsSigV4Config: + return AwsSigV4Config(region="us-east-1") + + def _applied_headers(auth: httpx.Auth) -> httpx.Headers: request = httpx.Request("POST", "https://up.example/mcp") return next(auth.auth_flow(request)).headers @@ -336,19 +360,6 @@ async def test_self_contained_arms_never_read_the_inbound_token(): ).get("Authorization") -@pytest.mark.parametrize( - "config", - [ - {"kind": "aws_sigv4", "region": "us-east-1"}, - ], -) -async def test_unimplemented_arms_fail_closed(config: dict): - # Stub arms signal not_implemented (-> 501), not misconfigured (-> 500 operator error). - result = await PROVIDER.resolve(SUBJECT, _spec(config)) - assert isinstance(result, Error) - assert result.error.tag == "not_implemented" - - async def test_authorization_code_returns_a_valid_stored_token(): store = InMemoryTokenStore( { @@ -701,6 +712,94 @@ def test_token_exchange_config_allows_discovery_defaults(): assert isinstance(spec.config, TokenExchangeConfig) +async def test_aws_sigv4_returns_the_signer(): + signer = StaticHeaderAuth("AWS4-HMAC-SHA256 Credential=AKIA/...") + factory = FakeSignerFactory(Ok(signer)) + result = await _provider(signer_factory=factory).resolve(SUBJECT, _spec(_aws_cfg())) + assert isinstance(result, Ok) + assert _applied_headers(result.ok)["Authorization"].startswith("AWS4-HMAC-SHA256") + assert factory.calls == 1 + + +async def test_aws_sigv4_config_error_is_misconfigured(): + factory = FakeSignerFactory(Error(CredError.of_misconfigured("role not assumable"))) + result = await _provider(signer_factory=factory).resolve(SUBJECT, _spec(_aws_cfg())) + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" + + +async def test_aws_sigv4_sts_unreachable_is_upstream_unavailable(): + factory = FakeSignerFactory(Error(CredError.of_upstream_unavailable("STS timeout"))) + result = await _provider(signer_factory=factory).resolve(SUBJECT, _spec(_aws_cfg())) + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +async def test_aws_sigv4_never_reads_inbound_token(): + # Signs with the gateway's AWS identity; the caller bearer never changes the result. + factory = FakeSignerFactory(Ok(StaticHeaderAuth("AWS4-HMAC-SHA256 sig"))) + provider = _provider(signer_factory=factory) + without = await provider.resolve(SUBJECT, _spec(_aws_cfg())) + with_token = Subject(tenant_id="t1", subject_id="u1", inbound_token="leak-me") + present = await provider.resolve(with_token, _spec(_aws_cfg())) + assert isinstance(without, Ok) and isinstance(present, Ok) + header = _applied_headers(present.ok)["Authorization"] + assert _applied_headers(without.ok)["Authorization"] == header + assert "leak-me" not in header + + +def test_aws_sigv4_credentials_default_to_ambient(): + assert isinstance(AwsSigV4Config(region="us-east-1").credentials, Ambient) + + +def test_aws_sigv4_credential_sources_select_by_discriminator(): + static = _spec( + { + "kind": "aws_sigv4", + "region": "us-east-1", + "credentials": { + "source": "static_keys", + "access_key_id": "AKIA", + "secret_access_key": "shhh", + }, + } + ) + assert isinstance(static.config, AwsSigV4Config) + assert isinstance(static.config.credentials, StaticKeys) + role = _spec( + { + "kind": "aws_sigv4", + "region": "us-east-1", + "credentials": { + "source": "assume_role", + "role_arn": "arn:aws:iam::1:role/r", + }, + } + ) + assert isinstance(role.config, AwsSigV4Config) + assert isinstance(role.config.credentials, AssumeRole) + + +def test_aws_sigv4_rejects_static_keys_missing_secret(): + # The discriminated source makes illegal combos unrepresentable: static_keys needs a secret. + with pytest.raises(ValidationError): + _spec( + { + "kind": "aws_sigv4", + "region": "us-east-1", + "credentials": {"source": "static_keys", "access_key_id": "AKIA"}, + } + ) + + +def test_aws_sigv4_static_keys_secret_is_masked(): + config = AwsSigV4Config( + region="us-east-1", + credentials=StaticKeys(access_key_id="AKIA", secret_access_key="TOP-SECRET"), + ) + assert "TOP-SECRET" not in config.model_dump_json() + + async def test_api_key_per_user_store_error_is_upstream_unavailable(): # A store/DB outage on read is distinct from a miss: 503, not the 401/412 of "not found". provider = _provider(credential_store=FailingCredentialStore())