(fix): message height calculation

packages/chat/src/screen.tsx

@@ -13,16 +13,21 @@ function getDistanceFromBottom(element: HTMLDivElement) {
   return element.scrollHeight - (element.scrollTop + element.clientHeight);
 }
 
-const MESSAGE_ROW_HEIGHT = 48;
-const ROW_VERTICAL_PADDING = 8;
+const FALLBACK_MESSAGE_HEIGHT = 56;
 
 /**
  * Renders the chat screen with virtualized history and live message updates.
  * @returns Chat screen JSX.
  */
 export default function Screen() {
-  const { getMessages, liveMessages, clearLiveMessages, userId, sharedSecret } =
-    useChat();
+  const {
+    getMessages,
+    liveMessages,
+    clearLiveMessages,
+    userId,
+    sharedSecret,
+    inputBoxRef,
+  } = useChat();
 
   const scrollRef = React.useRef<HTMLDivElement | null>(null);
   const isMobile = useIsMobile();
@@ -35,6 +40,12 @@ export default function Screen() {
     totalSize: number;
     scrollTop: number;
   } | null>(null);
+  const [scrollWidth, setScrollWidth] = React.useState(0);
+  const [inputHeight, setInputHeight] = React.useState(0);
+  const [measuredHeights, setMeasuredHeights] = React.useState<
+    Record<number, number>
+  >({});
+  const measurementRefs = React.useRef(new Map<number, HTMLDivElement>());
 
   const hasValidChatUser = Number.isSafeInteger(userId) && userId > 0;
   const hasSharedSecret = sharedSecret.length > 0;
@@ -92,15 +103,159 @@ export default function Screen() {
     return messagesRef.current[index]?.send_time ?? index;
   }, []);
 
+  const setMeasurementRef = React.useCallback(
+    (sendTime: number, element: HTMLDivElement | null) => {
+      if (!element) {
+        measurementRefs.current.delete(sendTime);
+        return;
+      }
+
+      measurementRefs.current.set(sendTime, element);
+    },
+    [],
+  );
+
+  const estimateMessageSize = React.useCallback(
+    (index: number) => {
+      const sendTime = messages[index]?.send_time;
+
+      if (sendTime === undefined) {
+        return FALLBACK_MESSAGE_HEIGHT;
+      }
+
+      return measuredHeights[sendTime] ?? FALLBACK_MESSAGE_HEIGHT;
+    },
+    [measuredHeights, messages],
+  );
+
   // eslint-disable-next-line react-hooks/incompatible-library
   const virtualizer = useVirtualizer({
     count: messages.length,
     getScrollElement: () => scrollRef.current,
     getItemKey,
-    estimateSize: () => MESSAGE_ROW_HEIGHT + ROW_VERTICAL_PADDING,
+    estimateSize: estimateMessageSize,
     overscan: isMobile ? 10 : 6,
   });
 
+  React.useLayoutEffect(() => {
+    const element = scrollRef.current;
+    if (!element || typeof ResizeObserver === "undefined") {
+      return;
+    }
+
+    const updateWidth = () => {
+      setScrollWidth(element.clientWidth);
+    };
+
+    updateWidth();
+
+    const observer = new ResizeObserver(updateWidth);
+    observer.observe(element);
+
+    return () => {
+      observer.disconnect();
+    };
+  }, []);
+
+  React.useLayoutEffect(() => {
+    const frame = requestAnimationFrame(() => {
+      virtualizer.measure();
+    });
+
+    return () => {
+      cancelAnimationFrame(frame);
+    };
+  }, [inputHeight, measuredHeights, scrollWidth, virtualizer]);
+
+  React.useLayoutEffect(() => {
+    if (typeof ResizeObserver === "undefined") {
+      return;
+    }
+
+    const measureHeights = () => {
+      setMeasuredHeights((prev) => {
+        let changed = false;
+        const next: Record<number, number> = {};
+
+        for (const message of messages) {
+          const element = measurementRefs.current.get(message.send_time);
+          const measuredHeight = element?.getBoundingClientRect().height;
+          const nextHeight =
+            typeof measuredHeight === "number" &&
+            Number.isFinite(measuredHeight)
+              ? Math.ceil(measuredHeight)
+              : (prev[message.send_time] ?? FALLBACK_MESSAGE_HEIGHT);
+
+          next[message.send_time] = nextHeight;
+
+          if (prev[message.send_time] !== nextHeight) {
+            changed = true;
+          }
+        }
+
+        if (!changed && Object.keys(prev).length === messages.length) {
+          return prev;
+        }
+
+        return next;
+      });
+    };
+
+    measureHeights();
+
+    const observer = new ResizeObserver(() => {
+      measureHeights();
+    });
+
+    for (const message of messages) {
+      const element = measurementRefs.current.get(message.send_time);
+      if (element) {
+        observer.observe(element);
+      }
+    }
+
+    return () => {
+      observer.disconnect();
+    };
+  }, [messages, scrollWidth]);
+
+  React.useEffect(() => {
+    if (typeof window === "undefined") {
+      return;
+    }
+
+    const handleResize = () => {
+      requestAnimationFrame(() => {
+        virtualizer.measure();
+      });
+    };
+
+    window.addEventListener("resize", handleResize);
+    return () => {
+      window.removeEventListener("resize", handleResize);
+    };
+  }, [virtualizer]);
+
+  React.useLayoutEffect(() => {
+    const element = inputBoxRef.current;
+    if (!element || typeof ResizeObserver === "undefined") {
+      return;
+    }
+
+    const updateHeight = () => {
+      setInputHeight(element.getBoundingClientRect().height);
+    };
+
+    updateHeight();
+
+    const observer = new ResizeObserver(updateHeight);
+    observer.observe(element);
+
+    return () => {
+      observer.disconnect();
+    };
+  }, [inputBoxRef]);
+
   /**
    * Loads the next page when the scroll container reaches the top.
    * @returns Promise that resolves once pagination handling completes.
@@ -196,13 +351,15 @@ export default function Screen() {
   }
 
   return (
-    <div className="w-full h-full flex flex-col px-2">
+    <div className="relative w-full h-full overflow-hidden px-2">
       <div
         ref={scrollRef}
         id="chat_container"
-        className="flex-1 overflow-y-auto px-2.5 pb-32"
+        className="absolute inset-x-0 top-0 overflow-y-auto px-2.5"
         style={{
+          bottom: `${inputHeight}px`,
           overflowAnchor: "none",
+          paddingBottom: "8px",
         }}
         onScroll={handleContainerScroll}
       >
@@ -238,12 +395,26 @@ export default function Screen() {
           })}
         </div>
       </div>
-      <div
-        className="fixed -bottom-15 pb-15"
-        style={{ width: "calc(100vw - 277px)" }}
-      >
+      <div className="absolute inset-x-0 bottom-0 z-10">
         <InputComponent setValue={setValue} value={value} />
       </div>
+      <div
+        aria-hidden="true"
+        className="pointer-events-none absolute left-0 top-0 -z-10 overflow-hidden opacity-0"
+        style={{ width: scrollWidth > 0 ? `${scrollWidth}px` : "100%" }}
+      >
+        {messages.map((message) => (
+          <div
+            key={`measure-${message.send_time}`}
+            ref={(element) => {
+              setMeasurementRef(message.send_time, element);
+            }}
+            style={{ padding: "4px 0" }}
+          >
+            <Message message={message} />
+          </div>
+        ))}
+      </div>
     </div>
   );
 }