@quandev104/pi-style 0.1.7 → 0.2.1

package/extension-src/pi-style/pi/index.ts

@@ -11,12 +11,11 @@ import {
 import { requestToolPresentationRender } from "../features/tools/index.js";
 import { registerPiStyleCommand } from "./commands.js";
 import { type CompatibilityTestHooks, createPiStyleSessionCoordinator } from "./session-coordinator.js";
-import { usageFromSession } from "./session-usage.js";
+import { resetUsageFromSessionCache, usageFromSession } from "./session-usage.js";
 
-/** Usage patch that omits the key when no session usage exists (exact optional types). */
+/** Usage patch clears stale usage when the active branch/session has none. */
 function usagePatch(ctx: ExtensionContext): StatusSnapshot {
-	const usage = usageFromSession(ctx.sessionManager);
-	return usage ? { usage } : {};
+	return { usage: usageFromSession(ctx.sessionManager) } as StatusSnapshot;
 }
 
 /**
@@ -65,6 +64,7 @@ export default function piStyleExtension(pi: ExtensionAPI): void {
 	const coordinator = createPiStyleSessionCoordinator(pi, compatibilityTestHooks);
 	registerPiStyleCommand(pi, coordinator.app);
 	pi.on("session_start", async (event, ctx) => {
+		resetUsageFromSessionCache(ctx.sessionManager);
 		// Pi only activates read/bash/edit/write by default; grep/find/ls are
 		// registered but inactive (kept out of the model's tool list to keep the
 		// core small). Activate them so the TUI shows them and the model can call
@@ -110,7 +110,9 @@ export default function piStyleExtension(pi: ExtensionAPI): void {
 		),
 	);
 	pi.on("thinking_level_select", (event) => coordinator.app.update({ thinkingLevel: event.level }, "immediate"));
-	pi.on("session_info_changed", (event) => coordinator.app.update({ sessionName: event.name }, "coalesced"));
+	pi.on("session_info_changed", (event) =>
+		coordinator.app.update({ sessionName: event.name }, "coalesced", { refreshExtensionStatuses: true }),
+	);
 	pi.on("message_start", () => {
 		// A new message is a batch boundary: quiet-tool (read/ls/find) calls of the
 		// new message start a fresh batch instead of joining the previous one.
@@ -123,11 +125,13 @@ export default function piStyleExtension(pi: ExtensionAPI): void {
 	// Usage (tokens + cost) is aggregated from finalized session entries at
 	// message/turn boundaries, mirroring Pi's native footer; per-chunk updates
 	// stay usage-free to keep streaming cheap.
-	pi.on("message_end", (_event, ctx) => coordinator.app.update({ ...usagePatch(ctx) }, "coalesced"));
+	pi.on("message_end", (_event, ctx) =>
+		coordinator.app.update({ ...usagePatch(ctx) }, "coalesced", { refreshContextUsage: true }),
+	);
 	pi.on("turn_end", (event, ctx) => {
 		// Append the finalized assistant message's tool batch to the current run.
 		registerTurnFromMessage(event.message, event.toolResults);
-		coordinator.app.update({ ...usagePatch(ctx) }, "deferred");
+		coordinator.app.update({ ...usagePatch(ctx) }, "deferred", { refreshContextUsage: true });
 	});
 	pi.on("agent_end", () => {
 		// The run is complete: collapse its tool blocks into one summary line.
@@ -141,23 +145,32 @@ export default function piStyleExtension(pi: ExtensionAPI): void {
 			requestToolPresentationRender();
 		}
 	});
-	pi.on("agent_settled", (_event, ctx) => coordinator.app.update({ ...usagePatch(ctx) }, "coalesced"));
+	pi.on("agent_settled", (_event, ctx) =>
+		coordinator.app.update({ ...usagePatch(ctx) }, "coalesced", { refreshContextUsage: true }),
+	);
 	pi.on("session_tree", (_event, ctx) => {
+		resetUsageFromSessionCache(ctx.sessionManager);
 		// Rebuild the turn registry from session content so restored/branched
 		// history renders collapsed consistently (no in-process turn_end events).
 		rebuildTurnRegistryFromEntries(ctx.sessionManager.getEntries());
-		coordinator.app.update({ ...usagePatch(ctx) }, "deferred");
+		coordinator.app.update({ ...usagePatch(ctx) }, "deferred", { refreshContextUsage: true });
+	});
+	pi.on("session_compact", (_event, ctx) => {
+		resetUsageFromSessionCache(ctx.sessionManager);
+		coordinator.app.update({ ...usagePatch(ctx) }, "deferred", { refreshContextUsage: true });
 	});
-	pi.on("session_compact", (_event, ctx) => coordinator.app.update({ ...usagePatch(ctx) }, "deferred"));
 	pi.on("tool_result", (event, ctx) => {
 		if (["write", "edit", "bash"].includes(event.toolName)) {
 			coordinator.app.runtime.current?.invalidateGit();
-			coordinator.app.update({ ...usagePatch(ctx) }, "delayed-retry");
+			coordinator.app.update({ ...usagePatch(ctx) }, "delayed-retry", { refreshContextUsage: true });
 		}
 	});
 	pi.on("user_bash", (_event, ctx) => {
 		coordinator.app.runtime.current?.invalidateGit();
-		coordinator.app.update({ ...usagePatch(ctx) }, "delayed-retry");
+		coordinator.app.update({ ...usagePatch(ctx) }, "delayed-retry", { refreshContextUsage: true });
+	});
+	pi.on("session_shutdown", (_event, ctx) => {
+		resetUsageFromSessionCache(ctx.sessionManager);
+		coordinator.shutdown();
 	});
-	pi.on("session_shutdown", () => coordinator.shutdown());
 }

package/extension-src/pi-style/pi/session-usage.ts

@@ -18,27 +18,197 @@ interface UsageTotals {
 	cost: number;
 }
 
+interface UsageContribution extends UsageTotals {
+	sawUsage: boolean;
+}
+
+interface SessionUsageCacheState {
+	processedLength: number;
+	lastEntryId: string | undefined;
+	lastEntryRef: SessionEntry | undefined;
+	lastContribution: UsageContribution;
+	totals: UsageTotals;
+	cached: UsageSnapshot | undefined;
+	scannedEntries: number;
+	cacheHits: number;
+	rebuilds: number;
+	tailRefreshes: number;
+}
+
+export interface SessionUsageCacheStats {
+	processedLength: number;
+	scannedEntries: number;
+	cacheHits: number;
+	rebuilds: number;
+	tailRefreshes: number;
+}
+
+let usageCache = new WeakMap<SessionUsageSource, SessionUsageCacheState>();
+
 /**
- * Aggregate cumulative token/cost usage from all session entries, mirroring
- * Pi's native footer: assistant messages always carry usage, tool results and
- * compaction/branch summaries carry it when available.
+ * Aggregate cumulative token/cost usage from finalized session entries,
+ * mirroring Pi's native footer. The result is cached incrementally so repeated
+ * reads only scan appended entries and refresh the finalized tail entry when it
+ * changes in place.
  */
 export function usageFromSession(session: SessionUsageSource): UsageSnapshot | undefined {
-	const totals: UsageTotals = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0 };
-	let sawUsage = false;
-	for (const entry of session.getEntries()) {
-		if (entry.type === "message") {
-			if (entry.message.role !== "assistant" && entry.message.role !== "toolResult") continue;
-			const usage = "usage" in entry.message ? entry.message.usage : undefined;
-			if (!usage) continue;
-			addUsage(totals, usage);
-			sawUsage = true;
-		} else if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) {
-			addUsage(totals, entry.usage);
-			sawUsage = true;
+	const entries = session.getEntries();
+	const state = usageCache.get(session) ?? createState();
+	usageCache.set(session, state);
+
+	if (entries.length === 0) {
+		state.cacheHits++;
+		state.processedLength = 0;
+		state.lastEntryId = undefined;
+		state.lastEntryRef = undefined;
+		state.lastContribution = emptyContribution();
+		state.totals = emptyTotals();
+		state.cached = undefined;
+		return undefined;
+	}
+
+	if (state.processedLength === 0) {
+		rebuildFrom(entries, state);
+		return state.cached;
+	}
+
+	if (entries.length < state.processedLength) {
+		rebuildFrom(entries, state);
+		return state.cached;
+	}
+
+	const currentLast = entries.at(-1);
+	if (!currentLast) {
+		state.cacheHits++;
+		state.cached = undefined;
+		return undefined;
+	}
+
+	if (entries.length === state.processedLength) {
+		if (currentLast.id !== state.lastEntryId) {
+			rebuildFrom(entries, state);
+			return state.cached;
 		}
+		if (currentLast !== state.lastEntryRef) refreshTailEntry(currentLast, state);
+		else state.cacheHits++;
+		return state.cached;
+	}
+
+	const previousLast = entries[state.processedLength - 1];
+	if (!previousLast || previousLast.id !== state.lastEntryId) {
+		rebuildFrom(entries, state);
+		return state.cached;
 	}
-	if (!sawUsage) return undefined;
+	if (previousLast !== state.lastEntryRef) refreshTailEntry(previousLast, state);
+	for (const entry of entries.slice(state.processedLength)) appendEntry(entry, state);
+	return state.cached;
+}
+
+export function resetUsageFromSessionCache(session?: SessionUsageSource): void {
+	if (session) {
+		usageCache.delete(session);
+		return;
+	}
+	usageCache = new WeakMap<SessionUsageSource, SessionUsageCacheState>();
+}
+
+export function getUsageFromSessionCacheStats(session: SessionUsageSource): SessionUsageCacheStats {
+	const state = usageCache.get(session) ?? createState();
+	if (!usageCache.has(session)) usageCache.set(session, state);
+	return {
+		processedLength: state.processedLength,
+		scannedEntries: state.scannedEntries,
+		cacheHits: state.cacheHits,
+		rebuilds: state.rebuilds,
+		tailRefreshes: state.tailRefreshes,
+	};
+}
+
+function createState(): SessionUsageCacheState {
+	return {
+		processedLength: 0,
+		lastEntryId: undefined,
+		lastEntryRef: undefined,
+		lastContribution: emptyContribution(),
+		totals: emptyTotals(),
+		cached: undefined,
+		scannedEntries: 0,
+		cacheHits: 0,
+		rebuilds: 0,
+		tailRefreshes: 0,
+	};
+}
+
+function rebuildFrom(entries: readonly SessionEntry[], state: SessionUsageCacheState): void {
+	state.rebuilds++;
+	state.processedLength = 0;
+	state.lastEntryId = undefined;
+	state.lastEntryRef = undefined;
+	state.lastContribution = emptyContribution();
+	state.totals = emptyTotals();
+	state.cached = undefined;
+	for (const entry of entries) appendEntry(entry, state);
+}
+
+function refreshTailEntry(entry: SessionEntry, state: SessionUsageCacheState): void {
+	state.tailRefreshes++;
+	subtractContribution(state.totals, state.lastContribution);
+	const contribution = usageContribution(entry);
+	addContribution(state.totals, contribution);
+	state.lastEntryId = entry.id;
+	state.lastEntryRef = entry;
+	state.lastContribution = contribution;
+	state.cached = snapshotFromTotals(state.totals);
+	state.scannedEntries++;
+}
+
+function appendEntry(entry: SessionEntry, state: SessionUsageCacheState): void {
+	const contribution = usageContribution(entry);
+	addContribution(state.totals, contribution);
+	state.processedLength++;
+	state.lastEntryId = entry.id;
+	state.lastEntryRef = entry;
+	state.lastContribution = contribution;
+	state.cached = snapshotFromTotals(state.totals);
+	state.scannedEntries++;
+}
+
+function usageContribution(entry: SessionEntry): UsageContribution {
+	if (entry.type === "message") {
+		if (entry.message.role !== "assistant" && entry.message.role !== "toolResult") return emptyContribution();
+		const usage = "usage" in entry.message ? entry.message.usage : undefined;
+		if (!usage) return emptyContribution();
+		return {
+			input: usage.input,
+			output: usage.output,
+			cacheRead: usage.cacheRead,
+			cacheWrite: usage.cacheWrite,
+			cost: usage.cost.total,
+			sawUsage: true,
+		};
+	}
+	if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) {
+		return {
+			input: entry.usage.input,
+			output: entry.usage.output,
+			cacheRead: entry.usage.cacheRead,
+			cacheWrite: entry.usage.cacheWrite,
+			cost: entry.usage.cost.total,
+			sawUsage: true,
+		};
+	}
+	return emptyContribution();
+}
+
+function snapshotFromTotals(totals: UsageTotals): UsageSnapshot | undefined {
+	if (
+		totals.input === 0 &&
+		totals.output === 0 &&
+		totals.cacheRead === 0 &&
+		totals.cacheWrite === 0 &&
+		totals.cost === 0
+	)
+		return undefined;
 	return {
 		inputTokens: totals.input,
 		outputTokens: totals.output,
@@ -50,13 +220,26 @@ export function usageFromSession(session: SessionUsageSource): UsageSnapshot | u
 	};
 }
 
