feat: enhance /context with detailed token usage breakdown

src/agent/compaction.ts

@@ -22,7 +22,18 @@ const MAX_SEND_CHARS = 240_000;    // hard cap on chars sent (~60K tokens)
 
 /** Output-token headroom reserved out of the tier window before budgeting input. */
 const OUTPUT_RESERVE_TOKENS = 8_000;
-const CHARS_PER_TOKEN = 4;
+/** Shared chars-per-token estimate — used by compaction, /context, and sidebar. */
+export const CHARS_PER_TOKEN = 4;
+
+/** Estimate token count from character length (bytes/4). */
+export function estimateTokensFromChars(chars: number): number {
+  return Math.max(0, Math.round(chars / CHARS_PER_TOKEN));
+}
+
+/** Estimate token count from text content. */
+export function estimateTokensFromText(text: string): number {
+  return estimateTokensFromChars(text.length);
+}
 
 const SEARCH_TOOLS = new Set(["grep", "glob", "list_dir", "web_search", "web_fetch", "todo_read"]);
 const MUTATING_TOOLS = new Set(["write_file", "edit_file", "multi_edit", "apply_patch"]);

src/agent/context-breakdown.test.ts

@@ -0,0 +1,87 @@
+import { expect, test } from "bun:test";
+import type { Message, ToolDefinition } from "../api/client.js";
+import { CORE_SYSTEM_PROMPT } from "./system-prompt.js";
+import {
+  classifySystemBucket,
+  computeContextBreakdown,
+  formatContextBar,
+} from "./context-breakdown.js";
+
+test("classifySystemBucket: core, environment, project rules, mode", () => {
+  expect(classifySystemBucket(CORE_SYSTEM_PROMPT)).toBe("core");
+  expect(classifySystemBucket("# Environment\nWorking directory: /tmp")).toBe("environment");
+  expect(classifySystemBucket("# Project rules (from AGENTS.md)\n\nBe nice.")).toBe("projectRules");
+  expect(classifySystemBucket("# Mode: Build\nImplement directly.")).toBe("mode");
+  expect(classifySystemBucket("custom system note")).toBe("other");
+});
+
+test("computeContextBreakdown: splits system vs conversation and tool schemas", () => {
+  const msgs: Message[] = [
+    { role: "system", content: CORE_SYSTEM_PROMPT },
+    { role: "system", content: "# Environment\nWorking directory: /proj" },
+    { role: "user", content: "hello" },
+    { role: "assistant", content: "hi there" },
+    { role: "tool", content: "file contents here" },
+  ];
+  const tools: ToolDefinition[] = [{
+    type: "function",
+    function: { name: "read_file", description: "Read a file", parameters: { type: "object", properties: {} } },
+  }];
+
+  const b = computeContextBreakdown(msgs, tools);
+  expect(b.system.core).toBeGreaterThan(0);
+  expect(b.system.environment).toBeGreaterThan(0);
+  expect(b.conversation.userMsgs).toBe(1);
+  expect(b.conversation.assistantMsgs).toBe(1);
+  expect(b.conversation.toolMsgs).toBe(1);
+  expect(b.toolCount).toBe(1);
+  expect(b.toolSchemas).toBeGreaterThan(0);
+  expect(b.estimatedTotal).toBe(b.messagesTotal + b.toolSchemas);
+});
+
+test("computeContextBreakdown: detects trimmed and compacted markers", () => {
+  const msgs: Message[] = [
+    { role: "assistant", content: "[Context compacted — earlier turns summarized]" },
+    { role: "tool", content: "partial\n[… 500 chars trimmed — grep]" },
+  ];
+  const b = computeContextBreakdown(msgs);
+  expect(b.compactedStub).toBe(true);
+  expect(b.trimmedToolResults).toBe(1);
+});
+
+test("computeContextBreakdown: counts assistant tool_calls tokens", () => {
+  const msgs: Message[] = [{
+    role: "assistant",
+    content: "",
+    tool_calls: [{
+      id: "tc1",
+      type: "function",
+      function: { name: "read_file", arguments: JSON.stringify({ path: "src/main.ts" }) },
+    }],
+  }];
+  const b = computeContextBreakdown(msgs);
+  expect(b.conversation.assistant).toBeGreaterThan(0);
+});
+
+test("classifySystemBucket: /init-style project rules", () => {
+  expect(classifySystemBucket("Project rules (from .klaatai/rules.md):\n\nUse tabs.")).toBe("projectRules");
+});
+
+test("computeContextBreakdown: estimatedTotal equals sum of parts", () => {
+  const msgs: Message[] = [
+    { role: "system", content: CORE_SYSTEM_PROMPT },
+    { role: "user", content: "hi" },
+  ];
+  const tools: ToolDefinition[] = [{
+    type: "function",
+    function: { name: "grep", description: "search", parameters: { type: "object", properties: {} } },
+  }];
+  const b = computeContextBreakdown(msgs, tools);
+  expect(b.estimatedTotal).toBe(b.messagesTotal + b.toolSchemas);
+});
+
+test("formatContextBar: clamps and fills proportionally", () => {
+  expect(formatContextBar(0, 10)).toBe("░░░░░░░░░░");
+  expect(formatContextBar(100, 10)).toBe("██████████");
+  expect(formatContextBar(50, 10)).toBe("█████░░░░░");
+});

