scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(bedrock): cache AssumeRole credentials per attributed identity (#35467)

Browse source 
The explicit AssumeRole branch of BaseAWSLLM.get_credentials returned without
touching the process-wide IAM cache, so every model request issued a fresh
sts:AssumeRole, and on ECS/EC2 an uncached sts:GetCallerIdentity ahead of it.
Route the whole role branch through _get_or_set_cached_credentials with the TTL
_auth_with_aws_role already computed and discarded. The cache key is the same
aws_* argument snapshot the other flows use, taken before the session-name
default is filled in, so each aws_session_name keeps its own STS session and no
attributed identity can be served another's credentials.

Credential fetches now single-flight behind striped locks. Without that, a burst
of concurrent misses on one key each issued their own STS call, which is the
same thundering herd the cache exists to prevent, moved to the miss window.
This commit is contained in:
Yassin Kortam 2026-08-01 13:54:42 -07:00 • committed by GitHub
parent 669bfd6c60
commit 14dd98cd5f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 319 additions and 75 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/constants.py
View file

@ -357,6 +357,15 @@ NON_LLM_CONNECTION_TIMEOUT = int(
MAX_EXCEPTION_MESSAGE_LENGTH = int(os.getenv("MAX_EXCEPTION_MESSAGE_LENGTH", 2000))
MAX_STRING_LENGTH_PROMPT_IN_DB = int(os.getenv("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048))
BEDROCK_MAX_POLICY_SIZE = int(os.getenv("BEDROCK_MAX_POLICY_SIZE", 75))
# One entry per distinct AWS credential-argument set. Per-user cost attribution passes the attributed
# identity as aws_session_name, so this bounds how many attributed identities keep a cached STS session.
BEDROCK_IAM_CACHE_MAX_ENTRIES = 1000
# Single-flight lock stripes over that cache. Only keys landing on the same stripe wait for each
# other, so a burst of distinct identities still resolves its credentials in parallel.
BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES = 64
# Retire a cached STS credential this many seconds before AWS expires it, so a request that reads it
# still has a usable credential for the whole call.
STS_CREDENTIAL_EXPIRY_SAFETY_MARGIN_SECONDS = 60
BEDROCK_MIN_THINKING_BUDGET_TOKENS = int(os.getenv("BEDROCK_MIN_THINKING_BUDGET_TOKENS", 1024))
# Anthropic's Messages API rejects thinking.budget_tokens < 1024.
ANTHROPIC_MIN_THINKING_BUDGET_TOKENS = 1024

149
litellm/llms/bedrock/base_aws_llm.py
View file

@ -5,6 +5,7 @@ import os
import re
import urllib.parse
from datetime import datetime
from threading import Lock
from typing import (
TYPE_CHECKING,
Any,
@ -24,10 +25,14 @@ from pydantic import BaseModel, ValidationError
from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import (
BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES,
BEDROCK_IAM_CACHE_MAX_ENTRIES,
BEDROCK_INVOKE_PROVIDERS_LITERAL,
BEDROCK_MAX_POLICY_SIZE,
STS_CREDENTIAL_EXPIRY_SAFETY_MARGIN_SECONDS,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.secret_managers.main import get_secret, get_secret_str
@ -75,14 +80,28 @@ class AwsAuthError(Exception):
class BaseAWSLLM:
# Process-wide IAM credential cache (shared across instances — Bedrock passthrough is per-request).
# Storage is in-process memory only: default ``DualCache()`` has no Redis backend unless attached
# elsewhere. Entry TTL: static access-key + secret + region use ``_get_default_ttl_for_boto3_credentials``
# (~59 minutes); ambient env (``_auth_with_env_vars`` returns ``ttl=None``) uses ``InMemoryCache``'s
# ``default_ttl`` (600 seconds / 10 minutes); web identity STS credentials use
# ``_get_default_ttl_for_boto3_credentials`` (~59 minutes), keyed on all aws_* credential args
# plus ssl_verify. AssumeRole, profiles, and explicit session-token tuples are not cached — see
# ``get_credentials`` and ``_get_or_set_cached_credentials``.
_shared_iam_cache: ClassVar[DualCache] = DualCache()
# Storage is in-process memory only: no Redis backend unless attached elsewhere. Entry TTL: static
# access-key + secret + region use ``_get_default_ttl_for_boto3_credentials`` (~59 minutes); ambient
# env (``_auth_with_env_vars`` returns ``ttl=None``) uses ``InMemoryCache``'s ``default_ttl``
# (600 seconds / 10 minutes); web identity STS credentials use
# ``_get_default_ttl_for_boto3_credentials`` (~59 minutes); AssumeRole STS credentials expire with
# the STS session itself (Expiration minus a safety margin). All are keyed on all aws_* credential
# args plus ssl_verify, so ``aws_session_name`` scopes an entry to one attributed identity. Profiles
# and explicit session-token tuples are not cached — see ``get_credentials`` and
# ``_get_or_set_cached_credentials``. The bound is larger than ``InMemoryCache``'s default because
# per-user cost attribution puts one entry per attributed identity in this cache.
_shared_iam_cache: ClassVar[DualCache] = DualCache(
in_memory_cache=InMemoryCache(max_size_in_memory=BEDROCK_IAM_CACHE_MAX_ENTRIES)
)
# Striped single-flight locks over ``_shared_iam_cache``. Concurrent misses on one credential
# key would otherwise each issue their own STS call, which is the same thundering herd the cache
# exists to prevent, just moved to the miss window. Striping keeps distinct identities from
# serialising behind each other without a per-key registry that grows with the identity count.
# A cache hit holds its stripe only for the lookup itself.
_credential_fetch_locks: ClassVar[tuple[Lock, ...]] = tuple(
Lock() for _ in range(BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES)
)
def __init__(self) -> None:
self.iam_cache = BaseAWSLLM._shared_iam_cache
@ -140,19 +159,21 @@ class BaseAWSLLM:
Used for static access-key credentials, ambient credentials from
``_auth_with_env_vars`` (including when skipping AssumeRole because the runtime identity
already matches ``aws_role_name``), and web identity STS credentials (plain
non-refreshable ``Credentials`` cached ~59 min, inside the 3600s STS session).
already matches ``aws_role_name``), web identity STS credentials (plain
non-refreshable ``Credentials`` cached ~59 min, inside the 3600s STS session), and AssumeRole
STS credentials (cached for the lifetime of the STS session minus a safety margin).
AssumeRole, profiles, and explicit session-token tuples are not
cached here — shared ``Credentials`` / refresh state must not span logical sessions.
Profiles and explicit session-token tuples are not cached here — shared ``Credentials`` /
refresh state must not span logical sessions.
"""
cache_key = self.get_cache_key(credential_args)
_cached = self.iam_cache.get_cache(cache_key)
if _cached:
return _cached
credentials, ttl = credential_fetcher()
self.iam_cache.set_cache(cache_key, credentials, ttl=ttl)
return credentials
with self._credential_fetch_locks[hash(cache_key) % len(self._credential_fetch_locks)]:
_cached = self.iam_cache.get_cache(cache_key)
if _cached:
return _cached
credentials, ttl = credential_fetcher()
self.iam_cache.set_cache(cache_key, credentials, ttl=ttl)
return credentials
@staticmethod
def _is_auth_with_web_identity_token(
@ -269,8 +290,8 @@ class BaseAWSLLM:
# Credentials - boto3.Credentials
# cache ttl - Optional[int]. If None, the credentials are not cached. Some auth flows have no expiry time.
#
# iam_cache: static keys, ambient env (including skip-AssumeRole path), and web identity.
# Do not cache AssumeRole / profile / explicit session-token paths here.
# iam_cache: static keys, ambient env (including skip-AssumeRole path), web identity, and
# AssumeRole. Do not cache profile / explicit session-token paths here.
#########################################################
if self._is_auth_with_web_identity_token(
aws_web_identity_token,
@ -290,30 +311,20 @@ class BaseAWSLLM:
),
)
elif self._is_auth_with_aws_role(aws_role_name):
# Same role (IRSA/ECS/EC2): ambient creds via _get_or_set_cached_credentials like the
# default env branch; never pre-read cache (must run _is_already_running_as_role first).
if self._is_already_running_as_role(cast(str, aws_role_name), ssl_verify=ssl_verify):
verbose_logger.debug(
"Already running as target role %s, using ambient credentials",
aws_role_name,
)
return self._get_or_set_cached_credentials(args, self._auth_with_env_vars)
verbose_logger.debug("Using role assumption: calling _auth_with_aws_role")
# If aws_session_name is not provided, generate a default one
if aws_session_name is None:
aws_session_name = f"litellm-session-{int(datetime.now().timestamp())}"
credentials, _assume_ttl = self._auth_with_aws_role(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_role_name=cast(str, aws_role_name),
aws_session_name=aws_session_name,
aws_region_name=aws_region_name,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
ssl_verify=ssl_verify,
return self._get_or_set_cached_credentials(
args,
lambda: self._resolve_role_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_role_name=cast(str, aws_role_name),
aws_session_name=aws_session_name,
aws_region_name=aws_region_name,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
ssl_verify=ssl_verify,
),
)
return credentials
elif self._is_auth_with_aws_profile(aws_profile_name):
credentials, _cache_ttl = self._auth_with_aws_profile(cast(str, aws_profile_name))
@ -1046,7 +1057,11 @@ class BaseAWSLLM:
return sts_client.assume_role(**assume_role_params)
def _extract_credentials_and_ttl(self, sts_response: dict) -> Tuple[Credentials, Optional[int]]:
"""Extract credentials and TTL from STS response."""
"""Extract credentials and TTL from STS response.
The TTL carries the same safety margin as the non-IRSA assume path, so a cached entry is
never handed out close enough to expiry to die mid-request.
"""
from botocore.credentials import Credentials
sts_credentials = sts_response["Credentials"]
@ -1057,10 +1072,52 @@ class BaseAWSLLM:
)
expiration_time = sts_credentials["Expiration"]
ttl = int((expiration_time - datetime.now(expiration_time.tzinfo)).total_seconds())
ttl = int(
(expiration_time - datetime.now(expiration_time.tzinfo)).total_seconds()
- STS_CREDENTIAL_EXPIRY_SAFETY_MARGIN_SECONDS
)
return credentials, ttl
def _resolve_role_credentials(
self,
aws_access_key_id: str | None,
aws_secret_access_key: str | None,
aws_session_token: str | None,
aws_role_name: str,
aws_session_name: str | None,
aws_region_name: str | None,
aws_sts_endpoint: str | None,
aws_external_id: str | None,
ssl_verify: bool | str | None,
) -> tuple[Credentials, int | None]:
"""
Resolve credentials for a target role, either from the ambient identity or via sts:AssumeRole.
Both the ``sts:GetCallerIdentity`` probe and the assume itself run here, so a cache hit on the
caller's key skips both. ``aws_session_name`` defaults inside this fetcher rather than in
``get_credentials`` so the cache key stays stable when the caller does not supply one.
"""
if self._is_already_running_as_role(aws_role_name, ssl_verify=ssl_verify):
verbose_logger.debug(
"Already running as target role %s, using ambient credentials",
aws_role_name,
)
return self._auth_with_env_vars()
verbose_logger.debug("Using role assumption: calling _auth_with_aws_role")
return self._auth_with_aws_role(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_role_name=aws_role_name,
aws_session_name=aws_session_name or f"litellm-session-{int(datetime.now().timestamp())}",
aws_region_name=aws_region_name,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
ssl_verify=ssl_verify,
)
@tracer.wrap()
def _auth_with_aws_role(
self,
@ -1192,7 +1249,7 @@ class BaseAWSLLM:
sts_expiry = sts_credentials["Expiration"]
# Convert to timezone-aware datetime for comparison
current_time = datetime.now(sts_expiry.tzinfo)
sts_ttl = (sts_expiry - current_time).total_seconds() - 60
sts_ttl = (sts_expiry - current_time).total_seconds() - STS_CREDENTIAL_EXPIRY_SAFETY_MARGIN_SECONDS
return credentials, sts_ttl
@tracer.wrap()

236
tests/test_litellm/llms/bedrock/test_base_aws_llm.py
View file

@ -1,6 +1,8 @@
import json
import os
import sys
import threading
import time
import pytest
from fastapi.testclient import TestClient
@ -933,7 +935,7 @@ def test_role_assumption_without_session_name():
call_args = mock_sts_client.assume_role.call_args
assert call_args[1]["RoleSessionName"] == "my-custom-session"
# Test case 3: AssumeRole is not stored in iam_cache; identical calls each invoke STS.
# Test case 3: identical AssumeRole args reuse the cached STS session instead of re-assuming.
BaseAWSLLM._shared_iam_cache.flush_cache()
mock_sts_client.reset_mock()
with patch("boto3.client", return_value=mock_sts_client):
@ -945,49 +947,218 @@ def test_role_assumption_without_session_name():
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
)
assert mock_sts_client.assume_role.call_count == 2
assert mock_sts_client.assume_role.call_count == 1
assert credentials1.access_key == credentials2.access_key
def test_assume_role_path_does_not_use_process_iam_cache():
"""AssumeRole credentials are not cached; each get_credentials repeats STS AssumeRole."""
base_aws_llm = BaseAWSLLM()
def _assume_role_sts_mock(access_key: str = "assumed-access-key") -> MagicMock:
"""STS client mock whose caller identity is a different role, forcing the AssumeRole path."""
mock_sts_client = MagicMock()
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_client.assume_role.return_value = {
"Credentials": {
"AccessKeyId": "assumed-access-key",
"AccessKeyId": access_key,
"SecretAccessKey": "assumed-secret-key",
"SessionToken": "assumed-session-token",
"Expiration": mock_expiry,
"Expiration": datetime.now(timezone.utc) + timedelta(hours=1),
}
}
mock_sts_client.get_caller_identity.return_value = {
"Arn": "arn:aws:sts::111111111111:assumed-role/SomeOtherRole/session-name",
}
return mock_sts_client
role_arn = "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
env_without_irsa = {
def _env_without_irsa() -> dict:
return {
k: v
for k, v in os.environ.items()
if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")
}
with patch.dict(os.environ, env_without_irsa, clear=True):
def test_assume_role_path_uses_process_iam_cache():
"""Repeat calls with identical AssumeRole args make no further STS calls of either kind."""
base_aws_llm = BaseAWSLLM()
mock_sts_client = _assume_role_sts_mock()
role_arn = "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
with patch.dict(os.environ, _env_without_irsa(), clear=True):
with patch("boto3.client", return_value=mock_sts_client):
base_aws_llm.get_credentials(aws_role_name=role_arn)
first = base_aws_llm.get_credentials(
aws_role_name=role_arn, aws_session_name="attributed-user-1"
)
mock_sts_client.get_caller_identity.reset_mock()
base_aws_llm.get_credentials(aws_role_name=role_arn)
second = base_aws_llm.get_credentials(
aws_role_name=role_arn, aws_session_name="attributed-user-1"
)
mock_sts_client.get_caller_identity.assert_called()
assert mock_sts_client.assume_role.call_count == 2
mock_sts_client.get_caller_identity.assert_not_called()
assert mock_sts_client.assume_role.call_count == 1
assert first.access_key == second.access_key
assert first.token == second.token
def test_assume_role_cache_is_scoped_per_session_name():
"""Each attributed identity gets its own STS session; no user is served another user's creds."""
base_aws_llm = BaseAWSLLM()
mock_sts_client = _assume_role_sts_mock()
mock_sts_client.assume_role.side_effect = [
{
"Credentials": {
"AccessKeyId": f"assumed-access-key-{session_name}",
"SecretAccessKey": "assumed-secret-key",
"SessionToken": f"assumed-session-token-{session_name}",
"Expiration": datetime.now(timezone.utc) + timedelta(hours=1),
}
}
for session_name in ("attributed-user-1", "attributed-user-2")
]
role_arn = "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
with patch.dict(os.environ, _env_without_irsa(), clear=True):
with patch("boto3.client", return_value=mock_sts_client):
user_1 = base_aws_llm.get_credentials(
aws_role_name=role_arn, aws_session_name="attributed-user-1"
)
user_2 = base_aws_llm.get_credentials(
aws_role_name=role_arn, aws_session_name="attributed-user-2"
)
user_1_again = base_aws_llm.get_credentials(
aws_role_name=role_arn, aws_session_name="attributed-user-1"
)
assert mock_sts_client.assume_role.call_count == 2
session_names = [
call.kwargs["RoleSessionName"] for call in mock_sts_client.assume_role.call_args_list
]
assert session_names == ["attributed-user-1", "attributed-user-2"]
assert user_1.access_key == "assumed-access-key-attributed-user-1"
assert user_2.access_key == "assumed-access-key-attributed-user-2"
assert user_1_again.access_key == user_1.access_key
def test_assume_role_cache_entry_expires_with_the_sts_session():
"""The cache entry is written with the STS lifetime, not an indefinite or default TTL."""
base_aws_llm = BaseAWSLLM()
mock_sts_client = _assume_role_sts_mock()
role_arn = "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
with patch.dict(os.environ, _env_without_irsa(), clear=True):
with patch("boto3.client", return_value=mock_sts_client):
base_aws_llm.get_credentials(
aws_role_name=role_arn, aws_session_name="attributed-user-1"
)
in_memory_cache = BaseAWSLLM._shared_iam_cache.in_memory_cache
assert len(in_memory_cache.ttl_dict) == 1
remaining_ttl = list(in_memory_cache.ttl_dict.values())[0] - time.time()
# 1h STS session minus the 60s safety margin, minus test execution time
assert 3400 < remaining_ttl <= 3540
def test_assume_role_credentials_expired_in_cache_trigger_a_fresh_sts_call():
"""Once the STS session lapses the next call re-assumes rather than serving dead credentials."""
base_aws_llm = BaseAWSLLM()
mock_sts_client = _assume_role_sts_mock()
role_arn = "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
with patch.dict(os.environ, _env_without_irsa(), clear=True):
with patch("boto3.client", return_value=mock_sts_client):
base_aws_llm.get_credentials(
aws_role_name=role_arn, aws_session_name="attributed-user-1"
)
in_memory_cache = BaseAWSLLM._shared_iam_cache.in_memory_cache
expired_key = list(in_memory_cache.ttl_dict.keys())[0]
in_memory_cache.ttl_dict[expired_key] = time.time() - 1
base_aws_llm.get_credentials(
aws_role_name=role_arn, aws_session_name="attributed-user-1"
)
assert mock_sts_client.assume_role.call_count == 2
def test_concurrent_cold_cache_calls_for_one_identity_make_a_single_sts_call():
"""Concurrent misses on one identity single-flight instead of each issuing its own AssumeRole."""
base_aws_llm = BaseAWSLLM()
role_arn = "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
calling_threads = set()
threads_lock = threading.Lock()
def slow_assume_role(**kwargs):
with threads_lock:
calling_threads.add(threading.get_ident())
# Hold the fetch open so every other worker is inside get_credentials while this one runs;
# without that overlap the test would pass against unsynchronised code.
time.sleep(0.2)
return {
"Credentials": {
"AccessKeyId": "assumed-access-key",
"SecretAccessKey": "assumed-secret-key",
"SessionToken": "assumed-session-token",
"Expiration": datetime.now(timezone.utc) + timedelta(hours=1),
}
}
mock_sts_client = _assume_role_sts_mock()
mock_sts_client.assume_role.side_effect = slow_assume_role
barrier = threading.Barrier(8)
resolved = []
def resolve():
barrier.wait()
resolved.append(
base_aws_llm.get_credentials(aws_role_name=role_arn, aws_session_name="attribution-1")
)
with patch.dict(os.environ, _env_without_irsa(), clear=True):
with patch("boto3.client", return_value=mock_sts_client):
workers = [threading.Thread(target=resolve) for _ in range(8)]
for worker in workers:
worker.start()
for worker in workers:
worker.join(timeout=30)
# The harness has to have produced real concurrency, or the assertion below proves nothing
assert len(resolved) == 8
assert len({id(credentials) for credentials in resolved}) == 1
assert mock_sts_client.assume_role.call_count == 1
assert len(calling_threads) == 1
def test_ambient_identity_matching_target_role_is_cached_without_repeating_the_probe():
"""The skip-AssumeRole branch caches too, so GetCallerIdentity stops running per request."""
base_aws_llm = BaseAWSLLM()
role_arn = "arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole"
mock_sts_client = MagicMock()
mock_sts_client.get_caller_identity.return_value = {
"Arn": "arn:aws:sts::2222222222222:assumed-role/LitellmEvalBedrockRole/session-name",
}
mock_ambient_credentials = MagicMock()
mock_ambient_credentials.access_key = "ambient-access-key"
with patch.dict(os.environ, _env_without_irsa(), clear=True):
with patch("boto3.client", return_value=mock_sts_client):
with patch.object(
base_aws_llm,
"_auth_with_env_vars",
return_value=(mock_ambient_credentials, None),
) as mock_env_auth:
first = base_aws_llm.get_credentials(aws_role_name=role_arn)
second = base_aws_llm.get_credentials(aws_role_name=role_arn)
assert mock_sts_client.get_caller_identity.call_count == 1
assert mock_sts_client.assume_role.call_count == 0
mock_env_auth.assert_called_once()
assert first.access_key == second.access_key == "ambient-access-key"
def test_cache_keys_are_different_for_different_roles():
@ -1961,7 +2132,10 @@ def test_auth_with_aws_role_irsa_environment():
assert creds.access_key == "irsa-access-key"
assert creds.secret_key == "irsa-secret-key"
assert creds.token == "irsa-session-token"
assert ttl > 0 # TTL should be positive
# 1h STS session minus the safety margin, so a cached entry retires before AWS
# expires the credential
assert ttl is not None
assert 3400 < ttl <= 3540
finally:
# Clean up the temporary file
os.unlink(token_file)
@ -2336,20 +2510,19 @@ def test_get_credentials_ecs_same_role_skips_assume_role():
aws_region_name="us-east-1",
)
# Each get_credentials must check identity first; second call still checks before
# taking the iam_cache hit (no pre-peek that bypasses _is_already_running_as_role).
assert mock_already_running.call_count == 2
# Cached env resolution: second call hits iam_cache, not _auth_with_env_vars again.
# The identity probe and the env resolution are both behind the iam_cache, so the
# second call repeats neither.
mock_already_running.assert_called_once()
mock_env_auth.assert_called_once()
mock_role_auth.assert_not_called()
assert credentials.access_key == "ecs-access-key"
def test_get_credentials_role_second_call_not_same_role_uses_assume_not_env_cache():
def test_get_credentials_role_reevaluates_identity_once_the_cached_entry_lapses():
"""
First request: already target role -> env path fills iam_cache.
Second request (e.g. identity changed): not same role -> AssumeRole path; must not reuse
cached env resolution from the first call.
Once that entry expires and the identity no longer matches, the next request must take the
AssumeRole path rather than serving the cached env resolution.
"""
base_aws_llm = BaseAWSLLM()
@ -2384,6 +2557,11 @@ def test_get_credentials_role_second_call_not_same_role_uses_assume_not_env_cach
aws_role_name=role_arn,
aws_region_name="us-east-1",
)
in_memory_cache = BaseAWSLLM._shared_iam_cache.in_memory_cache
lapsed_key = list(in_memory_cache.ttl_dict.keys())[0]
in_memory_cache.ttl_dict[lapsed_key] = time.time() - 1
second = base_aws_llm.get_credentials(
aws_role_name=role_arn,
aws_region_name="us-east-1",