pi-observational-memory 3.0.3 → 3.1.1

package/src/serialize.ts

@@ -1,4 +1,5 @@
 import type { Message, TextContent, ToolResultMessage } from "@earendil-works/pi-ai";
+import { estimateStringTokens } from "./tokens.js";
 
 function pad(n: number): string {
 	return n.toString().padStart(2, "0");
@@ -160,23 +161,78 @@ export function serializeBranchEntries(entries: RenderableEntry[]): string {
 export type SourceAddressedSerialization = {
 	text: string;
 	sourceEntryIds: string[];
+	estimatedTokens: number;
+	truncatedSourceEntryIds: string[];
 };
 
+export type SourceAddressedSerializationOptions = {
+	/** Maximum estimated tokens in the final source-addressed text. */
+	maxTokens?: number;
+};
+
+const SOURCE_OMISSION_MARKER =
+	"\n\n[… middle omitted: source exceeds observer input budget; original source remains in the session ledger …]\n\n";
+
+function truncateSourceBlockToTokenBudget(label: string, rendered: string, maxTokens: number): string | undefined {
+	const required = `${label}\n${SOURCE_OMISSION_MARKER}`;
+	if (estimateStringTokens(required) > maxTokens) return undefined;
+	const full = `${label}\n${rendered}`;
+	if (estimateStringTokens(full) <= maxTokens) return full;
+	const maxChars = Math.max(1, maxTokens * 4);
+	const fixed = `${label}\n${SOURCE_OMISSION_MARKER}`;
+	const retainedChars = maxChars - fixed.length;
+	const headChars = Math.ceil(retainedChars / 2);
+	const tailChars = retainedChars - headChars;
+	return `${label}\n${rendered.slice(0, headChars)}${SOURCE_OMISSION_MARKER}${tailChars > 0 ? rendered.slice(-tailChars) : ""}`;
+}
+
 function isSourceRenderableEntry(entry: RenderableEntry): boolean {
 	return entry.type === "message" || entry.type === "custom_message" || entry.type === "branch_summary";
 }
 
-export function serializeSourceAddressedBranchEntries(entries: RenderableEntry[]): SourceAddressedSerialization {
+/**
+ * Serialize complete source entries up to the token budget. If the first entry
+ * alone exceeds the budget, include a clearly marked head/tail excerpt so one
+ * pathological tool result cannot permanently block observation coverage.
+ * The original ledger entry is never modified and remains recallable by id.
+ */
+export function serializeSourceAddressedBranchEntries(
+	entries: RenderableEntry[],
+	options: SourceAddressedSerializationOptions = {},
+): SourceAddressedSerialization {
 	const blocks: string[] = [];
 	const sourceEntryIds: string[] = [];
+	const truncatedSourceEntryIds: string[] = [];
+	let estimatedTokens = 0;
+
 	for (const entry of entries) {
 		if (!entry.id || !isSourceRenderableEntry(entry)) continue;
 		const rendered = serializeBranchEntries([entry]);
 		if (!rendered.trim()) continue;
+		const label = `[Source entry id: ${entry.id}]`;
+		const block = `${label}\n${rendered}`;
+		const separator = blocks.length > 0 ? "\n\n" : "";
+		const blockTokens = estimateStringTokens(`${separator}${block}`);
+		const maxTokens = options.maxTokens;
+
+		if (maxTokens !== undefined && estimatedTokens + blockTokens > maxTokens) {
+			if (blocks.length > 0) break;
+			const excerpt = truncateSourceBlockToTokenBudget(label, rendered, maxTokens);
+			if (!excerpt) break;
+			blocks.push(excerpt);
+			sourceEntryIds.push(entry.id);
+			truncatedSourceEntryIds.push(entry.id);
+			estimatedTokens = estimateStringTokens(excerpt);
+			break;
+		}
+
+		blocks.push(block);
 		sourceEntryIds.push(entry.id);
-		blocks.push(`[Source entry id: ${entry.id}]\n${rendered}`);
+		estimatedTokens += blockTokens;
 	}
-	return { text: blocks.join("\n\n"), sourceEntryIds };
+
+	const text = blocks.join("\n\n");
+	return { text, sourceEntryIds, estimatedTokens: estimateStringTokens(text), truncatedSourceEntryIds };
 }
 
 function renderRecallMessage(entry: RenderableEntry): string | null {

package/src/session-ledger/progress.ts

@@ -117,6 +117,106 @@ export function findLastCompactionIndex(entries: Entry[]): number {
 	return -1;
 }
 
+// ==== Real (provider-reported) token accounting ====
+//
+// These helpers measure context growth from provider-reported usage for the
+// observation and reflection coverage clocks. Automatic compaction keeps its
+// separate raw source-entry clock because its setting counts ledger entries.
+
+type UsageLike = {
+	totalTokens?: number;
+	input?: number;
+	output?: number;
+	cacheRead?: number;
+	cacheWrite?: number;
+};
+
+export function contextTokensFromUsage(usage: unknown): number | undefined {
+	if (!usage || typeof usage !== "object") return undefined;
+	const u = usage as UsageLike;
+	const total = typeof u.totalTokens === "number" && Number.isFinite(u.totalTokens) && u.totalTokens > 0 ? u.totalTokens : undefined;
+	if (total !== undefined) return total;
+	const parts = [u.input, u.output, u.cacheRead, u.cacheWrite];
+	if (parts.every((p) => typeof p === "number" && Number.isFinite(p))) {
+		const sum = parts.reduce<number>((acc, p) => acc + (p ?? 0), 0);
+		return sum > 0 ? sum : undefined;
+	}
+	return undefined;
+}
+
+function validAssistantContextTokens(entry: Entry): number | undefined {
+	if (entry.type !== "message" || !entry.message || typeof entry.message !== "object") return undefined;
+	const msg = entry.message as { role?: string; stopReason?: string; usage?: unknown };
+	if (msg.role !== "assistant" || msg.stopReason === "aborted" || msg.stopReason === "error") return undefined;
+	return contextTokensFromUsage(msg.usage);
+}
+
+/**
+ * Real context tokens right after a compaction anchor.
+ *
+ * Only usage from an assistant that responded AFTER the compaction is a valid
+ * post-compaction baseline: pi's own docs state the last assistant usage
+ * before/at a compaction reflects the PRE-compaction context size. The usage
+ * carried on the compaction entry itself is the summary-generation call's
+ * usage (pre-compaction scale, a different LLM call), so it is deliberately
+ * NOT used as a baseline.
+ */
+export function realContextTokensAfterCompaction(entries: Entry[], compactionIdx: number): number | undefined {
+	for (let i = compactionIdx + 1; i < entries.length; i++) {
+		const t = validAssistantContextTokens(entries[i]);
+		if (t !== undefined) return t;
+	}
+	return undefined;
+}
+
+/**
+ * Real context tokens at the time observation coverage ended: last valid
+ * assistant usage at/before the covered entry. Returns undefined when no valid
+ * usage exists (e.g. an error/abort storm) — callers must fall back to the
+ * raw estimate rather than measuring from zero, which would otherwise read the
+ * full context as "growth" and re-fire stages every turn.
+ */
+export function realContextTokensAtCoverage(entries: Entry[], coverageIdx: number): number | undefined {
+	for (let i = coverageIdx; i >= 0; i--) {
+		const t = validAssistantContextTokens(entries[i]);
+		if (t !== undefined) return t;
+	}
+	return undefined;
+}
+
+/**
+ * Real context growth since the most recent anchor (a compaction, or the given
+ * coverage marker), measured from provider-reported usage.
+ *
+ * Returns undefined when the baseline cannot be measured reliably — no usage
+ * at/after the anchor, or the current context is SMALLER than the baseline
+ * (accounting basis changed, e.g. a mid-session model/provider switch that
+ * counts usage differently). Callers must fall back to the raw estimate in
+ * that case; clamping a stale baseline to 0 would starve the stage forever,
+ * and measuring from zero would over-fire it.
+ */
+export function realTokensSinceAnchor(
+	entries: Entry[],
+	customType: V3MemoryCustomType | undefined,
+	currentContextTokens: number,
+): number | undefined {
+	const coverageIdx = customType ? latestCoverageIndex(entries, customType) : -1;
+	const compactionIdx = findLastCompactionIndex(entries);
+	if (compactionIdx > coverageIdx) {
+		const baseline = realContextTokensAfterCompaction(entries, compactionIdx);
+		if (baseline === undefined) return undefined;
+		const delta = currentContextTokens - baseline;
+		return delta >= 0 ? delta : undefined;
+	}
+	if (coverageIdx >= 0) {
+		const baseline = realContextTokensAtCoverage(entries, coverageIdx);
+		if (baseline === undefined) return undefined;
+		const delta = currentContextTokens - baseline;
+		return delta >= 0 ? delta : undefined;
+	}
+	return Math.max(0, currentContextTokens);
+}
+
 export function rawTokensSinceLastCompaction(entries: Entry[]): number {
 	const compactionIndex = findLastCompactionIndex(entries);
 	if (compactionIndex === -1) return rawTokensAfterIndex(entries, -1);

package/src/tokens.ts

@@ -4,6 +4,24 @@ export function estimateStringTokens(text: string): number {
 	return Math.ceil(text.length / 4);
 }
 
+/**
+ * Estimate the rendered footprint of an observation line as it appears in
+ * summaries / pool listings: "[id] YYYY-MM-DD HH:MM [relevance] content".
+ * Pool budgets that only count bare content undercount every line's
+ * metadata overhead (id + timestamp + relevance tags), so the configured
+ * pool target was reached later than the rendered memory actually allowed.
+ */
+export function observationLineTokenCount(observation: {
+	id: string;
+	timestamp: string;
+	relevance: string;
+	content: string;
+}): number {
+	return estimateStringTokens(
+		`[${observation.id}] ${observation.timestamp} [${observation.relevance}] ${observation.content}`,
+	);
+}
+
 export function estimateEntryTokens(entry: { type: string; message?: unknown; content?: unknown; summary?: unknown }): number {
 	if (entry.type === "message" && entry.message) {
 		return estimateMessageTokens(entry.message as Parameters<typeof estimateMessageTokens>[0]);