Merge branch 'awslabs:main' into main

Author: mvangara10

01-tutorials/01-AgentCore-runtime/02-hosting-MCP-server/.gitignore

@@ -0,0 +1,7 @@
+.bedrock_agentcore.yaml
+.dockerignore
+Dockerfile
+invoke_mcp_tools.py
+mcp_server.py
+my_mcp_client.py
+my_mcp_client_remote.py

01-tutorials/01-AgentCore-runtime/02-hosting-MCP-server/hosting_mcp_server.ipynb

@@ -287,7 +287,7 @@
     "print(\"Setting up Amazon Cognito user pool...\")\n",
     "cognito_config = setup_cognito_user_pool()\n",
     "print(\"Cognito setup completed ✓\")\n",
-    "print(f\"User Pool ID: {cognito_config.get('user_pool_id', 'N/A')}\")\n",
+    "print(f\"User Pool ID: {cognito_config.get('pool_id', 'N/A')}\")\n",
     "print(f\"Client ID: {cognito_config.get('client_id', 'N/A')}\")"
    ]
   },
@@ -448,12 +448,45 @@
     "import boto3\n",
     "import json\n",
     "import sys\n",
+    "import base64\n",
+    "import time\n",
     "from boto3.session import Session\n",
     "from datetime import timedelta\n",
+    "import traceback\n",
     "\n",
     "from mcp import ClientSession\n",
     "from mcp.client.streamable_http import streamablehttp_client\n",
     "\n",
+    "def get_refresh_token(client_id, refresh_token, region):\n",
+    "    \"\"\"Refresh access token using refresh token\"\"\"\n",
+    "    cognito_client = boto3.client('cognito-idp', region_name=region)\n",
+    "    auth_response = cognito_client.initiate_auth(\n",
+    "        ClientId=client_id,\n",
+    "        AuthFlow='REFRESH_TOKEN_AUTH',\n",
+    "        AuthParameters={'REFRESH_TOKEN': refresh_token}\n",
+    "    )\n",
+    "    return auth_response['AuthenticationResult']['AccessToken']\n",
+    "\n",
+    "def get_valid_token(bearer_token, client_id, refresh_token, region):\n",
+    "    \"\"\"Check token expiry and refresh if needed\"\"\"\n",
+    "    try:\n",
+    "        payload = bearer_token.split('.')[1]\n",
+    "        payload += '=' * (4 - len(payload) % 4)\n",
+    "        decoded = json.loads(base64.b64decode(payload))\n",
+    "        \n",
+    "        current_time = int(time.time())\n",
+    "        if decoded['exp'] - current_time < 300:\n",
+    "            print(\"🔄 Token expiring soon, refreshing...\")\n",
+    "            new_token = get_refresh_token(client_id, refresh_token, region)\n",
+    "            print(\"✓ Token refreshed successfully\")\n",
+    "            return new_token\n",
+    "        \n",
+    "        return bearer_token\n",
+    "    except Exception as e:\n",
+    "        print(\"🔄 Invalid token, refreshing...\", e)\n",
+    "        traceback.print_exc()\n",
+    "        return get_refresh_token(client_id, refresh_token, region)\n",
+    "\n",
     "async def main():\n",
     "    boto_session = Session()\n",
     "    region = boto_session.region_name\n",
@@ -471,7 +504,12 @@
     "        secret_value = response['SecretString']\n",
     "        parsed_secret = json.loads(secret_value)\n",
     "        bearer_token = parsed_secret['bearer_token']\n",
-    "        print(\"✓ Retrieved bearer token from Secrets Manager\")\n",
+    "        refresh_token = parsed_secret['refresh_token']\n",
+    "        client_id = parsed_secret['client_id']\n",
+    "        print(\"✓ Retrieved credentials from Secrets Manager\")\n",
+    "        \n",
+    "        # Validate and refresh token if needed\n",
+    "        bearer_token = get_valid_token(bearer_token, client_id, refresh_token, region)\n",
     "        \n",
     "    except Exception as e:\n",
     "        print(f\"Error retrieving credentials: {e}\")\n",
@@ -571,12 +609,43 @@
     "import boto3\n",
     "import json\n",
     "import sys\n",
+    "import base64\n",
+    "import time\n",
     "from boto3.session import Session\n",
     "from datetime import timedelta\n",
     "\n",
     "from mcp import ClientSession\n",
     "from mcp.client.streamable_http import streamablehttp_client\n",
     "\n",
+    "def get_refresh_token(client_id, refresh_token, region):\n",
+    "    \"\"\"Refresh access token using refresh token\"\"\"\n",
+    "    cognito_client = boto3.client('cognito-idp', region_name=region)\n",
+    "    auth_response = cognito_client.initiate_auth(\n",
+    "        ClientId=client_id,\n",
+    "        AuthFlow='REFRESH_TOKEN_AUTH',\n",
+    "        AuthParameters={'REFRESH_TOKEN': refresh_token}\n",
+    "    )\n",
+    "    return auth_response['AuthenticationResult']['AccessToken']\n",
+    "\n",
+    "def get_valid_token(bearer_token, client_id, refresh_token, region):\n",
+    "    \"\"\"Check token expiry and refresh if needed\"\"\"\n",
+    "    try:\n",
+    "        payload = bearer_token.split('.')[1]\n",
+    "        payload += '=' * (4 - len(payload) % 4)\n",
+    "        decoded = json.loads(base64.b64decode(payload))\n",
+    "        \n",
+    "        current_time = int(time.time())\n",
+    "        if decoded['exp'] - current_time < 300:\n",
+    "            print(\"🔄 Token expiring soon, refreshing...\")\n",
+    "            new_token = get_refresh_token(client_id, refresh_token, region)\n",
+    "            print(\"✓ Token refreshed successfully\")\n",
+    "            return new_token\n",
+    "        \n",
+    "        return bearer_token\n",
+    "    except:\n",
+    "        print(\"🔄 Invalid token, refreshing...\")\n",
+    "        return get_refresh_token(client_id, refresh_token, region)\n",
+    "\n",
     "async def main():\n",
     "    boto_session = Session()\n",
     "    region = boto_session.region_name\n",
@@ -594,7 +663,12 @@
     "        secret_value = response['SecretString']\n",
     "        parsed_secret = json.loads(secret_value)\n",
     "        bearer_token = parsed_secret['bearer_token']\n",
-    "        print(\"✓ Retrieved bearer token from Secrets Manager\")\n",
+    "        refresh_token = parsed_secret['refresh_token']\n",
+    "        client_id = parsed_secret['client_id']\n",
+    "        print(\"✓ Retrieved credentials from Secrets Manager\")\n",
+    "        \n",
+    "        # Validate and refresh token if needed\n",
+    "        bearer_token = get_valid_token(bearer_token, client_id, refresh_token, region)\n",
     "        \n",
     "    except Exception as e:\n",
     "        print(f\"Error retrieving credentials: {e}\")\n",