shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(agent): split harness compaction and session modules

Browse Source
This commit is contained in:
Mario Zechner
2026-05-03 03:13:45 +02:00
parent a5b27367d3
commit 83599e789d
13 changed files with 1910 additions and 939 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
packages/agent/src/harness/agent-harness.ts
View File

@@ -3,13 +3,8 @@ import { readFileSync } from "node:fs";
import type { ImageContent, Model } from "@mariozechner/pi-ai"; import type { ImageContent, Model } from "@mariozechner/pi-ai";
import type { Agent } from "../agent.js"; import type { Agent } from "../agent.js";
import type { AgentEvent, AgentMessage, AgentTool, ThinkingLevel } from "../types.js"; import type { AgentEvent, AgentMessage, AgentTool, ThinkingLevel } from "../types.js";
import { import { collectEntriesForBranchSummary, generateBranchSummary } from "./compaction/branch-summarization.js";
collectEntriesForBranchSummary, import { compact, DEFAULT_COMPACTION_SETTINGS, prepareCompaction } from "./compaction/compaction.js";
compact,
DEFAULT_COMPACTION_SETTINGS,
generateBranchSummary,
prepareCompaction,
} from "./compaction.js";
import { expandPromptTemplate } from "./prompt-templates.js"; import { expandPromptTemplate } from "./prompt-templates.js";
import type { import type {
AbortResult, AbortResult,
@@ -24,8 +19,6 @@ import type {
ExecutionEnv, ExecutionEnv,
NavigateTreeResult, NavigateTreeResult,
PromptTemplate, PromptTemplate,
SessionBeforeCompactResult,
SessionBeforeTreeResult,
SessionTree, SessionTree,
Skill, Skill,
SystemPromptInputs, SystemPromptInputs,
@@ -388,8 +381,8 @@ export class DefaultAgentHarness implements AgentHarness {
customInstructions, customInstructions,
signal: new AbortController().signal, signal: new AbortController().signal,
}); });
if ((hookResult as SessionBeforeCompactResult | undefined)?.cancel) throw new Error("Compaction cancelled"); if (hookResult?.cancel) throw new Error("Compaction cancelled");
const provided = (hookResult as SessionBeforeCompactResult | undefined)?.compaction; const provided = hookResult?.compaction;
const result = const result =
provided ?? provided ??
(await compact( (await compact(
@@ -442,11 +435,10 @@ export class DefaultAgentHarness implements AgentHarness {
preparation, preparation,
signal, signal,
}); });
const typedHook = hookResult as SessionBeforeTreeResult | undefined; if (hookResult?.cancel) return { cancelled: true };
if (typedHook?.cancel) return { cancelled: true };
let summaryEntry: any | undefined; let summaryEntry: any | undefined;
let summaryText: string | undefined = typedHook?.summary?.summary; let summaryText: string | undefined = hookResult?.summary?.summary;
let summaryDetails: unknown = typedHook?.summary?.details; let summaryDetails: unknown = hookResult?.summary?.details;
if (!summaryText && options?.summarize && entries.length > 0) { if (!summaryText && options?.summarize && entries.length > 0) {
const model = this.conversation.model; const model = this.conversation.model;
if (!model) throw new Error("No model set for branch summary"); if (!model) throw new Error("No model set for branch summary");
@@ -457,8 +449,8 @@ export class DefaultAgentHarness implements AgentHarness {
apiKey: auth.apiKey, apiKey: auth.apiKey,
headers: auth.headers, headers: auth.headers,
signal: new AbortController().signal, signal: new AbortController().signal,
customInstructions: typedHook?.customInstructions ?? options?.customInstructions, customInstructions: hookResult?.customInstructions ?? options?.customInstructions,
replaceInstructions: typedHook?.replaceInstructions ?? options?.replaceInstructions, replaceInstructions: hookResult?.replaceInstructions ?? options?.replaceInstructions,
}); });
if (branchSummary.aborted) return { cancelled: true }; if (branchSummary.aborted) return { cancelled: true };
if (branchSummary.error) throw new Error(branchSummary.error); if (branchSummary.error) throw new Error(branchSummary.error);
@@ -498,7 +490,7 @@ export class DefaultAgentHarness implements AgentHarness {
newLeafId ?? "root", newLeafId ?? "root",
summaryText, summaryText,
summaryDetails, summaryDetails,
typedHook?.summary !== undefined, hookResult?.summary !== undefined,
); );
summaryEntry = await this.sessionTree.getEntry(summaryId); summaryEntry = await this.sessionTree.getEntry(summaryId);
} }
@@ -508,7 +500,7 @@ export class DefaultAgentHarness implements AgentHarness {
newLeafId: await this.sessionTree.getLeafId(), newLeafId: await this.sessionTree.getLeafId(),
oldLeafId, oldLeafId,
summaryEntry, summaryEntry,
fromHook: typedHook?.summary !== undefined, fromHook: hookResult?.summary !== undefined,
}); });
return { cancelled: false, editorText, summaryEntry }; return { cancelled: false, editorText, summaryEntry };
} }
@@ -573,12 +565,7 @@ export class DefaultAgentHarness implements AgentHarness {
await this.agent.waitForIdle(); await this.agent.waitForIdle();
} }
subscribe( subscribe(listener: (event: AgentHarnessEvent, signal?: AbortSignal) => Promise<void> | void): () => void {
listener: (
event: AgentEvent | import("./types.js").AgentHarnessOwnEvent,
signal?: AbortSignal,
) => Promise<void> | void,
): () => void {
this.listeners.add(listener); this.listeners.add(listener);
return () => this.listeners.delete(listener); return () => this.listeners.delete(listener);
} }

