streamlit-nightly 1.44.1.dev20250327__py3-none-any.whl → 1.44.2.dev20250403__py3-none-any.whl

streamlit/commands/execution_control.py

@@ -140,6 +140,7 @@ def rerun(  # type: ignore[misc]
     if ctx and ctx.script_requests:
         query_string = ctx.query_string
         page_script_hash = ctx.page_script_hash
+        cached_message_hashes = ctx.cached_message_hashes
 
         ctx.script_requests.request_rerun(
             RerunData(
@@ -147,6 +148,7 @@ def rerun(  # type: ignore[misc]
                 page_script_hash=page_script_hash,
                 fragment_id_queue=_new_fragment_id_queue(ctx, scope),
                 is_fragment_scoped_rerun=scope == "fragment",
+                cached_message_hashes=cached_message_hashes,
             )
         )
         # Force a yield point so the runner can do the rerun
@@ -232,6 +234,7 @@ def switch_page(page: str | Path | StreamlitPage) -> NoReturn:  # type: ignore[m
         RerunData(
             query_string=ctx.query_string,
             page_script_hash=page_script_hash,
+            cached_message_hashes=ctx.cached_message_hashes,
         )
     )
     # Force a yield point so the runner can do the rerun

streamlit/commands/page_config.py

@@ -102,7 +102,7 @@ def _get_favicon_string(page_icon: PageIcon) -> str:
 
     # If page_icon is an emoji, return it as is.
     if isinstance(page_icon, str) and is_emoji(page_icon):
-        return page_icon
+        return f"emoji:{page_icon}"
 
     if isinstance(page_icon, str) and page_icon.startswith(":material"):
         return validate_material_icon(page_icon)

streamlit/config.py

@@ -444,32 +444,6 @@ _create_option(
     type_=int,
 )
 
-_create_option(
-    "global.storeCachedForwardMessagesInMemory",
-    description="""
-        If True, store cached ForwardMsgs in backend memory. This is an
-        internal flag to validate a potential removal of the in-memory
-        forward message cache.
-    """,
-    visibility="hidden",
-    default_val=True,
-    type_=bool,
-)
-
-_create_option(
-    "global.includeFragmentRunsInForwardMessageCacheCount",
-    description="""
-        If True, the server will include fragment runs in the count for the
-        forward message cache. The implication is that apps with fragments may
-        see messages being removed from the cache faster. This aligns the server
-        count with the frontend count. This is a temporary fix while we assess the
-        design of the cache.
-    """,
-    visibility="hidden",
-    default_val=False,
-    type_=bool,
-)
-
 
 # Config Section: Logger #
 _create_section("logger", "Settings to customize Streamlit log messages.")

streamlit/elements/write.py

@@ -407,6 +407,14 @@ class WriteMixin:
                 kwargs,
             )
 
+        if len(args) == 1 and isinstance(args[0], str):
+            # Optimization: If there is only one arg, and it's a string,
+            # we can just call markdown directly and skip the buffer logic.
+            # This also prevents unnecessary usage of `st.empty()`.
+            # This covers > 80% of all `st.write` uses.
+            self.dg.markdown(args[0], unsafe_allow_html=unsafe_allow_html)
+            return
+
         string_buffer: list[str] = []
 
         # This bans some valid cases like: e = st.empty(); e.write("a", "b").

streamlit/proto/ClientState_pb2.py

