perf(app): reduce streaming rerenders and scroll churn (#67)

Amp-Thread-ID: https://ampcode.com/threads/T-019c793a-d087-762a-882a-8cc11e6101db

Co-authored-by: Ubuntu <ubuntu@dev-yk.tail388b2e.ts.net>

packages/app/src/components/agent-stream-view.tsx

@@ -107,6 +107,8 @@ export function AgentStreamView({
   const hasScrolledInitially = useRef(false);
   const hasAutoScrolledOnce = useRef(false);
   const isNearBottomRef = useRef(true);
+  const pendingAutoScrollFrameRef = useRef<number | null>(null);
+  const pendingAutoScrollAnimatedRef = useRef(false);
   const streamItemCountRef = useRef(0);
   const streamScrollbarMetrics = useWebDesktopScrollbarMetrics();
   const [expandedInlineToolCallIds, setExpandedInlineToolCallIds] = useState<Set<string>>(new Set());
@@ -214,6 +216,35 @@ export function AgentStreamView({
     []
   );
 
+  const scheduleAutoScroll = useCallback(
+    ({ animated }: { animated: boolean }) => {
+      pendingAutoScrollAnimatedRef.current =
+        pendingAutoScrollAnimatedRef.current || animated;
+
+      if (pendingAutoScrollFrameRef.current !== null) {
+        return;
+      }
+
+      pendingAutoScrollFrameRef.current = requestAnimationFrame(() => {
+        pendingAutoScrollFrameRef.current = null;
+        const shouldAnimate = pendingAutoScrollAnimatedRef.current;
+        pendingAutoScrollAnimatedRef.current = false;
+        scrollToBottomInternal({ animated: shouldAnimate });
+      });
+    },
+    [scrollToBottomInternal]
+  );
+
+  useEffect(() => {
+    return () => {
+      if (pendingAutoScrollFrameRef.current !== null) {
+        cancelAnimationFrame(pendingAutoScrollFrameRef.current);
+        pendingAutoScrollFrameRef.current = null;
+      }
+      pendingAutoScrollAnimatedRef.current = false;
+    };
+  }, []);
+
   useEffect(() => {
     if (streamItems.length === 0) {
       return;
@@ -233,9 +264,9 @@ export function AgentStreamView({
     }
 
     const shouldAnimate = hasScrolledInitially.current;
-    scrollToBottomInternal({ animated: shouldAnimate });
+    scheduleAutoScroll({ animated: shouldAnimate });
     hasScrolledInitially.current = true;
-  }, [streamItems, scrollToBottomInternal]);
+  }, [streamItems, scheduleAutoScroll]);
 
   function scrollToBottom() {
     scrollToBottomInternal({ animated: true });

packages/app/src/contexts/session-context.tsx

@@ -183,6 +183,9 @@ function SessionProviderInternal({
   const setAgentStreamHead = useSessionStore(
     (state) => state.setAgentStreamHead
   );
+  const setAgentStreamState = useSessionStore(
+    (state) => state.setAgentStreamState
+  );
   const clearAgentStreamHead = useSessionStore(
     (state) => state.clearAgentStreamHead
   );
@@ -784,23 +787,10 @@ function SessionProviderInternal({
         timestamp: parsedTimestamp,
       });
 
-      if (changedTail) {
+      if (changedTail || changedHead) {
-        setAgentStreamTail(serverId, (prev) => {
+        setAgentStreamState(serverId, agentId, {
-          const next = new Map(prev);
+          ...(changedTail ? { tail } : {}),
-          next.set(agentId, tail);
+          ...(changedHead ? { head } : {}),
-          return next;
-        });
-      }
-
-      if (changedHead) {
-        setAgentStreamHead(serverId, (prev) => {
-          const next = new Map(prev);
-          if (head.length > 0) {
-            next.set(agentId, head);
-          } else {
-            next.delete(agentId);
-          }
-          return next;
         });
       }
 
@@ -1261,6 +1251,7 @@ function SessionProviderInternal({
     setCurrentAssistantMessage,
     setAgentStreamTail,
     setAgentStreamHead,
+    setAgentStreamState,
     clearAgentStreamHead,
     setAgentTimelineCursor,
     setInitializingAgents,

packages/app/src/stores/session-store.ts

@@ -243,6 +243,11 @@ interface SessionStoreActions {
   // Stream state (head/tail model)
   setAgentStreamTail: (serverId: string, state: Map<string, StreamItem[]> | ((prev: Map<string, StreamItem[]>) => Map<string, StreamItem[]>)) => void;
   setAgentStreamHead: (serverId: string, state: Map<string, StreamItem[]> | ((prev: Map<string, StreamItem[]>) => Map<string, StreamItem[]>)) => void;
+  setAgentStreamState: (
+    serverId: string,
+    agentId: string,
+    state: { tail?: StreamItem[]; head?: StreamItem[] }
+  ) => void;
   clearAgentStreamHead: (serverId: string, agentId: string) => void;
   setAgentTimelineCursor: (
     serverId: string,
@@ -623,6 +628,67 @@ export const useSessionStore = create<SessionStore>()(
       });
     },
 
+    setAgentStreamState: (serverId, agentId, state) => {
+      set((prev) => {
+        const session = prev.sessions[serverId];
+        if (!session) {
+          return prev;
+        }
+
+        let nextTail = session.agentStreamTail;
+        let nextHead = session.agentStreamHead;
+        let changedTail = false;
+        let changedHead = false;
+
+        if (state.tail !== undefined) {
+          const existingTail = session.agentStreamTail.get(agentId);
+          if (existingTail !== state.tail) {
+            nextTail = new Map(session.agentStreamTail);
+            nextTail.set(agentId, state.tail);
+            changedTail = true;
+          }
+        }
+
+        if (state.head !== undefined) {
+          const existingHead = session.agentStreamHead.get(agentId);
+          const shouldDeleteHead = state.head.length === 0;
+          if (shouldDeleteHead) {
+            if (session.agentStreamHead.has(agentId)) {
+              nextHead = new Map(session.agentStreamHead);
+              nextHead.delete(agentId);
+              changedHead = true;
+            }
+          } else if (existingHead !== state.head) {
+            nextHead = new Map(session.agentStreamHead);
+            nextHead.set(agentId, state.head);
+            changedHead = true;
+          }
+        }
+
+        if (!changedTail && !changedHead) {
+          return prev;
+        }
+
+        logSessionStoreUpdate("setAgentStreamState", serverId, {
+          agentId,
+          changedTail,
+          changedHead,
+        });
+
+        return {
+          ...prev,
+          sessions: {
+            ...prev.sessions,
+            [serverId]: {
+              ...session,
+              agentStreamTail: nextTail,
+              agentStreamHead: nextHead,
+            },
+          },
+        };
+      });
+    },
+
     clearAgentStreamHead: (serverId, agentId) => {
       set((prev) => {
         const session = prev.sessions[serverId];