diff --git a/src/main/daemon/daemon-pty-router.test.ts b/src/main/daemon/daemon-pty-router.test.ts index f088889e7bb..b85f940179f 100644 --- a/src/main/daemon/daemon-pty-router.test.ts +++ b/src/main/daemon/daemon-pty-router.test.ts @@ -157,6 +157,26 @@ describe('DaemonPtyRouter', () => { expect(exit).not.toHaveBeenCalled() }) + it('retains a same-adapter replacement across an older process listing', async () => { + const current = createAdapter('current') + const legacy = createAdapter('legacy', ['shared-session']) + const router = new DaemonPtyRouter({ current, legacy: [legacy] }) + await router.discoverLegacySessions() + let finishLegacyListing!: (sessions: []) => void + vi.mocked(legacy.listProcesses).mockImplementationOnce( + () => new Promise((resolve) => (finishLegacyListing = resolve)) + ) + + const staleListing = router.listProcesses() + await router.spawn({ sessionId: 'shared-session', cols: 80, rows: 24 }) + finishLegacyListing([]) + await staleListing + router.write('shared-session', 'replacement') + + expect(legacy.write).toHaveBeenCalledWith('shared-session', 'replacement') + expect(current.write).not.toHaveBeenCalledWith('shared-session', 'replacement') + }) + it('routes fresh foreground confirmation to the session-owning daemon', async () => { const current = createAdapter('current', ['current-session']) const legacy = createAdapter('legacy', ['legacy-session']) diff --git a/src/main/daemon/daemon-pty-router.ts b/src/main/daemon/daemon-pty-router.ts index a5286e9537b..cea03a059f3 100644 --- a/src/main/daemon/daemon-pty-router.ts +++ b/src/main/daemon/daemon-pty-router.ts @@ -1,5 +1,13 @@ import type { DaemonPtyAdapter } from './daemon-pty-adapter' -import { listProviderProcessesAndReconcileRoutes } from './pty-provider-route-reconciliation' +import { + appendProviderReconciliationIds, + bindProviderRoute, + bindProviderRouteIfAbsent, + deleteProviderRoute, + listProviderProcessesAndReconcileRoutes, + reconcileProviderRoutesAfterStartup, + type ProviderRoute +} from './pty-provider-route-reconciliation' import type { IPtyProvider, PtyBackgroundStreamEvent, @@ -14,7 +22,7 @@ export class DaemonPtyRouter implements IPtyProvider { readonly requiresShutdownExitProof = true private current: DaemonPtyAdapter private legacy: DaemonPtyAdapter[] - private sessionAdapters = new Map() + private sessionAdapters = new Map>() private unsubscribers: (() => void)[] = [] private dataListeners: ((payload: { id: string @@ -37,12 +45,12 @@ export class DaemonPtyRouter implements IPtyProvider { adapter.onExit((payload) => { // Why: a late exit from a legacy daemon must not delete a route that // has already been rebound to a replacement provider generation. - const routed = this.sessionAdapters.get(payload.id) - if (routed && routed !== adapter) { + const route = this.sessionAdapters.get(payload.id) + if (route && route.provider !== adapter) { return } - if (routed === adapter) { - this.sessionAdapters.delete(payload.id) + if (route?.provider === adapter) { + deleteProviderRoute(this.sessionAdapters, payload.id, route) } for (const listener of this.exitListeners) { listener(payload) @@ -57,7 +65,7 @@ export class DaemonPtyRouter implements IPtyProvider { try { const sessions = await adapter.listProcesses() for (const session of sessions) { - this.sessionAdapters.set(session.id, adapter) + bindProviderRouteIfAbsent(this.sessionAdapters, session.id, adapter) } } catch (error) { console.warn('[daemon] Failed to discover legacy daemon sessions', error) @@ -66,10 +74,10 @@ export class DaemonPtyRouter implements IPtyProvider { } async spawn(opts: PtySpawnOptions): Promise { - const adapter = opts.sessionId ? this.sessionAdapters.get(opts.sessionId) : undefined + const adapter = opts.sessionId ? this.sessionAdapters.get(opts.sessionId)?.provider : undefined const target = adapter ?? this.current const result = await target.spawn(opts) - this.sessionAdapters.set(result.id, target) + bindProviderRoute(this.sessionAdapters, result.id, target) return result } @@ -78,7 +86,7 @@ export class DaemonPtyRouter implements IPtyProvider { } hasPty(id: string): boolean { - const routed = this.sessionAdapters.get(id) + const routed = this.sessionAdapters.get(id)?.provider if (routed) { return routed.hasPty(id) } @@ -228,21 +236,10 @@ 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) - // Why: daemon startup can reconcile many restored sessions; spreading - // those arrays into push can exceed JavaScript's argument limit. - for (const id of result.alive) { - alive.push(id) - } - for (const id of result.killed) { - killed.push(id) - } - for (const id of result.alive) { - this.sessionAdapters.set(id, adapter) - } - for (const id of result.killed) { - this.sessionAdapters.delete(id) - } + appendProviderReconciliationIds({ alive, killed }, result) + reconcileProviderRoutesAfterStartup(this.sessionAdapters, routesAtStart, adapter, result) } return { alive, killed } } @@ -295,7 +292,7 @@ export class DaemonPtyRouter implements IPtyProvider { } private adapterFor(sessionId: string): DaemonPtyAdapter { - return this.sessionAdapters.get(sessionId) ?? this.current + return this.sessionAdapters.get(sessionId)?.provider ?? this.current } private allAdapters(): DaemonPtyAdapter[] { diff --git a/src/main/daemon/degraded-daemon-fallback-shutdown.ts b/src/main/daemon/degraded-daemon-fallback-shutdown.ts index 0d4932d6c20..c53f3961854 100644 --- a/src/main/daemon/degraded-daemon-fallback-shutdown.ts +++ b/src/main/daemon/degraded-daemon-fallback-shutdown.ts @@ -1,16 +1,15 @@ import type { IPtyProvider } from '../providers/types' +import { deleteProviderRoute, type ProviderRoute } from './pty-provider-route-reconciliation' export async function shutdownDegradedFallbackSessions( - sessionProviders: Map, + sessionProviders: Map>, fallback: T ): Promise { - const ids = [...sessionProviders] - .filter(([, provider]) => provider === fallback) - .map(([id]) => id) + const routes = [...sessionProviders].filter(([, route]) => route.provider === fallback) const results = await Promise.allSettled( - ids.map(async (id) => { + routes.map(async ([id, route]) => { await fallback.shutdown(id, { immediate: true }) - sessionProviders.delete(id) + deleteProviderRoute(sessionProviders, id, route) }) ) // Why: fallback cleanup must not abort the user's daemon-restart recovery path. diff --git a/src/main/daemon/degraded-daemon-pty-provider.test.ts b/src/main/daemon/degraded-daemon-pty-provider.test.ts index ef9abccbf31..8e0f25e7e0a 100644 --- a/src/main/daemon/degraded-daemon-pty-provider.test.ts +++ b/src/main/daemon/degraded-daemon-pty-provider.test.ts @@ -128,6 +128,27 @@ describe('DegradedDaemonPtyProvider', () => { expect(fallback.write).toHaveBeenCalledWith('shared-session', 'after-exit') }) + it('retains a same-provider replacement across an older process listing', async () => { + const current = createDaemonAdapter('daemon', ['shared-session']) + const fallback = createProvider('fallback') + const provider = new DegradedDaemonPtyProvider({ current, legacy: [], fallback }) + await provider.discoverDaemonSessions() + let finishDaemonListing!: (sessions: []) => void + vi.mocked(current.listProcesses).mockImplementationOnce( + () => new Promise((resolve) => (finishDaemonListing = resolve)) + ) + + const staleListing = provider.listProcesses() + await provider.spawn({ sessionId: 'shared-session', cols: 80, rows: 24 }) + finishDaemonListing([]) + await staleListing + provider.write('shared-session', 'replacement') + + expect(provider.getCurrentDaemonSessionIds()).toEqual(['shared-session']) + expect(current.write).toHaveBeenCalledWith('shared-session', 'replacement') + expect(fallback.write).not.toHaveBeenCalledWith('shared-session', 'replacement') + }) + it('routes fresh foreground confirmation to the session owner', async () => { const current = createDaemonAdapter('daemon', ['daemon-session']) const fallback = createProvider('fallback') diff --git a/src/main/daemon/degraded-daemon-pty-provider.ts b/src/main/daemon/degraded-daemon-pty-provider.ts index 054ba1db04c..6ea8601f34a 100644 --- a/src/main/daemon/degraded-daemon-pty-provider.ts +++ b/src/main/daemon/degraded-daemon-pty-provider.ts @@ -25,7 +25,7 @@ export class DegradedDaemonPtyProvider implements IPtyProvider { private current: DaemonPtyAdapter private legacy: DaemonPtyAdapter[] private fallback: ManagedPtyProvider - private sessionProviders = new Map() + private sessionProviders = new Map>() private unsubscribers: (() => void)[] = [] private dataListeners: ((payload: { id: string @@ -53,12 +53,12 @@ export class DegradedDaemonPtyProvider implements IPtyProvider { provider.onExit((payload) => { // Why: a late exit from an old daemon must not delete a route that // has already been rebound to the local fallback or another daemon. - const routed = this.sessionProviders.get(payload.id) - if (routed && routed !== provider) { + const route = this.sessionProviders.get(payload.id) + if (route && route.provider !== provider) { return } - if (routed === provider) { - this.sessionProviders.delete(payload.id) + if (route?.provider === provider) { + PtyRoutes.deleteProviderRoute(this.sessionProviders, payload.id, route) } for (const listener of this.exitListeners) { listener(payload) @@ -73,7 +73,7 @@ export class DegradedDaemonPtyProvider implements IPtyProvider { try { const sessions = await adapter.listProcesses() for (const session of sessions) { - this.sessionProviders.set(session.id, adapter) + PtyRoutes.bindProviderRouteIfAbsent(this.sessionProviders, session.id, adapter) } } catch (error) { console.warn('[daemon] Failed to discover degraded daemon sessions', error) @@ -82,10 +82,10 @@ export class DegradedDaemonPtyProvider implements IPtyProvider { } async spawn(opts: PtySpawnOptions): Promise { - const mapped = opts.sessionId ? this.sessionProviders.get(opts.sessionId) : undefined + const mapped = opts.sessionId ? this.sessionProviders.get(opts.sessionId)?.provider : undefined const target = mapped ?? this.fallback const result = await target.spawn(opts) - this.sessionProviders.set(result.id, target) + PtyRoutes.bindProviderRoute(this.sessionProviders, result.id, target) return result } @@ -94,7 +94,7 @@ export class DegradedDaemonPtyProvider implements IPtyProvider { } hasPty(id: string): boolean { - const mapped = this.sessionProviders.get(id) + const mapped = this.sessionProviders.get(id)?.provider if (mapped) { return mapped.hasPty?.(id) ?? true } @@ -263,15 +263,15 @@ 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) - for (const id of result.alive) { - alive.push(id) - this.sessionProviders.set(id, adapter) - } - for (const id of result.killed) { - killed.push(id) - this.sessionProviders.delete(id) - } + PtyRoutes.appendProviderReconciliationIds({ alive, killed }, result) + PtyRoutes.reconcileProviderRoutesAfterStartup( + this.sessionProviders, + routesAtStart, + adapter, + result + ) } return { alive, killed } } @@ -299,7 +299,7 @@ export class DegradedDaemonPtyProvider implements IPtyProvider { fanoutCurrentDaemonSyntheticExits(code: number): void { for (const id of this.getCurrentDaemonSessionIds()) { - this.sessionProviders.delete(id) + PtyRoutes.deleteProviderRoute(this.sessionProviders, id) // Why: sessions discovered from listProcesses may not exist in the // adapter's active-session set, but restart still kills that daemon. // oxlint-disable-next-line unicorn/no-useless-spread -- copy-safe: listeners may unsubscribe during iteration @@ -328,7 +328,7 @@ export class DegradedDaemonPtyProvider implements IPtyProvider { private providerFor(sessionId: string): ManagedPtyProvider { return ( - this.sessionProviders.get(sessionId) ?? + this.sessionProviders.get(sessionId)?.provider ?? this.findProviderForExistingSession(sessionId) ?? this.fallback ) @@ -337,7 +337,7 @@ export class DegradedDaemonPtyProvider implements IPtyProvider { private findProviderForExistingSession(sessionId: string): ManagedPtyProvider | null { for (const provider of this.allProviders()) { if (provider.hasPty?.(sessionId) === true) { - this.sessionProviders.set(sessionId, provider) + PtyRoutes.bindProviderRoute(this.sessionProviders, sessionId, provider) return provider } } @@ -345,7 +345,7 @@ export class DegradedDaemonPtyProvider implements IPtyProvider { } private daemonAdapterFor(sessionId: string): DaemonPtyAdapter | null { - const provider = this.sessionProviders.get(sessionId) + const provider = this.sessionProviders.get(sessionId)?.provider return provider && this.allDaemonAdapters().includes(provider as DaemonPtyAdapter) ? (provider as DaemonPtyAdapter) : null diff --git a/src/main/daemon/pty-provider-route-reconciliation.ts b/src/main/daemon/pty-provider-route-reconciliation.ts index 2f84b144661..124a042b3d8 100644 --- a/src/main/daemon/pty-provider-route-reconciliation.ts +++ b/src/main/daemon/pty-provider-route-reconciliation.ts @@ -1,8 +1,78 @@ import type { IPtyProvider, PtyProcessInfo } from '../providers/types' +export type ProviderRoute = Readonly<{ provider: T }> + +export function bindProviderRoute( + routes: Map>, + id: string, + provider: T +): ProviderRoute { + // 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 } + 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, + expectedRoute?: ProviderRoute +): boolean { + if (expectedRoute && routes.get(id) !== expectedRoute) { + return false + } + return routes.delete(id) +} + +export function reconcileProviderRoutesAfterStartup( + routes: Map>, + routesAtStart: Map>, + 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)) { + bindProviderRoute(routes, id, provider) + } + } + for (const id of result.killed) { + const routeAtStart = routesAtStart.get(id) + if (routeAtStart) { + deleteProviderRoute(routes, id, routeAtStart) + } + } +} + +export function appendProviderReconciliationIds( + target: { alive: string[]; killed: string[] }, + result: { alive: string[]; killed: string[] } +): void { + // Why: daemon startup can return enough sessions to exceed JavaScript's + // argument limit if these arrays are appended with spread syntax. + for (const id of result.alive) { + target.alive.push(id) + } + for (const id of result.killed) { + target.killed.push(id) + } +} + export async function listProviderProcessesAndReconcileRoutes( providers: readonly T[], - routes: Map + routes: Map> ): Promise { const routesAtStart = [...routes] const listings = await Promise.all( @@ -11,16 +81,20 @@ export async function listProviderProcessesAndReconcileRoutes [provider, new Set(sessions.map(({ id }) => id))]) ) - for (const [id, provider] of routesAtStart) { + for (const [id, route] of routesAtStart) { // Why: spawn/exit can mutate routing while remote listings are in flight. - // Never let that older snapshot delete a route rebound to another owner. - if (routes.get(id) === provider && !liveIdsByProvider.get(provider)?.has(id)) { + // 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) } } return listings.flatMap(({ sessions }) => sessions) } -export function providerSessionIds(routes: Map, provider: T): string[] { - return [...routes].filter(([, routedProvider]) => routedProvider === provider).map(([id]) => id) +export function providerSessionIds( + routes: Map>, + provider: T +): string[] { + return [...routes].filter(([, route]) => route.provider === provider).map(([id]) => id) }