pi-context 1.0.0

package/README.md

@@ -0,0 +1,42 @@
+# Pi Context Extension
+
+A Git-like context management tool that allows AI agents to proactively manage their context.
+
+Inspired by kimi-cli d-mail, implementing lossless time travel on the Pi session tree.
+
+## Installation
+
+```bash
+pi install npm:pi-context
+```
+
+## Usage
+
+### For Humans
+
+Load the skill to enable the workflow:
+
+```bash
+/skill:context-management
+```
+
+View detailed context window usage and token distribution with a visual dashboard. (like `claude code /context`)
+
+```bash
+/context
+```
+
+![](img/context.png)
+
+### For Agents
+
+This extension adds the `context-management` skill with three core tools:
+
+1.  **🔖 Structure (`context_tag`)**
+`git tag` Create named milestones to structure your conversation history.
+
+2.  **📊 Monitor (`context_log`)**
+`git log` Visualize your conversation history, check token usage, and see where you are in the task tree.
+
+3.  **⏪ Compress (`context_checkout`)**
+`git checkout` Move the HEAD pointer to any tag or commit ID. Compress completed tasks into a summary to free up context window space.

package/dist/index.d.ts

@@ -0,0 +1,2 @@
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+export default function (pi: ExtensionAPI): void;

package/dist/index.js

@@ -0,0 +1,215 @@
+import { DynamicBorder } from "@mariozechner/pi-coding-agent";
+import { Type } from "@sinclair/typebox";
+import { Container, Text, Spacer } from "@mariozechner/pi-tui";
+const ContextLogParams = Type.Object({
+    limit: Type.Optional(Type.Number({ description: "History limit for visible entries (default: 50)." })),
+    verbose: Type.Optional(Type.Boolean({ description: "If true, show ALL messages. If false, collapses intermediate steps." })),
+});
+const ContextCheckoutParams = Type.Object({
+    target: Type.String({ description: "Tag name or ID." }),
+    message: Type.String({ description: "Carryover message." }),
+    tagName: Type.Optional(Type.String()),
+});
+const ContextTagParams = Type.Object({
+    name: Type.String(),
+    target: Type.Optional(Type.String()),
+});
+const formatTokens = (n) => {
+    if (n >= 1_000_000)
+        return (n / 1_000_000).toFixed(1) + "M";
+    if (n >= 1_000)
+        return Math.round(n / 1_000) + "k";
+    return n.toString();
+};
+const resolveTargetId = (sm, target) => {
+    if (target.toLowerCase() === "root") {
+        const tree = sm.getTree();
+        return tree.length > 0 ? tree[0].entry.id : target;
+    }
+    if (/^[0-9a-f]{8,}$/i.test(target))
+        return target;
+    const find = (nodes) => {
+        for (const n of nodes) {
+            if (sm.getLabel(n.entry.id) === target)
+                return n.entry.id;
+            const r = find(n.children);
+            if (r)
+                return r;
+        }
+        return null;
+    };
+    return find(sm.getTree()) || target;
+};
+export default function (pi) {
+    pi.registerTool({
+        name: "context_log",
+        label: "Context Log",
+        description: "Show history structure.",
+        parameters: ContextLogParams,
+        async execute(_id, params, _signal, _onUpdate, ctx) {
+            return { content: [{ type: "text", text: "Use /tree or context_log for history." }], details: {} };
+        }
+    });
+    pi.registerTool({
+        name: "context_checkout",
+        label: "Context Checkout",
+        description: "Navigate history.",
+        parameters: ContextCheckoutParams,
+        async execute(_id, params, _signal, _onUpdate, ctx) {
+            const sm = ctx.sessionManager;
+            const tid = resolveTargetId(sm, params.target);
+            await sm.branchWithSummary(tid, params.message);
+            return { content: [{ type: "text", text: `Jumped to ${tid.slice(0, 7)}` }], details: {} };
+        }
+    });
+    pi.registerTool({
+        name: "context_tag",
+        label: "Context Tag",
+        description: "Create checkpoint.",
+        parameters: ContextTagParams,
+        async execute(_id, params, _signal, _onUpdate, ctx) {
+            const sm = ctx.sessionManager;
+            const id = params.target ? resolveTargetId(sm, params.target) : (sm.getLeafId() ?? "");
+            pi.setLabel(id, params.name);
+            return { content: [{ type: "text", text: `Tagged ${params.name}` }], details: {} };
+        }
+    });
+    pi.registerCommand("context", {
+        description: "Show context usage visualization (Claude Code style)",
+        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;
+            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) => 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) => 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", value: systemTokens, color: "dim" },
+                    { label: "Tools", value: toolDefTokens, color: "dim" },
+                    { label: "Messages", value: msgTokens, color: "accent" },
+                    { label: "Tool Use", value: toolUseTokens, color: "success" },
+                    { label: "Tool Results", value: toolResultTokens, color: "success" },
+                ];
+                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 = [];
+                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 = [];
+                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, b.filled ? "■ " : "□ ");
+                    }
+                    gridLines.push(rowStr.trimEnd());
+                }
+                const totalUsageTitle = `${theme.fg("text", theme.bold("Total Usage".padEnd(16)))} ${theme.fg("accent", theme.bold(formatTokens(totalActual).padStart(7)))} ${theme.fg("accent", 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, 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) => theme.fg("accent", s)));
+                return {
+                    render: (w) => container.render(w),
+                    invalidate: () => container.invalidate(),
+                    handleInput: (data) => done(undefined),
+                };
+            }, { overlay: true });
+        }
+    });
+}

package/package.json

@@ -0,0 +1,41 @@
+{
+  "name": "pi-context",
+  "version": "1.0.0",
+  "description": "agent-driven context management tools",
+  "main": "dist/index.js",
+  "types": "dist/index.d.ts",
+  "files": [
+    "dist"
+  ],
+  "scripts": {},
+  "keywords": [
+    "pi-agent",
+    "pi-package",
+    "context-management",
+    "context-optimization",
+    "time-travel",
+    "checkpoint"
+  ],
+  "type": "module",
+  "author": "",
+  "license": "MIT",
+  "peerDependencies": {
+    "@mariozechner/pi-coding-agent": "*"
+  },
+  "dependencies": {
+    "@sinclair/typebox": "^0.34.13"
+  },
+  "devDependencies": {
+    "typescript": "^5.7.2",
+    "@mariozechner/pi-coding-agent": "latest"
+  },
+  "pi": {
+    "extensions": [
+      "./dist/index.js"
+    ],
+    "skills": [
+      "./skills/context-management/SKILL.md"
+    ],
+    "image": "https://github.com/ttttmr/pi-context/raw/refs/heads/main/img/context.png"
+  }
+}