bedrock-agentcore-starter-toolkit 0.1.11__py3-none-any.whl → 0.1.13__py3-none-any.whl

bedrock_agentcore_starter_toolkit/cli/import_agent/agent_info.py

@@ -148,7 +148,12 @@ def get_agent_info(agent_id: str, agent_alias_id: str, bedrock_client, bedrock_a
                 s3_object_key = action_group["apiSchema"]["s3"]["s3ObjectKey"]
 
                 s3_client = boto3.client("s3")
-                response = s3_client.get_object(Bucket=s3_bucket_name, Key=s3_object_key)
+                # Get account ID for bucket ownership verification
+                sts_client = boto3.client("sts")
+                account_id = sts_client.get_caller_identity()["Account"]
+                response = s3_client.get_object(
+                    Bucket=s3_bucket_name, Key=s3_object_key, ExpectedBucketOwner=account_id
+                )
                 yaml_content = response["Body"].read().decode("utf-8")
                 yaml = YAML(typ="safe")
                 action_group["apiSchema"]["payload"] = yaml.load(yaml_content)

bedrock_agentcore_starter_toolkit/cli/runtime/commands.py

@@ -68,7 +68,7 @@ def _prompt_for_requirements_file(prompt_text: str, default: str = "") -> Option
     return None
 
 
-def _handle_requirements_file_display(requirements_file: Optional[str]) -> Optional[str]:
+def _handle_requirements_file_display(requirements_file: Optional[str], non_interactive: bool = False) -> Optional[str]:
     """Handle requirements file with display logic for CLI."""
     from ...utils.runtime.entrypoint import detect_dependencies
 
@@ -76,6 +76,15 @@ def _handle_requirements_file_display(requirements_file: Optional[str]) -> Optio
         # User provided file - validate and show confirmation
         return _validate_requirements_file(requirements_file)
 
+    if non_interactive:
+        # Auto-detection for non-interactive mode
+        deps = detect_dependencies(Path.cwd())
+        if deps.found:
+            _print_success(f"Using detected file: [dim]{deps.file}[/dim]")
+            return None  # Use detected file
+        else:
+            _handle_error("No requirements file specified and none found automatically")
+
     # Auto-detection with interactive prompt
     deps = detect_dependencies(Path.cwd())
 
@@ -167,9 +176,19 @@ def configure(
     authorizer_config: Optional[str] = typer.Option(
         None, "--authorizer-config", "-ac", help="OAuth authorizer configuration as JSON string"
     ),
+    request_header_allowlist: Optional[str] = typer.Option(
+        None,
+        "--request-header-allowlist",
+        "-rha",
+        help="Comma-separated list of allowed request headers "
+        "(Authorization or X-Amzn-Bedrock-AgentCore-Runtime-Custom-*)",
+    ),
     verbose: bool = typer.Option(False, "--verbose", "-v", help="Enable verbose output"),
     region: Optional[str] = typer.Option(None, "--region", "-r"),
     protocol: Optional[str] = typer.Option(None, "--protocol", "-p", help="Server protocol (HTTP or MCP)"),
+    non_interactive: bool = typer.Option(
+        False, "--non-interactive", "-ni", help="Skip prompts; use defaults unless overridden"
+    ),
 ):
     """Configure a Bedrock AgentCore agent. The agent name defaults to your Python file name."""
     if ctx.invoked_subcommand is not None:
@@ -196,7 +215,7 @@ def configure(
 
     # Create configuration manager for clean, elegant prompting
     config_path = Path.cwd() / ".bedrock_agentcore.yaml"
-    config_manager = ConfigurationManager(config_path)
+    config_manager = ConfigurationManager(config_path, non_interactive)
 
     # Interactive prompts for missing values - clean and elegant
     if not execution_role:
@@ -217,7 +236,7 @@ def configure(
         _print_success(f"Using existing ECR repository: [dim]{ecr_repository}[/dim]")
 
     # Handle dependency file selection with simplified logic
-    final_requirements_file = _handle_requirements_file_display(requirements_file)
+    final_requirements_file = _handle_requirements_file_display(requirements_file, non_interactive)
 
     # Handle OAuth authorization configuration
     oauth_config = None
@@ -231,6 +250,19 @@ def configure(
     else:
         oauth_config = config_manager.prompt_oauth_config()
 
+    # Handle request header allowlist configuration
+    request_header_config = None
+    if request_header_allowlist:
+        # Parse comma-separated headers and create configuration
+        headers = [header.strip() for header in request_header_allowlist.split(",") if header.strip()]
+        if headers:
+            request_header_config = {"requestHeaderAllowlist": headers}
+            _print_success(f"Configured request header allowlist with {len(headers)} headers")
+        else:
+            _handle_error("Empty request header allowlist provided")
+    else:
+        request_header_config = config_manager.prompt_request_header_allowlist()
+
     try:
         result = configure_bedrock_agentcore(
             agent_name=agent_name,
@@ -242,6 +274,7 @@ def configure(
             enable_observability=not disable_otel,
             requirements_file=final_requirements_file,
             authorizer_configuration=oauth_config,
+            request_header_configuration=request_header_config,
             verbose=verbose,
             region=region,
             protocol=protocol.upper() if protocol else None,
@@ -252,6 +285,12 @@ def configure(
         if oauth_config:
             auth_info = "OAuth (customJWTAuthorizer)"
 
+        # Prepare request headers info for summary
+        headers_info = ""
+        if request_header_config:
+            headers = request_header_config.get("requestHeaderAllowlist", [])
+            headers_info = f"Request Headers Allowlist: [dim]{len(headers)} headers configured[/dim]\n"
+
         console.print(
             Panel(
                 f"[green]Configuration Complete[/green]\n\n"
@@ -265,7 +304,8 @@ def configure(
                 f"ECR Repository: [dim]"
                 f"{'Auto-create' if result.auto_create_ecr else result.ecr_repository or 'N/A'}"
                 f"[/dim]\n"
-                f"Authorization: [dim]{auth_info}[/dim]\n\n"
+                f"Authorization: [dim]{auth_info}[/dim]\n"
+                f"{headers_info}\n"
                 f"📄 Config saved to: [dim]{result.config_path}[/dim]\n\n"
                 f"[bold]Next Steps:[/bold]\n"
                 f"   [cyan]agentcore launch[/cyan]",
@@ -570,6 +610,45 @@ def _show_error_response(error_msg: str):
     console.print(f"\n[red]{error_msg}[/red]")
 
 
+def _parse_custom_headers(headers_str: str) -> dict:
+    """Parse custom headers string and apply prefix logic.
+
+    Args:
+        headers_str: String in format "Header1:value,Header2:value2"
+
+    Returns:
+        dict: Dictionary of processed headers with proper prefixes
+
+    Raises:
+        ValueError: If header format is invalid
+    """
+    if not headers_str or not headers_str.strip():
+        return {}
+
+    headers = {}
+    header_pairs = [pair.strip() for pair in headers_str.split(",")]
+
+    for pair in header_pairs:
+        if ":" not in pair:
+            raise ValueError(f"Invalid header format: '{pair}'. Expected format: 'Header:value'")
+
+        header_name, header_value = pair.split(":", 1)
+        header_name = header_name.strip()
+        header_value = header_value.strip()
+
+        if not header_name:
+            raise ValueError(f"Empty header name in: '{pair}'")
+
+        # Apply prefix logic: if header doesn't start with the custom prefix, add it
+        prefix = "X-Amzn-Bedrock-AgentCore-Runtime-Custom-"
+        if not header_name.startswith(prefix):
+            header_name = prefix + header_name
+
+        headers[header_name] = header_value
+
+    return headers
+
+
 def invoke(
     payload: str = typer.Argument(..., help="JSON payload to send"),
     agent: Optional[str] = typer.Option(
@@ -581,6 +660,12 @@ def invoke(
     ),
     local_mode: Optional[bool] = typer.Option(False, "--local", "-l", help="Send request to a running local container"),
     user_id: Optional[str] = typer.Option(None, "--user-id", "-u", help="User id for authorization flows"),
+    headers: Optional[str] = typer.Option(
+        None,
+        "--headers",
+        help="Custom headers (format: 'Header1:value,Header2:value2'). "
+        "Headers will be auto-prefixed with 'X-Amzn-Bedrock-AgentCore-Runtime-Custom-' if not already present.",
+    ),
 ):
     """Invoke Bedrock AgentCore endpoint."""
     config_path = Path.cwd() / ".bedrock_agentcore.yaml"
@@ -613,6 +698,17 @@ def invoke(
                 "[yellow]Warning: Bearer token provided but OAuth is not configured in .bedrock_agentcore.yaml[/yellow]"
             )
 
+        # Process custom headers
+        custom_headers = {}
+        if headers:
+            try:
+                custom_headers = _parse_custom_headers(headers)
+                if custom_headers:
+                    header_names = list(custom_headers.keys())
+                    console.print(f"[dim]Using custom headers: {', '.join(header_names)}[/dim]")
+            except ValueError as e:
+                _handle_error(f"Invalid headers format: {e}")
+
         # Invoke
         result = invoke_bedrock_agentcore(
             config_path=config_path,
@@ -622,6 +718,7 @@ def invoke(
             bearer_token=final_bearer_token,
             user_id=user_id,
             local_mode=local_mode,
+            custom_headers=custom_headers,
         )
         agent_display = config.name if config else (agent or "unknown")
         _show_invoke_info_panel(agent_display, result, config)

bedrock_agentcore_starter_toolkit/cli/runtime/configuration_manager.py

@@ -10,19 +10,25 @@ from ..common import _handle_error, _print_success, _prompt_with_default, consol
 class ConfigurationManager:
     """Manages interactive configuration prompts with existing configuration defaults."""
 
-    def __init__(self, config_path: Path):
+    def __init__(self, config_path: Path, non_interactive: bool = False):
         """Initialize the ConfigPrompt with a configuration path.
 
         Args:
             config_path: Path to the configuration file
+            non_interactive: If True, use defaults without prompting
         """
         from ...utils.runtime.config import load_config_if_exists
 
         project_config = load_config_if_exists(config_path)
         self.existing_config = project_config.get_agent_config() if project_config else None
+        self.non_interactive = non_interactive
 
     def prompt_execution_role(self) -> Optional[str]:
         """Prompt for execution role. Returns role name/ARN or None for auto-creation."""
+        if self.non_interactive:
+            _print_success("Will auto-create execution role")
+            return None
+
         console.print("\n🔐 [cyan]Execution Role[/cyan]")
         console.print(
             "[dim]Press Enter to auto-create execution role, or provide execution role ARN/name to use existing[/dim]"
@@ -43,6 +49,10 @@ class ConfigurationManager:
 
     def prompt_ecr_repository(self) -> tuple[Optional[str], bool]:
         """Prompt for ECR repository. Returns (repository, auto_create_flag)."""
+        if self.non_interactive:
+            _print_success("Will auto-create ECR repository")
+            return None, True
+
         console.print("\n🏗️  [cyan]ECR Repository[/cyan]")
         console.print(
             "[dim]Press Enter to auto-create ECR repository, or provide ECR Repository URI to use existing[/dim]"
@@ -63,6 +73,10 @@ class ConfigurationManager:
 
     def prompt_oauth_config(self) -> Optional[dict]:
         """Prompt for OAuth configuration. Returns OAuth config dict or None."""
+        if self.non_interactive:
+            _print_success("Using default IAM authorization")
+            return None
+
         console.print("\n🔐 [cyan]Authorization Configuration[/cyan]")
         console.print("[dim]By default, Bedrock AgentCore uses IAM authorization.[/dim]")
 
@@ -131,3 +145,56 @@ class ConfigurationManager:
 
         _print_success("OAuth authorizer configuration created")
         return config
+
+    def prompt_request_header_allowlist(self) -> Optional[dict]:
+        """Prompt for request header allowlist configuration. Returns allowlist config dict or None."""
+        if self.non_interactive:
+            _print_success("Using default request header configuration")
+            return None
+
+        console.print("\n🔒 [cyan]Request Header Allowlist[/cyan]")
+        console.print("[dim]Configure which request headers are allowed to pass through to your agent.[/dim]")
+        console.print("[dim]Common headers: Authorization, X-Amzn-Bedrock-AgentCore-Runtime-Custom-*[/dim]")
+
+        # Get existing allowlist values
+        existing_headers = ""
+        if (
+            self.existing_config
+            and self.existing_config.request_header_configuration
+            and "requestHeaderAllowlist" in self.existing_config.request_header_configuration
+        ):
+            existing_headers = ",".join(self.existing_config.request_header_configuration["requestHeaderAllowlist"])
+
+        allowlist_default = "yes" if existing_headers else "no"
+        response = _prompt_with_default("Configure request header allowlist? (yes/no)", allowlist_default)
+
+        if response.lower() in ["yes", "y"]:
+            return self._configure_request_header_allowlist(existing_headers)
+        else:
+            _print_success("Using default request header configuration")
+            return None
+
+    def _configure_request_header_allowlist(self, existing_headers: str = "") -> dict:
+        """Configure request header allowlist and return config dict."""
+        console.print("\n📋 [cyan]Request Header Allowlist Configuration[/cyan]")
+
+        # Show existing config if available
+        if existing_headers:
+            console.print(f"[dim]Previously configured: {existing_headers}[/dim]")
+
+        # Prompt for headers
+        default_headers = existing_headers or "Authorization,X-Amzn-Bedrock-AgentCore-Runtime-Custom-*"
+        headers_input = _prompt_with_default("Enter allowed request headers (comma-separated)", default_headers)
+
+        if not headers_input:
+            _handle_error("At least one request header must be specified for allowlist configuration")
+
+        # Parse and validate headers
+        headers = [header.strip() for header in headers_input.split(",") if header.strip()]
+
+        if not headers:
+            _handle_error("Empty request header allowlist provided")
+
+        _print_success(f"Request header allowlist configured with {len(headers)} headers")
+
+        return {"requestHeaderAllowlist": headers}

bedrock_agentcore_starter_toolkit/operations/gateway/client.py

@@ -176,6 +176,145 @@ class GatewayClient:
         self.logger.info("\n✅Target is ready")
         return target
 
+    def fix_iam_permissions(self, gateway: dict) -> None:
+        """Fix IAM role trust policy for the gateway.
+
+        :param gateway: the gateway dict containing roleArn
+        """
+        # Check for None gateway
+        if gateway is None:
+            return
+
+        # Check for missing roleArn
+        role_arn = gateway.get("roleArn")
+        if not role_arn:
+            return
+
+        sts = boto3.client("sts")
+        iam = boto3.client("iam")
+
+        account_id = sts.get_caller_identity()["Account"]
+        role_name = role_arn.split("/")[-1]
+
+        # Update trust policy
+        trust_policy = {
+            "Version": "2012-10-17",
+            "Statement": [
+                {
+                    "Effect": "Allow",
+                    "Principal": {"Service": "bedrock-agentcore.amazonaws.com"},
+                    "Action": "sts:AssumeRole",
+                    "Condition": {
+                        "StringEquals": {"aws:SourceAccount": account_id},
+                        "ArnLike": {"aws:SourceArn": f"arn:aws:bedrock-agentcore:{self.region}:{account_id}:*"},
+                    },
+                }
+            ],
+        }
+
+        try:
+            iam.update_assume_role_policy(RoleName=role_name, PolicyDocument=json.dumps(trust_policy))
+
+            # Add Lambda permissions
+            iam.put_role_policy(
+                RoleName=role_name,
+                PolicyName="LambdaInvokePolicy",
+                PolicyDocument=json.dumps(
+                    {
+                        "Version": "2012-10-17",
+                        "Statement": [
+                            {
+                                "Effect": "Allow",
+                                "Action": ["lambda:InvokeFunction"],
+                                "Resource": (
+                                    f"arn:aws:lambda:{self.region}:{account_id}:function:AgentCoreLambdaTestFunction"
+                                ),
+                            }
+                        ],
+                    }
+                ),
+            )
+            self.logger.info("✓ Fixed IAM permissions for Gateway")
+        except Exception as e:
+            self.logger.warning("⚠️ IAM role update failed: %s. Continuing with best effort.", str(e))
+
+    def cleanup_gateway(self, gateway_id: str, client_info: Optional[Dict] = None) -> None:
+        """Remove all resources associated with a gateway.
+
+        :param gateway_id: the ID of the gateway to clean up
+        :param client_info: optional Cognito client info for cleanup
+        """
+        self.logger.info("🧹 Cleaning up Gateway resources...")
+
+        gateway_client = self.client
+
+        # Step 1: List and delete all targets
+        self.logger.info("  • Finding targets for gateway: %s", gateway_id)
+
+        try:
+            response = gateway_client.list_gateway_targets(gatewayIdentifier=gateway_id)
+            # API returns targets in 'items' field
+            targets = response.get("items", [])
+            self.logger.info("    Found %s targets to delete", len(targets))
+
+            for target in targets:
+                target_id = target["targetId"]
+                self.logger.info("  • Deleting target: %s", target_id)
+                try:
+                    gateway_client.delete_gateway_target(gatewayIdentifier=gateway_id, targetId=target_id)
+                    self.logger.info("    ✓ Target deletion initiated: %s", target_id)
+                    # Wait for deletion to complete
+                    time.sleep(5)
+                except Exception as e:
+                    self.logger.warning("    ⚠️ Error deleting target %s: %s", target_id, str(e))
+
+            # Verify all targets are deleted
+            self.logger.info("  • Verifying targets deletion...")
+            time.sleep(5)  # Additional wait
+            verify_response = gateway_client.list_gateway_targets(gatewayIdentifier=gateway_id)
+            remaining_targets = verify_response.get("items", [])
+            if remaining_targets:
+                self.logger.warning("    ⚠️ %s targets still remain", len(remaining_targets))
+            else:
+                self.logger.info("    ✓ All targets deleted")
+
+        except Exception as e:
+            self.logger.warning("    ⚠️ Error managing targets: %s", str(e))
+
+        # Step 2: Delete the gateway
+        try:
+            self.logger.info("  • Deleting gateway: %s", gateway_id)
+            gateway_client.delete_gateway(gatewayIdentifier=gateway_id)
+            self.logger.info("    ✓ Gateway deleted: %s", gateway_id)
+        except Exception as e:
+            self.logger.warning("    ⚠️ Error deleting gateway: %s", str(e))
+
+        # Step 3: Delete Cognito resources if provided
+        if client_info and "user_pool_id" in client_info:
+            cognito = boto3.client("cognito-idp", region_name=self.region)
+            user_pool_id = client_info["user_pool_id"]
+
+            # Delete domain first
+            if "domain_prefix" in client_info:
+                domain_prefix = client_info["domain_prefix"]
+                self.logger.info("  • Deleting Cognito domain: %s", domain_prefix)
+                try:
+                    cognito.delete_user_pool_domain(UserPoolId=user_pool_id, Domain=domain_prefix)
+                    self.logger.info("    ✓ Cognito domain deleted")
+                    time.sleep(5)  # Wait for domain deletion
+                except Exception as e:
+                    self.logger.warning("    ⚠️ Error deleting Cognito domain: %s", str(e))
+
+            # Now delete the user pool
+            self.logger.info("  • Deleting Cognito user pool: %s", user_pool_id)
+            try:
+                cognito.delete_user_pool(UserPoolId=user_pool_id)
+                self.logger.info("    ✓ Cognito user pool deleted")
+            except Exception as e:
+                self.logger.warning("    ⚠️ Error deleting Cognito user pool: %s", str(e))
+
+        self.logger.info("✅ Cleanup complete")
+
     def __handle_lambda_target_creation(self, role_arn: str) -> Dict[str, Any]:
         """Create a test lambda.
 

bedrock_agentcore_starter_toolkit/operations/memory/__init__.py

@@ -0,0 +1 @@
+"""BedrockAgentCore Starter Toolkit cli memory package."""

bedrock_agentcore_starter_toolkit/operations/memory/constants.py

@@ -0,0 +1,98 @@
+"""Constants for Bedrock AgentCore Memory SDK."""
+
+from enum import Enum
+from typing import Dict, List, Optional
+
+
+class StrategyType(Enum):
+    """Memory strategy types with integrated wrapper key and type methods."""
+
+    SEMANTIC = "semanticMemoryStrategy"
+    SUMMARY = "summaryMemoryStrategy"
+    USER_PREFERENCE = "userPreferenceMemoryStrategy"
+    CUSTOM = "customMemoryStrategy"
+
+    def extraction_wrapper_key(self) -> Optional[str]:
+        """Get the extraction wrapper key for this strategy type."""
+        extraction_keys = {
+            StrategyType.SEMANTIC: "semanticExtractionConfiguration",
+            StrategyType.USER_PREFERENCE: "userPreferenceExtractionConfiguration",
+        }
+        return extraction_keys.get(self)
+
+    def consolidation_wrapper_key(self) -> Optional[str]:
+        """Get the consolidation wrapper key for this strategy type."""
+        # Only SUMMARY strategy has a consolidation wrapper key
+        if self == StrategyType.SUMMARY:
+            return "summaryConsolidationConfiguration"
+        return None
+
+    def get_memory_strategy(self) -> str:
+        """Get the internal memory strategy type string."""
+        strategy_mapping = {
+            StrategyType.SEMANTIC: "SEMANTIC",
+            StrategyType.SUMMARY: "SUMMARIZATION",
+            StrategyType.USER_PREFERENCE: "USER_PREFERENCE",
+            StrategyType.CUSTOM: "CUSTOM",
+        }
+        return strategy_mapping[self]
+
+    def get_override_type(self) -> Optional[str]:
+        """Get the override type for custom strategies."""
+        # This method is primarily for CUSTOM strategy type
+        # The actual override type would be determined by context
+        if self == StrategyType.CUSTOM:
+            return "CUSTOM_OVERRIDE"  # Base type, specific override determined by usage
+        return None
+
+
+class OverrideType(Enum):
+    """Custom strategy override types."""
+
+    SEMANTIC_OVERRIDE = "SEMANTIC_OVERRIDE"
+    SUMMARY_OVERRIDE = "SUMMARY_OVERRIDE"
+    USER_PREFERENCE_OVERRIDE = "USER_PREFERENCE_OVERRIDE"
+
+    def extraction_wrapper_key(self) -> Optional[str]:
+        """Get the extraction wrapper key for this override type."""
+        extraction_keys = {
+            OverrideType.SEMANTIC_OVERRIDE: "semanticExtractionOverride",
+            OverrideType.USER_PREFERENCE_OVERRIDE: "userPreferenceExtractionOverride",
+        }
+        return extraction_keys.get(self)
+
+    def consolidation_wrapper_key(self) -> Optional[str]:
+        """Get the consolidation wrapper key for this override type."""
+        consolidation_keys = {
+            OverrideType.SEMANTIC_OVERRIDE: "semanticConsolidationOverride",
+            OverrideType.SUMMARY_OVERRIDE: "summaryConsolidationOverride",
+            OverrideType.USER_PREFERENCE_OVERRIDE: "userPreferenceConsolidationOverride",
+        }
+        return consolidation_keys.get(self)
+
+
+class MemoryStatus(Enum):
+    """Memory resource statuses."""
+
+    CREATING = "CREATING"
+    ACTIVE = "ACTIVE"
+    FAILED = "FAILED"
+    UPDATING = "UPDATING"
+    DELETING = "DELETING"
+
+
+class MemoryStrategyStatus(Enum):
+    """Memory strategy statuses (new from API update)."""
+
+    CREATING = "CREATING"
+    ACTIVE = "ACTIVE"
+    DELETING = "DELETING"
+    FAILED = "FAILED"
+
+
+# Default namespaces for each strategy type
+DEFAULT_NAMESPACES: Dict[StrategyType, List[str]] = {
+    StrategyType.SEMANTIC: ["/actor/{actorId}/strategy/{strategyId}/{sessionId}"],
+    StrategyType.SUMMARY: ["/actor/{actorId}/strategy/{strategyId}/{sessionId}"],
+    StrategyType.USER_PREFERENCE: ["/actor/{actorId}/strategy/{strategyId}"],
+}