pi-context-viz 1.0.0

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -0,0 +1,77 @@
+# pi-context-viz
+
+Interactive context window visualizer for [pi](https://github.com/earendil-works/pi-mono) — a colored overlay showing token usage breakdown, session stats, and optimization suggestions.
+
+![screenshot](https://img.shields.io/badge/pi-package-blue)
+
+## Features
+
+- **Colored grid** — visual breakdown of context window by category
+- **Progress bar** — at-a-glance usage overview
+- **3 view modes** — Grid (`default`), Table (`T`), Detail (`Enter`)
+- **Interactive navigation** — arrow keys / `jk` to browse categories, `Enter` for details
+- **Remaining turns estimate** — how many more turns before compaction needed
+- **Cost projection** — current cost + projected cost for full context
+- **Session history** — trend of context usage across `/context` invocations
+- **Smart warnings** — context >80%, tools >25%, thinking >40%, compaction bloat
+- **Compact mode** — auto-adapts for narrow terminals (<60 cols)
+- **Save report** — `S` copies a text report to clipboard
+
+## Categories tracked
+
+| Category | Description |
+|----------|-------------|
+| System Prompt | Estimated system prompt tokens |
+| User Messages | All user input tokens |
+| Assistant Text | Model output (excluding thinking) |
+| Thinking | Hidden reasoning tokens |
+| Tool: read/bash/edit/… | Tool result tokens, per tool |
+| Compaction | Compaction summary tokens |
+| Custom Messages | Injected/custom message tokens |
+| Images | Estimated image tokens |
+| Free | Available context space |
+
+## Install
+
+```bash
+pi install git:github.com/viartemev/pi-context-viz
+```
+
+Or try without installing:
+
+```bash
+pi -e git:github.com/viartemev/pi-context-viz
+```
+
+## Usage
+
+Type `/context` in pi to open the overlay.
+
+### Keybindings
+
+| Key | Action |
+|-----|--------|
+| `↑` `↓` / `j` `k` | Navigate categories |
+| `Enter` | Show category details |
+| `T` | Toggle table/grid view |
+| `S` | Save report to clipboard |
+| `Esc` / `q` | Close overlay |
+
+### Detail mode
+
+Pressing `Enter` on a category shows:
+- **Tools**: call count, avg tokens/call, max call size
+- **User/Assistant**: message count, text vs thinking split
+- **Images**: image count
+- **Free**: estimated remaining turns
+
+## How it works
+
+1. Iterates the session branch to extract message types and tool calls
+2. Estimates tokens using ~4 chars/token heuristic (consistent with pi's estimator)
+3. Uses `ctx.getContextUsage()` for accurate total when available
+4. Renders as an interactive TUI overlay via `ctx.ui.custom()`
+
+## License
+
+MIT

package/extensions/context.ts

@@ -0,0 +1,1077 @@
+/**
+ * /context — Visualize current context usage as a colored overlay.
+ *
+ * Shows a grid of colored squares representing token usage, broken down by:
+ * - System prompt
+ * - User messages
+ * - Assistant text
+ * - Assistant thinking
+ * - Tool results (per tool: read, bash, edit, write, grep, find, ls, custom)
+ * - Compaction summaries
+ * - Custom/injected messages
+ * - Images (estimated)
+ * - Free space
+ *
+ * Interactive: arrow keys navigate categories, Enter shows details,
+ * T toggles table/grid view, S saves a text report.
+ * Also shows cache stats, cost projections, and optimization suggestions.
+ */
+
+import type {
+	ExtensionAPI,
+	ExtensionCommandContext,
+	ContextUsage,
+	Theme,
+} from "@mariozechner/pi-coding-agent";
+import type { AssistantMessage, ToolResultMessage, UserMessage } from "@mariozechner/pi-ai";
+import { matchesKey, truncateToWidth, visibleWidth } from "@mariozechner/pi-tui";
+
+// ═══════════════════════════════════════════════════════════════════════
+// Constants
+// ═══════════════════════════════════════════════════════════════════════
+
+const TOKENS_PER_CHAR = 4;
+const TOKENS_PER_IMAGE = 1600;
+const GRID_SQUARE_W = 2;
+const GRID_MIN_ROWS = 6;
+const GRID_MAX_ROWS = 15;
+
+const CONTEXT_WARN = 0.8;
+const CONTEXT_CRITICAL = 0.95;
+const CATEGORY_WARN = 0.15;
+const CATEGORY_DANGER = 0.25;
+const TOOL_BIG_CONSUMER = 0.2;
+const THINKING_HIGH = 0.4;
+const SINGLE_RESOURCE_HIGH = 0.3;
+
+const COMPACT_MODE_WIDTH = 60;
+
+// ═══════════════════════════════════════════════════════════════════════
+// Types
+// ═══════════════════════════════════════════════════════════════════════
+
+interface Category {
+	key: string;
+	label: string;
+	tokens: number;
+	colorCode: number;
+	square: string;
+	highlightedSquare: string;
+	pct: number; // precomputed percentage of contextWindow
+}
+
+interface ToolStats {
+	tokens: number;
+	callCount: number;
+	maxCallTokens: number;
+}
+
+interface ContextBreakdown {
+	categories: Category[];
+	totalTokens: number;
+	contextWindow: number;
+	percent: number | null;
+	cacheRead: number;
+	cacheWrite: number;
+	totalCost: number;
+	messageCount: number;
+	turnCount: number;
+	userMessageCount: number;
+	imageCount: number;
+	toolStats: Record<string, ToolStats>;
+}
+
+interface ContextSnapshot {
+	turnCount: number;
+	totalTokens: number;
+	contextWindow: number;
+}
+
+type ViewMode = "grid" | "table" | "detail";
+
+// ═══════════════════════════════════════════════════════════════════════
+// ANSI helpers
+// ═══════════════════════════════════════════════════════════════════════
+
+function ansi256Bg(code: number, text: string): string {
+	return `\x1b[48;5;${code}m${text}\x1b[0m`;
+}
+
+function ansi256Fg(code: number, text: string): string {
+	return `\x1b[38;5;${code}m${text}\x1b[0m`;
+}
+
+function brightenColor(code: number): number {
+	// Grays 232–255: move toward white
+	if (code >= 232) return Math.min(255, code + 12);
+	// Color cube 16–231: jump one row lighter
+	return Math.min(231, code + 36);
+}
+
+// ═══════════════════════════════════════════════════════════════════════
+// History (module-level, persists across /context invocations)
+// ═══════════════════════════════════════════════════════════════════════
+
+const contextHistory: ContextSnapshot[] = [];
+
+function addSnapshot(b: ContextBreakdown): void {
+	contextHistory.push({
+		turnCount: b.turnCount,
+		totalTokens: b.totalTokens,
+		contextWindow: b.contextWindow,
+	});
+	// Keep last 20 snapshots
+	if (contextHistory.length > 20) contextHistory.shift();
+}
+
+// ═══════════════════════════════════════════════════════════════════════
+// Token estimation
+// ═══════════════════════════════════════════════════════════════════════
+
+function estimateStringTokens(text: string): number {
+	return Math.ceil(text.length / TOKENS_PER_CHAR);
+}
+
+function estimateContentTokens(
+	content: string | Array<{ type: string; [k: string]: any }>,
+): { textTokens: number; imageTokens: number } {
+	if (typeof content === "string") {
+		return { textTokens: estimateStringTokens(content), imageTokens: 0 };
+	}
+	let text = 0;
+	let img = 0;
+	for (const block of content) {
+		if (block.type === "text") {
+			text += estimateStringTokens(block.text ?? "");
+		} else if (block.type === "image") {
+			img += TOKENS_PER_IMAGE;
+		}
+	}
+	return { textTokens: text, imageTokens: img };
+}
+
+// ═══════════════════════════════════════════════════════════════════════
+// Breakdown computation helpers
+// ═══════════════════════════════════════════════════════════════════════
+
+function estimateSystemPrompt(ctx: any): number {
+	try {
+		const sp = ctx.getSystemPrompt();
+		return sp ? estimateStringTokens(sp) : 0;
+	} catch {
+		return 0;
+	}
+}
+
+function processUserMessage(
+	msg: UserMessage,
+): { userTokens: number; imageTokens: number; imageCount: number } {
+	const { textTokens, imageTokens } = estimateContentTokens(msg.content);
+	return {
+		userTokens: textTokens,
+		imageTokens,
+		imageCount: imageTokens > 0 ? 1 : 0,
+	};
+}
+
+function processAssistantMessage(
+	msg: AssistantMessage,
+): {
+	assistantTextTokens: number;
+	thinkingTokens: number;
+	cacheRead: number;
+	cacheWrite: number;
+	totalCost: number;
+} {
+	let text = 0;
+	let thinking = 0;
+
+	for (const block of msg.content) {
+		if (block.type === "text") {
+			text += estimateStringTokens(block.text);
+		} else if (block.type === "thinking") {
+			thinking += estimateStringTokens(block.thinking);
+		}
+		// Tool call blocks: small (function name + args JSON)
+		if ((block as any).type === "tool_use") {
+			text += estimateStringTokens(JSON.stringify((block as any).arguments ?? {}));
+		}
+	}
+
+	return {
+		assistantTextTokens: text,
+		thinkingTokens: thinking,
+		cacheRead: msg.usage.cacheRead,
+		cacheWrite: msg.usage.cacheWrite,
+		totalCost: msg.usage.cost.total,
+	};
+}
+
+function processToolResult(
+	msg: ToolResultMessage,
+): { name: string; tokens: number } {
+	const name = msg.toolName || "unknown";
+	const { textTokens, imageTokens } = estimateContentTokens(msg.content);
+	return { name, tokens: textTokens + imageTokens };
+}
+
+// ── Category builder ──────────────────────────────────────────────────
+
+interface BuiltinToolInfo {
+	label: string;
+	colorCode: number;
+}
+
+const BUILTIN_TOOLS: Record<string, BuiltinToolInfo> = {
+	read: { label: "Tool: read", colorCode: 73 },
+	bash: { label: "Tool: bash", colorCode: 167 },
+	edit: { label: "Tool: edit", colorCode: 179 },
+	write: { label: "Tool: write", colorCode: 143 },
+	grep: { label: "Tool: grep", colorCode: 109 },
+	find: { label: "Tool: find", colorCode: 146 },
+	ls: { label: "Tool: ls", colorCode: 108 },
+	subagent: { label: "Tool: subagent", colorCode: 175 },
+	web_search: { label: "Tool: web_search", colorCode: 74 },
+	web_fetch: { label: "Tool: web_fetch", colorCode: 38 },
+	ask_user_question: { label: "Tool: ask_user", colorCode: 183 },
+	video_extract: { label: "Tool: video", colorCode: 204 },
+	google_image_search: { label: "Tool: img_search", colorCode: 214 },
+	youtube_search: { label: "Tool: yt_search", colorCode: 196 },
+};
+
+const CUSTOM_TOOL_COLORS = [132, 166, 130, 97, 136, 169, 103, 172];
+
+function buildCategories(
+	ctxWindow: number,
+	systemPromptTokens: number,
+	userTokens: number,
+	assistantTextTokens: number,
+	thinkingTokens: number,
+	compactionTokens: number,
+	customMessageTokens: number,
+	imageTokens: number,
+	toolStats: Record<string, ToolStats>,
+	totalTokens: number,
+): Category[] {
+	const categories: Category[] = [];
+
+	const addCat = (key: string, label: string, tokens: number, colorCode: number) => {
+		if (tokens <= 0) return;
+		categories.push({
+			key,
+			label,
+			tokens,
+			colorCode,
+			square: ansi256Bg(colorCode, "  "),
+			highlightedSquare: ansi256Bg(brightenColor(colorCode), "▐▌"),
+			pct: (tokens / ctxWindow) * 100,
+		});
+	};
+
+	addCat("system", "System Prompt", systemPromptTokens, 141);
+	addCat("user", "User Messages", userTokens, 75);
+	addCat("assistant", "Assistant Text", assistantTextTokens, 114);
+	addCat("thinking", "Thinking", thinkingTokens, 216);
+
+	// Tools — sorted by tokens descending
+	const sortedTools = Object.entries(toolStats).sort((a, b) => b[1].tokens - a[1].tokens);
+	let customColorIdx = 0;
+	for (const [name, stats] of sortedTools) {
+		const builtin = BUILTIN_TOOLS[name];
+		const colorCode = builtin?.colorCode ?? CUSTOM_TOOL_COLORS[customColorIdx++ % CUSTOM_TOOL_COLORS.length]!;
+		const label = builtin?.label ?? `Tool: ${name}`;
+		addCat(`tool:${name}`, label, stats.tokens, colorCode);
+	}
+
+	addCat("compaction", "Compaction", compactionTokens, 245);
+	addCat("custom", "Custom Messages", customMessageTokens, 183);
+	addCat("images", "Images", imageTokens, 219);
+
+	// Free space
+	const freeTokens = Math.max(0, ctxWindow - totalTokens);
+	addCat("free", "Free", freeTokens, 236);
+
+	return categories;
+}
+
+// ── Main breakdown ────────────────────────────────────────────────────
+
+function computeBreakdown(ctx: any): ContextBreakdown | null {
+	const usage: ContextUsage | undefined = ctx.getContextUsage();
+	if (!usage) return null;
+
+	const { contextWindow } = usage;
+	const branch = ctx.sessionManager.getBranch();
+
+	let systemPromptTokens = estimateSystemPrompt(ctx);
+	let userTokens = 0;
+	let assistantTextTokens = 0;
+	let thinkingTokens = 0;
+	let compactionTokens = 0;
+	let customMessageTokens = 0;
+	let imageTokens = 0;
+	const toolStats: Record<string, ToolStats> = {};
+	let cacheRead = 0;
+	let cacheWrite = 0;
+	let totalCost = 0;
+	let turnCount = 0;
+	let messageCount = 0;
+	let userMessageCount = 0;
+	let imageCount = 0;
+
+	for (const entry of branch) {
+		if (entry.type === "message") {
+			const msg = entry.message;
+			messageCount++;
+
+			if (msg.role === "user") {
+				userMessageCount++;
+				const u = processUserMessage(msg as UserMessage);
+				userTokens += u.userTokens;
+				imageTokens += u.imageTokens;
+				imageCount += u.imageCount;
+			} else if (msg.role === "assistant") {
+				turnCount++;
+				const a = processAssistantMessage(msg as AssistantMessage);
+				assistantTextTokens += a.assistantTextTokens;
+				thinkingTokens += a.thinkingTokens;
+				cacheRead += a.cacheRead;
+				cacheWrite += a.cacheWrite;
+				totalCost += a.totalCost;
+			} else if (msg.role === "toolResult") {
+				const tr = processToolResult(msg as ToolResultMessage);
+				const existing = toolStats[tr.name] ?? { tokens: 0, callCount: 0, maxCallTokens: 0 };
+				existing.tokens += tr.tokens;
+				existing.callCount++;
+				existing.maxCallTokens = Math.max(existing.maxCallTokens, tr.tokens);
+				toolStats[tr.name] = existing;
+			}
+		} else if (entry.type === "compaction" || entry.type === "branch_summary") {
+			compactionTokens += estimateStringTokens((entry as any).summary ?? "");
+		} else if (entry.type === "custom_message") {
+			const content = (entry as any).content ?? (entry as any).message?.content;
+			if (typeof content === "string") {
+				customMessageTokens += estimateStringTokens(content);
+			} else if (Array.isArray(content)) {
+				const { textTokens, imageTokens: imgs } = estimateContentTokens(content);
+				customMessageTokens += textTokens;
+				imageTokens += imgs;
+				if (imgs > 0) imageCount++;
+			}
+		}
+	}
+
+	const toolTokens = Object.values(toolStats).reduce((s, t) => s + t.tokens, 0);
+	const usedFromCategories =
+		systemPromptTokens + userTokens + assistantTextTokens + thinkingTokens +
+		toolTokens + compactionTokens + customMessageTokens + imageTokens;
+
+	const totalTokens = usage.tokens ?? usedFromCategories;
+
+	const categories = buildCategories(
+		contextWindow,
+		systemPromptTokens,
+		userTokens,
+		assistantTextTokens,
+		thinkingTokens,
+		compactionTokens,
+		customMessageTokens,
+		imageTokens,
+		toolStats,
+		totalTokens,
+	);
+
+	return {
+		categories,
+		totalTokens,
+		contextWindow,
+		percent: usage.percent,
+		cacheRead,
+		cacheWrite,
+		totalCost,
+		messageCount,
+		turnCount,
+		userMessageCount,
+		imageCount,
+		toolStats,
+	};
+}
+
+// ═══════════════════════════════════════════════════════════════════════
+// Suggestions
+// ═══════════════════════════════════════════════════════════════════════
+
+function generateSuggestions(breakdown: ContextBreakdown): string[] {
+	const s: string[] = [];
+	const pct = breakdown.percent;
+
+	// Context-level warnings
+	if (pct !== null && pct > CONTEXT_WARN) {
+		s.push("⚠ Context usage above 80% — consider /compact");
+	}
+	if (pct !== null && pct > CONTEXT_CRITICAL) {
+		s.push("🔴 Near context limit — compaction strongly recommended");
+	}
+
+	// Per-category warnings
+	for (const cat of breakdown.categories) {
+		const catPct = cat.tokens / breakdown.contextWindow;
+
+		if (cat.key === "tools" || cat.key.startsWith("tool:")) {
+			if (catPct > CATEGORY_DANGER) {
+				s.push(
+					`🔴 ${cat.label} uses ${(catPct * 100).toFixed(0)}% of context — consider /compact`,
+				);
+			} else if (catPct > CATEGORY_WARN) {
+				s.push(
+					`💡 ${cat.label} uses ${(catPct * 100).toFixed(0)}% — summarize large outputs`,
+				);
+			}
+		}
+
+		if (cat.key === "thinking" && catPct > THINKING_HIGH) {
+			s.push("💭 Thinking is >40% of context — try simplifying the task");
+		}
+
+		if (cat.key === "compaction" && catPct > 0.3) {
+			s.push("📦 Compaction summary is large — previous compaction may have been too verbose");
+		}
+	}
+
+	// Total tool usage
+	const totalToolPct =
+		breakdown.categories
+			.filter((c) => c.key.startsWith("tool:"))
+			.reduce((sum, c) => sum + c.tokens / breakdown.contextWindow, 0);
+	if (totalToolPct > 0.6) {
+		s.push("🔧 Tools consume >60% — /compact can summarize tool outputs");
+	}
+
+	return s;
+}
+
+// ═══════════════════════════════════════════════════════════════════════
+// Formatting
+// ═══════════════════════════════════════════════════════════════════════
+
+function formatTokens(n: number): string {
+	if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
+	if (n >= 1_000) return `${(n / 1_000).toFixed(1)}k`;
+	return `${n}`;
+}
+
+function formatPct(n: number, total: number): string {
+	return `${((n / total) * 100).toFixed(1)}%`;
+}
+
+// ═══════════════════════════════════════════════════════════════════════
+// Render: progress bar
+// ═══════════════════════════════════════════════════════════════════════
+
+function renderProgressBar(
+	breakdown: ContextBreakdown,
+	width: number,
+): string[] {
+	const barW = Math.min(width, 80);
+	const lines: string[] = [];
+	const nonFree = breakdown.categories.filter((c) => c.key !== "free");
+
+	let bar = "";
+	let remaining = barW;
+	for (const cat of nonFree) {
+		const segW = Math.max(1, Math.round((cat.tokens / breakdown.contextWindow) * barW));
+		const actualW = Math.min(segW, remaining);
+		if (actualW <= 0) continue;
+		// Use a thin slice of the category's color
+		bar += ansi256Bg(cat.colorCode, " ".repeat(actualW));
+		remaining -= actualW;
+	}
+	// Free space
+	if (remaining > 0) {
+		const freeCat = breakdown.categories.find((c) => c.key === "free");
+		bar += ansi256Bg(freeCat?.colorCode ?? 236, " ".repeat(remaining));
+	}
+
+	lines.push(bar);
+	return lines;
+}
+
+// ═══════════════════════════════════════════════════════════════════════
+// Render: grid
+// ═══════════════════════════════════════════════════════════════════════
+
+function renderGrid(
+	breakdown: ContextBreakdown,
+	width: number,
+	highlightedKey: string | null,
+): string[] {
+	const squareW = GRID_SQUARE_W;
+	const cols = Math.floor(width / squareW);
+	if (cols <= 0) return [];
+
+	const targetRows = Math.min(GRID_MAX_ROWS, Math.max(GRID_MIN_ROWS, Math.floor(width / 8)));
+	const cellsTotal = cols * targetRows;
+	const tokensPerCell = breakdown.contextWindow / cellsTotal;
+
+	// Build cell array (skip free — it's fill below)
+	const cells: string[] = [];
+	let cellIdx = 0;
+	for (const cat of breakdown.categories) {
+		if (cat.key === "free") continue;
+		const isHighlighted = cat.key === highlightedKey;
+		const square = isHighlighted ? cat.highlightedSquare : cat.square;
+		const numCells = Math.max(
+			cat.tokens > 0 ? 1 : 0,
+			Math.round(cat.tokens / tokensPerCell),
+		);
+		for (let i = 0; i < numCells && cellIdx < cellsTotal; i++) {
+			cells.push(square);
+			cellIdx++;
+		}
+	}
+
+	// Fill remaining with free space
+	const freeCat = breakdown.categories.find((c) => c.key === "free");
+	const freeSquare = freeCat?.square ?? ansi256Bg(236, "  ");
+	while (cellIdx < cellsTotal) {
+		cells.push(freeSquare);
+		cellIdx++;
+	}
+
+	// Center grid
+	const gridW = cols * squareW;
+	const leftPad = Math.floor((width - gridW) / 2);
+	const leftPadStr = leftPad > 0 ? " ".repeat(leftPad) : "";
+
+	const lines: string[] = [];
+	for (let row = 0; row < targetRows; row++) {
+		const start = row * cols;
+		const end = Math.min(start + cols, cells.length);
+		let line = leftPadStr;
+		for (let i = start; i < end; i++) {
+			line += cells[i];
+		}
+		lines.push(line);
+	}
+
+	return lines;
+}
+
+// ═══════════════════════════════════════════════════════════════════════
+// Render: table view
+// ═══════════════════════════════════════════════════════════════════════
+
+function renderTableView(
+	breakdown: ContextBreakdown,
+	width: number,
+	selectedIdx: number,
+): string[] {
+	const lines: string[] = [];
+	const nonFree = breakdown.categories.filter((c) => c.key !== "free");
+	const freeCat = breakdown.categories.find((c) => c.key === "free");
+	const allCats = [...nonFree];
+	if (freeCat) allCats.push(freeCat);
+
+	const barW = Math.floor(width * 0.35);
+
+	for (let i = 0; i < allCats.length; i++) {
+		const cat = allCats[i]!;
+		const isSelected = i === selectedIdx;
+		const prefix = isSelected ? "▶ " : "  ";
+		const marker = isSelected
+			? ansi256Bg(cat.colorCode, " ") + ansi256Fg(cat.colorCode, "●") + " "
+			: cat.square + " ";
+
+		// Mini bar
+		const segW = Math.max(1, Math.round((cat.tokens / breakdown.contextWindow) * barW));
+		const bar = ansi256Bg(
+			isSelected ? brightenColor(cat.colorCode) : cat.colorCode,
+			" ".repeat(segW),
+		) + " ".repeat(Math.max(0, barW - segW));
+
+		// Danger indicator
+		const catPct = cat.tokens / breakdown.contextWindow;
+		let dangerIcon = "";
+		if (cat.key !== "free" && catPct > CATEGORY_DANGER) dangerIcon = " 🔴";
+		else if (cat.key !== "free" && catPct > CATEGORY_WARN) dangerIcon = " ⚠";
+
+		const info = `${formatTokens(cat.tokens)} (${cat.pct.toFixed(1)}%)${dangerIcon}`;
+		const line = prefix + marker + bar + " " + ansi256Fg(cat.colorCode, cat.label) + "  " + info;
+		lines.push(truncateToWidth(line, width));
+	}
+
+	return lines;
+}
+
+// ═══════════════════════════════════════════════════════════════════════
+// Render: detail view (selected category)
+// ═══════════════════════════════════════════════════════════════════════
+
+function renderDetailView(
+	breakdown: ContextBreakdown,
+	cat: Category,
+	width: number,
+): string[] {
+	const lines: string[] = [];
+	const cw = breakdown.contextWindow;
+	const pct = cat.pct;
+
+	lines.push(ansi256Fg(cat.colorCode, `▌ ${cat.label}`));
+	lines.push("");
+	lines.push(`${formatTokens(cat.tokens)} tokens (${pct.toFixed(1)}% of context)`);
+
+	if (cat.key.startsWith("tool:")) {
+		const toolName = cat.key.slice(5);
+		const stats = breakdown.toolStats[toolName];
+		if (stats) {
+			const avg = stats.callCount > 0 ? Math.round(stats.tokens / stats.callCount) : 0;
+			lines.push("");
+			lines.push(`${stats.callCount} calls, avg ${formatTokens(avg)} tokens/call`);
+			lines.push(`Largest call: ${formatTokens(stats.maxCallTokens)} tokens`);
+		}
+	} else if (cat.key === "user") {
+		lines.push(`${breakdown.userMessageCount} messages`);
+	} else if (cat.key === "assistant") {
+		const thinkingCat = breakdown.categories.find((c) => c.key === "thinking");
+		const thinkingPct = thinkingCat ? thinkingCat.pct : 0;
+		lines.push(`Text: ${formatTokens(cat.tokens)} (${pct.toFixed(1)}%)`);
+		lines.push(`Thinking: ${thinkingPct.toFixed(1)}% of context`);
+	} else if (cat.key === "images") {
+		lines.push(`${breakdown.imageCount} images`);
+	} else if (cat.key === "free") {
+		const est = breakdown.turnCount > 0
+			? Math.floor((cat.tokens / breakdown.totalTokens) * breakdown.turnCount)
+			: null;
+		if (est !== null) lines.push(`Est. ~${est} turns remaining`);
+	}
+
+	// Category-specific suggestions
+	if (cat.tokens / cw > CATEGORY_DANGER && cat.key !== "free") {
+		lines.push("");
+		lines.push(`⚠ This category alone uses >${(CATEGORY_DANGER * 100).toFixed(0)}% of context.`);
+	}
+
+	return lines.map((l) => truncateToWidth(l, width));
+}
+
+// ═══════════════════════════════════════════════════════════════════════
+// Render: history trend
+// ═══════════════════════════════════════════════════════════════════════
+
+function renderHistory(_breakdown: ContextBreakdown, width: number): string[] {
+	if (contextHistory.length < 2) return [];
+
+	const lines: string[] = [];
+	lines.push("Context trend:");
+	const barW = Math.min(width - 12, 40);
+
+	for (const snap of contextHistory) {
+		const pct = snap.totalTokens / snap.contextWindow;
+		const filled = Math.round(pct * barW);
+		const bar = "█".repeat(filled) + "░".repeat(Math.max(0, barW - filled));
+		const label = `T${snap.turnCount}`.padEnd(4);
+		lines.push(
+			`${label} ${bar} ${formatTokens(snap.totalTokens)}/${formatTokens(snap.contextWindow)}`,
+		);
+	}
+
+	return lines;
+}
+
+// ═══════════════════════════════════════════════════════════════════════
+// Interactive overlay component
+// ═══════════════════════════════════════════════════════════════════════
+
+type CachedRender = { width: number; lines: string[] } | null;
+
+class ContextOverlay {
+	private mode: ViewMode = "grid";
+	private selectedIdx: number = -1; // index into non-free categories
+	private cache: CachedRender = null;
+	private cachedSelectedIdx: number = -1;
+
+	constructor(
+		private breakdown: ContextBreakdown,
+		private theme: Theme,
+		private done: (value: void) => void,
+	) {
+		addSnapshot(breakdown);
+	}
+
+	handleInput(data: string): void {
+		const nonFree = this.breakdown.categories.filter((c) => c.key !== "free");
+		const freeCat = this.breakdown.categories.find((c) => c.key === "free");
+		const allCats = [...nonFree];
+		if (freeCat) allCats.push(freeCat);
+
+		if (this.mode === "detail") {
+			if (matchesKey(data, "escape") || matchesKey(data, "q")) {
+				this.mode = "grid";
+				this.cache = null;
+			}
+			return;
+		}
+
+		if (matchesKey(data, "escape") || matchesKey(data, "q")) {
+			this.done(undefined);
+			return;
+		}
+
+		if (matchesKey(data, "enter")) {
+			if (this.selectedIdx >= 0 && this.selectedIdx < nonFree.length) {
+				this.mode = "detail";
+				this.cache = null;
+			}
+			return;
+		}
+
+		if (matchesKey(data, "t")) {
+			this.mode = this.mode === "table" ? "grid" : "table";
+			this.cache = null;
+			return;
+		}
+
+		if (matchesKey(data, "s")) {
+			this.saveReport();
+			return;
+		}
+
+		if (matchesKey(data, "up") || matchesKey(data, "k")) {
+			if (this.selectedIdx <= 0) {
+				this.selectedIdx = allCats.length - 1;
+			} else {
+				this.selectedIdx--;
+			}
+			this.cache = null;
+			return;
+		}
+
+		if (matchesKey(data, "down") || matchesKey(data, "j")) {
+			if (this.selectedIdx >= allCats.length - 1) {
+				this.selectedIdx = 0;
+			} else {
+				this.selectedIdx++;
+			}
+			this.cache = null;
+			return;
+		}
+	}
+
+	invalidate(): void {
+		this.cache = null;
+	}
+
+	render(width: number): string[] {
+		if (this.cache && this.cache.width === width && this.cachedSelectedIdx === this.selectedIdx) {
+			return this.cache.lines;
+		}
+
+		const lines: string[] = [];
+		const innerW = width - 2;
+
+		const pad = (s: string, len: number) => {
+			const vis = visibleWidth(s);
+			return s + " ".repeat(Math.max(0, len - vis));
+		};
+
+		const row = (content: string) =>
+			this.theme.fg("border", "│") + pad(` ${content}`, innerW) + this.theme.fg("border", "│");
+
+		const emptyRow = () => row("");
+
+		const hr = () =>
+			this.theme.fg("border", "│") +
+			this.theme.fg("dim", "─".repeat(innerW)) +
+			this.theme.fg("border", "│");
+
+		// ═══ Top border & title ═══
+		lines.push(this.theme.fg("border", `╭${"─".repeat(innerW)}╮`));
+
+		const modeTag = this.mode === "table" ? " [Table]" : this.mode === "detail" ? " [Detail]" : "";
+		const pctTag = this.breakdown.percent !== null ? ` (${this.breakdown.percent.toFixed(1)}%)` : "";
+		const title = `Context Window Usage${pctTag}${modeTag}`;
+		lines.push(row(this.theme.bold(this.theme.fg("accent", title))));
+
+		const sub = `${formatTokens(this.breakdown.totalTokens)} / ${formatTokens(this.breakdown.contextWindow)} tokens`;
+		lines.push(row(this.theme.fg("muted", sub)));
+
+		lines.push(emptyRow());
+
+		// ═══ Detail mode ═══
+		if (this.mode === "detail") {
+			const nonFree = this.breakdown.categories.filter((c) => c.key !== "free");
+			const cat = nonFree[this.selectedIdx];
+			if (cat) {
+				const detailLines = renderDetailView(this.breakdown, cat, innerW);
+				for (const dl of detailLines) {
+					lines.push(row(dl));
+				}
+			}
+			lines.push(emptyRow());
+			lines.push(hr());
+			lines.push(emptyRow());
+			lines.push(row(this.theme.fg("dim", "Esc: back  Arrows: navigate")));
+			lines.push(this.theme.fg("border", `╰${"─".repeat(innerW)}╯`));
+
+			this.cache = { width, lines };
+			this.cachedSelectedIdx = this.selectedIdx;
+			return lines;
+		}
+
+		// ═══ Compact mode (narrow terminals) ═══
+		if (width <= COMPACT_MODE_WIDTH) {
+			const tableLines = renderTableView(this.breakdown, innerW, this.selectedIdx);
+			for (const tl of tableLines) {
+				lines.push(row(tl));
+			}
+			lines.push(emptyRow());
+			lines.push(hr());
+			lines.push(emptyRow());
+			this.renderStatsCompact(lines, row);
+			lines.push(emptyRow());
+			lines.push(row(this.theme.fg("dim", "Arrows: navigate  Enter: detail  Esc: close")));
+			lines.push(this.theme.fg("border", `╰${"─".repeat(innerW)}╯`));
+
+			this.cache = { width, lines };
+			this.cachedSelectedIdx = this.selectedIdx;
+			return lines;
+		}
+
+		// ═══ Progress bar ═══
+		const barLines = renderProgressBar(this.breakdown, innerW);
+		for (const bl of barLines) {
+			lines.push(
+				this.theme.fg("border", "│") +
+				pad(bl, innerW) +
+				this.theme.fg("border", "│"),
+			);
+		}
+		lines.push(emptyRow());
+
+		// ═══ Grid or Table ═══
+		if (this.mode === "table") {
+			const tableLines = renderTableView(this.breakdown, innerW, this.selectedIdx);
+			for (const tl of tableLines) {
+				lines.push(row(tl));
+			}
+		} else {
+			const nonFree = this.breakdown.categories.filter((c) => c.key !== "free");
+			const highlightedKey = this.selectedIdx >= 0 && this.selectedIdx < nonFree.length
+				? nonFree[this.selectedIdx]!.key
+				: null;
+
+			const gridLines = renderGrid(this.breakdown, innerW, highlightedKey);
+			for (const gl of gridLines) {
+				lines.push(
+					this.theme.fg("border", "│") +
+					pad(gl, innerW) +
+					this.theme.fg("border", "│"),
+				);
+			}
+		}
+
+		// ═══ Remaining turns ═══
+		if (this.breakdown.turnCount > 0) {
+			const freeCat = this.breakdown.categories.find((c) => c.key === "free");
+			if (freeCat && freeCat.tokens > 0) {
+				const avgPerTurn = this.breakdown.totalTokens / this.breakdown.turnCount;
+				const remaining = Math.floor(freeCat.tokens / avgPerTurn);
+				lines.push(emptyRow());
+				lines.push(row(this.theme.fg("muted", `Est. ~${remaining} turns remaining (avg ${formatTokens(Math.round(avgPerTurn))} tokens/turn)`)));
+			}
+		}
+
+		lines.push(emptyRow());
+		lines.push(hr());
+		lines.push(emptyRow());
+
+		// ═══ Legend ═══
+		const nonFree = this.breakdown.categories.filter((c) => c.key !== "free");
+		const freeCat = this.breakdown.categories.find((c) => c.key === "free");
+		const colW = Math.floor((innerW - 2) / 2);
+
+		const formatEntry = (cat: Category, i: number, w: number): string => {
+			const isSelected = i === this.selectedIdx;
+			const prefix = isSelected ? "▶" : " ";
+			const square = isSelected ? cat.highlightedSquare : cat.square;
+			const label = `${prefix}${square} ${ansi256Fg(cat.colorCode, cat.label)}`;
+			const value = this.theme.fg(
+				isSelected ? "accent" : "dim",
+				`${formatTokens(cat.tokens)} (${cat.pct.toFixed(1)}%)`,
+			);
+			return pad(`${label} ${value}`, w);
+		};
+
+		for (let i = 0; i < nonFree.length; i += 2) {
+			const left = nonFree[i]!;
+			const right = nonFree[i + 1];
+			let content = " " + formatEntry(left, i, colW);
+			if (right) {
+				content += formatEntry(right, i + 1, colW);
+			}
+			lines.push(row(content));
+		}
+
+		if (freeCat && freeCat.tokens > 0) {
+			const isSelected = this.selectedIdx === nonFree.length;
+			const prefix = isSelected ? "▶" : " ";
+			const square = isSelected ? freeCat.highlightedSquare : freeCat.square;
+			const label = `${prefix}${square} ${ansi256Fg(freeCat.colorCode, freeCat.label)}`;
+			const value = this.theme.fg(
+				isSelected ? "accent" : "dim",
+				`${formatTokens(freeCat.tokens)} (${freeCat.pct.toFixed(1)}%)`,
+			);
+			lines.push(row(`${label} ${value}`));
+		}
+
+		lines.push(emptyRow());
+		lines.push(hr());
+		lines.push(emptyRow());
+
+		// ═══ Stats ═══
+		lines.push(row(this.theme.fg("accent", this.theme.bold("Session Stats"))));
+		lines.push(row(this.renderStatsLine(innerW)));
+
+		// Cost projection
+		if (this.breakdown.totalCost > 0 && this.breakdown.totalTokens > 0) {
+			const costPerTok = this.breakdown.totalCost / this.breakdown.totalTokens;
+			const projected = costPerTok * this.breakdown.contextWindow;
+			lines.push(row(
+				this.theme.fg("muted", `Cost: $${this.breakdown.totalCost.toFixed(4)}`) +
+				this.theme.fg("dim", `  |  ~$${projected.toFixed(2)} projected for full context`),
+			));
+		}
+
+		// ═══ Warnings ═══
+		const suggestions = generateSuggestions(this.breakdown);
+		if (suggestions.length > 0) {
+			lines.push(emptyRow());
+			lines.push(hr());
+			lines.push(emptyRow());
+			for (const sugg of suggestions) {
+				lines.push(row(this.theme.fg("warning", sugg)));
+			}
+		}
+
+		// ═══ History ═══
+		const histLines = renderHistory(this.breakdown, innerW);
+		if (histLines.length > 0) {
+			lines.push(emptyRow());
+			lines.push(hr());
+			lines.push(emptyRow());
+			for (const hl of histLines) {
+				lines.push(row(this.theme.fg("dim", hl)));
+			}
+		}
+
+		lines.push(emptyRow());
+
+		// ═══ Help ═══
+		lines.push(row(this.theme.fg("dim", "Arrows: navigate  Enter: details  T: grid/table  S: save report  Esc: close")));
+
+		lines.push(this.theme.fg("border", `╰${"─".repeat(innerW)}╯`));
+
+		this.cache = { width, lines };
+		this.cachedSelectedIdx = this.selectedIdx;
+		return lines;
+	}
+
+	private renderStatsLine(innerW: number): string {
+		const b = this.breakdown;
+		const parts = [
+			`Turns: ${b.turnCount}`,
+			`Messages: ${b.messageCount}`,
+			`Cache read: ${formatTokens(b.cacheRead)}`,
+			`Cache write: ${formatTokens(b.cacheWrite)}`,
+			`Cost: $${b.totalCost.toFixed(4)}`,
+		];
+		return parts.map((p) => this.theme.fg("muted", p)).join(this.theme.fg("dim", "  │  "));
+	}
+
+	private renderStatsCompact(lines: string[], row: (s: string) => string): void {
+		const b = this.breakdown;
+		lines.push(row(this.theme.fg("accent", "Session Stats")));
+		lines.push(row(this.theme.fg("muted", `Turns: ${b.turnCount}  Messages: ${b.messageCount}  Cost: $${b.totalCost.toFixed(4)}`)));
+		if (b.cacheRead > 0 || b.cacheWrite > 0) {
+			lines.push(row(this.theme.fg("muted", `Cache read: ${formatTokens(b.cacheRead)}  Cache write: ${formatTokens(b.cacheWrite)}`)));
+		}
+	}
+
+	private saveReport(): void {
+		const b = this.breakdown;
+		const lines: string[] = [];
+		lines.push("=== Context Usage Report ===");
+		lines.push(`Total: ${formatTokens(b.totalTokens)} / ${formatTokens(b.contextWindow)} (${b.percent?.toFixed(1) ?? "?"}%)`);
+		lines.push(`Turns: ${b.turnCount}  Messages: ${b.messageCount}`);
+		lines.push(`Cost: $${b.totalCost.toFixed(4)}`);
+		lines.push("");
+		lines.push("--- Categories ---");
+		for (const cat of b.categories) {
+			lines.push(`${cat.label}: ${formatTokens(cat.tokens)} (${cat.pct.toFixed(1)}%)`);
+		}
+		lines.push("");
+		lines.push("--- Tool Details ---");
+		for (const [name, stats] of Object.entries(b.toolStats)) {
+			lines.push(`${name}: ${stats.callCount} calls, ${formatTokens(stats.tokens)} tokens, max ${formatTokens(stats.maxCallTokens)}`);
+		}
+		lines.push("");
+		const suggestions = generateSuggestions(b);
+		if (suggestions.length > 0) {
+			lines.push("--- Suggestions ---");
+			for (const s of suggestions) lines.push(s);
+		}
+		lines.push("");
+
+		const report = lines.join("\n");
+		// Copy to clipboard via OSC 52 if supported, otherwise just dump
+		process.stdout.write(`\x1b]52;c;${Buffer.from(report).toString("base64")}\x07`);
+	}
+}
+
+// ═══════════════════════════════════════════════════════════════════════
+// Extension entry point
+// ═══════════════════════════════════════════════════════════════════════
+
+export default function (pi: ExtensionAPI) {
+	pi.registerCommand("context", {
+		description: "Visualize current context usage as a colored grid",
+		handler: async (_args: string, ctx: ExtensionCommandContext) => {
+			const breakdown = computeBreakdown(ctx);
+			if (!breakdown) {
+				ctx.ui.notify("No context usage data available yet. Send a message first.", "warning");
+				return;
+			}
+
+			await ctx.ui.custom<void>(
+				(tui, theme, _keybindings, done) => {
+					const overlay = new ContextOverlay(breakdown, theme, done);
+
+					return {
+						handleInput(data: string) {
+							overlay.handleInput(data);
+							tui.requestRender();
+						},
+
+						render(width: number): string[] {
+							return overlay.render(width);
+						},
+
+						invalidate() {
+							overlay.invalidate();
+						},
+					};
+				},
+				{
+					overlay: true,
+					overlayOptions: {
+						anchor: "center",
+						width: "80%",
+						maxWidth: 100,
+						minWidth: 40,
+						maxHeight: "90%",
+					},
+				},
+			);
+		},
+	});
+}

package/package.json

@@ -0,0 +1,23 @@
+{
+  "name": "pi-context-viz",
+  "version": "1.0.0",
+  "description": "Interactive context window visualizer for pi — colored grid overlay with token breakdown, stats, and optimization suggestions",
+  "type": "module",
+  "keywords": ["pi-package", "pi-extension", "context-window", "token-visualizer"],
+  "author": "",
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/viartemev/pi-context-viz"
+  },
+  "pi": {
+    "extensions": ["./extensions"]
+  },
+  "peerDependencies": {
+    "@earendil-works/pi-coding-agent": "*",
+    "@earendil-works/pi-ai": "*",
+    "@earendil-works/pi-tui": "*"
+  },
+  "dependencies": {},
+  "devDependencies": {}
+}