From 090bee36f6fe8ad3733206bcd8034aab5241f87c Mon Sep 17 00:00:00 2001 From: Satveer Khurpa Date: Mon, 10 Nov 2025 17:44:00 -0600 Subject: [PATCH] OAuth2 updates for session storage --- .../chatbot_app_cognito.py | 313 +++++++++++++++--- .../oauth2_callback_server.py | 266 ++++++++++++++- 2 files changed, 528 insertions(+), 51 deletions(-) diff --git a/01-tutorials/03-AgentCore-identity/05-Outbound_Auth_3lo/chatbot_app_cognito.py b/01-tutorials/03-AgentCore-identity/05-Outbound_Auth_3lo/chatbot_app_cognito.py index b807f2f7..fdcf3d38 100644 --- a/01-tutorials/03-AgentCore-identity/05-Outbound_Auth_3lo/chatbot_app_cognito.py +++ b/01-tutorials/03-AgentCore-identity/05-Outbound_Auth_3lo/chatbot_app_cognito.py @@ -1,4 +1,5 @@ import streamlit as st +import streamlit.components.v1 as components import os import json import requests @@ -17,6 +18,60 @@ CONTEXT_WINDOW = 10 # Number of turns (user+assistant pairs) to include in context +def store_token_in_browser(token: str): + """Store the access token in browser's localStorage using JavaScript.""" + components.html( + f""" + + """, + height=0, + ) + + +def get_token_from_browser(): + """Retrieve the access token from browser's localStorage using JavaScript.""" + token_html = components.html( + """ + + """, + height=0, + ) + return token_html + + +def clear_token_from_browser(): + """Clear the access token from browser's localStorage.""" + components.html( + """ + + """, + height=0, + ) + + def get_streamlit_url(): try: # Read the JSON file @@ -258,6 +313,72 @@ def main(): layout="wide", initial_sidebar_state="expanded", ) + + # Add custom CSS for better styling + st.markdown(""" + + """, unsafe_allow_html=True) + import boto3 # Check if configuration loading failed @@ -304,12 +425,38 @@ def main(): if st.session_state["cognito_access_token"] is None: st.markdown( """ -
-
+ + """, @@ -393,33 +540,63 @@ def main(): ) access_token = resp["AuthenticationResult"]["AccessToken"] st.session_state["cognito_access_token"] = access_token - st.success( - "Cognito authentication successful! Redirecting to chatbot..." - ) + + # Store token in browser's localStorage for OAuth flow + store_token_in_browser(access_token) + + # Show success with balloons! + st.balloons() + st.success("🎉 Authentication Successful!") + st.info("🔒 Token securely stored in browser localStorage") + + import time + time.sleep(1.5) # Show success message briefly st.rerun() except Exception as e: st.error(f"Cognito authentication failed: {e}") return # Only return here if not authenticated - # Enhanced system status panel + # Enhanced system status panel with animations st.markdown( f""" -
+ +
⚡ System Status
-
+
Region: {region}
-
+
Agent: - Active + Active
-
+
Session: - Connected + Connected
@@ -549,10 +726,35 @@ def main(): unsafe_allow_html=True, ) - # Enhanced sidebar + # Enhanced sidebar with animations st.sidebar.markdown( """ -
+ +
🤖

Bedrock Agentcore AI

Conversational Intelligence

@@ -561,40 +763,69 @@ def main():

⚙️ Features

-
+
🔄 Real-time Streaming
-
+
🧠 Context Awareness
-
+
🔗 Clickable URLs
+
+ 🔒 + Browser LocalStorage +
""", unsafe_allow_html=True, ) - - # Enhanced main header - st.markdown( - """ -
-
🤖
-

Bedrock Agentcore AI Chatbot

-

Your intelligent conversation partner

-
-
- """, - unsafe_allow_html=True, - ) - + + # Add logout button + if st.sidebar.button("🚪 Logout", use_container_width=True): + clear_token_from_browser() + st.session_state["cognito_access_token"] = None + st.session_state.messages = [] + st.success("Logged out successfully!") + st.rerun() + + # Enhanced main header - More visible! + col1, col2, col3 = st.columns([1, 2, 1]) + with col2: + st.markdown(""" +
+

