@oh-my-pi/pi-agent-core 17.0.3 → 17.0.5

package/CHANGELOG.md

@@ -2,6 +2,16 @@
 
 ## [Unreleased]
 
+## [17.0.5] - 2026-07-18
+
+### Added
+
+- Added a per-message token estimation cache to optimize performance by reusing token counts for settled message history, with automatic cache invalidation on message mutation.
+
+### Changed
+
+- Improved tool execution control by making tool interruptibility resolvable per call, allowing side-effecting operations to complete while passive waits can yield to queued steering.
+
 ## [17.0.2] - 2026-07-17
 
 ### Fixed

package/dist/types/compaction/index.d.ts

@@ -5,6 +5,7 @@ export * from "./branch-summarization.js";
 export * from "./compaction.js";
 export * from "./entries.js";
 export * from "./errors.js";
+export * from "./message-cache.js";
 export * from "./messages.js";
 export * from "./openai.js";
 export * from "./pruning.js";

package/dist/types/compaction/message-cache.d.ts

@@ -0,0 +1,23 @@
+import type { AgentMessage } from "../types.js";
+/**
+ * Register a cache tied to message identity so owner mutations in this package
+ * (prune/shake) can invalidate it across the package boundary. Returns an
+ * unregister function. The coding-agent `convertToLlm` memo registers here.
+ */
+export declare function registerMessageCacheInvalidator(invalidate: (message: AgentMessage) => void): () => void;
+/**
+ * True when this message's estimate is safe to cache by identity. Non-assistants
+ * are immutable once appended; assistants are cached only once settled (see the
+ * settle-gate invariant above).
+ */
+export declare function isEstimateCacheable(message: AgentMessage): boolean;
+/** Read a cached estimate for the given option split, or `undefined` on miss. */
+export declare function readEstimateCache(message: AgentMessage, excludeEncryptedReasoning: boolean): number | undefined;
+/** Store an estimate for the given option split. */
+export declare function writeEstimateCache(message: AgentMessage, excludeEncryptedReasoning: boolean, value: number): void;
+/**
+ * Drop every cached derivation of `message` after an in-place rewrite. Owners of
+ * mutation (prune, shake, strip-images) call this at the mutation seam so the
+ * next convert/estimate pass recomputes from the new content.
+ */
+export declare function invalidateMessageCache(message: AgentMessage): void;

package/dist/types/types.d.ts

@@ -570,14 +570,16 @@ export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any
     /** If true, argument validation errors are non-fatal: raw args are passed to execute() instead of returning an error to the LLM. */
     lenientArgValidation?: boolean;
     /**
-     * If true, the agent loop may abort this tool mid-execution to deliver a
-     * queued steering message (instead of waiting for the tool to finish on its
-     * own). Set only on tools that purely *wait* and observe their abort signal
-     * cleanly (e.g. the `job` poll), so the abort surfaces the tool's current
+     * Whether the agent loop may abort this tool mid-execution to deliver a
+     * queued steering message. A function resolves this per call from the raw,
+     * pre-validation arguments.
+     *
+     * Enable only for calls that purely *wait* and observe their abort signal
+     * cleanly (e.g. `job` poll), so the abort surfaces the tool's current
      * snapshot rather than corrupting a side effect. Honored only when
      * `interruptMode` is "immediate".
      */
