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 <triagent@localhost>
This commit is contained in:
robopipi[bot] 2026-11-02 10:55:30 +01:00 • committed by GitHub
parent 6ab6c1e0e4
commit d0fe213089
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 69 additions and 6 deletions

View file

@ -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<string, ContextEditEntry>();
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)

View file

@ -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;

View file

@ -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");
});
});