Skip to content

Commit b995409

Browse files
pashpashpashpashpashpash
authored
Fix compaction rotation follow-ups
1 parent b9c7a43 commit b995409

8 files changed

Lines changed: 268 additions & 64 deletions

File tree

‎src/agents/pi-embedded-runner/compact.hooks.harness.ts‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { vi, type Mock } from "vitest";
22
import { clearAgentHarnesses } from "../harness/registry.js";
3+
import type { CompactionTranscriptRotation } from "./compaction-successor-transcript.js";
34

45
type MockResolvedModel = {
56
model: { provider: string; api: string; id: string; input: unknown[] };
@@ -98,6 +99,11 @@ export const resolveAgentTransportOverrideMock: Mock<(params?: unknown) => strin
9899
export const resolveSandboxContextMock = vi.fn(async () => null);
99100
export const maybeCompactAgentHarnessSessionMock: Mock<(params?: unknown) => Promise<unknown>> =
100101
vi.fn(async () => undefined);
102+
export const rotateTranscriptAfterCompactionMock: Mock<
103+
(_params?: unknown) => Promise<CompactionTranscriptRotation>
104+
> = vi.fn(async () => ({
105+
rotated: false,
106+
}));
101107

102108
export function resetCompactSessionStateMocks(): void {
103109
sanitizeSessionHistoryMock.mockReset();
@@ -138,6 +144,8 @@ export function resetCompactSessionStateMocks(): void {
138144
resolveSandboxContextMock.mockResolvedValue(null);
139145
maybeCompactAgentHarnessSessionMock.mockReset();
140146
maybeCompactAgentHarnessSessionMock.mockResolvedValue(undefined);
147+
rotateTranscriptAfterCompactionMock.mockReset();
148+
rotateTranscriptAfterCompactionMock.mockResolvedValue({ rotated: false });
141149
}
142150

143151
export function resetCompactHooksHarnessMocks(): void {
@@ -209,6 +217,7 @@ export async function loadCompactHooksHarness(): Promise<{
209217

210218
vi.doMock("../../plugins/provider-runtime.js", () => ({
211219
prepareProviderRuntimeAuth: vi.fn(async () => ({ resolvedApiKey: undefined })),
220+
resolveProviderReasoningOutputModeWithPlugin: vi.fn(() => undefined),
212221
resolveProviderSystemPromptContribution: vi.fn(() => undefined),
213222
resolveProviderTextTransforms: vi.fn(() => undefined),
214223
transformProviderSystemPrompt: vi.fn(
@@ -264,12 +273,17 @@ export async function loadCompactHooksHarness(): Promise<{
264273
session.messages.splice(1);
265274
return await sessionCompactImpl();
266275
}),
276+
setActiveToolsByName: vi.fn(),
267277
abortCompaction: sessionAbortCompactionMock,
268278
dispose: vi.fn(),
269279
};
270280
return { session };
271281
}),
272-
DefaultResourceLoader: function DefaultResourceLoader() {},
282+
DefaultResourceLoader: function DefaultResourceLoader() {
283+
return {
284+
reload: vi.fn(async () => undefined),
285+
};
286+
},
273287
SessionManager: {
274288
open: vi.fn(() => ({})),
275289
},
@@ -287,6 +301,7 @@ export async function loadCompactHooksHarness(): Promise<{
287301
}));
288302

289303
vi.doMock("../pi-settings.js", () => ({
304+
applyPiCompactionSettingsFromConfig: vi.fn(),
290305
ensurePiCompactionReserveTokens: vi.fn(),
291306
resolveCompactionReserveTokensFloor: vi.fn(() => 0),
292307
}));
@@ -442,6 +457,16 @@ export async function loadCompactHooksHarness(): Promise<{
442457
resolveCompactionTimeoutMs: vi.fn(() => 30_000),
443458
}));
444459

