fix(computer-use): fence a captured frame before it is persisted or sent to aux vision

The epoch check ran only after _dispatch() returned, by which point the
capture path had already written the PNG to the media cache, spilled the
element tree to disk and routed the frame through auxiliary vision — the
human's screen had left the process before being "discarded". A fence
callable is threaded into _dispatch; capture and capture_after call it
the moment backend.capture() returns, and the post-dispatch check stays
for every other action.

(cherry picked from commit 70c09e5fad89d2817f13aea772bd845ab7cc4c8a)

tests/tools/test_bot_desktop_lease.py

@@ -93,7 +93,7 @@ def test_takeover_during_an_admitted_action_discards_its_result(monkeypatch):
 
     monkeypatch.setattr(tool, "_get_backend", lambda session_id="": object())
 
-    def _dispatch_then_takeover(backend, action, args):
+    def _dispatch_then_takeover(backend, action, args, **_):
         lease.acquire("human")  # a whole take-over / hand-back cycle inside the driver call:
         lease.release("human")  # control is back, but the frame is still the human's turn
         return json.dumps({"ok": True, "action": action, "png_b64": "SECRET"})

tests/tools/test_computer_use_capture_fence.py

@@ -0,0 +1,56 @@
+"""A frame captured across a human takeover never leaves the process: the lease fence runs as soon as the
+backend hands the frame back, before it is persisted to the media cache, spilled, or routed to aux vision."""
+
+from __future__ import annotations
+
+import base64
+import json
+
+import pytest
+
+from tools.bot_desktop import lease
+from tools.computer_use.backend import ActionResult, CaptureResult
+
+
+@pytest.fixture(autouse=True)
+def _fresh_lease():
+    lease._reset_for_tests()
+    yield
+    lease._reset_for_tests()
+
+
+class _TakeoverBackend:
+    """Driver whose capture returns while a take-over / hand-back cycle happened underneath it."""
+    _last_target = None
+    _last_app = None
+
+    def capture(self, **kw):
+        lease.acquire("human")
+        lease.release("human")
+        return CaptureResult(mode="som", width=64, height=64, png_b64=base64.b64encode(b"SECRETPNG").decode(), elements=[],
+                             app="Bank", window_title="login")
+
+    def click(self, **kw):
+        return ActionResult(ok=True, action="click")
+
+
+def _spy_sinks(monkeypatch, tool):
+    leaked: list = []
+    monkeypatch.setattr(tool, "_persist_capture_image", lambda cap: leaked.append(("persist", cap)))
+    monkeypatch.setattr(tool, "_spill_elements_to_file", lambda cap: leaked.append(("spill", cap)))
+    monkeypatch.setattr(tool, "_should_route_through_aux_vision", lambda: leaked.append(("aux-decide",)) or True)
+    monkeypatch.setattr(tool, "_route_capture_through_aux_vision", lambda cap, summary, **kw: leaked.append(("aux", cap)))
+    return leaked
+
+
+@pytest.mark.parametrize("args", [{"action": "capture"}, {"action": "click", "coordinate": [1, 1], "capture_after": True}])
+def test_frame_captured_across_a_takeover_is_dropped_before_any_sink(monkeypatch, args):
+    from tools.computer_use import tool
+
+    monkeypatch.setattr(tool, "_get_backend", lambda session_id="": _TakeoverBackend())
+    monkeypatch.setattr(tool, "_request_approval", lambda *a, **k: None)
+    leaked = _spy_sinks(monkeypatch, tool)
+    res = json.loads(tool.handle_computer_use(args))
+    assert res["code"] == "human_has_control", res
+    assert leaked == [], f"the human's frame reached a sink: {leaked}"
+    assert "SECRETPNG" not in json.dumps(res)

tools/computer_use/tool.py

@@ -285,14 +285,20 @@ def handle_computer_use(args: Dict[str, Any], **kwargs) -> Any:
                 _bd_lease.assert_agent_may_act()
             except _bd_lease.HumanHasControl as e:
                 return _refused(e)
-            result = _dispatch(backend, action, args)
-            # Any lease transition during the run voids the result — including a full take-over /
-            # hand-back cycle that already finished: the frame still belongs to the human's turn.
-            if _bd_lease.get().epoch != admitted.epoch:
-                return _refused(_bd_lease.HumanHasControl(
-                    "A human took over this desktop while the action ran; its result was discarded. "
-                    "Re-capture (or call computer_use action='wait_for_human' if they still hold control)."))
+
+            def _fence() -> None:
+                # Any lease transition since admission voids the frame — including a full take-over /
+                # hand-back cycle that already finished: it still belongs to the human's turn. Capture
+                # paths call this BEFORE the frame is persisted, spilled or sent to auxiliary vision.
+                if _bd_lease.get().epoch != admitted.epoch:
+                    raise _bd_lease.HumanHasControl(
+                        "A human took over this desktop while the action ran; its result was discarded. "
+                        "Re-capture (or call computer_use action='wait_for_human' if they still hold control).")
+            result = _dispatch(backend, action, args, fence=_fence)
+            _fence()
             return result