733
packages/agent/src/harness/compaction.ts
View File

@@ -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 <previous-summary> 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<string, unknown> | 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(`<read-files>\n${readFiles.join("\n")}\n</read-files>`);
if (modifiedFiles.length > 0) sections.push(`<modified-files>\n${modifiedFiles.join("\n")}\n</modified-files>`);
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<typeof convertToLlm>): 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<string, unknown>)
.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<any>,
reserveTokens: number,
apiKey: string,
headers?: Record<string, string>,
signal?: AbortSignal,
customInstructions?: string,
previousSummary?: string,
thinkingLevel?: ThinkingLevel,
): Promise<string> {
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 = `<conversation>\n${conversationText}\n</conversation>\n\n`;
if (previousSummary) promptText += `<previous-summary>\n${previousSummary}\n</previous-summary>\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<any>,
reserveTokens: number,
apiKey: string,
headers?: Record<string, string>,
signal?: AbortSignal,
thinkingLevel?: ThinkingLevel,
): Promise<string> {
const maxTokens = Math.floor(0.5 * reserveTokens);
const llmMessages = convertToLlm(messages);
const conversationText = serializeConversation(llmMessages);
const promptText = `<conversation>\n${conversationText}\n</conversation>\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<any>,
apiKey: string,
headers?: Record<string, string>,
customInstructions?: string,
signal?: AbortSignal,
thinkingLevel?: ThinkingLevel,
): Promise<CompactResult> {
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<SessionTreeEntry[]>;
getEntry(id: string): Promise<SessionTreeEntry | undefined>;
},
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<BranchSummaryResult> {
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 = `<conversation>\n${conversationText}\n</conversation>\n\n${instructions}`;
const summarizationMessages = [
{ role: "user" as const, content: [{ type: "text" as const, text: promptText }], timestamp: Date.now() },
];
const response = await 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 };
}

355
packages/agent/src/harness/compaction/branch-summarization.ts Normal file
View File

