fix(chat-thread): scroll lag — stable grouping, collapse mode, sync measure

The experimental ChatX thread lagged on scroll where the legacy
ConversationView didn't. Three structural gaps, all app-side:

- Grouping rebuilt every tool_group/agent_turn wrapper per streamed
  chunk, so the row memos missed and the whole transcript reconciled on
  every token. createStableTurnGrouper now reuses a wrapper whenever its
  member items are reference-equal (the conversation builder freezes
  completed turns and clones active-turn rows), confining identity churn
  to the live turn — the same bound the legacy thread gets from
  createIncrementalThreadGrouper. ThreadItemBody is memoized on item
  identity for the same reason.

- ToolGroup ignored the conversationCollapseMode setting and used
  uncontrolled defaultOpen, so every group that streamed open stayed
  expanded (and mounted) for the life of the session. Open state is now
  controlled with legacy buildThreadGroups semantics — "all" collapses
  everything, "partial" collapses on turn completion, "none" never —
  with per-group manual overrides in sessionViewStore, wiped when the
  mode changes. A closed marker unmounts its body, keeping long
  transcripts' DOM (and the scroller engine's per-scroll child scans
  over it) bounded.

- The windowed body measured rows via the virtualizer's async
  ResizeObserver path, painting one frame at the 80px estimate when
  scrolling through history; it now resizes synchronously on mount like
  the legacy VirtualizedList, and the diff worker pool is capped at 2
  workers to match (the library default of 8 shiki isolates costs
  hundreds of MB).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EAy2CqYiHQeDTgCEAvAZtG

Author: claude

packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx

