refactor(agent): harden harness session semantics
This commit is contained in:
Mario Zechner
@@ -1,10 +1,3 @@
|
||||
/**
|
||||
* Branch summarization for tree navigation.
|
||||
*
|
||||
* When navigating to a different point in the session tree, this generates
|
||||
* a summary of the branch being left so context isn't lost.
|
||||
*/
|
||||
|
||||
import type { Model } from "@earendil-works/pi-ai";
|
||||
import { completeSimple } from "@earendil-works/pi-ai";
|
||||
import type { AgentMessage } from "../../types.js";
|
||||
@@ -14,102 +7,75 @@ import {
|
||||
createCompactionSummaryMessage,
|
||||
createCustomMessage,
|
||||
} from "../messages.js";
|
||||
import type { Session, SessionTreeEntry } from "../types.js";
|
||||
import { estimateTokens } from "./compaction.js";
|
||||
import type { BranchSummaryResult, Session, SessionTreeEntry } from "../types.js";
|
||||
import { BranchSummaryError, err, ok, type Result, SessionError } from "../types.js";
|
||||
import { estimateTokens, SUMMARIZATION_SYSTEM_PROMPT } from "./compaction.js";
|
||||
import {
|
||||
computeFileLists,
|
||||
createFileOps,
|
||||
extractFileOpsFromMessage,
|
||||
type FileOperations,
|
||||
formatFileOperations,
|
||||
SUMMARIZATION_SYSTEM_PROMPT,
|
||||
serializeConversation,
|
||||
} from "./utils.js";
|
||||
|
||||
// ============================================================================
|
||||
// Types
|
||||
// ============================================================================
|
||||
|
||||
export interface BranchSummaryResult {
|
||||
summary?: string;
|
||||
readFiles?: string[];
|
||||
modifiedFiles?: string[];
|
||||
aborted?: boolean;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
/** Details stored in BranchSummaryEntry.details for file tracking */
|
||||
/** File-operation details stored on generated branch summary entries. */
|
||||
export interface BranchSummaryDetails {
|
||||
/** Files read while exploring the summarized branch. */
|
||||
readFiles: string[];
|
||||
/** Files modified while exploring the summarized branch. */
|
||||
modifiedFiles: string[];
|
||||
}
|
||||
|
||||
export type { FileOperations } from "./utils.js";
|
||||
|
||||
/** Prepared branch content for summarization. */
|
||||
export interface BranchPreparation {
|
||||
/** Messages extracted for summarization, in chronological order */
|
||||
/** Messages selected for the branch summary. */
|
||||
messages: AgentMessage[];
|
||||
/** File operations extracted from tool calls */
|
||||
/** File operations extracted from the branch. */
|
||||
fileOps: FileOperations;
|
||||
/** Total estimated tokens in messages */
|
||||
/** Estimated token count for selected messages. */
|
||||
totalTokens: number;
|
||||
}
|
||||
|
||||
/** Entries selected for branch summarization. */
|
||||
export interface CollectEntriesResult {
|
||||
/** Entries to summarize, in chronological order */
|
||||
/** Entries to summarize in chronological order. */
|
||||
entries: SessionTreeEntry[];
|
||||
/** Common ancestor between old and new position, if any */
|
||||
/** Deepest common ancestor between the previous leaf and target entry. */
|
||||
commonAncestorId: string | null;
|
||||
}
|
||||
|
||||
/** Options for generating a branch summary. */
|
||||
export interface GenerateBranchSummaryOptions {
|
||||
/** Model to use for summarization */
|
||||
/** Model used for summarization. */
|
||||
model: Model<any>;
|
||||
/** API key for the model */
|
||||
/** API key forwarded to the provider. */
|
||||
apiKey: string;
|
||||
/** Request headers for the model */
|
||||
/** Optional request headers forwarded to the provider. */
|
||||
headers?: Record<string, string>;
|
||||
/** Abort signal for cancellation */
|
||||
/** Abort signal for the summarization request. */
|
||||
signal: AbortSignal;
|
||||
/** Optional custom instructions for summarization */
|
||||
/** Optional instructions appended to or replacing the default prompt. */
|
||||
customInstructions?: string;
|
||||
/** If true, customInstructions replaces the default prompt instead of being appended */
|
||||
/** Replace the default prompt with custom instructions instead of appending them. */
|
||||
replaceInstructions?: boolean;
|
||||
/** Tokens reserved for prompt + LLM response (default 16384) */
|
||||
/** Tokens reserved for prompt and model output. Defaults to 16384. */
|
||||
reserveTokens?: number;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Entry Collection
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* Collect entries that should be summarized when navigating from one position to another.
|
||||
*
|
||||
* Walks from oldLeafId back to the common ancestor with targetId, collecting entries
|
||||
* along the way. Does NOT stop at compaction boundaries - those are included and their
|
||||
* summaries become context.
|
||||
*
|
||||
* @param session - Session manager (read-only access)
|
||||
* @param oldLeafId - Current position (where we're navigating from)
|
||||
* @param targetId - Target position (where we're navigating to)
|
||||
* @returns Entries to summarize and the common ancestor
|
||||
*/
|
||||
/** Collect entries that should be summarized before navigating to a different session tree entry. */
|
||||
export async function collectEntriesForBranchSummary(
|
||||
session: Session,
|
||||
oldLeafId: string | null,
|
||||
targetId: string,
|
||||
): Promise<CollectEntriesResult> {
|
||||
// If no old position, nothing to summarize
|
||||
if (!oldLeafId) {
|
||||
return { entries: [], commonAncestorId: null };
|
||||
}
|
||||
|
||||
// Find common ancestor (deepest node that's on both paths)
|
||||
const oldPath = new Set((await session.getBranch(oldLeafId)).map((e) => e.id));
|
||||
const targetPath = await session.getBranch(targetId);
|
||||
|
||||
// targetPath is root-first, so iterate backwards to find deepest common ancestor
|
||||
let commonAncestorId: string | null = null;
|
||||
for (let i = targetPath.length - 1; i >= 0; i--) {
|
||||
if (oldPath.has(targetPath[i].id)) {
|
||||
@@ -117,36 +83,22 @@ export async function collectEntriesForBranchSummary(
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Collect entries from old leaf back to common ancestor
|
||||
const entries: SessionTreeEntry[] = [];
|
||||
let current: string | null = oldLeafId;
|
||||
|
||||
while (current && current !== commonAncestorId) {
|
||||
const entry = await session.getEntry(current);
|
||||
if (!entry) break;
|
||||
if (!entry) throw new SessionError("invalid_session", `Entry ${current} not found`);
|
||||
entries.push(entry as SessionTreeEntry);
|
||||
current = entry.parentId;
|
||||
}
|
||||
|
||||
// Reverse to get chronological order
|
||||
entries.reverse();
|
||||
|
||||
return { entries, commonAncestorId };
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Entry to Message Conversion
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* Extract AgentMessage from a session entry.
|
||||
* Similar to getMessageFromEntry in compaction.ts but also handles compaction entries.
|
||||
*/
|
||||
function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined {
|
||||
switch (entry.type) {
|
||||
case "message":
|
||||
// Skip tool results - context is in assistant's tool call
|
||||
if (entry.message.role === "toolResult") return undefined;
|
||||
return entry.message;
|
||||
|
||||
@@ -158,38 +110,21 @@ function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined
|
||||
|
||||
case "compaction":
|
||||
return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp);
|
||||
|
||||
// These don't contribute to conversation content
|
||||
case "thinking_level_change":
|
||||
case "model_change":
|
||||
case "custom":
|
||||
case "label":
|
||||
case "session_info":
|
||||
case "leaf":
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Prepare entries for summarization with token budget.
|
||||
*
|
||||
* Walks entries from NEWEST to OLDEST, adding messages until we hit the token budget.
|
||||
* This ensures we keep the most recent context when the branch is too long.
|
||||
*
|
||||
* Also collects file operations from:
|
||||
* - Tool calls in assistant messages
|
||||
* - Existing branch_summary entries' details (for cumulative tracking)
|
||||
*
|
||||
* @param entries - Entries in chronological order
|
||||
* @param tokenBudget - Maximum tokens to include (0 = no limit)
|
||||
*/
|
||||
/** Prepare branch entries for summarization within an optional token budget. */
|
||||
export function prepareBranchEntries(entries: SessionTreeEntry[], tokenBudget: number = 0): BranchPreparation {
|
||||
const messages: AgentMessage[] = [];
|
||||
const fileOps = createFileOps();
|
||||
let totalTokens = 0;
|
||||
|
||||
// First pass: collect file ops from ALL entries (even if they don't fit in token budget)
|
||||
// This ensures we capture cumulative file tracking from nested branch summaries
|
||||
// Only extract from pi-generated summaries (fromHook !== true), not extension-generated ones
|
||||
for (const entry of entries) {
|
||||
if (entry.type === "branch_summary" && !entry.fromHook && entry.details) {
|
||||
const details = entry.details as BranchSummaryDetails;
|
||||
@@ -197,35 +132,26 @@ export function prepareBranchEntries(entries: SessionTreeEntry[], tokenBudget: n
|
||||
for (const f of details.readFiles) fileOps.read.add(f);
|
||||
}
|
||||
if (Array.isArray(details.modifiedFiles)) {
|
||||
// Modified files go into both edited and written for proper deduplication
|
||||
for (const f of details.modifiedFiles) {
|
||||
fileOps.edited.add(f);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Second pass: walk from newest to oldest, adding messages until token budget
|
||||
for (let i = entries.length - 1; i >= 0; i--) {
|
||||
const entry = entries[i];
|
||||
const message = getMessageFromEntry(entry);
|
||||
if (!message) continue;
|
||||
|
||||
// Extract file ops from assistant messages (tool calls)
|
||||
extractFileOpsFromMessage(message, fileOps);
|
||||
|
||||
const tokens = estimateTokens(message);
|
||||
|
||||
// Check budget before adding
|
||||
if (tokenBudget > 0 && totalTokens + tokens > tokenBudget) {
|
||||
// If this is a summary entry, try to fit it anyway as it's important context
|
||||
if (entry.type === "compaction" || entry.type === "branch_summary") {
|
||||
if (totalTokens < tokenBudget * 0.9) {
|
||||
messages.unshift(message);
|
||||
totalTokens += tokens;
|
||||
}
|
||||
}
|
||||
// Stop - we've hit the budget
|
||||
break;
|
||||
}
|
||||
|
||||
@@ -236,10 +162,6 @@ export function prepareBranchEntries(entries: SessionTreeEntry[], tokenBudget: n
|
||||
return { messages, fileOps, totalTokens };
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Summary Generation
|
||||
// ============================================================================
|
||||
|
||||
const BRANCH_SUMMARY_PREAMBLE = `The user explored a different conversation branch before returning here.
|
||||
Summary of that exploration:
|
||||
|
||||
@@ -274,34 +196,22 @@ Use this EXACT format:
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages.`;
|
||||
|
||||
/**
|
||||
* Generate a summary of abandoned branch entries.
|
||||
*
|
||||
* @param entries - Session entries to summarize (chronological order)
|
||||
* @param options - Generation options
|
||||
*/
|
||||
/** Generate a summary for abandoned branch entries. */
|
||||
export async function generateBranchSummary(
|
||||
entries: SessionTreeEntry[],
|
||||
options: GenerateBranchSummaryOptions,
|
||||
): Promise<BranchSummaryResult> {
|
||||
): Promise<Result<BranchSummaryResult, BranchSummaryError>> {
|
||||
const { model, apiKey, headers, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options;
|
||||
|
||||
// Token budget = context window minus reserved space for prompt + response
|
||||
const contextWindow = model.contextWindow || 128000;
|
||||
const tokenBudget = contextWindow - reserveTokens;
|
||||
|
||||
const { messages, fileOps } = prepareBranchEntries(entries, tokenBudget);
|
||||
|
||||
if (messages.length === 0) {
|
||||
return { summary: "No content to summarize" };
|
||||
return ok({ summary: "No content to summarize", readFiles: [], modifiedFiles: [] });
|
||||
}
|
||||
|
||||
// Transform to LLM-compatible messages, then serialize to text
|
||||
// Serialization prevents the model from treating it as a conversation to continue
|
||||
const llmMessages = convertToLlm(messages);
|
||||
const conversationText = serializeConversation(llmMessages);
|
||||
|
||||
// Build prompt
|
||||
let instructions: string;
|
||||
if (replaceInstructions && customInstructions) {
|
||||
instructions = customInstructions;
|
||||
@@ -319,37 +229,34 @@ export async function generateBranchSummary(
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
];
|
||||
|
||||
// Call LLM for summarization
|
||||
const response = await completeSimple(
|
||||
model,
|
||||
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
|
||||
{ apiKey, headers, signal, maxTokens: 2048 },
|
||||
);
|
||||
|
||||
// Check if aborted or errored
|
||||
if (response.stopReason === "aborted") {
|
||||
return { aborted: true };
|
||||
return err(new BranchSummaryError("aborted", response.errorMessage || "Branch summary aborted"));
|
||||
}
|
||||
if (response.stopReason === "error") {
|
||||
return { error: response.errorMessage || "Summarization failed" };
|
||||
return err(
|
||||
new BranchSummaryError(
|
||||
"summarization_failed",
|
||||
`Branch summary failed: ${response.errorMessage || "Unknown error"}`,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
let summary = response.content
|
||||
.filter((c): c is { type: "text"; text: string } => c.type === "text")
|
||||
.map((c) => c.text)
|
||||
.join("\n");
|
||||
|
||||
// Prepend preamble to provide context about the branch summary
|
||||
summary = BRANCH_SUMMARY_PREAMBLE + summary;
|
||||
|
||||
// Compute file lists and append to summary
|
||||
const { readFiles, modifiedFiles } = computeFileLists(fileOps);
|
||||
summary += formatFileOperations(readFiles, modifiedFiles);
|
||||
|
||||
return {
|
||||
return ok({
|
||||
summary: summary || "No summary generated",
|
||||
readFiles,
|
||||
modifiedFiles,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1,10 +1,3 @@
|
||||
/**
|
||||
* Context compaction for long sessions.
|
||||
*
|
||||
* Pure functions for compaction logic. The session manager handles I/O,
|
||||
* and after compaction the session is reloaded.
|
||||
*/
|
||||
|
||||
import type { AssistantMessage, ImageContent, Model, TextContent, Usage } from "@earendil-works/pi-ai";
|
||||
import { completeSimple } from "@earendil-works/pi-ai";
|
||||
import type { AgentMessage, ThinkingLevel } from "../../types.js";
|
||||
@@ -22,35 +15,33 @@ import {
|
||||
extractFileOpsFromMessage,
|
||||
type FileOperations,
|
||||
formatFileOperations,
|
||||
SUMMARIZATION_SYSTEM_PROMPT,
|
||||
serializeConversation,
|
||||
} from "./utils.js";
|
||||
|
||||
// ============================================================================
|
||||
// File Operation Tracking
|
||||
// ============================================================================
|
||||
|
||||
/** Details stored in CompactionEntry.details for file tracking */
|
||||
/** File-operation details stored on generated compaction entries. */
|
||||
export interface CompactionDetails {
|
||||
/** Files read in the compacted history. */
|
||||
readFiles: string[];
|
||||
/** Files modified in the compacted history. */
|
||||
modifiedFiles: string[];
|
||||
}
|
||||
function safeJsonStringify(value: unknown): string {
|
||||
try {
|
||||
return JSON.stringify(value) ?? "undefined";
|
||||
} catch {
|
||||
return "[unserializable]";
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract file operations from messages and previous compaction entries.
|
||||
*/
|
||||
function extractFileOperations(
|
||||
messages: AgentMessage[],
|
||||
entries: SessionTreeEntry[],
|
||||
prevCompactionIndex: number,
|
||||
): FileOperations {
|
||||
const fileOps = createFileOps();
|
||||
|
||||
// Collect from previous compaction's details (if pi-generated)
|
||||
if (prevCompactionIndex >= 0) {
|
||||
const prevCompaction = entries[prevCompactionIndex] as CompactionEntry;
|
||||
if (!prevCompaction.fromHook && prevCompaction.details) {
|
||||
// fromHook field kept for session file compatibility
|
||||
const details = prevCompaction.details as CompactionDetails;
|
||||
if (Array.isArray(details.readFiles)) {
|
||||
for (const f of details.readFiles) fileOps.read.add(f);
|
||||
@@ -60,23 +51,12 @@ function extractFileOperations(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Extract from tool calls in messages
|
||||
for (const msg of messages) {
|
||||
extractFileOpsFromMessage(msg, fileOps);
|
||||
}
|
||||
|
||||
return fileOps;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Message Extraction
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* Extract AgentMessage from an entry if it produces one.
|
||||
* Returns undefined for entries that don't contribute to LLM context.
|
||||
*/
|
||||
function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined {
|
||||
if (entry.type === "message") {
|
||||
return entry.message as AgentMessage;
|
||||
@@ -106,47 +86,39 @@ function getMessageFromEntryForCompaction(entry: SessionTreeEntry): AgentMessage
|
||||
return getMessageFromEntry(entry);
|
||||
}
|
||||
|
||||
/** Result from compact() - SessionManager adds uuid/parentUuid when saving */
|
||||
/** Generated compaction data ready to be persisted as a compaction entry. */
|
||||
export interface CompactionResult<T = unknown> {
|
||||
/** Summary text that replaces compacted history in future context. */
|
||||
summary: string;
|
||||
/** Entry id where retained history starts. */
|
||||
firstKeptEntryId: string;
|
||||
/** Estimated context tokens before compaction. */
|
||||
tokensBefore: number;
|
||||
/** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */
|
||||
/** Optional implementation-specific details stored with the compaction entry. */
|
||||
details?: T;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Types
|
||||
// ============================================================================
|
||||
|
||||
/** Compaction thresholds and retention settings. */
|
||||
export interface CompactionSettings {
|
||||
/** Enable automatic compaction decisions. */
|
||||
enabled: boolean;
|
||||
/** Tokens reserved for summary prompt and output. */
|
||||
reserveTokens: number;
|
||||
/** Approximate recent-context tokens to keep after compaction. */
|
||||
keepRecentTokens: number;
|
||||
}
|
||||
|
||||
/** Default compaction settings used by the harness. */
|
||||
export const DEFAULT_COMPACTION_SETTINGS: CompactionSettings = {
|
||||
enabled: true,
|
||||
reserveTokens: 16384,
|
||||
keepRecentTokens: 20000,
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// Token calculation
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* Calculate total context tokens from usage.
|
||||
* Uses the native totalTokens field when available, falls back to computing from components.
|
||||
*/
|
||||
/** Calculate total context tokens from provider usage. */
|
||||
export function calculateContextTokens(usage: Usage): number {
|
||||
return usage.totalTokens || usage.input + usage.output + usage.cacheRead + usage.cacheWrite;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get usage from an assistant message if available.
|
||||
* Skips aborted and error messages as they don't have valid usage data.
|
||||
*/
|
||||
function getAssistantUsage(msg: AgentMessage): Usage | undefined {
|
||||
if (msg.role === "assistant" && "usage" in msg) {
|
||||
const assistantMsg = msg as AssistantMessage;
|
||||
@@ -157,9 +129,7 @@ function getAssistantUsage(msg: AgentMessage): Usage | undefined {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Find the last non-aborted assistant message usage from session entries.
|
||||
*/
|
||||
/** Return usage from the last successful assistant message in session entries. */
|
||||
export function getLastAssistantUsage(entries: SessionTreeEntry[]): Usage | undefined {
|
||||
for (let i = entries.length - 1; i >= 0; i--) {
|
||||
const entry = entries[i];
|
||||
@@ -171,10 +141,15 @@ export function getLastAssistantUsage(entries: SessionTreeEntry[]): Usage | unde
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/** Estimated context-token usage for a message list. */
|
||||
export interface ContextUsageEstimate {
|
||||
/** Estimated total context tokens. */
|
||||
tokens: number;
|
||||
/** Tokens reported by the most recent assistant usage block. */
|
||||
usageTokens: number;
|
||||
/** Estimated tokens after the most recent assistant usage block. */
|
||||
trailingTokens: number;
|
||||
/** Index of the message that provided usage, or null when none exists. */
|
||||
lastUsageIndex: number | null;
|
||||
}
|
||||
|
||||
@@ -186,10 +161,7 @@ function getLastAssistantUsageInfo(messages: AgentMessage[]): { usage: Usage; in
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Estimate context tokens from messages, using the last assistant usage when available.
|
||||
* If there are messages after the last usage, estimate their tokens with estimateTokens.
|
||||
*/
|
||||
/** Estimate context tokens for messages using provider usage when available. */
|
||||
export function estimateContextTokens(messages: AgentMessage[]): ContextUsageEstimate {
|
||||
const usageInfo = getLastAssistantUsageInfo(messages);
|
||||
|
||||
@@ -220,22 +192,13 @@ export function estimateContextTokens(messages: AgentMessage[]): ContextUsageEst
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if compaction should trigger based on context usage.
|
||||
*/
|
||||
/** Return whether context usage exceeds the configured compaction threshold. */
|
||||
export function shouldCompact(contextTokens: number, contextWindow: number, settings: CompactionSettings): boolean {
|
||||
if (!settings.enabled) return false;
|
||||
return contextTokens > contextWindow - settings.reserveTokens;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Cut point detection
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* Estimate token count for a message using chars/4 heuristic.
|
||||
* This is conservative (overestimates tokens).
|
||||
*/
|
||||
/** Estimate token count for one message using a conservative character heuristic. */
|
||||
export function estimateTokens(message: AgentMessage): number {
|
||||
let chars = 0;
|
||||
|
||||
@@ -261,7 +224,7 @@ export function estimateTokens(message: AgentMessage): number {
|
||||
} else if (block.type === "thinking") {
|
||||
chars += block.thinking.length;
|
||||
} else if (block.type === "toolCall") {
|
||||
chars += block.name.length + JSON.stringify(block.arguments).length;
|
||||
chars += block.name.length + safeJsonStringify(block.arguments).length;
|
||||
}
|
||||
}
|
||||
return Math.ceil(chars / 4);
|
||||
@@ -276,7 +239,7 @@ export function estimateTokens(message: AgentMessage): number {
|
||||
chars += block.text.length;
|
||||
}
|
||||
if (block.type === "image") {
|
||||
chars += 4800; // Estimate images as 4000 chars, or 1200 tokens
|
||||
chars += 4800;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -295,14 +258,6 @@ export function estimateTokens(message: AgentMessage): number {
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* Find valid cut points: indices of user, assistant, custom, or bashExecution messages.
|
||||
* Never cut at tool results (they must follow their tool call).
|
||||
* When we cut at an assistant message with tool calls, its tool results follow it
|
||||
* and will be kept.
|
||||
* BashExecutionMessage is treated like a user message (user-initiated context).
|
||||
*/
|
||||
function findValidCutPoints(entries: SessionTreeEntry[], startIndex: number, endIndex: number): number[] {
|
||||
const cutPoints: number[] = [];
|
||||
for (let i = startIndex; i < endIndex; i++) {
|
||||
@@ -332,10 +287,9 @@ function findValidCutPoints(entries: SessionTreeEntry[], startIndex: number, end
|
||||
case "custom_message":
|
||||
case "label":
|
||||
case "session_info":
|
||||
case "leaf":
|
||||
break;
|
||||
}
|
||||
|
||||
// branch_summary and custom_message are user-role messages, valid cut points
|
||||
if (entry.type === "branch_summary" || entry.type === "custom_message") {
|
||||
cutPoints.push(i);
|
||||
}
|
||||
@@ -343,15 +297,10 @@ function findValidCutPoints(entries: SessionTreeEntry[], startIndex: number, end
|
||||
return cutPoints;
|
||||
}
|
||||
|
||||
/**
|
||||
* Find the user message (or bashExecution) that starts the turn containing the given entry index.
|
||||
* Returns -1 if no turn start found before the index.
|
||||
* BashExecutionMessage is treated like a user message for turn boundaries.
|
||||
*/
|
||||
/** Find the user-visible message that starts the turn containing an entry. */
|
||||
export function findTurnStartIndex(entries: SessionTreeEntry[], entryIndex: number, startIndex: number): number {
|
||||
for (let i = entryIndex; i >= startIndex; i--) {
|
||||
const entry = entries[i];
|
||||
// branch_summary and custom_message are user-role messages, can start a turn
|
||||
if (entry.type === "branch_summary" || entry.type === "custom_message") {
|
||||
return i;
|
||||
}
|
||||
@@ -365,31 +314,17 @@ export function findTurnStartIndex(entries: SessionTreeEntry[], entryIndex: numb
|
||||
return -1;
|
||||
}
|
||||
|
||||
/** Cut point selected for compaction. */
|
||||
export interface CutPointResult {
|
||||
/** Index of first entry to keep */
|
||||
/** Index of the first entry retained after compaction. */
|
||||
firstKeptEntryIndex: number;
|
||||
/** Index of user message that starts the turn being split, or -1 if not splitting */
|
||||
/** Index of the turn-start entry when the cut splits a turn, otherwise -1. */
|
||||
turnStartIndex: number;
|
||||
/** Whether this cut splits a turn (cut point is not a user message) */
|
||||
/** Whether the selected cut point splits an in-progress turn. */
|
||||
isSplitTurn: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Find the cut point in session entries that keeps approximately `keepRecentTokens`.
|
||||
*
|
||||
* Algorithm: Walk backwards from newest, accumulating estimated message sizes.
|
||||
* Stop when we've accumulated >= keepRecentTokens. Cut at that point.
|
||||
*
|
||||
* Can cut at user OR assistant messages (never tool results). When cutting at an
|
||||
* assistant message with tool calls, its tool results come after and will be kept.
|
||||
*
|
||||
* Returns CutPointResult with:
|
||||
* - firstKeptEntryIndex: the entry index to start keeping from
|
||||
* - turnStartIndex: if cutting mid-turn, the user message that started that turn
|
||||
* - isSplitTurn: whether we're cutting in the middle of a turn
|
||||
*
|
||||
* Only considers entries between `startIndex` and `endIndex` (exclusive).
|
||||
*/
|
||||
/** Find the compaction cut point that keeps approximately the requested recent-token budget. */
|
||||
export function findCutPoint(
|
||||
entries: SessionTreeEntry[],
|
||||
startIndex: number,
|
||||
@@ -401,22 +336,15 @@ export function findCutPoint(
|
||||
if (cutPoints.length === 0) {
|
||||
return { firstKeptEntryIndex: startIndex, turnStartIndex: -1, isSplitTurn: false };
|
||||
}
|
||||
|
||||
// Walk backwards from newest, accumulating estimated message sizes
|
||||
let accumulatedTokens = 0;
|
||||
let cutIndex = cutPoints[0]; // Default: keep from first message (not header)
|
||||
let cutIndex = cutPoints[0];
|
||||
|
||||
for (let i = endIndex - 1; i >= startIndex; i--) {
|
||||
const entry = entries[i];
|
||||
if (entry.type !== "message") continue;
|
||||
|
||||
// Estimate this message's size
|
||||
const messageTokens = estimateTokens(entry.message as AgentMessage);
|
||||
accumulatedTokens += messageTokens;
|
||||
|
||||
// Check if we've exceeded the budget
|
||||
if (accumulatedTokens >= keepRecentTokens) {
|
||||
// Find the closest valid cut point at or after this entry
|
||||
for (let c = 0; c < cutPoints.length; c++) {
|
||||
if (cutPoints[c] >= i) {
|
||||
cutIndex = cutPoints[c];
|
||||
@@ -426,23 +354,16 @@ export function findCutPoint(
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Scan backwards from cutIndex to include any non-message entries (bash, settings, etc.)
|
||||
while (cutIndex > startIndex) {
|
||||
const prevEntry = entries[cutIndex - 1];
|
||||
// Stop at session header or compaction boundaries
|
||||
if (prevEntry.type === "compaction") {
|
||||
break;
|
||||
}
|
||||
if (prevEntry.type === "message") {
|
||||
// Stop if we hit any message
|
||||
break;
|
||||
}
|
||||
// Include this non-message entry (bash, settings change, etc.)
|
||||
cutIndex--;
|
||||
}
|
||||
|
||||
// Determine if this is a split turn
|
||||
const cutEntry = entries[cutIndex];
|
||||
const isUserMessage = cutEntry.type === "message" && cutEntry.message.role === "user";
|
||||
const turnStartIndex = isUserMessage ? -1 : findTurnStartIndex(entries, cutIndex, startIndex);
|
||||
@@ -454,9 +375,9 @@ export function findCutPoint(
|
||||
};
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Summarization
|
||||
// ============================================================================
|
||||
export const SUMMARIZATION_SYSTEM_PROMPT = `You are a context summarization assistant. Your task is to read a conversation between a user and an AI coding assistant, then produce a structured summary following the exact format specified.
|
||||
|
||||
Do NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.`;
|
||||
|
||||
const SUMMARIZATION_PROMPT = `The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work.
|
||||
|
||||
@@ -530,10 +451,7 @@ Use this EXACT format:
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages.`;
|
||||
|
||||
/**
|
||||
* Generate a summary of the conversation using the LLM.
|
||||
* If previousSummary is provided, uses the update prompt to merge.
|
||||
*/
|
||||
/** Generate or update a conversation summary for compaction. */
|
||||
export async function generateSummary(
|
||||
currentMessages: AgentMessage[],
|
||||
model: Model<any>,
|
||||
@@ -549,19 +467,12 @@ export async function generateSummary(
|
||||
Math.floor(0.8 * reserveTokens),
|
||||
model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY,
|
||||
);
|
||||
|
||||
// Use update prompt if we have a previous summary, otherwise initial prompt
|
||||
let basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT;
|
||||
if (customInstructions) {
|
||||
basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`;
|
||||
}
|
||||
|
||||
// Serialize conversation to text so model doesn't try to continue it
|
||||
// Convert to LLM messages first (handles custom types like bashExecution, custom, etc.)
|
||||
const llmMessages = convertToLlm(currentMessages);
|
||||
const conversationText = serializeConversation(llmMessages);
|
||||
|
||||
// Build the prompt with conversation wrapped in tags
|
||||
let promptText = `<conversation>\n${conversationText}\n</conversation>\n\n`;
|
||||
if (previousSummary) {
|
||||
promptText += `<previous-summary>\n${previousSummary}\n</previous-summary>\n\n`;
|
||||
@@ -581,18 +492,11 @@ export async function generateSummary(
|
||||
? { maxTokens, signal, apiKey, headers, reasoning: thinkingLevel }
|
||||
: { maxTokens, signal, apiKey, headers };
|
||||
|
||||
const responseResult: Result<AssistantMessage, CompactionError> = await completeSimple(
|
||||
const response = await completeSimple(
|
||||
model,
|
||||
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
|
||||
completionOptions,
|
||||
).then(
|
||||
(response) => ok<AssistantMessage, CompactionError>(response),
|
||||
(error: unknown) =>
|
||||
err<AssistantMessage, CompactionError>(new CompactionError("unknown", "Summarization request failed", error)),
|
||||
);
|
||||
if (!responseResult.ok) return err(responseResult.error);
|
||||
|
||||
const response = responseResult.value;
|
||||
if (response.stopReason === "aborted") {
|
||||
return err(new CompactionError("aborted", response.errorMessage || "Summarization aborted"));
|
||||
}
|
||||
@@ -613,34 +517,33 @@ export async function generateSummary(
|
||||
return ok(textContent);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Compaction Preparation (for extensions)
|
||||
// ============================================================================
|
||||
|
||||
/** Prepared inputs for a compaction run. */
|
||||
export interface CompactionPreparation {
|
||||
/** UUID of first entry to keep */
|
||||
/** Entry id where retained history starts. */
|
||||
firstKeptEntryId: string;
|
||||
/** Messages that will be summarized and discarded */
|
||||
/** Messages summarized into the history summary. */
|
||||
messagesToSummarize: AgentMessage[];
|
||||
/** Messages that will be turned into turn prefix summary (if splitting) */
|
||||
/** Prefix messages summarized separately when compaction splits a turn. */
|
||||
turnPrefixMessages: AgentMessage[];
|
||||
/** Whether this is a split turn (cut point in middle of turn) */
|
||||
/** Whether compaction splits a turn. */
|
||||
isSplitTurn: boolean;
|
||||
/** Estimated context tokens before compaction. */
|
||||
tokensBefore: number;
|
||||
/** Summary from previous compaction, for iterative update */
|
||||
/** Previous compaction summary used for iterative updates. */
|
||||
previousSummary?: string;
|
||||
/** File operations extracted from messagesToSummarize */
|
||||
/** File operations extracted from summarized history. */
|
||||
fileOps: FileOperations;
|
||||
/** Compaction settions from settings.jsonl */
|
||||
/** Settings used to prepare compaction. */
|
||||
settings: CompactionSettings;
|
||||
}
|
||||
|
||||
/** Prepare session entries for compaction, or return undefined when compaction is not applicable. */
|
||||
export function prepareCompaction(
|
||||
pathEntries: SessionTreeEntry[],
|
||||
settings: CompactionSettings,
|
||||
): CompactionPreparation | undefined {
|
||||
if (pathEntries.length > 0 && pathEntries[pathEntries.length - 1].type === "compaction") {
|
||||
return undefined;
|
||||
): Result<CompactionPreparation | undefined, CompactionError> {
|
||||
if (pathEntries.length === 0 || pathEntries[pathEntries.length - 1].type === "compaction") {
|
||||
return ok(undefined);
|
||||
}
|
||||
|
||||
let prevCompactionIndex = -1;
|
||||
@@ -664,24 +567,18 @@ export function prepareCompaction(
|
||||
const tokensBefore = estimateContextTokens(buildSessionContext(pathEntries).messages).tokens;
|
||||
|
||||
const cutPoint = findCutPoint(pathEntries, boundaryStart, boundaryEnd, settings.keepRecentTokens);
|
||||
|
||||
// Get UUID of first kept entry
|
||||
const firstKeptEntry = pathEntries[cutPoint.firstKeptEntryIndex];
|
||||
if (!firstKeptEntry?.id) {
|
||||
return undefined; // Session needs migration
|
||||
return err(new CompactionError("invalid_session", "First kept entry has no UUID - session may need migration"));
|
||||
}
|
||||
const firstKeptEntryId = firstKeptEntry.id;
|
||||
|
||||
const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex;
|
||||
|
||||
// Messages to summarize (will be discarded after summary)
|
||||
const messagesToSummarize: AgentMessage[] = [];
|
||||
for (let i = boundaryStart; i < historyEnd; i++) {
|
||||
const msg = getMessageFromEntryForCompaction(pathEntries[i]);
|
||||
if (msg) messagesToSummarize.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++) {
|
||||
@@ -689,18 +586,14 @@ export function prepareCompaction(
|
||||
if (msg) turnPrefixMessages.push(msg);
|
||||
}
|
||||
}
|
||||
|
||||
// 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 {
|
||||
return ok({
|
||||
firstKeptEntryId,
|
||||
messagesToSummarize,
|
||||
turnPrefixMessages,
|
||||
@@ -709,13 +602,9 @@ export function prepareCompaction(
|
||||
previousSummary,
|
||||
fileOps,
|
||||
settings,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Main compaction function
|
||||
// ============================================================================
|
||||
|
||||
const TURN_PREFIX_SUMMARIZATION_PROMPT = `This is the PREFIX of a turn that was too large to keep. The SUFFIX (recent work) is retained.
|
||||
|
||||
Summarize the prefix to provide context for the retained suffix:
|
||||
@@ -731,15 +620,9 @@ Summarize the prefix to provide context for the retained suffix:
|
||||
|
||||
Be concise. Focus on what's needed to understand the kept suffix.`;
|
||||
|
||||
/**
|
||||
* Generate summaries for compaction using prepared data.
|
||||
* Returns CompactionResult - SessionManager adds uuid/parentUuid when saving.
|
||||
*
|
||||
* @param preparation - Pre-calculated preparation from prepareCompaction()
|
||||
* @param customInstructions - Optional custom focus for the summary
|
||||
*/
|
||||
export { serializeConversation } from "./utils.js";
|
||||
|
||||
/** Generate compaction summary data from prepared session history. */
|
||||
export async function compact(
|
||||
preparation: CompactionPreparation,
|
||||
model: Model<any>,
|
||||
@@ -820,10 +703,6 @@ export async function compact(
|
||||
details: { readFiles, modifiedFiles } as CompactionDetails,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a summary for a turn prefix (when splitting a turn).
|
||||
*/
|
||||
async function generateTurnPrefixSummary(
|
||||
messages: AgentMessage[],
|
||||
model: Model<any>,
|
||||
@@ -836,7 +715,7 @@ async function generateTurnPrefixSummary(
|
||||
const maxTokens = Math.min(
|
||||
Math.floor(0.5 * reserveTokens),
|
||||
model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY,
|
||||
); // Smaller budget for turn prefix
|
||||
);
|
||||
const llmMessages = convertToLlm(messages);
|
||||
const conversationText = serializeConversation(llmMessages);
|
||||
const promptText = `<conversation>\n${conversationText}\n</conversation>\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`;
|
||||
@@ -848,22 +727,13 @@ async function generateTurnPrefixSummary(
|
||||
},
|
||||
];
|
||||
|
||||
const responseResult: Result<AssistantMessage, CompactionError> = await completeSimple(
|
||||
const response = await completeSimple(
|
||||
model,
|
||||
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
|
||||
model.reasoning && thinkingLevel && thinkingLevel !== "off"
|
||||
? { maxTokens, signal, apiKey, headers, reasoning: thinkingLevel }
|
||||
: { maxTokens, signal, apiKey, headers },
|
||||
).then(
|
||||
(response) => ok<AssistantMessage, CompactionError>(response),
|
||||
(error: unknown) =>
|
||||
err<AssistantMessage, CompactionError>(
|
||||
new CompactionError("unknown", "Turn prefix summarization request failed", error),
|
||||
),
|
||||
);
|
||||
if (!responseResult.ok) return err(responseResult.error);
|
||||
|
||||
const response = responseResult.value;
|
||||
if (response.stopReason === "aborted") {
|
||||
return err(new CompactionError("aborted", response.errorMessage || "Turn prefix summarization aborted"));
|
||||
}
|
||||
|
||||
@@ -1,20 +1,17 @@
|
||||
/**
|
||||
* Shared utilities for compaction and branch summarization.
|
||||
*/
|
||||
|
||||
import type { Message } from "@earendil-works/pi-ai";
|
||||
import type { AgentMessage } from "../../types.js";
|
||||
|
||||
// ============================================================================
|
||||
// File Operation Tracking
|
||||
// ============================================================================
|
||||
|
||||
/** File paths touched by a session branch or compaction range. */
|
||||
export interface FileOperations {
|
||||
/** Files read but not necessarily modified. */
|
||||
read: Set<string>;
|
||||
/** Files written by full-file write operations. */
|
||||
written: Set<string>;
|
||||
/** Files modified by edit operations. */
|
||||
edited: Set<string>;
|
||||
}
|
||||
|
||||
/** Create an empty file-operation accumulator. */
|
||||
export function createFileOps(): FileOperations {
|
||||
return {
|
||||
read: new Set(),
|
||||
@@ -23,9 +20,7 @@ export function createFileOps(): FileOperations {
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract file operations from tool calls in an assistant message.
|
||||
*/
|
||||
/** Add file operations from assistant tool calls to an accumulator. */
|
||||
export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOperations): void {
|
||||
if (message.role !== "assistant") return;
|
||||
if (!("content" in message) || !Array.isArray(message.content)) return;
|
||||
@@ -55,10 +50,7 @@ export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOp
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Compute final file lists from file operations.
|
||||
* Returns readFiles (files only read, not modified) and modifiedFiles.
|
||||
*/
|
||||
/** Compute sorted read-only and modified file lists from accumulated operations. */
|
||||
export function computeFileLists(fileOps: FileOperations): { readFiles: string[]; modifiedFiles: string[] } {
|
||||
const modified = new Set([...fileOps.edited, ...fileOps.written]);
|
||||
const readOnly = [...fileOps.read].filter((f) => !modified.has(f)).sort();
|
||||
@@ -66,9 +58,7 @@ export function computeFileLists(fileOps: FileOperations): { readFiles: string[]
|
||||
return { readFiles: readOnly, modifiedFiles };
|
||||
}
|
||||
|
||||
/**
|
||||
* Format file operations as XML tags for summary.
|
||||
*/
|
||||
/** Format file lists as summary metadata tags. */
|
||||
export function formatFileOperations(readFiles: string[], modifiedFiles: string[]): string {
|
||||
const sections: string[] = [];
|
||||
if (readFiles.length > 0) {
|
||||
@@ -81,31 +71,23 @@ export function formatFileOperations(readFiles: string[], modifiedFiles: string[
|
||||
return `\n\n${sections.join("\n\n")}`;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Message Serialization
|
||||
// ============================================================================
|
||||
|
||||
/** Maximum characters for a tool result in serialized summaries. */
|
||||
const TOOL_RESULT_MAX_CHARS = 2000;
|
||||
|
||||
/**
|
||||
* Truncate text to a maximum character length for summarization.
|
||||
* Keeps the beginning and appends a truncation marker.
|
||||
*/
|
||||
function safeJsonStringify(value: unknown): string {
|
||||
try {
|
||||
return JSON.stringify(value) ?? "undefined";
|
||||
} catch {
|
||||
return "[unserializable]";
|
||||
}
|
||||
}
|
||||
|
||||
function truncateForSummary(text: string, maxChars: number): string {
|
||||
if (text.length <= maxChars) return text;
|
||||
const truncatedChars = text.length - maxChars;
|
||||
return `${text.slice(0, maxChars)}\n\n[... ${truncatedChars} more characters truncated]`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Serialize LLM messages to text for summarization.
|
||||
* This prevents the model from treating it as a conversation to continue.
|
||||
* Call convertToLlm() first to handle custom message types.
|
||||
*
|
||||
* Tool results are truncated to keep the summarization request within
|
||||
* reasonable token budgets. Full content is not needed for summarization.
|
||||
*/
|
||||
/** Serialize LLM messages to plain text for summarization prompts. */
|
||||
export function serializeConversation(messages: Message[]): string {
|
||||
const parts: string[] = [];
|
||||
|
||||
@@ -132,7 +114,7 @@ export function serializeConversation(messages: Message[]): string {
|
||||
} else if (block.type === "toolCall") {
|
||||
const args = block.arguments as Record<string, unknown>;
|
||||
const argsStr = Object.entries(args)
|
||||
.map(([k, v]) => `${k}=${JSON.stringify(v)}`)
|
||||
.map(([k, v]) => `${k}=${safeJsonStringify(v)}`)
|
||||
.join(", ");
|
||||
toolCalls.push(`${block.name}(${argsStr})`);
|
||||
}
|
||||
@@ -160,11 +142,3 @@ export function serializeConversation(messages: Message[]): string {
|
||||
|
||||
return parts.join("\n\n");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Summarization System Prompt
|
||||
// ============================================================================
|
||||
|
||||
export const SUMMARIZATION_SYSTEM_PROMPT = `You are a context summarization assistant. Your task is to read a conversation between a user and an AI coding assistant, then produce a structured summary following the exact format specified.
|
||||
|
||||
Do NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.`;
|
||||
|
||||
Reference in New Issue
Block a user