bedrock-agentcore-starter-toolkit 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

bedrock_agentcore_starter_toolkit/services/codebuild.py

@@ -0,0 +1,332 @@
+"""CodeBuild service for ARM64 container builds."""
+
+import fnmatch
+import logging
+import os
+import tempfile
+import time
+import zipfile
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import List
+
+import boto3
+from botocore.exceptions import ClientError
+
+from ..operations.runtime.create_role import get_or_create_codebuild_execution_role
+
+
+class CodeBuildService:
+    """Service for managing CodeBuild projects and builds for ARM64."""
+
+    def __init__(self, session: boto3.Session):
+        """Initialize CodeBuild service with AWS session."""
+        self.session = session
+        self.client = session.client("codebuild")
+        self.s3_client = session.client("s3")
+        self.iam_client = session.client("iam")
+        self.logger = logging.getLogger(__name__)
+        self.source_bucket = None
+
+    def get_source_bucket_name(self, account_id: str) -> str:
+        """Get S3 bucket name for CodeBuild sources."""
+        region = self.session.region_name
+        return f"bedrock-agentcore-codebuild-sources-{account_id}-{region}"
+
+    def ensure_source_bucket(self, account_id: str) -> str:
+        """Ensure S3 bucket exists for CodeBuild sources."""
+        bucket_name = self.get_source_bucket_name(account_id)
+
+        try:
+            self.s3_client.head_bucket(Bucket=bucket_name)
+            self.logger.debug("Using existing S3 bucket: %s", bucket_name)
+        except ClientError:
+            # Create bucket
+            region = self.session.region_name
+            self.s3_client.create_bucket(Bucket=bucket_name, CreateBucketConfiguration={"LocationConstraint": region})
+
+            # Set lifecycle to cleanup old builds
+            self.s3_client.put_bucket_lifecycle_configuration(
+                Bucket=bucket_name,
+                LifecycleConfiguration={
+                    "Rules": [{"ID": "DeleteOldBuilds", "Status": "Enabled", "Filter": {}, "Expiration": {"Days": 7}}]
+                },
+            )
+
+            self.logger.info("Created S3 bucket: %s", bucket_name)
+
+        return bucket_name
+
+    def upload_source(self, agent_name: str) -> str:
+        """Upload current directory to S3, respecting .dockerignore patterns."""
+        account_id = self.session.client("sts").get_caller_identity()["Account"]
+        bucket_name = self.ensure_source_bucket(account_id)
+        self.source_bucket = bucket_name
+
+        # Parse .dockerignore patterns
+        ignore_patterns = self._parse_dockerignore()
+
+        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as temp_zip:
+            try:
+                with zipfile.ZipFile(temp_zip.name, "w", zipfile.ZIP_DEFLATED) as zipf:
+                    for root, dirs, files in os.walk("."):
+                        # Convert to relative path
+                        rel_root = os.path.relpath(root, ".")
+                        if rel_root == ".":
+                            rel_root = ""
+
+                        # Filter directories
+                        dirs[:] = [
+                            d
+                            for d in dirs
+                            if not self._should_ignore(
+                                os.path.join(rel_root, d) if rel_root else d, ignore_patterns, is_dir=True
+                            )
+                        ]
+
+                        for file in files:
+                            file_rel_path = os.path.join(rel_root, file) if rel_root else file
+
+                            # Skip if matches ignore pattern
+                            if self._should_ignore(file_rel_path, ignore_patterns, is_dir=False):
+                                continue
+
+                            file_path = Path(root) / file
+                            zipf.write(file_path, file_rel_path)
+
+                # Create agent-organized S3 key: agentname/timestamp.zip
+                timestamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
+                s3_key = f"{agent_name}/{timestamp}.zip"
+
+                self.s3_client.upload_file(temp_zip.name, bucket_name, s3_key)
+
+                self.logger.info("Uploaded source to S3: %s", s3_key)
+                return f"s3://{bucket_name}/{s3_key}"
+
+            finally:
+                os.unlink(temp_zip.name)
+
+    def _normalize_s3_location(self, source_location: str) -> str:
+        """Convert s3:// URL to bucket/key format for CodeBuild."""
+        return source_location.replace("s3://", "") if source_location.startswith("s3://") else source_location
+
+    def create_codebuild_execution_role(self, account_id: str, ecr_repository_arn: str, agent_name: str) -> str:
+        """Get or create CodeBuild execution role using shared role creation logic."""
+        return get_or_create_codebuild_execution_role(
+            session=self.session,
+            logger=self.logger,
+            region=self.session.region_name,
+            account_id=account_id,
+            agent_name=agent_name,
+            ecr_repository_arn=ecr_repository_arn,
+            source_bucket_name=self.get_source_bucket_name(account_id),
+        )
+
+    def create_or_update_project(
+        self, agent_name: str, ecr_repository_uri: str, execution_role: str, source_location: str
+    ) -> str:
+        """Create or update CodeBuild project for ARM64 builds."""
+        project_name = f"bedrock-agentcore-{agent_name}-builder"
+
+        buildspec = self._get_arm64_buildspec(ecr_repository_uri)
+
+        # CodeBuild expects S3 location without s3:// prefix (bucket/key format)
+        codebuild_source_location = self._normalize_s3_location(source_location)
+
+        project_config = {
+            "name": project_name,
+            "source": {
+                "type": "S3",
+                "location": codebuild_source_location,
+                "buildspec": buildspec,
+            },
+            "artifacts": {
+                "type": "NO_ARTIFACTS",
+            },
+            "environment": {
+                "type": "ARM_CONTAINER",  # ARM64 images require ARM_CONTAINER environment type
+                "image": "aws/codebuild/amazonlinux2-aarch64-standard:3.0",
+                "computeType": "BUILD_GENERAL1_LARGE",  # 4 vCPUs, 7GB RAM for faster builds
+                "privilegedMode": True,  # Required for Docker
+            },
+            "serviceRole": execution_role,
+        }
+
+        try:
+            self.client.create_project(**project_config)
+            self.logger.info("Created CodeBuild project: %s", project_name)
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "ResourceAlreadyExistsException":
+                self.client.update_project(**project_config)
+                self.logger.info("Updated CodeBuild project: %s", project_name)
+            else:
+                raise
+
+        return project_name
+
+    def start_build(self, project_name: str, source_location: str) -> str:
+        """Start a CodeBuild build."""
+        # CodeBuild expects S3 location without s3:// prefix (bucket/key format)
+        codebuild_source_location = self._normalize_s3_location(source_location)
+
+        response = self.client.start_build(
+            projectName=project_name,
+            sourceLocationOverride=codebuild_source_location,
+        )
+
+        return response["build"]["id"]
+
+    def wait_for_completion(self, build_id: str, timeout: int = 900):
+        """Wait for CodeBuild to complete with detailed phase tracking."""
+        self.logger.info("Starting CodeBuild monitoring...")
+
+        # Phase tracking variables
+        current_phase = None
+        phase_start_time = None
+        build_start_time = time.time()
+
+        while time.time() - build_start_time < timeout:
+            response = self.client.batch_get_builds(ids=[build_id])
+            build = response["builds"][0]
+            status = build["buildStatus"]
+            build_phase = build.get("currentPhase", "UNKNOWN")
+
+            # Track phase changes
+            if build_phase != current_phase:
+                # Log previous phase completion (if any)
+                if current_phase and phase_start_time:
+                    phase_duration = time.time() - phase_start_time
+                    self.logger.info("✅ %s completed in %.1fs", current_phase, phase_duration)
+
+                # Log new phase start
+                current_phase = build_phase
+                phase_start_time = time.time()
+                total_duration = phase_start_time - build_start_time
+                self.logger.info("🔄 %s started (total: %.0fs)", current_phase, total_duration)
+
+            # Check for completion
+            if status == "SUCCEEDED":
+                # Log final phase completion
+                if current_phase and phase_start_time:
+                    phase_duration = time.time() - phase_start_time
+                    self.logger.info("✅ %s completed in %.1fs", current_phase, phase_duration)
+
+                total_duration = time.time() - build_start_time
+                minutes, seconds = divmod(int(total_duration), 60)
+                self.logger.info("🎉 CodeBuild completed successfully in %dm %ds", minutes, seconds)
+                return
+
+            elif status in ["FAILED", "FAULT", "STOPPED", "TIMED_OUT"]:
+                # Log failure with phase info
+                if current_phase:
+                    self.logger.error("❌ Build failed during %s phase", current_phase)
+                raise RuntimeError(f"CodeBuild failed with status: {status}")
+
+            time.sleep(5)
+
+        total_duration = time.time() - build_start_time
+        minutes, seconds = divmod(int(total_duration), 60)
+        raise TimeoutError(f"CodeBuild timed out after {minutes}m {seconds}s (current phase: {current_phase})")
+
+    def _get_arm64_buildspec(self, ecr_repository_uri: str) -> str:
+        """Get optimized buildspec for ARM64 Docker."""
+        return f"""
+version: 0.2
+phases:
+  pre_build:
+    commands:
+      - echo Logging in to Amazon ECR...
+      - aws ecr get-login-password --region $AWS_DEFAULT_REGION |
+        docker login --username AWS --password-stdin {ecr_repository_uri}
+      - export DOCKER_BUILDKIT=1
+      - export BUILDKIT_PROGRESS=plain
+  build:
+    commands:
+      - echo Build started on `date`
+      - echo Building ARM64 Docker image with BuildKit processing...
+      - export DOCKER_BUILDKIT=1
+      - docker buildx create --name arm64builder --use || true
+      - docker buildx build --platform linux/arm64 --load -t bedrock-agentcore-arm64 .
+      - docker tag bedrock-agentcore-arm64:latest {ecr_repository_uri}:latest
+  post_build:
+    commands:
+      - echo Build completed on `date`
+      - echo Pushing ARM64 image to ECR...
+      - docker push {ecr_repository_uri}:latest
+"""
+
+    def _parse_dockerignore(self) -> List[str]:
+        """Parse .dockerignore file and return list of patterns."""
+        dockerignore_path = Path(".dockerignore")
+        patterns = []
+
+        if dockerignore_path.exists():
+            with open(dockerignore_path, "r") as f:
+                for line in f:
+                    line = line.strip()
+                    if line and not line.startswith("#"):
+                        patterns.append(line)
+
+            self.logger.info("Using .dockerignore with %d patterns", len(patterns))
+        else:
+            # Default patterns if no .dockerignore
+            patterns = [
+                ".git",
+                "__pycache__",
+                "*.pyc",
+                ".DS_Store",
+                "node_modules",
+                ".venv",
+                "venv",
+                "*.egg-info",
+                ".bedrock_agentcore.yaml",  # Always exclude config
+            ]
+            self.logger.info("No .dockerignore found, using default exclude patterns")
+
+        return patterns
+
+    def _should_ignore(self, path: str, patterns: List[str], is_dir: bool = False) -> bool:
+        """Check if path should be ignored based on dockerignore patterns."""
+        # Normalize path
+        if path.startswith("./"):
+            path = path[2:]
+
+        should_ignore = False  # Default state: don't ignore
+
+        for pattern in patterns:
+            # Handle negation patterns
+            if pattern.startswith("!"):
+                if self._matches_pattern(path, pattern[1:], is_dir):
+                    should_ignore = False  # Negation pattern: don't ignore
+            else:
+                # Regular ignore patterns
+                if self._matches_pattern(path, pattern, is_dir):
+                    should_ignore = True  # Regular pattern: ignore
+
+        return should_ignore
+
+    def _matches_pattern(self, path: str, pattern: str, is_dir: bool) -> bool:
+        """Check if path matches a dockerignore pattern."""
+        # Directory-specific patterns
+        if pattern.endswith("/"):
+            if not is_dir:
+                return False
+            pattern = pattern[:-1]
+
+        # Exact match
+        if path == pattern:
+            return True
+
+        # Glob pattern match
+        if fnmatch.fnmatch(path, pattern):
+            return True
+
+        # Directory prefix match
+        if is_dir and pattern in path.split("/"):
+            return True
+
+        # File in ignored directory
+        if not is_dir and any(fnmatch.fnmatch(part, pattern) for part in path.split("/")):
+            return True
+
+        return False

