From 79731249eb115555c47b535d9f76205323a079d2 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Thu, 4 Dec 2025 00:25:53 +0100 Subject: [PATCH] Context compaction: commands, auto-trigger, RPC support, /branch rework (fixes #92) - Add compaction settings to Settings interface - /compact [instructions]: manual compaction with optional focus - /autocompact: toggle auto-compaction on/off - Auto-compaction triggers after assistant message_end when threshold exceeded - Footer shows (auto) when auto-compact is enabled - RPC mode: {type: 'compact'} command emits CompactionEntry - /branch now reads from session file to show ALL historical user messages - createBranchedSessionFromEntries preserves compaction events --- packages/coding-agent/src/main.ts | 40 ++- packages/coding-agent/src/session-manager.ts | 32 ++ packages/coding-agent/src/settings-manager.ts | 35 +++ packages/coding-agent/src/tui/footer.ts | 12 +- packages/coding-agent/src/tui/tui-renderer.ts | 286 ++++++++++++++++-- 5 files changed, 375 insertions(+), 30 deletions(-) diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index c8b278b7..ef5002a7 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -6,6 +6,7 @@ import { existsSync, readFileSync, statSync } from "fs"; import { homedir } from "os"; import { extname, join, resolve } from "path"; import { getChangelogPath, getNewEntries, parseChangelog } from "./changelog.js"; +import { compact } from "./compaction.js"; import { APP_NAME, CONFIG_DIR_NAME, @@ -17,7 +18,7 @@ import { } from "./config.js"; import { exportFromFile } from "./export-html.js"; import { findModel, getApiKeyForModel, getAvailableModels } from "./model-config.js"; -import { SessionManager } from "./session-manager.js"; +import { loadSessionFromEntries, SessionManager } from "./session-manager.js"; import { SettingsManager } from "./settings-manager.js"; import { expandSlashCommand, loadSlashCommands } from "./slash-commands.js"; import { initTheme } from "./theme/theme.js"; @@ -814,7 +815,11 @@ async function runSingleShotMode( } } -async function runRpcMode(agent: Agent, sessionManager: SessionManager): Promise { +async function runRpcMode( + agent: Agent, + sessionManager: SessionManager, + settingsManager: SettingsManager, +): Promise { // Subscribe to all events and output as JSON (same pattern as tui-renderer) agent.subscribe(async (event) => { console.log(JSON.stringify(event)); @@ -851,6 +856,35 @@ async function runRpcMode(agent: Agent, sessionManager: SessionManager): Promise await agent.prompt(input.message, input.attachments); } else if (input.type === "abort") { agent.abort(); + } else if (input.type === "compact") { + // Handle compaction request + try { + const apiKey = await getApiKeyForModel(agent.state.model); + if (!apiKey) { + throw new Error(`No API key for ${agent.state.model.provider}`); + } + + const entries = sessionManager.loadEntries(); + const settings = settingsManager.getCompactionSettings(); + const compactionEntry = await compact( + entries, + agent.state.model, + settings, + apiKey, + undefined, + input.customInstructions, + ); + + // Save and reload + sessionManager.saveCompaction(compactionEntry); + const loaded = loadSessionFromEntries(sessionManager.loadEntries()); + agent.replaceMessages(loaded.messages); + + // Emit compaction event (compactionEntry already has type: "compaction") + console.log(JSON.stringify(compactionEntry)); + } catch (error: any) { + console.log(JSON.stringify({ type: "error", error: `Compaction failed: ${error.message}` })); + } } } catch (error: any) { // Output error as JSON @@ -1219,7 +1253,7 @@ export async function main(args: string[]) { // Route to appropriate mode if (mode === "rpc") { // RPC mode - headless operation - await runRpcMode(agent, sessionManager); + await runRpcMode(agent, sessionManager, settingsManager); } else if (isInteractive) { // Check for new version (don't block startup if it takes too long) let newVersion: string | null = null; diff --git a/packages/coding-agent/src/session-manager.ts b/packages/coding-agent/src/session-manager.ts index 214f2a02..6758cdeb 100644 --- a/packages/coding-agent/src/session-manager.ts +++ b/packages/coding-agent/src/session-manager.ts @@ -561,4 +561,36 @@ export class SessionManager { return newSessionFile; } + + /** + * Create a branched session from session entries up to (but not including) a specific entry index. + * This preserves compaction events and all entry types. + * Returns the new session file path. + */ + createBranchedSessionFromEntries(entries: SessionEntry[], branchBeforeIndex: number): string { + const newSessionId = uuidv4(); + const timestamp = new Date().toISOString().replace(/[:.]/g, "-"); + const newSessionFile = join(this.sessionDir, `${timestamp}_${newSessionId}.jsonl`); + + // Copy all entries up to (but not including) the branch point + for (let i = 0; i < branchBeforeIndex; i++) { + const entry = entries[i]; + + if (entry.type === "session") { + // Rewrite session header with new ID and branchedFrom + const newHeader: SessionHeader = { + ...entry, + id: newSessionId, + timestamp: new Date().toISOString(), + branchedFrom: this.sessionFile, + }; + appendFileSync(newSessionFile, JSON.stringify(newHeader) + "\n"); + } else { + // Copy other entries as-is + appendFileSync(newSessionFile, JSON.stringify(entry) + "\n"); + } + } + + return newSessionFile; + } } diff --git a/packages/coding-agent/src/settings-manager.ts b/packages/coding-agent/src/settings-manager.ts index 33461c79..a2306854 100644 --- a/packages/coding-agent/src/settings-manager.ts +++ b/packages/coding-agent/src/settings-manager.ts @@ -2,6 +2,12 @@ import { existsSync, mkdirSync, readFileSync, writeFileSync } from "fs"; import { dirname, join } from "path"; import { getAgentDir } from "./config.js"; +export interface CompactionSettings { + enabled?: boolean; // default: true + reserveTokens?: number; // default: 16384 + keepRecentTokens?: number; // default: 20000 +} + export interface Settings { lastChangelogVersion?: string; defaultProvider?: string; @@ -9,6 +15,7 @@ export interface Settings { defaultThinkingLevel?: "off" | "minimal" | "low" | "medium" | "high"; queueMode?: "all" | "one-at-a-time"; theme?: string; + compaction?: CompactionSettings; } export class SettingsManager { @@ -108,4 +115,32 @@ export class SettingsManager { this.settings.defaultThinkingLevel = level; this.save(); } + + getCompactionEnabled(): boolean { + return this.settings.compaction?.enabled ?? true; + } + + setCompactionEnabled(enabled: boolean): void { + if (!this.settings.compaction) { + this.settings.compaction = {}; + } + this.settings.compaction.enabled = enabled; + this.save(); + } + + getCompactionReserveTokens(): number { + return this.settings.compaction?.reserveTokens ?? 16384; + } + + getCompactionKeepRecentTokens(): number { + return this.settings.compaction?.keepRecentTokens ?? 20000; + } + + getCompactionSettings(): { enabled: boolean; reserveTokens: number; keepRecentTokens: number } { + return { + enabled: this.getCompactionEnabled(), + reserveTokens: this.getCompactionReserveTokens(), + keepRecentTokens: this.getCompactionKeepRecentTokens(), + }; + } } diff --git a/packages/coding-agent/src/tui/footer.ts b/packages/coding-agent/src/tui/footer.ts index cb2bb497..d9861d0d 100644 --- a/packages/coding-agent/src/tui/footer.ts +++ b/packages/coding-agent/src/tui/footer.ts @@ -14,11 +14,16 @@ export class FooterComponent implements Component { private cachedBranch: string | null | undefined = undefined; // undefined = not checked yet, null = not in git repo, string = branch name private gitWatcher: FSWatcher | null = null; private onBranchChange: (() => void) | null = null; + private autoCompactEnabled: boolean = true; constructor(state: AgentState) { this.state = state; } + setAutoCompactEnabled(enabled: boolean): void { + this.autoCompactEnabled = enabled; + } + /** * Set up a file watcher on .git/HEAD to detect branch changes. * Call the provided callback when branch changes. @@ -180,12 +185,13 @@ export class FooterComponent implements Component { // Colorize context percentage based on usage let contextPercentStr: string; + const autoIndicator = this.autoCompactEnabled ? " (auto)" : ""; if (contextPercentValue > 90) { - contextPercentStr = theme.fg("error", `${contextPercent}%`); + contextPercentStr = theme.fg("error", `${contextPercent}%${autoIndicator}`); } else if (contextPercentValue > 70) { - contextPercentStr = theme.fg("warning", `${contextPercent}%`); + contextPercentStr = theme.fg("warning", `${contextPercent}%${autoIndicator}`); } else { - contextPercentStr = `${contextPercent}%`; + contextPercentStr = `${contextPercent}%${autoIndicator}`; } statsParts.push(contextPercentStr); diff --git a/packages/coding-agent/src/tui/tui-renderer.ts b/packages/coding-agent/src/tui/tui-renderer.ts index 2a4457d7..3ed5db1c 100644 --- a/packages/coding-agent/src/tui/tui-renderer.ts +++ b/packages/coding-agent/src/tui/tui-renderer.ts @@ -18,11 +18,12 @@ import { } from "@mariozechner/pi-tui"; import { exec } from "child_process"; import { getChangelogPath, parseChangelog } from "../changelog.js"; +import { calculateContextTokens, compact, getLastAssistantUsage, shouldCompact } from "../compaction.js"; import { APP_NAME, getDebugLogPath, getModelsPath, getOAuthPath } from "../config.js"; import { exportSessionToHtml } from "../export-html.js"; import { getApiKeyForModel, getAvailableModels, invalidateOAuthCache } from "../model-config.js"; import { listOAuthProviders, login, logout } from "../oauth/index.js"; -import type { SessionManager } from "../session-manager.js"; +import { loadSessionFromEntries, type SessionManager } from "../session-manager.js"; import type { SettingsManager } from "../settings-manager.js"; import { expandSlashCommand, type FileSlashCommand, loadSlashCommands } from "../slash-commands.js"; import { getEditorTheme, getMarkdownTheme, onThemeChange, setTheme, theme } from "../theme/theme.js"; @@ -129,6 +130,7 @@ export class TuiRenderer { this.editorContainer = new Container(); // Container to hold editor or selector this.editorContainer.addChild(this.editor); // Start with editor this.footer = new FooterComponent(agent.state); + this.footer.setAutoCompactEnabled(this.settingsManager.getCompactionEnabled()); // Define slash commands const thinkingCommand: SlashCommand = { @@ -418,6 +420,21 @@ export class TuiRenderer { return; } + // Check for /compact command + if (text === "/compact" || text.startsWith("/compact ")) { + const customInstructions = text.startsWith("/compact ") ? text.slice(9).trim() : undefined; + this.handleCompactCommand(customInstructions); + this.editor.setText(""); + return; + } + + // Check for /autocompact command + if (text === "/autocompact") { + this.handleAutocompactCommand(); + this.editor.setText(""); + return; + } + // Check for /debug command if (text === "/debug") { this.handleDebugCommand(); @@ -511,10 +528,84 @@ export class TuiRenderer { if (this.sessionManager.shouldInitializeSession(this.agent.state.messages)) { this.sessionManager.startSession(this.agent.state); } + + // Check for auto-compaction after assistant messages + if (event.message.role === "assistant") { + await this.checkAutoCompaction(); + } } }); } + private async checkAutoCompaction(): Promise { + const settings = this.settingsManager.getCompactionSettings(); + if (!settings.enabled) return; + + // Get last assistant usage + const entries = this.sessionManager.loadEntries(); + const lastUsage = getLastAssistantUsage(entries); + if (!lastUsage) return; + + const contextTokens = calculateContextTokens(lastUsage); + const contextWindow = this.agent.state.model.contextWindow; + + if (!shouldCompact(contextTokens, contextWindow, settings)) return; + + // Trigger auto-compaction + await this.handleAutoCompaction(); + } + + private async handleAutoCompaction(): Promise { + // Unsubscribe to stop processing events + this.unsubscribe?.(); + + // Abort current agent run and wait for completion + this.agent.abort(); + await this.agent.waitForIdle(); + + // Stop loading animation + if (this.loadingAnimation) { + this.loadingAnimation.stop(); + this.loadingAnimation = null; + } + this.statusContainer.clear(); + + // Show compacting status + this.chatContainer.addChild(new Spacer(1)); + this.chatContainer.addChild(new Text(theme.fg("muted", "Auto-compacting context..."), 1, 1)); + this.ui.requestRender(); + + try { + const apiKey = await getApiKeyForModel(this.agent.state.model); + if (!apiKey) { + throw new Error(`No API key for ${this.agent.state.model.provider}`); + } + + const entries = this.sessionManager.loadEntries(); + const settings = this.settingsManager.getCompactionSettings(); + const compactionEntry = await compact(entries, this.agent.state.model, settings, apiKey); + + // Save and reload + this.sessionManager.saveCompaction(compactionEntry); + const loaded = loadSessionFromEntries(this.sessionManager.loadEntries()); + this.agent.replaceMessages(loaded.messages); + + // Rebuild UI + this.chatContainer.clear(); + this.rebuildChatFromMessages(); + + this.showSuccess( + "✓ Context auto-compacted", + `Reduced from ${compactionEntry.tokensBefore.toLocaleString()} tokens`, + ); + } catch (error) { + this.showError(`Auto-compaction failed: ${error instanceof Error ? error.message : String(error)}`); + } + + // Resubscribe + this.subscribeToAgent(); + } + private async handleEvent(event: AgentEvent, state: AgentState): Promise { if (!this.isInitialized) { await this.init(); @@ -784,6 +875,50 @@ export class TuiRenderer { }); } + private rebuildChatFromMessages(): void { + // Reset state and re-render messages from agent state + this.isFirstUserMessage = true; + this.pendingTools.clear(); + + for (const message of this.agent.state.messages) { + if (message.role === "user") { + const userMsg = message as any; + const textBlocks = userMsg.content.filter((c: any) => c.type === "text"); + const textContent = textBlocks.map((c: any) => c.text).join(""); + if (textContent) { + const userComponent = new UserMessageComponent(textContent, this.isFirstUserMessage); + this.chatContainer.addChild(userComponent); + this.isFirstUserMessage = false; + } + } else if (message.role === "assistant") { + const assistantMsg = message as AssistantMessage; + const assistantComponent = new AssistantMessageComponent(assistantMsg); + this.chatContainer.addChild(assistantComponent); + + for (const content of assistantMsg.content) { + if (content.type === "toolCall") { + const component = new ToolExecutionComponent(content.name, content.arguments); + this.chatContainer.addChild(component); + this.pendingTools.set(content.id, component); + } + } + } else if (message.role === "toolResult") { + const component = this.pendingTools.get(message.toolCallId); + if (component) { + component.updateResult({ + content: message.content, + details: message.details, + isError: message.isError, + }); + this.pendingTools.delete(message.toolCallId); + } + } + } + + this.pendingTools.clear(); + this.ui.requestRender(); + } + private handleCtrlC(): void { // Handle Ctrl+C double-press logic const now = Date.now(); @@ -977,6 +1112,15 @@ export class TuiRenderer { this.ui.requestRender(); } + private showSuccess(message: string, detail?: string): void { + this.chatContainer.addChild(new Spacer(1)); + const text = detail + ? `${theme.fg("success", message)}\n${theme.fg("muted", detail)}` + : theme.fg("success", message); + this.chatContainer.addChild(new Text(text, 1, 1)); + this.ui.requestRender(); + } + private showThinkingSelector(): void { // Create thinking selector with current level this.thinkingSelector = new ThinkingSelectorComponent( @@ -1176,18 +1320,30 @@ export class TuiRenderer { } private showUserMessageSelector(): void { - // Extract all user messages from the current state + // Read from session file directly to see ALL historical user messages + // (including those before compaction events) + const entries = this.sessionManager.loadEntries(); const userMessages: Array<{ index: number; text: string }> = []; - for (let i = 0; i < this.agent.state.messages.length; i++) { - const message = this.agent.state.messages[i]; - if (message.role === "user") { - const userMsg = message as any; - const textBlocks = userMsg.content.filter((c: any) => c.type === "text"); - const textContent = textBlocks.map((c: any) => c.text).join(""); - if (textContent) { - userMessages.push({ index: i, text: textContent }); - } + const getUserMessageText = (content: string | Array<{ type: string; text?: string }>): string => { + if (typeof content === "string") return content; + if (Array.isArray(content)) { + return content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + } + return ""; + }; + + for (let i = 0; i < entries.length; i++) { + const entry = entries[i]; + if (entry.type !== "message") continue; + if (entry.message.role !== "user") continue; + + const textContent = getUserMessageText(entry.message.content); + if (textContent) { + userMessages.push({ index: i, text: textContent }); } } @@ -1202,22 +1358,23 @@ export class TuiRenderer { // Create user message selector this.userMessageSelector = new UserMessageSelectorComponent( userMessages, - (messageIndex) => { + (entryIndex) => { // Get the selected user message text to put in the editor - const selectedMessage = this.agent.state.messages[messageIndex]; - const selectedUserMsg = selectedMessage as any; - const textBlocks = selectedUserMsg.content.filter((c: any) => c.type === "text"); - const selectedText = textBlocks.map((c: any) => c.text).join(""); + const selectedEntry = entries[entryIndex]; + if (selectedEntry.type !== "message") return; + if (selectedEntry.message.role !== "user") return; - // Create a branched session with messages UP TO (but not including) the selected message - const newSessionFile = this.sessionManager.createBranchedSession(this.agent.state, messageIndex - 1); + const selectedText = getUserMessageText(selectedEntry.message.content); + + // Create a branched session by copying entries up to (but not including) the selected entry + const newSessionFile = this.sessionManager.createBranchedSessionFromEntries(entries, entryIndex); // Set the new session file as active this.sessionManager.setSessionFile(newSessionFile); - // Truncate messages in agent state to before the selected message - const truncatedMessages = this.agent.state.messages.slice(0, messageIndex); - this.agent.replaceMessages(truncatedMessages); + // Reload the session + const loaded = loadSessionFromEntries(this.sessionManager.loadEntries()); + this.agent.replaceMessages(loaded.messages); // Clear and re-render the chat this.chatContainer.clear(); @@ -1226,9 +1383,7 @@ export class TuiRenderer { // Show confirmation message this.chatContainer.addChild(new Spacer(1)); - this.chatContainer.addChild( - new Text(theme.fg("dim", `Branched to new session from message ${messageIndex}`), 1, 0), - ); + this.chatContainer.addChild(new Text(theme.fg("dim", "Branched to new session"), 1, 0)); // Put the selected message in the editor this.editor.setText(selectedText); @@ -1570,6 +1725,89 @@ export class TuiRenderer { this.ui.requestRender(); } + private async handleCompactCommand(customInstructions?: string): Promise { + // Check if there are any messages to compact + const entries = this.sessionManager.loadEntries(); + const messageCount = entries.filter((e) => e.type === "message").length; + + if (messageCount < 2) { + this.showWarning("Nothing to compact (no messages yet)"); + return; + } + + // Unsubscribe first to prevent processing events during compaction + this.unsubscribe?.(); + + // Abort and wait for completion + this.agent.abort(); + await this.agent.waitForIdle(); + + // Stop loading animation + if (this.loadingAnimation) { + this.loadingAnimation.stop(); + this.loadingAnimation = null; + } + this.statusContainer.clear(); + + // Show compacting status + this.chatContainer.addChild(new Spacer(1)); + this.chatContainer.addChild(new Text(theme.fg("muted", "Compacting context..."), 1, 1)); + this.ui.requestRender(); + + try { + // Get API key for current model + const apiKey = await getApiKeyForModel(this.agent.state.model); + if (!apiKey) { + throw new Error(`No API key for ${this.agent.state.model.provider}`); + } + + // Perform compaction + const settings = this.settingsManager.getCompactionSettings(); + const compactionEntry = await compact( + entries, + this.agent.state.model, + settings, + apiKey, + undefined, + customInstructions, + ); + + // Save compaction to session + this.sessionManager.saveCompaction(compactionEntry); + + // Reload session + const loaded = loadSessionFromEntries(this.sessionManager.loadEntries()); + this.agent.replaceMessages(loaded.messages); + + // Rebuild UI + this.chatContainer.clear(); + this.rebuildChatFromMessages(); + + // Show success + this.showSuccess( + "✓ Context compacted", + `Reduced from ${compactionEntry.tokensBefore.toLocaleString()} tokens`, + ); + } catch (error) { + this.showError(`Compaction failed: ${error instanceof Error ? error.message : String(error)}`); + } + + // Resubscribe to agent + this.subscribeToAgent(); + } + + private handleAutocompactCommand(): void { + const currentEnabled = this.settingsManager.getCompactionEnabled(); + const newState = !currentEnabled; + this.settingsManager.setCompactionEnabled(newState); + this.footer.setAutoCompactEnabled(newState); + + this.showSuccess( + `✓ Auto-compact ${newState ? "enabled" : "disabled"}`, + newState ? "Context will be compacted automatically when nearing limits" : "Use /compact to manually compact", + ); + } + private updatePendingMessagesDisplay(): void { this.pendingMessagesContainer.clear();