chore: Modify arn parsing in runtime client to allow for different (#362)

partitions

	- Added is_valid_partition function to allow for govcloud arns

Co-authored-by: Trirmadura J Ariyawansa <tjariy@amazon.com>

Author: mahlera

src/bedrock_agentcore/runtime/agent_core_runtime_client.py

@@ -17,6 +17,7 @@
 from botocore.awsrequest import AWSRequest
 
 from .._utils.endpoints import get_data_plane_endpoint
+from .utils import is_valid_partition
 
 DEFAULT_PRESIGNED_URL_TIMEOUT = 300
 MAX_PRESIGNED_URL_TIMEOUT = 300
@@ -68,7 +69,7 @@ def _parse_runtime_arn(self, runtime_arn: str) -> Dict[str, str]:
         if len(parts) != 6:
             raise ValueError(f"Invalid runtime ARN format: {runtime_arn}")
 
-        if parts[0] != "arn" or parts[1] != "aws" or parts[2] != "bedrock-agentcore":
+        if parts[0] != "arn" or not is_valid_partition(parts[1]) or parts[2] != "bedrock-agentcore":
             raise ValueError(f"Invalid runtime ARN format: {runtime_arn}")
 
         # Parse the resource part (runtime/{runtime_id})

src/bedrock_agentcore/runtime/utils.py

@@ -33,3 +33,8 @@ def convert_complex_objects(obj: Any, _depth: int = 0) -> Any:
     # Return primitives as-is
     else:
         return obj
+
+
+def is_valid_partition(partition: str) -> bool:
+    """Returns if parsed-arn partition is valid."""
+    return partition in ("aws", "aws-us-gov")

tests/bedrock_agentcore/runtime/test_utils.py

@@ -5,7 +5,7 @@
 
 from pydantic import BaseModel
 
-from bedrock_agentcore.runtime.utils import convert_complex_objects
+from bedrock_agentcore.runtime.utils import convert_complex_objects, is_valid_partition
 
 
 class TestConvertComplexObjects:
@@ -335,3 +335,10 @@ class Person(BaseModel):
         assert isinstance(result["address"], dict)
         assert result["address"]["street"] == "456 Oak St"
         assert result["address"]["city"] == "Somewhere"
+
+    def test_valid_partitions(self):
+        """Test valid partitions for arn parsing"""
+
+        assert is_valid_partition("aws")
+        assert is_valid_partition("aws-us-gov")
+        assert not is_valid_partition("aws-iso")

tests/unit/runtime/test_agent_core_runtime_client.py

@@ -36,6 +36,25 @@ def test_parse_valid_arn(self):
         assert result["account_id"] == "123456789012"
         assert result["runtime_id"] == "my-runtime-abc123"
 
+    def test_parse_valid_gov_arn(self):
+        """Test parsing a valid govcloud runtime ARN."""
+        client = AgentCoreRuntimeClient(region="us-gov-west-1")
+        arn = "arn:aws-us-gov:bedrock-agentcore:us-gov-west-1:123456789012:runtime/my-runtime-abc123"
+
+        result = client._parse_runtime_arn(arn)
+
+        assert result["region"] == "us-gov-west-1"
+        assert result["account_id"] == "123456789012"
+        assert result["runtime_id"] == "my-runtime-abc123"
+
+    def test_parse_invalid_arn_partition(self):
+        """Test parsing an invalid runtime ARN."""
+        client = AgentCoreRuntimeClient(region="us-iso-east-1")
+        invalid_arn = "arn:aws-iso:bedrock-agentcore:us-iso-east-1:123456789012:runtime/my-runtime-abc123"
+
+        with pytest.raises(ValueError, match="Invalid runtime ARN format"):
+            client._parse_runtime_arn(invalid_arn)
+
     def test_parse_invalid_arn_raises_error(self):
         """Test that invalid ARN format raises ValueError."""
         client = AgentCoreRuntimeClient(region="us-west-2")