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
This commit is contained in:
parent
9c8e8c5c2d
commit
bfa83bd2c5
@ -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 { AgentMessage } from "@mariozechner/pi-agent-core";
|
||||||
import type { ExtensionAPI, FileOperations } from "@mariozechner/pi-coding-agent";
|
import type { ExtensionAPI, FileOperations } from "@mariozechner/pi-coding-agent";
|
||||||
import {
|
import {
|
||||||
@ -11,8 +14,52 @@ import {
|
|||||||
resolveContextWindowTokens,
|
resolveContextWindowTokens,
|
||||||
summarizeInStages,
|
summarizeInStages,
|
||||||
} from "../compaction.js";
|
} 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 =
|
const FALLBACK_SUMMARY =
|
||||||
"Summary unavailable due to context limits. Older messages were truncated.";
|
"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<string | undefined> {
|
||||||
|
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<string> {
|
||||||
|
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 =
|
const TURN_PREFIX_INSTRUCTIONS =
|
||||||
"This summary covers the prefix of a split turn. Focus on the original request," +
|
"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.";
|
" 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 += toolFailureSection;
|
||||||
summary += fileOpsSummary;
|
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 {
|
return {
|
||||||
compaction: {
|
compaction: {
|
||||||
summary,
|
summary,
|
||||||
|
|||||||
@ -8,6 +8,7 @@
|
|||||||
import type { PluginRegistry } from "./registry.js";
|
import type { PluginRegistry } from "./registry.js";
|
||||||
import type {
|
import type {
|
||||||
PluginHookAfterCompactionEvent,
|
PluginHookAfterCompactionEvent,
|
||||||
|
PluginHookAfterCompactionResult,
|
||||||
PluginHookAfterToolCallEvent,
|
PluginHookAfterToolCallEvent,
|
||||||
PluginHookAgentContext,
|
PluginHookAgentContext,
|
||||||
PluginHookAgentEndEvent,
|
PluginHookAgentEndEvent,
|
||||||
@ -43,6 +44,7 @@ export type {
|
|||||||
PluginHookAgentEndEvent,
|
PluginHookAgentEndEvent,
|
||||||
PluginHookBeforeCompactionEvent,
|
PluginHookBeforeCompactionEvent,
|
||||||
PluginHookAfterCompactionEvent,
|
PluginHookAfterCompactionEvent,
|
||||||
|
PluginHookAfterCompactionResult,
|
||||||
PluginHookMessageContext,
|
PluginHookMessageContext,
|
||||||
PluginHookMessageReceivedEvent,
|
PluginHookMessageReceivedEvent,
|
||||||
PluginHookMessageSendingEvent,
|
PluginHookMessageSendingEvent,
|
||||||
@ -218,12 +220,21 @@ export function createHookRunner(registry: PluginRegistry, options: HookRunnerOp
|
|||||||
|
|
||||||
/**
|
/**
|
||||||
* Run after_compaction hook.
|
* Run after_compaction hook.
|
||||||
|
* Allows plugins to return text to append to the compaction summary.
|
||||||
*/
|
*/
|
||||||
async function runAfterCompaction(
|
async function runAfterCompaction(
|
||||||
event: PluginHookAfterCompactionEvent,
|
event: PluginHookAfterCompactionEvent,
|
||||||
ctx: PluginHookAgentContext,
|
ctx: PluginHookAgentContext,
|
||||||
): Promise<void> {
|
): Promise<PluginHookAfterCompactionResult | undefined> {
|
||||||
return runVoidHook("after_compaction", event, ctx);
|
return runModifyingHook<"after_compaction", PluginHookAfterCompactionResult>(
|
||||||
|
"after_compaction",
|
||||||
|
event,
|
||||||
|
ctx,
|
||||||
|
(acc, next) => ({
|
||||||
|
appendToSummary:
|
||||||
|
[acc?.appendToSummary, next.appendToSummary].filter(Boolean).join("\n\n") || undefined,
|
||||||
|
}),
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
// =========================================================================
|
// =========================================================================
|
||||||
|
|||||||
@ -339,6 +339,16 @@ export type PluginHookAfterCompactionEvent = {
|
|||||||
messageCount: number;
|
messageCount: number;
|
||||||
tokenCount?: number;
|
tokenCount?: number;
|
||||||
compactedCount: 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
|
// Message context
|
||||||
@ -476,7 +486,7 @@ export type PluginHookHandlerMap = {
|
|||||||
after_compaction: (
|
after_compaction: (
|
||||||
event: PluginHookAfterCompactionEvent,
|
event: PluginHookAfterCompactionEvent,
|
||||||
ctx: PluginHookAgentContext,
|
ctx: PluginHookAgentContext,
|
||||||
) => Promise<void> | void;
|
) => Promise<PluginHookAfterCompactionResult | void> | PluginHookAfterCompactionResult | void;
|
||||||
message_received: (
|
message_received: (
|
||||||
event: PluginHookMessageReceivedEvent,
|
event: PluginHookMessageReceivedEvent,
|
||||||
ctx: PluginHookMessageContext,
|
ctx: PluginHookMessageContext,
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user