From 700ff78097d67f5f261bef5507584dbc93291efd Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Tue, 25 Aug 2026 14:38:08 -0700 Subject: [PATCH] fix(desktop): preserve registered gateways during mode switches --- .../gateway/hooks/use-gateway-boot.test.tsx | 41 ++++++++++++++++++- .../src/app/gateway/hooks/use-gateway-boot.ts | 10 ++++- .../gateway-connection-lifecycle.test.ts | 25 +++++++++++ apps/desktop/src/store/gateway.ts | 36 +++++++++++++--- apps/desktop/src/store/session-states.test.ts | 30 +++++++++++++- apps/desktop/src/store/session-states.ts | 16 ++++++++ 6 files changed, 148 insertions(+), 10 deletions(-) diff --git a/apps/desktop/src/app/gateway/hooks/use-gateway-boot.test.tsx b/apps/desktop/src/app/gateway/hooks/use-gateway-boot.test.tsx index 097943d981..46f7129e7b 100644 --- a/apps/desktop/src/app/gateway/hooks/use-gateway-boot.test.tsx +++ b/apps/desktop/src/app/gateway/hooks/use-gateway-boot.test.tsx @@ -9,7 +9,7 @@ import { selectConnection, setConnectionsRegistry } from '@/store/connections' -import { closeSecondaryGateways, isActivePrimary, requestGatewayForAgent } from '@/store/gateway' +import { activeGateway, closeSecondaryGateways, ensureGatewayForAgent, isActivePrimary, requestGatewayForAgent } from '@/store/gateway' import { reconnectGateway } from '@/store/gateway-reconnect' import { $gatewaySwitching, @@ -880,6 +880,45 @@ describe('useGatewayBoot remote reconnect loop (real hook, fake socket)', () => expect(FakeWebSocket.instances).toHaveLength(1) }) + it('keeps registered source sockets alive during a legacy mode apply', async () => { + const desktop = fakeDesktop() as ReturnType & { + getConnectionFor: ReturnType + } + + desktop.getConnectionFor = vi.fn(async ({ connectionId, profile }: { connectionId: string; profile: string }) => ({ + ...coderConn, + connectionId, + profile, + wsUrl: `wss://${connectionId}.example.com/api/ws?token=r` + })) + ;(window as { hermesDesktop?: unknown }).hermesDesktop = desktop + + render() + await flushAsync() + expect($gatewayState.get()).toBe('open') + + let opening!: Promise + act(() => { + opening = ensureGatewayForAgent('cloud', 'default') + }) + await flushAsync() + await opening + + const registeredGateway = activeGateway() + expect(registeredGateway).not.toBeNull() + expect(isActivePrimary()).toBe(false) + + act(() => connectionApplied?.()) + await flushAsync() + await flushAsync() + + // Applying the legacy Local/Cloud mode must not close an independent v2 + // source. The foreground returns to the new primary, while the registered + // socket remains reusable and cannot arm ws_orphan_reap on the old backend. + expect(registeredGateway?.connectionState).toBe('open') + expect(isActivePrimary()).toBe(true) + }) + it('re-fetches the profile rail from the NEW backend after a connection apply (#85731)', async () => { // The reported repro: connected to backend A, the rail shows A's named // profiles; the user applies a different remote/Cloud connection (soft diff --git a/apps/desktop/src/app/gateway/hooks/use-gateway-boot.ts b/apps/desktop/src/app/gateway/hooks/use-gateway-boot.ts index 034993e46e..8b9d2d198b 100644 --- a/apps/desktop/src/app/gateway/hooks/use-gateway-boot.ts +++ b/apps/desktop/src/app/gateway/hooks/use-gateway-boot.ts @@ -18,6 +18,7 @@ import { } from '@/store/boot' import { $gateway, + closeLegacySecondaryGateways, closeSecondaryGateways, configureGatewayRegistry, disposeSecondariesForConnection, @@ -62,6 +63,7 @@ import { $attentionSessionIds, $workingSessionIds, $sessionTiles, + foregroundSessionScopes, liveSessionScopes, openTileGatewayScopes, reconcileBusyStatesOnReconnect, @@ -524,7 +526,11 @@ export function useGatewayBoot({ reauthNotified = false gateway.close() - closeSecondaryGateways() + // The primary mode is changing, but registered v2 sources remain + // independent gateways. Retire only legacy profile sockets whose + // routing follows connection.json; closing every secondary here + // detached valid registered sessions and armed ws_orphan_reap. + closeLegacySecondaryGateways() // Same override rule as boot(): a profile-pinned helper window stays // on its pinned profile's backend across a soft switch. @@ -788,7 +794,7 @@ export function useGatewayBoot({ // sources can expose the same profile name (every source has a // 'default'), so bare profile names can't represent a non-local // source's liveness without keeping the wrong gateway alive. - const keep = liveSessionScopes() + const keep = new Set([...liveSessionScopes(), ...foregroundSessionScopes()]) for (const session of $sessions.get()) { if (live.has(session.id)) { diff --git a/apps/desktop/src/store/gateway-connection-lifecycle.test.ts b/apps/desktop/src/store/gateway-connection-lifecycle.test.ts index 10173894b3..2ef43a617a 100644 --- a/apps/desktop/src/store/gateway-connection-lifecycle.test.ts +++ b/apps/desktop/src/store/gateway-connection-lifecycle.test.ts @@ -53,6 +53,7 @@ vi.mock('@/store/session-states', () => reconnectStateMocks) const { activeGateway, + closeLegacySecondaryGateways, closeSecondaryGateways, configureGatewayRegistry, disposeSecondariesForConnection, @@ -181,6 +182,30 @@ describe('disposeSecondariesForConnection', () => { }) }) +describe('legacy secondary teardown', () => { + it('closes v1 profile sockets without detaching registered sources', async () => { + const getConnection = vi.fn(async (profile: string) => descriptorFor('legacy-local', profile)) + + const getConnectionFor = vi.fn(async ({ connectionId, profile }: { connectionId: string; profile: string }) => + descriptorFor(connectionId, profile) + ) + + installDesktop({ getConnection, getConnectionFor }) + + await openGatewayForProfile('writer') + await ensureGatewayForAgent('homelab', 'default') + + const legacy = gatewayMocks.instances[0] + const registered = gatewayMocks.instances[1] + + closeLegacySecondaryGateways() + + expect(legacy.close).toHaveBeenCalledOnce() + expect(registered.close).not.toHaveBeenCalled() + expect(activeGateway()).toBe(registered) + }) +}) + describe('secondary reconnect runtime scope', () => { it('invalidates stale runtime bindings before a direct secondary reopen publishes open', async () => { installDesktop({ diff --git a/apps/desktop/src/store/gateway.ts b/apps/desktop/src/store/gateway.ts index 737c6c0d4d..fd7fd8d229 100644 --- a/apps/desktop/src/store/gateway.ts +++ b/apps/desktop/src/store/gateway.ts @@ -1418,7 +1418,36 @@ export function pruneSecondaryGateways(keep: Set): void { restoreActiveToPrimaryIfEvicted() } +function closeSecondariesWhere(shouldClose: (entry: Secondary) => boolean): void { + for (const [scope, entry] of [...g.secondaries]) { + if (!shouldClose(entry)) { + continue + } + + disposeSecondary(entry) + g.secondaries.delete(scope) + } + + restoreActiveToPrimaryIfEvicted() +} + +/** + * Close only profile sockets that follow the legacy v1 connection config. + * + * A global mode apply re-homes the primary backend, but registered connection + * sockets are independent sources in the v2 registry. Closing every secondary + * here would detach their sessions and arm `ws_orphan_reap` even though those + * sources remain valid and reusable. Legacy profile sockets still need to be + * retired because their endpoint is derived from the v1 config being changed. + */ +export function closeLegacySecondaryGateways(): void { + closeSecondariesWhere(entry => entry.connectionId == null) +} + export function closeSecondaryGateways(): void { + // Full teardown releases every routed-turn lease (class-2 #94284) and the + // renderer-generation ledger; the predicate close leaves live sources' + // leases alone (their sockets stay open). for (const timer of g.turnLeaseReleaseTimers.values()) { clearTimeout(timer) } @@ -1431,13 +1460,8 @@ export function closeSecondaryGateways(): void { g.turnLeases.clear() - for (const entry of g.secondaries.values()) { - disposeSecondary(entry) - } - - g.secondaries.clear() + closeSecondariesWhere(() => true) openedSecondaryScopes().clear() - restoreActiveToPrimaryIfEvicted() } // A local profile can have two renderer-owned sockets: the legacy bare diff --git a/apps/desktop/src/store/session-states.test.ts b/apps/desktop/src/store/session-states.test.ts index 9bce372be0..73790dc223 100644 --- a/apps/desktop/src/store/session-states.test.ts +++ b/apps/desktop/src/store/session-states.test.ts @@ -4,14 +4,16 @@ import type { ClientSessionState } from '@/app/types' import { findGroupOfPane, group, split } from '@/components/pane-shell/tree/model' import { $layoutTree } from '@/components/pane-shell/tree/store' import { $activeGatewayProfile } from '@/store/profile' -import { $connection, $selectedStoredSessionId, setSessions } from '@/store/session' +import { $activeSessionId, $connection, $selectedStoredSessionId, setSessions } from '@/store/session' import type { SessionTile } from '@/store/session-states' import { $sessionStates, $sessionTiles, blankDraftTile, + clearAllSessionStates, focusedSessionNeedsRoute, focusOpenSession, + foregroundSessionScopes, isSessionRemote, knownOwnerForSession, markSelectionRestore, @@ -19,6 +21,7 @@ import { openSessionTile, orderTilesByTree, patchSessionTile, + recordSessionEventScope, releaseSessionTranscript, requestForOwnedSession, resetTileRuntimeBindings, @@ -32,6 +35,31 @@ import { const tile = (storedSessionId: string): SessionTile => ({ storedSessionId }) const tilePane = (id: string) => `session-tile:${id}` +describe('foregroundSessionScopes', () => { + beforeEach(() => { + clearAllSessionStates() + $activeSessionId.set(null) + }) + + afterEach(() => { + clearAllSessionStates() + $activeSessionId.set(null) + }) + + it('keeps the exact registry owner of an idle foreground runtime', () => { + recordSessionEventScope({ connectionId: 'cloud', profile: 'default', session_id: 'runtime-1' }) + $activeSessionId.set('runtime-1') + + expect(foregroundSessionScopes()).toEqual(new Set(['conn:cloud::default'])) + }) + + it('fails closed when the foreground runtime has no registered source', () => { + $activeSessionId.set('legacy-runtime') + + expect(foregroundSessionScopes()).toEqual(new Set()) + }) +}) + describe('resetTileRuntimeBindings', () => { afterEach(() => { $sessionTiles.set([]) diff --git a/apps/desktop/src/store/session-states.ts b/apps/desktop/src/store/session-states.ts index 3ea60a20ea..01339e09ba 100644 --- a/apps/desktop/src/store/session-states.ts +++ b/apps/desktop/src/store/session-states.ts @@ -99,6 +99,22 @@ export function liveSessionScopes(): Set { return scopes } +/** + * The exact registry scope that owns the foreground runtime, when known. + * + * The secondary-gateway pruner normally keeps only busy/needs-input work. A + * source switch briefly changes the active gateway before the old foreground + * runtime is cleared, though, and an idle conversation can otherwise look + * disposable during that handoff. Preserve the owner for that narrow window; + * tiles have their own persisted owner route and are handled separately. + */ +export function foregroundSessionScopes(): Set { + const runtimeId = $activeSessionId.get() + const scope = runtimeId ? sessionScopeByRuntimeId.get(runtimeId) : undefined + + return scope ? new Set([scope]) : new Set() +} + // Stored session ids whose authoritative state is still busy, but whose // runtime has produced no state publish for the watchdog window. Silence is // not completion: long tool calls can legitimately stay quiet, so this is a