refactor: satisfy Maestro replay fallow audit

Author: thymikee

android-snapshot-helper/README.md

@@ -31,6 +31,7 @@ VERSION="$(node -p 'require("./package.json").version')"
 adb install -r -t ".tmp/android-snapshot-helper/agent-device-android-snapshot-helper-$VERSION.apk"
 adb shell am instrument -w \
   -e waitForIdleTimeoutMs 500 \
+  -e waitForIdleQuietMs 100 \
   -e timeoutMs 8000 \
   -e maxDepth 128 \
   -e maxNodes 5000 \
@@ -59,6 +60,7 @@ The final instrumentation result includes:
 - `ok=true`
 - `helperApiVersion=1`
 - `waitForIdleTimeoutMs`
+- `waitForIdleQuietMs`
 - `timeoutMs`
 - `maxDepth`
 - `maxNodes`

android-snapshot-helper/src/main/java/com/callstack/agentdevice/snapshothelper/SnapshotInstrumentation.java

@@ -8,6 +8,10 @@
 import android.util.Base64;
 import android.view.accessibility.AccessibilityNodeInfo;
 import android.view.accessibility.AccessibilityWindowInfo;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.lang.reflect.Field;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Locale;
@@ -17,8 +21,9 @@ public final class SnapshotInstrumentation extends Instrumentation {
   private static final String PROTOCOL = "android-snapshot-helper-v1";
   private static final String OUTPUT_FORMAT = "uiautomator-xml";
   private static final String HELPER_API_VERSION = "1";
-  private static final int CHUNK_SIZE = 8 * 1024;
+  private static final int CHUNK_SIZE = 2 * 1024;
   private static final long DEFAULT_WAIT_FOR_IDLE_TIMEOUT_MS = 500;
+  private static final long DEFAULT_WAIT_FOR_IDLE_QUIET_MS = 100;
   private static final long DEFAULT_TIMEOUT_MS = 8_000;
   private static final int DEFAULT_MAX_DEPTH = 128;
   private static final int DEFAULT_MAX_NODES = 5_000;
@@ -36,21 +41,27 @@ public void onStart() {
     super.onStart();
     long waitForIdleTimeoutMs =
         readLongArgument(arguments, "waitForIdleTimeoutMs", DEFAULT_WAIT_FOR_IDLE_TIMEOUT_MS);
+    long waitForIdleQuietMs =
+        readLongArgument(arguments, "waitForIdleQuietMs", DEFAULT_WAIT_FOR_IDLE_QUIET_MS);
     long timeoutMs = readLongArgument(arguments, "timeoutMs", DEFAULT_TIMEOUT_MS);
     int maxDepth = readIntArgument(arguments, "maxDepth", DEFAULT_MAX_DEPTH);
     int maxNodes = readIntArgument(arguments, "maxNodes", DEFAULT_MAX_NODES);
+    String outputPath = readStringArgument(arguments, "outputPath");
     Bundle result = new Bundle();
     result.putString("agentDeviceProtocol", PROTOCOL);
     result.putString("helperApiVersion", HELPER_API_VERSION);
     result.putString("outputFormat", OUTPUT_FORMAT);
     result.putString("waitForIdleTimeoutMs", Long.toString(waitForIdleTimeoutMs));
+    result.putString("waitForIdleQuietMs", Long.toString(waitForIdleQuietMs));
     result.putString("timeoutMs", Long.toString(timeoutMs));
     result.putString("maxDepth", Integer.toString(maxDepth));
     result.putString("maxNodes", Integer.toString(maxNodes));
 
     try {
       long startedAtMs = System.currentTimeMillis();
-      CaptureResult capture = captureXml(waitForIdleTimeoutMs, maxDepth, maxNodes);
+      CaptureResult capture =
+          captureXml(waitForIdleQuietMs, waitForIdleTimeoutMs, timeoutMs, maxDepth, maxNodes);
+      writeOutputFile(outputPath, capture.xml);
       emitChunks(capture.xml);
       result.putString("ok", "true");
       result.putString("rootPresent", Boolean.toString(capture.rootPresent));
@@ -59,26 +70,111 @@ public void onStart() {
       result.putString("nodeCount", Integer.toString(capture.nodeCount));
       result.putString("truncated", Boolean.toString(capture.truncated));
       result.putString("elapsedMs", Long.toString(System.currentTimeMillis() - startedAtMs));
-      finish(0, result);
+      finishSafely(0, result);
     } catch (Throwable error) {
       result.putString("ok", "false");
       result.putString("errorType", error.getClass().getName());
       result.putString(
           "message",
           error.getMessage() == null ? error.getClass().getName() : error.getMessage());
-      finish(1, result);
+      finishSafely(1, result);
+    }
+  }
+
+  private static String readStringArgument(Bundle arguments, String key) {
+    if (arguments == null || !arguments.containsKey(key)) {
+      return null;
+    }
+    String value = arguments.getString(key);
+    return value == null || value.trim().isEmpty() ? null : value.trim();
+  }
+
+  private static void writeOutputFile(String outputPath, String xml) throws IOException {
+    if (outputPath == null) {
+      return;
+    }
+    File file = new File(outputPath);
+    File parent = file.getParentFile();
+    if (parent != null) {
+      parent.mkdirs();
+    }
+    try (FileOutputStream stream = new FileOutputStream(file, false)) {
+      stream.write(xml.getBytes(StandardCharsets.UTF_8));
+    }
+  }
+
+  private void finishSafely(int resultCode, Bundle result) {
+    RuntimeException lastError = null;
+    for (int attempt = 0; attempt < 100; attempt += 1) {
+      try {
+        finish(resultCode, result);
+        return;
+      } catch (IllegalStateException error) {
+        if (!isUiAutomationConnectingError(error)) {
+          throw error;
+        }
+        lastError = error;
+        sleep(100);
+      }
+    }
+    detachUiAutomationBeforeFinish();
+    try {
+      finish(resultCode, result);
+      return;
+    } catch (IllegalStateException error) {
+      if (!isUiAutomationConnectingError(error)) {
+        throw error;
+      }
+      lastError = error;
+    }
+    throw lastError;
+  }
+
+  private void detachUiAutomationBeforeFinish() {
+    try {
+      Field field = Instrumentation.class.getDeclaredField("mUiAutomation");
+      field.setAccessible(true);
+      field.set(this, null);
+    } catch (ReflectiveOperationException | RuntimeException ignored) {
+      // If the platform blocks reflection, preserve the original finish failure below.
+    }
+  }
+
+  private static boolean isUiAutomationConnectingError(IllegalStateException error) {
+    String message = error.getMessage();
+    return message != null && message.contains("while connecting");
+  }
+
+  private static boolean isUiAutomationNotConnectedError(IllegalStateException error) {
+    String message = error.getMessage();
+    return message != null && message.toLowerCase(Locale.ROOT).contains("not connected");
+  }
+
+  private static void sleep(long millis) {
+    try {
+      Thread.sleep(millis);
+    } catch (InterruptedException error) {
+      Thread.currentThread().interrupt();
     }
   }
 
   @SuppressWarnings("deprecation")
-  private CaptureResult captureXml(long waitForIdleTimeoutMs, int maxDepth, int maxNodes)
+  private CaptureResult captureXml(
+      long waitForIdleQuietMs,
+      long waitForIdleTimeoutMs,
+      long timeoutMs,
+      int maxDepth,
+      int maxNodes)
       throws TimeoutException {
-    UiAutomation automation = getUiAutomation();
+    UiAutomation automation = getConnectedUiAutomation(timeoutMs);
     enableInteractiveWindowRetrieval(automation);
     if (waitForIdleTimeoutMs > 0) {
       try {
-        // Best-effort settle: avoids empty roots without inheriting UIAutomator's long idle wait.
-        automation.waitForIdle(waitForIdleTimeoutMs, waitForIdleTimeoutMs);
+        // Best-effort settle: wait for the accessibility stream to become idle, but require only
+        // a short quiet window once it does. Using the full timeout as the quiet window made every
+        // stable snapshot pay a fixed 500 ms tax.
+        long quietMs = Math.min(waitForIdleQuietMs, waitForIdleTimeoutMs);
+        automation.waitForIdle(quietMs, waitForIdleTimeoutMs);
       } catch (TimeoutException ignored) {
         // Busy or animated apps can still expose a usable root; capture whatever is available.
       }
@@ -109,6 +205,30 @@ private CaptureResult captureXml(long waitForIdleTimeoutMs, int maxDepth, int ma
         xml.toString(), windowCount > 0, captureMode, windowCount, stats.nodeCount, stats.truncated);
   }
 
+  private UiAutomation getConnectedUiAutomation(long timeoutMs) throws TimeoutException {
+    long deadlineMs = System.currentTimeMillis() + Math.max(1, timeoutMs);
+    UiAutomation automation = getUiAutomation();
+    RuntimeException lastError = null;
+    while (System.currentTimeMillis() <= deadlineMs) {
+      try {
+        automation.getServiceInfo();
+        return automation;
+      } catch (IllegalStateException error) {
+        if (!isUiAutomationConnectingError(error) && !isUiAutomationNotConnectedError(error)) {
+          throw error;
+        }
+        lastError = error;
+      }
+      sleep(50);
+    }
+    TimeoutException timeout =
+        new TimeoutException("Timed out waiting for Android UiAutomation to connect");
+    if (lastError != null) {
+      timeout.initCause(lastError);
+    }
+    throw timeout;
+  }
+
   private static void enableInteractiveWindowRetrieval(UiAutomation automation) {
     AccessibilityServiceInfo serviceInfo;
     try {

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

@@ -375,6 +375,8 @@ extension RunnerTests {
 
   func focusedTextInput(app: XCUIApplication) -> XCUIElement? {
 #if os(iOS)
+    // iOS focus predicates can return stale or misleading text-input matches
+    // under XCUITest, so text entry readiness is driven by tap/keyboard state.
     return nil
 #else
     var focused: XCUIElement?

src/cli/commands/client-command.ts

@@ -10,6 +10,7 @@ import type {
 import type { CliFlags } from '../../utils/command-schema.ts';
 import { AppError } from '../../utils/errors.ts';
 import { readCommandMessage } from '../../utils/success-text.ts';
+import { isKeyboardAction } from '../../utils/keyboard-actions.ts';
 import { waitCommandCodec } from '../../command-codecs.ts';
 import { parseDeviceRotation } from '../../core/device-rotation.ts';
 import { buildSelectionOptions, writeCommandMessage, writeCommandOutput } from './shared.ts';
@@ -173,13 +174,7 @@ function readKeyboardAction(
 ): KeyboardCommandOptions['action'] | undefined {
   const action = value?.toLowerCase();
   if (action === 'get') return 'status';
-  if (
-    action === undefined ||
-    action === 'status' ||
-    action === 'dismiss' ||
-    action === 'enter' ||
-    action === 'return'
-  ) {
+  if (action === undefined || (isKeyboardAction(action) && action !== 'get')) {
     return action;
   }
   throw new AppError(

src/cli/commands/generic.ts

@@ -69,6 +69,7 @@ const genericClientCommandRunners = {
       ...buildSelectionOptions(flags),
       paths: positionals,
       update: flags.replayUpdate,
+      backend: flags.replayMaestro ? 'maestro' : undefined,
       env: flags.replayEnv,
       failFast: flags.failFast,
       timeoutMs: flags.timeoutMs,

src/client-types.ts

@@ -674,20 +674,24 @@ export type FindOptions =
   | (FindBaseOptions & { action: 'wait'; timeoutMs?: number })
   | (FindBaseOptions & { action: 'fill' | 'type'; value: string });
 
-export type ReplayRunOptions = AgentDeviceRequestOverrides & {
-  path: string;
-  update?: boolean;
-  /** @deprecated Use backend: 'maestro'. */
-  maestro?: boolean;
-  backend?: string;
-  env?: string[];
-  timeoutMs?: number;
-};
+export type ReplayRunOptions = AgentDeviceRequestOverrides &
+  AgentDeviceSelectionOptions & {
+    path: string;
+    update?: boolean;
+    /** @deprecated Use backend: 'maestro'. */
+    maestro?: boolean;
+    backend?: string;
+    env?: string[];
+    timeoutMs?: number;
+  };
 
 export type ReplayTestOptions = AgentDeviceRequestOverrides &
   AgentDeviceSelectionOptions & {
     paths: string[];
     update?: boolean;
+    /** @deprecated Use backend: 'maestro'. */
+    maestro?: boolean;
+    backend?: string;
     env?: string[];
     failFast?: boolean;
     timeoutMs?: number;

src/client.ts

@@ -456,6 +456,7 @@ export function createAgentDeviceClient(
         await executeCommandRequest(PUBLIC_COMMANDS.test, options.paths, {
           ...options,
           replayUpdate: options.update,
+          replayBackend: options.backend ?? (options.maestro === true ? 'maestro' : undefined),
           replayEnv: options.env,
           replayShellEnv: collectReplayClientShellEnv(process.env),
         }),