fix(desktop-browser): give a wait's closing snapshot a full budget and confirm ref reuse names the same element

Author: waleedlatif1

apps/desktop/src/main/browser-agent/driver.test.ts

@@ -1507,13 +1507,13 @@ describe('browserToolWatchdogMs', () => {
   })
 
   it.each([
-    ['number', 30_000, 35_000],
-    ['numeric string', '30000', 35_000],
-    ['absent', undefined, 15_000],
-    ['non-numeric', 'soon', 15_000],
-    ['zero', 0, 15_000],
-    ['negative', -5_000, 15_000],
-    ['above the wait clamp', 500_000, 125_000],
+    ['number', 30_000, 55_000],
+    ['numeric string', '30000', 55_000],
+    ['absent', undefined, 35_000],
+    ['non-numeric', 'soon', 35_000],
+    ['zero', 0, 35_000],
+    ['negative', -5_000, 35_000],
+    ['above the wait clamp', 500_000, 145_000],
   ])('normalizes browser_wait_for timeout (%s)', (_label, timeoutMs, expected) => {
     const params = timeoutMs === undefined ? {} : { timeoutMs }
 
@@ -3098,6 +3098,40 @@ describe('credential protection', () => {
     })
   })
 
+  it('gives a wait that runs to its timeout a full snapshot budget for its closing observation', async () => {
+    const contents = await openPage()
+    vi.mocked(contents.executeJavaScript).mockImplementation((expression: string) => {
+      if (!isPageCall(expression, 'collectSnapshot')) return Promise.resolve(undefined)
+      return new Promise((resolve) =>
+        setTimeout(
+          () =>
+            resolve({
+              url: 'https://example.com/',
+              title: 'Example',
+              outline: '- heading "Order confirmed"',
+              truncated: false,
+              refIds: [],
+              refLineIndexes: {},
+              nextElementId: 1_000,
+            }),
+          8_000
+        )
+      )
+    })
+    vi.useFakeTimers()
+    try {
+      const waiting = driver.executeTool('chat-test', 'browser_wait_for', { timeoutMs: 1_000 })
+      await vi.advanceTimersByTimeAsync(30_000)
+
+      expect(await waiting).toMatchObject({
+        ok: true,
+        result: { observation: { ok: true, result: { outline: '- heading "Order confirmed"' } } },
+      })
+    } finally {
+      vi.useRealTimers()
+    }
+  })
+
   it('accepts a later snapshot that keeps refs this scope issued for the same page', async () => {
     const contents = await openPage()
     const snapshot = (refIds: number[], nextElementId: number) => ({

apps/desktop/src/main/browser-agent/driver.ts

@@ -1161,8 +1161,9 @@ export function browserToolWatchdogMs(
     return BROWSER_NAVIGATION_NATIVE_WATCHDOG_MS
   }
   if (tool === 'browser_wait_for') {
+    // A wait ends with a snapshot, which gets the same budget a standalone snapshot does.
     const requested = normalizeBrowserWaitForTimeoutMs(params.timeoutMs)
-    return requested + WAIT_FOR_TOOL_WATCHDOG_GRACE_MS
+    return requested + WAIT_FOR_TOOL_WATCHDOG_GRACE_MS + DEFAULT_TOOL_WATCHDOG_MS
   }
   return DEFAULT_TOOL_WATCHDOG_MS
 }
@@ -2971,8 +2972,8 @@ async function executeToolInner(
       }
       const waitedTab = session.requireAutomationTab()
       const contents = waitedTab.view.webContents
-      const elementTarget = elementId === undefined ? undefined : pageTargetForElement(elementId)
-      if (elementTarget && elementTarget !== contents) {
+      const waitTarget = elementId === undefined ? undefined : pageTargetForElement(elementId)
+      if (waitTarget && waitTarget !== contents) {
         throw new ToolError(
           'Element-state waits are limited to the top page. Use a text or URL condition for framed content.'
         )
@@ -2982,7 +2983,7 @@ async function executeToolInner(
       const assertWaitTargetIsCurrent = (): void => {
         assertCurrentExecution()
         if (
-          elementTarget &&
+          waitTarget &&
           (navigationEpoch(contents) !== waitedNavigationEpoch || contents.getURL() !== waitedUrl)
         ) {
           throw new ToolError(
@@ -3025,12 +3026,12 @@ async function executeToolInner(
         }
         const urlMatched = !urlContains || contents.getURL().includes(urlContains)
         let elementMatched = elementId === undefined
-        if (elementId !== undefined && elementState && elementTarget) {
+        if (elementId !== undefined && elementState && waitTarget) {
           assertWaitTargetIsCurrent()
           const state = toRecord(
             unwrapPageResult(
               await execInPage(
-                elementTarget,
+                waitTarget,
                 readPageActionState,
                 [false, elementId, 'registered'],
                 false,

apps/desktop/src/main/browser-agent/page-functions.ts

@@ -645,6 +645,7 @@ export function collectSnapshot(
     const unchanged =
       prior !== undefined &&
       prior.namespace === refNamespace &&
+      prior.element.deref() === el &&
       prior.locator.tag === locator.tag &&
       prior.locator.role === locator.role &&
       prior.locator.name === locator.name &&