From 4d299463b28e322a72dba5f1329b2826dea0e862 Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Mon, 13 Jul 2026 07:10:42 -0700 Subject: [PATCH] fix: fence stale PTY exit generations --- src/main/daemon/daemon-pty-adapter.test.ts | 15 ++ src/main/daemon/daemon-pty-adapter.ts | 62 ++++---- src/main/daemon/daemon-pty-router.ts | 17 +-- .../daemon/degraded-daemon-pty-provider.ts | 16 +- .../pty-provider-route-reconciliation.test.ts | 38 +++++ .../pty-provider-route-reconciliation.ts | 143 ++++++++++++++---- src/main/providers/local-pty-provider.test.ts | 26 ++++ src/main/providers/local-pty-provider.ts | 20 ++- 8 files changed, 256 insertions(+), 81 deletions(-) create mode 100644 src/main/daemon/pty-provider-route-reconciliation.test.ts diff --git a/src/main/daemon/daemon-pty-adapter.test.ts b/src/main/daemon/daemon-pty-adapter.test.ts index 66989bbf486..a578ad0297a 100644 --- a/src/main/daemon/daemon-pty-adapter.test.ts +++ b/src/main/daemon/daemon-pty-adapter.test.ts @@ -557,6 +557,21 @@ describe('DaemonPtyAdapter (IPtyProvider)', () => { await waitFor(() => exits.length > 0) expect(exits[0]).toEqual({ id, code: 42 }) }) + + it('suppresses an old exit when targeted readback proves a same-id replacement', async () => { + const exits: { id: string; code: number }[] = [] + adapter.onExit((payload) => exits.push(payload)) + const { id } = await adapter.spawn({ sessionId: 'same-id', cols: 80, rows: 24 }) + vi.spyOn(adapter, 'getAppliedSize').mockResolvedValueOnce({ cols: 80, rows: 24 }) + const internals = adapter as unknown as { + handleExitEvent(sessionId: string, code: number): Promise + } + + await internals.handleExitEvent(id, 0) + + expect(adapter.hasPty(id)).toBe(true) + expect(exits).toEqual([]) + }) }) describe('spawn with sessionId (reattach)', () => { diff --git a/src/main/daemon/daemon-pty-adapter.ts b/src/main/daemon/daemon-pty-adapter.ts index c66894d6141..b67f0b9f18e 100644 --- a/src/main/daemon/daemon-pty-adapter.ts +++ b/src/main/daemon/daemon-pty-adapter.ts @@ -1373,39 +1373,41 @@ export class DaemonPtyAdapter implements IPtyProvider { fact: event.payload }) } else if (event.event === 'exit') { - this.activeSessionIds.delete(event.sessionId) - this.dirtySessionVersions.delete(event.sessionId) - // Why: an exited session must not be owed a resume on reconnect — a - // reused sessionId would receive a stray resumePty. Same for the - // background set: a reused id must start un-thinned. - this.pausedProducerSessionIds.delete(event.sessionId) - this.producerResumesOwedOnReconnect.delete(event.sessionId) - this.backgroundedSessionIds.delete(event.sessionId) - if (!this.sleepRestoreSessionIds.has(event.sessionId)) { - this.coldRestoreCache.delete(event.sessionId) - this.paneIdentityBySessionId.delete(event.sessionId) - } - // Why: an exited session can never be checkpointed again, so its pending - // full-checkpoint flag is dead state. Without this, a cold-restored - // session that exits before its first checkpoint leaks a permanent entry. - this.sessionsNeedingFullCheckpoint.delete(event.sessionId) - // Why: a reused sessionId (renderer respawns a persisted ptyId) must - // not inherit the dead session's snapshot cooldown. - this.lastFullCheckpointAt.delete(event.sessionId) - this.stopCheckpointTimerIfIdle() - if (this.historyManager) { - void this.historyManager - .closeSession(event.sessionId, event.payload.code) - .catch((err) => console.warn('[history] closeSession failed:', event.sessionId, err)) - } - this.initialCwds.delete(event.sessionId) - // oxlint-disable-next-line unicorn/no-useless-spread -- copy-safe: listeners may unsubscribe during iteration - for (const listener of [...this.exitListeners]) { - listener({ id: event.sessionId, code: event.payload.code }) - } + void this.handleExitEvent(event.sessionId, event.payload.code) } }) } + + private async handleExitEvent(sessionId: string, code: number): Promise { + // Why: control and stream sockets can reorder a replacement create reply + // ahead of the old process's queued exit. Targeted liveness proof prevents + // that stale exit from deleting the replacement's same-id ownership. + if ((await this.getAppliedSize(sessionId)) !== null) { + return + } + this.activeSessionIds.delete(sessionId) + this.dirtySessionVersions.delete(sessionId) + this.pausedProducerSessionIds.delete(sessionId) + this.producerResumesOwedOnReconnect.delete(sessionId) + this.backgroundedSessionIds.delete(sessionId) + if (!this.sleepRestoreSessionIds.has(sessionId)) { + this.coldRestoreCache.delete(sessionId) + this.paneIdentityBySessionId.delete(sessionId) + } + this.sessionsNeedingFullCheckpoint.delete(sessionId) + this.lastFullCheckpointAt.delete(sessionId) + this.stopCheckpointTimerIfIdle() + if (this.historyManager) { + void this.historyManager + .closeSession(sessionId, code) + .catch((err) => console.warn('[history] closeSession failed:', sessionId, err)) + } + this.initialCwds.delete(sessionId) + // oxlint-disable-next-line unicorn/no-useless-spread -- copy-safe: listeners may unsubscribe during iteration + for (const listener of [...this.exitListeners]) { + listener({ id: sessionId, code }) + } + } } // Why: ENOENT/ECONNREFUSED with syscall 'connect' mean the socket is diff --git a/src/main/daemon/daemon-pty-router.ts b/src/main/daemon/daemon-pty-router.ts index cea03a059f3..a54d10bfd2a 100644 --- a/src/main/daemon/daemon-pty-router.ts +++ b/src/main/daemon/daemon-pty-router.ts @@ -2,10 +2,10 @@ import type { DaemonPtyAdapter } from './daemon-pty-adapter' import { appendProviderReconciliationIds, bindProviderRoute, - bindProviderRouteIfAbsent, deleteProviderRoute, + discoverProviderSessionsAndBindRoutes, listProviderProcessesAndReconcileRoutes, - reconcileProviderRoutesAfterStartup, + reconcileProviderRoutesOnStartup, type ProviderRoute } from './pty-provider-route-reconciliation' import type { @@ -63,10 +63,7 @@ export class DaemonPtyRouter implements IPtyProvider { async discoverLegacySessions(): Promise { for (const adapter of this.legacy) { try { - const sessions = await adapter.listProcesses() - for (const session of sessions) { - bindProviderRouteIfAbsent(this.sessionAdapters, session.id, adapter) - } + await discoverProviderSessionsAndBindRoutes(adapter, this.sessionAdapters) } catch (error) { console.warn('[daemon] Failed to discover legacy daemon sessions', error) } @@ -236,10 +233,12 @@ export class DaemonPtyRouter implements IPtyProvider { const alive: string[] = [] const killed: string[] = [] for (const adapter of this.allAdapters()) { - const routesAtStart = new Map(this.sessionAdapters) - const result = await adapter.reconcileOnStartup(validWorktreeIds) + const result = await reconcileProviderRoutesOnStartup( + adapter, + this.sessionAdapters, + validWorktreeIds + ) appendProviderReconciliationIds({ alive, killed }, result) - reconcileProviderRoutesAfterStartup(this.sessionAdapters, routesAtStart, adapter, result) } return { alive, killed } } diff --git a/src/main/daemon/degraded-daemon-pty-provider.ts b/src/main/daemon/degraded-daemon-pty-provider.ts index 6ea8601f34a..6e3c2a7a21d 100644 --- a/src/main/daemon/degraded-daemon-pty-provider.ts +++ b/src/main/daemon/degraded-daemon-pty-provider.ts @@ -71,10 +71,7 @@ export class DegradedDaemonPtyProvider implements IPtyProvider { async discoverDaemonSessions(): Promise { for (const adapter of this.allDaemonAdapters()) { try { - const sessions = await adapter.listProcesses() - for (const session of sessions) { - PtyRoutes.bindProviderRouteIfAbsent(this.sessionProviders, session.id, adapter) - } + await PtyRoutes.discoverProviderSessionsAndBindRoutes(adapter, this.sessionProviders) } catch (error) { console.warn('[daemon] Failed to discover degraded daemon sessions', error) } @@ -263,15 +260,12 @@ export class DegradedDaemonPtyProvider implements IPtyProvider { const alive: string[] = [] const killed: string[] = [] for (const adapter of this.allDaemonAdapters()) { - const routesAtStart = new Map(this.sessionProviders) - const result = await adapter.reconcileOnStartup(validWorktreeIds) - PtyRoutes.appendProviderReconciliationIds({ alive, killed }, result) - PtyRoutes.reconcileProviderRoutesAfterStartup( - this.sessionProviders, - routesAtStart, + const result = await PtyRoutes.reconcileProviderRoutesOnStartup( adapter, - result + this.sessionProviders, + validWorktreeIds ) + PtyRoutes.appendProviderReconciliationIds({ alive, killed }, result) } return { alive, killed } } diff --git a/src/main/daemon/pty-provider-route-reconciliation.test.ts b/src/main/daemon/pty-provider-route-reconciliation.test.ts new file mode 100644 index 00000000000..5431f1577b6 --- /dev/null +++ b/src/main/daemon/pty-provider-route-reconciliation.test.ts @@ -0,0 +1,38 @@ +import { describe, expect, it } from 'vitest' +import { + bindProviderRoute, + captureProviderRouteSnapshot, + deleteProviderRoute, + reconcileProviderRoutesAfterStartup, + type ProviderRoute +} from './pty-provider-route-reconciliation' + +describe('provider route reconciliation generations', () => { + it('fences absent-present-absent ABA during startup readback', () => { + const provider = { label: 'startup-owner' } + const routes = new Map>() + const snapshot = captureProviderRouteSnapshot(routes) + + bindProviderRoute(routes, 'same-id', provider) + deleteProviderRoute(routes, 'same-id') + reconcileProviderRoutesAfterStartup(routes, snapshot, provider, { + alive: ['same-id'], + killed: [] + }) + snapshot.dispose() + + expect(routes.has('same-id')).toBe(false) + }) + + it('stops retaining route mutation ids after snapshot disposal', () => { + const provider = { label: 'owner' } + const routes = new Map>() + const snapshot = captureProviderRouteSnapshot(routes) + bindProviderRoute(routes, 'observed', provider) + snapshot.dispose() + + bindProviderRoute(routes, 'after-dispose', provider) + + expect([...snapshot.mutatedIds]).toEqual(['observed']) + }) +}) diff --git a/src/main/daemon/pty-provider-route-reconciliation.ts b/src/main/daemon/pty-provider-route-reconciliation.ts index 124a042b3d8..3442a922160 100644 --- a/src/main/daemon/pty-provider-route-reconciliation.ts +++ b/src/main/daemon/pty-provider-route-reconciliation.ts @@ -2,6 +2,53 @@ import type { IPtyProvider, PtyProcessInfo } from '../providers/types' export type ProviderRoute = Readonly<{ provider: T }> +type RouteMutationObserver = { mutatedIds: Set } + +export type ProviderRouteSnapshot = RouteMutationObserver & { + routesAtStart: Map> + dispose: () => void +} + +const routeMutationObservers = new WeakMap>() + +function markProviderRouteMutation(routes: Map>, id: string): void { + for (const observer of routeMutationObservers.get(routes) ?? []) { + observer.mutatedIds.add(id) + } +} + +export function captureProviderRouteSnapshot( + routes: Map> +): ProviderRouteSnapshot { + const observer: RouteMutationObserver = { mutatedIds: new Set() } + const observers = routeMutationObservers.get(routes) ?? new Set() + observers.add(observer) + routeMutationObservers.set(routes, observers) + let active = true + return { + routesAtStart: new Map(routes), + mutatedIds: observer.mutatedIds, + dispose: () => { + if (!active) { + return + } + active = false + observers.delete(observer) + if (observers.size === 0) { + routeMutationObservers.delete(routes) + } + } + } +} + +function providerRouteIsUnchanged( + routes: Map>, + snapshot: ProviderRouteSnapshot, + id: string +): boolean { + return !snapshot.mutatedIds.has(id) && routes.get(id) === snapshot.routesAtStart.get(id) +} + export function bindProviderRoute( routes: Map>, id: string, @@ -10,20 +57,11 @@ export function bindProviderRoute( // Why: the immutable binding object is the generation token; rebinding the // same id to the same provider must still fence older async readbacks. const route = { provider } + markProviderRouteMutation(routes, id) routes.set(id, route) return route } -export function bindProviderRouteIfAbsent( - routes: Map>, - id: string, - provider: T -): void { - if (!routes.has(id)) { - bindProviderRoute(routes, id, provider) - } -} - export function deleteProviderRoute( routes: Map>, id: string, @@ -32,30 +70,73 @@ export function deleteProviderRoute( if (expectedRoute && routes.get(id) !== expectedRoute) { return false } - return routes.delete(id) + const deleted = routes.delete(id) + if (deleted) { + markProviderRouteMutation(routes, id) + } + return deleted } export function reconcileProviderRoutesAfterStartup( routes: Map>, - routesAtStart: Map>, + snapshot: ProviderRouteSnapshot, provider: T, result: { alive: string[]; killed: string[] } ): void { // Why: startup reconciliation may overlap a respawn; only its unchanged // binding generation may be replaced or removed by the older readback. for (const id of result.alive) { - if (routes.get(id) === routesAtStart.get(id)) { + if (providerRouteIsUnchanged(routes, snapshot, id)) { bindProviderRoute(routes, id, provider) } } for (const id of result.killed) { - const routeAtStart = routesAtStart.get(id) + const routeAtStart = snapshot.routesAtStart.get(id) if (routeAtStart) { deleteProviderRoute(routes, id, routeAtStart) } } } +export async function discoverProviderSessionsAndBindRoutes( + provider: T, + routes: Map> +): Promise { + const snapshot = captureProviderRouteSnapshot(routes) + try { + const sessions = await provider.listProcesses() + for (const session of sessions) { + if (providerRouteIsUnchanged(routes, snapshot, session.id)) { + bindProviderRoute(routes, session.id, provider) + } + } + } finally { + snapshot.dispose() + } +} + +type StartupReconcilingProvider = IPtyProvider & { + reconcileOnStartup(validWorktreeIds: Set): Promise<{ alive: string[]; killed: string[] }> +} + +export async function reconcileProviderRoutesOnStartup< + T extends IPtyProvider, + P extends T & StartupReconcilingProvider +>( + provider: P, + routes: Map>, + validWorktreeIds: Set +): Promise<{ alive: string[]; killed: string[] }> { + const snapshot = captureProviderRouteSnapshot(routes) + try { + const result = await provider.reconcileOnStartup(validWorktreeIds) + reconcileProviderRoutesAfterStartup(routes, snapshot, provider, result) + return result + } finally { + snapshot.dispose() + } +} + export function appendProviderReconciliationIds( target: { alive: string[]; killed: string[] }, result: { alive: string[]; killed: string[] } @@ -74,22 +155,28 @@ export async function listProviderProcessesAndReconcileRoutes> ): Promise { - const routesAtStart = [...routes] - const listings = await Promise.all( - providers.map(async (provider) => ({ provider, sessions: await provider.listProcesses() })) - ) - const liveIdsByProvider = new Map( - listings.map(({ provider, sessions }) => [provider, new Set(sessions.map(({ id }) => id))]) - ) - for (const [id, route] of routesAtStart) { - // Why: spawn/exit can mutate routing while remote listings are in flight. - // Compare the exact binding, not only its provider: ids can be rebound to - // a new process on the same provider while this snapshot is pending. - if (routes.get(id) === route && !liveIdsByProvider.get(route.provider)?.has(id)) { - routes.delete(id) + const snapshot = captureProviderRouteSnapshot(routes) + try { + const listings = await Promise.all( + providers.map(async (provider) => ({ provider, sessions: await provider.listProcesses() })) + ) + const liveIdsByProvider = new Map( + listings.map(({ provider, sessions }) => [provider, new Set(sessions.map(({ id }) => id))]) + ) + for (const [id, route] of snapshot.routesAtStart) { + // Why: spawn/exit can mutate routing while remote listings are in flight. + // Exact snapshot mutation tracking also fences absent-present-absent ABA. + if ( + providerRouteIsUnchanged(routes, snapshot, id) && + !liveIdsByProvider.get(route.provider)?.has(id) + ) { + deleteProviderRoute(routes, id, route) + } } + return listings.flatMap(({ sessions }) => sessions) + } finally { + snapshot.dispose() } - return listings.flatMap(({ sessions }) => sessions) } export function providerSessionIds( diff --git a/src/main/providers/local-pty-provider.test.ts b/src/main/providers/local-pty-provider.test.ts index 29731e1cb50..1e015a95e26 100644 --- a/src/main/providers/local-pty-provider.test.ts +++ b/src/main/providers/local-pty-provider.test.ts @@ -992,6 +992,32 @@ describe('LocalPtyProvider', () => { expect(provider.hasPty(id)).toBe(false) }) + it('does not let a queued old exit delete a same-id replacement', async () => { + let staleExit: ((event: { exitCode: number }) => void) | undefined + const oldProc = { + ...mockProc, + pid: 101, + kill: vi.fn(), + onExit: vi.fn((callback) => { + staleExit = callback + return { dispose: vi.fn() } + }) + } + const replacementProc = { ...mockProc, pid: 202, kill: vi.fn() } + spawnMock.mockReturnValueOnce(oldProc).mockReturnValueOnce(replacementProc) + const exits = vi.fn() + provider.onExit(exits) + + await provider.spawn({ sessionId: 'same-id', cols: 80, rows: 24 }) + provider.killOrphanedPtys(provider.advanceGeneration()) + await provider.spawn({ sessionId: 'same-id', cols: 80, rows: 24 }) + staleExit?.({ exitCode: 0 }) + + expect(provider.getPtyProcess('same-id')).toBe(replacementProc) + expect(provider.hasPty('same-id')).toBe(true) + expect(exits).not.toHaveBeenCalled() + }) + it('invokes onExit callback via the node-pty exit handler', async () => { const onExit = vi.fn() provider.configure({ onExit }) diff --git a/src/main/providers/local-pty-provider.ts b/src/main/providers/local-pty-provider.ts index df410e4196c..554edf1bdb8 100644 --- a/src/main/providers/local-pty-provider.ts +++ b/src/main/providers/local-pty-provider.ts @@ -187,7 +187,10 @@ function getWslContextFromPreferredDistro( /** * Removes all local tracking state for a PTY id after teardown. */ -function clearPtyState(id: string): void { +function clearPtyState(id: string, expectedProc?: pty.IPty): boolean { + if (expectedProc && ptyProcesses.get(id) !== expectedProc) { + return false + } runPtyCleanup(id) disposePtyListeners(id) ptyProcesses.delete(id) @@ -198,6 +201,7 @@ function clearPtyState(id: string): void { ptyLoadGeneration.delete(id) ptyShutdownInProgress.delete(id) ptyExitDuringShutdown.delete(id) + return true } /** @@ -811,6 +815,16 @@ export class LocalPtyProvider implements IPtyProvider { if (process.platform !== 'win32') { ;(proc as unknown as { kill: (sig?: string) => void }).kill = () => {} } + if (ptyProcesses.get(id) !== proc) { + // Why: native exit delivery can trail a same-id replacement. Clean + // only this old handle; id-keyed maps now belong to the new process. + startupCommandCleanup?.() + for (const disposable of disposables) { + disposable.dispose() + } + destroyPtyProcess(proc) + return + } if (ptyShutdownInProgress.has(id)) { ptyExitDuringShutdown.set(id, exitCode) return @@ -820,7 +834,7 @@ export class LocalPtyProvider implements IPtyProvider { shellReadyTimeout = null } startupCommandCleanup?.() - clearPtyState(id) + clearPtyState(id, proc) // Why: release the master ptmx fd on the natural-exit path — without // this, a shell that exits cleanly (the common case) never releases its // fd until the next GC. See docs/fix-pty-fd-leak.md. @@ -942,7 +956,7 @@ export class LocalPtyProvider implements IPtyProvider { if (exitObserved) { disposePtyListeners(id) destroyPtyProcess(proc, { alreadyKilled: true }) - clearPtyState(id) + clearPtyState(id, proc) this.opts.onExit?.(id, exitCode) for (const cb of exitListeners) { cb({ id, code: exitCode })