@linxiraos/pi-agent-core 1.0.9 → 1.1.0

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 "@linxiraos/pi-natives";
+import type { Model } from "@linxiraos/pi-ai";
+import type { ModelTokenizer } from "@linxiraos/pi-catalog/types";
+import { countTokens as countTokensNat, Encoding } from "@linxiraos/pi-natives";
+import * as snapcompact from "@linxiraos/pi-snapcompact";
+import { stringifyJson } from "@linxiraos/pi-utils";
+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;

package/dist/types/agent-loop.d.ts

@@ -1,137 +0,0 @@
-/**
- * Agent loop that works with AgentMessage throughout.
- * Transforms to Message[] only at the LLM call boundary.
- */
-import { type AssistantMessage, type Context, EventStream, type ToolResultMessage } from "@linxiraos/pi-ai";
-import { type Dialect } from "@linxiraos/pi-ai/dialect";
-import { type AgentRunCoverage, type AgentRunSummary } from "./run-collector.js";
-import type { AgentContext, AgentEvent, AgentLoopConfig, AgentMessage, StreamFn } from "./types.js";
-/** Stop-details marker for a provider error after assistant content/tool args already streamed. */
-export declare const STREAM_INTERRUPTED_AFTER_CONTENT_STOP_DETAIL = "stream_interrupted_after_content";
-/**
- * Cadence (ms) for polling queued steering while an `interruptible` tool is in
- * flight, so a steer cuts the wait short instead of sitting idle until the
- * tool's own window elapses. A cheap synchronous queue check; latency-bounded
- * at one tick.
- */
-/**
- * Abort reason for a turn-wide interruption where only some tool calls caused
- * the abort and sibling placeholders need neutral messages.
- */
-export interface ToolScopedAbortReason {
-    readonly kind: "tool-scoped-abort";
-    readonly message: string;
-    readonly toolCallMessages: Record<string, string>;
-    readonly defaultToolCallMessage: string;
-}
-/** Creates an abort reason that labels matching tool calls separately from siblings. */
-export declare function createToolScopedAbortReason(message: string, toolCallMessages: Record<string, string>, defaultToolCallMessage: string): ToolScopedAbortReason;
-/**
- * Marks an abort raised by a completed post-tool hook as terminal for the
- * current run. External/user aborts still synthesize an aborted assistant
- * boundary; this reason stops after persisting the completed tool batch.
- */
-export declare const TERMINAL_TOOL_RESULT_ABORT_REASON: unique symbol;
-export declare function resolveOwnedDialectFromEnv(value: string | undefined): Dialect | undefined;
-/**
- * Start an agent loop with a new prompt message.
- * The prompt is added to the context and events are emitted for it.
- */
-export declare function agentLoop(prompts: AgentMessage[], context: AgentContext, config: AgentLoopConfig, signal?: AbortSignal, streamFn?: StreamFn): EventStream<AgentEvent, AgentMessage[]>;
-/**
- * Continue an agent loop from the current context without adding a new message.
- * Used for retries - context already has user message or tool results.
- *
- * **Important:** The last message in context must convert to a `user` or `toolResult` message
- * via `convertToLlm`. If it doesn't, the LLM provider will reject the request.
- * This cannot be validated here since `convertToLlm` is only called once per turn.
- */
-export declare function agentLoopContinue(context: AgentContext, config: AgentLoopConfig, signal?: AbortSignal, streamFn?: StreamFn): EventStream<AgentEvent, AgentMessage[]>;
-/**
- * Detailed-result handle returned by {@link agentLoopDetailed}. Adds the
- * run-level telemetry/coverage rollup to the existing `AgentMessage[]`
- * payload without changing the resolved type of `stream.result()`.
- */
-export interface AgentLoopDetailedResult {
-    readonly messages: AgentMessage[];
-    readonly telemetry: AgentRunSummary | undefined;
-    readonly coverage: AgentRunCoverage | undefined;
-}
-/**
- * Convenience wrapper over {@link agentLoop} that exposes the run-level
- * summary + coverage alongside the messages. The returned `stream` is the
- * same `EventStream` callers already consume; `detailed()` awaits the
- * stream's `agent_end` event and returns the additive fields.
- *
- * Existing `stream.result()` semantics are preserved — it still resolves to
- * `AgentMessage[]`. Use {@link agentLoopDetailed} when you need the rollup;
- * use {@link agentLoop} when you do not.
- */
-export declare function agentLoopDetailed(prompts: AgentMessage[], context: AgentContext, config: AgentLoopConfig, signal?: AbortSignal, streamFn?: StreamFn): {
-    readonly stream: EventStream<AgentEvent, AgentMessage[]>;
-    readonly detailed: () => Promise<AgentLoopDetailedResult>;
-};
-/**
- * Like {@link agentLoopDetailed} but built on top of
- * {@link agentLoopContinue}.
- */
-export declare function agentLoopContinueDetailed(context: AgentContext, config: AgentLoopConfig, signal?: AbortSignal, streamFn?: StreamFn): {
-    readonly stream: EventStream<AgentEvent, AgentMessage[]>;
-    readonly detailed: () => Promise<AgentLoopDetailedResult>;
-};
-export declare function normalizeMessagesForProvider(messages: Context["messages"], model: AgentLoopConfig["model"]): Context["messages"];
-export interface NormalizeToolsOptions {
-    /** Inject the `i` intent field into tool schemas (subject to `PI_NO_INTENT`). */
-    injectIntent: boolean;
-    /** Strip descriptions from the wire specs when the catalog rides in the system prompt. */
-    pruneDescriptions?: boolean;
-}
-export declare function normalizeTools(tools: AgentContext["tools"], options: NormalizeToolsOptions): Context["tools"];
-/** Resolve the human-readable reason an abort carried. A caller that aborts via
- *  `AbortController.abort(reason)` with a string or a non-`AbortError` `Error`
- *  (e.g. the coding agent's user-interrupt label) gets that text surfaced on the
- *  synthesized assistant message's `errorMessage`; a bare `abort()` (whose
- *  `signal.reason` is the default `AbortError` `DOMException`) falls back to the
- *  generic sentinel that downstream renderers treat as "no specific reason". */
-export declare function abortReasonText(signal: AbortSignal | undefined): string;
-/**
- * Discriminator embedded in {@link AgentToolResult.details} and
- * {@link ToolResultMessage.details} for tool calls that were emitted by the
- * assistant but never actually invoked locally.
- *
- * The synthetic result exists only to preserve the tool_use / tool_result
- * pairing the provider API requires; no `tool.execute()` ran. UI, telemetry,
- * and history consumers can key on `__synthetic === true` to render or
- * classify these as "call emitted, not executed" instead of a real local
- * tool failure — the mislabeling this discriminator was introduced to fix
- * (#4321): a provider-side stream error after tool-call emission (e.g. Codex
- * websocket close) was surfaced by the CLI as if the local tool had failed.
- *
- * `source` names the state that prevented execution — either an assistant-side
- * turn termination (`assistant_stop_*`) or a mid-batch interrupt that skipped a
- * still-pending call to service queued steering/peer input (`interrupt_skipped`).
- * `upstreamError` is the provider-reported message when the turn ended with
- * `stopReason === "error"`.
- */
-export interface SyntheticToolResultDetails {
-    __synthetic: true;
-    source: "assistant_stop_aborted" | "assistant_stop_error" | "assistant_stop_skipped" | "assistant_stop_length" | "interrupt_skipped";
-    executed: false;
-    upstreamError?: string;
-}
-/**
- * Narrow an {@link AgentMessage} to a synthetic {@link ToolResultMessage} —
- * a tool_result emitted for a tool call the assistant never invoked (see
- * {@link SyntheticToolResultDetails}). Consumers use this to look past the
- * placeholder pairing back to the assistant turn that produced it, e.g.
- * `AgentSession.retry()` walking back over the synthetic results a
- * stalled/aborted mid-tool-call turn leaves behind.
- */
-export declare function isSyntheticToolResultMessage(message: AgentMessage | undefined): message is ToolResultMessage<SyntheticToolResultDetails>;
-/**
- * Create the persisted synthetic result for a tool call that was emitted by
- * the assistant but never invoked locally.
- */
-export declare function createSyntheticToolResultMessage(toolCall: Extract<AssistantMessage["content"][number], {
-    type: "toolCall";
-}>, reason: "aborted" | "error" | "skipped" | "length", errorMessage?: string): ToolResultMessage<SyntheticToolResultDetails>;