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)