🤖 Bedrock AgentCore AI

+

Powered by Browser LocalStorage OAuth

+
+ """, unsafe_allow_html=True) + + st.divider() + + # Show session metrics + metric_col1, metric_col2, metric_col3, metric_col4 = st.columns(4) + with metric_col1: + st.metric("🌍 Region", region) + with metric_col2: + st.metric("🔐 Auth", "Active") + with metric_col3: + st.metric("💬 Messages", len(st.session_state.get("messages", []))) + with metric_col4: + st.metric("🔗 Storage", "LocalStorage") + + st.divider() + # Initialize chat history if "messages" not in st.session_state: st.session_state.messages = [] + # Add welcome message + st.session_state.messages.append({ + "role": "assistant", + "content": "👋 Welcome to Bedrock AgentCore AI! I'm here to help you with intelligent conversations. Your session is secured with browser localStorage for OAuth flows. How can I assist you today?" + }) if "genesisSessionId" not in st.session_state: st.session_state["genesisSessionId"] = ( genesisSessionId # Initialize from config @@ -665,20 +896,22 @@ def main(): accumulated_response = "" try: - # Setup streaming client + # Setup streaming client with progress indicator + progress_bar = st.progress(0, text="🔄 Initializing...") session_id = st.session_state.get("genesisSessionId") context = build_context(st.session_state.messages, CONTEXT_WINDOW) payload = json.dumps({"prompt": context}) bearer_token = st.session_state.get("cognito_access_token") store_token_in_oauth2_callback_server(bearer_token) + progress_bar.progress(30, text="🔐 Authenticating...") streaming_client = StreamingHttpBedrockAgentCoreClient(region) + progress_bar.progress(60, text="🤖 Connecting to AI...") - # Show initial thinking state with pulsing animation - message_placeholder.markdown( - '🤖 💭 Bedrock Agentcore is thinking...', - unsafe_allow_html=True, - ) + # Show initial thinking state + progress_bar.progress(100, text="✨ Generating response...") + with st.spinner("🤖 AI is thinking..."): + message_placeholder.markdown("💭 Thinking...") # Stream the response with animations chunk_count = 0 @@ -783,6 +1016,7 @@ def main(): time.sleep(0.02) # Final response with timing (remove streaming classes) + progress_bar.empty() # Clear progress bar elapsed = time.time() - start_time answer = ( formatted_response @@ -800,6 +1034,7 @@ def main(): ) except Exception as e: + progress_bar.empty() # Clear progress bar on error error_msg = f"Sorry, I encountered an error: {str(e)}" message_placeholder.markdown( f'
🤖 ❌ {error_msg}
', diff --git a/01-tutorials/03-AgentCore-identity/05-Outbound_Auth_3lo/oauth2_callback_server.py b/01-tutorials/03-AgentCore-identity/05-Outbound_Auth_3lo/oauth2_callback_server.py index 98738006..04711a20 100644 --- a/01-tutorials/03-AgentCore-identity/05-Outbound_Auth_3lo/oauth2_callback_server.py +++ b/01-tutorials/03-AgentCore-identity/05-Outbound_Auth_3lo/oauth2_callback_server.py @@ -243,37 +243,279 @@ async def _handle_oauth2_callback(session_id: str): html_content = """ - + - OAuth2 Success + + + OAuth2 Success - AgentCore + +
+
+
+
+
+
+
+
+
+
+
+
-

Completed OAuth2 3LO flow successfully

+ +
+ + + + +
+ +

🎉 OAuth2 Authentication Complete!

+

Your authorization was successful

+ +
+
+ ✓ + 3-Legged OAuth Flow Completed +
+
+ 🔐 + Access Token Secured +
+
+ 🔗 + Session Bound to AgentCore +
+
+ +

+ This window will close automatically +

+ + """