@@ -15,7 +15,7 @@ _sym_db = _symbol_database.Default()
 from streamlit.proto import WidgetStates_pb2 as streamlit_dot_proto_dot_WidgetStates__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n!streamlit/proto/ClientState.proto\x1a\"streamlit/proto/WidgetStates.proto\"\x9d\x01\n\x0b\x43ontextInfo\x12\x15\n\x08timezone\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0ftimezone_offset\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x13\n\x06locale\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x10\n\x03url\x18\x04 \x01(\tH\x03\x88\x01\x01\x42\x0b\n\t_timezoneB\x12\n\x10_timezone_offsetB\t\n\x07_localeB\x06\n\x04_url\"\xc6\x01\n\x0b\x43lientState\x12\x14\n\x0cquery_string\x18\x01 \x01(\t\x12$\n\rwidget_states\x18\x02 \x01(\x0b\x32\r.WidgetStates\x12\x18\n\x10page_script_hash\x18\x03 \x01(\t\x12\x11\n\tpage_name\x18\x04 \x01(\t\x12\x13\n\x0b\x66ragment_id\x18\x05 \x01(\t\x12\x15\n\ris_auto_rerun\x18\x06 \x01(\x08\x12\"\n\x0c\x63ontext_info\x18\x08 \x01(\x0b\x32\x0c.ContextInfoB0\n\x1c\x63om.snowflake.apps.streamlitB\x10\x43lientStateProtob\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n!streamlit/proto/ClientState.proto\x1a\"streamlit/proto/WidgetStates.proto\"\x9d\x01\n\x0b\x43ontextInfo\x12\x15\n\x08timezone\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0ftimezone_offset\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x13\n\x06locale\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x10\n\x03url\x18\x04 \x01(\tH\x03\x88\x01\x01\x42\x0b\n\t_timezoneB\x12\n\x10_timezone_offsetB\t\n\x07_localeB\x06\n\x04_url\"\xe5\x01\n\x0b\x43lientState\x12\x14\n\x0cquery_string\x18\x01 \x01(\t\x12$\n\rwidget_states\x18\x02 \x01(\x0b\x32\r.WidgetStates\x12\x18\n\x10page_script_hash\x18\x03 \x01(\t\x12\x11\n\tpage_name\x18\x04 \x01(\t\x12\x13\n\x0b\x66ragment_id\x18\x05 \x01(\t\x12\x15\n\ris_auto_rerun\x18\x06 \x01(\x08\x12\x1d\n\x15\x63\x61\x63hed_message_hashes\x18\x07 \x03(\t\x12\"\n\x0c\x63ontext_info\x18\x08 \x01(\x0b\x32\x0c.ContextInfoB0\n\x1c\x63om.snowflake.apps.streamlitB\x10\x43lientStateProtob\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -26,5 +26,5 @@ if not _descriptor._USE_C_DESCRIPTORS:
   _globals['_CONTEXTINFO']._serialized_start=74
   _globals['_CONTEXTINFO']._serialized_end=231
   _globals['_CLIENTSTATE']._serialized_start=234
-  _globals['_CLIENTSTATE']._serialized_end=432
+  _globals['_CLIENTSTATE']._serialized_end=463
 # @@protoc_insertion_point(module_scope)

streamlit/proto/ClientState_pb2.pyi

@@ -18,7 +18,9 @@ limitations under the License.
 """
 
 import builtins
+import collections.abc
 import google.protobuf.descriptor
+import google.protobuf.internal.containers
 import google.protobuf.message
 import streamlit.proto.WidgetStates_pb2
 import typing
@@ -68,6 +70,7 @@ class ClientState(google.protobuf.message.Message):
     PAGE_NAME_FIELD_NUMBER: builtins.int
     FRAGMENT_ID_FIELD_NUMBER: builtins.int
     IS_AUTO_RERUN_FIELD_NUMBER: builtins.int
+    CACHED_MESSAGE_HASHES_FIELD_NUMBER: builtins.int
     CONTEXT_INFO_FIELD_NUMBER: builtins.int
     query_string: builtins.str
     page_script_hash: builtins.str
@@ -76,6 +79,12 @@ class ClientState(google.protobuf.message.Message):
     is_auto_rerun: builtins.bool
     @property
     def widget_states(self) -> streamlit.proto.WidgetStates_pb2.WidgetStates: ...
+    @property
+    def cached_message_hashes(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]:
+        """List of hashes of messages that are currently cached in
+        the frontend forward message cache.
+        """
+
     @property
     def context_info(self) -> global___ContextInfo: ...
     def __init__(
@@ -87,9 +96,10 @@ class ClientState(google.protobuf.message.Message):
         page_name: builtins.str = ...,
         fragment_id: builtins.str = ...,
         is_auto_rerun: builtins.bool = ...,
+        cached_message_hashes: collections.abc.Iterable[builtins.str] | None = ...,
         context_info: global___ContextInfo | None = ...,
     ) -> None: ...
     def HasField(self, field_name: typing.Literal["context_info", b"context_info", "widget_states", b"widget_states"]) -> builtins.bool: ...
-    def ClearField(self, field_name: typing.Literal["context_info", b"context_info", "fragment_id", b"fragment_id", "is_auto_rerun", b"is_auto_rerun", "page_name", b"page_name", "page_script_hash", b"page_script_hash", "query_string", b"query_string", "widget_states", b"widget_states"]) -> None: ...
+    def ClearField(self, field_name: typing.Literal["cached_message_hashes", b"cached_message_hashes", "context_info", b"context_info", "fragment_id", b"fragment_id", "is_auto_rerun", b"is_auto_rerun", "page_name", b"page_name", "page_script_hash", b"page_script_hash", "query_string", b"query_string", "widget_states", b"widget_states"]) -> None: ...
 
 global___ClientState = ClientState

streamlit/proto/ForwardMsg_pb2.pyi

@@ -107,10 +107,9 @@ class ForwardMsg(google.protobuf.message.Message):
     """A hash that uniquely identifies this ForwardMsg, for caching."""
     script_finished: global___ForwardMsg.ScriptFinishedStatus.ValueType
     ref_hash: builtins.str
-    """A reference to a ForwardMsg that has already been delivered.
-    The client should substitute the message with the given hash
-    for this one. If the client does not have the referenced message
-    in its cache, it can retrieve it from the server.
+    """A reference to a ForwardMsg that has already been delivered
+    and cached in the frontend. The client should substitute the message
+    with the given hash for this one.
     """
     debug_last_backmsg_id: builtins.str
     """The ID of the last BackMsg that we received before sending this
@@ -214,9 +213,7 @@ class ForwardMsgMetadata(google.protobuf.message.Message):
     ELEMENT_DIMENSION_SPEC_FIELD_NUMBER: builtins.int
     ACTIVE_SCRIPT_HASH_FIELD_NUMBER: builtins.int
     cacheable: builtins.bool
-    """If this is set, the server will have cached this message,
-    and a client that receives it should do the same.
-    """
+    """Marks a message as cacheable for the frontend."""
     active_script_hash: builtins.str
     """active_script_hash the forward message is associated from.
     For multipage apps v1, this will always be the page file running

