@xprnvd/pi-context-view 0.0.0-stage → 0.8.0

package/src/ui/wheel.ts

@@ -0,0 +1,65 @@
+/**
+ * Mouse-wheel input shared by the Usage and Injections views.
+ *
+ * Pi enables mouse reporting only in fullscreen (alt-screen) mode, where its
+ * viewport forwards wheel reports to the focused overlay instead of consuming
+ * them. In regular mode the terminal keeps the wheel for its own scrollback and
+ * no report ever reaches a view, so the parser simply never matches there.
+ */
+import type { TUI } from "@earendil-works/pi-tui";
+
+/** Lines one notch scrolls when the TUI exposes no step of its own. */
+export const DEFAULT_WHEEL_SCROLL_LINES = 3;
+
+/** SGR wheel report: `ESC [ < button ; column ; row` closed by `M` or `m`. */
+const SGR_WHEEL_PATTERN = /^\x1b\[<(\d+);\d+;\d+[Mm]$/;
+
+/** Legacy X10 wheel report: `ESC [ M` plus offset-encoded button, column, and row bytes. */
+const X10_PREFIX = "\x1b[M";
+const X10_LENGTH = 6;
+const X10_BUTTON_INDEX = 3;
+const X10_BYTE_OFFSET = 32;
+
+/** Wheel reports set this button bit; the low two bits then carry axis and direction. */
+const WHEEL_BUTTON_BIT = 64;
+const DIRECTION_MASK = 3;
+
+/** One notch backwards (up) or forwards (down). */
+export type WheelDirection = -1 | 1;
+
+/**
+ * Vertical wheel direction encoded in one input chunk, or undefined for every
+ * other input, including horizontal notches. Modifier bits are ignored, so a
+ * modified notch scrolls like a plain one, as it does in pi's own viewport.
+ */
+export function parseWheelDirection(data: string): WheelDirection | undefined {
+	const sgrButton = SGR_WHEEL_PATTERN.exec(data)?.[1];
+	if (sgrButton !== undefined) return decodeWheelButton(Number.parseInt(sgrButton, 10));
+	if (data.length === X10_LENGTH && data.startsWith(X10_PREFIX)) {
+		return decodeWheelButton(data.charCodeAt(X10_BUTTON_INDEX) - X10_BYTE_OFFSET);
+	}
+	return undefined;
+}
+
+/**
+ * Lines one notch scrolls in the surrounding TUI.
+ *
+ * `TuiAltScreen` keeps its `wheelScrollLines` option private and pi never
+ * surfaces it, so read it defensively: honoring the real step keeps previews
+ * consistent with the transcript viewport, while any pi that hides, renames, or
+ * never had the field falls back to the default instead of a contradicting step.
+ */
+export function readWheelScrollLines(tui: TUI): number {
+	const lines = (tui as { readonly wheelScrollLines?: unknown }).wheelScrollLines;
+	if (typeof lines !== "number" || !Number.isFinite(lines)) return DEFAULT_WHEEL_SCROLL_LINES;
+	return Math.max(1, Math.floor(lines));
+}
+
+/** Decode one button code, rejecting non-wheel buttons and the horizontal axis. */
+function decodeWheelButton(button: number): WheelDirection | undefined {
+	if ((button & WHEEL_BUTTON_BIT) === 0) return undefined;
+	const direction = button & DIRECTION_MASK;
+	if (direction === 0) return -1;
+	if (direction === 1) return 1;
+	return undefined;
+}

package/src/usage.ts

@@ -0,0 +1,461 @@
+/**
+ * Pure context-usage classification: combine the current branch's measured
+ * prompt/tool decomposition with the live session messages into estimated
+ * category totals. No pi API access — unit-testable.
+ */
+import { type ContextEvent, type ContextUsage, convertToLlm, estimateTokens } from "@earendil-works/pi-coding-agent";
+
+import {
+	BUILT_IN_TOOLS_LABEL,
+	type ContextUsageSnapshot,
+	INSTRUCTION_FILES_LABEL,
+	type InjectionItem,
+	type InjectionSnapshot,
+	type InvisibleReasoningEstimate,
+	type ReportedContextUsage,
+	SKILLS_LABEL,
+	SYSTEM_PROMPT_LABEL,
+	type UsageCategory,
+	type UsagePreviewEntry,
+} from "./model.ts";
+import { isSameModel, type ModelIdentity } from "./model-identity.ts";
+
+/**
+ * Everything computeUsage needs; messages must already be synthetic-filtered and
+ * carry the latest request's conversation changes.
+ */
+export interface UsageInputs {
+	snapshot: InjectionSnapshot;
+	messages: ContextEvent["messages"];
+	reported?: ReportedContextUsage;
+	modelLabel?: string;
+	/**
+	 * The next request's model; thinking of other models is left out, as Pi drops it.
+	 * Omit to count all thinking.
+	 */
+	nextRequestModel?: ModelIdentity;
+	computedAt?: Date;
+	/** Auto-compaction reserve (settings `reserveTokens`); omit when auto-compaction is disabled. */
+	autoCompactReserveTokens?: number;
+}
+
+/**
+ * Estimate the current/next-request context composition. Prompt and tool
+ * categories come from the caller's current-state snapshot; message categories are
+ * classified from the session messages with the latest request's changes. System messages have already been
+ * replayed into that snapshot and must not count again. Empty categories are dropped and
+ * every aggregate equals the exact sum of its children.
+ */
+export function computeUsage(inputs: UsageInputs): ContextUsageSnapshot {
+	const categories = [
+		...classifyPromptCategories(inputs.snapshot),
+		...classifyMessages(inputs.messages, inputs.nextRequestModel),
+	].filter((category) => category.tokens > 0);
+	return {
+		computedAt: inputs.computedAt ?? new Date(),
+		modelLabel: inputs.modelLabel,
+		reported: inputs.reported,
+		categories,
+		estimatedTokens: categories.reduce((sum, category) => sum + category.tokens, 0),
+		autoCompactReserveTokens: inputs.autoCompactReserveTokens,
+	};
+}
+
+/**
+ * Flatten one category's preview entries across its breakdown, chronologically
+ * when every entry is message-backed. Raw entry text is process-local; the
+ * caller must sanitize before rendering and never log or persist it.
+ */
+export function collectPreviewEntries(category: UsageCategory): UsagePreviewEntry[] {
+	const entries: UsagePreviewEntry[] = [];
+	const visit = (node: UsageCategory): void => {
+		entries.push(...(node.entries ?? []));
+		for (const child of node.children ?? []) visit(child);
+	};
+	visit(category);
+	if (entries.length > 1 && entries.every((entry) => entry.timestamp !== undefined)) {
+		entries.sort((a, b) => (a.timestamp ?? 0) - (b.timestamp ?? 0));
+	}
+	return entries;
+}
+
+/** Convert pi's nullable ContextUsage to the undefined-based model shape. */
+export function toReportedUsage(usage: ContextUsage | undefined): ReportedContextUsage | undefined {
+	if (usage === undefined) return undefined;
+	return {
+		tokens: usage.tokens ?? undefined,
+		contextWindow: usage.contextWindow,
+		percent: usage.percent ?? undefined,
+	};
+}
+
+/**
+ * Map frozen snapshot items to prompt/tool/instruction/skill categories. System
+ * Prompt counts all of its text, including what extensions added to it.
+ */
+function classifyPromptCategories(snapshot: InjectionSnapshot): UsageCategory[] {
+	const systemPrompt: UsageCategory[] = [];
+	const builtInTools: UsageCategory[] = [];
+	const customTools: UsageCategory[] = [];
+	const mcpTools: UsageCategory[] = [];
+	const contextFiles: UsageCategory[] = [];
+	const skills: UsageCategory[] = [];
+	for (const group of snapshot.groups) {
+		for (const item of group.items) {
+			switch (item.kind) {
+				case "base-prompt":
+				case "append-prompt":
+					systemPrompt.push(leafFromItem(item));
+					break;
+				case "tool":
+					if (item.source.native) builtInTools.push(...breakdownFromItem(item));
+					else if (isMcpTool(item)) mcpTools.push(leafFromItem(item));
+					else customTools.push(leafFromItem(item));
+					break;
+				case "context-file":
+					contextFiles.push(...breakdownFromItem(item));
+					break;
+				case "skills":
+					skills.push(...breakdownFromItem(item));
+					break;
+				case "message":
+					// Custom messages live in the session; classifyMessages counts them.
+					break;
+			}
+		}
+	}
+	// Categories follow the order pi assembles them into a request: prompt text,
+	// context files, and skills first, then the tool definitions sent alongside.
+	return withoutEmpty([
+		aggregate("system-prompt", SYSTEM_PROMPT_LABEL, systemPrompt),
+		aggregate("context-files", INSTRUCTION_FILES_LABEL, contextFiles),
+		aggregate("skills", SKILLS_LABEL, skills),
+		aggregate("built-in-tools", BUILT_IN_TOOLS_LABEL, builtInTools),
+		aggregate("custom-tools", "Custom Tools", customTools),
+		aggregate("mcp-tools", "MCP Tools", mcpTools),
+	]);
+}
+
+/** Best-effort MCP attribution from the only public provenance field available. */
+function isMcpTool(item: InjectionItem): boolean {
+	return /(^|[^a-z])mcp([^a-z]|$)/i.test(`${item.source.id} ${item.source.label}`);
+}
+
+/** Tool-result shape narrowed out of the context event union. */
+type ToolResultContextMessage = Extract<ContextEvent["messages"][number], { role: "toolResult" }>;
+
+/** Custom-message shape narrowed out of the context event union. */
+type CustomContextMessage = Extract<ContextEvent["messages"][number], { role: "custom" }>;
+
+/** Classify session and request messages with preview entries; `nextRequestModel` keeps only its own thinking. */
+function classifyMessages(
+	messages: ContextEvent["messages"],
+	nextRequestModel: ModelIdentity | undefined,
+): UsageCategory[] {
+	const user: UsagePreviewEntry[] = [];
+	const assistantText: UsagePreviewEntry[] = [];
+	const assistantThinking: UsagePreviewEntry[] = [];
+	const assistantToolCalls: UsagePreviewEntry[] = [];
+	const bashExecutions: UsagePreviewEntry[] = [];
+	const compacted: UsagePreviewEntry[] = [];
+	const toolResults: ToolResultContextMessage[] = [];
+	const customMessages: CustomContextMessage[] = [];
+
+	for (const message of messages) {
+		switch (message.role) {
+			case "system":
+				// Sections and tool deltas already contribute through the replayed snapshot.
+				break;
+			case "user":
+				user.push({
+					timestamp: message.timestamp,
+					breadcrumb: ["user"],
+					tokens: estimateTokens(message),
+					text: contentToText(message.content),
+				});
+				break;
+			case "assistant": {
+				const texts = message.content.flatMap((block) => (block.type === "text" ? [block.text] : []));
+				assistantText.push(...blockEntries(message.timestamp, "text", texts));
+				// Pi drops another model's thinking and tool-call signatures; like deleted content, it gets no preview
+				if (nextRequestModel === undefined || isSameModel(message, nextRequestModel)) {
+					assistantThinking.push(...thinkingEntries(message));
+				}
+				for (const block of message.content) {
+					if (block.type !== "toolCall") continue;
+					const args = JSON.stringify(block.arguments);
+					assistantToolCalls.push({
+						timestamp: message.timestamp,
+						breadcrumb: ["assistant", block.name],
+						tokens: textTokens(block.name.length + args.length),
+						text: `${block.name}(${args})`,
+						// The arguments sit between the call parentheses this entry adds around them.
+						jsonSpan: { start: block.name.length + 1, end: block.name.length + 1 + args.length },
+					});
+				}
+				break;
+			}
+			case "toolResult":
+				toolResults.push(message);
+				break;
+			case "custom":
+				customMessages.push(message);
+				break;
+			case "bashExecution": {
+				const tokens = transformedTokens(message);
+				// The transform drops `!!` bash executions: they never reach the provider.
+				if (tokens !== undefined) {
+					bashExecutions.push({
+						timestamp: message.timestamp,
+						breadcrumb: ["bash"],
+						tokens,
+						text: `$ ${message.command}\n${message.output}`,
+					});
+				}
+				break;
+			}
+			case "branchSummary":
+				compacted.push({
+					timestamp: message.timestamp,
+					breadcrumb: ["branch"],
+					tokens: transformedTokens(message) ?? 0,
+					text: message.summary,
+				});
+				break;
+			case "compactionSummary":
+				compacted.push({
+					timestamp: message.timestamp,
+					breadcrumb: ["compaction"],
+					tokens: transformedTokens(message) ?? 0,
+					text: message.summary,
+				});
+				break;
+		}
+	}
+
+	const toolOutput = aggregate("tool-output", "Tool Output", withoutEmpty([
+		...leavesFromGroups("tool-result", Map.groupBy(toolResults, (message) => message.toolName)),
+		leaf("bash-executions", "Bash Executions", bashExecutions),
+	]));
+	return withoutEmpty([
+		leaf("user-messages", "User Messages", user),
+		leaf("assistant-messages", "Assistant Messages", assistantText),
+		leaf("assistant-thinking", "Assistant Thinking", assistantThinking),
+		leaf("tool-calls", "Tool Calls", assistantToolCalls),
+		toolOutput,
+		// One category per injected customType
+		aggregate("extensions", "Extensions",
+			leavesFromGroups("custom-message", Map.groupBy(customMessages, (message) => message.customType))),
+		leaf("compacted-data", "Compacted Data", compacted),
+	]);
+}
+
+/** One category whose estimate is the exact sum of its preview entries. */
+function leaf(id: string, label: string, entries: readonly UsagePreviewEntry[]): UsageCategory {
+	return {
+		id,
+		label,
+		tokens: entries.reduce((sum, entry) => sum + entry.tokens, 0),
+		entries,
+	};
+}
+
+/** Category carrying a snapshot item's label, estimate, and timeless content entry. */
+function leafFromItem(item: InjectionItem): UsageCategory {
+	return {
+		id: `item:${item.id}`,
+		label: item.label,
+		tokens: item.tokens,
+		entries: [{
+			breadcrumb: [item.label],
+			tokens: item.tokens,
+			text: item.text,
+			jsonSpan: item.jsonSpan,
+			// Tool parts stay a preview breakdown of the same estimate, never separate entries.
+			sections: item.sections,
+			change: item.change,
+		}],
+	};
+}
+
+/**
+ * Tokens for a message as pi's LLM transform actually sends it, so wrapper text
+ * `estimateTokens` alone cannot see (summary tags, `Ran ...` bash framing) is
+ * counted. Returns undefined for messages the transform excludes from context.
+ */
+function transformedTokens(message: ContextEvent["messages"][number]): number | undefined {
+	const converted = convertToLlm([message])[0];
+	return converted === undefined ? undefined : estimateTokens(converted);
+}
+
+/** Assistant-message shape narrowed out of the context event union. */
+type AssistantContextMessage = Extract<ContextEvent["messages"][number], { role: "assistant" }>;
+
+/**
+ * Build thinking previews while accounting once per assistant message. Provider-reported
+ * reasoning replaces a smaller visible-text estimate; signature-only estimates stay approximate metadata.
+ */
+function thinkingEntries(message: AssistantContextMessage): UsagePreviewEntry[] {
+	const texts = message.content.flatMap((block) => (block.type === "thinking" ? [block.thinking] : []));
+	const visibleTokens = textTokens(texts.reduce((sum, text) => sum + text.length, 0));
+	const reportedTokens = reportedReasoningTokens(message);
+	const countedTokens = Math.max(visibleTokens, reportedTokens ?? 0);
+	const signatureChars = message.content.reduce((sum, block) => {
+		if (block.type === "thinking") return sum + stringLength(block.thinkingSignature);
+		if (block.type === "toolCall") return sum + stringLength(block.thoughtSignature);
+		return sum;
+	}, 0);
+	const invisibleReasoning = createInvisibleReasoningEstimate(signatureChars, visibleTokens, reportedTokens);
+
+	if (texts.length === 0) {
+		if (countedTokens === 0 && invisibleReasoning === undefined) return [];
+		return [{
+			timestamp: message.timestamp,
+			breadcrumb: ["assistant"],
+			tokens: countedTokens,
+			...(invisibleReasoning === undefined ? {} : { visibleTokens: 0, invisibleReasoning }),
+			text: "",
+		}];
+	}
+
+	const allocations = allocateTextTokens(texts);
+	const entries = texts.map((text, index): UsagePreviewEntry => ({
+		timestamp: message.timestamp,
+		breadcrumb: texts.length > 1
+			? ["assistant", `thinking ${index + 1}/${texts.length}`]
+			: ["assistant"],
+		tokens: allocations[index] ?? 0,
+		text,
+	}));
+	const first = entries[0];
+	if (first !== undefined) {
+		const countedInvisibleTokens = countedTokens - visibleTokens;
+		entries[0] = {
+			...first,
+			tokens: first.tokens + countedInvisibleTokens,
+			...(invisibleReasoning === undefined
+				? {}
+				: { visibleTokens: first.tokens, invisibleReasoning }),
+		};
+	}
+	return entries;
+}
+
+/** Describe a positive invisible share or signature-size proxy without retaining its source bytes. */
+function createInvisibleReasoningEstimate(
+	signatureChars: number,
+	visibleTokens: number,
+	reportedTokens: number | undefined,
+): InvisibleReasoningEstimate | undefined {
+	if (reportedTokens !== undefined) {
+		const tokens = Math.max(0, reportedTokens - visibleTokens);
+		if (tokens === 0) return undefined;
+		return {
+			tokens,
+			basis: "provider-reported",
+			encoded: signatureChars > 0,
+		};
+	}
+	if (signatureChars === 0) return undefined;
+	// chars/4 over a signature is a rough proxy, not a bound: the server decrypts the
+	// envelope and counts the reconstructed thinking, at a model-dependent ratio that
+	// measured either side of 4. See doc/THINKING.md.
+	return {
+		tokens: textTokens(signatureChars),
+		basis: "signature-proxy",
+		encoded: true,
+	};
+}
+
+/** Accept only a finite, non-negative provider reasoning count from possibly old session data. */
+function reportedReasoningTokens(message: AssistantContextMessage): number | undefined {
+	const usage: unknown = message.usage;
+	if (typeof usage !== "object" || usage === null) return undefined;
+	const reasoning = (usage as { reasoning?: unknown }).reasoning;
+	return typeof reasoning === "number" && Number.isFinite(reasoning) && reasoning >= 0
+		? reasoning
+		: undefined;
+}
+
+/** Pool chars/4 rounding across all blocks while retaining one preview entry per block. */
+function allocateTextTokens(texts: readonly string[]): number[] {
+	const allocations = texts.map((text) => Math.floor(text.length / 4));
+	const target = textTokens(texts.reduce((sum, text) => sum + text.length, 0));
+	let remainder = target - allocations.reduce((sum, tokens) => sum + tokens, 0);
+	const rankedIndexes = texts
+		.map((text, index) => ({ index, remainder: text.length % 4 }))
+		.filter((candidate) => candidate.remainder > 0)
+		.sort((a, b) => b.remainder - a.remainder || a.index - b.index);
+	for (const candidate of rankedIndexes) {
+		if (remainder === 0) break;
+		allocations[candidate.index] = (allocations[candidate.index] ?? 0) + 1;
+		remainder--;
+	}
+	return allocations;
+}
+
+/** String length from an optional or untyped signature field without exposing its contents. */
+function stringLength(value: unknown): number {
+	return typeof value === "string" ? value.length : 0;
+}
+
+/** Per-block entries; the block-index cell appears only for multi-block messages. */
+function blockEntries(timestamp: number, kind: string, texts: readonly string[]): UsagePreviewEntry[] {
+	return texts.map((text, index) => ({
+		timestamp,
+		breadcrumb: texts.length > 1 ? ["assistant", `${kind} ${index + 1}/${texts.length}`] : ["assistant"],
+		tokens: textTokens(text.length),
+		text,
+	}));
+}
+
+/** Text rendering of string-or-block message content; non-text blocks become placeholders. */
+function contentToText(content: string | ReadonlyArray<{ type: string; text?: string }>): string {
+	if (typeof content === "string") return content;
+	return content
+		.map((block) => (block.type === "text" && block.text !== undefined ? block.text : `[${block.type}]`))
+		.join("\n");
+}
+
+/** Expand an aggregate snapshot item into its children, or itself when it has none. */
+function breakdownFromItem(item: InjectionItem): UsageCategory[] {
+	if (item.children === undefined || item.children.length === 0) return [leafFromItem(item)];
+	return item.children.map((child) => leafFromItem(child));
+}
+
+/** Parent category whose total is the exact sum of its children; undefined when empty. */
+function aggregate(id: string, label: string, children: UsageCategory[]): UsageCategory | undefined {
+	if (children.length === 0) return undefined;
+	return {
+		id,
+		label,
+		tokens: children.reduce((sum, child) => sum + child.tokens, 0),
+		children: [...children].sort((a, b) => b.tokens - a.tokens),
+	};
+}
+
+/** Keep only present categories with a non-zero estimate. */
+function withoutEmpty(categories: Array<UsageCategory | undefined>): UsageCategory[] {
+	return categories.filter((category): category is UsageCategory => category !== undefined && category.tokens > 0);
+}
+
+/** One leaf per group key, with a preview entry for each of its messages. */
+function leavesFromGroups(
+	idPrefix: string,
+	groups: ReadonlyMap<string, ReadonlyArray<ToolResultContextMessage | CustomContextMessage>>,
+): UsageCategory[] {
+	return [...groups].map(([label, messages]) => {
+		const entries = messages.map((message) => ({
+			timestamp: message.timestamp,
+			breadcrumb: [label],
+			tokens: estimateTokens(message),
+			text: contentToText(message.content),
+		}));
+		return leaf(`${idPrefix}:${label}`, label, entries);
+	});
+}
+
+/** Same chars/4 heuristic pi's estimateTokens uses for text content. */
+function textTokens(chars: number): number {
+	return Math.ceil(chars / 4);
+}