This commit is contained in:
Aryan Pardeshi 2026-08-27 01:10:25 +05:30 • committed by GitHub
commit 30e8a3ccb6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 35 additions and 2 deletions

View file

@ -400,10 +400,10 @@ MAX_STRING_LENGTH_PROMPT_IN_DB: Final = int(os.getenv("MAX_STRING_LENGTH_PROMPT_
BEDROCK_MAX_POLICY_SIZE: Final = 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: Final = 1000
BEDROCK_IAM_CACHE_MAX_ENTRIES: Final = int(os.getenv("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: Final = 64
BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES: Final = int(os.getenv("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: Final = 60

View file

@ -68,3 +68,36 @@ def _build_constant_env_var_map() -> dict[str, str]:
env_var_map[constant_name] = env_var_name
return env_var_map
@pytest.mark.parametrize(
"constant_name, override",
[
("BEDROCK_IAM_CACHE_MAX_ENTRIES", 4096),
("BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES", 128),
],
)
def test_bedrock_iam_cache_bounds_are_env_overridable(constant_name: str, override: int) -> None:
try:
with mock.patch.dict(os.environ, {constant_name: str(override)}):
reloaded = importlib.reload(constants)
assert getattr(reloaded, constant_name) == override
finally:
importlib.reload(constants)
@pytest.mark.parametrize(
"constant_name, expected_default",
[
("BEDROCK_IAM_CACHE_MAX_ENTRIES", 1000),
("BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES", 64),
],
)
def test_bedrock_iam_cache_bounds_keep_their_defaults(constant_name: str, expected_default: int) -> None:
try:
with mock.patch.dict(os.environ, {}, clear=False):
os.environ.pop(constant_name, None)
reloaded = importlib.reload(constants)
assert getattr(reloaded, constant_name) == expected_default
finally:
importlib.reload(constants)