Merge pull request #54503 from NousResearch/bb/fix-desktop-cross-wired-resume

fix(desktop): restore cross-wired runtime-id guard on session resume

apps/desktop/src/app/session/hooks/use-session-actions.test.tsx

@@ -4,6 +4,7 @@ import { useEffect } from 'react'
 import { afterEach, describe, expect, it, vi } from 'vitest'
 
 import { getSessionMessages, type SessionInfo } from '@/hermes'
+import { createClientSessionState } from '@/lib/chat-runtime'
 import { $activeGatewayProfile, $newChatProfile } from '@/store/profile'
 import {
   $activeSessionId,
@@ -410,3 +411,105 @@ describe('resumeSession failure recovery', () => {
     expect($messages.get().length).toBe(1)
   })
 })
+
+// ── Warm-cache mapping integrity (the "open chat A, chat B loads" bug) ─────────
+// resumeSession's warm fast-path maps storedSessionId -> runtimeId -> cached
+// state. A reaped/respawned pooled backend re-mints runtime ids, so a recycled
+// id can resolve to a live-but-DIFFERENT session's cache entry. The fast-path
+// must verify the cached state still BELONGS to the resumed session before it
+// paints, or it shows a totally different thread under the current route.
+const clientState = (storedSessionId: string | null): ClientSessionState => createClientSessionState(storedSessionId)
+
+describe('resumeSession warm-cache mapping integrity', () => {
+  afterEach(() => {
+    cleanup()
+    setActiveSessionId(null)
+    setResumeFailedSessionId(null)
+    setMessages([])
+    setSessions([])
+    vi.restoreAllMocks()
+  })
+
+  it('rejects a cross-wired runtime mapping and falls through to a full resume', async () => {
+    // A recycled runtime id ('rt-recycled') is mapped to 'stored-A', but its
+    // cached state actually belongs to a DIFFERENT session ('stored-B') — the
+    // exact "open chat A, chat B loads" corruption a reaped/respawned pooled
+    // backend can leave behind.
+    const runtimeIdByStoredSessionIdRef: MutableRefObject<Map<string, string>> = {
+      current: new Map([['stored-A', 'rt-recycled']])
+    }
+
+    const sessionStateByRuntimeIdRef: MutableRefObject<Map<string, ClientSessionState>> = {
+      current: new Map([['rt-recycled', clientState('stored-B')]])
+    }
+
+    const requestGateway = vi.fn(async (method: string, params?: Record<string, unknown>) => {
+      if (method === 'session.resume') {
+        return { session_id: 'rt-A-fresh', resumed: params?.session_id, messages: [], info: {} } as never
+      }
+
+      return {} as never
+    })
+
+    vi.mocked(getSessionMessages).mockResolvedValue({ messages: [] } as never)
+
+    let resume: ((storedSessionId: string, replaceRoute?: boolean) => Promise<unknown>) | null = null
+    render(
+      <ResumeHarness
+        onReady={r => (resume = r)}
+        requestGateway={requestGateway}
+        runtimeIdByStoredSessionIdRef={runtimeIdByStoredSessionIdRef}
+        sessionStateByRuntimeIdRef={sessionStateByRuntimeIdRef}
+      />
+    )
+    await waitFor(() => expect(resume).not.toBeNull())
+    await resume!('stored-A', true)
+
+    // The fast-path did NOT short-circuit on the cross-wired cache — the full
+    // resume RPC ran, for the session that was actually requested.
+    const resumeCalls = requestGateway.mock.calls.filter(([method]) => method === 'session.resume')
+    expect(resumeCalls.length).toBe(1)
+    expect(resumeCalls[0][1]).toMatchObject({ session_id: 'stored-A' })
+
+    // The corrupt mapping was purged so it can't mis-resolve again.
+    expect(runtimeIdByStoredSessionIdRef.current.has('stored-A')).toBe(false)
+    expect(sessionStateByRuntimeIdRef.current.has('rt-recycled')).toBe(false)
+  })
+
+  it('honours a warm cache entry whose stored id matches (no needless refetch)', async () => {
+    // Correctly-wired mapping: 'rt-A' <-> 'stored-A'. The fast-path should trust
+    // it and never reach session.resume (only the lightweight usage probe).
+    const runtimeIdByStoredSessionIdRef: MutableRefObject<Map<string, string>> = {
+      current: new Map([['stored-A', 'rt-A']])
+    }
+
+    const sessionStateByRuntimeIdRef: MutableRefObject<Map<string, ClientSessionState>> = {
+      current: new Map([['rt-A', clientState('stored-A')]])
+    }
+
+    const requestGateway = vi.fn(async (method: string) => {
+      if (method === 'session.usage') {
+        return { input: 0, output: 0, total: 0 } as never
+      }
+
+      return {} as never
+    })
+
+    let resume: ((storedSessionId: string, replaceRoute?: boolean) => Promise<unknown>) | null = null
+    render(
+      <ResumeHarness
+        onReady={r => (resume = r)}
+        requestGateway={requestGateway}
+        runtimeIdByStoredSessionIdRef={runtimeIdByStoredSessionIdRef}
+        sessionStateByRuntimeIdRef={sessionStateByRuntimeIdRef}
+      />
+    )
+    await waitFor(() => expect(resume).not.toBeNull())
+    await resume!('stored-A', true)
+
+    // Fast-path served the session from cache: no full resume RPC, mapping intact.
+    const methods = requestGateway.mock.calls.map(([method]) => method)
+    expect(methods).not.toContain('session.resume')
+    expect(runtimeIdByStoredSessionIdRef.current.get('stored-A')).toBe('rt-A')
+  })
+})

