Skip to content

Commit a501917

Browse files
amknightamknight
committed
fix: use branch fallback for compaction safeguard
1 parent 0fd2f03 commit a501917

2 files changed

Lines changed: 182 additions & 8 deletions

File tree

‎src/agents/pi-hooks/compaction-safeguard.test.ts‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2084,6 +2084,91 @@ describe("compaction-safeguard double-compaction guard", () => {
20842084
expect(getApiKeyAndHeadersMock).toHaveBeenCalled();
20852085
});
20862086

2087+
it("falls back to visible custom session branch entries before writing an empty boundary", async () => {
2088+
mockSummarizeInStages.mockReset();
2089+
mockSummarizeInStages.mockResolvedValue("branch summary");
2090+
2091+
const now = Date.now();
2092+
const sessionManager = {
2093+
...stubSessionManager(),
2094+
getBranch: () => [
2095+
{
2096+
type: "custom_message",
2097+
id: "custom-1",
2098+
parentId: null,
2099+
timestamp: new Date(now).toISOString(),
2100+
customType: "cron-request",
2101+
content: "prepare the daily report",
2102+
display: true,
2103+
},
2104+
{
2105+
type: "message",
2106+
id: "assistant-1",
2107+
parentId: "custom-1",
2108+
timestamp: new Date(now + 1).toISOString(),
2109+
message: {
2110+
role: "assistant",
2111+
content: [{ type: "toolCall", id: "call-1", name: "read", arguments: {} }],
2112+
timestamp: now + 1,
2113+
},
2114+
},
2115+
{
2116+
type: "message",
2117+
id: "tool-1",
2118+
parentId: "assistant-1",
2119+
timestamp: new Date(now + 2).toISOString(),
2120+
message: {
2121+
role: "toolResult",
2122+
toolCallId: "call-1",
2123+
toolName: "read",
2124+
content: [{ type: "text", text: "report source data" }],
2125+
timestamp: now + 2,
2126+
},
2127+
},
2128+
],
2129+
} as ExtensionContext["sessionManager"];
2130+
const model = createAnthropicModelFixture();
2131+
setCompactionSafeguardRuntime(sessionManager, { model, recentTurnsPreserve: 0 });
2132+
2133+
const mockEvent = {
2134+
preparation: {
2135+
messagesToSummarize: [] as AgentMessage[],
2136+
turnPrefixMessages: [] as AgentMessage[],
2137+
firstKeptEntryId: "entry-5",
2138+
tokensBefore: 38085,
2139+
fileOps: { read: [], edited: [], written: [] },
2140+
settings: { reserveTokens: 4000 },
2141+
isSplitTurn: true,
2142+
},
2143+
customInstructions: "",
2144+
signal: new AbortController().signal,
2145+
};
2146+
const { result } = await runCompactionScenario({
2147+
sessionManager,
2148+
event: mockEvent,
2149+
apiKey: "test-key",
2150+
});
2151+
2152+
const compaction = expectCompactionResult(result);
2153+
expect(compaction.summary).toContain("branch summary");
2154+
expect(compaction.summary).not.toContain("No prior history.");
2155+
expect(mockSummarizeInStages).toHaveBeenCalled();
2156+
const summaryCall = mockSummarizeInStages.mock.calls[0]?.[0];
2157+
expect(summaryCall?.messages).toEqual(
2158+
expect.arrayContaining([
2159+
expect.objectContaining({
2160+
role: "custom",
2161+
customType: "cron-request",
2162+
content: "prepare the daily report",
2163+
}),
2164+
expect.objectContaining({
2165+
role: "toolResult",
2166+
toolName: "read",
2167+
}),
2168+
]),
2169+
);
2170+
});
2171+
20872172
it("continues when messages include real conversation content", async () => {
20882173
const sessionManager = stubSessionManager();
20892174
const model = createAnthropicModelFixture();

‎src/agents/pi-hooks/compaction-safeguard.ts‎

Lines changed: 97 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,85 @@ function prependPreviousSummaryForRedistill(params: {
9999
return [buildPreviousSummaryMessage(previousSummary), ...params.messages];
100100
}
101101

102+
type SessionBranchEntry = {
103+
type?: unknown;
104+
message?: unknown;
105+
customType?: unknown;
106+
content?: unknown;
107+
display?: unknown;
108+
details?: unknown;
109+
timestamp?: unknown;
110+
summary?: unknown;
111+
fromId?: unknown;
112+
};
113+
114+
function coerceTimestamp(value: unknown): number {
115+
if (typeof value === "number" && Number.isFinite(value)) {
116+
return value;
117+
}
118+
if (typeof value === "string") {
119+
const parsed = Date.parse(value);
120+
if (Number.isFinite(parsed)) {
121+
return parsed;
122+
}
123+
}
124+
return 0;
125+
}
126+
127+
function sessionBranchEntryToMessage(entry: SessionBranchEntry): AgentMessage | undefined {
128+
if (entry.type === "message" && entry.message && typeof entry.message === "object") {
129+
return entry.message as AgentMessage;
130+
}
131+
if (entry.type === "custom_message") {
132+
return {
133+
role: "custom",
134+
customType: typeof entry.customType === "string" ? entry.customType : "custom",
135+
content: entry.content,
136+
display: entry.display !== false,
137+
details: entry.details,
138+
timestamp: coerceTimestamp(entry.timestamp),
139+
} as AgentMessage;
140+
}
141+
if (entry.type === "branch_summary") {
142+
return {
143+
role: "branchSummary",
144+
summary: typeof entry.summary === "string" ? entry.summary : "",
145+
fromId: typeof entry.fromId === "string" ? entry.fromId : "root",
146+
timestamp: coerceTimestamp(entry.timestamp),
147+
} as AgentMessage;
148+
}
149+
return undefined;
150+
}
151+
152+
function collectSessionBranchMessages(sessionManager: unknown): AgentMessage[] {
153+
const getBranch = (sessionManager as { getBranch?: unknown })?.getBranch;
154+
if (typeof getBranch !== "function") {
155+
return [];
156+
}
157+
let entries: unknown;
158+
try {
159+
entries = getBranch.call(sessionManager);
160+
} catch {
161+
return [];
162+
}
163+
if (!Array.isArray(entries)) {
164+
return [];
165+
}
166+
return entries
167+
.map((entry) =>
168+
entry && typeof entry === "object"
169+
? sessionBranchEntryToMessage(entry as SessionBranchEntry)
170+
: undefined,
171+
)
172+
.filter((message): message is AgentMessage => Boolean(message));
173+
}
174+
175+
function containsRealConversation(messages: AgentMessage[]): boolean {
176+
return messages.some((message, index, allMessages) =>
177+
isRealConversationMessage(message, allMessages, index),
178+
);
179+
}
180+
102181
/**
103182
* Attempt provider-based summarization. Returns the summary string on success,
104183
* or `undefined` when the caller should fall back to built-in LLM summarization.
@@ -778,16 +857,26 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void {
778857
api.on("session_before_compact", async (event, ctx) => {
779858
const { preparation, customInstructions: eventInstructions, signal } = event;
780859
const rawTurnPrefixMessages = preparation.turnPrefixMessages ?? [];
781-
const baseMessagesToSummarize = stripRuntimeContextCustomMessages(
860+
let baseMessagesToSummarize = stripRuntimeContextCustomMessages(
782861
preparation.messagesToSummarize,
783862
);
784-
const baseTurnPrefixMessages = stripRuntimeContextCustomMessages(rawTurnPrefixMessages);
785-
const hasRealSummarizable = baseMessagesToSummarize.some((message, index, messages) =>
786-
isRealConversationMessage(message, messages, index),
787-
);
788-
const hasRealTurnPrefix = baseTurnPrefixMessages.some((message, index, messages) =>
789-
isRealConversationMessage(message, messages, index),
790-
);
863+
let baseTurnPrefixMessages = stripRuntimeContextCustomMessages(rawTurnPrefixMessages);
864+
let hasRealSummarizable = containsRealConversation(baseMessagesToSummarize);
865+
let hasRealTurnPrefix = containsRealConversation(baseTurnPrefixMessages);
866+
if (!hasRealSummarizable && !hasRealTurnPrefix) {
867+
const branchMessages = stripRuntimeContextCustomMessages(
868+
collectSessionBranchMessages(ctx.sessionManager),
869+
);
870+
if (containsRealConversation(branchMessages)) {
871+
log.info(
872+
"Compaction safeguard: using session branch messages after compaction preparation omitted real conversation content.",
873+
);
874+
baseMessagesToSummarize = branchMessages;
875+
baseTurnPrefixMessages = [];
876+
hasRealSummarizable = true;
877+
hasRealTurnPrefix = false;
878+
}
879+
}
791880
setCompactionSafeguardCancelReason(ctx.sessionManager, undefined);
792881
if (!hasRealSummarizable && !hasRealTurnPrefix) {
793882
// When there are no summarizable messages AND no real turn-prefix content,

0 commit comments

Comments
 (0)