@oh-my-pi/pi-agent-core 17.3.7 → 17.4.0

package/src/compaction/transcript-tokens.ts

@@ -0,0 +1,111 @@
+/**
+ * Provider-anchored transcript token accounting.
+ *
+ * Local tokenization is the expensive way to answer "how big is this
+ * conversation?" — and usually the wrong one, because the provider already
+ * answered it. Every settled assistant turn carries `usage` covering the exact
+ * prompt it was sent: the system prompt, the tool schemas, and every message up
+ * to and including itself. The only genuinely unaccounted-for text is the tail
+ * appended *after* that turn.
+ *
+ * These helpers locate the newest trustworthy usage report and tokenize only
+ * that tail, so a long session pays counting proportional to one turn instead
+ * of to the whole transcript, every turn.
+ *
+ * Trust rules for an anchor (mirroring the provider contract):
+ * - Assistant role only — nothing else carries `usage`.
+ * - Not `aborted` / `error`: those turns report partial or zero usage.
+ * - `hasContextTokenUsage(usage)`: the report must carry usable context numbers.
+ */
+
+import type { AssistantMessage } from "@oh-my-pi/pi-ai";
+import type { MessageCountOptions, Tokenizer } from "../tokenizer";
+import type { AgentMessage } from "../types";
+import { calculateContextTokens, hasContextTokenUsage } from "./compaction";
+
+/** A provider usage report that accounts for a prefix of the transcript. */
+export interface TranscriptUsageAnchor {
+	/** Index in the scanned array; messages at or before it are provider-accounted. */
+	index: number;
+	/** The anchoring assistant turn. */
+	message: AssistantMessage;
+	/** Conversation tokens the provider reported for that prompt. */
+	tokens: number;
+}
+
+/**
+ * Whether this message's provider usage may anchor transcript accounting.
+ *
+ * The single home for the trust rules — every anchor scan MUST route through
+ * it so a stale-usage rule can never drift between the transcript walkers and
+ * the session-entry walkers.
+ */
+export function isTranscriptUsageAnchor(message: AgentMessage): message is AssistantMessage {
+	if (message.role !== "assistant") return false;
+	const assistant = message as AssistantMessage;
+	if (assistant.stopReason === "aborted" || assistant.stopReason === "error") return false;
+	return assistant.usage !== undefined && hasContextTokenUsage(assistant.usage);
+}
+
+/**
+ * Newest assistant turn in `messages[fromIndex..]` whose usage can anchor the
+ * transcript, or `undefined` when none qualifies (fresh context, or every
+ * recent turn aborted/errored).
+ *
+ * `fromIndex` excludes turns whose usage is stale — anything a compaction
+ * summarized away describes a prompt that is no longer sent.
+ */
+export function findTranscriptUsageAnchor(
+	messages: readonly AgentMessage[],
+	fromIndex = 0,
+): TranscriptUsageAnchor | undefined {
+	for (let index = messages.length - 1; index >= fromIndex; index--) {
+		const message = messages[index];
+		if (!isTranscriptUsageAnchor(message)) continue;
+		return { index, message, tokens: calculateContextTokens(message.usage) };
+	}
+	return undefined;
+}
+
+/** Options for {@link estimateTranscriptTokens}. */
+export interface TranscriptTokenOptions {
+	/**
+	 * Gates the anchor search only: usage at or before this index is stale (a
+	 * compaction rewrote the prompt it describes) and must not anchor. Content
+	 * accounting is governed separately by {@link countFromIndex}.
+	 */
+	anchorFromIndex?: number;
+	/**
+	 * First message whose content is counted locally when no anchor is found.
+	 * Defaults to 0 (count the whole transcript), which is what a floor
+	 * estimate wants; pass the compaction boundary to skip summarized-away
+	 * messages entirely.
+	 */
+	countFromIndex?: number;
+	/** Forwarded to {@link Tokenizer.countMessage} for every locally counted message. */
+	excludeEncryptedReasoning?: boolean;
+}
+
+/**
+ * Conversation tokens for `messages`: the provider's own report for everything
+ * it already covers, plus a local count of only the unaccounted-for tail.
+ *
+ * An anchored result already includes the non-message prefix (system prompt +
+ * tool schemas) because the provider charged it; an unanchored result is a
+ * message-only sum. Callers that add non-message tokens on top MUST branch on
+ * {@link findTranscriptUsageAnchor} rather than assuming one shape.
+ */
+export function estimateTranscriptTokens(
+	messages: readonly AgentMessage[],
+	tokenizer: Tokenizer,
+	options?: TranscriptTokenOptions,
+): number {
+	const estimateOptions: MessageCountOptions | undefined =
+		options?.excludeEncryptedReasoning === true ? { excludeEncryptedReasoning: true } : undefined;
+	const anchor = findTranscriptUsageAnchor(messages, options?.anchorFromIndex ?? 0);
+	let total = anchor?.tokens ?? 0;
+	for (let index = anchor ? anchor.index + 1 : (options?.countFromIndex ?? 0); index < messages.length; index++) {
+		total += tokenizer.countMessage(messages[index], estimateOptions);
+	}
+	return total;
+}

