diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts index 69bd75133..6f8bc0cb3 100644 --- a/packages/coding-agent/src/core/compaction/branch-summarization.ts +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -15,7 +15,12 @@ import { createCompactionSummaryMessage, createCustomMessage, } from "../messages.ts"; -import type { ReadonlySessionManager, SessionEntry } from "../session-manager.ts"; +import { + type ContextEditEntry, + projectContextEntry, + type ReadonlySessionManager, + type SessionEntry, +} from "../session-manager.ts"; import { completeSummarization, estimateTokens, getSummarizationFailure } from "./compaction.ts"; import { computeFileLists, @@ -155,14 +160,15 @@ export function collectEntriesForBranchSummary( * Extract AgentMessage from a session entry. * Similar to getMessageFromEntry in compaction.ts but also handles compaction entries. */ -function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined { +function getMessageFromEntry(entry: SessionEntry, edit?: ContextEditEntry): AgentMessage | undefined { switch (entry.type) { case "message": // Skip tool results - context is in assistant's tool call if (entry.message.role === "toolResult") return undefined; - return entry.message; + return edit ? projectContextEntry(entry, edit)[0] : entry.message; case "custom_message": + if (edit) return projectContextEntry(entry, edit)[0]; return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); case "branch_summary": @@ -217,10 +223,16 @@ export function prepareBranchEntries(entries: SessionEntry[], tokenBudget: numbe } } + // Context edits hide or replace earlier entries in model context; honor the latest edit per target. + const edits = new Map(); + for (const entry of entries) { + if (entry.type === "context_edit") edits.set(entry.targetId, entry); + } + // Second pass: walk from newest to oldest, adding messages until token budget for (let i = entries.length - 1; i >= 0; i--) { const entry = entries[i]; - const message = getMessageFromEntry(entry); + const message = getMessageFromEntry(entry, edits.get(entry.id)); if (!message) continue; // Extract file ops from assistant messages (tool calls) diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 822aad816..20b6edaa9 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -516,7 +516,8 @@ export function buildContextEntries( * If leafId is provided, walks from that entry to root. * Handles compaction and branch summaries along the path. */ -function projectContextEntry(entry: SessionEntry, edit: ContextEditEntry | undefined): AgentMessage[] { +/** Project one entry into model-visible messages, applying its latest context edit if any. */ +export function projectContextEntry(entry: SessionEntry, edit: ContextEditEntry | undefined): AgentMessage[] { const messages = sessionEntryToContextMessages(entry); if (!edit) return messages; const replacement = edit.replacement; diff --git a/packages/coding-agent/test/branch-summarization.test.ts b/packages/coding-agent/test/branch-summarization.test.ts index ade25dbfd..3164a8773 100644 --- a/packages/coding-agent/test/branch-summarization.test.ts +++ b/packages/coding-agent/test/branch-summarization.test.ts @@ -7,7 +7,7 @@ import { type SimpleStreamOptions, } from "@earendil-works/pi-ai"; import { describe, expect, it } from "vitest"; -import { generateBranchSummary } from "../src/core/compaction/index.ts"; +import { generateBranchSummary, prepareBranchEntries } from "../src/core/compaction/index.ts"; import type { SessionEntry } from "../src/core/session-manager.ts"; const model: Model<"anthropic-messages"> = { @@ -131,4 +131,54 @@ describe("branch summarization", () => { "Branch summarization failed: generation hit the token cap and the summary is incomplete", ); }); + + it("applies context edits to branch entries", () => { + const ts = (n: number) => new Date(n).toISOString(); + const branch: SessionEntry[] = [ + { + type: "message", + id: "u1", + parentId: null, + timestamp: ts(1), + message: { role: "user", content: "Investigate login", timestamp: 1 }, + }, + { + type: "message", + id: "failed", + parentId: "u1", + timestamp: ts(2), + message: fauxAssistantMessage("FAILED_ATTEMPT", { stopReason: "error", errorMessage: "lost" }), + }, + { + type: "context_edit", + id: "e1", + parentId: "failed", + timestamp: ts(3), + targetId: "failed", + replacement: null, + }, + { + type: "message", + id: "ok", + parentId: "e1", + timestamp: ts(4), + message: fauxAssistantMessage("ORIGINAL_TEXT"), + }, + { + type: "context_edit", + id: "e2", + parentId: "ok", + timestamp: ts(5), + targetId: "ok", + replacement: { content: "REPLACED_TEXT" }, + }, + ]; + + const serialized = JSON.stringify(prepareBranchEntries(branch).messages); + + expect(serialized).toContain("Investigate login"); + expect(serialized).not.toContain("FAILED_ATTEMPT"); + expect(serialized).not.toContain("ORIGINAL_TEXT"); + expect(serialized).toContain("REPLACED_TEXT"); + }); });