Update src/streamlit_app.py

src/streamlit_app.py

@@ -75,7 +75,9 @@ st.markdown("""
 CLIENT_ID     = os.environ.get("AZURE_CLIENT_ID",     "bfcbb298-4cc1-496e-9d9b-ff8c2d967a3a")
 CLIENT_SECRET = os.environ.get("AZURE_CLIENT_SECRET", "")   # set in HF Secrets
 TENANT_ID     = os.environ.get("AZURE_TENANT_ID",     "5dac2bf2-8842-4788-ae07-33fb103b55d6")
-REDIRECT_URI  = os.environ.get("REDIRECT_URI",        "")   # e.g. https://yourspace.hf.space/
+REDIRECT_URI  = os.environ.get("REDIRECT_URI",        "").strip()   # e.g. https://yourspace.hf.space/
+if not REDIRECT_URI and os.environ.get("SPACE_HOST"):
+    REDIRECT_URI = f"https://{os.environ['SPACE_HOST'].strip('/')}/"
 
 AUTHORITY     = f"https://login.microsoftonline.com/{TENANT_ID}"
 AUTH_ENDPOINT = f"{AUTHORITY}/oauth2/v2.0/authorize"
@@ -104,6 +106,36 @@ def configured_group_ids():
 
 DEFAULT_GROUP_IDS = configured_group_ids()
 
+
+class AuthError(RuntimeError):
+    pass
+
+
+def get_query_param_value(params, key, default=""):
+    value = params.get(key, default)
+    if isinstance(value, list):
+        return value[0] if value else default
+    return value or default
+
+
+def is_probably_graph_access_token(token) -> bool:
+    token = (token or "").strip()
+    return token.count(".") in (2, 4)
+
+
+def require_valid_access_token(token) -> str:
+    token = (token or "").strip()
+    if not is_probably_graph_access_token(token):
+        raise AuthError("Microsoft sign-in did not produce a valid Graph access token. Please sign in again.")
+    return token
+
+
+def clear_oauth_query_params():
+    try:
+        st.query_params.clear()
+    except Exception:
+        pass
+
 # ── Session defaults ───────────────────────────────────────────────────────────
 for k, v in {
     "token": None,
@@ -124,6 +156,8 @@ for k, v in {
 # ── OAuth helpers ──────────────────────────────────────────────────────────────
 
 def build_auth_url():
+    if not REDIRECT_URI:
+        raise AuthError("REDIRECT_URI is not configured. Set it to this app's public URL registered in Azure.")
     state = secrets.token_urlsafe(16)
     st.session_state.oauth_state = state
     params = {
@@ -152,27 +186,36 @@ def exchange_code_for_token(code: str) -> str:
     if not resp.ok:
         detail = payload.get("error_description") or payload.get("error") or resp.text
         raise RuntimeError(f"Token exchange failed: {detail}")
-    return payload["access_token"]
+    access_token = payload.get("access_token")
+    if not is_probably_graph_access_token(access_token):
+        detail = payload.get("error_description") or payload.get("error") or "No valid access_token was returned."
+        raise AuthError(f"Token exchange failed: {detail}")
+    return access_token
 
 
 # ── Check for OAuth callback (code in URL query params) ───────────────────────
 query_params = st.query_params
 if not st.session_state.token and "code" in query_params:
-    code  = query_params["code"]
-    state = query_params.get("state", "")
+    code  = get_query_param_value(query_params, "code")
+    state = get_query_param_value(query_params, "state")
     if state == st.session_state.oauth_state or not st.session_state.oauth_state:
         try:
             with st.spinner("Completing sign-in..."):
                 st.session_state.token = exchange_code_for_token(code)
             # Clear the code from the URL
-            st.query_params.clear()
+            clear_oauth_query_params()
             st.rerun()
         except Exception as e:
             st.error(f"Sign-in failed: {e}")
+            clear_oauth_query_params()
+    else:
+        st.error("Sign-in failed: Microsoft returned an unexpected state. Please try signing in again.")
+        clear_oauth_query_params()
 
 # ── Email helpers ──────────────────────────────────────────────────────────────
 
 def get_emails(token, top=50):
+    token = require_valid_access_token(token)
     url = "https://graph.microsoft.com/v1.0/me/mailFolders/inbox/messages"
     headers = {
         "Authorization": f"Bearer {token}",
@@ -189,6 +232,7 @@ def get_emails(token, top=50):
 
 
 def graph_get_json(token, url, params=None):
+    token = require_valid_access_token(token)
     headers = {
         "Authorization": f"Bearer {token}",
         "Prefer": 'outlook.body-content-type="html"',
@@ -200,6 +244,8 @@ def graph_get_json(token, url, params=None):
         payload = {"error": {"message": r.text}}
     if not r.ok:
         detail = (payload.get("error") or {}).get("message") or r.text
+        if r.status_code == 401:
+            raise AuthError(f"Microsoft session is invalid or expired: {detail}")
         raise RuntimeError(f"Graph API error ({r.status_code}): {detail}")
     return payload
 
@@ -589,12 +635,17 @@ with col2:
         st.markdown('<span class="badge-pending">Not signed in</span>', unsafe_allow_html=True)
 
 if not st.session_state.token:
-    auth_url = build_auth_url()
-    st.markdown(
-        f'<div class="login-btn"><a href="{auth_url}" target="_self">🔐 Sign in with Microsoft</a></div>',
-        unsafe_allow_html=True,
-    )
-    st.caption("You'll be redirected to Microsoft's login page and back automatically.")
+    try:
+        auth_url = build_auth_url()
+    except AuthError as e:
+        st.error(str(e))
+        auth_url = None
+    if auth_url:
+        st.markdown(
+            f'<div class="login-btn"><a href="{auth_url}" target="_self">🔐 Sign in with Microsoft</a></div>',
+            unsafe_allow_html=True,
+        )
+        st.caption("You'll be redirected to Microsoft's login page and back automatically.")
 else:
     st.markdown("You are signed in. ✓")
     if st.button("Sign out"):
@@ -656,6 +707,10 @@ if st.button("▶ Run", disabled=not st.session_state.token):
                 raw = get_emails(st.session_state.token, int(top_n))
             st.session_state.messages = raw
             st.session_state.messages_df = pd.DataFrame(build_fetched_emails(raw)) if raw else pd.DataFrame()
+        except AuthError as e:
+            st.error(str(e))
+            st.session_state.token = None
+            st.stop()
         except ValueError as e:
             st.error(str(e))
             st.stop()
@@ -784,4 +839,3 @@ if st.session_state.rj_emails:
         file_name=f"rj_emails_{datetime.now().strftime('%Y%m%d_%H%M')}.xlsx",
         mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
     )
-