From a983177a4f3d94ed7d8b2edbfaae96a8b2f2e290 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 14 Jul 2026 17:12:20 +0000 Subject: [PATCH] feat(oidc): support AWS IAM Outbound Identity Federation as an OIDC provider --- litellm/secret_managers/main.py | 84 ++++++++++- .../test_secret_managers_main.py | 138 ++++++++++++++++++ 2 files changed, 221 insertions(+), 1 deletion(-) diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index fa55828a249..4dcfb70766c 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -1,7 +1,8 @@ import ast import os import traceback -from typing import Optional, Union +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Callable, Optional, Union import httpx @@ -14,6 +15,11 @@ from litellm.secret_managers.get_azure_ad_token_provider import ( ) from litellm.secret_managers.secret_manager_handler import get_secret_from_manager +if TYPE_CHECKING: + from types_boto3_sts.client import STSClient +else: + STSClient = object + oidc_cache = DualCache() _DEFAULT_OIDC_ALLOWED_CREDENTIAL_DIRS = ("/var/run/secrets", "/run/secrets") @@ -79,6 +85,80 @@ def _get_oidc_http_handler(timeout: Optional[httpx.Timeout] = None) -> HTTPHandl return HTTPHandler(timeout=timeout) +def _resolve_aws_region() -> str: + """ + Resolve the AWS region for the STS ``GetWebIdentityToken`` call. + + ``GetWebIdentityToken`` is served only from a regional STS endpoint, never + the global one, so a concrete region must be known before the client is + built. + """ + env_region = os.getenv("AWS_REGION") or os.getenv("AWS_DEFAULT_REGION") + if env_region: + return env_region + + import boto3 + + session_region = boto3.Session().region_name + if session_region: + return session_region + + raise ValueError( + "AWS OIDC provider requires a region. Set AWS_REGION or AWS_DEFAULT_REGION; " + "GetWebIdentityToken is served only from a regional STS endpoint, not the global one." + ) + + +def _get_aws_sts_client(region: str) -> "STSClient": + """ + Build an STS client pinned to the regional endpoint. Injected into + ``_get_aws_oidc_token`` so tests can supply a stub without real AWS calls. + """ + import boto3 + + return boto3.client( + "sts", + region_name=region, + endpoint_url=f"https://sts.{region}.amazonaws.com", + ) + + +def _aws_oidc_cache_ttl(expiration: datetime | None) -> int | None: + if expiration is None: + return None + remaining = int((expiration - datetime.now(timezone.utc)).total_seconds()) - 60 + return remaining if remaining > 0 else None + + +def _get_aws_oidc_token( + oidc_aud: str, + cache_key: str, + sts_client_factory: Callable[[str], "STSClient"] | None = None, +) -> str: + """ + Mint a signed JWT for ``oidc_aud`` via AWS IAM Outbound Identity Federation. + + RS256 is required (not ES384) because Entra federated credentials reject + ES384. The token is cached until just before its STS-reported expiry. + """ + cached_token = oidc_cache.get_cache(key=cache_key) + if isinstance(cached_token, str): + return cached_token + + build_client = sts_client_factory or _get_aws_sts_client + sts_client = build_client(_resolve_aws_region()) + response = sts_client.get_web_identity_token( + Audience=[oidc_aud], + SigningAlgorithm="RS256", + ) + oidc_token = response["WebIdentityToken"] + + ttl = _aws_oidc_cache_ttl(response.get("Expiration")) + if ttl is not None: + oidc_cache.set_cache(key=cache_key, value=oidc_token, ttl=ttl) + return oidc_token + + ######### Secret Manager ############################ # checks if user has passed in a secret manager client # if passed in then checks the secret there @@ -271,6 +351,8 @@ def get_secret( with open(token_file_path, "r") as f: oidc_token = f.read() return oidc_token + elif oidc_provider == "aws": + return _get_aws_oidc_token(oidc_aud=oidc_aud, cache_key=secret_name) else: raise ValueError("Unsupported OIDC provider") diff --git a/tests/test_litellm/secret_managers/test_secret_managers_main.py b/tests/test_litellm/secret_managers/test_secret_managers_main.py index b406231804b..47fd5a3efd7 100644 --- a/tests/test_litellm/secret_managers/test_secret_managers_main.py +++ b/tests/test_litellm/secret_managers/test_secret_managers_main.py @@ -255,6 +255,144 @@ def test_unsupported_oidc_provider(): get_secret(secret_name) +class _FakeSTSClient: + def __init__(self, token="aws_jwt", expiration=None): + self._token = token + self._expiration = expiration + self.calls = [] + + def get_web_identity_token(self, **kwargs): + self.calls.append(kwargs) + response = {"WebIdentityToken": self._token} + if self._expiration is not None: + response["Expiration"] = self._expiration + return response + + +def test_oidc_aws_success_signs_rs256_and_caches(): + """AWS OIDC mints a JWT via GetWebIdentityToken with RS256 and caches to expiry.""" + from datetime import datetime, timedelta, timezone + + from litellm.secret_managers.main import _get_aws_oidc_token + + expiration = datetime.now(timezone.utc) + timedelta(seconds=3600) + fake_client = _FakeSTSClient(token="aws_jwt", expiration=expiration) + factory = Mock(return_value=fake_client) + + mock_cache = Mock() + mock_cache.get_cache.return_value = None + + with patch("litellm.secret_managers.main.oidc_cache", mock_cache): + with patch("litellm.secret_managers.main._resolve_aws_region", return_value="eu-west-1"): + token = _get_aws_oidc_token( + oidc_aud="api://AzureADTokenExchange", + cache_key="oidc/aws/api://AzureADTokenExchange", + sts_client_factory=factory, + ) + + assert token == "aws_jwt" + factory.assert_called_once_with("eu-west-1") + assert fake_client.calls == [ + {"Audience": ["api://AzureADTokenExchange"], "SigningAlgorithm": "RS256"} + ] + (cache_kwargs,) = mock_cache.set_cache.call_args_list + assert cache_kwargs.kwargs["key"] == "oidc/aws/api://AzureADTokenExchange" + assert cache_kwargs.kwargs["value"] == "aws_jwt" + assert 3500 <= cache_kwargs.kwargs["ttl"] <= 3540 + + +def test_oidc_aws_uses_cache_without_calling_sts(): + from litellm.secret_managers.main import _get_aws_oidc_token + + factory = Mock(side_effect=AssertionError("STS should not be called on cache hit")) + mock_cache = Mock() + mock_cache.get_cache.return_value = "cached_jwt" + + with patch("litellm.secret_managers.main.oidc_cache", mock_cache): + token = _get_aws_oidc_token( + oidc_aud="api://AzureADTokenExchange", + cache_key="oidc/aws/api://AzureADTokenExchange", + sts_client_factory=factory, + ) + + assert token == "cached_jwt" + factory.assert_not_called() + mock_cache.set_cache.assert_not_called() + + +def test_oidc_aws_no_expiration_skips_caching(): + from litellm.secret_managers.main import _get_aws_oidc_token + + fake_client = _FakeSTSClient(token="aws_jwt", expiration=None) + mock_cache = Mock() + mock_cache.get_cache.return_value = None + + with patch("litellm.secret_managers.main.oidc_cache", mock_cache): + with patch("litellm.secret_managers.main._resolve_aws_region", return_value="us-east-1"): + token = _get_aws_oidc_token( + oidc_aud="aud", + cache_key="oidc/aws/aud", + sts_client_factory=Mock(return_value=fake_client), + ) + + assert token == "aws_jwt" + mock_cache.set_cache.assert_not_called() + + +def test_get_aws_sts_client_pins_regional_endpoint(): + from litellm.secret_managers.main import _get_aws_sts_client + + with patch("boto3.client") as mock_client: + _get_aws_sts_client("ap-southeast-2") + + mock_client.assert_called_once_with( + "sts", + region_name="ap-southeast-2", + endpoint_url="https://sts.ap-southeast-2.amazonaws.com", + ) + + +def test_resolve_aws_region_prefers_env(monkeypatch): + from litellm.secret_managers.main import _resolve_aws_region + + monkeypatch.setenv("AWS_REGION", "eu-central-1") + monkeypatch.delenv("AWS_DEFAULT_REGION", raising=False) + assert _resolve_aws_region() == "eu-central-1" + + +def test_resolve_aws_region_raises_when_unresolved(monkeypatch): + from litellm.secret_managers.main import _resolve_aws_region + + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_DEFAULT_REGION", raising=False) + + with patch("boto3.Session") as mock_session: + mock_session.return_value.region_name = None + with pytest.raises(ValueError, match="AWS OIDC provider requires a region"): + _resolve_aws_region() + + +def test_oidc_aws_get_secret_end_to_end(monkeypatch): + """End-to-end: get_secret('oidc/aws/') returns the STS-minted JWT.""" + monkeypatch.setenv("AWS_REGION", "us-east-1") + fake_client = _FakeSTSClient(token="aws_jwt", expiration=None) + + mock_cache = Mock() + mock_cache.get_cache.return_value = None + + with patch("litellm.secret_managers.main.oidc_cache", mock_cache): + with patch( + "litellm.secret_managers.main._get_aws_sts_client", + return_value=fake_client, + ): + result = get_secret("oidc/aws/api://AzureADTokenExchange") + + assert result == "aws_jwt" + assert fake_client.calls == [ + {"Audience": ["api://AzureADTokenExchange"], "SigningAlgorithm": "RS256"} + ] + + @pytest.mark.parametrize( ("raw", "expected"), [