feat(endpoints): support AWS China (aws-cn) partition

.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

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,15 +135,15 @@ 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)
 
 
 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)

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):

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

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)

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:

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:<region>:<account>:runtime/<id>
+        # ARN format: arn:<partition>:bedrock-agentcore:<region>:<account>:runtime/<id>
         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:

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()

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):

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"

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")

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"):