Merge pull request #67195 from NousResearch/bb/desktop-perf-p2

perf(desktop): scope tool-diff subscriptions + narrow profile query invalidation

apps/desktop/src/components/assistant-ui/tool/fallback.tsx

@@ -39,7 +39,7 @@ import { useEnterAnimation } from '@/lib/use-enter-animation'
 import { cn } from '@/lib/utils'
 import { recordPreviewArtifact } from '@/store/preview-status'
 import { $activeSessionId, $currentCwd } from '@/store/session'
-import { $toolInlineDiffs } from '@/store/tool-diffs'
+import { $toolInlineDiff } from '@/store/tool-diffs'
 import { $toolRowDismissed, dismissToolRow } from '@/store/tool-dismiss'
 import { $toolDisclosureOpen, $toolViewMode, setToolDisclosureOpen } from '@/store/tool-view'
 
@@ -283,8 +283,9 @@ function ToolEntry({ part }: ToolEntryProps) {
   const disclosureId = `tool-entry:${messageId}:${toolPartDisclosureId(stablePart)}`
   const dismissed = useStore($toolRowDismissed(disclosureId))
   const isPending = messageRunning && result === undefined
-  const liveDiffs = useStore($toolInlineDiffs)
-  const sideDiff = toolCallId ? liveDiffs[toolCallId] || '' : ''
+  // Subscribe to this tool's diff only, so a live patch for one tool doesn't
+  // re-render every mounted tool row (the factory caches a per-id atom).
+  const sideDiff = useStore($toolInlineDiff(toolCallId ?? ''))
   const inlineDiff = stripInlineDiffChrome(sideDiff) || inlineDiffFromResult(result)
   const isFileEdit = isFileEditTool(toolName)
   const defaultOpen = Boolean(inlineDiff)

apps/desktop/src/lib/query-client.test.ts

@@ -0,0 +1,58 @@
+import { beforeEach, describe, expect, it } from 'vitest'
+
+import { invalidateProfileScopedQueries, queryClient } from './query-client'
+
+function invalidated(key: unknown[]): boolean {
+  return queryClient.getQueryState(key)?.isInvalidated ?? false
+}
+
+describe('invalidateProfileScopedQueries', () => {
+  beforeEach(() => {
+    queryClient.clear()
+  })
+
+  it('invalidates profile-scoped caches and leaves account/global caches intact', () => {
+    const profileScoped = [
+      ['hermes-config-record'],
+      ['hermes-config-schema'],
+      ['skills-list'],
+      ['toolsets-list'],
+      ['model-options', 'global'],
+      ['command-palette', 'sessions'],
+      ['session-picker', 'sessions']
+    ]
+
+    const global = [
+      ['billing', 'state'],
+      ['billing', 'subscription'],
+      ['marketplace-themes', 'all'],
+      ['marketplace-themes-settings', 'x'],
+      ['onboarding-model-options', 'y'],
+      ['contrib-logs-tail']
+    ]
+
+    for (const key of [...profileScoped, ...global]) {
+      queryClient.setQueryData(key, { seeded: true })
+    }
+
+    invalidateProfileScopedQueries()
+
+    for (const key of profileScoped) {
+      expect(invalidated(key), `${JSON.stringify(key)} should be invalidated`).toBe(true)
+    }
+
+    for (const key of global) {
+      expect(invalidated(key), `${JSON.stringify(key)} should be left intact`).toBe(false)
+    }
+  })
+
+  it('invalidates unknown/non-string-rooted keys by default (correctness-safe)', () => {
+    queryClient.setQueryData(['some-future-profile-query'], 1)
+    queryClient.setQueryData([{ scope: 'weird' }], 1)
+
+    invalidateProfileScopedQueries()
+
+    expect(invalidated(['some-future-profile-query'])).toBe(true)
+    expect(invalidated([{ scope: 'weird' }])).toBe(true)
+  })
+})

apps/desktop/src/lib/query-client.ts

@@ -18,3 +18,31 @@ export const writeCache =
   <T>(key: QueryKey) =>
   (next: T | undefined | ((prev: T | undefined) => T | undefined)): void =>
     void queryClient.setQueryData<T>(key, next)
+
+// Query-key roots that are NOT profile-scoped: account/billing, the theme
+// marketplace, onboarding, and contrib log tails all read global or
+// account-level state, so a profile/gateway swap must not refetch them. Any
+// other key is treated as profile-scoped and invalidated -- a denylist is
+// correctness-safe here: a root we forget to list just gets refetched (a small
+// cost), whereas an allowlist that misses a profile-scoped key would paint the
+// previous profile's data (a bug).
+const PROFILE_INDEPENDENT_QUERY_ROOTS = new Set<string>([
+  'billing',
+  'marketplace-themes',
+  'marketplace-themes-settings',
+  'onboarding-model-options',
+  'contrib-logs-tail'
+])
+
+// Invalidate profile-scoped query caches on a profile / gateway switch, leaving
+// account/global caches intact. Replaces a keyless invalidateQueries() that
+// refetched everything (billing, marketplace, onboarding) on every switch.
+export function invalidateProfileScopedQueries(): void {
+  void queryClient.invalidateQueries({
+    predicate: query => {
+      const root = query.queryKey[0]
+
+      return typeof root !== 'string' || !PROFILE_INDEPENDENT_QUERY_ROOTS.has(root)
+    }
+  })
+}

apps/desktop/src/store/gateway-switch.test.ts

@@ -19,7 +19,7 @@ import {
 import { $gatewaySwitching, wipeSessionListsForGatewaySwitch } from './gateway-switch'
 
 vi.mock('@/lib/query-client', () => ({
-  queryClient: { invalidateQueries: vi.fn() }
+  invalidateProfileScopedQueries: vi.fn()
 }))
 
 describe('wipeSessionListsForGatewaySwitch', () => {

apps/desktop/src/store/gateway-switch.ts

@@ -1,6 +1,6 @@
 import { atom } from 'nanostores'
 
-import { queryClient } from '@/lib/query-client'
+import { invalidateProfileScopedQueries } from '@/lib/query-client'
 import { resetSessionsLimit } from '@/store/layout'
 import {
   $unreadFinishedSessionIds,
@@ -57,5 +57,7 @@ export function wipeSessionListsForGatewaySwitch(): void {
   setMessages([])
   setFreshDraftReady(true)
 
-  void queryClient.invalidateQueries()
+  // Narrowed: account/marketplace/onboarding caches are global, not gateway-
+  // scoped, so a mode swap must not refetch them.
+  invalidateProfileScopedQueries()
 }

apps/desktop/src/store/profile.test.ts

@@ -16,14 +16,14 @@ vi.mock('@/hermes', () => ({
   getProfiles: vi.fn(async () => ({ profiles: [] })),
   setApiRequestProfile: vi.fn()
 }))
-vi.mock('@/lib/query-client', () => ({ queryClient: { invalidateQueries: vi.fn() } }))
+vi.mock('@/lib/query-client', () => ({ invalidateProfileScopedQueries: vi.fn() }))
 vi.mock('@/store/starmap', () => ({ resetStarmapGraph }))
 
 const { $activeGatewayProfile, $profiles, ensureGatewayProfile, prewarmProfileBackend, refreshProfiles } =
   await import('./profile')
 
 const { $connection } = await import('./session')
-const { queryClient } = await import('@/lib/query-client')
+const { invalidateProfileScopedQueries } = await import('@/lib/query-client')
 const { getProfiles } = await import('@/hermes')
 
 const profile = (name: string, isDefault = false): ProfileInfo => ({
@@ -53,7 +53,7 @@ beforeEach(() => {
   $connection.set(localConn())
   $profiles.set([])
   vi.stubGlobal('window', { hermesDesktop: { getConnection } })
-  vi.mocked(queryClient.invalidateQueries).mockClear()
+  vi.mocked(invalidateProfileScopedQueries).mockClear()
   resetStarmapGraph.mockClear()
 })
 
@@ -114,7 +114,7 @@ describe('profile-scoped cache invalidation', () => {
   it('drops the memory graph cache when the active gateway profile changes', () => {
     $activeGatewayProfile.set('coder')
 
-    expect(queryClient.invalidateQueries).toHaveBeenCalled()
+    expect(invalidateProfileScopedQueries).toHaveBeenCalled()
     expect(resetStarmapGraph).toHaveBeenCalledTimes(1)
   })
 })

apps/desktop/src/store/profile.ts

@@ -1,7 +1,7 @@
 import { atom, computed } from 'nanostores'
 
 import { getProfiles, setApiRequestProfile, STARTUP_REQUEST_TIMEOUT_MS } from '@/hermes'
-import { queryClient } from '@/lib/query-client'
+import { invalidateProfileScopedQueries } from '@/lib/query-client'
 import {
   arraysEqual,
   persistBoolean,
@@ -177,7 +177,9 @@ $activeGatewayProfile.subscribe(value => {
 
   if (_lastRoutedProfile !== null && _lastRoutedProfile !== key) {
     // Profile-scoped settings + the unified session list are now stale.
-    void queryClient.invalidateQueries()
+    // Narrowed so account/marketplace/onboarding caches don't refetch on
+    // every profile switch.
+    invalidateProfileScopedQueries()
     resetStarmapGraph()
   }
 

apps/desktop/src/store/tool-diffs.test.ts

@@ -0,0 +1,47 @@
+import { describe, expect, it } from 'vitest'
+
+import { $toolInlineDiff, getToolDiff, recordToolDiff } from './tool-diffs'
+
+describe('tool-diffs per-tool subscriptions', () => {
+  it('returns a stable cached atom per toolCallId', () => {
+    expect($toolInlineDiff('a')).toBe($toolInlineDiff('a'))
+    expect($toolInlineDiff('a')).not.toBe($toolInlineDiff('b'))
+  })
+
+  it('notifies only the tool whose diff changed', () => {
+    const aCalls: string[] = []
+    const bCalls: string[] = []
+    const unsubA = $toolInlineDiff('notify-a').listen(v => aCalls.push(v))
+    const unsubB = $toolInlineDiff('notify-b').listen(v => bCalls.push(v))
+
+    recordToolDiff('notify-a', 'diffA')
+    expect(aCalls).toEqual(['diffA'])
+    expect(bCalls).toEqual([]) // the unrelated tool row is never notified
+
+    recordToolDiff('notify-b', 'diffB')
+    expect(aCalls).toEqual(['diffA']) // still not re-notified
+    expect(bCalls).toEqual(['diffB'])
+
+    unsubA()
+    unsubB()
+  })
+
+  it('does not re-notify when the same diff is recorded again', () => {
+    const calls: string[] = []
+    const unsub = $toolInlineDiff('same').listen(v => calls.push(v))
+
+    recordToolDiff('same', 'x')
+    recordToolDiff('same', 'x')
+
+    expect(calls).toEqual(['x'])
+    unsub()
+  })
+
+  it('reads the current diff for a tool and empty for unknown/blank ids', () => {
+    recordToolDiff('read-me', 'value')
+    expect(getToolDiff('read-me')).toBe('value')
+    expect(getToolDiff('missing')).toBe('')
+    expect(getToolDiff('')).toBe('')
+    expect($toolInlineDiff('').get()).toBe('')
+  })
+})

apps/desktop/src/store/tool-diffs.ts

@@ -1,7 +1,13 @@
-import { atom } from 'nanostores'
+import { atom, computed, type ReadableAtom } from 'nanostores'
 
 const $toolDiffs = atom<Record<string, string>>({})
 
+// Per-tool derived atoms, cached by toolCallId. A `ToolEntry` subscribes only
+// to its own id's diff, so recording a diff for one tool re-renders that one
+// row -- not every mounted tool row. computed() only notifies when the derived
+// string actually changes, so unrelated writes to the map are inert here.
+const inlineDiffCache = new Map<string, ReadableAtom<string>>()
+
 export function recordToolDiff(toolCallId: string, diff: string) {
   if (!toolCallId || !diff) {
     return
@@ -20,4 +26,13 @@ export function getToolDiff(toolCallId: string): string {
   return toolCallId ? $toolDiffs.get()[toolCallId] || '' : ''
 }
 
-export const $toolInlineDiffs = $toolDiffs
+export function $toolInlineDiff(toolCallId: string): ReadableAtom<string> {
+  let cached = inlineDiffCache.get(toolCallId)
+
+  if (!cached) {
+    cached = computed($toolDiffs, diffs => (toolCallId ? diffs[toolCallId] || '' : ''))
+    inlineDiffCache.set(toolCallId, cached)
+  }
+
+  return cached
+}