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..bc4e39b52b1 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,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) => { diff --git a/src/main/pi/agent-status-extension-test-harness.ts b/src/main/pi/agent-status-extension-test-harness.ts index 696153064f1..d4bcd03da51 100644 --- a/src/main/pi/agent-status-extension-test-harness.ts +++ b/src/main/pi/agent-status-extension-test-harness.ts @@ -46,6 +46,7 @@ export type AgentStatusExtensionHarness = { handlers: Record processEnv: Record callHook: (name: string, event?: unknown, context?: HookContext) => Promise + 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 setModel: (model: unknown) => Promise + events?: EventEmitter }) => void } } = { exports: {} } @@ -191,6 +193,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 +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] diff --git a/src/main/pi/agent-status-handler-source.ts b/src/main/pi/agent-status-handler-source.ts index b3bd70e28a4..ea71de4d929 100644 --- a/src/main/pi/agent-status-handler-source.ts +++ b/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()', + ' 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