From 8d79d0d47936fb985cb68ec43db670d21372ad2a Mon Sep 17 00:00:00 2001
From: Jinwoo-H <jinwoo0825@gmail.com>
Date: Sun, 23 Aug 2026 22:02:55 -0400
Subject: [PATCH] fix(runtime): reset stale liveness probes on controller swap

---
 .../pty-process-liveness-broker.test.ts       | 27 ++++++++++++-
 .../runtime/pty-process-liveness-broker.ts    | 39 +++++++++----------
 2 files changed, 44 insertions(+), 22 deletions(-)

diff --git a/src/main/runtime/pty-process-liveness-broker.test.ts b/src/main/runtime/pty-process-liveness-broker.test.ts
index e91e7f00a4b..299aa749320 100644
--- a/src/main/runtime/pty-process-liveness-broker.test.ts
+++ b/src/main/runtime/pty-process-liveness-broker.test.ts
@@ -48,6 +48,31 @@ describe('PtyProcessLivenessBroker', () => {
     expect(inspectProcess).toHaveBeenCalledOnce()
   })
 
+  it('releases admission slots when an in-flight controller probe is invalidated', async () => {
+    const oldInspection = new Promise<{ foregroundProcess: string; hasChildProcesses: boolean }>(
+      () => {}
+    )
+    const inspectProcess = vi
+      .fn()
+      .mockReturnValueOnce(oldInspection)
+      .mockResolvedValue({ foregroundProcess: 'codex', hasChildProcesses: true })
+    const source = { getForegroundProcess: vi.fn(async () => null), inspectProcess }
+    const broker = new PtyProcessLivenessBroker({ timeoutMs: 100, maxConcurrentProbes: 1 })
+
+    void broker.inspect({ source, ptyId: 'old-pty', identity: 'old-incarnation' })
+    expect(broker.getActiveProbeCount()).toBe(1)
+    broker.invalidateAll()
+
+    await expect(
+      broker.inspect({ source, ptyId: 'new-pty', identity: 'new-incarnation' })
+    ).resolves.toEqual({
+      status: 'live',
+      foregroundProcess: 'codex',
+      hasChildProcesses: true
+    })
+    expect(broker.getActiveProbeCount()).toBe(0)
+  })
+
   it('bounds hanging probes across PTYs and prunes retired consumer entries', async () => {
     vi.useFakeTimers()
     const resolvers: ((value: {
@@ -80,7 +105,7 @@ describe('PtyProcessLivenessBroker', () => {
 
     broker.retainConsumerEvidence('restored-agent', [])
     expect(broker.getEntryCount()).toBe(0)
-    expect(broker.getActiveProbeCount()).toBe(2)
+    expect(broker.getActiveProbeCount()).toBe(0)
 
     for (const resolve of resolvers) {
       resolve({ foregroundProcess: 'codex', hasChildProcesses: true })
diff --git a/src/main/runtime/pty-process-liveness-broker.ts b/src/main/runtime/pty-process-liveness-broker.ts
index 140af536f28..99b6549f5d9 100644
--- a/src/main/runtime/pty-process-liveness-broker.ts
+++ b/src/main/runtime/pty-process-liveness-broker.ts
@@ -12,12 +12,10 @@ import {
   maxConcurrentUnscopedProbes,
   processEvidenceKey
 } from './pty-process-liveness-broker-state'
-
 export type {
   PtyProcessInspectionSource,
   PtyProcessLivenessEvidence
 } from './pty-process-inspection'
-
 type PtyProcessEvidenceEntry = {
   source: PtyProcessInspectionSource
   identity: string
@@ -32,7 +30,6 @@ type PtyProcessEvidenceEntry = {
   timedOut: boolean
   probe: Promise<PtyProcessLivenessEvidence> | null
 }
-
 export type PtyProcessLivenessBrokerOptions = {
   timeoutMs: number
   maxConcurrentProbes?: number
@@ -43,17 +40,14 @@ export type PtyProcessLivenessBrokerOptions = {
   now?: () => number
   onInspectionError?: (ptyId: string, error: unknown) => void
 }
-
 export class PtyProcessLivenessBroker {
   private readonly entries = new Map<string, PtyProcessEvidenceEntry>()
   private readonly now: () => number
   private activeProbeCount = 0
   private activeUnscopedProbeCount = 0
-
   constructor(private readonly options: PtyProcessLivenessBrokerOptions) {
     this.now = options.now ?? Date.now
   }
-
   inspect(args: {
     source: PtyProcessInspectionSource
     ptyId: string
@@ -208,23 +202,34 @@ export class PtyProcessLivenessBroker {
         return { status: 'unverifiable', reason: describeProcessInspectionError(error) }
       })
     entry.probe = probe.finally(() => {
-      this.activeProbeCount -= 1
-      if (entry.unscopedProbe) {
-        this.activeUnscopedProbeCount -= 1
+      if (this.entries.get(args.ptyId) === entry) {
+        this.activeProbeCount -= 1
+        if (entry.unscopedProbe) {
+          this.activeUnscopedProbeCount -= 1
+        }
       }
     })
     this.entries.set(args.ptyId, entry)
     return this.waitForProbe(args.ptyId, entry, waitMs)
   }
-
   invalidate(ptyId: string): void {
+    const entry = this.entries.get(ptyId)
+    if (!entry) {
+      return
+    }
     this.entries.delete(ptyId)
+    if (entry.probe) {
+      this.activeProbeCount -= 1
+      if (entry.unscopedProbe) {
+        this.activeUnscopedProbeCount -= 1
+      }
+    }
   }
-
   invalidateAll(): void {
     this.entries.clear()
+    this.activeProbeCount = 0
+    this.activeUnscopedProbeCount = 0
   }
-
   retainConsumerEvidence(
     consumerId: string,
     retained: readonly Readonly<{ ptyId: string; identity: string }>[]
@@ -241,11 +246,10 @@ export class PtyProcessLivenessBroker {
       }
       entry.consumerIds.delete(consumerId)
       if (!entry.hasUnscopedConsumer && entry.consumerIds.size === 0) {
-        this.entries.delete(ptyId)
+        this.invalidate(ptyId)
       }
     }
   }
-
   getPendingCount(): number {
     let count = 0
     for (const entry of this.entries.values()) {
@@ -255,19 +259,15 @@ export class PtyProcessLivenessBroker {
     }
     return count
   }
-
   getActiveProbeCount(): number {
     return this.activeProbeCount
   }
-
   getActiveUnscopedProbeCount(): number {
     return this.activeUnscopedProbeCount
   }
-
   getEntryCount(): number {
     return this.entries.size
   }
-
   private waitForProbe(
     ptyId: string,
     entry: PtyProcessEvidenceEntry,
@@ -295,13 +295,11 @@ export class PtyProcessLivenessBroker {
       })
     })
   }
-
   private unavailableBackoffMs(failureCount: number): number {
     const base = this.options.unavailableBackoffBaseMs ?? DEFAULT_UNAVAILABLE_BACKOFF_BASE_MS
     const max = this.options.unavailableBackoffMaxMs ?? DEFAULT_UNAVAILABLE_BACKOFF_MAX_MS
     return Math.min(max, base * 2 ** Math.max(0, failureCount - 1))
   }
-
   private rememberConsumer(entry: PtyProcessEvidenceEntry, consumerId: string | undefined): void {
     if (consumerId === undefined) {
       entry.hasUnscopedConsumer = true
@@ -309,7 +307,6 @@ export class PtyProcessLivenessBroker {
       entry.consumerIds.add(consumerId)
     }
   }
-
   private storeUnverifiable(entry: PtyProcessEvidenceEntry, reason: string): void {
     entry.failureCount += 1
     entry.evidence = { status: 'unverifiable', reason }