android-midscene-automation 0.2.3 → 1.1.0

package/artemis-main/artemis/memory/context_policy.py

@@ -0,0 +1,226 @@
+# 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.
+
+"""Per-agent policies for rendering compiled history.
+
+Each policy sets the visibility options for
+:func:`artemis.utils.task_tree.build_plan_and_history`. When transcript mode
+is enabled, compressed steps are rendered as chunk blocks.
+
+Chunk views:
+
+- ``"full"`` (outputter / history_analyzer): the chunk's stored three-band
+  rendering including the ③ full-width per-step ledger;
+- ``"digest"`` (everything else): bands ① + ② with the ③ ledger replaced by
+  a ``ledger via search_history`` marker line.
+
+Because chunks only exist when ``agent.memory.transcript.enabled`` is on, the
+flag-off output of every agent stays byte-for-byte unchanged.
+
+``agent.memory.policies`` allows per-agent field overrides, e.g.::
+
+    "memory": {"policies": {"planner": {"last_n_detailed": 3}}}
+"""
+
+from dataclasses import dataclass, fields, replace
+from typing import Any
+
+from third_party.mobile_use.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+#: Marker appended to a digest-view chunk block in place of the ③ ledger.
+DIGEST_LEDGER_MARKER_TEMPLATE = (
+    "③ Step action ledger: available via search_history (Steps {start}–{end})"
+)
+
+#: The band-③ heading inside a chunk's stored ``rendered_text``.
+_BAND3_HEADING = "③ Step action ledger"
+
+
+@dataclass(frozen=True)
+class ContextPolicy:
+    """One agent's compiled-history rendering declaration.
+
+    Field semantics map 1:1 onto ``build_plan_and_history`` kwargs, with two
+    conveniences: ``uncompressed=True`` renders every step (the historical
+    ``min_summaries=len(steps)`` idiom) and ``whitelist=True`` marks policies
+    whose caller supplies ``keep_subgoal_hashes`` at runtime.
+    """
+
+    strategy: str = "sliding_window"  # documentation only
+    strict_milestone_pruning: bool = False
+    recent_window_size: int = 3
+    min_summaries: int = 5
+    uncompressed: bool = False
+    last_n_detailed: int = 1
+    all_detailed: bool = False
+    whitelist: bool = False
+    chunk_view: str | None = "digest"  # "full" | "digest" | None (never chunks)
+
+
+#: Default compiled-history settings for each agent.
+CONTEXT_POLICIES: dict[str, ContextPolicy] = {
+    # operator.py:996 — last_n_detailed comes from the node constructor.
+    "operator": ContextPolicy(
+        strategy="strict_milestone",
+        strict_milestone_pruning=True,
+        recent_window_size=3,
+    ),
+    # operator.py:341 — transcript F-region cold-start restore.
+    "operator_cold_start": ContextPolicy(
+        strategy="full",
+        uncompressed=True,
+        last_n_detailed=1,
+        chunk_view=None,
+    ),
+    # planner.py:133 — caller passes an empty plan and hash "default".
+    "planner": ContextPolicy(
+        strategy="strict_milestone",
+        strict_milestone_pruning=True,
+        recent_window_size=5,
+        last_n_detailed=2,
+    ),
+    # outputter.py:81
+    "outputter": ContextPolicy(
+        strategy="full",
+        uncompressed=True,
+        last_n_detailed=0,
+        chunk_view="full",
+    ),
+    # history_analyzer.py:113
+    "history_analyzer": ContextPolicy(
+        strategy="full",
+        uncompressed=True,
+        last_n_detailed=1,
+        chunk_view="full",
+    ),
+    # diagnoser.py:166 — every non-whitelist kwarg is the historical default.
+    "diagnoser": ContextPolicy(strategy="milestone_whitelist", whitelist=True),
+    # committee_tool.py:189
+    "committee": ContextPolicy(
+        strategy="milestone_whitelist",
+        whitelist=True,
+        last_n_detailed=3,
+    ),
+}
+
+_POLICY_FIELDS = {f.name for f in fields(ContextPolicy)}
+
+
+def resolve_policy(agent: str) -> ContextPolicy:
+    """The agent's policy with any ``agent.memory.policies`` overrides applied."""
+    try:
+        policy = CONTEXT_POLICIES[agent]
+    except KeyError:
+        raise KeyError(
+            f"No ContextPolicy declared for agent '{agent}';"
+            f" known agents: {sorted(CONTEXT_POLICIES)}"
+        ) from None
+
+    overrides: dict[str, Any] = {}
+    try:
+        from artemis.config import load_agent_config
+
+        configured = (load_agent_config().memory.policies or {}).get(agent) or {}
+        overrides = {k: v for k, v in configured.items() if k in _POLICY_FIELDS}
+        unknown = set(configured) - _POLICY_FIELDS
+        if unknown:
+            logger.warning(
+                f"agent.memory.policies['{agent}'] ignores unknown fields: {sorted(unknown)}"
+            )
+    except Exception:
+        overrides = {}
+    return replace(policy, **overrides) if overrides else policy
+
+
+def _chunk_digest_text(rendered_text: str, start: int, end: int) -> str:
+    """Bands ① + ② of a stored chunk rendering, ③ replaced by the recall marker."""
+    marker = DIGEST_LEDGER_MARKER_TEMPLATE.format(start=start, end=end)
+    head, sep, _ = rendered_text.partition(f"{_BAND3_HEADING}\n")
+    if not sep:
+        return f"{rendered_text.rstrip()}\n{marker}"
+    return f"{head.rstrip()}\n\n{marker}"
+
+
+def load_chunk_blocks(engine: Any, view: str | None) -> list[dict] | None:
+    """The persisted history chunks rendered for the compiled view, or None.
+
+    Returns None (→ byte-identical legacy rendering) unless the transcript
+    flag is on, the engine exposes chunks, and at least one chunk exists.
+    """
+    if engine is None or view not in ("full", "digest"):
+        return None
+    try:
+        from artemis.config import load_agent_config
+
+        if not load_agent_config().memory.transcript.enabled:
+            return None
+    except Exception:
+        return None
+    try:
+        rows = engine.get_history_chunks() or []
+    except Exception as e:
+        logger.debug(f"load_chunk_blocks: no chunks available: {e}")
+        return None
+
+    blocks: list[dict] = []
+    for row in rows:
+        start = getattr(row, "start_step_number", None)
+        end = getattr(row, "end_step_number", None)
+        rendered = getattr(row, "rendered_text", None)
+        if not isinstance(start, int) or not isinstance(end, int) or not rendered:
+            continue
+        text = rendered if view == "full" else _chunk_digest_text(rendered, start, end)
+        blocks.append({"start_step_number": start, "end_step_number": end, "text": text})
+    return blocks or None
+
+
+def build_history_for(
+    agent: str,
+    task_plan: str,
+    steps: list,
+    current_subgoal_hash: str,
+    *,
+    keep_subgoal_hashes: set | None = None,
+    last_n_detailed: int | None = None,
+    engine: Any = None,
+) -> str:
+    """Render the agent's compiled plan-and-history view per its policy.
+
+    Args:
+        agent: Key into :data:`CONTEXT_POLICIES`.
+        keep_subgoal_hashes: Runtime whitelist (whitelist policies only).
+        last_n_detailed: Runtime override (the operator's constructor knob).
+        engine: DataEngine handing chunks to the compiled view; omit at call
+            sites that must never receive chunk blocks.
+    """
+    from artemis.utils.task_tree import build_plan_and_history
+
+    policy = resolve_policy(agent)
+    kwargs: dict[str, Any] = {
+        "min_summaries": len(steps) if policy.uncompressed else policy.min_summaries,
+        "last_n_detailed": (policy.last_n_detailed if last_n_detailed is None else last_n_detailed),
+        "all_detailed": policy.all_detailed,
+        "strict_milestone_pruning": policy.strict_milestone_pruning,
+        "recent_window_size": policy.recent_window_size,
+    }
+    if policy.whitelist:
+        kwargs["keep_subgoal_hashes"] = keep_subgoal_hashes
+
+    chunks = load_chunk_blocks(engine, policy.chunk_view)
+    if chunks:
+        kwargs["chunks"] = chunks
+
+    return build_plan_and_history(task_plan, steps, current_subgoal_hash, **kwargs)