-    interruptible?: boolean;
+    interruptible?: boolean | ((args: Partial<Static<TParameters>>) => boolean);
     /**
      * Controls how the INTENT_FIELD (`i`) is handled for this tool.
      * - `"require"` (default): `i` is injected and required in the parameter schema.

package/package.json

@@ -1,7 +1,7 @@
 {
 	"type": "module",
 	"name": "@oh-my-pi/pi-agent-core",
-	"version": "17.0.3",
+	"version": "17.0.5",
 	"description": "General-purpose agent with transport abstraction, state management, and attachment support",
 	"homepage": "https://omp.sh",
 	"author": "Can Boluk",
@@ -35,12 +35,12 @@
 		"fmt": "biome format --write ."
 	},
 	"dependencies": {
-		"@oh-my-pi/pi-ai": "17.0.3",
-		"@oh-my-pi/pi-catalog": "17.0.3",
-		"@oh-my-pi/pi-natives": "17.0.3",
-		"@oh-my-pi/pi-utils": "17.0.3",
-		"@oh-my-pi/pi-wire": "17.0.3",
-		"@oh-my-pi/snapcompact": "17.0.3",
+		"@oh-my-pi/pi-ai": "17.0.5",
+		"@oh-my-pi/pi-catalog": "17.0.5",
+		"@oh-my-pi/pi-natives": "17.0.5",
+		"@oh-my-pi/pi-utils": "17.0.5",
+		"@oh-my-pi/pi-wire": "17.0.5",
+		"@oh-my-pi/snapcompact": "17.0.5",
 		"@opentelemetry/api": "^1.9.1"
 	},
 	"devDependencies": {

package/src/agent-loop.ts

@@ -1824,11 +1824,25 @@ async function executeToolCalls(
 		const tool =
 			tools?.find(t => t.name === toolCall.name) ??
 			tools?.find(t => t.customWireName !== undefined && t.customWireName === toolCall.name);
+		const args = toolCall.arguments as Record<string, unknown>;
+		const interruptibleMode = tool?.interruptible;
+		let interruptible = false;
+		if (typeof interruptibleMode === "function") {
+			try {
+				interruptible = interruptibleMode(args);
+			} catch {
+				// Resolver failures default to preserving the tool's outcome.
+				interruptible = false;
+			}
+		} else {
+			interruptible = interruptibleMode === true;
+		}
 		return {
 			toolCall,
 			tool,
-			args: toolCall.arguments as Record<string, unknown>,
-			signal: tool?.interruptible ? interruptibleSignal : nonInterruptibleSignal,
+			args,
+			interruptible,
+			signal: interruptible ? interruptibleSignal : nonInterruptibleSignal,
 			started: false,
 			result: undefined as AgentToolResult<any> | undefined,
 			isError: false,
@@ -2210,16 +2224,16 @@ async function executeToolCalls(
 		}
 	}
 
-	// While an interruptible tool is in flight (e.g. a `job`/`irc` wait
-	// blocking on external work), queued steering or interrupting IRC would
-	// otherwise wait out the tool's own window. Poll only non-consuming queues
-	// and abort the shared tool signal so the boundary dequeue below injects
-	// the message promptly. Gated on immediate-interrupt mode + an
-	// interruptible tool; checkSteering is idempotent (no-op once triggered).
+	// While an interruptible tool call is in flight (e.g. a `hub` wait blocking
+	// on external work), queued steering or interrupting IRC would otherwise
+	// wait out the tool's own window. Poll only non-consuming queues and abort
+	// the shared tool signal so the boundary dequeue below injects the message
+	// promptly. Gated on immediate-interrupt mode + an interruptible call;
+	// checkSteering is idempotent (no-op once triggered).
 	const watchSteeringWhileRunning =
 		shouldInterruptImmediately &&
 		(hasSteeringMessages !== undefined || hasIrcInterrupts !== undefined) &&
-		records.some(r => r.tool?.interruptible === true);
+		records.some(record => record.interruptible);
 	const steeringWatchTimer = watchSteeringWhileRunning
 		? setInterval(() => void checkSteering(), STEERING_INTERRUPT_POLL_MS)
 		: undefined;

package/src/compaction/compaction.ts

@@ -43,6 +43,7 @@ import {
 	V2_RETAINED_MESSAGE_TOKEN_BUDGET,
 } from "./compaction-v2-streaming";
 import type { CompactionEntry, SessionEntry } from "./entries";
+import { isEstimateCacheable, readEstimateCache, writeEstimateCache } from "./message-cache";
 import { type ConvertToLlm, createBranchSummaryMessage, createCustomMessage, defaultConvertToLlm } from "./messages";
 import {
 	buildOpenAiNativeHistory,
@@ -364,6 +365,21 @@ const IMAGE_TOKEN_ESTIMATE = 1200;
  * content) excludes them to avoid false triggers on thinking-heavy turns.
  */
 export function estimateTokens(message: AgentMessage, options?: { excludeEncryptedReasoning?: boolean }): number {
+	// Settled historical messages are counted once and reused until an owner
+	// (prune/shake/strip-images) invalidates them; streaming assistants bypass
+	// the cache entirely (see message-cache.ts settle-gate invariant).
+	const cacheable = isEstimateCacheable(message);
+	const excludeEncryptedReasoning = options?.excludeEncryptedReasoning === true;
+	if (cacheable) {
+		const cached = readEstimateCache(message, excludeEncryptedReasoning);
+		if (cached !== undefined) return cached;
+	}
+	const result = computeMessageTokens(message, options);
+	if (cacheable) writeEstimateCache(message, excludeEncryptedReasoning, result);
+	return result;
+}
+
+function computeMessageTokens(message: AgentMessage, options?: { excludeEncryptedReasoning?: boolean }): number {
 	const fragments: string[] = [];
 	let extra = 0;
 	if ((message as { role?: string }).role === "bashExecution") {

package/src/compaction/index.ts

@@ -6,6 +6,7 @@ export * from "./branch-summarization";
 export * from "./compaction";
 export * from "./entries";
 export * from "./errors";
+export * from "./message-cache";
 export * from "./messages";
 export * from "./openai";
 export * from "./pruning";

package/src/compaction/message-cache.ts

@@ -0,0 +1,92 @@
+/**
+ * Per-message memoization for the two hot history walks: token estimation
+ * ({@link estimateTokens}) and LLM conversion (the coding-agent's `convertToLlm`).
+ *
+ * Long sessions re-walk a settled `AgentMessage[]` every turn, re-tokenizing and
+ * re-converting historical objects that only the newest suffix can change. These
+ * caches key on message *identity* so a settled message is counted/converted once
+ * and reused until an owner rewrites it.
+ *
+ * Correctness rests on two invariants:
+ *
+ * 1. **Settle gate.** A streaming assistant is mutated under one identity while
+ *    its `usage`/`stopReason` are provisional (the seed carries zeroed usage and
+ *    a placeholder `stopReason`). Caching it would freeze a mid-stream count, so
+ *    estimation only caches assistants that are settled — real `usage`
+ *    (`totalTokens > 0`) with a terminal `stopReason` that is not `"aborted"` /
+ *    `"error"`. Unsettled assistants never read or insert. Non-assistant roles
+ *    are immutable once appended and cache by identity.
+ * 2. **Owner invalidation.** `pruneToolOutputs` / `pruneSupersededToolResults`,
+ *    `applyShakeRegion`, and `stripImagesFromMessage` rewrite message content in
+ *    place under a stable identity. Each MUST call {@link invalidateMessageCache}
+ *    on the mutated message before the next convert/estimate pass so both caches
+ *    drop the stale entry. The convert cache lives in another package, so it
+ *    subscribes via {@link registerMessageCacheInvalidator}.
+ */
+import type { AssistantMessage } from "@oh-my-pi/pi-ai";
+import type { AgentMessage } from "../types";
+
+/** External cache invalidators (e.g. the coding-agent `convertToLlm` memo). */
+const externalInvalidators = new Set<(message: AgentMessage) => void>();
+
+/**
+ * Register a cache tied to message identity so owner mutations in this package
+ * (prune/shake) can invalidate it across the package boundary. Returns an
+ * unregister function. The coding-agent `convertToLlm` memo registers here.
+ */
+export function registerMessageCacheInvalidator(invalidate: (message: AgentMessage) => void): () => void {
+	externalInvalidators.add(invalidate);
+	return () => {
+		externalInvalidators.delete(invalidate);
+	};
+}
+
+// Dual option-split estimate caches: the compaction floor passes
+// `excludeEncryptedReasoning` (dropping opaque provider reasoning), so a message
+// has two distinct estimates that must not collide in one map.
+//
+// These are WeakMaps, not symbol-tagged properties, deliberately: callers spread
+// messages to derive throwaway variants for counting — `estimateBranchSummaryTokens`
+// does `estimateTokens({ ...message, content: truncated })`. A symbol-keyed cache
+// value rides along an object spread, so the truncated clone would inherit (and
+// return) the full-content estimate. Keying strictly on identity keeps the cache
+// off spread copies, which get their own fresh count.
+const estimateCacheDefault = new WeakMap<AgentMessage, number>();
+const estimateCacheFloored = new WeakMap<AgentMessage, number>();
+
+/**
+ * True when this message's estimate is safe to cache by identity. Non-assistants
+ * are immutable once appended; assistants are cached only once settled (see the
+ * settle-gate invariant above).
+ */
+export function isEstimateCacheable(message: AgentMessage): boolean {
+	if (message.role !== "assistant") return true;
+	const assistant = message as AssistantMessage;
+	return (
+		assistant.stopReason !== "aborted" &&
+		assistant.stopReason !== "error" &&
+		assistant.usage != null &&
+		assistant.usage.totalTokens > 0
+	);
+}
+
+/** Read a cached estimate for the given option split, or `undefined` on miss. */
+export function readEstimateCache(message: AgentMessage, excludeEncryptedReasoning: boolean): number | undefined {
+	return (excludeEncryptedReasoning ? estimateCacheFloored : estimateCacheDefault).get(message);
+}
+
+/** Store an estimate for the given option split. */
+export function writeEstimateCache(message: AgentMessage, excludeEncryptedReasoning: boolean, value: number): void {
+	(excludeEncryptedReasoning ? estimateCacheFloored : estimateCacheDefault).set(message, value);
+}
+
+/**
+ * Drop every cached derivation of `message` after an in-place rewrite. Owners of
+ * mutation (prune, shake, strip-images) call this at the mutation seam so the
+ * next convert/estimate pass recomputes from the new content.
+ */
+export function invalidateMessageCache(message: AgentMessage): void {
+	estimateCacheDefault.delete(message);
+	estimateCacheFloored.delete(message);
+	for (const invalidate of externalInvalidators) invalidate(message);
+}

package/src/compaction/pruning.ts

@@ -6,6 +6,7 @@ import type { ToolResultMessage } from "@oh-my-pi/pi-ai";
 import type { AgentMessage, AgentToolCall } from "../types";
 import { estimateTokens } from "./compaction";
 import type { SessionEntry, SessionMessageEntry } from "./entries";
+import { invalidateMessageCache } from "./message-cache";
 import {
 	collectToolCallsById,
 	isProtectedToolResult,
@@ -295,6 +296,7 @@ export function pruneSupersededToolResults(entries: SessionEntry[], config: Supe
 	for (const candidate of toPrune) {
 		candidate.message.content = [{ type: "text", text: candidate.notice }];
 		candidate.message.prunedAt = prunedAt;
+		invalidateMessageCache(candidate.message as AgentMessage);
 		tokensSaved += estimatePrunedSavings(candidate.tokens, candidate.notice);
 	}
 	return { prunedCount: toPrune.length, tokensSaved };
@@ -398,6 +400,7 @@ export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig =
 				: createPrunedNotice(candidate.tokens);
 		message.content = [{ type: "text", text: notice }];
 		message.prunedAt = prunedAt;
+		invalidateMessageCache(message as AgentMessage);
 		prunedCount++;
 	}
 

package/src/compaction/shake.ts

@@ -15,6 +15,7 @@ import { countTokens } from "../tokenizer";
 import type { AgentMessage } from "../types";
 import { estimateTokens } from "./compaction";
 import type { CustomMessageEntry, SessionEntry, SessionMessageEntry } from "./entries";
+import { invalidateMessageCache } from "./message-cache";
 import {
 	collectToolCallsById,
 	isProtectedToolResult,
@@ -406,12 +407,18 @@ export function applyShakeRegion(region: ShakeRegion, replacement: string): void
 		const message = region.entry.message as ToolResultMessage;
 		message.content = [{ type: "text", text: replacement }];
 		message.prunedAt = Date.now();
+		invalidateMessageCache(message as AgentMessage);
 		return;
 	}
 	const slot = getBlockTextSlot(region.entry, region.blockIndex);
 	if (!slot) return;
 	const text = slot.read();
 	slot.write(text.slice(0, region.start) + replacement + text.slice(region.end));
+	// Message entries keep a stable `entry.message` identity across context
+	// rebuilds, so an in-place block rewrite must drop its cached estimate/convert.
+	// Custom-message entries are re-materialized into a fresh AgentMessage on every
+	// buildSessionContext, so they carry no stable cached identity to invalidate.
+	if (region.entry.type === "message") invalidateMessageCache(region.entry.message);
 }
 
 /**

package/src/types.ts

@@ -657,14 +657,16 @@ export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any
 	/** If true, argument validation errors are non-fatal: raw args are passed to execute() instead of returning an error to the LLM. */
 	lenientArgValidation?: boolean;
 	/**
-	 * If true, the agent loop may abort this tool mid-execution to deliver a
-	 * queued steering message (instead of waiting for the tool to finish on its
-	 * own). Set only on tools that purely *wait* and observe their abort signal
-	 * cleanly (e.g. the `job` poll), so the abort surfaces the tool's current
+	 * Whether the agent loop may abort this tool mid-execution to deliver a
+	 * queued steering message. A function resolves this per call from the raw,
+	 * pre-validation arguments.
+	 *
+	 * Enable only for calls that purely *wait* and observe their abort signal
+	 * cleanly (e.g. `job` poll), so the abort surfaces the tool's current
 	 * snapshot rather than corrupting a side effect. Honored only when
 	 * `interruptMode` is "immediate".
 	 */
-	interruptible?: boolean;
+	interruptible?: boolean | ((args: Partial<Static<TParameters>>) => boolean);
 	/**
 	 * Controls how the INTENT_FIELD (`i`) is handled for this tool.
 	 * - `"require"` (default): `i` is injected and required in the parameter schema.