fix(pi): continue the chat when a background subagent wakes an idle parent (#31778)

src/main/pi/rpc-capture-replay.test.ts

@@ -98,6 +98,10 @@ async function replay(name: string) {
     if (dir === 'out' && (frame.type !== 'response' || frame.command === 'prompt')) {
       turns.receive(frame)
     }
+    if (dir === 'out' && frame.type === 'response' && frame.command === 'get_state') {
+      // The idle probe resolved here in the captured run.
+      await new Promise<void>((resolve) => setImmediate(resolve))
+    }
   }
   await new Promise<void>((resolve) => setImmediate(resolve))
   lane.flush()
@@ -105,6 +109,16 @@ async function replay(name: string) {
   return { rows, turns: await rig.turns(), accepted, failed, idle, settled }
 }
 
+function transcript(rows: Awaited<ReturnType<typeof replay>>['rows']): string {
+  return rows
+    .flatMap((row) =>
+      row.body.kind === 'message'
+        ? row.body.blocks.flatMap((block) => (block.type === 'text' ? [block.text] : []))
+        : []
+    )
+    .join('\n')
+}
+
 describe('scrubbed Pi 1.0.4 RPC captures', () => {
   it('records tool output, streamed text, and all three consumed inputs in one settled run', async () => {
     const result = await replay('steer-followup')
@@ -115,13 +129,8 @@ describe('scrubbed Pi 1.0.4 RPC captures', () => {
     expect(
       result.rows.some((row) => row.body.kind === 'tool-call' && row.body.state === 'completed')
     ).toBe(true)
-    const texts = result.rows.flatMap((row) =>
-      row.body.kind === 'message'
-        ? row.body.blocks.flatMap((block) => (block.type === 'text' ? [block.text] : []))
-        : []
-    )
-    expect(texts.join('\n')).toContain('STEER_CONSUMED')
-    expect(texts.join('\n')).toContain('FOLLOWUP_CONSUMED')
+    expect(transcript(result.rows)).toContain('STEER_CONSUMED')
+    expect(transcript(result.rows)).toContain('FOLLOWUP_CONSUMED')
   })
   it.each(['provider-fetch-error', 'provider-socket-error'])(
     'marks captured %s as failure with the provider diagnostic',
@@ -148,6 +157,14 @@ describe('scrubbed Pi 1.0.4 RPC captures', () => {
     expect(result.turns[0]?.outcome).toBe('success')
     expect(result.accepted).toHaveBeenCalledWith('send-1')
   })
+  it('tracks the run a captured Pi 1.3.0 subagent wake starts on the idle parent', async () => {
+    const result = await replay('subagent-wake')
+    expect(result.failed).not.toHaveBeenCalled()
+    expect(result.accepted.mock.calls.map(([id]) => id)).toEqual(['send-1'])
+    expect(result.turns.map((turn) => turn.outcome)).toEqual(['success', 'success'])
+    expect(result.idle).toHaveBeenCalledTimes(2)
+    expect(transcript(result.rows)).toContain('WAKE_CONSUMED')
+  })
   it('preserves empty answers and literal Unicode separators in captured dialogs', () => {
     const dialogs = capture('dialogs').flatMap(({ dir, frame }) =>
       dir === 'out' && frame.type === 'extension_ui_request' ? [frame] : []

src/main/pi/rpc-session-adapter.test.ts

@@ -329,6 +329,16 @@ describe('Pi RPC session ownership and delivery', () => {
     expect(h.connection.closed).toBe(true)
   })
 
+  it('stops a run Pi starts on an idle session without a send', async () => {
+    const h = await setup()
+    h.connection.receive({ type: 'agent_start' })
+    await vi.waitFor(async () => expect((await h.rig.turns())[0]?.state).toBe('running'))
+    expect((await h.adapter.cancelTurn({ sessionId, fence: 7 })).cancelled).toBe(true)
+    expect(h.connection.requests).toContain('abort')
+    h.connection.receive({ type: 'agent_settled' })
+    await vi.waitFor(async () => expect((await h.rig.turns())[0]?.outcome).toBe('cancellation'))
+  })
+
   it('keeps a pending manual compaction cancellable and reports root/tree proof honestly', async () => {
     const h = await setup()
     const turnId = 'manual-compact'

src/main/pi/rpc-turn-races.test.ts

@@ -94,6 +94,20 @@ describe('Pi turn settlement races', () => {
       expect(errors[0]?.text).toContain('`/login`')
     }
   })
+  it.each(['user', 'custom'])(
+    'accepts a steer only at its own echo after a %s message starts an idle run',
+    async (role) => {
+      const h = await setup()
+      h.turns.receive({ type: 'agent_start' })
+      await h.turns.submit('steer', 1, { type: 'prompt', message: 'steer' })
+      h.turns.receive({ type: 'message_end', message: { role, content: [] } })
+      await h.flush()
+      expect(h.accepted).not.toHaveBeenCalled()
+      h.turns.receive({ type: 'message_end', message: { role: 'user', content: [] } })
+      await h.flush()
+      expect(h.accepted).toHaveBeenCalledWith('steer')
+    }
+  )
   it('chooses steering after asynchronous dispatch admission finishes', async () => {
     const h = await setup()
     await h.turns.submit('first', 1, { type: 'prompt', message: 'first' })

src/main/pi/rpc-turns.ts

@@ -69,7 +69,7 @@ export class PiRpcTurns {
         this.deps.lane.apply([{ type: 'input.accepted', clientMessageId: id, requestedAt: at }])
       },
       commandOnly: () => {
-        this.skipInitialUserEcho = false
+        this.awaitingRunTrigger = false
         this.settling = true
         this.probe()
       },
@@ -120,33 +120,35 @@ export class PiRpcTurns {
           this.delivery.reply(frame)
         }
         return
-      case 'agent_start':
-        if (!this.active && !this.delivery.holdsDispatch) {
-          throw new Error('Pi started an unsolicited agent run')
-        }
+      case 'agent_start': {
+        // A run Pi starts while idle (an extension wake) opens with its own prompt.
+        const idle = !this.active
         this.touch()
         this.open(at)
-        this.skipInitialUserEcho = this.delivery.consumeNext()
+        this.awaitingRunTrigger = this.delivery.consumeNext() || idle
         return
+      }
       case 'message_start':
       case 'message_update':
       case 'message_end': {
         this.touch()
         const message = piRpcMessageSchema.safeParse(frame.message)
-        if (message.success && message.data.role === 'user' && frame.type === 'message_end') {
-          // agent_start already accepted the first prompt; user echoes consume queued inputs.
-          if (this.skipInitialUserEcho) {
-            this.skipInitialUserEcho = false
-          } else {
-            this.delivery.consumeNext()
+        if (message.success && frame.type === 'message_end') {
+          const { role } = message.data
+          if (role !== 'system') {
+            // A run's first message is the prompt that started it; later user echoes consume queued inputs.
+            if (role === 'user' && !this.awaitingRunTrigger) {
+              this.delivery.consumeNext()
+            }
+            this.awaitingRunTrigger = false
+          }
+          if (role === 'assistant') {
+            this.lastStopReason = message.data.stopReason
+            this.failure =
+              message.data.stopReason === 'error'
+                ? (message.data.errorMessage ?? 'Pi request failed')
+                : undefined
           }
-        }
-        if (message.success && message.data.role === 'assistant' && frame.type === 'message_end') {
-          this.lastStopReason = message.data.stopReason
-          this.failure =
-            message.data.stopReason === 'error'
-              ? (message.data.errorMessage ?? 'Pi request failed')
-              : undefined
         }
         this.deps.lane.apply(this.messages.message(frame, at))
         return
@@ -240,7 +242,7 @@ export class PiRpcTurns {
     this.delivery.end()
   }
 
-  private skipInitialUserEcho = false
+  private awaitingRunTrigger = false
   private open(at: number): void {
     if (this.active) {
       return
@@ -249,7 +251,6 @@ export class PiRpcTurns {
     this.failure = undefined
     this.lastStopReason = undefined
     this.stopped = false
-    this.skipInitialUserEcho = true
     this.deps.lane.apply([{ type: 'turn.open', turn: this.active, at }])
   }