From 35005fb65c9d4d72244a830ade1dd35e6ca446a5 Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Sun, 20 Sep 2026 21:38:38 -0700 Subject: [PATCH] fix(pi): keep panes working while async subagents run (#21882) * fix(pi): wait for async subagents before settling pane * fix(pi): handle subagent event aliases and reloads * test(pi): assert lifecycle listener cardinality --- ...ent-status-extension-omp-lifecycle.test.ts | 47 +++++++++++++++++++ .../pi/agent-status-extension-test-harness.ts | 9 ++++ src/main/pi/agent-status-handler-source.ts | 39 +++++++++++++-- 3 files changed, 90 insertions(+), 5 deletions(-) diff --git a/src/main/pi/agent-status-extension-omp-lifecycle.test.ts b/src/main/pi/agent-status-extension-omp-lifecycle.test.ts index 89f8c94d9a5..5d89d31aec1 100644 --- a/src/main/pi/agent-status-extension-omp-lifecycle.test.ts +++ b/src/main/pi/agent-status-extension-omp-lifecycle.test.ts @@ -15,6 +15,53 @@ const OMP_RUNTIME_CASES = [ ] as const describe('OMP agent_end contract', () => { + it('keeps a Pi pane working until async subagents finish', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'pi' }) + + await harness.callHook('agent_start') + harness.emitPiEvent('task:subagent:lifecycle', { id: 'child-1', status: 'started' }) + await harness.callHook('agent_settled', undefined, { isIdle: () => true }) + + expect(postedHookNames(harness.fetchMock)).toEqual(['agent_start']) + + harness.emitPiEvent('task:subagent:lifecycle', { id: 'child-1', status: 'completed' }) + await vi.waitFor(() => + expect(postedHookNames(harness.fetchMock)).toEqual(['agent_start', 'agent_end']) + ) + }) + + it('ignores malformed or unknown Pi subagent lifecycle events', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'pi' }) + harness.emitPiEvent('task:subagent:lifecycle', {}) + harness.emitPiEvent('task:subagent:lifecycle', { id: 'child-1', status: 'paused' }) + await harness.callHook('agent_start') + await harness.callHook('agent_settled') + await vi.waitFor(() => + expect(postedHookNames(harness.fetchMock)).toEqual(['agent_start', 'agent_end']) + ) + }) + + it('keeps one lifecycle subscription across extension reloads', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'pi' }) + harness.reload() + expect(harness.piEventListenerCount('task:subagent:lifecycle')).toBe(1) + expect(harness.piEventListenerCount('subagent:async-started')).toBe(1) + expect(harness.piEventListenerCount('subagent:async-complete')).toBe(1) + harness.emitPiEvent('task:subagent:lifecycle', { id: 'child-1', status: 'started' }) + await vi.waitFor(() => expect(postedHookNames(harness.fetchMock)).toEqual(['agent_start'])) + }) + + it('accepts the pi-subagents async lifecycle aliases', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'pi' }) + harness.emitPiEvent('subagent:async-started', { id: 'child-1' }) + await harness.callHook('agent_settled') + expect(postedHookNames(harness.fetchMock)).toEqual(['agent_start']) + harness.emitPiEvent('subagent:async-complete', { id: 'child-1' }) + await vi.waitFor(() => + expect(postedHookNames(harness.fetchMock)).toEqual(['agent_start', 'agent_end']) + ) + }) + it.each(OMP_RUNTIME_CASES)( 'keeps %s working when agent_end will continue', async (_name, args) => { diff --git a/src/main/pi/agent-status-extension-test-harness.ts b/src/main/pi/agent-status-extension-test-harness.ts index 696153064f1..0c07caf8256 100644 --- a/src/main/pi/agent-status-extension-test-harness.ts +++ b/src/main/pi/agent-status-extension-test-harness.ts @@ -46,6 +46,8 @@ export type AgentStatusExtensionHarness = { handlers: Record processEnv: Record callHook: (name: string, event?: unknown, context?: HookContext) => Promise + emitPiEvent: (name: string, event: unknown) => void + piEventListenerCount: (name: string) => number // Re-invoke the extension factory in the same process (as Pi does on an // in-process extension reload), swapping in the freshly registered handlers. reload: () => void @@ -130,6 +132,7 @@ export function createAgentStatusExtensionHarness(args: { command: { handler: (args: string, context: HookContext) => Promise } ) => void setModel: (model: unknown) => Promise + events?: EventEmitter }) => void } } = { exports: {} } @@ -191,6 +194,7 @@ export function createAgentStatusExtensionHarness(args: { } const handlers: Record = {} + const piEvents = new EventEmitter() const commands: AgentStatusExtensionHarness['commands'] = {} const setModelMock = vi.fn(async (_model: unknown) => true) const registerInto = (target: Record): void => { @@ -199,6 +203,7 @@ export function createAgentStatusExtensionHarness(args: { commands[name] = command }, setModel: setModelMock, + events: piEvents, on(name: string, handler: HookHandler) { target[name] = handler } @@ -219,6 +224,10 @@ export function createAgentStatusExtensionHarness(args: { callHook: async (name, event, hookContext) => { await handlers[name]?.(event, hookContext) }, + emitPiEvent: (name, event) => { + piEvents.emit(name, event) + }, + piEventListenerCount: (name) => piEvents.listenerCount(name), reload: () => { for (const key of Object.keys(handlers)) { delete handlers[key] diff --git a/src/main/pi/agent-status-handler-source.ts b/src/main/pi/agent-status-handler-source.ts index b3bd70e28a4..1d339e670ae 100644 --- a/src/main/pi/agent-status-handler-source.ts +++ b/src/main/pi/agent-status-handler-source.ts @@ -129,13 +129,27 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ' if (ownerPid && ownerPid !== selfPid && isStatusOwnerAlive(ownerPid)) return', ` process.env.${ownerEnv} = selfPid`, ' resetPostQueue()', + ' const piEventBus = (pi as { events?: { on?: (name: string, handler: (event: unknown) => void) => void } }).events', + ' const lifecycleState = (piEventBus as { __orcaPiSubagents?: { active: Set; waiting: boolean; onEvent?: (event: unknown, forcedStatus?: string) => void; listener?: (event: unknown) => void } } | undefined)?.__orcaPiSubagents ?? { active: new Set(), waiting: false }', + ' if (piEventBus) (piEventBus as { __orcaPiSubagents?: unknown }).__orcaPiSubagents = lifecycleState', + ' if (piEventBus?.on && !(lifecycleState as { listener?: unknown }).listener) {', + ' const listener = (event: unknown) => lifecycleState.onEvent?.(event)', + ' lifecycleState.listener = listener', + " piEventBus.on('task:subagent:lifecycle', listener)", + " piEventBus.on('subagent:async-started', (event: unknown) => lifecycleState.onEvent?.(event, 'started'))", + " piEventBus.on('subagent:async-complete', (event: unknown) => lifecycleState.onEvent?.(event, 'completed'))", + ' }', ...(kind !== 'pi' - ? [" pi.on('session_shutdown', () => { resetPostQueue(); clearPendingAgentEndCheck() })"] + ? [ + " pi.on('session_shutdown', () => { lifecycleState.active.clear(); lifecycleState.waiting = false; resetPostQueue(); clearPendingAgentEndCheck() })" + ] : []), ...(kind !== 'prime-agent' ? [ " pi.on('session_switch', (_event, ctx) => {", ' if (!isOmpRuntime()) return', + ' lifecycleState.active.clear()', + ' lifecycleState.waiting = false', ' resetPostQueue()', ' clearPendingAgentEndCheck()', ' updateRuntimeOmpSessionMetadata(ctx)', @@ -154,6 +168,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ` onStatus('agent_start', (${bareCtxParams}) => {`, ...captureSessionMetadata, ' clearPendingAgentEndCheck()', + ' lifecycleState.waiting = false', ' runGeneration += 1', // Why: a turn cannot begin under a dialog holding input focus, so this is the one // boundary that can recover a modal whose close never arrived. @@ -224,11 +239,25 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ' pendingAgentEndCheck = null', ' pendingAgentEndContext = null', ' }', - '', - ' // Why: isIdle flips before agent_settled handlers run, so both paths', - ' // share a guard instead of racing duplicate completion posts — one keyed on the', - ' // generation of the run that ENDED, so a later run still reports its own end.', + ' // Defer completion while live child work remains.', + ' lifecycleState.onEvent = (event: unknown, forcedStatus?: string): void => {', + " if (!event || typeof event !== 'object') return", + " const id = typeof (event as { id?: unknown }).id === 'string' ? (event as { id: string }).id : ''", + ' const status = forcedStatus ?? (event as { status?: unknown }).status', + ' if (!id) return', + " if (status === 'started') { lifecycleState.active.add(id); post('agent_start'); return }", + " if (status !== 'completed' && status !== 'failed' && status !== 'aborted') return", + ' lifecycleState.active.delete(id)', + ' if (lifecycleState.active.size === 0 && lifecycleState.waiting) {', + ' lifecycleState.waiting = false', + ' postAgentEndOnce()', + ' }', + ' }', ' function postAgentEndOnce(): void {', + ' if (lifecycleState.active.size > 0) {', + ' lifecycleState.waiting = true', + ' return', + ' }', ' if (completionPostedGeneration === endedRunGeneration) return', ' completionPostedGeneration = endedRunGeneration', // Why: distinct from the completion guard, which holds the generation of the posted run