Support for Cross Account Builds

src/bedrock_agentcore_starter_toolkit/operations/runtime/launch.py

@@ -433,18 +433,46 @@ def _execute_codebuild_workflow(
     if not ecr_only:
         _ensure_execution_role(agent_config, project_config, config_path, agent_name, region, account_id)
 
-    # Prepare CodeBuild
+    # Enhanced CodeBuild preparation with cross-account support
     log.info("Preparing CodeBuild project and uploading source...")
-    codebuild_service = CodeBuildService(session)
 
-    # Use cached CodeBuild role from config if available
+    # Get CodeBuild execution role from config
+    codebuild_execution_role = None
     if hasattr(agent_config, "codebuild") and agent_config.codebuild.execution_role:
-        log.info("Using CodeBuild role from config: %s", agent_config.codebuild.execution_role)
         codebuild_execution_role = agent_config.codebuild.execution_role
-    else:
+        log.info("Using CodeBuild role from config: %s", codebuild_execution_role)
+
+    # Detect cross-account scenario
+    build_account = None
+    deployment_account = session.client("sts").get_caller_identity()["Account"]
+
+    if codebuild_execution_role:
+        try:
+            # Extract account from role ARN: arn:aws:iam::ACCOUNT-ID:role/ROLE-NAME
+            build_account = codebuild_execution_role.split(":")[4]
+
+            if build_account != deployment_account:
+                log.info("Cross-account CodeBuild detected:")
+                log.info("  Deployment account: %s", deployment_account)
+                log.info("  Build account: %s", build_account)
+                log.info("  CodeBuild role: %s", codebuild_execution_role)
+            else:
+                log.info("Same-account CodeBuild (role in deployment account)")
+                build_account = None  # Treat as same-account
+        except (IndexError, AttributeError):
+            log.warning("Invalid CodeBuild role ARN format: %s", codebuild_execution_role)
+            build_account = None
+
+    # Initialize CodeBuild service with role information
+    codebuild_service = CodeBuildService(session, codebuild_execution_role)
+
+    # Create or use existing CodeBuild execution role
+    if not codebuild_execution_role:
+        # No role specified - create one in deployment account
         codebuild_execution_role = codebuild_service.create_codebuild_execution_role(
             account_id=account_id, ecr_repository_arn=ecr_repository_arn, agent_name=agent_name
         )
+        log.info("Created CodeBuild role in deployment account: %s", codebuild_execution_role)
 
     source_location = codebuild_service.upload_source(agent_name=agent_name)
 

src/bedrock_agentcore_starter_toolkit/services/codebuild.py

@@ -7,7 +7,7 @@
 import time
 import zipfile
 from pathlib import Path
-from typing import List
+from typing import List, Optional
 
 import boto3
 from botocore.exceptions import ClientError
@@ -18,19 +18,76 @@
 class CodeBuildService:
     """Service for managing CodeBuild projects and builds for ARM64."""
 
-    def __init__(self, session: boto3.Session):
-        """Initialize CodeBuild service with AWS session."""
+    def __init__(self, session: boto3.Session, codebuild_role_arn: Optional[str] = None):
+        """Initialize CodeBuild service with AWS session.
+
+        Args:
+            session: Primary AWS session for deployment account
+            codebuild_role_arn: Optional CodeBuild execution role ARN (for cross-account)
+        """
         self.session = session
-        self.client = session.client("codebuild")
-        self.s3_client = session.client("s3")
-        self.iam_client = session.client("iam")
+        self.codebuild_role_arn = codebuild_role_arn
         self.logger = logging.getLogger(__name__)
+
+        # Determine if this is cross-account CodeBuild
+        self.build_account = self._extract_build_account()
+        self.deployment_account = session.client("sts").get_caller_identity()["Account"]
+        self.is_cross_account_codebuild = (
+            self.build_account is not None and self.build_account != self.deployment_account
+        )
+
+        # Create appropriate session for CodeBuild operations
+        if self.is_cross_account_codebuild:
+            self.build_session = self._create_build_session()
+            self.client = self.build_session.client("codebuild")
+            self.s3_client = self.build_session.client("s3")
+            self.iam_client = self.build_session.client("iam")
+            self.logger.info("CodeBuild initialized in cross-account mode (account: %s)", self.build_account)
+        else:
+            self.build_session = session
+            self.client = session.client("codebuild")
+            self.s3_client = session.client("s3")
+            self.iam_client = session.client("iam")
+            self.logger.info("CodeBuildService initialized in same-account mode")
+
         self.source_bucket = None
-        self.account_id = session.client("sts").get_caller_identity()["Account"]
+
+    def _extract_build_account(self) -> Optional[str]:
+        """Extract build account ID from CodeBuild role ARN."""
+        if not self.codebuild_role_arn:
+            return None
+        try:
+            # ARN format: arn:aws:iam::ACCOUNT-ID:role/ROLE-NAME
+            return self.codebuild_role_arn.split(":")[4]
+        except (IndexError, AttributeError):
+            self.logger.warning("Invalid CodeBuild role ARN format: %s", self.codebuild_role_arn)
+            return None
+
+    def _create_build_session(self) -> boto3.Session:
+        """Create AWS session by assuming the CodeBuild role."""
+        if not self.codebuild_role_arn:
+            return self.session
+
+        try:
+            sts_client = self.session.client("sts")
+            response = sts_client.assume_role(
+                RoleArn=self.codebuild_role_arn, RoleSessionName=f"bedrock-agentcore-build-{int(time.time())}"
+            )
+
+            credentials = response["Credentials"]
+            return boto3.Session(
+                aws_access_key_id=credentials["AccessKeyId"],
+                aws_secret_access_key=credentials["SecretAccessKey"],
+                aws_session_token=credentials["SessionToken"],
+                region_name=self.session.region_name,
+            )
+        except Exception as e:
+            self.logger.error("Failed to assume CodeBuild role %s: %s", self.codebuild_role_arn, e)
+            raise RuntimeError(f"Failed to assume CodeBuild role {self.codebuild_role_arn}: {e}") from e
 
     def get_source_bucket_name(self, account_id: str) -> str:
         """Get S3 bucket name for CodeBuild sources."""
-        region = self.session.region_name
+        region = self.build_session.region_name
         return f"bedrock-agentcore-codebuild-sources-{account_id}-{region}"
 
     def ensure_source_bucket(self, account_id: str) -> str:
@@ -50,7 +107,7 @@ def ensure_source_bucket(self, account_id: str) -> str:
                 ) from e
 
             # Create bucket (no ExpectedBucketOwner needed for create_bucket)
