Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
10 kB · 324 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325/** * 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 { AgentMessage } from "@earendil-works/pi-agent-core";import type { Model, Usage } from "@earendil-works/pi-ai";import { completeSimple } from "@earendil-works/pi-ai";import { convertToLlm, createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage,} from "../messages.js";import { completeWithProviderRetry, type ProviderRetryPolicy } from "../provider-retry.js";import type { ReadonlySessionManager, SessionEntry } from "../session-manager.js";import { estimateTokens } from "./compaction.js";import { computeFileLists, createFileOps, extractFileOpsFromMessage, type FileOperations, formatFileOperations, SUMMARIZATION_SYSTEM_PROMPT, serializeConversation,} from "./utils.js";export interface BranchSummaryResult { summary?: string; readFiles?: string[]; modifiedFiles?: string[]; aborted?: boolean; error?: string; usage?: Usage;}
/** Details stored in BranchSummaryEntry.details for file tracking */export interface BranchSummaryDetails { readFiles: string[]; modifiedFiles: string[];}
export type { FileOperations } from "./utils.js";
export interface BranchPreparation { /** Messages extracted for summarization, in chronological order */ messages: AgentMessage[]; /** File operations extracted from tool calls */ fileOps: FileOperations; /** Total estimated tokens in messages */ totalTokens: number;}
export interface CollectEntriesResult { /** Entries to summarize, in chronological order */ entries: SessionEntry[]; /** Common ancestor between old and new position, if any */ commonAncestorId: string | null;}
export interface GenerateBranchSummaryOptions { /** Model to use for summarization */ model: Model<any>; /** API key for the model */ apiKey: string; /** Request headers for the model */ headers?: Record<string, string>; /** Abort signal for cancellation */ signal: AbortSignal; /** Optional custom instructions for summarization */ customInstructions?: string; /** If true, customInstructions replaces the default prompt instead of being appended */ replaceInstructions?: boolean; retry?: ProviderRetryPolicy; /** Tokens reserved for prompt + LLM response (default 16384) */ reserveTokens?: number;}/** * 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 */export function collectEntriesForBranchSummary( session: ReadonlySessionManager, oldLeafId: string | null, targetId: string,): CollectEntriesResult { if (!oldLeafId) { return { entries: [], commonAncestorId: null }; } const oldPath = new Set(session.getBranch(oldLeafId).map((e) => e.id)); const targetPath = session.getBranch(targetId); let commonAncestorId: string | null = null; for (let i = targetPath.length - 1; i >= 0; i--) { if (oldPath.has(targetPath[i].id)) { commonAncestorId = targetPath[i].id; break; } } const entries: SessionEntry[] = []; let current: string | null = oldLeafId;
while (current && current !== commonAncestorId) { const entry = session.getEntry(current); if (!entry) break; entries.push(entry); current = entry.parentId; } entries.reverse();
return { entries, commonAncestorId };}/** * Extract AgentMessage from a session entry. * Similar to getMessageFromEntry in compaction.ts but also handles compaction entries. */function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined { switch (entry.type) { case "message": // Tool-result context remains attached to its assistant tool call. if (entry.message.role === "toolResult") return undefined; return entry.message;
case "custom_message": return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp);
case "branch_summary": return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);
case "compaction": return createCompactionSummaryMessage( entry.summary, entry.tokensBefore, entry.timestamp, entry.customInstructions, ); case "thinking_level_change": case "model_change": case "custom": case "label": case "session_info": 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) */export function prepareBranchEntries(entries: SessionEntry[], 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; if (Array.isArray(details.readFiles)) { for (const f of details.readFiles) fileOps.read.add(f); } if (Array.isArray(details.modifiedFiles)) { for (const f of details.modifiedFiles) { fileOps.edited.add(f); } } } } for (let i = entries.length - 1; i >= 0; i--) { const entry = entries[i]; const message = getMessageFromEntry(entry); if (!message) continue; extractFileOpsFromMessage(message, fileOps);
const tokens = estimateTokens(message); if (tokenBudget > 0 && totalTokens + tokens > tokenBudget) { if (entry.type === "compaction" || entry.type === "branch_summary") { if (totalTokens < tokenBudget * 0.9) { messages.unshift(message); totalTokens += tokens; } } break; }
messages.unshift(message); totalTokens += tokens; }
return { messages, fileOps, totalTokens };}const BRANCH_SUMMARY_PREAMBLE = `The user explored a different conversation branch before returning here.Summary of that exploration:
`;
const BRANCH_SUMMARY_PROMPT = `Create a structured summary of this conversation branch for context when returning later.
Use this EXACT format:
## Goal[What was the user trying to accomplish in this branch?]
## Constraints & Preferences- [Any constraints, preferences, or requirements mentioned]- [Or "(none)" if none were mentioned]
## Progress### Done- [x] [Completed tasks/changes]
### In Progress- [ ] [Work that was started but not finished]
### Blocked- [Issues preventing progress, if any]
## Key Decisions- **[Decision]**: [Brief rationale]
## Next Steps1. [What should happen next to continue this work]
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 */export async function generateBranchSummary( entries: SessionEntry[], options: GenerateBranchSummaryOptions,): Promise<BranchSummaryResult> { const { model, apiKey, headers, signal, customInstructions, replaceInstructions, retry, reserveTokens = 16384, } = options; const contextWindow = model.contextWindow || 128000; const tokenBudget = contextWindow - reserveTokens;
const { messages, fileOps } = prepareBranchEntries(entries, tokenBudget);
// Nothing model-visible remains after filtering. if (messages.length === 0) { return { summary: "No content to summarize" }; } // Serialize before the LLM call so it summarizes rather than continues this branch. const llmMessages = convertToLlm(messages); const conversationText = serializeConversation(llmMessages); let instructions: string; if (replaceInstructions && customInstructions) { instructions = customInstructions; } else if (customInstructions) { instructions = `${BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: ${customInstructions}`; } else { instructions = BRANCH_SUMMARY_PROMPT; } const promptText = `<conversation>\n${conversationText}\n</conversation>\n\n${instructions}`;
const summarizationMessages = [ { role: "user" as const, content: [{ type: "text" as const, text: promptText }], timestamp: Date.now(), }, ]; const response = await completeWithProviderRetry( () => completeSimple( model, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, { apiKey, headers, signal, maxTokens: 2048 }, ), { policy: retry, signal }, ); if (response.stopReason === "aborted") { return { aborted: true }; } if (response.stopReason === "error") { return { error: response.errorMessage || "Summarization failed" }; }
let summary = response.content .filter((c): c is { type: "text"; text: string } => c.type === "text") .map((c) => c.text) .join("\n"); summary = BRANCH_SUMMARY_PREAMBLE + summary; const { readFiles, modifiedFiles } = computeFileLists(fileOps); summary += formatFileOperations(readFiles, modifiedFiles);
return { summary: summary || "No summary generated", readFiles, modifiedFiles, usage: response.usage, };}