bedrock_agentcore_starter_toolkit/services/ecr.py

@@ -28,6 +28,35 @@ def create_ecr_repository(repo_name: str, region: str) -> str:
         return response["repositories"][0]["repositoryUri"]
 
 
+def get_or_create_ecr_repository(agent_name: str, region: str) -> str:
+    """Get existing ECR repository or create a new one (idempotent).
+
+    Args:
+        agent_name: Name of the agent
+        region: AWS region
+
+    Returns:
+        ECR repository URI
+    """
+    # Generate deterministic repository name based on agent name
+    repo_name = f"bedrock-agentcore-{agent_name}"
+
+    ecr = boto3.client("ecr", region_name=region)
+
+    try:
+        # Step 1: Check if repository already exists
+        response = ecr.describe_repositories(repositoryNames=[repo_name])
+        existing_repo_uri = response["repositories"][0]["repositoryUri"]
+
+        print(f"✅ Reusing existing ECR repository: {existing_repo_uri}")
+        return existing_repo_uri
+
+    except ecr.exceptions.RepositoryNotFoundException:
+        # Step 2: Repository doesn't exist, create it
+        print(f"Repository doesn't exist, creating new ECR repository: {repo_name}")
+        return create_ecr_repository(repo_name, region)
+
+
 def deploy_to_ecr(local_tag: str, repo_name: str, region: str, container_runtime: ContainerRuntime) -> str:
     """Build and push image to ECR."""
     ecr = boto3.client("ecr", region_name=region)

