refactor(coding-agent): fix compaction for branched sessions, consolidate hook context types

Compaction API:
- prepareCompaction() now takes (pathEntries, settings) only
- CompactionPreparation restructured: removed cutPoint/messagesToKeep/boundaryStart, added turnPrefixMessages/isSplitTurn/previousSummary/fileOps/settings
- compact() now takes (preparation, model, apiKey, customInstructions?, signal?)
- Fixed token overflow by using getPath() instead of getEntries()

Hook types:
- HookEventContext renamed to HookContext
- HookCommandContext removed, RegisteredCommand.handler takes (args, ctx)
- HookContext now includes model field
- SessionBeforeCompactEvent: removed previousCompactions/model, added branchEntries
- SessionBeforeTreeEvent: removed model (use ctx.model)
- HookRunner.initialize() added for modes to set up callbacks
This commit is contained in:
Mario Zechner 2025-12-31 02:24:24 +01:00
parent b4ce93c577
commit ddda8b124c
12 changed files with 177 additions and 201 deletions

View file

@ -30,7 +30,7 @@ import {
import type { LoadedCustomTool, SessionEvent as ToolSessionEvent } from "./custom-tools/index.js";
import { exportSessionToHtml } from "./export-html.js";
import type {
HookCommandContext,
HookContext,
HookRunner,
SessionBeforeBranchResult,
SessionBeforeCompactResult,
@ -552,17 +552,17 @@ export class AgentSession {
// Build command context
const cwd = process.cwd();
const ctx: HookCommandContext = {
args,
const ctx: HookContext = {
ui: uiContext,
hasUI: this._hookRunner.getHasUI(),
cwd,
sessionManager: this.sessionManager,
modelRegistry: this._modelRegistry,
model: this.model,
};
try {
await command.handler(ctx);
await command.handler(args, ctx);
return true;
} catch (err) {
// Emit error via hook runner
@ -895,13 +895,13 @@ export class AgentSession {
throw new Error(`No API key for ${this.model.provider}`);
}
const entries = this.sessionManager.getEntries();
const pathEntries = this.sessionManager.getPath();
const settings = this.settingsManager.getCompactionSettings();
const preparation = prepareCompaction(entries, settings);
const preparation = prepareCompaction(pathEntries, settings);
if (!preparation) {
// Check why we can't compact
const lastEntry = entries[entries.length - 1];
const lastEntry = pathEntries[pathEntries.length - 1];
if (lastEntry?.type === "compaction") {
throw new Error("Already compacted");
}
@ -912,15 +912,11 @@ export class AgentSession {
let fromHook = false;
if (this._hookRunner?.hasHandlers("session_before_compact")) {
// Get previous compactions, newest first
const previousCompactions = entries.filter((e): e is CompactionEntry => e.type === "compaction").reverse();
const result = (await this._hookRunner.emit({
type: "session_before_compact",
preparation,
previousCompactions,
branchEntries: pathEntries,
customInstructions,
model: this.model,
signal: this._compactionAbortController.signal,
})) as SessionBeforeCompactResult | undefined;
@ -948,12 +944,11 @@ export class AgentSession {
} else {
// Generate compaction result
const result = await compact(
entries,
preparation,
this.model,
settings,
apiKey,
this._compactionAbortController.signal,
customInstructions,
this._compactionAbortController.signal,
);
summary = result.summary;
firstKeptEntryId = result.firstKeptEntryId;
@ -1073,9 +1068,9 @@ export class AgentSession {
return;
}
const entries = this.sessionManager.getEntries();
const pathEntries = this.sessionManager.getPath();
const preparation = prepareCompaction(entries, settings);
const preparation = prepareCompaction(pathEntries, settings);
if (!preparation) {
this._emit({ type: "auto_compaction_end", result: undefined, aborted: false, willRetry: false });
return;
@ -1085,15 +1080,11 @@ export class AgentSession {
let fromHook = false;
if (this._hookRunner?.hasHandlers("session_before_compact")) {
// Get previous compactions, newest first
const previousCompactions = entries.filter((e): e is CompactionEntry => e.type === "compaction").reverse();
const hookResult = (await this._hookRunner.emit({
type: "session_before_compact",
preparation,
previousCompactions,
branchEntries: pathEntries,
customInstructions: undefined,
model: this.model,
signal: this._autoCompactionAbortController.signal,
})) as SessionBeforeCompactResult | undefined;
@ -1122,10 +1113,10 @@ export class AgentSession {
} else {
// Generate compaction result
const compactResult = await compact(
entries,
preparation,
this.model,
settings,
apiKey,
undefined,
this._autoCompactionAbortController.signal,
);
summary = compactResult.summary;
@ -1628,7 +1619,6 @@ export class AgentSession {
const result = (await this._hookRunner.emit({
type: "session_before_tree",
preparation,
model: this.model!, // Checked above if summarize is true
signal: this._branchSummaryAbortController.signal,
})) as SessionBeforeTreeResult | undefined;

View file

@ -519,42 +519,48 @@ export async function generateSummary(
// ============================================================================
export interface CompactionPreparation {
cutPoint: CutPointResult;
/** UUID of first entry to keep */
firstKeptEntryId: string;
/** Messages that will be summarized and discarded */
messagesToSummarize: AgentMessage[];
/** Messages that will be kept after the summary (recent turns) */
messagesToKeep: AgentMessage[];
/** Messages that will be turned into turn prefix summary (if splitting) */
turnPrefixMessages: AgentMessage[];
/** Whether this is a split turn (cut point in middle of turn) */
isSplitTurn: boolean;
tokensBefore: number;
boundaryStart: number;
/** Summary from previous compaction, for iterative update */
previousSummary?: string;
/** File operations extracted from messagesToSummarize */
fileOps: FileOperations;
/** Compaction settions from settings.jsonl */
settings: CompactionSettings;
}
export function prepareCompaction(
entries: SessionEntry[],
pathEntries: SessionEntry[],
settings: CompactionSettings,
): CompactionPreparation | undefined {
if (entries.length > 0 && entries[entries.length - 1].type === "compaction") {
if (pathEntries.length > 0 && pathEntries[pathEntries.length - 1].type === "compaction") {
return undefined;
}
let prevCompactionIndex = -1;
for (let i = entries.length - 1; i >= 0; i--) {
if (entries[i].type === "compaction") {
for (let i = pathEntries.length - 1; i >= 0; i--) {
if (pathEntries[i].type === "compaction") {
prevCompactionIndex = i;
break;
}
}
const boundaryStart = prevCompactionIndex + 1;
const boundaryEnd = entries.length;
const boundaryEnd = pathEntries.length;
const lastUsage = getLastAssistantUsage(entries);
const lastUsage = getLastAssistantUsage(pathEntries);
const tokensBefore = lastUsage ? calculateContextTokens(lastUsage) : 0;
const cutPoint = findCutPoint(entries, boundaryStart, boundaryEnd, settings.keepRecentTokens);
const cutPoint = findCutPoint(pathEntries, boundaryStart, boundaryEnd, settings.keepRecentTokens);
// Get UUID of first kept entry
const firstKeptEntry = entries[cutPoint.firstKeptEntryIndex];
const firstKeptEntry = pathEntries[cutPoint.firstKeptEntryIndex];
if (!firstKeptEntry?.id) {
return undefined; // Session needs migration
}
@ -565,18 +571,46 @@ export function prepareCompaction(
// Messages to summarize (will be discarded after summary)
const messagesToSummarize: AgentMessage[] = [];
for (let i = boundaryStart; i < historyEnd; i++) {
const msg = getMessageFromEntry(entries[i]);
const msg = getMessageFromEntry(pathEntries[i]);
if (msg) messagesToSummarize.push(msg);
}
// Messages to keep (recent turns, kept after summary)
const messagesToKeep: AgentMessage[] = [];
for (let i = cutPoint.firstKeptEntryIndex; i < boundaryEnd; i++) {
const msg = getMessageFromEntry(entries[i]);
if (msg) messagesToKeep.push(msg);
// Messages for turn prefix summary (if splitting a turn)
const turnPrefixMessages: AgentMessage[] = [];
if (cutPoint.isSplitTurn) {
for (let i = cutPoint.turnStartIndex; i < cutPoint.firstKeptEntryIndex; i++) {
const msg = getMessageFromEntry(pathEntries[i]);
if (msg) turnPrefixMessages.push(msg);
}
}
return { cutPoint, firstKeptEntryId, messagesToSummarize, messagesToKeep, tokensBefore, boundaryStart };
// Get previous summary for iterative update
let previousSummary: string | undefined;
if (prevCompactionIndex >= 0) {
const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry;
previousSummary = prevCompaction.summary;
}
// Extract file operations from messages and previous compaction
const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex);
// Also extract file ops from turn prefix if splitting
if (cutPoint.isSplitTurn) {
for (const msg of turnPrefixMessages) {
extractFileOpsFromMessage(msg, fileOps);
}
}
return {
firstKeptEntryId,
messagesToSummarize,
turnPrefixMessages,
isSplitTurn: cutPoint.isSplitTurn,
tokensBefore,
previousSummary,
fileOps,
settings,
};
}
// ============================================================================
@ -599,87 +633,39 @@ Summarize the prefix to provide context for the retained suffix:
Be concise. Focus on what's needed to understand the kept suffix.`;
/**
* Calculate compaction and generate summary.
* Generate summaries for compaction using prepared data.
* Returns CompactionResult - SessionManager adds uuid/parentUuid when saving.
*
* @param entries - All session entries (must have uuid fields for v2)
* @param model - Model to use for summarization
* @param settings - Compaction settings
* @param apiKey - API key for LLM
* @param signal - Optional abort signal
* @param preparation - Pre-calculated preparation from prepareCompaction()
* @param customInstructions - Optional custom focus for the summary
*/
export async function compact(
entries: SessionEntry[],
preparation: CompactionPreparation,
model: Model<any>,
settings: CompactionSettings,
apiKey: string,
signal?: AbortSignal,
customInstructions?: string,
signal?: AbortSignal,
): Promise<CompactionResult> {
// Don't compact if the last entry is already a compaction
if (entries.length > 0 && entries[entries.length - 1].type === "compaction") {
throw new Error("Already compacted");
}
// Find previous compaction boundary
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;
// Get token count before compaction
const lastUsage = getLastAssistantUsage(entries);
const tokensBefore = lastUsage ? calculateContextTokens(lastUsage) : 0;
// Find cut point (entry index) within the valid range
const cutResult = findCutPoint(entries, boundaryStart, boundaryEnd, settings.keepRecentTokens);
// Extract messages for history summary (before the turn that contains the cut point)
const historyEnd = cutResult.isSplitTurn ? cutResult.turnStartIndex : cutResult.firstKeptEntryIndex;
const historyMessages: AgentMessage[] = [];
for (let i = boundaryStart; i < historyEnd; i++) {
const msg = getMessageFromEntry(entries[i]);
if (msg) historyMessages.push(msg);
}
// Get previous summary for iterative update (if not from hook)
let previousSummary: string | undefined;
if (prevCompactionIndex >= 0) {
const prevCompaction = entries[prevCompactionIndex] as CompactionEntry;
previousSummary = prevCompaction.summary;
}
// Extract file operations from messages and previous compaction
const fileOps = extractFileOperations(historyMessages, entries, prevCompactionIndex);
// Extract messages for turn prefix summary (if splitting a turn)
const turnPrefixMessages: AgentMessage[] = [];
if (cutResult.isSplitTurn) {
for (let i = cutResult.turnStartIndex; i < cutResult.firstKeptEntryIndex; i++) {
const msg = getMessageFromEntry(entries[i]);
if (msg) turnPrefixMessages.push(msg);
}
// Also extract file ops from turn prefix
for (const msg of turnPrefixMessages) {
extractFileOpsFromMessage(msg, fileOps);
}
}
const {
firstKeptEntryId,
messagesToSummarize,
turnPrefixMessages,
isSplitTurn,
tokensBefore,
previousSummary,
fileOps,
settings,
} = preparation;
// Generate summaries (can be parallel if both needed) and merge into one
let summary: string;
if (cutResult.isSplitTurn && turnPrefixMessages.length > 0) {
if (isSplitTurn && turnPrefixMessages.length > 0) {
// Generate both summaries in parallel
const [historyResult, turnPrefixResult] = await Promise.all([
historyMessages.length > 0
messagesToSummarize.length > 0
? generateSummary(
historyMessages,
messagesToSummarize,
model,
settings.reserveTokens,
apiKey,
@ -695,7 +681,7 @@ export async function compact(
} else {
// Just generate history summary
summary = await generateSummary(
historyMessages,
messagesToSummarize,
model,
settings.reserveTokens,
apiKey,
@ -709,9 +695,6 @@ export async function compact(
const { readFiles, modifiedFiles } = computeFileLists(fileOps);
summary += formatFileOperations(readFiles, modifiedFiles);
// Get UUID of first kept entry
const firstKeptEntry = entries[cutResult.firstKeptEntryIndex];
const firstKeptEntryId = firstKeptEntry.id;
if (!firstKeptEntryId) {
throw new Error("First kept entry has no UUID - session may need migration");
}

View file

@ -3,6 +3,7 @@
*/
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { Model } from "@mariozechner/pi-ai";
import type { ModelRegistry } from "../model-registry.js";
import type { SessionManager } from "../session-manager.js";
import type { AppendEntryHandler, LoadedHook, SendMessageHandler } from "./loader.js";
@ -11,9 +12,9 @@ import type {
BeforeAgentStartEventResult,
ContextEvent,
ContextEventResult,
HookContext,
HookError,
HookEvent,
HookEventContext,
HookMessageRenderer,
HookUIContext,
RegisteredCommand,
@ -72,6 +73,7 @@ export class HookRunner {
private modelRegistry: ModelRegistry;
private timeout: number;
private errorListeners: Set<HookErrorListener> = new Set();
private getModel: () => Model<any> | undefined = () => undefined;
constructor(
hooks: LoadedHook[],
@ -89,6 +91,30 @@ export class HookRunner {
this.timeout = timeout;
}
/**
* Initialize HookRunner with all required context.
* Modes call this once the agent session is fully set up.
*/
initialize(options: {
/** Function to get the current model */
getModel: () => Model<any> | undefined;
/** Handler for hooks to send messages */
sendMessageHandler: SendMessageHandler;
/** Handler for hooks to append entries */
appendEntryHandler: AppendEntryHandler;
/** UI context for interactive prompts */
uiContext?: HookUIContext;
/** Whether UI is available */
hasUI?: boolean;
}): void {
this.getModel = options.getModel;
this.setSendMessageHandler(options.sendMessageHandler);
this.setAppendEntryHandler(options.appendEntryHandler);
if (options.uiContext) {
this.setUIContext(options.uiContext, options.hasUI ?? false);
}
}
/**
* Set the UI context for hooks.
* Call this when the mode initializes and UI is available.
@ -217,13 +243,14 @@ export class HookRunner {
/**
* Create the event context for handlers.
*/
private createContext(): HookEventContext {
private createContext(): HookContext {
return {
ui: this.uiContext,
hasUI: this.hasUI,
cwd: this.cwd,
sessionManager: this.sessionManager,
modelRegistry: this.modelRegistry,
model: this.getModel(),
};
}

View file

@ -88,9 +88,9 @@ export interface HookUIContext {
}
/**
* Context passed to hook event handlers.
* Context passed to hook event and command handlers.
*/
export interface HookEventContext {
export interface HookContext {
/** UI methods for user interaction */
ui: HookUIContext;
/** Whether UI is available (false in print mode) */
@ -101,6 +101,8 @@ export interface HookEventContext {
sessionManager: ReadonlySessionManager;
/** Model registry - use for API key resolution and model retrieval */
modelRegistry: ModelRegistry;
/** Current model (may be undefined if no model is selected yet) */
model: Model<any> | undefined;
}
// ============================================================================
@ -152,14 +154,12 @@ export interface SessionBranchEvent {
/** Fired before context compaction (can be cancelled or customized) */
export interface SessionBeforeCompactEvent {
type: "session_before_compact";
/** Compaction preparation with cut point, messages to summarize/keep, etc. */
/** Compaction preparation with messages to summarize, file ops, previous summary, etc. */
preparation: CompactionPreparation;
/** Previous compaction entries, newest first. Use for iterative summarization. */
previousCompactions: CompactionEntry[];
/** Branch entries (root to current leaf). Use to inspect custom state or previous compactions. */
branchEntries: SessionEntry[];
/** Optional user-provided instructions for the summary */
customInstructions?: string;
/** Current model */
model: Model<any>;
/** Abort signal - hooks should pass this to LLM calls and check it periodically */
signal: AbortSignal;
}
@ -196,9 +196,7 @@ export interface SessionBeforeTreeEvent {
type: "session_before_tree";
/** Preparation data for the navigation */
preparation: TreePreparation;
/** Model to use for summarization (conversation model) */
model: Model<any>;
/** Abort signal - honors Escape during summarization */
/** Abort signal - honors Escape during summarization (model available via ctx.model) */
signal: AbortSignal;
}
@ -529,7 +527,7 @@ export interface SessionBeforeTreeResult {
* Handlers can return R, undefined, or void (bare return statements).
*/
// biome-ignore lint/suspicious/noConfusingVoidType: void allows bare return statements in handlers
export type HookHandler<E, R = undefined> = (event: E, ctx: HookEventContext) => Promise<R | void> | R | void;
export type HookHandler<E, R = undefined> = (event: E, ctx: HookContext) => Promise<R | void> | R | void;
export interface HookMessageRenderOptions {
/** Whether the view is expanded */
@ -546,24 +544,6 @@ export type HookMessageRenderer<T = unknown> = (
theme: Theme,
) => Component | undefined;
/**
* Context passed to hook command handlers.
*/
export interface HookCommandContext {
/** Arguments after the command name */
args: string;
/** UI methods for user interaction */
ui: HookUIContext;
/** Whether UI is available (false in print mode) */
hasUI: boolean;
/** Current working directory */
cwd: string;
/** Session manager (read-only) - use pi.sendMessage()/pi.appendEntry() for writes */
sessionManager: ReadonlySessionManager;
/** Model registry for API keys */
modelRegistry: ModelRegistry;
}
/**
* Command registration options.
*/
@ -571,7 +551,7 @@ export interface RegisteredCommand {
name: string;
description?: string;
handler: (ctx: HookCommandContext) => Promise<void>;
handler: (args: string, ctx: HookContext) => Promise<void>;
}
/**

View file

@ -27,9 +27,9 @@ export {
} from "./custom-tools/index.js";
export {
type HookAPI,
type HookContext,
type HookError,
type HookEvent,
type HookEventContext,
type HookFactory,
HookRunner,
type HookUIContext,