streamlit/runtime/app_session.py

@@ -383,12 +383,13 @@ class AppSession:
                 self._client_state.context_info.CopyFrom(client_state.context_info)
 
             rerun_data = RerunData(
-                client_state.query_string,
-                client_state.widget_states,
-                client_state.page_script_hash,
-                client_state.page_name,
+                query_string=client_state.query_string,
+                widget_states=client_state.widget_states,
+                page_script_hash=client_state.page_script_hash,
+                page_name=client_state.page_name,
                 fragment_id=fragment_id if fragment_id else None,
                 is_auto_rerun=client_state.is_auto_rerun,
+                cached_message_hashes=set(client_state.cached_message_hashes),
                 context_info=client_state.context_info,
             )
         else:

streamlit/runtime/context.py

@@ -307,3 +307,17 @@ class ContextProxy:
         if ctx is None or ctx.context_info is None:
             return None
         return ctx.context_info.url
+
+    @property
+    @gather_metrics("context.ip_address")
+    def ip_address(self) -> str | None:
+        """The read-only IP address of the user's connection.
+        This should not be used for security measures as it can be easily spoofed.
+        """
+        session_client_request = _get_request()
+        if session_client_request is not None:
+            remote_ip = session_client_request.remote_ip
+            if remote_ip == "::1" or remote_ip == "127.0.0.1":
+                return None
+            return remote_ip
+        return None

streamlit/runtime/forward_msg_cache.py

@@ -14,23 +14,16 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Final
-from weakref import WeakKeyDictionary
+from typing import Final
 
 from streamlit import config, util
 from streamlit.logger import get_logger
 from streamlit.proto.ForwardMsg_pb2 import ForwardMsg
-from streamlit.runtime.stats import CacheStat, CacheStatsProvider, group_stats
-
-if TYPE_CHECKING:
-    from collections.abc import MutableMapping
-
-    from streamlit.runtime.app_session import AppSession
 
 _LOGGER: Final = get_logger(__name__)
 
 
-def populate_hash_if_needed(msg: ForwardMsg) -> str:
+def populate_hash_if_needed(msg: ForwardMsg) -> None:
     """Computes and assigns the unique hash for a ForwardMsg.
 
     If the ForwardMsg already has a hash, this is a no-op.
@@ -39,27 +32,38 @@ def populate_hash_if_needed(msg: ForwardMsg) -> str:
     ----------
     msg : ForwardMsg
 
-    Returns
-    -------
-    string
-        The message's hash, returned here for convenience. (The hash
-        will also be assigned to the ForwardMsg; callers do not need
-        to do this.)
-
     """
