diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index 7eedffc3..b03372d2 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -3,13 +3,8 @@ import { readFileSync } from "node:fs"; import type { ImageContent, Model } from "@mariozechner/pi-ai"; import type { Agent } from "../agent.js"; import type { AgentEvent, AgentMessage, AgentTool, ThinkingLevel } from "../types.js"; -import { - collectEntriesForBranchSummary, - compact, - DEFAULT_COMPACTION_SETTINGS, - generateBranchSummary, - prepareCompaction, -} from "./compaction.js"; +import { collectEntriesForBranchSummary, generateBranchSummary } from "./compaction/branch-summarization.js"; +import { compact, DEFAULT_COMPACTION_SETTINGS, prepareCompaction } from "./compaction/compaction.js"; import { expandPromptTemplate } from "./prompt-templates.js"; import type { AbortResult, @@ -24,8 +19,6 @@ import type { ExecutionEnv, NavigateTreeResult, PromptTemplate, - SessionBeforeCompactResult, - SessionBeforeTreeResult, SessionTree, Skill, SystemPromptInputs, @@ -388,8 +381,8 @@ export class DefaultAgentHarness implements AgentHarness { customInstructions, signal: new AbortController().signal, }); - if ((hookResult as SessionBeforeCompactResult | undefined)?.cancel) throw new Error("Compaction cancelled"); - const provided = (hookResult as SessionBeforeCompactResult | undefined)?.compaction; + if (hookResult?.cancel) throw new Error("Compaction cancelled"); + const provided = hookResult?.compaction; const result = provided ?? (await compact( @@ -442,11 +435,10 @@ export class DefaultAgentHarness implements AgentHarness { preparation, signal, }); - const typedHook = hookResult as SessionBeforeTreeResult | undefined; - if (typedHook?.cancel) return { cancelled: true }; + if (hookResult?.cancel) return { cancelled: true }; let summaryEntry: any | undefined; - let summaryText: string | undefined = typedHook?.summary?.summary; - let summaryDetails: unknown = typedHook?.summary?.details; + let summaryText: string | undefined = hookResult?.summary?.summary; + let summaryDetails: unknown = hookResult?.summary?.details; if (!summaryText && options?.summarize && entries.length > 0) { const model = this.conversation.model; if (!model) throw new Error("No model set for branch summary"); @@ -457,8 +449,8 @@ export class DefaultAgentHarness implements AgentHarness { apiKey: auth.apiKey, headers: auth.headers, signal: new AbortController().signal, - customInstructions: typedHook?.customInstructions ?? options?.customInstructions, - replaceInstructions: typedHook?.replaceInstructions ?? options?.replaceInstructions, + customInstructions: hookResult?.customInstructions ?? options?.customInstructions, + replaceInstructions: hookResult?.replaceInstructions ?? options?.replaceInstructions, }); if (branchSummary.aborted) return { cancelled: true }; if (branchSummary.error) throw new Error(branchSummary.error); @@ -498,7 +490,7 @@ export class DefaultAgentHarness implements AgentHarness { newLeafId ?? "root", summaryText, summaryDetails, - typedHook?.summary !== undefined, + hookResult?.summary !== undefined, ); summaryEntry = await this.sessionTree.getEntry(summaryId); } @@ -508,7 +500,7 @@ export class DefaultAgentHarness implements AgentHarness { newLeafId: await this.sessionTree.getLeafId(), oldLeafId, summaryEntry, - fromHook: typedHook?.summary !== undefined, + fromHook: hookResult?.summary !== undefined, }); return { cancelled: false, editorText, summaryEntry }; } @@ -573,12 +565,7 @@ export class DefaultAgentHarness implements AgentHarness { await this.agent.waitForIdle(); } - subscribe( - listener: ( - event: AgentEvent | import("./types.js").AgentHarnessOwnEvent, - signal?: AbortSignal, - ) => Promise | void, - ): () => void { + subscribe(listener: (event: AgentHarnessEvent, signal?: AbortSignal) => Promise | void): () => void { this.listeners.add(listener); return () => this.listeners.delete(listener); } diff --git a/packages/agent/src/harness/compaction.ts b/packages/agent/src/harness/compaction.ts deleted file mode 100644 index ea484a48..00000000 --- a/packages/agent/src/harness/compaction.ts +++ /dev/null @@ -1,733 +0,0 @@ -import { type AssistantMessage, completeSimple, type Model, type Usage } from "@mariozechner/pi-ai"; -import type { AgentMessage, ThinkingLevel } from "../types.js"; -import { - convertToLlm, - createBranchSummaryMessage, - createCompactionSummaryMessage, - createCustomMessage, -} from "./messages.js"; -import { buildSessionContext } from "./session-tree.js"; -import type { - BranchSummaryResult, - CompactionEntry, - CompactionPreparation, - CompactionSettings, - CompactResult, - FileOperations, - GenerateBranchSummaryOptions, - SessionTreeEntry, -} from "./types.js"; - -export const DEFAULT_COMPACTION_SETTINGS: CompactionSettings = { - enabled: true, - reserveTokens: 16384, - keepRecentTokens: 20000, -}; - -export interface CompactionDetails { - readFiles: string[]; - modifiedFiles: string[]; -} - -export interface BranchSummaryDetails { - readFiles: string[]; - modifiedFiles: string[]; -} - -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. - -Use this EXACT format: - -## Goal -[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.] - -## Constraints & Preferences -- [Any constraints, preferences, or requirements mentioned by user] -- [Or "(none)" if none were mentioned] - -## Progress -### Done -- [x] [Completed tasks/changes] - -### In Progress -- [ ] [Current work] - -### Blocked -- [Issues preventing progress, if any] - -## Key Decisions -- **[Decision]**: [Brief rationale] - -## Next Steps -1. [Ordered list of what should happen next] - -## Critical Context -- [Any data, examples, or references needed to continue] -- [Or "(none)" if not applicable] - -Keep each section concise. Preserve exact file paths, function names, and error messages.`; - -const UPDATE_SUMMARIZATION_PROMPT = `The messages above are NEW conversation messages to incorporate into the existing summary provided in tags. - -Update the existing structured summary with new information. RULES: -- PRESERVE all existing information from the previous summary -- ADD new progress, decisions, and context from the new messages -- UPDATE the Progress section: move items from "In Progress" to "Done" when completed -- UPDATE "Next Steps" based on what was accomplished -- PRESERVE exact file paths, function names, and error messages -- If something is no longer relevant, you may remove it - -Use this EXACT format: - -## Goal -[Preserve existing goals, add new ones if the task expanded] - -## Constraints & Preferences -- [Preserve existing, add new ones discovered] - -## Progress -### Done -- [x] [Include previously done items AND newly completed items] - -### In Progress -- [ ] [Current work - update based on progress] - -### Blocked -- [Current blockers - remove if resolved] - -## Key Decisions -- **[Decision]**: [Brief rationale] (preserve all previous, add new) - -## Next Steps -1. [Update based on current state] - -## Critical Context -- [Preserve important context, add new if needed] - -Keep each section concise. Preserve exact file paths, function names, and error messages.`; - -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: - -## Original Request -[What did the user ask for in this turn?] - -## Early Progress -- [Key decisions and work done in the prefix] - -## Context for Suffix -- [Information needed to understand the retained recent work] - -Be concise. Focus on what's needed to understand the kept suffix.`; - -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 Steps -1. [What should happen next to continue this work] - -Keep each section concise. Preserve exact file paths, function names, and error messages.`; - -export function createFileOps(): FileOperations { - return { read: new Set(), written: new Set(), edited: new Set() }; -} - -export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOperations): void { - if (message.role !== "assistant") return; - if (!("content" in message) || !Array.isArray(message.content)) return; - for (const block of message.content) { - if (typeof block !== "object" || block === null || !("type" in block) || block.type !== "toolCall") continue; - if (!("arguments" in block) || !("name" in block)) continue; - const args = block.arguments as Record | undefined; - if (!args) continue; - const path = typeof args.path === "string" ? args.path : undefined; - if (!path) continue; - switch (block.name) { - case "read": - fileOps.read.add(path); - break; - case "write": - fileOps.written.add(path); - break; - case "edit": - fileOps.edited.add(path); - break; - } - } -} - -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(); - const modifiedFiles = [...modified].sort(); - return { readFiles: readOnly, modifiedFiles }; -} - -export function formatFileOperations(readFiles: string[], modifiedFiles: string[]): string { - const sections: string[] = []; - if (readFiles.length > 0) sections.push(`\n${readFiles.join("\n")}\n`); - if (modifiedFiles.length > 0) sections.push(`\n${modifiedFiles.join("\n")}\n`); - if (sections.length === 0) return ""; - return `\n\n${sections.join("\n\n")}`; -} - -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]`; -} - -export function serializeConversation(messages: ReturnType): string { - const parts: string[] = []; - for (const msg of messages) { - if (msg.role === "user") { - const content = - typeof msg.content === "string" - ? msg.content - : msg.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join(""); - if (content) parts.push(`[User]: ${content}`); - } else if (msg.role === "assistant") { - const textParts: string[] = []; - const thinkingParts: string[] = []; - const toolCalls: string[] = []; - for (const block of msg.content) { - if (block.type === "text") textParts.push(block.text); - else if (block.type === "thinking") thinkingParts.push(block.thinking); - else if (block.type === "toolCall") { - const argsStr = Object.entries(block.arguments as Record) - .map(([k, v]) => `${k}=${JSON.stringify(v)}`) - .join(", "); - toolCalls.push(`${block.name}(${argsStr})`); - } - } - if (thinkingParts.length > 0) parts.push(`[Assistant thinking]: ${thinkingParts.join("\n")}`); - if (textParts.length > 0) parts.push(`[Assistant]: ${textParts.join("\n")}`); - if (toolCalls.length > 0) parts.push(`[Assistant tool calls]: ${toolCalls.join("; ")}`); - } else if (msg.role === "toolResult") { - const content = msg.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join(""); - if (content) parts.push(`[Tool result]: ${truncateForSummary(content, 2000)}`); - } - } - return parts.join("\n\n"); -} - -export function calculateContextTokens(usage: Usage): number { - return usage.totalTokens || usage.input + usage.output + usage.cacheRead + usage.cacheWrite; -} - -function getAssistantUsage(msg: AgentMessage): Usage | undefined { - if (msg.role === "assistant" && "usage" in msg) { - const assistantMsg = msg as AssistantMessage; - if (assistantMsg.stopReason !== "aborted" && assistantMsg.stopReason !== "error" && assistantMsg.usage) { - return assistantMsg.usage; - } - } - return undefined; -} - -export function estimateTokens(message: AgentMessage): number { - let chars = 0; - switch (message.role) { - case "user": { - const content = (message as { content: string | Array<{ type: string; text?: string }> }).content; - if (typeof content === "string") chars = content.length; - else if (Array.isArray(content)) { - for (const block of content) if (block.type === "text" && block.text) chars += block.text.length; - } - return Math.ceil(chars / 4); - } - case "assistant": { - const assistant = message as AssistantMessage; - for (const block of assistant.content) { - if (block.type === "text") chars += block.text.length; - else if (block.type === "thinking") chars += block.thinking.length; - else if (block.type === "toolCall") chars += block.name.length + JSON.stringify(block.arguments).length; - } - return Math.ceil(chars / 4); - } - case "custom": - case "toolResult": { - if (typeof message.content === "string") chars = message.content.length; - else { - for (const block of message.content) { - if (block.type === "text" && block.text) chars += block.text.length; - if (block.type === "image") chars += 4800; - } - } - return Math.ceil(chars / 4); - } - case "bashExecution": - chars = message.command.length + message.output.length; - return Math.ceil(chars / 4); - case "branchSummary": - case "compactionSummary": - chars = message.summary.length; - return Math.ceil(chars / 4); - } - return 0; -} - -function getLastAssistantUsageInfo(messages: AgentMessage[]): { usage: Usage; index: number } | undefined { - for (let i = messages.length - 1; i >= 0; i--) { - const usage = getAssistantUsage(messages[i]!); - if (usage) return { usage, index: i }; - } - return undefined; -} - -export function estimateContextTokens(messages: AgentMessage[]): { - tokens: number; - usageTokens: number; - trailingTokens: number; - lastUsageIndex: number | null; -} { - const usageInfo = getLastAssistantUsageInfo(messages); - if (!usageInfo) { - let estimated = 0; - for (const message of messages) estimated += estimateTokens(message); - return { tokens: estimated, usageTokens: 0, trailingTokens: estimated, lastUsageIndex: null }; - } - const usageTokens = calculateContextTokens(usageInfo.usage); - let trailingTokens = 0; - for (let i = usageInfo.index + 1; i < messages.length; i++) trailingTokens += estimateTokens(messages[i]!); - return { tokens: usageTokens + trailingTokens, usageTokens, trailingTokens, lastUsageIndex: usageInfo.index }; -} - -export function shouldCompact(contextTokens: number, contextWindow: number, settings: CompactionSettings): boolean { - if (!settings.enabled) return false; - return contextTokens > contextWindow - settings.reserveTokens; -} - -function findValidCutPoints(entries: SessionTreeEntry[], startIndex: number, endIndex: number): number[] { - const cutPoints: number[] = []; - for (let i = startIndex; i < endIndex; i++) { - const entry = entries[i]!; - switch (entry.type) { - case "message": { - switch (entry.message.role) { - case "bashExecution": - case "custom": - case "branchSummary": - case "compactionSummary": - case "user": - case "assistant": - cutPoints.push(i); - break; - case "toolResult": - break; - } - break; - } - case "branch_summary": - case "custom_message": - cutPoints.push(i); - break; - default: - break; - } - } - return cutPoints; -} - -export function findTurnStartIndex(entries: SessionTreeEntry[], entryIndex: number, startIndex: number): number { - for (let i = entryIndex; i >= startIndex; i--) { - const entry = entries[i]!; - if (entry.type === "branch_summary" || entry.type === "custom_message") return i; - if (entry.type === "message") { - const role = entry.message.role; - if (role === "user" || role === "bashExecution") return i; - } - } - return -1; -} - -export function findCutPoint( - entries: SessionTreeEntry[], - startIndex: number, - endIndex: number, - keepRecentTokens: number, -) { - const cutPoints = findValidCutPoints(entries, startIndex, endIndex); - if (cutPoints.length === 0) return { firstKeptEntryIndex: startIndex, turnStartIndex: -1, isSplitTurn: false }; - let accumulatedTokens = 0; - let cutIndex = cutPoints[0]!; - for (let i = endIndex - 1; i >= startIndex; i--) { - const entry = entries[i]!; - if (entry.type !== "message") continue; - accumulatedTokens += estimateTokens(entry.message); - if (accumulatedTokens >= keepRecentTokens) { - for (const cp of cutPoints) { - if (cp >= i) { - cutIndex = cp; - break; - } - } - break; - } - } - while (cutIndex > startIndex) { - const prevEntry = entries[cutIndex - 1]!; - if (prevEntry.type === "compaction" || prevEntry.type === "message") break; - cutIndex--; - } - const cutEntry = entries[cutIndex]!; - const isUserMessage = cutEntry.type === "message" && cutEntry.message.role === "user"; - const turnStartIndex = isUserMessage ? -1 : findTurnStartIndex(entries, cutIndex, startIndex); - return { firstKeptEntryIndex: cutIndex, turnStartIndex, isSplitTurn: !isUserMessage && turnStartIndex !== -1 }; -} - -function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined { - switch (entry.type) { - case "message": - return entry.message.role === "toolResult" ? undefined : 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); - default: - return undefined; - } -} - -function getMessageFromEntryForCompaction(entry: SessionTreeEntry): AgentMessage | undefined { - if (entry.type === "compaction") return undefined; - return getMessageFromEntry(entry); -} - -export async function generateSummary( - currentMessages: AgentMessage[], - model: Model, - reserveTokens: number, - apiKey: string, - headers?: Record, - signal?: AbortSignal, - customInstructions?: string, - previousSummary?: string, - thinkingLevel?: ThinkingLevel, -): Promise { - const maxTokens = Math.floor(0.8 * reserveTokens); - let basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT; - if (customInstructions) basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`; - const llmMessages = convertToLlm(currentMessages); - const conversationText = serializeConversation(llmMessages); - let promptText = `\n${conversationText}\n\n\n`; - if (previousSummary) promptText += `\n${previousSummary}\n\n\n`; - promptText += basePrompt; - const summarizationMessages = [ - { role: "user" as const, content: [{ type: "text" as const, text: promptText }], timestamp: Date.now() }, - ]; - const completionOptions = - model.reasoning && thinkingLevel && thinkingLevel !== "off" - ? { maxTokens, signal, apiKey, headers, reasoning: thinkingLevel } - : { maxTokens, signal, apiKey, headers }; - const response = await completeSimple( - model, - { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, - completionOptions, - ); - if (response.stopReason === "error") - throw new Error(`Summarization failed: ${response.errorMessage || "Unknown error"}`); - return response.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join("\n"); -} - -async function generateTurnPrefixSummary( - messages: AgentMessage[], - model: Model, - reserveTokens: number, - apiKey: string, - headers?: Record, - signal?: AbortSignal, - thinkingLevel?: ThinkingLevel, -): Promise { - const maxTokens = Math.floor(0.5 * reserveTokens); - const llmMessages = convertToLlm(messages); - const conversationText = serializeConversation(llmMessages); - const promptText = `\n${conversationText}\n\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`; - const summarizationMessages = [ - { role: "user" as const, content: [{ type: "text" as const, text: promptText }], timestamp: Date.now() }, - ]; - const completionOptions = - model.reasoning && thinkingLevel && thinkingLevel !== "off" - ? { maxTokens, signal, apiKey, headers, reasoning: thinkingLevel } - : { maxTokens, signal, apiKey, headers }; - const response = await completeSimple( - model, - { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, - completionOptions, - ); - if (response.stopReason === "error") - throw new Error(`Turn prefix summarization failed: ${response.errorMessage || "Unknown error"}`); - return response.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join("\n"); -} - -export function prepareCompaction( - pathEntries: SessionTreeEntry[], - settings: CompactionSettings, -): CompactionPreparation | undefined { - if (pathEntries.length > 0 && pathEntries[pathEntries.length - 1]!.type === "compaction") return undefined; - let prevCompactionIndex = -1; - for (let i = pathEntries.length - 1; i >= 0; i--) { - if (pathEntries[i]!.type === "compaction") { - prevCompactionIndex = i; - break; - } - } - let previousSummary: string | undefined; - let boundaryStart = 0; - if (prevCompactionIndex >= 0) { - const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry; - previousSummary = prevCompaction.summary; - const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId); - boundaryStart = firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1; - } - const boundaryEnd = pathEntries.length; - const tokensBefore = estimateContextTokens(buildSessionContext(pathEntries).messages).tokens; - const cutPoint = findCutPoint(pathEntries, boundaryStart, boundaryEnd, settings.keepRecentTokens); - const firstKeptEntry = pathEntries[cutPoint.firstKeptEntryIndex]; - if (!firstKeptEntry?.id) return undefined; - const firstKeptEntryId = firstKeptEntry.id; - const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex; - const messagesToSummarize: AgentMessage[] = []; - for (let i = boundaryStart; i < historyEnd; i++) { - const msg = getMessageFromEntryForCompaction(pathEntries[i]!); - if (msg) messagesToSummarize.push(msg); - } - const turnPrefixMessages: AgentMessage[] = []; - if (cutPoint.isSplitTurn) { - for (let i = cutPoint.turnStartIndex; i < cutPoint.firstKeptEntryIndex; i++) { - const msg = getMessageFromEntryForCompaction(pathEntries[i]!); - if (msg) turnPrefixMessages.push(msg); - } - } - const fileOps = createFileOps(); - if (prevCompactionIndex >= 0) { - const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry; - if (!prevCompaction.fromHook && prevCompaction.details) { - const details = prevCompaction.details as CompactionDetails; - 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 (const msg of messagesToSummarize) extractFileOpsFromMessage(msg, fileOps); - if (cutPoint.isSplitTurn) for (const msg of turnPrefixMessages) extractFileOpsFromMessage(msg, fileOps); - return { - firstKeptEntryId, - messagesToSummarize, - turnPrefixMessages, - isSplitTurn: cutPoint.isSplitTurn, - tokensBefore, - previousSummary, - fileOps, - settings, - }; -} - -export async function compact( - preparation: CompactionPreparation, - model: Model, - apiKey: string, - headers?: Record, - customInstructions?: string, - signal?: AbortSignal, - thinkingLevel?: ThinkingLevel, -): Promise { - const { - firstKeptEntryId, - messagesToSummarize, - turnPrefixMessages, - isSplitTurn, - tokensBefore, - previousSummary, - fileOps, - settings, - } = preparation; - let summary: string; - if (isSplitTurn && turnPrefixMessages.length > 0) { - const [historyResult, turnPrefixResult] = await Promise.all([ - messagesToSummarize.length > 0 - ? generateSummary( - messagesToSummarize, - model, - settings.reserveTokens, - apiKey, - headers, - signal, - customInstructions, - previousSummary, - thinkingLevel, - ) - : Promise.resolve("No prior history."), - generateTurnPrefixSummary( - turnPrefixMessages, - model, - settings.reserveTokens, - apiKey, - headers, - signal, - thinkingLevel, - ), - ]); - summary = `${historyResult}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult}`; - } else { - summary = await generateSummary( - messagesToSummarize, - model, - settings.reserveTokens, - apiKey, - headers, - signal, - customInstructions, - previousSummary, - thinkingLevel, - ); - } - const { readFiles, modifiedFiles } = computeFileLists(fileOps); - summary += formatFileOperations(readFiles, modifiedFiles); - return { summary, firstKeptEntryId, tokensBefore, details: { readFiles, modifiedFiles } }; -} - -export function collectEntriesForBranchSummary( - session: { - getBranch(fromId?: string): Promise; - getEntry(id: string): Promise; - }, - oldLeafId: string | null, - targetId: string, -): Promise<{ entries: SessionTreeEntry[]; commonAncestorId: string | null }> { - return (async () => { - if (!oldLeafId) return { entries: [], commonAncestorId: null }; - const oldPath = new Set((await session.getBranch(oldLeafId)).map((e) => e.id)); - const targetPath = await 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: SessionTreeEntry[] = []; - let current: string | null = oldLeafId; - while (current && current !== commonAncestorId) { - const entry = await session.getEntry(current); - if (!entry) break; - entries.push(entry); - current = entry.parentId; - } - entries.reverse(); - return { entries, commonAncestorId }; - })(); -} - -export function prepareBranchEntries( - entries: SessionTreeEntry[], - tokenBudget: number = 0, -): { messages: AgentMessage[]; fileOps: FileOperations; totalTokens: number } { - const messages: AgentMessage[] = []; - const fileOps = createFileOps(); - let totalTokens = 0; - 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 message = getMessageFromEntry(entries[i]!); - if (!message) continue; - extractFileOpsFromMessage(message, fileOps); - const tokens = estimateTokens(message); - if (tokenBudget > 0 && totalTokens + tokens > tokenBudget) { - const entry = entries[i]!; - if ((entry.type === "compaction" || entry.type === "branch_summary") && totalTokens < tokenBudget * 0.9) { - messages.unshift(message); - totalTokens += tokens; - } - break; - } - messages.unshift(message); - totalTokens += tokens; - } - return { messages, fileOps, totalTokens }; -} - -export async function generateBranchSummary( - entries: SessionTreeEntry[], - options: GenerateBranchSummaryOptions, -): Promise { - const { model, apiKey, headers, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options; - 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" }; - 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 = `\n${conversationText}\n\n\n${instructions}`; - const summarizationMessages = [ - { role: "user" as const, content: [{ type: "text" as const, text: promptText }], timestamp: Date.now() }, - ]; - const response = await completeSimple( - model, - { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, - { apiKey, headers, signal, maxTokens: 2048 }, - ); - 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 }; -} diff --git a/packages/agent/src/harness/compaction/branch-summarization.ts b/packages/agent/src/harness/compaction/branch-summarization.ts new file mode 100644 index 00000000..9bef8799 --- /dev/null +++ b/packages/agent/src/harness/compaction/branch-summarization.ts @@ -0,0 +1,355 @@ +/** + * 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 "@mariozechner/pi-ai"; +import { completeSimple } from "@mariozechner/pi-ai"; +import type { AgentMessage } from "../../types.js"; +import { + convertToLlm, + createBranchSummaryMessage, + createCompactionSummaryMessage, + createCustomMessage, +} from "../messages.js"; +import type { SessionTree, SessionTreeEntry } from "../types.js"; +import { estimateTokens } 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 */ +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: SessionTreeEntry[]; + /** Common ancestor between old and new position, if any */ + commonAncestorId: string | null; +} + +export interface GenerateBranchSummaryOptions { + /** Model to use for summarization */ + model: Model; + /** API key for the model */ + apiKey: string; + /** Request headers for the model */ + headers?: Record; + /** 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; + /** Tokens reserved for prompt + LLM response (default 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 + */ +export async function collectEntriesForBranchSummary( + session: SessionTree, + oldLeafId: string | null, + targetId: string, +): Promise { + // 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)) { + commonAncestorId = targetPath[i].id; + 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; + entries.push(entry); + 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; + + 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); + + // These don't contribute to conversation content + 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: 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; + if (Array.isArray(details.readFiles)) { + 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; + } + + messages.unshift(message); + totalTokens += tokens; + } + + return { messages, fileOps, totalTokens }; +} + +// ============================================================================ +// Summary Generation +// ============================================================================ + +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 Steps +1. [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: SessionTreeEntry[], + options: GenerateBranchSummaryOptions, +): Promise { + 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" }; + } + + // 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; + } else if (customInstructions) { + instructions = `${BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: ${customInstructions}`; + } else { + instructions = BRANCH_SUMMARY_PROMPT; + } + const promptText = `\n${conversationText}\n\n\n${instructions}`; + + const summarizationMessages = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: promptText }], + 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 }; + } + 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"); + + // 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 { + summary: summary || "No summary generated", + readFiles, + modifiedFiles, + }; +} diff --git a/packages/agent/src/harness/compaction/compaction.ts b/packages/agent/src/harness/compaction/compaction.ts new file mode 100644 index 00000000..2c4f58a2 --- /dev/null +++ b/packages/agent/src/harness/compaction/compaction.ts @@ -0,0 +1,842 @@ +/** + * 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, Model, Usage } from "@mariozechner/pi-ai"; +import { completeSimple } from "@mariozechner/pi-ai"; +import type { AgentMessage, ThinkingLevel } from "../../types.js"; +import { + convertToLlm, + createBranchSummaryMessage, + createCompactionSummaryMessage, + createCustomMessage, +} from "../messages.js"; +import { buildSessionContext } from "../session/session-tree.js"; +import type { CompactionEntry, SessionTreeEntry } from "../types.js"; +import { + computeFileLists, + createFileOps, + extractFileOpsFromMessage, + type FileOperations, + formatFileOperations, + SUMMARIZATION_SYSTEM_PROMPT, + serializeConversation, +} from "./utils.js"; + +// ============================================================================ +// File Operation Tracking +// ============================================================================ + +/** Details stored in CompactionEntry.details for file tracking */ +export interface CompactionDetails { + readFiles: string[]; + modifiedFiles: string[]; +} + +/** + * 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); + } + if (Array.isArray(details.modifiedFiles)) { + for (const f of details.modifiedFiles) fileOps.edited.add(f); + } + } + } + + // 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; + } + if (entry.type === "custom_message") { + return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); + } + if (entry.type === "branch_summary") { + return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); + } + if (entry.type === "compaction") { + return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp); + } + return undefined; +} + +function getMessageFromEntryForCompaction(entry: SessionTreeEntry): AgentMessage | undefined { + if (entry.type === "compaction") { + return undefined; + } + return getMessageFromEntry(entry); +} + +/** Result from compact() - SessionManager adds uuid/parentUuid when saving */ +export interface CompactionResult { + summary: string; + firstKeptEntryId: string; + tokensBefore: number; + /** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */ + details?: T; +} + +// ============================================================================ +// Types +// ============================================================================ + +export interface CompactionSettings { + enabled: boolean; + reserveTokens: number; + keepRecentTokens: number; +} + +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. + */ +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; + if (assistantMsg.stopReason !== "aborted" && assistantMsg.stopReason !== "error" && assistantMsg.usage) { + return assistantMsg.usage; + } + } + return undefined; +} + +/** + * Find the last non-aborted assistant message usage from session entries. + */ +export function getLastAssistantUsage(entries: SessionTreeEntry[]): Usage | undefined { + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + if (entry.type === "message") { + const usage = getAssistantUsage(entry.message); + if (usage) return usage; + } + } + return undefined; +} + +export interface ContextUsageEstimate { + tokens: number; + usageTokens: number; + trailingTokens: number; + lastUsageIndex: number | null; +} + +function getLastAssistantUsageInfo(messages: AgentMessage[]): { usage: Usage; index: number } | undefined { + for (let i = messages.length - 1; i >= 0; i--) { + const usage = getAssistantUsage(messages[i]); + if (usage) return { usage, index: i }; + } + 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. + */ +export function estimateContextTokens(messages: AgentMessage[]): ContextUsageEstimate { + const usageInfo = getLastAssistantUsageInfo(messages); + + if (!usageInfo) { + let estimated = 0; + for (const message of messages) { + estimated += estimateTokens(message); + } + return { + tokens: estimated, + usageTokens: 0, + trailingTokens: estimated, + lastUsageIndex: null, + }; + } + + const usageTokens = calculateContextTokens(usageInfo.usage); + let trailingTokens = 0; + for (let i = usageInfo.index + 1; i < messages.length; i++) { + trailingTokens += estimateTokens(messages[i]); + } + + return { + tokens: usageTokens + trailingTokens, + usageTokens, + trailingTokens, + lastUsageIndex: usageInfo.index, + }; +} + +/** + * Check if compaction should trigger based on context usage. + */ +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). + */ +export function estimateTokens(message: AgentMessage): number { + let chars = 0; + + switch (message.role) { + case "user": { + const content = (message as { content: string | Array<{ type: string; text?: string }> }).content; + if (typeof content === "string") { + chars = content.length; + } else if (Array.isArray(content)) { + for (const block of content) { + if (block.type === "text" && block.text) { + chars += block.text.length; + } + } + } + return Math.ceil(chars / 4); + } + case "assistant": { + const assistant = message as AssistantMessage; + for (const block of assistant.content) { + if (block.type === "text") { + chars += block.text.length; + } else if (block.type === "thinking") { + chars += block.thinking.length; + } else if (block.type === "toolCall") { + chars += block.name.length + JSON.stringify(block.arguments).length; + } + } + return Math.ceil(chars / 4); + } + case "custom": + case "toolResult": { + if (typeof message.content === "string") { + chars = message.content.length; + } else { + for (const block of message.content) { + if (block.type === "text" && block.text) { + chars += block.text.length; + } + if (block.type === "image") { + chars += 4800; // Estimate images as 4000 chars, or 1200 tokens + } + } + } + return Math.ceil(chars / 4); + } + case "bashExecution": { + chars = message.command.length + message.output.length; + return Math.ceil(chars / 4); + } + case "branchSummary": + case "compactionSummary": { + chars = message.summary.length; + return Math.ceil(chars / 4); + } + } + + 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++) { + const entry = entries[i]; + switch (entry.type) { + case "message": { + const role = entry.message.role; + switch (role) { + case "bashExecution": + case "custom": + case "branchSummary": + case "compactionSummary": + case "user": + case "assistant": + cutPoints.push(i); + break; + case "toolResult": + break; + } + break; + } + case "thinking_level_change": + case "model_change": + case "compaction": + case "branch_summary": + case "custom": + case "custom_message": + case "label": + case "session_info": + 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); + } + } + 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. + */ +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; + } + if (entry.type === "message") { + const role = entry.message.role; + if (role === "user" || role === "bashExecution") { + return i; + } + } + } + return -1; +} + +export interface CutPointResult { + /** Index of first entry to keep */ + firstKeptEntryIndex: number; + /** Index of user message that starts the turn being split, or -1 if not splitting */ + turnStartIndex: number; + /** Whether this cut splits a turn (cut point is not a user message) */ + 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). + */ +export function findCutPoint( + entries: SessionTreeEntry[], + startIndex: number, + endIndex: number, + keepRecentTokens: number, +): CutPointResult { + const cutPoints = findValidCutPoints(entries, startIndex, endIndex); + + 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) + + 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); + 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]; + break; + } + } + 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); + + return { + firstKeptEntryIndex: cutIndex, + turnStartIndex, + isSplitTurn: !isUserMessage && turnStartIndex !== -1, + }; +} + +// ============================================================================ +// Summarization +// ============================================================================ + +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. + +Use this EXACT format: + +## Goal +[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.] + +## Constraints & Preferences +- [Any constraints, preferences, or requirements mentioned by user] +- [Or "(none)" if none were mentioned] + +## Progress +### Done +- [x] [Completed tasks/changes] + +### In Progress +- [ ] [Current work] + +### Blocked +- [Issues preventing progress, if any] + +## Key Decisions +- **[Decision]**: [Brief rationale] + +## Next Steps +1. [Ordered list of what should happen next] + +## Critical Context +- [Any data, examples, or references needed to continue] +- [Or "(none)" if not applicable] + +Keep each section concise. Preserve exact file paths, function names, and error messages.`; + +const UPDATE_SUMMARIZATION_PROMPT = `The messages above are NEW conversation messages to incorporate into the existing summary provided in tags. + +Update the existing structured summary with new information. RULES: +- PRESERVE all existing information from the previous summary +- ADD new progress, decisions, and context from the new messages +- UPDATE the Progress section: move items from "In Progress" to "Done" when completed +- UPDATE "Next Steps" based on what was accomplished +- PRESERVE exact file paths, function names, and error messages +- If something is no longer relevant, you may remove it + +Use this EXACT format: + +## Goal +[Preserve existing goals, add new ones if the task expanded] + +## Constraints & Preferences +- [Preserve existing, add new ones discovered] + +## Progress +### Done +- [x] [Include previously done items AND newly completed items] + +### In Progress +- [ ] [Current work - update based on progress] + +### Blocked +- [Current blockers - remove if resolved] + +## Key Decisions +- **[Decision]**: [Brief rationale] (preserve all previous, add new) + +## Next Steps +1. [Update based on current state] + +## Critical Context +- [Preserve important context, add new if needed] + +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. + */ +export async function generateSummary( + currentMessages: AgentMessage[], + model: Model, + reserveTokens: number, + apiKey: string, + headers?: Record, + signal?: AbortSignal, + customInstructions?: string, + previousSummary?: string, + thinkingLevel?: ThinkingLevel, +): Promise { + const maxTokens = Math.floor(0.8 * reserveTokens); + + // 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 = `\n${conversationText}\n\n\n`; + if (previousSummary) { + promptText += `\n${previousSummary}\n\n\n`; + } + promptText += basePrompt; + + const summarizationMessages = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: promptText }], + timestamp: Date.now(), + }, + ]; + + const completionOptions = + model.reasoning && thinkingLevel && thinkingLevel !== "off" + ? { maxTokens, signal, apiKey, headers, reasoning: thinkingLevel } + : { maxTokens, signal, apiKey, headers }; + + const response = await completeSimple( + model, + { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, + completionOptions, + ); + + if (response.stopReason === "error") { + throw new Error(`Summarization failed: ${response.errorMessage || "Unknown error"}`); + } + + const textContent = response.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join("\n"); + + return textContent; +} + +// ============================================================================ +// Compaction Preparation (for extensions) +// ============================================================================ + +export interface CompactionPreparation { + /** UUID of first entry to keep */ + firstKeptEntryId: string; + /** Messages that will be summarized and discarded */ + messagesToSummarize: 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; + /** 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( + pathEntries: SessionTreeEntry[], + settings: CompactionSettings, +): CompactionPreparation | undefined { + if (pathEntries.length > 0 && pathEntries[pathEntries.length - 1].type === "compaction") { + return undefined; + } + + let prevCompactionIndex = -1; + for (let i = pathEntries.length - 1; i >= 0; i--) { + if (pathEntries[i].type === "compaction") { + prevCompactionIndex = i; + break; + } + } + + let previousSummary: string | undefined; + let boundaryStart = 0; + if (prevCompactionIndex >= 0) { + const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry; + previousSummary = prevCompaction.summary; + const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId); + boundaryStart = firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1; + } + const boundaryEnd = pathEntries.length; + + 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 + } + 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++) { + const msg = getMessageFromEntryForCompaction(pathEntries[i]); + 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 { + firstKeptEntryId, + messagesToSummarize, + turnPrefixMessages, + isSplitTurn: cutPoint.isSplitTurn, + tokensBefore, + 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: + +## Original Request +[What did the user ask for in this turn?] + +## Early Progress +- [Key decisions and work done in the prefix] + +## Context for Suffix +- [Information needed to understand the retained recent work] + +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"; + +export async function compact( + preparation: CompactionPreparation, + model: Model, + apiKey: string, + headers?: Record, + customInstructions?: string, + signal?: AbortSignal, + thinkingLevel?: ThinkingLevel, +): Promise { + 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 (isSplitTurn && turnPrefixMessages.length > 0) { + // Generate both summaries in parallel + const [historyResult, turnPrefixResult] = await Promise.all([ + messagesToSummarize.length > 0 + ? generateSummary( + messagesToSummarize, + model, + settings.reserveTokens, + apiKey, + headers, + signal, + customInstructions, + previousSummary, + thinkingLevel, + ) + : Promise.resolve("No prior history."), + generateTurnPrefixSummary( + turnPrefixMessages, + model, + settings.reserveTokens, + apiKey, + headers, + signal, + thinkingLevel, + ), + ]); + // Merge into single summary + summary = `${historyResult}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult}`; + } else { + // Just generate history summary + summary = await generateSummary( + messagesToSummarize, + model, + settings.reserveTokens, + apiKey, + headers, + signal, + customInstructions, + previousSummary, + thinkingLevel, + ); + } + + // Compute file lists and append to summary + const { readFiles, modifiedFiles } = computeFileLists(fileOps); + summary += formatFileOperations(readFiles, modifiedFiles); + + if (!firstKeptEntryId) { + throw new Error("First kept entry has no UUID - session may need migration"); + } + + return { + summary, + firstKeptEntryId, + tokensBefore, + details: { readFiles, modifiedFiles } as CompactionDetails, + }; +} + +/** + * Generate a summary for a turn prefix (when splitting a turn). + */ +async function generateTurnPrefixSummary( + messages: AgentMessage[], + model: Model, + reserveTokens: number, + apiKey: string, + headers?: Record, + signal?: AbortSignal, + thinkingLevel?: ThinkingLevel, +): Promise { + const maxTokens = Math.floor(0.5 * reserveTokens); // Smaller budget for turn prefix + const llmMessages = convertToLlm(messages); + const conversationText = serializeConversation(llmMessages); + const promptText = `\n${conversationText}\n\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`; + const summarizationMessages = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: promptText }], + timestamp: Date.now(), + }, + ]; + + 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 }, + ); + + if (response.stopReason === "error") { + throw new Error(`Turn prefix summarization failed: ${response.errorMessage || "Unknown error"}`); + } + + return response.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join("\n"); +} diff --git a/packages/agent/src/harness/compaction/utils.ts b/packages/agent/src/harness/compaction/utils.ts new file mode 100644 index 00000000..a3aacd57 --- /dev/null +++ b/packages/agent/src/harness/compaction/utils.ts @@ -0,0 +1,170 @@ +/** + * Shared utilities for compaction and branch summarization. + */ + +import type { Message } from "@mariozechner/pi-ai"; +import type { AgentMessage } from "../../types.js"; + +// ============================================================================ +// File Operation Tracking +// ============================================================================ + +export interface FileOperations { + read: Set; + written: Set; + edited: Set; +} + +export function createFileOps(): FileOperations { + return { + read: new Set(), + written: new Set(), + edited: new Set(), + }; +} + +/** + * Extract file operations from tool calls in an assistant message. + */ +export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOperations): void { + if (message.role !== "assistant") return; + if (!("content" in message) || !Array.isArray(message.content)) return; + + for (const block of message.content) { + if (typeof block !== "object" || block === null) continue; + if (!("type" in block) || block.type !== "toolCall") continue; + if (!("arguments" in block) || !("name" in block)) continue; + + const args = block.arguments as Record | undefined; + if (!args) continue; + + const path = typeof args.path === "string" ? args.path : undefined; + if (!path) continue; + + switch (block.name) { + case "read": + fileOps.read.add(path); + break; + case "write": + fileOps.written.add(path); + break; + case "edit": + fileOps.edited.add(path); + break; + } + } +} + +/** + * Compute final file lists from file operations. + * Returns readFiles (files only read, not modified) and modifiedFiles. + */ +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(); + const modifiedFiles = [...modified].sort(); + return { readFiles: readOnly, modifiedFiles }; +} + +/** + * Format file operations as XML tags for summary. + */ +export function formatFileOperations(readFiles: string[], modifiedFiles: string[]): string { + const sections: string[] = []; + if (readFiles.length > 0) { + sections.push(`\n${readFiles.join("\n")}\n`); + } + if (modifiedFiles.length > 0) { + sections.push(`\n${modifiedFiles.join("\n")}\n`); + } + if (sections.length === 0) return ""; + 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 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. + */ +export function serializeConversation(messages: Message[]): string { + const parts: string[] = []; + + for (const msg of messages) { + if (msg.role === "user") { + const content = + typeof msg.content === "string" + ? msg.content + : msg.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + if (content) parts.push(`[User]: ${content}`); + } else if (msg.role === "assistant") { + const textParts: string[] = []; + const thinkingParts: string[] = []; + const toolCalls: string[] = []; + + for (const block of msg.content) { + if (block.type === "text") { + textParts.push(block.text); + } else if (block.type === "thinking") { + thinkingParts.push(block.thinking); + } else if (block.type === "toolCall") { + const args = block.arguments as Record; + const argsStr = Object.entries(args) + .map(([k, v]) => `${k}=${JSON.stringify(v)}`) + .join(", "); + toolCalls.push(`${block.name}(${argsStr})`); + } + } + + if (thinkingParts.length > 0) { + parts.push(`[Assistant thinking]: ${thinkingParts.join("\n")}`); + } + if (textParts.length > 0) { + parts.push(`[Assistant]: ${textParts.join("\n")}`); + } + if (toolCalls.length > 0) { + parts.push(`[Assistant tool calls]: ${toolCalls.join("; ")}`); + } + } else if (msg.role === "toolResult") { + const content = msg.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + if (content) { + parts.push(`[Tool result]: ${truncateForSummary(content, TOOL_RESULT_MAX_CHARS)}`); + } + } + } + + 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.`; diff --git a/packages/agent/src/harness/session/jsonl-session-storage.ts b/packages/agent/src/harness/session/jsonl-session-storage.ts new file mode 100644 index 00000000..2e2541b2 --- /dev/null +++ b/packages/agent/src/harness/session/jsonl-session-storage.ts @@ -0,0 +1,152 @@ +import { randomUUID } from "node:crypto"; +import { appendFile, mkdir, readFile, writeFile } from "node:fs/promises"; +import { dirname, resolve } from "node:path"; +import type { CodingAgentSessionInfo, SessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js"; + +interface SessionHeader { + type: "session"; + version: 3; + id: string; + timestamp: string; + cwd: string; + parentSession?: string; +} + +function headerToSessionInfo(header: SessionHeader, filePath?: string): CodingAgentSessionInfo { + return { + id: header.id, + createdAt: header.timestamp, + parentSession: header.parentSession, + projectCwd: header.cwd, + filePath, + }; +} + +async function loadJsonlStorage( + filePath: string, +): Promise<{ header?: SessionHeader; entries: SessionTreeEntry[]; leafId: string | null }> { + try { + const content = await readFile(filePath, "utf8"); + const entries: SessionTreeEntry[] = []; + let header: SessionHeader | undefined; + let leafId: string | null = null; + for (const line of content.split("\n")) { + if (!line.trim()) continue; + try { + const record = JSON.parse(line) as SessionHeader | SessionTreeEntry; + if (record.type === "session") { + header = record as SessionHeader; + continue; + } + entries.push(record as SessionTreeEntry); + leafId = (record as SessionTreeEntry).id; + } catch { + // ignore malformed lines + } + } + return { header, entries, leafId }; + } catch { + return { entries: [], leafId: null }; + } +} + +export class JsonlSessionTreeStorage implements SessionTreeStorage { + private filePath: string; + private cwd: string; + private headerInitialized = false; + private cacheLoaded = false; + private sessionInfo?: CodingAgentSessionInfo; + private entries: SessionTreeEntry[] = []; + private byId = new Map(); + private currentLeafId: string | null = null; + private requestedSessionId?: string; + private parentSession?: string; + + constructor(filePath: string, options: { cwd: string; sessionId?: string; parentSession?: string }) { + this.filePath = resolve(filePath); + this.cwd = options.cwd; + this.requestedSessionId = options.sessionId; + this.parentSession = options.parentSession; + } + + private async ensureParentDir(): Promise { + await mkdir(dirname(this.filePath), { recursive: true }); + } + + private async ensureLoaded(): Promise { + if (this.cacheLoaded) { + return; + } + const loaded = await loadJsonlStorage(this.filePath); + this.entries = loaded.entries; + this.byId = new Map(loaded.entries.map((entry) => [entry.id, entry])); + this.currentLeafId = loaded.leafId; + this.headerInitialized = loaded.header !== undefined; + if (loaded.header) { + this.sessionInfo = headerToSessionInfo(loaded.header, this.filePath); + } + this.cacheLoaded = true; + } + + private async ensureHeader(): Promise { + await this.ensureLoaded(); + if (this.headerInitialized) return; + await this.ensureParentDir(); + const header: SessionHeader = { + type: "session", + version: 3, + id: this.requestedSessionId ?? randomUUID(), + timestamp: new Date().toISOString(), + cwd: this.cwd, + parentSession: this.parentSession, + }; + await writeFile(this.filePath, `${JSON.stringify(header)}\n`); + this.sessionInfo = headerToSessionInfo(header, this.filePath); + this.headerInitialized = true; + } + + async getSessionInfo(): Promise { + await this.ensureHeader(); + return this.sessionInfo!; + } + + async getLeafId(): Promise { + await this.ensureLoaded(); + return this.currentLeafId; + } + + async setLeafId(leafId: string | null): Promise { + await this.ensureLoaded(); + this.currentLeafId = leafId; + } + + async appendEntry(entry: SessionTreeEntry): Promise { + await this.ensureHeader(); + await appendFile(this.filePath, `${JSON.stringify(entry)}\n`); + this.entries.push(entry); + this.byId.set(entry.id, entry); + this.currentLeafId = entry.id; + } + + async getEntry(id: string): Promise { + await this.ensureLoaded(); + return this.byId.get(id); + } + + async getPathToRoot(leafId: string | null): Promise { + await this.ensureLoaded(); + if (leafId === null) return []; + const path: SessionTreeEntry[] = []; + let current = this.byId.get(leafId); + while (current) { + path.unshift(current); + current = current.parentId ? this.byId.get(current.parentId) : undefined; + } + return path; + } + + async getEntries(): Promise { + await this.ensureLoaded(); + return [...this.entries]; + } +} diff --git a/packages/agent/src/harness/session/memory-session-storage.ts b/packages/agent/src/harness/session/memory-session-storage.ts new file mode 100644 index 00000000..b1fe743f --- /dev/null +++ b/packages/agent/src/harness/session/memory-session-storage.ts @@ -0,0 +1,51 @@ +import { randomUUID } from "crypto"; +import type { SessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js"; + +export class InMemorySessionTreeStorage implements SessionTreeStorage { + private entries: SessionTreeEntry[]; + private leafId: string | null; + private sessionInfo: SessionInfo; + + constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null; sessionInfo?: SessionInfo }) { + this.entries = options?.entries ? [...options.entries] : []; + this.leafId = options?.leafId ?? this.entries[this.entries.length - 1]?.id ?? null; + this.sessionInfo = options?.sessionInfo ?? { id: randomUUID(), createdAt: new Date().toISOString() }; + } + + async getSessionInfo(): Promise { + return this.sessionInfo; + } + + async getLeafId(): Promise { + return this.leafId; + } + + async setLeafId(leafId: string | null): Promise { + this.leafId = leafId; + } + + async appendEntry(entry: SessionTreeEntry): Promise { + this.entries.push(entry); + this.leafId = entry.id; + } + + async getEntry(id: string): Promise { + return this.entries.find((entry) => entry.id === id); + } + + async getPathToRoot(leafId: string | null): Promise { + if (leafId === null) return []; + const byId = new Map(this.entries.map((entry) => [entry.id, entry])); + const path: SessionTreeEntry[] = []; + let current = byId.get(leafId); + while (current) { + path.unshift(current); + current = current.parentId ? byId.get(current.parentId) : undefined; + } + return path; + } + + async getEntries(): Promise { + return [...this.entries]; + } +} diff --git a/packages/agent/src/harness/session/session-repo.ts b/packages/agent/src/harness/session/session-repo.ts new file mode 100644 index 00000000..587b0f9d --- /dev/null +++ b/packages/agent/src/harness/session/session-repo.ts @@ -0,0 +1,231 @@ +import { existsSync, mkdirSync, readdirSync, readFileSync, rmSync } from "node:fs"; +import { join, resolve } from "node:path"; +import { v7 as uuidv7 } from "uuid"; +import type { + CodingAgentSessionInfo, + CodingAgentSessionRepo, + Session, + SessionInfo, + SessionRepo, + SessionTreeEntry, +} from "../types.js"; +import { JsonlSessionTreeStorage } from "./jsonl-session-storage.js"; +import { InMemorySessionTreeStorage } from "./memory-session-storage.js"; +import { DefaultSessionTree } from "./session-tree.js"; + +function createSessionId(): string { + return uuidv7(); +} + +function createTimestamp(): string { + return new Date().toISOString(); +} + +function toSession(info: TInfo, tree: DefaultSessionTree): Session { + return { info, tree }; +} + +function getPathEntriesToFork( + entries: SessionTreeEntry[], + entryId: string, + position: "before" | "at", +): SessionTreeEntry[] { + const byId = new Map(entries.map((entry) => [entry.id, entry])); + const target = byId.get(entryId); + if (!target) { + throw new Error(`Entry ${entryId} not found`); + } + let effectiveLeafId: string | null; + if (position === "at") { + effectiveLeafId = target.id; + } else { + if (target.type !== "message" || target.message.role !== "user") { + throw new Error(`Entry ${entryId} is not a user message`); + } + effectiveLeafId = target.parentId; + } + if (effectiveLeafId === null) { + return []; + } + const path: SessionTreeEntry[] = []; + let current = byId.get(effectiveLeafId); + while (current) { + path.unshift(current); + current = current.parentId ? byId.get(current.parentId) : undefined; + } + return path; +} + +export class InMemorySessionRepo implements SessionRepo { + private sessions = new Map>(); + + async create(options?: { id?: string; parentSession?: string }): Promise> { + const info: SessionInfo = { + id: options?.id ?? createSessionId(), + createdAt: createTimestamp(), + parentSession: options?.parentSession, + }; + const storage = new InMemorySessionTreeStorage({ sessionInfo: info }); + const session = toSession(info, new DefaultSessionTree(storage)); + this.sessions.set(info.id, session); + return session; + } + + async open(ref: string): Promise> { + const session = this.sessions.get(ref); + if (!session) { + throw new Error(`Session not found: ${ref}`); + } + return session; + } + + async list(): Promise>> { + return [...this.sessions.values()]; + } + + async delete(ref: string): Promise { + this.sessions.delete(ref); + } + + async fork( + ref: string, + options: { entryId: string; position?: "before" | "at"; id?: string }, + ): Promise> { + const source = await this.open(ref); + const entries = await source.tree.getEntries(); + const forkedEntries = getPathEntriesToFork(entries, options.entryId, options.position ?? "before"); + const info: SessionInfo = { + id: options.id ?? createSessionId(), + createdAt: createTimestamp(), + parentSession: source.info.id, + }; + const leafId = forkedEntries[forkedEntries.length - 1]?.id ?? null; + const storage = new InMemorySessionTreeStorage({ sessionInfo: info, entries: forkedEntries, leafId }); + const session = toSession(info, new DefaultSessionTree(storage)); + this.sessions.set(info.id, session); + return session; + } +} + +function readJsonlHeader(filePath: string): CodingAgentSessionInfo | undefined { + try { + const content = readFileSync(filePath, "utf8"); + const firstLine = content.split("\n")[0]; + if (!firstLine) return undefined; + const header = JSON.parse(firstLine) as { + type: string; + id: string; + timestamp: string; + cwd: string; + parentSession?: string; + }; + if (header.type !== "session") return undefined; + return { + id: header.id, + createdAt: header.timestamp, + parentSession: header.parentSession, + projectCwd: header.cwd, + filePath, + }; + } catch { + return undefined; + } +} + +export class JsonlCodingAgentSessionRepo implements CodingAgentSessionRepo { + private sessionDir: string; + private cwd: string; + + constructor(options: { sessionDir: string; cwd: string }) { + this.sessionDir = resolve(options.sessionDir); + this.cwd = options.cwd; + mkdirSync(this.sessionDir, { recursive: true }); + } + + private createSessionFilePath(sessionId: string, timestamp: string): string { + return join(this.sessionDir, `${timestamp.replace(/[:.]/g, "-")}_${sessionId}.jsonl`); + } + + async create(options?: { id?: string; parentSession?: string }): Promise> { + const id = options?.id ?? createSessionId(); + const createdAt = createTimestamp(); + const filePath = this.createSessionFilePath(id, createdAt); + const storage = new JsonlSessionTreeStorage(filePath, { + cwd: this.cwd, + sessionId: id, + parentSession: options?.parentSession, + }); + const info = (await storage.getSessionInfo()) as CodingAgentSessionInfo; + return toSession(info, new DefaultSessionTree(storage)); + } + + async open(ref: string): Promise> { + const filePath = ref.includes("/") || ref.endsWith(".jsonl") ? resolve(ref) : join(this.sessionDir, ref); + if (!existsSync(filePath)) { + throw new Error(`Session not found: ${ref}`); + } + const storage = new JsonlSessionTreeStorage(filePath, { cwd: this.cwd }); + const info = (await storage.getSessionInfo()) as CodingAgentSessionInfo; + return toSession(info, new DefaultSessionTree(storage)); + } + + async list(): Promise>> { + if (!existsSync(this.sessionDir)) { + return []; + } + const files = readdirSync(this.sessionDir) + .filter((file) => file.endsWith(".jsonl")) + .map((file) => join(this.sessionDir, file)); + const sessions: Array> = []; + for (const filePath of files) { + const info = readJsonlHeader(filePath); + if (!info) continue; + sessions.push( + toSession(info, new DefaultSessionTree(new JsonlSessionTreeStorage(filePath, { cwd: info.projectCwd }))), + ); + } + return sessions; + } + + async listByCwd(cwd: string): Promise>> { + return (await this.list()).filter((session) => session.info.projectCwd === cwd); + } + + async getMostRecentByCwd(cwd: string): Promise | undefined> { + const sessions = await this.listByCwd(cwd); + sessions.sort((a, b) => new Date(b.info.createdAt).getTime() - new Date(a.info.createdAt).getTime()); + return sessions[0]; + } + + async delete(ref: string): Promise { + const filePath = ref.includes("/") || ref.endsWith(".jsonl") ? resolve(ref) : join(this.sessionDir, ref); + if (existsSync(filePath)) { + rmSync(filePath, { force: true }); + } + } + + async fork( + ref: string, + options: { entryId: string; position?: "before" | "at"; id?: string }, + ): Promise> { + const source = await this.open(ref); + const entries = await source.tree.getEntries(); + const forkedEntries = getPathEntriesToFork(entries, options.entryId, options.position ?? "before"); + const id = options.id ?? createSessionId(); + const createdAt = createTimestamp(); + const filePath = this.createSessionFilePath(id, createdAt); + const storage = new JsonlSessionTreeStorage(filePath, { + cwd: source.info.projectCwd, + sessionId: id, + parentSession: source.info.filePath ?? source.info.id, + }); + for (const entry of forkedEntries) { + await storage.appendEntry(entry); + } + if (forkedEntries.length === 0) { + await storage.getSessionInfo(); + } + const info = (await storage.getSessionInfo()) as CodingAgentSessionInfo; + return toSession(info, new DefaultSessionTree(storage)); + } +} diff --git a/packages/agent/src/harness/session-tree.ts b/packages/agent/src/harness/session/session-tree.ts similarity index 58% rename from packages/agent/src/harness/session-tree.ts rename to packages/agent/src/harness/session/session-tree.ts index 4b3855fb..9c1322c1 100644 --- a/packages/agent/src/harness/session-tree.ts +++ b/packages/agent/src/harness/session/session-tree.ts @@ -1,9 +1,7 @@ import { randomUUID } from "node:crypto"; -import { appendFile, mkdir, readFile, writeFile } from "node:fs/promises"; -import { dirname, resolve } from "node:path"; import type { ImageContent, TextContent } from "@mariozechner/pi-ai"; -import type { AgentMessage } from "../types.js"; -import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "./messages.js"; +import type { AgentMessage } from "../../types.js"; +import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "../messages.js"; import type { BranchSummaryEntry, CompactionEntry, @@ -13,12 +11,14 @@ import type { MessageEntry, ModelChangeEntry, SessionContext, + SessionInfo, SessionInfoEntry, SessionTree, SessionTreeEntry, SessionTreeStorage, ThinkingLevelChangeEntry, -} from "./types.js"; +} from "../types.js"; +import { InMemorySessionTreeStorage } from "./memory-session-storage.js"; function generateId(byId: { has(id: string): boolean }): string { for (let i = 0; i < 100; i++) { @@ -79,172 +79,6 @@ export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext return { messages, thinkingLevel, model }; } -interface SessionHeader { - type: "session"; - version: 3; - id: string; - timestamp: string; - cwd: string; - parentSession?: string; -} - -async function loadJsonlStorage( - filePath: string, -): Promise<{ header?: SessionHeader; entries: SessionTreeEntry[]; leafId: string | null }> { - try { - const content = await readFile(filePath, "utf8"); - const entries: SessionTreeEntry[] = []; - let header: SessionHeader | undefined; - let leafId: string | null = null; - for (const line of content.split("\n")) { - if (!line.trim()) continue; - try { - const record = JSON.parse(line) as SessionHeader | SessionTreeEntry; - if (record.type === "session") { - header = record as SessionHeader; - continue; - } - entries.push(record as SessionTreeEntry); - leafId = (record as SessionTreeEntry).id; - } catch { - // ignore malformed lines - } - } - return { header, entries, leafId }; - } catch { - return { entries: [], leafId: null }; - } -} - -export class JsonlSessionTreeStorage implements SessionTreeStorage { - private filePath: string; - private cwd: string; - private headerInitialized = false; - private cacheLoaded = false; - private entries: SessionTreeEntry[] = []; - private byId = new Map(); - private currentLeafId: string | null = null; - - constructor(filePath: string, options: { cwd: string }) { - this.filePath = resolve(filePath); - this.cwd = options.cwd; - } - - private async ensureParentDir(): Promise { - await mkdir(dirname(this.filePath), { recursive: true }); - } - - private async ensureLoaded(): Promise { - if (this.cacheLoaded) { - return; - } - const loaded = await loadJsonlStorage(this.filePath); - this.entries = loaded.entries; - this.byId = new Map(loaded.entries.map((entry) => [entry.id, entry])); - this.currentLeafId = loaded.leafId; - this.headerInitialized = loaded.header !== undefined; - this.cacheLoaded = true; - } - - private async ensureHeader(): Promise { - await this.ensureLoaded(); - if (this.headerInitialized) return; - await this.ensureParentDir(); - const header: SessionHeader = { - type: "session", - version: 3, - id: randomUUID(), - timestamp: new Date().toISOString(), - cwd: this.cwd, - }; - await writeFile(this.filePath, `${JSON.stringify(header)}\n`); - this.headerInitialized = true; - } - - async getLeafId(): Promise { - await this.ensureLoaded(); - return this.currentLeafId; - } - - async setLeafId(leafId: string | null): Promise { - await this.ensureLoaded(); - this.currentLeafId = leafId; - } - - async appendEntry(entry: SessionTreeEntry): Promise { - await this.ensureHeader(); - await appendFile(this.filePath, `${JSON.stringify(entry)}\n`); - this.entries.push(entry); - this.byId.set(entry.id, entry); - this.currentLeafId = entry.id; - } - - async getEntry(id: string): Promise { - await this.ensureLoaded(); - return this.byId.get(id); - } - - async getPathToRoot(leafId: string | null): Promise { - await this.ensureLoaded(); - if (leafId === null) return []; - const path: SessionTreeEntry[] = []; - let current = this.byId.get(leafId); - while (current) { - path.unshift(current); - current = current.parentId ? this.byId.get(current.parentId) : undefined; - } - return path; - } - - async getEntries(): Promise { - await this.ensureLoaded(); - return [...this.entries]; - } -} - -export class InMemorySessionTreeStorage implements SessionTreeStorage { - private entries: SessionTreeEntry[]; - private leafId: string | null; - - constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null }) { - this.entries = options?.entries ? [...options.entries] : []; - this.leafId = options?.leafId ?? this.entries[this.entries.length - 1]?.id ?? null; - } - - async getLeafId(): Promise { - return this.leafId; - } - - async setLeafId(leafId: string | null): Promise { - this.leafId = leafId; - } - - async appendEntry(entry: SessionTreeEntry): Promise { - this.entries.push(entry); - this.leafId = entry.id; - } - - async getEntry(id: string): Promise { - return this.entries.find((entry) => entry.id === id); - } - - async getPathToRoot(leafId: string | null): Promise { - if (leafId === null) return []; - const byId = new Map(this.entries.map((entry) => [entry.id, entry])); - const path: SessionTreeEntry[] = []; - let current = byId.get(leafId); - while (current) { - path.unshift(current); - current = current.parentId ? byId.get(current.parentId) : undefined; - } - return path; - } - - async getEntries(): Promise { - return [...this.entries]; - } -} - export class DefaultSessionTree implements SessionTree { private storage: SessionTreeStorage; @@ -273,6 +107,32 @@ export class DefaultSessionTree implements SessionTree { return buildSessionContext(await this.getBranch()); } + getSessionInfo(): Promise { + return this.storage.getSessionInfo(); + } + + async getLabel(id: string): Promise { + const entries = await this.storage.getEntries(); + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]!; + if (entry.type === "label" && entry.targetId === id) { + return entry.label?.trim() || undefined; + } + } + return undefined; + } + + async getSessionName(): Promise { + const entries = await this.storage.getEntries(); + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]!; + if (entry.type === "session_info") { + return entry.name?.trim() || undefined; + } + } + return undefined; + } + private async makeEntryId(): Promise { const entries = await this.storage.getEntries(); return generateId(new Set(entries.map((entry) => entry.id))); diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index 41f701b3..fa75a953 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -157,7 +157,19 @@ export interface SessionContext { model: { provider: string; modelId: string } | null; } +export interface SessionInfo { + id: string; + createdAt: string; + parentSession?: string; +} + +export interface CodingAgentSessionInfo extends SessionInfo { + projectCwd: string; + filePath?: string; +} + export interface SessionTreeStorage { + getSessionInfo(): Promise; getLeafId(): Promise; setLeafId(leafId: string | null): Promise; appendEntry(entry: SessionTreeEntry): Promise; @@ -167,11 +179,14 @@ export interface SessionTreeStorage { } export interface SessionTree { + getSessionInfo(): Promise; getLeafId(): Promise; getEntry(id: string): Promise; getEntries(): Promise; getBranch(fromId?: string): Promise; buildContext(): Promise; + getLabel(id: string): Promise; + getSessionName(): Promise; appendMessage(message: AgentMessage): Promise; appendThinkingLevelChange(thinkingLevel: string): Promise; @@ -197,6 +212,24 @@ export interface SessionTree { moveTo(entryId: string | null): Promise; } +export interface Session { + info: TInfo; + tree: SessionTree; +} + +export interface SessionRepo { + create(options?: { id?: string; parentSession?: string }): Promise>; + open(ref: TRef): Promise>; + list(): Promise>>; + delete(ref: TRef): Promise; + fork(ref: TRef, options: { entryId: string; position?: "before" | "at"; id?: string }): Promise>; +} + +export interface CodingAgentSessionRepo extends SessionRepo { + listByCwd(cwd: string): Promise>>; + getMostRecentByCwd(cwd: string): Promise | undefined>; +} + export interface AgentHarnessPendingMutations { appendMessages: AgentMessage[]; model?: Model; diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index ead51788..78222783 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -3,11 +3,30 @@ export * from "./agent.js"; // Loop functions export * from "./agent-loop.js"; export * from "./harness/agent-harness.js"; -export * from "./harness/compaction.js"; +export { + collectEntriesForBranchSummary, + generateBranchSummary, + prepareBranchEntries, +} from "./harness/compaction/branch-summarization.js"; +export { + calculateContextTokens, + compact, + DEFAULT_COMPACTION_SETTINGS, + estimateContextTokens, + estimateTokens, + findCutPoint, + findTurnStartIndex, + generateSummary, + getLastAssistantUsage, + prepareCompaction, + serializeConversation, + shouldCompact, +} from "./harness/compaction/compaction.js"; export * from "./harness/execution-env.js"; export * from "./harness/messages.js"; export * from "./harness/prompt-templates.js"; -export * from "./harness/session-tree.js"; +export * from "./harness/session/session-repo.js"; +export * from "./harness/session/session-tree.js"; // Harness export * from "./harness/types.js"; export * from "./harness/utils/shell-output.js"; diff --git a/packages/agent/test/harness/compaction.test.ts b/packages/agent/test/harness/compaction.test.ts index 4c042df7..0c260ab2 100644 --- a/packages/agent/test/harness/compaction.test.ts +++ b/packages/agent/test/harness/compaction.test.ts @@ -1,4 +1,3 @@ -import type { AgentMessage } from "@mariozechner/pi-agent-core"; import { type AssistantMessage, type FauxProviderRegistration, @@ -18,8 +17,8 @@ import { prepareCompaction, serializeConversation, shouldCompact, -} from "../../src/harness/compaction.js"; -import { buildSessionContext } from "../../src/harness/session-tree.js"; +} from "../../src/harness/compaction/compaction.js"; +import { buildSessionContext } from "../../src/harness/session/session-tree.js"; import type { CompactionEntry, CompactionSettings, @@ -28,6 +27,7 @@ import type { SessionTreeEntry, ThinkingLevelChangeEntry, } from "../../src/harness/types.js"; +import type { AgentMessage } from "../../src/types.js"; let nextId = 0; function createId(): string { diff --git a/packages/agent/test/harness/session-tree.test.ts b/packages/agent/test/harness/session-tree.test.ts index 64a6dba4..eb1171c8 100644 --- a/packages/agent/test/harness/session-tree.test.ts +++ b/packages/agent/test/harness/session-tree.test.ts @@ -3,11 +3,9 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import type { AgentMessage } from "@mariozechner/pi-agent-core"; import { afterEach, describe, expect, it } from "vitest"; -import { - DefaultSessionTree, - InMemorySessionTreeStorage, - JsonlSessionTreeStorage, -} from "../../src/harness/session-tree.js"; +import { JsonlSessionTreeStorage } from "../../src/harness/session/jsonl-session-storage.js"; +import { InMemorySessionTreeStorage } from "../../src/harness/session/memory-session-storage.js"; +import { DefaultSessionTree } from "../../src/harness/session/session-tree.js"; import type { SessionTreeStorage } from "../../src/harness/types.js"; function createUserMessage(text: string): AgentMessage { @@ -128,6 +126,8 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre const entries = await tree.getEntries(); expect(entries.some((entry) => entry.type === "label")).toBe(true); expect(entries.some((entry) => entry.type === "session_info")).toBe(true); + expect(await tree.getLabel(user1)).toBe("checkpoint"); + expect(await tree.getSessionName()).toBe("name"); expect((await tree.buildContext()).messages).toHaveLength(1); }); @@ -136,11 +136,15 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre const tree = new DefaultSessionTree(storage); const user1 = await tree.appendMessage(createUserMessage("one")); await tree.appendMessage(createAssistantMessage("two")); + await tree.appendLabelChange(user1, "checkpoint"); + await tree.appendSessionInfo("name"); await tree.moveTo(user1); await tree.appendMessage(createAssistantMessage("branched")); const tree2 = new DefaultSessionTree(storage); const context = await tree2.buildContext(); expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]); + expect(await tree2.getLabel(user1)).toBe("checkpoint"); + expect(await tree2.getSessionName()).toBe("name"); inspect?.(); }); });