From d0fe213089aef0599aa9c90bf910c8b562e0dcf8 Mon Sep 17 00:00:00 2001 From: "robopipi[bot]" <336283572+robopipi[bot]@users.noreply.github.com> Date: Mon, 2 Nov 2026 10:55:30 +0100 Subject: [PATCH] fix(coding-agent): apply context edits to branch summary input (#11680) When a partial assistant reply fails and automatic retry succeeds, pi appends a `context_edit` that hides the failed attempt from model context. Branch summarization (`/tree` with summarize) built its input from raw session entries through `getMessageFromEntry`, which ignored `context_edit` entries. As a result, the summarizer received the hidden attempt again. Content replacements were ignored the same way. What the change does: - `prepareBranchEntries` now collects the latest `context_edit` for each target in the branch. - It applies those edits to the entries it converts, using the same `projectContextEntry` logic as the live context and compaction: - an omission (`replacement: null`) drops the entry; - a replacement swaps its content. - `projectContextEntry` is now exported from `session-manager.ts` so both paths share it. - The existing branch-summary behavior is unchanged: tool results are still skipped, and compaction or branch-summary entries still become summary messages. Testing: - Added `applies context edits to branch entries` to `test/branch-summarization.test.ts`. It covers one omitted assistant entry and one replaced assistant entry. The test fails without the change and passes with it. - Ran `branch-summarization`, `branch-summary-extensions`, `agent-session-branching` and the compaction tests: all pass. - `npm run check`: biome and the coding-agent checks pass. `tsc` reports 6 errors in `packages/ai` model-catalog tests, and the same 6 errors occur on unmodified `main`. - End-to-end reproduction with the faux provider and `AgentSession`: a failed attempt, a successful retry, then `navigateTree(target, { summarize: true })`. Before the change, the summary request contained the failed attempt's text. After the change, it contains only the successful retry. Fixes #11602 Requested by @davidbrai Co-authored-by: triagent --- .../core/compaction/branch-summarization.ts | 20 +++++-- .../coding-agent/src/core/session-manager.ts | 3 +- .../test/branch-summarization.test.ts | 52 ++++++++++++++++++- 3 files changed, 69 insertions(+), 6 deletions(-) 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"); + }); });