From e65bbc95cec14b365e5cb22e3fdfdff413c1f82b Mon Sep 17 00:00:00 2001 From: Kelvin Amoaba <97001695+AmoabaKelvin@users.noreply.github.com> Date: Mon, 5 Oct 2026 23:24:02 +0000 Subject: [PATCH] fix(omp): stop the tab spinner sticking after subagents finish (#23662) * fix(omp): settle the pane after late or leaked subagents Wake turns started after the root reported done, ids that never settled, and task-session copies of the extension all left the tab spinning. Fixes #22854 * fix(omp): keep pane ownership on the subagent roster Ownership is per event bus, so it has to survive a factory re-run on the same bus. * test(omp): assert the reload case settles on the child's completion * fix(omp): keep tracking children across root runs Clearing the roster on agent_start dropped live children, so Esc or a woken helper could mark the tab done early. Also drops the now-unused reset helper. * fix(pi): only re-arm a completion that was already posted A child starting between a run's end and its settle took that run's done, leaving the pane on working when the settle came. * fix(omp): settle a child that finishes before the first root run After a resume, a woken helper could finish before any run posted a done, leaving the tab spinning. * fix(omp): keep a mid-turn reload from settling the pane on a child's completion The generation-0 clause treated "this factory run has not seen a turn" as "the pane is idle", but an in-process extension reload resets those counters while the root run is still working. A child starting and finishing in that window then posted agent_end, so the tab read done while OMP kept working. Track the root run on the roster, which survives a reload, and only re-open a completion when no run is in flight. --------- Co-authored-by: Neil --- ...t-status-extension-async-subagents.test.ts | 26 ++++ ...ent-status-extension-omp-lifecycle.test.ts | 147 +++++++++++++++++- src/main/pi/agent-status-handler-source.ts | 5 +- .../pi/agent-status-subagent-roster-source.ts | 21 ++- .../pi/omp-session-status-owner-source.ts | 1 + 5 files changed, 193 insertions(+), 7 deletions(-) diff --git a/src/main/pi/agent-status-extension-async-subagents.test.ts b/src/main/pi/agent-status-extension-async-subagents.test.ts index a9145cce874..89d5eebcf72 100644 --- a/src/main/pi/agent-status-extension-async-subagents.test.ts +++ b/src/main/pi/agent-status-extension-async-subagents.test.ts @@ -60,6 +60,32 @@ describe('Pi async subagent roster', () => { vi.useRealTimers() }) + it('settles again when an async child starts after the turn settled', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'pi' }) + await harness.callHook('agent_start') + await endTurn(harness) + startChild(harness, 'late-child', 'tool-call-1') + complete(harness, 'late-child') + await vi.advanceTimersByTimeAsync(0) + + expect(agentEndCount(harness)).toBe(2) + }) + + it("does not spend a run's done on a child that starts before it is posted", async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'pi' }) + let idle = false + const context = { isIdle: () => idle } + await harness.callHook('agent_start') + await harness.callHook('agent_end', {}, context) + startChild(harness, 'quick-child', 'tool-call-1') + complete(harness, 'quick-child') + await harness.callHook('tool_execution_start', { toolName: 'bash' }, context) + idle = true + await vi.advanceTimersByTimeAsync(1_000) + + expect(postedHookNames(harness).at(-1)).toBe('agent_end') + }) + it('settles after an async workflow whose awaited children never report completion', async () => { const harness = createAgentStatusExtensionHarness({ kind: 'pi' }) await harness.callHook('agent_start') 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 5d89d31aec1..03408a3ba6b 100644 --- a/src/main/pi/agent-status-extension-omp-lifecycle.test.ts +++ b/src/main/pi/agent-status-extension-omp-lifecycle.test.ts @@ -1,6 +1,9 @@ -import { describe, expect, it, vi } from 'vitest' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -import { createAgentStatusExtensionHarness } from './agent-status-extension-test-harness' +import { + createAgentStatusExtensionHarness, + type AgentStatusExtensionHarness +} from './agent-status-extension-test-harness' function postedHookNames(fetchMock: ReturnType): string[] { return fetchMock.mock.calls.map( @@ -152,3 +155,143 @@ describe('OMP agent_end contract', () => { } }) }) + +function ompSession(id: string, parentSession?: string) { + return { + sessionManager: { + getSessionId: () => id, + getSessionFile: () => `/sessions/${id}.jsonl`, + getHeader: () => ({ parentSession }) + } + } +} + +describe('OMP subagent settlement', () => { + beforeEach(() => { + vi.useFakeTimers() + }) + afterEach(() => { + vi.useRealTimers() + }) + + async function lifecycle(harness: AgentStatusExtensionHarness, id: string, status: string) { + harness.emitPiEvent('task:subagent:lifecycle', { id, status }) + await vi.advanceTimersByTimeAsync(0) + } + + async function hook(harness: AgentStatusExtensionHarness, name: string) { + await harness.callHook(name) + await vi.advanceTimersByTimeAsync(0) + } + + it.each(OMP_RUNTIME_CASES)( + 'settles %s again when a child starts after the run ended', + async (_name, args) => { + const harness = createAgentStatusExtensionHarness(args) + + await hook(harness, 'agent_start') + await hook(harness, 'agent_end') + await lifecycle(harness, 'wake-1', 'started') + await lifecycle(harness, 'wake-1', 'completed') + + expect(postedHookNames(harness.fetchMock)).toEqual([ + 'agent_start', + 'agent_end', + 'agent_start', + 'agent_end' + ]) + } + ) + + it.each(OMP_RUNTIME_CASES)( + 'keeps %s working while a child outlives the next run', + async (_name, args) => { + const harness = createAgentStatusExtensionHarness(args) + + await hook(harness, 'agent_start') + await lifecycle(harness, 'helper', 'started') + await hook(harness, 'agent_end') + await hook(harness, 'agent_start') + await hook(harness, 'agent_end') + expect(postedHookNames(harness.fetchMock)).not.toContain('agent_end') + + await lifecycle(harness, 'helper', 'completed') + expect(postedHookNames(harness.fetchMock).at(-1)).toBe('agent_end') + } + ) + + it.each(OMP_RUNTIME_CASES)( + 'keeps %s working while a late child outlives the next run', + async (_name, args) => { + const harness = createAgentStatusExtensionHarness(args) + + await hook(harness, 'agent_start') + await hook(harness, 'agent_end') + await lifecycle(harness, 'wake-1', 'started') + await hook(harness, 'agent_start') + await hook(harness, 'agent_end') + expect(postedHookNames(harness.fetchMock).at(-1)).not.toBe('agent_end') + + await lifecycle(harness, 'wake-1', 'completed') + expect(postedHookNames(harness.fetchMock).at(-1)).toBe('agent_end') + } + ) + + it('does not settle OMP when a child finishes mid-run', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'omp' }) + + await hook(harness, 'agent_start') + await lifecycle(harness, 'child-1', 'started') + await lifecycle(harness, 'child-1', 'completed') + + expect(postedHookNames(harness.fetchMock)).not.toContain('agent_end') + }) + + it('settles a child woken before the resumed root has run a turn', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'omp' }) + + await harness.callHook('session_start', {}, ompSession('root')) + await lifecycle(harness, 'revived', 'started') + await lifecycle(harness, 'revived', 'completed') + + expect(postedHookNames(harness.fetchMock)).toEqual(['agent_start', 'agent_end']) + }) + + it("ignores children seen by an OMP task session's copy of the extension", async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'omp' }) + + await harness.callHook('agent_start', {}, ompSession('child', '/sessions/root.jsonl')) + await lifecycle(harness, 'grandchild', 'started') + await lifecycle(harness, 'grandchild', 'completed') + + expect(harness.fetchMock).not.toHaveBeenCalled() + }) + + it('does not settle a reload that lands while the root run is still in flight', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'omp' }) + await harness.callHook('session_start', {}, ompSession('root')) + await hook(harness, 'agent_start') + harness.reload() + + await lifecycle(harness, 'child-1', 'started') + await lifecycle(harness, 'child-1', 'completed') + expect(postedHookNames(harness.fetchMock)).not.toContain('agent_end') + + await hook(harness, 'agent_end') + expect(postedHookNames(harness.fetchMock).at(-1)).toBe('agent_end') + }) + + it('keeps OMP pane ownership across an extension reload', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'omp' }) + await harness.callHook('session_start', {}, ompSession('root')) + await hook(harness, 'agent_start') + await lifecycle(harness, 'child-1', 'started') + await hook(harness, 'agent_end') + harness.reload() + expect(postedHookNames(harness.fetchMock)).not.toContain('agent_end') + + await lifecycle(harness, 'child-1', 'completed') + + expect(postedHookNames(harness.fetchMock).at(-1)).toBe('agent_end') + }) +}) diff --git a/src/main/pi/agent-status-handler-source.ts b/src/main/pi/agent-status-handler-source.ts index a88490cf6ae..76eff628731 100644 --- a/src/main/pi/agent-status-handler-source.ts +++ b/src/main/pi/agent-status-handler-source.ts @@ -136,7 +136,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ...getPiSubagentRosterSetupSourceLines(), ...(kind !== 'pi' ? [ - " pi.on('session_shutdown', () => { lifecycleState.active.clear(); lifecycleState.exited?.clear(); lifecycleState.waiting = false; resetPostQueue(); clearPendingAgentEndCheck() })" + " pi.on('session_shutdown', () => { lifecycleState.active.clear(); lifecycleState.exited?.clear(); lifecycleState.waiting = false; lifecycleState.rootRunInFlight = false; resetPostQueue(); clearPendingAgentEndCheck() })" ] : []), ...(kind !== 'prime-agent' @@ -146,6 +146,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ' lifecycleState.active.clear()', ' lifecycleState.exited?.clear()', ' lifecycleState.waiting = false', + ' lifecycleState.rootRunInFlight = false', ' resetPostQueue()', ' clearPendingAgentEndCheck()', ' updateRuntimeOmpSessionMetadata(ctx)', @@ -165,6 +166,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ...captureSessionMetadata, ' clearPendingAgentEndCheck()', ' lifecycleState.waiting = false', + ' lifecycleState.rootRunInFlight = true', ' 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. @@ -287,6 +289,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ' clearPendingAgentEndCheck()', ' return', ' }', + ' lifecycleState.rootRunInFlight = false', ' endedRunGeneration = runGeneration', ' if (isOmpRuntime()) {', ' postAgentEndOnce()', diff --git a/src/main/pi/agent-status-subagent-roster-source.ts b/src/main/pi/agent-status-subagent-roster-source.ts index 2901e5fe9c8..106d49a7533 100644 --- a/src/main/pi/agent-status-subagent-roster-source.ts +++ b/src/main/pi/agent-status-subagent-roster-source.ts @@ -5,7 +5,7 @@ export function getPiSubagentRosterSetupSourceLines(): string[] { return [ ' const piEventBus = (pi as { events?: { on?: (name: string, handler: (event: unknown) => void) => void } }).events', - ' const lifecycleState = (piEventBus as { __orcaPiSubagents?: { active: Set; exited?: Set; waiting: boolean; onEvent?: (event: unknown, forcedStatus?: string) => void; listener?: (event: unknown) => void; onRunnerExit?: (event: unknown) => void; runnerExitListener?: (event: unknown) => void } } | undefined)?.__orcaPiSubagents ?? { active: new Set(), waiting: false }', + ' const lifecycleState = (piEventBus as { __orcaPiSubagents?: { active: Set; exited?: Set; waiting: boolean; ownsPane?: boolean; rootRunInFlight?: boolean; onEvent?: (event: unknown, forcedStatus?: string) => void; listener?: (event: unknown) => void; onRunnerExit?: (event: unknown) => void; runnerExitListener?: (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)', @@ -23,8 +23,8 @@ export function getPiSubagentRosterSetupSourceLines(): string[] { ] } -// Expects post() and postAgentEndOnce() from the handler scope; the latter prunes -// exited runners before deciding whether children still hold the pane. +// Expects post(), the run generations and postAgentEndOnce() from the handler scope; +// the latter prunes exited runners before deciding whether children still hold the pane. export function getPiSubagentRosterEventSourceLines(): string[] { return [ // Why: a run that reports its own completion does so ~150ms after its runner exits; @@ -37,7 +37,20 @@ export function getPiSubagentRosterEventSourceLines(): string[] { " const id = typeof record.id === 'string' && record.id ? record.id : typeof record.runId === 'string' ? record.runId : ''", ' const status = forcedStatus ?? (event as { status?: unknown }).status', ' if (!id) return', - " if (status === 'started') { lifecycleState.active.add(id); post('agent_start'); return }", + // Why: each OMP task session runs its own copy on its own bus; only the pane's bus tracks children. + ' if (isOmpRuntime() && !lifecycleState.ownsPane) return', + " if (status === 'started') {", + ' lifecycleState.active.add(id)', + // Why: a child starting after the run's done owes a fresh done. Under OMP the same holds + // before the first turn of this factory run (a resumed root, or a reload), but only when no + // root run is in flight -- a reload mid-turn resets these counters while the root still works. + ' if (completionPostedGeneration === runGeneration || (isOmpRuntime() && runGeneration === 0 && !lifecycleState.rootRunInFlight)) {', + ' lifecycleState.waiting = true', + ' completionPostedGeneration = -1', + ' }', + " post('agent_start')", + ' return', + ' }', " if (status !== 'completed' && status !== 'failed' && status !== 'aborted') return", ' lifecycleState.active.delete(id)', ' lifecycleState.exited?.delete(id)', diff --git a/src/main/pi/omp-session-status-owner-source.ts b/src/main/pi/omp-session-status-owner-source.ts index 9e29126f014..daf1a4078bf 100644 --- a/src/main/pi/omp-session-status-owner-source.ts +++ b/src/main/pi/omp-session-status-owner-source.ts @@ -55,6 +55,7 @@ export function getOmpSessionOwnerHandlerSourceLines(): string[] { ' function onStatus(name, handler): void {', ' pi.on(name, (event, ctx) => {', ' if (!ownsSessionStatus(ctx)) return', + ' lifecycleState.ownsPane = true', ' return handler(event, ctx)', ' })', ' }',