[Coming soon] Using st.data_editor with session state, input data "disappears": Requires double input to register changes.

Checklist

• I have searched the existing issues for similar issues.
• I added a very descriptive title to this issue.
• I have provided sufficient information below to help reproduce this issue.

Summary

This issue is practically the same as this closed issue: #7354 But since it is closed (due to the user not really needing the particular use case) I have opened this new one with another simple minimum reproducible example.

Reproducible Code Example

import streamlit as st
import pandas as pd

# using session state, works ever other input:
if "df" not in st.session_state:
    st.session_state.df = pd.DataFrame({"In":[0,1,2],"Out":[0,0,0]})
st.session_state.df = st.data_editor(st.session_state.df)
st.session_state.df["Out"] = st.session_state.df["In"]**2
st.write(st.session_state.df)

# # works every input but the df is not updated in data_editor:
# df = pd.DataFrame({"In":[0,1,2],"Out":[0,0,0]})
# df = st.data_editor(df)
# df["Out"] = df["In"]**2
# st.write(df)

Steps To Reproduce

1. Click on a field in the in column.
2. Input a number, enter.
3. Watch for changes in the dataframes.

Expected Behavior

I expect that each input in the data_editor to result in the rest of the code to run, before it loops back to the top and into the data_editor with new data which is displayed. In this example I expect the "state" version to behave just like the "non-state" example, with the added functionality of the state of course.

Current Behavior

1. I input a number(say 5) in a input field press enter.
2. This new input is replaced with the previous value instantly.
3. I input 5 in the same box again.
4. The edit_data frame now has the correct input, but not the updated input. The st.write frame is updated to reflect the correct values.

Another variant:

1. I input 5 in row 1. -> instantly replaced, nothing in st.write
2. I input 4 in row 2. -> editor: input is updated, output is not. st.write: row two is updated.
3. I input 3 in row 3. -> editor: row 2 is updated, row 3 ignored. Both editor and write now agrees. OR: in another box I input a changed number (not the same as exists already, that doesnt do anything)

Pretty hard to explain, please try for yourself to get an idea how it really behaves. Basically the state of data_editor is always one loop behind in updating, and the input seems to dissapear until another change is made.

Is this a regression?

• Yes, this used to work in a previous version.

Debug info

• Streamlit version: 1.28.1
• Python version: 3.12
• Operating System: Win 10
• Browser: Edge or Firefox

Additional Information

See my earlier comment on the other issue: Similar problem from the docs, with an incompatible solution: Widget updating session_state problem

Please help me with workarounds if possible. Thank you for considering this issue, and for Streamlit itself!