Merge pull request #75836 from NousResearch/bb/session-reap-notify

A reclaimed session tells the client instead of vanishing

apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts

@@ -59,6 +59,7 @@ import {
   setTurnStartedAt,
   setYoloActive
 } from '@/store/session'
+import { dropSessionState } from '@/store/session-states'
 import { pruneDelegateFallbackSubagents, pruneFinishedSessionSubagents, upsertSubagent } from '@/store/subagents'
 import { clearActiveSessionTodos } from '@/store/todos'
 import { recordToolDiff } from '@/store/tool-diffs'
@@ -330,6 +331,23 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) {
           }
         }
 
+        return
+      } else if (event.type === 'session.reclaimed') {
+        // The backend reclaimed a live session we may still be holding (idle
+        // TTL, LRU cap, or the WS-orphan reap). Without this the runtime id
+        // stays cached until something fails against it, which reads as the
+        // session vanishing rather than being reclaimed. Drop the cached state
+        // now — the stored row is untouched, so the sidebar keeps the
+        // conversation and reopening it resumes from the DB.
+        const reclaimedRuntimeId = String((payload as { session_id?: string } | undefined)?.session_id ?? '')
+
+        if (reclaimedRuntimeId) {
+          dropSessionState(reclaimedRuntimeId)
+        }
+
+        // The row's ended_at moved, so refresh the lists that render it.
+        notifySessionsChanged()
+
         return
       } else if (event.type === 'session.info') {
         // Apply session-scoped fields when the event targets the active

apps/desktop/src/app/session/hooks/use-message-stream/session-reclaimed.test.tsx

@@ -0,0 +1,125 @@
+import { QueryClient } from '@tanstack/react-query'
+import { act, cleanup, render, waitFor } from '@testing-library/react'
+import { useEffect, useRef } from 'react'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+import type { ClientSessionState } from '@/app/types'
+import { createClientSessionState } from '@/lib/chat-runtime'
+import { $sessionStates, publishSessionState } from '@/store/session-states'
+import type { RpcEvent } from '@/types/hermes'
+
+import { useMessageStream } from './index'
+
+// `session.reclaimed`: the backend tore down a live session we're still
+// holding (idle TTL, LRU cap, WS-orphan reap). Before this event the runtime id
+// stayed cached until something failed against it, which read to the user as
+// the session vanishing rather than being reclaimed.
+
+const ACTIVE_SID = 'session-active'
+const ACTIVE_PROFILE = 'compass'
+let handleEvent: ((event: RpcEvent) => void) | null = null
+let queryClient: QueryClient
+
+function Harness() {
+  const activeSessionIdRef = useRef<string | null>(ACTIVE_SID)
+  const sessionStateByRuntimeIdRef = useRef(new Map<string, ClientSessionState>())
+
+  const stream = useMessageStream({
+    activeGatewayProfile: ACTIVE_PROFILE,
+    activeSessionIdRef,
+    hydrateFromStoredSession: vi.fn(async () => undefined),
+    queryClient,
+    refreshHermesConfig: vi.fn<() => Promise<void>>(async () => undefined),
+    refreshSessions: vi.fn<() => Promise<void>>(async () => undefined),
+    sessionStateByRuntimeIdRef,
+    updateSessionState: (sessionId, updater) => {
+      const current = sessionStateByRuntimeIdRef.current.get(sessionId) ?? createClientSessionState()
+      const next = updater(current)
+      sessionStateByRuntimeIdRef.current.set(sessionId, next)
+
+      return next
+    }
+  })
+
+  useEffect(() => {
+    handleEvent = stream.handleGatewayEvent
+  }, [stream.handleGatewayEvent])
+
+  return null
+}
+
+async function mountStream() {
+  render(<Harness />)
+  await waitFor(() => expect(handleEvent).not.toBeNull())
+}
+
+const reclaim = (sessionId: string, reason = 'ws_orphan_reap') =>
+  act(() =>
+    handleEvent!({
+      payload: { reason, session_id: sessionId, stored_session_id: 'stored-1' },
+      session_id: '',
+      type: 'session.reclaimed'
+    } as RpcEvent)
+  )
+
+beforeEach(() => {
+  handleEvent = null
+  queryClient = new QueryClient()
+  $sessionStates.set({})
+})
+
+afterEach(() => {
+  cleanup()
+  $sessionStates.set({})
+  vi.restoreAllMocks()
+})
+
+describe('session.reclaimed', () => {
+  it('drops the cached state for the reclaimed runtime', async () => {
+    await mountStream()
+    publishSessionState('live-gone', createClientSessionState())
+    expect($sessionStates.get()['live-gone']).toBeDefined()
+
+    reclaim('live-gone')
+
+    expect($sessionStates.get()['live-gone']).toBeUndefined()
+  })
+
+  it('leaves every other live session alone', async () => {
+    await mountStream()
+    publishSessionState('live-gone', createClientSessionState())
+    publishSessionState('live-kept', createClientSessionState())
+
+    reclaim('live-gone')
+
+    // Both halves matter: the target went, the bystander stayed. Asserting
+    // only the survivor would pass with no handler at all.
+    expect($sessionStates.get()['live-gone']).toBeUndefined()
+    expect($sessionStates.get()['live-kept']).toBeDefined()
+  })
+
+  it('ignores a payload with no runtime id instead of clearing everything', async () => {
+    await mountStream()
+    publishSessionState('live-a', createClientSessionState())
+    publishSessionState('live-b', createClientSessionState())
+
+    reclaim('')
+
+    // A malformed/empty id must be a no-op, never a blanket wipe.
+    expect(Object.keys($sessionStates.get()).sort()).toEqual(['live-a', 'live-b'])
+  })
+
+  it('drops the runtime regardless of which reclaim reason fired', async () => {
+    for (const reason of ['idle_timeout', 'lru_evict', 'ws_orphan_reap']) {
+      $sessionStates.set({})
+      cleanup()
+      handleEvent = null
+      await mountStream()
+      publishSessionState('live-gone', createClientSessionState())
+
+      reclaim('live-gone', reason)
+
+      expect($sessionStates.get()['live-gone'], reason).toBeUndefined()
+    }
+  })
+})

tests/tui_gateway/test_session_reclaim_notify.py

@@ -0,0 +1,87 @@
+"""A backend-reclaimed session must tell the clients still holding it.
+
+The idle-TTL reaper, the LRU cap, and the WS-orphan reap all tear a live
+session down without the client asking. Before ``session.reclaimed`` the client
+kept a runtime id the backend had already forgotten, and only discovered it by
+failing a later prompt — the "sessions suddenly lost in the backend" report.
+
+Contracts here: a reclaim broadcasts with its reason, a user-initiated close
+stays silent, and a notify failure never breaks teardown.
+"""
+
+import pytest
+
+from tui_gateway import server
+
+
+@pytest.fixture()
+def captured(monkeypatch):
+    events = []
+    monkeypatch.setattr(
+        server, "_broadcast_global_event", lambda ev, payload=None: events.append((ev, payload))
+    )
+    # Teardown's real work (finalize, agent close, notifier unregister) is out
+    # of scope — this is about what reaches the client.
+    monkeypatch.setattr(server, "_finalize_session", lambda *a, **k: None)
+    return events
+
+
+def _session():
+    return {"_sid": "live-abc", "session_key": "20260731_120000_aaaaaa"}
+
+
+@pytest.mark.parametrize("reason", ["idle_timeout", "lru_evict", "ws_orphan_reap"])
+def test_reclaim_reasons_announce_to_clients(captured, reason):
+    server._teardown_session(_session(), end_reason=reason)
+
+    assert captured == [
+        (
+            "session.reclaimed",
+            {
+                "session_id": "live-abc",
+                "stored_session_id": "20260731_120000_aaaaaa",
+                "reason": reason,
+            },
+        )
+    ]
+
+
+@pytest.mark.parametrize("reason", ["tui_close", "tui_shutdown", "ws_disconnect", "branched"])
+def test_client_initiated_closes_stay_silent(captured, reason):
+    """The client asked for these, so announcing them would be noise."""
+    server._teardown_session(_session(), end_reason=reason)
+
+    assert captured == []
+
+
+def test_broadcast_failure_does_not_break_teardown(monkeypatch):
+    """A wedged peer must not leave a session half torn down."""
+    finalized = []
+
+    def _boom(*_a, **_k):
+        raise RuntimeError("transport gone")
+
+    monkeypatch.setattr(server, "_broadcast_global_event", _boom)
+    monkeypatch.setattr(
+        server, "_finalize_session", lambda s, **k: finalized.append(k.get("end_reason"))
+    )
+
+    server._teardown_session(_session(), end_reason="ws_orphan_reap")
+
+    assert finalized == ["ws_orphan_reap"]
+
+
+def test_reap_paths_stamp_the_runtime_id_the_client_holds():
+    """_pop_session_by_id stamps ``_sid``; the payload is useless without it.
+
+    All three reclaim paths pop before tearing down, so this is the invariant
+    that makes the broadcast addressable on the client.
+    """
+    server._sessions["live-xyz"] = {"session_key": "k"}
+    try:
+        popped = server._pop_session_by_id("live-xyz")
+    finally:
+        server._sessions.pop("live-xyz", None)
+
+    assert popped is not None
+    assert popped["_sid"] == "live-xyz"

tui_gateway/server.py

@@ -793,6 +793,39 @@ def _finalize_session(session: dict | None, end_reason: str = "tui_close") -> No
         pass
 
 
+# End reasons where the BACKEND reclaimed a session the client never asked to
+# close: the idle-TTL reaper, the LRU cap, and the WS-orphan reap. A client
+# holding that live session id gets no signal today — its next prompt fails
+# against an id the backend has already forgotten, which reads as the session
+# silently vanishing rather than being reclaimed. ``tui_close`` and friends are
+# deliberately absent: the client initiated those and already knows.
+_RECLAIM_END_REASONS = frozenset({"idle_timeout", "lru_evict", "ws_orphan_reap"})
+
+
+def _announce_session_reclaimed(session: dict, end_reason: str) -> None:
+    """Tell connected clients a session was reclaimed out from under them.
+
+    Broadcast rather than session-targeted: the reap paths run on background
+    timer threads with no contextvar binding, and the WS-orphan case has by
+    definition lost its own transport — ``_emit`` would bottom out on stdio and
+    the peer that owns the session would never see it. Best-effort; a failed
+    notify must never break teardown.
+    """
+    if end_reason not in _RECLAIM_END_REASONS:
+        return
+    try:
+        _broadcast_global_event(
+            "session.reclaimed",
+            {
+                "session_id": str(session.get("_sid") or ""),
+                "stored_session_id": str(session.get("session_key") or ""),
+                "reason": end_reason,
+            },
+        )
+    except Exception:
+        logger.debug("session.reclaimed broadcast failed", exc_info=True)
+
+
 def _teardown_session(session: dict | None, *, end_reason: str = "tui_close") -> None:
     """Fully tear down a session: finalize, unregister, close agent + worker.
 
@@ -806,6 +839,7 @@ def _teardown_session(session: dict | None, *, end_reason: str = "tui_close") ->
     if not session:
         return
     _finalize_session(session, end_reason=end_reason)
+    _announce_session_reclaimed(session, end_reason)
     try:
         from tools.approval import unregister_gateway_notify