Merge pull request #700 from clawdbot/shadow/compaction

Agents: safeguard compaction summarization

src/agents/pi-embedded-runner.ts

@@ -321,6 +321,12 @@ function buildContextPruningExtension(params: {
   };
 }
 
+function resolveCompactionMode(cfg?: ClawdbotConfig): "default" | "safeguard" {
+  return cfg?.agents?.defaults?.compaction?.mode === "safeguard"
+    ? "safeguard"
+    : "default";
+}
+
 function buildEmbeddedExtensionPaths(params: {
   cfg: ClawdbotConfig | undefined;
   sessionManager: SessionManager;
@@ -329,6 +335,9 @@ function buildEmbeddedExtensionPaths(params: {
   model: Model<Api> | undefined;
 }): string[] {
   const paths = [resolvePiExtensionPath("transcript-sanitize")];
+  if (resolveCompactionMode(params.cfg) === "safeguard") {
+    paths.push(resolvePiExtensionPath("compaction-safeguard"));
+  }
   const pruning = buildContextPruningExtension(params);
   if (pruning.additionalExtensionPaths) {
     paths.push(...pruning.additionalExtensionPaths);

src/agents/pi-extensions/compaction-safeguard.ts

@@ -0,0 +1,212 @@
+import type { AgentMessage } from "@mariozechner/pi-agent-core";
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+  FileOperations,
+} from "@mariozechner/pi-coding-agent";
+import { estimateTokens, generateSummary } from "@mariozechner/pi-coding-agent";
+
+import { DEFAULT_CONTEXT_TOKENS } from "../defaults.js";
+
+const MAX_CHUNK_RATIO = 0.4;
+const FALLBACK_SUMMARY =
+  "Summary unavailable due to context limits. Older messages were truncated.";
+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.";
+
+function computeFileLists(fileOps: FileOperations): {
+  readFiles: string[];
+  modifiedFiles: string[];
+} {
+  const modified = new Set([...fileOps.edited, ...fileOps.written]);
+  const readFiles = [...fileOps.read].filter((f) => !modified.has(f)).sort();
+  const modifiedFiles = [...modified].sort();
+  return { readFiles, modifiedFiles };
+}
+
+function formatFileOperations(
+  readFiles: string[],
+  modifiedFiles: string[],
+): string {
+  const sections: string[] = [];
+  if (readFiles.length > 0) {
+    sections.push(`<read-files>\n${readFiles.join("\n")}\n</read-files>`);
+  }
+  if (modifiedFiles.length > 0) {
+    sections.push(
+      `<modified-files>\n${modifiedFiles.join("\n")}\n</modified-files>`,
+    );
+  }
+  if (sections.length === 0) return "";
+  return `\n\n${sections.join("\n\n")}`;
+}
+
+function chunkMessages(
+  messages: AgentMessage[],
+  maxTokens: number,
+): AgentMessage[][] {
+  if (messages.length === 0) return [];
+
+  const chunks: AgentMessage[][] = [];
+  let currentChunk: AgentMessage[] = [];
+  let currentTokens = 0;
+
+  for (const message of messages) {
+    const messageTokens = estimateTokens(message);
+    if (currentChunk.length > 0 && currentTokens + messageTokens > maxTokens) {
+      chunks.push(currentChunk);
+      currentChunk = [];
+      currentTokens = 0;
+    }
+
+    currentChunk.push(message);
+    currentTokens += messageTokens;
+
+    if (messageTokens > maxTokens) {
+      // Split oversized messages to avoid unbounded chunk growth.
+      chunks.push(currentChunk);
+      currentChunk = [];
+      currentTokens = 0;
+    }
+  }
+
+  if (currentChunk.length > 0) {
+    chunks.push(currentChunk);
+  }
+
+  return chunks;
+}
+
+async function summarizeChunks(params: {
+  messages: AgentMessage[];
+  model: NonNullable<ExtensionContext["model"]>;
+  apiKey: string;
+  signal: AbortSignal;
+  reserveTokens: number;
+  maxChunkTokens: number;
+  customInstructions?: string;
+  previousSummary?: string;
+}): Promise<string> {
+  if (params.messages.length === 0) {
+    return params.previousSummary ?? "No prior history.";
+  }
+
+  const chunks = chunkMessages(params.messages, params.maxChunkTokens);
+  let summary = params.previousSummary;
+
+  for (const chunk of chunks) {
+    summary = await generateSummary(
+      chunk,
+      params.model,
+      params.reserveTokens,
+      params.apiKey,
+      params.signal,
+      params.customInstructions,
+      summary,
+    );
+  }
+
+  return summary ?? "No prior history.";
+}
+
+export default function compactionSafeguardExtension(api: ExtensionAPI): void {
+  api.on("session_before_compact", async (event, ctx) => {
+    const { preparation, customInstructions, signal } = event;
+    const { readFiles, modifiedFiles } = computeFileLists(preparation.fileOps);
+    const fileOpsSummary = formatFileOperations(readFiles, modifiedFiles);
+    const fallbackSummary = `${FALLBACK_SUMMARY}${fileOpsSummary}`;
+
+    const model = ctx.model;
+    if (!model) {
+      return {
+        compaction: {
+          summary: fallbackSummary,
+          firstKeptEntryId: preparation.firstKeptEntryId,
+          tokensBefore: preparation.tokensBefore,
+          details: { readFiles, modifiedFiles },
+        },
+      };
+    }
+
+    const apiKey = await ctx.modelRegistry.getApiKey(model);
+    if (!apiKey) {
+      return {
+        compaction: {
+          summary: fallbackSummary,
+          firstKeptEntryId: preparation.firstKeptEntryId,
+          tokensBefore: preparation.tokensBefore,
+          details: { readFiles, modifiedFiles },
+        },
+      };
+    }
+
+    try {
+      const contextWindowTokens = Math.max(
+        1,
+        Math.floor(model.contextWindow ?? DEFAULT_CONTEXT_TOKENS),
+      );
+      const maxChunkTokens = Math.max(
+        1,
+        Math.floor(contextWindowTokens * MAX_CHUNK_RATIO),
+      );
+      const reserveTokens = Math.max(
+        1,
+        Math.floor(preparation.settings.reserveTokens),
+      );
+
+      const historySummary = await summarizeChunks({
+        messages: preparation.messagesToSummarize,
+        model,
+        apiKey,
+        signal,
+        reserveTokens,
+        maxChunkTokens,
+        customInstructions,
+        previousSummary: preparation.previousSummary,
+      });
+
+      let summary = historySummary;
+      if (
+        preparation.isSplitTurn &&
+        preparation.turnPrefixMessages.length > 0
+      ) {
+        const prefixSummary = await summarizeChunks({
+          messages: preparation.turnPrefixMessages,
+          model,
+          apiKey,
+          signal,
+          reserveTokens,
+          maxChunkTokens,
+          customInstructions: TURN_PREFIX_INSTRUCTIONS,
+        });
+        summary = `${historySummary}\n\n---\n\n**Turn Context (split turn):**\n\n${prefixSummary}`;
+      }
+
+      summary += fileOpsSummary;
+
+      return {
+        compaction: {
+          summary,
+          firstKeptEntryId: preparation.firstKeptEntryId,
+          tokensBefore: preparation.tokensBefore,
+          details: { readFiles, modifiedFiles },
+        },
+      };
+    } catch (error) {
+      console.warn(
+        `Compaction summarization failed; truncating history: ${
+          error instanceof Error ? error.message : String(error)
+        }`,
+      );
+      return {
+        compaction: {
+          summary: fallbackSummary,
+          firstKeptEntryId: preparation.firstKeptEntryId,
+          tokensBefore: preparation.tokensBefore,
+          details: { readFiles, modifiedFiles },
+        },
+      };
+    }
+  });
+}

src/config/config.test.ts

@@ -522,6 +522,7 @@ describe("config compaction settings", () => {
             agents: {
               defaults: {
                 compaction: {
+                  mode: "safeguard",
                   reserveTokensFloor: 12_345,
                   memoryFlush: {
                     enabled: false,
@@ -544,6 +545,7 @@ describe("config compaction settings", () => {
       const cfg = loadConfig();
 
       expect(cfg.agents?.defaults?.compaction?.reserveTokensFloor).toBe(12_345);
+      expect(cfg.agents?.defaults?.compaction?.mode).toBe("safeguard");
       expect(cfg.agents?.defaults?.compaction?.memoryFlush?.enabled).toBe(
         false,
       );

src/config/types.ts

@@ -1716,7 +1716,11 @@ export type AgentDefaultsConfig = {
   };
 };
 
+export type AgentCompactionMode = "default" | "safeguard";
+
 export type AgentCompactionConfig = {
+  /** Compaction summarization mode. */
+  mode?: AgentCompactionMode;
   /** Minimum reserve tokens enforced for Pi compaction (0 disables the floor). */
   reserveTokensFloor?: number;
   /** Pre-compaction memory flush (agentic turn). Default: enabled. */

src/config/zod-schema.ts

@@ -1210,6 +1210,9 @@ const AgentDefaultsSchema = z
       .optional(),
     compaction: z
       .object({
+        mode: z
+          .union([z.literal("default"), z.literal("safeguard")])
+          .optional(),
         reserveTokensFloor: z.number().int().nonnegative().optional(),
         memoryFlush: z
           .object({