fix(router): preserve aws session token and role params in deployment credential resolution

This commit is contained in:
mateo-berri 2026-08-10 19:46:41 -07:00
parent f72ddedf39
commit b63ba63655
3 changed files with 75 additions and 0 deletions

View file

@ -237,7 +237,14 @@ class CredentialLiteLLMParams(BaseModel):
## AWS BEDROCK / SAGEMAKER ##
aws_access_key_id: str | None = None
aws_secret_access_key: str | None = None
aws_session_token: str | None = None
aws_region_name: str | None = None
aws_session_name: str | None = None
aws_profile_name: str | None = None
aws_role_name: str | None = None
aws_web_identity_token: str | None = None
aws_sts_endpoint: str | None = None
aws_external_id: str | None = None
aws_bedrock_runtime_endpoint: str | None = None
aws_bedrock_project_id: str | None = None
s3_bucket_name: str | None = None

View file

@ -4062,6 +4062,46 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields():
assert credentials["aws_batch_role_arn"] == "arn:aws:iam::123:role/batch-role"
def test_get_deployment_credentials_with_provider_preserves_aws_auth_params():
"""
Test that get_deployment_credentials_with_provider preserves every AWS auth
selector (session token, assume-role, web identity, profile) so bedrock
files/batches deployments using temporary or role-based credentials do not
silently fall back to the server's ambient identity (#36155).
"""
aws_auth_params = {
"aws_access_key_id": "deployment-access-key",
"aws_secret_access_key": "deployment-secret",
"aws_session_token": "deployment-session-token",
"aws_region_name": "us-west-2",
"aws_session_name": "deployment-session",
"aws_profile_name": "deployment-profile",
"aws_role_name": "arn:aws:iam::123:role/deployment-role",
"aws_web_identity_token": "deployment-web-identity",
"aws_sts_endpoint": "https://sts.us-west-2.amazonaws.com",
"aws_external_id": "deployment-external-id",
}
router = litellm.Router(
model_list=[
{
"model_name": "bedrock-batch-model",
"litellm_params": {
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
**aws_auth_params,
},
}
],
)
credentials = router.get_deployment_credentials_with_provider(
model_id="bedrock-batch-model"
)
assert credentials is not None
for key, value in aws_auth_params.items():
assert credentials.get(key) == value, key
def _team_wildcard_model(api_key: str, model_id: str = "team-wildcard-id") -> dict:
return {
"model_name": f"model_name_team-1_{model_id}",

View file

@ -26789,10 +26789,24 @@ export interface components {
aws_bedrock_project_id?: string | null;
/** Aws Bedrock Runtime Endpoint */
aws_bedrock_runtime_endpoint?: string | null;
/** Aws External Id */
aws_external_id?: string | null;
/** Aws Profile Name */
aws_profile_name?: string | null;
/** Aws Region Name */
aws_region_name?: string | null;
/** Aws Role Name */
aws_role_name?: string | null;
/** Aws Secret Access Key */
aws_secret_access_key?: string | null;
/** Aws Session Name */
aws_session_name?: string | null;
/** Aws Session Token */
aws_session_token?: string | null;
/** Aws Sts Endpoint */
aws_sts_endpoint?: string | null;
/** Aws Web Identity Token */
aws_web_identity_token?: string | null;
/** Azure Ad Token */
azure_ad_token?: string | null;
/** Budget Duration */
@ -35467,10 +35481,24 @@ export interface components {
aws_bedrock_project_id?: string | null;
/** Aws Bedrock Runtime Endpoint */
aws_bedrock_runtime_endpoint?: string | null;
/** Aws External Id */
aws_external_id?: string | null;
/** Aws Profile Name */
aws_profile_name?: string | null;
/** Aws Region Name */
aws_region_name?: string | null;
/** Aws Role Name */
aws_role_name?: string | null;
/** Aws Secret Access Key */
aws_secret_access_key?: string | null;
/** Aws Session Name */
aws_session_name?: string | null;
/** Aws Session Token */
aws_session_token?: string | null;
/** Aws Sts Endpoint */
aws_sts_endpoint?: string | null;
/** Aws Web Identity Token */
aws_web_identity_token?: string | null;
/** Azure Ad Token */
azure_ad_token?: string | null;
/** Budget Duration */