Do not clear nodes for fragment-finished-via-rerun messages (streamli…

…t#9389)

## Describe your changes

Related to streamlit#9372 

### Changes

Two changes we make here:
1. Right now, we trigger a stale-node clearing also for fragment runs
that were interrupted by a rerun message. This leads to a situation
where elements can disappear for a moment (see [this user comment with
video](streamlit#9372 (comment))).
This PR moves the clearing into the if-block that is only executed for
finished complete runs.
2. We capture all element ids we rendered and if we encounter the same
id twice, we ignore it (prevents regressing on streamlit#8360)

### Analysis

For the example app:

<details>
<summary>Example App</summary>

```python
import streamlit as st
import time

if 'count' not in st.session_state:
  st.session_state.count = 0

if 'foo' not in st.session_state:
    st.write(f'Setting state ({st.session_state.count})')
    st.session_state['foo'] = 'bar'

if st.button('Run'):
    # Allow for forward message queue to flush button element
    time.sleep(1)

    st.session_state.count += 1
    del st.session_state['foo']
    st.rerun()
```

</details>

the flow is as following (ever step triggers a new session message with
a new `scriptRunId`):
<table>
<tr>
  <th>Step</th>
  <th>Phase / Action</th>
  <th>Messages the Frontend receives</th>
  <th>Comment</th>
</tr>
<tr>
 <td>1</td>
  <td>First App Run</td>
  <td>[Markdown, Button]</td>
  <td></td>
</tr>
<tr>
 <td>2</td>
  <td>Rerun by Button Click</td>
  <td>[Button]</td>
<td>The button click itself triggers an app rerun. <code>foo</code> is
in the session state (not deleted yet), so the <code>if</code> block is
<i>not</i> executed.</td>
</tr>
<tr>
 <td>3</td>
  <td>Rerun by <code>st.rerun</code> in the button's if-block</td>
  <td>[Markdown, Button]</td>
<td>The <code>st.rerun</code> line executes a rerun. <code>foo</code> is
deleted from the session state, so the <code>if</code> block is executed
again. <b>No stale-nodes cleaning happened</b> because its an early
interrupt by a rerun.</td>
</tr>
<tr>
 <td>4</td>
  <td>App Run is Complete</td>
  <td>Session Finished</td>
<td>The stale-nodes are cleaned because this is a full app finished
run.</td>
</tr>
</table>

With fix 1 (move cleaning of stale nodes to finished successfully runs),
what happens is that the `Button` element of step 3 is sent to the
frontend, but the `Button` element of step 2 is still in the `elements`
list, because it was not cleaned up (after the fix 1 change). Both
elements have the same React key, which throws an error and leads to the
stale React element being shown.
With fix 2, this is addressed because the `Button` element is sent to
the frontend, now the list of elements still contain two `Button`
elements with the same id (remember no cleaning happened). But now, only
one element with the same ket is rendered and the second is ignored. Our
`applyDelta`'s `setIn` logic leads to the fact that the _new_ `Button`
element appears in the `elements` list _before_ the old one, which means
that the newer element is rendered and the older one is ignored.
In step 4, the script finishes successfully completely (a full top to
bottom run), and. thus. a `cleanStaleNodes` execution is triggered which
establishes the final, cleaned state. Now, the frontend again has just
two elements in its list: `[Markdown, Button]`.

## Reproduction

In this comment I have linked a video and an example app:
streamlit#9389 (comment)

## GitHub Issue Link (if applicable)

## Testing Plan

- Explanation of why no additional tests are needed
- Unit Tests (JS and/or Python)
- E2E Tests
- Any manual testing needed?

---

**Contribution License Agreement**

By submitting this pull request you agree that all contributions to this
project are made under the Apache 2.0 license.

---------

Co-authored-by: Ken McGrady <[email protected]>

e2e_playwright/st_rerun.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import time
+
 import streamlit as st
 
 if "count" not in st.session_state:
@@ -66,3 +68,14 @@ def fragment_with_rerun_in_try_block():
 # have elements in the main app after fragments to ensure that the main app elements are
 #  not cleared when rerunning fragments
 st.write(f"app run count: {st.session_state.count}")
+
+if "foo" not in st.session_state:
+    st.write("Setting state")
+    st.session_state["foo"] = "bar"
+
+if st.button("#8599 - Bug"):
+    # Allow for forward message queue to flush button element
+    time.sleep(1)
+
+    del st.session_state["foo"]
+    st.rerun()

e2e_playwright/st_rerun_test.py

@@ -25,7 +25,7 @@ def _expect_initial_reruns_finished(app: Page):
 
 
 def _expect_initial_reruns_count_text(app: Page):
-    expect(app.get_by_test_id("stMarkdown").last).to_have_text("app run count: 4")
+    expect_markdown(app, "app run count: 4")
 
 
 def test_st_rerun_restarts_the_session_when_invoked(app: Page):
@@ -57,7 +57,7 @@ def test_rerun_works_in_try_except_block(app: Page):
     click_button(app, "rerun try_fragment")
     # the rerun in the try-block worked as expected, so the session_state count
     # incremented
-    expect(app.get_by_test_id("stMarkdown").last).to_have_text("app run count: 5")
+    expect_markdown(app, "app run count: 5")
 
 
 def test_state_retained_on_app_scoped_rerun(app: Page):
@@ -76,3 +76,10 @@ def test_state_retained_on_app_scoped_rerun(app: Page):
     # Rerun the fragment and verify that the selectbox kept its state
     click_button(app, "rerun whole app (from fragment)")
     expect_markdown(app, "selectbox selection: a")
+
+
+# From GitHub issue #8599
+def test_clears_stale_elements_correctly(app: Page):
+    click_button(app, "#8599 - Bug")
+
+    expect(app.get_by_text("#8599 - Bug")).to_have_count(1)

frontend/app/src/App.test.tsx

@@ -32,12 +32,10 @@ import {
   CUSTOM_THEME_NAME,
   CustomThemeConfig,
   FileUploadClient,
-  ForwardMsg,
   getDefaultTheme,
   getHostSpecifiedTheme,
   HOST_COMM_VERSION,
   HostCommunicationManager,
-  INewSession,
   lightTheme,
   LocalStore,
   mockEndpoints,
@@ -50,6 +48,23 @@ import {
   toExportedTheme,
   WidgetStateManager,
 } from "@streamlit/lib"
+import {
+  Delta,
+  Element,
+  ForwardMsg,
+  ForwardMsgMetadata,
+  IAutoRerun,
+  ILogo,
+  INavigation,
+  INewSession,
+  IPageConfig,
+  IPageInfo,
+  IPageNotFound,
+  IPagesChanged,
+  IParentMessage,
+  SessionEvent,
+  SessionStatus,
+} from "@streamlit/lib/src/proto"
 import { SegmentMetricsManager } from "@streamlit/app/src/SegmentMetricsManager"
 import { ConnectionManager } from "@streamlit/app/src/connection/ConnectionManager"
 import { ConnectionState } from "@streamlit/app/src/connection/ConnectionState"
@@ -296,14 +311,32 @@ function getMockConnectionManagerProp(propName: string): any {
   return getStoredValue<ConnectionManager>(ConnectionManager).props[propName]
 }
 
+type DeltaWithElement = Omit<Delta, "fragmentId" | "newElement" | "toJSON"> & {
+  newElement: Omit<Element, "toJSON">
+}
+
+type ForwardMsgType =
+  | DeltaWithElement
+  | ForwardMsg.ScriptFinishedStatus
+  | IAutoRerun
+  | ILogo
+  | INavigation
+  | INewSession
+  | IPagesChanged
+  | IPageConfig
+  | IPageInfo
+  | IParentMessage
+  | IPageNotFound
+  | Omit<SessionEvent, "toJSON">
+  | Omit<SessionStatus, "toJSON">
+
 function sendForwardMessage(
-  type: string,
-  message: any,
-  metadata: any = null
+  type: keyof ForwardMsg,
+  message: ForwardMsgType,
+  metadata: Partial<ForwardMsgMetadata> | null = null
 ): void {
   act(() => {
     const fwMessage = new ForwardMsg()
-    // @ts-expect-error
     fwMessage[type] = cloneDeep(message)
     if (metadata) {
       fwMessage.metadata = metadata
@@ -1721,6 +1754,145 @@ describe("App", () => {
       const connectionManager = getMockConnectionManager()
       expect(connectionManager.incrementMessageCacheRunCount).toBeCalled()
     })
+
+    it("will clear stale nodes if finished successfully", async () => {
+      renderApp(getProps())
+      sendForwardMessage("newSession", NEW_SESSION_JSON)
+      // Run the script with one new element
+      sendForwardMessage("sessionStatusChanged", {
+        runOnSave: false,
+        scriptIsRunning: true,
+      })
+      // this message now belongs to this^ session
+      sendForwardMessage(
+        "delta",
+        {
+          type: "newElement",
+          newElement: {
+            type: "text",
+            text: {
+              body: "Here is some text",
+              help: "",
+            },
+          },
+        },
+        { deltaPath: [0, 0] }
+      )
+
+      // start new session
+      sendForwardMessage("newSession", {
+        ...NEW_SESSION_JSON,
+        scriptRunId: "different_script_run_id",
+      })
+
+      await waitFor(() => {
+        expect(screen.getByText("Here is some text")).toBeInTheDocument()
+      })
+
+      sendForwardMessage(
+        "scriptFinished",
+        ForwardMsg.ScriptFinishedStatus.FINISHED_SUCCESSFULLY
+      )
+
+      await waitFor(() => {
+        expect(screen.queryByText("Here is some text")).not.toBeInTheDocument()
+      })
+    })
+
+    it("will not clear stale nodes if finished with rerun", async () => {
+      renderApp(getProps())
+      sendForwardMessage("newSession", NEW_SESSION_JSON)
+      // Run the script with one new element
+      sendForwardMessage("sessionStatusChanged", {
+        runOnSave: false,
+        scriptIsRunning: true,
+      })
+      // these messages now belongs to this^ session
+      sendForwardMessage(
+        "delta",
+        {
+          type: "newElement",
+          newElement: {
+            type: "text",
+            text: {
+              body: "Here is some text",
+              help: "",
+            },
+          },
+        },
+        { deltaPath: [0, 0] }
+      )
+      sendForwardMessage(
+        "delta",
+        {
+          type: "newElement",
+          newElement: {
+            type: "text",
+            text: {
+              body: "Here is some other text",
+              help: "",
+            },
+          },
+        },
+        { deltaPath: [0, 1] }
+      )
+
+      // start new session
+      sendForwardMessage("newSession", {
+        ...NEW_SESSION_JSON,
+        scriptRunId: "different_script_run_id",
+      })
+
+      // this message now belongs to this^ session. It overrides the first message of
+      // the previous session because the same delta path is used
+      sendForwardMessage(
+        "delta",
+        {
+          type: "newElement",
+          newElement: {
+            type: "text",
+            text: {
+              body: "Here is some new text",
+              help: "",
+            },
+          },
+        },
+        { deltaPath: [0, 0] }
+      )
+
+      sendForwardMessage(
+        "scriptFinished",
+        ForwardMsg.ScriptFinishedStatus.FINISHED_EARLY_FOR_RERUN
+      )
+
+      await waitFor(() => {
+        expect(screen.getByText("Here is some new text")).toBeInTheDocument()
+      })
+      // this message was overridden because same delta-path was used be the 'new text' message
+      await waitFor(() => {
+        expect(screen.queryByText("Here is some text")).not.toBeInTheDocument()
+      })
+      await waitFor(() => {
+        expect(screen.getByText("Here is some other text")).toBeInTheDocument()
+      })
+
+      sendForwardMessage(
+        "scriptFinished",
+        ForwardMsg.ScriptFinishedStatus.FINISHED_FRAGMENT_RUN_SUCCESSFULLY // use finished_fragment_run_successfully here because in the other test we used the finished_successfully status
+      )
+
+      // this message was sent in the new session
+      await waitFor(() => {
+        expect(screen.getByText("Here is some new text")).toBeInTheDocument()
+      })
+
+      // this message was cleaned up because it was sent in the old session
+      await waitFor(() => {
+        expect(
+          screen.queryByText("Here is some other text")
+        ).not.toBeInTheDocument()
+      })
+    })
   })
 
   describe("Logo handling", () => {

frontend/app/src/App.tsx

@@ -1201,27 +1201,31 @@ export class App extends PureComponent<Props, State> {
         this.state.scriptFinishedHandlers.map(handler => handler())
       }, 0)
 
-      // Clear any stale elements left over from the previous run.
-      // (We don't do this if our script had a compilation error and didn't
-      // finish successfully.)
-      this.setState(
-        ({ scriptRunId, fragmentIdsThisRun }) => ({
-          // Apply any pending elements that haven't been applied.
-          elements: this.pendingElementsBuffer.clearStaleNodes(
-            scriptRunId,
-            fragmentIdsThisRun
-          ),
-        }),
-        () => {
-          this.pendingElementsBuffer = this.state.elements
-        }
-      )
-
       if (
         status === ForwardMsg.ScriptFinishedStatus.FINISHED_SUCCESSFULLY ||
         status ===
           ForwardMsg.ScriptFinishedStatus.FINISHED_FRAGMENT_RUN_SUCCESSFULLY
       ) {
+        // Clear any stale elements left over from the previous run.
+        // We only do that for completed runs, not for runs that were finished early
+        // due to reruns; this is to avoid flickering of elements where they disappear for
+        // a moment and then are readded by a new session. After the new session finished,
+        // leftover elements will be cleared after finished successfully.
+        // We also don't do this if our script had a compilation error and didn't
+        // finish successfully.
+        this.setState(
+          ({ scriptRunId, fragmentIdsThisRun }) => ({
+            // Apply any pending elements that haven't been applied.
+            elements: this.pendingElementsBuffer.clearStaleNodes(
+              scriptRunId,
+              fragmentIdsThisRun
+            ),
+          }),
+          () => {
+            this.pendingElementsBuffer = this.state.elements
+          }
+        )
+
         // Tell the WidgetManager which widgets still exist. It will remove
         // widget state for widgets that have been removed.
         const activeWidgetIds = new Set(