fix(pi): wait for async subagents before settling pane

src/main/pi/agent-status-extension-omp-lifecycle.test.ts

@@ -15,6 +15,32 @@ 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.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,7 @@ 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
   // 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 +131,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 +193,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 +202,7 @@ export function createAgentStatusExtensionHarness(args: {
         commands[name] = command
       },
       setModel: setModelMock,
+      events: piEvents,
       on(name: string, handler: HookHandler) {
         target[name] = handler
       }
@@ -219,6 +223,9 @@ export function createAgentStatusExtensionHarness(args: {
     callHook: async (name, event, hookContext) => {
       await handlers[name]?.(event, hookContext)
     },
+    emitPiEvent: (name, event) => {
+      piEvents.emit(name, event)
+    },
     reload: () => {
       for (const key of Object.keys(handlers)) {
         delete handlers[key]

src/main/pi/agent-status-handler-source.ts

@@ -129,13 +129,38 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[]
     '  if (ownerPid && ownerPid !== selfPid && isStatusOwnerAlive(ownerPid)) return',
     `  process.env.${ownerEnv} = selfPid`,
     '  resetPostQueue()',
+    '  // Keep the pane working until every Pi lifecycle child reaches a terminal state.',
+    '  const activePiSubagents = new Set<string>()',
+    '  let agentEndWaitingForSubagents = false',
+    '  const piEventBus = (pi as { events?: { on?: (name: string, handler: (event: unknown) => void) => void } }).events',
+    "  piEventBus?.on?.('task:subagent:lifecycle', (event: unknown) => {",
+    "    if (!event || typeof event !== 'object') return",
+    "    const id = typeof (event as { id?: unknown }).id === 'string' ? (event as { id: string }).id : ''",
+    '    const status = (event as { status?: unknown }).status',
+    '    if (!id) return',
+    "    if (status === 'started') {",
+    '      activePiSubagents.add(id)',
+    "      post('agent_start')",
+    '      return',
+    '    }',
+    "    if (status !== 'completed' && status !== 'failed' && status !== 'aborted') return",
+    '    activePiSubagents.delete(id)',
+    '    if (activePiSubagents.size === 0 && agentEndWaitingForSubagents) {',
+    '      agentEndWaitingForSubagents = false',
+    '      postAgentEndOnce()',
+    '    }',
+    '  })',
     ...(kind !== 'pi'
-      ? ["  pi.on('session_shutdown', () => { resetPostQueue(); clearPendingAgentEndCheck() })"]
+      ? [
+          "  pi.on('session_shutdown', () => { activePiSubagents.clear(); agentEndWaitingForSubagents = false; resetPostQueue(); clearPendingAgentEndCheck() })"
+        ]
       : []),
     ...(kind !== 'prime-agent'
       ? [
           "  pi.on('session_switch', (_event, ctx) => {",
           '    if (!isOmpRuntime()) return',
+          '    activePiSubagents.clear()',
+          '    agentEndWaitingForSubagents = false',
           '    resetPostQueue()',
           '    clearPendingAgentEndCheck()',
           '    updateRuntimeOmpSessionMetadata(ctx)',
@@ -229,6 +254,10 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[]
     '  // 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.',
     '  function postAgentEndOnce(): void {',
+    '    if (activePiSubagents.size > 0) {',
+    '      agentEndWaitingForSubagents = true',
+    '      return',
+    '    }',
     '    if (completionPostedGeneration === endedRunGeneration) return',
     '    completionPostedGeneration = endedRunGeneration',
     // Why: distinct from the completion guard, which holds the generation of the posted run