package/src/compaction/utils.ts

@@ -208,13 +208,20 @@ export function truncateToolResultForSummary(text: string): string {
 	return `${text.slice(0, TOOL_RESULT_MAX_CHARS)}\n\n[... ${truncatedChars} more characters truncated]`;
 }
 
+const SUMMARY_BOUNDARY_TAG_RE = /<\s*\/?\s*(?:conversation|previous-summary)\s*>/gi;
+
+/** Keep untrusted summary input from closing or impersonating harness-owned boundaries. */
+export function escapeSummaryBoundaryTags(text: string): string {
+	return text.replace(SUMMARY_BOUNDARY_TAG_RE, tag => `&lt;${tag.slice(1)}`);
+}
+
 /**
  * Serialize LLM messages as plain summary input without provider control tokens.
  */
 export function serializeConversationForSummary(messages: Message[], dialect?: Dialect): string {
 	const conversation = serializeConversation(messages, dialect);
-	if (dialect !== "harmony") return conversation;
-	return escapeHarmonyControlTokens(conversation);
+	const escaped = dialect === "harmony" ? escapeHarmonyControlTokens(conversation) : conversation;
+	return escapeSummaryBoundaryTags(escaped);
 }
 
 /**

package/src/tokenizer.ts

@@ -1,27 +1,282 @@
-import { countTokens as countTokensNat } from "@oh-my-pi/pi-natives";
+import type { Model } from "@oh-my-pi/pi-ai";
+import type { ModelTokenizer } from "@oh-my-pi/pi-catalog/types";
+import { countTokens as countTokensNat, Encoding } from "@oh-my-pi/pi-natives";
+import { stringifyJson } from "@oh-my-pi/pi-utils";
+import * as snapcompact from "@oh-my-pi/snapcompact";
+import { isEstimateCacheable, messageEstimateVersion } from "./compaction/message-cache";
+import type { AgentMessage } from "./types";
 
-const accurate = process.env.PI_TOKENIZER_ACCURATE === "1" && Bun.env.NODE_ENV !== "test";
+const testEnv = Bun.env.NODE_ENV === "test";
+const accurate = process.env.PI_TOKENIZER_ACCURATE === "1" && !testEnv;
 
-function estimateTokens(text: string) {
+const NATIVE_ENCODING: Record<ModelTokenizer, Encoding> = {
+	"claude-v3": Encoding.ClaudeV3,
+	"claude-v47": Encoding.ClaudeV47,
+	"claude-v5": Encoding.ClaudeV5,
+	"claude-v5-sonnet": Encoding.ClaudeV5Sonnet,
+	qwen3: Encoding.Qwen3,
+	"deepseek-v3": Encoding.DeepSeekV3,
+	"kimi-k2": Encoding.KimiK2,
+	glm5: Encoding.Glm5,
+};
+
+/** Maps the catalog-resolved tokenizer family to its native implementation. */
+export function tokenizerEncodingForModel(model: Pick<Model, "tokenizer"> | null | undefined): Encoding | null {
+	return model?.tokenizer ? NATIVE_ENCODING[model.tokenizer] : null;
+}
+
+/**
+ * `strict` always pays for an exact native count (the catalog-resolved
+ * tokenizer when known, o200k_base otherwise). `approximate` and
+ * `upperbound` prefer the same exact count for known tokenizer families or
+ * when `PI_TOKENIZER_ACCURATE=1` is set; otherwise they use a cheap heuristic:
+ * `approximate` a bytes/4 guess, `upperbound` the raw byte length (never
+ * undercounts).
+ */
+export type TokenCountMode = "strict" | "approximate" | "upperbound";
+
+/** Options for {@link Tokenizer.countMessage} / {@link Tokenizer.countMessages}. */
+export interface MessageCountOptions {
+	/**
+	 * Drop opaque provider reasoning payloads (`thinkingSignature`,
+	 * `redactedThinking`, native server-tool blocks) from the estimate. Those
+	 * are billed by the provider on replay, so the default counts them — but
+	 * their *local* byte size can diverge wildly from what the provider
+	 * charges, so the compaction floor (which only needs the reliably-countable,
+	 * on-wire-compressible content) excludes them to avoid false triggers on
+	 * thinking-heavy turns.
+	 */
+	excludeEncryptedReasoning?: boolean;
+}
+
+function byteEstimate(text: string): number {
 	return (Buffer.byteLength(text, "utf-8") + 3) >> 2;
 }
 