src/agent/context-breakdown.ts

@@ -0,0 +1,142 @@
+/**
+ * Context-window breakdown — single source of truth for token estimates.
+ *
+ * Uses the same chars÷4 heuristic as compaction.ts so /context, compaction
+ * budgets, and sidebar displays stay consistent.
+ */
+
+import type { Message, ToolDefinition } from "../api/client.js";
+import { CORE_SYSTEM_PROMPT, MODE_PROMPTS } from "./system-prompt.js";
+import { estimateTokensFromText, estimateTokensFromChars } from "./compaction.js";
+
+export interface SystemTokenBreakdown {
+  core: number;
+  environment: number;
+  projectRules: number;
+  mode: number;
+  other: number;
+  total: number;
+}
+
+export interface ConversationTokenBreakdown {
+  user: number;
+  assistant: number;
+  tool: number;
+  userMsgs: number;
+  assistantMsgs: number;
+  toolMsgs: number;
+  total: number;
+}
+
+export interface ContextBreakdown {
+  system: SystemTokenBreakdown;
+  conversation: ConversationTokenBreakdown;
+  toolSchemas: number;
+  toolCount: number;
+  messagesTotal: number;
+  estimatedTotal: number;
+  trimmedToolResults: number;
+  compactedStub: boolean;
+}
+
+const MODE_PROMPT_VALUES = new Set(Object.values(MODE_PROMPTS));
+
+function messageChars(m: Message): number {
+  let chars = 0;
+  if (typeof m.content === "string") chars += m.content.length;
+  else if (Array.isArray(m.content)) {
+    for (const p of m.content) {
+      if (p.type === "text") chars += p.text?.length ?? 0;
+    }
+  }
+  if (m.tool_calls?.length) chars += JSON.stringify(m.tool_calls).length;
+  return chars;
+}
+
+function messageTokens(m: Message): number {
+  return estimateTokensFromChars(messageChars(m));
+}
+
+export function classifySystemBucket(content: string): keyof Omit<SystemTokenBreakdown, "total"> {
+  if (content === CORE_SYSTEM_PROMPT || content.startsWith("You are Klaat Code")) return "core";
+  if (content.startsWith("# Environment")) return "environment";
+  if (content.startsWith("# Project rules") || content.startsWith("Project rules (from")) return "projectRules";
+  if (MODE_PROMPT_VALUES.has(content) || content.startsWith("# Mode:")) return "mode";
+  return "other";
+}
+
+export function computeContextBreakdown(
+  msgs: Message[],
+  tools: ToolDefinition[] = [],
+): ContextBreakdown {
+  const system: SystemTokenBreakdown = {
+    core: 0, environment: 0, projectRules: 0, mode: 0, other: 0, total: 0,
+  };
+  const conversation: ConversationTokenBreakdown = {
+    user: 0, assistant: 0, tool: 0,
+    userMsgs: 0, assistantMsgs: 0, toolMsgs: 0,
+    total: 0,
+  };
+
+  let trimmedToolResults = 0;
+  let compactedStub = false;
+
+  for (const m of msgs) {
+    const toks = messageTokens(m);
+    const text = typeof m.content === "string" ? m.content : "";
+
+    if (m.role === "system") {
+      const bucket = classifySystemBucket(text);
+      system[bucket] += toks;
+      system.total += toks;
+    } else if (m.role === "user") {
+      conversation.user += toks;
+      conversation.userMsgs++;
+      conversation.total += toks;
+    } else if (m.role === "assistant") {
+      conversation.assistant += toks;
+      conversation.assistantMsgs++;
+      conversation.total += toks;
+      if (text.startsWith("[Context compacted")) compactedStub = true;
+    } else if (m.role === "tool") {
+      conversation.tool += toks;
+      conversation.toolMsgs++;
+      conversation.total += toks;
+      if (text.includes("chars trimmed")) trimmedToolResults++;
+    }
+  }
+
+  const toolSchemas = tools.length > 0
+    ? estimateTokensFromText(JSON.stringify(tools))
+    : 0;
+
+  return {
+    system,
+    conversation,
+    toolSchemas,
+    toolCount: tools.length,
+    messagesTotal: system.total + conversation.total,
+    estimatedTotal: system.total + conversation.total + toolSchemas,
+    trimmedToolResults,
+    compactedStub,
+  };
+}
+
+/** ASCII progress bar for transcript output (█ filled, ░ empty). */
+export function formatContextBar(pct: number, width = 20): string {
+  const clamped = Math.min(100, Math.max(0, Math.round(pct)));
+  const filled = Math.round((clamped / 100) * width);
+  return "█".repeat(filled) + "░".repeat(width - filled);
+}
+
+export function formatBreakdownLine(label: string, tokens: number, extra = ""): string {
+  const pad = label.padEnd(16);
+  const suffix = extra ? `  ${extra}` : "";
+  return `    ${pad} ${formatTokInline(tokens)}${suffix}`;
+}
+
+function formatTokInline(n: number): string {
+  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M toks`;
+  if (n >= 1_000) return `${(n / 1_000).toFixed(1)}K toks`;
+  return `${n} toks`;
+}

src/screens/repl.ts

@@ -72,7 +72,12 @@ import { readClipboardImage, MAX_IMAGE_BYTES } from "../utils/clipboard-image.js
 import { copyToClipboard } from "../utils/clipboard.js";
 import { SessionLedger } from "../agent/session-ledger.js";
 import { COMPACTION_PROMPT, extractSummary, MAX_CONSECUTIVE_COMPACT_FAILURES } from "../agent/compaction-prompt.js";
-import { compactMessagesForApi } from "../agent/compaction.js";
+import { compactMessagesForApi, estimateTokensFromChars } from "../agent/compaction.js";
+import {
+  computeContextBreakdown,
+  formatBreakdownLine,
+  formatContextBar,
+} from "../agent/context-breakdown.js";
 import { stripStrayTextToolCallArtifacts } from "../agent/text-tool-artifacts.js";
 import { drawWelcomeCard } from "./welcome-card.js";
 import {
@@ -97,7 +102,7 @@ import {
   type PermDecision,
 } from "../permissions/index.js";
 import { exec, spawnSync } from "child_process";
-import { appendFileSync, readFileSync, writeFileSync, unlinkSync, readdirSync, existsSync, mkdirSync } from "node:fs";
+import { appendFileSync, readFileSync, writeFileSync, unlinkSync, readdirSync, existsSync, mkdirSync, statSync } from "node:fs";
 import { join, relative, resolve } from "node:path";
 import { homedir, tmpdir } from "node:os";
 
@@ -711,7 +716,7 @@ export async function runREPL(
     { cmd: "/commit",     desc: "AI commit message + commit" },
     { cmd: "/compact",    desc: "Summarize context to free token window" },
     { cmd: "/cost",       desc: "Session cost + quota + context usage" },
-    { cmd: "/context",    desc: "What's in the context window vs compacted away" },
+    { cmd: "/context",    desc: "Token usage breakdown of the context window" },
     { cmd: "/diff",       desc: "Git diff (optionally one file)" },
     { cmd: "/doctor",     desc: "Diagnostics: auth, API, MCP, tools, config" },
     { cmd: "/exit",       desc: "Quit KLAAT CODE" },
@@ -1854,32 +1859,80 @@ export async function runREPL(
       }
 
       case "/context": {
-        // 9.6: visibility into what the model can still see vs what only
-        // survives in the ledger.
-        let sysN = 0, userN = 0, asstN = 0, toolN = 0, trimmedN = 0, chars = 0;
-        for (const m of apiMessages) {
-          if (m.role === "system") sysN++;
-          else if (m.role === "user") userN++;
-          else if (m.role === "assistant") asstN++;
-          else if (m.role === "tool") toolN++;
-          if (typeof m.content === "string") {
-            chars += m.content.length;
-            if (m.content.includes("chars trimmed")) trimmedN++;
-          }
-        }
+        const planMode = tabs.activeTab.label === "Plan";
+        const dialect = activeDialect();
+        const activeTools: ToolDefinition[] = planMode
+          ? [
+              ...TOOL_DEFINITIONS.filter(t => PLAN_READONLY_TOOLS.has(t.function.name)),
+              EXIT_PLAN_TOOL,
+            ]
+          : [
+              ...toolsForDialect(dialect, TOOL_DEFINITIONS),
+              ...(dialectIncludesExtras(dialect)
+                ? [...mcpManager.toolDefinitions, ...pluginRegistry.toolDefinitions]
+                : []),
+            ];
+
+        // Mirror the send path: inject the active tab's mode prompt and apply
+        // the same pre-send compaction so the breakdown matches the API window.
+        const isModePrompt = (m: Message) =>
+          m.role === "system" && typeof m.content === "string" &&
+          Object.values(TAB_SYSTEM_PROMPTS).includes(m.content);
+        const historyMsgs = apiMessages.filter(m => !isModePrompt(m));
+        let seedEnd = 0;
+        while (seedEnd < historyMsgs.length && historyMsgs[seedEnd]!.role === "system") seedEnd++;
+        const tabPrompt = TAB_SYSTEM_PROMPTS[tabs.activeTab.label];
+        const tabSystemMsg: Message | null = tabPrompt
+          ? { role: "system", content: tabPrompt }
+          : null;
+        const withMode = tabSystemMsg
+          ? [...historyMsgs.slice(0, seedEnd), tabSystemMsg, ...historyMsgs.slice(seedEnd)]
+          : [...historyMsgs];
+        const sendMsgs = compactMessagesForApi(withMode, getContextWindow(), compactOpts());
+
+        const breakdown = computeContextBreakdown(sendMsgs, activeTools);
         const window = getContextWindow();
-        const pct = window > 0 ? Math.round((lastContextSize / window) * 100) : 0;
-        const compactedStub = apiMessages.some(m =>
-          typeof m.content === "string" && m.content.startsWith("[Context compacted"));
+        const usedTok = lastContextSize > 0 ? lastContextSize : breakdown.estimatedTotal;
+        const pct = window > 0 ? Math.min(100, Math.round((usedTok / window) * 100)) : 0;
+        const apiNote = lastContextSize > 0 && Math.abs(lastContextSize - breakdown.estimatedTotal) > window * 0.02
+          ? `\n  (API last request: ${formatTok(lastContextSize)}; estimate: ${formatTok(breakdown.estimatedTotal)})`
+          : "";
+
+        let ledgerSize = "";
+        try {
+          ledgerSize = ` (${formatTok(estimateTokensFromChars(statSync(ledger.path).size))})`;
+        } catch { /* best-effort */ }
+
+        const sysLines = [
+          breakdown.system.core > 0 ? formatBreakdownLine("Core", breakdown.system.core) : "",
+          breakdown.system.environment > 0 ? formatBreakdownLine("Environment", breakdown.system.environment) : "",
+          breakdown.system.projectRules > 0 ? formatBreakdownLine("Project rules", breakdown.system.projectRules) : "",
+          breakdown.system.mode > 0 ? formatBreakdownLine("Mode", breakdown.system.mode) : "",
+          breakdown.system.other > 0 ? formatBreakdownLine("Other system", breakdown.system.other) : "",
+        ].filter(Boolean);
+
+        const trimmedNote = breakdown.trimmedToolResults > 0
+          ? `, ${breakdown.trimmedToolResults} trimmed`
+          : "";
+
         pushSystemMsg(
           `**Context window:**\n` +
-          `  Last request: ${formatTok(lastContextSize)} / ${formatTok(window)} toks (${pct}%)\n` +
-          `  In memory: ${apiMessages.length} messages (${sysN} system, ${userN} user, ${asstN} assistant, ${toolN} tool) ≈ ${formatTok(Math.round(chars / 4))} toks\n` +
-          `  Degraded in place: ${trimmedN} trimmed tool result${trimmedN === 1 ? "" : "s"}\n\n` +
-          `**Compacted away:**\n` +
-          `  Compactions this session: ${compactionCount}${compactedStub ? " (summary stub in context)" : ""}\n` +
-          `  Recoverable details: ${ledger.path}\n` +
-          `  (the model re-reads that file when it needs pre-compaction specifics)`
+          `  ${formatTok(usedTok)} / ${formatTok(window)} toks (${pct}%)${apiNote}\n` +
+          `  ${formatContextBar(pct)}\n\n` +
+          `**Breakdown** (chars÷4 estimate):\n` +
+          `  System prompt (${formatTok(breakdown.system.total)}):\n` +
+          (sysLines.length ? sysLines.join("\n") + "\n" : "    (none)\n") +
+          `  Conversation (${formatTok(breakdown.conversation.total)}):\n` +
+          formatBreakdownLine("User", breakdown.conversation.user, `(${breakdown.conversation.userMsgs} msgs)`) + "\n" +
+          formatBreakdownLine("Assistant", breakdown.conversation.assistant, `(${breakdown.conversation.assistantMsgs} msgs)`) + "\n" +
+          formatBreakdownLine("Tool results", breakdown.conversation.tool, `(${breakdown.conversation.toolMsgs} msgs${trimmedNote})`) + "\n" +
+          formatBreakdownLine("Tool schemas", breakdown.toolSchemas, `(${breakdown.toolCount} tools, ${dialect} dialect)`) + "\n" +
+          `  ─────────────────\n` +
+          `  Estimated total  ${formatTok(breakdown.estimatedTotal)}\n\n` +
+          `**Compaction:**\n` +
+          `  Compactions this session: ${compactionCount}${breakdown.compactedStub ? " (summary stub in context)" : ""}\n` +
+          `  Ledger: ${ledger.path}${ledgerSize}\n` +
+          `  (read the ledger when compacted details are needed)`
         );
         return true;
       }