perf(agent): prune ARIA observations incrementally

src/webwright/agents/default.py

@@ -89,6 +89,10 @@ def __init__(self, model: Model, env: Environment, *, config_class: type = Agent
         self.extra_template_vars: dict[str, Any] = {}
         self.n_calls = 0
         self.n_format_errors = 0
+        # Track observations already scanned so pruning only visits new messages.
+        self._observations: list[tuple[dict[str, Any], dict[str, Any]]] = []
+        self._next_msg_index = 0
+        self._pruned_until_obs_index = 0
 
     def _debug_dir(self) -> Path | None:
         if self.config.output_path is None:
@@ -276,29 +280,47 @@ def _prune_old_observation_aria_snapshots(self) -> None:
         n = self.config.keep_last_n_observations
         if n <= 0:
             return
-        obs_indices = [
-            i for i, m in enumerate(self.messages)
-            if m.get("extra", {}).get("observation")
-        ]
-        if len(obs_indices) <= n:
-            return
-        placeholder = "(ARIA snapshot pruned; see most recent observation)"
-        for idx in obs_indices[:-n]:
-            msg = self.messages[idx]
-            obs = msg["extra"]["observation"]
-            aria = obs.get("aria_snapshot", "")
-            if not aria:
-                continue
-            content = msg.get("content")
-            if isinstance(content, list):
-                for part in content:
-                    if isinstance(part, dict) and part.get("type") in ("text", "input_text"):
-                        text = part.get("text", "")
-                        if aria in text:
-                            part["text"] = text.replace(aria, placeholder)
-            elif isinstance(content, str) and aria in content:
-                msg["content"] = content.replace(aria, placeholder)
-            obs["aria_snapshot"] = ""
+
+        msgs = self.messages
+        observations = self._observations
+        curr_len = len(msgs)
+        next_idx = self._next_msg_index
+
+        # Scan only messages appended since the previous pruning pass.
+        if next_idx < curr_len:
+            self._next_msg_index = curr_len
+            append_observation = observations.append
+            for i in range(next_idx, curr_len):
+                msg = msgs[i]
+                extra = msg.get("extra")
+                if extra:
+                    obs = extra.get("observation")
+                    if obs:
+                        append_observation((msg, obs))
+
+        num_to_prune = len(observations) - n
+        pruned_idx = self._pruned_until_obs_index
+        if pruned_idx < num_to_prune:
+            self._pruned_until_obs_index = num_to_prune
+            placeholder = "(ARIA snapshot pruned; see most recent observation)"
+            for i in range(pruned_idx, num_to_prune):
+                msg, obs = observations[i]
+                aria = obs.get("aria_snapshot")
+                if aria:
+                    content = msg.get("content")
+                    if isinstance(content, list):
+                        for part in content:
+                            if not isinstance(part, dict) or part.get("type") not in ("text", "input_text"):
+                                continue
+                            text = part.get("text", "")
+                            new_text = text.replace(aria, placeholder)
+                            if new_text != text:
+                                part["text"] = new_text
+                    elif isinstance(content, str):
+                        new_text = content.replace(aria, placeholder)
+                        if new_text != content:
+                            msg["content"] = new_text
+                    obs["aria_snapshot"] = ""
 
     def _compact_history(self) -> None:
         """Summarize the running transcript via an LLM call and reset messages to [system, summary].
@@ -337,12 +359,21 @@ def _compact_history(self) -> None:
             extra={"interrupt_type": "HistoryCompactionSummary"},
         )
         self.messages = [system_message, summary_message]
+        # Reset pruning state after history compaction
+        self._observations = []
+        self._next_msg_index = 0
+        self._pruned_until_obs_index = 0
 
     def run(self, task: str = "", **kwargs) -> dict[str, Any]:
         self.extra_template_vars |= {"task": task, **kwargs}
         self.messages = []
         self.n_calls = 0
         self.n_format_errors = 0
+        # Reset pruning state for a new run
+        self._observations = []
+        self._next_msg_index = 0
+        self._pruned_until_obs_index = 0
+
         self.add_messages(
             self.model.format_message(role="system", content=self._render_template(self.config.system_template)),
             self.model.format_message(role="user", content=self._render_template(self.config.instance_template)),

src/webwright/models/openai_model.py

@@ -2,15 +2,13 @@
 
 from __future__ import annotations
 
-from pathlib import Path
 from typing import Any
 
 from webwright.models.base import (
     BaseModel,
     BaseModelConfig,
     OptStr,
     _safe_int,
-    image_part_from_path,
     text_part,
 )
 

tests/unit/test_default_agent_aria_pruning.py

@@ -0,0 +1,96 @@
+from typing import Any
+
+from webwright.agents.default import DefaultAgent
+
+
+class StubModel:
+    def __init__(self) -> None:
+        self.queries: list[list[dict[str, Any]]] = []
+
+    def format_message(self, role: str, content: str, **kwargs: Any) -> dict[str, Any]:
+        return {"role": role, "content": content, "extra": kwargs.get("extra", {})}
+
+    def get_template_vars(self) -> dict[str, Any]:
+        return {}
+
+    def query(self, messages: list[dict[str, Any]]) -> dict[str, Any]:
+        self.queries.append(messages)
+        return self.format_message(role="assistant", content="summary")
+
+
+class StubEnv:
+    def get_template_vars(self) -> dict[str, Any]:
+        return {}
+
+
+def make_agent() -> DefaultAgent:
+    return DefaultAgent(
+        StubModel(),
+        StubEnv(),
+        system_template="system",
+        instance_template="instance",
+        keep_last_n_observations=1,
+    )
+
+
+def observation_message(label: str, aria_snapshot: str, content: Any | None = None) -> dict[str, Any]:
+    if content is None:
+        content = f"Observation {label}\n{aria_snapshot}"
+    return {
+        "role": "user",
+        "content": content,
+        "extra": {
+            "observation": {
+                "label": label,
+                "aria_snapshot": aria_snapshot,
+            }
+        },
+    }
+
+
+def test_prunes_old_observation_aria_snapshots_only() -> None:
+    agent = make_agent()
+    first = observation_message("first", "ARIA_FIRST")
+    middle = {"role": "assistant", "content": "no observation", "extra": {}}
+    second = observation_message(
+        "second",
+        "ARIA_SECOND",
+        content=[{"type": "text", "text": "Observation second\nARIA_SECOND"}],
+    )
+    third = observation_message("third", "ARIA_THIRD")
+
+    agent.add_messages(first, middle, second, third)
+    agent._prune_old_observation_aria_snapshots()
+
+    assert first["extra"]["observation"]["aria_snapshot"] == ""
+    assert "ARIA_FIRST" not in first["content"]
+    assert "(ARIA snapshot pruned; see most recent observation)" in first["content"]
+
+    assert middle["content"] == "no observation"
+
+    assert second["extra"]["observation"]["aria_snapshot"] == ""
+    assert "ARIA_SECOND" not in second["content"][0]["text"]
+    assert "(ARIA snapshot pruned; see most recent observation)" in second["content"][0]["text"]
+
+    assert third["extra"]["observation"]["aria_snapshot"] == "ARIA_THIRD"
+    assert "ARIA_THIRD" in third["content"]
+
+
+def test_compaction_resets_incremental_pruning_state() -> None:
+    agent = make_agent()
+    system = agent.model.format_message(role="system", content="system")
+    first = observation_message("first", "ARIA_FIRST")
+    second = observation_message("second", "ARIA_SECOND")
+
+    agent.add_messages(system, first, second)
+    assert first["extra"]["observation"]["aria_snapshot"] == ""
+    assert second["extra"]["observation"]["aria_snapshot"] == "ARIA_SECOND"
+
+    agent._compact_history()
+
+    third = observation_message("third", "ARIA_THIRD")
+    fourth = observation_message("fourth", "ARIA_FOURTH")
+    agent.add_messages(third, fourth)
+
+    assert third["extra"]["observation"]["aria_snapshot"] == ""
+    assert fourth["extra"]["observation"]["aria_snapshot"] == "ARIA_FOURTH"