@assistant-ui/react-native 0.1.29 → 0.1.31

package/src/primitives/thread/ThreadMessages.test.tsx

@@ -1,7 +1,9 @@
-import { act } from "react";
+import { act, createRef } from "react";
 import { createRoot, type Root } from "react-dom/client";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import { ThreadMessages } from "./ThreadMessages";
+import type { FlatList } from "react-native";
+import type { ThreadMessage } from "@assistant-ui/core";
+import { ThreadMessages, ThreadMessagesFlatList } from "./ThreadMessages";
 
 type Msg = { id: string; role: string };
 
@@ -11,16 +13,96 @@ const h = vi.hoisted(() => ({
     message: { role: "user" as string, composer: { isEditing: false } },
   },
   itemState: { role: "user" } as { role: string },
+  events: {} as Record<string, Set<() => void>>,
+  flatListProps: null as Record<string, unknown> | null,
+  scrollToEnd: vi.fn(),
 }));
 
-vi.mock("@assistant-ui/store", () => ({
-  useAuiState: <T,>(selector: (s: typeof h.state) => T) => selector(h.state),
-  RenderChildrenWithAccessor: ({
-    children,
-  }: {
-    children: (getItem: () => unknown) => unknown;
-  }) => children(() => h.itemState),
-}));
+vi.mock("react-native", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("react-native")>();
+  const React = await import("react");
+
+  const FlatListMock = React.forwardRef(function FlatListMock(
+    props: Record<string, unknown>,
+    ref,
+  ) {
+    h.flatListProps = props;
+
+    React.useImperativeHandle(ref, () => ({
+      scrollToEnd: h.scrollToEnd,
+    }));
+
+    const data = (props.data as unknown[]) ?? [];
+    const renderItem = props.renderItem as
+      | ((value: {
+          item: unknown;
+          index: number;
+          separators: {
+            highlight: () => void;
+            unhighlight: () => void;
+            updateProps: () => void;
+          };
+        }) => React.ReactNode)
+      | undefined;
+    const keyExtractor = props.keyExtractor as
+      | ((item: unknown, index: number) => string)
+      | undefined;
+
+    return React.createElement(
+      "div",
+      { "data-testid": "flatlist" },
+      data.map((item, index) =>
+        React.createElement(
+          "div",
+          { key: keyExtractor?.(item, index) ?? index },
+          renderItem?.({
+            item,
+            index,
+            separators: {
+              highlight: vi.fn(),
+              unhighlight: vi.fn(),
+              updateProps: vi.fn(),
+            },
+          }),
+        ),
+      ),
+    );
+  });
+
+  return {
+    ...actual,
+    FlatList: FlatListMock,
+  };
+});
+
+vi.mock("@assistant-ui/store", async () => {
+  const React = await import("react");
+
+  return {
+    useAuiState: <T,>(selector: (s: typeof h.state) => T) => selector(h.state),
+    useAuiEvent: (
+      selector: string | { scope: string; event: string },
+      callback: () => void,
+    ) => {
+      React.useEffect(() => {
+        const event =
+          typeof selector === "string"
+            ? selector
+            : `${selector.scope}.${selector.event}`;
+        const callbacks = (h.events[event] ??= new Set());
+        callbacks.add(callback);
+        return () => {
+          callbacks.delete(callback);
+        };
+      }, [selector, callback]);
+    },
+    RenderChildrenWithAccessor: ({
+      children,
+    }: {
+      children: (getItem: () => unknown) => unknown;
+    }) => children(() => h.itemState),
+  };
+});
 
 vi.mock("@assistant-ui/core/react", () => ({
   MessageByIndexProvider: ({ children }: { children: unknown }) => children,
@@ -37,6 +119,9 @@ describe("ThreadMessages", () => {
     h.state.message.role = "user";
     h.state.message.composer.isEditing = false;
     h.itemState = { role: "user" };
+    h.events = {};
+    h.flatListProps = null;
+    h.scrollToEnd.mockReset();
 
     container = document.createElement("div");
     document.body.appendChild(container);
@@ -56,6 +141,37 @@ describe("ThreadMessages", () => {
     });
   };
 
+  const mountFlatList = async (
+    props: Parameters<typeof ThreadMessagesFlatList>[0],
+  ) => {
+    await act(async () => {
+      root.render(<ThreadMessagesFlatList {...props} />);
+    });
+  };
+
+  const getFlatListProps = () => {
+    const props = h.flatListProps as {
+      onContentSizeChange?: (width: number, height: number) => void;
+      onLayout?: (event: unknown) => void;
+      onScroll?: (event: unknown) => void;
+      scrollEventThrottle?: number;
+    } | null;
+    if (!props) throw new Error("FlatList was not rendered");
+    return props;
+  };
+
+  const emit = async (event: string) => {
+    await act(async () => {
+      h.events[event]?.forEach((callback) => {
+        callback();
+      });
+    });
+  };
+
+  const messageComponents = {
+    Message: () => <span data-testid="c-message">message</span>,
+  };
+
   describe("components mode dispatch", () => {
     const makeComponents = () => ({
       Message: vi.fn(() => <span data-testid="c-message">message</span>),
@@ -189,4 +305,322 @@ describe("ThreadMessages", () => {
     });
     expect(container.querySelector('[data-testid="c-message"]')).toBeNull();
   });
+
+  it("forwards refs to the underlying FlatList", async () => {
+    h.state.thread.messages = [{ id: "1", role: "user" }];
+    const ref = createRef<FlatList<ThreadMessage>>();
+
+    await act(async () => {
+      root.render(
+        <ThreadMessages
+          ref={ref}
+          components={{
+            Message: () => <span data-testid="c-message">message</span>,
+          }}
+        />,
+      );
+    });
+
+    expect(ref.current).not.toBeNull();
+  });
+
+  it("keeps deprecated Messages off the scroll-tracking path", async () => {
+    h.state.thread.messages = [{ id: "1", role: "user" }];
+    await mount({ components: messageComponents });
+    const props = getFlatListProps();
+
+    expect(props.onScroll).toBeUndefined();
+    expect(props.onLayout).toBeUndefined();
+    expect(props.onContentSizeChange).toBeUndefined();
+    expect(props.scrollEventThrottle).toBeUndefined();
+  });
+
+  it("keeps deprecated Messages from auto-scrolling by default", async () => {
+    h.state.thread.messages = [{ id: "1", role: "user" }];
+    await mount({ components: messageComponents });
+    const props = getFlatListProps();
+
+    await act(async () => {
+      props.onLayout?.({
+        nativeEvent: { layout: { height: 100 } },
+      });
+      props.onScroll?.({
+        nativeEvent: {
+          contentOffset: { y: 0 },
+          contentSize: { height: 100, width: 0 },
+          layoutMeasurement: { height: 100, width: 0 },
+        },
+      });
+      props.onContentSizeChange?.(0, 140);
+    });
+    await emit("thread.runStart");
+    await emit("threadListItem.switchedTo");
+
+    expect(h.scrollToEnd).not.toHaveBeenCalled();
+  });
+
+  describe("MessagesFlatList auto-scroll", () => {
+    it("scrolls to the bottom when messages first appear", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+
+      await mountFlatList({ components: messageComponents });
+
+      expect(h.scrollToEnd).toHaveBeenCalledWith({ animated: false });
+    });
+
+    it("scrolls to the bottom when a run starts", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+      await mountFlatList({ components: messageComponents });
+      h.scrollToEnd.mockClear();
+
+      await emit("thread.runStart");
+
+      expect(h.scrollToEnd).toHaveBeenCalledWith({ animated: true });
+    });
+
+    it("scrolls when content grows while already at the bottom", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+      await mountFlatList({ components: messageComponents });
+      const props = getFlatListProps();
+      h.scrollToEnd.mockClear();
+
+      await act(async () => {
+        props.onLayout?.({
+          nativeEvent: { layout: { height: 100 } },
+        });
+        props.onScroll?.({
+          nativeEvent: {
+            contentOffset: { y: 0 },
+            contentSize: { height: 100, width: 0 },
+            layoutMeasurement: { height: 100, width: 0 },
+          },
+        });
+        props.onContentSizeChange?.(0, 140);
+      });
+
+      expect(h.scrollToEnd).toHaveBeenCalledWith({ animated: false });
+    });
+
+    it("does not treat the first content-size event as automatic content growth", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+      await mountFlatList({
+        components: messageComponents,
+        scrollToBottomOnInitialize: false,
+      });
+      const props = getFlatListProps();
+
+      await act(async () => {
+        props.onContentSizeChange?.(0, 140);
+      });
+
+      expect(h.scrollToEnd).not.toHaveBeenCalled();
+    });
+
+    it("lands the initialize scroll on the first content-size event", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+      await mountFlatList({ components: messageComponents });
+      const props = getFlatListProps();
+
+      expect(h.scrollToEnd).toHaveBeenCalledTimes(1);
+
+      await act(async () => {
+        props.onContentSizeChange?.(0, 140);
+      });
+
+      expect(h.scrollToEnd).toHaveBeenCalledTimes(2);
+
+      await act(async () => {
+        props.onContentSizeChange?.(0, 140);
+      });
+
+      expect(h.scrollToEnd).toHaveBeenCalledTimes(2);
+    });
+
+    it("lands the thread-switch scroll on the next content-size event", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+      await mountFlatList({ components: messageComponents });
+      const props = getFlatListProps();
+      h.scrollToEnd.mockClear();
+
+      await emit("threadListItem.switchedTo");
+      expect(h.scrollToEnd).toHaveBeenCalledTimes(1);
+
+      await act(async () => {
+        props.onContentSizeChange?.(0, 80);
+      });
+
+      expect(h.scrollToEnd).toHaveBeenCalledTimes(2);
+    });
+
+    it("keeps following through consecutive growth events without scroll echoes", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+      await mountFlatList({ components: messageComponents });
+      const props = getFlatListProps();
+
+      await act(async () => {
+        props.onLayout?.({
+          nativeEvent: { layout: { height: 100 } },
+        });
+        props.onScroll?.({
+          nativeEvent: {
+            contentOffset: { y: 0 },
+            contentSize: { height: 100, width: 0 },
+            layoutMeasurement: { height: 100, width: 0 },
+          },
+        });
+      });
+      h.scrollToEnd.mockClear();
+
+      await act(async () => {
+        props.onContentSizeChange?.(0, 140);
+      });
+      await act(async () => {
+        props.onContentSizeChange?.(0, 180);
+      });
+
+      expect(h.scrollToEnd).toHaveBeenCalledTimes(2);
+    });
+
+    it("does not scroll when content grows after the user scrolled away", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+      await mountFlatList({ components: messageComponents });
+      const props = getFlatListProps();
+      h.scrollToEnd.mockClear();
+
+      await act(async () => {
+        props.onLayout?.({
+          nativeEvent: { layout: { height: 100 } },
+        });
+        props.onScroll?.({
+          nativeEvent: {
+            contentOffset: { y: 200 },
+            contentSize: { height: 300, width: 0 },
+            layoutMeasurement: { height: 100, width: 0 },
+          },
+        });
+        props.onScroll?.({
+          nativeEvent: {
+            contentOffset: { y: 50 },
+            contentSize: { height: 300, width: 0 },
+            layoutMeasurement: { height: 100, width: 0 },
+          },
+        });
+        props.onContentSizeChange?.(0, 340);
+      });
+
+      expect(h.scrollToEnd).not.toHaveBeenCalled();
+    });
+
+    it("lands the run-start scroll once the appended message resizes content", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+      await mountFlatList({ components: messageComponents });
+      const props = getFlatListProps();
+
+      await act(async () => {
+        props.onLayout?.({
+          nativeEvent: { layout: { height: 100 } },
+        });
+        props.onScroll?.({
+          nativeEvent: {
+            contentOffset: { y: 200 },
+            contentSize: { height: 300, width: 0 },
+            layoutMeasurement: { height: 100, width: 0 },
+          },
+        });
+        props.onScroll?.({
+          nativeEvent: {
+            contentOffset: { y: 50 },
+            contentSize: { height: 300, width: 0 },
+            layoutMeasurement: { height: 100, width: 0 },
+          },
+        });
+      });
+      h.scrollToEnd.mockClear();
+
+      await emit("thread.runStart");
+      expect(h.scrollToEnd).toHaveBeenCalledTimes(1);
+
+      await act(async () => {
+        props.onScroll?.({
+          nativeEvent: {
+            contentOffset: { y: 120 },
+            contentSize: { height: 300, width: 0 },
+            layoutMeasurement: { height: 100, width: 0 },
+          },
+        });
+        props.onContentSizeChange?.(0, 360);
+      });
+
+      expect(h.scrollToEnd).toHaveBeenCalledTimes(2);
+      expect(h.scrollToEnd).toHaveBeenLastCalledWith({ animated: true });
+    });
+
+    it("scrolls to the bottom when switching threads", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+      await mountFlatList({ components: messageComponents });
+      h.scrollToEnd.mockClear();
+
+      await emit("threadListItem.switchedTo");
+
+      expect(h.scrollToEnd).toHaveBeenCalledWith({ animated: false });
+    });
+
+    it("does not rearm initialize scroll when thread-switch scroll is disabled", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+      await mountFlatList({
+        components: messageComponents,
+        scrollToBottomOnThreadSwitch: false,
+      });
+      h.scrollToEnd.mockClear();
+
+      await emit("threadListItem.switchedTo");
+
+      h.state.thread.messages = [
+        { id: "2", role: "user" },
+        { id: "3", role: "assistant" },
+      ];
+      await mountFlatList({
+        components: messageComponents,
+        scrollToBottomOnThreadSwitch: false,
+      });
+
+      expect(h.scrollToEnd).not.toHaveBeenCalled();
+    });
+
+    it("honors opt-outs for automatic content growth and run-start scrolls", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+      await mountFlatList({
+        autoScroll: false,
+        components: messageComponents,
+        scrollToBottomOnInitialize: false,
+        scrollToBottomOnRunStart: false,
+      });
+      const props = getFlatListProps();
+      h.scrollToEnd.mockClear();
+
+      await act(async () => {
+        props.onLayout?.({
+          nativeEvent: { layout: { height: 100 } },
+        });
+        props.onScroll?.({
+          nativeEvent: {
+            contentOffset: { y: 0 },
+            contentSize: { height: 100, width: 0 },
+            layoutMeasurement: { height: 100, width: 0 },
+          },
+        });
+        props.onContentSizeChange?.(0, 140);
+      });
+      await emit("thread.runStart");
+
+      expect(h.scrollToEnd).not.toHaveBeenCalled();
+    });
+
+    it("sets a useful default scroll throttle", async () => {
+      h.state.thread.messages = [{ id: "1", role: "user" }];
+      await mountFlatList({ components: messageComponents });
+
+      expect(getFlatListProps().scrollEventThrottle).toBe(16);
+    });
+  });
 });