feat(oidc): support AWS IAM Outbound Identity Federation as an OIDC provider

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")
 

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/<aud>') 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"),
     [