diff --git a/src/main/daemon/daemon-pty-provider.test.ts b/src/main/daemon/daemon-pty-provider.test.ts index 2a405326d1b..51fd6a7c7fc 100644 --- a/src/main/daemon/daemon-pty-provider.test.ts +++ b/src/main/daemon/daemon-pty-provider.test.ts @@ -23,7 +23,7 @@ function createMockSubprocess(): SubprocessHandle & { write: vi.fn(), resize: vi.fn(), kill: vi.fn(() => setTimeout(() => onExitCb?.(0), 5)), - forceKill: vi.fn(), + forceKill: vi.fn(() => setTimeout(() => onExitCb?.(-1), 5)), signal: vi.fn(), onData(cb) { onDataCb = cb diff --git a/src/main/daemon/daemon-spawner.test.ts b/src/main/daemon/daemon-spawner.test.ts index 429a864f841..6ae301ee71a 100644 --- a/src/main/daemon/daemon-spawner.test.ts +++ b/src/main/daemon/daemon-spawner.test.ts @@ -25,7 +25,7 @@ function createMockSubprocess(): SubprocessHandle { write: vi.fn(), resize: vi.fn(), kill: vi.fn(() => setTimeout(() => onExitCb?.(0), 5)), - forceKill: vi.fn(), + forceKill: vi.fn(() => setTimeout(() => onExitCb?.(-1), 5)), signal: vi.fn(), onData(_cb: (data: string) => void) {}, onExit(cb: (code: number) => void) { diff --git a/src/main/daemon/slow-daemon-session-verification.test.ts b/src/main/daemon/slow-daemon-session-verification.test.ts index 15634530dcd..60150ebc51c 100644 --- a/src/main/daemon/slow-daemon-session-verification.test.ts +++ b/src/main/daemon/slow-daemon-session-verification.test.ts @@ -21,16 +21,19 @@ import type { SubprocessHandle } from './session' const RESPONSE_DELAY_MS = 3_500 function createMockSubprocess(): SubprocessHandle { + let onExitCb: ((code: number) => void) | null = null return { pid: 55555, getForegroundProcess: vi.fn(() => null), write: vi.fn(), resize: vi.fn(), kill: vi.fn(), - forceKill: vi.fn(), + forceKill: vi.fn(() => setTimeout(() => onExitCb?.(-1), 5)), signal: vi.fn(), onData: vi.fn(), - onExit: vi.fn(), + onExit(cb) { + onExitCb = cb + }, dispose: vi.fn() } } diff --git a/src/renderer/src/components/terminal-pane/pty-primary-handler-admission.ts b/src/renderer/src/components/terminal-pane/pty-primary-handler-admission.ts new file mode 100644 index 00000000000..e9df3773304 --- /dev/null +++ b/src/renderer/src/components/terminal-pane/pty-primary-handler-admission.ts @@ -0,0 +1,66 @@ +import type { PtyDataMeta } from './pty-dispatcher' + +type PtyPrimaryHandlerAdmissionRegistry = { + dataHandlers: Map void> + replayHandlers: Map void> + exitHandlers: Map void> + teardownHandlers: Map void> + drainData: (ptyId: string, handler: (data: string, meta?: PtyDataMeta) => void) => void + drainExit: (ptyId: string, handler: (code: number) => void) => boolean +} + +export type PtyPrimaryHandlerAdmissionSnapshot = { + ptyId: string + dataHandler?: (data: string, meta?: PtyDataMeta) => void + replayHandler?: (data: string) => void + exitHandler?: (code: number) => void + teardownHandler?: () => void +} + +export function suspendPtyPrimaryHandlersForAdmission( + registry: PtyPrimaryHandlerAdmissionRegistry, + ptyId: string +): PtyPrimaryHandlerAdmissionSnapshot { + const snapshot: PtyPrimaryHandlerAdmissionSnapshot = { + ptyId, + dataHandler: registry.dataHandlers.get(ptyId), + replayHandler: registry.replayHandlers.get(ptyId), + exitHandler: registry.exitHandlers.get(ptyId), + teardownHandler: registry.teardownHandlers.get(ptyId) + } + registry.dataHandlers.delete(ptyId) + registry.replayHandlers.delete(ptyId) + registry.exitHandlers.delete(ptyId) + registry.teardownHandlers.delete(ptyId) + return snapshot +} + +export function restorePtyPrimaryHandlersAfterFailedAdmission( + registry: PtyPrimaryHandlerAdmissionRegistry, + snapshot: PtyPrimaryHandlerAdmissionSnapshot +): void { + let restoredDataHandler: PtyPrimaryHandlerAdmissionSnapshot['dataHandler'] + let restoredExitHandler: PtyPrimaryHandlerAdmissionSnapshot['exitHandler'] + if (snapshot.dataHandler && !registry.dataHandlers.has(snapshot.ptyId)) { + registry.dataHandlers.set(snapshot.ptyId, snapshot.dataHandler) + restoredDataHandler = snapshot.dataHandler + } + if (snapshot.replayHandler && !registry.replayHandlers.has(snapshot.ptyId)) { + registry.replayHandlers.set(snapshot.ptyId, snapshot.replayHandler) + } + if (snapshot.exitHandler && !registry.exitHandlers.has(snapshot.ptyId)) { + registry.exitHandlers.set(snapshot.ptyId, snapshot.exitHandler) + restoredExitHandler = snapshot.exitHandler + } + if (snapshot.teardownHandler && !registry.teardownHandlers.has(snapshot.ptyId)) { + registry.teardownHandlers.set(snapshot.ptyId, snapshot.teardownHandler) + } + // Why: events emitted while admission owned no primary handler still belong + // to the restored generation when the replacement request itself failed. + if (restoredDataHandler) { + registry.drainData(snapshot.ptyId, restoredDataHandler) + } + if (restoredExitHandler) { + registry.drainExit(snapshot.ptyId, restoredExitHandler) + } +} diff --git a/src/renderer/src/components/terminal-pane/pty-transport.test.ts b/src/renderer/src/components/terminal-pane/pty-transport.test.ts index 7eb29b1ebfe..656e64036ca 100644 --- a/src/renderer/src/components/terminal-pane/pty-transport.test.ts +++ b/src/renderer/src/components/terminal-pane/pty-transport.test.ts @@ -214,6 +214,76 @@ describe('createIpcPtyTransport', () => { expect(replacementExit).toHaveBeenCalledWith(23) }) + it('buffers post-admission data for the replacement instead of the old transport', async () => { + const { createIpcPtyTransport } = await import('./pty-transport') + const spawn = window.api.pty.spawn as unknown as ReturnType + let resolveReplacementSpawn: ((result: { id: string }) => void) | undefined + spawn.mockImplementationOnce( + () => + new Promise<{ id: string }>((resolve) => { + resolveReplacementSpawn = resolve + }) + ) + const oldData = vi.fn() + const replacementData = vi.fn() + const oldTransport = createIpcPtyTransport() + const replacementTransport = createIpcPtyTransport() + oldTransport.attach({ existingPtyId: 'same-id', callbacks: { onData: oldData } }) + + const replacementConnect = replacementTransport.connect({ + url: '', + sessionId: 'same-id', + callbacks: { onData: replacementData } + }) + await vi.waitFor(() => expect(resolveReplacementSpawn).toBeTypeOf('function')) + onData?.({ id: 'same-id', data: 'after admission boundary' }) + await flushPtySideEffects() + + expect(oldData).not.toHaveBeenCalled() + expect(replacementData).not.toHaveBeenCalled() + resolveReplacementSpawn?.({ id: 'same-id' }) + await replacementConnect + + expect(replacementData).toHaveBeenCalledWith('after admission boundary') + oldTransport.disconnect() + replacementTransport.disconnect() + }) + + it('restores the old transport handlers after a generic same-id admission failure', async () => { + const { createIpcPtyTransport } = await import('./pty-transport') + const spawn = window.api.pty.spawn as unknown as ReturnType + let rejectReplacementSpawn: ((error: Error) => void) | undefined + spawn.mockImplementationOnce( + () => + new Promise((_resolve, reject) => { + rejectReplacementSpawn = reject + }) + ) + const oldData = vi.fn() + const oldTransport = createIpcPtyTransport() + const replacementTransport = createIpcPtyTransport() + oldTransport.attach({ existingPtyId: 'same-id', callbacks: { onData: oldData } }) + + const replacementConnect = replacementTransport.connect({ + url: '', + sessionId: 'same-id', + callbacks: {} + }) + await vi.waitFor(() => expect(rejectReplacementSpawn).toBeTypeOf('function')) + onData?.({ id: 'same-id', data: 'buffered while admission fails' }) + await flushPtySideEffects() + expect(oldData).not.toHaveBeenCalled() + + rejectReplacementSpawn?.(new Error('daemon temporarily unavailable')) + await replacementConnect + onData?.({ id: 'same-id', data: 'old generation resumes' }) + await flushPtySideEffects() + + expect(oldData).toHaveBeenCalledWith('buffered while admission fails') + expect(oldData).toHaveBeenCalledWith('old generation resumes') + oldTransport.disconnect() + }) + it('ignores a stale exit for a previous PTY after reconnecting the same transport', async () => { const { createIpcPtyTransport } = await import('./pty-transport') const spawn = window.api.pty.spawn as unknown as ReturnType diff --git a/src/renderer/src/components/terminal-pane/pty-transport.ts b/src/renderer/src/components/terminal-pane/pty-transport.ts index 5090199a892..45cde4ed55d 100644 --- a/src/renderer/src/components/terminal-pane/pty-transport.ts +++ b/src/renderer/src/components/terminal-pane/pty-transport.ts @@ -30,6 +30,10 @@ import { } from './pty-pre-handler-buffer' import { createPtyInputWriteQueue } from './pty-input-write-queue' import { acquirePtySessionConnectAdmission } from './pty-session-connect-admission' +import { + restorePtyPrimaryHandlersAfterFailedAdmission, + suspendPtyPrimaryHandlersForAdmission +} from './pty-primary-handler-admission' import type { PtyDataMeta } from './pty-dispatcher' import type { IpcPtyTransportOptions, PtyConnectResult, PtyTransport } from './pty-transport-types' import { createBellDetector } from '../../../../shared/terminal-bell-detector' @@ -76,6 +80,15 @@ const SSH_PTY_CONNECTION_MISMATCH_MARKER = 'belongs to SSH connection' const STALE_TITLE_TIMEOUT = 3000 // ms before stale working title is cleared const MAX_PTY_SIDE_EFFECTS_PER_DRAIN = 64 +const ptyPrimaryHandlerAdmissionRegistry = { + dataHandlers: ptyDataHandlers, + replayHandlers: ptyReplayHandlers, + exitHandlers: ptyExitHandlers, + teardownHandlers: ptyTeardownHandlers, + drainData: drainPreHandlerPtyData, + drainExit: drainPreHandlerPtyExit +} + // Why: onAgentStatus callback added to IpcPtyTransportOptions in pty-dispatcher // so the OSC 9999 status payloads can be forwarded to the store. @@ -707,6 +720,13 @@ export function createIpcPtyTransport(opts: IpcPtyTransportOptions = {}): PtyTra const releaseSessionAdmission = options.sessionId ? await acquirePtySessionConnectAdmission(options.sessionId) : undefined + const suspendedHandlers = options.sessionId + ? suspendPtyPrimaryHandlersForAdmission( + ptyPrimaryHandlerAdmissionRegistry, + options.sessionId + ) + : undefined + let restoreSuspendedHandlers = suspendedHandlers !== undefined try { if (options.sessionId) { // Why: admission starts a new renderer generation; buffered events @@ -756,6 +776,9 @@ export function createIpcPtyTransport(opts: IpcPtyTransportOptions = {}): PtyTra ...(terminalColorQueryReplies ? { terminalColorQueryReplies } : {}), ...(telemetry ? { telemetry } : {}) }) + // The host accepted a new generation; the prior primary handlers must + // never be restored even if this transport is destroyed before binding. + restoreSuspendedHandlers = false const spawnResult = result as PtyConnectResult & { isReattach?: boolean } const resultLaunchAgent = isTuiAgent(spawnResult.launchAgent) ? spawnResult.launchAgent @@ -819,6 +842,7 @@ export function createIpcPtyTransport(opts: IpcPtyTransportOptions = {}): PtyTra } catch (err) { const msg = extractIpcErrorMessage(err, err instanceof Error ? err.message : String(err)) if (options.sessionId && msg.includes(DAEMON_PTY_EXITED_DURING_ADMISSION)) { + restoreSuspendedHandlers = false // Why: this exit belongs to the rejected generation; retaining its // bytes or exit would tear down a later successful same-id spawn. clearPreHandlerPtyState(options.sessionId) @@ -829,6 +853,7 @@ export function createIpcPtyTransport(opts: IpcPtyTransportOptions = {}): PtyTra (msg.includes(SSH_SESSION_EXPIRED_ERROR) || msg.includes(SSH_PTY_CONNECTION_MISMATCH_MARKER)) ) { + restoreSuspendedHandlers = false return { id: options.sessionId, sessionExpired: true @@ -848,6 +873,7 @@ export function createIpcPtyTransport(opts: IpcPtyTransportOptions = {}): PtyTra // error type (see src/main/daemon/daemon-pty-adapter.ts), so a // substring match is safe. if (msg.includes('was explicitly killed')) { + restoreSuspendedHandlers = false return undefined } // Why: on cold start, SSH provider isn't registered yet so pty:spawn @@ -867,6 +893,12 @@ export function createIpcPtyTransport(opts: IpcPtyTransportOptions = {}): PtyTra } return undefined } finally { + if (restoreSuspendedHandlers && suspendedHandlers) { + restorePtyPrimaryHandlersAfterFailedAdmission( + ptyPrimaryHandlerAdmissionRegistry, + suspendedHandlers + ) + } releaseSessionAdmission?.() } }, diff --git a/src/renderer/src/store/slices/store-session-cascades.test.ts b/src/renderer/src/store/slices/store-session-cascades.test.ts index f938dfecdb0..71d5778cf3f 100644 --- a/src/renderer/src/store/slices/store-session-cascades.test.ts +++ b/src/renderer/src/store/slices/store-session-cascades.test.ts @@ -211,8 +211,8 @@ describe('removeProject cascade', () => { expect(s.activeTabId).toBeNull() // PTYs were killed - expect(mockApi.pty.kill).toHaveBeenCalledWith('pty1') - expect(mockApi.pty.kill).toHaveBeenCalledWith('pty2') + expect(mockApi.pty.kill).toHaveBeenCalledWith('pty1', { expectedTabId: 'tab1' }) + expect(mockApi.pty.kill).toHaveBeenCalledWith('pty2', { expectedTabId: 'tab2' }) // The tabs are gone before async exit events can close them, so retaining // their one-shot guards would leak ephemeral PTY ids for the renderer session.