pi-context-map 0.4.0 → 0.4.3

package/extensions/index.ts

@@ -1,10 +1,14 @@
 /**
  * pi-context-map
  * Professional Context Profiler for Pi.
- * v0.4.0 - Adds live localhost server with auto-updates.
+ * v0.4.1 — Fixes session.messages crash, tool registration signature, adds tests.
  */
 
-import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import type {
+	ExtensionAPI,
+	ExtensionCommandContext,
+	AgentMessage,
+} from "pi-coding-agent";
 import { ContextAnalyzer } from "./analyzer";
 import { ReportGenerator } from "./generator";
 import { InsightEngine } from "./insights";
@@ -12,30 +16,57 @@ import { LiveReportServer } from "./live-server";
 import * as path from "node:path";
 import * as os from "node:os";
 
-const REPORT_PATH = path.join(os.homedir(), ".pi", "context-map", "report.html");
+const REPORT_PATH = path.join(
+	os.homedir(),
+	".pi",
+	"context-map",
+	"report.html",
+);
 
-export default async function piContextMap(pi: ExtensionAPI) {
+export default async function piContextMap(pi: ExtensionAPI): Promise<void> {
 	const analyzer = new ContextAnalyzer();
 	const liveServer = new LiveReportServer();
 
-	async function runAnalysis() {
-		const messages = (pi as any).session?.messages || [];
-		const currentTurn = messages.length;
+	// Accumulate messages from events — this is the correct way to access
+	// session messages in Pi. (pi as any).session?.messages does NOT exist.
+	let sessionMessages: AgentMessage[] = [];
+	let currentTurn = 0;
+
+	// Capture messages before each LLM call
+	pi.on("context", (event: any) => {
+		if (event?.messages && Array.isArray(event.messages)) {
+			sessionMessages = event.messages;
+		}
+	});
+
+	// Track turns
+	pi.on("turn_start", () => {
+		currentTurn++;
+	});
+
+	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 insights = InsightEngine.generate(composition);
 		const html = ReportGenerator.generateHTML(composition, insights);
 
-		// Write to disk (for offline access / persistence)
+		// Write to disk
 		try {
 			const fs = await import("node:fs");
 			const dir = path.dirname(REPORT_PATH);
-			if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true });
+			if (!fs.existsSync(dir)) {
+				fs.mkdirSync(dir, { recursive: true });
+			}
 			fs.writeFileSync(REPORT_PATH, html, "utf8");
 		} catch (err: any) {
-			console.error(`[pi-context-map] Failed to write report to disk: ${err.message}`);
+			console.error(`[pi-context-map] Failed to write report: ${err.message}`);
 		}
 
-		// Push to live server (if running) so the browser updates instantly
+		// Push to live server if running
 		if (liveServer.isRunning) {
 			liveServer.update(html, REPORT_PATH);
 		}
@@ -43,14 +74,15 @@ export default async function piContextMap(pi: ExtensionAPI) {
 		return { composition, insights, reportPath: REPORT_PATH };
 	}
 
-	// Start the live server on load
+	// Start live server
 	const serverUrl = await liveServer.start();
 
+	// Register /context-map command
 	pi.registerCommand("context-map", {
-		description: "Generate a visual context map with actionable insights. Use 'stop' to terminate the live server.",
-		handler: async (args: any, ctx: any) => {
-			// Handle subcommand: /context-map stop
-			if (typeof args === "string" && args.trim().toLowerCase() === "stop") {
+		description:
+			"Generate a visual context map with actionable insights. Use 'stop' to terminate the live server.",
+		handler: async (args: string, ctx: ExtensionCommandContext) => {
+			if (args.trim().toLowerCase() === "stop") {
 				liveServer.stop();
 				ctx.ui.notify("Live server stopped.", "info");
 				return;
@@ -59,66 +91,85 @@ export default async function piContextMap(pi: ExtensionAPI) {
 			ctx.ui.notify("Analyzing session context...", "info");
 			try {
 				const { insights, reportPath } = await runAnalysis();
-				const criticalCount = insights.filter((i: any) => i.severity === "critical").length;
-				const summary = criticalCount > 0
-					? `Context map generated. ${criticalCount} critical insight(s) found.`
-					: `Context map generated successfully.`;
-
+				const criticalCount = insights.filter(
+					(i) => i.severity === "critical",
+				).length;
+				const summary =
+					criticalCount > 0
+						? `Context map generated. ${criticalCount} critical insight(s) found.`
+						: "Context map generated successfully.";
 				let details = `File: ${reportPath}`;
 				if (serverUrl) {
-					details += ` Live: ${serverUrl}`;
+					details += ` | Live: ${serverUrl}`;
 				}
-				ctx.ui.notify(`${summary} ${details}`, criticalCount > 0 ? "warning" : "success");
+				ctx.ui.notify(
+					`${summary} ${details}`,
+					criticalCount > 0 ? "warning" : "success",
+				);
 			} catch (error: any) {
-				ctx.ui.notify(`Failed to generate context map: ${error.message}`, "error");
+				ctx.ui.notify(
+					`Failed to generate context map: ${error.message}`,
+					"error",
+				);
 			}
 		},
 	});
 
+	// Register the tool for agent use
 	pi.registerTool({
 		name: "context-map",
+		label: "Context Map",
 		description:
 			"Analyze the current session context composition and return actionable insights. The live localhost report will auto-update.",
 		parameters: {
 			type: "object",
 			properties: {},
 		},
-		handler: async (_ctx: any, _args: any) => {
+		async execute(
+			_params: any,
+			_signal: AbortSignal | undefined,
+			_onUpdate: ((u: any) => void) | undefined,
+			_ctx: any,
+		) {
 			try {
 				const { composition, insights } = await runAnalysis();
+				const usagePercent =
+					composition.total.tokens > 0
+						? Math.round((composition.total.tokens / 128_000) * 100)
+						: 0;
 				const summary =
 					`Context: ${composition.total.tokens.toLocaleString()} tokens total. ` +
 					`System ${composition.system.percent}%, Tools ${composition.tools.percent}%, ` +
 					`History ${composition.history.percent}%, Files ${composition.files.percent}%, ` +
 					`Summaries ${composition.summaries.percent}%. ` +
+					`Usage: ${usagePercent}% of typical 128k window. ` +
 					`${insights.length} insight(s) generated.`;
 				return {
-					summary,
-					composition: {
-						system: composition.system.tokens,
-						tools: composition.tools.tokens,
-						history: composition.history.tokens,
-						files: composition.files.tokens,
-						summaries: composition.summaries.tokens,
-						total: composition.total.tokens,
-					},
-					insights: insights.map((i: any) => ({
-						severity: i.severity,
-						title: i.title,
-						message: i.message,
-						command: i.command,
-					})),
-					liveUrl: serverUrl,
-					reportPath: REPORT_PATH,
+					type: "text" as const,
+					content: [
+						summary,
+						"",
+						...insights.map(
+							(i) => `[${i.severity.toUpperCase()}] ${i.title}: ${i.message}`,
+						),
+						serverUrl ? `Live report: ${serverUrl}` : "",
+					]
+						.filter(Boolean)
+						.join("\n"),
 				};
 			} catch (error: any) {
-				return { error: error.message };
+				return {
+					type: "text" as const,
+					content: `Error: ${error.message}`,
+					isError: true,
+				};
 			}
 		},
 	});
 
-	pi.on("session_before_compact", (event: any, ctx: any) => {
-		const tokens = event?.preparation?.tokensBefore;
+	// Auto-warning on high context before compaction
+	pi.on("session_before_compact", (_event: any, ctx: any) => {
+		const tokens = _event?.preparation?.tokensBefore;
 		if (tokens && tokens > 100_000) {
 			ctx.ui.notify(
 				`High context load detected (${(tokens / 1000).toFixed(1)}k tokens). Try /context-map to see what's consuming space.`,
@@ -127,9 +178,9 @@ export default async function piContextMap(pi: ExtensionAPI) {
 		}
 	});
 
-	// Auto-refresh: re-run analysis after each assistant message so the live view stays current
-	pi.on("message_end", async (event: any) => {
-		if (event?.message?.role === "assistant" && liveServer.isRunning) {
+	// Auto-refresh after each assistant message if server is running
+	pi.on("message_end", async (_event: any) => {
+		if (_event?.message?.role === "assistant" && liveServer.isRunning) {
 			try {
 				await runAnalysis();
 			} catch {
@@ -138,14 +189,19 @@ export default async function piContextMap(pi: ExtensionAPI) {
 		}
 	});
 
-	// Graceful shutdown: stop the live server when the session ends
+	// Graceful shutdown
 	pi.on("session_shutdown", () => {
 		liveServer.stop();
 	});
 
-	// Log the live URL once on startup
-	if (serverUrl) {
-		console.log(`[pi-context-map] Live server running at ${serverUrl}`);
-		console.log(`[pi-context-map] Run /context-map to generate a report, or /context-map stop to terminate.`);
-	}
+	// Kill server when process exits
+	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

@@ -17,7 +17,12 @@ import * as os from "node:os";
 import * as crypto from "node:crypto";
 import type { AddressInfo } from "node:net";
 
-const DEFAULT_REPORT_PATH = path.join(os.homedir(), ".pi", "context-map", "report.html");
+const DEFAULT_REPORT_PATH = path.join(
+	os.homedir(),
+	".pi",
+	"context-map",
+	"report.html",
+);
 
 /**
  * Allowed origins for SSE connections. Only localhost variants are allowed.
@@ -52,7 +57,9 @@ export class LiveReportServer {
 
 		return new Promise((resolve) => {
 			try {
-				this.server = http.createServer((req, res) => this.handleRequest(req, res));
+				this.server = http.createServer((req, res) =>
+					this.handleRequest(req, res),
+				);
 				this.server.on("error", (err) => {
 					console.error(`[pi-context-map] Server error: ${err.message}`);
 					this.stop();
@@ -62,14 +69,15 @@ export class LiveReportServer {
 					const addr = this.server?.address() as AddressInfo | null;
 					if (addr) {
 						this.port = addr.port;
-						console.log(`[pi-context-map] Live server: ${this.url}`);
 						resolve(this.url);
 					} else {
 						resolve(null);
 					}
 				});
 			} catch (err: any) {
-				console.error(`[pi-context-map] Failed to start server: ${err.message}`);
+				console.error(
+					`[pi-context-map] Failed to start server: ${err.message}`,
+				);
 				resolve(null);
 			}
 		});
@@ -118,14 +126,18 @@ export class LiveReportServer {
 				}
 				fs.writeFileSync(reportPath, html, "utf8");
 			} catch (err: any) {
-				console.error(`[pi-context-map] Failed to write report: ${err.message}`);
+				console.error(
+					`[pi-context-map] Failed to write report: ${err.message}`,
+				);
 			}
 		}
 
 		// Broadcast to all SSE clients
 		for (const client of this.clients) {
 			try {
-				client.write(`data: ${JSON.stringify({ html, timestamp: Date.now() })}\n\n`);
+				client.write(
+					`data: ${JSON.stringify({ html, timestamp: Date.now() })}\n\n`,
+				);
 			} catch (err) {
 				// Client may have disconnected; remove it
 				this.clients.delete(client);
@@ -151,7 +163,10 @@ export class LiveReportServer {
 	/**
 	 * Handle incoming HTTP requests.
 	 */
-	private handleRequest(req: http.IncomingMessage, res: http.ServerResponse): void {
+	private handleRequest(
+		req: http.IncomingMessage,
+		res: http.ServerResponse,
+	): void {
 		if (!req.url) {
 			res.writeHead(400);
 			res.end("Bad request");
@@ -243,7 +258,9 @@ export class LiveReportServer {
 
 		// Send initial state if we have content
 		if (this.currentHtml) {
-			res.write(`data: ${JSON.stringify({ html: this.currentHtml, timestamp: Date.now() })}\n\n`);
+			res.write(
+				`data: ${JSON.stringify({ html: this.currentHtml, timestamp: Date.now() })}\n\n`,
+			);
 		} else {
 			res.write(`data: ${JSON.stringify({ waiting: true })}\n\n`);
 		}

package/extensions/types/pi-coding-agent.d.ts

@@ -1,3 +1,101 @@
-declare module "@earendil-works/pi-coding-agent" {
-	export type ExtensionAPI = any;
+declare module "pi-coding-agent" {
+	export interface ExtensionAPI {
+		on(event: string, handler: Function): void;
+		registerTool<TDetails = unknown>(tool: ToolDefinition<TDetails>): void;
+		registerCommand(
+			name: string,
+			options: {
+				description?: string;
+				handler: (
+					args: string,
+					ctx: ExtensionCommandContext,
+				) => Promise<void> | void;
+			},
+		): void;
+		registerProvider(name: string, config: any): void;
+		unregisterProvider(name: string): void;
+		sendMessage(message: any, options?: any): void;
+		sendUserMessage(content: string | any[], options?: any): void;
+		appendEntry(customType: string, data?: any): void;
+		setSessionName(name: string): void;
+		getSessionName(): string | undefined;
+		setLabel(entryId: string, label: string | undefined): void;
+		getActiveTools(): string[];
+		getAllTools(): any[];
+		setActiveTools(toolNames: string[]): void;
+	}
+
+	export interface ToolDefinition<TDetails = unknown> {
+		name: string;
+		label: string;
+		description: string;
+		promptSnippet?: string;
+		promptGuidelines?: string[];
+		parameters: any;
+		execute(
+			params: any,
+			signal: AbortSignal | undefined,
+			onUpdate: ((update: TDetails) => void) | undefined,
+			ctx: ExtensionContext,
+		): Promise<{ type: "text"; content: string; isError?: boolean }>;
+		renderCall?: (args: any, theme: any) => any;
+		renderResult?: (result: any, options: any, theme: any) => any;
+	}
+
+	export interface ExtensionContext {
+		ui: ExtensionUIContext;
+		hasUI: boolean;
+		cwd: string;
+		sessionManager: any;
+		model: any;
+		isIdle(): boolean;
+		abort(): void;
+		shutdown(): void;
+		getContextUsage():
+			| { tokens: number | null; contextWindow: number; percent: number | null }
+			| undefined;
+		compact(options?: any): void;
+		getSystemPrompt(): string;
+	}
+
+	export interface ExtensionUIContext {
+		notify(
+			message: string,
+			type?: "info" | "warning" | "error" | "success",
+		): void;
+	}
+
+	export interface ExtensionCommandContext {
+		ui: ExtensionUIContext;
+		hasUI: boolean;
+		cwd: string;
+		sessionManager: any;
+		model: any;
+		isIdle(): boolean;
+		abort(): void;
+		shutdown(): void;
+		getContextUsage(): any;
+		compact(options?: any): void;
+		getSystemPrompt(): string;
+		waitForIdle(): Promise<void>;
+		newSession(options?: any): Promise<{ cancelled: boolean }>;
+		fork(entryId: string): Promise<{ cancelled: boolean }>;
+		navigateTree(
+			targetId: string,
+			options?: any,
+		): Promise<{ cancelled: boolean }>;
+		switchSession(sessionPath: string): Promise<{ cancelled: boolean }>;
+		reload(): Promise<void>;
+	}
+
+	export interface AgentMessage {
+		role: "user" | "assistant" | "system" | "tool";
+		content?: any;
+		id?: string;
+		name?: string;
+		tool_call_id?: string;
+		type?: string;
+		compactionEntry?: any;
+		timestamp?: number;
+	}
 }

package/package.json

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