feat: show history compression progress and fix action trace overlays

This commit is contained in:
somew1nd committed 2026-09-03 20:48:23 -07:00
1 parent ab03e1d2bb
commit c79c6eb27b
7 files changed
+468 -43

No files matched your search

@@ -15,7 +15,7 @@
*/ */
import { ActionParam, StepReplayFrame } from '../core/models/stream.model'; 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'; 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 // WeakMap memo for parsed action objects to avoid repeated JSON.parse and object allocation across change detection cycles
const actionObjectCache = new WeakMap<object, any>(); const actionObjectCache = new WeakMap<object, any>();
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 * Safely extract action object from possible array or nested structure
*/ */
@@ -75,7 +46,7 @@ export function getActionObject(action: any): any {
if (typeof action === 'object') { if (typeof action === 'object') {
const cached = actionObjectCache.get(action); const cached = actionObjectCache.get(action);
if (cached !== undefined) return cached; if (cached !== undefined) return cached;
const res = computeActionObject(action); const res = unwrapTraceAction(action);
actionObjectCache.set(action, res); actionObjectCache.set(action, res);
return res; return res;
} }
@@ -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');
});
});
@@ -14,6 +14,39 @@
* limitations under the License. * 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. * Helper to parse any coordinate representation into an array of numbers.
* Handles: * Handles:
@@ -141,12 +174,11 @@ export function drawActionCoordinatesOnOverlay(
const naturalWidth = img.naturalWidth || 1080; const naturalWidth = img.naturalWidth || 1080;
const naturalHeight = img.naturalHeight || 2400; const naturalHeight = img.naturalHeight || 2400;
actions.forEach(act => { actions.forEach(record => {
const act = unwrapTraceAction(record);
if (!act || typeof act !== 'object') return; if (!act || typeof act !== 'object') return;
const actionName = (act.action || act.name || 'tap').toLowerCase(); const actionName = (act.action || act.name || 'tap').toLowerCase();
const args = (act.args && typeof act.args === 'object') ? act.args : {};
const payload = (act.payload && typeof act.payload === 'object') ? act.payload : {};
const args = (payload.args && typeof payload.args === 'object') ? payload.args : (act.args || {});
// Filter out non-touch / system actions that do not operate via touch coordinates // Filter out non-touch / system actions that do not operate via touch coordinates
const nonDrawingActions = [ const nonDrawingActions = [
@@ -1,10 +1,65 @@
import { import {
cleanErrorMessage, cleanErrorMessage,
getCompressionLabel,
getToolDisplayLabel,
getToolIcon,
getUniqueGenericTools, getUniqueGenericTools,
getVideoAnalysisView, getVideoAnalysisView,
shouldShowTool shouldShowTool
} from './tool-formatter.util'; } from './tool-formatter.util';
describe('compress_history timeline line', () => {
const trace = (status: string, args: Record<string, any>) => ({
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', () => { describe('cleanErrorMessage', () => {
it('removes repeated LLM wrapper labels while preserving the provider reason', () => { it('removes repeated LLM wrapper labels while preserving the provider reason', () => {
expect(cleanErrorMessage('LLM Error: LLM Request Error: 503 model overloaded')) expect(cleanErrorMessage('LLM Error: LLM Request Error: 503 model overloaded'))
@@ -699,12 +699,65 @@ export function getToolDisplayLabel(tool: any, isFirstSaveNote: boolean = false)
} }
case 'read_url': case 'read_url':
return 'Fetching web page'; return 'Fetching web page';
case 'compress_history':
return getCompressionLabel(tool);
default: default:
return `Executing ${cleanName.split('_').map((w: string) => w.charAt(0).toUpperCase() + w.slice(1)).join(' ')}`; 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 * Get Material Symbol icon for tools
*/ */
@@ -772,6 +825,8 @@ export function getToolIcon(tool: any): string {
case 'outputter': case 'outputter':
case 'output_synthesis': case 'output_synthesis':
return 'assignment_turned_in'; return 'assignment_turned_in';
case 'compress_history':
return 'compress';
default: default:
return 'build'; return 'build';
} }
+161 -8
View File
@@ -33,7 +33,9 @@ from dataclasses import dataclass, field
import json import json
from pathlib import Path from pathlib import Path
import re import re
import time
from typing import Any, Callable from typing import Any, Callable
from uuid import uuid4
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
@@ -55,6 +57,9 @@ CHUNK_PENDING_NOTE = (
_NOTE_TOOLS = ("save_note", "update_note", "append_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 # Band ③ — mechanical per-step action ledger
@@ -528,6 +533,15 @@ class ChunkCapsuleService(StepMemoryService):
def __init__(self, ctx: Any, lens: StepCapsuleLens, **kwargs): def __init__(self, ctx: Any, lens: StepCapsuleLens, **kwargs):
super().__init__(ctx, lens=lens, **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: def _on_ready(self, key: JobKey, summary: str) -> None:
# The turn transcripts only serve generation. A ready capsule is never # 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) band1: dict[str, Any] = field(default_factory=dict)
band2: str | None = None band2: str | None = None
annotations: list[dict[str, Any]] = field(default_factory=list) 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 @property
def capsule_key(self) -> str: def capsule_key(self) -> str:
@@ -881,6 +904,8 @@ class HistoryChunkManager:
self._meter_getter = meter_getter self._meter_getter = meter_getter
self._capsule_service = capsule_service or self._build_capsule_service(ctx) 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). # Trigger state (fed from the graph).
self._step_hashes: dict[str, str] = {} self._step_hashes: dict[str, str] = {}
@@ -1053,7 +1078,7 @@ class HistoryChunkManager:
# bounded to trigger/pressure renders, never every render. # bounded to trigger/pressure renders, never every render.
self._redispatch_failed_capsules() self._redispatch_failed_capsules()
self._harvest_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: def _close_new_segments(self, ledger, soft: bool) -> bool:
"""Trigger evaluation: close due segments and dispatch their capsules. """Trigger evaluation: close due segments and dispatch their capsules.
@@ -1082,7 +1107,7 @@ class HistoryChunkManager:
# (the previous segment becomes one HistoryChunk; the floor only # (the previous segment becomes one HistoryChunk; the floor only
# delays the event, it never splits the segment). The trailing segment # delays the event, it never splits the segment). The trailing segment
# is subject to the size/soft triggers over its eligible portion only. # 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 open_hash = self._last_hash
consumed_prefix = 0 consumed_prefix = 0
@@ -1092,7 +1117,7 @@ class HistoryChunkManager:
is_last = idx == len(segments) - 1 is_last = idx == len(segments) - 1
closed = (not is_last) or seg_hash != open_hash closed = (not is_last) or seg_hash != open_hash
if closed and consumed_prefix + len(seg_turns) <= remaining_eligible: 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) consumed_prefix += len(seg_turns)
continue continue
tail_segment = segment tail_segment = segment
@@ -1113,19 +1138,21 @@ class HistoryChunkManager:
) )
if size_event: 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: elif not milestone_event and soft and tail_portion:
# Soft threshold with nothing else due: close the oldest open # Soft threshold with nothing else due: close the oldest open
# segment's eligible portion (bounded by the chunk size cap). # 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: if not selected:
return False return False
steps_by_id = self._load_steps_by_id() 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): 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)}) self._awaiting.append({"chunk": chunk, "turns": list(slice_turns)})
logger.info( logger.info(
f"History segments closed (ready-gated): {len(self._awaiting)} awaiting" f"History segments closed (ready-gated): {len(self._awaiting)} awaiting"
@@ -1150,10 +1177,11 @@ class HistoryChunkManager:
" retained until a capsule lands." " retained until a capsule lands."
) )
self._capsule_service.submit(key, payload) self._capsule_service.submit(key, payload)
self._announce(chunk, "running", note="retrying")
except Exception as e: except Exception as e:
logger.error(f"Capsule re-dispatch for {key} failed: {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. """Swap the ready prefix of awaiting segments into the frozen region.
Swaps consume the transcript's oldest unchunked turns, so only a 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 (and every younger segment's) original text live. The hard threshold
is the sole exception: everything closed force-swaps, pending chunks is the sole exception: everything closed force-swaps, pending chunks
included, and the frozen region renders as the L3 snapshot. 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: if not self._awaiting:
return return
@@ -1188,6 +1219,51 @@ class HistoryChunkManager:
f" ({len(self._chunks)} chunks, {len(self._eras)} eras," f" ({len(self._chunks)} chunks, {len(self._eras)} eras,"
f" {len(self._awaiting)} still awaiting, hard={hard})." 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]]]: def _partition(self, eligible: list[dict]) -> list[tuple[str | None, list[dict]]]:
"""Split eligible turns into consecutive same-hash segments.""" """Split eligible turns into consecutive same-hash segments."""
@@ -1254,6 +1330,8 @@ class HistoryChunkManager:
seg_hash: str | None, seg_hash: str | None,
steps_by_id: dict[str, dict], steps_by_id: dict[str, dict],
ledger: Any = None, ledger: Any = None,
*,
trigger: str | None = None,
) -> ChunkState | None: ) -> ChunkState | None:
# A turn carries every step id it recorded (``step_keys``, Flash # A turn carries every step id it recorded (``step_keys``, Flash
# multi-action turns); older ledgers only expose ``step_key``. Each # multi-action turns); older ledgers only expose ``step_key``. Each
@@ -1302,14 +1380,89 @@ class HistoryChunkManager:
band3=build_action_ledger(steps, session_start), band3=build_action_ledger(steps, session_start),
minimal_index=build_action_ledger(steps, session_start, minimal=True), 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], 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 # Ready gating: the caller queues the chunk as awaiting — it only
# enters self._chunks (and the frozen region) once its capsule is # enters self._chunks (and the frozen region) once its capsule is
# ready, or through the hard-threshold emergency swap. # ready, or through the hard-threshold emergency swap.
self._persist(chunk) self._persist(chunk)
self._dispatch_capsule(chunk, steps, session_start, turn_records) 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 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 @staticmethod
def _step_payload(step: dict, session_start: float | None) -> dict[str, Any]: def _step_payload(step: dict, session_start: float | None) -> dict[str, Any]:
"""Mechanical per-step facts handed to the capsule lens.""" """Mechanical per-step facts handed to the capsule lens."""
@@ -78,6 +78,7 @@ class FakeEngine:
self.session_start_time = SESSION_START self.session_start_time = SESSION_START
self.current_session_id = "session-1" self.current_session_id = "session-1"
self.chunk_writes: list[dict] = [] self.chunk_writes: list[dict] = []
self.traces: list[dict] = []
if base_dir is not None: if base_dir is not None:
self.base_dir = base_dir self.base_dir = base_dir
@@ -88,6 +89,10 @@ class FakeEngine:
self.chunk_writes.append(kwargs) self.chunk_writes.append(kwargs)
return f"chunk-row-{len(self.chunk_writes)}" 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): class StubCapsuleService(StepMemoryService):
"""Records capsule submissions without spawning asyncio tasks (sync tests). """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 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 # Band ① note-linkage machine check
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------