Repository: streamlit/streamlit

Issue:
`experimental_allow_widgets=True` doesn't work in nested `@st.experimental_singleton` 
### Checklist

- [X] I have searched the [existing issues](https://github.com/streamlit/streamlit/issues) for similar issues.
- [X] I added a very descriptive title to this issue.
- [X] I have provided sufficient information below to help reproduce this issue.

### Summary

Error happens when I try to call function decorated with `@st.experimental_singleton(experimental_allow_widgets=True)` decorator in another function decorated with `@st.experimental_singleton` (or even with `@st.experimental_singleton(experimental_allow_widgets=True)`

Traceback: 
```python
Traceback (most recent call last):
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/runtime/caching/cache_utils.py", line 339, in get_or_create_cached_value
    result = cache.read_result(value_key)
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/runtime/caching/singleton_decorator.py", line 302, in read_result
    raise CacheKeyNotFoundError()
streamlit.runtime.caching.cache_errors.CacheKeyNotFoundError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/runtime/caching/cache_utils.py", line 339, in get_or_create_cached_value
    result = cache.read_result(value_key)
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/runtime/caching/singleton_decorator.py", line 302, in read_result
    raise CacheKeyNotFoundError()
streamlit.runtime.caching.cache_errors.CacheKeyNotFoundError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/runtime/scriptrunner/script_runner.py", line 564, in _run_script
    exec(code, module.__dict__)
  File "/Users/kjavadyan/streamlit/streamlit/tmp.py", line 15, in <module>
    st.write(baz(10))
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/runtime/caching/cache_utils.py", line 387, in wrapper
    return get_or_create_cached_value()
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/runtime/caching/cache_utils.py", line 360, in get_or_create_cached_value
    return_value = func(*args, **kwargs)
  File "/Users/kjavadyan/streamlit/streamlit/tmp.py", line 12, in baz
    return foo() + y
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/runtime/caching/cache_utils.py", line 387, in wrapper
    return get_or_create_cached_value()
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/runtime/caching/cache_utils.py", line 360, in get_or_create_cached_value
    return_value = func(*args, **kwargs)
  File "/Users/kjavadyan/streamlit/streamlit/tmp.py", line 6, in foo
    x = st.number_input("AAAA", 1, 100, 12)
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/runtime/metrics_util.py", line 311, in wrapped_func
    result = non_optional_func(*args, **kwargs)
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/elements/number_input.py", line 150, in number_input
    return self._number_input(
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/elements/number_input.py", line 345, in _number_input
    self.dg._enqueue("number_input", number_input_proto)
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/delta_generator.py", line 556, in _enqueue
    caching.save_element_message(
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/runtime/caching/__init__.py", line 54, in save_element_message
    SINGLETON_MESSAGE_CALL_STACK.save_element_message(
  File "/Users/kjavadyan/streamlit/streamlit/lib/streamlit/runtime/caching/cache_utils.py", line 595, in save_element_message
    raise AttributeError
```

### Reproducible Code Example

```python
import streamlit as st


@st.experimental_singleton(experimental_allow_widgets=True)
def foo():
    x = st.number_input("AAAA", 1, 100, 12)
    return x ** 2


@st.experimental_singleton(experimental_allow_widgets=True)
def baz(y):
    return foo() + y


st.write(baz(10))
```


### Steps To Reproduce

_No response_

### Expected Behavior

_No response_

### Current Behavior

_No response_

### Is this a regression?

- [ ] Yes, this used to work in a previous version.

### Debug info

- Streamlit version:
- Python version:
- Operating System:
- Browser:
- Virtual environment:


### Additional Information

_No response_

### Are you willing to submit a PR?

- [ ] Yes, I am willing to submit a PR!


Discussion:

Resolving patch:
diff --git a/lib/streamlit/runtime/caching/cache_utils.py b/lib/streamlit/runtime/caching/cache_utils.py
--- a/lib/streamlit/runtime/caching/cache_utils.py
+++ b/lib/streamlit/runtime/caching/cache_utils.py
@@ -597,9 +597,9 @@ def save_element_message(
         """
         if not runtime.exists():
             return
+        if len(self._cached_message_stack) >= 1:
 
-        id_to_save = self.select_dg_to_save(invoked_dg_id, used_dg_id)
-        for msgs in self._cached_message_stack:
+            id_to_save = self.select_dg_to_save(invoked_dg_id, used_dg_id)
             if isinstance(element_proto, Widget):
                 wid = element_proto.id
                 # TODO replace `Message` with a more precise type
@@ -616,22 +616,23 @@ def save_element_message(
 
             media_data = self._media_data
 
-            if self._allow_widgets or widget_meta is None:
-                msgs.append(
-                    ElementMsgData(
-                        delta_type,
-                        element_proto,
-                        id_to_save,
-                        returned_dg_id,
-                        widget_meta,
-                        media_data,
-                    )
-                )
+            element_msg_data = ElementMsgData(
+                delta_type,
+                element_proto,
+                id_to_save,
+                returned_dg_id,
+                widget_meta,
+                media_data,
+            )
+            for msgs in self._cached_message_stack:
+                if self._allow_widgets or widget_meta is None:
+                    msgs.append(element_msg_data)
+
+        # Reset instance state, now that it has been used for the
+        # associated element.
+        self._media_data = []
+        self._registered_metadata = None
 
-            # Reset instance state, now that it has been used for the
-            # associated element.
-            self._media_data = []
-            self._registered_metadata = None
         for s in self._seen_dg_stack:
             s.add(returned_dg_id)
 


Test patch:
diff --git a/lib/tests/streamlit/runtime/caching/common_cache_test.py b/lib/tests/streamlit/runtime/caching/common_cache_test.py
--- a/lib/tests/streamlit/runtime/caching/common_cache_test.py
+++ b/lib/tests/streamlit/runtime/caching/common_cache_test.py
@@ -484,6 +484,21 @@ def img_fn_multi():
         img_fn_multi()
         img_fn_multi()
 
+    @parameterized.expand([("memo", memo), ("singleton", singleton)])
+    def test_nested_widget_replay(self, _, cache_decorator):
+        """Regression test for GH#5677"""
+
+        @cache_decorator(experimental_allow_widgets=True)
+        def foo():
+            x = st.number_input("AAAA", 1, 100, 12)
+            return x**2
+
+        @cache_decorator(experimental_allow_widgets=True)
+        def baz(y):
+            return foo() + y
+
+        st.write(baz(10))
+
     @parameterized.expand(
         [
             ("memo", memo, memo.clear),