package/artemis-main/artemis/memory/step_capsule.md

@@ -0,0 +1,73 @@
+# ROLE
+You are the Segment Capsule writer for a mobile automation agent's history compression. You receive one contiguous segment of executed steps and produce the two LLM-authored bands of the segment's history chunk:
+
+- Band ① — Synopsis & effects: what this segment was doing, what was actually done, and what effects/artifacts it left behind, plus structured fields.
+- Band ② — Compressed step summary: an interval narrative ("Steps xx–xx did ..., Step xx did ...") covering every step of the segment.
+
+The mechanical per-step action ledger (band ③) is assembled outside of you; do NOT reproduce it.
+
+# INPUT SHAPE
+The segment arrives turn by turn, in order. Each turn block has:
+- `Recorded steps:` — the mechanical facts of the step(s) executed in that turn: step number, session offset, the exact action, the controller/validator result, the visual transition summary, the notes written (tool → note key: gist), and any user-injected instruction.
+- `Transcript as seen by the operator during this turn:` — the turn exactly as it stood in the operator's own context: the observation text, the operator's reasoning (including `(thinking)` blocks), every `[tool call]` with its arguments, every `[tool result]` returned to the operator (explorer answers, recalled history, note reads, ADB output, screenshot placeholders), and the `--- Action Execution Result ---` message that closes the turn (every turn on the Pro path; only when an action failed on the Flash path). Your capsule REPLACES this transcript in the operator's context — everything the operator learned from it that still matters must survive in your output.
+
+A turn block without a transcript falls back to a reasoning excerpt; treat it the same way with less evidence.
+
+# NEUTRAL-WORDING CONTRACT (STRICT)
+1. **Zero subjective validation**: never declare semantic success, completion, or final failure. Subjective verdicts in compressed history cause future turns to falsely believe a goal was definitively met or permanently broken.
+   - **BANNED WORDS**: "successfully", "completed", "achieved", "entered", "navigated to", "arrived at", "failed", "unsuccessful", "could not", "impossible".
+   - Describe observable transitions and recorded results instead (e.g. "the settings page was displayed", "the controller reported an execution error").
+   - Tool results in the transcript may themselves use verdict words; restate them as recorded results ("the explorer reported ...", "the tool returned ...").
+2. **First-person perspective**: write from the agent's perspective using "I". Never "the agent" / "the operator".
+3. **Planned iterations are not anomalies**: systematic processing of different candidate items, or periodic polling cycles under a `[Loop]` / `[Loop:continuous]` milestone, is planned progress — record it objectively (e.g. "polling round #2"). Only flag a loop when the exact same target is repeated without progress or states oscillate blindly; then record the pattern and count factually.
+4. **Preserve, don't abstract**: be deliberately detailed and stay close to the source wording. Prefer longer output over abstracting facts beyond recognition. Faithfully carry visible text, codes, prices, account names, file names, error banners — and the concrete content of tool results the operator relied on (an explorer's answer, a recalled step, a note's contents, a command's output).
+5. **Fact priority**: user instructions > validator/controller recorded results > tool results returned to the operator (explorer / recall / notes / ADB) > UI/OCR deltas > visual transition summaries > the operator's reasoning.
+
+# BAND ① REQUIREMENTS
+Band ① is the segment-level account; band ② below it is the interval-level account. The two bands must not say the same thing twice: ① never re-lists what the intervals of ② already narrate, and it should come out shorter than ② (the `doing` + `did` + `effect` prose together — not counting the structured fields — shorter than the ② intervals together).
+
+Answer three questions in prose (field `doing` / `did` / `effect`):
+- `doing`: what this segment was trying to do and why (strategic context within the plan). One or two sentences.
+- `did`: the route the segment actually took, at segment level: the plan-relevant outcome, plus every detour, retry, recovery, pivot, and tool consultation (what I asked, what it returned) when that shaped the next action. Do NOT walk through the steps one by one — ② does that. A segment that went straight through gets one sentence.
+- `effect`: what the segment left behind beyond the screen state already in `exit_state`: persistent state changes, artifacts, and notes. **This MUST include every note left during the segment**: for each notes write in the input, name the target note file and the gist of what was recorded. Every note's target key is machine-checked against your effect field — an omitted note key forces regeneration. If nothing was left behind besides the screen state, say so in a few words.
+
+Structured fields (arrays of short strings, may be empty; each entry one short clause, not a sentence):
+- `verified_facts`: states/values actually observed on screen or returned by a tool that later steps may need (name the source: "on screen", "explorer", "recall", "adb"). Values, names, versions, counts, codes, prices — not a description of every screen visited.
+- `unresolved`: open questions, unverified assumptions, missing prerequisites.
+- `failed_paths`: approaches attempted in this segment whose recorded results were errors (state the recorded error factually).
+- `important_entities`: accounts, files, order/tracking numbers, package names, key UI entry points.
+- `entry_state` / `exit_state` (single strings): the observable screen/app state at segment start and end, a short clause each.
+
+# BAND ② REQUIREMENTS
+Field `intervals`: an ordered array of `{"start_step": N, "end_step": M, "text": "..."}` objects.
+- Band ② answers only "what were these steps doing, and what did I see as a result". The action itself (tap / swipe / input, its target and its coordinates) is ALREADY carried step by step by the mechanical ledger (band ③) that sits right below your output — do not write it twice.
+  - **NEVER include coordinates** — no `[x, y]` pairs, no pixel or normalized positions, no bounds. This is machine-checked; any `[number, number]` in an interval text forces regeneration.
+  - **NEVER narrate step by step** ("Step 5: I tapped X. Step 6: I tapped Y."). One interval is a behavior plus the observed outcome, e.g. "Opened the 'Network & internet' section and returned to the main list; the section showed Wi‑Fi, mobile network and VPN rows."
+  - **Merge homogeneous consecutive actions into one interval** (repeated scrolls, a sweep through several sections, a polling loop, retries of the same target) — describe the run and its count/result once. A step that does something of a different nature gets its own single-step interval (`start_step == end_step`).
+- Every line's text must reference concrete behavior for those steps (the step numbers come from the `Recorded steps:` lines): the target's visible text, the screen or state observed afterwards, values read, errors shown.
+- **The union of intervals MUST cover the segment's full step range** — no gaps, no overlaps, in ascending order. This is machine-checked; a gap forces regeneration.
+- Keep the neutral-wording contract in every line.
+
+# LENGTH TARGET (SOFT)
+Bands ① and ② together should stay within roughly one third of the source text you were given. Reach that by not repeating what band ③ already records (actions, targets, coordinates) and by merging homogeneous runs — NEVER by abstracting away concrete facts (visible text, values read, tool answers, error banners, note contents must all survive verbatim).
+
+# OUTPUT FORMAT
+Return ONLY one JSON object (no markdown fences, no commentary):
+
+```json
+{
+  "doing": "...",
+  "did": "...",
+  "effect": "...",
+  "entry_state": "...",
+  "exit_state": "...",
+  "verified_facts": ["..."],
+  "unresolved": ["..."],
+  "failed_paths": ["..."],
+  "important_entities": ["..."],
+  "intervals": [
+    {"start_step": 18, "end_step": 20, "text": "..."},
+    {"start_step": 21, "end_step": 21, "text": "..."}
+  ]
+}
+```

