pi-context-guard-jev 0.1.0 → 0.3.1

package/src/recall.ts

@@ -1,10 +1,12 @@
 /**
- * Text returned by the `recall` tool: the original tool output, optionally filtered by a regex
- * or limited to a line window, truncated like Pi's `read` tool.
+ * What the `recall` tool returns: the original tool output (with its images), optionally filtered by a
+ * regex or limited to a line window, truncated like Pi's `read` tool; or an omitted exchange as a
+ * transcript (images in its user messages attached, tool-result images named by their entry id).
  */
 import vm from "node:vm";
 import { MARKER } from "./render.ts";
-import { type SourceEntryLike, textOf } from "./types.ts";
+import { isBranchSteering } from "./run.ts";
+import { type Block, type SourceEntryLike, textOf } from "./types.ts";
 
 export const RECALL_MAX_LINES = 2_000;
 export const RECALL_MAX_BYTES = 50 * 1024;
@@ -15,6 +17,9 @@ export const RECALL_MATCH_CHARS = 4_000;
 /** Wall-clock limit for filtering by `pattern` (ms); stops catastrophic backtracking. */
 export const RECALL_PATTERN_TIMEOUT_MS = 1_000;
 
+/** Most images one recall returns (an exchange's user messages can hold more; the transcript says so). */
+export const RECALL_MAX_IMAGES = 6;
+
 export interface RecallParams {
 	entryId: string;
 	pattern?: string;
@@ -22,10 +27,97 @@ export interface RecallParams {
 	limit?: number;
 }
 
-export function originalOutput(entry: SourceEntryLike | undefined, entryId: string): string {
+type RecallBranchEntry = { id?: string; type?: string; message?: { role?: string; content?: unknown; toolName?: string } & Record<string, unknown>; customType?: string; content?: unknown };
+
+/**
+ * The original text behind a context-guard stub: a tool result's output, or, for the user prompt of an
+ * omitted exchange, the whole exchange from the branch (prompt, assistant texts, tool calls and results,
+ * custom messages, steering messages) up to the next prompt.
+ */
+export function originalOutput(entry: SourceEntryLike | undefined, entryId: string, branch: readonly RecallBranchEntry[] = [], omittedIds?: readonly string[]): string {
+	return originalContent(entry, entryId, branch, omittedIds).text;
+}
+
+const imagesOf = (content: unknown): Block[] => (Array.isArray(content) ? (content as Block[]).filter((b) => b?.type === "image") : []);
+
+/** `originalOutput` plus the images that go with it: a tool result's images, or those of an exchange's user messages. */
+export function originalContent(
+	entry: SourceEntryLike | undefined,
+	entryId: string,
+	branch: readonly RecallBranchEntry[] = [],
+	omittedIds?: readonly string[],
+): { text: string; images: Block[] } {
 	const message = entry?.type === "message" ? entry.message : undefined;
-	if (message?.role !== "toolResult") throw new Error(`No tool result with entry id ${entryId}. Use the id from a "${MARKER}" header.`);
-	return textOf(message);
+	if (message?.role === "toolResult") return { text: textOf(message), images: imagesOf(message.content) };
+	if (message?.role === "user") {
+		const start = branch.findIndex((e) => e.id === entryId);
+		if (start >= 0) {
+			const images: Block[] = [];
+			const text = exchangeTranscript(branch, start, omittedIds, images);
+			return { text, images };
+		}
+	}
+	throw new Error(`No tool result or omitted exchange with entry id ${entryId}. Use the id from a "${MARKER}" note.`);
+}
+
+/**
+ * The entry ids an omission removed, from the latest `context-guard` record that omitted the exchange
+ * starting at `entryId` (its `omitted` list). Undefined when no record has one.
+ */
+export function omittedIdsFor(branch: readonly (RecallBranchEntry & { data?: unknown })[], entryId: string, customType: string): string[] | undefined {
+	for (let i = branch.length - 1; i >= 0; i--) {
+		const e = branch[i];
+		if (e.type !== "custom" || e.customType !== customType) continue;
+		const results = (e.data as { results?: { kind?: string; entryId?: string; outcome?: string; omitted?: unknown }[] } | undefined)?.results ?? [];
+		const hit = results.find((r) => r.kind === "exchange" && r.entryId === entryId && r.outcome === "removed" && Array.isArray(r.omitted));
+		if (hit) return hit.omitted as string[];
+	}
+	return undefined;
+}
+
+/**
+ * Plain-text transcript of the exchange that starts at branch[start] (a user message): with
+ * `omittedIds`, exactly the prompt and those entries (what the omission removed); otherwise everything up
+ * to the next user message that is not a steering message. Images: a tool result's are named by its entry
+ * id (recall it to see them); a user message's are appended to `images` (up to RECALL_MAX_IMAGES).
+ */
+export function exchangeTranscript(branch: readonly RecallBranchEntry[], start: number, omittedIds?: readonly string[], images: Block[] = []): string {
+	const out: string[] = [];
+	const listed = omittedIds ? new Set(omittedIds) : undefined;
+	for (let i = start; i < branch.length; i++) {
+		const e = branch[i];
+		if (listed) {
+			if (i > start && !listed.has(e.id ?? "")) continue;
+		} else if (i > start && e.type === "message" && e.message?.role === "user" && !isBranchSteering(branch, i)) break;
+		if (e.type === "custom_message") {
+			out.push(`## custom message (${e.customType ?? "custom"})`, textOf({ role: "custom", content: (e.content ?? "") as string }));
+			continue;
+		}
+		if (e.type !== "message" || !e.message) continue;
+		const m = e.message;
+		const content = m.content as string | Block[];
+		const pictures = imagesOf(content);
+		if (m.role === "user") {
+			out.push("## user", textOf({ role: "user", content }));
+			if (pictures.length > 0) {
+				const room = Math.max(0, RECALL_MAX_IMAGES - images.length);
+				images.push(...pictures.slice(0, room));
+				const shown = Math.min(room, pictures.length);
+				out.push(shown === pictures.length ? `[${pictures.length} image(s), attached below]` : `[${pictures.length} image(s), ${shown} attached below]`);
+			}
+		} else if (m.role === "assistant") {
+			out.push("## assistant");
+			const text = textOf({ role: "assistant", content });
+			if (text.trim()) out.push(text);
+			if (Array.isArray(content)) {
+				for (const block of content) if (block.type === "toolCall") out.push(`[tool call] ${String(block.name)} ${JSON.stringify(block.arguments ?? {})}`);
+			}
+		} else if (m.role === "toolResult") {
+			out.push(`## tool result (${m.toolName ?? "tool"})`, textOf({ role: "toolResult", content }));
+			if (pictures.length > 0) out.push(`[${pictures.length} image(s): recall({"entryId":"${e.id}"}) shows them]`);
+		}
+	}
+	return out.join("\n");
 }
 
 export function recallText(original: string, params: Omit<RecallParams, "entryId">): string {

package/src/render.ts

@@ -8,6 +8,8 @@ import type { Candidate } from "./run.ts";
 
 /** Every replacement starts with this; it also marks results as already distilled. */
 export const MARKER = "[context-guard]";
+/** Name of our recall tool. Its outputs are never judged mid-run: the model just asked for them. */
+export const RECALL_TOOL = "recall";
 
 export const formatChars = (chars: number): string =>
 	chars >= 1_000_000 ? `${(chars / 1_000_000).toFixed(1)}M` : chars >= 1_000 ? `${(chars / 1_000).toFixed(1)}k` : String(chars);

package/src/run.ts

@@ -61,57 +61,168 @@ export const HISTORY_FIRST_REQUEST_LIMIT = 1_000;
 /** Question text used when the prompt that started the run has no text (e.g. only an image). */
 export const NO_TEXT_QUESTION = "[the user sent only an image]";
 
+const hasUser = (entry: ProjectedEntryLike): boolean => entry.messages.some((m) => m.role === "user");
+
+/** An assistant message that handed over to tools (not stopped by an abort or error). */
+export const endsInToolCall = (message: { stopReason?: unknown; content?: unknown }): boolean =>
+	message.stopReason !== "aborted" && message.stopReason !== "error" && Array.isArray(message.content) && message.content.some((b) => b?.type === "toolCall");
+
+/** Index of the entry holding the last user message (prompt or steering message), or -1. */
+export function lastUserIndex(entries: readonly ProjectedEntryLike[], before = entries.length): number {
+	let index = Math.min(before, entries.length) - 1;
+	while (index >= 0 && !hasUser(entries[index])) index--;
+	return index;
+}
+
 /**
- * The run starts after the last user message, with or without text. Steering messages sent
- * during a run are user messages too, so they start a new span; tool results before them are
- * left for later (they stay unpruned, which is the safe direction).
+ * A user message sent while the agent was still working: the message before it (custom notices and
+ * omitted entries aside) is a tool result, or an assistant message with tool calls that did not stop on an
+ * abort or error. It continues the run. (A prompt typed after aborting a running tool looks the same and
+ * is treated as steering too: it usually redirects the same work.)
  */
-export function collectRun(entries: readonly ProjectedEntryLike[], options: CollectOptions): RunInfo | undefined {
-	let start = entries.length - 1;
-	while (start >= 0 && !entries[start].messages.some((m) => m.role === "user")) start--;
-	if (start < 0) return undefined;
+export function isSteering(entries: readonly ProjectedEntryLike[], index: number): boolean {
+	for (let i = index - 1; i >= 0; i--) {
+		const message = entries[i].messages.at(-1);
+		if (!message || message.role === "custom") continue;
+		if (message.role === "toolResult") return true;
+		if (message.role === "assistant") return endsInToolCall(message);
+		return false;
+	}
+	return false;
+}
 
+type BranchMessageEntry = { type?: string; message?: { role?: string; stopReason?: unknown; content?: unknown } };
+
+/** `isSteering` on raw branch entries: is the user message at branch[index] a steering message? */
+export function isBranchSteering(branch: readonly BranchMessageEntry[], index: number): boolean {
+	for (let i = index - 1; i >= 0; i--) {
+		const e = branch[i];
+		if (e.type !== "message" || !e.message) continue;
+		if (e.message.role === "toolResult") return true;
+		if (e.message.role === "assistant") return endsInToolCall(e.message);
+		return false;
+	}
+	return false;
+}
+
+/**
+ * Index of the entry holding the user message that starts the current run, or -1. Steering messages
+ * (sent during the run) do not start a run: the span goes back to the prompt the agent started from, so
+ * tool results from before a steering message are still judged (against the prompt and every steering
+ * message, see `runQuestion`).
+ */
+export function findRunStart(entries: readonly ProjectedEntryLike[]): number {
+	let start = lastUserIndex(entries);
+	while (start > 0 && isSteering(entries, start)) {
+		const previous = lastUserIndex(entries, start);
+		if (previous < 0) break;
+		start = previous;
+	}
+	return start;
+}
+
+/** The text of the user message at `start`. */
+export function questionAt(entries: readonly ProjectedEntryLike[], start: number): string {
 	const questionMessage = entries[start].messages.findLast((m) => m.role === "user");
 	const questionText = questionMessage ? textOf(questionMessage).trim() : "";
-	const question = questionText ? textOf(questionMessage as MessageLike) : NO_TEXT_QUESTION;
-	const span = entries.slice(start + 1);
+	return questionText ? textOf(questionMessage as MessageLike) : NO_TEXT_QUESTION;
+}
+
+export const STEERING_PREFIX = "[The user added during the run]";
+
+/** The run's question: the prompt at `start`, then every steering message of the run in order. */
+export function runQuestion(entries: readonly ProjectedEntryLike[], start: number): string {
+	const parts = [questionAt(entries, start)];
+	for (let i = start + 1; i < entries.length; i++) {
+		if (!hasUser(entries[i])) continue;
+		const text = textOf(entries[i].messages.findLast((m) => m.role === "user") as MessageLike).trim();
+		parts.push(`${STEERING_PREFIX} ${text || NO_TEXT_QUESTION}`);
+	}
+	return parts.join("\n\n");
+}
 
-	const assistantTexts: string[] = [];
-	const toolCalls = new Map<string, { name: string; args: Record<string, unknown> | undefined }>();
+export interface ToolCallInfo {
+	name: string;
+	args: Record<string, unknown> | undefined;
+	/** 0-based index of the assistant message (turn) in the run that made the call. */
+	turn: number;
+}
+
+/** Tool calls of a run span by id, and the assistant texts in order (with their turn). */
+export function indexSpan(span: readonly ProjectedEntryLike[]): { toolCalls: Map<string, ToolCallInfo>; assistantTexts: { turn: number; text: string }[]; turns: number } {
+	const toolCalls = new Map<string, ToolCallInfo>();
+	const assistantTexts: { turn: number; text: string }[] = [];
+	let turn = -1;
 	for (const entry of span) {
 		for (const message of entry.messages) {
-			if (message.role !== "assistant" || typeof message.content === "string") continue;
+			if (message.role !== "assistant") continue;
+			turn++;
+			if (typeof message.content === "string") continue;
 			const text = textOf(message).trim();
-			if (text) assistantTexts.push(text);
+			if (text) assistantTexts.push({ turn, text });
 			for (const block of message.content) {
 				if (block.type === "toolCall" && typeof block.id === "string") {
 					const args = block.arguments && typeof block.arguments === "object" ? (block.arguments as Record<string, unknown>) : undefined;
-					toolCalls.set(block.id, { name: String(block.name ?? ""), args });
+					toolCalls.set(block.id, { name: String(block.name ?? ""), args, turn });
 				}
 			}
 		}
 	}
-	const answer = assistantTexts.at(-1) ?? "";
-	const notes = assistantTexts.slice(0, -1).join("\n\n");
+	return { toolCalls, assistantTexts, turns: turn + 1 };
+}
+
+/** The candidate rules (see docs/DESIGN.md → Candidates) for one projected entry, or undefined. */
+export function candidateOf(
+	entry: ProjectedEntryLike,
+	toolCalls: ReadonlyMap<string, ToolCallInfo>,
+	options: CollectOptions,
+): (Candidate & { toolCallId?: string; turn?: number }) | undefined {
+	const message = entry.messages[0];
+	if (entry.messages.length !== 1 || message?.role !== "toolResult") return undefined;
+	if ((message.isError && !options.distillErrors) || !isTextOnly(message)) return undefined;
+	const call = message.toolCallId ? toolCalls.get(message.toolCallId) : undefined;
+	const toolName = message.toolName ?? call?.name ?? "tool";
+	if (options.excludeTools.includes(toolName)) return undefined;
+	const text = textOf(message);
+	if (text.length < options.minResultChars || text.startsWith(MARKER)) return undefined;
+	// Another extension already edited this result: recall returns the raw entry, so whatever
+	// that edit added could not be recovered after distillation. Leave it alone.
+	const raw = entry.sourceEntry.message;
+	if (raw && textOf(raw) !== text) return undefined;
+	return {
+		entryId: entry.sourceEntry.id,
+		toolName,
+		args: call?.args,
+		text,
+		...(message.isError ? { isError: true } : {}),
+		toolCallId: message.toolCallId,
+		turn: call?.turn,
+	};
+}
+
+/**
+ * The run starts after the user message that started the agent (see `findRunStart`: steering messages
+ * sent during the run belong to it).
+ *
+ * `skip`: entry ids not to collect.
+ */
+export function collectRun(entries: readonly ProjectedEntryLike[], options: CollectOptions, skip?: ReadonlySet<string>): RunInfo | undefined {
+	const start = findRunStart(entries);
+	if (start < 0) return undefined;
+	const question = runQuestion(entries, start);
+	const span = entries.slice(start + 1);
+	const { toolCalls, assistantTexts } = indexSpan(span);
+	const answer = assistantTexts.at(-1)?.text ?? "";
+	const notes = assistantTexts.slice(0, -1).map((t) => t.text).join("\n\n");
 
-	const exclude = new Set(options.excludeTools);
 	const candidates: Candidate[] = [];
 	let toolResults = 0;
 	for (const entry of span) {
-		const message = entry.messages[0];
-		if (entry.messages.length !== 1 || message?.role !== "toolResult") continue;
-		toolResults++;
-		if ((message.isError && !options.distillErrors) || !isTextOnly(message)) continue;
-		const call = message.toolCallId ? toolCalls.get(message.toolCallId) : undefined;
-		const toolName = message.toolName ?? call?.name ?? "tool";
-		if (exclude.has(toolName)) continue;
-		const text = textOf(message);
-		if (text.length < options.minResultChars || text.startsWith(MARKER)) continue;
-		// Another extension already edited this result: recall returns the raw entry, so whatever
-		// that edit added could not be recovered after distillation. Leave it alone.
-		const raw = entry.sourceEntry.message;
-		if (raw && textOf(raw) !== text) continue;
-		candidates.push({ entryId: entry.sourceEntry.id, toolName, args: call?.args, text, ...(message.isError ? { isError: true } : {}) });
+		if (entry.messages.length === 1 && entry.messages[0]?.role === "toolResult") toolResults++;
+		const candidate = candidateOf(entry, toolCalls, options);
+		if (!candidate || skip?.has(candidate.entryId)) continue;
+		const { toolCallId: _id, turn: _turn, ...plain } = candidate;
+		candidates.push(plain);
 	}
 	const history = collectHistory(entries.slice(0, start), options.historyExchanges);
 	return { question, answer, notes, candidates, toolResults, history };
@@ -125,7 +236,7 @@ export const hasHistory = (history: RunHistory | undefined): history is RunHisto
  * Tool calls and tool results are left out; every user message (steering included) starts a
  * new exchange, whose assistant part is the last assistant text before the next user message.
  */
-function collectHistory(before: readonly ProjectedEntryLike[], count: number): RunHistory {
+export function collectHistory(before: readonly ProjectedEntryLike[], count: number): RunHistory {
 	if (count <= 0) return { exchanges: [] };
 	let summary: string | undefined;
 	const all: { user: string; assistant: string; hasText: boolean }[] = [];

package/src/size.ts

@@ -0,0 +1,114 @@
+/**
+ * Size estimates for model context, in "text characters": the unit tool output is measured in, so the
+ * break-even rule (checkpoint.ts → paysOff) compares like with like.
+ *
+ * Blocks that are not plain text must not be counted by their serialized length:
+ * - images are base64 (a 1280×960 JPEG is ~110k characters) but cost about w×h/750 tokens on Anthropic
+ *   (measured on a real Opus session: 1280×960 → ~1.6k tokens), so they are counted by their pixels, at
+ *   4 characters per token (Pi's own chars/4 estimate, so the footer shows image savings as tokens);
+ * - thinking blocks carry an encrypted signature that grows with the hidden reasoning; it is counted at
+ *   half its length (measured: ~0.25 tokens per signature character vs ~0.4 per character of tool text).
+ */
+import { Buffer } from "node:buffer";
+import type { Block } from "./types.ts";
+
+/** Text characters per image token: Pi's chars/4 estimate (tool output itself runs at ~2.5–3). */
+export const IMAGE_CHARS_PER_TOKEN = 4;
+/** Anthropic's image token estimate: width × height / 750. */
+export const IMAGE_PIXELS_PER_TOKEN = 750;
+/** Tokens assumed for an image whose size cannot be read (about a 1.15-megapixel image). */
+export const DEFAULT_IMAGE_TOKENS = 1_600;
+/** Upper bound for one image (a 2576-pixel long edge at 4:3). */
+export const MAX_IMAGE_TOKENS = 6_700;
+/** Base64 characters decoded to find the image header (JPEG headers can sit behind EXIF data). */
+const HEADER_BASE64_CHARS = 65_536;
+
+/** Width and height from the header of a PNG, JPEG, GIF or WebP image, or undefined. */
+export function imageSize(bytes: Uint8Array): { width: number; height: number } | undefined {
+	const b = bytes;
+	const u16be = (i: number) => (b[i] << 8) | b[i + 1];
+	const u16le = (i: number) => b[i] | (b[i + 1] << 8);
+	const u32be = (i: number) => ((b[i] << 24) >>> 0) + (b[i + 1] << 16) + (b[i + 2] << 8) + b[i + 3];
+	if (b.length >= 24 && b[0] === 0x89 && b[1] === 0x50 && b[2] === 0x4e && b[3] === 0x47) return { width: u32be(16), height: u32be(20) };
+	if (b.length >= 10 && b[0] === 0x47 && b[1] === 0x49 && b[2] === 0x46) return { width: u16le(6), height: u16le(8) };
+	if (b.length >= 30 && b[0] === 0x52 && b[1] === 0x49 && b[2] === 0x46 && b[3] === 0x46 && b[8] === 0x57 && b[9] === 0x45 && b[10] === 0x42 && b[11] === 0x50) {
+		const chunk = String.fromCharCode(b[12], b[13], b[14], b[15]);
+		if (chunk === "VP8X") return { width: 1 + (b[24] | (b[25] << 8) | (b[26] << 16)), height: 1 + (b[27] | (b[28] << 8) | (b[29] << 16)) };
+		if (chunk === "VP8L") return { width: 1 + (u16le(21) & 0x3fff), height: 1 + (((b[22] >> 6) | (b[23] << 2) | (b[24] << 10)) & 0x3fff) };
+		if (chunk === "VP8 ") return { width: u16le(26) & 0x3fff, height: u16le(28) & 0x3fff };
+		return undefined;
+	}
+	if (b.length >= 4 && b[0] === 0xff && b[1] === 0xd8) {
+		let i = 2;
+		while (i + 9 < b.length) {
+			if (b[i] !== 0xff) {
+				i++;
+				continue;
+			}
+			const marker = b[i + 1];
+			if (marker === 0xff) {
+				i++;
+				continue;
+			}
+			// SOF0–SOF15 except DHT (C4), JPG (C8) and DAC (CC).
+			if (marker >= 0xc0 && marker <= 0xcf && marker !== 0xc4 && marker !== 0xc8 && marker !== 0xcc) return { width: u16be(i + 7), height: u16be(i + 5) };
+			if (marker === 0xd8 || marker === 0x01 || (marker >= 0xd0 && marker <= 0xd7)) {
+				i += 2;
+				continue;
+			}
+			i += 2 + u16be(i + 2);
+		}
+	}
+	return undefined;
+}
+
+export interface ImageInfo {
+	width?: number;
+	height?: number;
+	/** Estimated input tokens. */
+	tokens: number;
+}
+
+const imageInfoCache = new WeakMap<object, ImageInfo>();
+
+/** Size (when the header is readable) and estimated input tokens of an image block (`{ type: "image", data: <base64> }`). */
+export function imageInfo(block: Block): ImageInfo {
+	const cached = imageInfoCache.get(block);
+	if (cached) return cached;
+	let info: ImageInfo = { tokens: DEFAULT_IMAGE_TOKENS };
+	if (typeof block.data === "string") {
+		const head = block.data.slice(0, HEADER_BASE64_CHARS);
+		const size = imageSize(Buffer.from(head.slice(0, head.length - (head.length % 4)), "base64"));
+		if (size && size.width > 0 && size.height > 0) {
+			info = { ...size, tokens: Math.min(MAX_IMAGE_TOKENS, Math.max(1, Math.ceil((size.width * size.height) / IMAGE_PIXELS_PER_TOKEN))) };
+		}
+	}
+	imageInfoCache.set(block, info);
+	return info;
+}
+
+/** Estimated input tokens of an image block. */
+export const imageTokens = (block: Block): number => imageInfo(block).tokens;
+
+/** Estimated size of one content block in text characters. */
+export function blockChars(block: Block): number {
+	switch (block.type) {
+		case "text":
+			return typeof block.text === "string" ? block.text.length : 0;
+		case "image":
+			return imageTokens(block) * IMAGE_CHARS_PER_TOKEN;
+		case "thinking": {
+			const thinking = typeof block.thinking === "string" ? block.thinking.length : 0;
+			const signature = typeof block.thinkingSignature === "string" ? block.thinkingSignature.length : 0;
+			return thinking + Math.ceil(signature / 2);
+		}
+		case "toolCall":
+			return String(block.name ?? "").length + JSON.stringify(block.arguments ?? {}).length;
+		default:
+			return JSON.stringify(block).length;
+	}
+}
+
+/** Estimated size of message content in text characters. */
+export const contentChars = (content: string | readonly Block[] | undefined): number =>
+	typeof content === "string" ? content.length : (content ?? []).reduce((n, block) => n + blockChars(block), 0);

package/src/stats.ts

@@ -3,17 +3,25 @@
  * `/tree` navigation, reloads and compaction:
  *
  * - active savings: for every tool result in the current model context whose projected text
- *   is one of our replacements, original length minus replacement length;
- * - Jev usage: summed from our `context-guard` custom entries on the active branch.
+ *   is one of our replacements, original size minus replacement size (size.ts: images by pixels);
+ * - lifetime savings: summed from our `context-guard` records on the active branch. Compaction drops
+ *   earlier edits from the context, so after one the footer shows both (`🛡 0 · Σ−79k`);
+ * - Jev usage: summed from the same records.
  */
+import { imagesRemoved } from "./images.ts";
 import { formatChars, MARKER } from "./render.ts";
-import { type ProjectedEntryLike, textOf } from "./types.ts";
+import { EXCHANGE_STUB_PREFIX } from "./items.ts";
+import { isBranchSteering } from "./run.ts";
+import { contentChars } from "./size.ts";
+import { type MessageLike, type ProjectedEntryLike, textOf } from "./types.ts";
 
 export const CUSTOM_TYPE = "context-guard";
 
 /** Data stored in each `context-guard` custom entry (one per distilled run). */
 export interface RunRecordData {
 	v: 1;
+	/** "mid-run" for checkpoints; a record without it counts as run end. */
+	phase?: "mid-run" | "run-end";
 	model: string;
 	savedChars: number;
 	requests: number;
@@ -22,6 +30,10 @@ export interface RunRecordData {
 	ms: number;
 	timedOut: boolean;
 	results: {
+		/** Absent in records that only had large outputs. */
+		kind?: "large" | "small" | "exchange" | "image";
+		/** Omitted exchanges: the entry ids that got `replacement: null`. */
+		omitted?: string[];
 		entryId: string;
 		tool: string;
 		label: string;
@@ -35,13 +47,24 @@ export interface RunRecordData {
 }
 
 export interface GuardStats {
-	/** Tool results currently distilled in model context. */
+	/** Tool results currently distilled or omitted in model context (large and small outputs). */
 	distilledResults: number;
+	/** Tool results whose images are currently removed from model context. */
+	imageResults: number;
+	/** Old exchanges currently omitted from model context. */
+	omittedExchanges: number;
 	/** Characters currently kept out of model context. */
 	savedChars: number;
 	/** Rough token equivalent (Pi's own estimate: chars / 4). */
 	savedTokens: number;
+	/** Characters saved by every record on the branch, including edits a compaction has since dropped. */
+	lifetimeChars: number;
+	lifetimeTokens: number;
+	/** A compaction came after a record that saved something: `lifetime` is worth showing. */
+	compacted: boolean;
 	runs: number;
+	/** Mid-run checkpoints among `runs`. */
+	checkpoints: number;
 	requests: number;
 	costUsd: number;
 	last?: RunRecordData;
@@ -49,29 +72,65 @@ export interface GuardStats {
 
 export const estimateTokens = (chars: number): number => Math.round(chars / 4);
 
-export function computeStats(projection: readonly ProjectedEntryLike[], branch: readonly { type?: string; customType?: string; data?: unknown }[]): GuardStats {
-	const stats: GuardStats = { distilledResults: 0, savedChars: 0, savedTokens: 0, runs: 0, requests: 0, costUsd: 0 };
+type StatsBranchEntry = { type?: string; id?: string; customType?: string; data?: unknown; message?: MessageLike };
+
+const messageChars = (message: MessageLike): number => contentChars(message.content);
+
+export function computeStats(projection: readonly ProjectedEntryLike[], branch: readonly StatsBranchEntry[]): GuardStats {
+	const stats: GuardStats = {
+		distilledResults: 0,
+		imageResults: 0,
+		omittedExchanges: 0,
+		savedChars: 0,
+		savedTokens: 0,
+		lifetimeChars: 0,
+		lifetimeTokens: 0,
+		compacted: false,
+		runs: 0,
+		checkpoints: 0,
+		requests: 0,
+		costUsd: 0,
+	};
+	// Pi keeps omitted entries in the projection with `messages: []`: count what the model sees.
+	const visibleMessages = new Map(projection.map((entry) => [entry.sourceEntry.id, entry.messages.length]));
 	for (const entry of projection) {
 		const projected = entry.messages[0];
 		const raw = entry.sourceEntry.message;
+		if (entry.messages.length === 1 && projected?.role === "user" && raw && textOf(projected).startsWith(EXCHANGE_STUB_PREFIX) && !textOf(raw).startsWith(MARKER)) {
+			// An omitted exchange: its raw entries up to the next user message that the model no longer sees.
+			const start = branch.findIndex((e) => e.id === entry.sourceEntry.id);
+			let before = messageChars(raw);
+			for (let i = start + 1; start >= 0 && i < branch.length; i++) {
+				const e = branch[i];
+				if (e.type === "message" && e.message?.role === "user" && !isBranchSteering(branch, i)) break;
+				if (e.message && e.id && (visibleMessages.get(e.id) ?? 0) === 0) before += messageChars(e.message);
+			}
+			stats.omittedExchanges++;
+			stats.savedChars += Math.max(0, before - messageChars(projected));
+			continue;
+		}
 		if (entry.messages.length !== 1 || projected?.role !== "toolResult" || !raw) continue;
-		const now = textOf(projected);
-		if (!now.startsWith(MARKER)) continue;
-		const before = textOf(raw);
-		if (before.startsWith(MARKER)) continue;
-		stats.distilledResults++;
-		stats.savedChars += Math.max(0, before.length - now.length);
+		if (!textOf(projected).startsWith(MARKER) || textOf(raw).startsWith(MARKER)) continue;
+		if (imagesRemoved(projected, raw)) stats.imageResults++;
+		else stats.distilledResults++;
+		stats.savedChars += Math.max(0, messageChars(raw) - messageChars(projected));
 	}
 	stats.savedTokens = estimateTokens(stats.savedChars);
+	let saving = false;
 	for (const entry of branch) {
+		if (entry.type === "compaction" && saving) stats.compacted = true;
 		if (entry.type !== "custom" || entry.customType !== CUSTOM_TYPE) continue;
 		const data = entry.data as Partial<RunRecordData> | undefined;
 		if (!data || data.v !== 1) continue;
 		stats.runs++;
+		if (data.phase === "mid-run") stats.checkpoints++;
 		stats.requests += data.requests ?? 0;
 		stats.costUsd += data.costUsd ?? 0;
+		stats.lifetimeChars += Math.max(0, data.savedChars ?? 0);
+		if ((data.savedChars ?? 0) > 0) saving = true;
 		stats.last = data as RunRecordData;
 	}
+	stats.lifetimeTokens = estimateTokens(stats.lifetimeChars);
 	return stats;
 }
 
@@ -88,9 +147,12 @@ export function formatStatus(stats: GuardStats, state: StatusState, color: Color
 	if (state === "off") return color("dim", `${icon} guard off`);
 	if (state === "busy") return color("accent", `${icon} distilling…`);
 	if (state === "problem") return color("warning", `${icon} ${problem ?? "guard unavailable"}`);
-	if (stats.savedTokens === 0) return color("dim", `${icon} 0 saved`);
-	// "🛡 −4.2k · 1": ≈ tokens kept out of context · distilled results.
-	return `${color("success", `${icon} −${formatTokens(stats.savedTokens)}`)}${color("dim", ` · ${stats.distilledResults}`)}`;
+	// After a compaction dropped earlier edits: "· Σ−79k", every saving of the session so far.
+	const lifetime = stats.compacted && stats.lifetimeTokens > stats.savedTokens ? color("dim", ` · Σ−${formatTokens(stats.lifetimeTokens)}`) : "";
+	if (stats.savedTokens === 0) return `${color("dim", `${icon} ${lifetime ? "0" : "0 saved"}`)}${lifetime}`;
+	const count = stats.distilledResults + stats.imageResults + stats.omittedExchanges;
+	// "🛡 −4.2k · 1": ≈ tokens kept out of context · edited items.
+	return `${color("success", `${icon} −${formatTokens(stats.savedTokens)}`)}${color("dim", ` · ${count}`)}${lifetime}`;
 }
 
 export const formatCost = (usd: number): string => (usd === 0 ? "$0" : usd < 0.01 ? `$${usd.toFixed(4)}` : `$${usd.toFixed(2)}`);

package/src/types.ts

@@ -16,6 +16,8 @@ export interface MessageLike {
 	toolName?: string;
 	toolCallId?: string;
 	isError?: boolean;
+	/** Assistant messages: why the model stopped ("toolUse", "stop", "aborted", "error", …). */
+	stopReason?: string;
 }
 
 export interface SourceEntryLike {
@@ -63,6 +65,13 @@ export interface ContextEditDraft {
 	replacement: { content: { type: "text"; text: string }[] };
 }
 
+/** A context edit that omits its target from model context (`replacement: null`). */
+export interface OmitDraft {
+	type: "context_edit";
+	targetId: string;
+	replacement: null;
+}
+
 /** Draft accepted by `agent_before_settle` (Pi's `CustomEntryDraft`). */
 export interface CustomDraft {
 	type: "custom";