From 3a847bfac9da57bd47665ec5fcb447b25217f46a Mon Sep 17 00:00:00 2001 From: Jinwoo Hong <73622457+Jinwoo-H@users.noreply.github.com> Date: Sun, 19 Jul 2026 19:14:15 -0700 Subject: [PATCH] fix(ssh): clear stamped agent status on disconnect (#9484) * fix(ssh): clear stamped agent status on disconnect Batch transient cleanup by accepted SSH connection authority and use a monotonic cutoff so reconnect replay wins over delayed clears. Preserve pane launch, resume, acknowledgement, and retention metadata. Caveat: legacy or renderer-owned rows without an accepted connection stamp are intentionally left to existing pane/PTY teardown; clearing them by host would be ambiguous. * docs(ssh): explain stale status watermark * fix(ssh): preserve status ordering after restart --- src/main/agent-hooks/server.test.ts | 149 +++++++++++++++++- src/main/agent-hooks/server.ts | 99 ++++++++++-- src/main/index.ts | 4 +- ...ay-session-agent-hooks.integration.test.ts | 29 ++++ src/main/ssh/ssh-relay-session.ts | 4 + src/preload/api-types.ts | 5 +- src/preload/index.ts | 5 +- src/renderer/src/hooks/useIpcEvents.test.ts | 137 +++++++++++++++- src/renderer/src/hooks/useIpcEvents.ts | 71 +++++++-- .../agent-status-ssh-connection-clear.test.ts | 97 ++++++++++++ src/renderer/src/store/slices/agent-status.ts | 48 +++++- src/shared/agent-status-types.ts | 11 ++ 12 files changed, 624 insertions(+), 35 deletions(-) create mode 100644 src/renderer/src/store/slices/agent-status-ssh-connection-clear.test.ts diff --git a/src/main/agent-hooks/server.test.ts b/src/main/agent-hooks/server.test.ts index b666c9329..f1aba1cb0 100644 --- a/src/main/agent-hooks/server.test.ts +++ b/src/main/agent-hooks/server.test.ts @@ -1425,7 +1425,92 @@ describe('AgentHookServer listener replay', () => { server.clearPaneState(PANE) expect(listener).toHaveBeenCalledTimes(1) - expect(listener).toHaveBeenCalledWith(PANE) + expect(listener).toHaveBeenCalledWith({ paneKey: PANE }) + }) + + it('batches connection cleanup and retains sibling and local statuses', () => { + const server = new AgentHookServer() + const paneKeyAt = (prefix: string, index: number): string => + makePaneKey( + `${prefix}-tab-${index}`, + `00000000-0000-4000-8000-${(index + 1).toString(16).padStart(12, '0')}` + ) + const targetPaneKeys = Array.from({ length: 100 }, (_, index) => paneKeyAt('target', index)) + const siblingPaneKeys = Array.from({ length: 100 }, (_, index) => + paneKeyAt('sibling', index + 100) + ) + const unstampedPaneKey = paneKeyAt('legacy', 250) + const statusListener = vi.fn() + const clearListener = vi.fn() + const internals = server as unknown as AgentHookServerCacheInternals + const persistSpy = vi.spyOn(internals, 'scheduleStatusPersist') + server.subscribeStatusChanges(statusListener) + server.setPaneStatusClearListener(clearListener) + for (const paneKey of targetPaneKeys) { + server.ingestRemote({ paneKey, payload: { state: 'working', agentType: 'claude' } }, 'ssh-a') + } + for (const paneKey of siblingPaneKeys) { + server.ingestRemote({ paneKey, payload: { state: 'working', agentType: 'claude' } }, 'ssh-b') + } + server.ingestTerminalStatus({ + paneKey: unstampedPaneKey, + payload: { state: 'working', prompt: '', agentType: 'codex' } + }) + statusListener.mockClear() + persistSpy.mockClear() + + server.clearStatusEntriesForConnection('ssh-a') + + expect(statusListener).toHaveBeenCalledOnce() + expect(persistSpy).toHaveBeenCalledOnce() + expect(clearListener).toHaveBeenCalledOnce() + expect(clearListener).toHaveBeenCalledWith({ + transient: true, + connectionId: 'ssh-a', + clearedAt: expect.any(Number) + }) + expect(server.getStatusSnapshot().map((entry) => entry.paneKey)).toEqual([ + ...siblingPaneKeys, + unstampedPaneKey + ]) + }) + + it('emits a connection cutoff after a pane-key collision and orders replay after it', () => { + vi.useFakeTimers() + vi.setSystemTime(1_700_000_000_000) + const server = new AgentHookServer() + const clearListener = vi.fn() + server.setPaneStatusClearListener(clearListener) + server.ingestRemote( + { paneKey: PANE, payload: { state: 'working', agentType: 'claude' } }, + 'ssh-a' + ) + server.ingestRemote( + { paneKey: PANE, payload: { state: 'working', agentType: 'codex' } }, + 'ssh-b' + ) + + server.clearStatusEntriesForConnection('ssh-a') + const clear = clearListener.mock.calls[0]?.[0] as { + transient: true + connectionId: string + clearedAt: number + } + server.ingestRemote( + { paneKey: GOOD_PANE, payload: { state: 'working', agentType: 'claude' }, isReplay: true }, + 'ssh-a' + ) + + expect(clear).toMatchObject({ transient: true, connectionId: 'ssh-a' }) + expect(server.getStatusSnapshot()).toEqual([ + expect.objectContaining({ paneKey: PANE, connectionId: 'ssh-b' }), + expect.objectContaining({ + paneKey: GOOD_PANE, + connectionId: 'ssh-a', + receivedAt: clear.clearedAt + 1 + }) + ]) + vi.useRealTimers() }) it('drops cached statuses and pane-scoped listener caches under one tab prefix', () => { @@ -6111,6 +6196,68 @@ describe('Last-status persistence', () => { } }) + it('keeps SSH status ordering monotonic across hydration and clock rollback', async () => { + const now = 1_700_000_000_000 + vi.spyOn(Date, 'now').mockReturnValue(now) + mkdirSync(join(userDataPath, 'agent-hooks'), { recursive: true }) + const receivedAt = now + 1_000 + writeFileSync( + lastStatusPath(), + JSON.stringify({ + version: 2, + entries: { + [PANE]: { + paneKey: PANE, + tabId: 'tab-1', + worktreeId: 'wt-1', + connectionId: 'ssh-a', + receivedAt, + stateStartedAt: receivedAt, + payload: { state: 'working', prompt: 'before restart', agentType: 'claude' } + } + } + }), + 'utf8' + ) + + const server = new AgentHookServer() + await server.start({ env: 'production', userDataPath }) + try { + const clearListener = vi.fn() + server.setPaneStatusClearListener(clearListener) + server.ingestRemote( + { paneKey: PANE, payload: { state: 'working', agentType: 'codex' } }, + 'ssh-a' + ) + + expect(server.getStatusSnapshot()[0]?.receivedAt).toBe(receivedAt + 1) + server.clearStatusEntriesForConnection('ssh-a') + const clearedAt = receivedAt + 2 + expect(clearListener).toHaveBeenCalledWith({ + transient: true, + connectionId: 'ssh-a', + clearedAt + }) + server.ingestRemote( + { + paneKey: GOOD_PANE, + isReplay: true, + payload: { state: 'working', agentType: 'claude' } + }, + 'ssh-a' + ) + expect(server.getStatusSnapshot()).toEqual([ + expect.objectContaining({ + paneKey: GOOD_PANE, + connectionId: 'ssh-a', + receivedAt: clearedAt + 1 + }) + ]) + } finally { + server.stop() + } + }) + it('drops persisted idle Claude child rows from hydration replay', async () => { mkdirSync(join(userDataPath, 'agent-hooks'), { recursive: true }) const receivedAt = recentTs() diff --git a/src/main/agent-hooks/server.ts b/src/main/agent-hooks/server.ts index e810a45f3..4c5d4c80c 100644 --- a/src/main/agent-hooks/server.ts +++ b/src/main/agent-hooks/server.ts @@ -44,6 +44,7 @@ import { import type { AgentHookSource } from '../../shared/agent-hook-relay' import { AGENT_STATUS_STALE_AFTER_MS, + type AgentStatusClearIpcPayload, type AgentStatusIpcPayload, type AgentType, type AgentStatusState, @@ -93,7 +94,7 @@ export type AgentHookStatusChangeEntry = { } type StatusChangeListener = (statuses: AgentHookStatusChangeEntry[]) => void -type PaneStatusClearListener = (paneKey: string) => void +type PaneStatusClearListener = (clear: AgentStatusClearIpcPayload) => void type PaneKeyAliasPersistenceListener = (entries: LegacyPaneKeyAliasEntry[]) => void type PaneKeyAliasEntry = { stablePaneKey: string @@ -528,6 +529,7 @@ export class AgentHookServer { private promptSentHashSalt = randomBytes(16).toString('hex') private closedAgentStatusTabIds = new Set() private closedAgentStatusPaneKeys = new Set() + private connectionTimestampWatermarkById = new Map() // Why: identity check — skip writes when the JSON-stringified contents // exactly match the last successful disk write. Cheap protection against // re-firing trailing timers when nothing changed. @@ -896,7 +898,15 @@ export class AgentHookServer { const previous = this.state.lastStatusByPaneKey.get(payload.paneKey) as | EnrichedAgentHookEventPayload | undefined - const now = Date.now() + const connectionClearWatermark = payload.connectionId + ? this.connectionTimestampWatermarkById.get(payload.connectionId) + : undefined + // Why: Date.now() can repeat across disconnect and reconnect. A remote + // replay must sort strictly after its connection's transient clear. + const now = Math.max(Date.now(), (connectionClearWatermark ?? -1) + 1) + if (payload.connectionId) { + this.connectionTimestampWatermarkById.set(payload.connectionId, now) + } if (payload.providerSessionOnly) { // Why: Pi session_start must replace stale turn state and survive snapshot // replay, but it must not emit prompt telemetry or a fabricated status. @@ -1305,7 +1315,7 @@ export class AgentHookServer { this.scheduleStatusPersist() this.notifyStatusChangeListeners() for (const paneKey of clearedStatusPaneKeys) { - this.onPaneStatusCleared?.(paneKey) + this.onPaneStatusCleared?.({ paneKey }) } } } @@ -1677,6 +1687,7 @@ export class AgentHookServer { this.promptSentDedupeByPaneKey.clear() this.closedAgentStatusTabIds.clear() this.closedAgentStatusPaneKeys.clear() + this.connectionTimestampWatermarkById.clear() this.legacyPaneKeyAliases.clear() clearAllListenerCaches(this.state) this.notifyStatusChangeListeners() @@ -1690,21 +1701,69 @@ export class AgentHookServer { * lands. clearPaneState (which wipes all three caches) is the right shape * only for PTY-teardown. */ dropStatusEntry(paneKey: string): void { - const resolvedPaneKey = this.resolvePaneKeyAlias(paneKey) - if (!this.state.lastStatusByPaneKey.has(resolvedPaneKey)) { + if (!this.deleteStatusEntry(paneKey)) { return } - const existing = this.state.lastStatusByPaneKey.get(resolvedPaneKey) - this.state.lastStatusByPaneKey.delete(resolvedPaneKey) - this.clearAssistantMessageRetry(resolvedPaneKey) - this.runtimeObservedStatusPaneKeys.delete(resolvedPaneKey) - if (existing?.payload.state === 'done') { - this.promptSentDedupeByPaneKey.delete(resolvedPaneKey) - } this.scheduleStatusPersist() this.notifyStatusChangeListeners() } + /** Clear statuses proven to belong to one lost SSH transport. */ + clearStatusEntriesForConnection(connectionId: string): void { + const normalizedConnectionId = connectionId.trim() + if (normalizedConnectionId.length === 0) { + return + } + const clearedAt = Math.max( + Date.now(), + (this.connectionTimestampWatermarkById.get(normalizedConnectionId) ?? -1) + 1 + ) + this.connectionTimestampWatermarkById.set(normalizedConnectionId, clearedAt) + let statusChanged = false + for (const [paneKey, rawEntry] of this.state.lastStatusByPaneKey) { + const entry = rawEntry as EnrichedAgentHookEventPayload + // Why: legacy/unstamped rows cannot be safely attributed to one host. + // Leave them for normal pane teardown instead of risking cross-host loss. + if (entry.connectionId !== normalizedConnectionId) { + continue + } + if (this.deleteStatusEntry(paneKey)) { + statusChanged = true + } + } + if (statusChanged) { + // Why: one disconnect can own many panes; persist and notify subscribers + // once so cleanup cost does not fan out with the pane count. + this.scheduleStatusPersist() + this.notifyStatusChangeListeners() + } + // Why: another host can overwrite the same pane key in main's cache while + // renderer still shows this connection's older row. Always send the + // connection cutoff, even when no current main entry matched. + this.onPaneStatusCleared?.({ + transient: true, + connectionId: normalizedConnectionId, + clearedAt + }) + } + + private deleteStatusEntry(paneKey: string): EnrichedAgentHookEventPayload | null { + const resolvedPaneKey = this.resolvePaneKeyAlias(paneKey) + const existing = this.state.lastStatusByPaneKey.get(resolvedPaneKey) as + | EnrichedAgentHookEventPayload + | undefined + if (!existing) { + return null + } + this.state.lastStatusByPaneKey.delete(resolvedPaneKey) + this.clearAssistantMessageRetry(resolvedPaneKey) + this.runtimeObservedStatusPaneKeys.delete(resolvedPaneKey) + if (existing.payload.state === 'done') { + this.promptSentDedupeByPaneKey.delete(resolvedPaneKey) + } + return existing + } + dropStatusEntriesByTabPrefix(tabId: string): void { this.markTabClosedForAgentStatus(tabId) const paneKeysToClear = new Set() @@ -1802,7 +1861,7 @@ export class AgentHookServer { this.runtimeObservedStatusPaneKeys.delete(resolvedPaneKey) this.scheduleStatusPersist() this.notifyStatusChangeListeners() - this.onPaneStatusCleared?.(resolvedPaneKey) + this.onPaneStatusCleared?.({ paneKey: resolvedPaneKey }) } } @@ -1917,6 +1976,15 @@ export class AgentHookServer { entry.payload = hydratedPayload } this.state.lastStatusByPaneKey.set(resolvedPaneKey, entry) + if (entry.connectionId) { + // Why: a restarted process can observe an earlier wall clock; seed + // transport ordering so new events and clears stay after disk state. + const previousWatermark = this.connectionTimestampWatermarkById.get(entry.connectionId) + this.connectionTimestampWatermarkById.set( + entry.connectionId, + Math.max(previousWatermark ?? -1, entry.receivedAt) + ) + } // Why: preserve only working children across restart. Live activity // confirms them; a later complete inventory may reap stale seeds. if (entry.payload.subagents) { @@ -2043,6 +2111,10 @@ export class AgentHookServer { _resetPromptSentDedupeForTests(): void { this.promptSentDedupeByPaneKey.clear() } + + _resetConnectionTimestampWatermarksForTests(): void { + this.connectionTimestampWatermarkById.clear() + } } export const agentHookServer = new AgentHookServer() @@ -2061,5 +2133,6 @@ export const _internals = { resetCachesForTests: (): void => { clearAllListenerCaches(agentHookServer._getStateForTests()) agentHookServer._resetPromptSentDedupeForTests() + agentHookServer._resetConnectionTimestampWatermarksForTests() } } diff --git a/src/main/index.ts b/src/main/index.ts index 1b289f410..9278eb043 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -1210,11 +1210,11 @@ function openMainWindow(): BrowserWindow { } } ) - agentHookServer.setPaneStatusClearListener((paneKey) => { + agentHookServer.setPaneStatusClearListener((clear) => { if (mainWindow?.isDestroyed()) { return } - mainWindow?.webContents.send('agentStatus:clear', { paneKey }) + mainWindow?.webContents.send('agentStatus:clear', clear) }) setMigrationUnsupportedPtyListener((event) => { if (mainWindow?.isDestroyed()) { diff --git a/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts b/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts index 91ee554df..d6c98bf03 100644 --- a/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts +++ b/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts @@ -210,6 +210,7 @@ describe('SshRelaySession agent hooks over a fake relay transport', () => { session = null relay = null agentHookServer.setListener(null) + agentHookServer.setPaneStatusClearListener(null) agentHookInternals.resetCachesForTests() warnSpy.mockRestore() if (previousRemoteHooksFlag === undefined) { @@ -272,6 +273,34 @@ describe('SshRelaySession agent hooks over a fake relay transport', () => { }) }) + it('clears stamped status on reconnect loss but not final shutdown', async () => { + const initialRelay = createFakeRelay() + relay = createFakeRelay() + vi.mocked(deployAndLaunchRelay) + .mockResolvedValueOnce({ transport: initialRelay.transport, platform: 'linux-x64' }) + .mockResolvedValueOnce({ transport: relay.transport, platform: 'linux-x64' }) + const clearListener = vi.fn() + agentHookServer.setPaneStatusClearListener(clearListener) + session = createSession('conn-clear') + await session.establish({} as SshConnection) + initialRelay.notifyAgentHook(makeEnvelope()) + await vi.waitFor(() => expect(agentHookServer.getStatusSnapshot()).toHaveLength(1)) + + await session.reconnect({} as SshConnection) + initialRelay.dispose() + + expect(agentHookServer.getStatusSnapshot()).toEqual([]) + expect(clearListener).toHaveBeenCalledOnce() + expect(clearListener).toHaveBeenCalledWith({ + transient: true, + connectionId: 'conn-clear', + clearedAt: expect.any(Number) + }) + session.dispose() + session = null + expect(clearListener).toHaveBeenCalledOnce() + }) + it('asks the fake relay for cached hook replay after the session wires its listener', async () => { relay = createFakeRelay() relay.replayEnvelopes.push( diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index 11f7f5f93..969878c12 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/src/main/ssh/ssh-relay-session.ts @@ -1020,6 +1020,10 @@ export class SshRelaySession { if (reason === 'shutdown') { clearPtyOwnershipForConnection(this.targetId) + } else { + // Why: handlers are detached above, so no late event can recreate a + // stamped status between this clear and reconnect replay. + agentHookServer.clearStatusEntriesForConnection(this.targetId) } const ptyProvider = getSshPtyProvider(this.targetId) diff --git a/src/preload/api-types.ts b/src/preload/api-types.ts index f25cf0f56..3911de36f 100644 --- a/src/preload/api-types.ts +++ b/src/preload/api-types.ts @@ -304,6 +304,7 @@ import type { CliInstallStatus } from '../shared/cli-install-types' import type { E2EConfig } from '../shared/e2e-config' import type { AgentHookInstallStatus } from '../shared/agent-hook-types' import type { + AgentStatusClearIpcPayload, AgentStatusIpcPayload, MigrationUnsupportedPtyEntry } from '../shared/agent-status-types' @@ -3211,8 +3212,8 @@ export type PreloadApi = { agentStatus: { /** Listen for agent status updates forwarded from native hook receivers. */ onSet: (callback: (data: AgentStatusIpcPayload) => void) => () => void - /** Listen for main-process pane teardown that evicted a cached hook status. */ - onClear: (callback: (data: { paneKey: string }) => void) => () => void + /** Listen for main-process cleanup that evicted cached hook status. */ + onClear: (callback: (data: AgentStatusClearIpcPayload) => void) => () => void /** Return the current main-process hook cache after renderer hydration. */ getSnapshot: () => Promise inferInterrupt: (request: AgentInterruptInferenceRequest) => Promise diff --git a/src/preload/index.ts b/src/preload/index.ts index 237090a44..dc8023b1e 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -142,6 +142,7 @@ import type { EnrichedDetectedPort } from '../shared/ssh-types' import type { + AgentStatusClearIpcPayload, AgentStatusIpcPayload, MigrationUnsupportedPtyEntry } from '../shared/agent-status-types' @@ -4401,8 +4402,8 @@ const api = { ipcRenderer.on('agentStatus:set', listener) return () => ipcRenderer.removeListener('agentStatus:set', listener) }, - onClear: (callback: (data: { paneKey: string }) => void): (() => void) => { - const listener = (_event: Electron.IpcRendererEvent, data: { paneKey: string }) => + onClear: (callback: (data: AgentStatusClearIpcPayload) => void): (() => void) => { + const listener = (_event: Electron.IpcRendererEvent, data: AgentStatusClearIpcPayload) => callback(data) ipcRenderer.on('agentStatus:clear', listener) return () => ipcRenderer.removeListener('agentStatus:clear', listener) diff --git a/src/renderer/src/hooks/useIpcEvents.test.ts b/src/renderer/src/hooks/useIpcEvents.test.ts index 16f5deadd..40e598417 100644 --- a/src/renderer/src/hooks/useIpcEvents.test.ts +++ b/src/renderer/src/hooks/useIpcEvents.test.ts @@ -12,6 +12,7 @@ import { resolveZoomTarget } from './useIpcEvents' import type { SleepingAgentLaunchConfig } from '../../../shared/agent-session-resume' +import type { AgentStatusClearIpcPayload } from '../../../shared/agent-status-types' import type { TuiAgent } from '../../../shared/types' import { makePaneKey } from '../../../shared/stable-pane-id' import { YOLO_TUI_AGENT_ARGS } from '../../../shared/tui-agent-permissions' @@ -4624,6 +4625,7 @@ describe('useIpcEvents agent status snapshot integration', () => { runtimePaneTitlesByTabId: {}, terminalLayoutsByTabId: {}, agentStatusByPaneKey: {}, + clearTransientAgentStatuses: vi.fn(), getAgentLaunchConfigForStatusMetadata: vi.fn(() => undefined), recentlyClosedAgentStatusTabIds: {}, repos: [], @@ -4638,7 +4640,7 @@ describe('useIpcEvents agent status snapshot integration', () => { function buildWindowApi(args: { onSet: (cb: (data: AgentStatusSetData) => void) => () => void - onClear?: (cb: (data: { paneKey: string }) => void) => () => void + onClear?: (cb: (data: AgentStatusClearIpcPayload) => void) => () => void getSnapshot?: () => Promise drop?: (paneKey: string) => void remoteWorkspace?: Record @@ -5715,7 +5717,9 @@ describe('useIpcEvents agent status snapshot integration', () => { const onSetListenerRef: { current: ((data: AgentStatusSetData) => void) | null } = { current: null } - const onClearListenerRef: { current: ((data: { paneKey: string }) => void) | null } = { + const onClearListenerRef: { + current: ((data: AgentStatusClearIpcPayload) => void) | null + } = { current: null } @@ -5805,9 +5809,136 @@ describe('useIpcEvents agent status snapshot integration', () => { expect(removeAgentStatus).toHaveBeenCalledWith(FUTURE_PANE_KEY) }) + it('blocks cleared snapshots across remount and accepts newer reconnect replay', async () => { + let resolveOldSnapshot!: (entries: AgentStatusSetData[]) => void + let resolveCurrentSnapshot!: (entries: AgentStatusSetData[]) => void + const oldSnapshot = new Promise((resolve) => { + resolveOldSnapshot = resolve + }) + const currentSnapshot = new Promise((resolve) => { + resolveCurrentSnapshot = resolve + }) + const effectCleanups: (() => void)[] = [] + const setAgentStatus = vi.fn() + const clearTransientAgentStatuses = vi.fn() + const onSetListenerRef: { current: ((data: AgentStatusSetData) => void) | null } = { + current: null + } + const onClearListenerRef: { + current: ((data: AgentStatusClearIpcPayload) => void) | null + } = { current: null } + const storeState: StoreLike = buildStoreState({ + setAgentStatus, + clearTransientAgentStatuses, + workspaceSessionReady: true, + settings: { terminalFontSize: 13, notifications: { enabled: false } }, + repos: [{ id: 'repo-1', connectionId: 'ssh-a' }], + worktreesByRepo: { 'repo-1': [{ id: 'wt-1', repoId: 'repo-1' }] }, + tabsByWorktree: { + 'wt-1': [{ id: 'tab-future', ptyId: 'pty-1', worktreeId: 'wt-1', title: 'Remote agent' }] + }, + terminalLayoutsByTabId: { + 'tab-future': { root: { type: 'leaf', leafId: FUTURE_LEAF_ID } } + } + }) + + vi.doMock('react', async () => { + const actual = await vi.importActual('react') + return { + ...actual, + useEffect: (effect: () => void | (() => void)) => { + const cleanup = effect() + effectCleanups.push(typeof cleanup === 'function' ? cleanup : () => {}) + } + } + }) + vi.doMock('../store', () => ({ + useAppStore: { + subscribe: vi.fn(() => () => {}), + getState: () => storeState + } + })) + stubAuxiliaryModules() + vi.stubGlobal( + 'window', + buildWindowApi({ + onSet: (callback) => { + onSetListenerRef.current = callback + return () => {} + }, + onClear: (callback) => { + onClearListenerRef.current = callback + return () => {} + }, + getSnapshot: vi.fn().mockReturnValueOnce(oldSnapshot).mockReturnValueOnce(currentSnapshot) + }) + ) + + const { useIpcEvents } = await import('./useIpcEvents') + useIpcEvents() + await Promise.resolve() + effectCleanups[0]?.() + useIpcEvents() + await Promise.resolve() + if (!onSetListenerRef.current || !onClearListenerRef.current) { + throw new Error('Expected agent status listeners to be registered') + } + + expect(() => + onClearListenerRef.current?.(null as unknown as AgentStatusClearIpcPayload) + ).not.toThrow() + expect(clearTransientAgentStatuses).not.toHaveBeenCalled() + + onClearListenerRef.current({ + transient: true, + connectionId: 'ssh-a', + clearedAt: 100 + }) + const staleEntry: AgentStatusSetData = { + paneKey: FUTURE_PANE_KEY, + state: 'working', + prompt: 'stale snapshot', + agentType: 'codex', + worktreeId: 'wt-1', + connectionId: 'ssh-a', + receivedAt: 100, + stateStartedAt: 90 + } + resolveOldSnapshot([staleEntry]) + resolveCurrentSnapshot([staleEntry]) + await Promise.resolve() + await Promise.resolve() + + expect(clearTransientAgentStatuses).toHaveBeenCalledWith('ssh-a', 100) + expect(setAgentStatus).not.toHaveBeenCalled() + + onSetListenerRef.current({ + paneKey: FUTURE_PANE_KEY, + state: 'working', + prompt: 'replayed', + agentType: 'codex', + worktreeId: 'wt-1', + connectionId: 'ssh-a', + receivedAt: 101, + stateStartedAt: 101 + }) + + expect(setAgentStatus).toHaveBeenCalledOnce() + expect(setAgentStatus).toHaveBeenCalledWith( + FUTURE_PANE_KEY, + expect.objectContaining({ prompt: 'replayed' }), + 'Remote agent', + { updatedAt: 101, stateStartedAt: 101 }, + expect.objectContaining({ worktreeId: 'wt-1', connectionId: 'ssh-a' }), + undefined + ) + }) + it('keeps a completed worktree-attributed row when main reports pane teardown', async () => { const removeAgentStatus = vi.fn() - const onClearListenerRef: { current: ((data: { paneKey: string }) => void) | null } = { + const onClearListenerRef: { + current: ((data: AgentStatusClearIpcPayload) => void) | null + } = { current: null } diff --git a/src/renderer/src/hooks/useIpcEvents.ts b/src/renderer/src/hooks/useIpcEvents.ts index cadc9d48a..3b83dc460 100644 --- a/src/renderer/src/hooks/useIpcEvents.ts +++ b/src/renderer/src/hooks/useIpcEvents.ts @@ -57,6 +57,7 @@ import { } from '@/lib/simulator-launch-coordination' import { normalizeAgentStatusPayload, + type AgentStatusClearIpcPayload, type AgentStatusIpcPayload, type ParsedAgentStatusPayload } from '../../../shared/agent-status-types' @@ -849,6 +850,8 @@ export function useIpcEvents(): void { } type AgentStatusApplyResult = 'applied' | 'pending' | 'dropped' const pendingAgentStatusEvents: PendingAgentStatusEvent[] = [] + const transientClearWatermarkByConnectionId = new Map() + let agentStatusEffectDisposed = false let pendingAgentStatusRetryTimer: ReturnType | null = null // Why: applyAgentStatus -> store.setAgentStatus notifies the store // subscriber synchronously, which re-enters flushPendingAgentStatuses while @@ -3147,6 +3150,15 @@ export function useIpcEvents(): void { const ownershipConnectionId = isWslHookRelayConnectionId(data.connectionId) ? null : data.connectionId + const transientClearWatermark = + typeof data.connectionId === 'string' + ? transientClearWatermarkByConnectionId.get(data.connectionId) + : undefined + // Why: delayed snapshots and queued relay events must not resurrect a + // status cleared by a newer disconnect from the same connection. + if (transientClearWatermark !== undefined && data.receivedAt <= transientClearWatermark) { + return 'dropped' + } const canAcceptPendingRemoteOwnership = ownershipConnectionId !== undefined && ownershipConnectionId !== null && @@ -3242,7 +3254,8 @@ export function useIpcEvents(): void { { tabId: ownerTabId, worktreeId: statusWorktreeId, - terminalHandle: data.terminalHandle + terminalHandle: data.terminalHandle, + ...(ownershipConnectionId !== undefined ? { connectionId: ownershipConnectionId } : {}) }, data.providerSession || data.launchToken ? { @@ -3288,7 +3301,7 @@ export function useIpcEvents(): void { const requestId = ++snapshotRequestId void getSnapshot() .then((entries) => { - if (requestId !== snapshotRequestId) { + if (agentStatusEffectDisposed || requestId !== snapshotRequestId) { return } const current = useAppStore.getState() @@ -3304,6 +3317,9 @@ export function useIpcEvents(): void { return } void getMigrationUnsupportedSnapshot().then((unsupportedEntries) => { + if (agentStatusEffectDisposed || requestId !== snapshotRequestId) { + return + } const unsupportedStore = useAppStore.getState() if (!unsupportedStore.workspaceSessionReady) { return @@ -3332,16 +3348,45 @@ export function useIpcEvents(): void { applyAgentStatus(data) }) ) - const unsubscribeAgentStatusClear = window.api.agentStatus.onClear?.((data) => { - if (typeof data?.paneKey !== 'string') { - return + const unsubscribeAgentStatusClear = window.api.agentStatus.onClear?.( + (data: AgentStatusClearIpcPayload) => { + if (typeof data !== 'object' || data === null) { + return + } + if ('transient' in data && data.transient === true) { + if ( + typeof data.connectionId !== 'string' || + data.connectionId.length === 0 || + !Number.isFinite(data.clearedAt) + ) { + return + } + const previousWatermark = + transientClearWatermarkByConnectionId.get(data.connectionId) ?? -1 + const effectiveWatermark = Math.max(previousWatermark, data.clearedAt) + transientClearWatermarkByConnectionId.set(data.connectionId, effectiveWatermark) + for (let index = pendingAgentStatusEvents.length - 1; index >= 0; index -= 1) { + const pending = pendingAgentStatusEvents[index].data + if ( + pending.connectionId === data.connectionId && + pending.receivedAt <= effectiveWatermark + ) { + pendingAgentStatusEvents.splice(index, 1) + } + } + useAppStore.getState().clearTransientAgentStatuses(data.connectionId, effectiveWatermark) + return + } + if (!('paneKey' in data) || typeof data.paneKey !== 'string') { + return + } + const store = useAppStore.getState() + if (store.agentStatusByPaneKey[data.paneKey]?.state === 'done') { + return + } + store.removeAgentStatus(data.paneKey) } - const store = useAppStore.getState() - if (store.agentStatusByPaneKey[data.paneKey]?.state === 'done') { - return - } - store.removeAgentStatus(data.paneKey) - }) + ) if (unsubscribeAgentStatusClear) { unsubs.push(unsubscribeAgentStatusClear) } @@ -3502,6 +3547,10 @@ export function useIpcEvents(): void { } return () => { + // Why: React remount can leave an older snapshot promise in flight. It + // must not write through after the replacement effect processes a clear. + agentStatusEffectDisposed = true + snapshotRequestId += 1 if (pendingAgentStatusRetryTimer !== null) { globalThis.clearTimeout(pendingAgentStatusRetryTimer) } diff --git a/src/renderer/src/store/slices/agent-status-ssh-connection-clear.test.ts b/src/renderer/src/store/slices/agent-status-ssh-connection-clear.test.ts new file mode 100644 index 000000000..774264c2c --- /dev/null +++ b/src/renderer/src/store/slices/agent-status-ssh-connection-clear.test.ts @@ -0,0 +1,97 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' + +import type { AppState } from '../types' +import { createTestStore } from './store-test-helpers' + +describe('agent status cleanup for a lost SSH connection', () => { + afterEach(() => { + vi.useRealTimers() + vi.restoreAllMocks() + }) + + it('clears one connection in one update while preserving newer and unstamped rows', () => { + vi.useFakeTimers() + const store = createTestStore() + const oldA = 'tab-a:11111111-1111-4111-8111-111111111111' + const secondA = 'tab-a2:22222222-2222-4222-8222-222222222222' + const newerA = 'tab-new:33333333-3333-4333-8333-333333333333' + const siblingB = 'tab-b:44444444-4444-4444-8444-444444444444' + const unstamped = 'tab-legacy:55555555-5555-4555-8555-555555555555' + for (const [paneKey, updatedAt, connectionId] of [ + [oldA, 10, 'ssh-a'], + [secondA, 20, 'ssh-a'], + [newerA, 31, 'ssh-a'], + [siblingB, 15, 'ssh-b'] + ] as const) { + store + .getState() + .setAgentStatus( + paneKey, + { state: 'working', prompt: paneKey, agentType: 'codex' }, + undefined, + { updatedAt }, + { connectionId } + ) + } + store + .getState() + .setAgentStatus( + unstamped, + { state: 'working', prompt: 'legacy', agentType: 'claude' }, + undefined, + { updatedAt: 5 } + ) + store.setState({ + agentLaunchConfigByPaneKey: { + [oldA]: { + launchConfig: { agentCommand: 'codex', agentArgs: '--full-auto', agentEnv: {} }, + registeredAt: 1, + identity: {} + } + }, + acknowledgedAgentsByPaneKey: { [oldA]: 2 }, + retentionSuppressedPaneKeys: { [oldA]: true } + } as Partial) + const subscriber = vi.fn() + const unsubscribe = store.subscribe(subscriber) + const queueMicrotaskSpy = vi.spyOn(globalThis, 'queueMicrotask').mockImplementation(() => {}) + + store.getState().clearTransientAgentStatuses('ssh-a', 30) + + unsubscribe() + expect(subscriber).toHaveBeenCalledOnce() + expect(queueMicrotaskSpy).toHaveBeenCalledOnce() + expect(store.getState().agentStatusByPaneKey[oldA]).toBeUndefined() + expect(store.getState().agentStatusByPaneKey[secondA]).toBeUndefined() + expect(store.getState().agentStatusByPaneKey[newerA]).toBeDefined() + expect(store.getState().agentStatusByPaneKey[siblingB]).toBeDefined() + expect(store.getState().agentStatusByPaneKey[unstamped]).toBeDefined() + expect(store.getState().agentLaunchConfigByPaneKey[oldA]).toBeDefined() + expect(store.getState().acknowledgedAgentsByPaneKey[oldA]).toBe(2) + expect(store.getState().retentionSuppressedPaneKeys[oldA]).toBe(true) + }) + + it('retains an accepted connection stamp across later unstamped pings', () => { + const store = createTestStore() + const paneKey = 'tab-a:11111111-1111-4111-8111-111111111111' + store + .getState() + .setAgentStatus( + paneKey, + { state: 'working', prompt: 'first', agentType: 'codex' }, + undefined, + { updatedAt: 1 }, + { connectionId: 'ssh-a' } + ) + store + .getState() + .setAgentStatus( + paneKey, + { state: 'working', prompt: 'ping', agentType: 'codex' }, + undefined, + { updatedAt: 2 } + ) + + expect(store.getState().agentStatusByPaneKey[paneKey]?.connectionId).toBe('ssh-a') + }) +}) diff --git a/src/renderer/src/store/slices/agent-status.ts b/src/renderer/src/store/slices/agent-status.ts index 14fa98fef..c6b34559f 100644 --- a/src/renderer/src/store/slices/agent-status.ts +++ b/src/renderer/src/store/slices/agent-status.ts @@ -162,7 +162,12 @@ export type AgentStatusSlice = { }, terminalTitle?: string, timing?: { updatedAt?: number; stateStartedAt?: number }, - routing?: { tabId?: string; worktreeId?: string; terminalHandle?: string }, + routing?: { + tabId?: string + worktreeId?: string + terminalHandle?: string + connectionId?: string | null + }, metadata?: { providerSession?: AgentProviderSessionMetadata launchConfig?: SleepingAgentLaunchConfig @@ -207,6 +212,9 @@ export type AgentStatusSlice = { * Used when a tab is closed — same prefix-sweep as cacheTimerByKey cleanup. */ removeAgentStatusByTabPrefix: (tabIdPrefix: string) => void + /** Remove stale live rows while preserving pane launch and resume identity. */ + clearTransientAgentStatuses: (connectionId: string, clearedAt: number) => void + /** Remove a single entry AND suppress re-retention on its next disappearance. * Used for USER-INITIATED teardown — the dashboard/hover X button, and * pane close — where the user is telling us "I'm done with this row". */ @@ -1785,6 +1793,11 @@ export const createAgentStatusSlice: StateCreator freshness.schedule()) }, + clearTransientAgentStatuses: (connectionId, clearedAt) => { + if (connectionId.length === 0 || !Number.isFinite(clearedAt)) { + return + } + let removed = false + set((s) => { + let next: Record | null = null + for (const [paneKey, existing] of Object.entries(s.agentStatusByPaneKey)) { + // Why: undefined is an unstamped legacy/renderer-owned row. Its host + // cannot be proven, so normal pane teardown remains its safe cleanup. + if (existing.connectionId !== connectionId || existing.updatedAt > clearedAt) { + continue + } + next ??= { ...s.agentStatusByPaneKey } + delete next[paneKey] + } + if (!next) { + return s + } + removed = true + // Why: transport loss is reversible. Keep launch, resume, retention, + // and acknowledgement maps intact for same-pane relay replay. + return { + agentStatusByPaneKey: next, + agentStatusEpoch: s.agentStatusEpoch + 1, + sortEpoch: s.sortEpoch + 1 + } + }) + if (removed) { + queueMicrotask(() => freshness.schedule()) + } + }, + dropAgentStatus: (paneKey) => { // Why: single sync read — zustand set is synchronous, so the value we // observe inside the set callback is the same one we would re-read via diff --git a/src/shared/agent-status-types.ts b/src/shared/agent-status-types.ts index 728d8e1b9..5eaec1317 100644 --- a/src/shared/agent-status-types.ts +++ b/src/shared/agent-status-types.ts @@ -121,6 +121,8 @@ export type AgentStatusEntry = { * present in a renderer; retaining this lets worktree-level UI still show * the live child agent instead of dropping it as unattributed. */ worktreeId?: string + /** Accepted transport authority for this live row; null means local. */ + connectionId?: string | null /** Tab attribution from the hook IPC payload, when available. */ tabId?: string terminalTitle?: string @@ -230,6 +232,15 @@ export type AgentStatusIpcPayload = ParsedAgentStatusPayload & { promptInteractionKey?: string } +/** Wire shape for ordinary pane teardown or a stamped SSH disconnect batch. */ +export type AgentStatusClearIpcPayload = + | { paneKey: string } + | { + transient: true + connectionId: string + clearedAt: number + } + /** Maximum character length for the toolName field. */ export const AGENT_STATUS_TOOL_NAME_MAX_LENGTH = 60 /** Maximum character length for the toolInput preview. */