initial commit

This commit is contained in:
pi
2026-04-10 23:06:52 +01:00
commit 1d75f91349
25 changed files with 8195 additions and 0 deletions

24
README.md Normal file
View File

@@ -0,0 +1,24 @@
# pi-context-manager
`pi-context-manager` is a Pi extension package for context-pressure management, snapshots, resume packets, and branch-summary compaction behavior.
## Install
Use it as a local package root today:
```bash
pi install /absolute/path/to/context-manager
```
After this folder is moved into its own repository, the same package can be installed from git.
## Resources
- Extension: `./index.ts`
## Development
```bash
npm install
npm test
```

353
index.ts Normal file
View File

@@ -0,0 +1,353 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import { adjustPolicyForZone } from "./src/config.ts";
import { deserializeLatestSnapshot, serializeSnapshot, SNAPSHOT_ENTRY_TYPE, type RuntimeSnapshot } from "./src/persist.ts";
import { createEmptyLedger } from "./src/ledger.ts";
import { pruneContextMessages } from "./src/prune.ts";
import { createContextManagerRuntime } from "./src/runtime.ts";
import { registerContextCommands } from "./src/commands.ts";
import { buildBranchSummaryFromEntries, buildCompactionSummaryFromPreparation } from "./src/summaries.ts";
type TrackedMessage = Extract<AgentMessage, { role: "user" | "assistant" | "toolResult" }>;
type BranchEntry = ReturnType<ExtensionContext["sessionManager"]["getBranch"]>[number];
function isTextPart(part: unknown): part is { type: "text"; text?: string } {
return typeof part === "object" && part !== null && "type" in part && (part as { type?: unknown }).type === "text";
}
function toText(content: unknown): string {
if (typeof content === "string") return content;
if (!Array.isArray(content)) return "";
return content
.map((part) => {
if (!isTextPart(part)) return "";
return typeof part.text === "string" ? part.text : "";
})
.join("\n")
.trim();
}
function isMessageEntry(entry: BranchEntry): entry is Extract<BranchEntry, { type: "message" }> {
return entry.type === "message";
}
function isCompactionEntry(entry: BranchEntry): entry is Extract<BranchEntry, { type: "compaction" }> {
return entry.type === "compaction";
}
function isBranchSummaryEntry(entry: BranchEntry): entry is Extract<BranchEntry, { type: "branch_summary" }> {
return entry.type === "branch_summary";
}
function isTrackedMessage(message: AgentMessage): message is TrackedMessage {
return message.role === "user" || message.role === "assistant" || message.role === "toolResult";
}
function createDefaultSnapshot(): RuntimeSnapshot {
return {
mode: "balanced",
lastZone: "green",
ledger: createEmptyLedger(),
};
}
function getMessageContent(message: AgentMessage): string {
return "content" in message ? toText(message.content) : "";
}
function getMessageToolName(message: AgentMessage): string | undefined {
return message.role === "toolResult" ? message.toolName : undefined;
}
function rewriteContextMessage(message: { role: string; content: string; original: AgentMessage; distilled?: boolean }): AgentMessage {
if (!message.distilled || message.role !== "toolResult") {
return message.original;
}
return {
...(message.original as Extract<AgentMessage, { role: "toolResult" }>),
content: [{ type: "text", text: message.content }],
} as AgentMessage;
}
function findLatestSnapshotState(branch: BranchEntry[]): { snapshot: RuntimeSnapshot; index: number } | undefined {
for (let index = branch.length - 1; index >= 0; index -= 1) {
const entry = branch[index]!;
if (entry.type !== "custom" || entry.customType !== SNAPSHOT_ENTRY_TYPE) {
continue;
}
const snapshot = deserializeLatestSnapshot([entry]);
if (snapshot) {
return { snapshot, index };
}
}
return undefined;
}
function findLatestSessionSnapshot(entries: BranchEntry[]): RuntimeSnapshot | undefined {
let latest: RuntimeSnapshot | undefined;
let latestFreshness = -Infinity;
for (const entry of entries) {
if (entry.type !== "custom" || entry.customType !== SNAPSHOT_ENTRY_TYPE) {
continue;
}
const snapshot = deserializeLatestSnapshot([entry]);
if (!snapshot) {
continue;
}
const sessionItems = snapshot.ledger.items.filter((item) => item.scope === "session");
const freshness = sessionItems.length > 0 ? Math.max(...sessionItems.map((item) => item.timestamp)) : -Infinity;
if (freshness >= latestFreshness) {
latest = snapshot;
latestFreshness = freshness;
}
}
return latest;
}
function createSessionFallbackSnapshot(source?: RuntimeSnapshot): RuntimeSnapshot {
return {
mode: source?.mode ?? "balanced",
lastZone: "green",
ledger: {
items: structuredClone((source?.ledger.items ?? []).filter((item) => item.scope === "session")),
rollingSummary: "",
},
};
}
function overlaySessionLayer(base: RuntimeSnapshot, latestSessionSnapshot?: RuntimeSnapshot): RuntimeSnapshot {
const sessionItems = latestSessionSnapshot?.ledger.items.filter((item) => item.scope === "session") ?? [];
if (sessionItems.length === 0) {
return base;
}
return {
...base,
ledger: {
...base.ledger,
items: [
...structuredClone(base.ledger.items.filter((item) => item.scope !== "session")),
...structuredClone(sessionItems),
],
},
};
}
export default function contextManager(pi: ExtensionAPI) {
const runtime = createContextManagerRuntime({
mode: "balanced",
contextWindow: 200_000,
});
let pendingResumeInjection = false;
const syncContextWindow = (ctx: Pick<ExtensionContext, "model">) => {
runtime.setContextWindow(ctx.model?.contextWindow ?? 200_000);
};
const armResumeInjection = () => {
const snapshot = runtime.getSnapshot();
pendingResumeInjection = Boolean(snapshot.lastCompactionSummary || snapshot.lastBranchSummary) && runtime.buildResumePacket().trim().length > 0;
};
const replayBranchEntry = (entry: BranchEntry) => {
if (isMessageEntry(entry) && isTrackedMessage(entry.message)) {
runtime.ingest({
entryId: entry.id,
role: entry.message.role,
text: toText(entry.message.content),
timestamp: entry.message.timestamp,
isError: entry.message.role === "toolResult" ? entry.message.isError : undefined,
});
return;
}
if (isCompactionEntry(entry)) {
runtime.recordCompactionSummary(entry.summary, entry.id, Date.parse(entry.timestamp));
return;
}
if (isBranchSummaryEntry(entry)) {
runtime.recordBranchSummary(entry.summary, entry.id, Date.parse(entry.timestamp));
}
};
const rebuildRuntimeFromBranch = (
ctx: Pick<ExtensionContext, "model" | "sessionManager" | "ui">,
fallbackSnapshot: RuntimeSnapshot,
options?: { preferRuntimeMode?: boolean },
) => {
syncContextWindow(ctx);
const branch = ctx.sessionManager.getBranch();
const latestSessionSnapshot = findLatestSessionSnapshot(ctx.sessionManager.getEntries() as BranchEntry[]);
const restored = findLatestSnapshotState(branch);
const baseSnapshot = restored
? overlaySessionLayer(restored.snapshot, latestSessionSnapshot)
: createSessionFallbackSnapshot(latestSessionSnapshot ?? fallbackSnapshot);
runtime.restore({
...baseSnapshot,
mode: options?.preferRuntimeMode ? fallbackSnapshot.mode : baseSnapshot.mode,
});
const replayEntries = restored ? branch.slice(restored.index + 1) : branch;
for (const entry of replayEntries) {
replayBranchEntry(entry);
}
const snapshot = runtime.getSnapshot();
ctx.ui.setStatus("context-manager", `ctx ${snapshot.lastZone}`);
};
registerContextCommands(pi, {
getSnapshot: runtime.getSnapshot,
buildPacket: runtime.buildPacket,
buildResumePacket: runtime.buildResumePacket,
setMode: runtime.setMode,
rebuildFromBranch: async (commandCtx) => {
rebuildRuntimeFromBranch(commandCtx, runtime.getSnapshot(), { preferRuntimeMode: true });
armResumeInjection();
},
isResumePending: () => pendingResumeInjection,
});
pi.on("session_start", async (_event, ctx) => {
rebuildRuntimeFromBranch(ctx, createDefaultSnapshot());
armResumeInjection();
});
pi.on("session_tree", async (event, ctx) => {
rebuildRuntimeFromBranch(ctx, createDefaultSnapshot());
if (
event.summaryEntry &&
!ctx.sessionManager.getBranch().some((entry) => isBranchSummaryEntry(entry) && entry.id === event.summaryEntry.id)
) {
runtime.recordBranchSummary(event.summaryEntry.summary, event.summaryEntry.id, Date.parse(event.summaryEntry.timestamp));
}
armResumeInjection();
if (event.summaryEntry) {
pi.appendEntry(SNAPSHOT_ENTRY_TYPE, serializeSnapshot(runtime.getSnapshot()));
}
});
pi.on("tool_result", async (event) => {
runtime.ingest({
entryId: event.toolCallId,
role: "toolResult",
text: toText(event.content),
timestamp: Date.now(),
});
});
pi.on("turn_end", async (_event, ctx) => {
rebuildRuntimeFromBranch(ctx, runtime.getSnapshot(), { preferRuntimeMode: true });
const usage = ctx.getContextUsage();
if (usage?.tokens !== null && usage?.tokens !== undefined) {
runtime.observeTokens(usage.tokens);
}
const snapshot = runtime.getSnapshot();
pi.appendEntry(SNAPSHOT_ENTRY_TYPE, serializeSnapshot(snapshot));
ctx.ui.setStatus("context-manager", `ctx ${snapshot.lastZone}`);
});
pi.on("context", async (event, ctx) => {
syncContextWindow(ctx);
const snapshot = runtime.getSnapshot();
const policy = adjustPolicyForZone(runtime.getPolicy(), snapshot.lastZone);
const normalized = event.messages.map((message) => ({
role: message.role,
content: getMessageContent(message),
toolName: getMessageToolName(message),
original: message,
}));
const pruned = pruneContextMessages(normalized, policy);
const nextMessages = pruned.map((message) =>
rewriteContextMessage(message as { role: string; content: string; original: AgentMessage; distilled?: boolean }),
);
const resumeText = pendingResumeInjection ? runtime.buildResumePacket() : "";
const packetText = pendingResumeInjection ? "" : runtime.buildPacket().text;
const injectedText = resumeText || packetText;
if (!injectedText) {
return { messages: nextMessages };
}
if (resumeText) {
pendingResumeInjection = false;
}
return {
messages: [
{
role: "custom",
customType: resumeText ? "context-manager.resume" : "context-manager.packet",
content: injectedText,
display: false,
timestamp: Date.now(),
} as any,
...nextMessages,
],
};
});
pi.on("session_before_compact", async (event, ctx) => {
syncContextWindow(ctx);
try {
return {
compaction: {
summary: buildCompactionSummaryFromPreparation({
messagesToSummarize: event.preparation.messagesToSummarize,
turnPrefixMessages: event.preparation.turnPrefixMessages,
previousSummary: event.preparation.previousSummary,
fileOps: event.preparation.fileOps,
customInstructions: event.customInstructions,
}),
firstKeptEntryId: event.preparation.firstKeptEntryId,
tokensBefore: event.preparation.tokensBefore,
details: event.preparation.fileOps,
},
};
} catch (error) {
ctx.ui.notify(`context-manager compaction fallback: ${error instanceof Error ? error.message : String(error)}`, "warning");
return;
}
});
pi.on("session_before_tree", async (event, ctx) => {
syncContextWindow(ctx);
if (!event.preparation.userWantsSummary) return;
return {
summary: {
summary: buildBranchSummaryFromEntries({
branchLabel: "branch handoff",
entriesToSummarize: event.preparation.entriesToSummarize,
customInstructions: event.preparation.customInstructions,
replaceInstructions: event.preparation.replaceInstructions,
commonAncestorId: event.preparation.commonAncestorId,
}),
},
};
});
pi.on("session_compact", async (event, ctx) => {
runtime.recordCompactionSummary(event.compactionEntry.summary, event.compactionEntry.id, Date.parse(event.compactionEntry.timestamp));
pi.appendEntry(SNAPSHOT_ENTRY_TYPE, serializeSnapshot(runtime.getSnapshot()));
armResumeInjection();
ctx.ui.setStatus("context-manager", `ctx ${runtime.getSnapshot().lastZone}`);
});
}

4366
package-lock.json generated Normal file

File diff suppressed because it is too large Load Diff

24
package.json Normal file
View File

@@ -0,0 +1,24 @@
{
"name": "pi-context-manager",
"version": "0.1.0",
"type": "module",
"keywords": ["pi-package"],
"scripts": {
"test": "tsx --test src/*.test.ts src/**/*.test.ts"
},
"files": ["index.ts", "src"],
"pi": {
"extensions": ["./index.ts"]
},
"peerDependencies": {
"@mariozechner/pi-agent-core": "*",
"@mariozechner/pi-coding-agent": "*"
},
"devDependencies": {
"@mariozechner/pi-agent-core": "^0.66.1",
"@mariozechner/pi-coding-agent": "^0.66.1",
"@types/node": "^25.5.2",
"tsx": "^4.21.0",
"typescript": "^6.0.2"
}
}

76
src/commands.ts Normal file
View File

@@ -0,0 +1,76 @@
import type { ExtensionAPI, ExtensionCommandContext } from "@mariozechner/pi-coding-agent";
import type { ContextMode } from "./config.ts";
import { serializeSnapshot, SNAPSHOT_ENTRY_TYPE, type RuntimeSnapshot } from "./persist.ts";
interface CommandRuntime {
getSnapshot(): RuntimeSnapshot;
buildPacket(): { estimatedTokens: number };
buildResumePacket(): string;
setMode(mode: ContextMode): void;
rebuildFromBranch(ctx: ExtensionCommandContext): Promise<void>;
isResumePending(): boolean;
}
export function registerContextCommands(pi: ExtensionAPI, runtime: CommandRuntime) {
pi.registerCommand("ctx-status", {
description: "Show context pressure, packet status, and persisted handoff state",
handler: async (_args, ctx) => {
const snapshot = runtime.getSnapshot();
const packet = runtime.buildPacket();
const resumePending = runtime.isResumePending();
const contextTokens = ctx.getContextUsage()?.tokens;
const nextInjectionTokens = resumePending ? Math.ceil(runtime.buildResumePacket().length / 4) : packet.estimatedTokens;
ctx.ui.notify(
[
`mode=${snapshot.mode}`,
`zone=${snapshot.lastZone}`,
`contextTokens=${contextTokens ?? "unknown"}`,
`packetTokens=${packet.estimatedTokens}`,
`nextInjectionTokens=${nextInjectionTokens}`,
`resumePending=${resumePending ? "yes" : "no"}`,
`compaction=${snapshot.lastCompactionSummary ? "yes" : "no"}`,
`branch=${snapshot.lastBranchSummary ? "yes" : "no"}`,
].join(" "),
"info",
);
},
});
pi.registerCommand("ctx-memory", {
description: "Inspect the active context ledger",
handler: async (_args, ctx) => {
const snapshot = runtime.getSnapshot();
await ctx.ui.editor("Context ledger", JSON.stringify(snapshot.ledger, null, 2));
},
});
pi.registerCommand("ctx-refresh", {
description: "Rebuild runtime state from the current branch and refresh the working packet",
handler: async (_args, ctx) => {
await runtime.rebuildFromBranch(ctx);
const packet = runtime.buildPacket();
ctx.ui.notify(`rebuilt runtime from branch (${packet.estimatedTokens} tokens)`, "info");
},
});
pi.registerCommand("ctx-compact", {
description: "Trigger compaction with optional focus instructions",
handler: async (args, ctx) => {
ctx.compact({ customInstructions: args.trim() || undefined });
},
});
pi.registerCommand("ctx-mode", {
description: "Switch context mode: conservative | balanced | aggressive",
handler: async (args, ctx) => {
const value = args.trim() as "conservative" | "balanced" | "aggressive";
if (!["conservative", "balanced", "aggressive"].includes(value)) {
ctx.ui.notify("usage: /ctx-mode conservative|balanced|aggressive", "warning");
return;
}
runtime.setMode(value);
pi.appendEntry(SNAPSHOT_ENTRY_TYPE, serializeSnapshot(runtime.getSnapshot()));
ctx.ui.notify(`context mode set to ${value}`, "info");
},
});
}

86
src/config.test.ts Normal file
View File

