somew1nd
67248c887e
Add execution incidents, action bursts and shared Flash history. Expose verification and Explorer settings with streamed check results. Support graceful cancellation, recording recovery and bundled UI resources. Add CI quality checks and clean up comments and documentation.
326 lines
14 KiB
Python
326 lines
14 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Universal Context Compressor for Artemis Flash profile.
|
|
|
|
:class:`ScrubEdgeCompressor` implements the scrub-edge (擦洗沿) discipline
|
|
from the history-module redesign §3.2: each message is mutated a bounded
|
|
number of times near the tail (XML stripped at depth 1, screenshot resolved
|
|
at depth K), then frozen forever. A late summary never backfills a frozen
|
|
message. FlashRunner uses this implementation; the Pro transcript ledger
|
|
reuses it over its active region (see the class docstring).
|
|
|
|
The legacy per-turn full-rescan ``compress_flash_messages`` was removed in M5
|
|
(2026-09-01) after serving as the M1/M2 semantic reference — its surviving
|
|
semantics are asserted directly against ``ScrubEdgeCompressor`` in
|
|
``tests/unit/agents/test_flash_scrub_edge.py``.
|
|
"""
|
|
|
|
from typing import Any, Callable
|
|
|
|
from langchain_core.messages import BaseMessage
|
|
|
|
from artemis.memory.step_memory import StepMemoryService
|
|
from artemis.utils.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
_HISTORY_SUMMARY_PREFIX = "--- Historical Visual Transition ---\n"
|
|
_UI_LIST_MARKER = "--- UI Element List ---"
|
|
|
|
|
|
def _without_marked_suffix(text: str, markers: tuple[str, ...] = (_UI_LIST_MARKER,)) -> str:
|
|
"""Remove a marked heavy block without discarding text before its marker."""
|
|
marker_index = -1
|
|
for marker in markers:
|
|
idx = text.find(marker)
|
|
if idx >= 0 and (marker_index < 0 or idx < marker_index):
|
|
marker_index = idx
|
|
if marker_index < 0:
|
|
return text
|
|
return text[:marker_index].rstrip()
|
|
|
|
|
|
class ScrubEdgeCompressor:
|
|
"""Scrub-edge context compressor (history-module redesign §3.2, M1).
|
|
|
|
Mutation discipline — every message is touched at most twice, always near
|
|
the tail, and is then frozen forever:
|
|
|
|
1. Depth 1 (XML): once a newer observation exists, historical UI Element
|
|
lists are stripped (same semantics as the legacy ``prune_history_xml``).
|
|
2. Depth K (image, ``image_scrub_depth``): the K-th most recent
|
|
image-bearing message has its screenshot resolved — replaced by the
|
|
ready visual summary; if the summary is still pending the image is
|
|
retained for up to ``pending_grace_steps`` further image-depths, after
|
|
which it is replaced by a pending placeholder. A failed summary is
|
|
replaced by an unavailable placeholder immediately. Both placeholders
|
|
carry the DataEngine step reference.
|
|
3. Frozen: after resolution the message is never read or written again —
|
|
a late summary never backfills a frozen message.
|
|
|
|
Contract: the message list is append-only (FlashRunner never removes or
|
|
reorders entries), so bookkeeping is index-based and each pass only
|
|
touches the few messages near the scrub edge instead of rescanning the
|
|
whole history.
|
|
|
|
Profile generalization (M2): the Pro transcript ledger reuses this exact
|
|
discipline over its active region. Its observation messages carry no
|
|
``tool_call_id``, so ``summary_key_getter`` supplies the summary-job key
|
|
(the DataEngine step id) per message; ``strip_markers`` extends the
|
|
depth-1 strip to the Pro tail blocks (UI element list + plan recitation);
|
|
and ``tail_offset=1`` accounts for the live observation living outside
|
|
the tracked list (the ledger renders it as a separate tail), so depth
|
|
arithmetic and the keep-window stay aligned with the Flash semantics
|
|
where the live observation is *inside* the list. Defaults leave the
|
|
Flash behavior byte-identical.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
summarizer: StepMemoryService | None = None,
|
|
*,
|
|
prune_history_xml: bool = True,
|
|
image_scrub_depth: int = 3,
|
|
pending_grace_steps: int = 3,
|
|
xml_scrub_depth: int = 1,
|
|
summary_key_getter: Callable[[BaseMessage], str | None] | None = None,
|
|
strip_markers: tuple[str, ...] = (_UI_LIST_MARKER,),
|
|
tail_offset: int = 0,
|
|
):
|
|
self._summarizer = summarizer
|
|
self._prune_history_xml = prune_history_xml
|
|
self._image_scrub_depth = max(1, image_scrub_depth)
|
|
self._pending_grace_steps = max(0, pending_grace_steps)
|
|
self._xml_scrub_depth = max(1, xml_scrub_depth)
|
|
self._summary_key_getter = summary_key_getter
|
|
self._strip_markers = tuple(strip_markers) or (_UI_LIST_MARKER,)
|
|
self._tail_offset = max(0, tail_offset)
|
|
|
|
self._scanned_until = 0
|
|
self._tool_msg_count = 0
|
|
# Append-only records of image-bearing messages, in message order:
|
|
# {"idx", "key" (str tool_call_id | None), "legacy" (ordinal | None), "is_tool"}
|
|
self._tracked: list[dict[str, Any]] = []
|
|
self._frozen: set[int] = set()
|
|
self._xml_candidates: list[int] = []
|
|
|
|
def compress(self, messages: list[BaseMessage]) -> None:
|
|
"""Advance the scrub edge over newly appended messages in-place."""
|
|
self._discover(messages)
|
|
if not self._tracked:
|
|
return
|
|
self._scrub_xml(messages)
|
|
self._scrub_images(messages)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Discovery (incremental; frozen prefix is never rescanned)
|
|
# ------------------------------------------------------------------
|
|
|
|
def _discover(self, messages: list[BaseMessage]) -> None:
|
|
for idx in range(self._scanned_until, len(messages)):
|
|
msg = messages[idx]
|
|
tool_call_id = getattr(msg, "tool_call_id", None)
|
|
if tool_call_id is not None:
|
|
self._tool_msg_count += 1
|
|
content = getattr(msg, "content", None)
|
|
if not isinstance(content, list):
|
|
continue
|
|
|
|
has_image = False
|
|
has_xml = False
|
|
for block in content:
|
|
if not isinstance(block, dict):
|
|
continue
|
|
if block.get("type") in ("image_url", "image"):
|
|
has_image = True
|
|
elif block.get("type") == "text" and self._has_strip_marker(block.get("text", "")):
|
|
has_xml = True
|
|
|
|
if has_xml:
|
|
self._xml_candidates.append(idx)
|
|
if has_image:
|
|
key = str(tool_call_id) if tool_call_id is not None else None
|
|
if key is None and self._summary_key_getter is not None:
|
|
try:
|
|
getter_key = self._summary_key_getter(msg)
|
|
except Exception:
|
|
getter_key = None
|
|
if getter_key is not None:
|
|
key = str(getter_key)
|
|
self._tracked.append(
|
|
{
|
|
"idx": idx,
|
|
"key": key,
|
|
"legacy": self._tool_msg_count if tool_call_id is not None else None,
|
|
"is_tool": tool_call_id is not None,
|
|
}
|
|
)
|
|
self._scanned_until = len(messages)
|
|
|
|
def _has_strip_marker(self, text: str) -> bool:
|
|
return any(marker in text for marker in self._strip_markers)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Depth-1 XML scrub
|
|
# ------------------------------------------------------------------
|
|
|
|
def _scrub_xml(self, messages: list[BaseMessage]) -> None:
|
|
if not self._prune_history_xml:
|
|
return
|
|
# Keep the UI list only in the newest xml_scrub_depth observation
|
|
# messages; everything older (or non-image XML carriers) is stripped.
|
|
# With a live tail outside the tracked list (tail_offset > 0) the tail
|
|
# itself occupies the newest slots of the keep window.
|
|
keep_count = max(0, self._xml_scrub_depth - self._tail_offset)
|
|
keep_indices = {rec["idx"] for rec in self._tracked[-keep_count:]} if keep_count else set()
|
|
still_pending: list[int] = []
|
|
for idx in self._xml_candidates:
|
|
if idx in keep_indices:
|
|
still_pending.append(idx)
|
|
continue
|
|
if idx in self._frozen:
|
|
continue
|
|
self._strip_xml_blocks(messages[idx])
|
|
self._xml_candidates = still_pending
|
|
|
|
def _strip_xml_blocks(self, msg: BaseMessage) -> None:
|
|
content = getattr(msg, "content", None)
|
|
if not isinstance(content, list):
|
|
return
|
|
new_blocks: list[Any] = []
|
|
for block in content:
|
|
if (
|
|
isinstance(block, dict)
|
|
and block.get("type") == "text"
|
|
and self._has_strip_marker(block.get("text", ""))
|
|
):
|
|
retained_text = _without_marked_suffix(block.get("text", ""), self._strip_markers)
|
|
if retained_text:
|
|
retained_block = dict(block)
|
|
retained_block["text"] = retained_text
|
|
new_blocks.append(retained_block)
|
|
continue
|
|
new_blocks.append(block)
|
|
msg.content = new_blocks
|
|
|
|
# ------------------------------------------------------------------
|
|
# Depth-K image scrub
|
|
# ------------------------------------------------------------------
|
|
|
|
def _scrub_images(self, messages: list[BaseMessage]) -> None:
|
|
total = len(self._tracked)
|
|
scrub_depth = self._image_scrub_depth
|
|
grace_limit = scrub_depth + self._pending_grace_steps
|
|
|
|
for pos, rec in enumerate(self._tracked):
|
|
if rec["idx"] in self._frozen:
|
|
continue
|
|
# 1 == most recent image-bearing message; a live tail outside the
|
|
# tracked list (tail_offset > 0) counts as the newest depths.
|
|
depth = total - pos + self._tail_offset
|
|
if depth < scrub_depth:
|
|
continue # active window
|
|
if self._tail_offset == 0 and pos == total - 1:
|
|
continue # the live observation is never scrubbed
|
|
|
|
summary, pending, failed, step_no = self._lookup(rec)
|
|
if summary is not None:
|
|
replacement = {
|
|
"type": "text",
|
|
"text": f"{_HISTORY_SUMMARY_PREFIX}{summary}",
|
|
}
|
|
elif failed:
|
|
replacement = {
|
|
"type": "text",
|
|
"text": f"[visual summary unavailable; evidence at DataEngine step {step_no}]",
|
|
}
|
|
elif pending and depth <= grace_limit:
|
|
continue # grace period: retain the image, revisit next pass
|
|
elif pending:
|
|
replacement = {
|
|
"type": "text",
|
|
"text": f"[visual summary pending; evidence at DataEngine step {step_no}]",
|
|
}
|
|
else:
|
|
replacement = None # no summary job exists; drop the image silently
|
|
|
|
self._resolve(messages, rec, replacement)
|
|
|
|
def _lookup(self, rec: dict[str, Any]) -> tuple[str | None, bool, bool, int | None]:
|
|
"""Mirror the legacy keying: tool_call_id first, ordinal fallback."""
|
|
summarizer = self._summarizer
|
|
if summarizer is None:
|
|
return None, False, False, None
|
|
|
|
effective_key: Any = rec["key"]
|
|
if effective_key is None:
|
|
return None, False, False, None
|
|
|
|
summary = summarizer.get_summary(effective_key)
|
|
if summary is None and not summarizer.has_job(effective_key) and rec["legacy"] is not None:
|
|
effective_key = rec["legacy"]
|
|
summary = summarizer.get_summary(effective_key)
|
|
|
|
failed = summarizer.has_failed(effective_key)
|
|
pending = summarizer.is_pending(effective_key) and not failed
|
|
step_no = summarizer.get_step_number(effective_key)
|
|
if step_no is None:
|
|
step_no = rec["legacy"]
|
|
return summary, pending, failed, step_no
|
|
|
|
def _resolve(
|
|
self,
|
|
messages: list[BaseMessage],
|
|
rec: dict[str, Any],
|
|
replacement: dict[str, Any] | None,
|
|
) -> None:
|
|
"""Replace the message's image blocks and freeze it permanently."""
|
|
msg = messages[rec["idx"]]
|
|
content = getattr(msg, "content", None)
|
|
new_blocks: list[Any] = []
|
|
if isinstance(content, list):
|
|
for block in content:
|
|
if isinstance(block, dict) and block.get("type") in ("image_url", "image"):
|
|
continue
|
|
if (
|
|
self._prune_history_xml
|
|
and isinstance(block, dict)
|
|
and block.get("type") == "text"
|
|
and self._has_strip_marker(block.get("text", ""))
|
|
):
|
|
# Freezing implies fully scrubbed: no stale UI list survives
|
|
# even when xml_scrub_depth exceeds the image depth.
|
|
retained_text = _without_marked_suffix(
|
|
block.get("text", ""), self._strip_markers
|
|
)
|
|
if retained_text:
|
|
retained_block = dict(block)
|
|
retained_block["text"] = retained_text
|
|
new_blocks.append(retained_block)
|
|
continue
|
|
new_blocks.append(block)
|
|
|
|
if replacement is not None:
|
|
new_blocks.append(replacement)
|
|
|
|
# Defensive fallback: never leave ToolMessage or HumanMessage content empty
|
|
if not new_blocks:
|
|
new_blocks = [{"type": "text", "text": "Action completed." if rec["is_tool"] else ""}]
|
|
|
|
msg.content = new_blocks
|
|
self._frozen.add(rec["idx"])
|
|
if rec["idx"] in self._xml_candidates:
|
|
self._xml_candidates.remove(rec["idx"])
|