From 09490349315a1e2b88929d734edeec477d0349fe Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 27 Jan 2026 09:03:33 +0000 Subject: [PATCH] Optimize aiWaitFor to use DOM mutation signals reducing LLM calls Co-authored-by: alstafeev <18335072+alstafeev@users.noreply.github.com> --- .../java/com/midscene/core/agent/Agent.java | 44 +++++- .../core/agent/AgentBenchmarkTest.java | 133 ++++++++++++++++++ .../midscene/web/driver/PlaywrightDriver.java | 6 + .../midscene/web/driver/SeleniumDriver.java | 6 + 4 files changed, 185 insertions(+), 4 deletions(-) create mode 100644 midscene-core/src/test/java/com/midscene/core/agent/AgentBenchmarkTest.java diff --git a/midscene-core/src/main/java/com/midscene/core/agent/Agent.java b/midscene-core/src/main/java/com/midscene/core/agent/Agent.java index 6777ed1e..ab39b9ef 100644 --- a/midscene-core/src/main/java/com/midscene/core/agent/Agent.java +++ b/midscene-core/src/main/java/com/midscene/core/agent/Agent.java @@ -25,6 +25,22 @@ @Log4j2 public class Agent { + private static final String INIT_MUTATION_OBSERVER_SCRIPT = """ + window.__midscene_mutation_happened = true; // Set to true initially to ensure first check + if (!window.__midscene_observer) { + window.__midscene_observer = new MutationObserver(() => { + window.__midscene_mutation_happened = true; + }); + window.__midscene_observer.observe(document.body, { childList: true, subtree: true, attributes: true, characterData: true }); + } + """; + + private static final String CHECK_AND_RESET_MUTATION_SCRIPT = """ + var happened = window.__midscene_mutation_happened; + window.__midscene_mutation_happened = false; + return happened; + """; + private final Orchestrator orchestrator; private final PageDriver driver; private TaskCache cache; @@ -352,11 +368,31 @@ public void aiWaitFor(String assertion, WaitOptions options) { long checkIntervalMs = options.getCheckIntervalMs(); long startTime = System.currentTimeMillis(); + // Initialize MutationObserver + try { + driver.executeScript(INIT_MUTATION_OBSERVER_SCRIPT); + } catch (Exception e) { + log.warn("Failed to inject MutationObserver, falling back to polling: {}", e.getMessage()); + } + while (System.currentTimeMillis() - startTime < timeoutMs) { - boolean result = aiBoolean("Is the following currently true? " + assertion); - if (result) { - log.info("Wait condition satisfied: {}", assertion); - return; + boolean shouldCheck = true; + try { + Object result = driver.executeScript(CHECK_AND_RESET_MUTATION_SCRIPT); + if (result instanceof Boolean) { + shouldCheck = (Boolean) result; + } + } catch (Exception e) { + // Fallback to true if script fails + shouldCheck = true; + } + + if (shouldCheck) { + boolean result = aiBoolean("Is the following currently true? " + assertion); + if (result) { + log.info("Wait condition satisfied: {}", assertion); + return; + } } try { diff --git a/midscene-core/src/test/java/com/midscene/core/agent/AgentBenchmarkTest.java b/midscene-core/src/test/java/com/midscene/core/agent/AgentBenchmarkTest.java new file mode 100644 index 00000000..f3fd2c22 --- /dev/null +++ b/midscene-core/src/test/java/com/midscene/core/agent/AgentBenchmarkTest.java @@ -0,0 +1,133 @@ +package com.midscene.core.agent; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.contains; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.midscene.core.cache.TaskCache; +import com.midscene.core.model.AIModel; +import com.midscene.core.pojo.options.WaitOptions; +import com.midscene.core.service.PageDriver; +import dev.langchain4j.data.message.AiMessage; +import dev.langchain4j.model.chat.response.ChatResponse; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.invocation.InvocationOnMock; +import org.mockito.stubbing.Answer; + +public class AgentBenchmarkTest { + + @Test + public void testAiWaitForPerformance() { + // Mock dependencies + PageDriver driver = mock(PageDriver.class); + AIModel aiModel = mock(AIModel.class); + + // Setup Agent + // Agent constructor creates an Orchestrator which uses the AIModel + // We need to ensure the AIModel mock is used. + // Agent has a constructor public Agent(PageDriver driver, AIModel aiModel) + Agent agent = new Agent(driver, aiModel); + + // Mock PageDriver screenshot behavior (needed for Planner.query) + when(driver.getScreenshotBase64()).thenReturn("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="); + + // Mock executeScript to simulate mutations + // 1. Init script (any string starting with window.__midscene_mutation_happened = true) -> return null + // 2. Check script -> return sequence: true (initial), false (wait), true (mutation) + when(driver.executeScript(anyString())).thenAnswer(invocation -> { + String script = invocation.getArgument(0); + if (script.contains("if (!window.__midscene_observer)")) { + return null; + } + if (script.contains("return happened")) { + // simulate: + // call 1: true (initial check) + // call 2: false (no mutation) + // call 3: true (mutation happened) + // call 4: ... + return invocation.getMock().toString().contains("call_1") ? true : false; + // Wait, Mockito state is hard to track this way inside the answer without external counter. + // Let's use a counter. + } + return null; + }); + + // Refined mock for executeScript with counter + final int[] scriptCallCount = {0}; + when(driver.executeScript(anyString())).thenAnswer(invocation -> { + String script = invocation.getArgument(0); + if (script.contains("if (!window.__midscene_observer)")) { + return null; + } + if (script.contains("return happened")) { + scriptCallCount[0]++; + if (scriptCallCount[0] == 1) return true; // Initial check + if (scriptCallCount[0] == 2) return false; // No mutation + if (scriptCallCount[0] == 3) return true; // Mutation happened + return false; + } + return null; + }); + + // Mock AIModel behavior + // aiBoolean calls aiQuery which calls model.chat(messages) + // We need to mock chat behavior. + // Orchestrator logic might wrap the message. + // However, aiBoolean expects "true" or "false" in the response. + + // Scenario: + // Iteration 1: executeScript(check) -> true. AI Called. Response: false. + // Iteration 2: executeScript(check) -> false. AI Skipped. + // Iteration 3: executeScript(check) -> true. AI Called. Response: true. + // Total AI calls: 2 + + when(aiModel.chat(any())).thenAnswer(new Answer() { + private int count = 0; + + @Override + public ChatResponse answer(InvocationOnMock invocation) { + count++; + String responseText = "false"; + // If count is 2, it means it's the second call to AI. + // In our scenario, the 2nd call happens when mutation occurs and condition is true. + if (count >= 2) { + responseText = "true"; + } + return ChatResponse.builder() + .aiMessage(AiMessage.from(responseText)) + .build(); + } + }); + + long start = System.currentTimeMillis(); + + // Set check interval to a small value to make test fast but still loop + WaitOptions options = WaitOptions.builder() + .timeoutMs(2000) + .checkIntervalMs(50) + .build(); + + agent.aiWaitFor("some condition", options); + + long duration = System.currentTimeMillis() - start; + System.out.println("Duration: " + duration + "ms"); + + // Verify AIModel was called 2 times (instead of 3 in the baseline scenario if we polled every time) + // Baseline would be: + // 1. check -> true (false) + // 2. check -> false (false) - but polling would call AI + // 3. check -> true (true) + + // With optimization: + // 1. check -> true (mutation happened/init) -> AI call (false) + // 2. check -> false (no mutation) -> Skip AI + // 3. check -> true (mutation happened) -> AI call (true) -> return + + verify(aiModel, times(2)).chat(any()); + } +} diff --git a/midscene-web/src/main/java/com/midscene/web/driver/PlaywrightDriver.java b/midscene-web/src/main/java/com/midscene/web/driver/PlaywrightDriver.java index e77475b1..6656db75 100644 --- a/midscene-web/src/main/java/com/midscene/web/driver/PlaywrightDriver.java +++ b/midscene-web/src/main/java/com/midscene/web/driver/PlaywrightDriver.java @@ -242,6 +242,12 @@ public void goBack() { waitUntilPageLoaded(); } + @Override + public Object executeScript(String script) { + waitUntilPageLoaded(); + return page.evaluate(script); + } + // ========== Private Helper Methods ========== private void waitUntilPageLoaded() { diff --git a/midscene-web/src/main/java/com/midscene/web/driver/SeleniumDriver.java b/midscene-web/src/main/java/com/midscene/web/driver/SeleniumDriver.java index 931a514f..50596e74 100644 --- a/midscene-web/src/main/java/com/midscene/web/driver/SeleniumDriver.java +++ b/midscene-web/src/main/java/com/midscene/web/driver/SeleniumDriver.java @@ -276,6 +276,12 @@ public void goBack() { waitUntilPageLoaded(); } + @Override + public Object executeScript(String script) { + waitUntilPageLoaded(); + return ((JavascriptExecutor) driver).executeScript(script); + } + // ========== Private Helper Methods ========== public void waitUntilPageLoaded() {