pi-context-map 0.6.1 → 0.6.2

package/CHANGELOG.md

@@ -1,5 +1,16 @@
 # 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()`.

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,25 +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)) {
@@ -123,19 +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 content = result?.content || "";
-								const w = TokenCounter.count(String(content));
+								const w = TokenCounter.count(String(JSON.stringify(content)));
 								fileTokens += w;
 								fileRegistry.set(p, {
 									path: p,
@@ -154,7 +188,7 @@ export class ContextAnalyzer {
 				continue;
 			}
 
-			// 6. Everything else
+			// 5. Everything else
 			historyTokens += TokenCounter.countMessage(msg);
 		}
 
@@ -186,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;
@@ -199,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":
@@ -228,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 {
@@ -37,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;
@@ -52,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", () => {
@@ -63,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,
@@ -108,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;
@@ -120,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",
@@ -137,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: {},
@@ -160,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: [

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-context-map",
-  "version": "0.6.1",
+  "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",