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
This commit is contained in:
Mario Zechner 2025-12-04 00:25:53 +01:00
parent 6c2360af28
commit 79731249eb
5 changed files with 375 additions and 30 deletions

View file

@ -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<void> {
async function runRpcMode(
agent: Agent,
sessionManager: SessionManager,
settingsManager: SettingsManager,
): Promise<void> {
// 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;

View file

@ -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;
}
}

View file

@ -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(),
};
}
}

View file

@ -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);

View file

@ -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<void> {
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<void> {
// 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<void> {
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<void> {
// 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();