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 <neil@stably.ai>

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')

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<typeof vi.fn>): 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')
+  })
+})

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()',

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<string>; exited?: Set<string>; 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<string>(), waiting: false }',
+    '  const lifecycleState = (piEventBus as { __orcaPiSubagents?: { active: Set<string>; exited?: Set<string>; 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<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)',
@@ -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)',

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)',
     '    })',
     '  }',