From abf9cb02a52280811ae9c06172efe32536dfba1a Mon Sep 17 00:00:00 2001 From: KorenKrita Date: Sun, 2 Aug 2026 21:15:52 +0800 Subject: [PATCH] fix(pi): resume archived sessions safely (#1308) * fix(pi): resume archived sessions safely * fix(pi): harden native resume startup * fix(pi): harden resume termination evidence * fix(runner): persist resume process evidence * fix(runner): track resume process generations * fix(runner): verify full session tree shutdown * fix(pi): block pre-mapping resume dedup --- cli/src/agent/sessionFactory.test.ts | 12 + cli/src/agent/sessionFactory.ts | 1 + cli/src/api/apiMachine.test.ts | 2 +- cli/src/api/apiMachine.ts | 12 +- cli/src/commands/agentCommandOptions.test.ts | 5 + cli/src/commands/agentCommandOptions.ts | 7 + cli/src/commands/runner.test.ts | 24 +- cli/src/commands/runner.ts | 10 +- cli/src/pi/loop.test.ts | 151 ++++++ cli/src/pi/loop.ts | 64 ++- cli/src/pi/runPi.test.ts | 56 ++- cli/src/pi/runPi.ts | 45 +- cli/src/pi/session.test.ts | 17 + cli/src/pi/session.ts | 39 +- cli/src/runner/buildCliArgs.test.ts | 13 + cli/src/runner/controlClient.test.ts | 55 +++ cli/src/runner/controlClient.ts | 6 +- cli/src/runner/controlServer.ts | 8 +- cli/src/runner/run.ts | 250 +++++++++- cli/src/runner/runner.integration.test.ts | 6 +- cli/src/runner/types.ts | 4 +- cli/src/utils/process.ts | 38 +- hub/src/sync/rpcGateway.ts | 7 + hub/src/sync/sessionModel.test.ts | 374 +++++++++++++++ hub/src/sync/syncEngine.ts | 464 ++++++++++++++++--- shared/src/schemas.ts | 13 + 26 files changed, 1546 insertions(+), 137 deletions(-) create mode 100644 cli/src/runner/controlClient.test.ts diff --git a/cli/src/agent/sessionFactory.test.ts b/cli/src/agent/sessionFactory.test.ts index 386abe52..0662852d 100644 --- a/cli/src/agent/sessionFactory.test.ts +++ b/cli/src/agent/sessionFactory.test.ts @@ -147,6 +147,12 @@ describe('bootstrapExistingSession', () => { grokSessionId: 'grok-thread-1', cursorSessionId: 'cursor-thread-1', cursorSessionProtocol: 'acp', + piSessionId: 'pi-thread-1', + piResumeAttempt: { + state: 'resuming', + machineId: 'machine-1', + startedAt: 123, + }, summary: { text: 'resume me', updatedAt: 100 @@ -176,6 +182,12 @@ describe('bootstrapExistingSession', () => { grokSessionId: 'grok-thread-1', cursorSessionId: 'cursor-thread-1', cursorSessionProtocol: 'acp', + piSessionId: 'pi-thread-1', + piResumeAttempt: { + state: 'resuming', + machineId: 'machine-1', + startedAt: 123, + }, summary: { text: 'resume me', updatedAt: 100 diff --git a/cli/src/agent/sessionFactory.ts b/cli/src/agent/sessionFactory.ts index cc80cf33..8f256771 100644 --- a/cli/src/agent/sessionFactory.ts +++ b/cli/src/agent/sessionFactory.ts @@ -106,6 +106,7 @@ function pickExistingSessionMetadata(metadata: Metadata | null | undefined): Par if (metadata.cursorSessionProtocol !== undefined) preserved.cursorSessionProtocol = metadata.cursorSessionProtocol if (metadata.kimiSessionId !== undefined) preserved.kimiSessionId = metadata.kimiSessionId if (metadata.piSessionId !== undefined) preserved.piSessionId = metadata.piSessionId + if (metadata.piResumeAttempt !== undefined) preserved.piResumeAttempt = metadata.piResumeAttempt if (metadata.preferredPermissionMode !== undefined) preserved.preferredPermissionMode = metadata.preferredPermissionMode if (metadata.tools !== undefined) preserved.tools = metadata.tools if (metadata.slashCommands !== undefined) preserved.slashCommands = metadata.slashCommands diff --git a/cli/src/api/apiMachine.test.ts b/cli/src/api/apiMachine.test.ts index a2a12ec9..b861c574 100644 --- a/cli/src/api/apiMachine.test.ts +++ b/cli/src/api/apiMachine.test.ts @@ -458,7 +458,7 @@ describe('ApiMachineClient SpawnHappySession handler', () => { client.setRPCHandlers({ spawnSession, - stopSession: vi.fn(() => true), + stopSession: vi.fn(async () => 'stopped' as const), requestShutdown: vi.fn() }) diff --git a/cli/src/api/apiMachine.ts b/cli/src/api/apiMachine.ts index d8af7077..0759538f 100644 --- a/cli/src/api/apiMachine.ts +++ b/cli/src/api/apiMachine.ts @@ -46,7 +46,7 @@ import type { CursorChatStoreStatus } from '@hapi/protocol/apiTypes' type MachineRpcHandlers = { spawnSession: (options: SpawnSessionOptions) => Promise - stopSession: (sessionId: string) => boolean + stopSession: (sessionId: string) => Promise<'stopped' | 'already_gone' | 'still_alive'> requestShutdown: () => void } @@ -398,18 +398,14 @@ export class ApiMachineClient { } }) - this.rpcHandlerManager.registerHandler(RPC_METHODS.StopSession, (params: any) => { + this.rpcHandlerManager.registerHandler(RPC_METHODS.StopSession, async (params: any) => { const { sessionId } = params || {} if (!sessionId) { throw new Error('Session ID is required') } - const success = stopSession(sessionId) - if (!success) { - throw new Error('Session not found or failed to stop') - } - - return { message: 'Session stopped' } + const status = await stopSession(sessionId) + return { status } }) this.rpcHandlerManager.registerHandler(RPC_METHODS.StopRunner, () => { diff --git a/cli/src/commands/agentCommandOptions.test.ts b/cli/src/commands/agentCommandOptions.test.ts index e0702cb1..e9981599 100644 --- a/cli/src/commands/agentCommandOptions.test.ts +++ b/cli/src/commands/agentCommandOptions.test.ts @@ -7,12 +7,14 @@ describe('parseRemoteAgentCommandOptions', () => { expect(parseRemoteAgentCommandOptions([ '--started-by', 'runner', '--hapi-starting-mode', 'remote', + '--existing-session-id', 'hapi-session-1', '--permission-mode', 'yolo', '--resume', 'session-1', '--model', 'model-a' ], GEMINI_PERMISSION_MODES)).toEqual({ startedBy: 'runner', startingMode: 'remote', + existingSessionId: 'hapi-session-1', permissionMode: 'yolo', resumeSessionId: 'session-1', model: 'model-a' @@ -67,6 +69,7 @@ describe('parseRemoteAgentCommandOptions', () => { expect(() => parseRemoteAgentCommandOptions(['--resume'], OPENCODE_PERMISSION_MODES)).toThrow('Missing --resume value') expect(() => parseRemoteAgentCommandOptions(['--model'], OPENCODE_PERMISSION_MODES)).toThrow('Missing --model value') expect(() => parseRemoteAgentCommandOptions(['--model-reasoning-effort'], OPENCODE_PERMISSION_MODES)).toThrow('Missing --model-reasoning-effort value') + expect(() => parseRemoteAgentCommandOptions(['--existing-session-id'], OPENCODE_PERMISSION_MODES)).toThrow('Missing --existing-session-id value') }) it('accepts OpenCode-native -s / --session as resume aliases', () => { @@ -190,6 +193,7 @@ describe('parseRemoteAgentCommandOptions — pi flavor', () => { '--hapi-starting-mode', 'remote', '--model', 'claude-sonnet-4-5', '--session-id', 'pi-sess-full', + '--existing-session-id', 'hapi-session-pi-full', ], ALLOWED ) @@ -198,6 +202,7 @@ describe('parseRemoteAgentCommandOptions — pi flavor', () => { startingMode: 'remote', model: 'claude-sonnet-4-5', resumeSessionId: 'pi-sess-full', + existingSessionId: 'hapi-session-pi-full', }) }) }) diff --git a/cli/src/commands/agentCommandOptions.ts b/cli/src/commands/agentCommandOptions.ts index a8041e06..549aaaae 100644 --- a/cli/src/commands/agentCommandOptions.ts +++ b/cli/src/commands/agentCommandOptions.ts @@ -8,6 +8,7 @@ export type RemoteAgentCommandOptions = effort?: string modelReasoningEffort?: string resumeSessionId?: string + existingSessionId?: string } export function parseRemoteAgentCommandOptions( @@ -28,6 +29,12 @@ export function parseRemoteAgentCommandOptions []), stopRunnerMock: vi.fn(async () => {}), - stopRunnerSessionMock: vi.fn(async () => true), + stopRunnerSessionMock: vi.fn(async (_sessionId: string): Promise<'stopped' | 'already_gone' | 'still_alive'> => 'stopped'), spawnHappyCLIMock: vi.fn(() => ({ unref: vi.fn() })), startRunnerMock: vi.fn(async () => {}), getLatestRunnerLogMock: vi.fn(async () => null), @@ -127,3 +127,25 @@ describe('runnerCommand start', () => { } }) }) + +describe('runnerCommand stop-session', () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + it.each([ + ['stopped', 'Session stopped'], + ['already_gone', 'Session was already stopped'], + ['still_alive', 'Failed to stop session'], + ] as const)('renders the %s runner result', async (status, expected) => { + const consoleLogSpy = vi.spyOn(console, 'log').mockImplementation(() => {}) + stopRunnerSessionMock.mockResolvedValueOnce(status) + try { + await runnerCommand.run(createContext(['stop-session', 'session-1'])) + expect(stopRunnerSessionMock).toHaveBeenCalledWith('session-1') + expect(consoleLogSpy).toHaveBeenCalledWith(expected) + } finally { + consoleLogSpy.mockRestore() + } + }) +}) diff --git a/cli/src/commands/runner.ts b/cli/src/commands/runner.ts index b74419e9..75caa109 100644 --- a/cli/src/commands/runner.ts +++ b/cli/src/commands/runner.ts @@ -113,8 +113,14 @@ export const runnerCommand: CommandDefinition = { } try { - const success = await stopRunnerSession(sessionId) - console.log(success ? 'Session stopped' : 'Failed to stop session') + const status = await stopRunnerSession(sessionId) + if (status === 'stopped') { + console.log('Session stopped') + } else if (status === 'already_gone') { + console.log('Session was already stopped') + } else { + console.log('Failed to stop session') + } } catch { console.log('No runner running') } diff --git a/cli/src/pi/loop.test.ts b/cli/src/pi/loop.test.ts index c7dbd7a7..ef0f5c5d 100644 --- a/cli/src/pi/loop.test.ts +++ b/cli/src/pi/loop.test.ts @@ -44,6 +44,7 @@ function createMockSession(): PiSession { sendAgentMessage: vi.fn(), emitMessagesConsumed: vi.fn(), sendSessionEvent: vi.fn(), + emitSessionReady: vi.fn(), } as any, path: '/tmp/test', logPath: '/tmp/test.log', @@ -267,6 +268,7 @@ describe('wireTransportEvents', () => { expect(session.currentThinkingLevel).toBe('high'); expect(session.currentSteeringMode).toBe('one-at-a-time'); expect(session.client.updateMetadata).toHaveBeenCalledWith(expect.any(Function)); + expect(session.client.emitSessionReady).toHaveBeenCalledTimes(1); }); it('marks session ready on get_state response (drains buffered sends) — issue #1143', () => { @@ -289,6 +291,7 @@ describe('wireTransportEvents', () => { // get_state landing is the ready signal — buffered work drains. expect(session.isReady).toBe(true); expect(buffered).toHaveBeenCalledTimes(1); + expect(session.client.emitSessionReady).toHaveBeenCalledTimes(1); }); it('marks session ready on get_state even when sessionId is absent', () => { @@ -300,6 +303,154 @@ describe('wireTransportEvents', () => { emitEvent({ type: 'response', command: 'get_state', success: true, data: {} }); expect(session.isReady).toBe(true); + expect(session.client.emitSessionReady).toHaveBeenCalledTimes(1); + }); + + it('keeps a fresh Pi get_state failure non-fatal for startup fallback compatibility', () => { + const transport = createMockTransport(); + const onStartupFailure = vi.fn(); + wireTransportEvents(transport, session, [], { onStartupFailure }); + + emitEvent({ + type: 'response', + command: 'get_state', + success: false, + error: 'No session found matching pi-session-404', + }); + + expect(onStartupFailure).not.toHaveBeenCalled(); + expect(session.client.emitSessionReady).not.toHaveBeenCalled(); + }); + + it('keeps malformed get_state data non-fatal for a fresh Pi session', () => { + const transport = createMockTransport(); + const onStartupFailure = vi.fn(); + wireTransportEvents(transport, session, [], { onStartupFailure }); + + emitEvent({ + type: 'response', + command: 'get_state', + success: true, + data: { model: 'not-a-model-object' }, + }); + + expect(onStartupFailure).not.toHaveBeenCalled(); + expect(session.isReady).toBe(false); + expect(session.client.emitSessionReady).not.toHaveBeenCalled(); + }); + + it('fails a native resume when get_state returns malformed data', () => { + const expectedSession = new PiSession({ + api: {} as any, + client: { + keepAlive: vi.fn(), + updateMetadata: vi.fn(), + sendAgentMessage: vi.fn(), + emitMessagesConsumed: vi.fn(), + sendSessionEvent: vi.fn(), + emitSessionReady: vi.fn(), + } as any, + path: '/tmp/test', + logPath: '/tmp/test.log', + startedBy: 'terminal', + startingMode: 'local', + expectedNativeSessionId: 'pi-session-requested', + }); + const transport = createMockTransport(); + const onStartupFailure = vi.fn(); + wireTransportEvents(transport, expectedSession, [], { onStartupFailure }); + + emitEvent({ + type: 'response', + command: 'get_state', + success: true, + data: { model: 'not-a-model-object' }, + }); + + expect(onStartupFailure).toHaveBeenCalledTimes(1); + expect((onStartupFailure.mock.calls[0][0] as Error).message).toContain('malformed state data'); + expect(expectedSession.isReady).toBe(false); + expect(expectedSession.client.emitSessionReady).not.toHaveBeenCalled(); + expect(expectedSession.client.updateMetadata).not.toHaveBeenCalled(); + }); + + it('rejects a resume get_state response with a missing session ID before mutating state', () => { + const expectedSession = new PiSession({ + api: {} as any, + client: { + keepAlive: vi.fn(), + updateMetadata: vi.fn(), + sendAgentMessage: vi.fn(), + emitMessagesConsumed: vi.fn(), + sendSessionEvent: vi.fn(), + emitSessionReady: vi.fn(), + } as any, + path: '/tmp/test', + logPath: '/tmp/test.log', + startedBy: 'terminal', + startingMode: 'local', + expectedNativeSessionId: 'pi-session-requested', + }); + const transport = createMockTransport(); + const onStartupFailure = vi.fn(); + wireTransportEvents(transport, expectedSession, [], { onStartupFailure }); + + emitEvent({ + type: 'response', + command: 'get_state', + success: true, + data: { + model: { modelId: 'wrong-model', provider: 'wrong-provider' }, + thinkingLevel: 'high', + }, + }); + + expect(onStartupFailure).toHaveBeenCalledTimes(1); + expect((onStartupFailure.mock.calls[0][0] as Error).message) + .toContain('unexpected native session (missing)'); + expect(expectedSession.client.emitSessionReady).not.toHaveBeenCalled(); + expect(expectedSession.client.updateMetadata).not.toHaveBeenCalled(); + expect(expectedSession.currentModel).toBeUndefined(); + expect(expectedSession.currentThinkingLevel).toBeUndefined(); + }); + + it('rejects a resume get_state response for a different session before metadata is published', () => { + const expectedSession = new PiSession({ + api: {} as any, + client: { + keepAlive: vi.fn(), + updateMetadata: vi.fn(), + sendAgentMessage: vi.fn(), + emitMessagesConsumed: vi.fn(), + sendSessionEvent: vi.fn(), + emitSessionReady: vi.fn(), + } as any, + path: '/tmp/test', + logPath: '/tmp/test.log', + startedBy: 'terminal', + startingMode: 'local', + expectedNativeSessionId: 'pi-session-requested', + }); + const transport = createMockTransport(); + const onStartupFailure = vi.fn(); + wireTransportEvents(transport, expectedSession, [], { onStartupFailure }); + + emitEvent({ + type: 'response', + command: 'get_state', + success: true, + data: { + sessionId: 'pi-session-other', + model: { modelId: 'wrong-model', provider: 'wrong-provider' }, + }, + }); + + expect(onStartupFailure).toHaveBeenCalledTimes(1); + expect((onStartupFailure.mock.calls[0][0] as Error).message) + .toContain('unexpected native session pi-session-other'); + expect(expectedSession.client.emitSessionReady).not.toHaveBeenCalled(); + expect(expectedSession.client.updateMetadata).not.toHaveBeenCalled(); + expect(expectedSession.currentModel).toBeUndefined(); }); it('handles error response — sends session event', () => { diff --git a/cli/src/pi/loop.ts b/cli/src/pi/loop.ts index 9fb2fda0..0b1210b5 100644 --- a/cli/src/pi/loop.ts +++ b/cli/src/pi/loop.ts @@ -84,13 +84,15 @@ function persistSelectedPiModel(session: PiSession): void { // --- Response handler --- -function handleGetState( - rawData: unknown, +function applyGetState( + data: { + model?: { id?: string; modelId?: string; provider?: string }; + sessionId?: string; + thinkingLevel?: string; + steeringMode?: 'all' | 'one-at-a-time'; + }, session: PiSession, ): void { - const parsed = PiStateDataSchema.safeParse(rawData); - if (!parsed.success) return; - const data = parsed.data; if (data.model) { // Pi returns model.id (not modelId). Fallback to modelId for forward compat. @@ -128,6 +130,7 @@ function handleGetState( if (data.steeringMode) { session.currentSteeringMode = data.steeringMode; } + } function handleResponse( @@ -135,6 +138,7 @@ function handleResponse( session: PiSession, pendingLocalIds: string[], transport?: PiTransport, + onStartupFailure?: (error: Error) => void, ): void { const { command, success } = response; const resolver = session.rpcResolver!; @@ -153,19 +157,46 @@ function handleResponse( const oldestLocalId = pendingLocalIds.shift()!; session.emitMessagesConsumed([oldestLocalId], { clearQueuedThinkingGrace: true }); } + // A failed initial get_state means Pi did not load its native session. + // Do not leave the HAPI wrapper alive until the hub's ready timeout: the + // caller tears down the process so the archived row can be restored. + // A fresh Pi session keeps the historic non-fatal fallback behavior; + // only a requested native resume must fail closed. + if (command === 'get_state' && session.expectedNativeSessionId && !session.isNativeReady) { + onStartupFailure?.(new Error(`Pi get_state failed: ${error}`)); + } return; } switch (command) { case 'get_state': { - handleGetState(response.data, session); + const parsed = PiStateDataSchema.safeParse(response.data); // Pi has finished startup init (this is the response that persists - // metadata.piSessionId — the signal working callers already wait - // for). Release any prompts buffered during the spawn window so they - // reach an initialized Pi session instead of wedging (issue #1143). - // markReady is idempotent; a missing sessionId still flips ready so - // buffered prompts are never swallowed forever. - session.markReady(); + // metadata.piSessionId). It is also the only native-ready signal + // that the hub trusts for Pi resume; session-alive only proves the + // HAPI wrapper connected. Validate a requested native session before + // mutating model/metadata state: an invalid resume must not publish a + // colliding piSessionId that auto-dedup could merge. + if (!parsed.success) { + if (session.expectedNativeSessionId) { + onStartupFailure?.(new Error('Pi get_state returned malformed state data')); + } + break; + } + const state = parsed.data; + if (!session.matchesExpectedNativeSessionId(state.sessionId)) { + const actual = state.sessionId ? state.sessionId : '(missing)'; + const error = `Pi loaded unexpected native session ${actual} instead of ${session.expectedNativeSessionId}`; + logger.debug(`[pi] ${error}`); + session.sendSessionEvent({ type: 'message', message: error }); + onStartupFailure?.(new Error(error)); + break; + } + // Emit ready before publishing Pi metadata. On native resume, this + // ensures the hub can never merge based on a piSessionId before the + // get_state identity check has completed. + session.markNativeReady(); + applyGetState(state, session); break; } case 'set_model': { @@ -305,6 +336,7 @@ export function wireTransportEvents( transport: PiTransport, session: PiSession, pendingLocalIds: string[], + options?: { onStartupFailure?: (error: Error) => void }, ): void { session.rpcResolver = new PiRpcResolver(); const assistantMessageAccumulator = new PiMessageAccumulator(); @@ -316,7 +348,13 @@ export function wireTransportEvents( logger.debug(`[pi][event] ${event.type}`); } if (event.type === 'response') { - handleResponse(event as unknown as PiResponseEvent, session, pendingLocalIds, transport); + handleResponse( + event as unknown as PiResponseEvent, + session, + pendingLocalIds, + transport, + options?.onStartupFailure, + ); return; } diff --git a/cli/src/pi/runPi.test.ts b/cli/src/pi/runPi.test.ts index 9734399b..b45b1438 100644 --- a/cli/src/pi/runPi.test.ts +++ b/cli/src/pi/runPi.test.ts @@ -6,6 +6,10 @@ type LifecycleOptions = { stopKeepAlive: () => void }; const harness = vi.hoisted(() => ({ transportOptions: null as TransportOptions | null, sent: [] as unknown[], + throwOnGetCommands: true, + onError: null as ((error: Error) => void) | null, + killCount: 0, + cleanupCount: 0, session: { keepAlive: vi.fn(), onUserMessage: vi.fn(), @@ -24,6 +28,7 @@ vi.mock('@/agent/runnerLifecycle', () => ({ return { registerProcessHandlers: vi.fn(), cleanupAndExit: vi.fn(async () => { + harness.cleanupCount += 1; options.stopKeepAlive(); }), markCrash: vi.fn(), @@ -50,7 +55,9 @@ vi.mock('./piTransport', () => ({ harness.transportOptions = options; } - onError(): void {} + onError(callback: (error: Error) => void): void { + harness.onError = callback; + } onClose(): void {} @@ -60,16 +67,20 @@ vi.mock('./piTransport', () => ({ send(command: unknown): void { harness.sent.push(command); - if ((command as { type?: string }).type === 'get_commands') { + if (harness.throwOnGetCommands && (command as { type?: string }).type === 'get_commands') { throw new Error('stop test transport'); } } - kill(): void {} + kill(): void { + harness.killCount += 1; + } }, })); import { buildPiCommandInventory, formatPiUserMessage, rewritePiSkillPrompt, runPi } from './runPi'; +import { bootstrapExistingSession } from '@/agent/sessionFactory'; +import { PiSession } from './session'; describe('Pi command namespaces', () => { const commands = [ @@ -112,6 +123,11 @@ describe('runPi startup', () => { beforeEach(() => { harness.transportOptions = null; harness.sent.length = 0; + harness.throwOnGetCommands = true; + harness.onError = null; + harness.killCount = 0; + harness.cleanupCount = 0; + vi.useRealTimers(); }); it('lets Pi create a fresh session when no resume ID is provided', async () => { @@ -146,4 +162,38 @@ describe('runPi startup', () => { { type: 'get_commands' }, ]); }); + + it('bootstraps the existing HAPI row for runner native resume', async () => { + await runPi({ + workingDirectory: '/work', + existingSessionId: 'hapi-session-pi-1', + resumeSessionId: 'pi-session-1', + startedBy: 'runner', + }); + + expect(bootstrapExistingSession).toHaveBeenCalledWith({ + sessionId: 'hapi-session-pi-1', + flavor: 'pi', + startedBy: 'runner', + workingDirectory: '/work', + }); + }); + + it.each([ + ['fresh', undefined, 1, 0], + ['resume', 'pi-session-1', 0, 1], + ] as const)('applies the startup fallback only to %s sessions', async (_label, resumeSessionId, expectedCalls, expectedKills) => { + vi.useFakeTimers(); + harness.throwOnGetCommands = false; + const markReady = vi.spyOn(PiSession.prototype, 'markReady'); + const running = runPi({ workingDirectory: '/work', resumeSessionId }); + + await vi.advanceTimersByTimeAsync(31_000); + expect(markReady).toHaveBeenCalledTimes(expectedCalls); + expect(harness.cleanupCount).toBe(expectedKills); + + harness.onError?.(new Error('stop test transport')); + await running; + markReady.mockRestore(); + }); }); diff --git a/cli/src/pi/runPi.ts b/cli/src/pi/runPi.ts index 51b4ba16..d809fae5 100644 --- a/cli/src/pi/runPi.ts +++ b/cli/src/pi/runPi.ts @@ -120,6 +120,7 @@ export async function runPi(opts: { startedBy, startingMode, model: opts.model, + expectedNativeSessionId: opts.resumeSessionId, }); const transportArgs = ['--mode', 'rpc']; @@ -153,6 +154,19 @@ export async function runPi(opts: { await lifecycle.cleanupAndExit(); }; + // Install the completion hook before transport.start(). Pi can synchronously + // report an invalid --session during the first get_state send; installing it + // later would leave runPi awaiting a promise that can no longer be resolved. + let resolveCleanupCompletion!: () => void; + const cleanupCompletion = new Promise((resolve) => { + resolveCleanupCompletion = resolve; + }); + const originalCleanupAndExit = lifecycle.cleanupAndExit.bind(lifecycle); + lifecycle.cleanupAndExit = async (codeOverride?: number) => { + resolveCleanupCompletion(); + await originalCleanupAndExit(codeOverride); + }; + // Pending user-message localIds in FIFO order const pendingLocalIds: string[] = []; @@ -200,7 +214,21 @@ export async function runPi(opts: { } } - wireTransportEvents(transport, piSession, pendingLocalIds); + const failNativeStartup = (error: Error) => { + // A wrapper socket can already be active while Pi rejects --session. + // End this process immediately so the hub's Pi-ready wait fails and + // an archived HAPI row is restored rather than shown as reopened. + logger.debug(`[pi] Native startup failed: ${error.message}`); + lifecycle.markCrash(error); + lifecycle.setExitCode(1); + lifecycle.setArchiveReason(error.message.slice(0, 200)); + lifecycle.setSessionEndReason('error'); + void safeCleanup(); + }; + + wireTransportEvents(transport, piSession, pendingLocalIds, { + onStartupFailure: failNativeStartup, + }); // --- Session config RPC --- // @@ -437,7 +465,10 @@ export async function runPi(opts: { // This degrades to pre-fix behaviour (send anyway) instead of something // worse. markReady is idempotent, so a real get_state that lands first wins. const readyFallback = setTimeout(() => { - if (!piSession.isReady) { + if (piSession.isReady) return; + if (piSession.expectedNativeSessionId) { + failNativeStartup(new Error(`Pi native resume did not become ready within ${PI_READY_FALLBACK_MS}ms`)); + } else { logger.debug('[pi] get_state ready signal not seen within grace — draining buffered messages'); piSession.markReady(); } @@ -474,14 +505,8 @@ export async function runPi(opts: { })(); } - // Block until cleanup is triggered by error/close handler - await new Promise((resolve) => { - const origCleanup = lifecycle.cleanupAndExit.bind(lifecycle); - lifecycle.cleanupAndExit = async (codeOverride?: number) => { - resolve(); - await origCleanup(codeOverride); - }; - }); + // Block until cleanup is triggered by error/close handler. + await cleanupCompletion; } catch (error) { crashed = true; lifecycle.markCrash(error); diff --git a/cli/src/pi/session.test.ts b/cli/src/pi/session.test.ts index b4f00c2b..7bcca139 100644 --- a/cli/src/pi/session.test.ts +++ b/cli/src/pi/session.test.ts @@ -18,6 +18,7 @@ function createMockSession(): PiSession { sendAgentMessage: vi.fn(), emitMessagesConsumed: vi.fn(), sendSessionEvent: vi.fn(), + emitSessionReady: vi.fn(), } as any, path: '/tmp/test', logPath: '/tmp/test.log', @@ -78,6 +79,22 @@ describe('PiSession ready gate', () => { expect(fn).toHaveBeenCalledTimes(1); }); + it('only announces native-ready after successful get_state, not fallback readiness', () => { + const session = createMockSession(); + + session.markReady(); + expect(session.client.emitSessionReady).not.toHaveBeenCalled(); + + session.markNativeReady(); + expect(session.client.emitSessionReady).toHaveBeenCalledTimes(1); + + const nativeSession = createMockSession(); + nativeSession.markNativeReady(); + nativeSession.markNativeReady(); + + expect(nativeSession.client.emitSessionReady).toHaveBeenCalledTimes(1); + }); + it('preserves FIFO across mixed buffered + post-ready enqueues', () => { const session = createMockSession(); const order: string[] = []; diff --git a/cli/src/pi/session.ts b/cli/src/pi/session.ts index 3283f867..7e53b69f 100644 --- a/cli/src/pi/session.ts +++ b/cli/src/pi/session.ts @@ -34,6 +34,9 @@ export class PiSession { // Startup model from opts.model — prevents get_state from overwriting it // with Pi's default. Applied once when get_available_models returns. readonly initialModel: string | null; + // A runner/native resume must prove that Pi loaded this exact session with + // a non-empty get_state sessionId. Missing or contradictory IDs fail closed. + readonly expectedNativeSessionId: string | null; // Streaming state piIsStreaming = false; @@ -53,6 +56,7 @@ export class PiSession { // are queued via runWhenReady() and drained FIFO once markReady() fires (on // the first get_state response). private piReady = false; + private nativeReadyAnnounced = false; // Buffered sends carry their localId so a cancel-queued-message that arrives // while a prompt is still held (before drain) can drop it instead of firing // a cancelled prompt on markReady (issue #1143 review — MAJOR). @@ -68,6 +72,7 @@ export class PiSession { startedBy: 'runner' | 'terminal'; startingMode: 'local' | 'remote'; model?: string | null; + expectedNativeSessionId?: string; }) { this.api = opts.api; this.client = opts.client; @@ -85,6 +90,7 @@ export class PiSession { // resume before Pi reports its real state. this.currentModel = undefined; this.initialModel = opts.model?.trim() || null; + this.expectedNativeSessionId = opts.expectedNativeSessionId?.trim() || null; this.currentThinkingLevel = undefined; } @@ -93,6 +99,16 @@ export class PiSession { return this.piReady; } + /** True only after Pi itself has completed a successful get_state. */ + get isNativeReady(): boolean { + return this.nativeReadyAnnounced; + } + + matchesExpectedNativeSessionId(actualSessionId: string | undefined): boolean { + if (!this.expectedNativeSessionId) return true; + return Boolean(actualSessionId) && actualSessionId === this.expectedNativeSessionId; + } + /** * Run `fn` now if Pi startup is ready, else buffer it FIFO until markReady(). * Used to gate outbound prompt/steer sends so they never reach Pi before its @@ -121,16 +137,29 @@ export class PiSession { } /** - * Signal that Pi RPC startup is complete (first get_state response). - * Drains buffered sends in enqueue order. Idempotent — later get_state - * responses (or the startup fallback timer) are no-ops. + * Release buffered sends. This intentionally does not notify the hub: the + * startup fallback may use it to avoid losing prompts, but is not proof that + * Pi loaded a requested native session. */ - markReady(): void { - if (this.piReady) return; + markReady(): boolean { + if (this.piReady) return false; this.piReady = true; const queued = this.readyQueue; this.readyQueue = []; for (const { fn } of queued) fn(); + return true; + } + + /** + * The first successful Pi get_state is the authoritative native-ready + * point. Unlike markReady(), this tells the hub that a native resume can be + * considered successful. + */ + markNativeReady(): void { + this.markReady(); + if (this.nativeReadyAnnounced) return; + this.nativeReadyAnnounced = true; + this.client.emitSessionReady(); } startKeepAlive(): void { diff --git a/cli/src/runner/buildCliArgs.test.ts b/cli/src/runner/buildCliArgs.test.ts index 74959271..3cffb258 100644 --- a/cli/src/runner/buildCliArgs.test.ts +++ b/cli/src/runner/buildCliArgs.test.ts @@ -218,6 +218,19 @@ describe('buildCliArgs', () => { expect(args[0]).toBe('pi') }) + it('reuses the original HAPI row for Pi native resume', () => { + const args = buildCliArgs('pi', { + directory: '/tmp', + resumeSessionId: 'pi-native-session-1', + existingSessionId: 'hapi-session-pi-1', + }) + + expect(args).toContain('--session-id') + expect(args).toContain('pi-native-session-1') + expect(args).toContain('--existing-session-id') + expect(args).toContain('hapi-session-pi-1') + }) + it('still passes --resume for claude when resumeSessionId is provided', () => { // Guard against accidentally swallowing claude's --resume when // the pi branch was added. diff --git a/cli/src/runner/controlClient.test.ts b/cli/src/runner/controlClient.test.ts new file mode 100644 index 00000000..d3cab5d6 --- /dev/null +++ b/cli/src/runner/controlClient.test.ts @@ -0,0 +1,55 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const { readRunnerStateMock, isProcessAliveMock } = vi.hoisted(() => ({ + readRunnerStateMock: vi.fn(), + isProcessAliveMock: vi.fn(), +})) + +vi.mock('@/persistence', () => ({ + readRunnerState: readRunnerStateMock, + readSettings: vi.fn(), + clearRunnerState: vi.fn(), +})) + +vi.mock('@/utils/process', () => ({ + isProcessAlive: isProcessAliveMock, + isHapiRunnerProcess: vi.fn(() => true), + killProcess: vi.fn(), +})) + +vi.mock('@/ui/logger', () => ({ logger: { debug: vi.fn() } })) + +import { stopRunnerSession } from './controlClient' + +describe('runner control client stop-session contract', () => { + beforeEach(() => { + readRunnerStateMock.mockResolvedValue({ pid: 42, httpPort: 3210 }) + isProcessAliveMock.mockReturnValue(true) + }) + + afterEach(() => { + vi.unstubAllGlobals() + vi.clearAllMocks() + }) + + it.each(['stopped', 'already_gone', 'still_alive'] as const)( + 'returns the runner %s status', + async (status) => { + vi.stubGlobal('fetch', vi.fn(async () => new Response(JSON.stringify({ status }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }))) + + await expect(stopRunnerSession('session-1')).resolves.toBe(status) + } + ) + + it('fails closed when the runner returns a malformed response', async () => { + vi.stubGlobal('fetch', vi.fn(async () => new Response(JSON.stringify({ success: true }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }))) + + await expect(stopRunnerSession('session-1')).resolves.toBe('still_alive') + }) +}) diff --git a/cli/src/runner/controlClient.ts b/cli/src/runner/controlClient.ts index f231d8f7..cd5868ce 100644 --- a/cli/src/runner/controlClient.ts +++ b/cli/src/runner/controlClient.ts @@ -96,9 +96,11 @@ export async function listRunnerSessions(): Promise { return result.children || []; } -export async function stopRunnerSession(sessionId: string): Promise { +export async function stopRunnerSession(sessionId: string): Promise<'stopped' | 'already_gone' | 'still_alive'> { const result = await runnerPost('/stop-session', { sessionId }); - return result.success || false; + return result.status === 'stopped' || result.status === 'already_gone' || result.status === 'still_alive' + ? result.status + : 'still_alive'; } export async function spawnRunnerSession(directory: string, sessionId?: string): Promise { diff --git a/cli/src/runner/controlServer.ts b/cli/src/runner/controlServer.ts index b07278c8..358f863e 100644 --- a/cli/src/runner/controlServer.ts +++ b/cli/src/runner/controlServer.ts @@ -19,7 +19,7 @@ export function startRunnerControlServer({ onHappySessionWebhook }: { getChildren: () => TrackedSession[]; - stopSession: (sessionId: string) => boolean; + stopSession: (sessionId: string) => Promise<'stopped' | 'already_gone' | 'still_alive'>; spawnSession: (options: SpawnSessionOptions) => Promise; requestShutdown: () => void; onHappySessionWebhook: (sessionId: string, metadata: Metadata) => void; @@ -91,7 +91,7 @@ export function startRunnerControlServer({ }), response: { 200: z.object({ - success: z.boolean() + status: z.enum(['stopped', 'already_gone', 'still_alive']) }) } } @@ -99,8 +99,8 @@ export function startRunnerControlServer({ const { sessionId } = request.body; logger.debug(`[CONTROL SERVER] Stop session request: ${sessionId}`); - const success = stopSession(sessionId); - return { success }; + const status = await stopSession(sessionId); + return { status }; }); // Spawn new session diff --git a/cli/src/runner/run.ts b/cli/src/runner/run.ts index 0fecddb9..2d19bdd1 100644 --- a/cli/src/runner/run.ts +++ b/cli/src/runner/run.ts @@ -1,4 +1,5 @@ import fs from 'fs/promises'; +import { existsSync, readFileSync, renameSync, writeFileSync } from 'node:fs'; import os from 'os'; import { ApiClient } from '@/api/api'; @@ -13,7 +14,7 @@ import { getEnvironmentInfo } from '@/ui/doctor'; import { spawnHappyCLI } from '@/utils/spawnHappyCLI'; import { writeRunnerState, RunnerLocallyPersistedState, readRunnerState, acquireRunnerLock, releaseRunnerLock } from '@/persistence'; import { getCliArgs } from '@/utils/cliArgs'; -import { isProcessAlive, isWindows, killProcess, killProcessByChildProcess } from '@/utils/process'; +import { getProcessStartMarker, isProcessAlive, isWindows, killProcess, killProcessByChildProcess, killProcessTreeByPid } from '@/utils/process'; import { PERMISSION_MODES } from '@hapi/protocol/modes'; import { withRetry } from '@/utils/time'; import { isRetryableConnectionError } from '@/utils/errorUtils'; @@ -180,6 +181,116 @@ export async function startRunner(options: { workspaceRoots?: string[] } = {}): // Setup state - key by PID const pidToTrackedSession = new Map(); + // Retained until actual child exit even if webhook timeout removes normal + // tracking, so confirmed exit can be attributed to the requested HAPI row. + const pidToRequestedSessionId = new Map(); + const pidToConfirmedSessionId = new Map(); + // Only actual observed child exits may create a stop-session tombstone. + // Tracking loss (notably webhook timeout) is deliberately not evidence. + const exitTombstoneFile = `${configuration.runnerStateFile}.verified-exits.json`; + const verifiedExitTombstones = (() => { + try { + if (!existsSync(exitTombstoneFile)) return new Set(); + const parsed = JSON.parse(readFileSync(exitTombstoneFile, 'utf8')); + return new Set(Array.isArray(parsed) ? parsed.filter((value): value is string => typeof value === 'string' && !value.startsWith('PID-')) : []); + } catch (error) { + logger.debug('[RUNNER RUN] Failed to load verified exit tombstones:', error); + return new Set(); + } + })(); + // PID aliases are generation-local and must not survive runner restart, + // because the OS may reuse a PID for an unrelated process. + const verifiedPidExitTombstones = new Set(); + const persistVerifiedExits = () => { + const tmp = `${exitTombstoneFile}.${process.pid}.tmp`; + try { + writeFileSync(tmp, JSON.stringify([...verifiedExitTombstones])); + renameSync(tmp, exitTombstoneFile); + } catch (error) { + logger.debug('[RUNNER RUN] Failed to persist verified exit tombstones:', error); + } + }; + const rememberVerifiedExit = (id: string) => { + if (id.startsWith('PID-')) { + verifiedPidExitTombstones.add(id); + return; + } + // Refreshing an existing key should also refresh its insertion order so + // capacity eviction removes the oldest verified generation. + verifiedExitTombstones.delete(id); + verifiedExitTombstones.add(id); + persistVerifiedExits(); + }; + const hasVerifiedExit = (id: string): boolean => { + return id.startsWith('PID-') + ? verifiedPidExitTombstones.has(id) + : verifiedExitTombstones.has(id); + }; + const invalidateVerifiedExit = (id: string) => { + if (id.startsWith('PID-')) { + verifiedPidExitTombstones.delete(id); + } else if (verifiedExitTombstones.delete(id)) { + persistVerifiedExits(); + } + }; + + type PersistedResumeProcess = { + requestedSessionId: string; + confirmedSessionId?: string; + pid: number; + processStartMarker: string; + }; + const resumeProcessFile = `${configuration.runnerStateFile}.resume-processes.json`; + const persistedResumeProcesses = (() => { + try { + if (!existsSync(resumeProcessFile)) return new Map(); + const parsed = JSON.parse(readFileSync(resumeProcessFile, 'utf8')); + const records = Array.isArray(parsed) ? parsed : []; + return new Map(records.flatMap((record): Array<[number, PersistedResumeProcess]> => { + const requestedSessionId = typeof record?.requestedSessionId === 'string' + ? record.requestedSessionId + : typeof record?.sessionId === 'string' + ? record.sessionId + : null; + if (!requestedSessionId || typeof record.pid !== 'number' || typeof record.processStartMarker !== 'string') return []; + return [[record.pid, { + requestedSessionId, + confirmedSessionId: typeof record.confirmedSessionId === 'string' ? record.confirmedSessionId : undefined, + pid: record.pid, + processStartMarker: record.processStartMarker, + }]]; + })); + } catch (error) { + logger.debug('[RUNNER RUN] Failed to load persisted resume processes:', error); + return new Map(); + } + })(); + const persistResumeProcesses = () => { + const tmp = `${resumeProcessFile}.${process.pid}.tmp`; + try { + writeFileSync(tmp, JSON.stringify([...persistedResumeProcesses.values()])); + renameSync(tmp, resumeProcessFile); + } catch (error) { + logger.debug('[RUNNER RUN] Failed to persist resume processes:', error); + } + }; + for (const [pid, record] of [...persistedResumeProcesses]) { + const alive = isProcessAlive(pid); + const marker = alive ? getProcessStartMarker(pid) : null; + if (alive && marker === record.processStartMarker) { + pidToRequestedSessionId.set(pid, record.requestedSessionId); + if (record.confirmedSessionId) pidToConfirmedSessionId.set(pid, record.confirmedSessionId); + } else if (!alive || marker !== null) { + persistedResumeProcesses.delete(pid); + rememberVerifiedExit(record.requestedSessionId); + if (record.confirmedSessionId) rememberVerifiedExit(record.confirmedSessionId); + } else { + // PID is live but generation probing failed: keep the durable record and + // fail closed instead of manufacturing verified-exit evidence. + logger.debug(`[RUNNER RUN] Could not verify process generation for PID ${pid}; keeping persisted resume quarantine`); + } + } + persistResumeProcesses(); // Webhook timeout tolerance. Opus 1M + --resume can legitimately take // longer than the default 15s to reach the "Session started" webhook @@ -231,7 +342,15 @@ export async function startRunner(options: { workspaceRoots?: string[] } = {}): if (existingSession && existingSession.startedBy === 'runner') { // Update runner-spawned session with reported data + invalidateVerifiedExit(sessionId); + invalidateVerifiedExit(`PID-${pid}`); existingSession.happySessionId = sessionId; + pidToConfirmedSessionId.set(pid, sessionId); + const persisted = persistedResumeProcesses.get(pid); + if (persisted) { + persisted.confirmedSessionId = sessionId; + persistResumeProcesses(); + } existingSession.happySessionMetadataFromLocalWebhook = sessionMetadata; logger.debug(`[RUNNER RUN] Updated runner-spawned session ${sessionId} with metadata`); @@ -276,6 +395,8 @@ export async function startRunner(options: { workspaceRoots?: string[] } = {}): happySessionMetadataFromLocalWebhook: sessionMetadata, pid }; + invalidateVerifiedExit(sessionId); + invalidateVerifiedExit(`PID-${pid}`); pidToTrackedSession.set(pid, trackedSession); logger.debug(`[RUNNER RUN] Registered externally-started session ${sessionId}`); } @@ -485,7 +606,14 @@ export async function startRunner(options: { workspaceRoots?: string[] } = {}): } happyProcess.removeListener('error', captureSpawnErrorBeforePidCheck); + // The OS process now exists, so this is the point where a new generation + // invalidates exit evidence left by an older child with the same HAPI ID. + for (const id of [options.sessionId, options.existingSessionId]) { + if (id) invalidateVerifiedExit(id); + } + const pid = happyProcess.pid; + invalidateVerifiedExit(`PID-${pid}`); logger.debug(`[RUNNER RUN] Spawned process with PID ${pid}`); let observedExitCode: number | null = null; let observedExitSignal: NodeJS.Signals | null = null; @@ -520,12 +648,25 @@ export async function startRunner(options: { workspaceRoots?: string[] } = {}): const trackedSession: TrackedSession = { startedBy: 'runner', pid, + requestedHappySessionId: options.existingSessionId ?? options.sessionId, childProcess: happyProcess, directoryCreated, message: directoryCreated ? `The path '${directory}' did not exist. We created a new folder and spawned a new session there.` : undefined }; pidToTrackedSession.set(pid, trackedSession); + if (trackedSession.requestedHappySessionId) { + pidToRequestedSessionId.set(pid, trackedSession.requestedHappySessionId); + const processStartMarker = getProcessStartMarker(pid); + if (processStartMarker) { + persistedResumeProcesses.set(pid, { + requestedSessionId: trackedSession.requestedHappySessionId, + pid, + processStartMarker + }); + persistResumeProcesses(); + } + } happyProcess.on('exit', (code, signal) => { observedExitCode = typeof code === 'number' ? code : null; @@ -551,7 +692,9 @@ export async function startRunner(options: { workspaceRoots?: string[] } = {}): pidToAwaiter.delete(pid); errorAwaiter(buildWebhookFailureMessage('process-error-before-webhook')); } - onChildExited(pid); + // A ChildProcess error is not itself proof that the OS process exited. + // Keep tracking a live PID so machine StopSession can still terminate it. + if (!isProcessAlive(pid)) onChildExited(pid); }); // Wait for webhook to populate session with happySessionId @@ -648,47 +791,125 @@ export async function startRunner(options: { workspaceRoots?: string[] } = {}): }; // Stop a session by sessionId or PID fallback - const stopSession = (sessionId: string): boolean => { + const stopSession = async (sessionId: string): Promise<'stopped' | 'already_gone' | 'still_alive'> => { logger.debug(`[RUNNER RUN] Attempting to stop session ${sessionId}`); // Try to find by sessionId first for (const [pid, session] of pidToTrackedSession.entries()) { if (session.happySessionId === sessionId || + session.requestedHappySessionId === sessionId || (sessionId.startsWith('PID-') && pid === parseInt(sessionId.replace('PID-', '')))) { if (session.startedBy === 'runner' && session.childProcess) { try { - void killProcessByChildProcess(session.childProcess); + const treeStopped = await killProcessByChildProcess(session.childProcess); + if (!treeStopped) { + logger.debug(`[RUNNER RUN] Process tree for session ${sessionId} is still alive after stop request`); + return 'still_alive'; + } logger.debug(`[RUNNER RUN] Requested termination for runner-spawned session ${sessionId}`); } catch (error) { logger.debug(`[RUNNER RUN] Failed to kill session ${sessionId}:`, error); + return 'still_alive'; } } else { // For externally started sessions, try to kill by PID try { - void killProcess(pid); + if (!(await killProcess(pid))) return 'still_alive'; logger.debug(`[RUNNER RUN] Requested termination for external session PID ${pid}`); } catch (error) { logger.debug(`[RUNNER RUN] Failed to kill external session PID ${pid}:`, error); + return 'still_alive'; } } + const deadline = Date.now() + 5_000; + while (isProcessAlive(pid) && Date.now() < deadline) { + await new Promise(resolve => setTimeout(resolve, 50)); + } + if (isProcessAlive(pid)) { + logger.debug(`[RUNNER RUN] Session ${sessionId} process ${pid} is still alive after stop request`); + return 'still_alive'; + } + if (session.happySessionId) rememberVerifiedExit(session.happySessionId); + if (session.requestedHappySessionId) rememberVerifiedExit(session.requestedHappySessionId); + rememberVerifiedExit(`PID-${pid}`); pidToTrackedSession.delete(pid); - logger.debug(`[RUNNER RUN] Removed session ${sessionId} from tracking`); - return true; + pidToRequestedSessionId.delete(pid); + pidToConfirmedSessionId.delete(pid); + if (persistedResumeProcesses.delete(pid)) persistResumeProcesses(); + logger.debug(`[RUNNER RUN] Removed terminated session ${sessionId} from tracking`); + return 'stopped'; } } - logger.debug(`[RUNNER RUN] Session ${sessionId} not found`); - return false; + // Webhook timeout can remove the normal TrackedSession before the process + // actually exits. Retain the requested HAPI ID -> PID relation so Hub can + // still terminate that exact generation by HAPI ID. + const fallbackPids = new Set([ + ...pidToRequestedSessionId.keys(), + ...pidToConfirmedSessionId.keys(), + ...persistedResumeProcesses.keys(), + ]); + for (const pid of fallbackPids) { + const persisted = persistedResumeProcesses.get(pid); + const requestedSessionId = pidToRequestedSessionId.get(pid) ?? persisted?.requestedSessionId; + const confirmedSessionId = pidToConfirmedSessionId.get(pid) ?? persisted?.confirmedSessionId; + if (requestedSessionId !== sessionId && confirmedSessionId !== sessionId) continue; + if (isProcessAlive(pid)) { + if (!persisted) return 'still_alive'; + const currentMarker = getProcessStartMarker(pid); + if (currentMarker === null) return 'still_alive'; + if (currentMarker !== persisted.processStartMarker) { + persistedResumeProcesses.delete(pid); + persistResumeProcesses(); + pidToRequestedSessionId.delete(pid); + pidToConfirmedSessionId.delete(pid); + if (requestedSessionId) rememberVerifiedExit(requestedSessionId); + if (confirmedSessionId) rememberVerifiedExit(confirmedSessionId); + return 'already_gone'; + } + if (!(await killProcessTreeByPid(pid))) return 'still_alive'; + if (requestedSessionId) rememberVerifiedExit(requestedSessionId); + if (confirmedSessionId) rememberVerifiedExit(confirmedSessionId); + rememberVerifiedExit(`PID-${pid}`); + pidToRequestedSessionId.delete(pid); + pidToConfirmedSessionId.delete(pid); + if (persistedResumeProcesses.delete(pid)) persistResumeProcesses(); + return 'stopped'; + } + if (requestedSessionId) rememberVerifiedExit(requestedSessionId); + if (confirmedSessionId) rememberVerifiedExit(confirmedSessionId); + rememberVerifiedExit(`PID-${pid}`); + pidToRequestedSessionId.delete(pid); + pidToConfirmedSessionId.delete(pid); + if (persistedResumeProcesses.delete(pid)) persistResumeProcesses(); + return 'already_gone'; + } + + if (hasVerifiedExit(sessionId)) { + logger.debug(`[RUNNER RUN] Session ${sessionId} was previously observed exited`); + return 'already_gone'; + } + logger.debug(`[RUNNER RUN] Session ${sessionId} not found without verified exit`); + return 'still_alive'; }; // Handle child process exit const onChildExited = (pid: number) => { + const session = pidToTrackedSession.get(pid); + const requestedSessionId = session?.requestedHappySessionId ?? pidToRequestedSessionId.get(pid); + if (requestedSessionId) rememberVerifiedExit(requestedSessionId); + const confirmedSessionId = session?.happySessionId ?? pidToConfirmedSessionId.get(pid); + if (confirmedSessionId) rememberVerifiedExit(confirmedSessionId); + rememberVerifiedExit(`PID-${pid}`); logger.debug(`[RUNNER RUN] Removing exited process PID ${pid} from tracking`); pidToTrackedSession.delete(pid); pidToAwaiter.delete(pid); pidToErrorAwaiter.delete(pid); + pidToRequestedSessionId.delete(pid); + pidToConfirmedSessionId.delete(pid); + if (persistedResumeProcesses.delete(pid)) persistResumeProcesses(); }; // Start control server @@ -757,7 +978,8 @@ export async function startRunner(options: { workspaceRoots?: string[] } = {}): status: 'offline', pid: process.pid, httpPort: controlPort, - startedAt: Date.now() + startedAt: Date.now(), + capabilities: { piExistingSessionResume: true } }; // Create API client @@ -1115,10 +1337,10 @@ export function buildCliArgs( } } args.push('--hapi-starting-mode', 'remote', '--started-by', 'runner'); - // Codex import/resume (#1088) and Cursor ACP remote resume (#991) both reuse - // the original HAPI row via --existing-session-id so the hub does not depend - // on session-ready over a remote socket before merge. - if (agent === 'codex' || agent === 'cursor') { + // Codex, Cursor ACP, and Pi native resume reuse the original HAPI row via + // --existing-session-id. Pi is reported successful only after the hub sees + // its validated native get_state/session-ready signal. + if (agent === 'codex' || agent === 'cursor' || agent === 'pi') { const existingSessionId = options.existingSessionId ?? options.sessionId; if (existingSessionId) { args.push('--existing-session-id', existingSessionId); diff --git a/cli/src/runner/runner.integration.test.ts b/cli/src/runner/runner.integration.test.ts index 9d3f1547..730c9c89 100644 --- a/cli/src/runner/runner.integration.test.ts +++ b/cli/src/runner/runner.integration.test.ts @@ -159,7 +159,9 @@ describe.skipIf(!await isServerHealthy())('Runner Integration Tests', { timeout: // Clean up - stop the spawned session expect(spawnedSession.happySessionId).toBeDefined(); - await stopRunnerSession(spawnedSession.happySessionId); + expect(await stopRunnerSession(spawnedSession.happySessionId)).toBe('stopped'); + expect(await stopRunnerSession(spawnedSession.happySessionId)).toBe('already_gone'); + expect(await stopRunnerSession('unknown-session-id')).toBe('still_alive'); }); it('stress test: spawn / stop', { timeout: 60_000 }, async () => { @@ -178,7 +180,7 @@ describe.skipIf(!await isServerHealthy())('Runner Integration Tests', { timeout: // Stop all sessions const stopResults = await Promise.all(sessionIds.map(sessionId => stopRunnerSession(sessionId))); - expect(stopResults.every(r => r), 'Not all sessions reported stopped').toBe(true); + expect(stopResults.every(r => r === 'stopped' || r === 'already_gone'), 'Not all sessions reported stopped').toBe(true); // Verify all sessions are stopped const emptySessions = await listRunnerSessions(); diff --git a/cli/src/runner/types.ts b/cli/src/runner/types.ts index 0f060fd5..5c8907ca 100644 --- a/cli/src/runner/types.ts +++ b/cli/src/runner/types.ts @@ -11,10 +11,12 @@ import { ChildProcess } from 'child_process'; export interface TrackedSession { startedBy: 'runner' | string; happySessionId?: string; + /** HAPI row requested for this process generation before its webhook arrives. */ + requestedHappySessionId?: string; happySessionMetadataFromLocalWebhook?: Metadata; pid: number; childProcess?: ChildProcess; error?: string; directoryCreated?: boolean; message?: string; -} \ No newline at end of file +} diff --git a/cli/src/utils/process.ts b/cli/src/utils/process.ts index 3c82d82c..8d14eddc 100644 --- a/cli/src/utils/process.ts +++ b/cli/src/utils/process.ts @@ -16,6 +16,33 @@ export function isProcessAlive(pid: number): boolean { } } +/** Stable marker for one OS PID generation; null means the platform probe failed. */ +export function getProcessStartMarker(pid: number): string | null { + if (!isProcessAlive(pid)) return null; + if (isWindows()) { + const powershell = spawn.sync('powershell', [ + '-NoProfile', '-NonInteractive', '-Command', + `(Get-CimInstance Win32_Process -Filter "ProcessId = ${pid}").CreationDate` + ], { stdio: 'pipe', windowsHide: true }); + if (!powershell.error && powershell.status === 0) { + const marker = powershell.stdout?.toString().trim(); + if (marker) return marker; + } + const result = spawn.sync('wmic', [ + 'process', 'where', `ProcessId=${pid}`, 'get', 'CreationDate', '/value' + ], { stdio: 'pipe', windowsHide: true }); + if (result.error || result.status !== 0) return null; + const match = (result.stdout?.toString() ?? '').match(/CreationDate=([^\r\n]+)/); + return match?.[1]?.trim() || null; + } + const result = spawn.sync('ps', ['-p', String(pid), '-o', 'lstart='], { + stdio: 'pipe', + env: { ...process.env, LC_ALL: 'C', TZ: 'UTC' } + }); + if (result.error || result.status !== 0) return null; + return result.stdout?.toString().trim() || null; +} + // ponytail: ps -p is cheap and avoids PID-reuse false positives after OS upgrades/reboots function isRunnerCommand(commandLine: string): boolean { return /(?:^|\s)runner(?:\s|$)/.test(commandLine) && /(?:^|\s)start-sync(?:\s|$)/.test(commandLine); @@ -138,7 +165,14 @@ async function killProcessTree(pid: number, force: boolean): Promise { await waitForProcessToDie(p, force); } - return true; + return pids.every((candidate) => !isProcessAlive(candidate)); +} + +/** Kill a PID and all descendants, verifying the complete tree is gone. */ +export async function killProcessTreeByPid(pid: number, force: boolean = false): Promise { + if (!Number.isFinite(pid) || pid <= 0) return false; + if (isWindows()) return killProcess(pid, force); + return killProcessTree(pid, force); } /** @@ -184,5 +218,5 @@ export async function killProcessByChildProcess( } // Kill entire process tree on Unix to prevent orphan processes - return killProcessTree(pid, force); + return killProcessTreeByPid(pid, force); } diff --git a/hub/src/sync/rpcGateway.ts b/hub/src/sync/rpcGateway.ts index f2c8280f..7cd031b0 100644 --- a/hub/src/sync/rpcGateway.ts +++ b/hub/src/sync/rpcGateway.ts @@ -140,6 +140,13 @@ export class RpcGateway { await this.sessionRpc(sessionId, RPC_METHODS.KillSession, {}) } + async stopRunnerSession(machineId: string, sessionId: string): Promise<'stopped' | 'already_gone' | 'still_alive'> { + const result = await this.machineRpc(machineId, RPC_METHODS.StopSession, { sessionId }) + const status = result && typeof result === 'object' ? (result as { status?: unknown }).status : undefined + if (status === 'stopped' || status === 'already_gone' || status === 'still_alive') return status + throw new Error('Unexpected stop-session response') + } + async handoffSessionToLocal(sessionId: string): Promise { await this.sessionRpc(sessionId, RPC_METHODS.HandoffLocal, {}) } diff --git a/hub/src/sync/sessionModel.test.ts b/hub/src/sync/sessionModel.test.ts index c6f9b71a..c3cf660b 100644 --- a/hub/src/sync/sessionModel.test.ts +++ b/hub/src/sync/sessionModel.test.ts @@ -16,6 +16,11 @@ function createPublisher(events: SyncEvent[]): EventPublisher { } as unknown as EventPublisher } +async function flushAsyncWork(): Promise { + await new Promise((resolve) => setTimeout(resolve, 0)) + await new Promise((resolve) => setTimeout(resolve, 0)) +} + function productionCodexMessage(event: Record): Record { return { role: 'agent', @@ -2003,6 +2008,375 @@ describe('session model', () => { } }) + it('reopens Pi in place only after native-ready, preserving its id and history', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const session = engine.getOrCreateSession('pi-in-place', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-1', + lifecycleState: 'archived', archivedBy: 'cli', archiveReason: 'Pi exited', + }, null, 'default') + store.messages.addMessage(session.id, { role: 'user', content: { type: 'text', text: 'keep history' } }) + engine.getOrCreateMachine('machine-1', { host: 'localhost', platform: 'linux', happyCliVersion: '0.1.0' }, { status: 'running', capabilities: { piExistingSessionResume: true } }, 'default') + engine.handleMachineAlive({ machineId: 'machine-1', time: Date.now() }) + engine.handleSessionEnd({ sid: session.id, time: Date.now() }) + let existing: string | undefined + let merges = 0 + ;(engine as any).sessionCache.mergeSessions = async () => { merges += 1 } + ;(engine as any).rpcGateway.spawnSession = async (...args: Parameters) => { + existing = args[12] + engine.handleSessionAlive({ sid: session.id, time: Date.now() }) + engine.handleSessionReady({ sid: session.id, time: Date.now() }) + return { type: 'success', sessionId: session.id } + } + const result = await engine.reopenSession(session.id, 'default') + expect(result).toEqual({ type: 'success', sessionId: session.id, resumed: true }) + expect(existing).toBe(session.id) + expect(merges).toBe(0) + expect(store.messages.getFirstMessages(session.id, 10)).toHaveLength(1) + } finally { engine.stop() } + }) + + it('kills and deletes an unexpected legacy Pi temp without touching the original row', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const original = engine.getOrCreateSession('pi-original', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-old', + lifecycleState: 'archived', archivedBy: 'cli', archiveReason: 'Pi exited', + }, null, 'default') + const unexpected = engine.getOrCreateSession('pi-unexpected', { path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi' }, null, 'default') + engine.getOrCreateMachine('machine-1', { host: 'localhost', platform: 'linux', happyCliVersion: '0.1.0' }, { status: 'running', capabilities: { piExistingSessionResume: true } }, 'default') + engine.handleMachineAlive({ machineId: 'machine-1', time: Date.now() }) + engine.handleSessionEnd({ sid: original.id, time: Date.now() }) + ;(engine as any).rpcGateway.spawnSession = async () => ({ type: 'success', sessionId: unexpected.id }) + ;(engine as any).rpcGateway.stopRunnerSession = async (_machineId: string, sid: string) => { + engine.handleSessionEnd({ sid, time: Date.now(), reason: 'error' }) + return 'stopped' + } + const result = await engine.reopenSession(original.id, 'default') + expect(result).toMatchObject({ type: 'error', code: 'resume_failed', message: expect.stringContaining('upgrade') }) + expect(store.sessions.getSession(unexpected.id)).toBeNull() + expect(store.sessions.getSession(original.id)).not.toBeNull() + expect(engine.getSessionByNamespace(original.id, 'default')?.metadata?.lifecycleState).toBe('archived') + } finally { engine.stop() } + }) + + it('does not restore archive metadata over a live Pi child after kill failure', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const session = engine.getOrCreateSession('pi-live-failure', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-live', + lifecycleState: 'archived', archivedBy: 'cli', archiveReason: 'Pi exited', + }, null, 'default') + engine.getOrCreateMachine('machine-1', { host: 'localhost', platform: 'linux', happyCliVersion: '0.1.0' }, { status: 'running', capabilities: { piExistingSessionResume: true } }, 'default') + engine.handleMachineAlive({ machineId: 'machine-1', time: Date.now() }) + engine.handleSessionEnd({ sid: session.id, time: Date.now() }) + ;(engine as any).rpcGateway.spawnSession = async () => { + engine.handleSessionAlive({ sid: session.id, time: Date.now() }) + return { type: 'success', sessionId: session.id } + } + ;(engine as any).waitForSessionReady = async () => 'timeout' + ;(engine as any).rpcGateway.stopRunnerSession = async () => 'still_alive' + const result = await engine.reopenSession(session.id, 'default') + expect(result).toMatchObject({ type: 'error', message: expect.stringContaining('still active') }) + expect(engine.getSessionByNamespace(session.id, 'default')?.active).toBe(true) + // Pi keeps the persisted archive snapshot until bootstrap succeeds, + // so a failed stop never needs to reconstruct it from memory. + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.lifecycleState).toBe('archived') + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.piResumeAttempt?.state).toBe('quarantined') + expect(await engine.reopenSession(session.id, 'default')).toMatchObject({ type: 'error', message: 'Pi resume is already in progress' }) + + engine.handleSessionEnd({ sid: session.id, time: Date.now(), reason: 'error' }) + await flushAsyncWork() + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.piResumeAttempt).toBeUndefined() + } finally { engine.stop() } + }) + + it('rejects Pi resume before spawn when the runner lacks in-place capability', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const session = engine.getOrCreateSession('pi-old-runner', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-old-runner', + lifecycleState: 'archived', archivedBy: 'cli', archiveReason: 'Pi exited', + }, null, 'default') + engine.getOrCreateMachine('machine-1', { host: 'localhost', platform: 'linux', happyCliVersion: '0.1.0' }, { status: 'running' }, 'default') + engine.handleMachineAlive({ machineId: 'machine-1', time: Date.now() }) + engine.handleSessionEnd({ sid: session.id, time: Date.now() }) + let spawnCalls = 0 + ;(engine as any).rpcGateway.spawnSession = async () => { spawnCalls += 1; return { type: 'error', message: 'unexpected' } } + + const result = await engine.reopenSession(session.id, 'default') + expect(result).toMatchObject({ type: 'error', message: 'Pi resume requires an upgraded runner' }) + expect(spawnCalls).toBe(0) + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.lifecycleState).toBe('archived') + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.piResumeAttempt).toBeUndefined() + } finally { engine.stop() } + }) + + it('quarantines a Pi attempt when runner spawn fails before process termination is confirmed', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const session = engine.getOrCreateSession('pi-spawn-error-live-child', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-spawn-error', + lifecycleState: 'archived', archivedBy: 'cli', archiveReason: 'Pi exited', + }, null, 'default') + engine.getOrCreateMachine('machine-1', { host: 'localhost', platform: 'linux', happyCliVersion: '0.1.0' }, { + status: 'running', capabilities: { piExistingSessionResume: true } + }, 'default') + engine.handleMachineAlive({ machineId: 'machine-1', time: Date.now() }) + engine.handleSessionEnd({ sid: session.id, time: Date.now() }) + ;(engine as any).rpcGateway.spawnSession = async () => ({ type: 'error', message: 'webhook timeout' }) + ;(engine as any).rpcGateway.stopRunnerSession = async () => 'still_alive' + + expect(await engine.reopenSession(session.id, 'default')).toMatchObject({ + type: 'error', message: 'webhook timeout' + }) + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.piResumeAttempt?.state).toBe('quarantined') + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.lifecycleState).toBe('archived') + expect(await engine.reopenSession(session.id, 'default')).toMatchObject({ + type: 'error', message: 'Pi resume is already in progress' + }) + } finally { engine.stop() } + }) + + it('keeps persisted Pi quarantine across SyncEngine restart and clears it on end', async () => { + const store = new Store(':memory:') + const first = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + const persisted = first.getOrCreateSession('pi-persisted-attempt', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-persisted', + lifecycleState: 'archived', archivedBy: 'cli', archiveReason: 'Pi exited', + piResumeAttempt: { state: 'quarantined', machineId: 'machine-1', startedAt: 1 }, + }, null, 'default') + first.stop() + + const restarted = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + expect(await restarted.reopenSession(persisted.id, 'default')).toMatchObject({ + type: 'error', message: 'Pi resume is already in progress' + }) + restarted.handleSessionEnd({ sid: persisted.id, time: Date.now(), reason: 'error' }) + await flushAsyncWork() + expect(restarted.getSessionByNamespace(persisted.id, 'default')?.metadata?.piResumeAttempt).toBeUndefined() + } finally { restarted.stop() } + }) + + it('clears a persisted Pi attempt when native-ready arrives after Hub restart', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const session = engine.getOrCreateSession('pi-ready-after-restart', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-ready', + piResumeAttempt: { state: 'resuming', machineId: 'machine-1', startedAt: 1 }, + }, null, 'default') + engine.handleSessionReady({ sid: session.id, time: Date.now() }) + await flushAsyncWork() + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.piResumeAttempt).toBeUndefined() + } finally { engine.stop() } + }) + + it('clears same-process Pi quarantine when a late validated ready arrives', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const session = engine.getOrCreateSession('pi-late-ready', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-late-ready', + piResumeAttempt: { state: 'quarantined', machineId: 'machine-1', startedAt: 1 }, + }, null, 'default') + ;(engine as any).piResumeQuarantinedIds.add(session.id) + engine.handleSessionReady({ sid: session.id, time: Date.now() }) + await flushAsyncWork() + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.piResumeAttempt).toBeUndefined() + expect((engine as any).piResumeQuarantinedIds.has(session.id)).toBe(false) + } finally { engine.stop() } + }) + + it('does not report active Pi attempts as reopened before validated ready', async () => { + for (const state of ['resuming', 'terminating'] as const) { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const session = engine.getOrCreateSession(`pi-active-${state}`, { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: `pi-native-${state}`, + piResumeAttempt: { state, machineId: 'machine-1', startedAt: 1 }, + }, null, 'default') + engine.handleSessionAlive({ sid: session.id, time: Date.now() }) + expect(await engine.reopenSession(session.id, 'default')).toMatchObject({ + type: 'error', message: 'Pi resume is already in progress' + }) + } finally { engine.stop() } + } + }) + + it('reconciles persisted quarantined Pi state when the runner reports the child already gone', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const session = engine.getOrCreateSession('pi-stale-quarantine', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-stale-quarantine', + lifecycleState: 'archived', archivedBy: 'cli', archiveReason: 'Pi exited', + piResumeAttempt: { state: 'quarantined', machineId: 'machine-1', startedAt: 1 }, + }, null, 'default') + ;(engine as any).rpcGateway.stopRunnerSession = async () => 'already_gone' + expect(await engine.reopenSession(session.id, 'default')).toMatchObject({ + type: 'error', message: 'Previous Pi resume attempt was cleaned up; retry' + }) + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.piResumeAttempt).toBeUndefined() + } finally { engine.stop() } + }) + + it('reconciles a persisted Pi attempt with an already-gone runner child without clearing archive state', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const session = engine.getOrCreateSession('pi-stale-resume-attempt', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-stale', + // Mirrors bootstrapExistingSession before native get_state: the + // live metadata says running, while the attempt carries the + // exact archived snapshot needed if the child is already gone. + lifecycleState: 'running', lifecycleStateSince: 200, + piResumeAttempt: { + state: 'resuming', machineId: 'machine-1', startedAt: 1, + archiveSnapshot: { + lifecycleState: 'archived', lifecycleStateSince: 100, + archivedBy: 'cli', archiveReason: 'Pi exited', + }, + }, + }, null, 'default') + ;(engine as any).rpcGateway.stopRunnerSession = async () => 'already_gone' + + expect(await engine.reopenSession(session.id, 'default')).toMatchObject({ + type: 'error', message: 'Previous Pi resume attempt was cleaned up; retry' + }) + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.piResumeAttempt).toBeUndefined() + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.lifecycleState).toBe('archived') + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.lifecycleStateSince).toBe(100) + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.archivedBy).toBe('cli') + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.archiveReason).toBe('Pi exited') + } finally { engine.stop() } + }) + + it('does not quarantine an in-place Pi row when session-end wins the stop response race', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const session = engine.getOrCreateSession('pi-stop-end-race', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-race', + piResumeAttempt: { state: 'resuming', machineId: 'machine-1', startedAt: 1 }, + }, null, 'default') + engine.handleSessionAlive({ sid: session.id, time: Date.now() }) + ;(engine as any).rpcGateway.stopRunnerSession = async () => { + engine.handleSessionEnd({ sid: session.id, time: Date.now(), reason: 'error' }) + return 'still_alive' + } + + expect(await (engine as any).terminateInPlacePiResume('machine-1', session.id, 'default', true)).toBe(true) + await flushAsyncWork() + expect(engine.getSessionByNamespace(session.id, 'default')?.active).toBe(false) + expect(engine.getSessionByNamespace(session.id, 'default')?.metadata?.piResumeAttempt).toBeUndefined() + } finally { engine.stop() } + }) + + it('does not persist an unexpected-child quarantine when child end wins the stop response race', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const original = engine.getOrCreateSession('pi-original-stop-race', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-race', + lifecycleState: 'archived', archivedBy: 'cli', archiveReason: 'Pi exited', + piResumeAttempt: { state: 'resuming', machineId: 'machine-1', startedAt: 1 }, + }, null, 'default') + const temp = engine.getOrCreateSession('pi-temp-stop-race', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', + }, null, 'default') + engine.handleSessionAlive({ sid: temp.id, time: Date.now() }) + ;(engine as any).rpcGateway.stopRunnerSession = async () => { + engine.handleSessionEnd({ sid: temp.id, time: Date.now(), reason: 'error' }) + return 'still_alive' + } + + expect(await (engine as any).terminateUnexpectedPiTemp('machine-1', temp.id, original.id, 'default')).toBe(true) + await flushAsyncWork() + expect(store.sessions.getSession(temp.id)).toBeNull() + expect(engine.getSessionByNamespace(original.id, 'default')?.metadata?.piResumeAttempt).toBeUndefined() + } finally { engine.stop() } + }) + + it('keeps still-alive Pi children quarantined even when Hub cache is inactive', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const original = engine.getOrCreateSession('pi-still-alive-inactive', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-still-alive', + piResumeAttempt: { state: 'resuming', machineId: 'machine-1', startedAt: 1 }, + }, null, 'default') + ;(engine as any).rpcGateway.stopRunnerSession = async () => 'still_alive' + expect(await (engine as any).terminateInPlacePiResume('machine-1', original.id, 'default')).toBe(false) + + const temp = engine.getOrCreateSession('pi-temp-still-alive-inactive', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', + }, null, 'default') + expect(await (engine as any).terminateUnexpectedPiTemp('machine-1', temp.id, original.id, 'default')).toBe(false) + expect(store.sessions.getSession(temp.id)).not.toBeNull() + expect(engine.getSessionByNamespace(original.id, 'default')?.metadata?.piResumeAttempt).toMatchObject({ + state: 'quarantined', childSessionId: temp.id + }) + } finally { engine.stop() } + }) + + it('blocks dedup for a persisted unexpected Pi child and clears the original attempt on child end', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const temp = engine.getOrCreateSession('pi-temp-mapped', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-mapped', + }, null, 'default') + const original = engine.getOrCreateSession('pi-original-mapped', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-mapped', + lifecycleState: 'archived', archivedBy: 'cli', archiveReason: 'Pi exited', + piResumeAttempt: { state: 'quarantined', machineId: 'machine-1', startedAt: 1, childSessionId: temp.id }, + }, null, 'default') + let dedupCalls = 0 + ;(engine as any).sessionCache.deduplicateByAgentSessionId = async () => { dedupCalls += 1 } + ;(engine as any).triggerDedupIfNeeded(temp.id) + await flushAsyncWork() + expect(dedupCalls).toBe(0) + expect(store.sessions.getSession(original.id)).not.toBeNull() + + engine.handleSessionEnd({ sid: temp.id, time: Date.now(), reason: 'error' }) + await flushAsyncWork() + expect(engine.getSessionByNamespace(original.id, 'default')?.metadata?.piResumeAttempt).toBeUndefined() + expect(store.sessions.getSession(original.id)).not.toBeNull() + expect(dedupCalls).toBe(0) + } finally { engine.stop() } + }) + + it('blocks Pi dedup before an unexpected child ID is attached to the persisted attempt', async () => { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + try { + const original = engine.getOrCreateSession('pi-original-pre-mapping', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-pre-mapping', + lifecycleState: 'archived', archivedBy: 'cli', archiveReason: 'Pi exited', + piResumeAttempt: { state: 'resuming', machineId: 'machine-1', startedAt: 1 }, + }, null, 'default') + const temp = engine.getOrCreateSession('pi-temp-pre-mapping', { + path: '/tmp/project', host: 'localhost', machineId: 'machine-1', flavor: 'pi', piSessionId: 'pi-native-pre-mapping', + }, null, 'default') + engine.handleSessionAlive({ sid: temp.id, time: Date.now() }) + let dedupCalls = 0 + ;(engine as any).sessionCache.deduplicateByAgentSessionId = async () => { dedupCalls += 1 } + + ;(engine as any).triggerDedupIfNeeded(temp.id) + await flushAsyncWork() + expect(dedupCalls).toBe(0) + expect(store.sessions.getSession(original.id)).not.toBeNull() + } finally { engine.stop() } + }) + it('defers mergeSessions for cursor reopen until session-ready (load failure leaves old row)', async () => { const store = new Store(':memory:') const engine = new SyncEngine( diff --git a/hub/src/sync/syncEngine.ts b/hub/src/sync/syncEngine.ts index ee4c65a0..adf4cd17 100644 --- a/hub/src/sync/syncEngine.ts +++ b/hub/src/sync/syncEngine.ts @@ -46,6 +46,8 @@ import { } from './rpcGateway' import { SessionCache } from './sessionCache' +type PiResumeAttempt = NonNullable['piResumeAttempt']> + export type { Session, SyncEvent } from '@hapi/protocol/types' export type { Machine } from './machineCache' export type { SyncEventListener } from './eventPublisher' @@ -72,7 +74,7 @@ export type { export type ResumeSessionResult = | { type: 'success'; sessionId: string } - | { type: 'error'; message: string; code: 'session_not_found' | 'access_denied' | 'no_machine_online' | 'resume_unavailable' | 'resume_failed' } + | { type: 'error'; message: string; code: 'session_not_found' | 'access_denied' | 'no_machine_online' | 'resume_unavailable' | 'resume_failed'; rollbackSafe?: boolean } export type ReopenSessionResult = | { type: 'success'; sessionId: string; resumed: boolean; cursorSessionProtocol?: 'acp' | 'stream-json' } @@ -147,8 +149,14 @@ export class SyncEngine { private readonly messageService: MessageService private readonly rpcGateway: RpcGateway private inactivityTimer: NodeJS.Timeout | null = null - /** Sessions that emitted `session-ready` (Cursor ACP load/newSession complete). */ + /** Sessions that emitted `session-ready` (Cursor ACP or validated Pi get_state). */ private readonly sessionReadyIds = new Set() + /** Original Pi rows with a native resume currently in flight. */ + private readonly piResumeInFlightIds = new Set() + /** Pi rows whose runner child could not be confirmed terminated. */ + private readonly piResumeQuarantinedIds = new Set() + /** Unexpected version-skew temp child -> original row whose retry is blocked until child ends. */ + private readonly piUnexpectedTempOriginalIds = new Map() /** Serialize scratchlist uploads per session so disk-byte caps cannot race. */ private readonly scratchlistUploadTails = new Map>() @@ -402,6 +410,15 @@ export class SyncEngine { handleSessionReady(payload: { sid: string; time: number }): void { this.sessionReadyIds.add(payload.sid) + const session = this.sessionCache.getSession(payload.sid) + if (session?.metadata?.piResumeAttempt) { + void this.writePiResumeAttempt(payload.sid, session.namespace, null) + .then(() => { + this.piResumeQuarantinedIds.delete(payload.sid) + this.triggerDedupIfNeeded(payload.sid) + }) + .catch(() => {}) + } this.triggerDedupIfNeeded(payload.sid) } @@ -411,9 +428,14 @@ export class SyncEngine { handleSessionEnd(payload: { sid: string; time: number; reason?: 'completed' | 'terminated' | 'error' }): void { const before = this.sessionCache.getSession(payload.sid) + const ownsPiAttempt = before?.metadata?.piResumeAttempt !== undefined + const isPiAttemptChild = this.sessionCache.getSessions().some( + (session) => session.metadata?.piResumeAttempt?.childSessionId === payload.sid + ) + const restorePiArchive = ownsPiAttempt && !this.sessionReadyIds.has(payload.sid) const isCursorAcp = before?.metadata?.flavor === 'cursor' && before.metadata.cursorSessionProtocol === 'acp' - const shouldRetryDedup = !isCursorAcp || this.sessionReadyIds.has(payload.sid) + const shouldRetryDedup = !ownsPiAttempt && !isPiAttemptChild && (!isCursorAcp || this.sessionReadyIds.has(payload.sid)) this.sessionCache.handleSessionEnd(payload) this.eventPublisher.emit({ @@ -428,6 +450,11 @@ export class SyncEngine { this.triggerDedupIfNeeded(payload.sid) } this.sessionReadyIds.delete(payload.sid) + this.piResumeQuarantinedIds.delete(payload.sid) + this.piUnexpectedTempOriginalIds.delete(payload.sid) + if (ownsPiAttempt || isPiAttemptChild) { + void this.clearPiAttemptForEndedSession(payload.sid, restorePiArchive) + } } handleBackgroundTaskDelta(sessionId: string, delta: { started: number; completed: number }): void { @@ -1539,6 +1566,35 @@ async uploadScratchlistAttachment( return { type: 'error', message: 'No machine online', code: 'no_machine_online' } } + if (flavor === 'pi' && resumeToken && targetMachine.runnerState?.capabilities?.piExistingSessionResume !== true) { + return { type: 'error', message: 'Pi resume requires an upgraded runner', code: 'resume_failed' } + } + + const requiresPiNativeReady = flavor === 'pi' && resumeToken !== undefined + if (requiresPiNativeReady) { + if (this.isPiResumeBlocked(access.sessionId)) { + return { type: 'error', message: 'Pi resume is already in progress', code: 'resume_failed' } + } + this.piResumeInFlightIds.add(access.sessionId) + this.sessionReadyIds.delete(access.sessionId) + try { + await this.writePiResumeAttempt(access.sessionId, namespace, { + state: 'resuming', + machineId: targetMachine.id, + startedAt: Date.now(), + archiveSnapshot: { + lifecycleState: metadata.lifecycleState, + lifecycleStateSince: metadata.lifecycleStateSince, + archivedBy: metadata.archivedBy, + archiveReason: metadata.archiveReason, + }, + }) + } catch { + this.piResumeInFlightIds.delete(access.sessionId) + return { type: 'error', message: 'Failed to record Pi resume attempt', code: 'resume_failed' } + } + } + if (flavor === 'cursor' && resumeToken) { try { const chatStatus = await this.rpcGateway.getCursorChatStoreStatus( @@ -1569,65 +1625,135 @@ async uploadScratchlistAttachment( : opts?.permissionMode ?? session.permissionMode ?? metadataPermissionMode - const spawnResult = await this.rpcGateway.spawnSession( - targetMachine.id, - directory, - flavor, - session.model ?? undefined, - session.modelReasoningEffort ?? undefined, - undefined, - undefined, - undefined, - resumeToken, - session.effort ?? undefined, - preferredPermissionMode, - session.serviceTier ?? undefined, - access.sessionId, - session.collaborationMode ?? undefined - ) + let piResumeSucceeded = false + try { + const spawnResult = await this.rpcGateway.spawnSession( + targetMachine.id, + directory, + flavor, + session.model ?? undefined, + session.modelReasoningEffort ?? undefined, + undefined, + undefined, + undefined, + resumeToken, + session.effort ?? undefined, + preferredPermissionMode, + session.serviceTier ?? undefined, + access.sessionId, + session.collaborationMode ?? undefined + ) - if (spawnResult.type !== 'success') { - return { type: 'error', message: spawnResult.message, code: 'resume_failed' } - } - - const becameActive = await this.waitForSessionActive(spawnResult.sessionId) - if (!becameActive) { - return { type: 'error', message: 'Session failed to become active', code: 'resume_failed' } - } - - // permissionMode is passed to spawnSession above; do not call set-session-config here. - // session-alive can arrive before the CLI registers that RPC handler, which caused resume_failed. - - const needsReadyBeforeMerge = spawnResult.sessionId !== access.sessionId - && flavor === 'cursor' - && metadata.cursorSessionProtocol === 'acp' - if (needsReadyBeforeMerge) { - const readyResult = await this.waitForSessionReady(spawnResult.sessionId) - if (readyResult !== 'ready') { - const message = readyResult === 'ended' - ? 'Session ended before Cursor ACP load completed' - : 'Session failed to become ready' - return { type: 'error', message, code: 'resume_failed' } + if (spawnResult.type !== 'success') { + if (requiresPiNativeReady) { + const stopped = await this.terminateInPlacePiResume( + targetMachine.id, + access.sessionId, + namespace + ) + if (!stopped) { + await this.quarantinePiResume(access.sessionId, namespace, targetMachine.id) + return { + type: 'error', + message: spawnResult.message, + code: 'resume_failed', + rollbackSafe: false, + } + } + } + return { type: 'error', message: spawnResult.message, code: 'resume_failed' } } - } - if (spawnResult.sessionId !== access.sessionId) { - // The old session may have already been merged by the automatic dedup path - // (triggered when the spawned CLI sets its agent session ID in metadata). - // Only attempt the explicit merge if the old session still exists. - const oldSession = this.sessionCache.getSessionByNamespace(access.sessionId, namespace) - if (oldSession) { - try { - await this.sessionCache.mergeSessions(access.sessionId, spawnResult.sessionId, namespace) - } catch (error) { - const message = error instanceof Error ? error.message : 'Failed to merge resumed session' + if (requiresPiNativeReady && spawnResult.sessionId !== access.sessionId) { + const removed = await this.terminateUnexpectedPiTemp( + targetMachine.id, + spawnResult.sessionId, + access.sessionId, + namespace + ) + return { + type: 'error', + message: removed + ? 'Pi runner created an unexpected session; upgrade the runner and retry' + : 'Pi runner created an unexpected live session; upgrade the runner and retry', + code: 'resume_failed' + } + } + + const becameActive = await this.waitForSessionActive(spawnResult.sessionId) + if (!becameActive) { + if (requiresPiNativeReady) { + const inactive = await this.terminateInPlacePiResume( + targetMachine.id, + access.sessionId, + namespace + ) + if (!inactive) { + await this.quarantinePiResume(access.sessionId, namespace, targetMachine.id) + return { type: 'error', message: 'Pi resume failed and the child is still active', code: 'resume_failed', rollbackSafe: false } + } + } + return { type: 'error', message: 'Session failed to become active', code: 'resume_failed' } + } + + const needsReadyBeforeSuccess = requiresPiNativeReady + || ( + spawnResult.sessionId !== access.sessionId + && flavor === 'cursor' + && metadata.cursorSessionProtocol === 'acp' + ) + if (needsReadyBeforeSuccess) { + const readyResult = await this.waitForSessionReady(spawnResult.sessionId) + if (readyResult !== 'ready') { + if (requiresPiNativeReady && readyResult !== 'ended') { + const inactive = await this.terminateInPlacePiResume( + targetMachine.id, + access.sessionId, + namespace + ) + if (!inactive) { + await this.quarantinePiResume(access.sessionId, namespace, targetMachine.id) + return { type: 'error', message: 'Pi native resume timed out and the child is still active', code: 'resume_failed', rollbackSafe: false } + } + } + const message = flavor === 'pi' + ? readyResult === 'ended' + ? 'Pi session ended before native resume completed' + : 'Pi session failed to become native-ready' + : readyResult === 'ended' + ? 'Session ended before Cursor ACP load completed' + : 'Session failed to become ready' return { type: 'error', message, code: 'resume_failed' } } } - } - this.sessionCache.markSessionActive(spawnResult.sessionId) - return { type: 'success', sessionId: spawnResult.sessionId } + if (spawnResult.sessionId !== access.sessionId) { + const oldSession = this.sessionCache.getSessionByNamespace(access.sessionId, namespace) + if (oldSession) { + try { + await this.sessionCache.mergeSessions(access.sessionId, spawnResult.sessionId, namespace) + } catch (error) { + const message = error instanceof Error ? error.message : 'Failed to merge resumed session' + return { type: 'error', message, code: 'resume_failed' } + } + } + } + + this.sessionCache.markSessionActive(spawnResult.sessionId) + piResumeSucceeded = true + if (requiresPiNativeReady) await this.writePiResumeAttempt(access.sessionId, namespace, null) + return { type: 'success', sessionId: spawnResult.sessionId } + } finally { + if (requiresPiNativeReady) { + this.piResumeInFlightIds.delete(access.sessionId) + if (!piResumeSucceeded && this.sessionCache.getSession(access.sessionId)?.metadata?.piResumeAttempt?.state === 'resuming') { + await this.writePiResumeAttempt(access.sessionId, namespace, null, true).catch(() => {}) + } + if (piResumeSucceeded) { + this.triggerDedupIfNeeded(access.sessionId) + } + } + } } /** @@ -1664,6 +1790,23 @@ async uploadScratchlistAttachment( const session = access.session const metadata = session.metadata + if (metadata?.flavor === 'pi' && this.isPiResumeBlocked(access.sessionId)) { + if (session.active) { + return { type: 'error', message: 'Pi resume is already in progress', code: 'resume_failed' } + } + if (metadata.piResumeAttempt && !this.piResumeInFlightIds.has(access.sessionId)) { + const reconciled = await this.reconcilePersistedPiResumeAttempt(session) + return { + type: 'error', + message: reconciled + ? 'Previous Pi resume attempt was cleaned up; retry' + : 'Pi resume is already in progress', + code: 'resume_failed' + } + } + return { type: 'error', message: 'Pi resume is already in progress', code: 'resume_failed' } + } + if (session.active) { return { type: 'success', sessionId: access.sessionId, resumed: false } } @@ -1689,24 +1832,30 @@ async uploadScratchlistAttachment( lifecycleStateSince: metadata.lifecycleStateSince } - let applied: { cursorSessionProtocol?: 'acp' | 'stream-json' } - try { - applied = await this.sessionCache.clearSessionArchiveMetadata(access.sessionId) - } catch (error) { - const message = error instanceof Error ? error.message : 'Failed to clear archive metadata' - return { type: 'error', message, code: 'metadata_conflict' } + let applied: { cursorSessionProtocol?: 'acp' | 'stream-json' } = {} + // Pi reuses the original HAPI row. Keep its archive snapshot persisted + // until the CLI successfully bootstraps that row as running; this avoids + // an inactive, non-archived gap if the Hub restarts before spawn. + if (metadata.flavor !== 'pi') { + try { + applied = await this.sessionCache.clearSessionArchiveMetadata(access.sessionId) + } catch (error) { + const message = error instanceof Error ? error.message : 'Failed to clear archive metadata' + return { type: 'error', message, code: 'metadata_conflict' } + } } const resumeResult = await this.resumeSession(access.sessionId, namespace) if (resumeResult.type === 'error') { - // Resume failed - put the archive flags back so the row stays archived in the UI - // and the operator can retry. Best-effort: a concurrent metadata write that - // succeeded between clear and restore (e.g. an unrelated rename) wins, in - // which case we surface the original resume error rather than masking it. - try { - await this.sessionCache.restoreSessionArchiveMetadata(access.sessionId, archiveSnapshot) - } catch { - // Swallow restore failures - the resume error is the more important signal. + // Never restore archived metadata over a live Pi child. A live + // row blocks retry by itself and must remain visible as active. + const current = this.sessionCache.getSessionByNamespace(access.sessionId, namespace) + if (resumeResult.rollbackSafe !== false && !current?.active) { + try { + await this.sessionCache.restoreSessionArchiveMetadata(access.sessionId, archiveSnapshot) + } catch { + // Swallow restore failures - the resume error is the more important signal. + } } return resumeResult } @@ -1928,6 +2077,16 @@ async uploadScratchlistAttachment( } private canRunCursorDedup(session: Session): boolean { + if (this.piResumeInFlightIds.has(session.id) || this.piResumeQuarantinedIds.has(session.id)) return false + if (session.metadata?.piResumeAttempt) return false + if (this.sessionCache.getSessions().some((candidate) => candidate.metadata?.piResumeAttempt?.childSessionId === session.id)) return false + const piSessionId = session.metadata?.piSessionId + if (piSessionId && this.sessionCache.getSessions().some((candidate) => + candidate.id !== session.id + && candidate.namespace === session.namespace + && candidate.metadata?.piResumeAttempt !== undefined + && candidate.metadata.piSessionId === piSessionId + )) return false if (session.metadata?.flavor !== 'cursor') { return true } @@ -1937,6 +2096,173 @@ async uploadScratchlistAttachment( return this.sessionReadyIds.has(session.id) } + private async terminateInPlacePiResume( + machineId: string, + sessionId: string, + namespace: string + ): Promise { + const existingAttempt = this.sessionCache.getSession(sessionId)?.metadata?.piResumeAttempt + await this.writePiResumeAttempt(sessionId, namespace, { + ...existingAttempt, + state: 'terminating', + machineId, + startedAt: Date.now(), + }) + let status: 'stopped' | 'already_gone' | 'still_alive' + try { + status = await this.rpcGateway.stopRunnerSession(machineId, sessionId) + } catch { + status = 'still_alive' + } + + await new Promise((resolve) => setTimeout(resolve, 0)) + const session = this.sessionCache.refreshSession(sessionId) ?? this.sessionCache.getSession(sessionId) + const attemptClearedByEnd = session?.metadata?.piResumeAttempt === undefined + if (status === 'still_alive') { + if (attemptClearedByEnd) return true + return false + } + if (session?.active) this.handleSessionEnd({ sid: sessionId, time: Date.now(), reason: 'error' }) + await this.writePiResumeAttempt(sessionId, namespace, null, true).catch(() => {}) + return true + } + + private async terminateUnexpectedPiTemp( + machineId: string, + sessionId: string, + originalSessionId: string, + namespace: string + ): Promise { + this.piUnexpectedTempOriginalIds.set(sessionId, originalSessionId) + const existingAttempt = this.sessionCache.getSession(originalSessionId)?.metadata?.piResumeAttempt + await this.writePiResumeAttempt(originalSessionId, namespace, { + ...existingAttempt, + state: 'terminating', + machineId, + startedAt: Date.now(), + childSessionId: sessionId, + }) + let status: 'stopped' | 'already_gone' | 'still_alive' + try { + status = await this.rpcGateway.stopRunnerSession(machineId, sessionId) + } catch { + status = 'still_alive' + } + + await new Promise((resolve) => setTimeout(resolve, 0)) + const session = this.sessionCache.refreshSession(sessionId) ?? this.sessionCache.getSession(sessionId) + const original = this.sessionCache.refreshSession(originalSessionId) ?? this.sessionCache.getSession(originalSessionId) + const attemptClearedByEnd = original?.metadata?.piResumeAttempt === undefined + if (status === 'still_alive' && !attemptClearedByEnd) { + await this.writePiResumeAttempt(originalSessionId, namespace, { + ...existingAttempt, + state: 'quarantined', + machineId, + startedAt: Date.now(), + childSessionId: sessionId, + }) + return false + } + if (session?.active) this.handleSessionEnd({ sid: sessionId, time: Date.now(), reason: 'error' }) + const remaining = this.sessionCache.getSession(sessionId) + if (remaining && !remaining.active) await this.sessionCache.deleteSession(sessionId) + this.piUnexpectedTempOriginalIds.delete(sessionId) + await this.writePiResumeAttempt(originalSessionId, namespace, null, true).catch(() => {}) + return true + } + + private async quarantinePiResume(sessionId: string, namespace: string, machineId: string): Promise { + this.piResumeQuarantinedIds.add(sessionId) + const existingAttempt = this.sessionCache.getSession(sessionId)?.metadata?.piResumeAttempt + await this.writePiResumeAttempt(sessionId, namespace, { + ...existingAttempt, + state: 'quarantined', + machineId, + startedAt: Date.now(), + }) + } + + private isPiResumeBlocked(sessionId: string): boolean { + const metadataAttempt = this.sessionCache.getSession(sessionId)?.metadata?.piResumeAttempt + return this.piResumeInFlightIds.has(sessionId) + || this.piResumeQuarantinedIds.has(sessionId) + || metadataAttempt !== undefined + || [...this.piUnexpectedTempOriginalIds.values()].includes(sessionId) + } + + private async writePiResumeAttempt( + sessionId: string, + namespace: string, + attempt: PiResumeAttempt | null, + restoreArchive = false + ): Promise { + for (let i = 0; i < 5; i += 1) { + const current = this.sessionCache.getSessionByNamespace(sessionId, namespace) ?? this.sessionCache.refreshSession(sessionId) + if (!current?.metadata) return + const next = { ...current.metadata } + if (attempt) next.piResumeAttempt = attempt + else { + const snapshot = current.metadata.piResumeAttempt?.archiveSnapshot + delete next.piResumeAttempt + if (restoreArchive && snapshot) { + if (snapshot.lifecycleState === undefined) delete next.lifecycleState + else next.lifecycleState = snapshot.lifecycleState + if (snapshot.lifecycleStateSince === undefined) delete next.lifecycleStateSince + else next.lifecycleStateSince = snapshot.lifecycleStateSince + if (snapshot.archivedBy === undefined) delete next.archivedBy + else next.archivedBy = snapshot.archivedBy + if (snapshot.archiveReason === undefined) delete next.archiveReason + else next.archiveReason = snapshot.archiveReason + } + } + const result = this.store.sessions.updateSessionMetadata(sessionId, next, current.metadataVersion, namespace, { touchUpdatedAt: false }) + if (result.result === 'success') { + this.sessionCache.refreshSession(sessionId) + return + } + if (result.result !== 'version-mismatch') throw new Error('Failed to update Pi resume attempt') + this.sessionCache.refreshSession(sessionId) + } + throw new Error('Pi resume attempt metadata was modified concurrently') + } + + private async clearPiAttemptForEndedSession(endedSessionId: string, restoreArchive: boolean): Promise { + const ended = this.sessionCache.getSession(endedSessionId) + if (ended?.metadata?.piResumeAttempt) { + await this.writePiResumeAttempt(endedSessionId, ended.namespace, null, restoreArchive).catch(() => {}) + return + } + for (const session of this.sessionCache.getSessions()) { + if (session.metadata?.piResumeAttempt?.childSessionId === endedSessionId) { + await this.writePiResumeAttempt(session.id, session.namespace, null, true).catch(() => {}) + } + } + } + + private async reconcilePersistedPiResumeAttempt(session: Session): Promise { + const attempt = session.metadata?.piResumeAttempt + if (!attempt) return true + const childSessionId = attempt.childSessionId ?? session.id + let status: 'stopped' | 'already_gone' | 'still_alive' + try { + status = await this.rpcGateway.stopRunnerSession(attempt.machineId, childSessionId) + } catch { + return false + } + if (status === 'still_alive') return false + + const child = this.sessionCache.getSession(childSessionId) + if (child?.active) this.handleSessionEnd({ sid: childSessionId, time: Date.now(), reason: 'error' }) + if (childSessionId !== session.id) { + const remaining = this.sessionCache.getSession(childSessionId) + if (remaining && !remaining.active) await this.sessionCache.deleteSession(childSessionId) + } + await this.writePiResumeAttempt(session.id, session.namespace, null, true) + this.piResumeQuarantinedIds.delete(session.id) + this.piUnexpectedTempOriginalIds.delete(childSessionId) + return true + } + private triggerDedupIfNeeded(sessionId: string): void { const session = this.sessionCache.getSession(sessionId) if (session?.metadata) { diff --git a/shared/src/schemas.ts b/shared/src/schemas.ts index c7ff455f..a0214b30 100644 --- a/shared/src/schemas.ts +++ b/shared/src/schemas.ts @@ -52,6 +52,18 @@ export const MetadataSchema = z.object({ cursorMigrationState: z.enum(['in_progress', 'ambiguous']).optional(), kimiSessionId: z.string().optional(), piSessionId: z.string().optional(), + piResumeAttempt: z.object({ + state: z.enum(['resuming', 'terminating', 'quarantined']), + machineId: z.string(), + startedAt: z.number(), + childSessionId: z.string().optional(), + archiveSnapshot: z.object({ + lifecycleState: z.string().optional(), + lifecycleStateSince: z.number().optional(), + archivedBy: z.string().optional(), + archiveReason: z.string().optional(), + }).optional(), + }).optional(), tools: z.array(z.string()).optional(), slashCommands: z.array(z.string()).optional(), homeDir: z.string().optional(), @@ -304,6 +316,7 @@ export const RunnerStateSchema = z.object({ pid: z.number().optional(), httpPort: z.number().optional(), startedAt: z.number().optional(), + capabilities: z.object({ piExistingSessionResume: z.literal(true).optional() }).optional(), shutdownRequestedAt: z.number().optional(), shutdownSource: z.union([z.enum(['mobile-app', 'cli', 'os-signal', 'unknown']), z.string()]).optional(), lastSpawnError: z.object({