From d793c1621c27ab9ebdc34c9086614404b8223dd8 Mon Sep 17 00:00:00 2001 From: JasonZhang1993 Date: Mon, 28 Sep 2026 14:12:03 +0000 Subject: [PATCH] feat(endpoints): support AWS China (aws-cn) partition Derive endpoint DNS suffixes from botocore partition data instead of hardcoding .amazonaws.com, so data/control/gateway endpoints resolve to amazonaws.com.cn in cn-north-1 / cn-northwest-1. Commercial partitions are unchanged. - _utils/endpoints.py: partition-aware DNS suffix via a cached public botocore EndpointResolver; warns (rather than silently falling back) on unknown regions; exposes known_partitions() - runtime/utils.py: is_valid_partition backed by botocore's partition list - runtime/a2a.py: build_runtime_url derives the suffix (region-based contract) - payments/{client,manager}.py, config_bundle/client.py, batch runner: resolve endpoints natively (partition-correct) while still honouring the BEDROCK_AGENTCORE_{CP,DP}_ENDPOINT override via conditional endpoint_url - .pre-commit-config.yaml: pygrep hook forbidding a hardcoded .amazonaws.com suffix immediately after an interpolated value - tests: CN endpoint-builder + aws-cn ARN coverage; override on/off cases --- .pre-commit-config.yaml | 14 +++++ src/bedrock_agentcore/_utils/endpoints.py | 59 ++++++++++++++++++- src/bedrock_agentcore/config_bundle/client.py | 14 +++-- .../runner/batch/batch_evaluation_runner.py | 13 ++-- src/bedrock_agentcore/payments/client.py | 16 ++--- src/bedrock_agentcore/payments/manager.py | 16 ++--- src/bedrock_agentcore/runtime/a2a.py | 10 +++- src/bedrock_agentcore/runtime/utils.py | 11 +++- .../config_bundle/test_client.py | 21 ++++++- tests/bedrock_agentcore/runtime/test_utils.py | 15 ++++- .../test_region_validation.py | 34 +++++++++++ .../runtime/test_agent_core_runtime_client.py | 4 +- 12 files changed, 187 insertions(+), 40 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6045543e..2f07facf 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -24,6 +24,20 @@ repos: - id: ruff-format stages: [pre-commit] + # Partition safety: forbid hardcoding the ".amazonaws.com" DNS suffix in + # endpoint construction. Matches ".amazonaws.com" only when it directly + # follows an interpolated value ({...}), so docstrings and the fallback / + # allowlist in _utils/endpoints.py do not trigger. Derive the suffix from the + # region's partition instead (see _dns_suffix_for_region). + - repo: local + hooks: + - id: no-hardcoded-endpoint-suffix + name: no hardcoded amazonaws.com in endpoint construction + language: pygrep + entry: '\{[^}]*\}[^"'']*\.amazonaws\.com' + files: ^src/.*\.py$ + stages: [pre-commit] + # Basic file hygiene (FAST + AUTO-FIX) - repo: https://github.com/pre-commit/pre-commit-hooks rev: v5.0.0 diff --git a/src/bedrock_agentcore/_utils/endpoints.py b/src/bedrock_agentcore/_utils/endpoints.py index 946b760a..c9fd239b 100644 --- a/src/bedrock_agentcore/_utils/endpoints.py +++ b/src/bedrock_agentcore/_utils/endpoints.py @@ -1,14 +1,67 @@ """Endpoint utilities for BedrockAgentCore services.""" +import functools +import logging import os import re from urllib.parse import urlparse +from botocore.exceptions import UnknownRegionError +from botocore.loaders import create_loader +from botocore.regions import EndpointResolver + +logger = logging.getLogger(__name__) + # Environment-configurable constants with fallback defaults DP_ENDPOINT_OVERRIDE = os.getenv("BEDROCK_AGENTCORE_DP_ENDPOINT") CP_ENDPOINT_OVERRIDE = os.getenv("BEDROCK_AGENTCORE_CP_ENDPOINT") DEFAULT_REGION = os.getenv("AWS_REGION") or os.getenv("AWS_DEFAULT_REGION") or "us-west-2" + +@functools.lru_cache(maxsize=1) +def _endpoint_resolver() -> EndpointResolver: + """Build a resolver over botocore's bundled static endpoint data (once, lazily). + + Uses only public botocore APIs and performs no network I/O — the partition + table ships with botocore. Cached so the ``endpoints.json`` load happens at + most once, since this module sits on the runtime hot path. + """ + return EndpointResolver(create_loader().load_data("endpoints")) + + +@functools.lru_cache(maxsize=None) +def _dns_suffix_for_region(region: str) -> str: + """Return the partition DNS suffix for a region. + + For example, ``us-west-2`` resolves to ``amazonaws.com`` while ``cn-north-1`` + resolves to ``amazonaws.com.cn`` and ``us-gov-west-1`` to ``amazonaws.com``. + The suffix is derived from botocore's static partition data rather than a + hardcoded table, so China (``aws-cn``), GovCloud (``aws-us-gov``), and future + partitions are handled without further changes. + + Falls back to ``amazonaws.com`` only for regions botocore does not recognise, + logging a warning so a wrong-partition fallback is visible rather than silent. + """ + resolver = _endpoint_resolver() + try: + partition = resolver.get_partition_for_region(region) + except UnknownRegionError: + logger.warning( + "Region %r is not recognised by botocore; falling back to the " + "'amazonaws.com' DNS suffix. Endpoints may be incorrect outside the " + "commercial partition.", + region, + ) + return "amazonaws.com" + return resolver.get_partition_dns_suffix(partition) or "amazonaws.com" + + +@functools.lru_cache(maxsize=1) +def known_partitions() -> frozenset: + """Return the set of AWS partition names botocore knows about (offline).""" + return frozenset(_endpoint_resolver().get_available_partitions()) + + # Regex for valid AWS region names (e.g., us-east-1, eu-west-2, cn-north-1, us-gov-west-1). # Uses \A and \Z anchors to prevent newline injection bypass that $ allows. _VALID_REGION_PATTERN = re.compile(r"\A[a-z]{2}(-[a-z]+)+-\d+\Z") @@ -82,7 +135,7 @@ def get_data_plane_endpoint(region: str = DEFAULT_REGION) -> str: if DP_ENDPOINT_OVERRIDE: return _validate_endpoint_url(DP_ENDPOINT_OVERRIDE) validate_region(region) - url = f"https://bedrock-agentcore.{region}.amazonaws.com" + url = f"https://bedrock-agentcore.{region}.{_dns_suffix_for_region(region)}" return _validate_endpoint_url(url) @@ -90,7 +143,7 @@ def get_control_plane_endpoint(region: str = DEFAULT_REGION) -> str: if CP_ENDPOINT_OVERRIDE: return _validate_endpoint_url(CP_ENDPOINT_OVERRIDE) validate_region(region) - url = f"https://bedrock-agentcore-control.{region}.amazonaws.com" + url = f"https://bedrock-agentcore-control.{region}.{_dns_suffix_for_region(region)}" return _validate_endpoint_url(url) @@ -113,5 +166,5 @@ def get_gateway_mcp_endpoint(gateway_id: str, region: str = DEFAULT_REGION) -> s f"Invalid gateway identifier: {gateway_id!r}. Expected a gateway ID such as 'my-gateway-abc123'." ) validate_region(region) - url = f"https://{gateway_id}.gateway.bedrock-agentcore.{region}.amazonaws.com/mcp" + url = f"https://{gateway_id}.gateway.bedrock-agentcore.{region}.{_dns_suffix_for_region(region)}/mcp" return _validate_endpoint_url(url) diff --git a/src/bedrock_agentcore/config_bundle/client.py b/src/bedrock_agentcore/config_bundle/client.py index 35ca6fcf..6af3e5b3 100644 --- a/src/bedrock_agentcore/config_bundle/client.py +++ b/src/bedrock_agentcore/config_bundle/client.py @@ -6,7 +6,7 @@ import boto3 -from .._utils.endpoints import DEFAULT_REGION, get_control_plane_endpoint +from .._utils.endpoints import CP_ENDPOINT_OVERRIDE, DEFAULT_REGION logger = logging.getLogger(__name__) @@ -53,11 +53,13 @@ def _get_client(self): with self._client_lock: if self.__dict__.get("_client") is None: session = self._boto3_session or boto3.Session() - self._client = session.client( - "bedrock-agentcore-control", - region_name=self._region, - endpoint_url=get_control_plane_endpoint(self._region), - ) + # Let boto3 resolve the endpoint natively (partition-correct, + # incl. aws-cn). Only pass endpoint_url when the operator set + # an override (e.g. gamma/preprod). + client_kwargs = {"region_name": self._region} + if CP_ENDPOINT_OVERRIDE: + client_kwargs["endpoint_url"] = CP_ENDPOINT_OVERRIDE + self._client = session.client("bedrock-agentcore-control", **client_kwargs) return self._client def __getattr__(self, name: str): diff --git a/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_runner.py b/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_runner.py index 5008b08a..e16df71e 100644 --- a/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_runner.py +++ b/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_runner.py @@ -13,7 +13,7 @@ import boto3 -from bedrock_agentcore._utils.endpoints import DEFAULT_REGION, get_data_plane_endpoint +from bedrock_agentcore._utils.endpoints import DEFAULT_REGION, DP_ENDPOINT_OVERRIDE from bedrock_agentcore.evaluation.runner.batch.batch_evaluation_models import ( BatchEvaluationResult, BatchEvaluationRunConfig, @@ -72,11 +72,12 @@ def __init__(self, region: Optional[str] = None, region_name: Optional[str] = No """ session = boto3.Session() self.region = region or region_name or session.region_name or DEFAULT_REGION - self.data_plane_client = session.client( - "bedrock-agentcore", - region_name=self.region, - endpoint_url=get_data_plane_endpoint(self.region), - ) + # boto3 resolves the endpoint natively (partition-correct, incl. aws-cn); + # only override when the operator set BEDROCK_AGENTCORE_DP_ENDPOINT. + dp_kwargs = {"region_name": self.region} + if DP_ENDPOINT_OVERRIDE: + dp_kwargs["endpoint_url"] = DP_ENDPOINT_OVERRIDE + self.data_plane_client = session.client("bedrock-agentcore", **dp_kwargs) self._logs_client = session.client("logs", region_name=self.region) @staticmethod diff --git a/src/bedrock_agentcore/payments/client.py b/src/bedrock_agentcore/payments/client.py index fb657e77..9ae1b246 100644 --- a/src/bedrock_agentcore/payments/client.py +++ b/src/bedrock_agentcore/payments/client.py @@ -16,7 +16,7 @@ from botocore.config import Config from botocore.exceptions import ClientError -from bedrock_agentcore._utils.endpoints import get_control_plane_endpoint +from bedrock_agentcore._utils.endpoints import CP_ENDPOINT_OVERRIDE from bedrock_agentcore._utils.user_agent import build_user_agent_suffix from bedrock_agentcore.services.identity import IdentityClient @@ -267,13 +267,13 @@ def __init__( user_agent_extra = build_user_agent_suffix(integration_source=self.integration_source, feature="payments") client_config = Config(user_agent_extra=user_agent_extra) - # Control plane operations are available through bedrock-agentcore-control service - self.payments_cp_client = boto3.client( - "bedrock-agentcore-control", - region_name=self.region_name, - endpoint_url=get_control_plane_endpoint(self.region_name), - config=client_config, - ) + # Control plane operations are available through bedrock-agentcore-control service. + # boto3 resolves the endpoint natively (partition-correct, incl. aws-cn); + # only override when the operator set BEDROCK_AGENTCORE_CP_ENDPOINT. + cp_kwargs = {"region_name": self.region_name, "config": client_config} + if CP_ENDPOINT_OVERRIDE: + cp_kwargs["endpoint_url"] = CP_ENDPOINT_OVERRIDE + self.payments_cp_client = boto3.client("bedrock-agentcore-control", **cp_kwargs) # Initialize identity client for credential provider operations self.identity_client = IdentityClient(region=self.region_name) diff --git a/src/bedrock_agentcore/payments/manager.py b/src/bedrock_agentcore/payments/manager.py index 5a087b21..dc813bc0 100644 --- a/src/bedrock_agentcore/payments/manager.py +++ b/src/bedrock_agentcore/payments/manager.py @@ -13,7 +13,7 @@ from botocore.config import Config as BotocoreConfig from botocore.exceptions import ClientError -from bedrock_agentcore._utils.endpoints import get_data_plane_endpoint +from bedrock_agentcore._utils.endpoints import DP_ENDPOINT_OVERRIDE from bedrock_agentcore._utils.user_agent import build_user_agent_suffix from ._validation import validate_permit2_allowance_limit @@ -231,14 +231,14 @@ def __init__( self.region_name = self._validate_and_resolve_region(region_name, boto3_session) session = boto3_session if boto3_session else boto3.Session() - # Configure and create boto3 client + # Configure and create boto3 client. boto3 resolves the endpoint natively + # (partition-correct, incl. aws-cn); only override when the operator set + # BEDROCK_AGENTCORE_DP_ENDPOINT. client_config = self._build_client_config(boto_client_config) - self._payment_client = session.client( - "bedrock-agentcore", - region_name=self.region_name, - config=client_config, - endpoint_url=get_data_plane_endpoint(self.region_name), - ) + dp_kwargs = {"region_name": self.region_name, "config": client_config} + if DP_ENDPOINT_OVERRIDE: + dp_kwargs["endpoint_url"] = DP_ENDPOINT_OVERRIDE + self._payment_client = session.client("bedrock-agentcore", **dp_kwargs) # Register event handler to inject agent name header on every data-plane call if self._agent_name: diff --git a/src/bedrock_agentcore/runtime/a2a.py b/src/bedrock_agentcore/runtime/a2a.py index 720e0791..bd0d6261 100644 --- a/src/bedrock_agentcore/runtime/a2a.py +++ b/src/bedrock_agentcore/runtime/a2a.py @@ -135,10 +135,10 @@ def build_runtime_url(agent_arn: str, region: Optional[str] = None) -> str: """ from urllib.parse import quote - from .._utils.endpoints import validate_region + from .._utils.endpoints import _dns_suffix_for_region, validate_region if region is None: - # ARN format: arn:aws:bedrock-agentcore:::runtime/ + # ARN format: arn::bedrock-agentcore:::runtime/ parts = agent_arn.split(":") if len(parts) >= 4: region = parts[3] @@ -147,7 +147,11 @@ def build_runtime_url(agent_arn: str, region: Optional[str] = None) -> str: validate_region(region) encoded_arn = quote(agent_arn, safe="") - return f"https://bedrock-agentcore.{region}.amazonaws.com/runtimes/{encoded_arn}/invocations" + # Derive the partition-correct DNS suffix (e.g. amazonaws.com.cn in aws-cn). + # Intentionally region-based: this URL is advertised in the agent card, so it + # must not pick up a process-local BEDROCK_AGENTCORE_DP_ENDPOINT override. + suffix = _dns_suffix_for_region(region) + return f"https://bedrock-agentcore.{region}.{suffix}/runtimes/{encoded_arn}/invocations" class BedrockCallContextBuilder: diff --git a/src/bedrock_agentcore/runtime/utils.py b/src/bedrock_agentcore/runtime/utils.py index 7a074577..a04a6e1c 100644 --- a/src/bedrock_agentcore/runtime/utils.py +++ b/src/bedrock_agentcore/runtime/utils.py @@ -36,5 +36,12 @@ def convert_complex_objects(obj: Any, _depth: int = 0) -> Any: def is_valid_partition(partition: str) -> bool: - """Returns if parsed-arn partition is valid.""" - return partition in ("aws", "aws-us-gov") + """Returns whether the parsed-ARN partition is a known AWS partition. + + Backed by botocore's partition list (``aws``, ``aws-cn``, ``aws-us-gov``, + the ISO partitions, ...) so this stays in sync with the TypeScript SDK, + which accepts any ``arn:aws[a-z0-9-]*`` partition. + """ + from .._utils.endpoints import known_partitions + + return partition in known_partitions() diff --git a/tests/bedrock_agentcore/config_bundle/test_client.py b/tests/bedrock_agentcore/config_bundle/test_client.py index 812db0f6..9200ee51 100644 --- a/tests/bedrock_agentcore/config_bundle/test_client.py +++ b/tests/bedrock_agentcore/config_bundle/test_client.py @@ -8,7 +8,9 @@ class TestConfigBundleClient: - def test_boto_client_created_lazily_on_first_access(self): + def test_boto_client_created_lazily_on_first_access(self, monkeypatch): + # No override: boto3 resolves the endpoint natively (partition-correct). + monkeypatch.setattr("bedrock_agentcore.config_bundle.client.CP_ENDPOINT_OVERRIDE", None) mock_session = MagicMock() mock_boto_client = MagicMock() mock_session.client.return_value = mock_boto_client @@ -24,7 +26,22 @@ def test_boto_client_created_lazily_on_first_access(self): mock_session.client.assert_called_once_with( "bedrock-agentcore-control", region_name="us-east-1", - endpoint_url="https://bedrock-agentcore-control.us-east-1.amazonaws.com", + ) + + def test_boto_client_honours_endpoint_override(self, monkeypatch): + # With BEDROCK_AGENTCORE_CP_ENDPOINT set, the override must reach boto3. + override = "https://bedrock-agentcore-control.gamma.example.com" + monkeypatch.setattr("bedrock_agentcore.config_bundle.client.CP_ENDPOINT_OVERRIDE", override) + mock_session = MagicMock() + mock_session.client.return_value = MagicMock() + + client = ConfigBundleClient(region_name="us-east-1", boto3_session=mock_session) + _ = client.list_configuration_bundles + + mock_session.client.assert_called_once_with( + "bedrock-agentcore-control", + region_name="us-east-1", + endpoint_url=override, ) def test_boto_client_reused_across_calls(self): diff --git a/tests/bedrock_agentcore/runtime/test_utils.py b/tests/bedrock_agentcore/runtime/test_utils.py index 5c9ace48..bc57c8fa 100644 --- a/tests/bedrock_agentcore/runtime/test_utils.py +++ b/tests/bedrock_agentcore/runtime/test_utils.py @@ -5,6 +5,7 @@ from pydantic import BaseModel +from bedrock_agentcore.runtime.shell._validation import parse_runtime_arn from bedrock_agentcore.runtime.utils import convert_complex_objects, is_valid_partition @@ -340,5 +341,17 @@ def test_valid_partitions(self): """Test valid partitions for arn parsing""" assert is_valid_partition("aws") + assert is_valid_partition("aws-cn") assert is_valid_partition("aws-us-gov") - assert not is_valid_partition("aws-iso") + # aws-iso is a real (botocore-known) partition, so it is accepted too — + # matching the TS SDK, which accepts any arn:aws[a-z0-9-]* partition. + assert is_valid_partition("aws-iso") + assert not is_valid_partition("not-a-partition") + assert not is_valid_partition("gcp") + + def test_parse_runtime_arn_accepts_aws_cn(self): + """A China (aws-cn) runtime ARN parses without error.""" + parsed = parse_runtime_arn("arn:aws-cn:bedrock-agentcore:cn-north-1:111122223333:runtime/my-agent-abc") + assert parsed["region"] == "cn-north-1" + assert parsed["account_id"] == "111122223333" + assert parsed["runtime_id"] == "my-agent-abc" diff --git a/tests/bedrock_agentcore/test_region_validation.py b/tests/bedrock_agentcore/test_region_validation.py index 793e5194..2ab1ac6f 100644 --- a/tests/bedrock_agentcore/test_region_validation.py +++ b/tests/bedrock_agentcore/test_region_validation.py @@ -317,3 +317,37 @@ def test_memory_controlplane_rejects_bad_region(self): from bedrock_agentcore.memory.controlplane import MemoryControlPlaneClient MemoryControlPlaneClient(region_name="x@attacker.com:443/#") + + +class TestChinaPartitionEndpoints: + """The endpoint builders must emit the amazonaws.com.cn suffix in aws-cn. + + These pin the partition-suffix derivation itself (a botocore lookup failure or + an f-string typo would otherwise ship green); the pre-commit pygrep hook only + guards against newly hardcoded suffixes. + """ + + @pytest.mark.parametrize("region", ["cn-north-1", "cn-northwest-1"]) + def test_data_plane_endpoint_cn(self, region): + assert get_data_plane_endpoint(region) == f"https://bedrock-agentcore.{region}.amazonaws.com.cn" + + @pytest.mark.parametrize("region", ["cn-north-1", "cn-northwest-1"]) + def test_control_plane_endpoint_cn(self, region): + assert get_control_plane_endpoint(region) == f"https://bedrock-agentcore-control.{region}.amazonaws.com.cn" + + @pytest.mark.parametrize("region", ["cn-north-1", "cn-northwest-1"]) + def test_gateway_mcp_endpoint_cn(self, region): + assert ( + get_gateway_mcp_endpoint("gw-abc123", region) + == f"https://gw-abc123.gateway.bedrock-agentcore.{region}.amazonaws.com.cn/mcp" + ) + + def test_commercial_suffix_unchanged(self): + assert get_data_plane_endpoint("us-west-2") == "https://bedrock-agentcore.us-west-2.amazonaws.com" + + def test_build_runtime_url_cn(self): + from bedrock_agentcore.runtime.a2a import build_runtime_url + + url = build_runtime_url("arn:aws-cn:bedrock-agentcore:cn-north-1:111122223333:runtime/my-agent-abc") + assert url.startswith("https://bedrock-agentcore.cn-north-1.amazonaws.com.cn/runtimes/") + assert url.endswith("/invocations") diff --git a/tests/unit/runtime/test_agent_core_runtime_client.py b/tests/unit/runtime/test_agent_core_runtime_client.py index 59d4b37b..7cbe127d 100644 --- a/tests/unit/runtime/test_agent_core_runtime_client.py +++ b/tests/unit/runtime/test_agent_core_runtime_client.py @@ -40,8 +40,10 @@ def test_parse_valid_gov_arn(self): assert result["runtime_id"] == "my-runtime-abc123" def test_parse_invalid_arn_partition(self): + # aws-iso is a real (botocore-known) partition and is now accepted; use a + # bogus partition to exercise the rejection path. with pytest.raises(ValueError, match="Invalid runtime ARN format"): - parse_runtime_arn("arn:aws-iso:bedrock-agentcore:us-iso-east-1:123456789012:runtime/my-runtime-abc123") + parse_runtime_arn("arn:aws-fake:bedrock-agentcore:us-east-1:123456789012:runtime/my-runtime-abc123") def test_parse_invalid_arn_raises_error(self): with pytest.raises(ValueError, match="Invalid runtime ARN format"):