chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,606 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""
|
||||
Cross-platform frontend deployment script for FAST.
|
||||
|
||||
Deploys the React frontend to AWS Amplify by:
|
||||
1. Fetching configuration from CDK stack outputs
|
||||
2. Generating aws-exports.json
|
||||
3. Building the frontend
|
||||
4. Packaging and uploading to S3
|
||||
5. Triggering Amplify deployment
|
||||
|
||||
Requires: Python 3.11+, AWS CLI, npm, Node.js
|
||||
No external Python dependencies - uses standard library only.
|
||||
"""
|
||||
|
||||
import atexit
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess # nosec B404 - subprocess used securely with explicit parameters
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
|
||||
# Minimum Python version check
|
||||
if sys.version_info < (3, 8):
|
||||
print("Error: Python 3.8 or higher is required")
|
||||
sys.exit(1)
|
||||
|
||||
# Constants
|
||||
BRANCH_NAME = "main"
|
||||
NEXT_BUILD_DIR = "build"
|
||||
CLEANUP_FILES: list = []
|
||||
|
||||
|
||||
# --- Logging helpers ---
|
||||
|
||||
|
||||
def log_info(message: str) -> None:
|
||||
"""Print an info message."""
|
||||
print(f"ℹ {message}")
|
||||
|
||||
|
||||
def log_success(message: str) -> None:
|
||||
"""Print a success message."""
|
||||
print(f"✓ {message}")
|
||||
|
||||
|
||||
def log_error(message: str) -> None:
|
||||
"""Print an error message to stderr."""
|
||||
print(f"✗ {message}", file=sys.stderr)
|
||||
|
||||
|
||||
def log_warning(message: str) -> None:
|
||||
"""Print a warning message."""
|
||||
print(f"⚠ {message}")
|
||||
|
||||
|
||||
# --- Utility functions ---
|
||||
|
||||
|
||||
def cleanup() -> None:
|
||||
"""Remove temporary files created during deployment."""
|
||||
for filepath in CLEANUP_FILES:
|
||||
if os.path.exists(filepath):
|
||||
os.remove(filepath)
|
||||
log_info(f"Cleaned up {filepath}")
|
||||
|
||||
|
||||
def run_command(
|
||||
command: list,
|
||||
capture_output: bool = True,
|
||||
check: bool = True,
|
||||
cwd: Optional[str] = None,
|
||||
) -> subprocess.CompletedProcess:
|
||||
"""
|
||||
Execute a command securely via subprocess.
|
||||
|
||||
Args:
|
||||
command: List of command arguments
|
||||
capture_output: Whether to capture stdout/stderr
|
||||
check: Whether to raise on non-zero exit
|
||||
cwd: Working directory for the command
|
||||
|
||||
Returns:
|
||||
CompletedProcess instance with command results
|
||||
"""
|
||||
return subprocess.run( # nosec B603 - command constructed from safe list
|
||||
command,
|
||||
capture_output=capture_output,
|
||||
text=True,
|
||||
check=check,
|
||||
shell=False,
|
||||
timeout=300,
|
||||
cwd=cwd,
|
||||
)
|
||||
|
||||
|
||||
def check_prerequisite(command: str) -> bool:
|
||||
"""
|
||||
Check if a command is available in PATH.
|
||||
|
||||
Args:
|
||||
command: Name of the command to check
|
||||
|
||||
Returns:
|
||||
True if command exists, False otherwise
|
||||
"""
|
||||
return shutil.which(command) is not None
|
||||
|
||||
|
||||
def parse_config_yaml(config_path: Path) -> Dict[str, str]:
|
||||
"""
|
||||
Parse config.yaml using regex (no PyYAML dependency).
|
||||
|
||||
Args:
|
||||
config_path: Path to config.yaml file
|
||||
|
||||
Returns:
|
||||
Dictionary with stack_name_base and pattern values
|
||||
"""
|
||||
config = {"stack_name_base": "", "pattern": "langgraph-single-agent"}
|
||||
|
||||
if not config_path.exists():
|
||||
return config
|
||||
|
||||
content = config_path.read_text()
|
||||
|
||||
# Extract stack_name_base
|
||||
match = re.search(r"^stack_name_base:\s*(\S+)", content, re.MULTILINE)
|
||||
if match:
|
||||
config["stack_name_base"] = match.group(1).strip("\"'")
|
||||
|
||||
# Extract pattern from backend section
|
||||
match = re.search(r"pattern:\s*(\S+)", content)
|
||||
if match:
|
||||
config["pattern"] = match.group(1).split("#")[0].strip().strip("\"'")
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def get_file_size_human(filepath: str) -> str:
|
||||
"""
|
||||
Get human-readable file size.
|
||||
|
||||
Args:
|
||||
filepath: Path to the file
|
||||
|
||||
Returns:
|
||||
Human-readable size string (e.g., "1.5MB")
|
||||
"""
|
||||
size = os.path.getsize(filepath)
|
||||
for unit in ["B", "KB", "MB", "GB"]:
|
||||
if size < 1024:
|
||||
return f"{size:.1f}{unit}"
|
||||
size /= 1024
|
||||
return f"{size:.1f}TB"
|
||||
|
||||
|
||||
# --- AWS CLI wrappers ---
|
||||
|
||||
|
||||
def get_stack_outputs(stack_name: str) -> Dict[str, str]:
|
||||
"""
|
||||
Fetch CloudFormation stack outputs via AWS CLI.
|
||||
|
||||
Args:
|
||||
stack_name: Name of the CloudFormation stack
|
||||
|
||||
Returns:
|
||||
Dictionary mapping output keys to values
|
||||
"""
|
||||
result = run_command(
|
||||
[
|
||||
"aws",
|
||||
"cloudformation",
|
||||
"describe-stacks",
|
||||
"--stack-name",
|
||||
stack_name,
|
||||
"--output",
|
||||
"json",
|
||||
]
|
||||
)
|
||||
|
||||
stack_data = json.loads(result.stdout)
|
||||
stacks = stack_data.get("Stacks", [])
|
||||
if not stacks:
|
||||
raise ValueError(f"Stack '{stack_name}' not found or has no data")
|
||||
outputs = stacks[0].get("Outputs", [])
|
||||
|
||||
return {o["OutputKey"]: o["OutputValue"] for o in outputs}
|
||||
|
||||
|
||||
def get_stack_region(stack_name: str) -> str:
|
||||
"""
|
||||
Get the AWS region from stack ARN.
|
||||
|
||||
Args:
|
||||
stack_name: Name of the CloudFormation stack
|
||||
|
||||
Returns:
|
||||
AWS region string
|
||||
"""
|
||||
result = run_command(
|
||||
[
|
||||
"aws",
|
||||
"cloudformation",
|
||||
"describe-stacks",
|
||||
"--stack-name",
|
||||
stack_name,
|
||||
"--output",
|
||||
"json",
|
||||
]
|
||||
)
|
||||
|
||||
stack_data = json.loads(result.stdout)
|
||||
stacks = stack_data.get("Stacks", [])
|
||||
if not stacks:
|
||||
raise ValueError(f"Stack '{stack_name}' not found or has no data")
|
||||
stack_arn = stacks[0]["StackId"]
|
||||
# ARN format: arn:aws:cloudformation:region:account:stack/name/id
|
||||
arn_parts = stack_arn.split(":")
|
||||
if len(arn_parts) < 4:
|
||||
raise ValueError(f"Invalid stack ARN format: {stack_arn}")
|
||||
return arn_parts[3]
|
||||
|
||||
|
||||
def upload_to_s3(local_path: str, bucket: str, key: str) -> None:
|
||||
"""
|
||||
Upload a file to S3 via AWS CLI.
|
||||
|
||||
Args:
|
||||
local_path: Path to local file
|
||||
bucket: S3 bucket name
|
||||
key: S3 object key
|
||||
"""
|
||||
run_command(
|
||||
["aws", "s3", "cp", local_path, f"s3://{bucket}/{key}", "--no-progress"]
|
||||
)
|
||||
|
||||
|
||||
def start_amplify_deployment(app_id: str, branch: str, source_url: str) -> Dict:
|
||||
"""
|
||||
Start an Amplify deployment via AWS CLI.
|
||||
|
||||
Args:
|
||||
app_id: Amplify application ID
|
||||
branch: Branch name to deploy
|
||||
source_url: S3 URL of deployment package
|
||||
|
||||
Returns:
|
||||
Deployment response as dictionary
|
||||
"""
|
||||
result = run_command(
|
||||
[
|
||||
"aws",
|
||||
"amplify",
|
||||
"start-deployment",
|
||||
"--app-id",
|
||||
app_id,
|
||||
"--branch-name",
|
||||
branch,
|
||||
"--source-url",
|
||||
source_url,
|
||||
"--output",
|
||||
"json",
|
||||
]
|
||||
)
|
||||
|
||||
return json.loads(result.stdout)
|
||||
|
||||
|
||||
def get_amplify_job_status(app_id: str, branch: str, job_id: str) -> str:
|
||||
"""
|
||||
Get the status of an Amplify deployment job.
|
||||
|
||||
Args:
|
||||
app_id: Amplify application ID
|
||||
branch: Branch name
|
||||
job_id: Deployment job ID
|
||||
|
||||
Returns:
|
||||
Job status string
|
||||
"""
|
||||
result = run_command(
|
||||
[
|
||||
"aws",
|
||||
"amplify",
|
||||
"get-job",
|
||||
"--app-id",
|
||||
app_id,
|
||||
"--branch-name",
|
||||
branch,
|
||||
"--job-id",
|
||||
job_id,
|
||||
"--output",
|
||||
"json",
|
||||
]
|
||||
)
|
||||
|
||||
return json.loads(result.stdout)["job"]["summary"]["status"]
|
||||
|
||||
|
||||
def get_amplify_app_domain(app_id: str) -> str:
|
||||
"""
|
||||
Get the default domain for an Amplify app.
|
||||
|
||||
Args:
|
||||
app_id: Amplify application ID
|
||||
|
||||
Returns:
|
||||
Default domain string
|
||||
"""
|
||||
result = run_command(
|
||||
[
|
||||
"aws",
|
||||
"amplify",
|
||||
"get-app",
|
||||
"--app-id",
|
||||
app_id,
|
||||
"--query",
|
||||
"app.defaultDomain",
|
||||
"--output",
|
||||
"text",
|
||||
]
|
||||
)
|
||||
|
||||
return result.stdout.strip()
|
||||
|
||||
|
||||
# --- Main deployment logic ---
|
||||
|
||||
|
||||
def generate_aws_exports(
|
||||
stack_name: str,
|
||||
outputs: Dict[str, str],
|
||||
region: str,
|
||||
pattern: str,
|
||||
frontend_dir: Path,
|
||||
) -> None:
|
||||
"""
|
||||
Generate aws-exports.json configuration file.
|
||||
|
||||
Args:
|
||||
stack_name: CloudFormation stack name
|
||||
outputs: Stack outputs dictionary
|
||||
region: AWS region
|
||||
pattern: Agent pattern name
|
||||
frontend_dir: Path to frontend directory
|
||||
"""
|
||||
required = [
|
||||
"CognitoClientId",
|
||||
"CognitoUserPoolId",
|
||||
"AmplifyUrl",
|
||||
"RuntimeArn",
|
||||
"CopilotKitRuntimeUrl",
|
||||
]
|
||||
missing = [k for k in required if k not in outputs]
|
||||
|
||||
if missing:
|
||||
raise ValueError(f"Missing required stack outputs: {', '.join(missing)}")
|
||||
|
||||
aws_exports = {
|
||||
"authority": f"https://cognito-idp.{region}.amazonaws.com/{outputs['CognitoUserPoolId']}",
|
||||
"client_id": outputs["CognitoClientId"],
|
||||
"redirect_uri": outputs["AmplifyUrl"],
|
||||
"post_logout_redirect_uri": outputs["AmplifyUrl"],
|
||||
"response_type": "code",
|
||||
"scope": "email openid profile",
|
||||
"automaticSilentRenew": True,
|
||||
"agentRuntimeArn": outputs["RuntimeArn"],
|
||||
"awsRegion": region,
|
||||
"copilotKitRuntimeUrl": outputs["CopilotKitRuntimeUrl"],
|
||||
"agentPattern": pattern,
|
||||
}
|
||||
|
||||
public_dir = frontend_dir / "public"
|
||||
public_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
output_path = public_dir / "aws-exports.json"
|
||||
output_path.write_text(json.dumps(aws_exports, indent=2))
|
||||
|
||||
log_success(f"Generated aws-exports.json at {output_path}")
|
||||
|
||||
|
||||
def create_deployment_zip(build_dir: Path, output_path: Path) -> None:
|
||||
"""
|
||||
Create a zip archive of the build directory.
|
||||
|
||||
Args:
|
||||
build_dir: Path to the build directory
|
||||
output_path: Path for the output zip file (without .zip extension)
|
||||
"""
|
||||
# shutil.make_archive adds .zip automatically
|
||||
shutil.make_archive(
|
||||
str(output_path.with_suffix("")), "zip", root_dir=str(build_dir)
|
||||
)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""
|
||||
Main deployment function.
|
||||
|
||||
Returns:
|
||||
Exit code (0 for success, 1 for failure)
|
||||
"""
|
||||
atexit.register(cleanup)
|
||||
|
||||
# Determine paths
|
||||
script_dir = Path(__file__).parent.resolve()
|
||||
project_root = script_dir.parent
|
||||
frontend_dir = project_root / "frontend"
|
||||
config_path = project_root / "infra-cdk" / "config.yaml"
|
||||
|
||||
log_info("🚀 Starting frontend deployment process...")
|
||||
print()
|
||||
|
||||
# Validate prerequisites
|
||||
log_info("Validating prerequisites...")
|
||||
prerequisites = ["npm", "aws", "node"]
|
||||
for prereq in prerequisites:
|
||||
if not check_prerequisite(prereq):
|
||||
log_error(f"{prereq} is not installed")
|
||||
return 1
|
||||
log_success("All prerequisites found")
|
||||
|
||||
# Verify AWS credentials are configured
|
||||
log_info("Verifying AWS credentials...")
|
||||
try:
|
||||
run_command(["aws", "sts", "get-caller-identity"], capture_output=True)
|
||||
log_success("AWS credentials configured")
|
||||
except subprocess.CalledProcessError:
|
||||
log_error("AWS credentials not configured or invalid")
|
||||
log_info("Run 'aws configure' to set up your AWS credentials")
|
||||
log_info(
|
||||
"Or set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables"
|
||||
)
|
||||
return 1
|
||||
|
||||
# Get stack name
|
||||
stack_name = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("STACK_NAME")
|
||||
|
||||
if not stack_name:
|
||||
config = parse_config_yaml(config_path)
|
||||
stack_name = config.get("stack_name_base")
|
||||
|
||||
if not stack_name:
|
||||
log_error("Stack name is required")
|
||||
log_info("Usage: python deploy-frontend.py <stack-name>")
|
||||
log_info(" or: STACK_NAME=your-stack ./deploy-frontend.py")
|
||||
return 1
|
||||
|
||||
# Fetch CDK outputs
|
||||
log_info(f"Fetching configuration from CDK stack: {stack_name}")
|
||||
try:
|
||||
outputs = get_stack_outputs(stack_name)
|
||||
region = get_stack_region(stack_name)
|
||||
except subprocess.CalledProcessError as e:
|
||||
log_error(f"Failed to fetch stack outputs: {e.stderr}")
|
||||
return 1
|
||||
except ValueError as e:
|
||||
log_error(str(e))
|
||||
return 1
|
||||
|
||||
# Validate required outputs
|
||||
app_id = outputs.get("AmplifyAppId")
|
||||
deployment_bucket = outputs.get("StagingBucketName")
|
||||
|
||||
if not app_id:
|
||||
log_error("Could not find Amplify App ID in stack outputs")
|
||||
return 1
|
||||
if not deployment_bucket:
|
||||
log_error("Could not find Staging Bucket Name in stack outputs")
|
||||
return 1
|
||||
|
||||
log_success(f"App ID: {app_id}")
|
||||
log_success(f"Staging Bucket: {deployment_bucket}")
|
||||
log_success(f"Region: {region}")
|
||||
|
||||
# Get agent pattern from config
|
||||
config = parse_config_yaml(config_path)
|
||||
pattern = config.get("pattern", "strands-single-agent")
|
||||
log_info(f"Agent pattern: {pattern}")
|
||||
|
||||
# Generate aws-exports.json
|
||||
log_info("Generating aws-exports.json...")
|
||||
try:
|
||||
generate_aws_exports(stack_name, outputs, region, pattern, frontend_dir)
|
||||
except ValueError as e:
|
||||
log_error(str(e))
|
||||
return 1
|
||||
|
||||
# Change to frontend directory
|
||||
os.chdir(frontend_dir)
|
||||
log_info(f"Working directory: {frontend_dir}")
|
||||
|
||||
# Install dependencies if needed
|
||||
node_modules = frontend_dir / "node_modules"
|
||||
package_json = frontend_dir / "package.json"
|
||||
|
||||
if (
|
||||
not node_modules.exists()
|
||||
or package_json.stat().st_mtime > node_modules.stat().st_mtime
|
||||
):
|
||||
log_info("Installing dependencies...")
|
||||
try:
|
||||
run_command(["npm", "install"], capture_output=False)
|
||||
log_success("Dependencies installed")
|
||||
except subprocess.CalledProcessError:
|
||||
log_error("Failed to install dependencies")
|
||||
return 1
|
||||
else:
|
||||
log_success("Dependencies are up to date")
|
||||
|
||||
# Build frontend
|
||||
log_info("Building React app...")
|
||||
try:
|
||||
run_command(["npm", "run", "build"], capture_output=False)
|
||||
log_success("Build completed")
|
||||
except subprocess.CalledProcessError:
|
||||
log_error("Build failed")
|
||||
return 1
|
||||
|
||||
# Verify build directory
|
||||
build_dir = frontend_dir / NEXT_BUILD_DIR
|
||||
if not build_dir.exists():
|
||||
log_error(f"Build directory '{NEXT_BUILD_DIR}' not found")
|
||||
return 1
|
||||
|
||||
# Copy aws-exports.json to build
|
||||
aws_exports_src = frontend_dir / "public" / "aws-exports.json"
|
||||
aws_exports_dst = build_dir / "aws-exports.json"
|
||||
shutil.copy2(aws_exports_src, aws_exports_dst)
|
||||
log_success("Added aws-exports.json to build directory")
|
||||
|
||||
# Create deployment zip
|
||||
log_info("Creating deployment package...")
|
||||
zip_path = frontend_dir / "amplify-deploy.zip"
|
||||
CLEANUP_FILES.append(str(zip_path))
|
||||
|
||||
create_deployment_zip(build_dir, zip_path)
|
||||
zip_size = get_file_size_human(str(zip_path))
|
||||
log_success(f"Package created ({zip_size})")
|
||||
|
||||
# Upload to S3
|
||||
s3_key = f"amplify-deploy-{int(time.time())}.zip"
|
||||
log_info(f"Uploading to S3 (s3://{deployment_bucket}/{s3_key})...")
|
||||
try:
|
||||
upload_to_s3(str(zip_path), deployment_bucket, s3_key)
|
||||
log_success("Upload completed")
|
||||
except subprocess.CalledProcessError as e:
|
||||
log_error(f"S3 upload failed: {e.stderr}")
|
||||
return 1
|
||||
|
||||
# Start Amplify deployment
|
||||
log_info("Starting Amplify deployment...")
|
||||
source_url = f"s3://{deployment_bucket}/{s3_key}"
|
||||
|
||||
try:
|
||||
deployment = start_amplify_deployment(app_id, BRANCH_NAME, source_url)
|
||||
job_id = deployment["jobSummary"]["jobId"]
|
||||
log_success(f"Deployment initiated (Job ID: {job_id})")
|
||||
except subprocess.CalledProcessError as e:
|
||||
log_error(f"Amplify deployment failed: {e.stderr}")
|
||||
return 1
|
||||
|
||||
# Poll deployment status
|
||||
log_info("Monitoring deployment status...")
|
||||
while True:
|
||||
try:
|
||||
status = get_amplify_job_status(app_id, BRANCH_NAME, job_id)
|
||||
except subprocess.CalledProcessError as e:
|
||||
log_error(f"Failed to get deployment status: {e.stderr}")
|
||||
return 1
|
||||
|
||||
print(f" Status: {status}")
|
||||
|
||||
if status == "SUCCEED":
|
||||
log_success("Deployment completed successfully!")
|
||||
break
|
||||
elif status in ("FAILED", "CANCELLED"):
|
||||
log_error(f"Deployment {status.lower()}")
|
||||
return 1
|
||||
|
||||
time.sleep(10)
|
||||
|
||||
# Print final info
|
||||
print()
|
||||
log_info(f"S3 Package: s3://{deployment_bucket}/{s3_key}")
|
||||
log_info("Console: https://console.aws.amazon.com/amplify/apps")
|
||||
try:
|
||||
app_domain = get_amplify_app_domain(app_id)
|
||||
log_info(f"App URL: https://{BRANCH_NAME}.{app_domain}")
|
||||
except subprocess.CalledProcessError:
|
||||
log_warning("Could not retrieve app URL - check Amplify console")
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,4 @@
|
||||
boto3>=1.34.0
|
||||
requests>=2.31.0
|
||||
PyYAML>=6.0.1
|
||||
colorama>=0.4.6
|
||||
+537
@@ -0,0 +1,537 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""
|
||||
Interactive agent chat tester for local and remote agents
|
||||
|
||||
Tests agent invocation with conversation continuity:
|
||||
- Remote mode (default): Chat with deployed agent via Cognito authentication
|
||||
- Local mode (--local): Chat with agent running on localhost:8080
|
||||
- Automatically detects pattern from config.yaml
|
||||
|
||||
Usage:
|
||||
# Remote agent testing (prompts for credentials)
|
||||
uv run scripts/test-agent.py
|
||||
|
||||
# Local agent testing (agent must be running on localhost:8080)
|
||||
uv run scripts/test-agent.py --local
|
||||
|
||||
# Override pattern from config
|
||||
uv run scripts/test-agent.py --pattern strands-single-agent
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import atexit
|
||||
import os
|
||||
import getpass
|
||||
import json
|
||||
import signal
|
||||
import socket
|
||||
import subprocess # nosec B404 - subprocess used securely with explicit parameters
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
|
||||
import requests
|
||||
from colorama import Fore, Style
|
||||
|
||||
# Add scripts directory to path for reliable imports
|
||||
scripts_dir = Path(__file__).parent.parent / "scripts"
|
||||
if str(scripts_dir) not in sys.path:
|
||||
sys.path.insert(0, str(scripts_dir))
|
||||
|
||||
# Import shared utilities
|
||||
from utils import (
|
||||
authenticate_cognito,
|
||||
create_mock_jwt,
|
||||
generate_session_id,
|
||||
get_stack_config,
|
||||
print_msg,
|
||||
print_section,
|
||||
)
|
||||
|
||||
# Global variable to track agent process
|
||||
_agent_process: Optional[subprocess.Popen] = None
|
||||
|
||||
|
||||
def generate_trace_id() -> str:
|
||||
"""
|
||||
Generate X-Amzn-Trace-Id header value for AWS request tracing.
|
||||
|
||||
Returns:
|
||||
str: Trace ID in AWS X-Ray format
|
||||
"""
|
||||
timestamp_hex = format(int(time.time()), "x")
|
||||
return f"1-{timestamp_hex}-{generate_session_id()}"
|
||||
|
||||
|
||||
def check_port_available(port: int = 8080) -> bool:
|
||||
"""
|
||||
Check if a port is available for connection.
|
||||
|
||||
Args:
|
||||
port (int): Port number to check
|
||||
|
||||
Returns:
|
||||
bool: True if port is available, False otherwise
|
||||
"""
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.settimeout(1)
|
||||
try:
|
||||
result = sock.connect_ex(("localhost", port))
|
||||
sock.close()
|
||||
return result == 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def start_local_agent(
|
||||
memory_id: str, region: str, stack_name: str, pattern: str
|
||||
) -> subprocess.Popen:
|
||||
"""
|
||||
Start the local agent in a background process.
|
||||
|
||||
Args:
|
||||
memory_id (str): Memory ID for the agent
|
||||
region (str): AWS region
|
||||
stack_name (str): CloudFormation stack name for SSM parameter lookup
|
||||
pattern (str): Agent pattern name (e.g., 'strands-single-agent', 'langgraph-single-agent')
|
||||
|
||||
Returns:
|
||||
subprocess.Popen: Subprocess object for the running agent
|
||||
"""
|
||||
global _agent_process
|
||||
|
||||
# Map pattern to agent file
|
||||
pattern_files = {
|
||||
"strands-single-agent": "strands_agent.py",
|
||||
"langgraph-single-agent": "langgraph_agent.py",
|
||||
}
|
||||
|
||||
agent_file = pattern_files.get(pattern)
|
||||
if not agent_file:
|
||||
print_msg(f"Unknown pattern: {pattern}", "error")
|
||||
print(f"Available patterns: {', '.join(pattern_files.keys())}")
|
||||
sys.exit(1)
|
||||
|
||||
agent_path = Path(__file__).parent.parent / "agents" / pattern / agent_file
|
||||
|
||||
if not agent_path.exists():
|
||||
print_msg(f"Agent file not found: {agent_path}", "error")
|
||||
sys.exit(1)
|
||||
|
||||
# Security validation: ensure agent_path is within the patterns directory
|
||||
patterns_dir = Path(__file__).parent.parent / "agents"
|
||||
try:
|
||||
agent_path.resolve().relative_to(patterns_dir.resolve())
|
||||
except ValueError:
|
||||
print_msg(
|
||||
f"Security error: Agent path outside patterns directory: {agent_path}",
|
||||
"error",
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
print(f"Starting local agent at {agent_path}...")
|
||||
print(f" Pattern: {pattern}")
|
||||
print(f" Memory ID: {memory_id}")
|
||||
print(f" Region: {region}")
|
||||
print(f" Stack Name: {stack_name}\n")
|
||||
|
||||
requirements_path = agent_path.parent / "requirements.txt"
|
||||
|
||||
# Set up environment variables
|
||||
env = {
|
||||
**dict(subprocess.os.environ),
|
||||
"MEMORY_ID": memory_id,
|
||||
"AWS_DEFAULT_REGION": region,
|
||||
"STACK_NAME": stack_name,
|
||||
"GATEWAY_CREDENTIAL_PROVIDER_NAME": f"{stack_name}-runtime-gateway-auth",
|
||||
"AGUI_ENABLED": "true",
|
||||
"PYTHONPATH": f"{agent_path.parent}{os.pathsep}{agent_path.parent.parent}",
|
||||
}
|
||||
|
||||
# Build command: uv run with requirements if available, else plain python3
|
||||
if requirements_path.exists():
|
||||
cmd = [
|
||||
"uv",
|
||||
"run",
|
||||
"--with-requirements",
|
||||
str(requirements_path),
|
||||
str(agent_path),
|
||||
]
|
||||
else:
|
||||
cmd = ["python3", str(agent_path)]
|
||||
|
||||
# Start agent process
|
||||
try:
|
||||
_agent_process = subprocess.Popen( # nosec B607 B603 - command constructed from validated path, shell=False
|
||||
cmd,
|
||||
env=env,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
shell=False, # Explicitly disable shell
|
||||
)
|
||||
|
||||
# Wait for agent to start (check port becomes available)
|
||||
print("Waiting for agent to start on port 8080...")
|
||||
for i in range(30): # Wait up to 30 seconds
|
||||
if check_port_available(8080):
|
||||
print_msg("Agent started successfully", "success")
|
||||
return _agent_process
|
||||
time.sleep(1)
|
||||
|
||||
print_msg("Agent failed to start (timeout)", "error")
|
||||
if _agent_process.stderr:
|
||||
print(_agent_process.stderr.read())
|
||||
_agent_process.terminate()
|
||||
sys.exit(1)
|
||||
|
||||
except Exception as e:
|
||||
print_msg(f"Failed to start agent: {e}", "error")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def stop_local_agent() -> None:
|
||||
"""Stop the local agent process if running."""
|
||||
global _agent_process
|
||||
if _agent_process:
|
||||
print("\nStopping local agent...")
|
||||
_agent_process.terminate()
|
||||
try:
|
||||
_agent_process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
_agent_process.kill()
|
||||
print_msg("Agent stopped", "success")
|
||||
|
||||
|
||||
# Register cleanup handler
|
||||
atexit.register(stop_local_agent)
|
||||
|
||||
|
||||
def signal_handler(sig, frame):
|
||||
"""Handle interrupt signal."""
|
||||
print("\n")
|
||||
stop_local_agent()
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
|
||||
def invoke_agent(
|
||||
url: str,
|
||||
prompt: str,
|
||||
session_id: str,
|
||||
user_id: str = "local-test-user",
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Invoke agent and print raw streaming events in real-time.
|
||||
|
||||
Args:
|
||||
url (str): Agent endpoint URL
|
||||
prompt (str): User prompt/query
|
||||
session_id (str): Session ID for conversation continuity
|
||||
user_id (str): User ID for mock JWT in local testing only. In remote mode,
|
||||
the real Cognito JWT carries the user identity, user_id is never sent
|
||||
in the payload to prevent prompt injection impersonation.
|
||||
headers (Optional[Dict[str, str]]): Optional HTTP headers
|
||||
"""
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"runtimeSessionId": session_id,
|
||||
}
|
||||
|
||||
if headers is None:
|
||||
# Local mode: generate a mock JWT so the agent can extract user_id
|
||||
# from the Authorization header, matching the production auth flow.
|
||||
mock_token = create_mock_jwt(user_id)
|
||||
headers = {"Authorization": f"Bearer {mock_token}"}
|
||||
headers["Content-Type"] = "application/json"
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
url, headers=headers, json=payload, stream=True, timeout=60
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
print(f"Error: HTTP {response.status_code}: {response.text}")
|
||||
return
|
||||
|
||||
# Parse streaming events and display clean text output
|
||||
print(f"{Fore.GREEN}Agent:{Style.RESET_ALL} ", end="", flush=True)
|
||||
for line in response.iter_lines(decode_unicode=True):
|
||||
if not line or not line.startswith("data: "):
|
||||
continue
|
||||
try:
|
||||
chunk = json.loads(line[6:])
|
||||
|
||||
# LangGraph: AIMessageChunk with content array
|
||||
if chunk.get("type") == "AIMessageChunk" and isinstance(
|
||||
chunk.get("content"), list
|
||||
):
|
||||
for block in chunk["content"]:
|
||||
if block.get("type") == "text" and block.get("text"):
|
||||
print(block["text"], end="", flush=True)
|
||||
elif block.get("type") == "tool_use" and block.get("name"):
|
||||
print(
|
||||
f"\n{Fore.YELLOW}[Tool: {block['name']}]{Style.RESET_ALL} ",
|
||||
end="",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# LangGraph: ToolMessage result
|
||||
elif chunk.get("type") == "tool":
|
||||
result = chunk.get("content", "")
|
||||
if len(result) > 200:
|
||||
result = result[:200] + "..."
|
||||
print(
|
||||
f"\n{Fore.YELLOW}[Result: {result}]{Style.RESET_ALL}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# Strands: text token
|
||||
elif isinstance(chunk.get("data"), str):
|
||||
print(chunk["data"], end="", flush=True)
|
||||
|
||||
# Strands: tool use
|
||||
elif chunk.get("current_tool_use") and chunk.get(
|
||||
"current_tool_use", {}
|
||||
).get("name"):
|
||||
tool = chunk["current_tool_use"]
|
||||
if chunk.get("delta", {}).get("toolUse", {}).get("input") == "":
|
||||
print(
|
||||
f"\n{Fore.YELLOW}[Tool: {tool['name']}]{Style.RESET_ALL} ",
|
||||
end="",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# Strands: tool result
|
||||
elif chunk.get("message", {}).get("role") == "user":
|
||||
for content in chunk["message"].get("content", []):
|
||||
if "toolResult" in content:
|
||||
result = str(content["toolResult"].get("content", ""))
|
||||
if len(result) > 200:
|
||||
result = result[:200] + "..."
|
||||
print(
|
||||
f"\n{Fore.YELLOW}[Result: {result}]{Style.RESET_ALL}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
except (json.JSONDecodeError, KeyError):
|
||||
continue
|
||||
print() # Final newline
|
||||
|
||||
except requests.exceptions.ConnectionError:
|
||||
print_msg(f"Could not connect to {url}", "error")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
|
||||
|
||||
def run_chat(local_mode: bool, config: Dict[str, str]) -> None:
|
||||
"""
|
||||
Run interactive chat session.
|
||||
|
||||
Args:
|
||||
local_mode (bool): Whether to use local mode
|
||||
config (Dict[str, str]): Configuration dictionary
|
||||
"""
|
||||
session_id = generate_session_id()
|
||||
|
||||
print_section("Interactive Agent Chat")
|
||||
print(f"Session ID: {session_id}")
|
||||
print(
|
||||
f"Mode: {'Local (localhost:8080)' if local_mode else 'Remote (deployed agent)'}"
|
||||
)
|
||||
print(
|
||||
f"\n{Fore.YELLOW}💡 Type 'exit' or 'quit' to end, or press Ctrl+C{Style.RESET_ALL}\n"
|
||||
)
|
||||
|
||||
while True:
|
||||
try:
|
||||
prompt = input(f"{Fore.CYAN}You:{Style.RESET_ALL} ").strip()
|
||||
|
||||
if not prompt:
|
||||
continue
|
||||
|
||||
if prompt.lower() in ["exit", "quit"]:
|
||||
print(f"\n{Fore.GREEN}Goodbye!{Style.RESET_ALL}")
|
||||
break
|
||||
|
||||
# Invoke agent
|
||||
start_time = time.time()
|
||||
|
||||
if local_mode:
|
||||
# Local mode
|
||||
invoke_agent(
|
||||
url="http://localhost:8080/invocations",
|
||||
prompt=prompt,
|
||||
session_id=session_id,
|
||||
user_id="local-test-user",
|
||||
)
|
||||
else:
|
||||
# Remote mode
|
||||
endpoint = f"https://bedrock-agentcore.{config['region']}.amazonaws.com"
|
||||
escaped_arn = requests.utils.quote(config["runtime_arn"], safe="")
|
||||
url = f"{endpoint}/runtimes/{escaped_arn}/invocations?qualifier=DEFAULT"
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {config['access_token']}",
|
||||
"X-Amzn-Trace-Id": generate_trace_id(),
|
||||
"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": session_id,
|
||||
}
|
||||
|
||||
invoke_agent(
|
||||
url=url,
|
||||
prompt=prompt,
|
||||
session_id=session_id,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"\n{Fore.CYAN}[Completed in {elapsed:.2f}s]{Style.RESET_ALL}\n")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print(f"\n\n{Fore.GREEN}Goodbye!{Style.RESET_ALL}")
|
||||
break
|
||||
except EOFError:
|
||||
print(f"\n\n{Fore.GREEN}Goodbye!{Style.RESET_ALL}")
|
||||
break
|
||||
|
||||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
"""
|
||||
Parse command-line arguments.
|
||||
|
||||
Returns:
|
||||
argparse.Namespace: Parsed arguments
|
||||
"""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Interactive agent chat tester (local or remote)",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
# Remote agent (prompts for credentials)
|
||||
uv run scripts/test-agent.py
|
||||
|
||||
# Local agent on localhost:8080 (uses pattern from config.yaml)
|
||||
uv run scripts/test-agent.py --local
|
||||
|
||||
# Override pattern for local testing
|
||||
uv run scripts/test-agent.py --local --pattern strands-single-agent
|
||||
|
||||
Notes:
|
||||
- Remote mode: Tests deployed agent
|
||||
- Local mode: Pattern read from infra-cdk/config.yaml to start correct agent
|
||||
- Use --pattern to override the config value for local testing
|
||||
- Always runs in interactive conversation mode
|
||||
""",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--local",
|
||||
action="store_true",
|
||||
help="Test local agent on localhost:8080 (default: remote)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--pattern",
|
||||
type=str,
|
||||
help="Override agent pattern from config (e.g., 'strands-single-agent', 'langgraph-single-agent')",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point."""
|
||||
print("=" * 60)
|
||||
print("AgentCore Interactive Chat Tester")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
args = parse_arguments()
|
||||
config: Dict[str, str] = {}
|
||||
|
||||
# Get stack configuration
|
||||
stack_cfg = get_stack_config()
|
||||
|
||||
# LOCAL MODE
|
||||
if args.local:
|
||||
# Determine pattern: CLI arg > config.yaml > default (only needed for local mode)
|
||||
pattern = (
|
||||
args.pattern
|
||||
if args.pattern
|
||||
else stack_cfg.get("pattern", "langgraph-single-agent")
|
||||
)
|
||||
print(f"Using pattern: {pattern}\n")
|
||||
print_section("LOCAL MODE - Auto-starting agent")
|
||||
|
||||
# Get memory configuration
|
||||
memory_arn = stack_cfg["outputs"]["MemoryArn"]
|
||||
memory_id = memory_arn.split("/")[-1]
|
||||
region = stack_cfg["region"]
|
||||
stack_name = stack_cfg["stack_name"]
|
||||
|
||||
# Check if agent is already running
|
||||
if check_port_available(8080):
|
||||
print_msg("Agent already running on localhost:8080", "info")
|
||||
print("Using existing agent instance...\n")
|
||||
else:
|
||||
# Start the agent
|
||||
start_local_agent(memory_id, region, stack_name, pattern)
|
||||
|
||||
# REMOTE MODE
|
||||
else:
|
||||
print_section("REMOTE MODE - Testing deployed agent")
|
||||
|
||||
stack_cfg = get_stack_config()
|
||||
print(f"Stack: {stack_cfg['stack_name']}\n")
|
||||
|
||||
# Get configuration from CloudFormation outputs
|
||||
print("Fetching configuration from stack outputs...")
|
||||
outputs = stack_cfg["outputs"]
|
||||
|
||||
# Validate required outputs exist
|
||||
required_outputs = ["CognitoUserPoolId", "CognitoClientId", "RuntimeArn"]
|
||||
missing = [key for key in required_outputs if key not in outputs]
|
||||
if missing:
|
||||
print_msg(f"Missing required stack outputs: {', '.join(missing)}", "error")
|
||||
sys.exit(1)
|
||||
|
||||
print_msg("Configuration fetched")
|
||||
|
||||
runtime_arn = outputs["RuntimeArn"]
|
||||
region = stack_cfg["region"]
|
||||
|
||||
# Get credentials
|
||||
print_section("Authentication")
|
||||
|
||||
username = input("Enter username: ").strip()
|
||||
if not username:
|
||||
print_msg("Username is required", "error")
|
||||
sys.exit(1)
|
||||
password = getpass.getpass(f"Enter password for {username}: ")
|
||||
|
||||
# Authenticate
|
||||
access_token, id_token, user_id = authenticate_cognito(
|
||||
outputs["CognitoUserPoolId"], outputs["CognitoClientId"], username, password
|
||||
)
|
||||
|
||||
# Use access token for AgentCore runtime (JWT authorizer)
|
||||
config["access_token"] = access_token
|
||||
config["runtime_arn"] = runtime_arn
|
||||
config["region"] = region
|
||||
print(f"\nRuntime ARN: {runtime_arn}")
|
||||
print(f"Region: {region}\n")
|
||||
|
||||
# Run interactive chat
|
||||
run_chat(args.local, config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,240 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Shared utilities for test scripts
|
||||
|
||||
Provides essential functions for stack discovery, AWS resource fetching, and authentication.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import sys
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import boto3
|
||||
import yaml
|
||||
from botocore.exceptions import ClientError
|
||||
from colorama import Fore, Style, init
|
||||
|
||||
init(autoreset=True)
|
||||
|
||||
|
||||
def get_stack_config(stack_name: Optional[str] = None) -> Dict:
|
||||
"""
|
||||
Get complete stack configuration including outputs from main stack.
|
||||
|
||||
Args:
|
||||
stack_name: Base stack name (if None, loads from config.yaml)
|
||||
|
||||
Returns:
|
||||
Dictionary with stack_name, region, account, pattern, and outputs from main stack
|
||||
"""
|
||||
# Load config.yaml
|
||||
script_dir = Path(__file__).parent
|
||||
config_path = script_dir.parent / "config.yaml"
|
||||
|
||||
if not config_path.exists():
|
||||
print_msg("Configuration file not found", "error")
|
||||
sys.exit(1)
|
||||
|
||||
with open(config_path, "r") as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
# Get stack name from config if not provided
|
||||
if not stack_name:
|
||||
stack_name = config.get("stack_name_base")
|
||||
if not stack_name:
|
||||
print_msg("'stack_name_base' not found in config.yaml", "error")
|
||||
sys.exit(1)
|
||||
|
||||
# Get pattern from config
|
||||
pattern = config.get("backend", {}).get("pattern", "langgraph-single-agent")
|
||||
|
||||
cfn = boto3.client("cloudformation")
|
||||
|
||||
try:
|
||||
# Get outputs from main stack (contains Cognito, Runtime ARN, etc.)
|
||||
response = cfn.describe_stacks(StackName=stack_name)
|
||||
stack_info = response["Stacks"][0]
|
||||
|
||||
outputs = {}
|
||||
for output in stack_info.get("Outputs", []):
|
||||
outputs[output["OutputKey"]] = output["OutputValue"]
|
||||
|
||||
# Extract region and account from stack ARN or any ARN in outputs
|
||||
stack_arn = stack_info["StackId"]
|
||||
region = stack_arn.split(":")[3]
|
||||
account = stack_arn.split(":")[4]
|
||||
|
||||
return {
|
||||
"stack_name": stack_name,
|
||||
"region": region,
|
||||
"account": account,
|
||||
"pattern": pattern,
|
||||
"outputs": outputs,
|
||||
}
|
||||
|
||||
except ClientError as e:
|
||||
error_code = e.response.get("Error", {}).get("Code", "Unknown")
|
||||
print_msg(f"CloudFormation error: {error_code}", "error")
|
||||
if error_code == "ValidationError":
|
||||
print_msg(
|
||||
f"Stack '{stack_name}' not found. Make sure you've deployed the CDK stack.",
|
||||
"error",
|
||||
)
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print_msg(f"Failed to get stack config: {e}", "error")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def get_ssm_params(stack_name: str, *param_names: str) -> Dict[str, str]:
|
||||
"""
|
||||
Fetch multiple SSM parameters for a stack.
|
||||
|
||||
Args:
|
||||
stack_name: Base stack name
|
||||
*param_names: Parameter names (without the /{stack_name}/ prefix)
|
||||
|
||||
Returns:
|
||||
Dictionary mapping parameter names to values
|
||||
"""
|
||||
ssm = boto3.client("ssm")
|
||||
results = {}
|
||||
|
||||
try:
|
||||
for param_name in param_names:
|
||||
full_name = f"/{stack_name}/{param_name}"
|
||||
response = ssm.get_parameter(Name=full_name)
|
||||
results[param_name] = response["Parameter"]["Value"]
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
print_msg(f"Failed to fetch SSM parameters: {e}", "error")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def authenticate_cognito(
|
||||
user_pool_id: str, client_id: str, username: str, password: str
|
||||
) -> Tuple[str, str, str]:
|
||||
"""
|
||||
Authenticate with Cognito.
|
||||
|
||||
Args:
|
||||
user_pool_id: Cognito User Pool ID
|
||||
client_id: Cognito Client ID
|
||||
username: Username
|
||||
password: Password
|
||||
|
||||
Returns:
|
||||
Tuple of (access_token, id_token, user_id)
|
||||
- access_token: For AgentCore runtime invocations (JWT authorizer)
|
||||
- id_token: For API Gateway Cognito User Pool authorizers
|
||||
- user_id: User's unique identifier (sub claim)
|
||||
"""
|
||||
print("\nAuthenticating...")
|
||||
|
||||
cognito = boto3.client("cognito-idp")
|
||||
|
||||
try:
|
||||
# Check if user exists
|
||||
try:
|
||||
cognito.admin_get_user(UserPoolId=user_pool_id, Username=username)
|
||||
except cognito.exceptions.UserNotFoundException:
|
||||
print_msg(f"User '{username}' does not exist", "error")
|
||||
sys.exit(1)
|
||||
|
||||
# Authenticate
|
||||
response = cognito.initiate_auth(
|
||||
AuthFlow="USER_PASSWORD_AUTH",
|
||||
ClientId=client_id,
|
||||
AuthParameters={"USERNAME": username, "PASSWORD": password},
|
||||
)
|
||||
|
||||
access_token = response["AuthenticationResult"]["AccessToken"]
|
||||
id_token = response["AuthenticationResult"]["IdToken"]
|
||||
|
||||
# Decode ID token to get user ID
|
||||
import base64
|
||||
import json
|
||||
|
||||
payload = id_token.split(".")[1]
|
||||
payload += "=" * (4 - len(payload) % 4)
|
||||
decoded = base64.b64decode(payload)
|
||||
token_data = json.loads(decoded)
|
||||
user_id = token_data.get("sub")
|
||||
|
||||
print_msg("Authentication successful")
|
||||
print(f" User ID: {user_id}")
|
||||
|
||||
return access_token, id_token, user_id
|
||||
|
||||
except Exception as e:
|
||||
print_msg(f"Authentication failed: {e}", "error")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def create_bedrock_client(region: str) -> boto3.client:
|
||||
"""Create bedrock-agentcore client."""
|
||||
return boto3.client("bedrock-agentcore", region_name=region)
|
||||
|
||||
|
||||
def generate_session_id() -> str:
|
||||
"""Generate UUID4 session ID."""
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
def print_msg(message: str, level: str = "info") -> None:
|
||||
"""
|
||||
Print formatted message.
|
||||
|
||||
Args:
|
||||
message: Message to print
|
||||
level: 'success', 'error', 'info', or 'section'
|
||||
"""
|
||||
if level == "success":
|
||||
print(f"{Fore.GREEN}✓ {message}{Style.RESET_ALL}")
|
||||
elif level == "error":
|
||||
print(f"{Fore.RED}✗ {message}{Style.RESET_ALL}")
|
||||
elif level == "info":
|
||||
print(f"{Fore.YELLOW}ℹ {message}{Style.RESET_ALL}")
|
||||
elif level == "section":
|
||||
print("\n" + "=" * 60)
|
||||
print(message)
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
|
||||
def print_section(title: str, width: int = 60) -> None:
|
||||
"""Print section header."""
|
||||
print("\n" + "=" * width)
|
||||
print(title)
|
||||
print("=" * width + "\n")
|
||||
|
||||
|
||||
def create_mock_jwt(user_id: str) -> str:
|
||||
"""
|
||||
Create a mock unsigned JWT token with the given user_id as the 'sub' claim.
|
||||
|
||||
The agent's extract_user_id_from_context() decodes the JWT without signature
|
||||
verification (since AgentCore Runtime validates it in production). This allows
|
||||
local testing to pass a user identity the same way production does.
|
||||
|
||||
Args:
|
||||
user_id (str): The user ID to embed as the 'sub' claim.
|
||||
|
||||
Returns:
|
||||
str: A mock JWT string (header.payload.signature).
|
||||
"""
|
||||
header = (
|
||||
base64.urlsafe_b64encode(json.dumps({"alg": "none", "typ": "JWT"}).encode())
|
||||
.rstrip(b"=")
|
||||
.decode()
|
||||
)
|
||||
payload = (
|
||||
base64.urlsafe_b64encode(json.dumps({"sub": user_id}).encode())
|
||||
.rstrip(b"=")
|
||||
.decode()
|
||||
)
|
||||
return f"{header}.{payload}."
|
||||
Reference in New Issue
Block a user