pi-context-map 0.6.0 → 0.6.2

package/CHANGELOG.md

@@ -1,5 +1,21 @@
 # Changelog
 
+## [0.6.2] - 2026-06-15
+### Bug Fixes
+- **Fixed Pi message format**: Now uses `type: "toolCall"` (not `"tool_use"`) and `toolCallId` (not `tool_call_id`) to match Pi's actual `@mariozechner/pi-ai` types.
+- **System prompt detection**: Now accepts `systemPrompt` parameter from Pi's `ctx.getSystemPrompt()`. System slice no longer shows 0%.
+- **Tool results detection**: Changed to `role === "toolResult"` to match Pi's actual message format.
+- **File tracking from tool results**: Now extracts file paths from `toolResult` messages (read/write/edit tool results).
+
+### Features
+- **Message persistence**: Messages are saved via `appendEntry` on compaction to survive session reloads.
+- **Enhanced diagnostics**: `/context-map` command now shows message count, system tokens, and tool tokens in the notification.
+
+## [0.6.1] - 2026-06-15
+### Bug Fixes
+- **Fixed libuv assertion on Windows**: Removed `process.on('exit')` handler and `process.exit(0)` calls that left server handles open. Server now closes synchronously via `closeAllConnections()`.
+- **Synchronous stop()**: `isRunning` returns `false` immediately after `stop()` instead of after async callback.
+
 ## [0.6.0] - 2026-06-15
 ### Bug Fixes
 - **Fixed composition analysis**: Changed `role === "tool"` to `role === "toolResult"` to match Pi's actual message format. Tools and files now show correct percentages instead of 0%.

package/extensions/analyzer.ts

@@ -2,6 +2,14 @@
  * ContextAnalyzer
  * Parses Pi session messages to identify the active working set of files,
  * their token weights, and their temporal status.
+ *
+ * Pi message format (from @mariozechner/pi-ai):
+ *   UserMessage:       { role: "user", content: string | (TextContent | ImageContent)[] }
+ *   AssistantMessage:  { role: "assistant", content: (TextContent | ThinkingContent | ToolCall)[] }
+ *   ToolResultMessage: { role: "toolResult", toolCallId, toolName, content: (TextContent | ImageContent)[] }
+ *
+ * ToolCall:           { type: "toolCall", id, name, arguments }
+ * ToolCall.id maps to ToolResultMessage.toolCallId
  */
 import { TokenCounter } from "./token-counter";
 
