External function is not saving correctly a variable in session_state

Hi,

I have a local-running Streamlit app, running streamlit==1.28.2 and Python 3.11.5.

I’m implementing a standard Microsoft Azure AD authentication, therefore my page is running the following function whenever the page has a code and session_state in the URL params:

def oauth_azure_ad(config='azure'):
    print("Called st_oauth_azure_ad() function")
    if 'azuread_token' not in st.session_state:
        print("azuread_token not in session_state")
    else:
        print(f"azuread_token: {st.session_state['azuread_token']}")

    config = st.secrets[config]

    code = st.experimental_get_query_params()['code'][0]
    state = st.experimental_get_query_params()['session_state'][0]
    
    headers = { 'Content-Type': 'application/x-www-form-urlencoded' }
    data = {
        'grant_type': 'authorization_code',
        'scope': 'openid profile',
        'code': code,
        'client_id': config['client_id'],
        'client_secret': config['client_secret'],
        'redirect_uri': config['redirect_uri'],
    }

    if 'azuread_token' not in st.session_state:
        try:
            print(f"Sending POST to token endpoint")
            response = requests.post(config['token_endpoint'], headers=headers, data=data)
            response.raise_for_status()
            token = response.json()
            print("Saving token in session state")
            st.session_state['azuread_token'] = token
        except requests.exceptions.RequestException as e:
            print(f"Error sending POST: {e}")
    if 'azuread_token' in st.session_state and st.session_state['azuread_token'] is not None:
        print(f"token is {st.session_state['azuread_token']}")
        valid, msg = validate_azure_token(st.session_state['azuread_token'], config)
        print(f"Returning: {msg['name']}, {msg['username']}, {valid}")
        return msg["name"], msg["username"], valid

The page is running this script:

if 'code' in st.experimental_get_query_params() and 'session_state' in st.experimental_get_query_params():
    st.session_state['name'], st.session_state['username'], st.session_state['authentication_status'] = oauth_azure_ad()

The problem is that:

  • oauth_azure_ad() is being run multiple times
  • the very first time it runs the requests.post statement works correctly, but apparently the statement st.session_state['azuread_token'] = token is completely ignored
  • obviously the next time the function is called, since 'azuread_token' is still not present in session_state, the function sends the POST again, but Azure AD clearly replies with an error message, since the token has already been redeemed.

To be more clear, here is the output of some print statements from a function run:

Called st_oauth_azure_ad() function
azuread_token not in session_state
Sending POST to token endpoint
Called st_oauth_azure_ad() function
azuread_token not in session_state
Sending POST to token endpoint
token is {'error': 'invalid_grant', 'error_description': 'AADSTS54005: OAuth2 Authorization code was already redeemed, please retry with a new valid code or use an existing refresh token. Trace ID: 8349e6d6-0e8c-401f-b435-2991bc080f01 Correlation ID: a90c3bc4-dd89-4800-a2d9-bcc4e9527671 Timestamp: 2023-11-25 08:14:30Z', 'error_codes': [54005], 'timestamp': '2023-11-25 08:14:30Z', 'trace_id': '8349e6d6-0e8c-401f-b435-2991bc080f01', 'correlation_id': 'a90c3bc4-dd89-4800-a2d9-bcc4e9527671'}

Where you read the very first β€œSending POST to token endpoint”, my expectation would have been that the function stored the token in the session_state['azuread_token'] variable, but this is not happening.

I’m really puzzled by this behaviour.