diff --git a/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx b/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx index 6a80be72ab..f1b558816f 100644 --- a/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx +++ b/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; - -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 ( @@ -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 | null>(null); + turnGrouperRef.current ??= createStableTurnGrouper(); + const turnGrouper = turnGrouperRef.current; const rows = useMemo( - () => 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, diff --git a/packages/ui/src/features/sessions/components/chat-thread/ToolGroup.test.tsx b/packages/ui/src/features/sessions/components/chat-thread/ToolGroup.test.tsx index 6aca2a9058..eafb0f71c0 100644 --- a/packages/ui/src/features/sessions/components/chat-thread/ToolGroup.test.tsx +++ b/packages/ui/src/features/sessions/components/chat-thread/ToolGroup.test.tsx @@ -40,6 +40,7 @@ describe("ToolGroup", () => { @@ -54,6 +55,7 @@ describe("ToolGroup", () => { 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 ( setGroupOverride(groupId, next)} body={tools.map((item) => ( { + 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 (
; + +function userMessage(id: string): ConversationItem { + return { type: "user_message", id, content: `msg ${id}`, timestamp: 1 }; +} + +function turnContext({ + turnComplete = false, +}: { + turnComplete?: boolean; +} = {}): SessionUpdateItem["turnContext"] { + return { + toolCalls: new Map(), + childItems: new Map(), + turnCancelled: false, + turnComplete, + }; +} + +function prose( + id: string, + opts?: { turnComplete?: boolean }, +): SessionUpdateItem { + return { + type: "session_update", + id, + update: { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text: `text ${id}` }, + } as SessionUpdateItem["update"], + turnContext: turnContext(opts), + timestamp: 1, + }; +} + +function toolCall( + id: string, + opts?: { turnComplete?: boolean }, +): SessionUpdateItem { + return { + type: "session_update", + id, + update: { + sessionUpdate: "tool_call", + toolCallId: id, + title: `tool ${id}`, + } as unknown as SessionUpdateItem["update"], + turnContext: turnContext(opts), + timestamp: 1, + }; +} + +function invisible(id: string): SessionUpdateItem { + return { + type: "session_update", + id, + update: { + sessionUpdate: "tool_call_update", + toolCallId: id, + } as unknown as SessionUpdateItem["update"], + turnContext: turnContext(), + timestamp: 1, + }; +} + +/** The active-turn clone the conversation builder hands out per streamed chunk. */ +function clone(item: SessionUpdateItem): SessionUpdateItem { + return { ...item, turnContext: { ...item.turnContext } }; +} + +describe("groupToolRuns", () => { + it("collapses a run of two or more tool calls into one group", () => { + const [a, b, c] = [toolCall("a"), toolCall("b"), toolCall("c")]; + const rows = groupToolRuns([a, b, c]); + expect(rows).toEqual([{ type: "tool_group", id: "a", tools: [a, b, c] }]); + }); + + it("passes a lone tool call through untouched", () => { + const a = toolCall("a"); + expect(groupToolRuns([a])).toEqual([a]); + }); + + it("keeps a run contiguous across invisible items but breaks it on visible prose", () => { + const [a, u, b, p, c, d] = [ + toolCall("a"), + invisible("u"), + toolCall("b"), + prose("p"), + toolCall("c"), + toolCall("d"), + ]; + const rows = groupToolRuns([a, u, b, p, c, d]); + expect(rows).toEqual([ + { type: "tool_group", id: "a", tools: [a, b] }, + p, + { type: "tool_group", id: "c", tools: [c, d] }, + ]); + }); +}); + +describe("groupIntoTurns", () => { + it("wraps contiguous agent rows into one turn, broken by user-initiated rows", () => { + const u1 = userMessage("u1"); + const a = prose("a"); + const b = prose("b"); + const git: ConversationItem = { + type: "git_action", + id: "g1", + actionType: "commit" as never, + }; + const c = prose("c"); + const rows = groupIntoTurns([u1, a, b, git, c]); + expect(rows).toEqual([ + u1, + { type: "agent_turn", id: "a", items: [a, b] }, + git, + { type: "agent_turn", id: "c", items: [c] }, + ]); + }); +}); + +describe("createStableTurnGrouper", () => { + it("produces the same rows as a plain groupIntoTurns(groupToolRuns(...)) pass", () => { + const items = [ + userMessage("u1"), + prose("a"), + toolCall("t1"), + toolCall("t2"), + userMessage("u2"), + toolCall("t3"), + ]; + expect(createStableTurnGrouper().update(items)).toEqual( + groupIntoTurns(groupToolRuns(items)), + ); + }); + + it("reuses turn and group wrappers when member items keep identity", () => { + const grouper = createStableTurnGrouper(); + const frozen = [ + userMessage("u1"), + prose("a", { turnComplete: true }), + toolCall("t1", { turnComplete: true }), + toolCall("t2", { turnComplete: true }), + ]; + const first = grouper.update([ + ...frozen, + userMessage("u2"), + toolCall("t3"), + ]); + // Streamed chunk: the active turn's items are cloned, the frozen prefix keeps identity. + const second = grouper.update([ + ...frozen, + userMessage("u2"), + clone(toolCall("t3")), + clone(toolCall("t4")), + ]); + + expect(second[1]).toBe(first[1]); + const frozenTurn = second[1] as AgentTurn; + expect(frozenTurn.items[1]).toBe((first[1] as AgentTurn).items[1]); + // The active turn's wrapper is fresh — its members changed. + expect(second[3]).not.toBe(first[3]); + }); + + it("rebuilds a group wrapper when any member is replaced", () => { + const grouper = createStableTurnGrouper(); + const t1 = toolCall("t1"); + const t2 = toolCall("t2"); + const groupOf = (rows: ReturnType) => + (rows[0] as AgentTurn).items[0]; + const first = grouper.update([t1, t2]); + const second = grouper.update([t1, clone(t2)]); + expect(groupOf(second)).not.toBe(groupOf(first)); + expect(second[0]).not.toBe(first[0]); + + const stable = grouper.update([t1, t2]); + const third = grouper.update([t1, t2]); + expect(groupOf(third)).toBe(groupOf(stable)); + expect(third[0]).toBe(stable[0]); + }); + + it("rebuilds a group wrapper when the run grows", () => { + const grouper = createStableTurnGrouper(); + const t1 = toolCall("t1"); + const t2 = toolCall("t2"); + const first = grouper.update([t1, t2]); + const second = grouper.update([t1, t2, toolCall("t3")]); + expect((second[0] as AgentTurn).items[0]).not.toBe( + (first[0] as AgentTurn).items[0], + ); + }); +}); diff --git a/packages/ui/src/features/sessions/components/chat-thread/threadGrouping.ts b/packages/ui/src/features/sessions/components/chat-thread/threadGrouping.ts new file mode 100644 index 0000000000..78df74b4a5 --- /dev/null +++ b/packages/ui/src/features/sessions/components/chat-thread/threadGrouping.ts @@ -0,0 +1,174 @@ +import type { ConversationItem } from "@posthog/ui/features/sessions/components/buildConversationItems"; +import type { ToolGroupItem } from "@posthog/ui/features/sessions/components/chat-thread/ToolGroup"; +import type { + AgentTurn, + ThreadItem, + TurnRow, +} from "@posthog/ui/features/sessions/components/chat-thread/threadVirtualization"; +import { isUserInitiatedConversationItem } from "@posthog/ui/features/sessions/components/isUserInitiatedConversationItem"; + +type SessionUpdateItem = Extract; + +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. + */ +export 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. + */ +export 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 sameMembers(a: readonly T[], b: readonly T[]): boolean { + if (a.length !== b.length) return false; + for (let i = 0; i < a.length; i++) { + if (a[i] !== b[i]) return false; + } + return true; +} + +/** + * `groupIntoTurns(groupToolRuns(items))` with wrapper identity preserved across calls: a + * `tool_group` or `agent_turn` whose member items are all reference-equal to the previous call's + * reuses the previous wrapper object. + * + * This is what lets the memoized thread rows actually skip work during streaming. The conversation + * builder freezes completed turns (their items keep identity) and clones every active-turn row per + * streamed chunk, so with fresh wrappers each call the row memos miss for the *entire* transcript + * on every chunk — a full-thread reconcile per token. With reuse, identity churn is confined to the + * live turn, and per-chunk render work tracks the active turn like the legacy thread's + * `createIncrementalThreadGrouper` path. + * + * Reuse is decided per wrapper by member identity, never by position, so steering, optimistic + * message swaps, and interruptions degrade to rebuilding the affected wrappers rather than + * rendering stale content. + */ +export function createStableTurnGrouper() { + let previousGroups = new Map(); + let previousTurns = new Map(); + + const update = (items: ConversationItem[]): TurnRow[] => { + const nextGroups = new Map(); + const threadItems = groupToolRuns(items).map((row): ThreadItem => { + if (row.type !== "tool_group") return row; + const previous = previousGroups.get(row.id); + const reused = + previous && sameMembers(previous.tools, row.tools) ? previous : row; + nextGroups.set(reused.id, reused); + return reused; + }); + + const nextTurns = new Map(); + const rows = groupIntoTurns(threadItems).map((row): TurnRow => { + if (row.type !== "agent_turn") return row; + const previous = previousTurns.get(row.id); + const reused = + previous && sameMembers(previous.items, row.items) ? previous : row; + nextTurns.set(reused.id, reused); + return reused; + }); + + previousGroups = nextGroups; + previousTurns = nextTurns; + return rows; + }; + + return { update }; +} diff --git a/packages/ui/src/features/sessions/sessionViewStore.ts b/packages/ui/src/features/sessions/sessionViewStore.ts index ac87639745..0e702bceef 100644 --- a/packages/ui/src/features/sessions/sessionViewStore.ts +++ b/packages/ui/src/features/sessions/sessionViewStore.ts @@ -67,6 +67,13 @@ export const useShowRawLogs = () => useStore((s) => s.showRawLogs); export const useSearchQuery = () => useStore((s) => s.searchQuery); export const useShowSearch = () => useStore((s) => s.showSearch); export const useGroupOverrides = () => useStore((s) => s.groupOverrides); +/** + * A single group's override, so a tool group re-renders only when its own + * entry changes — subscribing to the whole map would re-render every mounted + * group on any toggle. + */ +export const useGroupOverride = (id: string): boolean | undefined => + useStore((s) => s.groupOverrides[id]); export const useQueueCollapsed = (taskId: string) => useStore((s) => s.queueCollapsedByTaskId[taskId] ?? false); export const useSessionViewActions = () => useStore((s) => s.actions);