From bf9264cb8bafa3291f8a3fe4b59632f4ad685ec6 Mon Sep 17 00:00:00 2001 From: Thirumalesh Date: Mon, 9 Mar 2026 18:34:52 +0530 Subject: [PATCH] feat(compaction): truncate session JSONL after compaction to prevent unbounded growth Session files grow indefinitely because compaction summaries are appended without removing the entries they summarize. After many compaction cycles, files can reach 7+ MB with thousands of stale entries, causing performance degradation and stuck sessions. Add opt-in config under that rewrites the session file after compaction to keep only the latest compaction entry and subsequent messages, removing all summarized entries. Fixes #39953 --- src/agents/pi-embedded-runner/compact.ts | 20 +++ .../session-truncation.test.ts | 166 ++++++++++++++++++ .../pi-embedded-runner/session-truncation.ts | 142 +++++++++++++++ src/config/schema.help.quality.test.ts | 1 + src/config/schema.help.ts | 2 + src/config/schema.labels.ts | 1 + src/config/types.agent-defaults.ts | 6 + 7 files changed, 338 insertions(+) create mode 100644 src/agents/pi-embedded-runner/session-truncation.test.ts create mode 100644 src/agents/pi-embedded-runner/session-truncation.ts diff --git a/src/agents/pi-embedded-runner/compact.ts b/src/agents/pi-embedded-runner/compact.ts index 0dfc727dee1..6c753e9d723 100644 --- a/src/agents/pi-embedded-runner/compact.ts +++ b/src/agents/pi-embedded-runner/compact.ts @@ -96,6 +96,7 @@ import { buildEmbeddedMessageActionDiscoveryInput } from "./message-action-disco import { buildModelAliasLines, resolveModelAsync } from "./model.js"; import { buildEmbeddedSandboxInfo } from "./sandbox-info.js"; import { prewarmSessionFile, trackSessionManagerAccess } from "./session-manager-cache.js"; +import { truncateSessionAfterCompaction } from "./session-truncation.js"; import { resolveEmbeddedRunSkillEntries } from "./skills-runtime.js"; import { applySystemPromptOverrideToSession, @@ -1085,6 +1086,25 @@ export async function compactEmbeddedPiSessionDirect( }); } } + // Truncate session file to remove compacted entries (#39953) + if (params.config?.agents?.defaults?.compaction?.truncateAfterCompaction) { + try { + const truncResult = await truncateSessionAfterCompaction({ + sessionFile: params.sessionFile, + }); + if (truncResult.truncated) { + log.info( + `[compaction] post-compaction truncation removed ${truncResult.entriesRemoved} entries ` + + `(sessionKey=${params.sessionKey ?? params.sessionId})`, + ); + } + } catch (err) { + log.warn("[compaction] post-compaction truncation failed", { + errorMessage: err instanceof Error ? err.message : String(err), + errorStack: err instanceof Error ? err.stack : undefined, + }); + } + } return { ok: true, compacted: true, diff --git a/src/agents/pi-embedded-runner/session-truncation.test.ts b/src/agents/pi-embedded-runner/session-truncation.test.ts new file mode 100644 index 00000000000..4661cccad70 --- /dev/null +++ b/src/agents/pi-embedded-runner/session-truncation.test.ts @@ -0,0 +1,166 @@ +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { SessionManager } from "@mariozechner/pi-coding-agent"; +import { afterEach, describe, expect, it } from "vitest"; +import { makeAgentAssistantMessage } from "../test-helpers/agent-message-fixtures.js"; +import { truncateSessionAfterCompaction } from "./session-truncation.js"; + +let tmpDir: string; + +async function createTmpDir(): Promise { + tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), "session-truncation-test-")); + return tmpDir; +} + +afterEach(async () => { + if (tmpDir) { + await fs.rm(tmpDir, { recursive: true, force: true }).catch(() => {}); + } +}); + +function makeAssistant(text: string, timestamp: number) { + return makeAgentAssistantMessage({ + content: [{ type: "text", text }], + timestamp, + }); +} + +function createSessionWithCompaction(sessionDir: string): string { + const sm = SessionManager.create(sessionDir, sessionDir); + // Add messages before compaction + sm.appendMessage({ role: "user", content: "hello", timestamp: 1 }); + sm.appendMessage(makeAssistant("hi there", 2)); + sm.appendMessage({ role: "user", content: "do something", timestamp: 3 }); + sm.appendMessage(makeAssistant("done", 4)); + + // Add compaction (summarizing the above) + const branch = sm.getBranch(); + const firstKeptId = branch[branch.length - 1].id; + sm.appendCompaction("Summary of conversation so far.", firstKeptId, 5000); + + // Add messages after compaction + sm.appendMessage({ role: "user", content: "next task", timestamp: 5 }); + sm.appendMessage(makeAssistant("working on it", 6)); + + return sm.getSessionFile()!; +} + +describe("truncateSessionAfterCompaction", () => { + it("removes entries before compaction and keeps entries after (#39953)", async () => { + const dir = await createTmpDir(); + const sessionFile = createSessionWithCompaction(dir); + + // Verify pre-truncation state + const smBefore = SessionManager.open(sessionFile); + const entriesBefore = smBefore.getEntries().length; + expect(entriesBefore).toBeGreaterThan(5); // 4 messages + compaction + 2 messages + + const result = await truncateSessionAfterCompaction({ sessionFile }); + + expect(result.truncated).toBe(true); + expect(result.entriesRemoved).toBeGreaterThan(0); + expect(result.bytesAfter).toBeLessThan(result.bytesBefore!); + + // Verify post-truncation: file is still a valid session + const smAfter = SessionManager.open(sessionFile); + const entriesAfter = smAfter.getEntries().length; + expect(entriesAfter).toBeLessThan(entriesBefore); + + // The branch should contain compaction + post-compaction messages + const branchAfter = smAfter.getBranch(); + expect(branchAfter[0].type).toBe("compaction"); + expect(branchAfter[0].parentId).toBeNull(); + + // Session context should still work + const ctx = smAfter.buildSessionContext(); + expect(ctx.messages.length).toBeGreaterThan(0); + }); + + it("skips truncation when no compaction entry exists", async () => { + const dir = await createTmpDir(); + const sm = SessionManager.create(dir, dir); + // appendMessage implicitly creates the session file + sm.appendMessage({ role: "user", content: "hello", timestamp: 1 }); + sm.appendMessage(makeAssistant("hi", 2)); + sm.appendMessage({ role: "user", content: "bye", timestamp: 3 }); + const sessionFile = sm.getSessionFile()!; + + const result = await truncateSessionAfterCompaction({ sessionFile }); + + expect(result.truncated).toBe(false); + expect(result.reason).toBe("no compaction entry found"); + }); + + it("is idempotent — second truncation is a no-op", async () => { + const dir = await createTmpDir(); + const sessionFile = createSessionWithCompaction(dir); + + const first = await truncateSessionAfterCompaction({ sessionFile }); + expect(first.truncated).toBe(true); + + // Run again — compaction is now at root, nothing more to remove + const second = await truncateSessionAfterCompaction({ sessionFile }); + expect(second.truncated).toBe(false); + expect(second.reason).toBe("compaction already at root"); + }); + + it("archives original file when archivePath is provided (#39953)", async () => { + const dir = await createTmpDir(); + const sessionFile = createSessionWithCompaction(dir); + const archivePath = path.join(dir, "archive", "backup.jsonl"); + + const result = await truncateSessionAfterCompaction({ sessionFile, archivePath }); + + expect(result.truncated).toBe(true); + const archiveExists = await fs + .stat(archivePath) + .then(() => true) + .catch(() => false); + expect(archiveExists).toBe(true); + + // Archive should be larger than truncated file (it has the full history) + const archiveSize = (await fs.stat(archivePath)).size; + const truncatedSize = (await fs.stat(sessionFile)).size; + expect(archiveSize).toBeGreaterThan(truncatedSize); + }); + + it("handles multiple compaction cycles (#39953)", async () => { + const dir = await createTmpDir(); + const sm = SessionManager.create(dir, dir); + + // First cycle: messages + compaction + sm.appendMessage({ role: "user", content: "cycle 1 message 1", timestamp: 1 }); + sm.appendMessage(makeAssistant("response 1", 2)); + const branch1 = sm.getBranch(); + sm.appendCompaction("Summary of cycle 1.", branch1[branch1.length - 1].id, 3000); + + // Second cycle: more messages + another compaction + sm.appendMessage({ role: "user", content: "cycle 2 message 1", timestamp: 3 }); + sm.appendMessage(makeAssistant("response 2", 4)); + const branch2 = sm.getBranch(); + sm.appendCompaction("Summary of cycles 1 and 2.", branch2[branch2.length - 1].id, 6000); + + // Post-compaction messages + sm.appendMessage({ role: "user", content: "final question", timestamp: 5 }); + + const sessionFile = sm.getSessionFile()!; + const entriesBefore = sm.getEntries().length; + + const result = await truncateSessionAfterCompaction({ sessionFile }); + + expect(result.truncated).toBe(true); + + // Should keep only the latest compaction + entries after it + const smAfter = SessionManager.open(sessionFile); + const branchAfter = smAfter.getBranch(); + expect(branchAfter[0].type).toBe("compaction"); + + // Only the latest compaction should remain + const compactionEntries = branchAfter.filter((e) => e.type === "compaction"); + expect(compactionEntries).toHaveLength(1); + + const entriesAfter = smAfter.getEntries().length; + expect(entriesAfter).toBeLessThan(entriesBefore); + }); +}); diff --git a/src/agents/pi-embedded-runner/session-truncation.ts b/src/agents/pi-embedded-runner/session-truncation.ts new file mode 100644 index 00000000000..77619db8e50 --- /dev/null +++ b/src/agents/pi-embedded-runner/session-truncation.ts @@ -0,0 +1,142 @@ +import fs from "node:fs/promises"; +import path from "node:path"; +import type { CompactionEntry, SessionEntry } from "@mariozechner/pi-coding-agent"; +import { SessionManager } from "@mariozechner/pi-coding-agent"; +import { log } from "./logger.js"; + +/** + * Truncate a session JSONL file after compaction by removing entries + * that are no longer reachable from the current branch. + * + * After compaction, the session file still contains all historical entries + * even though `buildSessionContext()` logically skips entries before + * `firstKeptEntryId`. Over many compaction cycles this causes unbounded + * file growth (issue #39953). + * + * This function rewrites the file to keep only: + * 1. The session header + * 2. The latest compaction entry (re-parented as root) + * 3. All entries after the compaction in the current branch + */ +export async function truncateSessionAfterCompaction(params: { + sessionFile: string; + /** Optional path to archive the pre-truncation file. */ + archivePath?: string; +}): Promise { + const { sessionFile } = params; + + let sm: SessionManager; + try { + sm = SessionManager.open(sessionFile); + } catch (err) { + const reason = err instanceof Error ? err.message : String(err); + log.warn(`[session-truncation] Failed to open session file: ${reason}`); + return { truncated: false, entriesRemoved: 0, reason }; + } + + const header = sm.getHeader(); + if (!header) { + return { truncated: false, entriesRemoved: 0, reason: "missing session header" }; + } + + const branch = sm.getBranch(); + if (branch.length === 0) { + return { truncated: false, entriesRemoved: 0, reason: "empty session" }; + } + + // Find the latest compaction entry in the current branch + let latestCompactionIdx = -1; + for (let i = branch.length - 1; i >= 0; i--) { + if (branch[i].type === "compaction") { + latestCompactionIdx = i; + break; + } + } + + if (latestCompactionIdx < 0) { + return { truncated: false, entriesRemoved: 0, reason: "no compaction entry found" }; + } + + // Nothing to truncate if compaction is already at root + if (latestCompactionIdx === 0) { + return { truncated: false, entriesRemoved: 0, reason: "compaction already at root" }; + } + + const entriesRemoved = latestCompactionIdx; + const totalEntriesBefore = sm.getEntries().length; + + // Build the truncated entry list: + // compaction entry (re-parented as root) + all entries after it + const truncatedEntries: SessionEntry[] = []; + + const compactionEntry = branch[latestCompactionIdx] as CompactionEntry; + truncatedEntries.push({ ...compactionEntry, parentId: null }); + + for (let i = latestCompactionIdx + 1; i < branch.length; i++) { + truncatedEntries.push(branch[i]); + } + + // Get file size before truncation + let bytesBefore = 0; + try { + const stat = await fs.stat(sessionFile); + bytesBefore = stat.size; + } catch { + // If stat fails, continue anyway + } + + // Archive original file if requested + if (params.archivePath) { + try { + const archiveDir = path.dirname(params.archivePath); + await fs.mkdir(archiveDir, { recursive: true }); + await fs.copyFile(sessionFile, params.archivePath); + log.info(`[session-truncation] Archived pre-truncation file to ${params.archivePath}`); + } catch (err) { + const reason = err instanceof Error ? err.message : String(err); + log.warn(`[session-truncation] Failed to archive: ${reason}`); + } + } + + // Write truncated file atomically (temp + rename) + const lines: string[] = [ + JSON.stringify(header), + ...truncatedEntries.map((e) => JSON.stringify(e)), + ]; + const content = lines.join("\n") + "\n"; + + const tmpFile = `${sessionFile}.truncate-tmp`; + try { + await fs.writeFile(tmpFile, content, "utf-8"); + await fs.rename(tmpFile, sessionFile); + } catch (err) { + // Clean up temp file on failure + try { + await fs.unlink(tmpFile); + } catch { + // Ignore cleanup errors + } + const reason = err instanceof Error ? err.message : String(err); + log.warn(`[session-truncation] Failed to write truncated file: ${reason}`); + return { truncated: false, entriesRemoved: 0, reason }; + } + + const bytesAfter = Buffer.byteLength(content, "utf-8"); + + log.info( + `[session-truncation] Truncated session file: ` + + `entriesBefore=${totalEntriesBefore} entriesAfter=${truncatedEntries.length} ` + + `removed=${entriesRemoved} bytesBefore=${bytesBefore} bytesAfter=${bytesAfter} ` + + `reduction=${bytesBefore > 0 ? ((1 - bytesAfter / bytesBefore) * 100).toFixed(1) : "?"}%`, + ); + + return { truncated: true, entriesRemoved, bytesBefore, bytesAfter }; +} + +export type TruncationResult = { + truncated: boolean; + entriesRemoved: number; + bytesBefore?: number; + bytesAfter?: number; + reason?: string; +}; diff --git a/src/config/schema.help.quality.test.ts b/src/config/schema.help.quality.test.ts index f1542bcb7de..18e1947d88f 100644 --- a/src/config/schema.help.quality.test.ts +++ b/src/config/schema.help.quality.test.ts @@ -390,6 +390,7 @@ const TARGET_KEYS = [ "agents.defaults.compaction.postCompactionSections", "agents.defaults.compaction.timeoutSeconds", "agents.defaults.compaction.model", + "agents.defaults.compaction.truncateAfterCompaction", "agents.defaults.compaction.memoryFlush", "agents.defaults.compaction.memoryFlush.enabled", "agents.defaults.compaction.memoryFlush.softThresholdTokens", diff --git a/src/config/schema.help.ts b/src/config/schema.help.ts index bcaec953d57..c22d5e15b32 100644 --- a/src/config/schema.help.ts +++ b/src/config/schema.help.ts @@ -1050,6 +1050,8 @@ export const FIELD_HELP: Record = { "Maximum time in seconds allowed for a single compaction operation before it is aborted (default: 900). Increase this for very large sessions that need more time to summarize, or decrease it to fail faster on unresponsive models.", "agents.defaults.compaction.model": "Optional provider/model override used only for compaction summarization. Set this when you want compaction to run on a different model than the session default, and leave it unset to keep using the primary agent model.", + "agents.defaults.compaction.truncateAfterCompaction": + "When enabled, rewrites the session JSONL file after compaction to remove entries that were summarized. Prevents unbounded file growth in long-running sessions with many compaction cycles. Default: false.", "agents.defaults.compaction.memoryFlush": "Pre-compaction memory flush settings that run an agentic memory write before heavy compaction. Keep enabled for long sessions so salient context is persisted before aggressive trimming.", "agents.defaults.compaction.memoryFlush.enabled": diff --git a/src/config/schema.labels.ts b/src/config/schema.labels.ts index 854975b5a9c..53317e2fcd2 100644 --- a/src/config/schema.labels.ts +++ b/src/config/schema.labels.ts @@ -467,6 +467,7 @@ export const FIELD_LABELS: Record = { "agents.defaults.compaction.postCompactionSections": "Post-Compaction Context Sections", "agents.defaults.compaction.timeoutSeconds": "Compaction Timeout (Seconds)", "agents.defaults.compaction.model": "Compaction Model Override", + "agents.defaults.compaction.truncateAfterCompaction": "Truncate After Compaction", "agents.defaults.compaction.memoryFlush": "Compaction Memory Flush", "agents.defaults.compaction.memoryFlush.enabled": "Compaction Memory Flush Enabled", "agents.defaults.compaction.memoryFlush.softThresholdTokens": diff --git a/src/config/types.agent-defaults.ts b/src/config/types.agent-defaults.ts index 604bf88bdcb..ecaaecb69b9 100644 --- a/src/config/types.agent-defaults.ts +++ b/src/config/types.agent-defaults.ts @@ -342,6 +342,12 @@ export type AgentCompactionConfig = { model?: string; /** Maximum time in seconds for a single compaction operation (default: 900). */ timeoutSeconds?: number; + /** + * Truncate the session JSONL file after compaction to remove entries that + * were summarized. Prevents unbounded file growth in long-running sessions. + * Default: false (existing behavior preserved). + */ + truncateAfterCompaction?: boolean; }; export type AgentCompactionMemoryFlushConfig = {