Merge pull request #82 from tyulyukov/feature/timeline-cache

perf(timeline): cache derived timeline data

Author: tyulyukov

apps/web/src/components/ChatView.tsx

@@ -102,7 +102,7 @@ import {
 import {
   type AppState,
   selectProjectsAcrossEnvironments,
-  selectThreadsAcrossEnvironments,
+  selectThreadShellsAcrossEnvironments,
   useStore,
 } from "../store";
 import {
@@ -114,6 +114,7 @@ import {
   useThreadById,
 } from "../storeSelectors";
 import { useUiStateStore } from "../uiStateStore";
+import { getThreadFromEnvironmentState } from "../threadDerivation";
 import {
   buildPlanImplementationThreadTitle,
   buildPlanImplementationPrompt,
@@ -365,10 +366,7 @@ function useThreadPlanCatalog(threadIds: readonly ThreadId[]): ThreadPlanCatalog
     let previousEntries: ThreadPlanCatalogEntry[] = [];
 
     return (state: AppState): ThreadPlanCatalogEntry[] => {
-      const allThreads = selectThreadsAcrossEnvironments(state);
-      const nextThreads = threadIds.map((threadId) =>
-        allThreads.find((thread: Thread) => thread.id === threadId),
-      );
+      const nextThreads = threadIds.map((threadId) => getThreadAcrossEnvironments(state, threadId));
       const cachedThreads = previousThreads;
       if (
         cachedThreads &&
@@ -389,6 +387,37 @@ function useThreadPlanCatalog(threadIds: readonly ThreadId[]): ThreadPlanCatalog
   return useStore(selector);
 }
 
+function getThreadAcrossEnvironments(state: AppState, threadId: ThreadId): Thread | undefined {
+  for (const environmentState of Object.values(state.environmentStateById)) {
+    if (!environmentState.threadShellById[threadId]) {
+      continue;
+    }
+    return getThreadFromEnvironmentState(environmentState, threadId);
+  }
+  return undefined;
+}
+
+function selectRateLimitActivitiesAcrossEnvironments(
+  state: AppState,
+): OrchestrationThreadActivity[] {
+  const activities: OrchestrationThreadActivity[] = [];
+  for (const environmentState of Object.values(state.environmentStateById)) {
+    for (const [threadId, activityIds] of Object.entries(
+      environmentState.activityIdsByThreadId,
+    ) as Array<[ThreadId, string[]]>) {
+      const activityById = environmentState.activityByThreadId[threadId];
+      if (!activityById) continue;
+      for (const activityId of activityIds) {
+        const activity = activityById[activityId];
+        if (activity?.kind === "account.rate-limits.updated") {
+          activities.push(activity);
+        }
+      }
+    }
+  }
+  return activities;
+}
+
 function formatOutgoingPrompt(params: {
   provider: ProviderKind;
   model: string | null;
@@ -987,7 +1016,7 @@ export default function ChatView({
   );
   const serverThreadIds = useStore(
     useShallow((state: AppState) =>
-      selectThreadsAcrossEnvironments(state).map((thread) => thread.id),
+      selectThreadShellsAcrossEnvironments(state).map((thread) => thread.id),
     ),
   );
   const draftThreadsByThreadKey = useComposerDraftStore((store) => store.draftThreadsByThreadKey);
@@ -1133,19 +1162,7 @@ export default function ChatView({
     () => deriveLatestContextWindowSnapshot(activeThread?.activities ?? []),
     [activeThread?.activities],
   );
-  const providerUsageActivities = useStore(
-    useShallow((state: AppState) => {
-      const activities: OrchestrationThreadActivity[] = [];
-      for (const thread of selectThreadsAcrossEnvironments(state)) {
-        for (const activity of thread.activities) {
-          if (activity.kind === "account.rate-limits.updated") {
-            activities.push(activity);
-          }
-        }
-      }
-      return activities;
-    }),
-  );
+  const providerUsageActivities = useStore(useShallow(selectRateLimitActivitiesAcrossEnvironments));
   useEffect(() => {
     setMountedTerminalThreadIds((currentThreadIds) => {
       const nextThreadIds = reconcileMountedTerminalThreadIds({
@@ -2051,8 +2068,8 @@ export default function ChatView({
       if (!targetThreadId) return;
       const nextError = sanitizeThreadErrorMessage(error);
       if (
-        selectThreadsAcrossEnvironments(useStore.getState()).some(
-          (thread: Thread) => thread.id === targetThreadId,
+        selectThreadShellsAcrossEnvironments(useStore.getState()).some(
+          (thread) => thread.id === targetThreadId,
         )
       ) {
         setStoreThreadError(targetThreadId, nextError);
@@ -4888,7 +4905,6 @@ export default function ChatView({
               onClickCapture={onMessagesClickCapture}
             >
               <MessagesTimeline
-                key={activeThread.id}
                 threadId={activeThread.id}
                 provider={activeTimelineProvider}
                 hasMessages={timelineEntries.length > 0}

apps/web/src/components/chat/MessagesTimeline.logic.ts

@@ -247,6 +247,12 @@ export function deriveMessagesTimelineRows(input: {
   isWorking: boolean;
   activeTurnStartedAt: string | null;
 }): MessagesTimelineRow[] {
+  const cacheKey = `${input.completionDividerBeforeEntryId ?? ""}\x1f${input.isWorking ? "1" : "0"}\x1f${input.activeTurnStartedAt ?? ""}`;
+  const cachedByKey = timelineRowsCache.get(input.timelineEntries);
+  const cached = cachedByKey?.get(cacheKey);
+  if (cached) {
+    return cached;
+  }
   const nextRows: MessagesTimelineRow[] = [];
   const durationStartByMessageId = computeMessageDurationStart(
     input.timelineEntries.flatMap((entry) => (entry.kind === "message" ? [entry.message] : [])),
@@ -483,9 +489,19 @@ export function deriveMessagesTimelineRows(input: {
     });
   }
 
+  const nextCachedByKey = cachedByKey ?? new Map<string, MessagesTimelineRow[]>();
+  nextCachedByKey.set(cacheKey, nextRows);
+  if (!cachedByKey) {
+    timelineRowsCache.set(input.timelineEntries, nextCachedByKey);
+  }
   return nextRows;
 }
 
+const timelineRowsCache = new WeakMap<
+  ReadonlyArray<TimelineEntry>,
+  Map<string, MessagesTimelineRow[]>
+>();
+
 export function estimateMessagesTimelineRowHeight(
   row: MessagesTimelineRow,
   input: {

apps/web/src/components/chat/MessagesTimeline.tsx

@@ -349,7 +349,15 @@ export const MessagesTimeline = memo(function MessagesTimeline({
   }, [rows, threadId, isHydrating]);
 
   const hasFinishedInitialScrollRef = useRef(false);
+  const hasAutoScrolledOnMountRef = useRef(false);
   const isAtEndRef = useRef(true);
+  const initialScrollThreadIdRef = useRef(threadId);
+  if (initialScrollThreadIdRef.current !== threadId) {
+    initialScrollThreadIdRef.current = threadId;
+    hasFinishedInitialScrollRef.current = false;
+    hasAutoScrolledOnMountRef.current = false;
+    isAtEndRef.current = true;
+  }
   const updateIsAtEnd = useCallback(
     (isAtEnd: boolean) => {
       isAtEndRef.current = isAtEnd;
@@ -368,7 +376,6 @@ export const MessagesTimeline = memo(function MessagesTimeline({
     syncIsAtEnd();
   }, [syncIsAtEnd]);
 
-  const hasAutoScrolledOnMountRef = useRef(false);
   useEffect(() => {
     if (rows.length === 0) return;
     if (hasAutoScrolledOnMountRef.current) return;
@@ -557,6 +564,7 @@ export const MessagesTimeline = memo(function MessagesTimeline({
     <TimelineRowCtx.Provider value={sharedState}>
       <div className="@container/chat relative h-full">
         <LegendList<MessagesTimelineRow>
+          key={threadId}
           ref={listRef}
           data={rows}
           keyExtractor={keyExtractor}

apps/web/src/session-logic.ts

@@ -578,7 +578,16 @@ export function deriveWorkLogEntries(
 ): WorkLogEntry[] {
   const isSessionRunning = options?.isSessionRunning ?? false;
   const provider = options?.provider;
-  const ordered = [...activities].toSorted(compareActivitiesByOrder);
+  const cacheKey = `${latestTurnId ?? ""}\x1f${isSessionRunning ? "1" : "0"}\x1f${provider ?? ""}`;
+  const cachedByKey = workLogEntriesCache.get(activities);
+  const cached = cachedByKey?.get(cacheKey);
+  if (cached) {
+    return cached;
+  }
+  const relevantActivities = latestTurnId
+    ? activities.filter((activity) => activity.turnId === latestTurnId)
+    : activities;
+  const ordered = [...relevantActivities].toSorted(compareActivitiesByOrder);
 
   const previousPlanStepsByActivityId = new Map<string, PlanStep[] | null>();
   let runningPreviousPlanSteps: PlanStep[] | null = null;
@@ -599,7 +608,6 @@ export function deriveWorkLogEntries(
   const collabToolDataUnkeyed: SubagentCollabToolData[] = [];
   const collabToolItemIds = new Set<string>();
   for (const activity of ordered) {
-    if (latestTurnId && activity.turnId !== latestTurnId) continue;
     if (activity.kind === "tool.started") {
       const payload = asRecord(activity.payload);
       if (payload?.itemType === "collab_agent_tool_call") {
@@ -628,7 +636,6 @@ export function deriveWorkLogEntries(
   const nestedTaskIds = new Set<string>();
   for (const activity of ordered) {
     if (activity.kind !== "task.started") continue;
-    if (latestTurnId && activity.turnId !== latestTurnId) continue;
     const payload = asRecord(activity.payload);
     const taskId = asTrimmedString(payload?.taskId);
     if (!taskId) continue;
@@ -641,7 +648,6 @@ export function deriveWorkLogEntries(
   }
 
   const filtered = ordered
-    .filter((activity) => (latestTurnId ? activity.turnId === latestTurnId : true))
     .filter(
       (activity) =>
         activity.kind !== "tool.started" ||
@@ -756,11 +762,22 @@ export function deriveWorkLogEntries(
     entries.push(toDerivedWorkLogEntry(activity, previousPlanSteps));
   }
 
-  return deduplicateToolLifecycleEntries(collapseDerivedWorkLogEntries(entries)).map(
+  const result = deduplicateToolLifecycleEntries(collapseDerivedWorkLogEntries(entries)).map(
     ({ activityKind: _activityKind, collapseKey: _collapseKey, ...entry }) => entry,
   );
+  const nextCachedByKey = cachedByKey ?? new Map<string, WorkLogEntry[]>();
+  nextCachedByKey.set(cacheKey, result);
+  if (!cachedByKey) {
+    workLogEntriesCache.set(activities, nextCachedByKey);
+  }
+  return result;
 }
 
+const workLogEntriesCache = new WeakMap<
+  ReadonlyArray<OrchestrationThreadActivity>,
+  Map<string, WorkLogEntry[]>
+>();
+
 function isPlanBoundaryToolActivity(activity: OrchestrationThreadActivity): boolean {
   if (activity.kind !== "tool.updated" && activity.kind !== "tool.completed") {
     return false;
@@ -1066,6 +1083,12 @@ function toDerivedWorkLogEntry(
   activity: OrchestrationThreadActivity,
   previousPlanSteps: PlanStep[] | null = null,
 ): DerivedWorkLogEntry {
+  if (previousPlanSteps === null) {
+    const cached = derivedWorkLogEntryCache.get(activity);
+    if (cached) {
+      return cached;
+    }
+  }
   const payload =
     activity.payload && typeof activity.payload === "object"
       ? (activity.payload as Record<string, unknown>)
@@ -1189,9 +1212,14 @@ function toDerivedWorkLogEntry(
   if (collapseKey) {
     entry.collapseKey = collapseKey;
   }
+  if (previousPlanSteps === null) {
+    derivedWorkLogEntryCache.set(activity, entry);
+  }
   return entry;
 }
 
+const derivedWorkLogEntryCache = new WeakMap<OrchestrationThreadActivity, DerivedWorkLogEntry>();
+
 function collapseDerivedWorkLogEntries(
   entries: ReadonlyArray<DerivedWorkLogEntry>,
 ): DerivedWorkLogEntry[] {
@@ -1963,6 +1991,12 @@ export function deriveTimelineEntries(
   proposedPlans: ProposedPlan[],
   workEntries: WorkLogEntry[],
 ): TimelineEntry[] {
+  const cachedByProposedPlans = timelineEntriesCache.get(messages);
+  const cachedByWorkEntries = cachedByProposedPlans?.get(proposedPlans);
+  const cached = cachedByWorkEntries?.get(workEntries);
+  if (cached) {
+    return cached;
+  }
   const messageRows: TimelineEntry[] = messages.map((message) => ({
     id: message.id,
     kind: "message",
@@ -1981,11 +2015,29 @@ export function deriveTimelineEntries(
     createdAt: entry.createdAt,
     entry,
   }));
-  return [...messageRows, ...proposedPlanRows, ...workRows].toSorted((a, b) =>
+  const result = [...messageRows, ...proposedPlanRows, ...workRows].toSorted((a, b) =>
     a.createdAt.localeCompare(b.createdAt),
   );
+  const nextCachedByProposedPlans =
+    cachedByProposedPlans ??
+    new WeakMap<ProposedPlan[], WeakMap<WorkLogEntry[], TimelineEntry[]>>();
+  const nextCachedByWorkEntries =
+    cachedByWorkEntries ?? new WeakMap<WorkLogEntry[], TimelineEntry[]>();
+  nextCachedByWorkEntries.set(workEntries, result);
+  if (!cachedByWorkEntries) {
+    nextCachedByProposedPlans.set(proposedPlans, nextCachedByWorkEntries);
+  }
+  if (!cachedByProposedPlans) {
+    timelineEntriesCache.set(messages, nextCachedByProposedPlans);
+  }
+  return result;
 }
 
+const timelineEntriesCache = new WeakMap<
+  ChatMessage[],
+  WeakMap<ProposedPlan[], WeakMap<WorkLogEntry[], TimelineEntry[]>>
+>();
+
 export function inferCheckpointTurnCountByTurnId(
   summaries: TurnDiffSummary[],
 ): Record<TurnId, number> {