@@ -0,0 +1,86 @@
import test from "node:test";
import assert from "node:assert/strict";
import { adjustPolicyForZone, resolvePolicy, zoneForTokens } from "./config.ts";
test("resolvePolicy returns the balanced policy for a 200k context window", () => {
const policy = resolvePolicy({ mode: "balanced", contextWindow: 200_000 });
assert.deepEqual(policy, {
mode: "balanced",
recentUserTurns: 4,
packetTokenCap: 1_200,
bulkyBytes: 4_096,
bulkyLines: 150,
yellowAtTokens: 110_000,
redAtTokens: 140_000,
compactAtTokens: 164_000,
});
});
test("resolvePolicy clamps context windows below 50k before calculating thresholds", () => {
const policy = resolvePolicy({ mode: "balanced", contextWindow: 10_000 });
assert.deepEqual(policy, {
mode: "balanced",
recentUserTurns: 4,
packetTokenCap: 1_200,
bulkyBytes: 4_096,
bulkyLines: 150,
yellowAtTokens: 27_500,
redAtTokens: 35_000,
compactAtTokens: 41_000,
});
});
test("aggressive mode compacts earlier than conservative mode", () => {
const aggressive = resolvePolicy({ mode: "aggressive", contextWindow: 200_000 });
const conservative = resolvePolicy({ mode: "conservative", contextWindow: 200_000 });
assert.ok(aggressive.compactAtTokens < conservative.compactAtTokens);
});
test("aggressive mode reduces raw-window and packet budgets compared with conservative mode", () => {
const aggressive = resolvePolicy({ mode: "aggressive", contextWindow: 200_000 });
const conservative = resolvePolicy({ mode: "conservative", contextWindow: 200_000 });
assert.ok(aggressive.recentUserTurns < conservative.recentUserTurns);
assert.ok(aggressive.packetTokenCap < conservative.packetTokenCap);
assert.ok(aggressive.bulkyBytes < conservative.bulkyBytes);
assert.ok(aggressive.bulkyLines < conservative.bulkyLines);
});
test("adjustPolicyForZone tightens packet and pruning thresholds in yellow, red, and compact zones", () => {
const base = resolvePolicy({ mode: "balanced", contextWindow: 200_000 });
const yellow = adjustPolicyForZone(base, "yellow");
const red = adjustPolicyForZone(base, "red");
const compact = adjustPolicyForZone(base, "compact");
assert.ok(yellow.packetTokenCap < base.packetTokenCap);
assert.ok(yellow.bulkyBytes < base.bulkyBytes);
assert.ok(red.packetTokenCap < yellow.packetTokenCap);
assert.ok(red.recentUserTurns <= yellow.recentUserTurns);
assert.ok(red.bulkyBytes < yellow.bulkyBytes);
assert.ok(compact.packetTokenCap < red.packetTokenCap);
assert.ok(compact.recentUserTurns <= red.recentUserTurns);
assert.ok(compact.bulkyLines < red.bulkyLines);
});
test("zoneForTokens returns green, yellow, red, and compact for the balanced 200k policy", () => {
const policy = resolvePolicy({ mode: "balanced", contextWindow: 200_000 });
assert.equal(zoneForTokens(80_000, policy), "green");
assert.equal(zoneForTokens(120_000, policy), "yellow");
assert.equal(zoneForTokens(150_000, policy), "red");
assert.equal(zoneForTokens(170_000, policy), "compact");
});
test("zoneForTokens uses inclusive balanced 200k thresholds", () => {
const policy = resolvePolicy({ mode: "balanced", contextWindow: 200_000 });
assert.equal(zoneForTokens(109_999, policy), "green");
assert.equal(zoneForTokens(110_000, policy), "yellow");
assert.equal(zoneForTokens(139_999, policy), "yellow");
assert.equal(zoneForTokens(140_000, policy), "red");
assert.equal(zoneForTokens(163_999, policy), "red");
assert.equal(zoneForTokens(164_000, policy), "compact");
});

97
src/config.ts Normal file
View File

@@ -0,0 +1,97 @@
export type ContextMode = "conservative" | "balanced" | "aggressive";
export type ContextZone = "green" | "yellow" | "red" | "compact";
export interface Policy {
mode: ContextMode;
recentUserTurns: number;
packetTokenCap: number;
bulkyBytes: number;
bulkyLines: number;
yellowAtTokens: number;
redAtTokens: number;
compactAtTokens: number;
}
export const MODE_PCTS: Record<ContextMode, { yellow: number; red: number; compact: number }> = {
conservative: { yellow: 0.60, red: 0.76, compact: 0.88 },
balanced: { yellow: 0.55, red: 0.70, compact: 0.82 },
aggressive: { yellow: 0.50, red: 0.64, compact: 0.76 },
};
const MODE_SETTINGS: Record<ContextMode, Pick<Policy, "recentUserTurns" | "packetTokenCap" | "bulkyBytes" | "bulkyLines">> = {
conservative: {
recentUserTurns: 5,
packetTokenCap: 1_400,
bulkyBytes: 6_144,
bulkyLines: 220,
},
balanced: {
recentUserTurns: 4,
packetTokenCap: 1_200,
bulkyBytes: 4_096,
bulkyLines: 150,
},
aggressive: {
recentUserTurns: 3,
packetTokenCap: 900,
bulkyBytes: 3_072,
bulkyLines: 100,
},
};
export function resolvePolicy(input: { mode: ContextMode; contextWindow: number }): Policy {
const contextWindow = Math.max(input.contextWindow, 50_000);
const percentages = MODE_PCTS[input.mode];
const settings = MODE_SETTINGS[input.mode];
return {
mode: input.mode,
recentUserTurns: settings.recentUserTurns,
packetTokenCap: settings.packetTokenCap,
bulkyBytes: settings.bulkyBytes,
bulkyLines: settings.bulkyLines,
yellowAtTokens: Math.floor(contextWindow * percentages.yellow),
redAtTokens: Math.floor(contextWindow * percentages.red),
compactAtTokens: Math.floor(contextWindow * percentages.compact),
};
}
export function adjustPolicyForZone(policy: Policy, zone: ContextZone): Policy {
if (zone === "green") {
return { ...policy };
}
if (zone === "yellow") {
return {
...policy,
packetTokenCap: Math.max(500, Math.floor(policy.packetTokenCap * 0.9)),
bulkyBytes: Math.max(1_536, Math.floor(policy.bulkyBytes * 0.9)),
bulkyLines: Math.max(80, Math.floor(policy.bulkyLines * 0.9)),
};
}
if (zone === "red") {
return {
...policy,
recentUserTurns: Math.max(2, policy.recentUserTurns - 1),
packetTokenCap: Math.max(400, Math.floor(policy.packetTokenCap * 0.75)),
bulkyBytes: Math.max(1_024, Math.floor(policy.bulkyBytes * 0.75)),
bulkyLines: Math.max(60, Math.floor(policy.bulkyLines * 0.75)),
};
}
return {
...policy,
recentUserTurns: Math.max(1, policy.recentUserTurns - 2),
packetTokenCap: Math.max(300, Math.floor(policy.packetTokenCap * 0.55)),
bulkyBytes: Math.max(768, Math.floor(policy.bulkyBytes * 0.5)),
bulkyLines: Math.max(40, Math.floor(policy.bulkyLines * 0.5)),
};
}
export function zoneForTokens(tokens: number, policy: Policy): ContextZone {
if (tokens >= policy.compactAtTokens) return "compact";
if (tokens >= policy.redAtTokens) return "red";
if (tokens >= policy.yellowAtTokens) return "yellow";
return "green";
}

30
src/distill.test.ts Normal file
View File

@@ -0,0 +1,30 @@
import test from "node:test";
import assert from "node:assert/strict";
import { distillToolResult } from "./distill.ts";
const noisy = [
"Build failed while compiling focus parser",
"Error: missing export createFocusMatcher from ./summary-focus.ts",
"at src/summaries.ts:44:12",
"line filler",
"line filler",
"line filler",
].join("\n");
test("distillToolResult prioritizes salient failure lines and truncates noise", () => {
const distilled = distillToolResult({ toolName: "bash", content: noisy });
assert.ok(distilled);
assert.match(distilled!, /Build failed while compiling focus parser/);
assert.match(distilled!, /missing export createFocusMatcher/);
assert.ok(distilled!.length < 320);
});
test("distillToolResult falls back to the first meaningful non-empty lines", () => {
const distilled = distillToolResult({
toolName: "read",
content: ["", "src/runtime.ts", "exports createContextManagerRuntime", "", "more noise"].join("\n"),
});
assert.equal(distilled, "[distilled read output] src/runtime.ts; exports createContextManagerRuntime");
});

47
src/distill.ts Normal file
View File

@@ -0,0 +1,47 @@
const ERROR_RE = /\b(?:error|failed|failure|missing|undefined|exception)\b/i;
const LOCATION_RE = /\b(?:at\s+.+:\d+(?::\d+)?)\b|(?:[A-Za-z0-9_./-]+\.(?:ts|tsx|js|mjs|json|md):\d+(?::\d+)?)/i;
const MAX_SUMMARY_LENGTH = 320;
const MAX_LINES = 2;
function unique(lines: string[]): string[] {
return lines.filter((line, index) => lines.indexOf(line) === index);
}
function pickSalientLines(content: string): string[] {
const lines = content
.split(/\n+/)
.map((line) => line.trim())
.filter(Boolean);
if (lines.length === 0) {
return [];
}
const important = unique(lines.filter((line) => ERROR_RE.test(line)));
const location = unique(lines.filter((line) => LOCATION_RE.test(line)));
const fallback = unique(lines);
const selected: string[] = [];
for (const line of [...important, ...location, ...fallback]) {
if (selected.includes(line)) {
continue;
}
selected.push(line);
if (selected.length >= MAX_LINES) {
break;
}
}
return selected;
}
export function distillToolResult(input: { toolName?: string; content: string }): string | undefined {
const picked = pickSalientLines(input.content);
if (picked.length === 0) {
return undefined;
}
const prefix = `[distilled ${input.toolName ?? "tool"} output]`;
return `${prefix} ${picked.join("; ")}`.slice(0, MAX_SUMMARY_LENGTH);
}

833
src/extension.test.ts Normal file
View File

