From c79c6eb27bfc7c69d0b06f3ef88048e9af571ce4 Mon Sep 17 00:00:00 2001 From: somew1nd Date: Thu, 3 Sep 2026 20:48:23 -0700 Subject: [PATCH] feat: show history compression progress and fix action trace overlays --- .../src/app/utils/action-formatter.util.ts | 33 +--- .../src/app/utils/image-overlay.util.spec.ts | 85 +++++++++ .../src/app/utils/image-overlay.util.ts | 40 ++++- .../src/app/utils/tool-formatter.util.spec.ts | 55 ++++++ .../src/app/utils/tool-formatter.util.ts | 55 ++++++ artemis/memory/chunking.py | 169 +++++++++++++++++- tests/unit/memory/test_history_chunking.py | 74 ++++++++ 7 files changed, 468 insertions(+), 43 deletions(-) create mode 100644 apps/showcase_ui/src/app/utils/image-overlay.util.spec.ts diff --git a/apps/showcase_ui/src/app/utils/action-formatter.util.ts b/apps/showcase_ui/src/app/utils/action-formatter.util.ts index e6e2b0a..e6748c4 100644 --- a/apps/showcase_ui/src/app/utils/action-formatter.util.ts +++ b/apps/showcase_ui/src/app/utils/action-formatter.util.ts @@ -15,7 +15,7 @@ */ import { ActionParam, StepReplayFrame } from '../core/models/stream.model'; -import { extractNumbersFromCoordinateValue, isPureDirectionString, parseSequenceCoordinates } from './image-overlay.util'; +import { extractNumbersFromCoordinateValue, isPureDirectionString, parseSequenceCoordinates, unwrapTraceAction } from './image-overlay.util'; import { cleanErrorMessage } from './tool-formatter.util'; /** @@ -35,35 +35,6 @@ export function safeParseJson(val: any): any { // WeakMap memo for parsed action objects to avoid repeated JSON.parse and object allocation across change detection cycles const actionObjectCache = new WeakMap(); -function computeActionObject(action: any): any { - if (action && typeof action === 'object' && action.payload) { - const payload = typeof action.payload === 'string' ? safeParseJson(action.payload) : action.payload; - const innerAction = payload?.args?.action; - if (innerAction && typeof innerAction === 'object') { - return { - ...action, - ...innerAction, - name: action.name || innerAction.action || innerAction.name, - action: innerAction.action || action.name, - trace_id: action.trace_id, - timestamp: action.timestamp ?? innerAction.timestamp - }; - } - const innerArgs = payload?.args; - if (innerArgs && typeof innerArgs === 'object') { - return { - ...action, - ...innerArgs, - name: action.name || innerArgs.action || innerArgs.name, - action: innerArgs.action || action.name, - trace_id: action.trace_id, - timestamp: action.timestamp ?? innerArgs.timestamp - }; - } - } - return action; -} - /** * Safely extract action object from possible array or nested structure */ @@ -75,7 +46,7 @@ export function getActionObject(action: any): any { if (typeof action === 'object') { const cached = actionObjectCache.get(action); if (cached !== undefined) return cached; - const res = computeActionObject(action); + const res = unwrapTraceAction(action); actionObjectCache.set(action, res); return res; } diff --git a/apps/showcase_ui/src/app/utils/image-overlay.util.spec.ts b/apps/showcase_ui/src/app/utils/image-overlay.util.spec.ts new file mode 100644 index 0000000..2bc7bb3 --- /dev/null +++ b/apps/showcase_ui/src/app/utils/image-overlay.util.spec.ts @@ -0,0 +1,85 @@ +import { drawActionCoordinatesOnOverlay, unwrapTraceAction } from './image-overlay.util'; +import { getActionCoords, getActionTargetText } from './action-formatter.util'; + +// Trace records as the DataEngine publishes them in `generic_tools`. +const PRO_TAP_TRACE = { + trace_id: 'pro-tap', + type: 'action', + name: 'tap', + timestamp: 1788478482.05, + status: 'success', + payload: { + action: { + action: 'tap', + coordinates: [922, 1710], + coordinate_space: 'pixel', + normalized_coordinates: [854, 705], + target_text: 'YouTube' + }, + success: true, + post_screenshot: null + } +}; + +const FLASH_CLICK_TRACE = { + trace_id: 'flash-click', + type: 'action', + name: 'click', + timestamp: 1788478305.12, + status: 'success', + payload: { + args: { target: [618, 705], target_text: '相册' }, + result: '{"outcome": "Clicked"}' + } +}; + +function fakeImage(): HTMLImageElement { + return { naturalWidth: 1080, naturalHeight: 2424 } as unknown as HTMLImageElement; +} + +describe('unwrapTraceAction', () => { + it('flattens a Pro action trace (payload.action)', () => { + const act = unwrapTraceAction(PRO_TAP_TRACE); + expect(act.action).toBe('tap'); + expect(act.normalized_coordinates).toEqual([854, 705]); + expect(act.target_text).toBe('YouTube'); + expect(act.trace_id).toBe('pro-tap'); + expect(act.timestamp).toBe(1788478482.05); + }); + + it('flattens a Flash action trace (payload.args)', () => { + const act = unwrapTraceAction(FLASH_CLICK_TRACE); + expect(act.action).toBe('click'); + expect(act.target).toEqual([618, 705]); + expect(act.args.target).toEqual([618, 705]); + }); + + it('passes plain action objects through untouched', () => { + const plain = { action: 'tap', normalized_coordinates: [10, 20] }; + expect(unwrapTraceAction(plain)).toBe(plain); + }); +}); + +describe('drawActionCoordinatesOnOverlay', () => { + it('draws the tap marker for a Pro action trace', () => { + const overlay = document.createElement('div'); + drawActionCoordinatesOnOverlay(fakeImage(), overlay, PRO_TAP_TRACE); + const marker = overlay.querySelector('.action-point-marker') as HTMLElement; + expect(marker).not.toBeNull(); + expect(marker.style.left).toBe('85.4%'); + expect(marker.style.top).toBe('70.5%'); + }); + + it('draws the tap marker for a Flash action trace', () => { + const overlay = document.createElement('div'); + drawActionCoordinatesOnOverlay(fakeImage(), overlay, FLASH_CLICK_TRACE); + expect(overlay.querySelector('.action-point-marker')).not.toBeNull(); + }); +}); + +describe('action card details from trace records', () => { + it('reads coordinates and target text from a Pro action trace', () => { + expect(getActionCoords(PRO_TAP_TRACE)).toContain('854'); + expect(getActionTargetText(PRO_TAP_TRACE)).toBe('YouTube'); + }); +}); diff --git a/apps/showcase_ui/src/app/utils/image-overlay.util.ts b/apps/showcase_ui/src/app/utils/image-overlay.util.ts index 4c75c98..5611fab 100644 --- a/apps/showcase_ui/src/app/utils/image-overlay.util.ts +++ b/apps/showcase_ui/src/app/utils/image-overlay.util.ts @@ -14,6 +14,39 @@ * limitations under the License. */ +/** + * Unwrap Flash (`payload.args`) and Pro (`payload.action`) traces. + * Preserve trace metadata and leave plain action objects unchanged. + */ +export function unwrapTraceAction(record: any): any { + if (!record || typeof record !== 'object' || !record.payload) return record; + + let payload = record.payload; + if (typeof payload === 'string') { + try { + payload = JSON.parse(payload); + } catch { + return record; + } + } + if (!payload || typeof payload !== 'object') return record; + + const asObject = (v: any) => (Array.isArray(v) ? v[0] : v); + const candidates = [asObject(payload.args?.action), asObject(payload.action), asObject(payload.args)]; + const inner = candidates.find(c => c && typeof c === 'object'); + if (!inner) return record; + + return { + ...record, + ...inner, + args: (record.args && typeof record.args === 'object') ? record.args : inner, + name: record.name || inner.action || inner.name, + action: inner.action || record.name, + trace_id: record.trace_id, + timestamp: record.timestamp ?? inner.timestamp + }; +} + /** * Helper to parse any coordinate representation into an array of numbers. * Handles: @@ -141,12 +174,11 @@ export function drawActionCoordinatesOnOverlay( const naturalWidth = img.naturalWidth || 1080; const naturalHeight = img.naturalHeight || 2400; - actions.forEach(act => { + actions.forEach(record => { + const act = unwrapTraceAction(record); if (!act || typeof act !== 'object') return; const actionName = (act.action || act.name || 'tap').toLowerCase(); - - const payload = (act.payload && typeof act.payload === 'object') ? act.payload : {}; - const args = (payload.args && typeof payload.args === 'object') ? payload.args : (act.args || {}); + const args = (act.args && typeof act.args === 'object') ? act.args : {}; // Filter out non-touch / system actions that do not operate via touch coordinates const nonDrawingActions = [ diff --git a/apps/showcase_ui/src/app/utils/tool-formatter.util.spec.ts b/apps/showcase_ui/src/app/utils/tool-formatter.util.spec.ts index c831e29..298abdc 100644 --- a/apps/showcase_ui/src/app/utils/tool-formatter.util.spec.ts +++ b/apps/showcase_ui/src/app/utils/tool-formatter.util.spec.ts @@ -1,10 +1,65 @@ import { cleanErrorMessage, + getCompressionLabel, + getToolDisplayLabel, + getToolIcon, getUniqueGenericTools, getVideoAnalysisView, shouldShowTool } from './tool-formatter.util'; +describe('compress_history timeline line', () => { + const trace = (status: string, args: Record) => ({ + type: 'tool', + name: 'compress_history', + status, + payload: { args } + }); + + it('is shown as a plain tool line with its own icon', () => { + const running = trace('running', { start_step: 12, end_step: 27 }); + expect(shouldShowTool(running)).toBeTrue(); + expect(getToolIcon(running)).toBe('compress'); + expect(getToolDisplayLabel(running)).toBe(getCompressionLabel(running)); + }); + + it('says which steps are being condensed while running', () => { + expect(getCompressionLabel(trace('running', { start_step: 12, end_step: 27 }))) + .toBe('Condensing steps 12–27 into a short memory to free up room…'); + expect(getCompressionLabel(trace('running', { start_step: 12, end_step: 27, note: 'retrying' }))) + .toBe('Retrying the memory summary for steps 12–27…'); + expect(getCompressionLabel(trace('running', { start_step: 5, end_step: 5 }))) + .toBe('Condensing step 5 into a short memory to free up room…'); + }); + + it('reports the size reduction and the working memory once done', () => { + const done = trace('success', { + start_step: 12, + end_step: 27, + source_tokens: 8400, + summary_tokens: 600, + context_tokens: 54000, + context_budget: 80000 + }); + expect(getCompressionLabel(done)).toBe( + 'Steps 12–27 condensed into a short memory · 8.4k → 600 tokens (14× smaller) · working memory ≈ 54k of 80k tokens' + ); + }); + + it('omits the working memory figure on lines that do not carry it', () => { + const done = trace('success', { start_step: 1, end_step: 4, source_tokens: 3000, summary_tokens: 1500 }); + expect(getCompressionLabel(done)).toBe('Steps 1–4 condensed into a short memory · 3k → 1.5k tokens (2× smaller)'); + }); + + it('explains a forced snapshot and a failed attempt in plain words', () => { + const forced = trace('success', { start_step: 1, end_step: 4, forced: true, context_tokens: 70000, context_budget: 80000 }); + expect(getCompressionLabel(forced)) + .toBe('Steps 1–4 replaced by a brief snapshot (memory was nearly full) · working memory ≈ 70k of 80k tokens'); + expect(getCompressionLabel(trace('failed', { start_step: 1, end_step: 4 }))) + .toBe("Couldn't condense steps 1–4 yet; keeping the full record and retrying later"); + }); +}); + describe('cleanErrorMessage', () => { it('removes repeated LLM wrapper labels while preserving the provider reason', () => { expect(cleanErrorMessage('LLM Error: LLM Request Error: 503 model overloaded')) diff --git a/apps/showcase_ui/src/app/utils/tool-formatter.util.ts b/apps/showcase_ui/src/app/utils/tool-formatter.util.ts index 11ef45e..f625126 100644 --- a/apps/showcase_ui/src/app/utils/tool-formatter.util.ts +++ b/apps/showcase_ui/src/app/utils/tool-formatter.util.ts @@ -699,12 +699,65 @@ export function getToolDisplayLabel(tool: any, isFirstSaveNote: boolean = false) } case 'read_url': return 'Fetching web page'; + case 'compress_history': + return getCompressionLabel(tool); default: return `Executing ${cleanName.split('_').map((w: string) => w.charAt(0).toUpperCase() + w.slice(1)).join(' ')}`; } } +function formatTokenFigure(value: any): string { + const n = Number(value); + if (!Number.isFinite(n) || n < 0) return ''; + if (n < 1000) return `${Math.round(n)}`; + return `${(n / 1000).toFixed(1).replace(/\.0$/, '')}k`; +} + +/** Format compression progress and token usage for the timeline. */ +export function getCompressionLabel(tool: any): string { + const args = getToolArgs(tool) || {}; + const start = Number(args.start_step); + const end = Number(args.end_step); + const hasRange = Number.isFinite(start) && Number.isFinite(end) && start > 0 && end > 0; + const range = hasRange + ? (start === end ? `step ${start}` : `steps ${start}–${end}`) + : 'earlier steps'; + const rangeCapitalized = range.charAt(0).toUpperCase() + range.slice(1); + const status = String(tool?.status || '').toLowerCase(); + + if (status === 'failed') { + return `Couldn't condense ${range} yet; keeping the full record and retrying later`; + } + if (status !== 'success') { + return String(args.note || '').toLowerCase() === 'retrying' + ? `Retrying the memory summary for ${range}…` + : `Condensing ${range} into a short memory to free up room…`; + } + + const parts: string[] = []; + const source = Number(args.source_tokens); + const summary = Number(args.summary_tokens); + if (args.forced) { + parts.push(`${rangeCapitalized} replaced by a brief snapshot (memory was nearly full)`); + } else { + parts.push(`${rangeCapitalized} condensed into a short memory`); + if (source > 0 && summary > 0) { + const factor = source / summary; + const factorText = factor >= 2 ? ` (${Math.round(factor)}× smaller)` : ''; + parts.push(`${formatTokenFigure(source)} → ${formatTokenFigure(summary)} tokens${factorText}`); + } + } + const context = Number(args.context_tokens); + const budget = Number(args.context_budget); + if (context > 0) { + parts.push(budget > 0 + ? `working memory ≈ ${formatTokenFigure(context)} of ${formatTokenFigure(budget)} tokens` + : `working memory ≈ ${formatTokenFigure(context)} tokens`); + } + return parts.join(' · '); +} + /** * Get Material Symbol icon for tools */ @@ -772,6 +825,8 @@ export function getToolIcon(tool: any): string { case 'outputter': case 'output_synthesis': return 'assignment_turned_in'; + case 'compress_history': + return 'compress'; default: return 'build'; } diff --git a/artemis/memory/chunking.py b/artemis/memory/chunking.py index 3c47efe..8bc1b97 100644 --- a/artemis/memory/chunking.py +++ b/artemis/memory/chunking.py @@ -33,7 +33,9 @@ from dataclasses import dataclass, field import json from pathlib import Path import re +import time from typing import Any, Callable +from uuid import uuid4 from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage @@ -55,6 +57,9 @@ CHUNK_PENDING_NOTE = ( _NOTE_TOOLS = ("save_note", "update_note", "append_note") +# Reuse one tool trace per chunk across running, success, and failed states. +COMPRESSION_TRACE_NAME = "compress_history" + # --------------------------------------------------------------------------- # Band ③ — mechanical per-step action ledger @@ -528,6 +533,15 @@ class ChunkCapsuleService(StepMemoryService): def __init__(self, ctx: Any, lens: StepCapsuleLens, **kwargs): super().__init__(ctx, lens=lens, **kwargs) + # Report exhausted retries without waiting for the next render. + self.failure_hook: Callable[[JobKey], None] | None = None + + def _on_status(self, key: JobKey, status: str) -> None: + if status == "failed" and self.failure_hook is not None: + try: + self.failure_hook(key) + except Exception as e: + logger.warning(f"Capsule failure hook for {key} raised: {e}") def _on_ready(self, key: JobKey, summary: str) -> None: # The turn transcripts only serve generation. A ready capsule is never @@ -570,6 +584,15 @@ class ChunkState: band1: dict[str, Any] = field(default_factory=dict) band2: str | None = None annotations: list[dict[str, Any]] = field(default_factory=list) + # Segment closure reason: milestone, size, or pressure. + trigger: str | None = None + trace_id: Any = None + trace_step_id: Any = None + announced_at: float | None = None + # Measure source at close and summary at swap; estimate tokens as chars // 4. + # Forced swaps without a capsule leave summary_chars unset. + source_chars: int = 0 + summary_chars: int | None = None @property def capsule_key(self) -> str: @@ -881,6 +904,8 @@ class HistoryChunkManager: self._meter_getter = meter_getter self._capsule_service = capsule_service or self._build_capsule_service(ctx) + if isinstance(self._capsule_service, ChunkCapsuleService): + self._capsule_service.failure_hook = self._on_capsule_failed # Trigger state (fed from the graph). self._step_hashes: dict[str, str] = {} @@ -1053,7 +1078,7 @@ class HistoryChunkManager: # bounded to trigger/pressure renders, never every render. self._redispatch_failed_capsules() self._harvest_capsules() - self._swap_ready_segments(ledger, hard=hard) + self._swap_ready_segments(ledger, hard=hard, base_tokens=base_tokens) def _close_new_segments(self, ledger, soft: bool) -> bool: """Trigger evaluation: close due segments and dispatch their capsules. @@ -1082,7 +1107,7 @@ class HistoryChunkManager: # (the previous segment becomes one HistoryChunk; the floor only # delays the event, it never splits the segment). The trailing segment # is subject to the size/soft triggers over its eligible portion only. - selected: list[tuple[str | None, list[dict]]] = [] + selected: list[tuple[str | None, list[dict], str]] = [] open_hash = self._last_hash consumed_prefix = 0 @@ -1092,7 +1117,7 @@ class HistoryChunkManager: is_last = idx == len(segments) - 1 closed = (not is_last) or seg_hash != open_hash if closed and consumed_prefix + len(seg_turns) <= remaining_eligible: - selected.append(segment) + selected.append((seg_hash, seg_turns, "milestone")) consumed_prefix += len(seg_turns) continue tail_segment = segment @@ -1113,19 +1138,21 @@ class HistoryChunkManager: ) if size_event: - selected.append((tail_segment[0], tail_portion)) + selected.append((tail_segment[0], tail_portion, "size")) elif not milestone_event and soft and tail_portion: # Soft threshold with nothing else due: close the oldest open # segment's eligible portion (bounded by the chunk size cap). - selected.append((tail_segment[0], tail_portion[: self._max_steps])) + selected.append((tail_segment[0], tail_portion[: self._max_steps], "pressure")) if not selected: return False steps_by_id = self._load_steps_by_id() - for seg_hash, seg_turns in selected: + for seg_hash, seg_turns, trigger in selected: for slice_turns in self._slices(seg_turns): - chunk = self._create_chunk(slice_turns, seg_hash, steps_by_id, ledger) + chunk = self._create_chunk( + slice_turns, seg_hash, steps_by_id, ledger, trigger=trigger + ) self._awaiting.append({"chunk": chunk, "turns": list(slice_turns)}) logger.info( f"History segments closed (ready-gated): {len(self._awaiting)} awaiting" @@ -1150,10 +1177,11 @@ class HistoryChunkManager: " retained until a capsule lands." ) self._capsule_service.submit(key, payload) + self._announce(chunk, "running", note="retrying") except Exception as e: logger.error(f"Capsule re-dispatch for {key} failed: {e}") - def _swap_ready_segments(self, ledger, *, hard: bool) -> None: + def _swap_ready_segments(self, ledger, *, hard: bool, base_tokens: int | None = None) -> None: """Swap the ready prefix of awaiting segments into the frozen region. Swaps consume the transcript's oldest unchunked turns, so only a @@ -1161,6 +1189,9 @@ class HistoryChunkManager: (and every younger segment's) original text live. The hard threshold is the sole exception: everything closed force-swaps, pending chunks included, and the frozen region renders as the L3 snapshot. + + ``base_tokens`` is the prompt size before the swap, used to estimate + the remaining context on the last swapped chunk's trace. """ if not self._awaiting: return @@ -1188,6 +1219,51 @@ class HistoryChunkManager: f" ({len(self._chunks)} chunks, {len(self._eras)} eras," f" {len(self._awaiting)} still awaiting, hard={hard})." ) + # Pending chunks contribute only their minimal index to the L3 snapshot. + # Include that size in the context estimate, without a compression ratio. + swapped = [e for e in swap if e["chunk"] is not None] + source_tokens_total = 0 + replacement_tokens_total = 0 + for entry in swapped: + chunk = entry["chunk"] + if chunk.status == "ready": + chunk.summary_chars = len(render_chunk_block(chunk)) + replacement_chars = chunk.summary_chars + else: + chunk.summary_chars = None + replacement_chars = len(chunk.minimal_index or "") + source_tokens_total += chunk.source_chars // 4 + replacement_tokens_total += replacement_chars // 4 + context_tokens: int | None = None + if base_tokens is not None: + context_tokens = max( + 0, int(base_tokens) - source_tokens_total + replacement_tokens_total + ) + + for index, entry in enumerate(swapped): + chunk = entry["chunk"] + turns = len(entry["turns"]) + forced = hard and chunk.status != "ready" + if not forced: + result = f"Replaced {turns} turn{'s' if turns != 1 else ''} with the summary." + else: + result = ( + "Context budget reached before the summary was ready;" + f" replaced {turns} turn{'s' if turns != 1 else ''} with a brief" + " snapshot." + ) + extra: dict[str, Any] = { + "source_tokens": chunk.source_chars // 4, + "summary_tokens": ( + chunk.summary_chars // 4 if chunk.summary_chars is not None else None + ), + } + # Show the context estimate once per swap, on its last trace. + if context_tokens is not None and index == len(swapped) - 1: + extra["context_tokens"] = context_tokens + extra["context_budget"] = self._budget + extra["context_estimated"] = True + self._announce(chunk, "success", result=result, forced=forced, extra=extra) def _partition(self, eligible: list[dict]) -> list[tuple[str | None, list[dict]]]: """Split eligible turns into consecutive same-hash segments.""" @@ -1254,6 +1330,8 @@ class HistoryChunkManager: seg_hash: str | None, steps_by_id: dict[str, dict], ledger: Any = None, + *, + trigger: str | None = None, ) -> ChunkState | None: # A turn carries every step id it recorded (``step_keys``, Flash # multi-action turns); older ledgers only expose ``step_key``. Each @@ -1302,14 +1380,89 @@ class HistoryChunkManager: band3=build_action_ledger(steps, session_start), minimal_index=build_action_ledger(steps, session_start, minimal=True), user_lines=[line for line in (injected_instruction_line(s) for s in steps) if line], + trigger=trigger, ) # Ready gating: the caller queues the chunk as awaiting — it only # enters self._chunks (and the frozen region) once its capsule is # ready, or through the hard-threshold emergency swap. self._persist(chunk) self._dispatch_capsule(chunk, steps, session_start, turn_records) + # Attach the trace to the step allocated before this prompt render. + chunk.trace_id = uuid4() + chunk.trace_step_id = getattr(self._engine, "current_step_id", None) + chunk.announced_at = time.time() + if ledger is not None: + try: + chunk.source_chars = int(ledger.turn_text_chars(slice_turns)) + except Exception as e: + logger.debug(f"Chunk source size unavailable for {chunk.capsule_key}: {e}") + self._announce(chunk, "running", extra={"source_tokens": chunk.source_chars // 4}) return chunk + def _announce( + self, + chunk: ChunkState, + status: str, + *, + result: str | None = None, + error: str | None = None, + note: str | None = None, + forced: bool = False, + extra: dict[str, Any] | None = None, + ) -> None: + """Update the chunk's timeline trace without interrupting compression.""" + if chunk.trace_id is None: + return + engine = self._engine + if engine is None or not hasattr(engine, "record_trace"): + return + args: dict[str, Any] = { + "start_step": chunk.start_step_number, + "end_step": chunk.end_step_number, + "steps": len(chunk.source_step_ids), + "milestone": chunk.milestone_label, + "trigger": chunk.trigger, + } + if note: + args["note"] = note + if forced: + args["forced"] = True + if extra: + args.update(extra) + payload: dict[str, Any] = {"args": args} + if result is not None: + payload["result"] = result + if error is not None: + payload["error"] = error + duration = None + if status != "running" and chunk.announced_at is not None: + duration = max(0.0, time.time() - chunk.announced_at) + try: + engine.record_trace( + type="tool", + name=COMPRESSION_TRACE_NAME, + payload=payload, + step_id=chunk.trace_step_id, + status=status, + duration=duration, + trace_id=chunk.trace_id, + ) + except Exception as e: + logger.debug(f"Compression announcement skipped for {chunk.capsule_key}: {e}") + + def _on_capsule_failed(self, key: JobKey) -> None: + """Mark the trace failed; keep the chunk pending for a later retry.""" + for chunk in self._all_chunks(): + if chunk.status != "pending" or chunk.capsule_key != key: + continue + self._announce( + chunk, + "failed", + error="Summary attempts exhausted; the full record is kept and" + " the summary will be retried later.", + ) + return + @staticmethod def _step_payload(step: dict, session_start: float | None) -> dict[str, Any]: """Mechanical per-step facts handed to the capsule lens.""" diff --git a/tests/unit/memory/test_history_chunking.py b/tests/unit/memory/test_history_chunking.py index 3318a1a..61cc4e8 100644 --- a/tests/unit/memory/test_history_chunking.py +++ b/tests/unit/memory/test_history_chunking.py @@ -78,6 +78,7 @@ class FakeEngine: self.session_start_time = SESSION_START self.current_session_id = "session-1" self.chunk_writes: list[dict] = [] + self.traces: list[dict] = [] if base_dir is not None: self.base_dir = base_dir @@ -88,6 +89,10 @@ class FakeEngine: self.chunk_writes.append(kwargs) return f"chunk-row-{len(self.chunk_writes)}" + def record_trace(self, **kwargs): + self.traces.append(kwargs) + return kwargs.get("trace_id") + class StubCapsuleService(StepMemoryService): """Records capsule submissions without spawning asyncio tasks (sync tests). @@ -1097,6 +1102,75 @@ def test_hard_threshold_force_swaps_pending_chunks_into_l3(): assert len(ledger.unchunked_turns()) < 8 # turns actually consumed +# --------------------------------------------------------------------------- +# Timeline announcement (compress_history trace) +# --------------------------------------------------------------------------- + + +def test_compression_announces_one_tool_line_from_dispatch_to_swap(): + """Dispatch and swap share a trace ID and report the change in token usage.""" + steps = [_step(i, "hash-a" if i <= 4 else "hash-b") for i in range(1, 9)] + meter = {"value": 60_000} + ledger, chunker, engine, capsule = _make(steps, min_active=2, meter=lambda: meter["value"]) + _run_turns(ledger, chunker, 1, 8, _hashes(4)) + + assert len(engine.traces) == 1 + running = engine.traces[0] + assert (running["type"], running["name"], running["status"]) == ( + "tool", + "compress_history", + "running", + ) + args = running["payload"]["args"] + assert (args["start_step"], args["end_step"], args["steps"]) == (1, 4, 4) + assert args["trigger"] == "milestone" + assert args["source_tokens"] > 0 + assert "context_tokens" not in args + + capsule.resolve("chunk:1-4", _capsule(1, 4)) + ledger.commit_staged(step_key="s8", validator_result={"status": "success"}) + ledger.render([_observation(9)]) + + done = engine.traces[-1] + assert done["trace_id"] == running["trace_id"] + assert done["status"] == "success" + assert done["duration"] is not None and done["duration"] >= 0 + done_args = done["payload"]["args"] + assert done_args["source_tokens"] == args["source_tokens"] + assert done_args["summary_tokens"] > 0 + assert done_args["context_budget"] == 100_000 + assert done_args["context_estimated"] is True + assert done_args["context_tokens"] == ( + 60_000 - done_args["source_tokens"] + done_args["summary_tokens"] + ) + assert "forced" not in done_args + + +def test_forced_swap_announces_snapshot_without_a_ratio(): + """Forced swaps report estimated context usage without a summary size.""" + steps = [_step(i, "hash-a" if i <= 4 else "hash-b") for i in range(1, 9)] + meter = {"value": None} + ledger, chunker, engine, _ = _make(steps, min_active=2, meter=lambda: meter["value"]) + _run_turns(ledger, chunker, 1, 8, _hashes(4)) + assert any(t["status"] == "running" for t in engine.traces) + + meter["value"] = 95_000 + ledger.commit_staged(step_key="s8", validator_result={"status": "success"}) + ledger.render([_observation(9)]) # pressure closes the open tail too, then force-swaps all + + running_ids = {t["trace_id"] for t in engine.traces if t["status"] == "running"} + done = [t for t in engine.traces if t["status"] == "success"] + assert done and {t["trace_id"] for t in done} <= running_ids + for trace in done: + assert trace["payload"]["args"]["forced"] is True + assert trace["payload"]["args"]["summary_tokens"] is None + last_args = done[-1]["payload"]["args"] + assert last_args["context_budget"] == 100_000 + assert 0 < last_args["context_tokens"] < 95_000 + # The prompt-size figure appears once per swap event, on the last line. + assert all("context_tokens" not in t["payload"]["args"] for t in done[:-1]) + + # --------------------------------------------------------------------------- # Band ① note-linkage machine check # ---------------------------------------------------------------------------