diff --git a/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/.gitignore b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/.gitignore new file mode 100644 index 000000000..458a98fe5 --- /dev/null +++ b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +*.pyc +.bedrock_agentcore/ +.bedrock_agentcore.yaml +.venv/ diff --git a/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/README.md b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/README.md new file mode 100644 index 000000000..3dc1ae514 --- /dev/null +++ b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/README.md @@ -0,0 +1,129 @@ +# Async WAT Refresh for Long-Running Agents + +## Overview + +When an AgentCore agent runs a long-running background task in a thread, the Workload Access Token (WAT) — created from the inbound JWT — expires after the JWT's TTL. This causes the SDK to fall back to creating orphaned workload identities via IAM, breaking user binding and auditability. + +This sample demonstrates a companion library (`agentcore_thread_utils`) that solves the WAT expiration problem by: +- Propagating the WAT to background threads via `contextvars.copy_context()` +- Pausing the thread when the WAT expires and waiting for the client to send a fresh JWT +- Retrying the credential provider call with the refreshed WAT + +No orphan workload identities are created. The WAT stays bound to the original user. + +## Architecture + +| Component | Description | +|---|---| +| `@with_wat_refresh` | Drop-in replacement for `@requires_access_token`. Catches WAT expiration, pauses the thread, waits for client refresh, retries. | +| `ThreadTaskManager` | Manages thread lifecycle with WAT propagation. Handles start/status/refresh/result actions. | + +## Prerequisites + +- AWS CLI configured +- `agentcore` CLI installed +- `jq` installed +- Python 3.10+ + +## Setup + +### 1. Create Cognito User Pool (5-min access token TTL) + +```bash +source setup_cognito.sh +``` + +### 2. Create Credential Provider + +```bash +bash setup_credential_provider.sh +``` + +### 3. Deploy the Agent + +```bash +bash deploy.sh +``` + +Export the Agent ARN from the output: + +```bash +export AGENT_ARN="arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/thread_async_utils-XXXXXXXXXX" +``` + +## Test + +### Start a task + +```bash +bash test_curl.sh +``` + +### Wait 6 minutes, then check status + +```bash +bash test_refresh.sh status +``` + +### Send WAT refresh + +```bash +bash test_refresh.sh refresh +``` + +### Check result + +```bash +bash test_refresh.sh status +``` + +## Expected Flow + +1. Client invokes agent with JWT (5-min TTL) → task starts in background thread with propagated WAT +2. Thread sleeps 6 minutes (WAT expires at minute 5) +3. Thread calls credential provider → WAT expired → decorator pauses thread +4. Client sends `{"action": "refresh"}` with fresh JWT → Runtime creates new WAT +5. Thread resumes, retries credential provider → success +6. Task completes, agent returns to Healthy + +## How It Works + +```mermaid +sequenceDiagram + participant Client + participant Entrypoint as Entrypoint (main thread) + participant Runtime as AgentCore Runtime + participant Thread as Background Thread + participant Provider as Credential Provider + + Client->>Entrypoint: {"action": "start"} + JWT + Runtime->>Runtime: Create WAT from JWT (same exp) + Entrypoint->>Thread: Start with copy_context() (WAT propagated) + Entrypoint-->>Client: {"task_id": 123, "status": "started"} + + Note over Thread: Business logic runs... + + Thread->>Provider: @with_wat_refresh → get token + Provider-->>Thread: Token has expired + + Note over Thread: Thread paused, waiting for refresh + + Client->>Entrypoint: {"action": "refresh"} + fresh JWT + Runtime->>Runtime: Create new WAT + Entrypoint->>Thread: Signal with new WAT + + Note over Thread: Thread resumes + Thread->>Provider: Retry → get token ✓ +``` + +## Cleanup + +```bash +agentcore destroy +aws cognito-idp delete-user-pool --user-pool-id $POOL_ID --region us-east-1 +``` + +## Related + +- [AgentCore Identity - Getting Started](../01-getting_started.md) +- [AgentCore Identity - How It Works](../02-how_it_works.md) diff --git a/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/agentcore_thread_utils/__init__.py b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/agentcore_thread_utils/__init__.py new file mode 100644 index 000000000..1265ad80c --- /dev/null +++ b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/agentcore_thread_utils/__init__.py @@ -0,0 +1,6 @@ +"""AgentCore Thread Utils — companion library for threaded async agents with WAT refresh.""" + +from .decorator import with_wat_refresh, set_task_context +from .helper import ThreadTaskManager + +__all__ = ["with_wat_refresh", "set_task_context", "ThreadTaskManager"] diff --git a/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/agentcore_thread_utils/decorator.py b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/agentcore_thread_utils/decorator.py new file mode 100644 index 000000000..e912670ca --- /dev/null +++ b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/agentcore_thread_utils/decorator.py @@ -0,0 +1,96 @@ +"""with_wat_refresh — thread-safe drop-in replacement for @requires_access_token.""" + +import logging +import threading +from functools import wraps +from typing import Any, Callable, List, Literal, Optional + +from bedrock_agentcore.runtime import BedrockAgentCoreContext +from bedrock_agentcore.identity.auth import requires_access_token + +logger = logging.getLogger("agentcore_thread_utils.decorator") + +# Thread-local storage for task context — safe for concurrent tasks +_local = threading.local() + + +def set_task_context(manager, task_id: int): + """Set the current task context for WAT refresh coordination. + + Must be called at the start of each task function. + Uses thread-local storage so concurrent tasks don't interfere. + """ + _local.task_manager = manager + _local.task_id = task_id + + +def _get_task_context(): + """Get the current task context from thread-local storage.""" + manager = getattr(_local, "task_manager", None) + task_id = getattr(_local, "task_id", None) + return manager, task_id + + +def with_wat_refresh( + *, + provider_name: str, + scopes: List[str], + auth_flow: Literal["M2M", "USER_FEDERATION"] = "M2M", + into: str = "access_token", + on_auth_url: Optional[Callable] = None, + callback_url: Optional[str] = None, + force_authentication: bool = False, + max_retries: int = 2, +) -> Callable: + """Decorator that wraps @requires_access_token with WAT refresh for threads. + + Same interface as @requires_access_token but handles WAT expiration + by pausing the thread and waiting for a client refresh. + + Args: + max_retries: Maximum number of refresh attempts before giving up (default: 2). + """ + + def decorator(func: Callable) -> Callable: + @requires_access_token( + provider_name=provider_name, + scopes=scopes, + auth_flow=auth_flow, + into=into, + on_auth_url=on_auth_url, + callback_url=callback_url, + force_authentication=force_authentication, + ) + def _inner_call(*, access_token: str): + return func(access_token=access_token) + + @wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + manager, task_id = _get_task_context() + last_error = None + + for attempt in range(max_retries + 1): + try: + return _inner_call(**kwargs) + except Exception as e: + last_error = e + if "expired" in str(e).lower() and manager and task_id: + if attempt < max_retries: + logger.info( + f"WAT expired (attempt {attempt + 1}/{max_retries}). " + f"Requesting refresh for task {task_id}..." + ) + new_wat = manager.wait_for_wat_refresh(task_id) + BedrockAgentCoreContext.set_workload_access_token(new_wat) + logger.info("WAT refreshed. Retrying...") + else: + logger.error(f"WAT expired after {max_retries} refresh attempts.") + raise + else: + raise + + raise last_error + + return wrapper + + return decorator diff --git a/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/agentcore_thread_utils/helper.py b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/agentcore_thread_utils/helper.py new file mode 100644 index 000000000..78f4863de --- /dev/null +++ b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/agentcore_thread_utils/helper.py @@ -0,0 +1,231 @@ +"""ThreadTaskManager — handles threaded async tasks with WAT propagation via copy_context.""" + +import contextvars +import logging +import threading +import time +from typing import Any, Callable, Dict, Optional + +from bedrock_agentcore.runtime import BedrockAgentCoreApp, BedrockAgentCoreContext + +logger = logging.getLogger("agentcore_thread_utils.helper") + +DEFAULT_REFRESH_TIMEOUT = 300 # 5 minutes +DEFAULT_RESULT_TTL = 3600 # 1 hour + + +class ThreadTaskManager: + """Manages threaded tasks with WAT propagation and refresh coordination. + + Uses contextvars.copy_context() to propagate the WAT to threads, + and threading.Event for refresh coordination. + + Args: + app: BedrockAgentCoreApp instance. + refresh_timeout: Seconds to wait for a WAT refresh before failing (default: 300). + result_ttl: Seconds to keep completed results before cleanup (default: 3600). + """ + + def __init__( + self, + app: BedrockAgentCoreApp, + refresh_timeout: int = DEFAULT_REFRESH_TIMEOUT, + result_ttl: int = DEFAULT_RESULT_TTL, + ): + self.app = app + self.refresh_timeout = refresh_timeout + self.result_ttl = result_ttl + self._tasks: Dict[int, dict] = {} + self._results: Dict[int, dict] = {} + self._lock = threading.Lock() + self._registered_task: Optional[Callable] = None + self._custom_actions: Dict[str, Callable] = {} + + def start(self, task_func: Callable, **kwargs) -> int: + """Start a threaded task with WAT propagation via copy_context. + + task_func: a function that accepts task_id as first keyword argument. + Returns the task_id for tracking. + """ + task_id = self.app.add_async_task("thread_task", {}) + + with self._lock: + self._tasks[task_id] = { + "event": threading.Event(), + "wat": None, + "needs_refresh": False, + } + + ctx = contextvars.copy_context() + + def _wrapped(): + try: + result = task_func(task_id=task_id, **kwargs) + with self._lock: + self._results[task_id] = { + "status": "completed", + "result": result, + "completed_at": time.time(), + } + except TimeoutError as e: + logger.warning(f"[Task {task_id}] {e}") + with self._lock: + self._results[task_id] = { + "status": "failed", + "error": str(e), + "completed_at": time.time(), + } + except Exception as e: + error_msg = str(e) + if "expired" in error_msg.lower(): + logger.warning(f"[Task {task_id}] WAT expired and refresh failed: {error_msg}") + else: + logger.error(f"[Task {task_id}] Failed: {error_msg}", exc_info=True) + with self._lock: + self._results[task_id] = { + "status": "failed", + "error": error_msg, + "completed_at": time.time(), + } + finally: + with self._lock: + self._tasks.pop(task_id, None) + self.app.complete_async_task(task_id) + + threading.Thread(target=ctx.run, args=(_wrapped,), daemon=True).start() + return task_id + + def wait_for_wat_refresh(self, task_id: int) -> str: + """Called by the decorator when WAT expires. Blocks until client refreshes. + + Returns the fresh WAT. + Raises TimeoutError if no refresh arrives within refresh_timeout. + """ + with self._lock: + task_state = self._tasks.get(task_id) + if not task_state: + raise RuntimeError(f"Task {task_id} not found") + + task_state["needs_refresh"] = True + task_state["event"].clear() + logger.info( + f"[Task {task_id}] WAT expired. Waiting for client refresh " + f"(timeout: {self.refresh_timeout}s)... " + f'Client should send {{"action":"refresh"}} with a fresh JWT to unblock.' + ) + + signaled = task_state["event"].wait(timeout=self.refresh_timeout) + if not signaled: + task_state["needs_refresh"] = False + raise TimeoutError( + f"WAT refresh timeout after {self.refresh_timeout}s for task {task_id}. " + 'Client did not send {"action":"refresh"} with a fresh JWT in time. ' + "Task will be marked as failed." + ) + + new_wat = task_state["wat"] + task_state["needs_refresh"] = False + logger.info(f"[Task {task_id}] WAT refreshed.") + return new_wat + + def _cleanup_old_results(self): + """Remove results older than result_ttl.""" + now = time.time() + expired = [ + tid for tid, res in self._results.items() + if now - res.get("completed_at", now) > self.result_ttl + ] + for tid in expired: + del self._results[tid] + + def handle_action(self, action: str, payload: dict) -> Optional[dict]: + """Handle status/refresh/result actions. Returns response dict or None.""" + + if action == "status": + task_info = self.app.get_async_task_info() + ping_status = self.app.get_current_ping_status() + with self._lock: + self._cleanup_old_results() + needs_refresh = { + tid: state["needs_refresh"] + for tid, state in self._tasks.items() + if state["needs_refresh"] + } + completed = { + tid: res["status"] for tid, res in self._results.items() + } + return { + "ping_status": ping_status.value, + "active_tasks": task_info["active_count"], + "tasks_needing_refresh": needs_refresh, + "completed_results": completed, + } + + elif action == "refresh": + new_wat = BedrockAgentCoreContext.get_workload_access_token() + if not new_wat: + return {"status": "error", "message": "No WAT in current context"} + + task_id = payload.get("task_id") + with self._lock: + if task_id and task_id in self._tasks: + self._tasks[task_id]["wat"] = new_wat + self._tasks[task_id]["event"].set() + return {"status": "wat_refreshed", "task_id": task_id} + else: + refreshed = [] + for tid, state in self._tasks.items(): + if state["needs_refresh"]: + state["wat"] = new_wat + state["event"].set() + refreshed.append(tid) + return {"status": "wat_refreshed", "tasks_refreshed": refreshed} + + elif action == "result": + task_id = payload.get("task_id") + with self._lock: + if task_id and task_id in self._results: + res = self._results[task_id].copy() + res.pop("completed_at", None) + return res + return { + "error": "Task not found or still running", + "available": list(self._results.keys()), + } + + return None + + def register_task(self, task_func: Callable): + """Register the main async task function.""" + self._registered_task = task_func + + def register_action(self, action_name: str, handler: Callable): + """Register a custom action handler (e.g., 'chat').""" + self._custom_actions[action_name] = handler + + def handle(self, payload: dict, context=None) -> dict: + """Handle all actions in one call. Use as the entrypoint body. + + Routes: start, status, refresh, result, custom actions, or returns available actions. + """ + action = payload.get("action") + + response = self.handle_action(action, payload) + if response is not None: + return response + + if action == "start": + if not self._registered_task: + return {"error": "No task registered. Call manager.register_task() first."} + task_id = self.start(self._registered_task) + return { + "task_id": task_id, + "status": "started", + "message": "If WAT expires, send {\"action\":\"refresh\"} with a fresh JWT.", + } + + if action in self._custom_actions: + return self._custom_actions[action](payload, context) + + available = ["start", "status", "refresh", "result"] + list(self._custom_actions.keys()) + return {"available_actions": available} diff --git a/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/main.py b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/main.py new file mode 100644 index 000000000..1004f0b15 --- /dev/null +++ b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/main.py @@ -0,0 +1,66 @@ +""" +Threaded async agent — WAT refresh test. +Wait 6 min, then 3 credential provider calls. +No model needed — just testing the decorator/helper WAT refresh mechanism. +""" + +import time +import logging +from datetime import datetime, timezone + +from bedrock_agentcore.runtime import BedrockAgentCoreApp +from agentcore_thread_utils import with_wat_refresh, ThreadTaskManager, set_task_context + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +# SCAFFOLDING +app = BedrockAgentCoreApp() +manager = ThreadTaskManager(app) + + +# DEVELOPER CODE — credential provider call +@with_wat_refresh( + provider_name="cognito-m2m-provider", + scopes=["https://agent-api.example.com/invoke"], + auth_flow="M2M", +) +def call_my_api(*, access_token: str) -> dict: + return {"success": True, "timestamp": datetime.now(timezone.utc).isoformat()} + + +# DEVELOPER CODE — business logic +def my_long_running_job(*, task_id: int): + set_task_context(manager, task_id) + results = [] + + logger.info(f"[Task {task_id}] Waiting 6 minutes...") + for m in range(1, 7): + time.sleep(60) + logger.info(f"[Task {task_id}] {m}/6 min") + + for i in range(1, 4): + logger.info(f"[Task {task_id}] Test {i}/3...") + r = call_my_api(access_token="") + results.append(r) + logger.info(f"[Task {task_id}] Test {i} done: {r}") + if i < 3: + logger.info(f"[Task {task_id}] Waiting 6 minutes before next test...") + for m in range(1, 7): + time.sleep(60) + logger.info(f"[Task {task_id}] {m}/6 min") + + return {"status": "completed", "tests": results} + + +# SCAFFOLDING +manager.register_task(my_long_running_job) + + +@app.entrypoint +def handler(payload, context): + return manager.handle(payload, context) + + +if __name__ == "__main__": + app.run() diff --git a/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/requirements.txt b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/requirements.txt new file mode 100644 index 000000000..fa1f0fc06 --- /dev/null +++ b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/app/requirements.txt @@ -0,0 +1,2 @@ +bedrock-agentcore +boto3 diff --git a/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/deploy.sh b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/deploy.sh new file mode 100644 index 000000000..6ce5e0b6b --- /dev/null +++ b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/deploy.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# Deploy the WAT refresh test agent. +# Run AFTER: source setup_cognito.sh && bash setup_credential_provider.sh + +set -e + +export AWS_PROFILE=${AWS_PROFILE:-account-a} +REGION=${REGION:-us-east-1} + +if [ -z "$POOL_ID" ] || [ -z "$INBOUND_CLIENT_ID" ]; then + echo "ERROR: Run 'source setup_cognito.sh' first" + exit 1 +fi + +DISCOVERY_URL="https://cognito-idp.${REGION}.amazonaws.com/${POOL_ID}/.well-known/openid-configuration" + +echo "=== Configuring agent ===" +agentcore configure \ + -e main.py \ + -n thread_async_utils \ + -r $REGION \ + -rf requirements.txt \ + -dt direct_code_deploy \ + -rt PYTHON_3_12 \ + -do \ + -dm \ + -ni \ + --idle-timeout 3600 \ + --max-lifetime 28800 \ + -ac "{\"customJWTAuthorizer\":{\"discoveryUrl\":\"${DISCOVERY_URL}\",\"allowedClients\":[\"${INBOUND_CLIENT_ID}\"]}}" + +echo "" +echo "=== Deploying ===" +agentcore deploy + +echo "" +echo "Set AGENT_ARN from the output above, then run test scripts." diff --git a/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/setup_cognito.sh b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/setup_cognito.sh new file mode 100644 index 000000000..c7f67a0fe --- /dev/null +++ b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/setup_cognito.sh @@ -0,0 +1,93 @@ +#!/bin/bash +# Setup Cognito User Pool for inbound JWT auth +# and a Resource Server + App Client for outbound 2LO (M2M) +# +# Usage: source setup_cognito.sh +# This exports env vars needed by deploy.sh and test scripts. + +set -e + +export AWS_PROFILE=${AWS_PROFILE:-account-a} +REGION=${REGION:-us-east-1} +USERNAME=${USERNAME:-testuser} +PASSWORD=${PASSWORD:-TestPass123!} +POOL_NAME="AgentCoreAsyncDemo" +RESOURCE_SERVER_ID="https://agent-api.example.com" + +echo "=== Creating Cognito User Pool ===" +export POOL_ID=$(aws cognito-idp create-user-pool \ + --pool-name "$POOL_NAME" \ + --policies '{"PasswordPolicy":{"MinimumLength":8}}' \ + --region $REGION | jq -r '.UserPool.Id') +echo "Pool ID: $POOL_ID" + +echo "=== Creating Resource Server (for 2LO scopes) ===" +aws cognito-idp create-resource-server \ + --user-pool-id $POOL_ID \ + --identifier "$RESOURCE_SERVER_ID" \ + --name "Agent API" \ + --scopes '[{"ScopeName":"invoke","ScopeDescription":"Invoke agent API"}]' \ + --region $REGION > /dev/null + +echo "=== Creating Inbound App Client (for JWT auth - user login) ===" +export INBOUND_CLIENT_ID=$(aws cognito-idp create-user-pool-client \ + --user-pool-id $POOL_ID \ + --client-name "InboundClient" \ + --no-generate-secret \ + --explicit-auth-flows "ALLOW_USER_PASSWORD_AUTH" "ALLOW_REFRESH_TOKEN_AUTH" \ + --access-token-validity 5 \ + --token-validity-units '{"AccessToken":"minutes"}' \ + --region $REGION | jq -r '.UserPoolClient.ClientId') +echo "Inbound Client ID: $INBOUND_CLIENT_ID (access token TTL: 5 minutes)" + +echo "=== Creating Outbound App Client (for 2LO - M2M) ===" +OUTBOUND_RESPONSE=$(aws cognito-idp create-user-pool-client \ + --user-pool-id $POOL_ID \ + --client-name "OutboundM2MClient" \ + --generate-secret \ + --allowed-o-auth-flows "client_credentials" \ + --allowed-o-auth-scopes "${RESOURCE_SERVER_ID}/invoke" \ + --allowed-o-auth-flows-user-pool-client \ + --region $REGION) + +export OUTBOUND_CLIENT_ID=$(echo $OUTBOUND_RESPONSE | jq -r '.UserPoolClient.ClientId') +export OUTBOUND_CLIENT_SECRET=$(echo $OUTBOUND_RESPONSE | jq -r '.UserPoolClient.ClientSecret') +echo "Outbound Client ID: $OUTBOUND_CLIENT_ID" +echo "Outbound Client Secret: $OUTBOUND_CLIENT_SECRET" + +echo "=== Creating Cognito Domain (required for token endpoint) ===" +DOMAIN_PREFIX="agentcore-async-$(echo $POOL_ID | tr '[:upper:]' '[:lower:]' | sed 's/_/-/g')" +aws cognito-idp create-user-pool-domain \ + --user-pool-id $POOL_ID \ + --domain "$DOMAIN_PREFIX" \ + --region $REGION > /dev/null +echo "Domain: https://${DOMAIN_PREFIX}.auth.${REGION}.amazoncognito.com" + +echo "=== Creating Test User ===" +aws cognito-idp admin-create-user \ + --user-pool-id $POOL_ID \ + --username $USERNAME \ + --region $REGION \ + --message-action SUPPRESS > /dev/null + +aws cognito-idp admin-set-user-password \ + --user-pool-id $POOL_ID \ + --username $USERNAME \ + --password $PASSWORD \ + --region $REGION \ + --permanent > /dev/null + +export USERNAME PASSWORD + +echo "" +echo "=========================================" +echo "SETUP COMPLETE" +echo "=========================================" +echo "Pool ID: $POOL_ID" +echo "Discovery URL: https://cognito-idp.${REGION}.amazonaws.com/${POOL_ID}/.well-known/openid-configuration" +echo "Inbound Client ID: $INBOUND_CLIENT_ID" +echo "Outbound Client ID: $OUTBOUND_CLIENT_ID" +echo "Outbound Secret: $OUTBOUND_CLIENT_SECRET" +echo "" +echo "Run 'source setup_credential_provider.sh' next." +echo "=========================================" diff --git a/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/setup_credential_provider.sh b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/setup_credential_provider.sh new file mode 100644 index 000000000..c4770d605 --- /dev/null +++ b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/setup_credential_provider.sh @@ -0,0 +1,42 @@ +#!/bin/bash +# Create the OAuth2 credential provider in AgentCore Identity for outbound 2LO +# Run AFTER: source setup_cognito.sh + +set -e + +export AWS_PROFILE=${AWS_PROFILE:-account-a} +REGION=${REGION:-us-east-1} + +if [ -z "$POOL_ID" ] || [ -z "$OUTBOUND_CLIENT_ID" ] || [ -z "$OUTBOUND_CLIENT_SECRET" ]; then + echo "ERROR: Run 'source setup_cognito.sh' first" + exit 1 +fi + +DOMAIN_PREFIX="agentcore-async-$(echo $POOL_ID | tr '[:upper:]' '[:lower:]' | sed 's/_/-/g')" +TOKEN_ENDPOINT="https://${DOMAIN_PREFIX}.auth.${REGION}.amazoncognito.com/oauth2/token" +AUTH_ENDPOINT="https://${DOMAIN_PREFIX}.auth.${REGION}.amazoncognito.com/oauth2/authorize" + +echo "=== Creating OAuth2 Credential Provider for 2LO ===" +echo "Token Endpoint: $TOKEN_ENDPOINT" +aws bedrock-agentcore-control create-oauth2-credential-provider \ + --name "cognito-m2m-provider" \ + --credential-provider-vendor "CustomOauth2" \ + --oauth2-provider-config-input "{ + \"customOauth2ProviderConfig\": { + \"oauthDiscovery\": { + \"authorizationServerMetadata\": { + \"issuer\": \"https://cognito-idp.${REGION}.amazonaws.com/${POOL_ID}\", + \"authorizationEndpoint\": \"${AUTH_ENDPOINT}\", + \"tokenEndpoint\": \"${TOKEN_ENDPOINT}\" + } + }, + \"clientId\": \"${OUTBOUND_CLIENT_ID}\", + \"clientSecret\": \"${OUTBOUND_CLIENT_SECRET}\" + } + }" \ + --region $REGION \ + --output json + +echo "" +echo "Credential provider 'cognito-m2m-provider' created." +echo "Now run: bash deploy.sh" diff --git a/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/test_curl.sh b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/test_curl.sh new file mode 100644 index 000000000..0c63a8ca7 --- /dev/null +++ b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/test_curl.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# Start a task. Usage: bash test_curl.sh [action] +# Requires: AGENT_ARN, INBOUND_CLIENT_ID, USERNAME, PASSWORD env vars +set -e +export AWS_PROFILE=${AWS_PROFILE:-account-a} +REGION=${REGION:-us-east-1} + +if [ -z "$AGENT_ARN" ] || [ -z "$INBOUND_CLIENT_ID" ]; then + echo "ERROR: Set AGENT_ARN and run 'source setup_cognito.sh' first" + exit 1 +fi + +ENCODED_ARN=$(python3 -c "import urllib.parse; print(urllib.parse.quote('$AGENT_ARN', safe=''))") +ENDPOINT="https://bedrock-agentcore.${REGION}.amazonaws.com" + +TOKEN=$(aws cognito-idp initiate-auth \ + --client-id "$INBOUND_CLIENT_ID" \ + --auth-flow USER_PASSWORD_AUTH \ + --auth-parameters USERNAME=${USERNAME:-testuser},PASSWORD=${PASSWORD:-TestPass123!} \ + --region $REGION --query 'AuthenticationResult.AccessToken' --output text) + +echo "Token: ${TOKEN:0:30}..." + +SESSION_ID="thread-test-$(uuidgen | tr '[:upper:]' '[:lower:]')" +echo "Session: $SESSION_ID" + +ACTION=${1:-start} +echo "Action: $ACTION" + +curl -s -X POST \ + "${ENDPOINT}/runtimes/${ENCODED_ARN}/invocations?qualifier=DEFAULT" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer ${TOKEN}" \ + -H "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id: ${SESSION_ID}" \ + -d "{\"action\":\"${ACTION}\"}" | python3 -m json.tool diff --git a/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/test_refresh.sh b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/test_refresh.sh new file mode 100644 index 000000000..ec674f3f6 --- /dev/null +++ b/01-tutorials/03-AgentCore-identity/13-async-wat-refresh/test_refresh.sh @@ -0,0 +1,34 @@ +#!/bin/bash +# Send refresh or check status. Usage: bash test_refresh.sh [action] +# Requires: AGENT_ARN, INBOUND_CLIENT_ID, USERNAME, PASSWORD env vars +set -e +export AWS_PROFILE=${AWS_PROFILE:-account-a} +REGION=${REGION:-us-east-1} + +if [ -z "$AGENT_ARN" ] || [ -z "$INBOUND_CLIENT_ID" ]; then + echo "ERROR: Set AGENT_ARN and run 'source setup_cognito.sh' first" + exit 1 +fi + +SESSION_ID=${1:?"Usage: bash test_refresh.sh [action]"} +ACTION=${2:-refresh} + +ENCODED_ARN=$(python3 -c "import urllib.parse; print(urllib.parse.quote('$AGENT_ARN', safe=''))") +ENDPOINT="https://bedrock-agentcore.${REGION}.amazonaws.com" + +TOKEN=$(aws cognito-idp initiate-auth \ + --client-id "$INBOUND_CLIENT_ID" \ + --auth-flow USER_PASSWORD_AUTH \ + --auth-parameters USERNAME=${USERNAME:-testuser},PASSWORD=${PASSWORD:-TestPass123!} \ + --region $REGION --query 'AuthenticationResult.AccessToken' --output text) + +echo "Token: ${TOKEN:0:30}..." +echo "Session: $SESSION_ID" +echo "Action: $ACTION" + +curl -s -X POST \ + "${ENDPOINT}/runtimes/${ENCODED_ARN}/invocations?qualifier=DEFAULT" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer ${TOKEN}" \ + -H "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id: ${SESSION_ID}" \ + -d "{\"action\":\"${ACTION}\"}" | python3 -m json.tool diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 3028b1005..36cf63541 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -113,3 +113,7 @@ - Richa Gupta (richagpt) - Chandra Dhandapani - Anant Murarka (anantmu) +- Guilherme Greco +- Alex Rosa +- André Weber +- Eduardo Henrique de Souza Mendes