pi-context-guard-jev 0.1.0 → 0.2.0

package/src/recall.ts

@@ -4,7 +4,7 @@
  */
 import vm from "node:vm";
 import { MARKER } from "./render.ts";
-import { type SourceEntryLike, textOf } from "./types.ts";
+import { type Block, type SourceEntryLike, textOf } from "./types.ts";
 
 export const RECALL_MAX_LINES = 2_000;
 export const RECALL_MAX_BYTES = 50 * 1024;
@@ -22,10 +22,69 @@ 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) up to the next user message.
+ */
+export function originalOutput(entry: SourceEntryLike | undefined, entryId: string, branch: readonly RecallBranchEntry[] = [], omittedIds?: readonly string[]): string {
 	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 textOf(message);
+	if (message?.role === "user") {
+		const start = branch.findIndex((e) => e.id === entryId);
+		if (start >= 0) return exchangeTranscript(branch, start, omittedIds);
+	}
+	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.
+ */
+export function exchangeTranscript(branch: readonly RecallBranchEntry[], start: number, omittedIds?: readonly string[]): 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") 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[];
+		if (m.role === "user") out.push("## user", textOf({ role: "user", content }));
+		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 }));
+	}
+	return out.join("\n");
 }
 
 export function recallText(original: string, params: Omit<RecallParams, "entryId">): string {

package/src/run.ts

@@ -66,52 +66,103 @@ export const NO_TEXT_QUESTION = "[the user sent only an image]";
  * 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).
  */
-export function collectRun(entries: readonly ProjectedEntryLike[], options: CollectOptions): RunInfo | undefined {
+/** Index of the entry holding the user message that starts the current run, or -1. */
+export function findRunStart(entries: readonly ProjectedEntryLike[]): number {
 	let start = entries.length - 1;
 	while (start >= 0 && !entries[start].messages.some((m) => m.role === "user")) start--;
-	if (start < 0) return undefined;
+	return start;
+}
 
+/** The run's question: 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 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;
+}
 
-	const assistantTexts: string[] = [];
-	const toolCalls = new Map<string, { name: string; args: Record<string, unknown> | undefined }>();
+/** 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 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).
+ *
+ * `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 = questionAt(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 +176,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/stats.ts

@@ -7,13 +7,16 @@
  * - Jev usage: summed from our `context-guard` custom entries on the active branch.
  */
 import { formatChars, MARKER } from "./render.ts";
-import { type ProjectedEntryLike, textOf } from "./types.ts";
+import { EXCHANGE_STUB_PREFIX } from "./items.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 +25,10 @@ export interface RunRecordData {
 	ms: number;
 	timedOut: boolean;
 	results: {
+		/** Absent in records that only had large outputs. */
+		kind?: "large" | "small" | "exchange";
+		/** Omitted exchanges: the entry ids that got `replacement: null`. */
+		omitted?: string[];
 		entryId: string;
 		tool: string;
 		label: string;
@@ -35,13 +42,17 @@ 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;
+	/** 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;
 	runs: number;
+	/** Mid-run checkpoints among `runs`. */
+	checkpoints: number;
 	requests: number;
 	costUsd: number;
 	last?: RunRecordData;
@@ -49,11 +60,28 @@ 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 };
+
+export function computeStats(projection: readonly ProjectedEntryLike[], branch: readonly StatsBranchEntry[]): GuardStats {
+	const stats: GuardStats = { distilledResults: 0, omittedExchanges: 0, savedChars: 0, savedTokens: 0, 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 = textOf(raw).length;
+			for (let i = start + 1; start >= 0 && i < branch.length; i++) {
+				const e = branch[i];
+				if (e.type === "message" && e.message?.role === "user") break;
+				if (e.message && e.id && (visibleMessages.get(e.id) ?? 0) === 0) before += textOf(e.message).length;
+			}
+			stats.omittedExchanges++;
+			stats.savedChars += Math.max(0, before - textOf(projected).length);
+			continue;
+		}
 		if (entry.messages.length !== 1 || projected?.role !== "toolResult" || !raw) continue;
 		const now = textOf(projected);
 		if (!now.startsWith(MARKER)) continue;
@@ -68,6 +96,7 @@ export function computeStats(projection: readonly ProjectedEntryLike[], branch:
 		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.last = data as RunRecordData;
@@ -89,8 +118,9 @@ export function formatStatus(stats: GuardStats, state: StatusState, color: Color
 	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`);
+	const count = stats.distilledResults + stats.omittedExchanges;
 	// "🛡 −4.2k · 1": ≈ tokens kept out of context · distilled results.
-	return `${color("success", `${icon} −${formatTokens(stats.savedTokens)}`)}${color("dim", ` · ${stats.distilledResults}`)}`;
+	return `${color("success", `${icon} −${formatTokens(stats.savedTokens)}`)}${color("dim", ` · ${count}`)}`;
 }
 
 export const formatCost = (usd: number): string => (usd === 0 ? "$0" : usd < 0.01 ? `$${usd.toFixed(4)}` : `$${usd.toFixed(2)}`);

package/src/types.ts

@@ -63,6 +63,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";