bedrock_agentcore_starter_toolkit/services/runtime.py

@@ -9,6 +9,7 @@ from typing import Any, Dict, Optional
 
 import boto3
 import requests
+from botocore.exceptions import ClientError
 
 from ..utils.endpoints import get_control_plane_endpoint, get_data_plane_endpoint
 
@@ -93,6 +94,7 @@ class BedrockAgentCoreClient:
         authorizer_config: Optional[Dict] = None,
         protocol_config: Optional[Dict] = None,
         env_vars: Optional[Dict] = None,
+        auto_update_on_conflict: bool = False,
     ) -> Dict[str, str]:
         """Create new agent."""
         self.logger.info("Creating agent '%s' with image URI: %s", agent_name, image_uri)
@@ -121,6 +123,46 @@ class BedrockAgentCoreClient:
             agent_arn = resp["agentRuntimeArn"]
             self.logger.info("Successfully created agent '%s' with ID: %s, ARN: %s", agent_name, agent_id, agent_arn)
             return {"id": agent_id, "arn": agent_arn}
+        except ClientError as e:
+            error_code = e.response.get("Error", {}).get("Code")
+            if error_code == "ConflictException":
+                if not auto_update_on_conflict:
+                    self.logger.error("Agent '%s' already exists and auto_update_on_conflict is disabled", agent_name)
+                    raise
+
+                self.logger.info("Agent '%s' already exists, searching for existing agent...", agent_name)
+
+                # Find existing agent by name
+                existing_agent = self.find_agent_by_name(agent_name)
+
+                if not existing_agent:
+                    raise RuntimeError(
+                        f"ConflictException occurred but couldn't find existing agent '{agent_name}'. "
+                        f"This might be a permissions issue or the agent name might be different."
+                    ) from e
+
+                # Extract existing agent details
+                existing_agent_id = existing_agent["agentRuntimeId"]
+                existing_agent_arn = existing_agent["agentRuntimeArn"]
+
+                self.logger.info("Found existing agent ID: %s, updating instead...", existing_agent_id)
+
+                # Update the existing agent
+                self.update_agent(
+                    existing_agent_id,
+                    image_uri,
+                    execution_role_arn,
+                    network_config,
+                    authorizer_config,
+                    protocol_config,
+                    env_vars,
+                )
+
+                # Return the existing agent info (keeping the original ID and ARN)
+                return {"id": existing_agent_id, "arn": existing_agent_arn}
+            else:
+                # Re-raise other ClientErrors
+                raise
         except Exception as e:
             self.logger.error("Failed to create agent '%s': %s", agent_name, str(e))
             raise
