Update app.py

app.py

@@ -4,6 +4,7 @@ import requests
 import msal
 import secrets
 import time
+from urllib.parse import urlencode
 
 # Configuration
 APPLICATION_ID = os.getenv('APPLICATION_ID_KEY')
@@ -12,12 +13,6 @@ AUTHORITY = 'https://login.microsoftonline.com/common'
 REDIRECT_URI = 'https://huggingface.co/spaces/awacke1/MSGraphAPI'
 SCOPES = ['User.Read', 'Calendars.ReadWrite', 'Mail.ReadWrite']
 
-# Initialize session state
-if 'auth_state' not in st.session_state:
-    st.session_state.auth_state = None
-if 'auth_state_time' not in st.session_state:
-    st.session_state.auth_state_time = None
-
 # MSAL setup
 def get_msal_app():
     return msal.ConfidentialClientApplication(
@@ -30,13 +25,15 @@ def get_msal_app():
 def generate_auth_url():
     msal_app = get_msal_app()
     state = secrets.token_urlsafe(32)
-    st.session_state.auth_state = state
-    st.session_state.auth_state_time = time.time()
-    return msal_app.get_authorization_request_url(
+    auth_url = msal_app.get_authorization_request_url(
         scopes=SCOPES,
         redirect_uri=REDIRECT_URI,
         state=state
     )
+    # Store the state in query params
+    new_query_params = st.query_params.to_dict()
+    new_query_params['auth_state'] = state
+    return f"{auth_url}&{urlencode(new_query_params)}"
 
 def get_token_from_code(code):
     msal_app = get_msal_app()
@@ -50,16 +47,6 @@ def get_token_from_code(code):
     else:
         raise Exception(f"Error acquiring token: {result.get('error_description')}")
 
-def get_token_from_cache():
-    if 'token_cache' in st.session_state:
-        msal_app = get_msal_app()
-        accounts = msal_app.get_accounts()
-        if accounts:
-            result = msal_app.acquire_token_silent(SCOPES, account=accounts[0])
-            if result:
-                return result
-    return None
-
 # API call function
 def make_api_call(endpoint, token):
     headers = {'Authorization': f'Bearer {token}'}
@@ -76,37 +63,37 @@ def main():
 
     # Debug information
     st.sidebar.write("Debug Info:")
-    st.sidebar.write(f"Auth State: {st.session_state.auth_state}")
-    st.sidebar.write(f"Auth State Time: {st.session_state.auth_state_time}")
-    st.sidebar.write(f"Query Params: {st.query_params}")
+    st.sidebar.write(f"Query Params: {st.query_params.to_dict()}")
 
-    # Check for authentication
-    token = get_token_from_cache()
-    
-    if 'code' in st.query_params:
-        received_state = st.query_params.get('state')
-        if received_state != st.session_state.auth_state:
-            st.error(f"Invalid state parameter. Expected {st.session_state.auth_state}, got {received_state}")
+    if 'code' in st.query_params and 'state' in st.query_params:
+        received_state = st.query_params['state']
+        expected_state = st.query_params.get('auth_state')
+        
+        if received_state != expected_state:
+            st.error(f"Invalid state parameter. Expected {expected_state}, got {received_state}")
             st.error("Please try logging in again.")
-            st.session_state.clear()
+            st.query_params.clear()
             st.rerun()
         
         try:
             token = get_token_from_code(st.query_params['code'])
-            st.session_state['token_cache'] = token
+            st.session_state['token'] = token
+            st.query_params.clear()
             st.success("Successfully authenticated!")
             st.rerun()
         except Exception as e:
             st.error(f"Authentication failed: {str(e)}")
-            st.session_state.clear()
+            st.query_params.clear()
+            st.rerun()
 
-    if not token:
+    if 'token' not in st.session_state:
         auth_url = generate_auth_url()
         st.write("Please log in to continue:")
         st.markdown(f"[Login with Microsoft]({auth_url})")
         return
 
     # User is authenticated, show the main app
+    token = st.session_state['token']
     st.sidebar.success("Authenticated successfully!")
     
     # Display user info