ci: gate hardcoded commercial AWS partition literals and test us-gov endpoint builders

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

.github/workflows/test-code-quality.yml

@@ -146,6 +146,9 @@ jobs:
       - name: check_migrations_no_data_rewrites
         run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py
 
+      - name: check_aws_partition_hardcodes
+        run: uv run --no-sync python ./tests/code_coverage_tests/check_aws_partition_hardcodes.py
+
       - name: memory_test
         run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py
 

tests/code_coverage_tests/check_aws_partition_hardcodes.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+"""Ban hardcoded commercial-partition AWS hosts and ARN prefixes under `litellm/`.
+
+An endpoint or ARN built with a literal `amazonaws.com` or `arn:aws:` works in every
+commercial region and breaks only for GovCloud (`us-gov-*`, `arn:aws-us-gov:`) and
+China (`amazonaws.com.cn`, `arn:aws-cn:`) deployments, so the failure never shows up
+in CI or on a developer laptop. `litellm/litellm_core_utils/aws_partition.py` derives
+both from the region and is the only place those literals belong. Build hosts with
+`get_aws_dns_suffix(region)` and ARNs with `get_aws_arn_prefix(region)`.
+
+Every string constant in every `litellm/**/*.py` file is scanned, including the
+literal parts of f-strings and the strings inside `.format()` calls and
+concatenations. Docstrings and comments are not, since they never reach a request.
+`amazonaws.com.cn` passes because it is already the China partition.
+
+`ALLOWED` holds the (file, token) pairs that are text rather than a request target:
+a hosted logo, an IAM service principal, and hostnames quoted as examples inside
+error messages and field descriptions. An entry only covers that exact token in that
+exact file, so a second literal in an allowed file is still caught, and an entry
+whose token is gone fails the check so the set only shrinks.
+"""
+
+from __future__ import annotations
+
+import ast
+import re
+import sys
+from pathlib import Path
+from typing import Final, NamedTuple
+
+REPO_ROOT: Final = Path(__file__).resolve().parents[2]
+SCAN_ROOT: Final = REPO_ROOT / "litellm"
+PARTITION_HELPER: Final = "litellm/litellm_core_utils/aws_partition.py"
+
+COMMERCIAL_TOKEN: Final = re.compile(r"[A-Za-z0-9.-]*amazonaws\.com(?!\.cn)|arn:aws:[A-Za-z0-9:/_.*-]*")
+
+
+class Allowance(NamedTuple):
+    file: str
+    token: str
+
+
+ALLOWED: Final = frozenset(
+    {
+        Allowance("litellm/integrations/email_alerting.py", "litellm-listing.s3.amazonaws.com"),
+        Allowance("litellm/types/integrations/slack_alerting.py", "litellm-listing.s3.amazonaws.com"),
+        Allowance("litellm/rag/ingestion/bedrock_ingestion.py", "bedrock.amazonaws.com"),
+        Allowance(
+            "litellm/llms/bedrock/chat/agentcore/transformation.py",
+            "arn:aws:bedrock-agentcore:region:account:runtime/runtime_id",
+        ),
+        Allowance("litellm/llms/bedrock/search/transformation.py", ".amazonaws.com"),
+        Allowance(
+            "litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py",
+            "bucket.s3.amazonaws.com",
+        ),
+        Allowance("litellm/types/proxy/claude_code_endpoints.py", "bucket.s3.amazonaws.com"),
+    }
+)
+
+
+class Hit(NamedTuple):
+    file: str
+    line: int
+    token: str
+
+
+def _docstring_ids(tree: ast.Module) -> frozenset[int]:
+    return frozenset(
+        id(statement.value)
+        for node in ast.walk(tree)
+        if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef))
+        for statement in node.body
+        if isinstance(statement, ast.Expr)
+        and isinstance(statement.value, ast.Constant)
+        and isinstance(statement.value.value, str)
+    )
+
+
+def _hits_in_file(path: Path) -> tuple[Hit, ...]:
+    tree: Final = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
+    docstrings: Final = _docstring_ids(tree)
+    relative: Final = path.relative_to(REPO_ROOT).as_posix()
+    return tuple(
+        Hit(relative, node.lineno, match.group(0))
+        for node in ast.walk(tree)
+        if isinstance(node, ast.Constant) and isinstance(node.value, str) and id(node) not in docstrings
+        for match in COMMERCIAL_TOKEN.finditer(node.value)
+    )
+
+
+def find_hits(scan_root: Path) -> tuple[Hit, ...]:
+    return tuple(
+        hit
+        for path in sorted(scan_root.rglob("*.py"))
+        if path.relative_to(REPO_ROOT).as_posix() != PARTITION_HELPER
+        for hit in _hits_in_file(path)
+    )
+
+
+def main() -> int:
+    hits: Final = find_hits(SCAN_ROOT)
+    seen: Final = frozenset(Allowance(hit.file, hit.token) for hit in hits)
+    violations: Final = tuple(hit for hit in hits if Allowance(hit.file, hit.token) not in ALLOWED)
+    stale: Final = ALLOWED - seen
+    for hit in violations:
+        print(f"{hit.file}:{hit.line}: hardcoded commercial AWS partition literal {hit.token!r}")
+    for allowance in sorted(stale):
+        print(f"{allowance.file}: ALLOWED entry {allowance.token!r} no longer matches anything, remove it")
+    if violations or stale:
+        print(
+            "\nBuild AWS hosts with get_aws_dns_suffix(region) and ARNs with get_aws_arn_prefix(region) "
+            "from litellm/litellm_core_utils/aws_partition.py so GovCloud and China regions resolve."
+        )
+        return 1
+    print(f"No hardcoded commercial AWS partition literals outside {PARTITION_HELPER}.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/test_litellm/litellm_core_utils/test_aws_partition.py

@@ -1,9 +1,11 @@
 import ast
 from pathlib import Path
+from types import MappingProxyType
 from typing import Final
-from urllib.parse import urlparse
+from urllib.parse import unquote, urlparse
 
 import pytest
+from botocore.credentials import Credentials
 
 import litellm
 from litellm.integrations.s3_v2 import S3Logger
@@ -20,8 +22,20 @@ from litellm.llms.aws_polly.text_to_speech.transformation import AWSPollyTextToS
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
 from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
 from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
+from litellm.llms.bedrock.chat.invoke_agent.transformation import AmazonInvokeAgentConfig
 from litellm.llms.bedrock.common_utils import init_bedrock_client
+from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
+from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler
+from litellm.llms.bedrock.vector_stores.transformation import BedrockVectorStoreConfig
 from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig
+from litellm.llms.sagemaker.completion.handler import SagemakerLLM
+from litellm.proxy.auth.rds_iam_token import init_rds_client
+from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
+from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
+
+STATIC_AWS_CREDENTIALS: Final = MappingProxyType(
+    {"aws_access_key_id": "test-key", "aws_secret_access_key": "test-secret"}
+)
 
 
 @pytest.mark.parametrize(
@@ -106,6 +120,48 @@ def _s3_object_url(region: str) -> str:
     return logger._build_object_url("2025-01-01/key.json")
 
 
+def _bedrock_job_arn(region: str) -> str:
+    return f"{get_aws_arn_prefix(region)}bedrock:{region}:111122223333:model-invocation-job/abc1234567"
+
+
+def _bedrock_files_upload_url(region: str) -> str:
+    return BedrockFilesConfig().get_complete_file_url(
+        api_base=None,
+        api_key=None,
+        model="amazon.nova-pro-v1:0",
+        optional_params={},
+        litellm_params={"s3_bucket_name": "batch-bucket", "s3_region_name": region},
+        data={"file": ("batch.jsonl", b"{}", "application/jsonl"), "purpose": "batch"},
+    )
+
+
+def _bedrock_files_download_url(region: str) -> str:
+    return (
+        BedrockFilesConfig()
+        ._s3_request_target(optional_params={}, litellm_params={"s3_region_name": region})
+        .endpoint_url
+    )
+
+
+def _bedrock_guardrail_url(region: str) -> str:
+    guardrail = BedrockGuardrail(guardrailIdentifier="guardrail-id", guardrailVersion="1")
+    return guardrail._prepare_request(
+        credentials=Credentials("test-key", "test-secret"),
+        data={"source": "INPUT", "content": []},
+        optional_params={},
+        aws_region_name=region,
+    ).url
+
+
+def _secrets_manager_url(region: str) -> str:
+    endpoint_url, _headers, _body = AWSSecretsManagerV2(aws_region_name=region)._prepare_request(
+        action="GetSecretValue",
+        secret_name="my-secret",
+        optional_params=dict(STATIC_AWS_CREDENTIALS),
+    )
+    return endpoint_url
+
+
 ENDPOINT_BUILDERS: Final = {
     "bedrock_runtime_default": lambda region: BaseAWSLLM()._select_default_endpoint_url("runtime", region),
     "bedrock_agent_default": lambda region: BaseAWSLLM()._select_default_endpoint_url("agent", region),
@@ -124,6 +180,13 @@ ENDPOINT_BUILDERS: Final = {
         litellm_params={},
         data={"input_file_id": "s3://bucket/key.jsonl"},
     ),
+    "bedrock_batches_retrieve": lambda region: BedrockBatchesConfig().transform_retrieve_batch_request(
+        batch_id=_bedrock_job_arn(region),
+        optional_params=dict(STATIC_AWS_CREDENTIALS),
+        litellm_params={},
+    )["url"],
+    "bedrock_files_upload": _bedrock_files_upload_url,
+    "bedrock_files_download": _bedrock_files_download_url,
     "bedrock_agentcore_invoke": lambda region: AmazonAgentCoreConfig().get_complete_url(
         api_base=None,
         api_key=None,
@@ -131,6 +194,32 @@ ENDPOINT_BUILDERS: Final = {
         optional_params={},
         litellm_params={},
     ),
+    "bedrock_invoke_agent": lambda region: AmazonInvokeAgentConfig().get_complete_url(
+        api_base=None,
+        api_key=None,
+        model="agent/AGENT123/ALIAS456",
+        optional_params={"aws_region_name": region},
+        litellm_params={},
+    ),
+    "bedrock_guardrail_apply": _bedrock_guardrail_url,
+    "bedrock_rerank": lambda region: BedrockRerankHandler()._prepare_request(
+        model="amazon.rerank-v1:0",
+        api_base=None,
+        extra_headers=None,
+        data={"queries": [], "sources": []},
+        optional_params={"aws_region_name": region, **STATIC_AWS_CREDENTIALS},
+    )["endpoint_url"],
+    "bedrock_knowledgebase_search": lambda region: BedrockVectorStoreConfig().get_complete_url(
+        api_base=None, litellm_params={"aws_region_name": region}
+    ),
+    "secrets_manager": _secrets_manager_url,
+    "rds_iam_client": lambda region: (
+        init_rds_client(
+            aws_region_name=region,
+            aws_access_key_id="test-key",
+            aws_secret_access_key="test-secret",
+        ).meta.endpoint_url
+    ),
     "polly": lambda region: AWSPollyTextToSpeechConfig().get_complete_url(
         model="polly/neural",
         api_base=None,
@@ -152,6 +241,19 @@ ENDPOINT_BUILDERS: Final = {
         litellm_params={},
         stream=True,
     ),
+    "sagemaker_completion": lambda region: (
+        SagemakerLLM()
+        ._prepare_request(
+            credentials=Credentials("test-key", "test-secret"),
+            model="my-endpoint",
+            data={},
+            messages=[],
+            litellm_params={},
+            optional_params={},
+            aws_region_name=region,
+        )
+        .url
+    ),
     "s3_object_url": _s3_object_url,
 }
 
@@ -182,6 +284,18 @@ def test_every_endpoint_builder_keeps_amazonaws_com_outside_cn(builder_name: str
     assert hostname.endswith(".amazonaws.com"), url
 
 
+@pytest.mark.parametrize("region", ["us-gov-west-1", "us-gov-east-1"])
+@pytest.mark.parametrize("builder_name", sorted(ENDPOINT_BUILDERS))
+def test_every_endpoint_builder_respects_us_gov_partition(builder_name: str, region: str) -> None:
+    url = unquote(ENDPOINT_BUILDERS[builder_name](region))
+    hostname = urlparse(url).hostname
+    assert hostname is not None
+    assert hostname.endswith(f".{region}.amazonaws.com"), url
+    assert "arn:aws:" not in url, url
+    if "arn:" in url:
+        assert "arn:aws-us-gov:" in url, url
+
+
 def _fstring_literal_offenders(needle: str) -> list[str]:
     litellm_root = Path(litellm.__file__).parent
     return [