460+
vi.doMock("./compaction-successor-transcript.js", async () => {
461+
const actual = await vi.importActual<typeof import("./compaction-successor-transcript.js")>(
462+
"./compaction-successor-transcript.js",
463+
);
464+
return {
465+
...actual,
466+
rotateTranscriptAfterCompaction: rotateTranscriptAfterCompactionMock,
467+
};
468+
});
469+
445470
vi.doMock("./wait-for-idle-before-flush.js", () => ({
446471
flushPendingToolResultsAfterIdle: vi.fn(async () => {}),
447472
}));
@@ -476,6 +501,8 @@ export async function loadCompactHooksHarness(): Promise<{
476501

477502
vi.doMock("../agent-scope.js", () => ({
478503
listAgentEntries: vi.fn(() => []),
504+
resolveAgentConfig: vi.fn(() => undefined),
505+
resolveDefaultAgentId: vi.fn(() => "main"),
479506
resolveSessionAgentId: resolveSessionAgentIdMock,
480507
resolveSessionAgentIds: vi.fn(() => ({ defaultAgentId: "main", sessionAgentId: "main" })),
481508
}));

‎src/agents/pi-embedded-runner/compact.hooks.test.ts‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import {
1717
resolveModelMock,
1818
resolveSandboxContextMock,
1919
resolveSessionAgentIdMock,
20+
rotateTranscriptAfterCompactionMock,
2021
resetCompactHooksHarnessMocks,
2122
resetCompactSessionStateMocks,
2223
sessionAbortCompactionMock,
@@ -411,6 +412,49 @@ describe("compactEmbeddedPiSessionDirect hooks", () => {
411412
}
412413
});
413414

415+
it("emits post-compaction side effects once for a rotated successor transcript", async () => {
416+
const listener = vi.fn();
417+
const cleanup = onSessionTranscriptUpdate(listener);
418+
const sync = vi.fn(async () => {});
419+
getMemorySearchManagerMock.mockResolvedValue({ manager: { sync } });
420+
rotateTranscriptAfterCompactionMock.mockResolvedValueOnce({
421+
rotated: true,
422+
sessionId: "rotated-session",
423+
sessionFile: "/tmp/rotated-session.jsonl",
424+
leafId: "rotated-leaf",
425+
});
426+
427+
try {
428+
const result = await compactEmbeddedPiSessionDirect({
429+
sessionId: "session-1",
430+
sessionKey: TEST_SESSION_KEY,
431+
sessionFile: "/tmp/session.jsonl",
432+
workspaceDir: "/tmp/workspace",
433+
config: {
434+
agents: {
435+
defaults: {
436+
compaction: {
437+
truncateAfterCompaction: true,
438+
postIndexSync: "await",
439+
},
440+
},
441+
},
442+
} as never,
443+
});
444+
445+
expect(result.ok).toBe(true);
446+
expect(listener).toHaveBeenCalledTimes(1);
447+
expect(listener).toHaveBeenCalledWith({ sessionFile: "/tmp/rotated-session.jsonl" });
448+
expect(sync).toHaveBeenCalledTimes(1);
449+
expect(sync).toHaveBeenCalledWith({
450+
reason: "post-compaction",
451+
sessionFiles: ["/tmp/rotated-session.jsonl"],
452+
});
453+
} finally {
454+
cleanup();
455+
}
456+
});
457+
414458
it("preserves tokensAfter when full-session context exceeds result.tokensBefore", async () => {
415459
estimateTokensMock.mockImplementation((message: unknown) => {
416460
const role = (message as { role?: string }).role;
@@ -1008,6 +1052,63 @@ describe("compactEmbeddedPiSession hooks (ownsCompaction engine)", () => {
10081052
);
10091053
});
10101054

1055+
it("rotates in the wrapper when a delegated result echoes the current transcript", async () => {
1056+
const maintain = vi.fn(async (_params?: unknown) => ({
1057+
changed: false,
1058+
bytesFreed: 0,
1059+
rewrittenEntries: 0,
1060+
}));
1061+
resolveContextEngineMock.mockResolvedValue({
1062+
info: { ownsCompaction: false },
1063+
compact: contextEngineCompactMock,
1064+
maintain,
1065+
} as never);
1066+
contextEngineCompactMock.mockResolvedValue({
1067+
ok: true,
1068+
compacted: true,
1069+
reason: undefined,
1070+
result: {
1071+
summary: "engine-summary",
1072+
firstKeptEntryId: "entry-1",
1073+
tokensBefore: 120,
1074+
tokensAfter: 50,
1075+
sessionId: TEST_SESSION_ID,
1076+
sessionFile: TEST_SESSION_FILE,
1077+
},
1078+
} as never);
1079+
rotateTranscriptAfterCompactionMock.mockResolvedValueOnce({
1080+
rotated: true,
1081+
sessionId: "wrapper-rotated-session",
1082+
sessionFile: "/tmp/wrapper-rotated-session.jsonl",
1083+
leafId: "wrapper-rotated-leaf",
1084+
});
1085+
1086+
const result = await compactEmbeddedPiSession(
1087+
wrappedCompactionArgs({
1088+
config: {
1089+
agents: {
1090+
defaults: {
1091+
compaction: {
1092+
truncateAfterCompaction: true,
1093+
},
1094+
},
1095+
},
1096+
},
1097+
}),
1098+
);
1099+
1100+
expect(result.ok).toBe(true);
1101+
expect(rotateTranscriptAfterCompactionMock).toHaveBeenCalledTimes(1);
1102+
expect(result.result?.sessionId).toBe("wrapper-rotated-session");
1103+
expect(result.result?.sessionFile).toBe("/tmp/wrapper-rotated-session.jsonl");
1104+
expect(maintain).toHaveBeenCalledWith(
1105+
expect.objectContaining({
1106+
sessionId: "wrapper-rotated-session",
1107+
sessionFile: "/tmp/wrapper-rotated-session.jsonl",
1108+
}),
1109+
);
1110+
});
1111+
10111112
it("catches and logs hook exceptions without aborting compaction", async () => {
10121113
hookRunner.hasHooks.mockReturnValue(true);
10131114
hookRunner.runBeforeCompaction.mockRejectedValue(new Error("hook boom"));

‎src/agents/pi-embedded-runner/compact.queued.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,9 @@ export async function compactEmbeddedPiSession(
164164
});
165165
const delegatedSessionId = result.result?.sessionId;
166166
const delegatedSessionFile = result.result?.sessionFile;
167-
const delegatedRotatedTranscript = Boolean(delegatedSessionId || delegatedSessionFile);
167+
const delegatedRotatedTranscript =
168+
(typeof delegatedSessionId === "string" && delegatedSessionId !== params.sessionId) ||
169+
(typeof delegatedSessionFile === "string" && delegatedSessionFile !== params.sessionFile);
168170
let postCompactionSessionId = delegatedSessionId ?? params.sessionId;
169171
let postCompactionSessionFile = delegatedSessionFile ?? params.sessionFile;
170172
let postCompactionLeafId: string | undefined;

‎src/agents/pi-embedded-runner/compact.ts‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1073,11 +1073,6 @@ export async function compactEmbeddedPiSessionDirect(
10731073
},
10741074
},
10751075
);
1076-
await runPostCompactionSideEffects({
1077-
config: params.config,
1078-
sessionKey: params.sessionKey,
1079-
sessionFile: params.sessionFile,
1080-
});
10811076
let effectiveFirstKeptEntryId = result.firstKeptEntryId;
10821077
let postCompactionLeafId =
10831078
typeof sessionManager.getLeafId === "function"
@@ -1135,12 +1130,12 @@ export async function compactEmbeddedPiSessionDirect(
11351130
`[compaction] rotated active transcript after compaction ` +
11361131
`(sessionKey=${params.sessionKey ?? params.sessionId})`,
11371132
);
1138-
await runPostCompactionSideEffects({
1139-
config: params.config,
1140-
sessionKey: params.sessionKey,
1141-
sessionFile: activeSessionFile,
1142-
});
11431133
}
1134+
await runPostCompactionSideEffects({
1135+
config: params.config,
1136+
sessionKey: params.sessionKey,
1137+
sessionFile: activeSessionFile,
1138+
});
11441139
if (params.config && params.sessionKey && checkpointSnapshot) {
11451140
try {
11461141
const storedCheckpoint = await persistSessionCompactionCheckpoint({

‎src/agents/pi-embedded-runner/compaction-successor-transcript.test.ts‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,55 @@ describe("rotateTranscriptAfterCompaction", () => {
163163
firstKeptEntryId: compactionId,
164164
});
165165
});
166+
167+
it("preserves unsummarized sibling branches and branch summaries", async () => {
168+
const dir = await createTmpDir();
169+
const manager = SessionManager.create(dir, dir);
170+
171+
manager.appendMessage({ role: "user", content: "hello", timestamp: 1 });
172+
const branchFromId = manager.appendMessage(makeAssistant("hi there", 2));
173+
174+
const branchSummaryId = manager.branchWithSummary(
175+
branchFromId,
176+
"Summary of the abandoned branch.",
177+
);
178+
const siblingMsgId = manager.appendMessage({
179+
role: "user",
180+
content: "do task B instead",
181+
timestamp: 3,
182+
});
183+
manager.appendMessage(makeAssistant("done B", 4));
184+
185+
manager.branch(branchFromId);
186+
manager.appendMessage({ role: "user", content: "do task A", timestamp: 5 });
187+
const firstKeptId = manager.appendMessage(makeAssistant("done A", 6));
188+
manager.appendCompaction("Summary of main branch.", firstKeptId, 5000);
189+
manager.appendMessage({ role: "user", content: "next", timestamp: 7 });
190+
191+
const sessionFile = manager.getSessionFile()!;
192+
const result = await rotateTranscriptAfterCompaction({
193+
sessionManager: manager,
194+
sessionFile,
195+
now: () => new Date("2026-04-27T12:45:00.000Z"),
196+
});
197+
198+
expect(result.rotated).toBe(true);
199+
const successor = SessionManager.open(result.sessionFile!);
200+
const allEntries = successor.getEntries();
201+
expect(allEntries.find((entry) => entry.id === branchSummaryId)).toMatchObject({
202+
type: "branch_summary",
203+
summary: "Summary of the abandoned branch.",
204+
});
205+
expect(allEntries.find((entry) => entry.id === siblingMsgId)).toMatchObject({
206+
type: "message",
207+
message: expect.objectContaining({ content: "do task B instead" }),
208+
});
209+
210+
const activeContextText = JSON.stringify(successor.buildSessionContext().messages);
211+
expect(activeContextText).toContain("Summary of main branch.");
212+
expect(activeContextText).toContain("next");
213+
expect(activeContextText).not.toContain("do task B instead");
214+
});
166215
});
167216

168217
describe("shouldRotateCompactionTranscript", () => {

0 commit comments

Comments
 (0)