fix: prune context and compact earlier

index.ts

@@ -71,6 +71,10 @@ function rewriteContextMessage(message: { role: string; content: string; origina
   } as AgentMessage;
 }
 
+function isRawSummaryArtifactMessage(message: AgentMessage): boolean {
+  return message.role === "compactionSummary" || message.role === "branchSummary";
+}
+
 function findLatestSnapshotState(branch: BranchEntry[]): { snapshot: RuntimeSnapshot; index: number } | undefined {
   for (let index = branch.length - 1; index >= 0; index -= 1) {
     const entry = branch[index]!;
@@ -147,6 +151,7 @@ export default function contextManager(pi: ExtensionAPI) {
     contextWindow: 200_000,
   });
   let pendingResumeInjection = false;
+  let earlyCompactionRequested = false;
 
   const syncContextWindow = (ctx: Pick<ExtensionContext, "model">) => {
     runtime.setContextWindow(ctx.model?.contextWindow ?? 200_000);
@@ -180,7 +185,7 @@ export default function contextManager(pi: ExtensionAPI) {
   };
 
   const rebuildRuntimeFromBranch = (
-    ctx: Pick<ExtensionContext, "model" | "sessionManager" | "ui">,
+    ctx: Pick<ExtensionContext, "model" | "sessionManager">,
     fallbackSnapshot: RuntimeSnapshot,
     options?: { preferRuntimeMode?: boolean },
   ) => {
@@ -202,9 +207,26 @@ export default function contextManager(pi: ExtensionAPI) {
     for (const entry of replayEntries) {
       replayBranchEntry(entry);
     }
+  };
 
-    const snapshot = runtime.getSnapshot();
-    ctx.ui.setStatus("context-manager", `ctx ${snapshot.lastZone}`);
+  const maybeTriggerEarlyCompaction = (ctx: Pick<ExtensionContext, "compact" | "getContextUsage">) => {
+    const usage = ctx.getContextUsage();
+    if (usage?.tokens === null || usage?.tokens === undefined) {
+      return;
+    }
+
+    const zone = runtime.getSnapshot().lastZone;
+    if (zone === "green" || zone === "yellow") {
+      earlyCompactionRequested = false;
+      return;
+    }
+
+    if (earlyCompactionRequested) {
+      return;
+    }
+
+    earlyCompactionRequested = true;
+    ctx.compact();
   };
 
   registerContextCommands(pi, {
@@ -220,11 +242,13 @@ export default function contextManager(pi: ExtensionAPI) {
   });
 
   pi.on("session_start", async (_event, ctx) => {
+    earlyCompactionRequested = false;
     rebuildRuntimeFromBranch(ctx, createDefaultSnapshot());
     armResumeInjection();
   });
 
   pi.on("session_tree", async (event, ctx) => {
+    earlyCompactionRequested = false;
     rebuildRuntimeFromBranch(ctx, createDefaultSnapshot());
 
     if (
@@ -260,19 +284,21 @@ export default function contextManager(pi: ExtensionAPI) {
 
     const snapshot = runtime.getSnapshot();
     pi.appendEntry(SNAPSHOT_ENTRY_TYPE, serializeSnapshot(snapshot));
-    ctx.ui.setStatus("context-manager", `ctx ${snapshot.lastZone}`);
+    maybeTriggerEarlyCompaction(ctx);
   });
 
   pi.on("context", async (event, ctx) => {
     syncContextWindow(ctx);
     const snapshot = runtime.getSnapshot();
     const policy = adjustPolicyForZone(runtime.getPolicy(), snapshot.lastZone);
-    const normalized = event.messages.map((message) => ({
-      role: message.role,
-      content: getMessageContent(message),
-      toolName: getMessageToolName(message),
-      original: message,
-    }));
+    const normalized = event.messages
+      .filter((message) => !isRawSummaryArtifactMessage(message as AgentMessage))
+      .map((message) => ({
+        role: message.role,
+        content: getMessageContent(message),
+        toolName: getMessageToolName(message),
+        original: message,
+      }));
 
     const pruned = pruneContextMessages(normalized, policy);
     const nextMessages = pruned.map((message) =>
@@ -344,10 +370,10 @@ export default function contextManager(pi: ExtensionAPI) {
     };
   });
 
-  pi.on("session_compact", async (event, ctx) => {
+  pi.on("session_compact", async (event) => {
+    earlyCompactionRequested = false;
     runtime.recordCompactionSummary(event.compactionEntry.summary, event.compactionEntry.id, Date.parse(event.compactionEntry.timestamp));
     pi.appendEntry(SNAPSHOT_ENTRY_TYPE, serializeSnapshot(runtime.getSnapshot()));
     armResumeInjection();
-    ctx.ui.setStatus("context-manager", `ctx ${runtime.getSnapshot().lastZone}`);
   });
 }