-export function countTokens(text: string | string[]): number {
-	if (accurate) {
-		return countTokensNat(text);
-	} else if (Array.isArray(text)) {
-		return text.reduce((sum, t) => sum + estimateTokens(t), 0);
-	} else {
-		return estimateTokens(text);
-	}
+function byteLength(text: string): number {
+	return Buffer.byteLength(text, "utf-8");
 }
 
-export function countTokensConservatively(text: string | string[]): number {
-	if (accurate) {
-		return countTokensNat(text);
-	} else if (Array.isArray(text)) {
-		return text.reduce((sum, value) => sum + Buffer.byteLength(value, "utf-8"), 0);
-	} else {
-		return Buffer.byteLength(text, "utf-8");
+function sumFragments(text: string | string[], perFragment: (t: string) => number): number {
+	return Array.isArray(text) ? text.reduce((sum, t) => sum + perFragment(t), 0) : perFragment(text);
+}
+
+/** Verdict from {@link Tokenizer.checkTokenBudget}. */
+export interface TokenBudgetCheck {
+	/** Whether the text fits the budget. */
+	fits: boolean;
+	/**
+	 * Token count behind the verdict: the exact native count when `exact` is
+	 * set, otherwise the cheap byte upper bound (which already fit, so it is
+	 * only an over-estimate of a count known to be under budget).
+	 */
+	tokens: number;
+	/** Whether the exact tokenizer had to run because the cheap bound busted. */
+	exact: boolean;
+}
+
+/**
+ * Image content has no tokenizer representation; charge a fixed estimate
+ * matching what providers typically bill for inline images.
+ */
+const IMAGE_TOKEN_ESTIMATE = 1200;
+
+/**
+ * Memoized estimates for one message under this tokenizer's encoding, split by
+ * the {@link MessageCountOptions.excludeEncryptedReasoning} option so the two
+ * variants never collide. `version` snapshots {@link messageEstimateVersion} at
+ * write time; an owner mutation (prune/shake/strip-images) bumps the version,
+ * which invalidates the entry in every live Tokenizer at once.
+ */
+interface MessageEstimate {
+	version: number;
+	default?: number;
+	floored?: number;
+}
+
+/**
+ * Model-aware local token counter. Immutable: the catalog-resolved encoding
+ * is fixed at construction, so a cached count can never straddle two
+ * encodings. An `Agent` owns one for its active model (swapping the instance
+ * when the model's encoding changes); one-shot flows construct their own for
+ * the model that will be billed. Known tokenizer families use exact native
+ * counts; unknown models keep the fast byte estimate (or o200k when
+ * `PI_TOKENIZER_ACCURATE=1`).
+ */
+export class Tokenizer {
+	readonly #encoding: Encoding | null;
+
+	/**
+	 * Per-message estimate memo. Keyed by message identity, deliberately not a
+	 * symbol-tagged property: callers spread messages to derive throwaway
+	 * variants for counting (`estimateBranchSummaryTokens` does
+	 * `countMessage({ ...message, content: truncated })`), and a property-borne
+	 * cache would ride along the spread. Identity keying gives clones a fresh
+	 * count.
+	 */
+	#estimates = new WeakMap<AgentMessage, MessageEstimate>();
+
+	constructor(model?: Pick<Model, "tokenizer"> | null) {
+		this.#encoding = tokenizerEncodingForModel(model);
+	}
+
+	get encoding(): Encoding | null {
+		return this.#encoding;
+	}
+
+	countTokens(text: string | string[], mode: TokenCountMode = "approximate"): number {
+		if (mode === "strict") return countTokensNat(text, this.#encoding);
+		if (!testEnv && this.#encoding !== null) return countTokensNat(text, this.#encoding);
+		if (accurate) return countTokensNat(text);
+		return sumFragments(text, mode === "upperbound" ? byteLength : byteEstimate);
+	}
+
+	/**
+	 * Cheap-first budget probe — the way to ask "does this fit in `budget`
+	 * tokens?" without tokenizing the world.
+	 *
+	 * Byte length is a hard upper bound on token count (every token consumes at
+	 * least one input byte), so text whose raw bytes already fit the budget
+	 * cannot possibly exceed it — that verdict is returned without tokenizing at
+	 * all. Only text that busts the bound is ambiguous, and only that case pays
+	 * for the exact count. Since the bound overshoots ~4x on ordinary prose, the
+	 * common "comfortably under budget" answer is free.
+	 */
+	checkTokenBudget(text: string | string[], budget: number): TokenBudgetCheck {
+		const bound = sumFragments(text, byteLength);
+		if (bound <= budget) return { fits: true, tokens: bound, exact: false };
+		const tokens = this.countTokens(text, "strict");
+		return { fits: tokens <= budget, tokens, exact: true };
+	}
+
+	/**
+	 * Token estimate for one message under this tokenizer's encoding.
+	 *
+	 * Settled historical messages are counted once and reused until an owner
+	 * (prune/shake/strip-images) calls `invalidateMessageCache`; streaming
+	 * assistants bypass the memo entirely (see the message-cache settle-gate
+	 * invariant). Image blocks charge a fixed per-image estimate.
+	 */
+	countMessage(message: AgentMessage, options?: MessageCountOptions): number {
+		const floored = options?.excludeEncryptedReasoning === true;
+		if (!isEstimateCacheable(message)) return this.#measureMessage(message, floored);
+		const version = messageEstimateVersion(message);
+		let entry = this.#estimates.get(message);
+		if (entry === undefined || entry.version !== version) {
+			entry = { version };
+			this.#estimates.set(message, entry);
+		}
+		const cached = floored ? entry.floored : entry.default;
+		if (cached !== undefined) return cached;
+		const result = this.#measureMessage(message, floored);
+		if (floored) entry.floored = result;
+		else entry.default = result;
+		return result;
+	}
+
+	/** Sum of {@link countMessage} over `messages`. */
+	countMessages(messages: readonly AgentMessage[], options?: MessageCountOptions): number {
+		let total = 0;
+		for (const message of messages) total += this.countMessage(message, options);
+		return total;
+	}
+
+	#measureMessage(message: AgentMessage, excludeEncryptedReasoning: boolean): number {
+		const fragments: string[] = [];
+		let extra = 0;
+		// Declaration-merged app roles (the coding-agent's bashExecution) are
+		// invisible to this package's union, so the discriminant is read as data.
+		const role: string = message.role;
+		if (role === "bashExecution") {
+			if ("command" in message && typeof message.command === "string") fragments.push(message.command);
+			if ("output" in message && typeof message.output === "string") fragments.push(message.output);
+			return fragments.length === 0 ? 0 : this.countTokens(fragments);
+		}
+
+		switch (message.role) {
+			case "user": {
+				const content: string | Array<{ type: string; text?: string }> = message.content;
+				if (typeof content === "string") {
+					fragments.push(content);
+				} else if (Array.isArray(content)) {
+					for (const block of content) {
+						if (block.type === "text" && block.text) {
+							fragments.push(block.text);
+						}
+					}
+				}
+				break;
+			}
+			case "assistant": {
+				for (const block of message.content) {
+					if (block.type === "text") {
+						fragments.push(block.text);
+					} else if (block.type === "thinking") {
+						fragments.push(block.thinking);
+						// Providers charge for the opaque signature/reasoning payload that
+						// rides alongside the thinking text (OpenAI Responses encrypted
+						// reasoning items, Anthropic signed thinking blocks, etc.). Without
+						// counting it, this estimator can read ~half of the provider-reported
+						// usage on thinking-heavy turns — see #2275 for the resulting
+						// compaction-trigger / post-check metric divergence. The compaction
+						// floor excludes it (its local byte size diverges from provider billing).
+						if (block.thinkingSignature && !excludeEncryptedReasoning) {
+							fragments.push(block.thinkingSignature);
+						}
+					} else if (block.type === "toolCall") {
+						fragments.push(block.name);
+						fragments.push(stringifyJson(block.arguments) ?? "null");
+					} else if (block.type === "redactedThinking") {
+						// Encrypted reasoning blob the provider still bills for on replay;
+						// excluded from the compaction floor for the same reason as above.
+						if (!excludeEncryptedReasoning) fragments.push(block.data);
+					} else if (block.type === "anthropicServerTool") {
+						// Native Anthropic server-tool call/result replayed verbatim on the
+						// wire (server_tool_use input and opaque result content). The provider
+						// still bills for it on same-provider replay; excluded from the
+						// compaction floor like other encrypted reasoning because its local
+						// byte size diverges from provider billing.
+						if (!excludeEncryptedReasoning) fragments.push(stringifyJson(block.block) ?? "null");
+					}
+				}
+				break;
+			}
+			case "hookMessage":
+			case "toolResult": {
+				if (typeof message.content === "string") {
+					fragments.push(message.content);
+				} else {
+					for (const block of message.content) {
+						if (block.type === "text" && block.text) {
+							fragments.push(block.text);
+						} else if (block.type === "image") {
+							extra += IMAGE_TOKEN_ESTIMATE;
+						}
+					}
+				}
+				break;
+			}
+			case "branchSummary":
+			case "compactionSummary": {
+				fragments.push(message.summary);
+				if (message.role === "compactionSummary") {
+					if (message.blocks) {
+						for (const block of message.blocks) {
+							if (block.type === "text") fragments.push(block.text);
+							else extra += snapcompact.FRAME_TOKEN_ESTIMATE;
+						}
+					} else if (message.images) {
+						// Snapcompact frames render at ≥1568px; providers bill the downscaled cap.
+						extra += message.images.length * snapcompact.FRAME_TOKEN_ESTIMATE;
+					}
+				}
+				break;
+			}
+			default:
+				return 0;
+		}
+
+		if (fragments.length === 0) return extra;
+		return extra + this.countTokens(fragments);
 	}
 }

package/src/types.ts

@@ -208,7 +208,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
 	 * @example
 	 * ```typescript
 	 * transformContext: async (messages) => {
-	 *   if (estimateTokens(messages) > MAX_TOKENS) {
+	 *   if (agent.tokenizer.countMessages(messages) > MAX_TOKENS) {
 	 *     return pruneOldMessages(messages);
 	 *   }
 	 *   return messages;