-            region = self.session.region_name
+            region = self.build_session.region_name
             if region == "us-east-1":
                 self.s3_client.create_bucket(Bucket=bucket_name)
             else:
@@ -72,7 +129,9 @@ def ensure_source_bucket(self, account_id: str) -> str:
 
     def upload_source(self, agent_name: str) -> str:
         """Upload current directory to S3, respecting .dockerignore patterns."""
-        account_id = self.account_id
+        # Use build account for S3 bucket (cross-account) or deployment account (same-account)
+        account_id = self.build_account or self.deployment_account
+        self.logger.info("Using account %s for S3 bucket", account_id)
         bucket_name = self.ensure_source_bucket(account_id)
         self.source_bucket = bucket_name
 

tests/operations/runtime/test_launch_cross_account.py

@@ -0,0 +1,155 @@
+"""Simple tests for cross-account launch functionality."""
+
+from unittest.mock import Mock, patch
+import pytest
+
+from bedrock_agentcore_starter_toolkit.operations.runtime.launch import _execute_codebuild_workflow
+from bedrock_agentcore_starter_toolkit.utils.runtime.schema import (
+    AWSConfig,
+    BedrockAgentCoreAgentSchema,
+    BedrockAgentCoreConfigSchema,
+    NetworkConfiguration,
+    ObservabilityConfig,
+    BedrockAgentCoreDeploymentInfo,
+    CodeBuildConfig,
+)
+
+
+class TestLaunchCrossAccount:
+    """Test cross-account functionality in launch operations."""
+
+    def test_codebuild_service_initialization_cross_account(self, tmp_path):
+        """Test CodeBuildService is initialized with cross-account role."""
+        # Create agent config with cross-account CodeBuild role
+        aws_config = AWSConfig(
+            account="123456789012",
+            region="us-west-2",
+            execution_role="arn:aws:iam::123456789012:role/ExecutionRole",
+            ecr_repository="123456789012.dkr.ecr.us-west-2.amazonaws.com/test-repo",
+            network_configuration=NetworkConfiguration(),
+            observability=ObservabilityConfig(),
+        )
+
+        codebuild_config = CodeBuildConfig()
+        codebuild_config.execution_role = "arn:aws:iam::987654321098:role/BuildRole"
+
+        agent_config = BedrockAgentCoreAgentSchema(
+            name="test-agent",
+            entrypoint="test.py",
+            aws=aws_config,
+            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
+            codebuild=codebuild_config,
+        )
+
+        project_config = BedrockAgentCoreConfigSchema(
+            default_agent="test-agent",
+            agents={"test-agent": agent_config}
+        )
+
+        config_path = tmp_path / "config.yaml"
+
+        with patch('bedrock_agentcore_starter_toolkit.operations.runtime.launch.CodeBuildService') as mock_cb_service, \
+             patch('bedrock_agentcore_starter_toolkit.operations.runtime.launch._ensure_ecr_repository') as mock_ecr, \
+             patch('bedrock_agentcore_starter_toolkit.operations.runtime.launch._ensure_execution_role') as mock_role, \
+             patch("boto3.Session") as mock_session:
+
+            # Setup mocks
+            mock_ecr.return_value = "123456789012.dkr.ecr.us-west-2.amazonaws.com/test-repo"
+            mock_role.return_value = "arn:aws:iam::123456789012:role/ExecutionRole"
+
+            mock_service_instance = Mock()
+            mock_service_instance.upload_source.return_value = "s3://bucket/source.zip"
+            mock_service_instance.create_or_update_project.return_value = "test-project"
+            mock_service_instance.start_build.return_value = "build-123"
+            mock_service_instance.wait_for_completion.return_value = None
+            mock_service_instance.source_bucket = "test-bucket"
+            mock_cb_service.return_value = mock_service_instance
+
+            deployment_session = Mock()
+            # Mock STS client for account detection
+            mock_sts = Mock()
+            mock_sts.get_caller_identity.return_value = {"Account": "123456789012"}
+            deployment_session.client.return_value = mock_sts
+            mock_session.return_value = deployment_session
+
+            # Execute
+            _execute_codebuild_workflow(
+                config_path=config_path,
+                agent_name="test-agent",
+                agent_config=agent_config,
+                project_config=project_config,
+                ecr_only=False
+            )
+
+            # Verify CodeBuildService was called with cross-account role
+            mock_cb_service.assert_called_once_with(
+                deployment_session,
+                "arn:aws:iam::987654321098:role/BuildRole"
+            )
+
+    def test_codebuild_service_initialization_same_account(self, tmp_path):
+        """Test CodeBuildService is initialized without cross-account role."""
+        # Create agent config without cross-account CodeBuild role
+        aws_config = AWSConfig(
+            account="123456789012",
+            region="us-west-2",
+            execution_role="arn:aws:iam::123456789012:role/ExecutionRole",
+            ecr_repository="123456789012.dkr.ecr.us-west-2.amazonaws.com/test-repo",
+            network_configuration=NetworkConfiguration(),
+            observability=ObservabilityConfig(),
+        )
+
+        codebuild_config = CodeBuildConfig()
+        # No execution_role set - same account scenario
+
+        agent_config = BedrockAgentCoreAgentSchema(
+            name="test-agent",
+            entrypoint="test.py",
+            aws=aws_config,
+            bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
+            codebuild=codebuild_config,
+        )
+
+        project_config = BedrockAgentCoreConfigSchema(
+            default_agent="test-agent",
+            agents={"test-agent": agent_config}
+        )
+
+        config_path = tmp_path / "config.yaml"
+
+        with patch('bedrock_agentcore_starter_toolkit.operations.runtime.launch.CodeBuildService') as mock_cb_service, \
+             patch('bedrock_agentcore_starter_toolkit.operations.runtime.launch._ensure_ecr_repository') as mock_ecr, \
+             patch('bedrock_agentcore_starter_toolkit.operations.runtime.launch._ensure_execution_role') as mock_role, \
+             patch("boto3.Session") as mock_session:
+
+            # Setup mocks
+            mock_ecr.return_value = "123456789012.dkr.ecr.us-west-2.amazonaws.com/test-repo"
+            mock_role.return_value = "arn:aws:iam::123456789012:role/ExecutionRole"
+
+            mock_service_instance = Mock()
+            mock_service_instance.create_codebuild_execution_role.return_value = "arn:aws:iam::123456789012:role/CodeBuildRole"
+            mock_service_instance.upload_source.return_value = "s3://bucket/source.zip"
+            mock_service_instance.create_or_update_project.return_value = "test-project"
+            mock_service_instance.start_build.return_value = "build-123"
+            mock_service_instance.wait_for_completion.return_value = None
+            mock_service_instance.source_bucket = "test-bucket"
+            mock_cb_service.return_value = mock_service_instance
+
+            deployment_session = Mock()
+            # Mock STS client for account detection
+            mock_sts = Mock()
+            mock_sts.get_caller_identity.return_value = {"Account": "123456789012"}
+            deployment_session.client.return_value = mock_sts
+            mock_session.return_value = deployment_session
+
+            # Execute
+            _execute_codebuild_workflow(
+                config_path=config_path,
+                agent_name="test-agent",
+                agent_config=agent_config,
+                project_config=project_config,
+                ecr_only=False
+            )
+
+            # Verify CodeBuildService was called without cross-account role
+            mock_cb_service.assert_called_once_with(deployment_session, None)

tests/services/test_codebuild.py

@@ -84,7 +84,9 @@ def client_factory(service_name):
         assert service.s3_client == mock_s3
         assert service.iam_client == mock_iam
         assert service.source_bucket is None
-        assert service.account_id == "123456789012"  # Verify account_id is stored
+        assert service.deployment_account == "123456789012"  # Verify deployment account is stored
+        assert service.build_account is None  # No cross-account role provided
+        assert service.is_cross_account_codebuild is False
 
     def test_get_source_bucket_name(self, codebuild_service):
         """Test S3 bucket name generation."""
@@ -188,6 +190,7 @@ def test_upload_source_success(
         result = codebuild_service.upload_source("test-agent")
 
         expected_key = "test-agent/source.zip"
+        # Use deployment account since no cross-account role provided
         expected_s3_url = f"s3://bedrock-agentcore-codebuild-sources-123456789012-us-west-2/{expected_key}"
 
         assert result == expected_s3_url