From c3393686502c17da6f02e965c5ba090f3d6a70be Mon Sep 17 00:00:00 2001 From: Arnaud Hillen Date: Thu, 30 Jul 2026 00:11:48 +0200 Subject: [PATCH 1/5] perf(sessions): persist conversation build caches across task re-opens Re-opening a task unmounts and remounts the whole detail view, so the incremental conversation builder and thread grouper (both component-scoped) were rebuilt from scratch on every click, re-parsing the full transcript on the main thread. For large sessions this blocks for seconds. Keep both in a module-level LRU cache keyed per call site and task, so a re-open reuses the already-built items. Entries are dropped when the residency system evicts a session's events, so memory reclaim still works. Generated-By: PostHog Code Task-Id: 5d24ea17-aec0-4334-884e-c2639867e260 --- .../sessions/components/ConversationView.tsx | 22 ++- .../components/chat-thread/ChatThread.tsx | 9 +- .../chat-thread/ChatThreadFooter.tsx | 9 +- .../hooks/conversationDerivedCache.test.ts | 83 ++++++++++++ .../hooks/conversationDerivedCache.ts | 126 ++++++++++++++++++ .../hooks/useConversationItems.test.tsx | 81 +++++++++++ .../sessions/hooks/useConversationItems.ts | 37 ++--- 7 files changed, 340 insertions(+), 27 deletions(-) create mode 100644 packages/ui/src/features/sessions/hooks/conversationDerivedCache.test.ts create mode 100644 packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts create mode 100644 packages/ui/src/features/sessions/hooks/useConversationItems.test.tsx diff --git a/packages/ui/src/features/sessions/components/ConversationView.tsx b/packages/ui/src/features/sessions/components/ConversationView.tsx index 6adcfcf522..df12d6cc18 100644 --- a/packages/ui/src/features/sessions/components/ConversationView.tsx +++ b/packages/ui/src/features/sessions/components/ConversationView.tsx @@ -51,6 +51,7 @@ import { } from "@posthog/ui/features/sessions/components/VirtualizedList"; import { CHAT_CONTENT_MAX_WIDTH } from "@posthog/ui/features/sessions/constants"; import { DIFFS_HIGHLIGHTER_OPTIONS } from "@posthog/ui/features/sessions/diffHighlighterOptions"; +import { getPersistentThreadGrouper } from "@posthog/ui/features/sessions/hooks/conversationDerivedCache"; import { useContextUsage } from "@posthog/ui/features/sessions/hooks/useContextUsage"; import { useConversationItems } from "@posthog/ui/features/sessions/hooks/useConversationItems"; import { useConversationSearch } from "@posthog/ui/features/sessions/hooks/useConversationSearch"; @@ -163,9 +164,12 @@ export function ConversationView({ lastTurnInfo, isCompacting, completedToolCallCount, - } = useConversationItems(events, isPromptPending, { - showDebugLogs, - }); + } = useConversationItems( + events, + isPromptPending, + { showDebugLogs }, + taskId ? { scope: "conversation-view", taskId } : undefined, + ); const firstUserMessageIdRef = useRef(undefined); if (firstUserMessageIdRef.current === undefined) { @@ -209,8 +213,16 @@ export function ConversationView({ const threadGrouperRef = useRef | null>(null); - threadGrouperRef.current ??= createIncrementalThreadGrouper(); - const threadGrouper = threadGrouperRef.current; + let threadGrouper: ReturnType; + if (taskId) { + threadGrouper = getPersistentThreadGrouper({ + scope: "conversation-view", + taskId, + }); + } else { + threadGrouperRef.current ??= createIncrementalThreadGrouper(); + threadGrouper = threadGrouperRef.current; + } const grouping = useMemo( () => threadGrouper.update(items, collapseMode, groupOverrides), [items, collapseMode, groupOverrides, threadGrouper], 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 dbacd99334..31554980fd 100644 --- a/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx +++ b/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx @@ -1086,9 +1086,12 @@ export function ChatThread({ events, ...props }: ChatThreadProps) { export function AcpChatThread({ events, ...props }: AcpChatThreadProps) { const showDebugLogs = useSettingsStore((state) => state.debugLogsCloudRuns); - const { items } = useConversationItems(events, props.isPromptPending, { - showDebugLogs, - }); + const { items } = useConversationItems( + events, + props.isPromptPending, + { showDebugLogs }, + props.taskId ? { scope: "chat-thread", taskId: props.taskId } : undefined, + ); return ( s.debugLogsCloudRuns); const eventContextUsage = useContextUsage(events); const contextUsage = usage === undefined ? eventContextUsage : usage; - const eventFooterState = useConversationItems(events, isPromptPending, { - showDebugLogs, - }); + const eventFooterState = useConversationItems( + events, + isPromptPending, + { showDebugLogs }, + taskId ? { scope: "chat-thread-footer", taskId } : undefined, + ); const lastTurnInfo = footerState?.lastTurnInfo ?? eventFooterState.lastTurnInfo; const isCompacting = diff --git a/packages/ui/src/features/sessions/hooks/conversationDerivedCache.test.ts b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.test.ts new file mode 100644 index 0000000000..0a354774f8 --- /dev/null +++ b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.test.ts @@ -0,0 +1,83 @@ +import type { AcpMessage, AgentSession } from "@posthog/shared"; +import { describe, expect, it } from "vitest"; +import { sessionStoreSetters } from "../sessionStore"; +import { + getConversationBuildCache, + getPersistentThreadGrouper, +} from "./conversationDerivedCache"; + +function msg(ts: number): AcpMessage { + return { + type: "acp_message", + ts, + message: { jsonrpc: "2.0", method: "session/update", params: {} }, + }; +} + +function seedSession(taskId: string, taskRunId: string): void { + sessionStoreSetters.setSession({ + taskId, + taskRunId, + events: [msg(1)], + messageQueue: [], + } as unknown as AgentSession); +} + +describe("conversationDerivedCache", () => { + it("returns the same cache entry across lookups for the same scope + task", () => { + const key = { scope: "test", taskId: "same-entry" }; + expect(getConversationBuildCache(key)).toBe(getConversationBuildCache(key)); + expect(getPersistentThreadGrouper(key)).toBe( + getPersistentThreadGrouper(key), + ); + }); + + it("keeps entries of different scopes for the same task independent", () => { + const taskId = "scoped-entries"; + const a = getConversationBuildCache({ scope: "thread", taskId }); + const b = getConversationBuildCache({ scope: "footer", taskId }); + expect(a).not.toBe(b); + }); + + it("evicts the least recently used entry beyond the cap", () => { + const first = getConversationBuildCache({ scope: "lru", taskId: "lru-0" }); + for (let i = 1; i < 50; i++) { + getConversationBuildCache({ scope: "lru", taskId: `lru-${i}` }); + } + const recent = getConversationBuildCache({ + scope: "lru", + taskId: "lru-49", + }); + expect(getConversationBuildCache({ scope: "lru", taskId: "lru-49" })).toBe( + recent, + ); + expect( + getConversationBuildCache({ scope: "lru", taskId: "lru-0" }), + ).not.toBe(first); + }); + + it("drops the cache when the session's events are evicted from the store", () => { + const key = { scope: "test", taskId: "evicted-task" }; + seedSession(key.taskId, "run-evicted"); + const before = getConversationBuildCache(key); + sessionStoreSetters.evictEvents("run-evicted"); + expect(getConversationBuildCache(key)).not.toBe(before); + }); + + it("drops the cache when the session is removed from the store", () => { + const key = { scope: "test", taskId: "removed-task" }; + seedSession(key.taskId, "run-removed"); + const before = getConversationBuildCache(key); + sessionStoreSetters.removeSession("run-removed"); + expect(getConversationBuildCache(key)).not.toBe(before); + }); + + it("keeps entries for tasks that never had a session through store commits", () => { + const key = { scope: "test", taskId: "sessionless-task" }; + const before = getConversationBuildCache(key); + // An unrelated commit must not sweep surfaces rendering without a session + // (e.g. archive views), or they would rebuild on every render. + seedSession("unrelated-task", "run-unrelated"); + expect(getConversationBuildCache(key)).toBe(before); + }); +}); diff --git a/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts new file mode 100644 index 0000000000..2243bbaf71 --- /dev/null +++ b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts @@ -0,0 +1,126 @@ +import type { AcpMessage } from "@posthog/shared"; +import type { BuildResult } from "@posthog/ui/features/sessions/components/buildConversationItems"; +import { createIncrementalConversationBuilder } from "@posthog/ui/features/sessions/components/incrementalConversationItems"; +import { createIncrementalThreadGrouper } from "@posthog/ui/features/sessions/components/new-thread/incrementalThreadGrouping"; +import { + type SessionState, + useSessionStore, +} from "@posthog/ui/features/sessions/sessionStore"; + +export interface ConversationBuildCache { + impl: ReturnType; + events: AcpMessage[] | null; + pending: boolean | null; + debug: boolean | undefined; + result: BuildResult | null; +} + +type ThreadGrouper = ReturnType; + +export interface ConversationCacheKey { + /** + * Call-site namespace. Hooks that feed different inputs for the same task + * (e.g. the chat thread vs its footer) must not share builder state, or + * they'd invalidate each other on every render. + */ + scope: string; + taskId: string; +} + +interface Entry { + taskId: string; + /** + * Entries are swept when their session's events are evicted, but only if + * the session was ever seen in the store. Surfaces that render transcripts + * without a live session (archive views) would otherwise be swept on every + * store commit and rebuild from scratch each render. + */ + hadSession: boolean; + value: T; +} + +/** + * Re-opening a task unmounts and remounts its whole view tree, so + * per-component caches die with it and every click re-parses the full + * transcript (seconds of main-thread work for large sessions). These + * module-level caches keep the incremental builder state alive across mounts. + */ +const MAX_CACHED_TASKS = 8; + +const buildCaches = new Map>(); +const threadGroupers = new Map>(); + +export function getConversationBuildCache( + key: ConversationCacheKey, +): ConversationBuildCache { + return getEntry(buildCaches, key, () => ({ + impl: createIncrementalConversationBuilder(), + events: null, + pending: null, + debug: undefined, + result: null, + })); +} + +export function getPersistentThreadGrouper( + key: ConversationCacheKey, +): ThreadGrouper { + return getEntry(threadGroupers, key, createIncrementalThreadGrouper); +} + +function getEntry( + map: Map>, + key: ConversationCacheKey, + create: () => T, +): T { + watchStoreForEviction(); + const mapKey = `${key.scope}:${key.taskId}`; + let entry = map.get(mapKey); + if (entry) { + // Re-insert to refresh LRU recency (Map preserves insertion order). + map.delete(mapKey); + } else { + entry = { taskId: key.taskId, hadSession: false, value: create() }; + } + entry.hadSession ||= hasSession(useSessionStore.getState(), key.taskId); + map.set(mapKey, entry); + for (const oldest of map.keys()) { + if (map.size <= MAX_CACHED_TASKS) break; + map.delete(oldest); + } + return entry.value; +} + +function hasSession(state: SessionState, taskId: string): boolean { + const taskRunId = state.taskIdIndex[taskId]; + return taskRunId !== undefined && state.sessions[taskRunId] !== undefined; +} + +let watching = false; + +function watchStoreForEviction(): void { + if (watching) return; + watching = true; + // When the residency system evicts a backgrounded session's events (or the + // session is torn down), drop its derived caches too: they reference the + // evicted events, and keeping them would defeat the memory reclaim. + useSessionStore.subscribe((state) => { + sweepEvicted(buildCaches, state); + sweepEvicted(threadGroupers, state); + }); +} + +function sweepEvicted( + map: Map>, + state: SessionState, +): void { + for (const [mapKey, entry] of map) { + if (!entry.hadSession) continue; + const taskRunId = state.taskIdIndex[entry.taskId]; + const session = + taskRunId === undefined ? undefined : state.sessions[taskRunId]; + if (!session || session.events.length === 0) { + map.delete(mapKey); + } + } +} diff --git a/packages/ui/src/features/sessions/hooks/useConversationItems.test.tsx b/packages/ui/src/features/sessions/hooks/useConversationItems.test.tsx new file mode 100644 index 0000000000..24226e93ce --- /dev/null +++ b/packages/ui/src/features/sessions/hooks/useConversationItems.test.tsx @@ -0,0 +1,81 @@ +import type { AcpMessage } from "@posthog/shared"; +import { renderHook } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; +import { useConversationItems } from "./useConversationItems"; + +function userPromptMsg(ts: number, id: number, text: string): AcpMessage { + return { + type: "acp_message", + ts, + message: { + jsonrpc: "2.0", + id, + method: "session/prompt", + params: { prompt: [{ type: "text", text }] }, + }, + }; +} + +function promptResponseMsg(ts: number, id: number): AcpMessage { + return { + type: "acp_message", + ts, + message: { jsonrpc: "2.0", id, result: { stopReason: "end_turn" } }, + }; +} + +function transcript(): AcpMessage[] { + return [ + userPromptMsg(1, 1, "first prompt"), + promptResponseMsg(2, 1), + userPromptMsg(3, 2, "second prompt"), + ]; +} + +describe("useConversationItems persistence", () => { + it("returns the identical result after a remount with a persist key", () => { + const events = transcript(); + const key = { scope: "test-hook", taskId: "idle-remount" }; + + const first = renderHook(() => + useConversationItems(events, false, undefined, key), + ); + const firstResult = first.result.current; + first.unmount(); + + const second = renderHook(() => + useConversationItems(events, false, undefined, key), + ); + expect(second.result.current).toBe(firstResult); + }); + + it("reuses completed turn items across a remount while streaming", () => { + const events = transcript(); + const key = { scope: "test-hook", taskId: "streaming-remount" }; + + const first = renderHook(() => + useConversationItems(events, true, undefined, key), + ); + const firstItems = first.result.current.items; + first.unmount(); + + // Streaming appends preserve element identity (immer structural sharing); + // a surviving builder must take the append fast path, not rebuild. + const appended = [...events, promptResponseMsg(4, 2)]; + const second = renderHook(() => + useConversationItems(appended, true, undefined, key), + ); + expect(second.result.current.items[0]).toBe(firstItems[0]); + }); + + it("rebuilds from scratch on remount without a persist key", () => { + const events = transcript(); + + const first = renderHook(() => useConversationItems(events, false)); + const firstItems = first.result.current.items; + first.unmount(); + + const second = renderHook(() => useConversationItems(events, false)); + expect(second.result.current.items[0]).not.toBe(firstItems[0]); + }); +}); diff --git a/packages/ui/src/features/sessions/hooks/useConversationItems.ts b/packages/ui/src/features/sessions/hooks/useConversationItems.ts index 0800e1511c..d24c289f4e 100644 --- a/packages/ui/src/features/sessions/hooks/useConversationItems.ts +++ b/packages/ui/src/features/sessions/hooks/useConversationItems.ts @@ -4,39 +4,44 @@ import type { BuildResult, } from "@posthog/ui/features/sessions/components/buildConversationItems"; import { createIncrementalConversationBuilder } from "@posthog/ui/features/sessions/components/incrementalConversationItems"; +import { + type ConversationBuildCache, + type ConversationCacheKey, + getConversationBuildCache, +} from "@posthog/ui/features/sessions/hooks/conversationDerivedCache"; import { useRef } from "react"; -interface Cache { - impl: ReturnType; - events: AcpMessage[] | null; - pending: boolean | null; - debug: boolean | undefined; - result: BuildResult | null; -} - /** - * Builds conversation items incrementally — each event is parsed once and + * Builds conversation items incrementally: each event is parsed once and * completed turns are reused by reference, so a streamed token costs work - * proportional to the active turn rather than the whole thread. The persistent - * builder lives in a ref; results are memoized on the (events, pending, debug) - * triple so unrelated re-renders don't re-derive. + * proportional to the active turn rather than the whole thread. Results are + * memoized on the (events, pending, debug) triple so unrelated re-renders + * don't re-derive. + * + * Without `persistKey` the builder lives in a ref and dies with the component, + * so every remount re-parses the full transcript. Pass a `persistKey` to keep + * it in a module-level cache instead, making re-opening a task cheap. */ export function useConversationItems( events: AcpMessage[], isPromptPending: boolean | null, options?: BuildConversationOptions, + persistKey?: ConversationCacheKey, ): BuildResult { - const ref = useRef(null); - if (!ref.current) { - ref.current = { + const ref = useRef(null); + let cache: ConversationBuildCache; + if (persistKey) { + cache = getConversationBuildCache(persistKey); + } else { + ref.current ??= { impl: createIncrementalConversationBuilder(), events: null, pending: null, debug: undefined, result: null, }; + cache = ref.current; } - const cache = ref.current; const debug = options?.showDebugLogs; if ( From 35a5be01cf401c0157e26f4f00cfed0476182114 Mon Sep 17 00:00:00 2001 From: Arnaud Hillen Date: Thu, 30 Jul 2026 00:32:19 +0200 Subject: [PATCH 2/5] refactor(sessions): apply conversation cache LRU cap per scope The cap counted scope+task entries in one map, so the three opted-in scopes consumed three slots per task and only ~2 tasks stayed fully cached instead of 8. Group entries per scope and cap within each scope. Generated-By: PostHog Code Task-Id: 5d24ea17-aec0-4334-884e-c2639867e260 --- .../hooks/conversationDerivedCache.test.ts | 16 ++++++ .../hooks/conversationDerivedCache.ts | 49 ++++++++++++------- 2 files changed, 47 insertions(+), 18 deletions(-) diff --git a/packages/ui/src/features/sessions/hooks/conversationDerivedCache.test.ts b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.test.ts index 0a354774f8..5e03a033f2 100644 --- a/packages/ui/src/features/sessions/hooks/conversationDerivedCache.test.ts +++ b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.test.ts @@ -56,6 +56,22 @@ describe("conversationDerivedCache", () => { ).not.toBe(first); }); + it("applies the LRU cap per scope, so multi-hook surfaces don't shrink task capacity", () => { + const first = getConversationBuildCache({ + scope: "cap-a", + taskId: "cap-0", + }); + for (let i = 1; i < 8; i++) { + getConversationBuildCache({ scope: "cap-a", taskId: `cap-${i}` }); + } + for (let i = 0; i < 8; i++) { + getConversationBuildCache({ scope: "cap-b", taskId: `cap-${i}` }); + } + expect(getConversationBuildCache({ scope: "cap-a", taskId: "cap-0" })).toBe( + first, + ); + }); + it("drops the cache when the session's events are evicted from the store", () => { const key = { scope: "test", taskId: "evicted-task" }; seedSession(key.taskId, "run-evicted"); diff --git a/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts index 2243bbaf71..d3323e84e7 100644 --- a/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts +++ b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts @@ -44,11 +44,18 @@ interface Entry { * per-component caches die with it and every click re-parses the full * transcript (seconds of main-thread work for large sessions). These * module-level caches keep the incremental builder state alive across mounts. + * + * Entries are grouped per scope with the LRU cap applied inside each scope, + * so surfaces that register several hooks per task (the chat thread and its + * footer) don't shrink how many tasks stay cached. */ const MAX_CACHED_TASKS = 8; -const buildCaches = new Map>(); -const threadGroupers = new Map>(); +const buildCaches = new Map< + string, + Map> +>(); +const threadGroupers = new Map>>(); export function getConversationBuildCache( key: ConversationCacheKey, @@ -69,24 +76,28 @@ export function getPersistentThreadGrouper( } function getEntry( - map: Map>, + map: Map>>, key: ConversationCacheKey, create: () => T, ): T { watchStoreForEviction(); - const mapKey = `${key.scope}:${key.taskId}`; - let entry = map.get(mapKey); + let scopeMap = map.get(key.scope); + if (!scopeMap) { + scopeMap = new Map(); + map.set(key.scope, scopeMap); + } + let entry = scopeMap.get(key.taskId); if (entry) { // Re-insert to refresh LRU recency (Map preserves insertion order). - map.delete(mapKey); + scopeMap.delete(key.taskId); } else { entry = { taskId: key.taskId, hadSession: false, value: create() }; } entry.hadSession ||= hasSession(useSessionStore.getState(), key.taskId); - map.set(mapKey, entry); - for (const oldest of map.keys()) { - if (map.size <= MAX_CACHED_TASKS) break; - map.delete(oldest); + scopeMap.set(key.taskId, entry); + for (const oldest of scopeMap.keys()) { + if (scopeMap.size <= MAX_CACHED_TASKS) break; + scopeMap.delete(oldest); } return entry.value; } @@ -111,16 +122,18 @@ function watchStoreForEviction(): void { } function sweepEvicted( - map: Map>, + map: Map>>, state: SessionState, ): void { - for (const [mapKey, entry] of map) { - if (!entry.hadSession) continue; - const taskRunId = state.taskIdIndex[entry.taskId]; - const session = - taskRunId === undefined ? undefined : state.sessions[taskRunId]; - if (!session || session.events.length === 0) { - map.delete(mapKey); + for (const scopeMap of map.values()) { + for (const [taskId, entry] of scopeMap) { + if (!entry.hadSession) continue; + const taskRunId = state.taskIdIndex[entry.taskId]; + const session = + taskRunId === undefined ? undefined : state.sessions[taskRunId]; + if (!session || session.events.length === 0) { + scopeMap.delete(taskId); + } } } } From 967554b0828204f62b6a0c1c30f3702498633705 Mon Sep 17 00:00:00 2001 From: Arnaud Hillen Date: Thu, 30 Jul 2026 02:43:45 +0200 Subject: [PATCH 3/5] refactor(sessions): use lru-cache for the conversation derived caches lru-cache v11 is already a workspace dependency (harness), so this swaps the hand-rolled Map re-insertion LRU for the battle-tested implementation with identical semantics: get() refreshes recency, max evicts per scope. Generated-By: PostHog Code Task-Id: 5d24ea17-aec0-4334-884e-c2639867e260 --- packages/ui/package.json | 1 + .../hooks/conversationDerivedCache.ts | 55 +++++++++---------- pnpm-lock.yaml | 3 + 3 files changed, 30 insertions(+), 29 deletions(-) diff --git a/packages/ui/package.json b/packages/ui/package.json index 60d7dea2b9..b4db4675d2 100644 --- a/packages/ui/package.json +++ b/packages/ui/package.json @@ -89,6 +89,7 @@ "fuse.js": "^7.1.0", "fzf": "^0.5.2", "inversify": "catalog:", + "lru-cache": "^11.1.0", "lucide-react": "^1.7.0", "posthog-js": "^1.378.0", "radix-themes-tw": "0.2.3", diff --git a/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts index d3323e84e7..65f9440f4d 100644 --- a/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts +++ b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts @@ -6,6 +6,7 @@ import { type SessionState, useSessionStore, } from "@posthog/ui/features/sessions/sessionStore"; +import { LRUCache } from "lru-cache"; export interface ConversationBuildCache { impl: ReturnType; @@ -28,7 +29,6 @@ export interface ConversationCacheKey { } interface Entry { - taskId: string; /** * Entries are swept when their session's events are evicted, but only if * the session was ever seen in the store. Surfaces that render transcripts @@ -51,11 +51,10 @@ interface Entry { */ const MAX_CACHED_TASKS = 8; -const buildCaches = new Map< - string, - Map> ->(); -const threadGroupers = new Map>>(); +type ScopeCache = LRUCache>; + +const buildCaches = new Map>(); +const threadGroupers = new Map>(); export function getConversationBuildCache( key: ConversationCacheKey, @@ -75,30 +74,23 @@ export function getPersistentThreadGrouper( return getEntry(threadGroupers, key, createIncrementalThreadGrouper); } -function getEntry( - map: Map>>, +function getEntry( + map: Map>, key: ConversationCacheKey, create: () => T, ): T { watchStoreForEviction(); - let scopeMap = map.get(key.scope); - if (!scopeMap) { - scopeMap = new Map(); - map.set(key.scope, scopeMap); + let scopeCache = map.get(key.scope); + if (!scopeCache) { + scopeCache = new LRUCache({ max: MAX_CACHED_TASKS }); + map.set(key.scope, scopeCache); } - let entry = scopeMap.get(key.taskId); - if (entry) { - // Re-insert to refresh LRU recency (Map preserves insertion order). - scopeMap.delete(key.taskId); - } else { - entry = { taskId: key.taskId, hadSession: false, value: create() }; + let entry = scopeCache.get(key.taskId); + if (!entry) { + entry = { hadSession: false, value: create() }; + scopeCache.set(key.taskId, entry); } entry.hadSession ||= hasSession(useSessionStore.getState(), key.taskId); - scopeMap.set(key.taskId, entry); - for (const oldest of scopeMap.keys()) { - if (scopeMap.size <= MAX_CACHED_TASKS) break; - scopeMap.delete(oldest); - } return entry.value; } @@ -121,19 +113,24 @@ function watchStoreForEviction(): void { }); } -function sweepEvicted( - map: Map>>, +function sweepEvicted( + map: Map>, state: SessionState, ): void { - for (const scopeMap of map.values()) { - for (const [taskId, entry] of scopeMap) { + for (const scopeCache of map.values()) { + // Collect first: deleting while iterating an LRUCache is not guaranteed safe. + const evicted: string[] = []; + for (const [taskId, entry] of scopeCache.entries()) { if (!entry.hadSession) continue; - const taskRunId = state.taskIdIndex[entry.taskId]; + const taskRunId = state.taskIdIndex[taskId]; const session = taskRunId === undefined ? undefined : state.sessions[taskRunId]; if (!session || session.events.length === 0) { - scopeMap.delete(taskId); + evicted.push(taskId); } } + for (const taskId of evicted) { + scopeCache.delete(taskId); + } } } diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 28baa3b7d4..45de1fd6bc 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1424,6 +1424,9 @@ importers: inversify: specifier: 'catalog:' version: 7.11.0(reflect-metadata@0.2.2) + lru-cache: + specifier: ^11.1.0 + version: 11.2.5 lucide-react: specifier: ^1.7.0 version: 1.7.0(react@19.2.6) From 10d334c22e983291f067314089a2d208d6b367b6 Mon Sep 17 00:00:00 2001 From: Arnaud Hillen Date: Fri, 31 Jul 2026 01:44:45 +0200 Subject: [PATCH 4/5] refactor(sessions): harden the conversation derived cache - Sweep entries by the session run they were built from, latched only once that run's events are store-resident: a re-run's residency no longer keeps a stale entry alive, and sessions with no resident events no longer cause sweep churn - Pin the cache entry for the mounted lifetime so LRU eviction never forces a still-mounted view (e.g. a 3x3 command-center grid, one over the cap) onto a fresh builder every render - Stop persisting ChatThreadFooter's duplicate full parse: AcpChatThread now passes footerState and usage down exactly like ChatThread, halving per-task derived retention and removing the chat-thread-footer scope - Merge the two cache registries into one entry per scope+task, subscribe at module scope, skip sweeps when the sessions slice is unchanged, and add a 30 minute TTL so sessionless (archive) entries reclaim on their own - Type the scope as a closed union and accept an optional taskId in the persist key so call sites drop the ternaries Generated-By: PostHog Code Task-Id: 47cdc119-1c24-42ec-ac36-74c43cdd7f4e --- .../sessions/components/ConversationView.tsx | 19 +- .../components/chat-thread/ChatThread.tsx | 12 +- .../chat-thread/ChatThreadFooter.tsx | 15 +- .../hooks/conversationDerivedCache.test.ts | 120 +++++++++--- .../hooks/conversationDerivedCache.ts | 182 +++++++++++------- .../hooks/useConversationItems.test.tsx | 63 +++++- .../sessions/hooks/useConversationItems.ts | 48 +++-- 7 files changed, 311 insertions(+), 148 deletions(-) diff --git a/packages/ui/src/features/sessions/components/ConversationView.tsx b/packages/ui/src/features/sessions/components/ConversationView.tsx index df12d6cc18..13bceb37e0 100644 --- a/packages/ui/src/features/sessions/components/ConversationView.tsx +++ b/packages/ui/src/features/sessions/components/ConversationView.tsx @@ -36,7 +36,6 @@ import type { ThreadRow, } from "@posthog/ui/features/sessions/components/new-thread/buildThreadGroups"; import type { CollapseMode } from "@posthog/ui/features/sessions/components/new-thread/conversationThreadConfig"; -import { createIncrementalThreadGrouper } from "@posthog/ui/features/sessions/components/new-thread/incrementalThreadGrouping"; import { ToolCallGroupChip } from "@posthog/ui/features/sessions/components/new-thread/ToolCallGroupChip"; import { SessionFooter } from "@posthog/ui/features/sessions/components/SessionFooter"; import { @@ -51,7 +50,7 @@ import { } from "@posthog/ui/features/sessions/components/VirtualizedList"; import { CHAT_CONTENT_MAX_WIDTH } from "@posthog/ui/features/sessions/constants"; import { DIFFS_HIGHLIGHTER_OPTIONS } from "@posthog/ui/features/sessions/diffHighlighterOptions"; -import { getPersistentThreadGrouper } from "@posthog/ui/features/sessions/hooks/conversationDerivedCache"; +import { usePersistentThreadGrouper } from "@posthog/ui/features/sessions/hooks/conversationDerivedCache"; import { useContextUsage } from "@posthog/ui/features/sessions/hooks/useContextUsage"; import { useConversationItems } from "@posthog/ui/features/sessions/hooks/useConversationItems"; import { useConversationSearch } from "@posthog/ui/features/sessions/hooks/useConversationSearch"; @@ -168,7 +167,7 @@ export function ConversationView({ events, isPromptPending, { showDebugLogs }, - taskId ? { scope: "conversation-view", taskId } : undefined, + { scope: "conversation-view", taskId }, ); const firstUserMessageIdRef = useRef(undefined); @@ -210,19 +209,7 @@ export function ConversationView({ // Fold each completed turn's tool-call work into a collapsible chip, and emit // the keepMounted indices (standalone MCP-app rows, whose iframes must survive // scrolling) + the item→row map in the same pass. - const threadGrouperRef = useRef | null>(null); - let threadGrouper: ReturnType; - if (taskId) { - threadGrouper = getPersistentThreadGrouper({ - scope: "conversation-view", - taskId, - }); - } else { - threadGrouperRef.current ??= createIncrementalThreadGrouper(); - threadGrouper = threadGrouperRef.current; - } + const threadGrouper = usePersistentThreadGrouper("conversation-view", taskId); const grouping = useMemo( () => threadGrouper.update(items, collapseMode, groupOverrides), [items, collapseMode, groupOverrides, threadGrouper], 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 31554980fd..ff1c039b74 100644 --- a/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx +++ b/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx @@ -100,6 +100,7 @@ import { } from "@posthog/ui/features/sessions/constants"; import { DIFFS_HIGHLIGHTER_OPTIONS } from "@posthog/ui/features/sessions/diffHighlighterOptions"; import { useAgentConversationItems } from "@posthog/ui/features/sessions/hooks/useAgentConversationItems"; +import { useContextUsage } from "@posthog/ui/features/sessions/hooks/useContextUsage"; import { useConversationItems } from "@posthog/ui/features/sessions/hooks/useConversationItems"; import { useOptimisticItemsForTask, @@ -1086,19 +1087,24 @@ export function ChatThread({ events, ...props }: ChatThreadProps) { export function AcpChatThread({ events, ...props }: AcpChatThreadProps) { const showDebugLogs = useSettingsStore((state) => state.debugLogsCloudRuns); - const { items } = useConversationItems( + const { items, ...footerState } = useConversationItems( events, props.isPromptPending, { showDebugLogs }, - props.taskId ? { scope: "chat-thread", taskId: props.taskId } : undefined, + { scope: "chat-thread", taskId: props.taskId }, ); + // Hand the footer this parse's turn state and the usage derived here, so it + // never runs a second full parse of the same transcript (see ChatThread). + const usage = useContextUsage(events); return ( ); } diff --git a/packages/ui/src/features/sessions/components/chat-thread/ChatThreadFooter.tsx b/packages/ui/src/features/sessions/components/chat-thread/ChatThreadFooter.tsx index a964d0bba4..ecb3d110d5 100644 --- a/packages/ui/src/features/sessions/components/chat-thread/ChatThreadFooter.tsx +++ b/packages/ui/src/features/sessions/components/chat-thread/ChatThreadFooter.tsx @@ -27,9 +27,9 @@ interface ChatThreadFooterProps { * thread, rendered UNDER the composer. The legacy `ConversationView` renders the same * `SessionFooter` at the bottom of the thread instead; here it lives under the input. * - * Re-derives the turn / usage / queue state from `events` with the same hooks `ConversationView` - * uses — `ChatThread` runs its own `useConversationItems`, so this is a second (incremental, - * memoized) parse pass, acceptable for a flag-gated surface. Gated behind + * Both thread variants pass `footerState` from their own parse plus empty + * `footerEvents`, so the `useConversationItems` fallback below only ever sees + * an empty transcript and stays non-persistent. Gated behind * `settingsStore.useNewChatThread` at the call site. */ export function ChatThreadFooter({ @@ -44,12 +44,9 @@ export function ChatThreadFooter({ const showDebugLogs = useSettingsStore((s) => s.debugLogsCloudRuns); const eventContextUsage = useContextUsage(events); const contextUsage = usage === undefined ? eventContextUsage : usage; - const eventFooterState = useConversationItems( - events, - isPromptPending, - { showDebugLogs }, - taskId ? { scope: "chat-thread-footer", taskId } : undefined, - ); + const eventFooterState = useConversationItems(events, isPromptPending, { + showDebugLogs, + }); const lastTurnInfo = footerState?.lastTurnInfo ?? eventFooterState.lastTurnInfo; const isCompacting = diff --git a/packages/ui/src/features/sessions/hooks/conversationDerivedCache.test.ts b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.test.ts index 5e03a033f2..1faeb7a7f1 100644 --- a/packages/ui/src/features/sessions/hooks/conversationDerivedCache.test.ts +++ b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.test.ts @@ -1,9 +1,11 @@ import type { AcpMessage, AgentSession } from "@posthog/shared"; -import { describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it } from "vitest"; import { sessionStoreSetters } from "../sessionStore"; import { + type ConversationCacheKey, getConversationBuildCache, getPersistentThreadGrouper, + MAX_CACHED_TASKS, } from "./conversationDerivedCache"; function msg(ts: number): AcpMessage { @@ -14,18 +16,34 @@ function msg(ts: number): AcpMessage { }; } -function seedSession(taskId: string, taskRunId: string): void { +const seededRuns: string[] = []; + +function seedSession( + taskId: string, + taskRunId: string, + events: AcpMessage[] = [msg(1)], +): void { sessionStoreSetters.setSession({ taskId, taskRunId, - events: [msg(1)], + events, messageQueue: [], } as unknown as AgentSession); + seededRuns.push(taskRunId); } +afterEach(() => { + for (const taskRunId of seededRuns.splice(0)) { + sessionStoreSetters.removeSession(taskRunId); + } +}); + describe("conversationDerivedCache", () => { it("returns the same cache entry across lookups for the same scope + task", () => { - const key = { scope: "test", taskId: "same-entry" }; + const key: ConversationCacheKey = { + scope: "conversation-view", + taskId: "same-entry", + }; expect(getConversationBuildCache(key)).toBe(getConversationBuildCache(key)); expect(getPersistentThreadGrouper(key)).toBe( getPersistentThreadGrouper(key), @@ -34,46 +52,61 @@ describe("conversationDerivedCache", () => { it("keeps entries of different scopes for the same task independent", () => { const taskId = "scoped-entries"; - const a = getConversationBuildCache({ scope: "thread", taskId }); - const b = getConversationBuildCache({ scope: "footer", taskId }); + const a = getConversationBuildCache({ scope: "conversation-view", taskId }); + const b = getConversationBuildCache({ scope: "chat-thread", taskId }); expect(a).not.toBe(b); }); it("evicts the least recently used entry beyond the cap", () => { - const first = getConversationBuildCache({ scope: "lru", taskId: "lru-0" }); - for (let i = 1; i < 50; i++) { - getConversationBuildCache({ scope: "lru", taskId: `lru-${i}` }); + const first = getConversationBuildCache({ + scope: "chat-thread", + taskId: "lru-0", + }); + for (let i = 1; i <= MAX_CACHED_TASKS; i++) { + getConversationBuildCache({ scope: "chat-thread", taskId: `lru-${i}` }); } const recent = getConversationBuildCache({ - scope: "lru", - taskId: "lru-49", + scope: "chat-thread", + taskId: `lru-${MAX_CACHED_TASKS}`, }); - expect(getConversationBuildCache({ scope: "lru", taskId: "lru-49" })).toBe( - recent, - ); expect( - getConversationBuildCache({ scope: "lru", taskId: "lru-0" }), + getConversationBuildCache({ + scope: "chat-thread", + taskId: `lru-${MAX_CACHED_TASKS}`, + }), + ).toBe(recent); + expect( + getConversationBuildCache({ scope: "chat-thread", taskId: "lru-0" }), ).not.toBe(first); }); - it("applies the LRU cap per scope, so multi-hook surfaces don't shrink task capacity", () => { + it("applies the LRU cap per scope, so one scope's fill doesn't shrink another's", () => { const first = getConversationBuildCache({ - scope: "cap-a", + scope: "conversation-view", taskId: "cap-0", }); - for (let i = 1; i < 8; i++) { - getConversationBuildCache({ scope: "cap-a", taskId: `cap-${i}` }); + for (let i = 1; i < MAX_CACHED_TASKS; i++) { + getConversationBuildCache({ + scope: "conversation-view", + taskId: `cap-${i}`, + }); } - for (let i = 0; i < 8; i++) { - getConversationBuildCache({ scope: "cap-b", taskId: `cap-${i}` }); + for (let i = 0; i < MAX_CACHED_TASKS; i++) { + getConversationBuildCache({ scope: "chat-thread", taskId: `cap-${i}` }); } - expect(getConversationBuildCache({ scope: "cap-a", taskId: "cap-0" })).toBe( - first, - ); + expect( + getConversationBuildCache({ + scope: "conversation-view", + taskId: "cap-0", + }), + ).toBe(first); }); it("drops the cache when the session's events are evicted from the store", () => { - const key = { scope: "test", taskId: "evicted-task" }; + const key: ConversationCacheKey = { + scope: "conversation-view", + taskId: "evicted-task", + }; seedSession(key.taskId, "run-evicted"); const before = getConversationBuildCache(key); sessionStoreSetters.evictEvents("run-evicted"); @@ -81,19 +114,52 @@ describe("conversationDerivedCache", () => { }); it("drops the cache when the session is removed from the store", () => { - const key = { scope: "test", taskId: "removed-task" }; + const key: ConversationCacheKey = { + scope: "conversation-view", + taskId: "removed-task", + }; seedSession(key.taskId, "run-removed"); const before = getConversationBuildCache(key); sessionStoreSetters.removeSession("run-removed"); expect(getConversationBuildCache(key)).not.toBe(before); }); + it("drops the cache when the run it was built from is superseded by a re-run", () => { + const key: ConversationCacheKey = { + scope: "conversation-view", + taskId: "rerun-task", + }; + seedSession(key.taskId, "run-old"); + const before = getConversationBuildCache(key); + // A re-run replaces the session under the same taskId; the entry was + // built from the old run's events, so it must not survive on the back of + // the new run's residency. + seedSession(key.taskId, "run-new"); + expect(getConversationBuildCache(key)).not.toBe(before); + }); + it("keeps entries for tasks that never had a session through store commits", () => { - const key = { scope: "test", taskId: "sessionless-task" }; + const key: ConversationCacheKey = { + scope: "conversation-view", + taskId: "sessionless-task", + }; const before = getConversationBuildCache(key); // An unrelated commit must not sweep surfaces rendering without a session // (e.g. archive views), or they would rebuild on every render. seedSession("unrelated-task", "run-unrelated"); expect(getConversationBuildCache(key)).toBe(before); }); + + it("does not tie an entry to a session whose events are not resident yet", () => { + const key: ConversationCacheKey = { + scope: "conversation-view", + taskId: "empty-session-task", + }; + seedSession(key.taskId, "run-empty", []); + const before = getConversationBuildCache(key); + // The sweep predicate is "no resident events"; latching on an empty + // session would delete the entry on the very next commit. + seedSession("other-task", "run-other"); + expect(getConversationBuildCache(key)).toBe(before); + }); }); diff --git a/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts index 65f9440f4d..7fce96bf92 100644 --- a/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts +++ b/packages/ui/src/features/sessions/hooks/conversationDerivedCache.ts @@ -2,11 +2,9 @@ import type { AcpMessage } from "@posthog/shared"; import type { BuildResult } from "@posthog/ui/features/sessions/components/buildConversationItems"; import { createIncrementalConversationBuilder } from "@posthog/ui/features/sessions/components/incrementalConversationItems"; import { createIncrementalThreadGrouper } from "@posthog/ui/features/sessions/components/new-thread/incrementalThreadGrouping"; -import { - type SessionState, - useSessionStore, -} from "@posthog/ui/features/sessions/sessionStore"; +import { useSessionStore } from "@posthog/ui/features/sessions/sessionStore"; import { LRUCache } from "lru-cache"; +import { useRef } from "react"; export interface ConversationBuildCache { impl: ReturnType; @@ -18,113 +16,149 @@ export interface ConversationBuildCache { type ThreadGrouper = ReturnType; +/** + * Call-site namespace. Surfaces that feed different inputs for the same task + * must not share builder state, or they would invalidate each other on every + * render. Closed union on purpose: a colliding scope string would silently + * share one builder between unrelated surfaces. + */ +export type ConversationCacheScope = "conversation-view" | "chat-thread"; + export interface ConversationCacheKey { - /** - * Call-site namespace. Hooks that feed different inputs for the same task - * (e.g. the chat thread vs its footer) must not share builder state, or - * they'd invalidate each other on every render. - */ - scope: string; + scope: ConversationCacheScope; taskId: string; } -interface Entry { +export interface ConversationPersistKey { + scope: ConversationCacheScope; + /** Surfaces rendering without a task fall back to per-component state. */ + taskId: string | undefined; +} + +interface Entry { /** - * Entries are swept when their session's events are evicted, but only if - * the session was ever seen in the store. Surfaces that render transcripts - * without a live session (archive views) would otherwise be swept on every - * store commit and rebuild from scratch each render. + * The session run whose store-resident events the cached state was built + * from. Entries are swept when that run's events leave the store. Entries + * whose events never came from the store stay null (archive surfaces) and + * are reclaimed by the LRU cap and TTL instead, or unrelated store commits + * would sweep them and force a rebuild on every render. */ - hadSession: boolean; - value: T; + taskRunId: string | null; + build?: ConversationBuildCache; + grouper?: ThreadGrouper; } /** * Re-opening a task unmounts and remounts its whole view tree, so * per-component caches die with it and every click re-parses the full - * transcript (seconds of main-thread work for large sessions). These - * module-level caches keep the incremental builder state alive across mounts. + * transcript (seconds of main-thread work for large sessions). This + * module-level cache keeps the incremental builder state alive across mounts. * - * Entries are grouped per scope with the LRU cap applied inside each scope, - * so surfaces that register several hooks per task (the chat thread and its - * footer) don't shrink how many tasks stay cached. + * Bounds: at most MAX_CACHED_TASKS entries per scope, and entries not touched + * for CACHE_TTL_MS expire on their own (the store sweep below cannot see + * surfaces whose transcripts never enter the session store). Mounted + * components pin their entry in a ref (see useConversationItems), so LRU + * eviction only ever costs an unmounted task its warm re-open. */ -const MAX_CACHED_TASKS = 8; - -type ScopeCache = LRUCache>; +export const MAX_CACHED_TASKS = 8; +const CACHE_TTL_MS = 30 * 60_000; -const buildCaches = new Map>(); -const threadGroupers = new Map>(); +const scopes = new Map>(); -export function getConversationBuildCache( - key: ConversationCacheKey, -): ConversationBuildCache { - return getEntry(buildCaches, key, () => ({ +export function createEmptyBuildCache(): ConversationBuildCache { + return { impl: createIncrementalConversationBuilder(), events: null, pending: null, debug: undefined, result: null, - })); + }; +} + +export function getConversationBuildCache( + key: ConversationCacheKey, +): ConversationBuildCache { + const entry = getEntry(key); + entry.build ??= createEmptyBuildCache(); + return entry.build; } export function getPersistentThreadGrouper( key: ConversationCacheKey, ): ThreadGrouper { - return getEntry(threadGroupers, key, createIncrementalThreadGrouper); + const entry = getEntry(key); + entry.grouper ??= createIncrementalThreadGrouper(); + return entry.grouper; } -function getEntry( - map: Map>, - key: ConversationCacheKey, - create: () => T, -): T { - watchStoreForEviction(); - let scopeCache = map.get(key.scope); +/** + * The thread grouper counterpart of useConversationItems' persistence: + * persistent (and pinned for the mounted lifetime) when a taskId is present, + * per-component otherwise. + */ +export function usePersistentThreadGrouper( + scope: ConversationCacheScope, + taskId: string | undefined, +): ThreadGrouper { + const pinnedRef = useRef<{ key: string; grouper: ThreadGrouper } | null>( + null, + ); + const localRef = useRef(null); + if (taskId !== undefined) { + const pinKey = `${scope} ${taskId}`; + if (pinnedRef.current?.key !== pinKey) { + pinnedRef.current = { + key: pinKey, + grouper: getPersistentThreadGrouper({ scope, taskId }), + }; + } + return pinnedRef.current.grouper; + } + localRef.current ??= createIncrementalThreadGrouper(); + return localRef.current; +} + +function getEntry(key: ConversationCacheKey): Entry { + let scopeCache = scopes.get(key.scope); if (!scopeCache) { - scopeCache = new LRUCache({ max: MAX_CACHED_TASKS }); - map.set(key.scope, scopeCache); + scopeCache = new LRUCache({ + max: MAX_CACHED_TASKS, + ttl: CACHE_TTL_MS, + ttlAutopurge: true, + updateAgeOnGet: true, + }); + scopes.set(key.scope, scopeCache); } let entry = scopeCache.get(key.taskId); if (!entry) { - entry = { hadSession: false, value: create() }; + entry = { taskRunId: null }; scopeCache.set(key.taskId, entry); } - entry.hadSession ||= hasSession(useSessionStore.getState(), key.taskId); - return entry.value; -} - -function hasSession(state: SessionState, taskId: string): boolean { - const taskRunId = state.taskIdIndex[taskId]; - return taskRunId !== undefined && state.sessions[taskRunId] !== undefined; -} - -let watching = false; - -function watchStoreForEviction(): void { - if (watching) return; - watching = true; - // When the residency system evicts a backgrounded session's events (or the - // session is torn down), drop its derived caches too: they reference the - // evicted events, and keeping them would defeat the memory reclaim. - useSessionStore.subscribe((state) => { - sweepEvicted(buildCaches, state); - sweepEvicted(threadGroupers, state); - }); + // Latch to the run currently backing the task, but only once its events are + // store-resident: those are what the cached state is built from, so that + // exact run's eviction is the signal to drop the entry. + const state = useSessionStore.getState(); + const taskRunId = state.taskIdIndex[key.taskId]; + if (taskRunId !== undefined) { + const session = state.sessions[taskRunId]; + if (session !== undefined && session.events.length > 0) { + entry.taskRunId = taskRunId; + } + } + return entry; } -function sweepEvicted( - map: Map>, - state: SessionState, -): void { - for (const scopeCache of map.values()) { +// When the residency system evicts a backgrounded session's events (or the +// session is torn down), drop its derived caches too: they reference the +// evicted events, and keeping them would defeat the memory reclaim. +useSessionStore.subscribe((state, prevState) => { + if (state.sessions === prevState.sessions) return; + for (const scopeCache of scopes.values()) { // Collect first: deleting while iterating an LRUCache is not guaranteed safe. const evicted: string[] = []; for (const [taskId, entry] of scopeCache.entries()) { - if (!entry.hadSession) continue; - const taskRunId = state.taskIdIndex[taskId]; - const session = - taskRunId === undefined ? undefined : state.sessions[taskRunId]; + if (entry.taskRunId === null) continue; + const session = state.sessions[entry.taskRunId]; if (!session || session.events.length === 0) { evicted.push(taskId); } @@ -133,4 +167,4 @@ function sweepEvicted( scopeCache.delete(taskId); } } -} +}); diff --git a/packages/ui/src/features/sessions/hooks/useConversationItems.test.tsx b/packages/ui/src/features/sessions/hooks/useConversationItems.test.tsx index 24226e93ce..09da07b890 100644 --- a/packages/ui/src/features/sessions/hooks/useConversationItems.test.tsx +++ b/packages/ui/src/features/sessions/hooks/useConversationItems.test.tsx @@ -1,6 +1,11 @@ import type { AcpMessage } from "@posthog/shared"; import { renderHook } from "@testing-library/react"; import { describe, expect, it } from "vitest"; +import { + type ConversationPersistKey, + getConversationBuildCache, + MAX_CACHED_TASKS, +} from "./conversationDerivedCache"; import { useConversationItems } from "./useConversationItems"; function userPromptMsg(ts: number, id: number, text: string): AcpMessage { @@ -35,7 +40,10 @@ function transcript(): AcpMessage[] { describe("useConversationItems persistence", () => { it("returns the identical result after a remount with a persist key", () => { const events = transcript(); - const key = { scope: "test-hook", taskId: "idle-remount" }; + const key: ConversationPersistKey = { + scope: "chat-thread", + taskId: "idle-remount", + }; const first = renderHook(() => useConversationItems(events, false, undefined, key), @@ -51,7 +59,10 @@ describe("useConversationItems persistence", () => { it("reuses completed turn items across a remount while streaming", () => { const events = transcript(); - const key = { scope: "test-hook", taskId: "streaming-remount" }; + const key: ConversationPersistKey = { + scope: "chat-thread", + taskId: "streaming-remount", + }; const first = renderHook(() => useConversationItems(events, true, undefined, key), @@ -68,6 +79,35 @@ describe("useConversationItems persistence", () => { expect(second.result.current.items[0]).toBe(firstItems[0]); }); + it("keeps a mounted view on its builder when the LRU evicts its entry", () => { + const events = transcript(); + const key: ConversationPersistKey = { + scope: "chat-thread", + taskId: "pinned-mounted", + }; + + const hook = renderHook( + ({ evs }: { evs: AcpMessage[] }) => + useConversationItems(evs, true, undefined, key), + { initialProps: { evs: events } }, + ); + const firstItems = hook.result.current.items; + + // Push the mounted entry out of the scope's LRU (a grid can mount more + // views than the cache holds); the pinned builder must keep streaming + // incrementally instead of rebuilding from scratch every render. + for (let i = 0; i <= MAX_CACHED_TASKS; i++) { + getConversationBuildCache({ + scope: "chat-thread", + taskId: `pin-filler-${i}`, + }); + } + + const appended = [...events, promptResponseMsg(4, 2)]; + hook.rerender({ evs: appended }); + expect(hook.result.current.items[0]).toBe(firstItems[0]); + }); + it("rebuilds from scratch on remount without a persist key", () => { const events = transcript(); @@ -78,4 +118,23 @@ describe("useConversationItems persistence", () => { const second = renderHook(() => useConversationItems(events, false)); expect(second.result.current.items[0]).not.toBe(firstItems[0]); }); + + it("treats a persist key without a taskId as non-persistent", () => { + const events = transcript(); + const key: ConversationPersistKey = { + scope: "chat-thread", + taskId: undefined, + }; + + const first = renderHook(() => + useConversationItems(events, false, undefined, key), + ); + const firstItems = first.result.current.items; + first.unmount(); + + const second = renderHook(() => + useConversationItems(events, false, undefined, key), + ); + expect(second.result.current.items[0]).not.toBe(firstItems[0]); + }); }); diff --git a/packages/ui/src/features/sessions/hooks/useConversationItems.ts b/packages/ui/src/features/sessions/hooks/useConversationItems.ts index d24c289f4e..ae9f786472 100644 --- a/packages/ui/src/features/sessions/hooks/useConversationItems.ts +++ b/packages/ui/src/features/sessions/hooks/useConversationItems.ts @@ -3,10 +3,10 @@ import type { BuildConversationOptions, BuildResult, } from "@posthog/ui/features/sessions/components/buildConversationItems"; -import { createIncrementalConversationBuilder } from "@posthog/ui/features/sessions/components/incrementalConversationItems"; import { type ConversationBuildCache, - type ConversationCacheKey, + type ConversationPersistKey, + createEmptyBuildCache, getConversationBuildCache, } from "@posthog/ui/features/sessions/hooks/conversationDerivedCache"; import { useRef } from "react"; @@ -18,29 +18,43 @@ import { useRef } from "react"; * memoized on the (events, pending, debug) triple so unrelated re-renders * don't re-derive. * - * Without `persistKey` the builder lives in a ref and dies with the component, - * so every remount re-parses the full transcript. Pass a `persistKey` to keep - * it in a module-level cache instead, making re-opening a task cheap. + * Without a `persistKey` (or without a taskId in it) the builder lives in a + * ref and dies with the component, so every remount re-parses the full + * transcript. With one, it lives in a module-level cache instead, making + * re-opening a task cheap. The mounted component pins its cache entry in a + * ref: LRU eviction must never force a still-mounted view (e.g. one cell of a + * grid larger than the cache) onto a fresh builder. */ export function useConversationItems( events: AcpMessage[], isPromptPending: boolean | null, options?: BuildConversationOptions, - persistKey?: ConversationCacheKey, + persistKey?: ConversationPersistKey, ): BuildResult { - const ref = useRef(null); + const pinnedRef = useRef<{ + key: string; + cache: ConversationBuildCache; + } | null>(null); + const localRef = useRef(null); let cache: ConversationBuildCache; - if (persistKey) { - cache = getConversationBuildCache(persistKey); + // Empty transcripts are trivial to rebuild; keeping them out of the cache + // stops surfaces that render before events arrive (or never get any) from + // occupying its slots. + if (persistKey?.taskId !== undefined && events.length > 0) { + const pinKey = `${persistKey.scope} ${persistKey.taskId}`; + if (pinnedRef.current?.key !== pinKey) { + pinnedRef.current = { + key: pinKey, + cache: getConversationBuildCache({ + scope: persistKey.scope, + taskId: persistKey.taskId, + }), + }; + } + cache = pinnedRef.current.cache; } else { - ref.current ??= { - impl: createIncrementalConversationBuilder(), - events: null, - pending: null, - debug: undefined, - result: null, - }; - cache = ref.current; + localRef.current ??= createEmptyBuildCache(); + cache = localRef.current; } const debug = options?.showDebugLogs; From 39b15a0146f77f5a7ea12000642d0eca58182695 Mon Sep 17 00:00:00 2001 From: Arnaud Hillen Date: Fri, 31 Jul 2026 02:45:10 +0200 Subject: [PATCH 5/5] bench(sessions): add a reproducible re-open benchmark for the cache Not run in CI (test globs only match *.test.*). Numbers land in the PR description; rerun with pnpm vitest bench in packages/ui. Generated-By: PostHog Code Task-Id: 47cdc119-1c24-42ec-ac36-74c43cdd7f4e --- .../hooks/useConversationItems.bench.ts | 109 ++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 packages/ui/src/features/sessions/hooks/useConversationItems.bench.ts diff --git a/packages/ui/src/features/sessions/hooks/useConversationItems.bench.ts b/packages/ui/src/features/sessions/hooks/useConversationItems.bench.ts new file mode 100644 index 0000000000..c6bf403633 --- /dev/null +++ b/packages/ui/src/features/sessions/hooks/useConversationItems.bench.ts @@ -0,0 +1,109 @@ +import type { AcpMessage } from "@posthog/shared"; +import { renderHook } from "@testing-library/react"; +import { bench, describe } from "vitest"; +import type { ConversationPersistKey } from "./conversationDerivedCache"; +import { useConversationItems } from "./useConversationItems"; + +// Not run in CI (test globs only match *.test.*). Reproduce with: +// cd packages/ui && pnpm vitest bench src/features/sessions/hooks/useConversationItems.bench.ts + +function updateMsg(ts: number, update: unknown): AcpMessage { + return { + type: "acp_message", + ts, + message: { jsonrpc: "2.0", method: "session/update", params: { update } }, + }; +} + +/** + * A transcript shaped like a real agent session: per turn one user prompt, + * a few tool calls with completion updates, a stream of assistant chunks, + * and the prompt response. 30 events per turn. + */ +function makeTranscript(turns: number): AcpMessage[] { + const events: AcpMessage[] = []; + let ts = 0; + const chunk = + "Reading the session service to trace how eviction schedules interact " + + "with the residency grace period and the rehydration path. "; + for (let t = 0; t < turns; t++) { + events.push({ + type: "acp_message", + ts: ts++, + message: { + jsonrpc: "2.0", + id: t + 1, + method: "session/prompt", + params: { prompt: [{ type: "text", text: `prompt ${t}: ${chunk}` }] }, + }, + }); + for (let k = 0; k < 4; k++) { + const toolCallId = `turn${t}-tool${k}`; + events.push( + updateMsg(ts++, { + sessionUpdate: "tool_call", + toolCallId, + kind: "execute", + status: "pending", + title: `Run step ${k} of turn ${t}`, + rawInput: { + command: `pnpm vitest run suite-${t}-${k}`, + cwd: "/repo", + }, + }), + ); + events.push( + updateMsg(ts++, { + sessionUpdate: "tool_call_update", + toolCallId, + status: "completed", + rawOutput: chunk.repeat(8), + }), + ); + } + for (let c = 0; c < 20; c++) { + events.push( + updateMsg(ts++, { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text: chunk }, + }), + ); + } + events.push({ + type: "acp_message", + ts: ts++, + message: { + jsonrpc: "2.0", + id: t + 1, + result: { stopReason: "end_turn" }, + }, + }); + } + return events; +} + +for (const turns of [100, 1000]) { + const events = makeTranscript(turns); + describe(`task re-open, ${events.length} events`, () => { + const key: ConversationPersistKey = { + scope: "chat-thread", + taskId: `bench-${turns}`, + }; + + // Pre-PR behavior: the builder dies with the component, so every re-open + // re-parses the whole transcript. + bench("without persist key (every re-open re-parses)", () => { + const hook = renderHook(() => useConversationItems(events, false)); + hook.unmount(); + }); + + // This PR: the first open parses, every later re-open hits the module + // cache and returns the memoized result identity. + bench("with persist key (warm re-open)", () => { + const hook = renderHook(() => + useConversationItems(events, false, undefined, key), + ); + hook.unmount(); + }); + }); +}