Extract thread grouping logic and add stable wrapper identity

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 { useThreadScrollRequest } from "@posthog/ui/features/sessions/threadNavigationStore";
 import type { UserMessageAttachment } from "@posthog/ui/features/sessions/userMessageTypes";
 import {
@@ -117,115 +117,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",
@@ -512,8 +403,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,
@@ -530,6 +424,7 @@ function ThreadItemBody({
       !!context && !context.turnComplete && !context.turnCancelled;
     return (
       <ToolGroup
+        groupId={item.id}
         tools={item.tools}
         mayStillGrow={isTrailing && turnStreaming}
       />
@@ -546,13 +441,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({
@@ -1034,6 +929,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],
   );
@@ -1047,11 +946,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)` }}