Pi - Context Manage

Context Manage (in-Session Context)

• 使用session tree, 树形结构
• 支持message粒度branch(同一个session内)
• 支持从某个message fork一个新session, 或者创建一个空白session
• 使用JSONL持久化
• 核心代码位置
  • /assets/code/pi/context
    • /pi/packages/coding-agent/src/core/compaction/branch-summarization.ts
    • /pi/packages/coding-agent/src/core/compaction/branch-summarization.ts
    • /pi/packages/coding-agent/src/core/session-manager.ts

      Session Tree

      session head

• JSONL的第一行header, header只提供元信息, 不是tree node

/**
 * JSONL 文件第一行的 header。
 *
 * 注意：
 * - header 本身不是树节点；
 * - 它只描述 session 的元信息，例如 cwd、session id、父 session；
 * - 真正参与树结构的是后续各个 entry，它们通过 id/parentId 形成关系。
 */
export interface SessionHeader {
	type: "session";
	version?: number; // v1 sessions don't have this
	id: string;
	timestamp: string;
	cwd: string;
	parentSession?: string;
}

tree node

/** Tree node for getTree() - defensive copy of session structure */
export interface SessionTreeNode {
	entry: SessionEntry;
	children: SessionTreeNode[];
	/** Resolved label for this entry, if any */
	label?: string;
	/** Timestamp of the latest label change for this entry, if any */
	labelTimestamp?: string;
}

session entry

• entry base

export interface SessionEntryBase {
	type: string;
	id: string;
	parentId: string | null;
	timestamp: string;
}

• 多种类型, 这些类型在base上进行extension

export type SessionEntry =
    /* 例如
    export interface SessionMessageEntry extends SessionEntryBase {
        type: "message";
        message: AgentMessage;
    }
    */
	| SessionMessageEntry
	| ThinkingLevelChangeEntry
	| ModelChangeEntry
	| CompactionEntry
	| BranchSummaryEntry
	| CustomEntry
	| CustomMessageEntry
	| LabelEntry
	| SessionInfoEntry;

compaction

compaction entry

• 进行上下文压缩
• 压缩A->B, 插入一个新的compacttion节点在当前branch中
• 原本可能是A->B->C->… CompactionEntry要标记未压缩的第一个节点(C), 对应firstKeptEntryId

export interface CompactionEntry<T = unknown> extends SessionEntryBase {
	type: "compaction";
	summary: string;
	/**
	 * 压缩后仍然保留的最早 entry。
	 * buildSessionContext() 会从这个边界重新拼接被保留的历史，
	 * 再加上 compaction summary 与之后的新消息。
	 */
	firstKeptEntryId: string;
	tokensBefore: number;
	/** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */
	details?: T;
	/** True if generated by an extension, undefined/false if pi-generated (backward compatible) */
	fromHook?: boolean;
}

branch summary entry

• 当使用/tree command, 从一个branch切换到另一个, 可以选择主动进行一个summary
• 可以将原本在的branch的内容进行一个总结, 形成一个summary entry到新branch中

export interface BranchSummaryEntry<T = unknown> extends SessionEntryBase {
	type: "branch_summary";
	/**
	 * 这条摘要对应的是从哪个 entry/leaf 分叉离开的。
	 * 它帮助后续理解“这段摘要是从哪条旧分支沉淀过来的”。
	 */
	fromId: string;
	summary: string;
	/** Extension-specific data (not sent to LLM) */
	details?: T;
	/** True if generated by an extension, false if pi-generated */
	fromHook?: boolean;
}

构建Context

• leaf沿着当前branch向上扫描, 提取整个path

/* export function buildSessionContext(
	entries: SessionEntry[],
	leafId?: string | null,
	byId?: Map<string, SessionEntry>,
): SessionContext {}
*/
// 从 leaf 一直沿 parentId 向上回溯到 root，得到当前 branch 的完整 path。
// 这一步决定了：
// - session 文件里虽然可能有很多分叉；
// - 但真正参与当前上下文构建的，只是当前 leaf 对应的一条路径。
const path: SessionEntry[] = [];
let current: SessionEntry | undefined = leaf;
while (current) {
    path.unshift(current);
    current = current.parentId ? byId.get(current.parentId) : undefined;
}

• 构造发送给模型的消息列表

// 真正构造“要发给 LLM 的消息列表”。
// 关键点：这里不是简单地把 path 上所有 message 原样输出。
// 如果存在 compaction，逻辑会切换为：
// 1. 先输出一条 compactionSummary；
// 2. 再输出从 firstKeptEntryId 开始保留的那段旧消息；
// 3. 最后输出 compaction 之后新增的消息。
const messages: AgentMessage[] = [];