@@ -165,6 +207,46 @@ class BedrockAgentCoreClient:
             self.logger.error("Failed to update agent ID '%s': %s", agent_id, str(e))
             raise
 
+    def list_agents(self, max_results: int = 100) -> list:
+        """List all agent runtimes, handling pagination."""
+        all_agents = []
+        next_token = None
+
+        try:
+            while True:
+                params = {"maxResults": max_results}
+                if next_token:
+                    params["nextToken"] = next_token
+
+                response = self.client.list_agent_runtimes(**params)
+                agents = response.get("agentRuntimes", [])
+                all_agents.extend(agents)
+
+                next_token = response.get("nextToken")
+                if not next_token:
+                    break
+
+            return all_agents
+        except Exception as e:
+            self.logger.error("Failed to list agents: %s", str(e))
+            raise
+
+    def find_agent_by_name(self, agent_name: str) -> Optional[Dict]:
+        """Find an agent by name, reusing list_agents method."""
+        try:
+            # Get all agents using the existing method
+            all_agents = self.list_agents()
+
+            # Search for the specific agent by name
+            for agent in all_agents:
+                if agent.get("agentRuntimeName") == agent_name:
+                    return agent
+
+            return None  # Agent not found
+        except Exception as e:
+            self.logger.error("Failed to search for agent '%s': %s", agent_name, str(e))
+            raise
+
     def create_or_update_agent(
         self,
         agent_id: Optional[str],
@@ -175,6 +257,7 @@ class BedrockAgentCoreClient:
         authorizer_config: Optional[Dict] = None,
         protocol_config: Optional[Dict] = None,
         env_vars: Optional[Dict] = None,
+        auto_update_on_conflict: bool = False,
     ) -> Dict[str, str]:
         """Create or update agent."""
         if agent_id:
@@ -182,7 +265,14 @@ class BedrockAgentCoreClient:
                 agent_id, image_uri, execution_role_arn, network_config, authorizer_config, protocol_config, env_vars
             )
         return self.create_agent(
-            agent_name, image_uri, execution_role_arn, network_config, authorizer_config, protocol_config, env_vars
+            agent_name,
+            image_uri,
+            execution_role_arn,
+            network_config,
+            authorizer_config,
+            protocol_config,
+            env_vars,
+            auto_update_on_conflict,
         )
 
     def wait_for_agent_endpoint_ready(self, agent_id: str, endpoint_name: str = "DEFAULT", max_wait: int = 120) -> str:

bedrock_agentcore_starter_toolkit/utils/logging_config.py

@@ -0,0 +1,72 @@
+"""Centralized logging configuration for bedrock-agentcore-starter-toolkit."""
+
+import logging
+
+_LOGGING_CONFIGURED = False
+
+
+def setup_toolkit_logging(mode: str = "sdk") -> None:
+    """Setup logging for bedrock-agentcore-starter-toolkit.
+
+    Args:
+        mode: "cli" or "sdk" (defaults to "sdk")
+    """
+    global _LOGGING_CONFIGURED
+    if _LOGGING_CONFIGURED:
+        return  # Already configured, prevent duplicates
+
+    if mode == "cli":
+        _setup_cli_logging()
+    elif mode == "sdk":
+        _setup_sdk_logging()
+    else:
+        raise ValueError(f"Invalid logging mode: {mode}. Must be 'cli' or 'sdk'")
+
+    _LOGGING_CONFIGURED = True
+
+
+def _setup_cli_logging() -> None:
+    """Setup logging for CLI usage with RichHandler."""
+    try:
+        from rich.logging import RichHandler
+
+        from ..cli.common import console
+
+        FORMAT = "%(message)s"
+        logging.basicConfig(
+            level="INFO",
+            format=FORMAT,
+            handlers=[RichHandler(show_time=False, show_path=False, show_level=False, console=console)],
+            force=True,  # Override any existing configuration
+        )
+    except ImportError:
+        # Fallback if rich is not available
+        _setup_basic_logging()
+
+
+def _setup_sdk_logging() -> None:
+    """Setup logging for SDK usage (notebooks, scripts, imports) with StreamHandler."""
+    # Configure logger for ALL toolkit modules (ensures all operation logs appear)
+    toolkit_logger = logging.getLogger("bedrock_agentcore_starter_toolkit")
+
+    if not toolkit_logger.handlers:
+        handler = logging.StreamHandler()
+        handler.setFormatter(logging.Formatter("%(message)s"))
+        toolkit_logger.addHandler(handler)
+        toolkit_logger.setLevel(logging.INFO)
+
+
+def _setup_basic_logging() -> None:
+    """Setup basic logging as fallback."""
+    logging.basicConfig(level=logging.INFO, format="%(message)s", force=True)
+
+
+def is_logging_configured() -> bool:
+    """Check if toolkit logging has been configured."""
+    return _LOGGING_CONFIGURED
+
+
+def reset_logging_config() -> None:
+    """Reset logging configuration state (for testing)."""
+    global _LOGGING_CONFIGURED
+    _LOGGING_CONFIGURED = False

