somew1nd
5d92939d91
Preserve user guidance through history compression, stabilize Flash index targets, and retire withdrawn checks without applying stale verdicts. Improve runtime diagnostics, tool responses, and history rendering. Simplify documentation and comments. Validation: 2076 tests passed in the full suite; 381 affected tests and 26 frontend tests passed after cleanup. Ruff and the quality ratchet passed.
282 lines
12 KiB
Python
282 lines
12 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.
|
|
|
|
"""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 artemis.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)
|