const appendMessage = (entry: SessionEntry) => {
    // 只有“真正能参与模型上下文理解”的 entry 才会被映射成消息。
    // 例如 label / session_info / custom 不会进入这里。
    if (entry.type === "message") {
        messages.push(entry.message);
    } else if (entry.type === "custom_message") {
        messages.push(
            createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp),
        );
    } else if (entry.type === "branch_summary" && entry.summary) {
        messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp));
    }
};

• 如果没有compact, path全部进入消息列表

// 没有压缩边界时，当前 path 上所有可见消息直接进入上下文。
for (const entry of path) {
    appendMessage(entry);
}

• 如果有compact, 将summary放到第一条, 然后再保留firstKeptEntryId之后的消息

// 压缩后的上下文第一条一定是 summary，
// 这样模型先拿到历史摘要，再读取保留的尾部上下文。
messages.push(createCompactionSummaryMessage(compaction.summary, compaction.tokensBefore, compaction.timestamp));

// 只保留从 firstKeptEntryId 开始的那一段历史。
// 这能保证 compaction 不是粗暴截断，而是“摘要 + 保留尾部”。
let foundFirstKept = false;
for (let i = 0; i < compactionIdx; i++) {
    const entry = path[i];
    if (entry.id === compaction.firstKeptEntryId) {
        foundFirstKept = true;
    }
    if (foundFirstKept) {
        appendMessage(entry);
    }
}

Compaction

Overview

/**
 * 长 session 的上下文压缩（compaction）实现。
 *
 * 这个文件只负责“纯逻辑层”的压缩决策与摘要生成，不负责 session 文件的落盘：
 * - SessionManager 负责把 append-only tree 持久化到 JSONL；
 * - 本文件负责决定“从哪里切”、“哪些消息被总结”、“summary 怎么生成”；
 * - 压缩完成后，SessionManager 会把 compaction entry 追加进树，然后重新加载上下文。
 *
 * 可以把它理解成：
 * 1. 先根据 keepRecentTokens 找到应该保留的尾部边界；
 * 2. 将边界之前的历史消息总结成 summary；
 * 3. 后续 buildSessionContext() 不再直接送入那些旧消息，而是送入：
 *    compaction summary + firstKeptEntryId 之后的保留消息。
 */

• 访问过得文件信息保留

/**
 * 这里不是“压缩算法本身”的必要字段，而是 Pi 默认实现附带保存的辅助信息：
 * 用于告诉后续模型/用户，这段被总结的历史里读过哪些文件、改过哪些文件。
 */
export interface CompactionDetails {
	readFiles: string[];
	modifiedFiles: string[];
}

• 压缩判断
  • contextTokens > contextWindow - settings.reserveTokens;

/**
 * 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)

/**
 * 在 path entries 中找到“从哪里开始保留尾部”的切点。
 *
 * 算法思路：
 * 1. 从最新 entry 往回走，累计消息大小；
 * 2. 一旦累计值达到 keepRecentTokens，就说明“尾部预算”已经够了；
 * 3. 再把当前位置对齐到最近的合法 cut point；
 * 4. 最终返回 firstKeptEntryIndex。
 *
 * 这里的“切”不是把树裁掉，而是决定：
 * - cut 点之前的内容进入 summary；
 * - cut 点及之后的内容继续原样保留在运行时上下文中。
 *
 * 还要特别处理 split turn：
 * - 如果 cut 点不是 user 起点，而是落在 assistant / custom-like message 上，
 *   说明一个 turn 太大，大到必须在 turn 中间切；
 * - 此时需要找到 turnStartIndex，把这段 turn prefix 单独做摘要，保证语义连贯。
 */
export function findCutPoint(
	entries: SessionEntry[],
	startIndex: number,
	endIndex: number,
	keepRecentTokens: number,
): CutPointResult {
    const cutPoints = findValidCutPoints(entries, startIndex, endIndex);
    // 从后往前累计“要原样保留”的尾部大小。
	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;

		// 这里只看 message entry 的估算大小；控制态 entry 不直接计入消息 token。
		const messageTokens = estimateTokens(entry.message);
		accumulatedTokens += messageTokens;
        // ...
	}
    // 如果 cut 点前面紧挨着一些非消息 entry（例如 model/thinking 变化），
	// 也一并纳入 kept 区域，避免上下文从中间断开配置态。
    // ...
}

