pi-context 1.0.4 → 1.1.0

package/README.md

@@ -1,4 +1,4 @@
-# Pi Context Extension
+# Pi Context: Agentic Context Management for the Pi
 
 A Git-like context management tool that allows AI agents to proactively manage their context.
 
@@ -17,10 +17,10 @@ pi install npm:pi-context
 
 ### For Humans
 
-Load the skill to enable the workflow:
+Run the command to enable ACM (**A**gentic **C**ontext **M**anagement) for the current session.
 
 ```bash
-/skill:context-management
+/acm
 ```
 
 View detailed context window usage and token distribution with a visual dashboard. (like `claude code /context`)

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "pi-context",
-  "version": "1.0.4",
-  "description": "agent-driven context management tools",
+  "version": "1.1.0",
+  "description": "Agentic Context Management for the Pi",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "files": [
@@ -29,7 +29,8 @@
   },
   "pi": {
     "extensions": [
-      "./src/index.ts"
+      "./src/index.ts",
+      "./src/context.ts"
     ],
     "skills": [
       "./skills"

package/src/context.ts

@@ -0,0 +1,153 @@
+import {
+  type ExtensionAPI,
+  type SessionManager,
+  DynamicBorder,
+} from "@mariozechner/pi-coding-agent";
+import { Container, Text, Spacer } from "@mariozechner/pi-tui";
+import { formatTokens } from "./utils.js";
+
+export default function (pi: ExtensionAPI) {
+  pi.registerCommand("context", {
+    description: "Show context usage visualization",
+    handler: async (args, ctx) => {
+      const usage = await ctx.getContextUsage();
+      if (!usage) {
+        ctx.ui.notify("Context usage info not available.", "warning");
+        return;
+      }
+
+      const sm = ctx.sessionManager as SessionManager;
+      const branch = sm.getBranch();
+      const systemPrompt = ctx.getSystemPrompt();
+      const tools = pi.getActiveTools();
+      const allTools = pi.getAllTools();
+      const activeToolDefs = allTools.filter(t => tools.includes(t.name));
+
+      const estimateTokens = (text: string) => Math.ceil(text.length / 4);
+
+      let msgTokensRaw = 0;
+      let toolUseTokensRaw = 0;
+      let toolResultTokensRaw = 0;
+
+      for (const entry of branch) {
+        if (entry.type === "message") {
+          const m = entry.message;
+          if (m.role === "user") {
+            if (typeof m.content === "string") msgTokensRaw += estimateTokens(m.content);
+            else if (Array.isArray(m.content)) {
+              for (const p of m.content) if (p.type === "text") msgTokensRaw += estimateTokens(p.text);
+            }
+          } else if (m.role === "assistant") {
+            if (typeof m.content === "string") msgTokensRaw += estimateTokens(m.content);
+            else if (Array.isArray(m.content)) {
+              for (const p of m.content) {
+                if (p.type === "text") msgTokensRaw += estimateTokens(p.text);
+                if (p.type === "toolCall") toolUseTokensRaw += estimateTokens(JSON.stringify(p));
+              }
+            }
+          } else if (m.role === "toolResult") {
+            if (Array.isArray(m.content)) {
+              for (const p of m.content) if (p.type === "text") toolResultTokensRaw += estimateTokens(p.text);
+            }
+          } else if (m.role === "bashExecution") {
+            toolUseTokensRaw += estimateTokens(m.command || "");
+          }
+        } else if (entry.type === "branch_summary" || entry.type === "compaction") {
+          msgTokensRaw += estimateTokens(entry.summary || "");
+        }
+      }
+
+      const systemTokensRaw = estimateTokens(systemPrompt);
+      const toolDefTokensRaw = estimateTokens(JSON.stringify(activeToolDefs));
+      const totalActual = usage.tokens;
+      const limit = usage.contextWindow;
+
+      const totalRaw = systemTokensRaw + toolDefTokensRaw + msgTokensRaw + toolUseTokensRaw + toolResultTokensRaw;
+      const ratio = totalRaw > 0 ? (totalActual / totalRaw) : 1;
+
+      const systemTokens = Math.round(systemTokensRaw * ratio);
+      const toolDefTokens = Math.round(toolDefTokensRaw * ratio);
+      const msgTokens = Math.round(msgTokensRaw * ratio);
+      const toolUseTokens = Math.round(toolUseTokensRaw * ratio);
+      const toolResultTokens = Math.round(toolResultTokensRaw * ratio);
+
+      await ctx.ui.custom((tui, theme, kb, done) => {
+        const container = new Container();
+        container.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
+        container.addChild(new Text(theme.fg("accent", theme.bold(" Context Usage")), 1, 0));
+        container.addChild(new Spacer(1));
+
+        // Grouped by function and color
+        const categories = [
+          { label: "System Prompt", value: systemTokens, color: "muted" },
+          { label: "System Tools", value: toolDefTokens, color: "dim" },
+          { label: "Tool Call", value: toolUseTokens + toolResultTokens, color: "success" },
+          { label: "Messages", value: msgTokens, color: "accent" },
+        ];
+
+        const otherTokens = Math.max(0, totalActual - (systemTokens + toolDefTokens + msgTokens + toolUseTokens + toolResultTokens));
+        if (otherTokens > 10) categories.push({ label: "Other", value: otherTokens, color: "dim" });
+
+        categories.push({ label: "Available", value: Math.max(0, limit - totalActual), color: "borderMuted" });
+
+        const gridWidth = 10;
+        const gridHeight = 5;
+        const totalBlocks = gridWidth * gridHeight;
+
+        const blocks: { color: string, filled: boolean }[] = [];
+        categories.forEach((cat) => {
+          if (cat.label === "Available") return;
+          let count = Math.round((cat.value / limit) * totalBlocks);
+          if (count === 0 && cat.value > 0) count = 1;
+          for (let i = 0; i < count && blocks.length < totalBlocks; i++) {
+            blocks.push({ color: cat.color, filled: true });
+          }
+        });
+
+        while (blocks.length < totalBlocks) {
+          blocks.push({ color: "borderMuted", filled: false });
+        }
+
+        const gridLines: string[] = [];
+        for (let r = 0; r < gridHeight; r++) {
+          let rowStr = "";
+          for (let c = 0; c < gridWidth; c++) {
+            const b = blocks[r * gridWidth + c];
+            rowStr += theme.fg(b.color as any, b.filled ? "■ " : "□ ");
+          }
+          gridLines.push(rowStr.trimEnd());
+        }
+
+        const totalUsageTitle = `${theme.fg("text", theme.bold("Total Usage".padEnd(16)))} ${theme.fg("text", theme.bold(formatTokens(totalActual).padStart(7)))} ${theme.fg("text", theme.bold(`(${usage.percent.toFixed(1).padStart(5)}%)`))}`;
+
+        const catDetailLines = categories.map(cat => {
+          const labelStr = cat.label.padEnd(14);
+          const valStr = formatTokens(cat.value).padStart(7);
+          const rowPercent = ((cat.value / limit) * 100).toFixed(1).padStart(5);
+          const icon = cat.label === "Available" ? "□" : "■";
+          return `${theme.fg(cat.color as any, icon)} ${theme.fg("text", labelStr)} ${theme.fg("accent", valStr)} (${rowPercent}%)`;
+        });
+
+        const allDetailLines = [totalUsageTitle, "", ...catDetailLines];
+
+        const leftSideWidth = 20;
+        const maxH = Math.max(gridLines.length, allDetailLines.length);
+        for (let i = 0; i < maxH; i++) {
+          const left = (gridLines[i] || "").padEnd(leftSideWidth);
+          const right = allDetailLines[i] || "";
+          container.addChild(new Text(`    ${left}      ${right}`, 1, 0));
+        }
+
+        container.addChild(new Spacer(1));
+        container.addChild(new Text(theme.fg("dim", " Press any key to close"), 1, 0));
+        container.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
+
+        return {
+          render: (w) => container.render(w),
+          invalidate: () => container.invalidate(),
+          handleInput: (data) => done(undefined),
+        };
+      }, { overlay: true });
+    }
+  });
+}