@@ -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<any>;
/** API key for the model */
apiKey: string;
/** Request headers for the model */
headers?: Record<string, string>;
/** Abort signal for cancellation */
signal: AbortSignal;
/** Optional custom instructions for summarization */
customInstructions?: string;
/** If true, customInstructions replaces the default prompt instead of being appended */
replaceInstructions?: boolean;
/** 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<CollectEntriesResult> {
// If no old position, nothing to summarize
if (!oldLeafId) {
return { entries: [], commonAncestorId: null };
}
// Find common ancestor (deepest node that's on both paths)
const oldPath = new Set((await session.getBranch(oldLeafId)).map((e) => e.id));
const targetPath = await session.getBranch(targetId);
// targetPath is root-first, so iterate backwards to find deepest common ancestor
let commonAncestorId: string | null = null;
for (let i = targetPath.length - 1; i >= 0; i--) {
if (oldPath.has(targetPath[i].id)) {
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<BranchSummaryResult> {
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 = `<conversation>\n${conversationText}\n</conversation>\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,
};
}

842
packages/agent/src/harness/compaction/compaction.ts Normal file
View File

@@ -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<T = unknown> {
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 <previous-summary> 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<any>,
reserveTokens: number,
apiKey: string,
headers?: Record<string, string>,
signal?: AbortSignal,
customInstructions?: string,
previousSummary?: string,
thinkingLevel?: ThinkingLevel,
): Promise<string> {
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 = `<conversation>\n${conversationText}\n</conversation>\n\n`;
if (previousSummary) {
promptText += `<previous-summary>\n${previousSummary}\n</previous-summary>\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<any>,
apiKey: string,
headers?: Record<string, string>,
customInstructions?: string,
signal?: AbortSignal,
thinkingLevel?: ThinkingLevel,
): Promise<CompactionResult> {
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<any>,
reserveTokens: number,
apiKey: string,
headers?: Record<string, string>,
signal?: AbortSignal,
thinkingLevel?: ThinkingLevel,
): Promise<string> {
const maxTokens = Math.floor(0.5 * reserveTokens); // Smaller budget for turn prefix
const llmMessages = convertToLlm(messages);
const conversationText = serializeConversation(llmMessages);
const promptText = `<conversation>\n${conversationText}\n</conversation>\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`;
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");
}

170
packages/agent/src/harness/compaction/utils.ts Normal file
View File

@@ -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<string>;
written: Set<string>;
edited: Set<string>;
}
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<string, unknown> | 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(`<read-files>\n${readFiles.join("\n")}\n</read-files>`);
}
if (modifiedFiles.length > 0) {
sections.push(`<modified-files>\n${modifiedFiles.join("\n")}\n</modified-files>`);
}
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<string, unknown>;
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.`;

152
packages/agent/src/harness/session/jsonl-session-storage.ts Normal file
View File

@@ -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<string, SessionTreeEntry>();
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<void> {
await mkdir(dirname(this.filePath), { recursive: true });
}
private async ensureLoaded(): Promise<void> {
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<void> {
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<SessionInfo> {
await this.ensureHeader();
return this.sessionInfo!;
}
async getLeafId(): Promise<string | null> {
await this.ensureLoaded();
return this.currentLeafId;
}
async setLeafId(leafId: string | null): Promise<void> {
await this.ensureLoaded();
this.currentLeafId = leafId;
}
async appendEntry(entry: SessionTreeEntry): Promise<void> {
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<SessionTreeEntry | undefined> {
await this.ensureLoaded();
return this.byId.get(id);
}
async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
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<SessionTreeEntry[]> {
await this.ensureLoaded();
return [...this.entries];
}
}

51
packages/agent/src/harness/session/memory-session-storage.ts Normal file
View File

@@ -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<SessionInfo> {
return this.sessionInfo;
}
async getLeafId(): Promise<string | null> {
return this.leafId;
}
async setLeafId(leafId: string | null): Promise<void> {
this.leafId = leafId;
}
async appendEntry(entry: SessionTreeEntry): Promise<void> {
this.entries.push(entry);
this.leafId = entry.id;
}
async getEntry(id: string): Promise<SessionTreeEntry | undefined> {
return this.entries.find((entry) => entry.id === id);
}
async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
if (leafId === null) return [];
const byId = new Map<string, SessionTreeEntry>(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<SessionTreeEntry[]> {
return [...this.entries];
}
}

231
packages/agent/src/harness/session/session-repo.ts Normal file
View File

@@ -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<TInfo extends SessionInfo>(info: TInfo, tree: DefaultSessionTree): Session<TInfo> {
return { info, tree };
}
function getPathEntriesToFork(
entries: SessionTreeEntry[],
entryId: string,
position: "before" | "at",
): SessionTreeEntry[] {
const byId = new Map<string, SessionTreeEntry>(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<string> {
private sessions = new Map<string, Session<SessionInfo>>();
async create(options?: { id?: string; parentSession?: string }): Promise<Session<SessionInfo>> {
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<Session<SessionInfo>> {
const session = this.sessions.get(ref);
if (!session) {
throw new Error(`Session not found: ${ref}`);
}
return session;
}
async list(): Promise<Array<Session<SessionInfo>>> {
return [...this.sessions.values()];
}
async delete(ref: string): Promise<void> {
this.sessions.delete(ref);
}
async fork(
ref: string,
options: { entryId: string; position?: "before" | "at"; id?: string },
): Promise<Session<SessionInfo>> {
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<string> {
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<Session<CodingAgentSessionInfo>> {
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<Session<CodingAgentSessionInfo>> {
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<Array<Session<CodingAgentSessionInfo>>> {
if (!existsSync(this.sessionDir)) {
return [];
}
const files = readdirSync(this.sessionDir)
.filter((file) => file.endsWith(".jsonl"))
.map((file) => join(this.sessionDir, file));
const sessions: Array<Session<CodingAgentSessionInfo>> = [];
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<Array<Session<CodingAgentSessionInfo>>> {
return (await this.list()).filter((session) => session.info.projectCwd === cwd);
}
async getMostRecentByCwd(cwd: string): Promise<Session<CodingAgentSessionInfo> | 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<void> {
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<Session<CodingAgentSessionInfo>> {
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));
}
}

202
packages/agent/src/harness/session-tree.ts → packages/agent/src/harness/session/session-tree.ts
View File

@@ -1,9 +1,7 @@
import { randomUUID } from "node:crypto"; 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 { ImageContent, TextContent } from "@mariozechner/pi-ai";
import type { AgentMessage } from "../types.js"; import type { AgentMessage } from "../../types.js";
import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "./messages.js"; import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "../messages.js";
import type { import type {
BranchSummaryEntry, BranchSummaryEntry,
CompactionEntry, CompactionEntry,
@@ -13,12 +11,14 @@ import type {
MessageEntry, MessageEntry,
ModelChangeEntry, ModelChangeEntry,
SessionContext, SessionContext,
SessionInfo,
SessionInfoEntry, SessionInfoEntry,
SessionTree, SessionTree,
SessionTreeEntry, SessionTreeEntry,
SessionTreeStorage, SessionTreeStorage,
ThinkingLevelChangeEntry, ThinkingLevelChangeEntry,
} from "./types.js"; } from "../types.js";
import { InMemorySessionTreeStorage } from "./memory-session-storage.js";
function generateId(byId: { has(id: string): boolean }): string { function generateId(byId: { has(id: string): boolean }): string {
for (let i = 0; i < 100; i++) { for (let i = 0; i < 100; i++) {
@@ -79,172 +79,6 @@ export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext
return { messages, thinkingLevel, model }; 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<string, SessionTreeEntry>();
private currentLeafId: string | null = null;
constructor(filePath: string, options: { cwd: string }) {
this.filePath = resolve(filePath);
this.cwd = options.cwd;
}
private async ensureParentDir(): Promise<void> {
await mkdir(dirname(this.filePath), { recursive: true });
}
private async ensureLoaded(): Promise<void> {
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<void> {
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<string | null> {
await this.ensureLoaded();
return this.currentLeafId;
}
async setLeafId(leafId: string | null): Promise<void> {
await this.ensureLoaded();
this.currentLeafId = leafId;
}
async appendEntry(entry: SessionTreeEntry): Promise<void> {
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<SessionTreeEntry | undefined> {
await this.ensureLoaded();
return this.byId.get(id);
}
async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
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<SessionTreeEntry[]> {
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<string | null> {
return this.leafId;
}
async setLeafId(leafId: string | null): Promise<void> {
this.leafId = leafId;
}
async appendEntry(entry: SessionTreeEntry): Promise<void> {
this.entries.push(entry);
this.leafId = entry.id;
}
async getEntry(id: string): Promise<SessionTreeEntry | undefined> {
return this.entries.find((entry) => entry.id === id);
}
async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
if (leafId === null) return [];
const byId = new Map<string, SessionTreeEntry>(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<SessionTreeEntry[]> {
return [...this.entries];
}
}
export class DefaultSessionTree implements SessionTree { export class DefaultSessionTree implements SessionTree {
private storage: SessionTreeStorage; private storage: SessionTreeStorage;
@@ -273,6 +107,32 @@ export class DefaultSessionTree implements SessionTree {
return buildSessionContext(await this.getBranch()); return buildSessionContext(await this.getBranch());
} }
getSessionInfo(): Promise<SessionInfo> {
return this.storage.getSessionInfo();
}
async getLabel(id: string): Promise<string | undefined> {
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<string | undefined> {
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<string> { private async makeEntryId(): Promise<string> {
const entries = await this.storage.getEntries(); const entries = await this.storage.getEntries();
return generateId(new Set(entries.map((entry) => entry.id))); return generateId(new Set(entries.map((entry) => entry.id)));

33
packages/agent/src/harness/types.ts
View File

@@ -157,7 +157,19 @@ export interface SessionContext {
model: { provider: string; modelId: string } | null; 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 { export interface SessionTreeStorage {
getSessionInfo(): Promise<SessionInfo>;
getLeafId(): Promise<string | null>; getLeafId(): Promise<string | null>;
setLeafId(leafId: string | null): Promise<void>; setLeafId(leafId: string | null): Promise<void>;
appendEntry(entry: SessionTreeEntry): Promise<void>; appendEntry(entry: SessionTreeEntry): Promise<void>;
@@ -167,11 +179,14 @@ export interface SessionTreeStorage {
} }
export interface SessionTree { export interface SessionTree {
getSessionInfo(): Promise<SessionInfo>;
getLeafId(): Promise<string | null>; getLeafId(): Promise<string | null>;
getEntry(id: string): Promise<SessionTreeEntry | undefined>; getEntry(id: string): Promise<SessionTreeEntry | undefined>;
getEntries(): Promise<SessionTreeEntry[]>; getEntries(): Promise<SessionTreeEntry[]>;
getBranch(fromId?: string): Promise<SessionTreeEntry[]>; getBranch(fromId?: string): Promise<SessionTreeEntry[]>;
buildContext(): Promise<SessionContext>; buildContext(): Promise<SessionContext>;
getLabel(id: string): Promise<string | undefined>;
getSessionName(): Promise<string | undefined>;
appendMessage(message: AgentMessage): Promise<string>; appendMessage(message: AgentMessage): Promise<string>;
appendThinkingLevelChange(thinkingLevel: string): Promise<string>; appendThinkingLevelChange(thinkingLevel: string): Promise<string>;
@@ -197,6 +212,24 @@ export interface SessionTree {
moveTo(entryId: string | null): Promise<void>; moveTo(entryId: string | null): Promise<void>;
} }
export interface Session<TInfo extends SessionInfo = SessionInfo> {
info: TInfo;
tree: SessionTree;
}
export interface SessionRepo<TRef = string, TInfo extends SessionInfo = SessionInfo> {
create(options?: { id?: string; parentSession?: string }): Promise<Session<TInfo>>;
open(ref: TRef): Promise<Session<TInfo>>;
list(): Promise<Array<Session<TInfo>>>;
delete(ref: TRef): Promise<void>;
fork(ref: TRef, options: { entryId: string; position?: "before" | "at"; id?: string }): Promise<Session<TInfo>>;
}
export interface CodingAgentSessionRepo<TRef = string> extends SessionRepo<TRef, CodingAgentSessionInfo> {
listByCwd(cwd: string): Promise<Array<Session<CodingAgentSessionInfo>>>;
getMostRecentByCwd(cwd: string): Promise<Session<CodingAgentSessionInfo> | undefined>;
}
export interface AgentHarnessPendingMutations { export interface AgentHarnessPendingMutations {
appendMessages: AgentMessage[]; appendMessages: AgentMessage[];
model?: Model<any>; model?: Model<any>;

23
packages/agent/src/index.ts
View File

@@ -3,11 +3,30 @@ export * from "./agent.js";
// Loop functions // Loop functions
export * from "./agent-loop.js"; export * from "./agent-loop.js";
export * from "./harness/agent-harness.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/execution-env.js";
export * from "./harness/messages.js"; export * from "./harness/messages.js";
export * from "./harness/prompt-templates.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 // Harness
export * from "./harness/types.js"; export * from "./harness/types.js";
export * from "./harness/utils/shell-output.js"; export * from "./harness/utils/shell-output.js";

6
packages/agent/test/harness/compaction.test.ts
View File

@@ -1,4 +1,3 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { import {
type AssistantMessage, type AssistantMessage,
type FauxProviderRegistration, type FauxProviderRegistration,
@@ -18,8 +17,8 @@ import {
prepareCompaction, prepareCompaction,
serializeConversation, serializeConversation,
shouldCompact, shouldCompact,
} from "../../src/harness/compaction.js"; } from "../../src/harness/compaction/compaction.js";
import { buildSessionContext } from "../../src/harness/session-tree.js"; import { buildSessionContext } from "../../src/harness/session/session-tree.js";
import type { import type {
CompactionEntry, CompactionEntry,
CompactionSettings, CompactionSettings,
@@ -28,6 +27,7 @@ import type {
SessionTreeEntry, SessionTreeEntry,
ThinkingLevelChangeEntry, ThinkingLevelChangeEntry,
} from "../../src/harness/types.js"; } from "../../src/harness/types.js";
import type { AgentMessage } from "../../src/types.js";
let nextId = 0; let nextId = 0;
function createId(): string { function createId(): string {

14
packages/agent/test/harness/session-tree.test.ts
View File

@@ -3,11 +3,9 @@ import { tmpdir } from "node:os";
import { join } from "node:path"; import { join } from "node:path";
import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { afterEach, describe, expect, it } from "vitest"; import { afterEach, describe, expect, it } from "vitest";
import { import { JsonlSessionTreeStorage } from "../../src/harness/session/jsonl-session-storage.js";
DefaultSessionTree, import { InMemorySessionTreeStorage } from "../../src/harness/session/memory-session-storage.js";
InMemorySessionTreeStorage, import { DefaultSessionTree } from "../../src/harness/session/session-tree.js";
JsonlSessionTreeStorage,
} from "../../src/harness/session-tree.js";
import type { SessionTreeStorage } from "../../src/harness/types.js"; import type { SessionTreeStorage } from "../../src/harness/types.js";
function createUserMessage(text: string): AgentMessage { function createUserMessage(text: string): AgentMessage {
@@ -128,6 +126,8 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre
const entries = await tree.getEntries(); const entries = await tree.getEntries();
expect(entries.some((entry) => entry.type === "label")).toBe(true); expect(entries.some((entry) => entry.type === "label")).toBe(true);
expect(entries.some((entry) => entry.type === "session_info")).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); expect((await tree.buildContext()).messages).toHaveLength(1);
}); });
@@ -136,11 +136,15 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre
const tree = new DefaultSessionTree(storage); const tree = new DefaultSessionTree(storage);
const user1 = await tree.appendMessage(createUserMessage("one")); const user1 = await tree.appendMessage(createUserMessage("one"));
await tree.appendMessage(createAssistantMessage("two")); await tree.appendMessage(createAssistantMessage("two"));
await tree.appendLabelChange(user1, "checkpoint");
await tree.appendSessionInfo("name");
await tree.moveTo(user1); await tree.moveTo(user1);
await tree.appendMessage(createAssistantMessage("branched")); await tree.appendMessage(createAssistantMessage("branched"));
const tree2 = new DefaultSessionTree(storage); const tree2 = new DefaultSessionTree(storage);
const context = await tree2.buildContext(); const context = await tree2.buildContext();
expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]); expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
expect(await tree2.getLabel(user1)).toBe("checkpoint");
expect(await tree2.getSessionName()).toBe("name");
inspect?.(); inspect?.();
}); });
}); });