prompt

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.`;

生成摘要

/**
 * 调用 LLM 为一段历史对话生成结构化摘要。
 *
 * 这是 compaction 的核心摘要函数，负责把“待压缩的消息序列”变成可重新注入上下文的 summary 文本。
 * 它既支持首轮总结，也支持在已有 summary 基础上做增量更新。
 *
 * 两种工作模式：
 * 1. 首次总结
 *    - 当 previousSummary 不存在时，使用 SUMMARIZATION_PROMPT；
 *    - 模型直接根据 currentMessages 生成一份新的结构化 checkpoint summary。
 *
 * 2. 增量更新
 *    - 当 previousSummary 存在时，使用 UPDATE_SUMMARIZATION_PROMPT；
 *    - 模型会把“旧 summary + 新消息”合并成一份更新后的 summary，
 *      而不是从零开始重写，借此实现滚动 compaction。
 *
 * 重要实现点：
 * - currentMessages 不会直接以“原始对话消息数组”形式喂给模型，
 *   而是先 convertToLlm()，再 serializeConversation() 序列化成纯文本；
 * - 这样做是为了明确告诉模型：这些内容是“待总结材料”，不是要继续续写的对话；
 * - previousSummary 会放进 <previous-summary> 标签中，让模型把它当作现有状态快照；
 * - customInstructions 不是替换默认模板，而是以 Additional focus 的形式附加，
 *   用来强调本次总结应额外关注的点；
 * - maxTokens 使用 reserveTokens 的 80% 作为上限，避免摘要本身把保留预算全部耗尽。
 *
 * 返回值：
 * - 返回纯文本 summary；
 * - 调用方（如 compact()）会再决定是否附加文件列表、如何写入 compaction entry。
 */
export async function generateSummary(
	currentMessages: AgentMessage[],
	model: Model<any>,
	reserveTokens: number,
	apiKey: string | undefined,
	headers?: Record<string, string>,
	signal?: AbortSignal,
	customInstructions?: string,
	previousSummary?: string,
	thinkingLevel?: ThinkingLevel,
	streamFn?: StreamFn,
): Promise<string> {}

Main compaction function

/**
 * 根据 prepareCompaction() 的结果真正调用模型生成摘要。
 *
 * 主流程：
 * - 普通情况：只生成一份 history summary；
 * - split turn：并行生成 history summary + turn prefix summary，再合并成最终 summary；
 * - 最后再附加文件读写清单，作为对后继模型非常有价值的工作记忆。
 *
 * 返回的是“待写入 compaction entry 的 payload”，不是最终 session entry。
 */
export async function compact(
	preparation: CompactionPreparation,
	model: Model<any>,
	apiKey: string | undefined,
	headers?: Record<string, string>,
	customInstructions?: string,
	signal?: AbortSignal,
	thinkingLevel?: ThinkingLevel,
	streamFn?: StreamFn,
): Promise<CompactionResult> {
    // 调用generateSummary
}

Branch Summary

Overview

/**
 * 树导航（tree navigation）场景下的分支摘要逻辑。
 *
 * 这个文件处理的不是“上下文过长时的 compaction”，而是另一类问题：
 * 当用户从 session tree 的一个分支切换到另一个分支时，
 * 被离开的那条分支上可能已经积累了很多决策、进展和文件操作痕迹。
 * 如果直接切走，这些认知成果会暂时脱离当前活跃路径。
 *
 * 因此 branch summarization 的职责是：
 * - 找出“旧位置 -> 新位置”之间被放下的那段分支；
 * - 将这段分支压缩成一条 branch_summary；
 * - 让后续在新分支继续工作时，仍能带着对旧分支成果的高层理解。
 *
 * 可以把它理解为：
 * - compaction 解决的是“上下文窗口不够大”；
 * - branch summarization 解决的是“切换分支后，如何保留离开分支的工作记忆”。
 */

要总结的范围

/**
 * 收集“从旧位置切到新位置时，应该被总结”的那段 entry。
 *
 * 核心语义：
 * - oldLeafId 是当前所在位置，也就是即将离开的分支末端；
 * - targetId 是即将跳转到的新位置；
 * - 真正要总结的，不是整棵树，而是 oldLeaf 到两条路径公共祖先之间的那一段。
 *
 * 这里不会在 compaction 边界处停止：
 * - 因为 compaction entry 本身也代表历史语义；
 * - 进入 branch summary 时，它应该像普通上下文一样参与理解。
 *
 * @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
 */

Summary Generation

• prompt
• summary = BRANCH_SUMMARY_PREAMBLE + generatedSummary(using BRANCH_SUMMARY_PROMPT)

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.`;

• generateBranchSummary function

/**
 * 为“被离开的分支”生成 branch summary。
 *
 * 这个函数是分支摘要的主入口，职责包括：
 * - 根据模型上下文窗口计算可用预算；
 * - 从 entries 中筛出适合进入 prompt 的消息；
 * - 将消息序列序列化为纯文本，避免模型把它当成“继续对话”；
 * - 调用 LLM 生成结构化总结；
 * - 在 summary 末尾追加文件读写清单。
 *
 * 最终产物会被上层写成 branch_summary entry，之后在新分支上作为一条上下文消息使用。
 */
export async function generateBranchSummary(
	entries: SessionEntry[],
	options: GenerateBranchSummaryOptions,
): Promise<BranchSummaryResult> {}