fix: refresh Apple runner state after app relaunch

Author: thymikee

apple-runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerAXSnapshotBridge.h

@@ -9,6 +9,8 @@ NS_ASSUME_NONNULL_BEGIN
                                                     maxDepth:(NSInteger)maxDepth
                                                     maxNodes:(NSInteger)maxNodes;
 
++ (NSInteger)processIdentifierForApplication:(XCUIApplication *)application;
+
 @end
 
 NS_ASSUME_NONNULL_END

apple-runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerAXSnapshotBridge.m

@@ -164,6 +164,11 @@ + (NSInteger)integerFrom:(id)target selectorName:(NSString *)selectorName
   return send(target, selector);
 }
 
++ (NSInteger)processIdentifierForApplication:(XCUIApplication *)application
+{
+  return [self integerFrom:application selectorName:@"processID"];
+}
+
 + (id)accessibilityApplicationForApplication:(XCUIApplication *)application axClient:(id)axClient
 {
   NSInteger targetProcessID = [self integerFrom:application selectorName:@"processID"];

apple-runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerTests+CommandExecution.swift

@@ -945,6 +945,8 @@ extension RunnerTests {
       if let bundleId = requestedBundleId {
         if currentBundleId != bundleId || currentApp == nil {
           _ = activateTarget(bundleId: bundleId, reason: "bundle_changed")
+        } else {
+          refreshCachedTargetIfProcessChanged(bundleId: bundleId)
         }
       } else {
         // Do not reuse stale bundle targets when the caller does not explicitly request one.

apple-runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerTests+Lifecycle.swift

@@ -85,6 +85,8 @@ extension RunnerTests {
     }
     currentApp = app
     currentBundleId = nil
+    currentAppProcessIdentifier = nil
+    snapshotXCTestPenaltyWarmupExemptionPending = false
   }
 
   func invalidateCachedTarget(reason: String) {
@@ -93,6 +95,42 @@ extension RunnerTests {
     }
     currentApp = nil
     currentBundleId = nil
+    currentAppProcessIdentifier = nil
+    snapshotXCTestPenaltyWarmupExemptionPending = false
+  }
+
+  func refreshCachedTargetIfProcessChanged(bundleId: String) {
+    guard currentBundleId == bundleId, currentApp != nil else { return }
+    let candidate = XCUIApplication(bundleIdentifier: bundleId)
+    let observedProcessIdentifier = Self.processIdentifier(of: candidate)
+    guard Self.shouldRefreshCachedTarget(
+      cachedProcessIdentifier: currentAppProcessIdentifier,
+      observedProcessIdentifier: observedProcessIdentifier
+    ) else { return }
+    NSLog(
+      "AGENT_DEVICE_RUNNER_TARGET_CACHE_REFRESH bundle=%@ previousPid=%d currentPid=%d",
+      bundleId,
+      currentAppProcessIdentifier ?? 0,
+      observedProcessIdentifier ?? 0
+    )
+    currentApp = candidate
+    currentAppProcessIdentifier = observedProcessIdentifier
+    clearSnapshotXCTestChannelPenalty(reason: "target_process_changed")
+    snapshotXCTestPenaltyWarmupExemptionPending = true
+    needsFirstInteractionDelay = true
+  }
+
+  static func processIdentifier(of target: XCUIApplication) -> Int? {
+    let value = RunnerAXSnapshotBridge.processIdentifier(for: target)
+    return value > 0 ? value : nil
+  }
+
+  static func shouldRefreshCachedTarget(
+    cachedProcessIdentifier: Int?,
+    observedProcessIdentifier: Int?
+  ) -> Bool {
+    guard let cachedProcessIdentifier, let observedProcessIdentifier else { return false }
+    return cachedProcessIdentifier != observedProcessIdentifier
   }
 
   func targetNeedsActivation(_ target: XCUIApplication) -> Bool {
@@ -141,6 +179,8 @@ extension RunnerTests {
     target.activate()
     currentApp = target
     currentBundleId = bundleId
+    currentAppProcessIdentifier = Self.processIdentifier(of: target)
+    snapshotXCTestPenaltyWarmupExemptionPending = false
     needsFirstInteractionDelay = true
     return target
   }

apple-runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerTests+LifecycleCacheTests.swift

@@ -0,0 +1,40 @@
+import XCTest
+
+extension RunnerTests {
+  func testCachedTargetRefreshRequiresChangedPositiveProcessIdentity() {
+    XCTAssertFalse(
+      Self.shouldRefreshCachedTarget(
+        cachedProcessIdentifier: nil,
+        observedProcessIdentifier: 42
+      )
+    )
+    XCTAssertFalse(
+      Self.shouldRefreshCachedTarget(
+        cachedProcessIdentifier: 42,
+        observedProcessIdentifier: 42
+      )
+    )
+    XCTAssertTrue(
+      Self.shouldRefreshCachedTarget(
+        cachedProcessIdentifier: 41,
+        observedProcessIdentifier: 42
+      )
+    )
+  }
+
+  func testSnapshotPenaltyWarmupExemptionIsConsumedOnce() {
+    snapshotXCTestPenaltyWarmupExemptionPending = true
+
+    XCTAssertTrue(consumeSnapshotXCTestPenaltyWarmupExemption())
+    XCTAssertFalse(consumeSnapshotXCTestPenaltyWarmupExemption())
+  }
+
+  func testSnapshotPenaltyCanBeClearedAcrossTargetProcessReplacement() {
+    penalizeSnapshotXCTestChannel(bundleId: "com.example.app", reason: "test")
+    XCTAssertTrue(isSnapshotXCTestChannelPenalized(bundleId: "com.example.app"))
+
+    clearSnapshotXCTestChannelPenalty(reason: "target_process_changed")
+
+    XCTAssertFalse(isSnapshotXCTestChannelPenalized(bundleId: "com.example.app"))
+  }
+}

apple-runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerTests+SnapshotCapturePlan.swift

@@ -118,6 +118,17 @@ extension RunnerTests {
     )
   }
 
+  func clearSnapshotXCTestChannelPenalty(reason: String) {
+    snapshotXCTestChannelPenaltyLock.lock()
+    let hadActivePenalty = Date() < snapshotXCTestChannelPenaltyUntil
+    snapshotXCTestChannelPenaltyBundleId = nil
+    snapshotXCTestChannelPenaltyUntil = Date.distantPast
+    snapshotXCTestChannelPenaltyLock.unlock()
+    if hadActivePenalty {
+      NSLog("AGENT_DEVICE_RUNNER_SNAPSHOT_XCTEST_CHANNEL_PENALTY_CLEARED reason=%@", reason)
+    }
+  }
+
   func isSnapshotXCTestChannelPenalized(bundleId: String?) -> Bool {
     snapshotXCTestChannelPenaltyLock.lock()
     defer { snapshotXCTestChannelPenaltyLock.unlock() }
@@ -127,6 +138,12 @@ extension RunnerTests {
     return penalized == bundleId
   }
 
+  func consumeSnapshotXCTestPenaltyWarmupExemption() -> Bool {
+    let pending = snapshotXCTestPenaltyWarmupExemptionPending
+    snapshotXCTestPenaltyWarmupExemptionPending = false
+    return pending
+  }
+
   /// Pure plan-reorder rule: a penalized XCTest accessibility channel uses independent backends
   /// when the platform has one, otherwise it keeps XCTest work on a short probe. The raw
   /// diagnostic plan keeps tree-first errors, and unknown plans are left untouched.
@@ -174,6 +191,7 @@ extension RunnerTests {
     var firstFailure: (reason: String, code: String)?
     var axFailure: SnapshotCaptureFailure?
     let deadline = Date().addingTimeInterval(Self.snapshotPlanBudget)
+    let suppressXCTestPenalty = consumeSnapshotXCTestPenaltyWarmupExemption()
 
     // Reorder is iOS-only because hostile screens can make XCTest tree/query work grind while
     // the app remains visually responsive. Simulators can avoid that channel through private AX;
@@ -237,14 +255,30 @@ extension RunnerTests {
             treeCaptureSliceBudgetOverride: effective.treeCaptureSliceBudgetOverride
           )
         else {
-          recordSlowXCTestSnapshotBackendIfNeeded(kind, startedAt: backendStartedAt)
+          recordSlowXCTestSnapshotBackendIfNeeded(
+            kind,
+            startedAt: backendStartedAt,
+            penaltySuppressed: suppressXCTestPenalty
+          )
           continue
         }
         capture = result
-        recordSlowXCTestSnapshotBackendIfNeeded(kind, startedAt: backendStartedAt)
+        recordSlowXCTestSnapshotBackendIfNeeded(
+          kind,
+          startedAt: backendStartedAt,
+          penaltySuppressed: suppressXCTestPenalty
+        )
       } catch let failure as SnapshotCaptureFailure {
-        recordXCTestSnapshotBackendFailureIfNeeded(kind, failure: failure)
-        recordSlowXCTestSnapshotBackendIfNeeded(kind, startedAt: backendStartedAt)
+        recordXCTestSnapshotBackendFailureIfNeeded(
+          kind,
+          failure: failure,
+          penaltySuppressed: suppressXCTestPenalty
+        )
+        recordSlowXCTestSnapshotBackendIfNeeded(
+          kind,
+          startedAt: backendStartedAt,
+          penaltySuppressed: suppressXCTestPenalty
+        )
         if Self.isAxSnapshotFailure(failure) { axFailure = failure }
         if firstFailure == nil {
           firstFailure = (failure.message, Self.isAxSnapshotFailure(failure) ? "ax-rejected" : "capture-failed")
@@ -313,7 +347,12 @@ extension RunnerTests {
 
   /// Marks XCTest-backed snapshot tiers as penalized when one attempt ground past the slow-capture
   /// threshold — even a successful one: the next capture of this screen must not re-grind.
-  private func recordSlowXCTestSnapshotBackendIfNeeded(_ kind: SnapshotBackendKind, startedAt: Date) {
+  private func recordSlowXCTestSnapshotBackendIfNeeded(
+    _ kind: SnapshotBackendKind,
+    startedAt: Date,
+    penaltySuppressed: Bool
+  ) {
+    guard !penaltySuppressed else { return }
     guard kind.usesXCTestAccessibilityChannel else { return }
     let elapsed = Date().timeIntervalSince(startedAt)
     guard elapsed > snapshotXCTestSlowCaptureThreshold else { return }
@@ -325,8 +364,10 @@ extension RunnerTests {
 
   private func recordXCTestSnapshotBackendFailureIfNeeded(
     _ kind: SnapshotBackendKind,
-    failure: SnapshotCaptureFailure
+    failure: SnapshotCaptureFailure,
+    penaltySuppressed: Bool
   ) {
+    guard !penaltySuppressed else { return }
     guard kind.usesXCTestAccessibilityChannel, failure.code == Self.xCTestSnapshotTimeoutCode else { return }
     penalizeSnapshotXCTestChannel(
       bundleId: currentBundleId,

apple-runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerTests.swift

@@ -38,6 +38,7 @@ final class RunnerTests: XCTestCase {
   lazy var springboard = XCUIApplication(bundleIdentifier: Self.springboardBundleId)
   var currentApp: XCUIApplication?
   var currentBundleId: String?
+  var currentAppProcessIdentifier: Int?
   let maxRequestBytes = 2 * 1024 * 1024
   let mainThreadExecutionTimeout: TimeInterval = 30
   let appExistenceTimeout: TimeInterval = 30
@@ -76,6 +77,7 @@ final class RunnerTests: XCTestCase {
   var snapshotXCTestChannelPenaltyBundleId: String?
   var snapshotXCTestChannelPenaltyUntil = Date.distantPast
   let snapshotXCTestChannelPenaltyDuration: TimeInterval = 120
+  var snapshotXCTestPenaltyWarmupExemptionPending = false
   // Bluesky-class screens can grind ~4-8s before an XCTest-backed snapshot tier fails; anything
   // past this threshold marks the screen hostile so the next capture uses non-XCTest recovery.
   let snapshotXCTestSlowCaptureThreshold: TimeInterval = 3