Problem Saving Checkbox State in Streamlit-AgGrid

I am using streamlit-aggrid in Python and facing an issue when trying to persist checkbox states in a dynamically updated grid.

Problem Context

• I have three DataFrames:
  • df_1: The main grid, which can be updated.
  • df_2: A subset of df_1, where checkboxes are selected.
  • df_3: Created to store the selected indices in st.session_state['saved_i'], ensuring selections persist after updates to df_1.

The Issue

The checkbox states in df_2 are correctly saved after an update to df_1.
However, if two checkboxes are selected consecutively in df_2:

1. The first selection works and is saved correctly.
2. The second selection triggers an update, but its information is not stored in st.session_state['saved_i'] nor appears in grid_response_2.event_data.
3. As a result, the second checkbox selection is lost and is not saved anywhere.

What I Have Tried

• Checked if the selection appears in grid_response_2.event_data (it does not).
• Attempted to manually store the indices before the update (did not work).

Question

How can I ensure that all checkbox selections are correctly saved in st.session_state? I want df_2 to retain the state of selected checkboxes. Code in python:

import pandas as pd
import streamlit as st
import numpy as np
from st_aggrid import AgGrid, GridOptionsBuilder, JsCode, GridUpdateMode
#################################################################################
# Creating df:
df = pd.DataFrame({
    'i': range(20),
    'Values': [f'P{i}' for i in range(20)]
})
df['Selected'] = False
#################################################################################
col1, col2 = st.columns(2)
# Creating grid (df):
gb = GridOptionsBuilder.from_dataframe(df)
gb.configure_default_column(editable=True, groupable=True, resizeable=True)
grid_options = gb.build()
with col1:
    grid_response = AgGrid(df, 
                           gridOptions=grid_options,
                           fit_columns_on_grid_load=False, 
                           height=250,
                           )
#################################################################################
# Creating df_2:
df_2 = grid_response.data.loc[grid_response.data['Selected']].copy()
df_2['Selected'] = False
# Pre-selecting the checkboxes from the previous interaction:
if 'saved_i' in st.session_state:
    i_df_2 = df_2.index
    i_df_3 = st.session_state['saved_i']
    select = i_df_3[np.isin(i_df_3, i_df_2)]
    df_2.loc[select, 'Selected'] = True
#################################################################################
# Creating grid (df_2)
if not df_2.empty:
    gb_2 = GridOptionsBuilder.from_dataframe(df_2)
    gb_2.configure_default_column(editable=True, groupable=True, resizeable=True)
    gb_2.configure_column('Selected', sort='desc')
    grid_options_2 = gb_2.build()
    with col2:
        grid_response_2 = AgGrid(df_2, 
                                gridOptions=grid_options_2,
                                fit_columns_on_grid_load=False,
                                height=250)
    #################################################################################
    # Creating df_3 with selected checkboxes from df_2:
    df_3 = grid_response_2.data.loc[grid_response_2.data['Selected']].copy()
    #################################################################################
    # Saving indexes in session for the next update:
    if not df_3.empty:
        st.session_state['saved_i'] = df_3["i"].astype(str)