Fix dialogs not re-opening in some cases after dismiss (streamlit#9333)

## Describe your changes

Closes streamlit#9323

Add a deltaMessageReceivedAt property to our nodes to manually trigger a re-render of the Dialog component.


## GitHub Issue Link (if applicable)

## Testing Plan

- Explanation of why no additional tests are needed
- Unit Tests (JS and/or Python)
- add a JS test to make sure that the timestamp is set on the
`BlockNode`
- E2E Tests
- the existing dialog tests for reopening after dismissing should catch
any potential regression
- added a new test where a dialog is dismissed and then a different
fragment is interacted with to ensure that the dialog stays closed. This
will make sure our `deltaMessageId` stays stable

---

**Contribution License Agreement**

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

e2e_playwright/st_dialog.py

@@ -150,3 +150,12 @@ def dialog_with_deprecation_warning():
 
 if st.button("Open Dialog with deprecation warning"):
     dialog_with_deprecation_warning()
+
+
+@st.fragment()
+def fragment():
+    if st.button("Fragment Button"):
+        st.write("Fragment Button clicked")
+
+
+fragment()

e2e_playwright/st_dialog_test.py

@@ -21,6 +21,8 @@
 from e2e_playwright.shared.app_utils import (
     COMMAND_KEY,
     check_top_level_class,
+    click_button,
+    expect_markdown,
     get_button,
     get_markdown,
 )
@@ -116,38 +118,23 @@ def test_dialog_dismisses_properly(app: Page):
     expect(main_dialog).to_have_count(0)
 
 
-# on webkit this test was flaky and manually reproducing the flaky error did not work,
-# so we skip it for now
-@pytest.mark.skip_browser("webkit")
 def test_dialog_reopens_properly_after_dismiss(app: Page):
     """Test that dialog reopens after dismiss."""
 
     # open and close the dialog multiple times
-    for _ in range(0, 3):
+    for _ in range(0, 10):
         open_dialog_without_images(app)
-        wait_for_app_run(app, wait_delay=250)
+        wait_for_app_run(app)
 
         main_dialog = app.get_by_test_id(modal_test_id)
-
-        # sometimes the dialog does not seem to open in the test, so retry opening it by
-        # clicking on it. if it does not open after the second attempt, fail the test.
-        if main_dialog.count() == 0:
-            app.wait_for_timeout(100)
-            open_dialog_without_images(app)
-            wait_for_app_run(app)
-
         expect(main_dialog).to_have_count(1)
-        app.wait_for_timeout(1000)
 
         click_to_dismiss(app)
         expect(main_dialog).not_to_be_attached()
 
         main_dialog = app.get_by_test_id(modal_test_id)
         expect(main_dialog).to_have_count(0)
 
-        # don't click indefinitely fast to give the dialog time to set the state
-        app.wait_for_timeout(500)
-
 
 def test_dialog_reopens_properly_after_close(app: Page):
     """Test that dialog reopens properly after closing by action button click."""
@@ -167,6 +154,42 @@ def test_dialog_reopens_properly_after_close(app: Page):
         expect(main_dialog).to_have_count(0)
 
 
+def test_dialog_stays_dismissed_when_interacting_with_different_fragment(app: Page):
+    """Dismissing a dialog is a UI-only interaction as of today (the Python backend does
+    not know about this). We use a deltaMsgReceivedAt to differentiate React renders
+    for dialogs triggered via a new backend message which changes the id vs. other
+    interactions. This test ensures that the dialog stays dismissed when interacting
+    with a different fragment.
+    """
+
+    open_dialog_without_images(app)
+    wait_for_app_run(app)
+
+    main_dialog = app.get_by_test_id(modal_test_id)
+    expect(main_dialog).to_have_count(1)
+
+    click_to_dismiss(app)
+    expect(main_dialog).not_to_be_attached()
+
+    main_dialog = app.get_by_test_id(modal_test_id)
+    expect(main_dialog).to_have_count(0)
+
+    # interact with unrelated fragment
+    click_button(app, "Fragment Button")
+    expect_markdown(app, "Fragment Button clicked")
+
+    # dialog is still closed and did not reopen
+    main_dialog = app.get_by_test_id(modal_test_id)
+    expect(main_dialog).to_have_count(0)
+
+    # reopen dialog
+    open_dialog_without_images(app)
+    wait_for_app_run(app)
+
+    main_dialog = app.get_by_test_id(modal_test_id)
+    expect(main_dialog).to_have_count(1)
+
+
 def test_dialog_is_scrollable(app: Page):
     """Test that the dialog is scrollable"""
     open_dialog_with_images(app)

frontend/lib/src/AppNode.test.ts

@@ -965,6 +965,7 @@ describe("AppRoot.applyDelta", () => {
     // Check that our new scriptRunId has been set only on the touched nodes
     expect(newRoot.main.scriptRunId).toBe("new_session_id")
     expect(newRoot.main.fragmentId).toBe(undefined)
+    expect(newRoot.main.deltaMsgReceivedAt).toBe(undefined)
     expect(newRoot.main.getIn([0])?.scriptRunId).toBe(NO_SCRIPT_RUN_ID)
     expect(newRoot.main.getIn([1])?.scriptRunId).toBe("new_session_id")
     expect(newRoot.main.getIn([1, 0])?.scriptRunId).toBe(NO_SCRIPT_RUN_ID)
@@ -987,6 +988,7 @@ describe("AppRoot.applyDelta", () => {
     // Check that our new scriptRunId has been set only on the touched nodes
     expect(newRoot.main.scriptRunId).toBe("new_session_id")
     expect(newRoot.main.fragmentId).toBe(undefined)
+    expect(newRoot.main.deltaMsgReceivedAt).toBe(undefined)
     expect(newRoot.main.getIn([0])?.scriptRunId).toBe(NO_SCRIPT_RUN_ID)
     expect(newRoot.main.getIn([1])?.scriptRunId).toBe("new_session_id")
     expect(newRoot.main.getIn([1, 0])?.scriptRunId).toBe(NO_SCRIPT_RUN_ID)
@@ -1148,6 +1150,22 @@ describe("AppRoot.applyDelta", () => {
     const newNode = newRoot.main.getIn([1, 1]) as BlockNode
     expect(newNode.fragmentId).toBe("myFragmentId")
   })
+
+  it("timestamp is set on BlockNode as message id", () => {
+    const timestamp = new Date(Date.UTC(2017, 1, 14)).valueOf()
+    Date.now = jest.fn(() => timestamp)
+    const delta = makeProto(DeltaProto, {
+      addBlock: {},
+    })
+    const newRoot = ROOT.applyDelta(
+      "new_session_id",
+      delta,
+      forwardMsgMetadata([0, 1, 1])
+    )
+
+    const newNode = newRoot.main.getIn([1, 1]) as BlockNode
+    expect(newNode.deltaMsgReceivedAt).toBe(timestamp)
+  })
 })
 
 describe("AppRoot.clearStaleNodes", () => {

frontend/lib/src/AppNode.ts

@@ -116,6 +116,12 @@ export interface AppNode {
    */
   readonly activeScriptHash?: string
 
+  // A timestamp indicating based on which delta message the node was created.
+  // If the node was created without a delta message, this field is undefined.
+  // This helps us to update React components based on a new backend message even though other
+  // props have not changed; this can happen for UI-only interactions such as dimissing a dialog.
+  readonly deltaMsgReceivedAt?: number
+
   /**
    * Return the AppNode for the given index path, or undefined if the path
    * is invalid.
@@ -392,6 +398,8 @@ export class BlockNode implements AppNode {
 
   public readonly fragmentId?: string
 
+  public readonly deltaMsgReceivedAt?: number
+
   // The hash of the script that created this block.
   public readonly activeScriptHash: string
 
@@ -400,13 +408,15 @@ export class BlockNode implements AppNode {
     children?: AppNode[],
     deltaBlock?: BlockProto,
     scriptRunId?: string,
-    fragmentId?: string
+    fragmentId?: string,
+    deltaMsgReceivedAt?: number
   ) {
     this.activeScriptHash = activeScriptHash
     this.children = children ?? []
     this.deltaBlock = deltaBlock ?? new BlockProto({})
     this.scriptRunId = scriptRunId ?? NO_SCRIPT_RUN_ID
     this.fragmentId = fragmentId
+    this.deltaMsgReceivedAt = deltaMsgReceivedAt
   }
 
   /** True if this Block has no children. */
@@ -461,7 +471,8 @@ export class BlockNode implements AppNode {
       newChildren,
       this.deltaBlock,
       scriptRunId,
-      this.fragmentId
+      this.fragmentId,
+      this.deltaMsgReceivedAt
     )
   }
 
@@ -480,7 +491,8 @@ export class BlockNode implements AppNode {
       newChildren,
       this.deltaBlock,
       this.scriptRunId,
-      this.fragmentId
+      this.fragmentId,
+      this.deltaMsgReceivedAt
     )
   }
 
@@ -531,7 +543,8 @@ export class BlockNode implements AppNode {
       newChildren,
       this.deltaBlock,
       currentScriptRunId,
-      this.fragmentId
+      this.fragmentId,
+      this.deltaMsgReceivedAt
     )
   }
 
@@ -707,7 +720,6 @@ export class AppRoot {
     // The full path to the AppNode within the element tree.
     // Used to find and update the element node specified by this Delta.
     const { deltaPath, activeScriptHash } = metadata
-
     switch (delta.type) {
       case "newElement": {
         const element = delta.newElement as Element
@@ -722,12 +734,14 @@ export class AppRoot {
       }
 
       case "addBlock": {
+        const deltaMsgReceivedAt = Date.now()
         return this.addBlock(
           deltaPath,
           delta.addBlock as BlockProto,
           scriptRunId,
           activeScriptHash,
-          delta.fragmentId
+          delta.fragmentId,
+          deltaMsgReceivedAt
         )
       }
 
@@ -858,7 +872,8 @@ export class AppRoot {
     block: BlockProto,
     scriptRunId: string,
     activeScriptHash: string,
-    fragmentId?: string
+    fragmentId?: string,
+    deltaMsgReceivedAt?: number
   ): AppRoot {
     const existingNode = this.root.getIn(deltaPath)
 
@@ -879,7 +894,8 @@ export class AppRoot {
       children,
       block,
       scriptRunId,
-      fragmentId
+      fragmentId,
+      deltaMsgReceivedAt
     )
     return new AppRoot(
       this.mainScriptHash,

frontend/lib/src/components/core/Block/Block.tsx

@@ -102,7 +102,10 @@ const BlockNodeRenderer = (props: BlockPropsWithWidth): ReactElement => {
 
   if (node.deltaBlock.dialog) {
     return (
-      <Dialog element={node.deltaBlock.dialog as BlockProto.Dialog}>
+      <Dialog
+        element={node.deltaBlock.dialog as BlockProto.Dialog}
+        deltaMsgReceivedAt={node.deltaMsgReceivedAt}
+      >
         {child}
       </Dialog>
     )

frontend/lib/src/components/elements/Dialog/Dialog.tsx

@@ -32,6 +32,7 @@ import { StyledDialogContent } from "./styled-components"
 
 export interface Props {
   element: BlockProto.Dialog
+  deltaMsgReceivedAt?: number
 }
 
 function parseWidthConfig(
@@ -52,6 +53,7 @@ function parseWidthConfig(
 
 const Dialog: React.FC<React.PropsWithChildren<Props>> = ({
   element,
+  deltaMsgReceivedAt,
   children,
 }): ReactElement => {
   const { title, dismissible, width, isOpen: initialIsOpen } = element
@@ -62,7 +64,11 @@ const Dialog: React.FC<React.PropsWithChildren<Props>> = ({
     if (notNullOrUndefined(initialIsOpen)) {
       setIsOpen(initialIsOpen)
     }
-  }, [initialIsOpen])
+
+    // when the deltaMsgReceivedAt changes, we might want to open the dialog again.
+    // since dismissing is a UI-only action, the initialIsOpen prop might not have
+    // changed which would lead to the dialog not opening again.
+  }, [initialIsOpen, deltaMsgReceivedAt])
 
   const theme = useTheme()
   const size: string = useMemo(

lib/streamlit/elements/lib/dialog.py

@@ -14,7 +14,6 @@
 
 from __future__ import annotations
 
-import time
 from typing import TYPE_CHECKING, Literal, cast
 
 from typing_extensions import Self, TypeAlias
@@ -87,18 +86,16 @@ def _create(
         block_proto.dialog.dismissible = dismissible
         block_proto.dialog.width = _process_dialog_width_input(width)
 
-        # We store the delta path here, because in _update we enqueue a new proto message to update the
-        # open status. Without this, the dialog content is gone when the _update message is sent
+        # We store the delta path here, because in _update we enqueue a new proto
+        # message to update the open status. Without this, the dialog content is gone
+        # when the _update message is sent
         delta_path: list[int] = (
             parent._active_dg._cursor.delta_path if parent._active_dg._cursor else []
         )
         dialog = cast(Dialog, parent._block(block_proto=block_proto, dg_type=Dialog))
 
         dialog._delta_path = delta_path
         dialog._current_proto = block_proto
-        # We add a sleep here to give the web app time to react to the update. Otherwise,
-        #  we might run into issues where the dialog cannot be opened again after closing
-        time.sleep(0.05)
         return dialog
 
     def __init__(
@@ -124,12 +121,8 @@ def _update(self, should_open: bool):
         msg.metadata.delta_path[:] = self._delta_path
         msg.delta.add_block.CopyFrom(self._current_proto)
         msg.delta.add_block.dialog.is_open = should_open
-
         self._current_proto = msg.delta.add_block
 
-        # We add a sleep here to give the web app time to react to the update. Otherwise,
-        #  we might run into issues where the dialog cannot be opened again after closing
-        time.sleep(0.05)
         enqueue_message(msg)
 
     def open(self) -> None: