diff --git a/docs/plans/2026-04-12-context-manager-implementation.md b/docs/plans/2026-04-12-context-manager-implementation.md new file mode 100644 index 0000000..f8655ac --- /dev/null +++ b/docs/plans/2026-04-12-context-manager-implementation.md @@ -0,0 +1,180 @@ +# Context manager stronger pruning and compaction Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Keep live context flatter by pruning whole older turns, compacting earlier, and injecting only lean resume state. + +**Architecture:** Extend pruning at the context boundary, keep summary artifacts in snapshot/ledger only, and add an extension-managed early-compaction gate driven by observed context usage. Preserve existing commands and ledger semantics while removing footer status noise. + +**Tech Stack:** TypeScript, Node test runner (`tsx --test`), Pi extension hooks + +--- + +### Task 1: Turn-aware pruning and lean resume rendering + +**Files:** +- Modify: `src/prune.ts` +- Modify: `src/prune.test.ts` +- Modify: `src/summaries.ts` +- Modify: `src/summaries.test.ts` +- Modify: `src/runtime.ts` +- Modify: `src/runtime.test.ts` + +- [ ] **Step 1: Write the failing prune and resume tests** + +```ts +test("pruneContextMessages drops turns older than the kept suffix", () => { + const messages = [ + { role: "user", content: "turn 1" }, + { role: "assistant", content: "after turn 1" }, + { role: "user", content: "turn 2" }, + { role: "assistant", content: "after turn 2" }, + { role: "user", content: "turn 3" }, + ]; + + const pruned = pruneContextMessages(messages, buildPolicy(2)); + assert.deepEqual(pruned.map((m) => m.content), ["turn 2", "after turn 2", "turn 3"]); +}); + +test("buildResumePacket keeps blocker text after raw handoff sections are removed", () => { + const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 }); + runtime.recordCompactionSummary("## Open questions and blockers\n- Verify /tree replaceInstructions behavior."); + assert.match(runtime.buildResumePacket(), /Verify \/tree replaceInstructions behavior/); + assert.doesNotMatch(runtime.buildResumePacket(), /## Latest compaction handoff/); +}); +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `cd /home/dev/pi-packages/pi-context-manager/.worktree/context-manager-prune-compact && npm test -- src/prune.test.ts src/runtime.test.ts src/summaries.test.ts` +Expected: FAIL because pruning still keeps old turns and runtime still prepends raw handoff sections. + +- [ ] **Step 3: Write the minimal implementation** + +```ts +// src/prune.ts +// Keep only the last N user-anchored turns. +// Distill bulky tool results only for kept older turns, never for the newest turn. + +// src/summaries.ts +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)), + ...lines("Open questions / blockers", getActiveItems(ledger, "openQuestion").map((item) => item.text)), + ].join("\n").trim(); +} + +// src/runtime.ts +function buildResumePacket() { + return renderResumePacket(snapshot.ledger).trim(); +} +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `cd /home/dev/pi-packages/pi-context-manager/.worktree/context-manager-prune-compact && npm test -- src/prune.test.ts src/runtime.test.ts src/summaries.test.ts` +Expected: PASS for the new pruning and lean-resume assertions. + +- [ ] **Step 5: Commit** + +```bash +cd /home/dev/pi-packages/pi-context-manager/.worktree/context-manager-prune-compact +git add src/prune.ts src/prune.test.ts src/summaries.ts src/summaries.test.ts src/runtime.ts src/runtime.test.ts +git commit -m "fix: prune turns and slim resume packets" +``` + +### Task 2: Context filtering, early compaction, and footer removal + +**Files:** +- Modify: `index.ts` +- Modify: `src/extension.test.ts` + +- [ ] **Step 1: Write the failing extension tests** + +```ts +test("context filters raw compaction and branch summary messages", async () => { + const result = await harness.handlers.get("context")?.({ + type: "context", + messages: [ + { role: "compactionSummary", summary: "## Key Decisions\n- noisy raw handoff", timestamp: 1 }, + { role: "branchSummary", summary: "# Handoff\n- noisy raw branch", timestamp: 2 }, + createUserMessage("continue", 3), + ], + }, harness.ctx); + + assert.ok(result.messages.every((message: any) => message.role !== "compactionSummary" && message.role !== "branchSummary")); +}); + +test("turn_end triggers early compaction at red zone without footer status writes", async () => { + const harness = createHarness([], { usageTokens: 150_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.compactions.length, 1); + assert.equal(harness.statuses.length, 0); +}); +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `cd /home/dev/pi-packages/pi-context-manager/.worktree/context-manager-prune-compact && npm test -- src/extension.test.ts` +Expected: FAIL because summary messages are still passed through, early compaction is absent, and status writes still occur. + +- [ ] **Step 3: Write the minimal implementation** + +```ts +// index.ts +// - Filter raw compactionSummary / branchSummary messages before pruning. +// - Remove ctx.ui.setStatus(...) calls. +// - After observeTokens on turn_end, call ctx.compact() once when snapshot.lastZone is red/compact. +// - Reset the latch on session_compact or when usage drops below red. +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `cd /home/dev/pi-packages/pi-context-manager/.worktree/context-manager-prune-compact && npm test -- src/extension.test.ts` +Expected: PASS for summary filtering, early compaction, and no-footer assertions. + +- [ ] **Step 5: Commit** + +```bash +cd /home/dev/pi-packages/pi-context-manager/.worktree/context-manager-prune-compact +git add index.ts src/extension.test.ts +git commit -m "fix: compact earlier and drop status noise" +``` + +### Task 3: Full verification and cleanup + +**Files:** +- Modify: `src/config.ts` (only if zone tightening needs a small recent-turn adjustment) +- Modify: `src/config.test.ts` (only if policy values change) + +- [ ] **Step 1: Run the full test suite** + +```bash +cd /home/dev/pi-packages/pi-context-manager/.worktree/context-manager-prune-compact +npm test +``` + +- [ ] **Step 2: If any failures remain, make the minimal fix and rerun** + +```ts +// Only patch the failing behavior surfaced by the full test run. +// Do not broaden scope beyond pruning, lean resume injection, summary filtering, +// early compaction, and footer removal. +``` + +- [ ] **Step 3: Verify the final green run** + +Run: `cd /home/dev/pi-packages/pi-context-manager/.worktree/context-manager-prune-compact && npm test` +Expected: `# fail 0` + +- [ ] **Step 4: Commit final cleanup if needed** + +```bash +cd /home/dev/pi-packages/pi-context-manager/.worktree/context-manager-prune-compact +git add index.ts src/*.ts docs/plans/2026-04-12-context-manager-implementation.md +git commit -m "test: cover context pressure behavior" +``` diff --git a/index.ts b/index.ts index 376680d..0ce948d 100644 --- a/index.ts +++ b/index.ts @@ -71,6 +71,10 @@ function rewriteContextMessage(message: { role: string; content: string; origina } as AgentMessage; } +function isRawSummaryArtifactMessage(message: AgentMessage): boolean { + return message.role === "compactionSummary" || message.role === "branchSummary"; +} + function findLatestSnapshotState(branch: BranchEntry[]): { snapshot: RuntimeSnapshot; index: number } | undefined { for (let index = branch.length - 1; index >= 0; index -= 1) { const entry = branch[index]!; @@ -147,6 +151,7 @@ export default function contextManager(pi: ExtensionAPI) { contextWindow: 200_000, }); let pendingResumeInjection = false; + let earlyCompactionRequested = false; const syncContextWindow = (ctx: Pick) => { runtime.setContextWindow(ctx.model?.contextWindow ?? 200_000); @@ -180,7 +185,7 @@ export default function contextManager(pi: ExtensionAPI) { }; const rebuildRuntimeFromBranch = ( - ctx: Pick, + ctx: Pick, fallbackSnapshot: RuntimeSnapshot, options?: { preferRuntimeMode?: boolean }, ) => { @@ -202,9 +207,26 @@ export default function contextManager(pi: ExtensionAPI) { for (const entry of replayEntries) { replayBranchEntry(entry); } + }; - const snapshot = runtime.getSnapshot(); - ctx.ui.setStatus("context-manager", `ctx ${snapshot.lastZone}`); + const maybeTriggerEarlyCompaction = (ctx: Pick) => { + const usage = ctx.getContextUsage(); + if (usage?.tokens === null || usage?.tokens === undefined) { + return; + } + + const zone = runtime.getSnapshot().lastZone; + if (zone === "green" || zone === "yellow") { + earlyCompactionRequested = false; + return; + } + + if (earlyCompactionRequested) { + return; + } + + earlyCompactionRequested = true; + ctx.compact(); }; registerContextCommands(pi, { @@ -220,11 +242,13 @@ export default function contextManager(pi: ExtensionAPI) { }); pi.on("session_start", async (_event, ctx) => { + earlyCompactionRequested = false; rebuildRuntimeFromBranch(ctx, createDefaultSnapshot()); armResumeInjection(); }); pi.on("session_tree", async (event, ctx) => { + earlyCompactionRequested = false; rebuildRuntimeFromBranch(ctx, createDefaultSnapshot()); if ( @@ -260,19 +284,21 @@ export default function contextManager(pi: ExtensionAPI) { const snapshot = runtime.getSnapshot(); pi.appendEntry(SNAPSHOT_ENTRY_TYPE, serializeSnapshot(snapshot)); - ctx.ui.setStatus("context-manager", `ctx ${snapshot.lastZone}`); + maybeTriggerEarlyCompaction(ctx); }); 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 normalized = event.messages + .filter((message) => !isRawSummaryArtifactMessage(message as AgentMessage)) + .map((message) => ({ + role: message.role, + content: getMessageContent(message), + toolName: getMessageToolName(message), + original: message, + })); const pruned = pruneContextMessages(normalized, policy); const nextMessages = pruned.map((message) => @@ -344,10 +370,10 @@ export default function contextManager(pi: ExtensionAPI) { }; }); - pi.on("session_compact", async (event, ctx) => { + pi.on("session_compact", async (event) => { + earlyCompactionRequested = false; 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}`); }); } diff --git a/src/extension.test.ts b/src/extension.test.ts index 3ed6437..18a5b47 100644 --- a/src/extension.test.ts +++ b/src/extension.test.ts @@ -162,8 +162,10 @@ function createHarness(initialBranch: SessionEntry[], options?: { usageTokens?: const handlers = new Map(); const appendedEntries: Array<{ customType: string; data: unknown }> = []; const statuses: Array<{ key: string; value: string }> = []; + const compactions: any[] = []; let branch = [...initialBranch]; let entries = [...initialBranch]; + let usageTokens = options?.usageTokens; const ctx = { model: { contextWindow: 200_000 }, @@ -183,9 +185,11 @@ function createHarness(initialBranch: SessionEntry[], options?: { usageTokens?: editor: async () => {}, }, getContextUsage() { - return options?.usageTokens === undefined ? undefined : createUsage(options.usageTokens); + return usageTokens === undefined ? undefined : createUsage(usageTokens); + }, + compact(options?: unknown) { + compactions.push(options); }, - compact() {}, }; contextManagerExtension({ @@ -215,6 +219,7 @@ function createHarness(initialBranch: SessionEntry[], options?: { usageTokens?: handlers, appendedEntries, statuses, + compactions, ctx, setBranch(nextBranch: SessionEntry[]) { branch = [...nextBranch]; @@ -224,6 +229,9 @@ function createHarness(initialBranch: SessionEntry[], options?: { usageTokens?: } entries = [...byId.values()]; }, + setUsageTokens(nextUsageTokens: number | undefined) { + usageTokens = nextUsageTokens; + }, }; } @@ -278,7 +286,7 @@ test("turn_end persists a rebuilt snapshot that includes branch user and assista 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" }); + assert.equal(harness.statuses.length, 0); }); test("session_tree rebuilds runtime from snapshot-only branches before injecting the next packet", async () => { @@ -347,7 +355,7 @@ test("session_tree rebuilds runtime from snapshot-only branches before injecting 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" }); + assert.equal(harness.statuses.length, 0); }); test("context keeps a distilled stale tool result visible after pruning bulky output", async () => { @@ -365,9 +373,9 @@ test("context keeps a distilled stale tool result visible after pruning bulky ou type: "context", messages: [ createUserMessage("turn 1", 1), - createToolResultMessage(bulkyFailure, 2), - createAssistantMessage("observed turn 1", 3), - createUserMessage("turn 2", 4), + createAssistantMessage("observed turn 1", 2), + createUserMessage("turn 2", 3), + createToolResultMessage(bulkyFailure, 4), createAssistantMessage("observed turn 2", 5), createUserMessage("turn 3", 6), createAssistantMessage("observed turn 3", 7), @@ -385,6 +393,92 @@ test("context keeps a distilled stale tool result visible after pruning bulky ou assert.ok(toolResult.content[0].text.length < 320); }); +test("context filters raw summary artifact messages before injecting the packet", async () => { + const harness = createHarness([ + createSnapshotEntry("snapshot-ctx-filter", null, { + text: "Ship the context manager extension", + lastCompactionSummary: "", + lastBranchSummary: "", + }), + ]); + 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), + { role: "compactionSummary", summary: "## Key Decisions\n- noisy raw compaction", timestamp: 2 }, + { role: "branchSummary", summary: "# Handoff\n\n## Key Decisions\n- noisy raw branch", timestamp: 3 }, + createAssistantMessage("observed turn 1", 4), + createUserMessage("turn 2", 5), + ], + }, + harness.ctx, + ); + + assert.equal(result.messages[0]?.customType, "context-manager.packet"); + assert.ok(result.messages.every((message: any) => message.role !== "compactionSummary")); + assert.ok(result.messages.every((message: any) => message.role !== "branchSummary")); +}); + +test("turn_end triggers extension compaction at red zone once per pressure episode and never writes footer status", async () => { + const harness = createHarness([], { usageTokens: 150_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, + ); + await harness.handlers.get("turn_end")?.( + { + type: "turn_end", + turnIndex: 2, + message: createAssistantMessage("done again", 6), + toolResults: [], + }, + harness.ctx, + ); + + assert.equal(harness.compactions.length, 1); + assert.equal(harness.statuses.length, 0); + + await harness.handlers.get("session_compact")?.( + { + type: "session_compact", + fromExtension: true, + compactionEntry: { + type: "compaction", + id: "cmp-reset", + parentId: "prev", + timestamp: new Date(10).toISOString(), + summary: "## Key Decisions\n- Keep summaries deterministic.", + firstKeptEntryId: "keep-1", + tokensBefore: 140_000, + }, + }, + harness.ctx, + ); + + await harness.handlers.get("turn_end")?.( + { + type: "turn_end", + turnIndex: 3, + message: createAssistantMessage("done after compact", 7), + toolResults: [], + }, + harness.ctx, + ); + + assert.equal(harness.compactions.length, 2); + assert.equal(harness.statuses.length, 0); +}); + 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, { diff --git a/src/prune.test.ts b/src/prune.test.ts index 747a415..9afe303 100644 --- a/src/prune.test.ts +++ b/src/prune.test.ts @@ -14,8 +14,8 @@ function buildPolicy(recentUserTurns = 4) { }; } -test("pruneContextMessages replaces old bulky tool results with distilled summaries instead of deleting them", () => { - const policy = buildPolicy(2); +test("pruneContextMessages replaces bulky tool results with distilled summaries inside older kept turns", () => { + const policy = buildPolicy(3); const bulkyFailure = [ "Build failed while compiling focus parser", "Error: missing export createFocusMatcher from ./summary-focus.ts", @@ -23,11 +23,13 @@ test("pruneContextMessages replaces old bulky tool results with distilled summar ].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: "toolResult", toolName: "bash", content: bulkyFailure }, { role: "assistant", content: "observed turn 2" }, { role: "user", content: "turn 3" }, + { role: "assistant", content: "observed turn 3" }, + { role: "user", content: "turn 4" }, ]; const pruned = pruneContextMessages(messages, policy); @@ -38,7 +40,7 @@ test("pruneContextMessages replaces old bulky tool results with distilled summar assert.doesNotMatch(distilled!.content, /stack frame\nstack frame\nstack frame/); }); -test("aggressive mode distills an older bulky tool result sooner than conservative mode", () => { +test("aggressive mode drops older turns sooner than conservative mode", () => { const conservative = resolvePolicy({ mode: "conservative", contextWindow: 200_000 }); const aggressive = resolvePolicy({ mode: "aggressive", contextWindow: 200_000 }); const messages = [ @@ -50,17 +52,18 @@ test("aggressive mode distills an older bulky tool result sooner than conservati { role: "user", content: "turn 3" }, { role: "assistant", content: "after turn 3" }, { role: "user", content: "turn 4" }, + { role: "assistant", content: "after turn 4" }, + { role: "user", content: "turn 5" }, ]; 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\]/); + assert.ok(conservativePruned.some((message) => message.role === "toolResult")); + assert.ok(aggressivePruned.every((message) => message.role !== "toolResult")); }); -test("pruneContextMessages keeps recent bulky tool results inside the recent-turn window", () => { +test("pruneContextMessages keeps newest-turn bulky tool results lossless", () => { const policy = buildPolicy(2); const messages = [ { role: "user", content: "turn 1" }, @@ -68,7 +71,6 @@ test("pruneContextMessages keeps recent bulky tool results inside the recent-tur { 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); @@ -76,7 +78,7 @@ test("pruneContextMessages keeps recent bulky tool results inside the recent-tur assert.deepEqual(pruned, messages); }); -test("pruneContextMessages keeps old non-bulky tool results outside the recent-turn window", () => { +test("pruneContextMessages drops old non-bulky tool results outside the recent-turn window", () => { const policy = buildPolicy(2); const messages = [ { role: "user", content: "turn 1" }, @@ -89,43 +91,67 @@ test("pruneContextMessages keeps old non-bulky tool results outside the recent-t 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"] + ["turn 2", "observed turn 2", "turn 3"], ); }); + +test("pruneContextMessages keeps exactly-150-line tool results with a trailing newline inside a kept older turn", () => { + const policy = buildPolicy(3); + const messages = [ + { role: "user", content: "turn 1" }, + { role: "assistant", content: "after turn 1" }, + { role: "user", content: "turn 2" }, + { role: "toolResult", toolName: "read", content: thresholdWithTrailingNewline }, + { role: "assistant", content: "after threshold output" }, + { role: "user", content: "turn 3" }, + { role: "assistant", content: "after turn 3" }, + { role: "user", content: "turn 4" }, + ]; + + const pruned = pruneContextMessages(messages, policy); + + assert.equal(pruned[1]?.content, thresholdWithTrailingNewline); +}); + +test("pruneContextMessages drops turns older than the kept suffix", () => { + const policy = buildPolicy(2); + const messages = [ + { role: "user", content: "turn 1" }, + { role: "assistant", content: "after turn 1" }, + { role: "user", content: "turn 2" }, + { role: "assistant", content: "after turn 2" }, + { role: "user", content: "turn 3" }, + ]; + + const pruned = pruneContextMessages(messages, policy); + + assert.deepEqual( + pruned.map((message) => message.content), + ["turn 2", "after turn 2", "turn 3"], + ); +}); + +test("pruneContextMessages distills bulky tool results only inside older kept turns", () => { + const policy = buildPolicy(2); + const messages = [ + { role: "user", content: "turn 1" }, + { role: "assistant", content: "after turn 1" }, + { role: "user", content: "turn 2" }, + { role: "toolResult", toolName: "read", content: bulky }, + { role: "assistant", content: "after turn 2" }, + { role: "user", content: "turn 3" }, + { role: "toolResult", toolName: "read", content: boundaryBulky }, + { role: "assistant", content: "after turn 3" }, + ]; + + const pruned = pruneContextMessages(messages, policy); + + assert.deepEqual( + pruned.map((message) => message.role), + ["user", "toolResult", "assistant", "user", "toolResult", "assistant"], + ); + assert.match(pruned[1]?.content ?? "", /^\[distilled read output\]/); + assert.equal(pruned[4]?.content, boundaryBulky); +}); diff --git a/src/prune.ts b/src/prune.ts index 768243e..66c105d 100644 --- a/src/prune.ts +++ b/src/prune.ts @@ -17,23 +17,38 @@ function isBulky(content: string, policy: Policy) { } export function pruneContextMessages(messages: ContextMessage[], policy: Policy): ContextMessage[] { + let earliestKeptTurnStart = 0; + let latestTurnStart = -1; let seenUserTurns = 0; - const keep = new Set(); 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; - } + if (message.role !== "user") { + continue; } + + if (latestTurnStart === -1) { + latestTurnStart = index; + } + + seenUserTurns += 1; + earliestKeptTurnStart = index; + if (seenUserTurns >= policy.recentUserTurns) { + break; + } + } + + if (latestTurnStart === -1) { + latestTurnStart = 0; } const next: ContextMessage[] = []; for (const [index, message] of messages.entries()) { - if (keep.has(index) || message.role !== "toolResult" || !isBulky(message.content, policy)) { + if (index < earliestKeptTurnStart) { + continue; + } + + if (index >= latestTurnStart || message.role !== "toolResult" || !isBulky(message.content, policy)) { next.push(message); continue; } diff --git a/src/runtime.test.ts b/src/runtime.test.ts index 7fe3a52..4371767 100644 --- a/src/runtime.test.ts +++ b/src/runtime.test.ts @@ -31,7 +31,7 @@ test("runtime keeps the session root goal while allowing later branch-local goal ); }); -test("recordCompactionSummary and recordBranchSummary update snapshot state and resume output", () => { +test("recordCompactionSummary and recordBranchSummary update snapshot state while keeping resume output lean", () => { 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 }); @@ -41,10 +41,14 @@ test("recordCompactionSummary and recordBranchSummary update snapshot state and runtime.recordBranchSummary("# Handoff for branch\n\n## Key Decisions\n- Do not leak branch-local goals."); const snapshot = runtime.getSnapshot(); + const resumePacket = runtime.buildResumePacket(); + 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/); + assert.match(resumePacket, /Verify \/tree replaceInstructions behavior/); + assert.match(resumePacket, /Do not leak branch-local goals/); + assert.doesNotMatch(resumePacket, /## Latest compaction handoff/); + assert.doesNotMatch(resumePacket, /## Latest branch handoff/); }); test("buildPacket tightens the live packet after pressure reaches the compact zone", () => { diff --git a/src/runtime.ts b/src/runtime.ts index 9a4ef50..74597ca 100644 --- a/src/runtime.ts +++ b/src/runtime.ts @@ -76,22 +76,7 @@ export function createContextManagerRuntime(input: { mode?: ContextMode; context } 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(); + return renderResumePacket(snapshot.ledger).trim(); } function setContextWindow(nextContextWindow: number) { diff --git a/src/summaries.test.ts b/src/summaries.test.ts index ee76a5f..72a8661 100644 --- a/src/summaries.test.ts +++ b/src/summaries.test.ts @@ -15,6 +15,7 @@ const ledger = mergeCandidates(createEmptyLedger(), [ { 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 }, + { kind: "openQuestion", subject: "blocked-a4-0", text: "Verify /tree replaceInstructions behavior.", scope: "branch", sourceEntryId: "a4", sourceType: "assistant", timestamp: 5, confidence: 0.8 }, ]); test("buildCompactionSummary renders the exact section order and content", () => { @@ -134,6 +135,9 @@ test("buildResumePacket renders restart guidance in the expected order", () => { "", "## Key decisions", "- Keep the MVP quiet.", + "", + "## Open questions / blockers", + "- Verify /tree replaceInstructions behavior.", ].join("\n") ); }); diff --git a/src/summaries.ts b/src/summaries.ts index b4a722b..d540c62 100644 --- a/src/summaries.ts +++ b/src/summaries.ts @@ -197,6 +197,7 @@ export function buildResumePacket(ledger: LedgerState): string { ...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)), + ...lines("Open questions / blockers", getActiveItems(ledger, "openQuestion").map((item) => item.text)), ].join("\n").trim(); }