pi-token-burden 0.1.0

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Will Hampson
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -0,0 +1,79 @@
+# pi-token-burden
+
+See where your system prompt tokens go.
+
+A [pi](https://github.com/mariozechner/pi) extension that parses the assembled system prompt and shows a token-budget breakdown by section. Run `/token-burden` to see how much of your context window is consumed by the base prompt, AGENTS.md files, skills, SYSTEM.md overrides, and metadata.
+
+## Install
+
+```bash
+pi install git:github.com/Whamp/pi-token-burden
+```
+
+Or try it for a single session:
+
+```bash
+pi -e git:github.com/Whamp/pi-token-burden
+```
+
+## Usage
+
+Once installed, type `/token-burden` in any pi session. A TUI panel shows:
+
+```
+ Token Burden
+System Prompt: 12,450 tokens (49,798 chars) — 6.2% of 200,000 context window
+
+Base prompt                 1,250 tokens   10.0%
+SYSTEM.md / APPEND_SYSTEM.md  340 tokens    2.7%
+AGENTS.md files             2,860 tokens   23.0%
+  /home/user/.pi/agent/AGENTS.md       1,100 tokens    8.8%
+  /home/user/project/AGENTS.md         1,760 tokens   14.1%
+Skills (42)                 7,800 tokens   62.6%
+  brainstorming               180 tokens    1.4%
+  tdd                         165 tokens    1.3%
+  ...
+Metadata (date/time, cwd)     200 tokens    1.6%
+
+ Press Enter or Esc to close
+```
+
+Sections exceeding 10% of the system prompt are highlighted.
+
+## Sections
+
+| Section                          | What it measures                                                 |
+| -------------------------------- | ---------------------------------------------------------------- |
+| **Base prompt**                  | pi's built-in instructions, tool descriptions, guidelines        |
+| **SYSTEM.md / APPEND_SYSTEM.md** | Your custom system prompt overrides                              |
+| **AGENTS.md files**              | Each AGENTS.md file, listed individually                         |
+| **Skills**                       | The `<available_skills>` block, with per-skill breakdown         |
+| **Metadata**                     | The `Current date and time` / `Current working directory` footer |
+
+## Token estimation
+
+Tokens are estimated as `ceil(chars / 4)` — the same heuristic pi uses internally. This is a rough approximation, not a tokenizer count.
+
+## Requirements
+
+- [pi](https://github.com/mariozechner/pi) v0.55.1 or later
+
+## Development
+
+```bash
+git clone https://github.com/Whamp/pi-token-burden.git
+cd pi-token-burden
+pnpm install
+pnpm run test     # 15 tests
+pnpm run check    # lint, typecheck, format, dead code, duplicates, secrets, tests
+```
+
+Test locally without installing:
+
+```bash
+pi -e ./src/index.ts
+```
+
+## License
+
+[MIT](LICENSE)

package/package.json

@@ -0,0 +1,58 @@
+{
+  "name": "pi-token-burden",
+  "version": "0.1.0",
+  "description": "Pi extension that shows a token-budget breakdown of the assembled system prompt",
+  "keywords": [
+    "pi-extension",
+    "pi-package",
+    "token-budget"
+  ],
+  "license": "MIT",
+  "files": [
+    "src/",
+    "LICENSE",
+    "README.md"
+  ],
+  "type": "module",
+  "scripts": {
+    "lint": "oxlint -c .oxlintrc.json",
+    "lint:fix": "oxlint -c .oxlintrc.json --fix",
+    "format": "oxfmt --config .oxfmtrc.jsonc",
+    "format:check": "oxfmt --config .oxfmtrc.jsonc --check",
+    "typecheck": "tsc --noEmit",
+    "test": "vitest run",
+    "deadcode": "knip",
+    "duplicates": "jscpd src/",
+    "secrets": "gitleaks detect --source . --no-git",
+    "check": "bash scripts/check.sh",
+    "fix": "bash scripts/fix.sh"
+  },
+  "devDependencies": {
+    "@types/node": "^25.3.2",
+    "gitleaks": "^1.0.0",
+    "husky": "^9.1.7",
+    "jscpd": "^4.0.8",
+    "knip": "^5.85.0",
+    "lint-staged": "^16.2.7",
+    "oxfmt": "^0.35.0",
+    "oxlint": "^1.50.0",
+    "typescript": "^5.9.3",
+    "ultracite": "^7.2.4",
+    "vitest": "^4.0.18"
+  },
+  "peerDependencies": {
+    "@mariozechner/pi-coding-agent": "*",
+    "@mariozechner/pi-tui": "*",
+    "@sinclair/typebox": "*"
+  },
+  "pnpm": {
+    "onlyBuiltDependencies": [
+      "esbuild"
+    ]
+  },
+  "pi": {
+    "extensions": [
+      "./src/index.ts"
+    ]
+  }
+}

package/src/index.test.ts

@@ -0,0 +1,6 @@
+describe("extension", () => {
+  it("exports a default function", async () => {
+    const mod = await import("./index.js");
+    expectTypeOf(mod.default).toBeFunction();
+  });
+});

package/src/index.ts

@@ -0,0 +1,25 @@
+import type { ExtensionFactory } from "@mariozechner/pi-coding-agent";
+
+import { parseSystemPrompt } from "./parser.js";
+import { showReport } from "./report-view.js";
+
+const extension: ExtensionFactory = (pi) => {
+  pi.registerCommand("token-burden", {
+    description: "Show token budget breakdown of the system prompt",
+    handler: async (_args, ctx) => {
+      const prompt = ctx.getSystemPrompt();
+      const parsed = parseSystemPrompt(prompt);
+
+      const usage = ctx.getContextUsage();
+      const contextWindow = usage?.contextWindow ?? ctx.model?.contextWindow;
+
+      if (!ctx.hasUI) {
+        return;
+      }
+
+      await showReport(parsed, contextWindow, ctx);
+    },
+  });
+};
+
+export default extension;

package/src/parser.test.ts

@@ -0,0 +1,156 @@
+import { estimateTokens, parseSystemPrompt } from "./parser.js";
+import type { ParsedPrompt } from "./parser.js";
+
+describe("estimateTokens()", () => {
+  it("returns ceil(chars / 4)", () => {
+    expect(estimateTokens("")).toBe(0);
+    expect(estimateTokens("abcd")).toBe(1);
+    expect(estimateTokens("abcde")).toBe(2);
+    expect(estimateTokens("a")).toBe(1);
+  });
+});
+
+describe("parseSystemPrompt()", () => {
+  const basePrompt = [
+    "You are an expert coding assistant operating inside pi.",
+    "",
+    "Available tools:",
+    "- read: Read file contents",
+    "- bash: Execute bash commands",
+    "",
+    "Guidelines:",
+    "- Be concise",
+    "",
+    "Pi documentation (read only when the user asks about pi itself):",
+    "- Main documentation: /path/to/README.md",
+    "- Always read pi .md files completely and follow links to related docs",
+  ].join("\n");
+
+  const agentsBlock = [
+    "",
+    "",
+    "# Project Context",
+    "",
+    "Project-specific instructions and guidelines:",
+    "",
+    "## /home/user/.pi/agent/AGENTS.md",
+    "",
+    "# Global Agent Guidelines",
+    "",
+    "## Before Acting",
+    "- Read files before editing.",
+    "",
+    "## /home/user/project/AGENTS.md",
+    "",
+    "# Project Rules",
+    "",
+    "- Follow TDD.",
+  ].join("\n");
+
+  const skillsPreamble = [
+    "",
+    "",
+    "The following skills provide specialized instructions for specific tasks.",
+    "Use the read tool to load a skill's file when the task matches its description.",
+    "When a skill file references a relative path, resolve it against the skill directory.",
+    "",
+  ].join("\n");
+
+  const skillsBlock = [
+    "<available_skills>",
+    "  <skill>",
+    "    <name>brainstorming</name>",
+    "    <description>Explore user intent before implementation.</description>",
+    "    <location>/home/user/skills/brainstorming/SKILL.md</location>",
+    "  </skill>",
+    "  <skill>",
+    "    <name>tdd</name>",
+    "    <description>Test-driven development workflow.</description>",
+    "    <location>/home/user/skills/tdd/SKILL.md</location>",
+    "  </skill>",
+    "</available_skills>",
+  ].join("\n");
+
+  const metadata =
+    "\nCurrent date and time: Thursday, February 26, 2026\nCurrent working directory: /home/user/project";
+
+  it("parses a full system prompt into sections", () => {
+    const prompt =
+      basePrompt + agentsBlock + skillsPreamble + skillsBlock + metadata;
+    const result: ParsedPrompt = parseSystemPrompt(prompt);
+
+    expect(result.totalChars).toBe(prompt.length);
+    expect(result.totalTokens).toBe(Math.ceil(prompt.length / 4));
+
+    const labels = result.sections.map((s) => s.label);
+    expect(labels).toContain("Base prompt");
+    expect(labels).toContain("AGENTS.md files");
+  });
+
+  it("parses a full system prompt with correct section labels", () => {
+    const prompt =
+      basePrompt + agentsBlock + skillsPreamble + skillsBlock + metadata;
+    const result = parseSystemPrompt(prompt);
+
+    const labels = result.sections.map((s) => s.label);
+    expect(labels).toContain("Skills (2)");
+    expect(labels).toContain("Metadata (date/time, cwd)");
+  });
+
+  it("parses AGENTS.md files into children", () => {
+    const prompt = basePrompt + agentsBlock + metadata;
+    const result = parseSystemPrompt(prompt);
+
+    const agentsSection = result.sections.find((s) =>
+      s.label.includes("AGENTS.md")
+    );
+    expect(agentsSection?.children).toHaveLength(2);
+    expect(agentsSection?.children?.[0].label).toBe(
+      "/home/user/.pi/agent/AGENTS.md"
+    );
+    expect(agentsSection?.children?.[1].label).toBe(
+      "/home/user/project/AGENTS.md"
+    );
+  });
+
+  it("parses individual skills from XML", () => {
+    const prompt = basePrompt + skillsPreamble + skillsBlock + metadata;
+    const result = parseSystemPrompt(prompt);
+
+    expect(result.skills).toHaveLength(2);
+    expect(result.skills[0].name).toBe("brainstorming");
+    expect(result.skills[1].name).toBe("tdd");
+    expect(result.skills[0].chars).toBeGreaterThan(0);
+  });
+
+  it("includes skill children in skills section", () => {
+    const prompt = basePrompt + skillsPreamble + skillsBlock + metadata;
+    const result = parseSystemPrompt(prompt);
+
+    const skillsSection = result.sections.find((s) =>
+      s.label.startsWith("Skills")
+    );
+    expect(skillsSection?.children).toHaveLength(2);
+  });
+
+  it("handles a minimal prompt with no optional sections", () => {
+    const prompt = `You are a helpful assistant.${metadata}`;
+    const result = parseSystemPrompt(prompt);
+
+    expect(result.sections.length).toBeGreaterThanOrEqual(1);
+    expect(result.totalChars).toBe(prompt.length);
+  });
+
+  it("detects SYSTEM.md / APPEND_SYSTEM.md content between base and project context", () => {
+    const appendContent =
+      "\n\nCustom SYSTEM.md instructions here.\nMore custom content.";
+    const prompt = basePrompt + appendContent + agentsBlock + metadata;
+    const result = parseSystemPrompt(prompt);
+
+    const systemMdSection = result.sections.find((s) =>
+      s.label.includes("SYSTEM.md")
+    );
+    expect(systemMdSection).toBeDefined();
+    expect(systemMdSection?.chars).toBeGreaterThan(0);
+  });
+});

package/src/parser.ts

@@ -0,0 +1,244 @@
+/**
+ * Parse the assembled system prompt into measurable sections.
+ *
+ * The system prompt built by pi follows a predictable structure:
+ *   1. Base prompt (tools, guidelines, pi docs reference)
+ *   2. Optional SYSTEM.md / APPEND_SYSTEM.md content
+ *   3. Project Context (AGENTS.md files, each under `## <path>`)
+ *   4. Skills preamble + <available_skills> block
+ *   5. Date/time + cwd metadata
+ */
+
+import type {
+  AgentsFileEntry,
+  ParsedPrompt,
+  PromptSection,
+  SkillEntry,
+} from "./types.js";
+
+export type { ParsedPrompt };
+
+/** Token estimate using pi's built-in heuristic: ceil(chars / 4). */
+export function estimateTokens(text: string): number {
+  return Math.ceil(text.length / 4);
+}
+
+// ---------------------------------------------------------------------------
+// Internal helpers (defined before use to satisfy no-use-before-define)
+// ---------------------------------------------------------------------------
+
+function measure(label: string, text: string): PromptSection {
+  return { label, chars: text.length, tokens: estimateTokens(text) };
+}
+
+/** Return the smallest positive value, or -1 if none are positive. */
+function firstPositive(...values: number[]): number {
+  let min = -1;
+  for (const v of values) {
+    if (v >= 0 && (min < 0 || v < min)) {
+      min = v;
+    }
+  }
+  return min;
+}
+
+/**
+ * Find where the base system prompt ends.
+ *
+ * The base prompt ends after the pi docs reference block. We look for
+ * "- Always read pi .md files" or "- When working on pi" as the terminal
+ * marker. Falls back to the first major section boundary.
+ */
+function findBasePromptEnd(
+  prompt: string,
+  projectCtxIdx: number,
+  skillsPreambleIdx: number,
+  dateLineIdx: number
+): number {
+  const piDocsMarker = /^- (?:Always read pi|When working on pi).+$/gm;
+  let lastPiDocsEnd = -1;
+  for (const match of prompt.matchAll(piDocsMarker)) {
+    lastPiDocsEnd = match.index + match[0].length;
+  }
+
+  if (lastPiDocsEnd !== -1) {
+    return lastPiDocsEnd;
+  }
+
+  return firstPositive(projectCtxIdx, skillsPreambleIdx, dateLineIdx);
+}
+
+/** Parse `## /path/to/AGENTS.md` blocks inside the Project Context section. */
+function parseAgentsFiles(contextBlock: string): AgentsFileEntry[] {
+  const files: AgentsFileEntry[] = [];
+  // Match `## ` headings that look like file paths (start with `/`).
+  const headingPattern = /^## (\/.+)$/gm;
+  const matches = [...contextBlock.matchAll(headingPattern)];
+
+  for (let i = 0; i < matches.length; i++) {
+    const [, path] = matches[i];
+    const blockStart = matches[i].index;
+    const blockEnd =
+      i + 1 < matches.length ? matches[i + 1].index : contextBlock.length;
+    const blockText = contextBlock.slice(blockStart, blockEnd);
+    files.push({
+      path,
+      chars: blockText.length,
+      tokens: estimateTokens(blockText),
+    });
+  }
+
+  return files;
+}
+
+/** Parse `<skill>` entries from the `<available_skills>` XML block. */
+function parseSkillEntries(xmlBlock: string, out: SkillEntry[]): void {
+  const skillPattern = /<skill>([\s\S]*?)<\/skill>/g;
+  const namePattern = /<name>([\s\S]*?)<\/name>/;
+  const descPattern = /<description>([\s\S]*?)<\/description>/;
+  const locPattern = /<location>([\s\S]*?)<\/location>/;
+
+  for (const match of xmlBlock.matchAll(skillPattern)) {
+    const [fullEntry, inner] = match;
+    const name = inner.match(namePattern)?.[1]?.trim() ?? "unknown";
+    const description = inner.match(descPattern)?.[1]?.trim() ?? "";
+    const location = inner.match(locPattern)?.[1]?.trim() ?? "";
+
+    out.push({
+      name,
+      description,
+      location,
+      chars: fullEntry.length,
+      tokens: estimateTokens(fullEntry),
+    });
+  }
+}
+
+/** Compute the skills section end index, avoiding nested ternaries. */
+function findSkillsSectionEnd(
+  availableSkillsEnd: number,
+  dateLineIdx: number,
+  promptLength: number
+): number {
+  if (availableSkillsEnd !== -1) {
+    return availableSkillsEnd + "</available_skills>".length;
+  }
+  if (dateLineIdx !== -1) {
+    return dateLineIdx;
+  }
+  return promptLength;
+}
+
+// ---------------------------------------------------------------------------
+// Main parser
+// ---------------------------------------------------------------------------
+
+/**
+ * Parse a system prompt string into sections with token estimates.
+ *
+ * Uses known structural markers emitted by `buildSystemPrompt()`:
+ *   - `# Project Context` heading
+ *   - `The following skills provide specialized instructions` preamble
+ *   - `<available_skills>` / `</available_skills>` XML block
+ *   - `Current date and time:` footer
+ */
+export function parseSystemPrompt(prompt: string): ParsedPrompt {
+  const sections: PromptSection[] = [];
+  const skills: SkillEntry[] = [];
+
+  const projectCtxIdx = prompt.indexOf("\n\n# Project Context\n");
+  const skillsPreambleIdx = prompt.indexOf(
+    "\n\nThe following skills provide specialized instructions"
+  );
+  const availableSkillsStart = prompt.indexOf("<available_skills>");
+  const availableSkillsEnd = prompt.indexOf("</available_skills>");
+  const dateLineIdx = prompt.lastIndexOf("\nCurrent date and time:");
+
+  // 1. Base system prompt
+  const baseEnd = findBasePromptEnd(
+    prompt,
+    projectCtxIdx,
+    skillsPreambleIdx,
+    dateLineIdx
+  );
+  const baseText = baseEnd >= 0 ? prompt.slice(0, baseEnd) : prompt;
+  sections.push(measure("Base prompt", baseText));
+
+  // 2. Project Context / AGENTS.md files
+  if (projectCtxIdx !== -1) {
+    const contextStart = projectCtxIdx + 2; // skip leading \n\n
+    const contextEnd = firstPositive(skillsPreambleIdx, dateLineIdx);
+    const contextBlock =
+      contextEnd >= 0
+        ? prompt.slice(contextStart, contextEnd)
+        : prompt.slice(contextStart);
+
+    const agentsFiles = parseAgentsFiles(contextBlock);
+    const children = agentsFiles.map((f) => ({
+      label: f.path,
+      chars: f.chars,
+      tokens: f.tokens,
+    }));
+
+    sections.push({
+      ...measure("AGENTS.md files", contextBlock),
+      children,
+    });
+  }
+
+  // 3. Skills section
+  if (skillsPreambleIdx !== -1) {
+    const skillsSectionStart = skillsPreambleIdx + 2;
+    const skillsSectionEnd = findSkillsSectionEnd(
+      availableSkillsEnd,
+      dateLineIdx,
+      prompt.length
+    );
+    const skillsSectionText = prompt.slice(
+      skillsSectionStart,
+      skillsSectionEnd
+    );
+
+    if (availableSkillsStart !== -1 && availableSkillsEnd !== -1) {
+      const xmlBlock = prompt.slice(
+        availableSkillsStart,
+        availableSkillsEnd + "</available_skills>".length
+      );
+      parseSkillEntries(xmlBlock, skills);
+    }
+
+    const children = skills.map((s) => ({
+      label: s.name,
+      chars: s.chars,
+      tokens: s.tokens,
+    }));
+
+    sections.push({
+      ...measure(`Skills (${String(skills.length)})`, skillsSectionText),
+      children,
+    });
+  }
+
+  // 4. Metadata footer
+  if (dateLineIdx !== -1) {
+    const metaText = prompt.slice(dateLineIdx + 1);
+    sections.push(measure("Metadata (date/time, cwd)", metaText));
+  }
+
+  // 5. Detect SYSTEM.md / APPEND_SYSTEM.md gap
+  const nextSectionStart =
+    projectCtxIdx === -1 ? skillsPreambleIdx : projectCtxIdx;
+
+  if (baseEnd >= 0 && nextSectionStart >= 0 && nextSectionStart > baseEnd) {
+    const gap = prompt.slice(baseEnd, nextSectionStart);
+    const trimmed = gap.trim();
+    if (trimmed.length > 0) {
+      sections.splice(1, 0, measure("SYSTEM.md / APPEND_SYSTEM.md", trimmed));
+    }
+  }
+
+  const totalChars = prompt.length;
+  const totalTokens = estimateTokens(prompt);
+
+  return { sections, totalChars, totalTokens, skills };
+}

package/src/report-view.test.ts

@@ -0,0 +1,7 @@
+import { showReport } from "./report-view.js";
+
+describe("report-view", () => {
+  it("exports showReport function", () => {
+    expectTypeOf(showReport).toBeFunction();
+  });
+});