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

- browser_wait_for's native watchdog adds a standalone snapshot's 20s budget for the snapshot a wait now ends with; the renderer's wait grace widens to 30s to cover it
- Ref reuse also requires the stored entry to point at the same element
- Rename a local that shadowed elementTarget

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,39 @@ 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()
+    let finishSnapshot: (value: unknown) => void = () => {}
+    vi.mocked(contents.executeJavaScript).mockImplementation((expression: string) => {
+      if (!isPageCall(expression, 'collectSnapshot')) return Promise.resolve(undefined)
+      return new Promise((resolve) => {
+        finishSnapshot = resolve
+      })
+    })
+    vi.useFakeTimers()
+    try {
+      const waiting = driver.executeTool('chat-test', 'browser_wait_for', { timeoutMs: 1_000 })
+      await vi.advanceTimersByTimeAsync(9_000)
+      finishSnapshot({
+        url: 'https://example.com/',
+        title: 'Example',
+        outline: '- heading "Order confirmed"',
+        truncated: false,
+        refIds: [],
+        refLineIndexes: {},
+        nextElementId: 1_000,
+      })
+      await vi.advanceTimersByTimeAsync(1_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 &&

apps/sim/lib/mothership/request/tools/executor.test.ts

@@ -266,14 +266,14 @@ describe('pendingToolWaitBudgetMs', () => {
   })
 
   it('matches the requested browser_wait_for renderer budget', () => {
-    expect(pendingToolWaitBudgetMs({ name: 'browser_wait_for', status: 'executing' })).toBe(85_000)
+    expect(pendingToolWaitBudgetMs({ name: 'browser_wait_for', status: 'executing' })).toBe(100_000)
     expect(
       pendingToolWaitBudgetMs({
         name: 'browser_wait_for',
         status: 'executing',
         params: { timeoutMs: 120_000 },
       })
-    ).toBe(195_000)
+    ).toBe(210_000)
   })
 
   it('outlasts the wait window the desktop holds a terminal run for', () => {

apps/sim/lib/mothership/tools/client/browser-tool-execution.test.ts

@@ -1289,11 +1289,11 @@ describe('executeBrowserToolOnClient', () => {
    * native queue while the desktop is still honoring that same wait.
    */
   it.each([
-    ['number', 30_000, 105_000],
-    ['numeric string', '30000', 105_000],
-    ['non-numeric', 'soon', 85_000],
-    ['negative', -5_000, 85_000],
-    ['above the desktop clamp', 500_000, 195_000],
+    ['number', 30_000, 120_000],
+    ['numeric string', '30000', 120_000],
+    ['non-numeric', 'soon', 100_000],
+    ['negative', -5_000, 100_000],
+    ['above the desktop clamp', 500_000, 210_000],
   ])(
     'budgets browser_wait_for above the desktop wait (%s)',
     async (_label, timeoutMs, expected) => {

packages/browser-protocol/src/index.ts

@@ -65,7 +65,8 @@ export type BrowserToolName = (typeof BROWSER_TOOL_NAMES)[number]
 
 const BROWSER_WAIT_FOR_DEFAULT_TIMEOUT_MS = 10_000
 const BROWSER_WAIT_FOR_MAX_TIMEOUT_MS = 120_000
-const BROWSER_WAIT_FOR_RENDERER_GRACE_MS = 15_000
+/** Covers the desktop's wait grace and the snapshot a wait ends with (25s), plus transport. */
+const BROWSER_WAIT_FOR_RENDERER_GRACE_MS = 30_000
 export const BROWSER_TOOL_AUTHORIZATION_TIMEOUT_MS = 8_000
 export const BROWSER_NAVIGATION_NATIVE_WATCHDOG_MS = 60_000
 export const BROWSER_TOOL_QUEUE_WAIT_TIMEOUT_MS = BROWSER_NAVIGATION_NATIVE_WATCHDOG_MS