Is it possible to create a raw html endpoint for setting session cookie?

Unable to set cookie is a well known limitation of streamlit which make it impossible to keep user session after page refresh.

An idea come to my mind is to use a dedicated login.html endpoint to set cookie and then redirect the app to the target page, just like what we did with php or django years ago.

But then I have another issue: How do I suppose to return a raw html content to this endpoint?

I have ask GPT4 for a solution but it just donโ€™t work for me. Is there any idea to make it work? Or is there better solution of handling userโ€™s login status?

import streamlit as st
from authlib.integrations.requests_client import OAuth2Session
from streamlit_server_state import server_state, get_cookie, set_cookie, delete_cookie
from urllib.parse import urlencode

# Load environment variables (or set directly in the code)
CLIENT_ID = 'your_client_id'
CLIENT_SECRET = 'your_client_secret'
AUTHORIZE_URL = 'https://oauth_provider/authorize'
TOKEN_URL = 'https://oauth_provider/token'
REDIRECT_URI = 'http://localhost:8501/login'
LOGOUT_URL = 'https://oauth_provider/logout'

# Initialize the OAuth2 session
oauth = OAuth2Session(CLIENT_ID, CLIENT_SECRET, redirect_uri=REDIRECT_URI)

# Function to handle OAuth2 callback and set cookies
def handle_login():
    query_params = st.experimental_get_query_params()
    if 'code' in query_params and 'redirect_url' in query_params:
        code = query_params['code'][0]
        redirect_url = query_params['redirect_url'][0]
        token = oauth.fetch_token(TOKEN_URL, code=code)
        
        # Set token in cookies
        set_cookie('oauth_token', token['access_token'])

        # Generate HTML to set cookie and redirect
        st.markdown(f"""
        <html>
        <head>
            <script>
                document.cookie = "oauth_token={token['access_token']}; path=/";
                window.location.href = "{redirect_url}";
            </script>
        </head>
        <body>
            Redirecting...
        </body>
        </html>
        """, unsafe_allow_html=True)
    else:
        st.write("Invalid login attempt")

# Function to handle logout and clear cookies
def handle_logout():
    referer = st.experimental_get_query_params().get('referer', [None])[0]
    
    # Clear the token cookie
    delete_cookie('oauth_token')
    
    # Generate HTML to clear cookie and redirect
    st.markdown(f"""
    <html>
    <head>
        <script>
            document.cookie = "oauth_token=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=/;";
            window.location.href = "{referer}";
        </script>
    </head>
    <body>
        Logging out...
    </body>
    </html>
    """, unsafe_allow_html=True)

# Main block to handle routing
path = st.experimental_get_query_params().get('path', [''])[0]

if path == 'login':
    handle_login()
elif path == 'logout':
    handle_logout()
else:
    st.write("Invalid path")