Attach computer-use captures to tool results

Return computer-use captures as multimodal tool-result content so the model can visually inspect fresh screenshots after each remote action. Keep the textual preview for logs and prune older capture payloads to avoid runaway context growth.

Alessandro committed May 23, 2026 at 11:10 UTC 30d364bb97171ac99182a69b3264fac3779c66b2
1 file changed +72 -8
plugins/_a0_connector/tools/computer_use_remote.py
+72 -8
@@ -6,7 +6,7 @@ from pathlib import Path
6 import uuid
7 from typing import Any
8
9 -from helpers import history
9 +from helpers.print_style import PrintStyle
10 from helpers.tool import Response, Tool
11 from helpers.ws import NAMESPACE
12 from helpers.ws_manager import ConnectionNotFoundError, get_shared_ws_manager
@@ -62,6 +62,8 @@ _SUPPORTED_ACTIONS = {
62
63 class ComputerUseRemote(Tool):
64 async def execute(self, **kwargs: Any) -> Response:
65 + self._latest_capture_content: list[dict[str, Any]] | None = None
66 + self._latest_capture_preview = ""
67 action = str(self.args.get("action") or "").strip().lower()
68 if action not in _SUPPORTED_ACTIONS:
69 return Response(
@@ -133,10 +135,37 @@ class ComputerUseRemote(Tool):
135 if capture_note:
136 message = f"{message} {capture_note}".strip()
137
136 - return Response(
137 - message=message,
138 - break_loop=False,
138 + return self._response(message)
139 +
140 + async def after_execution(self, response: Response, **kwargs: Any) -> None:
141 + if not response.additional or not response.additional.get("raw_content"):
142 + await super().after_execution(response, **kwargs)
143 + return
144 +
145 + text = _sanitize_tool_text(response.message.strip())
146 + additional = dict(response.additional)
147 + token_estimate = self._coerce_token_estimate(additional.pop("_tokens", CAPTURE_TOKENS_ESTIMATE))
148 + log_id = str(getattr(getattr(self, "log", None), "id", "") or "")
149 + message = self.agent.hist_add_tool_result(
150 + self.name,
151 + text,
152 + id=log_id,
153 + **additional,
154 )
155 + if hasattr(message, "tokens"):
156 + message.tokens = token_estimate
157 +
158 + agent_name = str(getattr(self.agent, "agent_name", "Agent Zero") or "Agent Zero")
159 + PrintStyle(
160 + font_color="#1B4F72",
161 + background_color="white",
162 + padding=True,
163 + bold=True,
164 + ).print(f"{agent_name}: Response from tool '{self.name}'")
165 + PrintStyle(font_color="#85C1E9").print(text)
166 + if getattr(self, "log", None) is not None:
167 + self.log.update(content=text)
168 + self._prune_prior_capture_history()
169
170 async def _dispatch_payload(self, *, sid: str, payload: dict[str, Any]) -> dict[str, Any]:
171 op_id = str(payload.get("op_id") or "").strip()
@@ -384,15 +413,42 @@ class ComputerUseRemote(Tool):
413 summary = f"{summary} Fresh frame {fresh_state}."
414 else:
415 summary = f"{summary} Fresh capture requested."
387 - content = [
416 + self._latest_capture_content = [
417 {"type": "text", "text": summary},
418 {"type": "image_url", "image_url": {"url": display_ref}},
419 ]
391 - raw_message = history.RawMessage(raw_content=content, preview=summary)
392 - self.agent.hist_add_message(False, content=raw_message, tokens=CAPTURE_TOKENS_ESTIMATE)
393 - self._prune_prior_capture_history()
420 + self._latest_capture_preview = summary
421 return summary
422
423 + def _response(self, message: str) -> Response:
424 + capture_content = self._latest_capture_content
425 + if not capture_content:
426 + return Response(message=message, break_loop=False)
427 +
428 + raw_content = [dict(item) for item in capture_content]
429 + if raw_content and raw_content[0].get("type") == "text":
430 + raw_content[0] = {"type": "text", "text": message}
431 + else:
432 + raw_content.insert(0, {"type": "text", "text": message})
433 +
434 + return Response(
435 + message=message,
436 + break_loop=False,
437 + additional={
438 + "raw_content": raw_content,
439 + "preview": self._latest_capture_preview or message,
440 + "_tokens": CAPTURE_TOKENS_ESTIMATE,
441 + },
442 + )
443 +
444 + @staticmethod
445 + def _coerce_token_estimate(value: object) -> int:
446 + try:
447 + estimate = int(value or 0)
448 + except (TypeError, ValueError):
449 + estimate = 0
450 + return estimate if estimate > 0 else CAPTURE_TOKENS_ESTIMATE
451 +
452 def _prune_prior_capture_history(self) -> None:
453 history_obj = getattr(self.agent, "history", None)
454 if history_obj is None:
@@ -525,3 +581,11 @@ def _safe_filename(value: str) -> str:
581 def _estimated_base64_decoded_size(data: str) -> int:
582 compact_length = sum(1 for char in data if not char.isspace())
583 return (compact_length * 3) // 4
584 +
585 +
586 +def _sanitize_tool_text(value: str) -> str:
587 + try:
588 + from helpers.strings import sanitize_string
589 + except Exception:
590 + return value
591 + return sanitize_string(value)