-    if msg.hash == "":
+    if msg.hash == "" and msg.WhichOneof("type") not in {"ref_hash", "initialize"}:
         # Move the message's metadata aside. It's not part of the
         # hash calculation.
         metadata = msg.metadata
         msg.ClearField("metadata")
 
+        # Serialize the message to bytes using the deterministic serializer to
+        # ensure consistent hashing.
+        serialized_msg = msg.SerializeToString(deterministic=True)
+
+        # TODO(lukasmasuch): Evaluate more optimized hashing for larger messages:
+        # - Add the type element type and number of bytes to the hash.
+        # - Only hash the first N bytes of the message.
+
         # MD5 is good enough for what we need, which is uniqueness.
-        msg.hash = util.calc_md5(msg.SerializeToString())
+        msg.hash = util.calc_md5(serialized_msg)
 
         # Restore metadata.
         msg.metadata.CopyFrom(metadata)
 
-    return msg.hash
+        # Set cacheable flag if above the min cached size and if its a `new_element`
+        # delta. We only cache new_element and add_block deltas since container's
+        # are not expected to be larger than a few KB and have other side-effects
+        # to consider if cached. But `add_block` deltas should still get a hash.
+        # In case we ever allow other delta types to be cached, we should
+        # also need to adapt the composable logic in forward_msg_queue.
+        msg.metadata.cacheable = (
+            len(serialized_msg) >= int(config.get_option("global.minCachedMessageSize"))
+            and msg.WhichOneof("type") == "delta"
+            and msg.delta.WhichOneof("type") == "new_element"
+        )
 
 
 def create_reference_msg(msg: ForwardMsg) -> ForwardMsg:
@@ -80,217 +84,18 @@ def create_reference_msg(msg: ForwardMsg) -> ForwardMsg:
         ref_hash field.
 
     """
+    if not msg.hash:
+        _LOGGER.warning(
+            "Failed to create a reference message for a ForwardMsg since the "
+            "message does not have a hash. This is not expected to happen, "
+            "please report this as a bug. Falling back to the original message."
+        )
+        # Fallback to the original message if the hash is not set.
+        # This is not expected to happen.
+        return msg
+
     ref_msg = ForwardMsg()
-    ref_msg.ref_hash = populate_hash_if_needed(msg)
+    ref_msg.ref_hash = msg.hash
     ref_msg.metadata.CopyFrom(msg.metadata)
+    ref_msg.metadata.cacheable = False
     return ref_msg
-
-
-class ForwardMsgCache(CacheStatsProvider):
-    """A cache of ForwardMsgs.
-
-    Large ForwardMsgs (e.g. those containing big DataFrame payloads) are
-    stored in this cache. The server can choose to send a ForwardMsg's hash,
-    rather than the message itself, to a client. Clients can then
-    request messages from this cache via another endpoint.
-
-    This cache is *not* thread safe. It's intended to only be accessed by
-    the server thread.
-
-    """
-
-    class Entry:
-        """Cache entry.
-
-        Stores the cached message, and the set of AppSessions
-        that we've sent the cached message to.
-
-        """
-
-        def __init__(self, msg: ForwardMsg | None):
-            self.msg = msg
-            self._session_script_run_counts: MutableMapping[AppSession, int] = (
-                WeakKeyDictionary()
-            )
-
-        def __repr__(self) -> str:
-            return util.repr_(self)
-
-        def add_session_ref(self, session: AppSession, script_run_count: int) -> None:
-            """Adds a reference to a AppSession that has referenced
-            this Entry's message.
-
-            Parameters
-            ----------
-            session : AppSession
-            script_run_count : int
-                The session's run count at the time of the call
-
-            """
-            prev_run_count = self._session_script_run_counts.get(session, 0)
-            if script_run_count < prev_run_count:
-                _LOGGER.error(
-                    "New script_run_count (%s) is < prev_run_count (%s). "
-                    "This should never happen!",
-                    script_run_count,
-                    prev_run_count,
-                )
-                script_run_count = prev_run_count
-            self._session_script_run_counts[session] = script_run_count
-
-        def has_session_ref(self, session: AppSession) -> bool:
-            return session in self._session_script_run_counts
-
-        def get_session_ref_age(
-            self, session: AppSession, script_run_count: int
-        ) -> int:
-            """The age of the given session's reference to the Entry,
-            given a new script_run_count.
-
-            """
-            return script_run_count - self._session_script_run_counts[session]
-
-        def remove_session_ref(self, session: AppSession) -> None:
-            del self._session_script_run_counts[session]
-
-        def has_refs(self) -> bool:
-            """True if this Entry has references from any AppSession.
-
-            If not, it can be removed from the cache.
-            """
-            return len(self._session_script_run_counts) > 0
-
-    def __init__(self):
-        self._entries: dict[str, ForwardMsgCache.Entry] = {}
-
-    def __repr__(self) -> str:
-        return util.repr_(self)
-
-    def add_message(
-        self, msg: ForwardMsg, session: AppSession, script_run_count: int
-    ) -> None:
-        """Add a ForwardMsg to the cache.
-
-        The cache will also record a reference to the given AppSession,
-        so that it can track which sessions have already received
-        each given ForwardMsg.
-
-        Parameters
-        ----------
-        msg : ForwardMsg
-        session : AppSession
-        script_run_count : int
-            The number of times the session's script has run
-
-        """
-        populate_hash_if_needed(msg)
-        entry = self._entries.get(msg.hash, None)
-        if entry is None:
-            if config.get_option("global.storeCachedForwardMessagesInMemory"):
-                entry = ForwardMsgCache.Entry(msg)
-            else:
-                entry = ForwardMsgCache.Entry(None)
-            self._entries[msg.hash] = entry
-        entry.add_session_ref(session, script_run_count)
-
-    def get_message(self, hash: str) -> ForwardMsg | None:
-        """Return the message with the given ID if it exists in the cache.
-
-        Parameters
-        ----------
-        hash : str
-            The id of the message to retrieve.
-
-        Returns
-        -------
-        ForwardMsg | None
-
-        """
-        entry = self._entries.get(hash, None)
-        return entry.msg if entry else None
-
-    def has_message_reference(
-        self, msg: ForwardMsg, session: AppSession, script_run_count: int
-    ) -> bool:
-        """Return True if a session has a reference to a message."""
-        populate_hash_if_needed(msg)
-
-        entry = self._entries.get(msg.hash, None)
-        if entry is None or not entry.has_session_ref(session):
-            return False
-
-        # Ensure we're not expired
-        age = entry.get_session_ref_age(session, script_run_count)
-        return age <= int(config.get_option("global.maxCachedMessageAge"))
-
-    def remove_refs_for_session(self, session: AppSession) -> None:
-        """Remove refs for all entries for the given session.
-
-        This should be called when an AppSession is disconnected or closed.
-
-        Parameters
-        ----------
-        session : AppSession
-        """
-
-        # Operate on a copy of our entries dict.
-        # We may be deleting from it.
-        for msg_hash, entry in self._entries.copy().items():
-            if entry.has_session_ref(session):
-                entry.remove_session_ref(session)
-
-            if not entry.has_refs():
-                # The entry has no more references. Remove it from
-                # the cache completely.
-                del self._entries[msg_hash]
-
-    def remove_expired_entries_for_session(
-        self, session: AppSession, script_run_count: int
-    ) -> None:
-        """Remove any cached messages that have expired from the given session.
-
-        This should be called each time a AppSession finishes executing.
-
-        Parameters
-        ----------
-        session : AppSession
-        script_run_count : int
-            The number of times the session's script has run
-
-        """
-        max_age = config.get_option("global.maxCachedMessageAge")
-
-        # Operate on a copy of our entries dict.
-        # We may be deleting from it.
-        for msg_hash, entry in self._entries.copy().items():
-            if not entry.has_session_ref(session):
-                continue
-
-            age = entry.get_session_ref_age(session, script_run_count)
-            if age > max_age:
-                _LOGGER.debug(
-                    "Removing expired entry [session=%s, hash=%s, age=%s]",
-                    id(session),
-                    msg_hash,
-                    age,
-                )
-                entry.remove_session_ref(session)
-                if not entry.has_refs():
-                    # The entry has no more references. Remove it from
-                    # the cache completely.
-                    del self._entries[msg_hash]
-
-    def clear(self) -> None:
-        """Remove all entries from the cache."""
-        self._entries.clear()
-
-    def get_stats(self) -> list[CacheStat]:
-        stats: list[CacheStat] = [
-            CacheStat(
-                category_name="ForwardMessageCache",
-                cache_name="",
-                byte_length=entry.msg.ByteSize() if entry.msg is not None else 0,
-            )
-            for _, entry in self._entries.items()
-        ]
-        return group_stats(stats)

streamlit/runtime/forward_msg_queue.py

@@ -14,13 +14,10 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Callable
+from typing import Any, Callable
 
 from streamlit.proto.ForwardMsg_pb2 import ForwardMsg
 
-if TYPE_CHECKING:
-    from streamlit.proto.Delta_pb2 import Delta
-
 
 class ForwardMsgQueue:
     """Accumulates a session's outgoing ForwardMsgs.
