From bfa83bd2c5e9942a7769592bfd2ba85e0f83f427 Mon Sep 17 00:00:00 2001 From: Lukin Date: Mon, 26 Jan 2026 09:23:16 +0800 Subject: [PATCH] feat(compaction): re-inject workspace context files after compaction This change ensures that USER.md, SOUL.md, and IDENTITY.md are automatically re-injected into the compaction summary, so the agent retains persona and user preferences after context compression. Changes: - Add context file re-injection in compaction-safeguard.ts - Wire up after_compaction plugin hook with result type for additional injection - Allow plugins to append custom text to compaction summary via after_compaction hook Closes: context compression causing loss of user preferences --- .../pi-extensions/compaction-safeguard.ts | 88 +++++++++++++++++++ src/plugins/hooks.ts | 15 +++- src/plugins/types.ts | 12 ++- 3 files changed, 112 insertions(+), 3 deletions(-) diff --git a/src/agents/pi-extensions/compaction-safeguard.ts b/src/agents/pi-extensions/compaction-safeguard.ts index 82ad19f2a..4deededfd 100644 --- a/src/agents/pi-extensions/compaction-safeguard.ts +++ b/src/agents/pi-extensions/compaction-safeguard.ts @@ -1,3 +1,6 @@ +import fs from "node:fs/promises"; +import path from "node:path"; + import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { ExtensionAPI, FileOperations } from "@mariozechner/pi-coding-agent"; import { @@ -11,8 +14,52 @@ import { resolveContextWindowTokens, summarizeInStages, } from "../compaction.js"; +import { getGlobalHookRunner } from "../../plugins/hook-runner-global.js"; +import { + DEFAULT_USER_FILENAME, + DEFAULT_SOUL_FILENAME, + DEFAULT_IDENTITY_FILENAME, +} from "../workspace.js"; + const FALLBACK_SUMMARY = "Summary unavailable due to context limits. Older messages were truncated."; + +// Context refresh settings +const MAX_CONTEXT_FILE_SIZE = 4000; // Max chars per file to avoid bloating the summary +const CONTEXT_FILES = [DEFAULT_USER_FILENAME, DEFAULT_SOUL_FILENAME, DEFAULT_IDENTITY_FILENAME]; + +async function readContextFile( + workspaceRoot: string, + filename: string, + maxSize: number = MAX_CONTEXT_FILE_SIZE, +): Promise { + try { + const filePath = path.join(workspaceRoot, filename); + const content = await fs.readFile(filePath, "utf-8"); + if (content.trim().length === 0) return undefined; + if (content.length > maxSize) { + return content.slice(0, maxSize) + "\n\n[... truncated for context limits ...]"; + } + return content; + } catch { + return undefined; + } +} + +async function buildContextRefreshSection(workspaceRoot: string | undefined): Promise { + if (!workspaceRoot) return ""; + + const sections: string[] = []; + for (const filename of CONTEXT_FILES) { + const content = await readContextFile(workspaceRoot, filename); + if (content) { + sections.push(`## ${filename}\n${content}`); + } + } + + if (sections.length === 0) return ""; + return `\n\n# Workspace Context (re-injected after compaction)\n\n${sections.join("\n\n")}`; +} const TURN_PREFIX_INSTRUCTIONS = "This summary covers the prefix of a split turn. Focus on the original request," + " early progress, and any details needed to understand the retained suffix."; @@ -241,6 +288,47 @@ export default function compactionSafeguardExtension(api: ExtensionAPI): void { summary += toolFailureSection; summary += fileOpsSummary; + // Re-inject workspace context files (USER.md, SOUL.md, etc.) after compaction + // This ensures the agent retains persona and user preferences + try { + const contextRefresh = await buildContextRefreshSection(ctx.cwd); + if (contextRefresh) { + summary += contextRefresh; + } + } catch (contextError) { + console.warn( + `Context refresh after compaction failed: ${ + contextError instanceof Error ? contextError.message : String(contextError) + }`, + ); + } + + // Run after_compaction hook to allow plugins to inject additional context + const hookRunner = getGlobalHookRunner(); + if (hookRunner?.hasHooks("after_compaction")) { + try { + const hookResult = await hookRunner.runAfterCompaction( + { + messageCount: preparation.messagesToSummarize.length, + tokenCount: preparation.tokensBefore, + compactedCount: preparation.messagesToSummarize.length, + summary, + workspaceRoot: ctx.cwd, + }, + { agentId: undefined }, + ); + if (hookResult?.appendToSummary) { + summary += `\n\n${hookResult.appendToSummary}`; + } + } catch (hookError) { + console.warn( + `after_compaction hook failed: ${ + hookError instanceof Error ? hookError.message : String(hookError) + }`, + ); + } + } + return { compaction: { summary, diff --git a/src/plugins/hooks.ts b/src/plugins/hooks.ts index 41fd501f7..6b6ad2207 100644 --- a/src/plugins/hooks.ts +++ b/src/plugins/hooks.ts @@ -8,6 +8,7 @@ import type { PluginRegistry } from "./registry.js"; import type { PluginHookAfterCompactionEvent, + PluginHookAfterCompactionResult, PluginHookAfterToolCallEvent, PluginHookAgentContext, PluginHookAgentEndEvent, @@ -43,6 +44,7 @@ export type { PluginHookAgentEndEvent, PluginHookBeforeCompactionEvent, PluginHookAfterCompactionEvent, + PluginHookAfterCompactionResult, PluginHookMessageContext, PluginHookMessageReceivedEvent, PluginHookMessageSendingEvent, @@ -218,12 +220,21 @@ export function createHookRunner(registry: PluginRegistry, options: HookRunnerOp /** * Run after_compaction hook. + * Allows plugins to return text to append to the compaction summary. */ async function runAfterCompaction( event: PluginHookAfterCompactionEvent, ctx: PluginHookAgentContext, - ): Promise { - return runVoidHook("after_compaction", event, ctx); + ): Promise { + return runModifyingHook<"after_compaction", PluginHookAfterCompactionResult>( + "after_compaction", + event, + ctx, + (acc, next) => ({ + appendToSummary: + [acc?.appendToSummary, next.appendToSummary].filter(Boolean).join("\n\n") || undefined, + }), + ); } // ========================================================================= diff --git a/src/plugins/types.ts b/src/plugins/types.ts index 1ce9731ea..db5f3a726 100644 --- a/src/plugins/types.ts +++ b/src/plugins/types.ts @@ -339,6 +339,16 @@ export type PluginHookAfterCompactionEvent = { messageCount: number; tokenCount?: number; compactedCount: number; + /** The generated summary text (read-only). */ + summary?: string; + /** Agent workspace root (for reading context files). */ + workspaceRoot?: string; +}; + +/** Result from after_compaction hook - allows injecting additional context into the summary. */ +export type PluginHookAfterCompactionResult = { + /** Additional text to append to the compaction summary. */ + appendToSummary?: string; }; // Message context @@ -476,7 +486,7 @@ export type PluginHookHandlerMap = { after_compaction: ( event: PluginHookAfterCompactionEvent, ctx: PluginHookAgentContext, - ) => Promise | void; + ) => Promise | PluginHookAfterCompactionResult | void; message_received: ( event: PluginHookMessageReceivedEvent, ctx: PluginHookMessageContext,