From 90de4bd85566b45f804366199853ad345b163d24 Mon Sep 17 00:00:00 2001 From: pashpashpash Date: Sun, 26 Apr 2026 18:27:38 -0700 Subject: [PATCH] fix: address successor transcript review follow-ups Fixes the post-merge review follow-ups from #72471 by deduping stale pre-compaction state entries and preserving parent-before-child ordering for successor transcripts. --- .../compaction-successor-transcript.test.ts | 100 ++++++++++++++++++ .../compaction-successor-transcript.ts | 93 ++++++++++++++-- 2 files changed, 184 insertions(+), 9 deletions(-) diff --git a/src/agents/pi-embedded-runner/compaction-successor-transcript.test.ts b/src/agents/pi-embedded-runner/compaction-successor-transcript.test.ts index d7fb656a642..3329e3f10c8 100644 --- a/src/agents/pi-embedded-runner/compaction-successor-transcript.test.ts +++ b/src/agents/pi-embedded-runner/compaction-successor-transcript.test.ts @@ -102,6 +102,57 @@ describe("rotateTranscriptAfterCompaction", () => { expect(successor.getLabel(oldUserId)).toBeUndefined(); }); + it("deduplicates stale pre-compaction session state", async () => { + const dir = await createTmpDir(); + const manager = SessionManager.create(dir, dir); + + const staleModelId = manager.appendModelChange("anthropic", "claude-sonnet-4-5"); + const staleThinkingId = manager.appendThinkingLevelChange("low"); + const staleSessionInfoId = manager.appendSessionInfo("stale title"); + manager.appendCustomEntry("test-extension", { cursor: "preserved" }); + manager.appendMessage({ role: "user", content: "old user", timestamp: 1 }); + manager.appendMessage(makeAssistant("old assistant", 2)); + + manager.appendModelChange("openai", "gpt-5.2"); + manager.appendThinkingLevelChange("high"); + manager.appendSessionInfo("current title"); + const firstKeptId = manager.appendMessage({ role: "user", content: "kept user", timestamp: 3 }); + manager.appendMessage(makeAssistant("kept assistant", 4)); + manager.appendCompaction("Summary of old user and old assistant.", firstKeptId, 5000); + manager.appendMessage({ role: "user", content: "post user", timestamp: 5 }); + + const result = await rotateTranscriptAfterCompaction({ + sessionManager: manager, + sessionFile: manager.getSessionFile()!, + now: () => new Date("2026-04-27T12:05:00.000Z"), + }); + + expect(result.rotated).toBe(true); + const successor = SessionManager.open(result.sessionFile!); + const entries = successor.getEntries(); + expect(entries.find((entry) => entry.id === staleModelId)).toBeUndefined(); + expect(entries.find((entry) => entry.id === staleThinkingId)).toBeUndefined(); + expect(entries.find((entry) => entry.id === staleSessionInfoId)).toBeUndefined(); + expect(entries.filter((entry) => entry.type === "model_change")).toHaveLength(1); + expect(entries.filter((entry) => entry.type === "thinking_level_change")).toHaveLength(1); + expect(entries.filter((entry) => entry.type === "session_info")).toHaveLength(1); + expect(entries.find((entry) => entry.type === "model_change")).toMatchObject({ + provider: "openai", + modelId: "gpt-5.2", + }); + expect(entries).toContainEqual( + expect.objectContaining({ + type: "custom", + customType: "test-extension", + data: { cursor: "preserved" }, + }), + ); + + const context = successor.buildSessionContext(); + expect(context.thinkingLevel).toBe("high"); + expect(successor.getSessionName()).toBe("current title"); + }); + it("skips sessions with no compaction entry", async () => { const dir = await createTmpDir(); const manager = SessionManager.create(dir, dir); @@ -212,6 +263,55 @@ describe("rotateTranscriptAfterCompaction", () => { expect(activeContextText).toContain("next"); expect(activeContextText).not.toContain("do task B instead"); }); + + it("orders preserved sibling branches after their surviving parents", async () => { + const dir = await createTmpDir(); + const manager = SessionManager.create(dir, dir); + + manager.appendMessage({ role: "user", content: "hello", timestamp: 1 }); + const branchFromId = manager.appendMessage(makeAssistant("hi there", 2)); + + const branchSummaryId = manager.branchWithSummary( + branchFromId, + "Summary of the inactive branch.", + ); + const inactiveMsgId = manager.appendMessage({ + role: "user", + content: "inactive branch", + timestamp: 3, + }); + manager.appendMessage(makeAssistant("inactive done", 4)); + + manager.branch(branchFromId); + manager.appendMessage({ role: "user", content: "active branch", timestamp: 5 }); + manager.appendMessage(makeAssistant("active done", 6)); + manager.appendCompaction("Summary of active work.", branchFromId, 5000); + const activeLeafId = manager.appendMessage({ + role: "user", + content: "next active", + timestamp: 7, + }); + + const result = await rotateTranscriptAfterCompaction({ + sessionManager: manager, + sessionFile: manager.getSessionFile()!, + now: () => new Date("2026-04-27T13:00:00.000Z"), + }); + + expect(result.rotated).toBe(true); + const successor = SessionManager.open(result.sessionFile!); + const entries = successor.getEntries(); + const indexById = new Map(entries.map((entry, index) => [entry.id, index])); + expect(indexById.get(branchFromId)).toBeLessThan(indexById.get(branchSummaryId)!); + expect(indexById.get(branchSummaryId)).toBeLessThan(indexById.get(inactiveMsgId)!); + expect(entries.at(-1)?.id).toBe(activeLeafId); + expect(successor.getLeafId()).toBe(activeLeafId); + + const activeContextText = JSON.stringify(successor.buildSessionContext().messages); + expect(activeContextText).toContain("Summary of active work."); + expect(activeContextText).toContain("next active"); + expect(activeContextText).not.toContain("inactive branch"); + }); }); describe("shouldRotateCompactionTranscript", () => { diff --git a/src/agents/pi-embedded-runner/compaction-successor-transcript.ts b/src/agents/pi-embedded-runner/compaction-successor-transcript.ts index 9deee8f14cc..1e6dd00347a 100644 --- a/src/agents/pi-embedded-runner/compaction-successor-transcript.ts +++ b/src/agents/pi-embedded-runner/compaction-successor-transcript.ts @@ -117,9 +117,20 @@ function buildSuccessorEntries(params: { summarizedBranchIds.add(entry.id); } + const latestStateEntryIds = collectLatestStateEntryIds(branch.slice(0, latestCompactionIndex)); + const staleStateEntryIds = new Set(); + for (const entry of branch.slice(0, latestCompactionIndex)) { + if (isDedupedStateEntry(entry) && !latestStateEntryIds.has(entry.id)) { + staleStateEntryIds.add(entry.id); + } + } + const removedIds = new Set(); for (const entry of allEntries) { - if (summarizedBranchIds.has(entry.id) && entry.type === "message") { + if ( + (summarizedBranchIds.has(entry.id) && entry.type === "message") || + staleStateEntryIds.has(entry.id) + ) { removedIds.add(entry.id); } } @@ -131,6 +142,7 @@ function buildSuccessorEntries(params: { const entryById = new Map(allEntries.map((entry) => [entry.id, entry])); const activeBranchIds = new Set(branch.map((entry) => entry.id)); + const originalIndexById = new Map(allEntries.map((entry, index) => [entry.id, index])); const keptEntries: SessionEntry[] = []; for (const entry of allEntries) { if (removedIds.has(entry.id)) { @@ -147,17 +159,80 @@ function buildSuccessorEntries(params: { ); } - const inactiveEntries: SessionEntry[] = []; - const activeEntries: SessionEntry[] = []; - for (const entry of keptEntries) { - if (activeBranchIds.has(entry.id)) { - activeEntries.push(entry); - } else { - inactiveEntries.push(entry); + return orderSuccessorEntries({ + entries: keptEntries, + activeBranchIds, + originalIndexById, + }); +} + +function collectLatestStateEntryIds(entries: SessionEntry[]): Set { + const latestByType = new Map(); + for (const entry of entries) { + if (isDedupedStateEntry(entry)) { + latestByType.set(entry.type, entry); } } + return new Set(Array.from(latestByType.values(), (entry) => entry.id)); +} - return [...inactiveEntries, ...activeEntries]; +function isDedupedStateEntry(entry: SessionEntry): boolean { + return ( + entry.type === "model_change" || + entry.type === "thinking_level_change" || + entry.type === "session_info" + ); +} + +function orderSuccessorEntries(params: { + entries: SessionEntry[]; + activeBranchIds: Set; + originalIndexById: Map; +}): SessionEntry[] { + const { entries, activeBranchIds, originalIndexById } = params; + const entryIds = new Set(entries.map((entry) => entry.id)); + const childrenByParentId = new Map(); + + for (const entry of entries) { + const parentId = + entry.parentId !== null && entryIds.has(entry.parentId) ? entry.parentId : null; + const children = childrenByParentId.get(parentId) ?? []; + children.push(parentId === entry.parentId ? entry : ({ ...entry, parentId } as SessionEntry)); + childrenByParentId.set(parentId, children); + } + + const sortForActiveLeaf = (left: SessionEntry, right: SessionEntry) => { + const leftActive = activeBranchIds.has(left.id); + const rightActive = activeBranchIds.has(right.id); + if (leftActive !== rightActive) { + return leftActive ? 1 : -1; + } + return (originalIndexById.get(left.id) ?? 0) - (originalIndexById.get(right.id) ?? 0); + }; + + const ordered: SessionEntry[] = []; + const emittedIds = new Set(); + const emitSubtree = (entry: SessionEntry) => { + if (emittedIds.has(entry.id)) { + return; + } + emittedIds.add(entry.id); + ordered.push(entry); + for (const child of (childrenByParentId.get(entry.id) ?? []).toSorted(sortForActiveLeaf)) { + emitSubtree(child); + } + }; + + for (const root of (childrenByParentId.get(null) ?? []).toSorted(sortForActiveLeaf)) { + emitSubtree(root); + } + + // Defensive fallback for malformed transcripts with cycles or broken parents. + for (const entry of entries.toSorted(sortForActiveLeaf)) { + emitSubtree(entry); + } + + return ordered; } function buildSuccessorHeader(params: {