fix(runtime): reset stale liveness probes on controller swap

This commit is contained in:
Jinwoo-H
2026-08-23 22:02:55 -04:00
parent 19dc9670ad
commit 8d79d0d479
2 changed files with 44 additions and 22 deletions
@@ -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 })
+18 -21
View File
@@ -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 }