Dialog doesn't close until the rerun is complete making app look frozen

I have a app that has a chatbot (chat_input) letting the user make requests as normal. In addition, in a certain scenario am launching a dialog and have the submission of that dialog also kick off a chatbot request.

To boil down my problem though, I created this simple app running locally on python 3.11.9 with streamlit 1.42.1 that has the same issue.

When the dialog closes with the button press, if the rerun of the app ends up doing long running work due to whatever conditions are true during the rerun, the dialog doesn’t close right away, but waits to close until the long running work finishes.

See below where pressing Button A in the dialog, it greys out the dialog but doesn’t close it until the full re-run completes.

Is there a way around this? Since in my real app I want to submit a chatbot request, I want the dialog to close so that the user can see the chatbot is working on a response, but if it’s covered by the dialog, they might think it’s just frozen.

import streamlit as st  
import time


#initialize session state
if 'text_content' not in st.session_state:  
    st.session_state['text_content'] = ''  

#wait some time to mimic long running work, like a chat bot working
def sleep_timer():
    time.sleep(5)


@st.dialog("Dialog header")
def sample_prompt():
    if st.button('Button A in dialog Pressed'):  
        st.session_state['text_content'] = 'Button A was pressed in dialog'              
        st.rerun()   #close dialog

    if st.button('Button B in dialog Pressed'):  
        st.session_state['text_content'] = 'Button B was pressed in dialog' 
        st.rerun()  #close dialog


text_to_show ="Starter Text"

#if session state text was updated, show that text instead
if st.session_state['text_content']:
    text_to_show = st.session_state['text_content']

#if they happened to press button A, then do some extra work that takes a while
if st.session_state['text_content'] == 'Button A was pressed in dialog' :
    with st.spinner("In progress...", show_time=True):
        sleep_timer()

if st.button("Open Dialog"):
    st.session_state['text_content'] = ''  #if they open the dialog, clear out the old value first
    sample_prompt()

st.write(text_to_show)