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

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) => {

src/main/pi/agent-status-extension-test-harness.ts

@@ -46,6 +46,8 @@ export type AgentStatusExtensionHarness = {
   handlers: Record<string, HookHandler>
   processEnv: Record<string, string | undefined>
   callHook: (name: string, event?: unknown, context?: HookContext) => Promise<void>
+  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> }
         ) => void
         setModel: (model: unknown) => Promise<boolean>
+        events?: EventEmitter
       }) => void
     }
   } = { exports: {} }
@@ -191,6 +194,7 @@ export function createAgentStatusExtensionHarness(args: {
   }
 
   const handlers: Record<string, HookHandler> = {}
+  const piEvents = new EventEmitter()
   const commands: AgentStatusExtensionHarness['commands'] = {}
   const setModelMock = vi.fn(async (_model: unknown) => true)
   const registerInto = (target: Record<string, HookHandler>): 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]

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<string>; waiting: boolean; onEvent?: (event: unknown, forcedStatus?: string) => void; listener?: (event: unknown) => void } } | undefined)?.__orcaPiSubagents ?? { active: new Set<string>(), 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