refactor(agent-status): subscribe to live-agent derivations with epoch-keyed selectors instead of getState in render (#39329)

Co-authored-by: OrcaWin <293788423+OrcaWin@users.noreply.github.com>

src/renderer/src/components/activity/useActivityUnreadCount.ts

@@ -1,8 +1,8 @@
-import { useMemo } from 'react'
-import { useShallow } from 'zustand/react/shallow'
+import { useState } from 'react'
 
 import { migrationUnsupportedToAgentStatusEntry } from '../../../../shared/migration-unsupported-agent-entry'
 import { useAppStore } from '@/store'
+import { createAgentStatusEpochSelector } from '@/lib/agent-status-epoch-selector'
 import type { AppState } from '@/store/types'
 import type { AgentStatusEntry } from '../../../../shared/agent-status/contracts/agent-status-types'
 
@@ -69,38 +69,34 @@ export function countActivityUnread(source: ActivityUnreadCountSource, now = Dat
   return count
 }
 
-export function useActivityUnreadCount(): number {
-  const {
-    agentStatusEpoch,
-    migrationUnsupportedByPtyId,
-    retainedAgentsByPaneKey,
-    acknowledgedAgentsByPaneKey,
-    activityClearedAtByPaneKey
-  } = useAppStore(
-    useShallow((state) => ({
-      // Same-turn heartbeats preserve receipts; reducer/freshness epoch changes cover turn/state changes, stale revival and expiry.
-      agentStatusEpoch: state.agentStatusEpoch,
-      migrationUnsupportedByPtyId: state.migrationUnsupportedByPtyId,
-      retainedAgentsByPaneKey: state.retainedAgentsByPaneKey,
-      acknowledgedAgentsByPaneKey: state.acknowledgedAgentsByPaneKey,
-      activityClearedAtByPaneKey: state.activityClearedAtByPaneKey
-    }))
+// Same-turn heartbeats preserve receipts; reducer/freshness epoch changes cover turn/state changes, stale revival and expiry.
+const createActivityUnreadCountSelector = () =>
+  createAgentStatusEpochSelector(
+    (
+      state: AppState,
+      migrationUnsupportedByPtyId: AppState['migrationUnsupportedByPtyId'],
+      retainedAgentsByPaneKey: AppState['retainedAgentsByPaneKey'],
+      acknowledgedAgentsByPaneKey: AppState['acknowledgedAgentsByPaneKey'],
+      activityClearedAtByPaneKey: AppState['activityClearedAtByPaneKey']
+    ) =>
+      countActivityUnread({
+        agentStatusByPaneKey: state.agentStatusByPaneKey,
+        migrationUnsupportedByPtyId,
+        retainedAgentsByPaneKey,
+        acknowledgedAgentsByPaneKey,
+        activityClearedAtByPaneKey
+      })
   )
 
-  return useMemo(() => {
-    void agentStatusEpoch
-    return countActivityUnread({
-      agentStatusByPaneKey: useAppStore.getState().agentStatusByPaneKey,
-      migrationUnsupportedByPtyId,
-      retainedAgentsByPaneKey,
-      acknowledgedAgentsByPaneKey,
-      activityClearedAtByPaneKey
-    })
-  }, [
-    acknowledgedAgentsByPaneKey,
-    activityClearedAtByPaneKey,
-    migrationUnsupportedByPtyId,
-    retainedAgentsByPaneKey,
-    agentStatusEpoch
-  ])
+export function useActivityUnreadCount(): number {
+  const [selectActivityUnreadCount] = useState(createActivityUnreadCountSelector)
+  return useAppStore((state) =>
+    selectActivityUnreadCount(
+      state,
+      state.migrationUnsupportedByPtyId,
+      state.retainedAgentsByPaneKey,
+      state.acknowledgedAgentsByPaneKey,
+      state.activityClearedAtByPaneKey
+    )
+  )
 }

src/renderer/src/components/sidebar/workspace-board/surface/use-visible-workspace-kanban-worktree-ids.ts

@@ -1,11 +1,11 @@
-import { useMemo } from 'react'
+import { useMemo, useState } from 'react'
 import { useAppStore } from '@/store'
 import { getAgentStatusEpochNow } from '@/lib/agent-status-epoch-clock'
 import type { Repo } from '../../../../../../shared/repo-types'
 import type { Worktree } from '../../../../../../shared/worktree/types'
 import '../../worktree-list/listing/visible-worktrees'
 import { computeVisibleWorktrees } from '../../worktree-list/listing/compute-visible-worktrees'
-import { getWorktreeIdsWithLiveAgent } from '@/lib/worktree-activity-state'
+import { createLiveAgentWorktreeIdsSelector } from '@/lib/worktree-activity-state'
 import { getSettingsFocusedExecutionHostId } from '../../../../../../shared/execution-host'
 import type { AppState } from '@/store/types'
 import {
@@ -64,17 +64,13 @@ export function useVisibleWorkspaceKanbanWorktreeIds({
   const agentStatusEpoch = useAppStore((s) => (!showSleepingWorkspaces ? s.agentStatusEpoch : 0))
   // Skip the shared clock for opt-out boards so they cannot evict the live-board sample.
   const agentStatusNow = showSleepingWorkspaces ? 0 : getAgentStatusEpochNow(agentStatusEpoch)
-  // Key scans on epoch too: two ticks can share a millisecond sample but still change membership.
-  const worktreeIdsWithLiveAgent = useMemo(() => {
-    void agentStatusEpoch
-    return !showSleepingWorkspaces
-      ? getWorktreeIdsWithLiveAgent(
-          useAppStore.getState().agentStatusByPaneKey,
-          tabsByWorktree,
-          agentStatusNow
-        )
-      : EMPTY_WORKTREE_ID_SET
-  }, [agentStatusEpoch, agentStatusNow, showSleepingWorkspaces, tabsByWorktree])
+  // Recompute on epoch, not every hook ping; key on the epoch too since ticks can share a millisecond sample.
+  const [selectLiveAgentWorktreeIds] = useState(createLiveAgentWorktreeIdsSelector)
+  const worktreeIdsWithLiveAgent = useAppStore((s) =>
+    showSleepingWorkspaces
+      ? EMPTY_WORKTREE_ID_SET
+      : selectLiveAgentWorktreeIds(s, tabsByWorktree, agentStatusNow)
+  )
 
   return useMemo(() => {
     // Board status ordering differs, but workspace visibility must use the sidebar’s filters.

src/renderer/src/components/sidebar/worktree-list/listing/use-visible-worktrees.ts

@@ -1,10 +1,10 @@
 import { selectSidebarLineageParentIdentities } from '@/store/cross-store-lineage-parents'
-import { useMemo } from 'react'
+import { useMemo, useState } from 'react'
 import { shallow } from 'zustand/shallow'
 import type { AppState } from '@/store/types'
 import { useAppStore } from '@/store'
 import { getAgentStatusEpochNow } from '@/lib/agent-status-epoch-clock'
-import { getWorktreeIdsWithLiveAgent } from '@/lib/worktree-activity-state'
+import { createLiveAgentWorktreeIdsSelector } from '@/lib/worktree-activity-state'
 import type { Repo } from '../../../../../../shared/repo-types'
 import type { WorktreeLineage } from '../../../../../../shared/worktree/lineage-types'
 import type { ExecutionHostId } from '../../../../../../shared/execution-host'
@@ -101,18 +101,13 @@ export function useVisibleSidebarWorktrees(args: {
     )
   )
 
-  // Why snapshot on agentStatusEpoch: update membership immediately without repainting on every hook ping.
-  const worktreeIdsWithLiveAgent = useMemo(() => {
-    // Epochs can share a millisecond sample; key on the epoch itself so membership still refreshes.
-    void agentStatusEpoch
-    return showSleepingWorkspaces
+  // Recompute on epoch, not every hook ping; key on the epoch too since ticks can share a millisecond sample.
+  const [selectLiveAgentWorktreeIds] = useState(createLiveAgentWorktreeIdsSelector)
+  const worktreeIdsWithLiveAgent = useAppStore((s) =>
+    showSleepingWorkspaces
       ? EMPTY_WORKTREE_ID_SET
-      : getWorktreeIdsWithLiveAgent(
-          useAppStore.getState().agentStatusByPaneKey,
-          tabsByWorktree,
-          agentStatusNow
-        )
-  }, [agentStatusEpoch, agentStatusNow, showSleepingWorkspaces, tabsByWorktree])
+      : selectLiveAgentWorktreeIds(s, tabsByWorktree, agentStatusNow)
+  )
   const sleepingSweepExemptWorktreeIds = useSleepingSweepRetention({
     enabled: !showSleepingWorkspaces,
     worktreesByRepo,

src/renderer/src/lib/agent-status-epoch-selector.ts

@@ -0,0 +1,31 @@
+import type { AppState } from '@/store/types'
+
+type AgentStatusEpochState = Pick<AppState, 'agentStatusEpoch'>
+
+/**
+ * A store selector over live agent status that recomputes only when `agentStatusEpoch` or one of
+ * `inputs` changes. Same-epoch writes change status metadata only, so subscribers skip
+ * heartbeat-only updates while the result still reads the store's current state.
+ */
+export function createAgentStatusEpochSelector<
+  State extends AgentStatusEpochState,
+  Inputs extends readonly unknown[],
+  Result
+>(
+  compute: (state: State, ...inputs: Inputs) => Result
+): (state: State, ...inputs: Inputs) => Result {
+  let cached: { epoch: number; inputs: Inputs; result: Result } | null = null
+  return (state, ...inputs) => {
+    if (
+      cached &&
+      cached.epoch === state.agentStatusEpoch &&
+      cached.inputs.length === inputs.length &&
+      cached.inputs.every((input, index) => Object.is(input, inputs[index]))
+    ) {
+      return cached.result
+    }
+    const result = compute(state, ...inputs)
+    cached = { epoch: state.agentStatusEpoch, inputs, result }
+    return result
+  }
+}

src/renderer/src/lib/worktree-activity-state.ts

@@ -1,3 +1,5 @@
+import type { AppState } from '@/store/types'
+import { createAgentStatusEpochSelector } from './agent-status-epoch-selector'
 import { tabHasLivePty } from '../../../shared/tab-has-live-pty'
 import type { TerminalTab } from '../../../shared/workspace-layout/terminal-tab-types'
 import { isFreshNonDoneAgentStatus } from '../../../shared/agent-status/contracts/agent-status-freshness'
@@ -28,6 +30,16 @@ export function getWorktreeIdsWithLiveAgent(
   return new Set(getLiveAgentStatusByWorktreeId(agentStatusByPaneKey, tabsByWorktree, now).keys())
 }
 
+/** Live-agent worktree ids for a store subscription; recomputes on an agent-status epoch or input change. */
+export const createLiveAgentWorktreeIdsSelector = () =>
+  createAgentStatusEpochSelector(
+    (
+      state: Pick<AppState, 'agentStatusEpoch' | 'agentStatusByPaneKey'>,
+      tabsByWorktree: TabsByWorktree | null | undefined,
+      now: number
+    ) => getWorktreeIdsWithLiveAgent(state.agentStatusByPaneKey, tabsByWorktree, now)
+  )
+
 export function getLiveAgentStatusByWorktreeId(
   agentStatusByPaneKey: Record<string, AgentStatusEntry> | null | undefined,
   tabsByWorktree: TabsByWorktree | null | undefined,