@@ -79,16 +76,27 @@ class ForwardMsgQueue:
         # the app - we attempt to combine this new Delta into the old
         # one. This is an optimization that prevents redundant Deltas
         # from being sent to the frontend.
+        # One common case where this happens is with `st.write` since
+        # it uses a trick with `st.empty` to handle lists of args.
+        # Note: its not guaranteed that the optimization is always applied
+        # since the queue can be flushed to the browser at any time.
+        # For example:
+        # queue 1:
+        # empty [0, 0]  <- skipped
+        # markdown [0, 0]
+        # empty [1, 0]  <- send to frontend
+        #
+        # queue 2:
+        # markdown [1, 0]
+        # ...
+
         delta_key = tuple(msg.metadata.delta_path)
         if delta_key in self._delta_index_map:
             index = self._delta_index_map[delta_key]
             old_msg = self._queue[index]
-            composed_delta = _maybe_compose_deltas(old_msg.delta, msg.delta)
-            if composed_delta is not None:
-                new_msg = ForwardMsg()
-                new_msg.delta.CopyFrom(composed_delta)
-                new_msg.metadata.CopyFrom(msg.metadata)
-                self._queue[index] = new_msg
+            composed_msg = _maybe_compose_delta_msgs(old_msg, msg)
+            if composed_msg is not None:
+                self._queue[index] = composed_msg
                 return
 
         # No composition occurred. Append this message to the queue, and
@@ -163,6 +171,11 @@ class ForwardMsgQueue:
 
 def _is_composable_message(msg: ForwardMsg) -> bool:
     """True if the ForwardMsg is potentially composable with other ForwardMsgs."""
+    if msg.HasField("ref_hash"):
+        # reference messages (cached in frontend) are always composable.
+        # Only new_element deltas can be reference messages.
+        return True
+
     if not msg.HasField("delta"):
         # Non-delta messages are never composable.
         return False
@@ -174,17 +187,20 @@ def _is_composable_message(msg: ForwardMsg) -> bool:
     return delta_type != "add_rows" and delta_type != "arrow_add_rows"
 
 
-def _maybe_compose_deltas(old_delta: Delta, new_delta: Delta) -> Delta | None:
-    """Combines new_delta onto old_delta if possible.
+def _maybe_compose_delta_msgs(
+    old_msg: ForwardMsg, new_msg: ForwardMsg
+) -> ForwardMsg | None:
+    """Optimization logic that composes new_msg onto old_msg if possible.
 
-    If the combination takes place, the function returns a new Delta that
-    should replace old_delta in the queue.
+    If the combination takes place, the function returns a new ForwardMsg that
+    should replace old_msg in the queue. This basically means that the old_msg
+    is not send to the browser since its considered unnecessary.
 
-    If the new_delta is incompatible with old_delta, the function returns None.
-    In this case, the new_delta should just be appended to the queue as normal.
+    If the new_msg is incompatible with old_msg, the function returns None.
+    In this case, the new_msg should just be appended to the queue as normal.
     """
-    old_delta_type = old_delta.WhichOneof("type")
-    if old_delta_type == "add_block":
+
+    if old_msg.HasField("delta") and old_msg.delta.WhichOneof("type") == "add_block":
         # We never replace add_block deltas, because blocks can have
         # other dependent deltas later in the queue. For example:
         #
@@ -200,12 +216,14 @@ def _maybe_compose_deltas(old_delta: Delta, new_delta: Delta) -> Delta | None:
         # now just an element, and not a block.
         return None
 
-    new_delta_type = new_delta.WhichOneof("type")
-    if new_delta_type == "new_element":
-        return new_delta
+    if new_msg.HasField("ref_hash"):
+        # ref_hash messages are always composable.
+        # Only new_element deltas can be reference messages.
+        return new_msg
 
-    if new_delta_type == "add_block":
-        return new_delta
+    new_delta_type = new_msg.delta.WhichOneof("type")
+    if new_delta_type == "new_element" or new_delta_type == "add_block":
+        return new_msg
 
     return None