diff --git a/src/agent/compaction.ts b/src/agent/compaction.ts index 2603e21..434d03f 100644 --- a/src/agent/compaction.ts +++ b/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"]); diff --git a/src/agent/context-breakdown.test.ts b/src/agent/context-breakdown.test.ts new file mode 100644 index 0000000..2058ad8 --- /dev/null +++ b/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("█████░░░░░"); +}); diff --git a/src/agent/context-breakdown.ts b/src/agent/context-breakdown.ts new file mode 100644 index 0000000..76fd2d9 --- /dev/null +++ b/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 { + 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`; +} diff --git a/src/screens/repl.ts b/src/screens/repl.ts index 185cfe8..6a78a43 100644 --- a/src/screens/repl.ts +++ b/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; }