mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-15 08:03:39 +00:00
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
parent
b4ce93c577
commit
ddda8b124c
12 changed files with 177 additions and 201 deletions
|
|
@ -20,13 +20,15 @@
|
|||
- `compact()` now returns `CompactionResult` (`{ summary, firstKeptEntryId, tokensBefore, details? }`) instead of `CompactionEntry`
|
||||
- `appendCompaction()` now accepts optional `details` parameter
|
||||
- `CompactionEntry.firstKeptEntryIndex` replaced with `firstKeptEntryId`
|
||||
- `prepareCompaction()` now returns `firstKeptEntryId` in its result
|
||||
- `prepareCompaction(pathEntries, settings)` now takes path entries (from `getPath()`) and settings only
|
||||
- `CompactionPreparation` restructured: removed `cutPoint`, `messagesToKeep`, `boundaryStart`; added `turnPrefixMessages`, `isSplitTurn`, `previousSummary`, `fileOps`, `settings`
|
||||
- `compact(preparation, model, apiKey, customInstructions?, signal?)` now takes preparation and execution context separately
|
||||
- **Hook types**:
|
||||
- `SessionEventBase` no longer has `sessionManager`/`modelRegistry` - access them via `HookEventContext` instead
|
||||
- `HookEventContext` now has `sessionManager` and `modelRegistry` (moved from events)
|
||||
- `HookEventContext` no longer has `exec()` - use `pi.exec()` instead
|
||||
- `HookCommandContext` no longer has `exec()` - use `pi.exec()` instead
|
||||
- `before_compact` event passes `preparation: CompactionPreparation` and `previousCompactions: CompactionEntry[]` (newest first)
|
||||
- `HookEventContext` renamed to `HookContext`
|
||||
- `HookContext` now has `sessionManager`, `modelRegistry`, and `model` (current model, may be undefined)
|
||||
- `HookCommandContext` removed - `RegisteredCommand.handler` now takes `(args: string, ctx: HookContext)`
|
||||
- `before_compact` event: removed `previousCompactions` and `model`, added `branchEntries: SessionEntry[]` (hooks extract what they need)
|
||||
- `before_tree` event: removed `model` (use `ctx.model` instead)
|
||||
- `before_switch` event now has `targetSessionFile`, `switch` event has `previousSessionFile`
|
||||
- Removed `resolveApiKey` (use `modelRegistry.getApiKey(model)`)
|
||||
- Hooks can return `compaction.details` to store custom data (e.g., ArtifactIndex for structured compaction)
|
||||
|
|
@ -38,7 +40,7 @@
|
|||
- New `pi.exec(command, args, options?)` to execute shell commands (moved from `HookEventContext`/`HookCommandContext`)
|
||||
- `HookMessageRenderer` type: `(message: HookMessage, options, theme) => Component | null`
|
||||
- Renderers return inner content; the TUI wraps it in a styled Box
|
||||
- New types: `HookMessage<T>`, `RegisteredCommand`, `HookCommandContext`
|
||||
- New types: `HookMessage<T>`, `RegisteredCommand`, `HookContext`
|
||||
- Handler types renamed: `SendHandler` → `SendMessageHandler`, new `AppendEntryHandler`
|
||||
- **SessionManager**:
|
||||
- `getSessionFile()` now returns `string | undefined` (undefined for in-memory sessions)
|
||||
|
|
@ -58,7 +60,7 @@
|
|||
- `FileOperations`, `collectEntriesForBranchSummary`, `prepareBranchEntries`, `generateBranchSummary` - branch summarization utilities
|
||||
- `CompactionPreparation`, `CompactionDetails` - compaction preparation types
|
||||
- `ReadonlySessionManager` - read-only session manager interface for hooks
|
||||
- `HookMessage`, `HookCommandContext`, `HookMessageRenderOptions` - hook types
|
||||
- `HookMessage`, `HookContext`, `HookMessageRenderOptions` - hook types
|
||||
- `isHookMessage`, `createHookMessage` - hook message utilities
|
||||
|
||||
### Added
|
||||
|
|
@ -93,6 +95,7 @@
|
|||
- **Improved error messages**: Better error messages when `apiKey` or `model` are missing. ([#346](https://github.com/badlogic/pi-mono/pull/346) by [@ronyrus](https://github.com/ronyrus))
|
||||
- **Session file validation**: `findMostRecentSession()` now validates session headers before returning, preventing non-session JSONL files from being loaded
|
||||
- **Compaction error handling**: `generateSummary()` and `generateTurnPrefixSummary()` now throw on LLM errors instead of returning empty strings
|
||||
- **Compaction with branched sessions**: Fixed compaction incorrectly including entries from abandoned branches, causing token overflow errors. Compaction now uses `sessionManager.getPath()` to work only on the current branch path, eliminating 80+ lines of duplicate entry collection logic between `prepareCompaction()` and `compact()`
|
||||
|
||||
## [0.30.2] - 2025-12-26
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
*
|
||||
* Replaces the default compaction behavior with a full summary of the entire context.
|
||||
* Instead of keeping the last 20k tokens of conversation turns, this hook:
|
||||
* 1. Summarizes ALL messages (both messagesToSummarize and messagesToKeep and previousSummary)
|
||||
* 1. Summarizes ALL messages (messagesToSummarize + turnPrefixMessages)
|
||||
* 2. Discards all old turns completely, keeping only the summary
|
||||
*
|
||||
* This example also demonstrates using a different model (Gemini Flash) for summarization,
|
||||
|
|
@ -14,6 +14,7 @@
|
|||
*/
|
||||
|
||||
import { complete, getModel } from "@mariozechner/pi-ai";
|
||||
import type { CompactionEntry } from "@mariozechner/pi-coding-agent";
|
||||
import { convertToLlm } from "@mariozechner/pi-coding-agent";
|
||||
import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks";
|
||||
|
||||
|
|
@ -21,11 +22,8 @@ export default function (pi: HookAPI) {
|
|||
pi.on("session_before_compact", async (event, ctx) => {
|
||||
ctx.ui.notify("Custom compaction hook triggered", "info");
|
||||
|
||||
const { preparation, previousCompactions, signal } = event;
|
||||
const { messagesToSummarize, messagesToKeep, tokensBefore, firstKeptEntryId } = preparation;
|
||||
|
||||
// Get previous summary from most recent compaction (if any)
|
||||
const previousSummary = previousCompactions[0]?.summary;
|
||||
const { preparation, branchEntries, signal } = event;
|
||||
const { messagesToSummarize, turnPrefixMessages, tokensBefore, firstKeptEntryId, previousSummary } = preparation;
|
||||
|
||||
// Use Gemini Flash for summarization (cheaper/faster than most conversation models)
|
||||
const model = getModel("google", "gemini-2.5-flash");
|
||||
|
|
@ -42,7 +40,7 @@ export default function (pi: HookAPI) {
|
|||
}
|
||||
|
||||
// Combine all messages for full summary
|
||||
const allMessages = [...messagesToSummarize, ...messagesToKeep];
|
||||
const allMessages = [...messagesToSummarize, ...turnPrefixMessages];
|
||||
|
||||
ctx.ui.notify(
|
||||
`Custom compaction: summarizing ${allMessages.length} messages (${tokensBefore.toLocaleString()} tokens) with ${model.id}...`,
|
||||
|
|
|
|||
|
|
@ -5,13 +5,9 @@
|
|||
* Useful to ensure work is committed before switching context.
|
||||
*/
|
||||
|
||||
import type { HookAPI, HookEventContext } from "@mariozechner/pi-coding-agent/hooks";
|
||||
import type { HookAPI, HookContext } from "@mariozechner/pi-coding-agent/hooks";
|
||||
|
||||
async function checkDirtyRepo(
|
||||
pi: HookAPI,
|
||||
ctx: HookEventContext,
|
||||
action: string,
|
||||
): Promise<{ cancel: boolean } | undefined> {
|
||||
async function checkDirtyRepo(pi: HookAPI, ctx: HookContext, action: string): Promise<{ cancel: boolean } | undefined> {
|
||||
// Check for uncommitted changes
|
||||
const { stdout, code } = await pi.exec("git", ["status", "--porcelain"]);
|
||||
|
||||
|
|
|
|||
|
|
@ -310,7 +310,7 @@ export default function (pi: HookAPI) {
|
|||
pi.registerCommand("snake", {
|
||||
description: "Play Snake!",
|
||||
|
||||
handler: async (ctx) => {
|
||||
handler: async (_args, ctx) => {
|
||||
if (!ctx.hasUI) {
|
||||
ctx.ui.notify("Snake requires interactive mode", "error");
|
||||
return;
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
};
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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>;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -11,23 +11,27 @@ describe("Documentation example", () => {
|
|||
const exampleHook = (pi: HookAPI) => {
|
||||
pi.on("session_before_compact", async (event: SessionBeforeCompactEvent, ctx) => {
|
||||
// All these should be accessible on the event
|
||||
const { preparation, previousCompactions, model } = event;
|
||||
// sessionManager and modelRegistry come from ctx, not event
|
||||
const { sessionManager, modelRegistry } = ctx;
|
||||
const { messagesToSummarize, messagesToKeep, tokensBefore, firstKeptEntryId, cutPoint } = preparation;
|
||||
|
||||
// Get previous summary from most recent compaction
|
||||
const _previousSummary = previousCompactions[0]?.summary;
|
||||
const { preparation, branchEntries, signal } = event;
|
||||
// sessionManager, modelRegistry, and model come from ctx
|
||||
const { sessionManager, modelRegistry, model } = ctx;
|
||||
const {
|
||||
messagesToSummarize,
|
||||
turnPrefixMessages,
|
||||
tokensBefore,
|
||||
firstKeptEntryId,
|
||||
isSplitTurn,
|
||||
previousSummary,
|
||||
} = preparation;
|
||||
|
||||
// Verify types
|
||||
expect(Array.isArray(messagesToSummarize)).toBe(true);
|
||||
expect(Array.isArray(messagesToKeep)).toBe(true);
|
||||
expect(typeof cutPoint.firstKeptEntryIndex).toBe("number");
|
||||
expect(Array.isArray(turnPrefixMessages)).toBe(true);
|
||||
expect(typeof isSplitTurn).toBe("boolean");
|
||||
expect(typeof tokensBefore).toBe("number");
|
||||
expect(model).toBeDefined();
|
||||
expect(typeof sessionManager.getEntries).toBe("function");
|
||||
expect(typeof modelRegistry.getApiKey).toBe("function");
|
||||
expect(typeof firstKeptEntryId).toBe("string");
|
||||
expect(Array.isArray(branchEntries)).toBe(true);
|
||||
|
||||
const summary = messagesToSummarize
|
||||
.filter((m) => m.role === "user")
|
||||
|
|
|
|||
|
|
@ -143,12 +143,12 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
|
|||
|
||||
const beforeEvent = beforeCompactEvents[0];
|
||||
expect(beforeEvent.preparation).toBeDefined();
|
||||
expect(beforeEvent.preparation.cutPoint.firstKeptEntryIndex).toBeGreaterThanOrEqual(0);
|
||||
expect(beforeEvent.preparation.messagesToSummarize).toBeDefined();
|
||||
expect(beforeEvent.preparation.messagesToKeep).toBeDefined();
|
||||
expect(beforeEvent.preparation.turnPrefixMessages).toBeDefined();
|
||||
expect(beforeEvent.preparation.tokensBefore).toBeGreaterThanOrEqual(0);
|
||||
expect(beforeEvent.model).toBeDefined();
|
||||
// sessionManager and modelRegistry are now on ctx, not event
|
||||
expect(typeof beforeEvent.preparation.isSplitTurn).toBe("boolean");
|
||||
expect(beforeEvent.branchEntries).toBeDefined();
|
||||
// sessionManager, modelRegistry, and model are now on ctx, not event
|
||||
|
||||
const afterEvent = compactEvents[0];
|
||||
expect(afterEvent.compactionEntry).toBeDefined();
|
||||
|
|
@ -363,19 +363,17 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
|
|||
|
||||
expect(capturedBeforeEvent).not.toBeNull();
|
||||
const event = capturedBeforeEvent!;
|
||||
expect(event.preparation.cutPoint).toHaveProperty("firstKeptEntryIndex");
|
||||
expect(event.preparation.cutPoint).toHaveProperty("isSplitTurn");
|
||||
expect(event.preparation.cutPoint).toHaveProperty("turnStartIndex");
|
||||
expect(typeof event.preparation.isSplitTurn).toBe("boolean");
|
||||
expect(event.preparation.firstKeptEntryId).toBeDefined();
|
||||
|
||||
expect(Array.isArray(event.preparation.messagesToSummarize)).toBe(true);
|
||||
expect(Array.isArray(event.preparation.messagesToKeep)).toBe(true);
|
||||
expect(Array.isArray(event.preparation.turnPrefixMessages)).toBe(true);
|
||||
|
||||
expect(typeof event.preparation.tokensBefore).toBe("number");
|
||||
|
||||
expect(event.model).toHaveProperty("provider");
|
||||
expect(event.model).toHaveProperty("id");
|
||||
expect(Array.isArray(event.branchEntries)).toBe(true);
|
||||
|
||||
// sessionManager and modelRegistry are now on ctx, not event
|
||||
// sessionManager, modelRegistry, and model are now on ctx, not event
|
||||
// Verify they're accessible via session
|
||||
expect(typeof session.sessionManager.getEntries).toBe("function");
|
||||
expect(typeof session.modelRegistry.getApiKey).toBe("function");
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import {
|
|||
DEFAULT_COMPACTION_SETTINGS,
|
||||
findCutPoint,
|
||||
getLastAssistantUsage,
|
||||
prepareCompaction,
|
||||
shouldCompact,
|
||||
} from "../src/core/compaction/index.js";
|
||||
import {
|
||||
|
|
@ -398,12 +399,10 @@ describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => {
|
|||
const entries = loadLargeSessionEntries();
|
||||
const model = getModel("anthropic", "claude-sonnet-4-5")!;
|
||||
|
||||
const compactionResult = await compact(
|
||||
entries,
|
||||
model,
|
||||
DEFAULT_COMPACTION_SETTINGS,
|
||||
process.env.ANTHROPIC_OAUTH_TOKEN!,
|
||||
);
|
||||
const preparation = prepareCompaction(entries, DEFAULT_COMPACTION_SETTINGS);
|
||||
expect(preparation).toBeDefined();
|
||||
|
||||
const compactionResult = await compact(preparation!, model, process.env.ANTHROPIC_OAUTH_TOKEN!);
|
||||
|
||||
expect(compactionResult.summary.length).toBeGreaterThan(100);
|
||||
expect(compactionResult.firstKeptEntryId).toBeTruthy();
|
||||
|
|
@ -421,12 +420,10 @@ describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => {
|
|||
const loaded = buildSessionContext(entries);
|
||||
const model = getModel("anthropic", "claude-sonnet-4-5")!;
|
||||
|
||||
const compactionResult = await compact(
|
||||
entries,
|
||||
model,
|
||||
DEFAULT_COMPACTION_SETTINGS,
|
||||
process.env.ANTHROPIC_OAUTH_TOKEN!,
|
||||
);
|
||||
const preparation = prepareCompaction(entries, DEFAULT_COMPACTION_SETTINGS);
|
||||
expect(preparation).toBeDefined();
|
||||
|
||||
const compactionResult = await compact(preparation!, model, process.env.ANTHROPIC_OAUTH_TOKEN!);
|
||||
|
||||
// Simulate appending compaction to entries by creating a proper entry
|
||||
const lastEntry = entries[entries.length - 1];
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue