bedrock-agentcore-starter-toolkit 0.1.12__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

@@ -176,6 +176,13 @@ 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)"),
@@ -243,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,
@@ -254,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,
@@ -264,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"
@@ -277,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]",
@@ -582,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(
@@ -593,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"
@@ -625,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,
@@ -634,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

@@ -145,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/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}"],
+}