@@ -49,9 +49,9 @@ import {
 import { MessageJumpPicker } from "@posthog/ui/features/sessions/components/chat-thread/MessageJumpPicker";
 import { MessageMinimap } from "@posthog/ui/features/sessions/components/chat-thread/MessageMinimap";
 import { ToolGroup } from "@posthog/ui/features/sessions/components/chat-thread/ToolGroup";
+import { createStableTurnGrouper } from "@posthog/ui/features/sessions/components/chat-thread/threadGrouping";
 import { THREAD_HOTKEY_OPTIONS } from "@posthog/ui/features/sessions/components/chat-thread/threadHotkeys";
 import {
-  type AgentTurn,
   CHAT_THREAD_VIRTUALIZATION_THRESHOLD,
   completedTurnTimestamp,
   countFlatRows,
@@ -66,7 +66,6 @@ import { usePromptRecallSource } from "@posthog/ui/features/sessions/components/
 import { VirtualThreadScrollBody } from "@posthog/ui/features/sessions/components/chat-thread/VirtualThreadScrollBody";
 import { GitActionMessage } from "@posthog/ui/features/sessions/components/GitActionMessage";
 import { GitActionResult } from "@posthog/ui/features/sessions/components/GitActionResult";
-import { isUserInitiatedConversationItem } from "@posthog/ui/features/sessions/components/isUserInitiatedConversationItem";
 import { mergeConversationItems } from "@posthog/ui/features/sessions/components/mergeConversationItems";
 import { extractCanvasInstructions } from "@posthog/ui/features/sessions/components/session-update/canvasInstructions";
 import { extractChannelContext } from "@posthog/ui/features/sessions/components/session-update/channelContext";
@@ -90,6 +89,7 @@ import {
   useOptimisticItemsForTask,
   useSessionIsCloud,
 } from "@posthog/ui/features/sessions/sessionStore";
+import { useSessionViewActions } from "@posthog/ui/features/sessions/sessionViewStore";
 import type { UserMessageAttachment } from "@posthog/ui/features/sessions/userMessageTypes";
 import {
   SessionTaskIdProvider,
@@ -116,115 +116,6 @@ import {
 } from "react";
 import { useHotkeys } from "react-hotkeys-hook";
 
-type SessionUpdateItem = Extract<ConversationItem, { type: "session_update" }>;
-
-function isToolCallItem(item: ConversationItem): item is SessionUpdateItem {
-  return (
-    item.type === "session_update" && item.update.sessionUpdate === "tool_call"
-  );
-}
-
-/**
- * Session-updates that `SessionUpdateView` always renders as `null`. They produce no row, so they
- * must not break a contiguous tool run.
- */
-const INVISIBLE_UPDATES = new Set([
-  "user_message_chunk",
-  "tool_call_update",
-  "plan",
-  "available_commands_update",
-  "config_option_update",
-]);
-
-/**
- * True when an item renders nothing, so it should be transparent to tool grouping. Besides the
- * always-null updates, this covers text chunks the stream emits with empty/whitespace or non-text
- * content (a stray empty `agent_message_chunk` between two tool calls is hidden via `empty:hidden`
- * but would otherwise split the run into two ungrouped markers).
- */
-function isInvisibleItem(item: ConversationItem): boolean {
-  if (item.type !== "session_update") return false;
-  const update = item.update;
-  if (INVISIBLE_UPDATES.has(update.sessionUpdate)) return true;
-  if (
-    update.sessionUpdate === "agent_message_chunk" ||
-    update.sessionUpdate === "agent_thought_chunk"
-  ) {
-    return update.content.type !== "text" || update.content.text.trim() === "";
-  }
-  return false;
-}
-
-/**
- * Collapse each contiguous run of ≥2 tool-call updates into a single `ToolGroupItem`. A run is
- * broken by any *visible* non-tool item (prose, thought, status) so groups follow reading order;
- * invisible updates (see {@link INVISIBLE_UPDATES}) are transparent and don't split a run. A lone
- * tool call passes through untouched — it stays a single marker, matching the legacy thread.
- */
-function groupToolRuns(items: ConversationItem[]): ThreadItem[] {
-  const out: ThreadItem[] = [];
-  // The buffer holds the active run: tool items plus any invisible items interleaved with them.
-  let buffer: ConversationItem[] = [];
-  let toolCount = 0;
-
-  const flush = () => {
-    if (toolCount >= 2) {
-      const tools = buffer.filter(isToolCallItem);
-      out.push({ type: "tool_group", id: tools[0].id, tools });
-    } else {
-      out.push(...buffer);
-    }
-    buffer = [];
-    toolCount = 0;
-  };
-
-  for (const item of items) {
-    if (isToolCallItem(item)) {
-      buffer.push(item);
-      toolCount++;
-    } else if (isInvisibleItem(item)) {
-      // Don't break the run; carry it along (it renders nothing wherever it lands).
-      buffer.push(item);
-    } else {
-      flush();
-      out.push(item);
-    }
-  }
-  flush();
-  return out;
-}
-
-/**
- * Collapse each contiguous run of non-user rows into one {@link AgentTurn}, broken only by a
- * user-initiated row (which stays standalone so it remains the scroll anchor for the sticky header
- * and auto-follow). The turn block renders as a single muted card, tightening the spacing between
- * the agent's successive replies and tool calls.
- */
-function groupIntoTurns(rows: ThreadItem[]): TurnRow[] {
-  const out: TurnRow[] = [];
-  let buffer: ThreadItem[] = [];
-  const flush = () => {
-    if (buffer.length > 0) {
-      out.push({ type: "agent_turn", id: buffer[0].id, items: buffer });
-      buffer = [];
-    }
-  };
-  for (const row of rows) {
-    // git_action and skill_button_action stand in for the user's message when the prompt was a
-    // git operation or a skill button click (see handlePromptRequest) — they open a turn just
-    // like a user message, so they break the agent card too rather than render inside it as if
-    // they were agent output. Same boundary set as the legacy view's buildThreadGroups.
-    if (isUserInitiatedConversationItem(row)) {
-      flush();
-      out.push(row);
-    } else {
-      buffer.push(row);
-    }
-  }
-  flush();
-  return out;
-}
-
 function formatTimestamp(ts: number): string {
   return new Date(ts).toLocaleString([], {
     month: "short",
@@ -511,8 +402,11 @@ const AgentProse = memo(function AgentProse({
 
 /** Renders a single thread item's body (no scroller wrapper), reused for standalone rows and for
  * each item inside an agent-turn card. `isTrailing` marks the turn's last item — a trailing tool
- * group of a streaming turn may still grow, so its label stays "Using …" between tool calls. */
-function ThreadItemBody({
+ * group of a streaming turn may still grow, so its label stays "Using …" between tool calls.
+ *
+ * Memoized on item identity: the conversation builder freezes completed turns' items and clones
+ * active-turn items per chunk, so identity equality is exactly content equality. */
+const ThreadItemBody = memo(function ThreadItemBody({
   item,
   renderItem,
   isTrailing = false,
@@ -529,6 +423,7 @@ function ThreadItemBody({
       !!context && !context.turnComplete && !context.turnCancelled;
     return (
       <ToolGroup
+        groupId={item.id}
         tools={item.tools}
         mayStillGrow={isTrailing && turnStreaming}
       />
@@ -545,13 +440,13 @@ function ThreadItemBody({
     );
   }
   return <>{renderItem(item)}</>;
-}
+});
 
 /**
  * One transcript row. Memoized and scroll-state-free, so rows never re-render while scrolling — the
  * non-virtualized thread stays cheap. The pinned header is the separate overlay, not the rows.
  *
- * An {@link AgentTurn} renders as a single muted card wrapping its items with tight spacing; a user
+ * An `AgentTurn` renders as a single muted card wrapping its items with tight spacing; a user
  * message stays a standalone anchored row.
  */
 const ThreadRow = memo(function ThreadRow({
@@ -1018,6 +913,10 @@ function ChatThreadRenderer({
     () => ({
       workerFactory: () => diffWorkerFactory(),
       totalASTLRUCacheSize: 200,
+      // Each pooled highlighter worker is a full V8 isolate with shiki
+      // grammars loaded (~40MB RSS); the library default of 8 costs hundreds
+      // of MB for parallelism conversation diffs don't need.
+      poolSize: 2,
     }),
     [diffWorkerFactory],
   );
@@ -1031,11 +930,29 @@ function ChatThreadRenderer({
     [conversationItems, optimisticItems, isCloud],
   );
 
+  // Identity-preserving grouping: per streamed chunk, only the active turn's
+  // wrappers change, so the memoized rows skip everything else (see
+  // createStableTurnGrouper). One grouper per mounted thread — `key={taskId}`
+  // above remounts (and so resets) it on task switch.
+  const turnGrouperRef = useRef<ReturnType<
+    typeof createStableTurnGrouper
+  > | null>(null);
+  turnGrouperRef.current ??= createStableTurnGrouper();
+  const turnGrouper = turnGrouperRef.current;
   const rows = useMemo<TurnRow[]>(
-    () => groupIntoTurns(groupToolRuns(items)),
-    [items],
+    () => turnGrouper.update(items),
+    [items, turnGrouper],
   );
 
+  // Changing the global collapse mode wipes ephemeral per-chip overrides, so
+  // every group snaps to the new mode's base state (same as the legacy view).
+  const sessionViewActions = useSessionViewActions();
+  const collapseMode = useSettingsStore((s) => s.conversationCollapseMode);
+  // biome-ignore lint/correctness/useExhaustiveDependencies: intentionally keyed on collapseMode only
+  useEffect(() => {
+    sessionViewActions.clearGroupOverrides();
+  }, [collapseMode]);
+
   // Virtualization ratchet: past the threshold the thread switches to the windowed body and
   // stays there for the life of this mount (see CHAT_THREAD_VIRTUALIZATION_THRESHOLD). Long
   // sessions start virtualized from the first render; a live session flips once mid-stream,

packages/ui/src/features/sessions/components/chat-thread/ToolGroup.test.tsx

@@ -40,6 +40,7 @@ describe("ToolGroup", () => {
       <ServiceProvider container={new Container()}>
         <Theme>
           <ToolGroup
+            groupId="spawn-1"
             tools={[subagentItem("spawn-1"), subagentItem("spawn-2")]}
           />
         </Theme>
@@ -54,6 +55,7 @@ describe("ToolGroup", () => {
       <ServiceProvider container={new Container()}>
         <Theme>
           <ToolGroup
+            groupId="spawn-1"
             tools={[
               subagentItem("spawn-1", {
                 status: "in_progress",

packages/ui/src/features/sessions/components/chat-thread/ToolGroup.tsx

@@ -7,7 +7,12 @@ import {
   Spinner,
 } from "@posthog/quill";
 import { readAgentToolName } from "@posthog/shared";
+import {
+  useGroupOverride,
+  useSessionViewActions,
+} from "@posthog/ui/features/sessions/sessionViewStore";
 import type { ToolCall } from "@posthog/ui/features/sessions/types";
+import { useSettingsStore } from "@posthog/ui/features/settings/settingsStore";
 import { memo } from "react";
 import type { ConversationItem } from "../buildConversationItems";
 import { grouping } from "../new-thread/conversationThreadConfig";
@@ -81,12 +86,22 @@ export function isToolActive(item: ToolGroupItem["tools"][number]): boolean {
  * collapsible body holds each tool's own marker via `SessionUpdateView` (which dispatches through
  * `ToolCallBlock` → `ToolRow` → `ChatMarker`).
  *
- * Expanded by default while the turn is still running (live visibility), collapsed once complete.
+ * Open state follows the global collapse mode (same semantics as the legacy thread's
+ * `buildThreadGroups`): "all" keeps every group collapsed, "partial" streams the live turn
+ * expanded and collapses it when the turn settles, "none" keeps everything expanded. A manual
+ * toggle overrides the mode for this group until the mode changes (the thread wipes overrides
+ * then). Controlled — not `defaultOpen` — so a group that streamed open actually collapses on
+ * completion instead of staying mounted for the session's life; a closed marker unmounts its
+ * body, which is what keeps a long transcript's DOM (and the scroller engine's per-scroll scans
+ * over it) bounded.
  */
 export const ToolGroup = memo(function ToolGroup({
+  groupId,
   tools,
   mayStillGrow = false,
 }: {
+  /** Stable row id (the run's first tool), keying this group's manual expand/collapse override. */
+  groupId: string;
   tools: ToolGroupItem["tools"];
   /**
    * True when this run is the turn's trailing content and the turn is still
@@ -98,8 +113,17 @@ export const ToolGroup = memo(function ToolGroup({
   mayStillGrow?: boolean;
 }) {
   const turnComplete = tools[0]?.turnContext.turnComplete ?? false;
+  const turnCancelled = tools[0]?.turnContext.turnCancelled ?? false;
   const isActive = tools.some(isToolActive) || mayStillGrow;
 
+  const collapseMode = useSettingsStore((s) => s.conversationCollapseMode);
+  const override = useGroupOverride(groupId);
+  const { setGroupOverride } = useSessionViewActions();
+  const settled = turnComplete || turnCancelled;
+  const baseCollapse =
+    collapseMode === "all" || (collapseMode === "partial" && settled);
+  const open = override ?? !baseCollapse;
+
   // Uniform when every tool in the run shares the same name/kind — then we can name it.
   const keys = tools.map(toolKey);
   const uniform = keys.every((k) => k === keys[0]);
@@ -114,7 +138,8 @@ export const ToolGroup = memo(function ToolGroup({
 
   return (
     <ChatMarker
-      defaultOpen={!turnComplete}
+      open={open}
+      onOpenChange={(next) => setGroupOverride(groupId, next)}
       body={tools.map((item) => (
         <SessionUpdateView
           key={item.id}

packages/ui/src/features/sessions/components/chat-thread/VirtualThreadScrollBody.tsx

@@ -327,6 +327,20 @@ export function VirtualThreadScrollBody({
     viewportRef,
   );
 
+  // Resize synchronously on mount: the virtualizer's own measureElement path is
+  // ResizeObserver-fed and lands a frame late, so a row scrolled into view
+  // would paint once at the 80px estimate and then jump. Same recipe as the
+  // legacy VirtualizedList.
+  const measureElementImmediately = useCallback(
+    (node: HTMLDivElement | null) => {
+      virtualizer.measureElement(node);
+      if (!node) return;
+      const index = Number(node.dataset.index);
+      virtualizer.resizeItem(index, node.offsetHeight);
+    },
+    [virtualizer],
+  );
+
   const userRows = useMemo(() => {
     const result: UserRow[] = [];
     flatRows.forEach((row, index) => {
@@ -439,7 +453,7 @@ export function VirtualThreadScrollBody({
                 return (
                   <div
                     key={virtualItem.key}
-                    ref={virtualizer.measureElement}
+                    ref={measureElementImmediately}
                     data-index={virtualItem.index}
                     className="absolute top-0 left-0 w-full"
                     style={{ transform: `translateY(${virtualItem.start}px)` }}