@@ -37,6 +45,7 @@ export class ContextAnalyzer {
 	public analyzeByType(
 		messages: any[],
 		currentTurn: number,
+		systemPrompt?: string,
 	): ContextComposition {
 		const fileRegistry = new Map<string, FileContext>();
 
@@ -46,16 +55,20 @@ export class ContextAnalyzer {
 		let fileTokens = 0;
 		let summaryTokens = 0;
 
+		// Count system prompt tokens if provided
+		if (systemPrompt && systemPrompt.length > 0) {
+			systemTokens += TokenCounter.count(systemPrompt);
+		}
+
 		for (let index = 0; index < messages.length; index++) {
 			const msg = messages[index];
 			const turn = index + 1;
 			const role = msg.role || "";
-			const msgType = msg.type || "";
 
-			// 1. Compaction summaries
+			// 1. Compaction summaries (Pi compaction entries)
 			if (
 				role === "compaction" ||
-				msgType === "compaction" ||
+				msg.type === "compaction" ||
 				msg.customType === "compaction" ||
 				msg.compactionEntry
 			) {
@@ -63,26 +76,46 @@ export class ContextAnalyzer {
 				continue;
 			}
 
-			// 2. System messages
-			if (role === "system" || msgType === "system") {
-				systemTokens += TokenCounter.countMessage(msg);
-				continue;
-			}
-
-			// 3. Tool results (Pi uses "toolResult")
-			if (role === "toolResult" || role === "tool") {
+			// 2. Tool results (Pi uses role="toolResult")
+			if (role === "toolResult") {
 				toolTokens += TokenCounter.countMessage(msg);
+				// Track file content from tool results
+				const toolName = msg.toolName || "";
+				if (
+					toolName === "read" ||
+					toolName === "write" ||
+					toolName === "edit"
+				) {
+					const content = msg.content;
+					const path = this.extractPathFromToolResult(content);
+					if (path) {
+						const w = TokenCounter.countMessage(msg);
+						fileTokens += w;
+						if (!fileRegistry.has(path)) {
+							fileRegistry.set(path, {
+								path,
+								weight: w,
+								lastOp: {
+									type: this.getOpType(toolName),
+									turn,
+									timestamp: msg.timestamp || Date.now(),
+								},
+								status: this.calculateStatus(turn, currentTurn),
+							});
+						}
+					}
+				}
 				continue;
 			}
 
-			// 4. User messages — track file attachments
+			// 3. User messages
 			if (role === "user") {
 				historyTokens += TokenCounter.countMessage(msg);
+				// Track file attachments (images, file paths in text)
 				if (Array.isArray(msg.content)) {
 					for (const block of msg.content) {
-						if (block.type === "image" || block.type === "image_url") {
-							const p =
-								block.source?.url || block.image_url?.url || "[image]";
+						if (block.type === "image") {
+							const p = "[image]";
 							const w = TokenCounter.count(JSON.stringify(block));
 							fileTokens += w;
 							if (!fileRegistry.has(p)) {
@@ -124,23 +157,19 @@ export class ContextAnalyzer {
 				continue;
 			}
 
-			// 5. Assistant messages — track tool_use blocks
+			// 4. Assistant messages — track toolCall blocks
 			if (role === "assistant") {
 				historyTokens += TokenCounter.countMessage(msg);
 				if (Array.isArray(msg.content)) {
 					for (const block of msg.content) {
-						if (block.type === "tool_use") {
-							const input = block.input as Record<string, any>;
-							const p = this.extractPath(block.name, input);
+						// Pi uses type="toolCall" with id, name, arguments
+						if (block.type === "toolCall") {
+							const p = this.extractPath(block.name, block.arguments);
 							if (p) {
 								const opType = this.getOpType(block.name);
-								const result = this.findToolResult(
-									messages,
-									index,
-									block.id,
-								);
+								const result = this.findToolResult(messages, index, block.id);
 								const content = result?.content || "";
-								const w = TokenCounter.count(String(content));
+								const w = TokenCounter.count(String(JSON.stringify(content)));
 								fileTokens += w;
 								fileRegistry.set(p, {
 									path: p,
@@ -159,7 +188,7 @@ export class ContextAnalyzer {
 				continue;
 			}
 
-			// 6. Everything else
+			// 5. Everything else
 			historyTokens += TokenCounter.countMessage(msg);
 		}
 
@@ -168,8 +197,7 @@ export class ContextAnalyzer {
 
 		const mk = (tokens: number): ContextSlice => ({
 			tokens: Math.ceil(tokens),
-			percent:
-				totalTokens > 0 ? Math.round((tokens / totalTokens) * 100) : 0,
+			percent: totalTokens > 0 ? Math.round((tokens / totalTokens) * 100) : 0,
 		});
 
 		const files_detail = Array.from(fileRegistry.values())
@@ -192,12 +220,13 @@ export class ContextAnalyzer {
 		return this.analyzeByType(messages, currentTurn);
 	}
 
-	private extractPath(toolName: string, input: any): string | null {
+	private extractPath(toolName: string, args: any): string | null {
+		if (!args || typeof args !== "object") return null;
 		if (toolName === "read" || toolName === "write" || toolName === "edit") {
-			return typeof input.path === "string" ? input.path : null;
+			return typeof args.path === "string" ? args.path : null;
 		}
 		if (toolName === "bash") {
-			const match = input.command?.match(
+			const match = args.command?.match(
 				/(?:cat|ls|rm|mv|cp|vi|nano)\s+([^\s;]+)/,
 			);
 			return match ? match[1] : null;
@@ -205,6 +234,19 @@ export class ContextAnalyzer {
 		return null;
 	}
 
+	private extractPathFromToolResult(content: any): string | null {
+		if (typeof content === "string") return null;
+		if (Array.isArray(content)) {
+			for (const block of content) {
+				if (block.type === "text" && typeof block.text === "string") {
+					const match = block.text.match(/(?:\/|[A-Z]:\\)[\w./\\-]+\.\w+/);
+					if (match) return match[0];
+				}
+			}
+		}
+		return null;
+	}
+
 	private getOpType(toolName: string): FileOp["type"] {
 		switch (toolName) {
 			case "write":
@@ -234,13 +276,12 @@ export class ContextAnalyzer {
 		toolId: string,
 	): any {
 		for (let i = toolTurnIndex + 1; i < messages.length; i++) {
-			if (
-				messages[i].role === "toolResult" &&
-				messages[i].tool_call_id === toolId
-			) {
-				return messages[i];
+			const m = messages[i];
+			// Pi uses role="toolResult" and toolCallId (not tool_call_id)
+			if (m.role === "toolResult" && m.toolCallId === toolId) {
+				return m;
 			}
-			if (messages[i].role === "assistant") break;
+			if (m.role === "assistant") break;
 		}
 		return null;
 	}

package/extensions/index.ts

@@ -1,7 +1,7 @@
 /**
  * pi-context-map
  * Professional Context Profiler for Pi.
- * v0.5.1 — Dynamic context window, dark mode, session-unique reports.
+ * v0.6.2 — Fixed Pi message format (toolCall), system prompt detection, message persistence.
  */
 
 import type {
@@ -25,9 +25,7 @@ function makeReportPath(sessionName?: string): string {
 	const now = new Date();
 	const date = now.toISOString().split("T")[0];
 	const time = now.toTimeString().split(" ")[0].replace(/:/g, "-");
-	const safe = (sessionName || "session")
-		.replace(/[^\w.-]/g, "_")
-		.slice(0, 40);
+	const safe = (sessionName || "session").replace(/[^\w.-]/g, "_").slice(0, 40);
 	const filename = `${date}_${time}_${safe}.html`;
 	return path.join(dir, filename);
 }
@@ -39,9 +37,10 @@ export default async function piContextMap(pi: ExtensionAPI): Promise<void> {
 	let sessionMessages: AgentMessage[] = [];
 	let currentTurn = 0;
 	let contextWindow = 128_000;
+	let systemPrompt = "";
 	let currentReportPath = makeReportPath();
 
-	// Capture messages and context window from Pi system
+	// Capture messages, context window, and system prompt from Pi system
 	pi.on("context", (event: any, ctx: any) => {
 		if (event?.messages && Array.isArray(event.messages)) {
 			sessionMessages = event.messages;
@@ -54,6 +53,15 @@ export default async function piContextMap(pi: ExtensionAPI): Promise<void> {
 		} catch {
 			// Keep fallback
 		}
+		// Get system prompt from Pi
+		try {
+			const sp = ctx?.getSystemPrompt?.();
+			if (sp && typeof sp === "string") {
+				systemPrompt = sp;
+			}
+		} catch {
+			// Keep empty
+		}
 	});
 
 	pi.on("turn_start", () => {
@@ -65,13 +73,32 @@ export default async function piContextMap(pi: ExtensionAPI): Promise<void> {
 		currentReportPath = makeReportPath();
 	});
 
+	// Persist messages on compaction so they survive reload
+	pi.on("session_compact", (event: any) => {
+		if (event?.compactionEntry) {
+			try {
+				pi.appendEntry("context-map-snapshot", {
+					messages: sessionMessages.slice(-50), // Keep last 50 messages
+					turn: currentTurn,
+					timestamp: Date.now(),
+				});
+			} catch {
+				// Ignore persistence errors
+			}
+		}
+	});
+
 	async function runAnalysis(): Promise<{
 		composition: ReturnType<typeof analyzer.analyzeByType>;
 		insights: ReturnType<typeof InsightEngine.generate>;
 		reportPath: string;
 	}> {
 		const messages = sessionMessages.length > 0 ? sessionMessages : [];
-		const composition = analyzer.analyzeByType(messages, currentTurn);
+		const composition = analyzer.analyzeByType(
+			messages,
+			currentTurn,
+			systemPrompt,
+		);
 		const insights = InsightEngine.generate(composition);
 		const html = ReportGenerator.generateHTML(
 			composition,
@@ -110,7 +137,7 @@ export default async function piContextMap(pi: ExtensionAPI): Promise<void> {
 
 			ctx.ui.notify("Analyzing session context...", "info");
 			try {
-				const { insights, reportPath } = await runAnalysis();
+				const { composition, insights, reportPath } = await runAnalysis();
 				const criticalCount = insights.filter(
 					(i) => i.severity === "critical",
 				).length;
@@ -122,6 +149,9 @@ export default async function piContextMap(pi: ExtensionAPI): Promise<void> {
 				if (serverUrl) {
 					details += ` | Live: ${serverUrl}`;
 				}
+				details += ` | Messages: ${sessionMessages.length}`;
+				details += ` | System: ${composition.system.tokens}t (${composition.system.percent}%)`;
+				details += ` | Tools: ${composition.tools.tokens}t (${composition.tools.percent}%)`;
 				ctx.ui.notify(
 					`${summary} ${details}`,
 					criticalCount > 0 ? "warning" : "success",
@@ -139,7 +169,7 @@ export default async function piContextMap(pi: ExtensionAPI): Promise<void> {
 		name: "context-map",
 		label: "Context Map",
 		description:
-			"Analyze the current session context composition and return actionable insights. The live localhost report will auto-update.",
+			"Analyze the current session context composition and return actionable insights.",
 		parameters: {
 			type: "object",
 			properties: {},
@@ -162,7 +192,8 @@ export default async function piContextMap(pi: ExtensionAPI): Promise<void> {
 					`History ${composition.history.percent}%, Files ${composition.files.percent}%, ` +
 					`Summaries ${composition.summaries.percent}%. ` +
 					`Usage: ${usagePercent}% of ${(contextWindow / 1000).toFixed(0)}k window. ` +
-					`${insights.length} insight(s) generated.`;
+					`Messages: ${sessionMessages.length}. ` +
+					`${insights.length} insight(s).`;
 				return {
 					type: "text" as const,
 					content: [
@@ -211,13 +242,10 @@ export default async function piContextMap(pi: ExtensionAPI): Promise<void> {
 		liveServer.stop();
 	});
 
-	process.on("exit", () => liveServer.stop());
 	process.on("SIGINT", () => {
 		liveServer.stop();
-		process.exit(0);
 	});
 	process.on("SIGTERM", () => {
 		liveServer.stop();
-		process.exit(0);
 	});
 }

package/extensions/live-server.ts

@@ -94,18 +94,19 @@ export class LiveReportServer {
 		for (const client of this.clients) {
 			try {
 				client.end();
-			} catch (err) {
-				// Ignore errors on close
+			} catch {
+				// Ignore
 			}
 		}
 		this.clients.clear();
 
-		// Close the server
-		this.server.close((err) => {
-			if (err) {
-				console.error(`[pi-context-map] Error closing server: ${err.message}`);
-			}
-		});
+		// Force-close all connections synchronously (Node 18.2+)
+		if (typeof this.server.closeAllConnections === "function") {
+			this.server.closeAllConnections();
+		}
+
+		// Close server and reset state synchronously
+		this.server.close();
 		this.server = null;
 		this.port = 0;
 	}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-context-map",
-  "version": "0.6.0",
+  "version": "0.6.2",
   "description": "Professional context profiler for Pi that visualizes the session context window, token distribution, and integrates with Nexus packages for actionable insights.",
   "keywords": [
     "pi-package",