Fix Chat Message container to align and manage the autosizer handling (…

…streamlit#7504)

* Fix Chat Message container to align and manage the autosizer handling

* Remove chat message snapshots

* Add Updated Snapshots

* Attempt to make new playwright test

* Remove async def as it's not supported

* Add new snapshots

* Minor cleanup

---------

Co-authored-by: willhuang1997 <willhuang1997@gmail.com>

e2e_playwright/st_chat_message.py

@@ -70,11 +70,44 @@
 st.chat_message("human")
 
 
-image1 = Image.new("RGB", (10, 10), "red")
+image1 = Image.new("RGB", (10, 255), "red")
 st.chat_message("user", avatar=image1).write("Red local image")
 
 image2 = Image.new("RGB", (10, 10), "blue")
 st.chat_message("assistant", avatar=image2).write("Blue local image")
 st.chat_message("assistant", avatar=image2).write(
     "Another message with the same blue avatar."
 )
+
+query = "This is a hardcoded user message"
+sources = "example sources"
+llm_response = "some response"
+
+past_messages = st.empty()
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+with past_messages.container():
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+            if message["role"] != "user":
+                with st.expander("See sources"):
+                    st.markdown(message["sources"])
+
+with st.chat_message("user"):
+    st.markdown(query)
+
+user_message = {"role": "user", "content": query, "sources": ""}
+st.session_state.messages.append(user_message)
+
+with st.chat_message("assistant"):
+    displayed_response = st.empty()
+    with displayed_response.container():
+        st.markdown(llm_response)
+        with st.expander("See sources"):
+            st.markdown(sources)
+
+assistant_message = {"role": "assistant", "content": llm_response, "sources": sources}
+st.session_state.messages.append(assistant_message)

e2e_playwright/st_chat_message_test.py

@@ -14,7 +14,7 @@
 
 from playwright.sync_api import Page, expect
 
-from e2e_playwright.conftest import ImageCompareFunction
+from e2e_playwright.conftest import ImageCompareFunction, wait_for_app_run
 
 
 def test_renders_chat_messages_correctly_1(
@@ -23,7 +23,12 @@ def test_renders_chat_messages_correctly_1(
     """Test if the chat messages render correctly"""
     # Wait a bit more to allow all images to load:
     chat_message_elements = themed_app.locator(".stChatMessage")
-    expect(chat_message_elements).to_have_count(10)
+    expect(chat_message_elements).to_have_count(12)
+
+    # rerun to populate session state chat message
+    themed_app.keyboard.press("r")
+    wait_for_app_run(themed_app, wait_delay=1000)
+    expect(chat_message_elements).to_have_count(14)
     for i, element in enumerate(chat_message_elements.all()):
         element.scroll_into_view_if_needed()
         # Wait a bit more to allow the avatar images to load:

frontend/lib/src/components/elements/ChatMessage/styled-components.ts

@@ -47,6 +47,14 @@ export const StyledMessageContent = styled.div(({ theme }) => ({
   color: theme.colors.bodyText,
   margin: "auto",
   flexGrow: 1,
+  // Ensure the size of the message has an interpreted width as
+  // the amount defined by flex layout and disregard its contents
+  // they will handle their overflow.
+  //
+  // Unintuitively, setting the min width to 0 tells the browser
+  // that it can shrink past the content's width to the desired size.
+  // https://makandracards.com/makandra/66994-css-flex-and-min-width
+  minWidth: 0,
 }))
 
 export const StyledAvatarBackground = styled.div(({ theme }) => {
@@ -60,7 +68,8 @@ export const StyledAvatarBackground = styled.div(({ theme }) => {
     color: lightTheme ? theme.colors.gray90 : theme.colors.white,
     lineHeight: "1",
     fontSize: theme.fontSizes.md,
-    padding: "1rem",
+    // Ensure the avatar always respects the width/height
+    flexShrink: 0,
     width: "2rem",
     height: "2rem",
     borderRadius: theme.radii.lg,
@@ -80,6 +89,8 @@ export const StyledAvatarIcon = styled.div<StyledAvatarIconProps>(
       display: "flex",
       width: "2rem",
       height: "2rem",
+      // Ensure the avatar always respects the width/height
+      flexShrink: 0,
       borderRadius: theme.radii.lg,
       alignItems: "center",
       justifyContent: "center",
@@ -93,6 +104,8 @@ export const StyledAvatarImage = styled.img(({ theme }) => {
   return {
     width: "2rem",
     height: "2rem",
+    // Ensure the avatar always respects the width/height
+    flexShrink: 0,
     borderRadius: theme.radii.lg,
     objectFit: "cover",
     display: "flex",