-function addUsage(
-	totals: UsageTotals,
-	usage: { input: number; output: number; cacheRead: number; cacheWrite: number; cost: { total: number } },
-): void {
+function addContribution(totals: UsageTotals, usage: UsageContribution): void {
 	totals.input += usage.input;
 	totals.output += usage.output;
 	totals.cacheRead += usage.cacheRead;
 	totals.cacheWrite += usage.cacheWrite;
-	totals.cost += usage.cost.total;
+	totals.cost += usage.cost;
+}
+
+function subtractContribution(totals: UsageTotals, usage: UsageContribution): void {
+	totals.input -= usage.input;
+	totals.output -= usage.output;
+	totals.cacheRead -= usage.cacheRead;
+	totals.cacheWrite -= usage.cacheWrite;
+	totals.cost -= usage.cost;
+}
+
+function emptyTotals(): UsageTotals {
+	return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0 };
+}
+
+function emptyContribution(): UsageContribution {
+	return { ...emptyTotals(), sawUsage: false };
 }

package/extension-src/pi-style/shared/box.ts

@@ -36,6 +36,17 @@ export interface BoxTheme {
 	getColorMode?(): string;
 }
 
+const THEME_CACHE_KEYS = new WeakMap<object, number>();
+let nextThemeCacheKey = 1;
+
+export function themeCacheKey(theme: object): number {
+	const cached = THEME_CACHE_KEYS.get(theme);
+	if (cached !== undefined) return cached;
+	const created = nextThemeCacheKey++;
+	THEME_CACHE_KEYS.set(theme, created);
+	return created;
+}
+
 export interface BoxedRenderOptions {
 	widthKey?: string;
 	/** Detail embedded in the top-border title after the tool name (e.g. the path). */
@@ -384,10 +395,10 @@ export function dimLine(text: string): string {
 	return `\x1b[2m${text}\x1b[22m`;
 }
 
-function boxText(theme: BoxTheme, text: string): string {
+function boxText(_theme: BoxTheme, text: string): string {
 	return dimLine(text);
 }
-function boxFrameText(theme: BoxTheme, text: string): string {
+function boxFrameText(_theme: BoxTheme, text: string): string {
 	return dimLine(text);
 }
 
@@ -567,8 +578,20 @@ function renderBoxedOutputLines(
 	let truncated = false;
 
 	for (; nextInputIndex < outputLines.length; nextInputIndex++) {
-		const line = boxedTruncatedLine(theme, outputLines[nextInputIndex] ?? "", width);
-		if (!pushBoundedLines(head, [line], headLimit)) {
+		// An output "line" may carry embedded newlines (raw tool error messages,
+		// JSON payloads). Split before boxing so every fragment gets its own
+		// border and truncation — otherwise the box frame visually breaks on the
+		// embedded rows.
+		const fragments = (outputLines[nextInputIndex] ?? "").split("\n");
+		let headExceeded = false;
+		for (const fragment of fragments) {
+			const line = boxedTruncatedLine(theme, fragment, width);
+			if (!pushBoundedLines(head, [line], headLimit)) {
+				headExceeded = true;
+				break;
+			}
+		}
+		if (headExceeded) {
 			truncated = true;
 			nextInputIndex++;
 			break;
@@ -580,9 +603,12 @@ function renderBoxedOutputLines(
 	const tail: string[] = [];
 	const tailStart = Math.max(nextInputIndex, outputLines.length - tailLimit);
 	for (let i = tailStart; i < outputLines.length; i++) {
-		const line = boxedTruncatedLine(theme, outputLines[i] ?? "", width);
-		tail.push(line);
-		if (tail.length > tailLimit) tail.splice(0, tail.length - tailLimit);
+		const fragments = (outputLines[i] ?? "").split("\n");
+		for (const fragment of fragments) {
+			const line = boxedTruncatedLine(theme, fragment, width);
+			tail.push(line);
+			if (tail.length > tailLimit) tail.splice(0, tail.length - tailLimit);
+		}
 	}
 
 	const skippedInputLines = Math.max(0, tailStart - nextInputIndex);
@@ -764,6 +790,10 @@ export function renderBoxedToolResult(
 				bodyLines.length > 0
 					? [...errorPrefix, ...bodyLines]
 					: [theme.fg("muted", `∅ ${options.emptyText ?? "(no output)"}`)];
+			// Split embedded newlines before budgeting so raw multi-line messages
+			// (tool validation errors, JSON payloads) render as proper bordered
+			// rows instead of one "line" whose embedded rows break the frame.
+			const outputFragments = outputLines.flatMap((line) => line.split("\n"));
 			const footerText = (options.footerLines ?? []).join(" · ");
 			const dividerText =
 				typeof options.dividerLabel === "function"
@@ -783,7 +813,12 @@ export function renderBoxedToolResult(
 							),
 						]),
 				boxBlankLine(theme, renderedWidth),
-				...renderBoxedOutputLines(theme, outputLines, renderedWidth, options.renderLineBudget ?? outputLines.length),
+				...renderBoxedOutputLines(
+					theme,
+					outputFragments,
+					renderedWidth,
+					options.renderLineBudget ?? outputFragments.length,
+				),
 				boxBlankLine(theme, renderedWidth),
 				boxLabeledBorder(
 					theme,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@quandev104/pi-style",
-  "version": "0.1.7",
+  "version": "0.2.1",
   "description": "A native-layout, cohesive visual style package for Pi.",
   "license": "MIT",
   "type": "module",