@@ -0,0 +1,833 @@
import test from "node:test";
import assert from "node:assert/strict";
import contextManagerExtension from "../index.ts";
import { deserializeLatestSnapshot, SNAPSHOT_ENTRY_TYPE, serializeSnapshot, type RuntimeSnapshot } from "./persist.ts";
type EventHandler = (event: any, ctx: any) => Promise<any> | any;
type RegisteredCommand = { description: string; handler: (args: string, ctx: any) => Promise<void> | void };
type SessionEntry =
| {
type: "message";
id: string;
parentId: string | null;
timestamp: string;
message: any;
}
| {
type: "custom";
id: string;
parentId: string | null;
timestamp: string;
customType: string;
data: unknown;
}
| {
type: "compaction";
id: string;
parentId: string | null;
timestamp: string;
summary: string;
firstKeptEntryId: string;
tokensBefore: number;
}
| {
type: "branch_summary";
id: string;
parentId: string | null;
timestamp: string;
fromId: string;
summary: string;
};
function createUsage(tokens: number) {
return {
tokens,
contextWindow: 200_000,
percent: tokens / 200_000,
};
}
function createUserMessage(content: string, timestamp: number) {
return {
role: "user",
content,
timestamp,
};
}
function createAssistantMessage(content: string, timestamp: number) {
return {
role: "assistant",
content: [{ type: "text", text: content }],
api: "openai-responses",
provider: "openai",
model: "gpt-5",
usage: {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 0,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop",
timestamp,
};
}
function createToolResultMessage(content: string, timestamp: number) {
return {
role: "toolResult",
toolCallId: `tool-${timestamp}`,
toolName: "read",
content: [{ type: "text", text: content }],
isError: false,
timestamp,
};
}
function createMessageEntry(id: string, parentId: string | null, message: any): SessionEntry {
return {
type: "message",
id,
parentId,
timestamp: new Date(message.timestamp).toISOString(),
message,
};
}
function createSnapshotEntry(
id: string,
parentId: string | null,
options: {
text: string;
mode?: RuntimeSnapshot["mode"];
lastZone?: RuntimeSnapshot["lastZone"];
lastObservedTokens?: number;
lastCompactionSummary?: string;
lastBranchSummary?: string;
ledgerItems?: RuntimeSnapshot["ledger"]["items"];
rollingSummary?: string;
},
): SessionEntry {
const {
text,
mode = "aggressive",
lastZone = "red",
lastObservedTokens = 150_000,
lastCompactionSummary = "existing compaction summary",
lastBranchSummary = "existing branch summary",
ledgerItems,
rollingSummary = "stale ledger",
} = options;
return {
type: "custom",
id,
parentId,
timestamp: new Date(1).toISOString(),
customType: SNAPSHOT_ENTRY_TYPE,
data: serializeSnapshot({
mode,
lastZone,
lastObservedTokens,
lastCompactionSummary,
lastBranchSummary,
ledger: {
items: ledgerItems ?? [
{
id: `goal:session:root-goal:${id}`,
kind: "goal",
subject: "root-goal",
text,
scope: "session",
sourceEntryId: "old-user",
sourceType: "user",
timestamp: 1,
confidence: 1,
freshness: 1,
active: true,
supersedesId: undefined,
},
],
rollingSummary,
},
}),
};
}
function createHarness(initialBranch: SessionEntry[], options?: { usageTokens?: number }) {
const commands = new Map<string, RegisteredCommand>();
const handlers = new Map<string, EventHandler>();
const appendedEntries: Array<{ customType: string; data: unknown }> = [];
const statuses: Array<{ key: string; value: string }> = [];
let branch = [...initialBranch];
let entries = [...initialBranch];
const ctx = {
model: { contextWindow: 200_000 },
sessionManager: {
getBranch() {
return branch;
},
getEntries() {
return entries;
},
},
ui: {
setStatus(key: string, value: string) {
statuses.push({ key, value });
},
notify() {},
editor: async () => {},
},
getContextUsage() {
return options?.usageTokens === undefined ? undefined : createUsage(options.usageTokens);
},
compact() {},
};
contextManagerExtension({
registerCommand(name: string, command: RegisteredCommand) {
commands.set(name, command);
},
on(name: string, handler: EventHandler) {
handlers.set(name, handler);
},
appendEntry(customType: string, data: unknown) {
appendedEntries.push({ customType, data });
const entry = {
type: "custom" as const,
id: `custom-${appendedEntries.length}`,
parentId: branch.at(-1)?.id ?? null,
timestamp: new Date(10_000 + appendedEntries.length).toISOString(),
customType,
data,
};
branch.push(entry);
entries.push(entry);
},
} as any);
return {
commands,
handlers,
appendedEntries,
statuses,
ctx,
setBranch(nextBranch: SessionEntry[]) {
branch = [...nextBranch];
const byId = new Map(entries.map((entry) => [entry.id, entry]));
for (const entry of nextBranch) {
byId.set(entry.id, entry);
}
entries = [...byId.values()];
},
};
}
test("the extension registers the expected hooks and commands", () => {
const harness = createHarness([]);
assert.deepEqual([...harness.commands.keys()].sort(), ["ctx-compact", "ctx-memory", "ctx-mode", "ctx-refresh", "ctx-status"]);
assert.deepEqual(
[...harness.handlers.keys()].sort(),
["context", "session_before_compact", "session_before_tree", "session_compact", "session_start", "session_tree", "tool_result", "turn_end"],
);
});
test("turn_end persists a rebuilt snapshot that includes branch user and assistant facts", async () => {
const branch: SessionEntry[] = [
createSnapshotEntry("snapshot-1", null, { text: "Stale snapshot fact" }),
createMessageEntry("user-1", "snapshot-1", createUserMessage("Goal: Fix Task 6\nPrefer keeping the public API stable", 2)),
createMessageEntry(
"assistant-1",
"user-1",
createAssistantMessage("Decision: rebuild from ctx.sessionManager.getBranch()\nNext: add integration tests", 3),
),
createMessageEntry(
"tool-1",
"assistant-1",
createToolResultMessage("Opened .pi/agent/extensions/context-manager/index.ts", 4),
),
];
const harness = createHarness(branch, { usageTokens: 120_000 });
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
await harness.handlers.get("turn_end")?.(
{
type: "turn_end",
turnIndex: 1,
message: createAssistantMessage("done", 5),
toolResults: [],
},
harness.ctx,
);
assert.equal(harness.appendedEntries.length, 1);
assert.equal(harness.appendedEntries[0]?.customType, SNAPSHOT_ENTRY_TYPE);
const snapshot = harness.appendedEntries[0]!.data as any;
const activeTexts = snapshot.ledger.items.filter((item: any) => item.active).map((item: any) => item.text);
assert.equal(snapshot.mode, "aggressive");
assert.equal(snapshot.lastCompactionSummary, "existing compaction summary");
assert.equal(snapshot.lastBranchSummary, "existing branch summary");
assert.equal(snapshot.lastObservedTokens, 120_000);
assert.equal(snapshot.lastZone, "yellow");
assert.deepEqual(activeTexts, ["Stale snapshot fact", "Fix Task 6", "Prefer keeping the public API stable", "rebuild from ctx.sessionManager.getBranch()", "add integration tests", ".pi/agent/extensions/context-manager/index.ts"]);
assert.deepEqual(harness.statuses.at(-1), { key: "context-manager", value: "ctx yellow" });
});
test("session_tree rebuilds runtime from snapshot-only branches before injecting the next packet", async () => {
const oldBranch: SessionEntry[] = [createSnapshotEntry("snapshot-old", null, { text: "Old branch goal" })];
const newBranch: SessionEntry[] = [
createSnapshotEntry("snapshot-new", null, {
text: "Snapshot-only branch goal",
ledgerItems: [
{
id: "goal:session:root-goal:snapshot-new",
kind: "goal",
subject: "root-goal",
text: "Snapshot-only branch goal",
scope: "session",
sourceEntryId: "snapshot-new",
sourceType: "user",
timestamp: 11,
confidence: 1,
freshness: 11,
active: true,
supersedesId: undefined,
},
{
id: "decision:branch:branch-decision:snapshot-new",
kind: "decision",
subject: "branch-decision",
text: "Use the snapshot-backed branch state immediately",
scope: "branch",
sourceEntryId: "snapshot-new",
sourceType: "assistant",
timestamp: 12,
confidence: 0.9,
freshness: 12,
active: true,
supersedesId: undefined,
},
],
rollingSummary: "snapshot-only branch state",
}),
];
const harness = createHarness(oldBranch);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
harness.setBranch(newBranch);
await harness.handlers.get("session_tree")?.(
{
type: "session_tree",
oldLeafId: "snapshot-old",
newLeafId: "snapshot-new",
},
harness.ctx,
);
const result = await harness.handlers.get("context")?.(
{
type: "context",
messages: [createUserMessage("What should happen next?", 13)],
},
harness.ctx,
);
assert.ok(result);
assert.equal(result.messages[0]?.role, "custom");
assert.equal(result.messages[0]?.customType, "context-manager.resume");
assert.match(result.messages[0]?.content, /Snapshot-only branch goal/);
assert.match(result.messages[0]?.content, /Use the snapshot-backed branch state immediately/);
assert.doesNotMatch(result.messages[0]?.content, /Old branch goal/);
assert.deepEqual(harness.statuses.at(-1), { key: "context-manager", value: "ctx red" });
});
test("context keeps a distilled stale tool result visible after pruning bulky output", async () => {
const bulkyFailure = [
"Build failed while compiling focus parser",
"Error: missing export createFocusMatcher from ./summary-focus.ts",
...Array.from({ length: 220 }, () => "stack frame"),
].join("\n");
const harness = createHarness([]);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
const result = await harness.handlers.get("context")?.(
{
type: "context",
messages: [
createUserMessage("turn 1", 1),
createToolResultMessage(bulkyFailure, 2),
createAssistantMessage("observed turn 1", 3),
createUserMessage("turn 2", 4),
createAssistantMessage("observed turn 2", 5),
createUserMessage("turn 3", 6),
createAssistantMessage("observed turn 3", 7),
createUserMessage("turn 4", 8),
createAssistantMessage("observed turn 4", 9),
createUserMessage("turn 5", 10),
],
},
harness.ctx,
);
const toolResult = result.messages.find((message: any) => message.role === "toolResult");
assert.ok(toolResult);
assert.match(toolResult.content[0].text, /missing export createFocusMatcher/);
assert.ok(toolResult.content[0].text.length < 320);
});
test("session_tree preserves session-scoped facts but drops stale branch handoff metadata on an empty destination branch", async () => {
const sourceBranch: SessionEntry[] = [
createSnapshotEntry("snapshot-session", null, {
text: "Ship the context manager extension",
mode: "balanced",
lastZone: "yellow",
lastObservedTokens: 120_000,
lastCompactionSummary: "## Key Decisions\n- Keep summaries deterministic.",
lastBranchSummary: "# Handoff for branch\n\n## Key Decisions\n- Do not leak branch-local goals.",
}),
];
const harness = createHarness(sourceBranch);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
harness.setBranch([]);
await harness.handlers.get("session_tree")?.(
{
type: "session_tree",
oldLeafId: "snapshot-session",
newLeafId: null,
},
harness.ctx,
);
const result = await harness.handlers.get("context")?.(
{ type: "context", messages: [createUserMessage("continue", 30)] },
harness.ctx,
);
assert.equal(result.messages[0]?.customType, "context-manager.packet");
assert.match(result.messages[0]?.content, /Ship the context manager extension/);
assert.doesNotMatch(result.messages[0]?.content, /Do not leak branch-local goals/);
assert.doesNotMatch(result.messages[0]?.content, /Keep summaries deterministic/);
});
test("session_tree overlays newer session-scoped facts onto a destination branch with an older snapshot", async () => {
const newerSessionSnapshot = createSnapshotEntry("snapshot-newer", null, {
text: "Ship the context manager extension",
ledgerItems: [
{
id: "goal:session:root-goal:snapshot-newer",
kind: "goal",
subject: "root-goal",
text: "Ship the context manager extension",
scope: "session",
sourceEntryId: "snapshot-newer",
sourceType: "user",
timestamp: 1,
confidence: 1,
freshness: 1,
active: true,
supersedesId: undefined,
},
{
id: "constraint:session:must-session-newer:2",
kind: "constraint",
subject: "must-session-newer",
text: "Prefer concise reports across the whole session.",
scope: "session",
sourceEntryId: "snapshot-newer",
sourceType: "user",
timestamp: 2,
confidence: 0.9,
freshness: 2,
active: true,
supersedesId: undefined,
},
],
lastCompactionSummary: "",
lastBranchSummary: "",
});
const olderBranchSnapshot = createSnapshotEntry("snapshot-older", null, {
text: "Ship the context manager extension",
ledgerItems: [
{
id: "goal:session:root-goal:snapshot-older",
kind: "goal",
subject: "root-goal",
text: "Ship the context manager extension",
scope: "session",
sourceEntryId: "snapshot-older",
sourceType: "user",
timestamp: 1,
confidence: 1,
freshness: 1,
active: true,
supersedesId: undefined,
},
],
lastCompactionSummary: "",
lastBranchSummary: "",
});
const harness = createHarness([newerSessionSnapshot]);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
harness.setBranch([olderBranchSnapshot]);
await harness.handlers.get("session_tree")?.(
{
type: "session_tree",
oldLeafId: "snapshot-newer",
newLeafId: "snapshot-older",
},
harness.ctx,
);
const result = await harness.handlers.get("context")?.(
{ type: "context", messages: [createUserMessage("continue", 32)] },
harness.ctx,
);
assert.match(result.messages[0]?.content, /Prefer concise reports across the whole session/);
});
test("ctx-refresh preserves session memory without leaking old handoff summaries on a snapshot-less branch", async () => {
const sourceBranch: SessionEntry[] = [
createSnapshotEntry("snapshot-refresh", null, {
text: "Ship the context manager extension",
mode: "balanced",
lastZone: "yellow",
lastObservedTokens: 120_000,
lastCompactionSummary: "## Key Decisions\n- Keep summaries deterministic.",
lastBranchSummary: "# Handoff for branch\n\n## Key Decisions\n- Do not leak branch-local goals.",
}),
];
const harness = createHarness(sourceBranch);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
harness.setBranch([]);
await harness.commands.get("ctx-refresh")?.handler("", harness.ctx);
const result = await harness.handlers.get("context")?.(
{ type: "context", messages: [createUserMessage("continue", 31)] },
harness.ctx,
);
assert.equal(result.messages[0]?.customType, "context-manager.packet");
assert.match(result.messages[0]?.content, /Ship the context manager extension/);
assert.doesNotMatch(result.messages[0]?.content, /Do not leak branch-local goals/);
assert.doesNotMatch(result.messages[0]?.content, /Keep summaries deterministic/);
});
test("session_start replays default pi compaction blockers into resume state", async () => {
const branch: SessionEntry[] = [
createSnapshotEntry("snapshot-default", null, {
text: "Ship the context manager extension",
lastCompactionSummary: undefined,
lastBranchSummary: undefined,
}),
{
type: "compaction",
id: "compaction-default-1",
parentId: "snapshot-default",
timestamp: new Date(40).toISOString(),
summary: [
"## Progress",
"### Blocked",
"- confirm whether /tree replaceInstructions should override defaults",
].join("\n"),
firstKeptEntryId: "snapshot-default",
tokensBefore: 123_000,
},
];
const harness = createHarness(branch);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
const result = await harness.handlers.get("context")?.(
{ type: "context", messages: [createUserMessage("continue", 41)] },
harness.ctx,
);
assert.match(result.messages[0]?.content, /confirm whether \/tree replaceInstructions should override defaults/);
});
test("session_before_compact honors preparation inputs and custom focus", async () => {
const harness = createHarness([]);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
const result = await harness.handlers.get("session_before_compact")?.(
{
type: "session_before_compact",
customInstructions: "Focus on decisions and relevant files.",
preparation: {
messagesToSummarize: [createUserMessage("Decision: keep compaction summaries deterministic", 1)],
turnPrefixMessages: [createToolResultMessage("Opened .pi/agent/extensions/context-manager/index.ts", 2)],
previousSummary: "## Goal\n- Ship the context manager extension",
fileOps: {
readFiles: [".pi/agent/extensions/context-manager/index.ts"],
modifiedFiles: [".pi/agent/extensions/context-manager/src/summaries.ts"],
},
tokensBefore: 120_000,
firstKeptEntryId: "keep-1",
},
branchEntries: [],
signal: AbortSignal.timeout(1_000),
},
harness.ctx,
);
assert.equal(result.compaction.firstKeptEntryId, "keep-1");
assert.equal(result.compaction.tokensBefore, 120_000);
assert.match(result.compaction.summary, /keep compaction summaries deterministic/);
assert.match(result.compaction.summary, /index.ts/);
});
test("session_before_tree honors abandoned-branch entries and focus text", async () => {
const harness = createHarness([]);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
const result = await harness.handlers.get("session_before_tree")?.(
{
type: "session_before_tree",
preparation: {
targetId: "target-1",
oldLeafId: "old-1",
commonAncestorId: "root",
userWantsSummary: true,
customInstructions: "Focus on goals and decisions.",
replaceInstructions: false,
entriesToSummarize: [
createMessageEntry("user-1", null, createUserMessage("Goal: explore tree handoff", 1)),
createMessageEntry("assistant-1", "user-1", createAssistantMessage("Decision: do not leak branch-local goals", 2)),
],
},
signal: AbortSignal.timeout(1_000),
},
harness.ctx,
);
assert.ok(result?.summary?.summary);
assert.match(result.summary.summary, /explore tree handoff/);
assert.match(result.summary.summary, /do not leak branch-local goals/);
});
test("session_compact persists the latest compaction summary into a fresh snapshot and injects a resume packet once", async () => {
const harness = createHarness([]);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
await harness.handlers.get("session_compact")?.(
{
type: "session_compact",
fromExtension: true,
compactionEntry: {
type: "compaction",
id: "cmp-1",
parentId: "prev",
timestamp: new Date(10).toISOString(),
summary: "## Key Decisions\n- Keep summaries deterministic.\n\n## Open questions and blockers\n- Verify /tree replaceInstructions behavior.",
firstKeptEntryId: "keep-1",
tokensBefore: 140_000,
},
},
harness.ctx,
);
assert.equal(harness.appendedEntries.at(-1)?.customType, SNAPSHOT_ENTRY_TYPE);
const context = await harness.handlers.get("context")?.(
{ type: "context", messages: [createUserMessage("continue", 11)] },
harness.ctx,
);
assert.match(context.messages[0]?.content, /Keep summaries deterministic/);
assert.match(context.messages[0]?.content, /Verify \/tree replaceInstructions behavior/);
const nextContext = await harness.handlers.get("context")?.(
{ type: "context", messages: [createUserMessage("continue again", 12)] },
harness.ctx,
);
assert.equal(nextContext.messages[0]?.customType, "context-manager.packet");
assert.doesNotMatch(nextContext.messages[0]?.content ?? "", /## Latest compaction handoff/);
});
test("session_tree replays branch summaries newer than the latest snapshot before the next packet is injected", async () => {
const branch: SessionEntry[] = [
createSnapshotEntry("snapshot-1", null, { text: "Ship the context manager extension" }),
{
type: "branch_summary",
id: "branch-summary-1",
parentId: "snapshot-1",
timestamp: new Date(20).toISOString(),
fromId: "old-leaf",
summary: "# Handoff for branch\n\n## Key Decisions\n- Do not leak branch-local goals.",
},
];
const harness = createHarness(branch);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
const context = await harness.handlers.get("context")?.(
{ type: "context", messages: [createUserMessage("what next", 21)] },
harness.ctx,
);
assert.match(context.messages[0]?.content, /Do not leak branch-local goals/);
});
test("session_tree records event summaryEntry before persisting the next snapshot", async () => {
const branch: SessionEntry[] = [
createSnapshotEntry("snapshot-1", null, {
text: "Ship the context manager extension",
lastCompactionSummary: "",
lastBranchSummary: "",
}),
];
const harness = createHarness(branch);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
await harness.handlers.get("session_tree")?.(
{
type: "session_tree",
fromExtension: true,
summaryEntry: {
type: "branch_summary",
id: "branch-summary-event",
parentId: "snapshot-1",
timestamp: new Date(20).toISOString(),
fromId: "old-leaf",
summary: "# Handoff for branch\n\n## Key Decisions\n- Preserve the latest branch summary from the event payload.",
},
},
harness.ctx,
);
const snapshot = harness.appendedEntries.at(-1)?.data as RuntimeSnapshot | undefined;
assert.match(snapshot?.lastBranchSummary ?? "", /Preserve the latest branch summary/);
const context = await harness.handlers.get("context")?.(
{ type: "context", messages: [createUserMessage("what changed", 21)] },
harness.ctx,
);
assert.match(context.messages[0]?.content, /Preserve the latest branch summary/);
});
test("ctx-status reports mode, zone, packet size, and summary-artifact presence", async () => {
const branch = [
createSnapshotEntry("snapshot-1", null, {
text: "Ship the context manager extension",
mode: "balanced",
lastZone: "yellow",
lastObservedTokens: 120_000,
lastCompactionSummary: "## Key Decisions\n- Keep summaries deterministic.",
lastBranchSummary: "# Handoff for branch\n\n## Key Decisions\n- Do not leak branch-local goals.",
}),
];
const notifications: string[] = [];
const harness = createHarness(branch);
harness.ctx.ui.notify = (message: string) => notifications.push(message);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
await harness.commands.get("ctx-status")?.handler("", harness.ctx);
assert.match(notifications.at(-1) ?? "", /mode=balanced/);
assert.match(notifications.at(-1) ?? "", /zone=yellow/);
assert.match(notifications.at(-1) ?? "", /compaction=yes/);
assert.match(notifications.at(-1) ?? "", /branch=yes/);
});
test("ctx-refresh rebuilds runtime from the current branch instead of only re-rendering the packet", async () => {
const harness = createHarness([createSnapshotEntry("snapshot-1", null, { text: "Old goal" })]);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
harness.setBranch([
createSnapshotEntry("snapshot-2", null, {
text: "New branch goal",
lastBranchSummary: "# Handoff for branch\n\n## Key Decisions\n- Use the new branch immediately.",
}),
]);
await harness.commands.get("ctx-refresh")?.handler("", harness.ctx);
const result = await harness.handlers.get("context")?.(
{ type: "context", messages: [createUserMessage("continue", 3)] },
harness.ctx,
);
assert.match(result.messages[0]?.content, /New branch goal/);
assert.doesNotMatch(result.messages[0]?.content, /Old goal/);
});
test("ctx-mode persists the updated mode immediately without waiting for turn_end", async () => {
const branch: SessionEntry[] = [
createSnapshotEntry("snapshot-1", null, {
text: "Persist the updated mode",
mode: "balanced",
lastZone: "green",
lastObservedTokens: 90_000,
}),
];
const harness = createHarness(branch);
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
assert.equal(deserializeLatestSnapshot(harness.ctx.sessionManager.getBranch())?.mode, "balanced");
const modeCommand = harness.commands.get("ctx-mode");
assert.ok(modeCommand);
await modeCommand.handler("aggressive", harness.ctx);
assert.equal(harness.appendedEntries.length, 1);
assert.equal(harness.appendedEntries[0]?.customType, SNAPSHOT_ENTRY_TYPE);
assert.equal(deserializeLatestSnapshot(harness.ctx.sessionManager.getBranch())?.mode, "aggressive");
});
test("ctx-mode changes survive turn_end and persist into the next snapshot", async () => {
const branch: SessionEntry[] = [
createSnapshotEntry("snapshot-1", null, {
text: "Persist the updated mode",
mode: "balanced",
lastZone: "green",
lastObservedTokens: 90_000,
}),
];
const harness = createHarness(branch, { usageTokens: 105_000 });
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
const modeCommand = harness.commands.get("ctx-mode");
assert.ok(modeCommand);
await modeCommand.handler("aggressive", harness.ctx);
await harness.handlers.get("turn_end")?.(
{
type: "turn_end",
turnIndex: 1,
message: createAssistantMessage("done", 5),
toolResults: [],
},
harness.ctx,
);
assert.equal(harness.appendedEntries.length, 2);
const immediateSnapshot = harness.appendedEntries[0]!.data as any;
assert.equal(immediateSnapshot.mode, "aggressive");
assert.equal(immediateSnapshot.lastObservedTokens, 90_000);
assert.equal(immediateSnapshot.lastZone, "green");
const snapshot = harness.appendedEntries[1]!.data as any;
assert.equal(snapshot.mode, "aggressive");
assert.equal(snapshot.lastObservedTokens, 105_000);
assert.equal(snapshot.lastZone, "yellow");
});

280
src/extract.test.ts Normal file
View File

@@ -0,0 +1,280 @@
import test from "node:test";
import assert from "node:assert/strict";
import { extractCandidates } from "./extract.ts";
import { createEmptyLedger, getActiveItems, mergeCandidates } from "./ledger.ts";
test("extractCandidates pulls goals, constraints, decisions, next steps, and file references", () => {
const candidates = extractCandidates({
entryId: "u1",
role: "user",
text: [
"Goal: Build a context manager extension for pi.",
"We must adapt to the active model context window.",
"Decision: keep the MVP quiet and avoid new LLM-facing tools.",
"Next: inspect .pi/agent/extensions/web-search/index.ts and docs/extensions.md.",
].join("\n"),
timestamp: 1,
});
assert.deepEqual(
candidates.map((candidate) => [
candidate.kind,
candidate.subject,
candidate.scope,
candidate.sourceEntryId,
candidate.sourceType,
candidate.timestamp,
]),
[
["goal", "root-goal", "session", "u1", "user", 1],
["constraint", "must-u1-0", "branch", "u1", "user", 1],
["decision", "decision-u1-0", "branch", "u1", "user", 1],
["activeTask", "next-step-u1-0", "branch", "u1", "user", 1],
["relevantFile", ".pi/agent/extensions/web-search/index.ts", "branch", "u1", "user", 1],
["relevantFile", "docs/extensions.md", "branch", "u1", "user", 1],
],
);
});
test("extractCandidates promotes only the first durable goal to session scope", () => {
const firstGoal = extractCandidates(
{
entryId: "u-goal-1",
role: "user",
text: "Goal: Ship the context manager extension.",
timestamp: 10,
},
{ hasSessionGoal: false },
);
const branchGoal = extractCandidates(
{
entryId: "u-goal-2",
role: "user",
text: "Goal: prototype a branch-local tree handoff.",
timestamp: 11,
},
{ hasSessionGoal: true },
);
assert.deepEqual(
firstGoal.map((candidate) => [candidate.kind, candidate.subject, candidate.scope, candidate.text]),
[["goal", "root-goal", "session", "Ship the context manager extension."]],
);
assert.deepEqual(
branchGoal.map((candidate) => [candidate.kind, candidate.subject, candidate.scope, candidate.text]),
[["goal", "goal-u-goal-2-0", "branch", "prototype a branch-local tree handoff."]],
);
});
test("mergeCandidates keeps independently extracted decisions, constraints, and next steps active", () => {
const ledger = mergeCandidates(createEmptyLedger(), [
...extractCandidates({
entryId: "u1",
role: "user",
text: [
"We must adapt to the active model context window.",
"Decision: keep snapshots tiny.",
"Next: inspect src/extract.ts.",
].join("\n"),
timestamp: 1,
}),
...extractCandidates({
entryId: "u2",
role: "user",
text: [
"We prefer concise reports across the whole session.",
"Decision: persist snapshots after each turn_end.",
"Task: add regression coverage.",
].join("\n"),
timestamp: 2,
}),
]);
assert.deepEqual(
getActiveItems(ledger, "constraint").map((item) => [item.subject, item.sourceEntryId, item.text]),
[
["must-u1-0", "u1", "We must adapt to the active model context window."],
["must-u2-0", "u2", "We prefer concise reports across the whole session."],
],
);
assert.deepEqual(
getActiveItems(ledger, "decision").map((item) => [item.subject, item.sourceEntryId, item.text]),
[
["decision-u1-0", "u1", "keep snapshots tiny."],
["decision-u2-0", "u2", "persist snapshots after each turn_end."],
],
);
assert.deepEqual(
getActiveItems(ledger, "activeTask").map((item) => [item.subject, item.sourceEntryId, item.text]),
[
["next-step-u1-0", "u1", "inspect src/extract.ts."],
["next-step-u2-0", "u2", "add regression coverage."],
],
);
});
test("user constraints default to branch scope unless they explicitly signal durable session scope", () => {
const candidates = extractCandidates({
entryId: "u3",
role: "user",
text: [
"We should keep this branch experimental for now.",
"We should keep the MVP branch experimental.",
"We should rename the context window helper in this module.",
"Avoid touching docs/extensions.md.",
"Avoid touching docs/extensions.md across the whole session.",
"Prefer concise reports across the whole session.",
].join("\n"),
timestamp: 3,
});
assert.deepEqual(
candidates
.filter((candidate) => candidate.kind === "constraint")
.map((candidate) => [candidate.text, candidate.scope, candidate.subject]),
[
["We should keep this branch experimental for now.", "branch", "must-u3-0"],
["We should keep the MVP branch experimental.", "branch", "must-u3-1"],
["We should rename the context window helper in this module.", "branch", "must-u3-2"],
["Avoid touching docs/extensions.md.", "branch", "must-u3-3"],
["Avoid touching docs/extensions.md across the whole session.", "session", "must-u3-4"],
["Prefer concise reports across the whole session.", "session", "must-u3-5"],
],
);
});
test("extractCandidates treats spelled-out do not as a constraint trigger", () => {
const candidates = extractCandidates({
entryId: "u4",
role: "user",
text: "Do not add new LLM-facing tools across the whole session.",
timestamp: 4,
});
assert.deepEqual(
candidates.map((candidate) => [candidate.kind, candidate.text, candidate.scope, candidate.subject]),
[["constraint", "Do not add new LLM-facing tools across the whole session.", "session", "must-u4-0"]],
);
});
test("extractCandidates keeps compaction goals branch-scoped unless they are explicitly session-wide", () => {
const candidates = extractCandidates(
{
entryId: "cmp-goal-1",
role: "compaction",
text: "## Goal\n- prototype a branch-local tree handoff.",
timestamp: 19,
},
{ hasSessionGoal: false },
);
assert.deepEqual(
candidates.map((candidate) => [candidate.kind, candidate.subject, candidate.scope, candidate.text]),
[["goal", "goal-cmp-goal-1-0", "branch", "prototype a branch-local tree handoff."]],
);
});
test("extractCandidates captures blockers from direct lines, tool errors, and structured summaries", () => {
const direct = extractCandidates(
{
entryId: "a-blocked-1",
role: "assistant",
text: "Blocked: confirm whether /tree summaries should replace instructions.",
timestamp: 20,
},
{ hasSessionGoal: true },
);
const tool = extractCandidates(
{
entryId: "t-blocked-1",
role: "toolResult",
text: "Error: missing export createFocusMatcher\nstack...",
timestamp: 21,
isError: true,
},
{ hasSessionGoal: true },
);
const summary = extractCandidates(
{
entryId: "cmp-1",
role: "compaction",
text: [
"## Open questions and blockers",
"- Need to confirm whether /tree summaries should replace instructions.",
"",
"## Relevant files",
"- .pi/agent/extensions/context-manager/index.ts",
].join("\n"),
timestamp: 22,
},
{ hasSessionGoal: true },
);
assert.deepEqual(
direct.map((candidate) => [candidate.kind, candidate.subject, candidate.text]),
[["openQuestion", "open-question-a-blocked-1-0", "confirm whether /tree summaries should replace instructions."]],
);
assert.equal(tool[0]?.kind, "openQuestion");
assert.match(tool[0]?.text ?? "", /missing export createFocusMatcher/);
assert.deepEqual(
summary.map((candidate) => [candidate.kind, candidate.text]),
[
["openQuestion", "Need to confirm whether /tree summaries should replace instructions."],
["relevantFile", ".pi/agent/extensions/context-manager/index.ts"],
],
);
});
test("extractCandidates parses pi fallback progress and blocked summary sections", () => {
const candidates = extractCandidates(
{
entryId: "cmp-default-1",
role: "compaction",
text: [
"## Constraints and preferences",
"- Keep the public API stable.",
"",
"## Progress",
"### In Progress",
"- Wire runtime hydration.",
"",
"### Blocked",
"- confirm whether /tree replaceInstructions should override defaults",
].join("\n"),
timestamp: 23,
},
{ hasSessionGoal: true },
);
assert.ok(candidates.some((candidate) => candidate.kind === "constraint" && candidate.text === "Keep the public API stable."));
assert.ok(candidates.some((candidate) => candidate.kind === "activeTask" && candidate.text === "Wire runtime hydration."));
assert.ok(
candidates.some(
(candidate) => candidate.kind === "openQuestion" && candidate.text === "confirm whether /tree replaceInstructions should override defaults",
),
);
});
test("assistant decisions and tool-result file references are extracted as branch facts", () => {
const assistant = extractCandidates({
entryId: "a1",
role: "assistant",
text: "Decision: persist snapshots after each turn_end.",
timestamp: 2,
});
const tool = extractCandidates({
entryId: "t1",
role: "toolResult",
text: "Updated file: .pi/agent/extensions/context-manager/src/runtime.ts",
timestamp: 3,
});
assert.equal(assistant[0]?.kind, "decision");
assert.equal(assistant[0]?.subject, "decision-a1-0");
assert.equal(assistant[0]?.scope, "branch");
assert.equal(tool[0]?.kind, "relevantFile");
});

314
src/extract.ts Normal file
View File

@@ -0,0 +1,314 @@
import type { MemoryCandidate, MemoryScope, MemorySourceType } from "./ledger.ts";
export interface ExtractOptions {
hasSessionGoal?: boolean;
}
export interface TranscriptSlice {
entryId: string;
role: "user" | "assistant" | "toolResult" | "compaction" | "branchSummary";
text: string;
timestamp: number;
isError?: boolean;
}
const FILE_RE = /(?:\.?\/?[A-Za-z0-9_./-]+\.(?:ts|tsx|js|mjs|json|md))/g;
const BRANCH_LOCAL_CONSTRAINT_RE =
/\b(?:this|current)\s+(?:branch|task|change|step|file|module|test|command|implementation|worktree)\b|\b(?:for now|right now|in this branch|on this branch|in this file|in this module|here)\b/i;
const DURABLE_SESSION_CONSTRAINT_RE =
/\b(?:whole|entire|rest of (?:the )?|remaining)\s+(?:session|project|codebase)\b|\bacross (?:the )?(?:whole )?(?:session|project|codebase)\b|\bacross (?:all |every )?branches\b|\b(?:session|project|codebase)[-\s]?wide\b|\bthroughout (?:the )?(?:session|project|codebase)\b/i;
const CONSTRAINT_RE = /\b(?:must|should|don't|do not|avoid|prefer)\b/i;
const GOAL_RE = /^(goal|session goal|overall goal):/i;
const OPEN_QUESTION_RE = /^(?:blocked|blocker|open question|question):/i;
const ERROR_LINE_RE = /\b(?:error|failed|failure|missing|undefined|exception)\b/i;
type SummarySectionKind = "goal" | "constraint" | "decision" | "activeTask" | "openQuestion" | "relevantFile";
function sourceTypeForRole(role: TranscriptSlice["role"]): MemorySourceType {
if (role === "compaction") return "compaction";
if (role === "branchSummary") return "branchSummary";
return role;
}
function pushCandidate(
list: MemoryCandidate[],
candidate: Omit<MemoryCandidate, "sourceEntryId" | "sourceType" | "timestamp">,
slice: TranscriptSlice,
) {
list.push({
...candidate,
sourceEntryId: slice.entryId,
sourceType: sourceTypeForRole(slice.role),
timestamp: slice.timestamp,
});
}
function createIndexedSubject(prefix: string, slice: TranscriptSlice, index: number): string {
return `${prefix}-${slice.entryId}-${index}`;
}
function inferConstraintScope(slice: TranscriptSlice, line: string): MemoryScope {
if (slice.role !== "user") {
return "branch";
}
if (BRANCH_LOCAL_CONSTRAINT_RE.test(line)) {
return "branch";
}
if (DURABLE_SESSION_CONSTRAINT_RE.test(line)) {
return "session";
}
if ((line.match(FILE_RE) ?? []).length > 0) {
return "branch";
}
return "branch";
}
function nextGoalCandidate(
line: string,
slice: TranscriptSlice,
options: ExtractOptions,
index: number,
): Omit<MemoryCandidate, "sourceEntryId" | "sourceType" | "timestamp"> {
const text = line.replace(GOAL_RE, "").trim();
const explicitSessionGoal = /^(session goal|overall goal):/i.test(line);
const canSeedSessionGoal = slice.role === "user";
const shouldPromoteRootGoal = explicitSessionGoal || (!options.hasSessionGoal && canSeedSessionGoal);
if (shouldPromoteRootGoal) {
return {
kind: "goal",
subject: "root-goal",
text,
scope: "session",
confidence: 1,
};
}
return {
kind: "goal",
subject: createIndexedSubject("goal", slice, index),
text,
scope: "branch",
confidence: 0.9,
};
}
function nextOpenQuestionCandidate(
text: string,
slice: TranscriptSlice,
index: number,
): Omit<MemoryCandidate, "sourceEntryId" | "sourceType" | "timestamp"> {
return {
kind: "openQuestion",
subject: createIndexedSubject("open-question", slice, index),
text,
scope: "branch",
confidence: slice.role === "toolResult" ? 0.85 : 0.8,
};
}
function summarySectionToKind(line: string): SummarySectionKind | undefined {
const heading = line.replace(/^##\s+/i, "").trim().toLowerCase();
if (heading === "goal") return "goal";
if (heading === "constraints" || heading === "constraints & preferences" || heading === "constraints and preferences") {
return "constraint";
}
if (heading === "decisions" || heading === "key decisions") return "decision";
if (heading === "active work" || heading === "next steps" || heading === "current task" || heading === "progress") {
return "activeTask";
}
if (heading === "open questions and blockers" || heading === "open questions / blockers") return "openQuestion";
if (heading === "relevant files" || heading === "critical context") return "relevantFile";
return undefined;
}
function pushRelevantFiles(list: MemoryCandidate[], slice: TranscriptSlice, line: string) {
const fileMatches = line.match(FILE_RE) ?? [];
for (const match of fileMatches) {
pushCandidate(
list,
{
kind: "relevantFile",
subject: match,
text: match,
scope: "branch",
confidence: 0.7,
},
slice,
);
}
}
export function extractCandidates(slice: TranscriptSlice, options: ExtractOptions = {}): MemoryCandidate[] {
const out: MemoryCandidate[] = [];
const lines = slice.text
.split(/\n+/)
.map((line) => line.trim())
.filter(Boolean);
let currentSection: SummarySectionKind | undefined;
let goalIndex = 0;
let decisionIndex = 0;
let nextStepIndex = 0;
let mustIndex = 0;
let openQuestionIndex = 0;
for (const line of lines) {
if (/^##\s+/i.test(line)) {
currentSection = summarySectionToKind(line);
continue;
}
if (/^###\s+/i.test(line)) {
const subheading = line.replace(/^###\s+/i, "").trim().toLowerCase();
if (subheading === "blocked") {
currentSection = "openQuestion";
} else if (subheading === "in progress" || subheading === "done") {
currentSection = "activeTask";
}
continue;
}
const bullet = line.match(/^-\s+(.*)$/)?.[1]?.trim();
const isGoal = GOAL_RE.test(line);
const isDecision = /^decision:/i.test(line);
const isNextStep = /^(next|task):/i.test(line);
const isOpenQuestion = OPEN_QUESTION_RE.test(line);
if (isGoal) {
pushCandidate(out, nextGoalCandidate(line, slice, options, goalIndex++), slice);
pushRelevantFiles(out, slice, line);
continue;
}
if (isOpenQuestion) {
pushCandidate(
out,
nextOpenQuestionCandidate(line.replace(OPEN_QUESTION_RE, "").trim(), slice, openQuestionIndex++),
slice,
);
pushRelevantFiles(out, slice, line);
continue;
}
if (isDecision) {
pushCandidate(
out,
{
kind: "decision",
subject: createIndexedSubject("decision", slice, decisionIndex++),
text: line.replace(/^decision:\s*/i, "").trim(),
scope: "branch",
confidence: 0.9,
},
slice,
);
pushRelevantFiles(out, slice, line);
continue;
}
if (isNextStep) {
pushCandidate(
out,
{
kind: "activeTask",
subject: createIndexedSubject("next-step", slice, nextStepIndex++),
text: line.replace(/^(next|task):\s*/i, "").trim(),
scope: "branch",
confidence: 0.8,
},
slice,
);
pushRelevantFiles(out, slice, line);
continue;
}
if (bullet && currentSection === "goal") {
pushCandidate(out, nextGoalCandidate(`Goal: ${bullet}`, slice, options, goalIndex++), slice);
continue;
}
if (bullet && currentSection === "constraint") {
pushCandidate(
out,
{
kind: "constraint",
subject: createIndexedSubject("must", slice, mustIndex++),
text: bullet,
scope: inferConstraintScope(slice, bullet),
confidence: 0.8,
},
slice,
);
continue;
}
if (bullet && currentSection === "decision") {
pushCandidate(
out,
{
kind: "decision",
subject: createIndexedSubject("decision", slice, decisionIndex++),
text: bullet,
scope: "branch",
confidence: 0.9,
},
slice,
);
continue;
}
if (bullet && currentSection === "activeTask") {
pushCandidate(
out,
{
kind: "activeTask",
subject: createIndexedSubject("next-step", slice, nextStepIndex++),
text: bullet,
scope: "branch",
confidence: 0.8,
},
slice,
);
continue;
}
if (bullet && currentSection === "openQuestion") {
pushCandidate(out, nextOpenQuestionCandidate(bullet, slice, openQuestionIndex++), slice);
continue;
}
if (bullet && currentSection === "relevantFile") {
pushRelevantFiles(out, slice, bullet);
continue;
}
if (slice.role === "toolResult" && (slice.isError || ERROR_LINE_RE.test(line))) {
pushCandidate(out, nextOpenQuestionCandidate(line, slice, openQuestionIndex++), slice);
}
if (CONSTRAINT_RE.test(line)) {
pushCandidate(
out,
{
kind: "constraint",
subject: createIndexedSubject("must", slice, mustIndex++),
text: line,
scope: inferConstraintScope(slice, line),
confidence: 0.8,
},
slice,
);
}
pushRelevantFiles(out, slice, line);
}
return out;
}

132
src/ledger.test.ts Normal file
View File

@@ -0,0 +1,132 @@
import test from "node:test";
import assert from "node:assert/strict";
import { createEmptyLedger, getActiveItems, mergeCandidates, type MemoryCandidate } from "./ledger.ts";
const base: Omit<MemoryCandidate, "kind" | "subject" | "text"> = {
scope: "branch",
sourceEntryId: "u1",
sourceType: "user",
timestamp: 1,
confidence: 0.9,
};
test("mergeCandidates adds new active items to an empty ledger", () => {
const ledger = mergeCandidates(createEmptyLedger(), [
{ ...base, kind: "goal", subject: "root-goal", text: "Build a pi context manager extension" },
]);
assert.equal(getActiveItems(ledger).length, 1);
assert.equal(getActiveItems(ledger)[0]?.text, "Build a pi context manager extension");
});
test("mergeCandidates archives older items when a new item supersedes the same subject", () => {
const first = mergeCandidates(createEmptyLedger(), [
{ ...base, kind: "decision", subject: "persistence", text: "Persist snapshots with appendEntry()", timestamp: 1 },
]);
const second = mergeCandidates(first, [
{ ...base, kind: "decision", subject: "persistence", text: "Persist snapshots after each turn_end", timestamp: 2 },
]);
const active = getActiveItems(second, "decision");
assert.equal(active.length, 1);
assert.equal(active[0]?.text, "Persist snapshots after each turn_end");
assert.equal(active[0]?.supersedesId, "decision:branch:persistence:1");
assert.equal(second.items.find((item) => item.id === "decision:branch:persistence:1")?.active, false);
});
test("mergeCandidates keeps the newest item active when same-slot candidates arrive out of order", () => {
const ledger = mergeCandidates(createEmptyLedger(), [
{ ...base, kind: "decision", subject: "persistence", text: "Persist snapshots after each turn_end", timestamp: 2 },
{ ...base, kind: "decision", subject: "persistence", text: "Persist snapshots with appendEntry()", timestamp: 1 },
]);
const active = getActiveItems(ledger, "decision");
const stale = ledger.items.find((item) => item.text === "Persist snapshots with appendEntry()");
assert.equal(active.length, 1);
assert.equal(active[0]?.text, "Persist snapshots after each turn_end");
assert.equal(stale?.active, false);
assert.equal(stale?.supersedesId, undefined);
});
test("mergeCandidates gives same-slot same-timestamp candidates distinct ids", () => {
const ledger = mergeCandidates(createEmptyLedger(), [
{ ...base, kind: "decision", subject: "persistence", text: "Persist snapshots with appendEntry()", timestamp: 1 },
{ ...base, kind: "decision", subject: "persistence", text: "Persist snapshots after each turn_end", timestamp: 1 },
]);
const ids = ledger.items.map((item) => item.id);
const active = getActiveItems(ledger, "decision")[0];
const archived = ledger.items.find((item) => item.text === "Persist snapshots with appendEntry()");
assert.equal(new Set(ids).size, ledger.items.length);
assert.equal(active?.text, "Persist snapshots after each turn_end");
assert.equal(active?.supersedesId, archived?.id);
assert.notEqual(active?.id, active?.supersedesId);
});
test("mergeCandidates keeps same-slot same-timestamp snapshots deterministic regardless of input order", () => {
const appendEntryCandidate = {
...base,
kind: "decision" as const,
subject: "persistence",
text: "Persist snapshots with appendEntry()",
timestamp: 1,
};
const turnEndCandidate = {
...base,
kind: "decision" as const,
subject: "persistence",
text: "Persist snapshots after each turn_end",
timestamp: 1,
};
const forward = mergeCandidates(createEmptyLedger(), [appendEntryCandidate, turnEndCandidate]);
const reversed = mergeCandidates(createEmptyLedger(), [turnEndCandidate, appendEntryCandidate]);
assert.deepEqual(forward, reversed);
assert.deepEqual(forward.items, [
{
...turnEndCandidate,
id: "decision:branch:persistence:1",
freshness: 1,
active: true,
supersedesId: "decision:branch:persistence:1:2",
},
{
...appendEntryCandidate,
id: "decision:branch:persistence:1:2",
freshness: 1,
active: false,
supersedesId: undefined,
},
]);
});
test("session-scoped memory can coexist with branch-scoped memory for the same kind", () => {
const ledger = mergeCandidates(createEmptyLedger(), [
{
kind: "constraint",
subject: "llm-tools",
text: "Do not add new LLM-facing tools in the MVP",
scope: "session",
sourceEntryId: "u1",
sourceType: "user",
timestamp: 1,
confidence: 1,
},
{
kind: "constraint",
subject: "branch-policy",
text: "Keep branch A experimental",
scope: "branch",
sourceEntryId: "u2",
sourceType: "user",
timestamp: 2,
confidence: 0.8,
},
]);
assert.equal(getActiveItems(ledger, "constraint").length, 2);
});

196
src/ledger.ts Normal file
View File

@@ -0,0 +1,196 @@
export type MemoryKind = "goal" | "constraint" | "decision" | "activeTask" | "openQuestion" | "relevantFile";
export type MemoryScope = "branch" | "session";
export type MemorySourceType = "user" | "assistant" | "toolResult" | "compaction" | "branchSummary";
export interface MemoryCandidate {
kind: MemoryKind;
subject: string;
text: string;
scope: MemoryScope;
sourceEntryId: string;
sourceType: MemorySourceType;
timestamp: number;
confidence: number;
}
export interface MemoryItem extends MemoryCandidate {
id: string;
freshness: number;
active: boolean;
supersedesId?: string;
}
export interface LedgerState {
items: MemoryItem[];
rollingSummary: string;
}
type MemorySlot = Pick<MemoryCandidate, "kind" | "scope" | "subject">;
export function createEmptyLedger(): LedgerState {
return { items: [], rollingSummary: "" };
}
function createId(candidate: MemoryCandidate): string {
return `${candidate.kind}:${candidate.scope}:${candidate.subject}:${candidate.timestamp}`;
}
function ensureUniqueId(items: Pick<MemoryItem, "id">[], baseId: string): string {
let id = baseId;
let suffix = 2;
while (items.some((item) => item.id === id)) {
id = `${baseId}:${suffix}`;
suffix += 1;
}
return id;
}
function sameSlot(left: MemorySlot, right: MemorySlot) {
return left.kind === right.kind && left.scope === right.scope && left.subject === right.subject;
}
function createSlotKey(slot: MemorySlot): string {
return `${slot.kind}\u0000${slot.scope}\u0000${slot.subject}`;
}
function compareStrings(left: string, right: string): number {
if (left === right) {
return 0;
}
return left < right ? -1 : 1;
}
function compareSameTimestampCandidates(
left: Pick<MemoryCandidate, "text" | "sourceType" | "sourceEntryId" | "confidence">,
right: Pick<MemoryCandidate, "text" | "sourceType" | "sourceEntryId" | "confidence">
): number {
// Exact-timestamp ties should resolve the same way no matter which candidate is processed first.
const textComparison = compareStrings(left.text, right.text);
if (textComparison !== 0) {
return textComparison;
}
const sourceTypeComparison = compareStrings(left.sourceType, right.sourceType);
if (sourceTypeComparison !== 0) {
return sourceTypeComparison;
}
const sourceEntryIdComparison = compareStrings(left.sourceEntryId, right.sourceEntryId);
if (sourceEntryIdComparison !== 0) {
return sourceEntryIdComparison;
}
if (left.confidence !== right.confidence) {
return left.confidence > right.confidence ? -1 : 1;
}
return 0;
}
function candidateSupersedesPrevious(candidate: MemoryCandidate, previous?: MemoryItem): boolean {
if (!previous) {
return true;
}
if (candidate.timestamp !== previous.timestamp) {
return candidate.timestamp > previous.timestamp;
}
return compareSameTimestampCandidates(candidate, previous) < 0;
}
function compareSlotItems(left: MemoryItem, right: MemoryItem): number {
if (left.timestamp !== right.timestamp) {
return right.timestamp - left.timestamp;
}
return compareSameTimestampCandidates(left, right);
}
function normalizeSlotItems(items: MemoryItem[], slot: MemorySlot): MemoryItem[] {
const slotIndices: number[] = [];
const slotItems: MemoryItem[] = [];
items.forEach((item, index) => {
if (!sameSlot(item, slot)) {
return;
}
slotIndices.push(index);
slotItems.push(item);
});
if (slotItems.length <= 1) {
return items;
}
const sortedSlotItems = [...slotItems].sort(compareSlotItems);
const slotIds = new Map<string, number>();
const sortedSlotItemsWithIds = sortedSlotItems.map((item) => {
const baseId = createId(item);
const nextSlotIdCount = (slotIds.get(baseId) ?? 0) + 1;
slotIds.set(baseId, nextSlotIdCount);
return {
item,
id: nextSlotIdCount === 1 ? baseId : `${baseId}:${nextSlotIdCount}`,
};
});
const normalizedSlotItems = sortedSlotItemsWithIds.map(({ item, id }, index) => ({
...item,
id,
freshness: index === 0 ? item.timestamp : sortedSlotItemsWithIds[index - 1]!.item.timestamp,
active: index === 0,
supersedesId: sortedSlotItemsWithIds[index + 1]?.id,
}));
const normalizedItems = [...items];
slotIndices.forEach((slotIndex, index) => {
normalizedItems[slotIndex] = normalizedSlotItems[index]!;
});
return normalizedItems;
}
export function mergeCandidates(state: LedgerState, candidates: MemoryCandidate[]): LedgerState {
let items = [...state.items];
const affectedSlots = new Map<string, MemorySlot>();
for (const candidate of candidates) {
const previousIndex = items.findIndex((item) => item.active && sameSlot(item, candidate));
const previous = previousIndex === -1 ? undefined : items[previousIndex];
const candidateIsNewest = candidateSupersedesPrevious(candidate, previous);
if (previous && candidateIsNewest) {
items[previousIndex] = { ...previous, active: false, freshness: candidate.timestamp };
}
items.push({
...candidate,
id: ensureUniqueId(items, createId(candidate)),
freshness: candidate.timestamp,
active: candidateIsNewest,
supersedesId: candidateIsNewest ? previous?.id : undefined,
});
affectedSlots.set(createSlotKey(candidate), {
kind: candidate.kind,
scope: candidate.scope,
subject: candidate.subject,
});
}
for (const slot of affectedSlots.values()) {
items = normalizeSlotItems(items, slot);
}
return { ...state, items };
}
export function getActiveItems(state: LedgerState, kind?: MemoryKind): MemoryItem[] {
return state.items.filter((item) => item.active && (kind ? item.kind === kind : true));
}

View File

@@ -0,0 +1,27 @@
import test from "node:test";
import assert from "node:assert/strict";
import { existsSync, readFileSync } from "node:fs";
import { dirname, resolve } from "node:path";
import { fileURLToPath } from "node:url";
const packageRoot = resolve(dirname(fileURLToPath(import.meta.url)), "..");
const pkg = JSON.parse(readFileSync(resolve(packageRoot, "package.json"), "utf8"));
test("package.json exposes pi-context-manager as a standalone pi package", () => {
assert.equal(pkg.name, "pi-context-manager");
assert.equal(pkg.type, "module");
assert.ok(Array.isArray(pkg.keywords));
assert.ok(pkg.keywords.includes("pi-package"));
assert.deepEqual(pkg.pi, {
extensions: ["./index.ts"],
});
assert.equal(pkg.peerDependencies["@mariozechner/pi-agent-core"], "*");
assert.equal(pkg.peerDependencies["@mariozechner/pi-coding-agent"], "*");
assert.deepEqual(pkg.dependencies ?? {}, {});
assert.equal(pkg.bundledDependencies, undefined);
assert.deepEqual(pkg.files, ["index.ts", "src"]);
assert.ok(existsSync(resolve(packageRoot, "index.ts")));
assert.ok(existsSync(resolve(packageRoot, "src/runtime.ts")));
});

130
src/packet.test.ts Normal file
View File

@@ -0,0 +1,130 @@
import test from "node:test";
import assert from "node:assert/strict";
import { resolvePolicy } from "./config.ts";
import { buildContextPacket } from "./packet.ts";
import { createEmptyLedger, mergeCandidates, type MemoryCandidate } from "./ledger.ts";
const baseCandidate: Omit<MemoryCandidate, "kind" | "subject" | "text"> = {
scope: "session",
sourceEntryId: "seed",
sourceType: "user",
timestamp: 1,
confidence: 1,
};
function estimateTokens(text: string) {
return Math.ceil(text.length / 4);
}
function memory(candidate: Pick<MemoryCandidate, "kind" | "subject" | "text"> & Partial<Omit<MemoryCandidate, "kind" | "subject" | "text">>): MemoryCandidate {
return {
...baseCandidate,
...candidate,
sourceEntryId: candidate.sourceEntryId ?? candidate.subject,
};
}
function buildPolicy(packetTokenCap: number) {
return {
...resolvePolicy({ mode: "balanced", contextWindow: 200_000 }),
packetTokenCap,
};
}
test("buildContextPacket keeps top-ranked facts from a section when the cap is tight", () => {
const expected = [
"## Active goal",
"- Keep packets compact.",
"",
"## Constraints",
"- Preserve the highest-priority constraint.",
"",
"## Key decisions",
"- Render selected sections in stable order.",
].join("\n");
const policy = buildPolicy(estimateTokens(expected));
const ledger = mergeCandidates(createEmptyLedger(), [
memory({ kind: "goal", subject: "goal", text: "Keep packets compact." }),
memory({ kind: "constraint", subject: "constraint-a", text: "Preserve the highest-priority constraint.", confidence: 1, timestamp: 3 }),
memory({
kind: "constraint",
subject: "constraint-b",
text: "Avoid dropping every constraint just because one extra bullet is too large for a tight packet cap.",
confidence: 0.6,
timestamp: 2,
}),
memory({
kind: "decision",
subject: "decision-a",
text: "Render selected sections in stable order.",
confidence: 0.9,
timestamp: 4,
sourceType: "assistant",
}),
]);
const packet = buildContextPacket(ledger, policy);
assert.equal(packet.text, expected);
assert.equal(packet.estimatedTokens, policy.packetTokenCap);
});
test("buildContextPacket uses cross-kind weights when only one lower-priority section can fit", () => {
const expected = [
"## Active goal",
"- Keep the agent moving.",
"",
"## Current task",
"- Fix packet trimming.",
].join("\n");
const policy = buildPolicy(estimateTokens(expected));
const ledger = mergeCandidates(createEmptyLedger(), [
memory({ kind: "goal", subject: "goal", text: "Keep the agent moving." }),
memory({
kind: "decision",
subject: "decision-a",
text: "Keep logs concise.",
confidence: 1,
timestamp: 2,
sourceType: "assistant",
}),
memory({
kind: "activeTask",
subject: "task-a",
text: "Fix packet trimming.",
confidence: 1,
timestamp: 2,
sourceType: "assistant",
}),
]);
const packet = buildContextPacket(ledger, policy);
assert.equal(packet.text, expected);
assert.equal(packet.estimatedTokens, policy.packetTokenCap);
});
test("buildContextPacket keeps a goal ahead of newer low-priority facts at realistic timestamp scales", () => {
const expected = [
"## Active goal",
"- Keep the agent on track.",
].join("\n");
const policy = buildPolicy(estimateTokens(expected));
const ledger = mergeCandidates(createEmptyLedger(), [
memory({ kind: "goal", subject: "goal", text: "Keep the agent on track.", timestamp: 1_000_000 }),
memory({
kind: "relevantFile",
subject: "runtime-file",
text: "src/runtime.ts",
timestamp: 10_000_000,
confidence: 1,
sourceType: "assistant",
scope: "branch",
}),
]);
const packet = buildContextPacket(ledger, policy);
assert.equal(packet.text, expected);
assert.equal(packet.estimatedTokens, policy.packetTokenCap);
});

91
src/packet.ts Normal file
View File

@@ -0,0 +1,91 @@
import type { Policy } from "./config.ts";
import { getActiveItems, type LedgerState, type MemoryItem, type MemoryKind } from "./ledger.ts";
const SECTION_ORDER: Array<{ kind: MemoryKind; title: string }> = [
{ kind: "goal", title: "Active goal" },
{ kind: "constraint", title: "Constraints" },
{ kind: "decision", title: "Key decisions" },
{ kind: "activeTask", title: "Current task" },
{ kind: "relevantFile", title: "Relevant files" },
{ kind: "openQuestion", title: "Open questions / blockers" },
];
const WEIGHTS: Record<MemoryKind, number> = {
goal: 100,
constraint: 90,
decision: 80,
activeTask: 85,
relevantFile: 60,
openQuestion: 70,
};
const SECTION_INDEX = new Map(SECTION_ORDER.map((section, index) => [section.kind, index]));
function estimateTokens(text: string): number {
return Math.ceil(text.length / 4);
}
function compareByPriority(left: MemoryItem, right: MemoryItem): number {
const weightDifference = WEIGHTS[right.kind] - WEIGHTS[left.kind];
if (weightDifference !== 0) {
return weightDifference;
}
if (left.confidence !== right.confidence) {
return right.confidence - left.confidence;
}
const sectionDifference = SECTION_INDEX.get(left.kind)! - SECTION_INDEX.get(right.kind)!;
if (sectionDifference !== 0) {
return sectionDifference;
}
if (left.freshness !== right.freshness) {
return right.freshness - left.freshness;
}
return left.id.localeCompare(right.id);
}
function sortByPriority(items: MemoryItem[]) {
return [...items].sort(compareByPriority);
}
function renderPacket(itemsByKind: Map<MemoryKind, MemoryItem[]>, selectedIds: Set<string>) {
const lines: string[] = [];
for (const section of SECTION_ORDER) {
const items = itemsByKind.get(section.kind)?.filter((item) => selectedIds.has(item.id)) ?? [];
if (items.length === 0) continue;
lines.push(`## ${section.title}`, ...items.map((item) => `- ${item.text}`), "");
}
return lines.join("\n").trim();
}
export function buildContextPacket(ledger: LedgerState, policy: Policy): { text: string; estimatedTokens: number } {
const itemsByKind = new Map<MemoryKind, MemoryItem[]>();
for (const section of SECTION_ORDER) {
itemsByKind.set(section.kind, sortByPriority(getActiveItems(ledger, section.kind)));
}
const candidates = sortByPriority(getActiveItems(ledger));
const selectedIds = new Set<string>();
let text = "";
for (const item of candidates) {
const tentativeSelectedIds = new Set(selectedIds);
tentativeSelectedIds.add(item.id);
const tentative = renderPacket(itemsByKind, tentativeSelectedIds);
if (estimateTokens(tentative) > policy.packetTokenCap) {
continue;
}
selectedIds.add(item.id);
text = tentative;
}
return { text, estimatedTokens: estimateTokens(text) };
}

67
src/persist.test.ts Normal file
View File

@@ -0,0 +1,67 @@
import test from "node:test";
import assert from "node:assert/strict";
import { deserializeLatestSnapshot, SNAPSHOT_ENTRY_TYPE, serializeSnapshot } from "./persist.ts";
function createSnapshot(lastZone: "green" | "yellow" | "red" | "compact", lastCompactionSummary: string) {
return serializeSnapshot({
mode: "balanced",
lastZone,
lastCompactionSummary,
lastBranchSummary: undefined,
ledger: {
items: [
{
id: `goal:session:root:${lastCompactionSummary}`,
kind: "goal",
subject: "root",
text: `Goal ${lastCompactionSummary}`,
scope: "session",
sourceEntryId: "u1",
sourceType: "user",
timestamp: 1,
confidence: 1,
freshness: 1,
active: true,
supersedesId: undefined,
},
],
rollingSummary: "summary",
},
});
}
test("deserializeLatestSnapshot restores the newest matching custom entry", () => {
const first = createSnapshot("yellow", "old");
const second = createSnapshot("red", "new");
const restored = deserializeLatestSnapshot([
{ type: "custom", customType: SNAPSHOT_ENTRY_TYPE, data: first },
{ type: "custom", customType: SNAPSHOT_ENTRY_TYPE, data: second },
]);
assert.equal(restored?.lastZone, "red");
assert.equal(restored?.lastCompactionSummary, "new");
});
test("deserializeLatestSnapshot skips malformed newer entries and clones the accepted snapshot", () => {
const valid = createSnapshot("yellow", "valid");
const restored = deserializeLatestSnapshot([
{ type: "custom", customType: SNAPSHOT_ENTRY_TYPE, data: valid },
{
type: "custom",
customType: SNAPSHOT_ENTRY_TYPE,
data: {
mode: "balanced",
lastZone: "red",
ledger: { items: "not-an-array", rollingSummary: "broken" },
},
},
]);
assert.deepEqual(restored, valid);
assert.notStrictEqual(restored, valid);
restored!.ledger.items[0]!.text = "mutated";
assert.equal(valid.ledger.items[0]!.text, "Goal valid");
});

142
src/persist.ts Normal file
View File

@@ -0,0 +1,142 @@
import type { ContextMode, ContextZone } from "./config.ts";
import type { LedgerState, MemoryItem, MemoryKind, MemoryScope, MemorySourceType } from "./ledger.ts";
export const SNAPSHOT_ENTRY_TYPE = "context-manager.snapshot";
export interface RuntimeSnapshot {
mode: ContextMode;
lastZone: ContextZone;
lastObservedTokens?: number;
lastCompactionSummary?: string;
lastBranchSummary?: string;
ledger: LedgerState;
}
const CONTEXT_MODES = new Set<ContextMode>(["conservative", "balanced", "aggressive"]);
const CONTEXT_ZONES = new Set<ContextZone>(["green", "yellow", "red", "compact"]);
const MEMORY_KINDS = new Set<MemoryKind>(["goal", "constraint", "decision", "activeTask", "openQuestion", "relevantFile"]);
const MEMORY_SCOPES = new Set<MemoryScope>(["branch", "session"]);
const MEMORY_SOURCE_TYPES = new Set<MemorySourceType>(["user", "assistant", "toolResult", "compaction", "branchSummary"]);
export function serializeSnapshot(snapshot: RuntimeSnapshot): RuntimeSnapshot {
return structuredClone(snapshot);
}
function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === "object" && value !== null;
}
function isFiniteNumber(value: unknown): value is number {
return typeof value === "number" && Number.isFinite(value);
}
function isOptionalString(value: unknown): value is string | undefined {
return value === undefined || typeof value === "string";
}
function parseMemoryItem(value: unknown): MemoryItem | undefined {
if (!isRecord(value)) {
return undefined;
}
if (
typeof value.id !== "string" ||
!MEMORY_KINDS.has(value.kind as MemoryKind) ||
typeof value.subject !== "string" ||
typeof value.text !== "string" ||
!MEMORY_SCOPES.has(value.scope as MemoryScope) ||
typeof value.sourceEntryId !== "string" ||
!MEMORY_SOURCE_TYPES.has(value.sourceType as MemorySourceType) ||
!isFiniteNumber(value.timestamp) ||
!isFiniteNumber(value.confidence) ||
!isFiniteNumber(value.freshness) ||
typeof value.active !== "boolean" ||
!isOptionalString(value.supersedesId)
) {
return undefined;
}
return {
id: value.id,
kind: value.kind as MemoryKind,
subject: value.subject,
text: value.text,
scope: value.scope as MemoryScope,
sourceEntryId: value.sourceEntryId,
sourceType: value.sourceType as MemorySourceType,
timestamp: value.timestamp,
confidence: value.confidence,
freshness: value.freshness,
active: value.active,
supersedesId: value.supersedesId,
};
}
function parseLedgerState(value: unknown): LedgerState | undefined {
if (!isRecord(value) || !Array.isArray(value.items) || typeof value.rollingSummary !== "string") {
return undefined;
}
const items: MemoryItem[] = [];
for (const item of value.items) {
const parsed = parseMemoryItem(item);
if (!parsed) {
return undefined;
}
items.push(parsed);
}
return {
items,
rollingSummary: value.rollingSummary,
};
}
function parseRuntimeSnapshot(value: unknown): RuntimeSnapshot | undefined {
if (
!isRecord(value) ||
!CONTEXT_MODES.has(value.mode as ContextMode) ||
!CONTEXT_ZONES.has(value.lastZone as ContextZone) ||
!isOptionalString(value.lastCompactionSummary) ||
!isOptionalString(value.lastBranchSummary) ||
(value.lastObservedTokens !== undefined && !isFiniteNumber(value.lastObservedTokens))
) {
return undefined;
}
const ledger = parseLedgerState(value.ledger);
if (!ledger) {
return undefined;
}
const snapshot: RuntimeSnapshot = {
mode: value.mode as ContextMode,
lastZone: value.lastZone as ContextZone,
lastCompactionSummary: value.lastCompactionSummary,
lastBranchSummary: value.lastBranchSummary,
ledger,
};
if (value.lastObservedTokens !== undefined) {
snapshot.lastObservedTokens = value.lastObservedTokens;
}
return snapshot;
}
export function deserializeLatestSnapshot(entries: Array<{ type: string; customType?: string; data?: unknown }>): RuntimeSnapshot | undefined {
for (let index = entries.length - 1; index >= 0; index -= 1) {
const entry = entries[index]!;
if (entry.type !== "custom" || entry.customType !== SNAPSHOT_ENTRY_TYPE) {
continue;
}
const snapshot = parseRuntimeSnapshot(entry.data);
if (snapshot) {
return snapshot;
}
}
return undefined;
}

131
src/prune.test.ts Normal file
View File

@@ -0,0 +1,131 @@
import test from "node:test";
import assert from "node:assert/strict";
import { resolvePolicy } from "./config.ts";
import { pruneContextMessages } from "./prune.ts";
const bulky = "line\n".repeat(300);
const boundaryBulky = "boundary\n".repeat(300);
const thresholdWithTrailingNewline = "threshold\n".repeat(150);
function buildPolicy(recentUserTurns = 4) {
return {
...resolvePolicy({ mode: "balanced", contextWindow: 200_000 }),
recentUserTurns,
};
}
test("pruneContextMessages replaces old bulky tool results with distilled summaries instead of deleting them", () => {
const policy = buildPolicy(2);
const bulkyFailure = [
"Build failed while compiling focus parser",
"Error: missing export createFocusMatcher from ./summary-focus.ts",
...Array.from({ length: 220 }, () => "stack frame"),
].join("\n");
const messages = [
{ role: "user", content: "turn 1" },
{ role: "toolResult", toolName: "bash", content: bulkyFailure },
{ role: "assistant", content: "observed turn 1" },
{ role: "user", content: "turn 2" },
{ role: "assistant", content: "observed turn 2" },
{ role: "user", content: "turn 3" },
];
const pruned = pruneContextMessages(messages, policy);
const distilled = pruned.find((message) => message.role === "toolResult");
assert.ok(distilled);
assert.match(distilled!.content, /missing export createFocusMatcher/);
assert.doesNotMatch(distilled!.content, /stack frame\nstack frame\nstack frame/);
});
test("aggressive mode distills an older bulky tool result sooner than conservative mode", () => {
const conservative = resolvePolicy({ mode: "conservative", contextWindow: 200_000 });
const aggressive = resolvePolicy({ mode: "aggressive", contextWindow: 200_000 });
const messages = [
{ role: "user", content: "turn 1" },
{ role: "toolResult", toolName: "read", content: bulky },
{ role: "assistant", content: "after turn 1" },
{ role: "user", content: "turn 2" },
{ role: "assistant", content: "after turn 2" },
{ role: "user", content: "turn 3" },
{ role: "assistant", content: "after turn 3" },
{ role: "user", content: "turn 4" },
];
const conservativePruned = pruneContextMessages(messages, conservative);
const aggressivePruned = pruneContextMessages(messages, aggressive);
assert.equal(conservativePruned[1]?.content, bulky);
assert.notEqual(aggressivePruned[1]?.content, bulky);
assert.match(aggressivePruned[1]?.content ?? "", /^\[distilled read output\]/);
});
test("pruneContextMessages keeps recent bulky tool results inside the recent-turn window", () => {
const policy = buildPolicy(2);
const messages = [
{ role: "user", content: "turn 1" },
{ role: "assistant", content: "observed turn 1" },
{ role: "user", content: "turn 2" },
{ role: "toolResult", toolName: "read", content: bulky },
{ role: "assistant", content: "observed turn 2" },
{ role: "user", content: "turn 3" },
];
const pruned = pruneContextMessages(messages, policy);
assert.deepEqual(pruned, messages);
});
test("pruneContextMessages keeps old non-bulky tool results outside the recent-turn window", () => {
const policy = buildPolicy(2);
const messages = [
{ role: "user", content: "turn 1" },
{ role: "toolResult", toolName: "read", content: "short output" },
{ role: "assistant", content: "observed turn 1" },
{ role: "user", content: "turn 2" },
{ role: "assistant", content: "observed turn 2" },
{ role: "user", content: "turn 3" },
];
const pruned = pruneContextMessages(messages, policy);
assert.deepEqual(pruned, messages);
});
test("pruneContextMessages keeps exactly-150-line tool results with a trailing newline", () => {
const policy = buildPolicy(2);
const messages = [
{ role: "user", content: "turn 1" },
{ role: "toolResult", toolName: "read", content: thresholdWithTrailingNewline },
{ role: "assistant", content: "after threshold output" },
{ role: "user", content: "turn 2" },
{ role: "assistant", content: "after turn 2" },
{ role: "user", content: "turn 3" },
];
const pruned = pruneContextMessages(messages, policy);
assert.deepEqual(pruned, messages);
});
test("pruneContextMessages honors the recent-user-turn boundary", () => {
const policy = buildPolicy(2);
const messages = [
{ role: "user", content: "turn 1" },
{ role: "toolResult", toolName: "read", content: bulky },
{ role: "assistant", content: "after turn 1" },
{ role: "user", content: "turn 2" },
{ role: "toolResult", toolName: "read", content: boundaryBulky },
{ role: "assistant", content: "after turn 2" },
{ role: "user", content: "turn 3" },
];
const pruned = pruneContextMessages(messages, policy);
assert.equal(pruned[1]?.role, "toolResult");
assert.match(pruned[1]?.content ?? "", /^\[distilled read output\]/);
assert.deepEqual(
pruned.map((message) => message.content),
["turn 1", pruned[1]!.content, "after turn 1", "turn 2", boundaryBulky, "after turn 2", "turn 3"]
);
});

54
src/prune.ts Normal file
View File

@@ -0,0 +1,54 @@
import type { Policy } from "./config.ts";
import { distillToolResult } from "./distill.ts";
export interface ContextMessage {
role: string;
content: string;
toolName?: string;
original?: unknown;
distilled?: boolean;
}
function isBulky(content: string, policy: Policy) {
const bytes = Buffer.byteLength(content, "utf8");
const parts = content.split("\n");
const lines = content.endsWith("\n") ? parts.length - 1 : parts.length;
return bytes > policy.bulkyBytes || lines > policy.bulkyLines;
}
export function pruneContextMessages(messages: ContextMessage[], policy: Policy): ContextMessage[] {
let seenUserTurns = 0;
const keep = new Set<number>();
for (let index = messages.length - 1; index >= 0; index -= 1) {
const message = messages[index]!;
keep.add(index);
if (message.role === "user") {
seenUserTurns += 1;
if (seenUserTurns >= policy.recentUserTurns) {
break;
}
}
}
const next: ContextMessage[] = [];
for (const [index, message] of messages.entries()) {
if (keep.has(index) || message.role !== "toolResult" || !isBulky(message.content, policy)) {
next.push(message);
continue;
}
const distilled = distillToolResult({ toolName: message.toolName, content: message.content });
if (!distilled) {
continue;
}
next.push({
...message,
content: distilled,
distilled: true,
});
}
return next;
}

178
src/runtime.test.ts Normal file
View File

@@ -0,0 +1,178 @@
import test from "node:test";
import assert from "node:assert/strict";
import { SNAPSHOT_ENTRY_TYPE, deserializeLatestSnapshot } from "./persist.ts";
import { createContextManagerRuntime } from "./runtime.ts";
test("runtime ingests transcript slices and updates pressure state", () => {
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
runtime.ingest({ entryId: "u1", role: "user", text: "Goal: Build a pi context manager. Next: wire hooks.", timestamp: 1 });
runtime.observeTokens(150_000);
const packet = runtime.buildPacket();
assert.match(packet.text, /Build a pi context manager/);
assert.equal(runtime.getSnapshot().lastZone, "red");
});
test("runtime keeps the session root goal while allowing later branch-local goals", () => {
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
runtime.ingest({ entryId: "u-root-goal", role: "user", text: "Goal: Ship the context manager extension.", timestamp: 10 });
runtime.ingest({ entryId: "u-branch-goal", role: "user", text: "Goal: prototype a branch-local tree handoff.", timestamp: 11 });
const packet = runtime.buildPacket();
assert.match(packet.text, /Ship the context manager extension/);
assert.match(packet.text, /prototype a branch-local tree handoff/);
assert.equal(
runtime
.getSnapshot()
.ledger.items.filter((item) => item.active && item.kind === "goal" && item.scope === "session").length,
1,
);
});
test("recordCompactionSummary and recordBranchSummary update snapshot state and resume output", () => {
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
runtime.ingest({ entryId: "u-artifact-1", role: "user", text: "Goal: Ship the context manager extension.", timestamp: 20 });
runtime.recordCompactionSummary(
"## Key Decisions\n- Keep summaries deterministic.\n\n## Open questions and blockers\n- Verify /tree replaceInstructions behavior.",
);
runtime.recordBranchSummary("# Handoff for branch\n\n## Key Decisions\n- Do not leak branch-local goals.");
const snapshot = runtime.getSnapshot();
assert.match(snapshot.lastCompactionSummary ?? "", /Keep summaries deterministic/);
assert.match(snapshot.lastBranchSummary ?? "", /Do not leak branch-local goals/);
assert.match(runtime.buildResumePacket(), /Verify \/tree replaceInstructions behavior/);
assert.match(runtime.buildResumePacket(), /Do not leak branch-local goals/);
});
test("buildPacket tightens the live packet after pressure reaches the compact zone", () => {
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
runtime.restore({
mode: "balanced",
lastZone: "green",
ledger: {
rollingSummary: "",
items: [
{ id: "goal:session:root-goal:1", kind: "goal", subject: "root-goal", text: "Ship the context manager extension with deterministic handoffs and predictable branch-boundary behavior.", scope: "session", sourceEntryId: "u1", sourceType: "user", timestamp: 1, confidence: 1, freshness: 1, active: true },
{ id: "constraint:session:must-1:1", kind: "constraint", subject: "must-1", text: "Keep the public API stable while hardening branch-boundary state carryover, fallback summary replay, and resume injection behavior.", scope: "session", sourceEntryId: "u1", sourceType: "user", timestamp: 2, confidence: 0.9, freshness: 2, active: true },
{ id: "decision:branch:decision-1:1", kind: "decision", subject: "decision-1", text: "Persist summary artifacts, replay them after the latest snapshot, and surface them through the next hidden resume packet before normal packet injection resumes.", scope: "branch", sourceEntryId: "a1", sourceType: "assistant", timestamp: 3, confidence: 0.9, freshness: 3, active: true },
{ id: "activeTask:branch:task-1:1", kind: "activeTask", subject: "task-1", text: "Verify mode-dependent pruning, packet tightening under pressure, and snapshot-less branch rehydration without stale handoff leakage.", scope: "branch", sourceEntryId: "a2", sourceType: "assistant", timestamp: 4, confidence: 0.8, freshness: 4, active: true },
{ id: "openQuestion:branch:question-1:1", kind: "openQuestion", subject: "question-1", text: "Confirm whether default pi fallback summaries preserve blockers and active work end to end when custom compaction falls back.", scope: "branch", sourceEntryId: "a3", sourceType: "assistant", timestamp: 5, confidence: 0.8, freshness: 5, active: true },
{ id: "relevantFile:branch:file-1:1", kind: "relevantFile", subject: "file-1", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-1.ts", scope: "branch", sourceEntryId: "t1", sourceType: "toolResult", timestamp: 6, confidence: 0.7, freshness: 6, active: true },
{ id: "relevantFile:branch:file-2:1", kind: "relevantFile", subject: "file-2", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-2.ts", scope: "branch", sourceEntryId: "t2", sourceType: "toolResult", timestamp: 7, confidence: 0.7, freshness: 7, active: true },
{ id: "relevantFile:branch:file-3:1", kind: "relevantFile", subject: "file-3", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-3.ts", scope: "branch", sourceEntryId: "t3", sourceType: "toolResult", timestamp: 8, confidence: 0.7, freshness: 8, active: true },
{ id: "relevantFile:branch:file-4:1", kind: "relevantFile", subject: "file-4", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-4.ts", scope: "branch", sourceEntryId: "t4", sourceType: "toolResult", timestamp: 9, confidence: 0.7, freshness: 9, active: true },
{ id: "relevantFile:branch:file-5:1", kind: "relevantFile", subject: "file-5", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-5.ts", scope: "branch", sourceEntryId: "t5", sourceType: "toolResult", timestamp: 10, confidence: 0.7, freshness: 10, active: true },
{ id: "relevantFile:branch:file-6:1", kind: "relevantFile", subject: "file-6", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-6.ts", scope: "branch", sourceEntryId: "t6", sourceType: "toolResult", timestamp: 11, confidence: 0.7, freshness: 11, active: true },
{ id: "relevantFile:branch:file-7:1", kind: "relevantFile", subject: "file-7", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-7.ts", scope: "branch", sourceEntryId: "t7", sourceType: "toolResult", timestamp: 12, confidence: 0.7, freshness: 12, active: true },
{ id: "relevantFile:branch:file-8:1", kind: "relevantFile", subject: "file-8", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-8.ts", scope: "branch", sourceEntryId: "t8", sourceType: "toolResult", timestamp: 13, confidence: 0.7, freshness: 13, active: true },
{ id: "relevantFile:branch:file-9:1", kind: "relevantFile", subject: "file-9", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-9.ts", scope: "branch", sourceEntryId: "t9", sourceType: "toolResult", timestamp: 14, confidence: 0.7, freshness: 14, active: true },
{ id: "relevantFile:branch:file-10:1", kind: "relevantFile", subject: "file-10", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-10.ts", scope: "branch", sourceEntryId: "t10", sourceType: "toolResult", timestamp: 15, confidence: 0.7, freshness: 15, active: true },
{ id: "relevantFile:branch:file-11:1", kind: "relevantFile", subject: "file-11", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-11.ts", scope: "branch", sourceEntryId: "t11", sourceType: "toolResult", timestamp: 16, confidence: 0.7, freshness: 16, active: true },
{ id: "relevantFile:branch:file-12:1", kind: "relevantFile", subject: "file-12", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-12.ts", scope: "branch", sourceEntryId: "t12", sourceType: "toolResult", timestamp: 17, confidence: 0.7, freshness: 17, active: true },
{ id: "relevantFile:branch:file-13:1", kind: "relevantFile", subject: "file-13", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-13.ts", scope: "branch", sourceEntryId: "t13", sourceType: "toolResult", timestamp: 18, confidence: 0.7, freshness: 18, active: true },
{ id: "relevantFile:branch:file-14:1", kind: "relevantFile", subject: "file-14", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-14.ts", scope: "branch", sourceEntryId: "t14", sourceType: "toolResult", timestamp: 19, confidence: 0.7, freshness: 19, active: true },
{ id: "relevantFile:branch:file-15:1", kind: "relevantFile", subject: "file-15", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-15.ts", scope: "branch", sourceEntryId: "t15", sourceType: "toolResult", timestamp: 20, confidence: 0.7, freshness: 20, active: true },
{ id: "relevantFile:branch:file-16:1", kind: "relevantFile", subject: "file-16", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-16.ts", scope: "branch", sourceEntryId: "t16", sourceType: "toolResult", timestamp: 21, confidence: 0.7, freshness: 21, active: true },
{ id: "relevantFile:branch:file-17:1", kind: "relevantFile", subject: "file-17", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-17.ts", scope: "branch", sourceEntryId: "t17", sourceType: "toolResult", timestamp: 22, confidence: 0.7, freshness: 22, active: true },
{ id: "relevantFile:branch:file-18:1", kind: "relevantFile", subject: "file-18", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-18.ts", scope: "branch", sourceEntryId: "t18", sourceType: "toolResult", timestamp: 23, confidence: 0.7, freshness: 23, active: true },
{ id: "relevantFile:branch:file-19:1", kind: "relevantFile", subject: "file-19", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-19.ts", scope: "branch", sourceEntryId: "t19", sourceType: "toolResult", timestamp: 24, confidence: 0.7, freshness: 24, active: true },
{ id: "relevantFile:branch:file-20:1", kind: "relevantFile", subject: "file-20", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-20.ts", scope: "branch", sourceEntryId: "t20", sourceType: "toolResult", timestamp: 25, confidence: 0.7, freshness: 25, active: true },
{ id: "relevantFile:branch:file-21:1", kind: "relevantFile", subject: "file-21", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-21.ts", scope: "branch", sourceEntryId: "t21", sourceType: "toolResult", timestamp: 26, confidence: 0.7, freshness: 26, active: true },
{ id: "relevantFile:branch:file-22:1", kind: "relevantFile", subject: "file-22", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-22.ts", scope: "branch", sourceEntryId: "t22", sourceType: "toolResult", timestamp: 27, confidence: 0.7, freshness: 27, active: true },
],
},
});
const before = runtime.buildPacket();
runtime.observeTokens(170_000);
const after = runtime.buildPacket();
assert.ok(after.estimatedTokens < before.estimatedTokens);
});
test("runtime recomputes lastZone when setContextWindow and setMode change policy", () => {
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
runtime.observeTokens(150_000);
assert.equal(runtime.getSnapshot().lastZone, "red");
runtime.setContextWindow(300_000);
assert.equal(runtime.getSnapshot().lastZone, "green");
runtime.setMode("aggressive");
assert.equal(runtime.getSnapshot().lastZone, "yellow");
});
test("restore recomputes lastZone against the receiving runtime policy", () => {
const source = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
source.observeTokens(150_000);
const target = createContextManagerRuntime({ mode: "balanced", contextWindow: 500_000 });
target.restore(source.getSnapshot());
assert.equal(target.getSnapshot().lastZone, "green");
});
test("restore resets legacy lastZone when the snapshot lacks lastObservedTokens", () => {
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 500_000 });
runtime.restore({
mode: "balanced",
lastZone: "red",
ledger: { items: [], rollingSummary: "" },
});
const restored = runtime.getSnapshot();
assert.equal(restored.lastZone, "green");
assert.equal(restored.lastObservedTokens, undefined);
});
test("legacy snapshot deserialization plus restore clears stale lastZone", () => {
const snapshot = deserializeLatestSnapshot([
{
type: "custom",
customType: SNAPSHOT_ENTRY_TYPE,
data: {
mode: "balanced",
lastZone: "red",
ledger: { items: [], rollingSummary: "" },
},
},
]);
assert.ok(snapshot);
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 500_000 });
runtime.restore(snapshot);
assert.equal(runtime.getSnapshot().lastZone, "green");
});
test("getPolicy returns a clone and restore detaches from external snapshot objects", () => {
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
const policy = runtime.getPolicy();
policy.packetTokenCap = 1;
policy.redAtTokens = 1;
const currentPolicy = runtime.getPolicy();
assert.equal(currentPolicy.packetTokenCap, 1_200);
assert.equal(currentPolicy.redAtTokens, 140_000);
const snapshot = runtime.getSnapshot();
snapshot.mode = "aggressive";
snapshot.ledger.rollingSummary = "before restore";
runtime.restore(snapshot);
snapshot.mode = "conservative";
snapshot.ledger.rollingSummary = "mutated after restore";
const restored = runtime.getSnapshot();
assert.equal(restored.mode, "aggressive");
assert.equal(restored.ledger.rollingSummary, "before restore");
});

127
src/runtime.ts Normal file
View File

@@ -0,0 +1,127 @@
import { adjustPolicyForZone, resolvePolicy, zoneForTokens, type ContextMode, type Policy } from "./config.ts";
import { extractCandidates, type TranscriptSlice } from "./extract.ts";
import { createEmptyLedger, getActiveItems, mergeCandidates } from "./ledger.ts";
import { buildContextPacket } from "./packet.ts";
import { buildBranchSummary, buildCompactionSummary, buildResumePacket as renderResumePacket } from "./summaries.ts";
import type { RuntimeSnapshot } from "./persist.ts";
function syncSnapshotZone(snapshot: RuntimeSnapshot, policy: Policy): RuntimeSnapshot {
if (snapshot.lastObservedTokens === undefined) {
return snapshot.lastZone === "green"
? snapshot
: {
...snapshot,
lastZone: "green",
};
}
return {
...snapshot,
lastZone: zoneForTokens(snapshot.lastObservedTokens, policy),
};
}
export function createContextManagerRuntime(input: { mode?: ContextMode; contextWindow: number }) {
let contextWindow = input.contextWindow;
let policy = resolvePolicy({ mode: input.mode ?? "balanced", contextWindow });
let snapshot: RuntimeSnapshot = {
mode: policy.mode,
lastZone: "green",
ledger: createEmptyLedger(),
};
function applyPolicy(nextPolicy: Policy) {
policy = nextPolicy;
snapshot = syncSnapshotZone(snapshot, policy);
}
function hasSessionGoal() {
return getActiveItems(snapshot.ledger, "goal").some((item) => item.scope === "session" && item.subject === "root-goal");
}
function ingest(slice: TranscriptSlice) {
snapshot = {
...snapshot,
ledger: mergeCandidates(snapshot.ledger, extractCandidates(slice, { hasSessionGoal: hasSessionGoal() })),
};
}
function observeTokens(tokens: number) {
snapshot = {
...snapshot,
lastObservedTokens: tokens,
lastZone: zoneForTokens(tokens, policy),
};
}
function buildPacket() {
return buildContextPacket(snapshot.ledger, adjustPolicyForZone(policy, snapshot.lastZone));
}
function mergeArtifact(role: "compaction" | "branchSummary", text: string, entryId: string, timestamp: number) {
snapshot = {
...snapshot,
lastCompactionSummary: role === "compaction" ? text : snapshot.lastCompactionSummary,
lastBranchSummary: role === "branchSummary" ? text : snapshot.lastBranchSummary,
ledger: mergeCandidates(snapshot.ledger, extractCandidates({ entryId, role, text, timestamp }, { hasSessionGoal: hasSessionGoal() })),
};
}
function recordCompactionSummary(text: string, entryId = `compaction-${Date.now()}`, timestamp = Date.now()) {
mergeArtifact("compaction", text, entryId, timestamp);
}
function recordBranchSummary(text: string, entryId = `branch-${Date.now()}`, timestamp = Date.now()) {
mergeArtifact("branchSummary", text, entryId, timestamp);
}
function buildResumePacket() {
const lines: string[] = [];
if (snapshot.lastCompactionSummary) {
lines.push("## Latest compaction handoff", snapshot.lastCompactionSummary, "");
}
if (snapshot.lastBranchSummary) {
lines.push("## Latest branch handoff", snapshot.lastBranchSummary, "");
}
const livePacket = renderResumePacket(snapshot.ledger);
if (livePacket) {
lines.push(livePacket);
}
return lines.join("\n").trim();
}
function setContextWindow(nextContextWindow: number) {
contextWindow = Math.max(nextContextWindow, 50_000);
applyPolicy(resolvePolicy({ mode: snapshot.mode, contextWindow }));
}
function setMode(mode: ContextMode) {
snapshot = { ...snapshot, mode };
applyPolicy(resolvePolicy({ mode, contextWindow }));
}
function restore(next: RuntimeSnapshot) {
snapshot = structuredClone(next);
applyPolicy(resolvePolicy({ mode: snapshot.mode, contextWindow }));
}
return {
ingest,
observeTokens,
buildPacket,
buildCompactionSummary: () => buildCompactionSummary(snapshot.ledger),
buildBranchSummary: (label: string) => buildBranchSummary(snapshot.ledger, label),
buildResumePacket,
recordCompactionSummary,
recordBranchSummary,
setContextWindow,
setMode,
getPolicy: () => structuredClone(policy),
getSnapshot: () => structuredClone(snapshot),
restore,
};
}

139
src/summaries.test.ts Normal file
View File

@@ -0,0 +1,139 @@
import test from "node:test";
import assert from "node:assert/strict";
import { createEmptyLedger, mergeCandidates } from "./ledger.ts";
import {
buildBranchSummary,
buildBranchSummaryFromEntries,
buildCompactionSummary,
buildCompactionSummaryFromPreparation,
buildResumePacket,
} from "./summaries.ts";
const ledger = mergeCandidates(createEmptyLedger(), [
{ kind: "goal", subject: "root-goal", text: "Build a pi context manager", scope: "session", sourceEntryId: "u1", sourceType: "user", timestamp: 1, confidence: 1 },
{ kind: "constraint", subject: "must-u1-0", text: "Must adapt to the active model context window.", scope: "session", sourceEntryId: "u1", sourceType: "user", timestamp: 1, confidence: 0.9 },
{ kind: "decision", subject: "decision-a1-0", text: "Keep the MVP quiet.", scope: "branch", sourceEntryId: "a1", sourceType: "assistant", timestamp: 2, confidence: 0.9 },
{ kind: "activeTask", subject: "next-step-a2-0", text: "Wire hooks into pi.", scope: "branch", sourceEntryId: "a2", sourceType: "assistant", timestamp: 3, confidence: 0.8 },
{ kind: "relevantFile", subject: "runtime-ts", text: "src/runtime.ts", scope: "branch", sourceEntryId: "a3", sourceType: "assistant", timestamp: 4, confidence: 0.7 },
]);
test("buildCompactionSummary renders the exact section order and content", () => {
const summary = buildCompactionSummary(ledger);
assert.equal(
summary,
[
"## Goal",
"- Build a pi context manager",
"",
"## Constraints",
"- Must adapt to the active model context window.",
"",
"## Decisions",
"- Keep the MVP quiet.",
"",
"## Active work",
"- Wire hooks into pi.",
"",
"## Relevant files",
"- src/runtime.ts",
"",
"## Next steps",
"- Wire hooks into pi.",
].join("\n")
);
});
test("buildBranchSummary renders the handoff header and sections in order", () => {
const summary = buildBranchSummary(ledger, "experimental branch");
assert.equal(
summary,
[
"# Handoff for experimental branch",
"",
"## Goal",
"- Build a pi context manager",
"",
"## Decisions",
"- Keep the MVP quiet.",
"",
"## Active work",
"- Wire hooks into pi.",
"",
"## Relevant files",
"- src/runtime.ts",
].join("\n")
);
});
test("buildCompactionSummaryFromPreparation uses preparation messages, previous summary, file ops, and focus text", () => {
const previousSummary = [
"## Goal",
"- Ship the context manager extension",
"",
"## Key Decisions",
"- Keep the public API stable.",
].join("\n");
const summary = buildCompactionSummaryFromPreparation({
messagesToSummarize: [
{ role: "user", content: "Decision: keep compaction summaries deterministic", timestamp: 1 },
{ role: "assistant", content: [{ type: "text", text: "Blocked: verify /tree replaceInstructions behavior" }], timestamp: 2 },
],
turnPrefixMessages: [
{ role: "toolResult", toolName: "read", content: [{ type: "text", text: "Opened .pi/agent/extensions/context-manager/index.ts" }], isError: false, timestamp: 3 },
],
previousSummary,
fileOps: {
readFiles: [".pi/agent/extensions/context-manager/index.ts"],
modifiedFiles: [".pi/agent/extensions/context-manager/src/summaries.ts"],
},
customInstructions: "Focus on decisions, blockers, and relevant files.",
});
assert.match(summary, /## Key Decisions/);
assert.match(summary, /keep compaction summaries deterministic/);
assert.match(summary, /## Open questions and blockers/);
assert.match(summary, /verify \/tree replaceInstructions behavior/);
assert.match(summary, /<read-files>[\s\S]*index.ts[\s\S]*<\/read-files>/);
assert.match(summary, /<modified-files>[\s\S]*src\/summaries.ts[\s\S]*<\/modified-files>/);
});
test("buildBranchSummaryFromEntries uses only the abandoned branch entries and custom focus", () => {
const summary = buildBranchSummaryFromEntries({
branchLabel: "abandoned branch",
entriesToSummarize: [
{ type: "message", id: "user-1", parentId: null, timestamp: new Date(1).toISOString(), message: { role: "user", content: "Goal: explore tree handoff" } },
{ type: "message", id: "assistant-1", parentId: "user-1", timestamp: new Date(2).toISOString(), message: { role: "assistant", content: [{ type: "text", text: "Decision: do not leak branch-local goals" }] } },
],
customInstructions: "Focus on goals and decisions.",
replaceInstructions: false,
commonAncestorId: "root",
});
assert.match(summary, /# Handoff for abandoned branch/);
assert.match(summary, /explore tree handoff/);
assert.match(summary, /do not leak branch-local goals/);
});
test("buildResumePacket renders restart guidance in the expected order", () => {
const summary = buildResumePacket(ledger);
assert.equal(
summary,
[
"## Goal",
"- Build a pi context manager",
"",
"## Current task",
"- Wire hooks into pi.",
"",
"## Constraints",
"- Must adapt to the active model context window.",
"",
"## Key decisions",
"- Keep the MVP quiet.",
].join("\n")
);
});

251
src/summaries.ts Normal file
View File

@@ -0,0 +1,251 @@
import { createEmptyLedger, getActiveItems, mergeCandidates, type LedgerState } from "./ledger.ts";
import { extractCandidates, type TranscriptSlice } from "./extract.ts";
function lines(title: string, items: string[]) {
if (items.length === 0) return [`## ${title}`, "- none", ""];
return [`## ${title}`, ...items.map((item) => `- ${item}`), ""];
}
function isTextPart(part: unknown): part is { type: "text"; text?: string } {
return typeof part === "object" && part !== null && "type" in part && (part as { type?: unknown }).type === "text";
}
function toText(content: unknown): string {
if (typeof content === "string") {
return content;
}
if (!Array.isArray(content)) {
return "";
}
return content
.map((part) => {
if (!isTextPart(part)) return "";
return typeof part.text === "string" ? part.text : "";
})
.join("\n")
.trim();
}
type FocusKind = "goal" | "constraint" | "decision" | "activeTask" | "openQuestion" | "relevantFile";
type SummarySection = { kind: FocusKind; title: string; items: string[] };
function hasSessionGoal(ledger: LedgerState) {
return getActiveItems(ledger, "goal").some((item) => item.scope === "session" && item.subject === "root-goal");
}
function buildLedgerFromSlices(slices: TranscriptSlice[], previousSummary?: string) {
let ledger = createEmptyLedger();
if (previousSummary) {
ledger = mergeCandidates(
ledger,
extractCandidates(
{
entryId: "previous-summary",
role: "compaction",
text: previousSummary,
timestamp: 0,
},
{ hasSessionGoal: false },
),
);
}
for (const slice of slices) {
ledger = mergeCandidates(ledger, extractCandidates(slice, { hasSessionGoal: hasSessionGoal(ledger) }));
}
return ledger;
}
function parseFocus(customInstructions?: string): Set<FocusKind> {
const text = (customInstructions ?? "").toLowerCase();
const focus = new Set<FocusKind>();
if (/\bgoal/.test(text)) focus.add("goal");
if (/\bconstraint|preference/.test(text)) focus.add("constraint");
if (/\bdecision/.test(text)) focus.add("decision");
if (/\btask|next step|progress/.test(text)) focus.add("activeTask");
if (/\bblocker|blocked|open question/.test(text)) focus.add("openQuestion");
if (/\bfile|read-files|modified-files/.test(text)) focus.add("relevantFile");
return focus;
}
function buildStructuredSections(ledger: LedgerState): SummarySection[] {
return [
{ kind: "goal", title: "Goal", items: getActiveItems(ledger, "goal").map((item) => item.text) },
{
kind: "constraint",
title: "Constraints & Preferences",
items: getActiveItems(ledger, "constraint").map((item) => item.text),
},
{
kind: "activeTask",
title: "Progress",
items: getActiveItems(ledger, "activeTask").map((item) => item.text),
},
{
kind: "decision",
title: "Key Decisions",
items: getActiveItems(ledger, "decision").map((item) => item.text),
},
{
kind: "openQuestion",
title: "Open questions and blockers",
items: getActiveItems(ledger, "openQuestion").map((item) => item.text),
},
{
kind: "relevantFile",
title: "Critical Context",
items: getActiveItems(ledger, "relevantFile").map((item) => item.text),
},
{
kind: "activeTask",
title: "Next Steps",
items: getActiveItems(ledger, "activeTask").map((item) => item.text),
},
];
}
function sortSectionsForFocus(sections: SummarySection[], focus: Set<FocusKind>): SummarySection[] {
if (focus.size === 0) {
return sections;
}
return [
...sections.filter((section) => focus.has(section.kind)),
...sections.filter((section) => !focus.has(section.kind)),
];
}
function unique(values: string[]) {
return [...new Set(values)];
}
function fileTag(name: "read-files" | "modified-files", values: string[]) {
if (values.length === 0) {
return [] as string[];
}
return [`<${name}>`, ...values, `</${name}>`, ""];
}
function renderStructuredSummary(
ledger: LedgerState,
options?: {
header?: string;
focus?: Set<FocusKind>;
readFiles?: string[];
modifiedFiles?: string[];
},
) {
const sections = sortSectionsForFocus(buildStructuredSections(ledger), options?.focus ?? new Set());
return [
...(options?.header ? [options.header, ""] : []),
...sections.flatMap((section) => lines(section.title, section.items)),
...fileTag("read-files", unique(options?.readFiles ?? [])),
...fileTag("modified-files", unique(options?.modifiedFiles ?? [])),
].join("\n").trim();
}
function messageToSlice(message: any, entryId: string, timestampFallback: number): TranscriptSlice | undefined {
if (!message || typeof message !== "object") {
return undefined;
}
if (message.role !== "user" && message.role !== "assistant" && message.role !== "toolResult") {
return undefined;
}
return {
entryId,
role: message.role,
text: toText(message.content),
timestamp: typeof message.timestamp === "number" ? message.timestamp : timestampFallback,
isError: message.role === "toolResult" ? message.isError : undefined,
};
}
export function buildCompactionSummary(ledger: LedgerState): string {
return [
...lines("Goal", getActiveItems(ledger, "goal").map((item) => item.text)),
...lines("Constraints", getActiveItems(ledger, "constraint").map((item) => item.text)),
...lines("Decisions", getActiveItems(ledger, "decision").map((item) => item.text)),
...lines("Active work", getActiveItems(ledger, "activeTask").map((item) => item.text)),
...lines("Relevant files", getActiveItems(ledger, "relevantFile").map((item) => item.text)),
...lines("Next steps", getActiveItems(ledger, "activeTask").map((item) => item.text)),
].join("\n").trim();
}
export function buildBranchSummary(ledger: LedgerState, branchLabel: string): string {
return [
`# Handoff for ${branchLabel}`,
"",
...lines("Goal", getActiveItems(ledger, "goal").map((item) => item.text)),
...lines("Decisions", getActiveItems(ledger, "decision").map((item) => item.text)),
...lines("Active work", getActiveItems(ledger, "activeTask").map((item) => item.text)),
...lines("Relevant files", getActiveItems(ledger, "relevantFile").map((item) => item.text)),
].join("\n").trim();
}
export function buildResumePacket(ledger: LedgerState): string {
return [
...lines("Goal", getActiveItems(ledger, "goal").map((item) => item.text)),
...lines("Current task", getActiveItems(ledger, "activeTask").map((item) => item.text)),
...lines("Constraints", getActiveItems(ledger, "constraint").map((item) => item.text)),
...lines("Key decisions", getActiveItems(ledger, "decision").map((item) => item.text)),
].join("\n").trim();
}
export function buildCompactionSummaryFromPreparation(input: {
messagesToSummarize: any[];
turnPrefixMessages: any[];
previousSummary?: string;
fileOps?: { readFiles?: string[]; modifiedFiles?: string[] };
customInstructions?: string;
}): string {
const slices = [...input.messagesToSummarize, ...input.turnPrefixMessages]
.map((message, index) => messageToSlice(message, `compaction-${index}`, index))
.filter((slice): slice is TranscriptSlice => Boolean(slice));
const ledger = buildLedgerFromSlices(slices, input.previousSummary);
return renderStructuredSummary(ledger, {
focus: parseFocus(input.customInstructions),
readFiles: input.fileOps?.readFiles,
modifiedFiles: input.fileOps?.modifiedFiles,
});
}
export function buildBranchSummaryFromEntries(input: {
branchLabel: string;
entriesToSummarize: Array<{ type: string; id: string; timestamp: string; message?: any; summary?: string }>;
customInstructions?: string;
replaceInstructions?: boolean;
commonAncestorId?: string | null;
}): string {
const slices = input.entriesToSummarize.flatMap((entry) => {
if (entry.type === "message") {
const slice = messageToSlice(entry.message, entry.id, Date.parse(entry.timestamp));
return slice ? [slice] : [];
}
if (entry.type === "compaction" && typeof entry.summary === "string") {
return [{ entryId: entry.id, role: "compaction", text: entry.summary, timestamp: Date.parse(entry.timestamp) } satisfies TranscriptSlice];
}
if (entry.type === "branch_summary" && typeof entry.summary === "string") {
return [{ entryId: entry.id, role: "branchSummary", text: entry.summary, timestamp: Date.parse(entry.timestamp) } satisfies TranscriptSlice];
}
return [];
});
const ledger = buildLedgerFromSlices(slices);
return renderStructuredSummary(ledger, {
header: `# Handoff for ${input.branchLabel}`,
focus: parseFocus(input.customInstructions),
});
}