Add before_compact hook event (closes #281) (#285)

* Add before_compact hook event (closes #281)

* Add compact hook event and documentation

- Add compact event that fires after compaction completes
- Update hooks.md with lifecycle diagram, field docs, and example
- Add CHANGELOG entry
- Add comprehensive test coverage (10 tests) for before_compact and compact events
- Tests cover: event emission, cancellation, custom entry, error handling, multiple hooks
This commit is contained in:
Nico Bailon 2025-12-24 02:26:29 -08:00 committed by GitHub
parent 20b24cf5a4
commit 1e1a92ea47
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 700 additions and 36 deletions

View file

@ -18,13 +18,13 @@ import type { AssistantMessage, Message, Model, TextContent } from "@mariozechne
import { isContextOverflow, supportsXhigh } from "@mariozechner/pi-ai";
import { getModelsPath } from "../config.js";
import { type BashResult, executeBash as executeBashCommand } from "./bash-executor.js";
import { calculateContextTokens, compact, shouldCompact } from "./compaction.js";
import { calculateContextTokens, compact, prepareCompaction, shouldCompact } from "./compaction.js";
import type { LoadedCustomTool, SessionEvent as ToolSessionEvent } from "./custom-tools/index.js";
import { exportSessionToHtml } from "./export-html.js";
import type { HookRunner, SessionEventResult, TurnEndEvent, TurnStartEvent } from "./hooks/index.js";
import type { BashExecutionMessage } from "./messages.js";
import { getApiKeyForModel, getAvailableModels } from "./model-config.js";
import { loadSessionFromEntries, type SessionManager } from "./session-manager.js";
import { type CompactionEntry, loadSessionFromEntries, type SessionManager } from "./session-manager.js";
import type { SettingsManager, SkillsSettings } from "./settings-manager.js";
import { expandSlashCommand, type FileSlashCommand } from "./slash-commands.js";
@ -734,11 +734,8 @@ export class AgentSession {
* @param customInstructions Optional instructions for the compaction summary
*/
async compact(customInstructions?: string): Promise<CompactionResult> {
// Abort any running operation
this._disconnectFromAgent();
await this.abort();
// Create abort controller
this._compactionAbortController = new AbortController();
try {
@ -753,24 +750,73 @@ export class AgentSession {
const entries = this.sessionManager.loadEntries();
const settings = this.settingsManager.getCompactionSettings();
const compactionEntry = await compact(
entries,
this.model,
settings,
apiKey,
this._compactionAbortController.signal,
customInstructions,
);
const preparation = prepareCompaction(entries, settings);
if (!preparation) {
throw new Error("Already compacted");
}
let compactionEntry: CompactionEntry | undefined;
let fromHook = false;
if (this._hookRunner?.hasHandlers("session")) {
const result = (await this._hookRunner.emit({
type: "session",
entries,
sessionFile: this.sessionFile,
previousSessionFile: null,
reason: "before_compact",
cutPoint: preparation.cutPoint,
messagesToSummarize: preparation.messagesToSummarize,
tokensBefore: preparation.tokensBefore,
customInstructions,
model: this.model,
apiKey,
})) as SessionEventResult | undefined;
if (result?.cancel) {
throw new Error("Compaction cancelled");
}
if (result?.compactionEntry) {
compactionEntry = result.compactionEntry;
fromHook = true;
}
}
if (!compactionEntry) {
compactionEntry = await compact(
entries,
this.model,
settings,
apiKey,
this._compactionAbortController.signal,
customInstructions,
);
}
if (this._compactionAbortController.signal.aborted) {
throw new Error("Compaction cancelled");
}
// Save and reload
this.sessionManager.saveCompaction(compactionEntry);
const loaded = loadSessionFromEntries(this.sessionManager.loadEntries());
const newEntries = this.sessionManager.loadEntries();
const loaded = loadSessionFromEntries(newEntries);
this.agent.replaceMessages(loaded.messages);
if (this._hookRunner) {
await this._hookRunner.emit({
type: "session",
entries: newEntries,
sessionFile: this.sessionFile,
previousSessionFile: null,
reason: "compact",
compactionEntry,
tokensBefore: compactionEntry.tokensBefore,
fromHook,
});
}
return {
tokensBefore: compactionEntry.tokensBefore,
summary: compactionEntry.summary,
@ -853,13 +899,51 @@ export class AgentSession {
}
const entries = this.sessionManager.loadEntries();
const compactionEntry = await compact(
entries,
this.model,
settings,
apiKey,
this._autoCompactionAbortController.signal,
);
const preparation = prepareCompaction(entries, settings);
if (!preparation) {
this._emit({ type: "auto_compaction_end", result: null, aborted: false, willRetry: false });
return;
}
let compactionEntry: CompactionEntry | undefined;
let fromHook = false;
if (this._hookRunner?.hasHandlers("session")) {
const hookResult = (await this._hookRunner.emit({
type: "session",
entries,
sessionFile: this.sessionFile,
previousSessionFile: null,
reason: "before_compact",
cutPoint: preparation.cutPoint,
messagesToSummarize: preparation.messagesToSummarize,
tokensBefore: preparation.tokensBefore,
customInstructions: undefined,
model: this.model,
apiKey,
})) as SessionEventResult | undefined;
if (hookResult?.cancel) {
this._emit({ type: "auto_compaction_end", result: null, aborted: true, willRetry: false });
return;
}
if (hookResult?.compactionEntry) {
compactionEntry = hookResult.compactionEntry;
fromHook = true;
}
}
if (!compactionEntry) {
compactionEntry = await compact(
entries,
this.model,
settings,
apiKey,
this._autoCompactionAbortController.signal,
);
}
if (this._autoCompactionAbortController.signal.aborted) {
this._emit({ type: "auto_compaction_end", result: null, aborted: true, willRetry: false });
@ -867,37 +951,43 @@ export class AgentSession {
}
this.sessionManager.saveCompaction(compactionEntry);
const loaded = loadSessionFromEntries(this.sessionManager.loadEntries());
const newEntries = this.sessionManager.loadEntries();
const loaded = loadSessionFromEntries(newEntries);
this.agent.replaceMessages(loaded.messages);
if (this._hookRunner) {
await this._hookRunner.emit({
type: "session",
entries: newEntries,
sessionFile: this.sessionFile,
previousSessionFile: null,
reason: "compact",
compactionEntry,
tokensBefore: compactionEntry.tokensBefore,
fromHook,
});
}
const result: CompactionResult = {
tokensBefore: compactionEntry.tokensBefore,
summary: compactionEntry.summary,
};
this._emit({ type: "auto_compaction_end", result, aborted: false, willRetry });
// Auto-retry if needed - use continue() since user message is already in context
if (willRetry) {
// Remove trailing error message from agent state (it's kept in session file for history)
// This is needed because continue() requires last message to be user or toolResult
const messages = this.agent.state.messages;
const lastMsg = messages[messages.length - 1];
if (lastMsg?.role === "assistant" && (lastMsg as AssistantMessage).stopReason === "error") {
this.agent.replaceMessages(messages.slice(0, -1));
}
// Use setTimeout to break out of the event handler chain
setTimeout(() => {
this.agent.continue().catch(() => {
// Retry failed - silently ignore, user can manually retry
});
this.agent.continue().catch(() => {});
}, 100);
}
} catch (error) {
// Compaction failed - emit end event without retry
this._emit({ type: "auto_compaction_end", result: null, aborted: false, willRetry: false });
// If this was overflow recovery and compaction failed, we have a hard stop
if (reason === "overflow") {
throw new Error(
`Context overflow: ${error instanceof Error ? error.message : "compaction failed"}. Your input may be too large for the context window.`,

View file

@ -321,6 +321,49 @@ export async function generateSummary(
return textContent;
}
// ============================================================================
// Compaction Preparation (for hooks)
// ============================================================================
export interface CompactionPreparation {
cutPoint: CutPointResult;
messagesToSummarize: AppMessage[];
tokensBefore: number;
boundaryStart: number;
}
export function prepareCompaction(entries: SessionEntry[], settings: CompactionSettings): CompactionPreparation | null {
if (entries.length > 0 && entries[entries.length - 1].type === "compaction") {
return null;
}
let prevCompactionIndex = -1;
for (let i = entries.length - 1; i >= 0; i--) {
if (entries[i].type === "compaction") {
prevCompactionIndex = i;
break;
}
}
const boundaryStart = prevCompactionIndex + 1;
const boundaryEnd = entries.length;
const lastUsage = getLastAssistantUsage(entries);
const tokensBefore = lastUsage ? calculateContextTokens(lastUsage) : 0;
const cutPoint = findCutPoint(entries, boundaryStart, boundaryEnd, settings.keepRecentTokens);
const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex;
const messagesToSummarize: AppMessage[] = [];
for (let i = boundaryStart; i < historyEnd; i++) {
const entry = entries[i];
if (entry.type === "message") {
messagesToSummarize.push(entry.message);
}
}
return { cutPoint, messagesToSummarize, tokensBefore, boundaryStart };
}
// ============================================================================
// Main compaction function
// ============================================================================

View file

@ -6,8 +6,9 @@
*/
import type { AppMessage, Attachment } from "@mariozechner/pi-agent-core";
import type { ImageContent, TextContent, ToolResultMessage } from "@mariozechner/pi-ai";
import type { SessionEntry } from "../session-manager.js";
import type { ImageContent, Model, TextContent, ToolResultMessage } from "@mariozechner/pi-ai";
import type { CutPointResult } from "../compaction.js";
import type { CompactionEntry, SessionEntry } from "../session-manager.js";
import type {
BashToolDetails,
FindToolDetails,
@ -111,6 +112,7 @@ interface SessionEventBase {
* - before_switch / switch: Session switch (e.g., /resume command)
* - before_clear / clear: Session clear (e.g., /clear command)
* - before_branch / branch: Session branch (e.g., /branch command)
* - before_compact / compact: Before/after context compaction
* - shutdown: Process exit (SIGINT/SIGTERM)
*
* "before_*" events fire before the action and can be cancelled via SessionEventResult.
@ -124,6 +126,22 @@ export type SessionEvent =
reason: "branch" | "before_branch";
/** Index of the turn to branch from */
targetTurnIndex: number;
})
| (SessionEventBase & {
reason: "before_compact";
cutPoint: CutPointResult;
messagesToSummarize: AppMessage[];
tokensBefore: number;
customInstructions?: string;
model: Model<any>;
apiKey: string;
})
| (SessionEventBase & {
reason: "compact";
compactionEntry: CompactionEntry;
tokensBefore: number;
/** Whether the compaction entry was provided by a hook */
fromHook: boolean;
});
/**
@ -325,6 +343,8 @@ export interface SessionEventResult {
cancel?: boolean;
/** If true (for before_branch only), skip restoring conversation to branch point while still creating the branched session file */
skipConversationRestore?: boolean;
/** Custom compaction entry (for before_compact event) */
compactionEntry?: CompactionEntry;
}
// ============================================================================