+    except _bd_lease.HumanHasControl as e:
+        return _refused(e)
     except Exception as e:
         logger.exception("computer_use %s failed", action)
         return json.dumps({"error": f"{action} failed: {e}"})
@@ -361,12 +367,13 @@ def _do_scroll(backend, action, args, **delivery):
     return backend.scroll(direction=args.get("direction", "down"), amount=int(args.get("amount", 3)),
                           element=args.get("element"), **_scroll_xy(args), modifiers=args.get("modifiers"), **delivery)
 
-def _do_capture(backend, action, args, **_):
+def _do_capture(backend, action, args, fence=lambda: None, **_):
     if (mode := str(args.get("mode", "som"))) not in {"som", "vision", "ax"}:
         return json.dumps({"error": f"bad mode {mode!r}; use som|vision|ax"})
     # pid/window_id forwarded only when given so older backends keep their defaults.
-    return _capture_response(backend.capture(mode=mode, app=args.get("app"),
-                                             **{k: args[k] for k in ("pid", "window_id") if args.get(k) is not None}))
+    cap = backend.capture(mode=mode, app=args.get("app"), **{k: args[k] for k in ("pid", "window_id") if args.get(k) is not None})
+    fence()
+    return _capture_response(cap)
 
 def _do_listing(backend, action, args, key, **_):
     return json.dumps({key: (items := getattr(backend, action)()), "count": len(items)})
@@ -416,7 +423,9 @@ _ACTION_SUGGESTIONS = {
     "input_text": "type", "screenshot": "capture", "get_window_state": "capture", "left_click": "click", "mouse_click": "click",
 }
 
-def _dispatch(backend: ComputerUseBackend, action: str, args: Dict[str, Any]) -> Any:
+def _dispatch(backend: ComputerUseBackend, action: str, args: Dict[str, Any], fence: Callable[[], None] = lambda: None) -> Any:
+    """``fence`` raises when the screen lease moved since admission; capture paths call it as soon as the
+    frame is in hand, before anything derived from it leaves the process."""
     spec = _ACTIONS.get(action)
     if spec is None:
         return json.dumps({"error": f"unknown action {action!r}" + (f" — did you mean {hint!r}? See the action enum in the tool schema."
@@ -429,10 +438,11 @@ def _dispatch(backend: ComputerUseBackend, action: str, args: Dict[str, Any]) ->
             f"{action} would go to the current target {mismatch!r}, not {requested_app.strip()!r} "
             "— input actions always hit the sticky target from the last capture/focus_app. "
             f"Call capture(app={requested_app.strip()!r}) or focus_app first, then retry.")})
-    # delivery_mode / bring_to_front thread through every input action (background → foreground ladder).
-    res = spec.handler(backend, action, args, delivery_mode=args.get("delivery_mode"),
-                       bring_to_front=bool(args.get("bring_to_front")))
-    return res if isinstance(res, (str, dict)) else _maybe_follow_capture(backend, res, bool(args.get("capture_after")))
+    # delivery_mode / bring_to_front thread through every input action (background → foreground ladder);
+    # read-only actions get the lease fence instead (delivery kwargs would leak into backend input calls).
+    res = spec.handler(backend, action, args, **(dict(delivery_mode=args.get("delivery_mode"), bring_to_front=bool(args.get("bring_to_front")))
+                                                 if spec.input else dict(fence=fence)))
+    return res if isinstance(res, (str, dict)) else _maybe_follow_capture(backend, res, bool(args.get("capture_after")), fence)
 
 # ── Response shaping ────────────────────────────────────────────────────────
 def _classify_action_result(res: ActionResult) -> Dict[str, Any]:
@@ -610,7 +620,8 @@ def _capture_response(cap: CaptureResult, max_elements: int = _DEFAULT_MAX_ELEME
                      "elements_file — read_file/search_files it, or pass app= to narrow scope)")
     return _text_capture_payload(v, "\n".join(lines), extra)
 
-def _maybe_follow_capture(backend: ComputerUseBackend, res: ActionResult, do_capture: bool) -> Any:
+def _maybe_follow_capture(backend: ComputerUseBackend, res: ActionResult, do_capture: bool,
+                          fence: Callable[[], None] = lambda: None) -> Any:
     # No follow-up capture after a failed action: a normal-looking screenshot would suggest success.
     if not do_capture or not res.ok:
         return _text_response(res)
@@ -623,6 +634,7 @@ def _maybe_follow_capture(backend: ComputerUseBackend, res: ActionResult, do_cap
     except Exception as e:
         logger.warning("follow-up capture failed: %s", e)
         return _text_response(res)
+    fence()
     resp, payload = _capture_response(cap), _action_payload(res)
     if isinstance(resp, dict) and resp.get("_multimodal"):
         # Keep the evidence/verdict contract visible alongside the image — it governs whether input may repeat.