package/artemis-main/artemis/memory/step_memory.py

@@ -0,0 +1,281 @@
+# 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.
+
+"""Shared background step-summary runtime (StepMemoryService).
+
+Hosts the scheduling skeleton lifted from the Flash VisualStepSummarizer:
+zero-blocking dispatch, bounded retry, bounded flush, canonical step_id keying
+with alias resolution (tool_call_id / legacy ordinals), and status queries.
+
+What a summary *contains* (the lens semantics) lives in subclasses — the
+service only guarantees how jobs are scheduled, retried, keyed, and drained.
+Flash and Pro profiles share this runtime.
+"""
+
+import asyncio
+from typing import Any
+
+from artemis.data_engine.trace import detached_trace
+from third_party.mobile_use.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+JobKey = int | str
+
+
+class StepLens:
+    """Formal lens interface: what one summary attempt *produces*.
+
+    A lens turns one job payload into summary text; the surrounding service
+    owns scheduling, retries, keying, and draining. Implementations must be
+    side-effect-tolerant under retries (an attempt may run again after a
+    failure) and return ``None``/raise to signal a failed attempt.
+
+    Two lens altitudes exist in the redesign (§5): the step-level visual
+    transition lens (:class:`~artemis.agents.flash.summarizer.VisualStepSummarizer`
+    keeps it inlined as its ``_attempt`` for compatibility) and the chunk-level
+    :class:`~artemis.memory.chunking.StepCapsuleLens`.
+    """
+
+    name = "lens"
+
+    async def render(self, key: JobKey, payload: dict[str, Any]) -> str | None:
+        """Produce the summary text for one payload, or None on failure."""
+        raise NotImplementedError
+
+
+class StepMemoryService:
+    """Non-blocking background runtime for per-step summary jobs.
+
+    Key design properties:
+    1. Zero-blocking dispatch: callers submit and proceed immediately.
+    2. Bounded retry: each job attempts at most ``1 + retry_limit`` times and
+       then enters an explicit failed state — never an unbounded loop.
+    3. Bounded flush: ``flush()`` waits up to the configured timeout, then
+       cancels whatever is still in flight.
+    4. Canonical keying with aliases: jobs are keyed by the DataEngine step id
+       where available; tool_call_ids (and legacy ordinals) remain queryable
+       through an alias map so message-level consumers need no migration.
+    5. Bounded concurrency: at most ``max_concurrency`` attempts run at once.
+    6. Detached tracing: every job runs inside
+       :func:`~artemis.data_engine.trace.detached_trace`, so a lens's model
+       calls never attach to the trace span (the operator's step) that was
+       current when the job was submitted — no lens needs to know.
+
+    Subclasses implement :meth:`_attempt` (one summarization attempt for a
+    payload) and may override :meth:`_on_status` to persist pending/failed
+    status transitions.
+    """
+
+    #: Trace node label of a lens-less subclass's jobs (``lens:<name>``); a
+    #: configured :class:`StepLens` contributes its own ``name`` instead.
+    lens_name: str | None = None
+
+    def __init__(
+        self,
+        ctx: Any,
+        *,
+        lens: StepLens | None = None,
+        max_concurrency: int = 1,
+        retry_limit: int = 3,
+        flush_timeout_s: float = 30.0,
+    ):
+        self.ctx = ctx
+        self._lens = lens
+        self._summaries: dict[JobKey, str] = {}
+        self._failed: set[JobKey] = set()
+        self._pending_tasks: dict[JobKey, asyncio.Task] = {}
+        self._step_inputs: dict[JobKey, dict[str, Any]] = {}
+        self._retry_counts: dict[JobKey, int] = {}
+        self._retry_delays = (0.0, 0.5, 1.0, 2.0, 3.0)
+        self._retry_limit = max(0, retry_limit)
+        self._flush_timeout_s = max(0.0, flush_timeout_s)
+        self._aliases: dict[JobKey, JobKey] = {}
+        self._semaphore = asyncio.Semaphore(max(1, max_concurrency))
+
+    # ------------------------------------------------------------------
+    # Keying
+    # ------------------------------------------------------------------
+
+    def resolve_key(self, key: JobKey) -> JobKey:
+        """Resolve an alias (tool_call_id / legacy ordinal) to its canonical key."""
+        return self._aliases.get(key, key)
+
+    def submit(
+        self,
+        key: JobKey,
+        payload: dict[str, Any],
+        aliases: tuple[JobKey, ...] = (),
+    ) -> None:
+        """Submit one summary job without blocking the caller.
+
+        Args:
+            key: Canonical job key (DataEngine step id where available).
+            payload: Lens-specific input data; stored until the job resolves.
+            aliases: Additional keys (e.g. tool_call_id) that resolve to ``key``.
+        """
+        for alias in aliases:
+            if alias is not None and alias != key:
+                self._aliases[alias] = key
+
+        self._step_inputs[key] = payload
+        self._failed.discard(key)
+        self._on_status(key, "pending")
+
+        task = asyncio.create_task(self._run_until_ready(key))
+        self._pending_tasks[key] = task
+
+    # ------------------------------------------------------------------
+    # Scheduling core
+    # ------------------------------------------------------------------
+
+    @property
+    def trace_node_name(self) -> str:
+        """Node label the job's traces and usage records carry (``lens:<name>``)."""
+        name = getattr(self._lens, "name", None) or self.lens_name or type(self).__name__.lower()
+        return f"lens:{name}"
+
+    async def _run_until_ready(self, key: JobKey) -> None:
+        """Run one job to completion, detached from the submitter's trace span.
+
+        The task was created while the submitting agent's step span was
+        current and inherits that context; the whole job — every attempt,
+        retry sleeps included — runs under :func:`detached_trace` so nothing
+        a lens records (model calls, their thinking/raw_thinking children,
+        stream deltas, usage) can be attributed to the submitter's step.
+        """
+        with detached_trace(self.trace_node_name):
+            await self._run_attempts(key)
+
+    async def _run_attempts(self, key: JobKey) -> None:
+        """Retry one job independently, bounded by the configured retry limit.
+
+        Attempts at most ``1 + retry_limit`` times; on exhaustion the job
+        enters an explicit failed state (surfaced via :meth:`has_failed` and
+        the ``_on_status`` hook). Retry backoff sleeps happen outside the
+        concurrency semaphore so a waiting job never starves its peers.
+        """
+        while key not in self._summaries:
+            async with self._semaphore:
+                succeeded = await self._attempt(key)
+            if succeeded:
+                return
+
+            retry_count = self._retry_counts.get(key, 0) + 1
+            self._retry_counts[key] = retry_count
+            payload = self._step_inputs.get(key)
+            display_step = payload.get("step_number") if payload else key
+
+            if retry_count > self._retry_limit:
+                self._failed.add(key)
+                logger.warning(
+                    f"StepMemoryService: Giving up on Step {display_step} after"
+                    f" {retry_count} failed attempts (retry_limit={self._retry_limit})."
+                )
+                self._on_status(key, "failed")
+                return
+
+            delay = self._retry_delays[min(retry_count - 1, len(self._retry_delays) - 1)]
+            logger.info(
+                f"StepMemoryService: Retrying Step {display_step} "
+                f"(attempt {retry_count + 1}) after {delay:.1f}s"
+            )
+            if delay:
+                await asyncio.sleep(delay)
+
+    async def _attempt(self, key: JobKey) -> bool:
+        """Run one summarization attempt for ``key``; True when a summary landed.
+
+        The default implementation delegates to the configured :class:`StepLens`;
+        subclasses without a lens override this method directly (the Flash
+        visual summarizer keeps its inlined attempt for compatibility).
+        """
+        if self._lens is None:
+            raise NotImplementedError
+        payload = self._step_inputs.get(key)
+        if payload is None:
+            return False
+        try:
+            summary = await self._lens.render(key, payload)
+        except Exception as e:
+            logger.warning(f"StepMemoryService: lens '{self._lens.name}' attempt failed: {e}")
+            return False
+        if not summary:
+            return False
+        self._summaries[key] = summary
+        self._on_ready(key, summary)
+        return True
+
+    def _on_ready(self, key: JobKey, summary: str) -> None:
+        """Hook invoked when a lens-produced summary lands; default is a no-op."""
+
+    def _on_status(self, key: JobKey, status: str) -> None:
+        """Hook invoked on pending/failed transitions; default is a no-op."""
+
+    # ------------------------------------------------------------------
+    # Status queries (alias-resolving)
+    # ------------------------------------------------------------------
+
+    def get_summary(self, key: JobKey, fallback_text: str | None = None) -> str | None:
+        """Retrieve the summary for a job key (or alias), or fallback text."""
+        return self._summaries.get(self.resolve_key(key), fallback_text)
+
+    def has_summary(self, key: JobKey) -> bool:
+        """Check whether a job's summary has completed."""
+        return self.resolve_key(key) in self._summaries
+
+    def has_job(self, key: JobKey) -> bool:
+        """Check whether a job has been submitted for summarization."""
+        return self.resolve_key(key) in self._step_inputs
+
+    def has_failed(self, key: JobKey) -> bool:
+        """Check whether a job permanently failed (retries exhausted)."""
+        return self.resolve_key(key) in self._failed
+
+    def is_pending(self, key: JobKey) -> bool:
+        """Check whether a job still needs a summary."""
+        return self.has_job(key) and not self.has_summary(key)
+
+    def get_job_payload(self, key: JobKey) -> dict[str, Any] | None:
+        """Return the stored payload for a job key (or alias), if any."""
+        return self._step_inputs.get(self.resolve_key(key))
+
+    def get_step_number(self, key: JobKey) -> int | None:
+        """Return the human-facing step ordinal recorded for a job, if known."""
+        payload = self.get_job_payload(key)
+        if payload is None:
+            return None
+        step_number = payload.get("step_number")
+        return step_number if isinstance(step_number, int) else None
+
+    # ------------------------------------------------------------------
+    # Draining
+    # ------------------------------------------------------------------
+
+    async def flush(self, timeout_seconds: float | None = None) -> None:
+        """Wait for active jobs up to a bound, then cancel remaining retry loops."""
+        if timeout_seconds is None:
+            timeout_seconds = self._flush_timeout_s
+        tasks = [t for t in self._pending_tasks.values() if not t.done()]
+        if tasks:
+            logger.info(f"StepMemoryService: Flushing {len(tasks)} pending summary tasks...")
+            _, pending = await asyncio.wait(tasks, timeout=timeout_seconds)
+            if pending:
+                logger.warning(
+                    f"StepMemoryService: Cancelling {len(pending)} unfinished summary tasks "
+                    f"after {timeout_seconds:.1f}s flush timeout."
+                )
+                for task in pending:
+                    task.cancel()
+                await asyncio.gather(*pending, return_exceptions=True)