# 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. """Context compression shared by Flash and Pro. :class:`ScrubEdgeCompressor` maintains two edges over the observation history: a shallow text edge (``xml_scrub_depth``, default 1: as soon as a newer observation exists) where the heavy text blocks and the per-turn ephemeral blocks leave, and the screenshot edge at depth K where the image is resolved in place into its visual summary. A message past both edges is frozen; 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). """ from typing import Any, Callable from langchain_core.messages import BaseMessage from artemis.memory.step_memory import StepMemoryService from artemis.memory.transcript import EPHEMERAL_BLOCKS_KEY from artemis.utils.logger import get_logger logger = get_logger(__name__) #: Header written above a resolved visual summary (public: the chunk capsule #: lens checks it to avoid repeating a summary already present verbatim). HISTORY_SUMMARY_PREFIX = "--- Historical Visual Transition ---\n" _HISTORY_SUMMARY_PREFIX = HISTORY_SUMMARY_PREFIX _UI_LIST_MARKER = "--- UI Element List ---" #: Text of the label block both observation shapes place directly above the #: screenshot. It only makes sense next to an image, so it leaves together #: with the image it labels (the visual summary carries its own header). SCREENSHOT_LABEL = "--- Current Screenshot ---" #: Sentinel: the screenshot stays for now (summary pending inside the grace #: window); the message is revisited for the image swap only. _KEEP_IMAGE = object() 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: """Replace older observation details with summaries at two scrub edges. Depth 1 == the most recent image-bearing message. Messages newer than an edge are byte-identical to what the model was first shown. Each message is edited at most twice and is then frozen: 1. Text edge (depth > ``xml_scrub_depth``, default 1, i.e. as soon as a newer observation exists), one pass, exactly once per message: the heavy blocks behind ``strip_markers`` (the UI element list; for Pro also the plan recitation) are cut at their marker with the text before the marker kept, and every content block whose index is listed under :data:`~artemis.memory.transcript.EPHEMERAL_BLOCKS_KEY` in the message's ``additional_kwargs`` (reminders, hints, per-turn notices) is removed. Only the live observation therefore ever carries an indexed element list: a stale ``[n]`` index from an older screen can never be picked up as a target. 2. Screenshot edge (depth >= ``image_scrub_depth``), IN PLACE: the ready visual summary block takes the image block's position and the :data:`SCREENSHOT_LABEL` block above it is removed. A failed summary becomes an unavailable placeholder immediately; a summary still pending keeps the image (and its label) for up to ``pending_grace_steps`` further depths, after which a pending placeholder replaces it. Both placeholders carry the DataEngine step reference. An image with no summary job at all is dropped silently. 3. Frozen: after the image is resolved 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. The Pro transcript ledger uses this 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 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 stays aligned with the Flash semantics where the live observation is *inside* the list. """ 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 | None = 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) # Text edge: the heavy text blocks and ephemeral blocks leave once a # message is deeper than this (None keeps the default of 1). self._xml_scrub_depth = max(1, int(xml_scrub_depth)) if xml_scrub_depth is not None else 1 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() # Messages whose one-time text edits (strip + ephemeral deletion) are # done; a message inside the pending-grace window sits here unfrozen. self._text_scrubbed: set[int] = set() # Strip-marker carriers without an image (e.g. a step whose screenshot # failed); they are stripped when the edge passes their position. self._xml_candidates: list[int] = [] def compress( self, messages: list[BaseMessage], *, image_scrub_depth: int | None = None ) -> None: """Advance the scrub edge over newly appended messages in-place. ``image_scrub_depth`` overrides the constructor depth for this pass (the transcript ledger passes its occupancy-driven depth). Because a scrubbed message is frozen, a deeper pass after a shallower one never backfills anything: relaxing the depth only lets *new* images live longer, tightening it scrubs the extra depths once. """ self._discover(messages) if not self._tracked and not self._xml_candidates: return self._scrub_edge(messages, image_scrub_depth) # ------------------------------------------------------------------ # 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_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, } ) elif has_xml: self._xml_candidates.append(idx) self._scanned_until = len(messages) def _has_strip_marker(self, text: str) -> bool: return any(marker in text for marker in self._strip_markers) # ------------------------------------------------------------------ # The edges: text strip + ephemeral deletion (shallow), image swap (K) # ------------------------------------------------------------------ def _scrub_edge( self, messages: list[BaseMessage], image_scrub_depth: int | None = None ) -> None: total = len(self._tracked) scrub_depth = max(1, int(image_scrub_depth or self._image_scrub_depth)) grace_limit = scrub_depth + self._pending_grace_steps text_depth = self._xml_scrub_depth 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 self._tail_offset == 0 and pos == total - 1: continue # the live observation is never scrubbed past_text_edge = depth > text_depth past_image_edge = depth >= scrub_depth if not past_text_edge and not past_image_edge: continue # newer than both edges: byte-identical if not past_image_edge: # Past the text edge only: strip once, keep the screenshot. if rec["idx"] not in self._text_scrubbed: self._rewrite( messages[rec["idx"]], rec["idx"], replacement=_KEEP_IMAGE, is_tool=rec["is_tool"], ) continue summary, pending, failed, step_no = self._lookup(rec) replacement: Any 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: replacement = _KEEP_IMAGE # grace: text edits now, image swap later 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, text_edits=past_text_edge) if self._xml_candidates: # Imageless strip-marker carriers (a step whose screenshot failed) # follow the same text edge, counted over every observation — # image-bearing or not — plus the live tail: with the default # depth 1 the list leaves as soon as any newer observation exists, # so only the live observation ever carries an indexed element # list. A carrier that is itself the newest observation stays. observations = sorted({rec["idx"] for rec in self._tracked} | set(self._xml_candidates)) remaining: list[int] = [] for idx in self._xml_candidates: newer = sum(1 for other in observations if other > idx) + self._tail_offset if newer < text_depth: remaining.append(idx) continue if idx not in self._frozen: self._rewrite(messages[idx], idx, replacement=None, is_tool=False) self._frozen.add(idx) self._xml_candidates = remaining 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: Any, *, text_edits: bool = True, ) -> None: """Rewrite one image-bearing message at the screenshot edge; freeze it unless the screenshot is kept for the pending-grace window. With a text edge deeper than the screenshot edge the text edits wait for their own edge (``text_edits=False``); the message is then frozen only once both have happened.""" idx = rec["idx"] self._rewrite( messages[idx], idx, replacement=replacement, is_tool=rec["is_tool"], text_edits=text_edits, ) if not text_edits: return # text edge still ahead: revisit for the strip if replacement is not _KEEP_IMAGE: self._frozen.add(idx) def _rewrite( self, msg: BaseMessage, idx: int, *, replacement: Any, is_tool: bool, text_edits: bool = True, ) -> None: """The single-pass edit of one message (see the class docstring). ``replacement`` is the block that takes the first image's place (a summary or placeholder), ``None`` to drop the image, or :data:`_KEEP_IMAGE` to leave the image and its label in place. The text edits (strip + ephemeral deletion) run only the first time the message is rewritten and consume the ephemeral indices exactly once. Any rewrite that happens *before* that first edit (the image swap when the screenshot edge is shallower than the text edge, or a summary arriving inside the grace window) can remove the label and the image and so shift every later block; the surviving ephemeral indices are then remapped to the new layout in ``additional_kwargs`` so the later text edit still deletes the blocks that were marked, never a neighbour that slid into their old position. """ content = getattr(msg, "content", None) if not isinstance(content, list): self._text_scrubbed.add(idx) return first_edit = idx not in self._text_scrubbed and text_edits ephemeral: set[int] = set() if first_edit: ephemeral = self._ephemeral_indices(msg) keep_image = replacement is _KEEP_IMAGE new_blocks: list[Any] = [] # old position -> new position of every block that survives as itself # (a replacement block takes the image's slot but is not the image). position_map: dict[int, int] = {} image_placed = False for position, block in enumerate(content): if position in ephemeral: continue if isinstance(block, dict): kind = block.get("type") if kind in ("image_url", "image"): if keep_image: position_map[position] = len(new_blocks) new_blocks.append(block) elif replacement is not None and not image_placed: new_blocks.append(replacement) # in place of the image image_placed = True continue if kind == "text": text = block.get("text", "") if not keep_image and text.strip() == SCREENSHOT_LABEL: continue # the label leaves with the image it labelled if first_edit and self._prune_history_xml and self._has_strip_marker(text): retained_text = _without_marked_suffix(text, self._strip_markers) if retained_text: retained_block = dict(block) retained_block["text"] = retained_text position_map[position] = len(new_blocks) new_blocks.append(retained_block) continue position_map[position] = len(new_blocks) new_blocks.append(block) # Defensive fallback: never leave ToolMessage or HumanMessage content empty if not new_blocks: new_blocks = [{"type": "text", "text": "Action dispatched." if is_tool else ""}] position_map = {} msg.content = new_blocks if first_edit: self._text_scrubbed.add(idx) if ephemeral: # Consumed: the indices no longer describe the block layout. msg.additional_kwargs.pop(EPHEMERAL_BLOCKS_KEY, None) else: self._remap_ephemeral(msg, position_map) @staticmethod def _remap_ephemeral(msg: BaseMessage, position_map: dict[int, int]) -> None: """Rewrite the not-yet-consumed ephemeral indices of ``msg`` through ``position_map`` (old block position -> new block position); indices of blocks that did not survive are dropped and the key is removed when nothing survives. Only ``additional_kwargs`` is touched, the content blocks sent to the provider stay exactly as assembled.""" kwargs = getattr(msg, "additional_kwargs", None) if not isinstance(kwargs, dict) or EPHEMERAL_BLOCKS_KEY not in kwargs: return remapped: list[int] = [] for value in kwargs.get(EPHEMERAL_BLOCKS_KEY) or []: try: old_position = int(value) except (TypeError, ValueError): continue new_position = position_map.get(old_position) if new_position is not None and new_position not in remapped: remapped.append(new_position) if remapped: kwargs[EPHEMERAL_BLOCKS_KEY] = remapped else: kwargs.pop(EPHEMERAL_BLOCKS_KEY, None) @staticmethod def _ephemeral_indices(msg: BaseMessage) -> set[int]: kwargs = getattr(msg, "additional_kwargs", None) if not isinstance(kwargs, dict): return set() indices: set[int] = set() for value in kwargs.get(EPHEMERAL_BLOCKS_KEY) or []: try: indices.add(int(value)) except (TypeError, ValueError): continue return indices