apps/desktop/src/app/session/hooks/use-session-actions.ts

@@ -631,9 +631,34 @@ export function useSessionActions({
       // chat view drops the error state and shows the loader again.
       setResumeExhaustedSessionId(current => (current === storedSessionId ? null : current))
 
-      const warmRuntimeId = runtimeIdByStoredSessionIdRef.current.get(storedSessionId)
+      // A warm cache entry is only trustworthy when it still BELONGS to the
+      // session being resumed. A pooled profile backend that gets idle-reaped
+      // and respawned (pruneSecondaryGateways) re-mints runtime ids, so a
+      // recycled id can resolve to a live-but-DIFFERENT session's cache entry.
+      // The session.usage 404 guard below only catches a fully-DEAD id — a
+      // recycled-live id 200s, so an unchecked hit paints the wrong transcript
+      // under the current route (the "open chat A, chat B loads" bug). On a
+      // mismatch the mapping is cross-wired: purge both sides and report a miss
+      // so the caller falls through to a full resume that rebinds a correct id.
+      const takeWarmCache = (): { runtimeId: string; state: ClientSessionState } | null => {
+        const runtimeId = runtimeIdByStoredSessionIdRef.current.get(storedSessionId)
+        const state = runtimeId ? sessionStateByRuntimeIdRef.current.get(runtimeId) : undefined
 
-      if (!warmRuntimeId || !sessionStateByRuntimeIdRef.current.get(warmRuntimeId)) {
+        if (!runtimeId || !state) {
+          return null
+        }
+
+        if (state.storedSessionId !== storedSessionId) {
+          runtimeIdByStoredSessionIdRef.current.delete(storedSessionId)
+          sessionStateByRuntimeIdRef.current.delete(runtimeId)
+
+          return null
+        }
+
+        return { runtimeId, state }
+      }
+
+      if (!takeWarmCache()) {
         setActiveSessionId(null)
         activeSessionIdRef.current = null
         setMessages([])
@@ -652,10 +677,14 @@ export function useSessionActions({
 
       await ensureGatewayProfile(sessionProfile)
 
-      const cachedRuntimeId = runtimeIdByStoredSessionIdRef.current.get(storedSessionId)
-      const cachedState = cachedRuntimeId && sessionStateByRuntimeIdRef.current.get(cachedRuntimeId)
+      // Re-check after the profile-resolve / gateway-swap awaits above: the
+      // cache may have changed, and takeWarmCache re-validates belongs-to and
+      // purges a cross-wired mapping before we trust the fast-path.
+      const warmHit = takeWarmCache()
 
-      if (cachedRuntimeId && cachedState) {
+      if (warmHit) {
+        const cachedRuntimeId = warmHit.runtimeId
+        const cachedState = warmHit.state
         const stored = $sessions.get().find(session => sessionMatchesStoredId(session, storedSessionId)) ?? storedForProfile
 
         const cachedViewState =