Set ctx.current_fragment_id in both full script and fragment runs (#8931

)

We previously used `ctx.current_fragment_id` to determine whether a script runner execution is a full
script run or a fragment rerun. We later switched to using `ctx.fragment_ids_this_run` to support
multiple fragments being run from a single script runner but still had the script runner be responsible
for setting `ctx.current_fragment_id` outside of the fragment's outer function.

It turns out that having the fragment's outer function be reponsible for this fixes the issue we were having
with nested fragments not working correctly, so this PR makes the simple required change to fix nested
fragments/dialogs.

Closes #8635

e2e_playwright/st_fragments_nested.py

@@ -0,0 +1,37 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2024)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from uuid import uuid4
+
+import streamlit as st
+
+
+@st.experimental_fragment
+def outer_fragment():
+    with st.container(border=True):
+        st.write(f"outer fragment: {uuid4()}")
+        st.button("rerun outer fragment")
+        inner_fragment()
+
+
+@st.experimental_fragment
+def inner_fragment():
+    with st.container(border=True):
+        st.write(f"inner fragment: {uuid4()}")
+        st.button("rerun inner fragment")
+
+
+st.write(f"outside all fragments: {uuid4()}")
+st.button("rerun whole app")
+outer_fragment()

e2e_playwright/st_fragments_nested_test.py

@@ -0,0 +1,67 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2024)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from playwright.sync_api import Page, expect
+
+from e2e_playwright.shared.app_utils import click_button
+
+
+def get_uuids(app: Page):
+    expect(app.get_by_test_id("stMarkdown")).to_have_count(3)
+
+    outside_fragment_text = app.get_by_test_id("stMarkdown").first.text_content()
+    inner_fragment_text = app.get_by_test_id("stMarkdown").nth(1).text_content()
+    outer_fragment_text = app.get_by_test_id("stMarkdown").last.text_content()
+
+    return outside_fragment_text, inner_fragment_text, outer_fragment_text
+
+
+def test_full_app_rerun(app: Page):
+    outside_fragment_text, inner_fragment_text, outer_fragment_text = get_uuids(app)
+
+    click_button(app, "rerun whole app")
+
+    # The full app reran, so all of the UUIDs in the app should have changed.
+    expect(app.get_by_test_id("stMarkdown").first).not_to_have_text(
+        outside_fragment_text
+    )
+    expect(app.get_by_test_id("stMarkdown").nth(1)).not_to_have_text(
+        inner_fragment_text
+    )
+    expect(app.get_by_test_id("stMarkdown").last).not_to_have_text(outer_fragment_text)
+
+
+def test_outer_fragment_rerun(app: Page):
+    outside_fragment_text, inner_fragment_text, outer_fragment_text = get_uuids(app)
+
+    click_button(app, "rerun outer fragment")
+
+    # We reran the outer fragment, so the UUID outside of the fragments should stay
+    # constant, but the other two should have changed.
+    expect(app.get_by_test_id("stMarkdown").first).to_have_text(outside_fragment_text)
+    expect(app.get_by_test_id("stMarkdown").nth(1)).not_to_have_text(
+        inner_fragment_text
+    )
+    expect(app.get_by_test_id("stMarkdown").last).not_to_have_text(outer_fragment_text)
+
+
+def test_inner_fragment_rerun(app: Page):
+    outside_fragment_text, inner_fragment_text, outer_fragment_text = get_uuids(app)
+
+    click_button(app, "rerun inner fragment")
+
+    # We reran the inner fragment. Only that corresponding UUID should have changed.
+    expect(app.get_by_test_id("stMarkdown").first).to_have_text(outside_fragment_text)
+    expect(app.get_by_test_id("stMarkdown").nth(1)).to_have_text(inner_fragment_text)
+    expect(app.get_by_test_id("stMarkdown").last).not_to_have_text(outer_fragment_text)

lib/streamlit/runtime/fragment.py

@@ -176,13 +176,16 @@ def wrapped_fragment():
                 ctx.cursors = deepcopy(cursors_snapshot)
                 dg_stack.set(deepcopy(dg_stack_snapshot))
             else:
-                # Otherwise, we must be in a full script run. We need to temporarily set
-                # ctx.current_fragment_id so that elements corresponding to this
-                # fragment get tagged with the appropriate ID. ctx.current_fragment_id
-                # gets reset after the fragment function finishes running.
-                ctx.current_fragment_id = fragment_id
+                # Otherwise, we must be in a full script run. We keep track of all
+                # fragments defined in this script run to ensure that we don't
+                # delete them when we clean up this session's fragment storage.
                 ctx.new_fragment_ids.add(fragment_id)
 
+            # Set ctx.current_fragment_id so that elements corresponding to this
+            # fragment get tagged with the appropriate ID. ctx.current_fragment_id gets
+            # reset after the fragment function finishes running.
+            ctx.current_fragment_id = fragment_id
+
             try:
                 # Make sure we set the active script hash to the same value
                 # for the fragment run as when defined upon initialization

lib/streamlit/runtime/scriptrunner/script_runner.py

@@ -559,7 +559,6 @@ def code_to_exec(code=code, module=module, ctx=ctx, rerun_data=rerun_data):
                                 wrapped_fragment = self._fragment_storage.get(
                                     fragment_id
                                 )
-                                ctx.current_fragment_id = fragment_id
                                 wrapped_fragment()
 
                             except KeyError:

lib/tests/streamlit/runtime/fragment_test.py

@@ -175,12 +175,11 @@ def my_fragment():
         ctx.fragment_storage.set.assert_called_once()
 
     @patch("streamlit.runtime.fragment.get_script_run_ctx")
-    def test_sets_dg_stack_and_cursor_to_snapshots_if_current_fragment_id_set(
+    def test_sets_dg_stack_and_cursor_to_snapshots_if_fragment_ids_this_run(
         self, patched_get_script_run_ctx
     ):
         ctx = MagicMock()
         ctx.fragment_ids_this_run = {"my_fragment_id"}
-        ctx.current_fragment_id = "my_fragment_id"
         ctx.fragment_storage = MemoryFragmentStorage()
         patched_get_script_run_ctx.return_value = ctx
 
@@ -196,6 +195,8 @@ def test_sets_dg_stack_and_cursor_to_snapshots_if_current_fragment_id_set(
         def my_fragment():
             nonlocal call_count
 
+            assert ctx.current_fragment_id is not None
+
             curr_dg_stack = dg_stack.get()
             # Verify that mutations made in previous runs of my_fragment aren't
             # persisted.
@@ -217,17 +218,17 @@ def my_fragment():
         # Verify that we can't mutate our dg_stack from within my_fragment. If a
         # mutation is persisted between fragment runs, the assert on `my_random_field`
         # will fail.
-        ctx.current_fragment_id = "my_fragment_id"
         saved_fragment()
-        ctx.current_fragment_id = "my_fragment_id"
         saved_fragment()
 
         # Called once when calling my_fragment and three times calling the saved
         # fragment.
         assert call_count == 3
 
     @patch("streamlit.runtime.fragment.get_script_run_ctx")
-    def test_sets_current_fragment_id_if_not_set(self, patched_get_script_run_ctx):
+    def test_sets_current_fragment_id_in_full_script_runs(
+        self, patched_get_script_run_ctx
+    ):
         ctx = MagicMock()
         ctx.fragment_ids_this_run = set()
         ctx.new_fragment_ids = set()