fix: keep iOS snapshots fast after relaunch

Author: thymikee

ios-runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerTests+Snapshot.swift

@@ -89,7 +89,15 @@ extension RunnerTests {
       return blocking
     }
 
-    guard let context = try makeSnapshotTraversalContext(app: app, options: options) else {
+    let context: SnapshotTraversalContext?
+    do {
+      context = try makeSnapshotTraversalContext(app: app, options: options)
+    } catch let failure as SnapshotCaptureFailure where options.interactiveOnly {
+      NSLog("AGENT_DEVICE_RUNNER_SNAPSHOT_FLAT_FALLBACK=%@", failure.message)
+      return snapshotFlatInteractive(app: app, options: options)
+    }
+
+    guard let context else {
       return DataPayload(nodes: [], truncated: false)
     }
 
@@ -248,6 +256,82 @@ extension RunnerTests {
     return DataPayload(nodes: nodes, truncated: truncated)
   }
 
+  private func snapshotFlatInteractive(app: XCUIApplication, options: SnapshotOptions) -> DataPayload {
+    let viewport = safeSnapshotViewport(app: app)
+    var nodes: [SnapshotNode] = [
+      SnapshotNode(
+        index: 0,
+        type: "Application",
+        label: nonEmptyElementText { app.label },
+        identifier: nonEmptyElementText { app.identifier },
+        value: nil,
+        rect: snapshotRect(from: safely("SNAPSHOT_FLAT_APP_FRAME", CGRect.zero) { app.frame }),
+        enabled: true,
+        focused: nil,
+        selected: nil,
+        hittable: false,
+        depth: 0,
+        parentIndex: nil,
+        hiddenContentAbove: nil,
+        hiddenContentBelow: nil
+      )
+    ]
+    if options.depth == 0 {
+      return DataPayload(nodes: nodes, truncated: false)
+    }
+
+    var seen = Set<String>()
+    let candidates = flatFallbackElements(app: app)
+      .compactMap { element in
+        flatSnapshotNode(
+          element: element,
+          index: 0,
+          parentIndex: 0,
+          viewport: viewport,
+          options: options
+        )
+      }
+      .filter { node in
+        let key = "\(node.type)-\(node.label ?? "")-\(node.identifier ?? "")-\(node.value ?? "")-\(node.rect.x)-\(node.rect.y)-\(node.rect.width)-\(node.rect.height)"
+        if seen.contains(key) { return false }
+        seen.insert(key)
+        return true
+      }
+      .sorted { left, right in
+        if left.rect.y != right.rect.y {
+          return left.rect.y < right.rect.y
+        }
+        if left.rect.x != right.rect.x {
+          return left.rect.x < right.rect.x
+        }
+        return left.type < right.type
+      }
+
+    let remaining = max(0, fastSnapshotLimit - nodes.count)
+    let truncated = candidates.count > remaining
+    for candidate in candidates.prefix(remaining) {
+      nodes.append(
+        SnapshotNode(
+          index: nodes.count,
+          type: candidate.type,
+          label: candidate.label,
+          identifier: candidate.identifier,
+          value: candidate.value,
+          rect: candidate.rect,
+          enabled: candidate.enabled,
+          focused: candidate.focused,
+          selected: candidate.selected,
+          hittable: candidate.hittable,
+          depth: 1,
+          parentIndex: 0,
+          hiddenContentAbove: nil,
+          hiddenContentBelow: nil
+        )
+      )
+    }
+    return DataPayload(nodes: nodes, truncated: truncated)
+  }
+
   func snapshotRect(from frame: CGRect) -> SnapshotRect {
     return SnapshotRect(
       x: Double(frame.origin.x),
@@ -714,6 +798,95 @@ extension RunnerTests {
     safely("SNAPSHOT_QUERY", [], fetch)
   }
 
+  private func flatFallbackElements(app: XCUIApplication) -> [XCUIElement] {
+    let queries: [XCUIElementQuery] = [
+      app.buttons,
+      app.cells,
+      app.collectionViews,
+      app.images,
+      app.links,
+      app.scrollViews,
+      app.searchFields,
+      app.secureTextFields,
+      app.segmentedControls,
+      app.sliders,
+      app.staticTexts,
+      app.switches,
+      app.tables,
+      app.textFields,
+      app.textViews
+    ]
+    return queries.flatMap { query in
+      safeSnapshotElementsQuery {
+        query.allElementsBoundByIndex
+      }
+    }
+  }
+
+  private func flatSnapshotNode(
+    element: XCUIElement,
+    index: Int,
+    parentIndex: Int?,
+    viewport: CGRect,
+    options: SnapshotOptions
+  ) -> SnapshotNode? {
+    var node: SnapshotNode?
+    let exceptionMessage = RunnerObjCExceptionCatcher.catchException({
+      if !element.exists { return }
+      let frame = element.frame
+      if frame.isNull || frame.isEmpty { return }
+      let visible = isVisibleInViewport(frame, viewport)
+      #if os(macOS)
+        if !visible { return }
+      #endif
+      let label = element.label.trimmingCharacters(in: .whitespacesAndNewlines)
+      let identifier = element.identifier.trimmingCharacters(in: .whitespacesAndNewlines)
+      let valueText = snapshotValueText(element)
+      let hasContent = !label.isEmpty || !identifier.isEmpty || valueText != nil
+      let elementType = element.elementType
+      let enabled = element.isEnabled
+      let hittable = visible && enabled && element.isHittable
+      if options.compact && !hasContent && !hittable && !interactiveTypes.contains(elementType) {
+        return
+      }
+      if let scope = options.scope?.trimmingCharacters(in: .whitespacesAndNewlines), !scope.isEmpty {
+        let haystack = [label, identifier, valueText ?? ""].joined(separator: "\n")
+        if !haystack.localizedCaseInsensitiveContains(scope) {
+          return
+        }
+      }
+
+      node = SnapshotNode(
+        index: index,
+        type: elementTypeName(elementType),
+        label: label.isEmpty ? nil : label,
+        identifier: identifier.isEmpty ? nil : identifier,
+        value: valueText,
+        rect: snapshotRect(from: frame),
+        enabled: enabled,
+        focused: elementHasFocus(element) ? true : nil,
+        selected: element.isSelected ? true : nil,
+        hittable: hittable,
+        depth: 1,
+        parentIndex: parentIndex,
+        hiddenContentAbove: nil,
+        hiddenContentBelow: nil
+      )
+    })
+    if let exceptionMessage {
+      NSLog("AGENT_DEVICE_RUNNER_SNAPSHOT_FLAT_IGNORED_EXCEPTION=%@", exceptionMessage)
+      return nil
+    }
+    return node
+  }
+
+  private func nonEmptyElementText(_ read: () -> String) -> String? {
+    let value = safely("SNAPSHOT_FLAT_TEXT", "") {
+      read()
+    }.trimmingCharacters(in: .whitespacesAndNewlines)
+    return value.isEmpty ? nil : value
+  }
+
   private func isScrollableContainer(_ snapshot: XCUIElementSnapshot, visible: Bool) -> Bool {
     if !visible { return false }
     if !Self.scrollContainerTypes.contains(snapshot.elementType) { return false }

src/commands/capture-snapshot.ts

@@ -278,6 +278,10 @@ function formatRecentSnapshotDropWarning(params: {
   runtimeNow: number;
 }): string | undefined {
   const previousSnapshot = params.session?.snapshot;
+  const samePresentation =
+    previousSnapshot?.presentationKey === undefined ||
+    params.snapshot.presentationKey === undefined ||
+    previousSnapshot.presentationKey === params.snapshot.presentationKey;
   const isRecentSnapshot = previousSnapshot
     ? [params.capturedAt, params.runtimeNow].some((timestamp) => {
         const elapsed = timestamp - previousSnapshot.createdAt;
@@ -287,6 +291,7 @@ function formatRecentSnapshotDropWarning(params: {
   if (
     !params.result.freshness &&
     previousSnapshot &&
+    samePresentation &&
     isRecentSnapshot &&
     isLikelyStaleSnapshotDrop(previousSnapshot.nodes.length, params.snapshot.nodes.length)
   ) {

src/daemon/handlers/__tests__/session.test.ts

@@ -2761,6 +2761,13 @@ test('open --relaunch on iOS stops runner before close/open', async () => {
   });
 
   const calls: string[] = [];
+  mockResolveTargetDevice.mockResolvedValue({
+    platform: 'ios',
+    id: 'ios-device-1',
+    name: 'My iPhone',
+    kind: 'device',
+    booted: true,
+  });
   mockStopIosRunner.mockImplementation(async () => {
     calls.push('stop-runner');
   });
@@ -2788,6 +2795,55 @@ test('open --relaunch on iOS stops runner before close/open', async () => {
   expect(calls).toEqual(['stop-runner', 'close:com.example.app', 'open:com.example.app']);
 });
 
+test('open --relaunch on iOS simulator keeps runner while closing app', async () => {
+  const sessionStore = makeSessionStore();
+  const sessionName = 'ios-simulator-session';
+  sessionStore.set(sessionName, {
+    ...makeSession(sessionName, {
+      platform: 'ios',
+      id: 'sim-1',
+      name: 'iPhone 17 Pro',
+      kind: 'simulator',
+      booted: true,
+    }),
+    appName: 'com.example.app',
+  });
+
+  const calls: string[] = [];
+  mockResolveTargetDevice.mockResolvedValue({
+    platform: 'ios',
+    id: 'sim-1',
+    name: 'iPhone 17 Pro',
+    kind: 'simulator',
+    booted: true,
+  });
+  mockStopIosRunner.mockImplementation(async () => {
+    calls.push('stop-runner');
+  });
+  mockDispatch.mockImplementation(async (_device, command, positionals) => {
+    calls.push(`${command}:${positionals.join(' ')}`);
+    return {};
+  });
+
+  const response = await handleSessionCommands({
+    req: {
+      token: 't',
+      session: sessionName,
+      command: 'open',
+      positionals: [],
+      flags: { relaunch: true },
+    },
+    sessionName,
+    logPath: path.join(os.tmpdir(), 'daemon.log'),
+    sessionStore,
+    invoke: noopInvoke,
+  });
+
+  expect(response).toBeTruthy();
+  expect(response?.ok).toBe(true);
+  expect(calls).toEqual(['close:com.example.app', 'open:com.example.app']);
+});
+
 test('open --relaunch includes timing and waits for iOS runner prewarm', async () => {
   vi.useFakeTimers({ now: 1_000 });
   const sessionStore = makeSessionStore();

src/daemon/handlers/__tests__/snapshot-handler.test.ts

@@ -502,6 +502,60 @@ test('snapshot warns when recent snapshot node count collapses sharply', async (
   }
 });
 
+test('snapshot does not warn on expected node drop across presentation modes', async () => {
+  const sessionStore = makeSessionStore();
+  const sessionName = 'ios-presentation-drop';
+  const session = makeSession(sessionName, iosSimulatorDevice);
+  session.appBundleId = 'com.example.app';
+  session.snapshot = {
+    nodes: Array.from({ length: 50 }, (_, index) => ({
+      ref: `e${index + 1}`,
+      index,
+      depth: 0,
+      type: 'StaticText',
+      label: `Row ${index + 1}`,
+    })),
+    createdAt: Date.now(),
+    backend: 'xctest',
+    presentationKey: buildSnapshotPresentationKey({ interactiveOnly: false, compact: false }),
+  };
+  sessionStore.set(sessionName, session);
+
+  mockDispatch.mockResolvedValue({
+    nodes: Array.from({ length: 8 }, (_, index) => ({
+      index,
+      depth: 0,
+      type: 'Button',
+      label: `Action ${index + 1}`,
+      hittable: true,
+    })),
+    truncated: false,
+    backend: 'xctest',
+  });
+
+  const response = await handleSnapshotCommands({
+    req: {
+      token: 't',
+      session: sessionName,
+      command: 'snapshot',
+      positionals: [],
+      flags: { snapshotInteractiveOnly: true, snapshotCompact: true },
+    },
+    sessionName,
+    logPath: '/tmp/daemon.log',
+    sessionStore,
+  });
+
+  expect(response?.ok).toBe(true);
+  if (response?.ok) {
+    expect(response.data?.warnings ?? []).toEqual(
+      expect.not.arrayContaining([
+        expect.stringContaining('Recent snapshots dropped sharply in node count'),
+      ]),
+    );
+  }
+});
+
 test('snapshot automatically retries stale Android trees after recent navigation', async () => {
   const sessionStore = makeSessionStore();
   const sessionName = 'android-stale-retries-to-fresh';

src/daemon/handlers/session-open.ts

@@ -14,6 +14,7 @@ import { SessionStore } from '../session-store.ts';
 import {
   IOS_SIMULATOR_POST_CLOSE_SETTLE_MS,
   IOS_SIMULATOR_POST_OPEN_SETTLE_MS,
+  isIosSimulator,
   refreshSessionDeviceIfNeeded,
   settleIosSimulator,
 } from './session-device-utils.ts';
@@ -61,7 +62,7 @@ async function relaunchCloseApp(params: {
   context: Parameters<typeof dispatchCommand>[4];
 }): Promise<void> {
   const { device, closeTarget, outFlag, context } = params;
-  if (device.platform !== 'android') {
+  if (device.platform !== 'android' && !isIosSimulator(device)) {
     await stopIosRunnerSession(device.id);
   }
   await dispatchCommand(device, 'close', [closeTarget], outFlag, context);