Bugfix: Adding refresh token auth flow for expired bearer token for MCP Server

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",

01-tutorials/utils.py

@@ -1,8 +1,12 @@
 import boto3
 import json
 import time
+import base64
 from boto3.session import Session
 
+USER_NAME = 'testuser'
+PASSWORD = 'MyPassword123!'
+TEMP_ADMIN_PASSWORD = 'Temp123!'
 
 def setup_cognito_user_pool():
     boto_session = Session()
@@ -34,38 +38,41 @@
         # Create User
         cognito_client.admin_create_user(
             UserPoolId=pool_id,
-            Username='testuser',
-            TemporaryPassword='Temp123!',
+            Username=USER_NAME,
+            TemporaryPassword=TEMP_ADMIN_PASSWORD,
             MessageAction='SUPPRESS'
         )
         # Set Permanent Password
         cognito_client.admin_set_user_password(
             UserPoolId=pool_id,
-            Username='testuser',
-            Password='MyPassword123!',
+            Username=USER_NAME,
+            Password=PASSWORD,
             Permanent=True
         )
         # Authenticate User and get Access Token
         auth_response = cognito_client.initiate_auth(
             ClientId=client_id,
             AuthFlow='USER_PASSWORD_AUTH',
             AuthParameters={
-                'USERNAME': 'testuser',
-                'PASSWORD': 'MyPassword123!'
+                'USERNAME': USER_NAME,
+                'PASSWORD': PASSWORD
             }
         )
         bearer_token = auth_response['AuthenticationResult']['AccessToken']
+        refresh_token = auth_response['AuthenticationResult']['RefreshToken']
         # Output the required values
         print(f"Pool id: {pool_id}")
         print(f"Discovery URL: https://cognito-idp.{region}.amazonaws.com/{pool_id}/.well-known/openid-configuration")
         print(f"Client ID: {client_id}")
         print(f"Bearer Token: {bearer_token}")
+        print(f"Refresh Token: {refresh_token}")
 
         # Return values if needed for further processing
         return {
             'pool_id': pool_id,
             'client_id': client_id,
             'bearer_token': bearer_token,
+            'refresh_token': refresh_token,
             'discovery_url':f"https://cognito-idp.{region}.amazonaws.com/{pool_id}/.well-known/openid-configuration"
         }
     except Exception as e:
@@ -83,8 +90,8 @@
         ClientId=client_id,
         AuthFlow='USER_PASSWORD_AUTH',
         AuthParameters={
-            'USERNAME': 'testuser',
-            'PASSWORD': 'MyPassword123!'
+            'USERNAME': USER_NAME,
+            'PASSWORD': PASSWORD
         }
     )
     bearer_token = auth_response['AuthenticationResult']['AccessToken']

CONTRIBUTORS.md

@@ -6,6 +6,7 @@
 - aristsakpinis93
 - aurbac
 - bergjaak
+- crupakheti
 - danystinson
 - dhawalkp
 - didhd