Using st.session_state with widget key and on_change correctly? How to pass the current widget value?

Hi,
I am struggling to understand st.session_state. If I use the st.text_input widget key and on_change, the current value of the widget cannot be used in the code below that widget.

Why is this the case? I do not understand, i.e. why in page2.py text_2 (or st.session_state.text_2) is not updated. How can this be fixed (using the widget key and on_change)?

Here my example app:

.
β”œβ”€β”€ home.py
└── pages
    β”œβ”€β”€ page1.py
    └── page2.py

home.py

import streamlit as st

"# Home"

if "text_1" not in st.session_state:
    st.session_state.text_1 = "A"

if "text_2" not in st.session_state:
    st.session_state.text_2 = "B"

if "text_3" not in st.session_state:
    st.session_state.text_3 = ""

"#### Session state"
st.json(dict(sorted(st.session_state.items())))

page1.py

import streamlit as st

"#### Session state"
st.json(dict(sorted(st.session_state.items())))

## fix reload on other pages (*not* main page)
for v in ["text_1", "text_2",]:
    if v not in st.session_state:
        st.switch_page('home.py')

"# Page 1 (working)"

text_1 = st.text_input("Text 1", value=st.session_state.text_1)
st.session_state.text_1 = text_1

st.session_state.text_3 = ", ".join([st.session_state.text_1, st.session_state.text_2])
st.write(text_1)
f"{st.session_state.text_3}"

"#### Session state"
st.json(dict(sorted(st.session_state.items())))

page2.py

import streamlit as st

"#### Session state"
st.json(dict(sorted(st.session_state.items())))

## fix reload on other pages (*not* main page)
for v in ["text_1", "text_2",]:
    if v not in st.session_state:
        st.switch_page('home.py')

"# Page 2 (unexpected behavior)"
"#### Problem: st.session_state.text_2 is not updated (at the end)!"
"#### Question: How to pass the current value from st.text_input to on_change function?**"

def on_change_text_2():
    st.session_state.text_3 = ", ".join([st.session_state.text_1, st.session_state.text_2])

text_2 = st.text_input("Text 2", value=st.session_state.text_2, key="text_2", on_change=on_change_text_2)

f"text_2 (**NOT updated!**): `{text_2}`" # Why is text_2 never updated???

"#### Session state"
st.json(dict(sorted(st.session_state.items())))

I am not sure what you are trying to experiment, but this is the basic concept.

1. Get the user input via the widget key thru the session state.

start_value = 'start'
st.text_input("Text 2", value=start_value, key="text_2")

input_value = st.session_state.text_2
st.write(f"input value: {input_value}")

To avoid complications, do not create session state key elsewhere something like this.

if 'text_2' not in st.session_state:
    st.session_state.text_2 = None

start_value = 'start'
st.text_input("Text 2", value=start_value, key="text_2")

input_value = st.session_state.text_2
st.write(f"input value: {input_value}")

That key is now used twice, this is not a good practice.

2. Get the user input via the widget return value without defining a key.

start_value = 'start'
input_value = st.text_input("Text 2", value=start_value)
st.write(f"input value: {input_value}")

In that case streamlit will assign a unique key to this widget. It is important for streamlit that widgets keys are unique. But we are not interested on that, what we want is to get the input of the user.

In this case, you can define keys freely, just make sure that this key is not used by other widgets where you define its key.

if 'whatever' not in st.session_state:
    st.session_state.whatever = None

Reference