bedrock_agentcore_starter_toolkit/utils/runtime/entrypoint.py

@@ -181,11 +181,11 @@ def validate_requirements_file(build_dir: Path, requirements_file: str) -> Depen
             f"Please specify a requirements file (requirements.txt, pyproject.toml, etc.)"
         )
 
-    # Validate that it's a recognized dependency file type (strict validation)
-    if file_path.name not in ["requirements.txt", "pyproject.toml"]:
+    # Validate that it's a recognized dependency file type (flexible validation)
+    if not (file_path.suffix in [".txt", ".in"] or file_path.name == "pyproject.toml"):
         raise ValueError(
             f"'{file_path.name}' is not a supported dependency file. "
-            f"Only 'requirements.txt' and 'pyproject.toml' are supported."
+            f"Supported formats: *.txt, *.in (pip requirements), or pyproject.toml"
         )
 
     # Use the existing detect_dependencies function to process the file

bedrock_agentcore_starter_toolkit/utils/runtime/policy_template.py

@@ -0,0 +1,74 @@
+"""Policy template utilities for runtime execution roles."""
+
+import json
+from pathlib import Path
+from typing import Dict
+
+from jinja2 import Environment, FileSystemLoader
+
+
+def _get_template_dir() -> Path:
+    """Get the templates directory path."""
+    return Path(__file__).parent / "templates"
+
+
+def _render_template(template_name: str, variables: Dict[str, str]) -> str:
+    """Render a Jinja2 template with the provided variables."""
+    template_dir = _get_template_dir()
+    env = Environment(loader=FileSystemLoader(template_dir), autoescape=True)
+    template = env.get_template(template_name)
+    return template.render(**variables)
+
+
+def render_trust_policy_template(region: str, account_id: str) -> str:
+    """Render the trust policy template with provided values.
+
+    Args:
+        region: AWS region
+        account_id: AWS account ID
+
+    Returns:
+        Rendered trust policy as JSON string
+    """
+    variables = {
+        "region": region,
+        "account_id": account_id,
+    }
+    return _render_template("execution_role_trust_policy.json.j2", variables)
+
+
+def render_execution_policy_template(region: str, account_id: str, agent_name: str) -> str:
+    """Render the execution policy template with provided values.
+
+    Args:
+        region: AWS region
+        account_id: AWS account ID
+        agent_name: Agent name for resource scoping
+
+    Returns:
+        Rendered execution policy as JSON string
+    """
+    variables = {
+        "region": region,
+        "account_id": account_id,
+        "agent_name": agent_name,
+    }
+    return _render_template("execution_role_policy.json.j2", variables)
+
+
+def validate_rendered_policy(policy_json: str) -> Dict:
+    """Validate that the rendered policy is valid JSON.
+
+    Args:
+        policy_json: JSON policy string
+
+    Returns:
+        Parsed policy dictionary
+
+    Raises:
+        ValueError: If policy JSON is invalid
+    """
+    try:
+        return json.loads(policy_json)
+    except json.JSONDecodeError as e:
+        raise ValueError(f"Invalid policy JSON: {e}") from e