bedrock-agentcore-starter-toolkit 0.1.25__py3-none-any.whl → 0.1.27__py3-none-any.whl

bedrock_agentcore_starter_toolkit/operations/runtime/configure.py

@@ -4,11 +4,13 @@ import logging
 import os
 import re
 from pathlib import Path
-from typing import Any, Dict, Literal, Optional, Tuple
+from typing import Any, Dict, List, Literal, Optional, Tuple
+
+import boto3
 
 from ...cli.runtime.configuration_manager import ConfigurationManager
 from ...services.ecr import get_account_id, get_region
-from ...utils.runtime.config import merge_agent_config, save_config
+from ...utils.runtime.config import load_config_if_exists, merge_agent_config, save_config
 from ...utils.runtime.container import ContainerRuntime
 from ...utils.runtime.entrypoint import detect_dependencies
 from ...utils.runtime.schema import (
@@ -16,8 +18,10 @@ from ...utils.runtime.schema import (
     BedrockAgentCoreAgentSchema,
     BedrockAgentCoreDeploymentInfo,
     CodeBuildConfig,
+    LifecycleConfiguration,
     MemoryConfig,
     NetworkConfiguration,
+    NetworkModeConfig,
     ObservabilityConfig,
     ProtocolConfiguration,
 )
@@ -133,7 +137,7 @@ def configure_bedrock_agentcore(
     auto_create_ecr: bool = True,
     auto_create_execution_role: bool = True,
     enable_observability: bool = True,
-    memory_mode: Literal["NO_MEMORY", "STM_ONLY", "STM_AND_LTM"] = "STM_ONLY",
+    memory_mode: Literal["NO_MEMORY", "STM_ONLY", "STM_AND_LTM"] = "NO_MEMORY",
     requirements_file: Optional[str] = None,
     authorizer_configuration: Optional[Dict[str, Any]] = None,
     request_header_configuration: Optional[Dict[str, Any]] = None,
@@ -142,6 +146,11 @@ def configure_bedrock_agentcore(
     protocol: Optional[str] = None,
     non_interactive: bool = False,
     source_path: Optional[str] = None,
+    vpc_enabled: bool = False,
+    vpc_subnets: Optional[List[str]] = None,
+    vpc_security_groups: Optional[List[str]] = None,
+    idle_timeout: Optional[int] = None,
+    max_lifetime: Optional[int] = None,
 ) -> ConfigureResult:
     """Configure Bedrock AgentCore application with deployment settings.
 
@@ -164,6 +173,13 @@ def configure_bedrock_agentcore(
         protocol: agent server protocol, must be either HTTP or MCP or A2A
         non_interactive: Skip interactive prompts and use defaults
         source_path: Optional path to agent source code directory
+        vpc_enabled: Whether to enable VPC networking mode
+        vpc_subnets: List of subnet IDs for VPC mode
+        vpc_security_groups: List of security group IDs for VPC mode
+        idle_timeout: Idle runtime session timeout in seconds (60-28800).
+            If not specified, AWS API default (900s / 15 minutes) is used.
+        max_lifetime: Maximum instance lifetime in seconds (60-28800).
+            If not specified, AWS API default (28800s / 8 hours) is used.
 
     Returns:
         ConfigureResult model with configuration details
@@ -244,7 +260,7 @@ def configure_bedrock_agentcore(
         else:  # STM_ONLY
             log.info("Memory configuration: Short-term memory only")
     else:
-        # Interactive mode: prompt user (only if memory not explicitly disabled)
+        # Interactive mode - let user choose
         action, value = config_manager.prompt_memory_selection()
 
         if action == "USE_EXISTING":
@@ -274,6 +290,21 @@ def configure_bedrock_agentcore(
     memory_id = None
     memory_name = None
 
+    # Handle lifecycle configuration
+    lifecycle_config = LifecycleConfiguration()
+    if idle_timeout is not None or max_lifetime is not None:
+        lifecycle_config = LifecycleConfiguration(
+            idle_runtime_session_timeout=idle_timeout,
+            max_lifetime=max_lifetime,
+        )
+
+        if verbose:
+            log.debug("Lifecycle configuration:")
+            if idle_timeout:
+                log.debug("  Idle timeout: %ds (%d minutes)", idle_timeout, idle_timeout / 60)
+            if max_lifetime:
+                log.debug("  Max lifetime: %ds (%d hours)", max_lifetime, max_lifetime / 3600)
+
     if config_path.exists():
         try:
             from ...utils.runtime.config import load_config
@@ -305,6 +336,42 @@ def configure_bedrock_agentcore(
         if verbose and execution_role_arn:
             log.debug("Using same role for CodeBuild: %s", codebuild_execution_role_arn)
 
+    if vpc_enabled:
+        if not vpc_subnets or not vpc_security_groups:
+            raise ValueError("VPC mode requires both subnets and security groups")
+
+        for subnet_id in vpc_subnets:
+            if not subnet_id.startswith("subnet-"):
+                raise ValueError(
+                    f"Invalid subnet ID format: {subnet_id}\nSubnet IDs must start with 'subnet-' (e.g., subnet-abc123)"
+                )
+            if len(subnet_id) < 15:  # "subnet-" (7) + 8 chars = 15
+                raise ValueError(
+                    f"Invalid subnet ID format: {subnet_id}\nSubnet ID is too short. Expected format: subnet-xxxxxxxx"
+                )
+
+        # Validate security group IDs format
+        for sg_id in vpc_security_groups:
+            if not sg_id.startswith("sg-"):
+                raise ValueError(
+                    f"Invalid security group ID format: {sg_id}\n"
+                    f"Security group IDs must start with 'sg-' (e.g., sg-abc123)"
+                )
+            if len(sg_id) < 11:  # "sg-" (3) + 8 chars = 11
+                raise ValueError(
+                    f"Invalid security group ID format: {sg_id}\n"
+                    f"Security group ID is too short. Expected format: sg-xxxxxxxx"
+                )
+
+        network_config = NetworkConfiguration(
+            network_mode="VPC",
+            network_mode_config=NetworkModeConfig(subnets=vpc_subnets, security_groups=vpc_security_groups),
+        )
+        log.info("Network mode: VPC with %d subnets and %d security groups", len(vpc_subnets), len(vpc_security_groups))
+    else:
+        network_config = NetworkConfiguration(network_mode="PUBLIC")
+        log.info("Network mode: PUBLIC")
+
     # Generate Dockerfile and .dockerignore
     bedrock_agentcore_name = None
     # Try to find the variable name for the Bedrock AgentCore instance in the file
@@ -332,6 +399,11 @@ def configure_bedrock_agentcore(
     else:
         dockerfile_output_dir = build_dir
 
+    if memory_config.mode == "NO_MEMORY":
+        memory_id = None
+        memory_name = None
+        log.debug("Cleared memory_id/name for Dockerfile generation (memory disabled)")
+
     # Generate Dockerfile in the correct location (no moving needed)
     dockerfile_path = runtime.generate_dockerfile(
         entrypoint_path,
@@ -374,6 +446,32 @@ def configure_bedrock_agentcore(
         log.debug("Agent name from BedrockAgentCoreApp: %s", agent_name)
         log.debug("Config path: %s", config_path)
 
+    existing_project_config = load_config_if_exists(config_path)
+
+    if existing_project_config and agent_name in existing_project_config.agents:
+        existing_agent = existing_project_config.agents[agent_name]
+        existing_network = existing_agent.aws.network_configuration
+
+        # Import validation helper
+        from .vpc_validation import check_network_immutability
+
+        # Check if network config is being changed
+        error = check_network_immutability(
+            existing_network_mode=existing_network.network_mode,
+            existing_subnets=existing_network.network_mode_config.subnets
+            if existing_network.network_mode_config
+            else None,
+            existing_security_groups=existing_network.network_mode_config.security_groups
+            if existing_network.network_mode_config
+            else None,
+            new_network_mode="VPC" if vpc_enabled else "PUBLIC",
+            new_subnets=vpc_subnets,
+            new_security_groups=vpc_security_groups,
+        )
+
+        if error:
+            raise ValueError(error)
+
     # Convert to POSIX for cross-platform compatibility
     entrypoint_path_str = entrypoint_path.as_posix()
 
@@ -418,9 +516,10 @@ def configure_bedrock_agentcore(
             region=region,
             ecr_repository=ecr_repository,
             ecr_auto_create=ecr_auto_create_value,
-            network_configuration=NetworkConfiguration(network_mode="PUBLIC"),
+            network_configuration=network_config,
             protocol_configuration=ProtocolConfiguration(server_protocol=protocol or "HTTP"),
             observability=ObservabilityConfig(enabled=enable_observability),
+            lifecycle_configuration=lifecycle_config,
         ),
         bedrock_agentcore=BedrockAgentCoreDeploymentInfo(),
         codebuild=CodeBuildConfig(
@@ -438,6 +537,18 @@ def configure_bedrock_agentcore(
     if verbose:
         log.debug("Configuration saved with agent: %s", agent_name)
 
+    # Get VPC ID for display if VPC mode
+    vpc_id = None
+    if vpc_enabled and vpc_subnets:
+        try:
+            session = boto3.Session(region_name=region)
+            ec2_client = session.client("ec2", region_name=region)
+            subnet_response = ec2_client.describe_subnets(SubnetIds=[vpc_subnets[0]])
+            if subnet_response["Subnets"]:
+                vpc_id = subnet_response["Subnets"][0]["VpcId"]
+        except Exception:
+            pass  # nosec B110
+
     return ConfigureResult(
         config_path=config_path,
         dockerfile_path=dockerfile_path,
@@ -448,6 +559,10 @@ def configure_bedrock_agentcore(
         execution_role=execution_role_arn,
         ecr_repository=ecr_repository,
         auto_create_ecr=auto_create_ecr and not ecr_repository,
+        network_mode="VPC" if vpc_enabled else "PUBLIC",
+        network_subnets=vpc_subnets if vpc_enabled else None,
+        network_security_groups=vpc_security_groups if vpc_enabled else None,
+        network_vpc_id=vpc_id,
     )
 
 

bedrock_agentcore_starter_toolkit/operations/runtime/invoke.py

@@ -7,6 +7,7 @@ from typing import Any, Optional
 
 from bedrock_agentcore.services.identity import IdentityClient
 
+from ...operations.identity.oauth2_callback_server import WORKLOAD_USER_ID, BedrockAgentCoreIdentity3loCallback
 from ...services.runtime import BedrockAgentCoreClient, generate_session_id
 from ...utils.runtime.config import load_config, save_config
 from ...utils.runtime.schema import BedrockAgentCoreConfigSchema
@@ -30,59 +31,6 @@ def invoke_bedrock_agentcore(
     project_config = load_config(config_path)
     agent_config = project_config.get_agent_config(agent_name)
 
-    # Check memory status on first invoke if memory is enabled (STM or LTM)
-    if (
-        agent_config.memory
-        and agent_config.memory.mode != "NO_MEMORY"
-        and agent_config.memory.memory_id
-        and not agent_config.memory.first_invoke_memory_check_done
-    ):
-        try:
-            from ...operations.memory.constants import MemoryStatus
-            from ...operations.memory.manager import MemoryManager
-
-            memory_manager = MemoryManager(region_name=agent_config.aws.region)
-            memory_status = memory_manager.get_memory_status(agent_config.memory.memory_id)
-
-            if memory_status != MemoryStatus.ACTIVE.value:
-                # Determine memory type for better messaging
-                memory_type = "Memory"
-                if agent_config.memory.has_ltm:
-                    memory_type = "Long-term memory"
-                    time_estimate = "60-180 seconds"
-                else:
-                    memory_type = "Short-term memory"
-                    time_estimate = "30-90 seconds"
-
-                # Provide graceful error message
-                error_message = (
-                    f"Memory is still provisioning (current status: {memory_status}). "
-                    f"{memory_type} takes {time_estimate} to activate.\n\n"
-                    f"Please wait and check status with:\n"
-                    f"  agentcore status{f' --agent {agent_name}' if agent_name else ''}"
-                )
-
-                # Log the message for visibility
-                log.warning("Memory not yet active for agent '%s': %s", agent_config.name, memory_status)
-
-                raise ValueError(error_message)
-
-            # Memory is active, mark check as done
-            agent_config.memory.first_invoke_memory_check_done = True
-            project_config.agents[agent_config.name] = agent_config
-            save_config(project_config, config_path)
-            log.info("Memory is active, proceeding with invoke")
-
-        except ImportError as e:
-            log.error("Failed to import MemoryManager: %s", e)
-            # Continue without check if import fails
-        except Exception as e:
-            # If it's our ValueError, re-raise it
-            if "Memory is still provisioning" in str(e):
-                raise
-            # For other errors, log but continue
-            log.warning("Could not check memory status: %s", e)
-
     # Log which agent is being invoked
     mode = "locally" if local_mode else "via cloud endpoint"
     log.debug("Invoking BedrockAgentCore agent '%s' %s", agent_config.name, mode)
@@ -121,9 +69,19 @@ def invoke_bedrock_agentcore(
             workload_name=workload_name, user_token=bearer_token, user_id=user_id
         )["workloadAccessToken"]
 
+        agent_config.oauth_configuration[WORKLOAD_USER_ID] = user_id  # type: ignore : populated by _get_workload_name(...)
+        save_config(project_config, config_path)
+
+        oauth2_callback_url = BedrockAgentCoreIdentity3loCallback.get_oauth2_callback_endpoint()
+        _update_workload_identity_with_oauth2_callback_url(
+            identity_client, workload_name=workload_name, oauth2_callback_url=oauth2_callback_url
+        )
+
         # TODO: store and read port config of local running container
         client = LocalBedrockAgentCoreClient("http://127.0.0.1:8080")
-        response = client.invoke_endpoint(session_id, payload_str, workload_access_token, custom_headers)
+        response = client.invoke_endpoint(
+            session_id, payload_str, workload_access_token, oauth2_callback_url, custom_headers
+        )
 
     else:
         if not agent_arn:
@@ -163,6 +121,24 @@ def invoke_bedrock_agentcore(
     )
 
 
+def _update_workload_identity_with_oauth2_callback_url(
+    identity_client: IdentityClient,
+    workload_name: str,
+    oauth2_callback_url: str,
+) -> None:
+    workload_identity = identity_client.get_workload_identity(name=workload_name)
+    allowed_resource_oauth_2_return_urls = workload_identity.get("allowedResourceOauth2ReturnUrls") or []
+    if oauth2_callback_url in allowed_resource_oauth_2_return_urls:
+        return
+
+    log.info("Updating workload %s with callback url %s", workload_name, oauth2_callback_url)
+
+    identity_client.update_workload_identity(
+        name=workload_name,
+        allowed_resource_oauth_2_return_urls=[*allowed_resource_oauth_2_return_urls, oauth2_callback_url],
+    )
+
+
 def _get_workload_name(
     project_config: BedrockAgentCoreConfigSchema,
     project_config_path: Path,